[YAMY1234]

Motivation

The TRTLLM MHA decode backend does not support decode context parallelism (DCP). This PR adds the missing DCP decode path while keeping prefill on a supported backend. The common GQA KV-head mapping fix is split into #32858.

Modifications

• Build DCP-local sequence lengths and token-indexed TRTLLM page tables.
• All-gather Q heads, request per-rank output/LSE, and merge partial attention results with the existing DCP collective.
• Add DCP metadata for target verify and draft extend, plus graph-stable Q gather and output/LSE buffers.
• Preserve the configured prefill backend for ordinary draft prompt prefill.
• Add targeted page-table, graph-metadata, LSE, Q-layout, and backend-routing tests.
• Reject unsupported DCP combinations: SWA KV pools, native FP4 KV cache, and XQA.

Accuracy Tests

Qwen3.5-397B-A17B-NVFP4-V2 was evaluated on one GB300 node with TP4, 200 GSM8K examples, 8-shot prompting, temperature 0.6, max_tokens=16384, and full decode CUDA Graph through batch size 200. The DCP4 runs include #32858.

Layout	Decode backend	GSM8K
TP4/DCP1	Triton	0.990
TP4/DCP4	Triton	0.980
TP4/DCP4	TRTLLM MHA	0.975

All runs completed 200/200 examples with no missing or unextractable answers. The TRTLLM MHA implementation also passed 154 targeted page-table and CUDA Graph metadata tests on ARM64/GB300 and x86/GB300, plus eager and full-graph DCP1/DCP4 serving smoke tests.

Speculative Decode Validation

Validated against merged FlashInfer #4518.

Test	Configuration	Result
SGLang targeted tests	DCP metadata/page table, base-2 LSE merge, graph-stable Q gather, and draft prefill/decode routing	10 passed
FlashInfer kernel tests	D256, FP8 KV, page size 64, Q/KV heads 16/1, CP4, q_len 1/2/4/7/8, CUDA Graph replay	6 passed
End-to-end serving	TP16/DCP4, MTP4, Triton prompt prefill, TRTLLM MHA target/draft attention, full target-verify/draft-decode/draft-extend CUDA Graph, two deterministic 8192-input/8-output requests	Passed; identical output IDs and speculative acceptance metadata

The end-to-end smoke used dummy weights, so it validates integration, shape routing, and CUDA Graph replay rather than model quality. This path requires the first stable FlashInfer release containing #4518; the current SGLang FlashInfer pin (0.6.17) does not contain it.

Speed Tests and Profiling

End-to-end serving was measured on 16 GB300 GPUs with FP8 weights, FP8 E4M3 KV cache, TP16, ISL 8192, OSL 1024, concurrency 1024, 1,024 warmup requests, 5,120 measured requests, and full decode CUDA Graph through batch size 256.

Layout	GQA prefill / decode	Output tok/s	Total tok/s	Total tok/s/GPU	Mean TTFT	Mean TPOT
TP16	TRTLLM MHA / TRTLLM MHA	1,326.19	11,935.67	745.98	502.17 s	231.85 ms
TP16/DCP1	Triton / TRTLLM MHA	1,304.11	11,736.98	733.56	512.69 s	233.69 ms
TP16/DCP4	Triton / TRTLLM MHA	1,822.46	16,402.15	1,025.13	168.89 s	397.20 ms

DCP1 and DCP4 differ only in dcp_size. At concurrency 1024, DCP4 improves output throughput by 39.75% over DCP1 and 37.42% over TP16, with a higher per-request TPOT. All 5,120 measured requests completed with exact 8K/1K lengths, zero request errors, zero empty outputs, and zero HTTP non-200 responses.

This CC1024 result was collected on a stacked validation branch containing this PR, #32858, #32800, and the high-concurrency output-buffer capacity guard found during the run; the last guard is not yet in this PR head.

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 Write documentations. No user-facing documentation change is required for this backend integration.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): ❌ Run #33147127468
Latest PR Test (Extra): ❌ Run #33147127407
Latest PR Test (AMD ROCm 7.2): ❌ Run #33147127409

[gemini-code-assist]

Caution

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

[gemini-code-assist]

Caution

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

[kpham-sgl]

Might be blocked by flashinfer-ai/flashinfer#4323

[nvpohanh]

cc @thanhhao98 for vis

[YAMY1234]

Depends on flashinfer-ai/flashinfer#4518, will rebase after flashinfer version bump