[ormandj]

Motivation

#33465 made DraftBlockProposer construct a CUDA int32 scalar on every speculative decoding step.

num_token_non_padded is only consumed when enable_num_token_non_padded() is true. For MoE EP size 1, the predicate is false, so the device allocation is unused.

Modifications

• Return None instead of constructing the device scalar when token padding is disabled.
• Preserve the existing int32 device scalar when token padding is enabled.
• Add unit coverage for both paths.

Accuracy Tests

This does not change model arithmetic. The focused test verifies:

• The disabled path returns None.
• The enabled path preserves the existing scalar value and torch.int32 dtype.

Passed during the Linux image build:

uv run --no-project python test/registered/spec/dspark/test_dspark_dp_tier.py

The changed files also pass the repository-pinned Black, Ruff, and isort checks.

Speed Tests and Profiling

Hardware and workload:

• 2× RTX PRO 6000 Blackwell (SM120), TP2
• DeepSeek-V4-Flash-0731 with DSpark
• Fixed 16,384-token inputs
• 2,048 requested output tokens
• Temperature 0
• 128 scheduler steps per capture
• Mean timeline duration across both ranks

All step times are milliseconds per scheduler step; lower is faster. The final column is the reduction from the after-#33465 control to the fixed build.

C	Before #33465	After #33465	Fixed	Step-time reduction
1	14.744 ms	15.681 ms	14.605 ms	6.86%
8	32.979 ms	35.399 ms	32.119 ms	9.27%

The post-fix capture used the same post-#33465 source stack, launch configuration, and fixed token inputs as the post-#33465 control, with only this patch added.

Every capture retained 127 complete target-plus-draft graph replay pairs. Graph kernel counts were identical between corresponding controls.

Checklist

• Format the changed files with the repository-pinned tools.
• Add focused unit tests.
• Provide matched profiling results.
• Follow the SGLang code-style guidance.
• Documentation is not required for this internal allocation change.

Prepared with AI assistance.

---

CI States

Latest PR Test (Base): ❌ Run #31754996272
Latest PR Test (Extra): ❌ Run #31754996074

[ormandj]

@McZyWu @Hexq0210 This fixes the CUDA EP1 decode performance regression caused by the unconditional DSpark device-scalar allocation introduced in #33465. The PR body includes matched before, after, and fixed profiling results.

[hnyls2002]

/rerun-test test_basic_sanity_dspark.py test_dspark_dp_tier.py test_dspark_draft_path_default.py

[github-actions]

Results for /rerun-test test_basic_sanity_dspark.py test_dspark_dp_tier.py test_dspark_draft_path_default.py:

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

cd test/ && python3 registered/core/test_basic_sanity_dspark.py

🚀 ubuntu-latest (2 tests): ✅ View workflow run

cd test/ && python3 registered/spec/dspark/test_dspark_dp_tier.py
cd test/ && python3 registered/spec/dspark/test_dspark_draft_path_default.py