[DSA] Fix decode context parallelism for RoPE sparse attention on Blackwell

[vincentzed]

Motivation

• Fixes [Bug] DCP fails for GLM-5.3-NVFP4 on B300 #41926. With --dcp-size > 1, GLM-5.3, GLM-5.2 and DeepSeek-V3.2 crash. These models use RoPE sparse attention (DSA) with the trtllm backend. There are three causes:
  • forward_mla.py expects the attention backend to return (attn_output, lse), but the DSA backend only returns the output.
  • multi_ctas_kv_counter_buffer is sized for the heads of one rank, but with DCP each rank runs the heads of the whole DCP group.
  • The EAGLE draft pool asserts page_size == 64, and DCP uses a larger page size.
• This makes a documented flag work. It does not make DCP the fastest option: TP alone is still fastest when the load fits in memory, and attention DP is still faster at high concurrency.

Modifications

• Index remap: a new Triton kernel takes the global top-k slots, keeps the ones this rank owns, converts them to local slot ids, moves them to the front of the row, and returns the count per row.
• Backend: sparse attention is called with return_lse=True. Rows that have no local KV are fixed up, and the partial results are merged with the existing ag_rs / a2a / fi_a2a code. The fused top-k path stays on.
• Memory: the target model's KV is split across the DCP group. The index keys are kept in full on every rank. KV sizing now reserves room for the DCP workspace, the counters, the merge buffers and the replicated Q weights.
• Decode latency:
  • Layers that share the previous layer's top-k reuse its remapped table. This cuts the remaps from 78 to 21 per step.
  • With a2a / fi_a2a, every rank keeps the full Q projection weights by default, so there is no Q all-gather per layer. --no-dcp-replicate-q-proj turns this off.
• EAGLE/MTP: works with the checkpoint's own draft as a single chain. The draft KV is kept in full on every rank.
• Launch validation: flag combinations that are not supported raise an error at startup.
• Fixes in shared code:
  • The LSE combine returns zero when every shard is empty.
  • Eager autotuning no longer allocates a logits buffer that it never uses.
• FlashInfer workaround: in 0.7.0.post1, the sparse RoPE kernel with return_lse returns wrong values for short-KV FP8 rows when there are more than 32 heads. To stay within 32, the heads are split into groups and each group runs as extra batch rows. I have a reproducer and will file the upstream issue.

Not in this PR

• PD disaggregation.
• Prefill CP together with DCP.

Both raise an error at launch.

Follow-ups, most important first

1. A fused merge kernel. Today each layer runs three kernels to merge the partial results: pack, exchange and LSE combine. That is 234 launches per step. Doing all three inside the FlashInfer all-to-all is the biggest remaining latency item.
2. Splitting the index cache across ranks.
   • Today every rank keeps all index keys: 2,772 bytes per token on every rank, next to 44,928 bytes per token of KV that is split.
   • This costs 5.8% of capacity at DCP2, 17% at DCP4 and 41% at DCP8.
   • Splitting the index cache gets that back, but each indexer layer then needs one more collective (21 per step) to merge the local top-k candidates.
   • It is worth doing for DCP8, but not for DCP2 or DCP4.
3. HiSparse and HiCache / L3.
4. Unified memory.
5. LoRA with replicated Q. It works today with --no-dcp-replicate-q-proj.
6. Hopper (flashmla_kv).
7. DCP groups that span several nodes.

These also raise an error at launch: sparse backends other than trtllm, KV dtypes other than FP8 E4M3 or BF16, and DCP groups that span more than one attention-TP group. EAGLE needs the checkpoint's own draft, top-k 1, and draft tokens = steps + 1.

Related work

• Replaces [Feature] Decode context parallelism (DCP) for DSA models (DeepSeek V3.2, GLM-5.x) #31821.
• Not built on top of [DCP] Restore GLM-5.3 Flash decode CP and fix virtual KV addressing #39117 (NoPE path) or [DSA][DCP] Decode context parallel for the DSA attention family (GLM-5.3 / DeepSeek-V3.2) on Hopper #39330 (Hopper). Whichever lands second will need a rebase.
• [Feature] Support DSA PD with prefill CP and decode DCP #40117 covers Hopper PD, DCP8 and a sharded index cache.
• DeepEP is not changed and not tested.

