Feat[DCP] Support decode context parallelism for DSA models with the flashmla_kv backend by Zhiy-Zhang · Pull Request #36990 · sgl-project/sglang

Zhiy-Zhang · 2026-08-29T07:51:10Z

co-author：https://github.com/Lai-JX

Motivation

Inspired by PR #31821, this PR builds on it to improve DCP support for NSA models.The main changes introduced in this PR are as follows:

DCP support is available for DSA models such as GLM-5.2. On Hopper-architecture hardware, only the flashmla_kv backend is currently supported.
DCP supports both TP-DCP (all_gather q)and CP-DCP (all_gather kv). CP-DCP is compatible with Decode Slice and outperforms TP-DCP when KV cache capacity is not the bottleneck.
Both FP16 and FP8 KV cache formats are supported. When the KV cache uses FP8, the DCP temporary buffer uses the same data type as the persistent KV cache—packed FP8—avoiding quantization and dequantization overhead.
Add unit tests covering various scenarios

Accuracy Tests

tp (baseline)
server command

python -m sglang.launch_server \
  --model-path /home/admin/GLM-5.2-W4AFP8 \
  --quantization w4afp8 \
  --disable-shared-experts-fusion \
  --tp 8 \
  --kv-cache-dtype fp8_e4m3 \
  --reasoning-parser glm45 \
  --tool-call-parser glm47 \
  --chunked-prefill-size 8192 \
  --max-running-requests 64 \
  --page-size 64 \
  --mem-fraction-static 0.70 \
  --trust-remote-code \
  --attention-backend nsa \
  --port 30100 \
  --dsa-decode-backend flashmla_kv \
  --dsa-prefill-backend flashmla_kv

accuracy result

100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1314/1314 [11:30<00:00,  1.90it/s]
Total latency: 690.054 s
Score: 0.948
Output throughput: 194.914 token/s
[METRIC] gsm8k_score=0.9482496194824962 labels={"model": "/home/admin/GLM-5.2-W4AFP8", "eval": "gsm8k"}
[METRIC] gsm8k_latency=690.0544659122825 labels={"model": "/home/admin/GLM-5.2-W4AFP8", "eval": "gsm8k"}
====================
Speculative decoding: no per-request spec_accept_length in responses (non-speculative server, or --api completion which lacks return_meta_info).
====================

cp-dcp
server command

export SGLANG_ENABLE_CP_V2=1
python -m sglang.launch_server \
  --model-path /home/admin/GLM-5.2-W4AFP8 \
  --quantization w4afp8 \
  --disable-shared-experts-fusion \
  --tp 8 \
  --attn-cp-size 8 \
  --enable-prefill-cp \
  --cp-strategy interleave \
  --dcp-size 8 \
  --enable-cp-decode-attn-tp \
  --kv-cache-dtype fp8_e4m3 \
  --reasoning-parser glm45 \
  --tool-call-parser glm47 \
  --chunked-prefill-size 8192 \
  --max-running-requests 64 \
  --page-size 64 \
  --mem-fraction-static 0.70 \
  --trust-remote-code \
  --attention-backend nsa \
  --port 30100 \
  --dsa-decode-backend flashmla_kv \
  --dsa-prefill-backend flashmla_kv

accuracy result

100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1314/1314 [23:31<00:00,  1.07s/it]
Total latency: 1411.064 s
Score: 0.951
Output throughput: 96.001 token/s
[METRIC] gsm8k_score=0.9512937595129376 labels={"model": "/home/admin/GLM-5.2-W4AFP8", "eval": "gsm8k"}
[METRIC] gsm8k_latency=1411.0640860423446 labels={"model": "/home/admin/GLM-5.2-W4AFP8", "eval": "gsm8k"}
====================
Speculative decoding: no per-request spec_accept_length in responses (non-speculative server, or --api completion which lacks return_meta_info).
====================

tp-dcp
server command

python -m sglang.launch_server \
  --model-path /home/admin/GLM-5.2-W4AFP8 \
  --quantization w4afp8 \
  --disable-shared-experts-fusion \
  --tp 8 \
  --dcp-size 8 \
  --kv-cache-dtype fp8_e4m3 \
  --reasoning-parser glm45 \
  --tool-call-parser glm47 \
  --chunked-prefill-size 8192 \
  --max-running-requests 64 \
  --page-size 64 \
  --mem-fraction-static 0.70 \
  --trust-remote-code \
  --attention-backend nsa \
  --port 30100 \
  --dsa-decode-backend flashmla_kv \
  --dsa-prefill-backend flashmla_kv

accuracy result

