[Feature] DSpark Support Prefill TP+CP

[zhangxiaolei123456]

Motivation

Add DSpark support for DeepSeek-V4 PD prefill with context parallelism.

Motivation

In the target deployment, prefill uses TP8/CP8 while decode uses TP8/DP8. Prefill only constructs and transfers target and draft KV caches, so configuring the Markov head against the TP-sharded LM head is unnecessary and fails because attn_tp_size=1 while lm_head.tp_size=8 .

Modifications

• Allow DSpark with Mooncake PD prefill when using PP=1, attention TP=1, and round-robin context parallelism.
• Pack captured auxiliary hidden states and restore global token order with a single CP all-gather.
• Validate hidden-state, position, and KV cache-location alignment before draft KV injection.
• Skip Decode-only embedding, LM-head, and Markov TP-shard initialization on PD prefill workers.
• Keep DSpark Decode and unsupported CP configurations unchanged.

Accuracy Tests

/sgl-workspace/sglang/benchmark/gsm8k# python3 bench_sglang.py --host http://localhost  --port 8090 --data-path /data00 --num-questions 5000 --parallel 128
100%|████| 1319/1319 [00:22<00:00, 58.35it/s]
Accuracy: 0.944
Invalid: 0.001
Latency: 22.607 s

Speed Tests and Profiling

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

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 #31099832671
Latest PR Test (Extra): ❌ Run #31099831690

[elmo2019]

@zhangxiaolei123456 hello~ any plan to merge it? Looking forward to this fea!