[Kernel] Rewrite JIT custom all-reduce (v2) with a decoupled kernel/storage design

[DarkSharpness]

Generated by Claude.

Motivation

The JIT custom all-reduce (v2) previously coupled compute and storage in one C++ object: CustomAllReduceBase owned a cudaMalloc arena (signal buffers + graph-param slots + pull/push buffers), shared it via raw cudaIpcMemHandle exchange, and tracked CUDA-graph capture inputs inside C++. Every kernel had to inherit from it, which made the kernels hard to reuse (e.g. tp_qknorm), the storage lifecycle opaque to Python, and graph-input registration split across three layers.

This PR rewrites it with a fully decoupled design: storage is owned in Python, kernels are pure functions of (input, Communicator, algo, pull_arg).

Design

Storage plane (Python-owned)

• All shared buffers live in one torch symmetric-memory allocation, sliced into [2 * world_size push buffers | pull buffer | pull semaphores]; push counters are a plain local tensor. Symmetric memory also provides the multicast address of the pull workspace for free.
• Communicator (C++) is a thin pointer holder: it validates the tensor views and records raw pointers plus grid-size knobs. No allocation, no IPC, no graph bookkeeping.

Kernels (one file, three algorithms)

• 1shot_push: lamport-style push to every peer's push workspace, then a local polling reduce (best at small sizes).
• 1shot_pull / 2shot_pull: semaphore-synchronized pull with three data sources: eager workspace, CUDA-graph pointer table, or multimem.ld_reduce multicast. 2shot fuses reduce-scatter with the broadcast by writing shards back to all peers.
• PDL is used throughout, including a dedicated PDL memcpy kernel for the eager staging copies.
• An f16x2 multimem.ld_reduce variant is added so fp16 instantiates in multicast mode; all multimem asm is preprocessed out below sm90, so the module still compiles for sm80 (verified).

CUDA-graph inputs

• During capture, each all-reduce consumes one row of a device-side graph_params pointer table; the captured kernel dereferences its row at replay time.
• After capture, Python exchanges the input pointers: batched cudaIpc handles via a new IPCManager for cudaMalloc-backed pointers, and the existing fabric/posix-fd VMM path (expandable segments) via VmmGraphInputManager, whose base-range walk moved to Python (compute_graph_capture_bases).

Dispatch config

• Tuned thresholds and block counts (H200/B200; separate graph vs eager heuristics; multicast windows) live in device_communicators/configs/custom_all_reduce_v2.py.
• Unlike the old implementation, 2shot_pull now has a tuned ceiling: above the NCCL crossover, should_custom_ar returns False. Benchmarks/tests that must keep every size on the custom-AR path call uncap_pull_thresholds().
• Sizing API: CustomAllReduceV2(group, device, max_size=..., *, max_pull_size=None, max_push_size=None, ...). Explicit per-direction sizes are respected verbatim; unset directions take min(tuned_bytes, max_size). max_size defaults to SGLANG_CUSTOM_ALL_REDUCE_V2_MAX_SIZE_KB (16 MB), keeping the previous default memory footprint.

Downstream

• tp_qknorm is ported onto the Communicator push plane (no more base-class inheritance); the MiniMax-M2 call sites are unchanged.
• The public CustomAllReduceV2 surface (capture(), should_custom_ar, custom_all_reduce, override_algo, close, .obj) is preserved; parallel_state / the v1 wrapper need no changes. The unused override_shot hook is removed.
• The all-reduce benchmark gains separate jit-eager / jit-graph providers, and both the benchmark and the tests now allocate their NCCL reference group independently (dist.new_group) instead of reaching into parallel_state internals.
• test_custom_all_reduce wall time drops from 213 s to ~3.4 s: triton.testing.assert_close compares via numpy on the host (~0.6 s per 32 MB tensor); torch.testing.assert_close compares on device (~2 ms) with better mismatch reporting.

Performance (B200, bf16, us per all-reduce)

world_size = 8:

message	nccl	aot (v1)	jit-graph	speedup vs v1
4 KB	26.6	6.8	3.9	1.7x
256 KB	31.1	26.0	6.6	3.9x
1 MB	32.1	26.9	12.4	2.2x
16 MB	108.1	147.1	53.2	2.8x
64 MB	276.7	512.6	189.9	2.7x

world_size = 2: 3.0 us at 4 KB (vs 10.3 nccl / 6.0 v1), 579 GB/s bus bandwidth at 64 MB. The new jit-eager line additionally exposes the Python dispatch floor (~36 us) separately from kernel cost.

Verification (8x B200)

• test_custom_all_reduce.py: 180/180 at world size 2 and 8 (all sizes x fp16/bf16/fp32 x 3 algos x eager/graph, exact-equality against NCCL).
• test_tp_qknorm.py: 42/42 at world sizes 2/4/8.
• Multicast smoke (heuristic dispatch, eager + graph windows) on 8 GPUs, bit-exact vs NCCL.
• E2E: TP=2 server with CUDA graphs (Qwen3-0.6B), 6k-token prefill and concurrent batches.
• sm80 cross-compile via override_jit_cuda_arch(8, 0) for all dtypes.
• bench_custom_all_reduce.py full sweep at world sizes 2/4/8.

---

CI States

Latest PR Test (Base): 🚫 Run #29559587415
Latest PR Test (Extra): 🚫 Run #29559587350

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!