[DCP] Optimize shared Output, Query, and Indexer/Top-K via peer memory by foraxe · Pull Request #32851 · sgl-project/sglang

foraxe · 2026-07-30T02:07:38Z

Summary

Depends on #31821.

This draft adds the complete opt-in Shared-DCP optimization for the DSA/MLA
decode path:

producer-direct Output/LSE stores into destination-owned peer mappings;
fused consumer-direct Query loads remote BF16 shards while applying RoPE
and FP8 conversion;
owner-sharded Indexer K storage and exact global Top-K, with direct peer
candidate reads for bounded decode.

The producer-fanout Query research route is not included. Large or ragged
prefill rows retain packed candidate AllGather after a direct ragged transport
experiment was measured, found substantially slower, and reverted.

Why

#31821 context-shards DCP attention, but it still:

materializes an Output/LSE collective;
AllGathers the BF16 Query before RoPE/FP8 conversion;
stores the complete Indexer K cache and repeats the same full token scan on
every DCP rank.

Shared-DCP makes these objects peer-addressable and aligns Indexer storage and
compute with logical-token ownership. The default remains the #31821 path.

Configuration

--dcp-comm-backend vmm
--dcp-query-backend vmm_direct
--dcp-indexer-backend owner_sharded
--dcp-topk-backend vmm

The selected routes require a same-host NVIDIA CUDA DCP group with peer access
and native peer atomics, and currently target the DeepSeek-V3.2/GLM-5.x DSA
MLA path.

Performance

Measured on four NVIDIA GB200 GPUs with GLM-5.2 NVFP4, FP8 KV cache, fixed
1800 MHz clocks, and the exact #31821 implementation as the same-lineage
control.

Three fresh server repetitions per DCP concurrency-1 cell:

Context	TP4	DCP2 #31821	DCP2 Shared-DCP	Delta	DCP4 #31821	DCP4 Shared-DCP	Delta
32K	7.476 ms	12.645 ms	12.408 ms	-0.238 ms (-1.88%)	13.013 ms	12.233 ms	-0.780 ms (-6.00%)
128K	7.612 ms	13.491 ms	12.837 ms	-0.654 ms (-4.85%)	13.914 ms	12.408 ms	-1.505 ms (-10.82%)

All DCP TPOT CVs are below 0.18%.

Ordinary NCCL A2A was the fastest existing collective on this host. The full
stack still lowers TPOT by 2.15%/5.14% for DCP2 at 32K/128K and
3.28%/8.27% for DCP4. FlashInfer fi_a2a was attempted at DCP2 and DCP4, but
this allocation rejected MNNVL fabric memory with
CUDA_ERROR_NOT_PERMITTED; no performance value is inferred.

KV capacity

physical tokens/rank is measured max_total_num_tokens; logical KV tokens
multiply by DCP size.

Parallelism	Native fraction	Physical tokens/rank	Logical KV tokens	Residual HBM	Versus TP4
TP4	0.85	912,832	912,832	14.85 GiB	1.00x
DCP2 #31821	0.82	795,456	1,590,912	10.30 GiB	1.74x
DCP2 Shared-DCP	0.82	795,456	1,590,912	18.27 GiB	1.74x
DCP4 #31821	0.70	368,192	1,472,768	28.02 GiB	1.61x
DCP4 Shared-DCP	0.70	368,192	1,472,768	39.20 GiB	1.61x

Measured frontier:

Parallelism	Control max	Control logical KV	Shared-DCP max	Shared-DCP logical KV	Gain	Versus TP4
DCP2	0.85	1,802,496	0.88	2,014,080	+211,584 (+11.74%)	2.21x
DCP4	0.76	2,317,568	0.85	3,585,280	+1,267,712 (+54.70%)	3.93x

Concurrency tradeoff

At 128K C4/C8, Shared-DCP improves output throughput by 0.49-2.30%, TPOT by
0.99-2.86%, and p99 ITL by 3.93-10.30%.

At 32K C16/C32 mixed prefill/decode, aggregate output throughput is
4.25-5.26% lower. Tail ITL improves by 1.11-11.97%, but TPOT is mixed. This
draft therefore claims long-context decode and KV capacity, not universal
throughput or prefill improvement.

Isolated mechanisms

