feat(cp): add opt-in NVSHMEM CP state handoff

[foraxe]

Summary

• add an opt-in NVSHMEM backend for KDA context-parallel recurrent-state communication
• support the CP=2 predecessor handoff with signal-based publication and selective V-state fetch
• support CP=4 multi-predecessor assembly with symmetric state storage, one visibility barrier, and a fused remote recurrence merge
• preserve the existing FLA preprocessing path as the default
• add deterministic distributed forward/backward checks, reproducible benchmarks, and machine-readable reports

Motivation

The general FLA CP preprocessing path communicates and stages complete intermediate states. The NVSHMEM backend specializes the recurrent-state boundary:

• At CP=2, rank 0 publishes its state in symmetric memory and rank 1 fetches only the V-state it consumes.
• At CP=4, each producer writes its affine recurrence state directly into symmetric memory. Ranks with multiple predecessors read those states directly and assemble the final state in one Triton kernel, avoiding peer-copy staging and a second local merge pass.

The backend is disabled by default. Existing behavior is unchanged unless both CULA_CP_OVERLAP=1 and CULA_CP_COMM_BACKEND=nvshmem are set.

Supported configuration

• one host with a global CP group of size two or four
• one global sequence and eager execution
• BF16 KDA inputs with FP32 recurrent state
• NVSHMEM 3.4.5 on GB200 with NVSHMEM_DISABLE_CUDA_VMM=1
• CP=2 signal path: CUDA_DEVICE_MAX_CONNECTIONS=1
• CP=4 optimized path: CUDA_DEVICE_MAX_CONNECTIONS=2, barrier publication, fused remote merge, and current-stream execution

The CP=4 optimization is limited to world size at most four and a non-transposed state layout. Other configurations retain the existing path unless strict opt-in settings request an error instead of fallback.

Correctness

The deterministic oracle compares the cuLA NVSHMEM path with FLA CP at the same output and backward boundary. It checks output and active Q/K/V/g/beta gradients.

CP size	Global sequence	Heads	Fresh launches	Maximum absolute difference
2	1,024	4	pass	0.0
4	1,024	4	3/3 pass	0.0
4	4,096	8	pass	0.0
4	8,192	16	pass	0.0

Performance

Measurements use BF16, 20 warmups, 60 iterations, five in-process repeats, and the median of each repeat's rank-max iteration median.

CP=2

Path	Latency	Improvement
cuLA + FLA CP preprocessing	0.6076 ms	reference
cuLA + NVSHMEM overlap	0.5636 ms	7.24%

CP=4

Shape	FLA CP preprocessing	NVSHMEM optimized	Improvement
4096x8	0.6317 ms	0.5680 ms	10.08%
8192x16	0.6112 ms	0.5662 ms	7.36%

These numbers isolate the cuLA KDA CP preprocessing boundary. They are not an integrated Megatron training result and do not claim a 20% end-to-end speedup.

Validation

• based on inclusionAI/cuLA commit 326a307
• FLA submodule commit 3a9ce1c8
• Ruff, shell syntax, Python compilation, and JSON validation passed
• opt-in CP=2 and CP=4 forward/backward suite: 3 passed
• CP=4 optimized path passed three fresh tiny-oracle launches and the 8192x16 oracle
• documented five-repeat CP=2 and CP=4 benchmark matrices completed

Reproduction commands and complete repeat data are under docs/cp_overlap/ and benchmarks/run_cp4_nvshmem_optimized.sh.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.