[ormandj]

Motivation

Without this change, an SPS cost table profiled on a server that has no table yet records full-width verify costs for every budget, so the planner that uses the table plans from wrong costs.

python -m sglang.benchmark.dspark_sps_profiler run --fracs ... profiles the off-diagonal SPS cost T(bs, M) by pinning dspark_force_budget_frac (through /set_internal_state) on a server running SGLANG_RAGGED_VERIFY_MODE=compact. A server started without --speculative-dspark-sps-table-path has the flat, uninitialized table, so DSparkVerifyPlanner is verify-all. This is the state when profiling the off-diagonal costs without a seed table. schedule_layout then returns the cached full-width layout before it looks at the budget. The budget planner does compute the pinned budget, but the layout ignores it. Pinned cells still run full-width verify, while the profiler records them under the requested M; its validation does not reject the mismatch.

I observed this on DeepSeek-V4.1-Flash, TP4 on 4x RTX PRO 6000 Blackwell (SM120), gamma=5: a --batch-size 1 2 4 8 --fracs 0.2 0.4 0.6 0.8 1.0 sweep on a server with no SPS table recorded only full-width steps (bs=4 always verified 24 tokens). The same sweep on a server started with a 1D table applied the budgets, for example bs=8 ran at 44.3 steps/s with 16 verify tokens vs 33.8 steps/s with 48.

#40567 changes how profiler control reaches the planner; it leaves this verify-all shortcut unchanged.

Modifications

• DSparkVerifyPlanner.schedule_layout: the verify-all shortcut (cached uniform layout) now applies only while no budget fraction is pinned. Clearing the pin (dspark_force_budget_frac=None through /set_internal_state) brings back the cached layout.
• /set_internal_state rejects pinning dspark_force_budget_frac while simulated acceptance is on (SGLANG_SIMULATE_ACC_LEN above 0), unless the length is exactly 1.0 with SGLANG_SIMULATE_ACC_METHOD=match-expected (the default), which always accepts exactly one token. multinomial samples the accept length from a normal distribution around SGLANG_SIMULATE_ACC_LEN, so even 1.0 can accept two or more tokens, past a trimmed compact verify window and onto target KV that verify never wrote. Now that a pin trims verify under verify-all, a constant simulated accept length could exceed the trimmed budget, which the DSpark worker already rejects at startup for profiled tables. Clearing the pin is still allowed.
• New TestForcedBudgetUnderVerifyAll (test/registered/spec/dspark/test_dspark_scheduler.py, CPU): with the uninitialized table, schedule_layout gives the full-width layout without a pin, a trimmed layout at frac=0.4 (12 verify tokens instead of 24 at bs=4, on main 24), and the cached layout again after the pin is cleared.
• New TestForcedBudgetRejectsSimulatedAcceptLength (same file) drives SchedulerInternalStateController.set_internal_state. With SGLANG_SIMULATE_ACC_LEN=3, or 1.0 with SGLANG_SIMULATE_ACC_METHOD=multinomial, the pin is rejected and the worker's budget setter is not called. Clearing the pin, and pinning with the variable set to -1.0 (disabled, with either method) or to 1.0 with match-expected, still work. On main, the rejection case fails.

Serving is unchanged unless a fraction is pinned. The pin is only set through /set_internal_state, which the profiler uses.

Accuracy Tests

This changes pinned layout selection and rejects incompatible simulated-accept-length settings; model outputs are unchanged when no fraction is pinned.
test_dspark_scheduler.py: 29 passed on CPU, including the four new tests.

Speed Tests and Profiling

No GPU speed measurements for this diff; serving is unchanged unless a fraction is pinned.

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.

Developed with AI assistance.

---

CI States

Latest PR Test (Base): ❌ Run #40961582812
Latest PR Test (Extra): ❌ Run #40961582767
Latest PR Test (AMD ROCm 10): ❌ Run #40961582732