Feat/cp nvshmem enhanced

[Knight-of-Thunder]

Description

To make the computation and communication become overlap, we create a new stream for communication.
To use NVSHMEM APIs easier, and make the code cleaner, we use NVSHMEM pybindings to replace original cpp code.

Fixes # (issue)

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

Please list the changes introduced in this PR:

• Change A
• Change B

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Greptile Summary

This PR introduces two new Python modules (context_parallel_nvshmem.py and context_parallel_nvshmem_enhanced.py) that replace raw C++ NVSHMEM calls with NVSHMEM pybindings to overlap computation and communication in Context Parallel attention. It also adds a new comm_gemm C++ backend that wraps cuBLASMp for fused AllGather/GEMM, GEMM/ReduceScatter, and GEMM/AllReduce operations, plus several quality-of-life fixes across the linear ops and CUDA graph support.

Key issues found:

• nvshmem_kv NameError in backward (context_parallel_nvshmem_enhanced.py ~line 2089): nvshmem_kv is allocated and freed in the forward pass and never saved to ctx. The backward method references it directly, resulting in NameError on every backward call.
• Duplicate function definition (context_parallel_nvshmem_enhanced.py ~lines 105–148): torchrun_uid_init_bcast_object_no_reinit is defined twice; the second copy silently shadows the first and contains Chinese-language development comments that indicate this is an unclean artifact from development.
• nvshmem_malloc return not checked (comm_gemm.cpp line ~431): The symmetric heap allocation for the cuBLASMp workspace is used without a null-check, risking a null-pointer dereference if the heap is exhausted.
• Debug print in setup.py: print("CMAKE_FLAGS:", ...) will be emitted on every NVTE_WITH_CUBLASMP=1 build.
• Numerous previously-flagged issues remain: inverted error message in nvshmem_comm.cpp, missing cross-stream synchronization for NVSHMEM gets, unconditional re-initialization of NVSHMEM on every forward pass, unused tensor_get_buffer import, and cp_global_ranks NameError in the enhanced backward path.

Confidence Score: 1/5

• Not safe to merge — multiple NameError crashes in the forward and backward passes of both new files.
• The two primary new files (context_parallel_nvshmem_enhanced.py and context_parallel_nvshmem.py) contain several confirmed runtime crashes: a NameError for nvshmem_kv in the backward pass, a NameError for cp_global_ranks in the backward pass (previously flagged), a duplicate function definition with Chinese development comments indicating the code is not production-ready, missing NVSHMEM re-initialization guards, and a null-pointer risk in the C++ workspace allocation. The non-NVSHMEM changes (graph GC workaround, Megatron-LM wgrad fusion, comm_gemm backend) appear solid, but the core new feature is not ready.
• transformer_engine/pytorch/attention/dot_product_attention/context_parallel_nvshmem_enhanced.py and context_parallel_nvshmem.py require the most attention — they contain the majority of the critical bugs. transformer_engine/common/comm_gemm/comm_gemm.cpp needs the nvshmem_malloc null-check.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/attention/dot_product_attention/context_parallel_nvshmem_enhanced.py	New 3223-line file implementing enhanced NVSHMEM context parallelism. Contains multiple critical bugs: duplicate function definition, nvshmem_kv used in backward but freed in forward without ctx save, and cp_global_ranks NameError in backward.
transformer_engine/pytorch/attention/dot_product_attention/context_parallel_nvshmem.py	New 4096-line file implementing NVSHMEM-based context parallelism. Contains unused import (tensor_get_buffer) and unconditional NVSHMEM re-initialization on every forward pass.
transformer_engine/pytorch/csrc/extensions/nvshmem_comm.cpp	Adds nvshmem_get_on_current_stream C++ binding. The #else error message is inverted — it fires when NVTE_ENABLE_NVSHMEM is NOT defined but states the opposite condition.
transformer_engine/common/comm_gemm/comm_gemm.cpp	New 519-line file implementing cuBLASMp-based comm+GEMM fusion (AG-GEMM, GEMM-RS, GEMM-AR). Workspace is allocated via nvshmem_malloc without checking for null return. Otherwise well-structured.
setup.py	Adds cuBLASMp/NVSHMEM cmake flags. Contains a debug print statement that should be removed before merging.
transformer_engine/pytorch/graph.py	Adds _graph_context_wrapper to temporarily disable GC during CUDA graph capture as a workaround for a PyTorch bug. Clean, well-documented change.
transformer_engine/pytorch/ops/basic/basic_linear.py	Improves Megatron-LM wgrad fusion to set grad_added_to_main_grad flag and return a dummy wgrad tensor when needed. Minor refactoring to use weight_param local variable.