Mechanism	DCP	Rows	Control median	Selected median	Mean delta
Output/LSE	2	1	354.18 us	108.34 us	-69.80%
Output/LSE	4	512	372.72 us	113.69 us	-69.31%
Query pipeline	2	1	230.19 us	108.26 us	-125.92 us (-53.84%)
Query pipeline	2	32	231.16 us	107.86 us	-125.18 us (-53.81%)
Top-K candidate merge	4	1	442.29 us	109.39 us	-334.03 us (-75.38%)
Top-K candidate merge	4	512	470.11 us	109.89 us	-359.37 us (-76.57%)

Validation

DCP2 and DCP4 Output/LSE eager and changing-input CUDA-graph equivalence;
DCP2 and DCP4 byte-exact FP8 Query, sparse MLA consumption, and
changing-input CUDA-graph replay;
DCP2 and DCP4 exact-set Top-K, deterministic lower-global-ID cutoff, and
changing-input CUDA-graph replay;
full integrated DCP2/DCP4 server CUDA-graph capture;
Black, Ruff F401,F821,UP037, py_compile, and git diff --check;
focused final-stack CPU suite: 109 passed, 4 skipped.

Accuracy

The final-current-main matrix is complete on the four-GPU allocation.

Configuration	GSM8K warmup	GSM8K counted	AIME25 warmup	AIME25 counted
TP4 reference	0.937	0.938	91.67% +/- 1.92%	94.17% +/- 1.67%
DCP2 #31821 control	0.943	0.946	95.83% +/- 3.19%	91.67% +/- 1.92%
DCP2 Shared-DCP	0.939	0.945	92.50% +/- 3.19%	92.50% +/- 4.19%
DCP4 #31821 control	0.944	0.940	90.83% +/- 3.19%	91.67% +/- 3.33%
DCP4 Shared-DCP	0.943	0.943	90.83% +/- 3.19%	92.50% +/- 7.39%

Pass zero is the warmup and is excluded from the counted result. AIME25 uses
30 problems x 4 stochastic repeats with a 64K reasoning cap. Every GSM8K pass
has zero invalid outputs and every AIME25 pass has a 0.00% error rate. Counted
AIME25 truncation rates at the cap are 0.83%/0%/0%/2.50%/0% for
TP4/DCP2-control/DCP2-Shared/DCP4-control/DCP4-Shared. The earlier interrupted
DCP2 counted AIME25 run is excluded; the table uses its fresh 120-sample
replacement restarted from sample zero.

Coverage limits

Measured: TP4, TP4/DCP2, TP4/DCP4 on four GPUs.
TP8: UNMEASURED; no eight-GPU allocation was available.
Direct candidate reads: bounded decode rows; ragged prefill uses candidate
AllGather.
Multi-node and non-DSA paths are unchanged.

Review structure

The branch keeps three explicit commit boundaries for independent review:

Output/LSE: d16fd9994..2a2b15629
Query: 2a2b15629..5e8618eaf
Indexer/Top-K: 5e8618eaf..6f984f7a3

CI States

Latest PR Test (Base): ❌ Missing run-ci label -- add it to run CI tests.
Latest PR Test (Extra): ❌ Blocked -- run-ci is required first.

…3.2, GLM-5.x) Based on the merged MLA DCP path (sgl-project#14194 lineage): latent KV stays interleaved across DCP ranks (masked/divided write already inside set_mla_kv_buffer's kernel), while the DSA indexer K cache is kept replicated (global slots, sized *dcp) so every rank scores the full sequence and produces identical top-k. Owned selections are filtered and mapped to local rows inside the index transforms; sparse flashmla kernels return their LSE and outputs are combined across the DCP group (existing cp_lse_ag_out_rs_mla for the head-widened decode, new cp_lse_ag_out_ar all-reduce variant for extend). v1 scope: --dsa-{decode,prefill}-backend flashmla_sparse, bf16 KV cache, no hisparse, unfused top-k (fused variants bake in full-sequence index math). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

… groups); contiguous LSE for the combine Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Under DCP decode the indexer (replicated, full-sequence paged scoring) is independent of the q-side pipeline, so fork it onto the existing alt stream and join after the cross-rank q all-gather is issued: comm hides behind indexer compute. GLM-5.2-NVFP4 tp4+dcp2 32k-ctx decode: +11% output throughput, -6% mean ITL; gsm8k parity unchanged. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

