perf(comm): LL128 protocol for decode_cp_a2a_lse_reduce, on any torch symmetric-memory backend

[thanhhao98]

📌 Description

decode_cp_a2a_lse_reduce (#4929) has two problems; each commit fixes one.

1. It spends most of each call synchronizing. On 8×B200 one call is slower than the MNNVL all-to-all plus the separate
   combine kernel it is meant to replace.
   • The first commit keeps the op, its Python API and its workspace contract, and replaces the kernel's protocol.
   • Data now carries its own readiness, as in NCCL's LL128 protocol, so there is no fence, readiness word or grid-wide
     barrier.
2. It selects torch's NCCL symmetric-memory backend for the whole process. create_workspace called
   symm_mem.set_backend("NCCL"), which is process-wide and frozen by the first symmetric allocation.
   • In SGLang that turns off its default custom all-reduce, which costs more than the kernel saves.
   • The second commit takes each rank's workspace address from whichever backend the process uses (CUDA by default,
     NCCL, or NVSHMEM).
   • The op therefore no longer needs NCCL's device API, NCCL ≥ 2.29 or torch 2.14.

Results on 8× B200:

Measurement	Result
FlashInfer's bench_dcp_lse_reduce.py, #4929 → this PR	2.0–2.4× faster per call in single-op CUDA graphs, 2.6–3.6× in 50-op graphs
DeepSeek-V2 decode shape, per layer (16 heads per rank, head_dim 512)	6.5 µs at batch 1, vs 27.4 for #4929 in the same job; 17–32% faster than SGLang's current path (pack + MNNVL all-to-all + combine) at every batch from 1 to 256 (19–33% on the NCCL backend)
Commit 2 on the NCCL backend	Same speed as commit 1 at every batch
Commit 2 on torch's default CUDA backend	Same at batch 1; 0.1–1.4 µs (at most ~3%) slower at batch ≥ 2, over two jobs (see Backend parity)
SGLang serving (DeepSeek-V2-236B, TP 8 = DCP 8, 256K context), TPOT at concurrency 1 vs the MNNVL path	−0.87% with SGLang's default settings. #4929 is +14.5%, measured with SGLang's v1 all-reduce in every arm, which #4929 needs. Commit 1 alone was +2.6% because of the backend switch.

The release candidates so far (0.7.1rc1–rc5, 0.7.2rc1) ship #4929's kernel and its set_backend("NCCL"), which SGLang
cannot adopt without turning off its default all-reduce. If this PR can be reviewed in time, it may be worth considering
for 0.7.1, or else 0.7.2.

Why the #4929 kernel is slow

We timed the kernel's phases with %globaltimer in block 0 on 8×B200, batch 1, 16 heads per rank, head_dim 512, µs:

epoch grid.sync	send	__threadfence_system in every thread	grid.sync	release flags	acquire-wait	grid.sync	merge	total
~1.2	~4.3	~5.8	~1.8	~2.8	0.5–3	~3	~5	~25

• About 14 µs is the protocol itself.
• Tweaking it recovers ~13% at most: one fence per block, no epoch grid.sync, busy spin (27.2 vs 31.1 µs).
• Most of the rest is a latency-bound send loop and a block-per-row merge with serial weights and scalar loads.

What changes — commit 1: LL128 protocol

Transport (include/flashinfer/comm/dcp_lse_reduce.cuh)

• A row's 8-byte data words are sent as 128-byte lines: 15 data words plus a flag word.
• A warp writes four lines per 16-byte vector store straight into the peer's workspace. Lane part 7 carries data word 14 and
  the flag.
• The row's LSE travels as one 16-byte LL line ({lse, flag, 0, flag}) per group of four data lines.
• Receivers poll the lines themselves, so data and readiness arrive together. As in NCCL's LL128, this relies on a
  warp-coalesced 128-byte store becoming visible atomically over NVLink. Every line is 128-byte aligned.

Work split

• One warp sends each (row, destination).
• One warp receives each (row, group of four lines). It polls every source's line together with the LSE lines, merges in
  fixed source order (deterministic), and writes the output.
• Sources are processed in batches of 8, so cp_size up to 64 is still supported.

Slot reuse

• Two slots alternate by call, selected by a device-side counter, so CUDA-graph replays advance it.
• Every line is read by exactly one lane, which zeroes its flag. Flags are never zero, so a stale line can never look
  ready, including after the 32-bit counter wraps.
• The workspace is therefore zeroed at creation.
• The launch stays cooperative, only to guarantee co-residency.

Strided input (csrc/dcp_lse_reduce.cu)