Accuracy Tests

• 4x B300 (SM103), FlashInfer 0.7.0.post1. The tables were measured with torch 2.13 cu130 and sglang-kernel 0.4.8.
• Rerun on the current head (torch 2.14.1, sglang-kernel 0.4.9):
  • Unit and GPU test suites: 143 passed.
  • GSM8K: DCP2 and DCP4 both score 198 of 200. TPOT at C1 is 7.63 ms for DCP2 and 8.05 ms for DCP4.
  • Registered test (MTP fixture, completion API, 20-shot, 500 questions): TP4 + DCP2 scores 0.904 and TP4 without DCP scores 0.914. The accept length is 4.05 in both.
• Rerun after merging main at 0f5bc4af4b: the same test suites give 146 passed. TP4 + DCP2 scores 198 of 200 on GSM8K, and TPOT at C1 with 32K input is 7.67 ms.
• Rerun after merging main at e6072935bf: the same test suites give 148 passed. TP4 + DCP2 scores 198 of 200 on GSM8K, and TPOT at C1 with 32K input is 7.66 ms.
• Model: nvidia/GLM-5.3-NVFP4, FP8 KV, temperature 0.

Configuration	GSM8K
TP4	1283/1319 (97.3%)
TP4 + DCP2	1279/1319 (97.0%)
TP4 + DCP4	1279/1319 (97.0%)
DCP2 ag_rs / ag_rs + symm-mem / a2a	198 / 198 / 197 of 200
DCP4 ag_rs / ag_rs + symm-mem / a2a	197 / 197 / 196 of 200
DCP2 + EAGLE (steps 5, top-k 1, draft 6)	198/200
DCP2 + EP4	196/200
DCP2 + attention DP2	198/200
DCP2, BF16 KV	197/200
DCP2, zai-org/GLM-5.3 (FP8 weights)	199/200

• On 65K and 131K retrieval prompts, TP4, DCP2, DCP4 and DCP2 + EAGLE return the same answers.
• No regression without DCP: at TP4, this branch and main both score 198/200 and have the same KV capacity. TPOT is 6.714 ms on the branch and 6.713 ms on main.
• The full-set rows and the variant rows were measured before the two decode-latency changes. After those changes, DCP2 and DCP4 both score 198 of 200.
• Not tested: BF16 weights, SM100, GLM-5.2 and DeepSeek-V3.2 checkpoints, TP8/DCP8.

Speed Tests and Profiling

• All configurations use the same prompts, 1,024 output tokens and default flags. Attention DP uses sticky routing.
• The prefix cache is warm. The numbers are steady-state output tokens per second.
• The DCP columns are from the current head on a smaller grid. · means the cell was not rerun. The number in parentheses is the earlier measurement.

Input	C	TP4	DP2	DP4	DCP2	DCP4
32K	1	149	77	81	131 (119)	124 (117)
32K	16	1,300	927	955	1,169 (1,129)	1,112 (1,065)
32K	64	3,079	2,552	2,829	2,882 (2,692)	2,757 (2,596)
32K	128	–	4,338	4,710	4,469 (4,181)	· (3,808)
32K	256	–	–	7,866	–	4,441 (4,181)
128K	1	147	76	80	130 (118)	123 (116)
128K	16	1,185	892	935	1,076 (1,034)	1,036 (1,001)
128K	32	–	1,469	1,589	1,644 (1,598)	· (1,557)
128K	64	–	–	2,605	–	2,391 (2,273)

– means the load does not fit in the KV cache.

• The two decode-latency changes make DCP 3-10% faster. The order relative to TP4 and attention DP does not change.
• DCP is faster than attention DP up to C64 at 32K (DCP2) and up to C32 at 128K. Above that, attention DP4 is faster (1.8x DCP4 at 32K C256).
• At 32K C256, DCP4 cannot run all requests at once. 163 requests were waiting at the peak.

	TP4	DP2	DP4	DCP2	DCP4