100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1314/1314 [17:29<00:00,  1.25it/s]
Total latency: 1049.677 s
Score: 0.944
Output throughput: 128.980 token/s
[METRIC] gsm8k_score=0.9444444444444444 labels={"model": "/home/admin/GLM-5.2-W4AFP8", "eval": "gsm8k"}
[METRIC] gsm8k_latency=1049.677041195333 labels={"model": "/home/admin/GLM-5.2-W4AFP8", "eval": "gsm8k"}
====================
Speculative decoding: no per-request spec_accept_length in responses (non-speculative server, or --api completion which lacks return_meta_info).
====================

Speed Tests and Profiling

The tests were conducted using an open-source agent workload dataset: Inferact/codex_swebenchpro_traces.
Here, Concurrency refers to the number of concurrent sessions, while Turns refers to the number of interaction rounds.

More detailed performance metrics, such as TTFT and TPOT. **Latency summary (Concurrency = 8， Turns = 10)**

**Latency summary (Concurrency = 6， Turns = 10)**

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.
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

Ping Merge Oncalls to start the process. See the PR Merge Process.
Get approvals from CODEOWNERS and other reviewers.
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
After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

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.
Latest PR Test (AMD ROCm 7.2): ➖ No AMD PR run found for this commit.

huangzhilin-hzl · 2026-09-01T09:42:43Z

+            if parallel.dcp_enabled:
+                # The indexer K cache is not sharded under DCP: every rank


The allocator can produce locations up to (size + page_size) * dcp_size - 1, but the indexer cache only covers size * dcp_size + page_size slots. Accessing the final widened page may cause an OOB GPU write. Please allocate the full (size + page_size) * dcp_size span.

huangzhilin-hzl · 2026-09-01T09:48:54Z