Sequence Diagram

sequenceDiagram
    participant FWD as Forward Pass
    participant CTX as ctx (saved state)
    participant BWD as Backward Pass
    participant NVSHMEM as NVSHMEM Heap
    participant P2P as P2P Fallback

    FWD->>NVSHMEM: nvshmem.tensor() → nvshmem_kv
    FWD->>NVSHMEM: nvshmem.tensor() → nvshmem_q

    loop cp_size iterations
        FWD->>NVSHMEM: nvshmem_get_on_stream(p2p_comm_buffers[i+1], nvshmem_kv, owner) [communicate_stream]
        Note over FWD: ⚠ No sync between communicate_stream and flash_attn_streams
        FWD->>FWD: flash attention compute on flash_attn_streams[i%2]
    end

    FWD->>CTX: ctx.cp_global_ranks = cp_global_ranks
    FWD->>NVSHMEM: nvshmem.free_tensor(nvshmem_kv)
    Note over NVSHMEM: ⚠ nvshmem_kv freed but NOT saved to ctx

    BWD->>CTX: read ctx.cp_global_ranks ✓
    BWD->>BWD: if nvshmem_kv is not None → NameError ✗
    Note over BWD: ⚠ nvshmem_kv undefined in backward scope

    alt NVSHMEM path
        BWD->>NVSHMEM: tex.nvshmem_get_on_current_stream(recv, nvshmem_kv, owner)
    else P2P fallback
        BWD->>P2P: flash_attn_p2p_communicate(...)
    end

Loading

_{Last reviewed commit: 1c88f8f}

[greptile-apps]

Function signature mismatch causes TypeError at runtime

The function _store_fa_nvshmem is defined with 2 parameters but called with 3 positional arguments at lines 932 and 2999:

# Definition (line 699): 2 params
def _store_fa_nvshmem(out_tensor, softmax_tensor):
    ...

# Call sites: 3 args
_store_fa_nvshmem(i, out_per_step[i], softmax_lse_per_step[i])  # lines 932, 2999

This will raise TypeError: _store_fa_nvshmem() takes 2 positional arguments but 3 were given at runtime.

The commented-out version at line 682 had the correct 3-parameter signature. Restore it with the idx parameter:

Suggested change

	def _store_fa_nvshmem(out_tensor, softmax_tensor):
	# allocate symmetric tensors lazily and copy
	if not causal:
	return
	try:
	nvshmem_fa_out = [tex.create_nvshmem_tensor(list(out_tensor.shape), out_tensor.dtype) for _ in range(cp_size)]
	nvshmem_fa_softmax_lse = [tex.create_nvshmem_tensor(
	list(softmax_tensor.shape), softmax_tensor.dtype
	) for _ in range(cp_size)]
	except Exception:
	nvshmem_fa_out = [None for _ in range(cp_size)]
	nvshmem_fa_softmax_lse = [None for _ in range(cp_size)]
	# if nvshmem_fa_out is not None:
	# for idx in range(cp_size):
	# if nvshmem_fa_out[idx] is not None:
	# nvshmem_fa_out[idx].copy_(out_tensor)
	# nvshmem_fa_softmax_lse[idx].copy_(softmax_tensor)
	# create two streams to resolve wave quantization issue of Flash Attn in each step
	flash_attn_streams = [torch.cuda.current_stream(), cp_stream]
	# synchronize fwd results correction across steps
	fwd_results_correction_done = torch.cuda.Event()
	def _store_fa_nvshmem(idx, out_tensor, softmax_tensor):
	# allocate symmetric tensors lazily and copy
	if not causal:
	return
	try:
	nvshmem_fa_out[idx] = tex.create_nvshmem_tensor(list(out_tensor.shape), out_tensor.dtype)
	nvshmem_fa_softmax_lse[idx] = tex.create_nvshmem_tensor(
	list(softmax_tensor.shape), softmax_tensor.dtype
	)
	except Exception:
	nvshmem_fa_out[idx] = None
	nvshmem_fa_softmax_lse[idx] = None
	if nvshmem_fa_out[idx] is not None:
	nvshmem_fa_out[idx].copy_(out_tensor)
	nvshmem_fa_softmax_lse[idx].copy_(softmax_tensor)

