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: three files and 122 additions. 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;
• dispatch the CUDA Driver stream-memory APIs through Transformer Engine's
  existing lazy driver loader, avoiding a hard libcuda link dependency.

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.
• Rebuilt after review commit f676f681; readelf reports no libcuda.so
  dependency. Wheel SHA256:
  7e6df6d2042fa14b6e9e72cad5e669aafa72e31809ad253a898f363d321f569f.
• 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 an experimental stream-ordered CP=4 gradient-return primitive for a Megatron NVSHMEM context-parallel attention backend. The implementation is narrow (three files) and deliberately does not add a hard NVSHMEM dependency — it uses caller-supplied peer tensor views and dispatches cuStreamWriteValue32/cuStreamWaitValue32 through TE's existing lazy CUDA driver loader.

• nvshmem_comm.cpp: Implements scatter → stream-write epoch → stream-wait epoch → accumulate, with cp_check_current_device guards on all tensor arguments, self-alias assertions for the rank-self slots, and the driver API routed through NVTE_CALL_CHECK_CUDA_DRIVER (previous concerns about direct driver calls and unconditional libcuda linkage were addressed in commit f676f68).
• extensions.h / pybind.cpp: Declare and expose the function unconditionally (no NVTE_ENABLE_NVSHMEM guard required), consistent with the design that the primitive works with any caller-allocated peer tensor views.
• Two minor internal-API gaps: cp_stream_wait_epochs does not call cp_check_current_device unlike its paired write function, and the static epoch map grows without pruning for dynamic buffer-allocation patterns.

Confidence Score: 5/5

Safe to merge; the new primitive is isolated to three files, does not touch existing fused-attention paths, and all previously raised concerns (lazy driver loading, self-alias assertion, current-device validation) have been addressed.

The change is narrow and well-guarded: every tensor argument is validated for CUDA device membership, the self-peer alias invariant is enforced with explicit data_ptr() comparisons, and driver symbol resolution goes through the lazy loader so there is no link-time libcuda dependency. The two remaining observations are minor API-consistency gaps in internal helper functions that have no effect on correctness at the current call site.

transformer_engine/pytorch/csrc/extensions/nvshmem_comm.cpp — specifically the cp_stream_wait_epochs helper and the static epoch map lifetime.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/csrc/extensions/nvshmem_comm.cpp	Adds stream-ordered CP=4 gradient return primitive using NVTE_CALL_CHECK_CUDA_DRIVER (lazy loader), with cp_check_current_device guards and self-alias assertions; two minor internal-API consistency gaps (cp_stream_wait_epochs device check, epoch map lifetime) but no correctness bugs
transformer_engine/pytorch/csrc/extensions.h	Adds declaration of nvshmem_cp_global_grad_return_execute outside any NVSHMEM guard, consistent with the implementation which does not require NVTE_ENABLE_NVSHMEM
transformer_engine/pytorch/csrc/extensions/pybind.cpp	Registers nvshmem_cp_global_grad_return_execute in the nvshmem section with py::call_guardpy::gil_scoped_release(); consistent with adjacent NVSHMEM function registrations

Sequence Diagram

sequenceDiagram
    participant PE0 as PE 0 (rank=0)
    participant PE1 as PE 1 (rank=1)
    participant PEN as PE N (rank=N)
    participant GPU as GPU Stream

    Note over PE0,PEN: All PEs call nvshmem_cp_global_grad_return_execute concurrently

    PE0->>GPU: "copy_() scatter dk_global to peer_grad_key_returns[owner][rank=0]"
    PE1->>GPU: "copy_() scatter dk_global to peer_grad_key_returns[owner][rank=1]"
    PEN->>GPU: "copy_() scatter dk_global to peer_grad_key_returns[owner][rank=N]"

    Note over PE0,PEN: Each PE writes its epoch to all owners' epoch buffers

    PE0->>GPU: "cuStreamWriteValue32 peer_grad_committed_epochs[owner][rank=0] = epoch"
    PE1->>GPU: "cuStreamWriteValue32 peer_grad_committed_epochs[owner][rank=1] = epoch"
    PEN->>GPU: "cuStreamWriteValue32 peer_grad_committed_epochs[owner][rank=N] = epoch"

    Note over PE0,PEN: Each PE waits until its local epoch buffer has all cp_size slots >= epoch

    PE0->>GPU: "cuStreamWaitValue32 grad_committed_epoch[0..cp_size-1] >= epoch"
    PE1->>GPU: "cuStreamWaitValue32 grad_committed_epoch[0..cp_size-1] >= epoch"
    PEN->>GPU: "cuStreamWaitValue32 grad_committed_epoch[0..cp_size-1] >= epoch"

    Note over PE0,PEN: Barrier satisfied - all scatter writes visible

    PE0->>GPU: accumulate grad_key_return[0..cp_size-1] to dk
    PE1->>GPU: accumulate grad_key_return[0..cp_size-1] to dk
    PEN->>GPU: accumulate grad_key_return[0..cp_size-1] to dk

Loading

_{Reviews (4): Last reviewed commit: "Validate CP gradient return tensor devic..." | Re-trigger Greptile}

[sudhakarsingh27]

@foraxe, could you elaborate what's the motivation behind this. I see that this is a part of a bigger effort here NVIDIA/Megatron-LM#5969 but why are we then hardcoding this for only CP=4, why not for CP=2 or 8?

[foraxe]

Hi @sudhakarsingh27

Thanks for asking. The motivation is to reduce Megatron CP attention communication overhead

• In Megatron: We try to Reduce CP attention communication overhead, by using
  • symmetric K/V storage and direct peer access in forward,
  • direct owner gradient return in backward.
• This TE PR try to Reduce CP backward dK/dV redistribution overhead through:
  • direct peer-buffer writes and stream-ordered GPU synchronization.

---

CP=4 came from the production workload we were targeting and the hardware used for validation. The end-to-end stack was tested at CP=4, S=262144 on four GB200 GPUs, including training correctness against TE p2p and integrated timing. I added the CP=4 check as a fail-closed validation boundary, rather than implying that the operation fundamentally requires CP=4.

That said, I agree that hardcoding the experimental integration constraint in the TE primitive is unnecessarily restrictive. The ownership rule itself is generic: for CP size N, owner o receives native sequence chunks o and (2 * N - 1 - o), and the implementation already sizes its peer lists, epoch vectors, and accumulation loop from cp_size.

---

I will update this PR to make the low-level primitive CP-size-generic and add focused CP=2/4/8 tests for its scatter, synchronization, and reduction contract. The dependent Megatron backend can remain explicitly gated to CP=4 until its full forward/backward path is validated at CP=2 and CP=8.
This keeps the TE API general while being precise about what has end-to-end performance and correctness evidence today.