• A 4-D partial_o (3-D partial_lse) may be any view whose last dimension is contiguous.
• A caller can pass its attention output [B, cp * H, D] as out.unflatten(1, (cp, H)).permute(0, 2, 1, 3) instead of
  packing it first.
• Other shapes are unchanged and must be contiguous.

Workspace (flashinfer/comm/dcp_lse_reduce.py)

• decode_cp_a2a_lse_reduce_workspace_size follows the new row layout: 128 bytes of metadata +
  2 · cp · rows · (lines · 128 + groups · 16).
• That is ~17% larger than before at head_dim 512 bf16 (1200 vs 1028 bytes per row per source per slot).
• Rows are limited to 15360 bytes.

What changes — commit 2: no forced symmetric-memory backend

Why. torch's symmetric-memory backend is process-wide and frozen by the first symmetric allocation, so
create_workspace had two bad outcomes:

• it failed with Backend can not be changed after use. once another library had allocated; or
• it moved every other user of torch symmetric memory onto NCCL.

SGLang's custom all-reduce v2 allocates group-named buffers, which torch's NCCL allocator rejects. SGLang therefore falls back
to its v1 all-reduce, which on its own costs +3.3% TPOT, more than this kernel saves.

What.

• The kernel needs only each rank's workspace address. Every backend exposes it as SymmetricMemory::get_buffer_ptrs_dev()
  plus get_offset(), with null for a rank it cannot map for load/store.
• The kernel takes that device array and the payload offset instead of an ncclWindow_t and ncclGetLsaPointer. The LL128
  protocol and the workspace layout are unchanged.
• The launcher looks the handle up through the base class. It checks the group size and rank against cp_size and
  cp_rank, and checks that every rank's workspace is mapped and 128-byte aligned.
• create_workspace allocates on the process's current backend and never changes it.
• Without NCCL's device API, the NCCL_HAS_DEVCOMM stub and the NCCL header/library discovery in the JIT spec are gone.
• bench_dcp_lse_reduce.py gains --symm-mem-backend.

Behavior and API.

• No signature or schema change. decode_cp_a2a_lse_reduce_workspace_size is as in commit 1.
• The workspace lives on the process's current backend: CUDA unless the application selected another.
  • An application that wants NCCL selects it itself, before its first symmetric allocation.
  • The NCCL backend still needs an existing communicator on the group, and cuMem.
• The first call no longer creates an NCCL device communicator, which was a collective.
• Requirements: torch ≥ 2.10, no NCCL floor. Tested on torch 2.13 and 2.14.1.

Results (8× B200, one HGX node, NVLink 5; torch 2.14.1, NCCL 2.30.7)

Backend parity. Commit 2 vs commit 1 at the DeepSeek-V2 decode shape:

• 16 heads per rank, head_dim 512, bf16, cp 8, strided input;
• 60 calls captured per CUDA graph, L2 hot;
• µs per layer: max over ranks of each rank's median, mean of 2 repeats;
• one node, palindromic process order (commit 1 → commit 2 on NCCL → commit 2 on CUDA ×2 → commit 2 on NCCL → commit 1).

batch	SGLang's current path (pack + MNNVL + combine)	commit 1 (NCCL)	commit 2, NCCL backend	commit 2, CUDA backend
1	8.08	6.51	6.53	6.55
2	8.21	6.66	6.64	6.85
4	8.82	7.16	7.15	7.27
8	10.68	8.05	8.07	8.20
16	15.05	10.40	10.49	10.50
32	20.60	13.81	13.88	13.97
64	32.38	21.71	21.95	22.38
128	54.47	38.42	38.74	39.23
256	98.61	73.33	73.74	74.34

• On the NCCL backend, commit 2 matches commit 1 at every batch, strided and packed, L2 hot and flushed. Mean deltas are
  +0.45% (strided, hot), +0.31% (strided, flushed) and ~0 (packed).
• On the CUDA backend the same code is 0.1–1.0 µs slower from batch 2 up (at most +3.1%, at batch 64). A second job
  measured +0.1–0.15 µs at batch 2–32 and +0.45–1.4 µs at batch 64–256.
  • Since commit 2 on NCCL matches commit 1, this is a property of the two backends' peer mappings, not of the kernel change.
  • The CUDA backend maps each peer's cuMemCreate block separately; NCCL maps its symmetric window into one flat range.
    We have not pinned down which property of the mappings matters.
  • In SGLang the CUDA backend is still the right choice: forcing NCCL costs it far more (see the serving results).
• With L2 flushed before every layer the picture is the same: 9.51 → 7.25 µs at batch 1, 98.67 → 74.36 at batch 256
  (current path → commit 2, CUDA backend).