…ff by default), env knobs - Row-slice the head-padded flash_mla_sparse call under DCP (SGLANG_DSA_DCP_SPARSE_MAX_ROWS, default 8192, ~5% extend overhead) — fixes the OOM at large chunked-prefill sizes without lowering --mem-fraction-static. - Fused top-k v2 under DCP: correct via a non-mutating owner filter on the fused output (IndexShare-safe), but the v2 plan/persistent-pool is decode-shaped and ~2x slower on 32k-row extend chunks with no measured decode win — disabled by default, SGLANG_DSA_DCP_FUSED_TOPK=1 to experiment. Follow-up: extend-shaped plan (work-proportional routing + request-blocked tiling). - Speculative decoding stays rejected with DCP: EAGLE draft-graph capture crashes on current main with this flag set even without DCP (upstream Triton IMA), so there is no baseline to validate against. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

trtllm-gen already surfaces the base-2 LSE (ComputeLSEFromMD) through flashinfer's return_lse, so both decode and the decode-ized extend flow through the existing cross-rank combine. The owner-filtered top-k rows are compacted valid-first (attention is permutation-invariant over the selected set) with per-row local counts as seq_lens — the trtllm-gen sparse kernel scans a counted prefix, unlike flashmla's -1 skipping. Zero-count rows contribute exp2(-inf)=0. The LSE request needs a larger softmax-stats workspace (2x under DCP). This makes the resolved sm100 defaults (trtllm + fp8_e4m3 KV) the primary DCP configuration; flashmla_sparse + bf16 KV remains supported. GLM-5.2-NVFP4 tp4+dcp2, default flags: gsm8k 0.640 vs 0.610 baseline. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

…lse_ag_out_ar SGLANG_DSA_DCP_SPARSE_MAX_ROWS and SGLANG_DSA_DCP_FUSED_TOPK move to envs.EnvInt/EnvBool (malformed values now warn and fall back instead of crashing at import). cp_lse_ag_out_ar was the combine for the abandoned same-query extend scheme (wrong across DCP head groups) — removed so nobody wires it back in; extend uses the gathered-q cp_lse_ag_out_rs_mla path. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