Logical KV tokens	2.97M	5.44M	9.35M	5.36M	9.11M
Cold C1 TTFT, 32K (s)	1.24	1.44	1.86	1.89	2.28
Cold C1 TTFT, 128K (s)	5.85	6.68	8.27	8.40	10.05

Effect of the two decode-latency changes (32K, C1, 512 output tokens):

Configuration	TPOT before (ms)	TPOT now (ms)	GSM8K (200)
TP4 (reference)	6.78	6.78	198
DCP2	8.39	7.61	198
DCP2, --no-dcp-replicate-q-proj	8.39	8.25	198
DCP4	8.51	8.08	198
DCP2 + EAGLE	1.99	1.94	196
DCP2, ag_rs	8.87	8.77	196

Time per DCP2 decode step, by component:

Component	Before	Now
Q all-gather	0.60 ms, 78 launches	removed
Q packing	0.19 ms, 78 launches	removed
Index remap	0.20 ms, 78 launches	0.05 ms, 21 launches
Output/LSE pack + combine + fixup	0.34 ms	0.32 ms
Graph span	8.45 ms	7.80 ms

• KV capacity compared to TP4: 1.80x at DCP2 and 3.07x at DCP4.
• TPOT by comm backend before the two changes (DCP2, 32K C1): fi_a2a 8.39 ms, ag_rs + symm-mem 8.65 ms, ag_rs 8.87 ms, a2a 9.58 ms.

Checklist

• Format with pre-commit.
• Add unit tests and a registered four-GPU model test.
• Update docs/docs/advanced_features/dcp.mdx.
• Provide accuracy and speed results.

---

CI States

Latest PR Test (Base): ❌ Run #38232310139
Latest PR Test (Extra): ❌ Run #38232309598
Latest PR Test (AMD ROCm 10): ❌ Run #38232310077

[sfbemerk]

Many thanks for this PR! I just deployed it on B300s (TP=8, DCP=2) and it seems to work really well within its scopes!

[vincentzed]

@sfbemerk thanks for testing this on TP8.

I merged main into the branch. There was one conflict with #42178 in memory_pool.py. After the merge I ran it again on 4x B300 with TP4 + DCP2: GSM8K is 198/200, TPOT at 32K is 7.67 ms, and the unit tests pass.

cc @kpham-sgl @lmyybh

[vincentzed]

@mmangkad @ch-wan could you take a look at the DSA and memory pool changes? The PR also needs the run-ci label.

[kpham-sgl]

@vincentzed I can review this PR this week

[vincentzed]

I merged main into this branch again. #42671 was merged and it conflicted with this PR.

Two small changes came out of the merge:

• [DSA] Clarify pooled-page naming and ownership #42671 renamed slots_per_page to index_page_size. This PR already used the name index_page_size for something else, so I renamed that one to index_alloc_page_size.
• One test used real_page_size. It now uses the new name, physical_page_size.

I ran the tests again on 4x B300 after the merge:

• The DCP and memory pool unit tests pass (146 tests).
• GSM8K with TP4 + DCP2 on nvidia/GLM-5.3-NVFP4: 198 of 200 correct.
• Time per output token with a 32K prompt: 7.67 ms.
• No errors in the server log.

[vincentzed]

I merged main into this branch again because kv_cache_configurator.py had a conflict.

• main now passes the DSA pool class into _build_dsa_kv_pool ([Platform] Thread vendor KV pool classes through the in-tree builders; add PlatformCapabilities #35594). The DCP pool sizing is unchanged and uses that class.
• The DCP pool unit test passes its CPU pool through the new argument.
• I ran the test suites and the TP4 + DCP2 check on the merged branch again. The results are in the PR description under "Rerun after merging main".

[vincentzed]

I merged main into this branch again because pool_configurator.py had a conflict.

• main now subtracts a fixed scratch size before sizing the pool. The DCP padding pages are subtracted on the same line, so the merged code subtracts both.
• I ran the test suites and the TP4 + DCP2 check on the merged branch again. The results are in the PR description under "Rerun after merging main at e6072935bf".

@kpham-sgl I saw #43735, thank you for carrying this forward.

[vincentzed]

Done in #43735