[greptile-apps]

Undefined variable cp_global_ranks in backward

The backward function accesses cp_global_ranks at lines 1805 and 1835 as a bare local variable, but this variable is never assigned locally. The code only saves and retrieves ctx.cp_global_ranks. At runtime, this will raise NameError: name 'cp_global_ranks' is not defined.

# Line 1598-1599: correct usage with ctx prefix
send_dst = ctx.cp_global_ranks[(rank - 1) % cp_size * cp_size_a2a + rank_a2a]
recv_src = ctx.cp_global_ranks[(rank + 1) % cp_size * cp_size_a2a + rank_a2a]

# Line 1805, 1835: incorrect usage without ctx prefix
owner_global = cp_global_ranks[owner_idx * cp_size_a2a + rank_a2a]  # NameError

Add this assignment near the top of the backward method alongside the other context unpacking:

Suggested change

	def backward(ctx, dout):
	# pylint: disable=missing-function-docstring
	nvtx_range_push("transformer_engine.AttnFuncWithCPAndKVP2P.backward")
	cp_size_a2a = ctx.cp_size_a2a
	rank_a2a = ctx.rank_a2a
	cp_size = get_distributed_world_size(ctx.cp_group)
	rank = get_distributed_rank(ctx.cp_group)
	send_dst = ctx.cp_global_ranks[(rank - 1) % cp_size * cp_size_a2a + rank_a2a]
	recv_src = ctx.cp_global_ranks[(rank + 1) % cp_size * cp_size_a2a + rank_a2a]
	device_compute_capability = get_device_compute_capability()
	batch_p2p_comm = int(os.getenv("NVTE_BATCH_MHA_P2P_COMM", "0")) or (
	device_compute_capability < (10, 0) and cp_size == 2
	)
	q, kv, out, softmax_lse, cu_seqlens_q_padded, cu_seqlens_kv_padded, *other_tensors = (
	restore_from_saved(ctx.tensor_objects, ctx.saved_tensors)
	)
	cu_seqlens_q_per_step = other_tensors[:cp_size]
	cu_seqlens_kv_per_step = other_tensors[cp_size : cp_size * 2]
	rng_states = other_tensors[cp_size * 2 : cp_size * 3]
	cp_size_a2a = ctx.cp_size_a2a
	rank_a2a = ctx.rank_a2a
	cp_global_ranks = ctx.cp_global_ranks

[greptile-apps]

Inverted error message in #else branch

The #else branch executes when NVTE_ENABLE_NVSHMEM is not defined, but the error message incorrectly states the function "cannot be initialized… when TE is compiled with NVTE_ENABLE_NVSHMEM=1" — the exact opposite condition.

Suggested change

	#else
	NVTE_ERROR(
	"Internal TE error: nvshmem_get_on_current_stream cannot be initialized with valid PyTorch ",
	"distributed process groups when TE is compiled with NVTE_ENABLE_NVSHMEM=1!");
	#endif
	NVTE_ERROR(
	"Internal TE error: nvshmem_get_on_current_stream requires TE to be compiled with ",
	"NVTE_ENABLE_NVSHMEM=1!");

[greptile-apps]

Missing cross-stream synchronization for NVSHMEM gets

In the forward loop, the NVSHMEM get for p2p_comm_buffers[i+1] is issued on stream flash_attn_streams[i % 2] (line 765). In the next iteration, the computation consumes p2p_comm_buffers[i+1] on stream flash_attn_streams[(i+1) % 2] — a different stream — with no event-based synchronization.

The existing P2P fallback updates send_recv_reqs with requests that are waited on (line 754-755), but when the NVSHMEM path is taken, send_recv_reqs[(i) % 2] remains empty, so the wait is a no-op:

# Line 765: NVSHMEM get on stream[i % 2]
tex.nvshmem_get_on_current_stream(p2p_comm_buffers[i + 1], nvshmem_kv, int(owner_global))
# send_recv_reqs[(i) % 2] is never populated with a request

