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 ~4000-line Python modules that replace point-to-point distributed collectives with NVSHMEM-based one-sided get operations for KV exchange, adds nvshmem_get_on_current_stream C++ binding, and brings in a new comm_gemm module using cuBLASMp. Several ancillary changes update FlashAttention version caps and cuDNN thresholds.

Critical issues identified:

1. context_parallel_nvshmem_enhanced.py: _store_fa_nvshmem is defined with 2 parameters but called with 3 at lines 932 and 2999 — this raises TypeError at runtime.

2. context_parallel_nvshmem_enhanced.py: The backward pass uses bare variable cp_global_ranks at lines 1805 and 1835, which is never assigned locally; only ctx.cp_global_ranks is persisted. This raises NameError on backward.

3. context_parallel_nvshmem_enhanced.py: NVSHMEM get operations (line 765) are issued on flash_attn_streams[i % 2], but data is consumed in the next iteration on flash_attn_streams[(i+1) % 2] — a different stream — with no cross-stream CUDA event synchronization. The send_recv_reqs list stays [[], []] on the NVSHMEM path, bypassing the intended wait. This risks silent data corruption.

4. nvshmem_comm.cpp: The #else error message (line 100) incorrectly states the function "cannot be initialized when TE is compiled with NVTE_ENABLE_NVSHMEM=1" — the exact opposite of the true condition.

5. context_parallel_nvshmem.py: torchrun_uid_init_bcast_object_no_reinit calls nvshmem.init() unconditionally on every forward pass (line 730); repeated initialization can cause undefined behavior.

6. context_parallel_nvshmem.py: tensor_get_buffer is imported at line 14 but never used.

Confidence Score: 1/5

• Critical — not safe to merge. Two new modules contain runtime crashes (TypeError, NameError) and a cross-stream data-race that produces silently incorrect results.
• Three separate code paths will fail before end-to-end testing can pass: (1) wrong-arity call to _store_fa_nvshmem raises TypeError, (2) undefined cp_global_ranks in backward raises NameError, and (3) missing cross-stream synchronization between NVSHMEM get stream and FlashAttention compute stream creates a data-race that can silently corrupt gradients. Additionally, NVSHMEM is re-initialized on every forward pass with undefined behavior, and an error message is logically inverted.
• context_parallel_nvshmem_enhanced.py (critical), context_parallel_nvshmem.py (critical), nvshmem_comm.cpp (logic fix required)

_{Last reviewed commit: f63d3c3}

[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