@@ -4520,6 +4473,12 @@ def __init__(
        self.index_head_dim = index_head_dim
        if index_buf_size is None:
            index_buf_size = size
+            parallel = get_parallel()
+            if parallel.dcp_enabled:


This multiplies the actual indexer allocation by dcp_size, but the pool configurator still charges for one copy. Auto-calculated max_total_tokens may therefore overcommit memory and OOM during pool initialization. ref: _compute_dsa_indexer_cell_size

huangzhilin-hzl · 2026-09-01T09:52:44Z

+        parallel = get_parallel()
+        self.dcp_enabled = parallel.dcp_enabled
+        self.dcp_size = parallel.attn_dcp_size if self.dcp_enabled else 1
+        self.dcp_rank = parallel.attn_dcp_rank if self.dcp_enabled else 0


This bucket is based on the pre-gather local head count. For a 128-head model with TP8/DCP8, it selects 64 from 16 local heads, but FlashMLA receives 128 heads after Q all-gather, causing a shape mismatch and incorrect metadata. Please select the bucket per phase using the effective head count: local_heads * dcp_size for Q-gather paths and local_heads for CP-DCP global-KV prefill.

huangzhilin-hzl · 2026-09-01T09:54:56Z

+            if model_runner.server_args.enable_dp_attention:
+                # Keep each DCP group inside one attention-DP shard so the
+                # replicated indexer sees identical requests group-wide.
+                if parallel.attn_tp_size % self.dcp_size != 0:


Checking attn_tp_size alone rejects the standard DSA CP topology, which intentionally uses attn_tp_size=1 (e.g. TP8/CP8/DCP8 fails as 1 % 8 != 0). The requirement is that a DCP group must not cross an attention-DP shard, whose span is attn_tp_size * attn_cp_size. Please validate that span and its group alignment instead.

huangzhilin-hzl · 2026-09-01T09:58:56Z

+            dcp_meta = forward_batch.attn_dcp_metadata
+            if dcp_meta is None or dcp_meta.dcp_kv_buffer is None:
+                raise RuntimeError("DSA CP+DCP prefill requires a gathered KV buffer.")
+            dcp_page_table_1 = self._build_dcp_prefill_page_table(


indexer_seq_lens_cpu has already been filtered by bs_idx, while dcp_meta.dcp_kv_indptr still describes the original full batch. If a CP rank owns no tokens for some requests, this either fails the batch-size check or maps the remaining request to the wrong KV segment

huangzhilin-hzl · 2026-09-01T10:00:53Z

-        if gpu_mem is not None and gpu_mem > 60 * 1024:
-            reserved_mem = max(reserved_mem, 10 * 1024)
-        return reserved_mem
+        if self.dcp_size > 1 and is_cuda() and self.speculative_algorithm is not None:


This global guard breaks existing non-DSA configurations, including the registered Kimi K3 DCP8 + DSPARK path. Please scope this restriction to DSA with flashmla_kv, preferably in the DSA backend’s model-aware validation, rather than rejecting every CUDA DCP configuration.

huangzhilin-hzl · 2026-09-01T10:03:00Z

Under CP-v2 zigzag, materialize_full_mla_kv() writes the gathered KV into the cache and returns None, but this caller unconditionally unpacks a tuple, causing a TypeError.

huangzhilin-hzl · 2026-09-01T10:07:15Z

If every shard has LSE=-inf, all weights are zero and this computes 0 / 0, reintroducing NaNs even when the partial outputs were cleared beforehand. Please use a safe denominator and explicitly return zero output (and -inf global LSE) when weight_sum == 0

huangzhilin-hzl · 2026-09-01T10:08:08Z

+            quantize_k_cache_separate,
+        )
+
+        current_nope, current_rope = quantize_k_cache_separate(k_nope, k_pe)


These rows are quantized here for the temporary gathered buffer, then the same k_nope/k_pe tensors are quantized again when writing the persistent KV cache.

…flashmla_kv backend Co-authored-by: Lai-JX <76830814+Lai-JX@users.noreply.github.com>

Zhiy-Zhang · 2026-09-06T11:34:17Z

@huangzhilin-hzl Thank you very much for your review comments. Most of the boundary issues have been fixed. The only remaining issues are related to the DCP and MTP implementations, which will be addressed after #31785 is merged.

Zhiy-Zhang · 2026-09-09T09:52:49Z

@kpham-sgl Hi, I noticed #36507 has already been merged. Could you share some feedback on the implementation in this PR?

github-actions Bot added deepseek jit-kernel memory-pool labels Aug 29, 2026

Zhiy-Zhang mentioned this pull request Aug 29, 2026

[Feature] Decode context parallelism (DCP) for DSA models (DeepSeek V3.2, GLM-5.x) #31821

Open

huangzhilin-hzl reviewed Sep 1, 2026

View reviewed changes

qiufan.zzy and others added 2 commits September 6, 2026 19:27

Feat[DCP] Support decode context parallelism for DSA models with the …

4b3fe69

…flashmla_kv backend Co-authored-by: Lai-JX <76830814+Lai-JX@users.noreply.github.com>

Fix[DCP]: Fix boundary issues

c6aeb8b

Zhiy-Zhang force-pushed the gitlab_dc_dcp_glm52 branch from e11018e to c6aeb8b Compare September 6, 2026 11:29

leoong55 mentioned this pull request Sep 10, 2026

GLM-5.3: CP8 + DCP4 on 0bcd822377da, v3 working baseline leoong55/sglang-lasolovev#1

Draft

kpham-sgl mentioned this pull request Sep 14, 2026

[Roadmap][DCP] Decode Context Parallelism & Helix Parallelism (2026 Q3) #29736

Open

36 tasks

Navigation Menu

Uh oh!

Feat[DCP] Support decode context parallelism for DSA models with the flashmla_kv backend - #36990

Feat[DCP] Support decode context parallelism for DSA models with the flashmla_kv backend#36990
Zhiy-Zhang wants to merge 2 commits into
sgl-project:mainfrom
Zhiy-Zhang:gitlab_dc_dcp_glm52

Zhiy-Zhang commented Aug 29, 2026 •

edited by github-actions Bot

Loading

Uh oh!

huangzhilin-hzl Sep 1, 2026

Uh oh!

huangzhilin-hzl Sep 1, 2026

Uh oh!

huangzhilin-hzl Sep 1, 2026

Uh oh!

huangzhilin-hzl Sep 1, 2026

Uh oh!

huangzhilin-hzl Sep 1, 2026

Uh oh!

huangzhilin-hzl Sep 1, 2026

Uh oh!

huangzhilin-hzl Sep 1, 2026

Uh oh!

huangzhilin-hzl Sep 1, 2026

Uh oh!

huangzhilin-hzl Sep 1, 2026

Uh oh!

Zhiy-Zhang commented Sep 6, 2026

Uh oh!

Zhiy-Zhang commented Sep 9, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		if parallel.dcp_enabled:
		# The indexer K cache is not sharded under DCP: every rank

Uh oh!

Conversation

Zhiy-Zhang commented Aug 29, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Accuracy Tests

Speed Tests and Profiling

Checklist

Review and Merge Process

CI States

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Zhiy-Zhang commented Sep 6, 2026

Uh oh!

Zhiy-Zhang commented Sep 9, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Zhiy-Zhang commented Aug 29, 2026 •

edited by github-actions Bot

Loading