benchmarks/bench_dcp_lse_reduce.py, bf16, cp 8, µs per call (max over ranks of each rank's median):

batch × local heads × head_dim	#4929, eager → graph 1 op → graph 50 ops	this PR (commit 2, CUDA backend), eager → graph 1 op → graph 50 ops
1 × 2 × 64	19.82 → 22.54 → 15.98	12.2 → 11.35 → 6.12
1 × 8 × 128	23.48 → 24.90 → 20.21	12.5 → 12.28 → 6.26
4 × 8 × 128	30.75 → 30.37 → 23.60	12.2 → 12.36 → 6.38
16 × 8 × 128	30.70 → 30.45 → 23.70	12.3 → 12.80 → 7.04

• The feat(comm): add fused decode CP A2A + LSE reduce #4929 column is from 2026-09-30 on the same node type, with the same bench.
• In the same jobs, commit 2 on the CUDA and NCCL backends stays within ±0.3 µs of commit 1 in both graph modes, at every
  shape (two jobs, palindromic order).
• Single eager samples occasionally read ~20 µs on either backend. That is host-side noise; the graph modes do not show it.

End to end (SGLang, DeepSeek-V2-236B, TP 8 = DCP 8, FlashInfer MLA, 256K context, ISL 260992 / OSL 1024).

• SGLang main e49e971131 vs the same main with its DCP fi_a2a backend calling this op (strided input, no pack kernel).
• SGLang's default settings in every arm: custom all-reduce v2, the CUDA symmetric-memory backend, NCCL_CUMEM_ENABLE=0.
• Each job runs four servers on one node, two per path, in A/B/B/A order. Median TPOT in ms; Δ is per job.

job (order)	concurrency 1: MNNVL a2a + combine	this PR	Δ	concurrency 2 Δ	median ITL Δ, conc. 1 / 2 / 4
1 (MNNVL, PR, PR, MNNVL)	13.271	13.144	−0.96%	−0.40%	−0.93 / −0.66 / −1.03%
2 (PR, MNNVL, MNNVL, PR)	13.227	13.120	−0.81%	−0.63%	−0.88 / −0.90 / −0.79%
3 (MNNVL, PR, PR, MNNVL)	13.246	13.127	−0.90%	−0.66%	−0.90 / −0.95 / −0.81%
4 (PR, MNNVL, MNNVL, PR)	13.244	13.137	−0.81%	−1.43%	−0.84 / −0.88 / −0.87%
mean			−0.87% (4/4)	−0.78%	−0.89 / −0.85 / −0.88%

• Control: the unchanged NCCL a2a backend, main vs this branch, is +0.05% at concurrency 1.
• Decode-only step time at 32K context, batch 8 / 16 / 32: −1.19 / −1.14 / −1.52%, negative in all 4 jobs.
• Concurrency-4 TPOT is noisy: only 9 requests finish per run, and TPOT includes stalls behind other requests'
  261K-token prefills.
  • Six of the eight runs of this PR are 0.2–0.5% below main.
  • Two read 46.9 and 58.2 ms (main: 44.5–44.7) while their median ITL, 21.56 and 21.51 ms, matches the other runs of
    this PR.
  • Median ITL at concurrency 4 is −0.88% (−0.79 to −1.03% per job).
• Accuracy:
  • GSM8K, 5-shot, 50 questions, --parallel 1: SGLang main 0.80 and this op 0.76 in every run (2 of 50 questions; greedy
    decoding repeats exactly).
  • The full 1319 questions at --parallel 32 with CUDA graphs only up to batch 4, so most batches run the op eagerly:
    0.772 (this op) vs 0.766 (main). Without decode CUDA graphs: 0.767.
  • Greedy outputs of four canary prompts are token-identical between the two paths: 8 of 8 cross-path run pairs.

nsys (20 decode steps at concurrency 1 per arm; two captures, opposite orders). Per layer, on the critical path from the
last rank's attention end to the last rank's reduce end:

	MNNVL a2a + combine	this PR
kernels per layer (median)	pack 1.50 + helixAllToAllKernel 5.44 + combine 1.70 µs	FusedKernel 7.26 µs
critical path per layer	9.41 µs (incl. 1.4 µs of gaps between the three kernels)	7.54 µs
× 60 layers, per decode step	566 µs (4.3% of 13.2 ms)	453 µs (3.5%)

• Per step: −113 µs (nsys) vs −115 µs measured in serving.
• The microbenchmark predicts −93 µs (L2 hot) to −141 µs (L2 flushed), from 60 × the per-layer difference.
• Custom all-reduce v2 runs in both arms: 121 kernels per step, no v1.

Against #4929 in SGLang. #4929 cannot start next to SGLang's default all-reduce (its set_backend("NCCL") fails once
that has allocated). So this comparison runs SGLang's v1 all-reduce in every arm. #4929 gets contiguous input (two torch
copies), since it does not accept strided views. Two jobs, opposite arm orders; nsys from a third.

	MNNVL a2a + combine	#4929	this PR
