Fix DSpark + DeepSeek V4 prefill CP compatibility by whybeyoung · Pull Request #33865 · sgl-project/sglang

whybeyoung · 2026-08-06T10:36:00Z

Summary

Two related fixes so DSpark speculative decoding works together with DeepSeek V4 prefill context parallelism (--enable-nsa-prefill-context-parallel --nsa-prefill-cp-mode round-robin-split) at PD co-located deployment.

Changes

deepseek_v4_dspark.py: use lm_head's shard group for markov_w2 TP-shard
configure_tp_shard hard-coded attn_tp_group for the per-step all-gather. Under prefill CP (attn_cp_size == tp_size) attn_tp_group degenerates to world_size == 1 while lm_head still shards over the full TP group, so the geometry check failed and the gather (if forced through) would be a no-op leaving half the vocab missing. Now we pick the group that matches lm_head's own choice: attn_tp_group when use_attn_tp_group=True, otherwise the full TP group.
deepseek_v4.py: enable DSpark aux capture under prefill CP
Drop the NotImplementedError blocking DSpark + CP and all-gather + rerange each dspark_aux_hidden_states tensor on the last PP rank with the same CP token split used for hidden_states, so downstream logits_processor indexing sees the global token layout.

Test

Verified with:

--tp 8 --speculative-algorithm DSPARK
--enable-nsa-prefill-context-parallel --nsa-prefill-cp-mode round-robin-split
--page-size 256 ...

No more NotImplementedError; markov_w2 TP-shard geometry check passes.

CI States

Latest PR Test (Base): ❌ Run #31455444987
Latest PR Test (Extra): ❌ Run #31455444860

CC @xu-yfei @ShangmingCai @Fridge003

whybeyoung · 2026-08-06T10:38:00Z

/tag-and-rerun-ci

Fridge003 · 2026-08-08T00:39:44Z

@whybeyoung We have shifted the CP implementation of DSV4 to CP_V2 in #33532. The older CP implementation will be deprecated soon.

Can you please test how it works with DSpark

Its usage should be:

--attn-cp-size 8 \
--enable-prefill-cp \
--cp-strategy interleave \

with flag SGLANG_ENABLE_CP_V2=1

…project#32035 sgl-project#33656 sgl-project#32183 sgl-project#33145) Applied PRs (latest from GitHub): sgl-project#33288 Indexer logits OOM fix sgl-project#30393 HiCache packed/sidecar draft caches sgl-project#31170 DPA prefix_affinity load balancing sgl-project#33795 DSpark compact ragged-verify CUDA graph JIT race sgl-project#32467 C128 plan-kernel warp barrier sgl-project#33865 DSpark x prefill CP unblock sgl-project#30371 SWA state pool sizing (storage page) sgl-project#33358 FlashMLA norm-rope K-tokens-per-block ILP sgl-project#33872 num_draft_tokens clamp + extend_len==0 skip (supersede sgl-project#32183) sgl-project#34002 Sidecar backup vacuously-successful fix (replaces sgl-project#33656, with tests) sgl-project#33862 Reclaim redundant host mirrors after storage backup sgl-project#31315 Avoid repeated Mooncake gets after stale hits sgl-project#32327 Q8KV8 sparse MLA prefill backend (flashmla_sparse_q8) sgl-project#31668 Fix sidecar pool life-time (use-after-free on prefetch abort) sgl-project#31195 TP0 verify-token-budget broadcast (adapted to get_schedule() API) Dropped (per user request or superseded): sgl-project#32771 IndexCache C4 top-k reuse — has bug sgl-project#32035 DSpark C128 online compressor — has bug sgl-project#33656 Superseded by sgl-project#34002 (same fix + unit tests) sgl-project#32183 Superseded by sgl-project#33872 (included in supersede PR) sgl-project#33145 Base f01f706 already has superior reasoning-effort profile system Conflicts resolved: sgl-project#31195: adapted to base get_schedule().disable_overlap_schedule API sgl-project#32327: path remapped jit_kernel/ -> kernels/jit/ and kernels/ops/attention/ sgl-project#31668: applied cleanly on top of sgl-project#30393+sgl-project#34002+sgl-project#33862 modifications

