Fix DSpark and DP/EP

[vladnosiv]

Motivation

Fix DSpark with DP attention: its draft ForwardBatch missed metadata required by CUDA graph admission and EP token accounting, causing startup and first-request crashes.

Modifications

Populate original_global_num_tokens_cpu and non-padded token counts when DSpark constructs the draft batch, and add regression coverage.

Accuracy Tests

verified on 4xH200 with:

sglang serve \
  --trust-remote-code \
  --model-path deepseek-ai/DeepSeek-V4-Flash-0731 \
  --tp 4 \
  --dp 4 \
  --enable-dp-attention \
  --enable-dp-lm-head \
  --moe-runner-backend marlin \
  --ep 4 \
  --host 0.0.0.0 \
  --port 30000 \
  --max-running-requests 512 \
  --swa-full-tokens-ratio 0.01 \
  --cuda-graph-max-bs-decode 128 \
  --speculative-algorithm DSPARK

python3 -m sglang.test.few_shot_gsm8k --num-questions 200 --port 30000

Accuracy: 0.960
Invalid: 0.000
Latency: 11.855 s
Output throughput: 1455.725 token/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 #30704262557
Latest PR Test (Extra): ❌ Run #30825602108

[gemini-code-assist]

Caution

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

[junliu-mde]

Same issue, this works for me.

Tested the cumulative diff at 9c37e64041aa501449dcd99f0c67ed22167a0795 with lmsysorg/sglang:v0.5.16 on 2x NVIDIA B300, DeepSeek-V4-Flash, --tp 2 --dp-size 2 --enable-dp-attention --enable-dp-lm-head --speculative-algorithm DSPARK. Both target and draft CUDA graph capture completed with backend=full. Concurrent generation requests returned HTTP 200, and both DP ranks reported cuda graph: True during decode. The pod remained Ready with zero restarts and no TypeError, traceback, or CUDA error.

[JustinTong0323]

/tag-and-rerun-ci

[JustinTong0323]

Validated this fix end-to-end on 4x GB300 with DeepSeek-V4-Flash-0731. I cherry-picked the three fix commits onto the exact v0.5.16 base (fdebc938) where the first request previously crashed because original_global_num_tokens_cpu was missing, then launched TP4/DP4 with DP attention, DP LM head, FlashInfer MXFP4, and bundled DSpark.

Results:

• The target and draft CUDA graphs both completed all 27/27 tiers; the warmup, streaming request, and 32-request concurrent smoke test all passed, with no scheduler exception or CUDA/OOM error.
• The added DSpark DP-tier test passed: 5/5.
• Full GSM8K: 96.66% (1275/1319), 99.92% stop rate, 0 request errors.
• 8k input / 1k output serving benchmark, with 64 warmups and a separate cache flush before every case:
  • c64: 16,248.70 total tok/s (4,062.18 tok/s/GPU), 1,857.75 output tok/s, 5.52 acceptance length. This is +10.73% versus TP4 + DSpark and +51.69% versus target-only DPA.
  • c256: 26,868.24 total tok/s (6,717.06 tok/s/GPU), 3,015.63 output tok/s, 5.47 acceptance length. This is +54.67% versus TP4 + DSpark and +33.55% versus target-only DPA.

This confirms that the patch fixes the previously broken DSpark + DP-attention path and makes it the strongest measured throughput-oriented configuration. TP4 + DSpark still has better latency at c64, and this result does not cover DSpark with PD disaggregation.

[JustinTong0323]

/rerun-failed-ci

[junliu-mde]

Same issue here. This patch works for me with DeepSeek-V4-Flash, DSpark, and DPA2 on 2x B300. Target and draft CUDA Graph capture completed with decode backend=full, two concurrent generation requests returned HTTP 200, both DP ranks reported cuda graph: True, and the pod stayed Ready with 0 restarts.

[vladnosiv]

@JustinTong0323 Hi !
Can we merge it?

[JustinTong0323]

@JustinTong0323 Hi ! Can we merge it?

Just pinged code owner to take a look

[Phoenix3334]

Follow-up integration regression: #34410 now pins the DSpark PD-decode draft-input handoff on current main.

This complements #33098 rather than duplicating it: #33098 fixed DSpark DP-attention draft ForwardBatch metadata, while #34410 covers the separate disaggregation handoff invariant that SpecInput and overlap relay state exist before the first decode step.