From 5c492ca316ca68d87b101a703d5b1390229aabcb Mon Sep 17 00:00:00 2001 From: David Orman Date: Sat, 3 Oct 2026 19:13:03 -0500 Subject: [PATCH 1/2] [DSpark] Apply a pinned verify budget while the SPS table is uninitialized dspark_sps_profiler --fracs pins dspark_force_budget_frac on a compact server to profile the off-diagonal SPS cost T(bs, M). Before a profiled table exists the planner is verify-all, and schedule_layout returned the cached full-width layout before consulting the budget, so every pinned cell measured full-width verify and the fitted 2D table was flat in M. Serve the cached layout only while no fraction is pinned. Since a pin now trims verify under verify-all, reject pinning it while SGLANG_SIMULATE_ACC_LEN is above 1, as the DSpark worker already does at startup for profiled tables. --- .../scheduler_components/internal_state.py | 10 ++ .../dspark_components/dspark_planner.py | 9 +- .../spec/dspark/test_dspark_scheduler.py | 122 +++++++++++++++++- 3 files changed, 138 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/managers/scheduler_components/internal_state.py b/python/sglang/srt/managers/scheduler_components/internal_state.py index 574f4a796968..6e1653d95e54 100644 --- a/python/sglang/srt/managers/scheduler_components/internal_state.py +++ b/python/sglang/srt/managers/scheduler_components/internal_state.py @@ -218,6 +218,16 @@ def _validate_dspark_budget_frac(self, key: str, value: Any) -> None: self._require_dspark_worker(key, "set_dspark_forced_budget_frac") if value is not None and not (0.0 < float(value) <= 1.0): raise UpdateRejected(f"{key} must be in (0, 1] or null, got {value}.") + simulate_acc_len = envs.SGLANG_SIMULATE_ACC_LEN.get() + if value is not None and simulate_acc_len > 0 and simulate_acc_len != 1.0: + # A trimmed verify budget breaks the cap-accept accounting of a + # constant simulated accept length, as the DSpark worker checks at + # startup for profiled SPS tables. + raise UpdateRejected( + f"{key} cannot be pinned with SGLANG_SIMULATE_ACC_LEN=" + f"{simulate_acc_len}; only SGLANG_SIMULATE_ACC_LEN=1.0 stays " + "within every verify budget." + ) def _apply_dspark_budget_frac(self, key: str, value: Any) -> None: # A worker command, not a server arg: keep it out of the override. diff --git a/python/sglang/srt/speculative/dspark_components/dspark_planner.py b/python/sglang/srt/speculative/dspark_components/dspark_planner.py index abeb2da68820..2de3f9e7cc3b 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_planner.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_planner.py @@ -386,10 +386,15 @@ def schedule_layout( ) -> Optional[RaggedVerifyLayout]: if self._ragged_verify_mode is RaggedVerifyMode.STATIC: return None - if self._is_verify_all and self._ragged_verify_mode is RaggedVerifyMode.COMPACT: + if ( + self._is_verify_all + and self._ragged_verify_mode is RaggedVerifyMode.COMPACT + and self._budget_planner.forced_budget_frac is None + ): # Verify-all: the uniform layout (or None, past the captured grid) # is constant per (bs, tier); serve it from cache instead of paying - # the per-step schedule and its host<->device round-trips. + # the per-step schedule and its host<->device round-trips. A pinned + # budget (dspark_sps_profiler --fracs) must still be scheduled. key = (int(req_pool_indices.shape[0]), global_num_reqs) if key not in self._uniform_layout_cache: self._uniform_layout_cache[key] = uniform_ragged_layout( diff --git a/test/registered/spec/dspark/test_dspark_scheduler.py b/test/registered/spec/dspark/test_dspark_scheduler.py index b7a782785378..5df8eb7f08e1 100644 --- a/test/registered/spec/dspark/test_dspark_scheduler.py +++ b/test/registered/spec/dspark/test_dspark_scheduler.py @@ -1,14 +1,17 @@ import functools import types import unittest +import unittest.mock import torch from sglang.kernels.ops.speculative.dspark.dspark_schedule import ( schedule_verify_lens_topk_from_survival, ) +from sglang.srt.environ import envs from sglang.srt.speculative.dspark_components.dspark_planner import ( DSparkScheduleConfig, + DSparkVerifyPlanner, HostConfidenceBudgetPlanner, VerifyBudgetDecision, compute_verify_token_budget, @@ -17,8 +20,9 @@ from sglang.srt.speculative.dspark_components.dspark_sps import ( SpsAdditiveCostTable, SpsCostTable, + build_uninitialized_sps_table, ) -from sglang.srt.speculative.ragged_verify import RaggedVerifyLayout +from sglang.srt.speculative.ragged_verify import RaggedVerifyLayout, RaggedVerifyMode from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -398,6 +402,48 @@ def test_feeding_budget_fills_topk_total_to_tier(self): self.assertEqual(total, min(graph_num_tokens, bs * cap)) +class TestForcedBudgetRejectsSimulatedAcceptLength(CustomTestCase): + """A pinned budget now trims verify under verify-all, so it must not combine + with a constant simulated accept length above 1, which can exceed it.""" + + def _set(self, frac): + from sglang.srt.managers.io_struct import SetInternalStateReq + from sglang.srt.managers.scheduler_components.internal_state import ( + SchedulerInternalStateController, + ) + + draft_worker = types.SimpleNamespace( + set_dspark_forced_budget_frac=unittest.mock.Mock() + ) + scheduler = types.SimpleNamespace( + spec_algorithm=types.SimpleNamespace(is_dspark=lambda: True), + draft_worker=draft_worker, + ) + controller = SchedulerInternalStateController(scheduler, False) + output = controller.set_internal_state( + SetInternalStateReq(server_args={"dspark_force_budget_frac": frac}) + ) + return output.updated, draft_worker.set_dspark_forced_budget_frac + + def test_pin_rejected_with_simulated_accept_length(self): + with envs.SGLANG_SIMULATE_ACC_LEN.override(3.0): + updated, apply = self._set(0.4) + self.assertFalse(updated) + apply.assert_not_called() + # Clearing the pin stays allowed. + updated, apply = self._set(None) + self.assertTrue(updated) + apply.assert_called_once_with(None) + + def test_pin_allowed_without_or_with_unit_simulated_accept_length(self): + for simulate_acc_len in (-1.0, 1.0): + with self.subTest(simulate_acc_len=simulate_acc_len): + with envs.SGLANG_SIMULATE_ACC_LEN.override(simulate_acc_len): + updated, apply = self._set(0.4) + self.assertTrue(updated) + apply.assert_called_once_with(0.4) + + class _FakeRaggedRunner(types.SimpleNamespace): pass @@ -481,5 +527,79 @@ def test_exceeds_gate_checks_slots_and_tier(self): ) +class TestForcedBudgetUnderVerifyAll(CustomTestCase): + """dspark_sps_profiler --fracs pins a budget fraction on a compact server whose + SPS table is still uninitialized (verify-all); the pinned budget must shape + the layout instead of the cached full-width one.""" + + BS = 4 + GAMMA = 5 + + def _planner(self): + num_draft_tokens = self.GAMMA + 1 + model_runner = _fake_model_runner( + [num_draft_tokens * n for n in range(1, 9)], max_bs=8 + ) + model_runner.req_to_token_pool = types.SimpleNamespace( + req_to_token=torch.zeros((16, 1), dtype=torch.int32) + ) + cfg = DSparkScheduleConfig(gamma=self.GAMMA) + planner = DSparkVerifyPlanner.__new__(DSparkVerifyPlanner) + planner.model_runner = model_runner + planner.verify_num_draft_tokens = num_draft_tokens + planner._ragged_verify_mode = RaggedVerifyMode.COMPACT + planner._schedule_cfg = cfg + planner._is_verify_all = True + planner._uniform_layout_cache = {} + planner._dynamic_graph_tier = True + planner._align_verify_tokens_to_graph_tier = False + planner._tp_sync = types.SimpleNamespace(sync=lambda *args: None) + planner._budget_planner = HostConfidenceBudgetPlanner( + sps_table=build_uninitialized_sps_table(max_batch_tokens=8 * 6), + cfg=cfg, + model_runner=model_runner, + ) + return planner + + def _schedule(self, planner): + req_pool_indices = torch.arange(self.BS) + generation = torch.ones(self.BS, dtype=torch.int64) + confidence = torch.full((self.BS, self.GAMMA), 0.9) + budget = planner._budget_planner.compute_budget( + confidence=confidence, + generation=generation, + current_generation=generation, + req_pool_indices_cpu=req_pool_indices, + ) + layout = planner.schedule_layout( + req_pool_indices=req_pool_indices, + prefix_lens=torch.zeros(self.BS, dtype=torch.int64), + device=torch.device("cpu"), + confidence=confidence, + budget=budget, + ) + return budget, layout + + def test_pinned_budget_shapes_the_layout(self): + planner = self._planner() + full_width = self.BS * (self.GAMMA + 1) + + _, uniform = self._schedule(planner) + self.assertEqual(uniform.verify_lens.tolist(), [self.GAMMA + 1] * self.BS) + self.assertEqual(uniform.graph_num_tokens, full_width) + + planner.set_forced_budget_frac(0.4) + budget, pinned = self._schedule(planner) + total = int(pinned.verify_lens.sum()) + self.assertEqual(budget, int(0.4 * self.BS * self.GAMMA)) + self.assertEqual(total, self.BS + budget) + self.assertLess(pinned.graph_num_tokens, full_width) + self.assertEqual(pinned.graph_num_tokens, 12) + + planner.set_forced_budget_frac(None) + _, cleared = self._schedule(planner) + self.assertIs(cleared, uniform) + + if __name__ == "__main__": unittest.main() From 0c2039fbd203307f23b1ef9ca0016f052b4328cc Mon Sep 17 00:00:00 2001 From: David Orman Date: Tue, 6 Oct 2026 21:17:10 -0500 Subject: [PATCH 2/2] [DSpark] Reject a pinned budget with multinomial simulated acceptance SGLANG_SIMULATE_ACC_METHOD=multinomial draws the accept length from a normal around SGLANG_SIMULATE_ACC_LEN, so a configured 1.0 can still accept past a trimmed compact verify window. Allow the pin only when the simulated length is exactly 1.0 with match-expected sampling. --- .../scheduler_components/internal_state.py | 18 +++++++++---- .../spec/dspark/test_dspark_scheduler.py | 27 ++++++++++++++++--- 2 files changed, 37 insertions(+), 8 deletions(-) diff --git a/python/sglang/srt/managers/scheduler_components/internal_state.py b/python/sglang/srt/managers/scheduler_components/internal_state.py index 6e1653d95e54..3e5c36e21f2f 100644 --- a/python/sglang/srt/managers/scheduler_components/internal_state.py +++ b/python/sglang/srt/managers/scheduler_components/internal_state.py @@ -219,14 +219,22 @@ def _validate_dspark_budget_frac(self, key: str, value: Any) -> None: if value is not None and not (0.0 < float(value) <= 1.0): raise UpdateRejected(f"{key} must be in (0, 1] or null, got {value}.") simulate_acc_len = envs.SGLANG_SIMULATE_ACC_LEN.get() - if value is not None and simulate_acc_len > 0 and simulate_acc_len != 1.0: + simulate_acc_method = envs.SGLANG_SIMULATE_ACC_METHOD.get() + if ( + value is not None + and simulate_acc_len > 0 + and (simulate_acc_len != 1.0 or simulate_acc_method != "match-expected") + ): # A trimmed verify budget breaks the cap-accept accounting of a - # constant simulated accept length, as the DSpark worker checks at - # startup for profiled SPS tables. + # simulated accept length above 1, as the DSpark worker checks at + # startup for profiled SPS tables. Only match-expected sampling of + # 1.0 always yields 1; multinomial draws around the mean. raise UpdateRejected( f"{key} cannot be pinned with SGLANG_SIMULATE_ACC_LEN=" - f"{simulate_acc_len}; only SGLANG_SIMULATE_ACC_LEN=1.0 stays " - "within every verify budget." + f"{simulate_acc_len} and SGLANG_SIMULATE_ACC_METHOD=" + f"{simulate_acc_method}; only SGLANG_SIMULATE_ACC_LEN=1.0 with " + "SGLANG_SIMULATE_ACC_METHOD=match-expected stays within every " + "verify budget." ) def _apply_dspark_budget_frac(self, key: str, value: Any) -> None: diff --git a/test/registered/spec/dspark/test_dspark_scheduler.py b/test/registered/spec/dspark/test_dspark_scheduler.py index 5df8eb7f08e1..b3d85947e9fe 100644 --- a/test/registered/spec/dspark/test_dspark_scheduler.py +++ b/test/registered/spec/dspark/test_dspark_scheduler.py @@ -435,10 +435,31 @@ def test_pin_rejected_with_simulated_accept_length(self): self.assertTrue(updated) apply.assert_called_once_with(None) + def test_pin_rejected_with_multinomial_unit_accept_length(self): + # multinomial draws from N(1, 1), so it can accept past the trimmed + # verify window even though the configured length is 1. + with ( + envs.SGLANG_SIMULATE_ACC_LEN.override(1.0), + envs.SGLANG_SIMULATE_ACC_METHOD.override("multinomial"), + ): + updated, apply = self._set(0.4) + self.assertFalse(updated) + apply.assert_not_called() + updated, apply = self._set(None) + self.assertTrue(updated) + apply.assert_called_once_with(None) + def test_pin_allowed_without_or_with_unit_simulated_accept_length(self): - for simulate_acc_len in (-1.0, 1.0): - with self.subTest(simulate_acc_len=simulate_acc_len): - with envs.SGLANG_SIMULATE_ACC_LEN.override(simulate_acc_len): + for simulate_acc_len, method in ( + (-1.0, "match-expected"), + (-1.0, "multinomial"), + (1.0, "match-expected"), + ): + with self.subTest(simulate_acc_len=simulate_acc_len, method=method): + with ( + envs.SGLANG_SIMULATE_ACC_LEN.override(simulate_acc_len), + envs.SGLANG_SIMULATE_ACC_METHOD.override(method), + ): updated, apply = self._set(0.4) self.assertTrue(updated) apply.assert_called_once_with(0.4)