[SPEC][5/N] feat: batchsize-aware support for adaptive speculative_num_steps (#24055)

Co-authored-by: 坤钧 <maoyuhan.myh@antgroup.co>
Co-authored-by: alphabetc1 <alphabetc1@users.noreply.github.com>
Co-authored-by: EanWang211123 <wangyiheng@sangfor.com.cn>
Co-authored-by: shuwenn <47200617+alphabetc1@users.noreply.github.com>
Co-authored-by: shuwenn <2508695655@qq.com>

test/registered/spec/eagle/test_adaptive_speculative.py

@@ -47,11 +47,13 @@ class TestAdaptiveSpeculativeServer(CustomTestCase):
         with tempfile.NamedTemporaryFile("w", suffix=".json", delete=False) as f:
             json.dump(
                 {
-                    "candidate_steps": [1, 3],
-                    "ema_alpha": 1.0,
-                    "warmup_batches": 1,
-                    "update_interval": 1,
-                    "up_hysteresis": 0.0,
+                    "1": {
+                        "candidate_steps": [1, 3],
+                        "ema_alpha": 1.0,
+                        "warmup_batches": 1,
+                        "update_interval": 1,
+                        "up_hysteresis": 0.0,
+                    },
                 },
                 f,
             )

test/registered/unit/spec/test_adaptive_spec_params.py

@@ -4,6 +4,9 @@ import unittest
 
 from sglang.srt.speculative.adaptive_spec_params import (
     AdaptiveSpeculativeParams,
+    AdaptiveStepSlot,
+    resolve_candidate_steps_from_config,
+    validate_adaptive_initial_steps,
 )
 from sglang.test.ci.ci_register import register_cpu_ci, register_xpu_ci
 
@@ -11,39 +14,16 @@ register_cpu_ci(est_time=6, suite="base-a-test-cpu")
 register_xpu_ci(est_time=10, suite="stage-a-test-1-gpu-xpu")
 
 
-class TestAdaptiveSpeculativeParams(unittest.TestCase):
-    def _make_params_from_config(self, initial_steps: int, config: dict[str, object]):
-        with tempfile.NamedTemporaryFile("w", suffix=".json") as f:
-            json.dump(config, f)
-            f.flush()
-            return AdaptiveSpeculativeParams(
-                initial_steps=initial_steps, cfg_path=f.name
-            )
+class TestAdaptiveStepSlot(unittest.TestCase):
+    def _make_params_from_config(self, initial_steps: int, config: dict):
+        return AdaptiveStepSlot(initial_steps=initial_steps, cfg=config)
 
-    def test_params_loads_config_path(self):
-        with tempfile.NamedTemporaryFile("w", suffix=".json") as f:
-            json.dump(
-                {
-                    "candidate_steps": [1, 5],
-                    "ema_alpha": 0.75,
-                    "warmup_batches": 2,
-                },
-                f,
-            )
-            f.flush()
-
-            params = AdaptiveSpeculativeParams(initial_steps=3, cfg_path=f.name)
-
-        self.assertEqual(params.candidate_steps, [1, 3, 5])
-        self.assertEqual(params.ema_alpha, 0.75)
-        self.assertEqual(params.warmup_batches, 2)
-
-    def test_initial_steps_added_to_candidates_when_missing(self):
+    def test_initial_steps_snaps_to_middle_when_missing(self):
         params = self._make_params_from_config(2, {"candidate_steps": [1, 3, 7]})
 
-        self.assertEqual(params.candidate_steps, [1, 2, 3, 7])
-        self.assertEqual(params.current_steps, 2)
-        self.assertEqual(params.ema_accept_len, 1.0)
+        self.assertEqual(params.candidate_steps, [1, 3, 7])
+        self.assertEqual(params.current_steps, 3)
+        self.assertEqual(params.ema_accept_len, 2.0)
 
     def test_update_respects_warmup_and_interval(self):
         params = self._make_params_from_config(
@@ -219,6 +199,229 @@ class TestAdaptiveSpeculativeParams(unittest.TestCase):
         self.assertEqual(params.current_steps, 1)
         self.assertEqual(params.ema_accept_len, 0.375)
 
+    def test_ceiling_coeff_caps_steps(self):
+        params = self._make_params_from_config(
+            7,
+            {
+                "candidate_steps": [1, 3, 7],
+                "ema_alpha": 1.0,
+                "warmup_batches": 0,
+                "update_interval": 1,
+                "ceiling_coeff": 1.0,
+            },
+        )
+        # Force low ema to trigger ceiling
+        params.ema_accept_len = 1.0
+        self.assertTrue(params.update([1, 1]))
+        # ceiling = ceil(1.0 * 1.0) = 1, target capped to 1
+        self.assertEqual(params.current_steps, 1)
+
+    def test_ceiling_disabled_by_default(self):
+        params = self._make_params_from_config(3, {"candidate_steps": [1, 3, 7]})
+        self.assertEqual(params.ceiling_coeff, 0)
+
+
+class TestAdaptiveSpeculativeParams(unittest.TestCase):
+    def test_default_config_loads(self):
+        params = AdaptiveSpeculativeParams(initial_steps=3)
+        self.assertEqual(params._bs_list, [1, 8, 32])
+        self.assertEqual(params._slots[1].candidate_steps, [1, 3, 7])
+        self.assertEqual(params._slots[8].candidate_steps, [1, 3])
+        self.assertEqual(params._slots[32].candidate_steps, [1])
+
+    def test_config_file(self):
+        with tempfile.NamedTemporaryFile("w", suffix=".json") as f:
+            json.dump(
+                {
+                    "1": {"candidate_steps": [1, 5], "up_hysteresis": 0.3},
+                    "32": {"candidate_steps": [1, 2]},
+                },
+                f,
+            )
+            f.flush()
+            params = AdaptiveSpeculativeParams(initial_steps=5, cfg_path=f.name)
+        self.assertEqual(params._bs_list, [1, 32])
+        # Slots are built straight from the config; the launch flag never pollutes
+        # them. initial_steps just selects the smallest slot's starting step.
+        self.assertEqual(params._slots[1].candidate_steps, [1, 5])
+        self.assertEqual(params._slots[1].current_steps, 5)
+        self.assertEqual(params._slots[1].up_hysteresis, 0.3)
+        self.assertEqual(params._slots[32].candidate_steps, [1, 2])
+
+    def test_launch_flag_not_injected_into_slots(self):
+        # initial_steps lives only in a larger slot. It must NOT be merged into
+        # any other slot's candidates: slots come straight from the config.
+        with tempfile.NamedTemporaryFile("w", suffix=".json") as f:
+            json.dump(
+                {
+                    "1": {"candidate_steps": [1, 5]},
+                    "8": {"candidate_steps": [1, 3, 7]},
+                },
+                f,
+            )
+            f.flush()
+            params = AdaptiveSpeculativeParams(initial_steps=7, cfg_path=f.name)
+        self.assertEqual(params._slots[1].candidate_steps, [1, 5])
+        self.assertEqual(params._slots[8].candidate_steps, [1, 3, 7])
+        # The slot that does not own initial_steps starts at its own median.
+        self.assertEqual(params._slots[1].current_steps, 5)
+        self.assertEqual(params._slots[8].current_steps, 7)
+
+    def test_invalid_config_raises(self):
+        with tempfile.NamedTemporaryFile("w", suffix=".json") as f:
+            json.dump({"not_a_bs": "bad"}, f)
+            f.flush()
+            with self.assertRaises(ValueError):
+                AdaptiveSpeculativeParams(initial_steps=3, cfg_path=f.name)
+
+    def test_invalid_steps_raises(self):
+        with tempfile.NamedTemporaryFile("w", suffix=".json") as f:
+            json.dump({"1": {"candidate_steps": "bad"}}, f)
+            f.flush()
+            with self.assertRaises(ValueError):
+                AdaptiveSpeculativeParams(initial_steps=3, cfg_path=f.name)
+
+    def test_empty_steps_raises(self):
+        with tempfile.NamedTemporaryFile("w", suffix=".json") as f:
+            json.dump({"1": {"candidate_steps": []}}, f)
+            f.flush()
+            with self.assertRaises(ValueError):
+                AdaptiveSpeculativeParams(initial_steps=3, cfg_path=f.name)
+
+    def test_zero_steps_raises(self):
+        with tempfile.NamedTemporaryFile("w", suffix=".json") as f:
+            json.dump({"1": {"candidate_steps": [0]}}, f)
+            f.flush()
+            with self.assertRaises(ValueError):
+                AdaptiveSpeculativeParams(initial_steps=3, cfg_path=f.name)
+
+    def test_global_hysteresis_inherited(self):
+        with tempfile.NamedTemporaryFile("w", suffix=".json") as f:
+            json.dump(
+                {
+                    "up_hysteresis": 0.5,
+                    "1": {"candidate_steps": [1, 3]},
+                },
+                f,
+            )
+            f.flush()
+            params = AdaptiveSpeculativeParams(initial_steps=3, cfg_path=f.name)
+        self.assertEqual(params._slots[1].up_hysteresis, 0.5)
+
+    def test_entry_hysteresis_overrides_global(self):
+        with tempfile.NamedTemporaryFile("w", suffix=".json") as f:
+            json.dump(
+                {
+                    "up_hysteresis": 0.5,
+                    "1": {"candidate_steps": [1, 3], "up_hysteresis": 0.1},
+                },
+                f,
+            )
+            f.flush()
+            params = AdaptiveSpeculativeParams(initial_steps=3, cfg_path=f.name)
+        self.assertEqual(params._slots[1].up_hysteresis, 0.1)
+
+
+class TestBatchSizeRouting(unittest.TestCase):
+    """BS-aware routing: batch size selects the slot, CUDA-graph BS pads first."""
+
+    def _params(self):
+        # Slots: bs=1 -> [1,3,7], bs=8 -> [1,3], bs=32 -> [1].
+        return AdaptiveSpeculativeParams(initial_steps=3)
+
+    def test_routes_to_floor_slot_without_cuda_graph(self):
+        params = self._params()
+        # A batch maps to the largest slot BS <= batch (floor), capped at the top slot.
+        self.assertEqual(params._route(1).candidate_steps, [1, 3, 7])
+        self.assertEqual(params._route(7).candidate_steps, [1, 3, 7])
+        self.assertEqual(params._route(8).candidate_steps, [1, 3])
+        self.assertEqual(params._route(31).candidate_steps, [1, 3])
+        self.assertEqual(params._route(32).candidate_steps, [1])
+        self.assertEqual(params._route(1000).candidate_steps, [1])
+
+    def test_cuda_graph_bs_pads_batch_up_before_routing(self):
+        params = self._params()
+        params.set_cuda_graph_bs([4, 8, 16, 32])
+        # bs=5 pads up to the captured graph BS 8 -> slot bs=8.
+        self.assertEqual(params._route(5).candidate_steps, [1, 3])
+        # bs=17 pads up to 32 -> slot bs=32.
+        self.assertEqual(params._route(17).candidate_steps, [1])
+        # A batch larger than every captured BS keeps its own value -> top slot.
+        self.assertEqual(params._route(100).candidate_steps, [1])
+
+    def test_cuda_graph_bs_for_step_prunes_unreachable_graphs(self):
+        params = self._params()
+        params.set_cuda_graph_bs([4, 8, 16, 32])
+        # step=1 is reachable from every slot.
+        self.assertEqual(params.cuda_graph_bs_for_step(1), [4, 8, 16, 32])
+        # step=3 lives in the bs=1 and bs=8 slots: graphs 4,8,16 floor into them.
+        self.assertEqual(params.cuda_graph_bs_for_step(3), [4, 8, 16])
+        # step=7 lives only in the bs=1 slot: only graph BS 4 floors into it.
+        self.assertEqual(params.cuda_graph_bs_for_step(7), [4])
+
+    def test_cuda_graph_bs_for_step_returns_none_when_disabled(self):
+        params = self._params()
+        self.assertIsNone(params.cuda_graph_bs_for_step(7))
+        params.set_cuda_graph_bs(None)
+        self.assertIsNone(params.cuda_graph_bs_for_step(7))
+
+    def test_observe_verify_feeds_the_routed_slot(self):
+        params = self._params()
+        # Drive the bs=1 slot up with perfect acceptance; the bs=32 slot is
+        # untouched and stays at its single candidate step.
+        for _ in range(40):
+            params.on_verify_complete([7, 7, 7], batch_size=1)
+        self.assertGreater(params.get_steps_for_batch(1), 1)
+        self.assertEqual(params.get_steps_for_batch(32), 1)
+
+
+class TestResolveCandidateSteps(unittest.TestCase):
+    def test_default_config(self):
+        steps = resolve_candidate_steps_from_config()
+        self.assertIn(1, steps)
+        self.assertIn(3, steps)
+        self.assertIn(7, steps)
+
+    def test_config_file(self):
+        with tempfile.NamedTemporaryFile("w", suffix=".json") as f:
+            json.dump({"1": {"candidate_steps": [2, 4]}}, f)
+            f.flush()
+            steps = resolve_candidate_steps_from_config(cfg_path=f.name)
+        self.assertEqual(steps, [2, 4])
+
+    def test_unions_and_dedups_across_slots(self):
+        with tempfile.NamedTemporaryFile("w", suffix=".json") as f:
+            json.dump(
+                {
+                    "1": {"candidate_steps": [1, 5]},
+                    "8": {"candidate_steps": [3, 5, 7]},
+                },
+                f,
+            )
+            f.flush()
+            steps = resolve_candidate_steps_from_config(cfg_path=f.name)
+        self.assertEqual(steps, [1, 3, 5, 7])
+
+
+class TestValidateAdaptiveInitialSteps(unittest.TestCase):
+    def test_accepts_value_from_any_slot(self):
+        with tempfile.NamedTemporaryFile("w", suffix=".json") as f:
+            json.dump(
+                {
+                    "1": {"candidate_steps": [1, 5]},
+                    "8": {"candidate_steps": [1, 3, 7]},
+                },
+                f,
+            )
+            f.flush()
+            # Membership in any slot is enough: 5 lives in the smallest slot,
+            # 7 only in a larger slot -- both accepted.
+            validate_adaptive_initial_steps(5, cfg_path=f.name)
+            validate_adaptive_initial_steps(7, cfg_path=f.name)
+            # 9 is in no slot -> rejected.
+            with self.assertRaises(ValueError):
+                validate_adaptive_initial_steps(9, cfg_path=f.name)
+
 
 if __name__ == "__main__":
     unittest.main()