Add stream-ordered CP gradient return primitive

[foraxe]

Description

Add one experimental CP=4 gradient-return primitive used by an opt-in Megatron
NVSHMEM context-parallel attention backend.

The primitive:

• scatters global dK/dV into symmetric owner buffers using Megatron's native
  two-chunk CP layout;
• publishes writer epochs with cuStreamWriteValue32;
• waits for all CP writers with cuStreamWaitValue32;
• accumulates four writer slots into the local owner's dK/dV;
• accepts peer tensor views supplied by the caller, so it does not add an
  NVSHMEM dependency to Transformer Engine.

This is intentionally narrow: four files and 122 additions after automated
formatting. It does not add a
Transformer Engine attention backend or modify existing fused-attention
behavior.

Type of change

• Documentation change
• Bug fix
• New feature
• Breaking change
• Infra/Build change
• Code refactoring

Changes

• add a CUDA implementation of stream-ordered CP owner-gradient return;
• expose the operation through the PyTorch extension;
• link the extension with libcuda for the CUDA Driver stream-memory APIs.

Validation

• Rebased the four-file patch onto current Transformer Engine main.
• Built the Python 3.12 / CUDA 13 PyTorch extension from the rebased source and
  verified that nvshmem_cp_global_grad_return_execute is present in the
  resulting extension.
• In the reviewed Transformer Engine 2.13 integration stack, ran four-GB200
  Megatron CP=4 training at S=262144, H=2560, 20 heads, head dimension 128,
  four layers, and vocabulary 157184.
• TE p2p and the NVSHMEM candidate matched all 73 checkpoint tensors with worst
  max-absolute difference 1.9073486328125e-6 (atol=2e-6, rtol=0).

For that single clean integrated comparison, median iteration time over updates
4-7 was 1707.45 ms for TE p2p and 1413.65 ms for the NVSHMEM candidate, 17.21%
lower. This is an integrated stack result, not a standalone speed claim for the
primitive in this PR.

Checklist

• I have read and followed the contributing guidelines
• The functionality is complete for the documented experimental contract
• I have commented hard-to-understand ownership and synchronization logic
• I have made corresponding user documentation changes
• My changes generate no new build warnings observed in local validation
• I have added an in-tree unit test
• The full existing test suite passes locally

The end-to-end correctness test currently lives in the dependent Megatron PR,
which owns symmetric allocation and peer tensor construction. The current-main
validation is a compile/export check; end-to-end runtime validation used the
reviewed Transformer Engine 2.13 integration branch.

Stack relationship

• Dao-AILab/flash-attention#2722
  provides two-section causal backward over shared K/V.
• NVIDIA/Megatron-LM#5969
  owns the opt-in backend, symmetric allocation, and peer tensor views passed
  to this primitive.

[greptile-apps]

Greptile Summary

This PR adds a narrow, experimental CP=4 stream-ordered gradient-return primitive (nvshmem_cp_global_grad_return_execute) that scatters global dK/dV into symmetric owner buffers using the two-chunk CP layout, coordinates completion via cuStreamWriteValue32/cuStreamWaitValue32 epoch fencing, and accumulates the result. The primitive is exposed through the PyTorch extension and deliberately avoids an NVSHMEM compile-time dependency by accepting peer tensor views from the caller.

• The scatter logic, epoch protocol ordering (copy → write-epoch → wait-epoch → accumulate), and the two-chunk index formula are all functionally correct for CP=4.
• The cuStreamWriteValue32 and cuStreamWaitValue32 calls are made as direct symbol references rather than through the repo's cuda_driver::call dynamic-dispatch mechanism, which requires adding \"cuda\" unconditionally to libraries in build_tools/pytorch.py, changing link-time behavior for ALL TE PyTorch extension builds.
• The scatter formula (7 - owner) * half encodes 2 * cp_size - 1 - owner for CP=4 and would benefit from an inline constant or comment to make the relationship self-evident.

Confidence Score: 3/5

The gradient accumulation and synchronization logic are correct, but the unconditional libcuda linkage changes the build contract for every TE PyTorch extension build, not just those using the new primitive.

The core synchronization protocol (scatter, write-epoch, wait-epoch, accumulate) is sound and matches the described two-chunk CP layout. The concern is in build_tools/pytorch.py: libraries = ["cuda"] is applied to all builds, and the driver calls that require it bypass the repo's own lazy-loading mechanism documented in cuda_driver.h. This is a build-wide side effect that should be addressed before merge.

build_tools/pytorch.py and the two driver-call sites in transformer_engine/pytorch/csrc/extensions/nvshmem_comm.cpp (cp_stream_write_epoch and cp_stream_wait_epochs) need the most attention.

Important Files Changed

