[Attention][DCP] Add symm_a2a: single-node symmetric-memory A2A backend for MLA decode

[Cerdore]

Motivation

SGLang's existing single-node DCP (Decode Context Parallel) attention-reduction backends both route the exchange through NCCL collectives:

• ag_rs (default): NCCL all_gather + Triton correction + reduce_scatter (two collectives per layer).
• a2a: pack output + fp32 LSE → NCCL all_to_all_single → local Triton combine (one collective + pack/unpack copies per layer).

At high concurrency (conc=64+) with Kimi-K2.5 TP8/DCP8, this NCCL exchange sits on the critical path of every decode step and becomes a measurable bottleneck. The fi_a2a backend avoids NCCL via a FlashInfer MNNVL kernel, but requires MNNVL fabric memory (GB200 NVL72), unavailable on the common single-node NVLink/NVSwitch clusters.

This adds a symm_a2a backend targeting single-node NVLink/NVSwitch domains (e.g. H20 ×8, A100 ×8, H100 ×8): it allocates a peer-mapped symmetric memory slab via PyTorch _SymmetricMemory, and each rank directly writes its partial output + LSE into peer buffers via P2P, then runs a lightweight CUDA combine kernel locally — no NCCL involved. Eligibility is detected at init via can_use_custom_all_reduce_v2 (runtime check, not a hardware whitelist); validated on H20 ×8.

Modifications

• dcp_direct_a2a.cu — four CUDA kernels exposed via the torch op direct_dcp_a2a_lse_reduce, launched in sequence on the compute stream (all CUDA-graph capturable):
  1. increment_epoch_kernel — bumps the epoch counter; parity = epoch & 1 selects one of two staging slots for lock-free double-buffering (no host sync between graph replays).
  2. dispatch_output_lse_kernel — each rank writes its partial attention output + LSE directly into every peer's region of the symmetric slab, using vectorized uint4 P2P stores over NVLink/NVSwitch.
  3. signal_kernel — each rank writes its epoch value into every peer's signal slot with release semantics (st_flag_release_u64), marking its data ready.
  4. wait_lse_combine_kernel<scalar_t> — each rank spin-waits (acquire + timeout trap) on its local received-signal slots until all peers signal the current epoch, then performs the LSE-weighted softmax reduction over all ranks' received partial outputs.
• layers/dcp/comm.py — SymmA2AWorkspace: symmetric slab layout, pointer tables, epoch double-buffering, init/destroy. Wired into the DCP merge dispatch as a new backend alongside ag_rs/a2a/fi_a2a.
• base_runner.py — _pre_initialize_symm_a2a_workspace() pre-allocates the workspace before CUDA graph capture. Also fixes a device-index bug on this path: mr.device resolves to "cuda" (index=None), which made can_use_custom_all_reduce_v2 / NVLink P2P access checks evaluate against the default device rather than the rank-local GPU — pass cuda:{mr.gpu_id}.
• forward_mla.py — MLA decode path routes to symm_a2a when selected.
• server_args.py — --dcp-comm-backend symm_a2a CLI arg + validation (requires CUDA, single-node NVLink, no --enable-two-batch-overlap). Compatible with --dcp-replicate-q-proj.

Accuracy Tests

• Kernel unit tests (test_dcp_direct_a2a_kernel.py, 2-GPU and 4-GPU, fp16/bf16, with/without CUDA graph) — passed. Validates the numerical correctness of direct dispatch + LSE-weighted combine, including the CUDA-graph replay path.
• E2E correctness parity (test_dcp_symm_a2a_tp2.py TP2, test_dcp_symm_a2a.py TP4, DeepSeek-V2-Lite) — identical token sequences vs ag_rs.

Speed Tests and Profiling

Kernel micro-benchmark (bench_dcp_direct_a2a.py)

Kimi-K2.5 geometry (64 heads, 16/rank, D=512, BF16), 4×H20, TP4/DCP4. 4 ranks × 4 paths × 6 token counts (1/8/16/32/64/128), 10 warmup + 100 iterations, latency of the slowest rank (μs).

tokens	Direct eager	Direct graph	A2A+Triton eager	A2A+Triton graph	Direct graph vs A2A graph
1	25.6	14.7	131.1	27.8	1.89× faster
8	24.7	15.2	341.3	25.7	1.69× faster
16	24.3	17.3	131.5	29.4	1.70× faster
32	24.5	20.1	130.3	35.1	1.75× faster
64	31.4	28.0	129.7	47.2	1.69× faster
128	52.3	49.9	129.0	71.7	1.44× faster

Direct CUDA-graph path is 1.44×–1.89× faster than A2A+Triton CUDA-graph across all token counts; single decode step (tokens=1) drops 27.8 → 14.7 μs (-47%). Direct-path latency scales linearly with tokens (P2P write + local combine), while A2A-graph grows faster (NCCL all_to_all + pack/unpack). Note: micro-bench runs on idle GPUs with full NVLink bandwidth; the E2E speedup below is smaller because A2A is only a fraction of the decode step.

E2E serving benchmark

Kimi-K2.5 (555GB, MoE, MLA), H20 ×8, TP=8, DCP=8. 6 scenarios × 3 concurrency (8/32/64) × 2 backends = 36 rounds; 4 synthetic (decode-heavy / balanced / prefill-heavy / heavy-load) + 2 real datasets (ShareGPT V3, online prompts).

conc	throughput Δ	ITL p50 Δ
8	+0.5~1.0%	-0.6~2.1%
32	+1.2~3.5%	-1.5~2.8%
64	+7.6~10.3%	-10.1~12.5%

conc=64 averages across 6 scenarios: throughput +9.0%, ITL p50 -11.0%. Prefill-heavy (ISL=8192) conc=64: TTFT p50 -15.5%. Low concurrency is flat — communication is not the bottleneck there.

Trend matches vLLM PR #48897 (gains scale with concurrency); the SGLang-side ITL improvement is larger, likely because TP8/DCP8 exchanges more data per step than vLLM's TP4/DCP4.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to the Write documentations.
• Provide accuracy and speed benchmark results according to the Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Closes #33355.

---

CI States

Latest PR Test (Base): ❌ Run #30799521351
Latest PR Test (Extra): ❌ Run #30799520968

[gemini-code-assist]

Caution

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

[foraxe]

Hi @Cerdore, is this a sglang version of vllm-project/vllm#48897?

[Cerdore]

vllm-project/vllm#48897

Hi @foraxe — yes, it's the SGLang port of GirasoleY's vLLM #48897.

[foraxe]

Hi @Cerdore, Thanks for confirming.

One related point: SGLang #32851 already adopts the faster Output/LSE implementation from vLLM #50009, which improves upon the original vLLM #48897 approach.

In our apples-to-apples vLLM benchmark, vLLM #50009 was 17.85% faster than vLLM #48897 at the c32 geometry.

If it's useful, we're happy to share our implementation experience or benchmark results.
It would be great to align on a common high-performance backend for the community.