- Coerce dsa_prefill_impl=flashmla_auto to flashmla_sparse under DCP instead of hard-failing before auto-select runs (non-sm100 defaults). - Gate the 2x LSE workspace on the trtllm impls actually being used (saves 384MB on flashmla_sparse+DCP runs). - Initialize lse=None in forward_absorb_core and assert in the DCP combine so a fusion path skipping the LSE fails loudly, not unbound. - Union return annotation on _forward_flashmla_sparse. - Unit tests for the DCP owner-filter branches of both index transforms (test_dsa_transform_index.py, 2 passed on B300). Remaining comments triaged as noise: negative-index guards (decode kernel's masked lanes never store computed values and C-remainder of -1 never matches a rank; prefill already guards via valid_topk_mask), backend-level spec assert (server_args already rejects spec+DCP at launch). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

…P combos - Revert test_dsa_transform_index.py additions (keep the PR code-only). - Skip building the dense-MLA DCP extend gather metadata for DSA models (unused by the sparse local-shard extend path). - Reject --enable-mixed-chunk and --enable-dp-attention under DSA DCP: MIXED batches would walk the DCP extend path with mixed rows, and dp-attention currently overflows the trtllm softmax workspace besides being structurally unvalidated (both are follow-up work). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

- Mixed chunk needs no new machinery: decode and extend rows share the row-wise DCP recipe (gathered q -> local-shard sparse attention -> LSE combine). Validated under concurrent load (gsm8k 0.620, flashmla dcp2). - dp-attention composes when each DCP group nests inside one attention-DP shard (attn_tp_size % dcp_size == 0, asserted): ranks in a group then share requests and the replicated indexer keeps producing identical top-k. The trtllm LSE workspace now scales with the dcp-widened head count, and dp idle-companion batches skip the gather/combine. tp4 dp2 dcp2 on the default trtllm+fp8 stack: gsm8k 0.648 (500q) vs 0.635 dp-only baseline. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

…TOPK (opt-in ENABLE_ verb per env-var conventions) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

…ecipe The earlier 'blocked upstream' conclusion was wrong: the EAGLE draft-capture IMA was the sgl-project#30209-class flashinfer_trtllm bf16 draft-MoE instability, avoided by the production draft-MoE pins (--speculative-moe-runner-backend triton --speculative-moe-a2a-backend none). With those pins, spec works on this branch's stacks, so the CUDA spec x DCP rejection becomes an experimental warning and the three DSA DCP mode gates (q gather, attention branch, LSE combine) include DRAFT_EXTEND_V2 — draft rows use the same row-wise recipe; the dense-MLA draft path stays unimplemented. GLM-5.2-NVFP4 tp4 dcp2 + EAGLE 5-1-6 + triton draft-MoE: serves with accept length 4.4-4.6; gsm8k 0.590 (500q) vs no-spec dcp2 0.622-0.648 — functional, verify-path numerics validation outstanding. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

…ft-MoE pin The triton draft-MoE pin itself costs ~4.5 gsm8k points (0.592 vs 0.636 at 500q, no DCP, otherwise identical) — the earlier 0.590 spec x DCP reading was entirely that recipe, not DCP. On the default stack (trtllm attention, fp8 KV, flashinfer_trtllm draft MoE) with flashinfer >= 0.6.15 (contains the flashinfer#3238 draft-MoE fix): spec no-DCP 0.636, spec + dcp2 0.650 (500q), accept length 4.35-4.61. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Same two dsa_backend.py conflicts as the earlier Helix-branch merge — upstream sgl-project#30514's Q8KV8 restructure vs the DCP init block and the extend flashmla_sparse call; resolved keep-both / concat-then-DCP. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

The DCP extend recipe all-gathers q across the DCP group, which requires every rank in the group to hold identical extend rows. Prefill-CP splits rows across ranks, so DCP-group peers hold different tokens and the gathered q mixes unrelated rows. Short prompts previously bypassed CP engagement, which masked this; the batch-prep rework in the upstream window (post sgl-project#30514) routes short extends through CP and surfaced it as zero accuracy. Reject the combination at launch until the CP-split extend path gathers position-ordered KV across the DCP group instead (all_gather_kv_cache_for_mla_extend-style, RAGGED indexing). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

… path The trtllm backends are the sm100 defaults and the only stack anyone should serve with; the flashmla_sparse DCP support existed only as a bring-up vehicle and carried real weight: an LSE tuple return + row chunking (with its SGLANG_DSA_DCP_SPARSE_MAX_ROWS env) to bound the 128-head padding transients, a flashmla_auto coercion, and dtype asserts. Remove all of it; DCP now requires trtllm decode+prefill at init. _forward_flashmla_sparse returns to its upstream shape. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

…rlap block, rename slot mapper - transform_index / dsa_backend / init: comments cut to one or two lines - delete the dangling flashmla row-slice comment - workspace multiplier computed in _dcp_workspace_multiplier instead of inline in the buffer lambda - _apply_dcp_owner_filter -> _dcp_global_slots_to_local_rows - forward_mla capture-overlap block: shared _indexer_topk closure instead of two duplicated arms - spec x DCP warning no longer talks about old flashinfer releases Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

- init guards reordered (backend requirement first, then composition, then toggles) with assert messages cut to the point - slot mapper uses scalar torch.where - extend gather branch restructured as two positive arms, no-op elif/pass arms dropped (idle needs nothing; dense draft-extend simply has no gather) - comments in memory_pool/environ/kernels cut to one or two lines Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

# Conflicts: # python/sglang/srt/layers/attention/dsa_backend.py

@b8zhong

…buffer for DCP The persistent multi-CTAS KV counter buffer (sgl-project#31607) is sized by head count; under DCP the trtllm decode runs with dcp-widened heads, so pre-size it by num_q_heads * dcp_size and let the grow check use the per-call head count. token_capacity in metrics and the load inquirer now reports logical tokens (physical * dcp_size), matching the allocator's accounting. Fix by @b8zhong, thanks. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

…fixup kernel Route the trtllm decode LSE output through a pre-registered symmetric-memory buffer so cp_lse_ag_out_rs_mla's all-gather needs no separate cast/copy, and replace the zero_rows torch.where masking with the fixup_zero_kv_rows kernel (rows where this rank owns none of the selected tokens were producing NaN, not 0, so the prior masking was incomplete).

# Conflicts: # python/sglang/srt/layers/attention/dsa_backend.py

The init-time disable turned fusion off for decode too, which costs decode perf for no reason: the per-mode force_unfused gate already keeps extends on the unfused transform (the v2 plan is decode-shaped), and the fused decode table goes through the owner filter. The SGLANG_DSA_DCP_ENABLE_FUSED_TOPK escape hatch forced fusion onto extends, which produces wrong indices; with the per-mode gate there is no configuration that reaches a fused extend, so the env var is gone. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Re-enabling fused decode (c192ef9) drops gsm8k from 0.920 to 0.000 at dcp2 with everything else fixed, so the v2 fused transform no longer composes with the DCP owner filter after the recent upstream churn. Keep the unfused path until that is root-caused; the incorrect SGLANG_DSA_DCP_ENABLE_FUSED_TOPK escape hatch stays removed. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

# Conflicts: # python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py # python/sglang/srt/server_args.py

…-latest

github-actions · 2026-09-28T01:33:30Z

Thanks @foraxe. Closing this because you have 11 PRs open and none updated in 7 days, over our soft cap of 5 for idle PRs.

Reopen the two or three you most want reviewed - a short active set moves faster than a long idle one.

Some directories moved recently, so an older branch may need retargeting:
sgl-kernel/ -> python/sglang/kernels/aot/, python/sglang/jit_kernel/
-> python/sglang/kernels/jit/, docs/ -> docs/docs/ (.mdx),
bench_serving.py -> benchmark/serving.py, test/srt/ -> test/registered/.

brayden and others added 30 commits July 19, 2026 03:46

fix: DSA extend under DCP all-gathers q (ranks hold different TP head…

ffd6236

… groups); contiguous LSE for the combine Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

env: rename SGLANG_DSA_DCP_FUSED_TOPK -> SGLANG_DSA_DCP_ENABLE_FUSED_…

aeef46e

…TOPK (opt-in ENABLE_ verb per env-var conventions) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Fix lint: ruff unused import + black formatting

eb88776

Merge remote-tracking branch 'origin/main' into dcp-dsa

59518f4

# Conflicts: # python/sglang/srt/layers/attention/dsa_backend.py

Merge remote-tracking branch 'vincentzed/dcp-dsa' into dcp-dsa

f465988

# Conflicts: # python/sglang/srt/layers/attention/dsa_backend.py

Merge remote-tracking branch 'origin/main' into dcp-dsa

a0aac34

Merge current main into SGLang DSA DCP base

181e3c8

# Conflicts: # python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py # python/sglang/srt/server_args.py

Merge remote-tracking branch 'origin/main' into agent/shared-dcp-base…

d16fd99

…-latest

[DCP] Add producer-direct VMM output reduction

51ebfb0

[DCP] Fix CUDA 13 VMM handles and add benchmark

a305bac

[DCP] Finalize Output landing validation

26371f4

b8zhong mentioned this pull request Aug 5, 2026

[NVIDIA] DeepSeek V4 Perf Tracking #33636

Open

15 of 41 tasks

foraxe marked this pull request as ready for review August 12, 2026 09:35

github-actions Bot closed this Sep 28, 2026

Navigation Menu

Repository navigation

[DCP] Optimize shared Output, Query, and Indexer/Top-K via peer memory - #32851

[DCP] Optimize shared Output, Query, and Indexer/Top-K via peer memory#32851
foraxe wants to merge 34 commits into
sgl-project:mainfrom
foraxe:agent/shared-dcp-integrated-latest

foraxe commented Jul 30, 2026 •

edited by github-actions Bot

Loading

Uh oh!

github-actions Bot commented Sep 28, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

foraxe commented Jul 30, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Why

Configuration

Performance

KV capacity

Concurrency tradeoff

Isolated mechanisms

Validation

Accuracy

Coverage limits

Review structure

CI States

Uh oh!

github-actions Bot commented Sep 28, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

foraxe commented Jul 30, 2026 •

edited by github-actions Bot

Loading