TPOT, concurrency 1 (ms)	13.678	15.658 (+14.5%)	13.570 (−0.79%)
median ITL, concurrency 1 / 2 / 4		+14.5 / +12.9 / +8.5%	−0.74 / −1.16 / −0.72%
DCP critical path per decode step (nsys)	568 µs	2534 µs	445 µs
decode-shape microbenchmark, batch 1, per layer	8.09 µs	27.39 µs	6.54 µs

• feat(comm): add fused decode CP A2A + LSE reduce #4929 adds 1980 µs per decode step in serving and 1966 µs of DCP time in nsys.
• The microbenchmark predicts 1414–2011 µs per step for feat(comm): add fused decode CP A2A + LSE reduce #4929 (L2 hot to flushed, 60 layers, including the two copies).

Tests

tests/comm/test_dcp_lse_reduce.py:

• runs on the process's default backend; FLASHINFER_TEST_SYMM_MEM_BACKEND=NCCL runs it on NCCL;
• commit 1 adds:
  • strided views vs packed inputs and vs the reference;
  • row-count sequences;
  • call-number reuse and the 32-bit wrap;
  • 50 replays of a 16-call graph;
• commit 2 adds:
  • a workspace created next to an earlier group-named symmetric allocation (what SGLang's all-reduce does);
  • the backend is left alone;
  • a group whose ranks differ from the global ones, and a mismatched cp_rank rejected on every rank before launch.

🔍 Related Issues

Follow-up to #4929.

🚀 Pull Request Checklist

✅ Pre-commit Checks

• I have installed pre-commit by running pip install pre-commit (or used your preferred method).
• I have installed the hooks with pre-commit install.
• I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

🧪 Tests

• Tests have been added or updated as needed.
• All tests are passing (unittest, etc.).

At commit 2, torchrun ... -m pytest tests/comm/test_dcp_lse_reduce.py: 34 passed, 0 skipped on every rank in every run
below.

Platform	Backend	Ranks
8× B200, torch 2.14.1	CUDA (NCCL_CUMEM_ENABLE=0)	8, 4, 2
8× B200, torch 2.14.1	NCCL, via FLASHINFER_TEST_SYMM_MEM_BACKEND=NCCL and via TORCH_SYMMMEM=NCCL	8
8× B200, torch 2.13 (FlashInfer 0.6.18 with this op)	CUDA	8, 4
2 × 4 GB300, the CP group spanning two trays of one NVL72 domain	CUDA (fabric handles) and NCCL	8

• Negative control: commit 1 fails exactly the two backend tests: assert 'NCCL' == 'CUDA', and
  RuntimeError: Backend can not be changed after use.
• SGLang's integration test of this op (CUDA backend, with its custom all-reduce v2 in the same process): passes at 2,
  4 and 8 ranks.

Not in this PR

• An NVLink topology check. LL128 readiness needs 128-byte stores to land atomically, which NVLink provides and PCIe does
  not. The op checks that every rank's workspace is mapped; like the NCCL LSA check it replaces, that does not tell NVLink
  from PCIe peer-to-peer. A topology check could reuse flashinfer/comm/ulysses_topology.py.
• The CUDA-backend gap. It is at most ~3% at batch ≥ 2 and comes from torch's CUDA-backend peer mappings, not the
  kernel.
• A 16-byte LL variant was faster only at batch 1 (5.8 vs 6.5 µs per layer with the same launch).
• Programmatic dependent launch hid a further ~0.5 µs per call on the LL variant.

Reviewer Notes

• Readiness relies on 128-byte atomicity. A line's data is trusted once its flag is seen, which requires a
  warp-coalesced 128-byte store to become visible atomically. NCCL enables LL128 on NVLink on the same premise, and the
  MNNVL helix kernel in FlashInfer already relies on it. Every rank's workspace must be mapped over NVLink.
• The call counter is still device-side. It is read by every block at entry and advanced by the last block to finish,
  so it stays correct under CUDA-graph replay.
• Moving off NCCL's device API. The kernel used only ncclGetLsaPointer, i.e. peer addressing. torch's symmetric
  memory gives the same addresses on every backend.

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

• Autopilot · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}