configure_tp_shard hard-coded attn_tp_group for the per-step all-gather, but under prefill CP (attn_cp_size == tp_size) attn_tp_group degenerates to world_size=1 while lm_head still shards over the full TP group, leaving the check failing and the gather a no-op. Pick the group that matches lm_head's own sharding: attn_tp_group when use_attn_tp_group is True, otherwise the full TP group. Store the group and use it in _apply_step_logits_sharded.

Drop the NotImplementedError blocking DSpark + CP and all-gather+rerange each dspark_aux_hidden_states tensor on the last PP rank with the same CP token split used for hidden_states, so downstream logits_processor indexing sees the global token layout.

TP=4 with attn_cp_size=4 and round-robin-split so attn_tp_group degenerates to size 1: exercises both fixes (markov_w2 TP-shard picking lm_head's shard group and DSpark aux hidden-state CP all-gather).

Under CP-v2 the model-level all-gather is skipped and hidden_states is gathered by _execute_extend_cp_v2 via cp_gather_after_forward. Do the same for aux_hidden_states so DSpark's per-layer captures land in the global token layout that logits_processor indexes with extend_seq_lens.

Per PR review: drop the standalone test file and add a DSPARK subtest inside test_deepseek_v4_flash_fp4_b200_cp.py using the CP-v2 flag set (--enable-prefill-cp --cp-strategy interleave, SGLANG_ENABLE_CP_V2=1).

whybeyoung · 2026-08-08T06:21:15Z

Addressed review comments and rebased on latest main:

@Fridge003 - CP_V2 support:

Extended _execute_extend_cp_v2 in eager_runner.py to also gather aux_hidden_states via cp_gather_after_forward, mirroring the hidden-states gather; without it DSpark aux tensors stay CP-local while logits_processor indexes them with global extend_seq_lens.
The in-model gather in deepseek_v4.py is now gated on use_prefill_cp and not cp_v2_active (matches upstream), so CP_V1 still works and CP_V2 uses the new external gather.

@Fridge003 - test consolidation:

Removed the standalone test_deepseek_v4_dspark_cp.py.
Added TestDSV4FlashFP4B200_CP_DSpark as a subtest in test_deepseek_v4_flash_fp4_b200_cp.py using the CP_V2 flag set (--enable-prefill-cp --cp-strategy interleave, SGLANG_ENABLE_CP_V2=1).

Rebased onto upstream/main; merge conflicts (with #33532) resolved.

whybeyoung · 2026-08-08T08:24:20Z

/tag-and-rerun-ci

Fridge003 · 2026-08-10T03:08:25Z

/rerun-test test/registered/cp/test_gqa_prefill_cp.py test/registered/cp/test_deepseek_v4_flash_fp4_b200_cp.py test/registered/cp/test_cp_strategy_unit.py test/registered/cp/test_dsa_prefill_cp.py test/registered/cp/test_deepseek_v3_cp_single_node.py

github-actions · 2026-08-10T03:08:51Z

Results for /rerun-test test/registered/cp/test_gqa_prefill_cp.py test/registered/cp/test_deepseek_v4_flash_fp4_b200_cp.py test/registered/cp/test_cp_strategy_unit.py test/registered/cp/test_dsa_prefill_cp.py test/registered/cp/test_deepseek_v3_cp_single_node.py:

🚀 4-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/cp/test_gqa_prefill_cp.py

🚀 4-gpu-b200 (1 test): ❌ View workflow run

cd test/ && python3 registered/cp/test_deepseek_v4_flash_fp4_b200_cp.py

🚀 ubuntu-latest (1 test): ✅ View workflow run

cd test/ && python3 registered/cp/test_cp_strategy_unit.py

🚀 8-gpu-h200 (2 tests): ✅ View workflow run

cd test/ && python3 registered/cp/test_dsa_prefill_cp.py
cd test/ && python3 registered/cp/test_deepseek_v3_cp_single_node.py

Fridge003 · 2026-08-10T03:09:30Z

/rerun-test test/registered/dcp/test_kimi_linear_dcp_dspark4.py

github-actions · 2026-08-10T03:10:03Z

Results for /rerun-test test/registered/dcp/test_kimi_linear_dcp_dspark4.py:

🚀 4-gpu-b200 (1 test): ✅ View workflow run

cd test/ && python3 registered/dcp/test_kimi_linear_dcp_dspark4.py

outdated

Under DSV4 prefill CP, the CP hook auto-sets enable_dp_attention=True but dp_size stays 1 (no real DP sharding). The DSpark dp-attention validation (needs --enable-dp-lm-head, forbids attn_cp_size>1) targets real DP deployments; skip it when dp_size==1 so DSpark + CP can launch.

whybeyoung · 2026-08-10T06:14:38Z

/rerun-test test/registered/cp/test_gqa_prefill_cp.py test/registered/cp/test_deepseek_v4_flash_fp4_b200_cp.py test/registered/cp/test_cp_strategy_unit.py test/registered/cp/test_dsa_prefill_cp.py test/registered/cp/test_deepseek_v3_cp_single_node.py

whybeyoung · 2026-08-10T06:15:09Z

/rerun-test test/registered/dcp/test_kimi_linear_dcp_dspark4.py

github-actions · 2026-08-10T06:15:11Z

Results for /rerun-test test/registered/cp/test_gqa_prefill_cp.py test/registered/cp/test_deepseek_v4_flash_fp4_b200_cp.py test/registered/cp/test_cp_strategy_unit.py test/registered/cp/test_dsa_prefill_cp.py test/registered/cp/test_deepseek_v3_cp_single_node.py:

🚀 4-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/cp/test_gqa_prefill_cp.py

🚀 4-gpu-b200 (1 test): ❌ View workflow run

cd test/ && python3 registered/cp/test_deepseek_v4_flash_fp4_b200_cp.py

🚀 ubuntu-latest (1 test): ✅ View workflow run

cd test/ && python3 registered/cp/test_cp_strategy_unit.py

🚀 8-gpu-h200 (2 tests): ✅ View workflow run

cd test/ && python3 registered/cp/test_dsa_prefill_cp.py
cd test/ && python3 registered/cp/test_deepseek_v3_cp_single_node.py

github-actions · 2026-08-10T06:15:28Z

Results for /rerun-test test/registered/dcp/test_kimi_linear_dcp_dspark4.py:

🚀 4-gpu-b200 (1 test): ✅ View workflow run

cd test/ && python3 registered/dcp/test_kimi_linear_dcp_dspark4.py

whybeyoung · 2026-08-11T02:19:51Z

Pushed a follow-up fix (64ac5ae) for the b200 CP failure.

Root cause: in _execute_extend_cp_v2 I was always forwarding pre_hc_head as hidden_states_before_norm to the logits processor. DeepseekV4ForCausalLM.forward only does that when aux capture is off; when DSpark aux capture is on, it passes hidden_states_before_norm=None so the logits processor stores the packed aux tensor ([T, K*hidden]) and hands it to the DSpark KV injector. Because _get_hidden_states_to_store prefers hidden_states_before_norm when present, my old CP-v2 code was silently swapping the packed aux for pre_hc_head ([T, mult*hidden]), which then tripped the A.shape[-1] == B.shape[-1] assertion inside main_proj's fp8 matmul.

The CP-v2 code now mirrors the wrapper's contract: pass hidden_states_before_norm only when aux_hidden_states is None.

/rerun-test test/registered/cp/test_deepseek_v4_flash_fp4_b200_cp.py test/registered/dcp/test_kimi_linear_dcp_dspark4.py

whybeyoung · 2026-08-11T02:21:34Z

/rerun-test test/registered/cp/test_gqa_prefill_cp.py test/registered/cp/test_deepseek_v4_flash_fp4_b200_cp.py test/registered/cp/test_cp_strategy_unit.py test/registered/cp/test_dsa_prefill_cp.py test/registered/cp/test_deepseek_v3_cp_single_node.py

whybeyoung · 2026-08-11T02:21:48Z

/rerun-test test/registered/dcp/test_kimi_linear_dcp_dspark4.py

github-actions · 2026-08-11T02:21:55Z

Results for /rerun-test test/registered/cp/test_gqa_prefill_cp.py test/registered/cp/test_deepseek_v4_flash_fp4_b200_cp.py test/registered/cp/test_cp_strategy_unit.py test/registered/cp/test_dsa_prefill_cp.py test/registered/cp/test_deepseek_v3_cp_single_node.py:

🚀 4-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/cp/test_gqa_prefill_cp.py

🚀 4-gpu-b200 (1 test): ✅ View workflow run

cd test/ && python3 registered/cp/test_deepseek_v4_flash_fp4_b200_cp.py

🚀 ubuntu-latest (1 test): ❌ View workflow run

cd test/ && python3 registered/cp/test_cp_strategy_unit.py

🚀 8-gpu-h200 (2 tests): ✅ View workflow run

cd test/ && python3 registered/cp/test_dsa_prefill_cp.py
cd test/ && python3 registered/cp/test_deepseek_v3_cp_single_node.py

github-actions · 2026-08-11T02:22:10Z

Results for /rerun-test test/registered/dcp/test_kimi_linear_dcp_dspark4.py:

🚀 4-gpu-b200 (1 test): ✅ View workflow run

cd test/ && python3 registered/dcp/test_kimi_linear_dcp_dspark4.py

When DSpark captures aux hidden states, DeepseekV4ForCausalLM.forward already passes hidden_states_before_norm=None so the logits processor returns packed aux tensors (shape [T, K*hidden]) to the DSpark KV injector. The CP-v2 path in _execute_extend_cp_v2 was still forwarding pre_hc_head (shape [T, mult*hidden]) as hidden_states_before_norm, which _get_hidden_states_to_store prefers over the packed aux buffer. The DSpark KV injector then handed main_proj a tensor whose last dim mismatched the ReplicatedLinear weight, tripping the fp8 matmul A.shape[-1] == B.shape[-1] assertion during prefill. Mirror the wrapper's contract: pass hidden_states_before_norm only when aux_hidden_states is None.

whybeyoung · 2026-08-11T03:09:17Z

/rerun-test test/registered/cp/test_gqa_prefill_cp.py test/registered/cp/test_deepseek_v4_flash_fp4_b200_cp.py test/registered/cp/test_cp_strategy_unit.py test/registered/cp/test_dsa_prefill_cp.py test/registered/cp/test_deepseek_v3_cp_single_node.py test/registered/dcp/test_kimi_linear_dcp_dspark4.py

github-actions · 2026-08-11T03:09:43Z

Results for /rerun-test test/registered/cp/test_gqa_prefill_cp.py test/registered/cp/test_deepseek_v4_flash_fp4_b200_cp.py test/registered/cp/test_cp_strategy_unit.py test/registered/cp/test_dsa_prefill_cp.py test/registered/cp/test_deepseek_v3_cp_single_node.py test/registered/dcp/test_kimi_linear_dcp_dspark4.py:

🚀 4-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/cp/test_gqa_prefill_cp.py

🚀 4-gpu-b200 (2 tests): ✅ View workflow run

cd test/ && python3 registered/cp/test_deepseek_v4_flash_fp4_b200_cp.py
cd test/ && python3 registered/dcp/test_kimi_linear_dcp_dspark4.py

🚀 ubuntu-latest (1 test): ✅ View workflow run

cd test/ && python3 registered/cp/test_cp_strategy_unit.py

🚀 8-gpu-h200 (2 tests): ✅ View workflow run

cd test/ && python3 registered/cp/test_dsa_prefill_cp.py
cd test/ && python3 registered/cp/test_deepseek_v3_cp_single_node.py

github-actions Bot added the deepseek label Aug 6, 2026

github-actions Bot added the run-ci label Aug 6, 2026

ShangmingCai assigned Fridge003 and ShangmingCai Aug 6, 2026

hzh0425 reviewed Aug 6, 2026

View reviewed changes

Comment thread python/sglang/srt/models/deepseek_v4.py

hzh0425 approved these changes Aug 6, 2026

View reviewed changes

Fridge003 previously requested changes Aug 8, 2026

View reviewed changes

Comment thread test/registered/cp/test_deepseek_v4_dspark_cp.py Outdated

ybyang added 5 commits August 8, 2026 14:14

Add DSPARK + DSA prefill CP e2e test for DeepSeek V4-Flash

b75ad6f

TP=4 with attn_cp_size=4 and round-robin-split so attn_tp_group degenerates to size 1: exercises both fixes (markov_w2 TP-shard picking lm_head's shard group and DSpark aux hidden-state CP all-gather).

Move DSPARK + prefill CP e2e to test_deepseek_v4_flash_fp4_b200_cp

174efc6

Per PR review: drop the standalone test file and add a DSPARK subtest inside test_deepseek_v4_flash_fp4_b200_cp.py using the CP-v2 flag set (--enable-prefill-cp --cp-strategy interleave, SGLANG_ENABLE_CP_V2=1).

whybeyoung force-pushed the fix-dspark-cp branch from f253193 to 174efc6 Compare August 8, 2026 06:20

whybeyoung requested review from Ying1123, hnyls2002, ispobock and merrymercy as code owners August 8, 2026 06:20

Merge branch 'main' into fix-dspark-cp

f20d497

Fridge003 approved these changes Aug 10, 2026

View reviewed changes

github-actions Bot added the speculative-decoding label Aug 10, 2026

whybeyoung force-pushed the fix-dspark-cp branch from 64ac5ae to f50e1f5 Compare August 11, 2026 03:06

Merge branch 'main' into fix-dspark-cp

c016bb2

chore: black-format eager_runner.py

1c5ec03

Fridge003 merged commit 9d4be40 into sgl-project:main Aug 11, 2026
91 of 130 checks passed

vstone-w pushed a commit to ClownBin/sglang that referenced this pull request Aug 12, 2026

Fix DSpark + DeepSeek V4 prefill CP compatibility (sgl-project#33865)

5ed5bd9

saturn-acc pushed a commit to saturn-acc/sglang that referenced this pull request Aug 16, 2026

Fix DSpark + DeepSeek V4 prefill CP compatibility (sgl-project#33865)

4b39c50

jakki-amd pushed a commit to jakki-amd/sglang that referenced this pull request Sep 9, 2026

Fix DSpark + DeepSeek V4 prefill CP compatibility (sgl-project#33865)

268c4f4

Atituiset pushed a commit to Atituiset/sglang that referenced this pull request Sep 10, 2026

Fix DSpark + DeepSeek V4 prefill CP compatibility (sgl-project#33865)

261cbe9

Uh oh!

Conversation

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

Uh oh!

Summary

Changes

Test

CI States

Uh oh!

whybeyoung commented Aug 6, 2026

Uh oh!

Uh oh!

Fridge003 commented Aug 8, 2026

Uh oh!

Uh oh!

whybeyoung commented Aug 8, 2026

Uh oh!

whybeyoung commented Aug 8, 2026

Uh oh!

Fridge003 commented Aug 10, 2026

Uh oh!

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

Uh oh!

Uh oh!

Fridge003 commented Aug 10, 2026

Uh oh!

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

Uh oh!

Uh oh!

whybeyoung commented Aug 10, 2026

Uh oh!

whybeyoung commented Aug 10, 2026

Uh oh!

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

Uh oh!

Uh oh!

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

Uh oh!

Uh oh!

whybeyoung commented Aug 11, 2026

Uh oh!

whybeyoung commented Aug 11, 2026

Uh oh!

whybeyoung commented Aug 11, 2026

Uh oh!

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

Uh oh!

Uh oh!

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

Uh oh!

Uh oh!

whybeyoung commented Aug 11, 2026

Uh oh!

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

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

whybeyoung commented Aug 6, 2026 •

edited by github-actions Bot

Loading

github-actions Bot commented Aug 10, 2026 •

edited

Loading

github-actions Bot commented Aug 10, 2026 •

edited

Loading

github-actions Bot commented Aug 10, 2026 •

edited

Loading

github-actions Bot commented Aug 10, 2026 •

edited

Loading

github-actions Bot commented Aug 11, 2026 •

edited

Loading

github-actions Bot commented Aug 11, 2026 •

edited

Loading

github-actions Bot commented Aug 11, 2026 •

edited

Loading