# Next iteration (i+1):
# Line 754-755: wait is a no-op because send_recv_reqs[(i+1) % 2] is empty
for req in send_recv_reqs[(i + 1) % 2]:
    req.wait()
# Line 779-782: computation accesses p2p_comm_buffers[i+1] on stream[(i+1) % 2]

Without synchronization, the computation kernel on stream [(i+1) % 2] may observe stale data. Add CUDA event-based synchronization after the NVSHMEM get and before subsequent iterations use the buffer.

[greptile-apps]

NVSHMEM initialized unconditionally on every forward pass

Despite its name torchrun_uid_init_bcast_object_no_reinit, this function calls nvshmem.init(...) unconditionally (line 457) with no guard against re-initialization. The function is invoked inside the forward method (line 730) on every forward pass. Calling nvshmem.init() multiple times may cause undefined behavior or errors depending on the NVSHMEM pybinding version.

Move NVSHMEM initialization to module-load time or wrap the call in a one-time initialization guard:

_nvshmem_initialized = False

def _maybe_init_nvshmem(cp_group):
    global _nvshmem_initialized
    if _nvshmem_initialized:
        return
    torchrun_uid_init_bcast_object_no_reinit(cp_group)
    _nvshmem_initialized = True

Then call _maybe_init_nvshmem(cp_group) instead of torchrun_uid_init_bcast_object_no_reinit(cp_group) at line 730.

[greptile-apps]

Unused import tensor_get_buffer

tensor_get_buffer is imported at line 14 but is not referenced anywhere else in the module. This creates an unnecessary dependency and will trigger linting warnings.

Remove this unused import:

Suggested change

	from nvshmem.core.interop.torch import tensor_get_buffer
	# Removed: from nvshmem.core.interop.torch import tensor_get_buffer

[greptile-apps]

Duplicate function definition silently discards first implementation

torchrun_uid_init_bcast_object_no_reinit is defined twice in this file (once at ~line 105 and again at ~line 143). In Python, the second definition completely replaces the first at module load time without any error or warning. The second definition also contains Chinese-language inline comments (# 设备已经在外面 set 过最好 and # 不要再 init_process_group !!!), strongly suggesting it is a copy-paste development artifact that was never cleaned up.

Remove the first definition and consolidate into a single, clean function — or if the two definitions are intentionally different, unify them.

[greptile-apps]

nvshmem_kv used in backward but never defined there — NameError at runtime

In the backward method, nvshmem_kv is referenced at multiple points (if nvshmem_kv is not None:, tex.nvshmem_get_on_current_stream(recv_tensor[0], nvshmem_kv, ...)), but it is never assigned within the backward function.

nvshmem_kv was allocated in the forward pass and then explicitly freed at the end of forward via nvshmem.free_tensor(nvshmem_kv). It is not saved to ctx at all. Accessing an unresolved name will raise NameError: name 'nvshmem_kv' is not defined on the first backward call.

The fix depends on the intended design:

• If NVSHMEM KV buffers should be re-allocated in the backward pass, add the allocation at the start of backward.
• If the backward should fall back to P2P when no NVSHMEM buffer is available, initialise it to None at the top of backward:

nvshmem_kv = None  # backward allocates its own buffer or falls back to P2P

[greptile-apps]

Debug print left in build script

This print("CMAKE_FLAGS:", cmake_flags[-2:]) is a debug statement that will be emitted on every build where NVTE_WITH_CUBLASMP=1. It should be removed before merging to keep build output clean.

Suggested change

cmake_flags.append(f"-DNVSHMEM_DIR={nvshmem_dir}")

[greptile-apps]

nvshmem_malloc return value not checked — silent null-pointer dereference

nvshmem_malloc can return nullptr on allocation failure (e.g. symmetric heap exhausted), but the result is stored without any null check. ctx->workspace is then passed directly to cublasMpMatmul, which will dereference the null pointer and produce undefined behaviour or a hard crash.

Add a null check after allocation:

  if (ctx->workspace_size < wrksp_size_device) {
    nvshmem_free(ctx->workspace);
    ctx->workspace = nvshmem_malloc(wrksp_size_device);
    NVTE_CHECK(ctx->workspace != nullptr,
               "nvshmem_malloc failed to allocate workspace of size ", wrksp_size_device);
    ctx->workspace_size = wrksp_size_device;
  }