Filename	Overview
build_tools/pytorch.py	Unconditionally adds "cuda" to libraries, linking libcuda.so into ALL TE PyTorch extension builds rather than only when the new primitive is used; bypasses the repo's existing lazy-loading driver dispatch.
transformer_engine/pytorch/csrc/extensions/nvshmem_comm.cpp	Adds the CP=4 gradient-return primitive; scatter logic, epoch protocol, and accumulation are logically sound, but driver API calls use direct symbol references instead of the repo's NVTE_CALL_CHECK_CUDA_DRIVER dynamic-dispatch macro, creating the link-time dependency in pytorch.py.
transformer_engine/pytorch/csrc/extensions.h	Adds declaration for nvshmem_cp_global_grad_return_execute; straightforward and consistent with existing declarations.
transformer_engine/pytorch/csrc/extensions/pybind.cpp	Registers the new function alongside existing NVSHMEM bindings; no guard differences from other registrations in the same section.

Sequence Diagram

sequenceDiagram
    participant R as Rank R (caller)
    participant OB as Owner O buffer
    participant EP as Owner O epoch buf
    participant LE as Local epoch buf
    participant LC as Local collect buf

    Note over R: for owner = 0..cp_size-1
    R->>OB: .copy_() scatter dK/dV slices (async, current stream)
    Note over R: fetch_add epoch (atomic)
    Note over R: for owner = 0..cp_size-1
    R->>EP: cuStreamWriteValue32(epoch) stream-ordered after copies

    Note over R: cuStreamWaitValue32 for all S in 0..cp_size-1
    LE-->>R: unblocks when all writers published epoch

    Note over R: accumulate grad_key_return
    LC->>R: dk.add_(grad_key_return.select(0, source))

Loading

_{Reviews (1): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile}

[greptile-apps]

Unconditional libcuda linkage bypasses TE's existing lazy-loading pattern

libraries = ["cuda"] is now applied to every TE PyTorch extension build, not just NVSHMEM-enabled ones. The existing pattern — established in common/util/cuda_driver.h and documented there — uses cuda_driver::call / NVTE_CALL_CHECK_CUDA_DRIVER to load driver symbols at runtime via dlopen/dlsym, precisely to avoid a hard link-time dependency on libcuda.so (e.g., on headless build servers that have only the CUDA toolkit stubs). The direct cuStreamWriteValue32 / cuStreamWaitValue32 calls in nvshmem_comm.cpp require this symbol to be resolved at link time; switching those two sites to NVTE_CALL_CHECK_CUDA_DRIVER would eliminate the need for this unconditional library entry and keep the new primitive consistent with the rest of the codebase.

[greptile-apps]

Direct driver calls instead of NVTE_CALL_CHECK_CUDA_DRIVER

cuStreamWriteValue32 and cuStreamWaitValue32 are called directly here (and in cp_stream_wait_epochs), requiring a link-time symbol from libcuda.so. The repo's cuda_driver.h comment explicitly says direct calls are the wrong approach: "The CUDA driver library … may be different at compile-time and run-time … Indirect function calls into a lazily-initialized library ensures we are accessing the correct version." Using NVTE_CALL_CHECK_CUDA_DRIVER(cuStreamWriteValue32, stream, addr, value, 0) at these two call sites would follow the established pattern and remove the need for the unconditional "cuda" entry in build_tools/pytorch.py.

[greptile-apps]

The magic number 7 is 2 * cp_size - 1 for CP=4. Expressing it as a named constant with a comment makes the two-chunk layout formula self-documenting and prevents silent breakage if the CP=4 constraint is ever relaxed.

Suggested change

	key_slot.narrow(0, 0, half).copy_(dk_global.narrow(0, owner * half, half));
	key_slot.narrow(0, half, half).copy_(dk_global.narrow(0, (7 - owner) * half, half));
	value_slot.narrow(0, 0, half).copy_(dv_global.narrow(0, owner * half, half));
	value_slot.narrow(0, half, half).copy_(dv_global.narrow(0, (7 - owner) * half, half));
	// Two-chunk CP layout: owner `o` owns half-chunks at index `o` and
	// `(2*cp_size - 1 - o)`. For CP=4 that resolves to `7 - o`.
	const int64_t mirror_owner = (2 * cp_size - 1 - owner);
	key_slot.narrow(0, 0, half).copy_(dk_global.narrow(0, owner * half, half));
	key_slot.narrow(0, half, half).copy_(dk_global.narrow(0, mirror_owner * half, half));
	value_slot.narrow(0, 0, half).copy_(dv_global.narrow(0, owner * half, half));
	value_slot.narrow(0, half, half).copy_(dv_global.narrow(0, mirror_owner * half, half));

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!