[spec decoding] supports step 0 in adaptive spec decoding (updating draft kv cache without draft decoding) (#23994)

Co-authored-by: shuwenn <2508695655@qq.com>

python/sglang/srt/environ.py

@@ -638,6 +638,10 @@ class Envs:
 
     # Spec Config
     SGLANG_SPEC_ENABLE_STRICT_FILTER_CHECK = EnvBool(True)
+    # Skip draft_extend while adaptive spec is at steps=0 (drafting disabled).
+    # Saves the per-step draft forward, but the draft KV goes stale: an upshift
+    # back to steps>0 starts from a cold draft state (low accept until it recovers).
+    SGLANG_SPEC_SKIP_ZERO_STEP_DRAFT_EXTEND = EnvBool(False)
     # Master switch for all async-asserted invariant probes (NaN, Inf, OOB,
     # page alignment). Off in prod; tests turn it on to fail-fast on
     # numerical / index violations instead of getting silent NaN cascades.

python/sglang/srt/speculative/adaptive_spec_params.py

@@ -19,7 +19,6 @@ if TYPE_CHECKING:
 
 logger = logging.getLogger(__name__)
 
-# TODO: add step=0 (nospec fallback) for BS>=8 once supported.
 DEFAULT_ADAPTIVE_CONFIG: dict[str, dict] = {
     "1": {
         "candidate_steps": [1, 3, 7],
@@ -28,13 +27,19 @@ DEFAULT_ADAPTIVE_CONFIG: dict[str, dict] = {
         "ceiling_coeff": 0,
     },
     "8": {
-        "candidate_steps": [1, 3],
+        "candidate_steps": [0, 1, 3],
         "up_hysteresis": 0.0,
         "down_hysteresis": 0.0,
         "ceiling_coeff": 0,
     },
     "32": {
-        "candidate_steps": [1],
+        "candidate_steps": [0, 1],
+        "up_hysteresis": 0.0,
+        "down_hysteresis": 0.0,
+        "ceiling_coeff": 0,
+    },
+    "64": {
+        "candidate_steps": [0],
         "up_hysteresis": 0.0,
         "down_hysteresis": 0.0,
         "ceiling_coeff": 0,
@@ -102,10 +107,11 @@ def _load_adaptive_config(
         if (
             not isinstance(steps, list)
             or not steps
-            or not all(isinstance(s, int) and s > 0 for s in steps)
+            or not all(isinstance(s, int) and s >= 0 for s in steps)
         ):
             raise ValueError(
-                f"BS {key}: candidate_steps must be a list of positive ints, got {steps!r}"
+                f"BS {key}: candidate_steps must be a list of non-negative ints, "
+                f"got {steps!r}"
             )
         bs_entries[int(key)] = entry
 
@@ -172,10 +178,13 @@ class AdaptiveStepSlot:
         if not num_correct_drafts_per_req:
             return False
 
-        batch_avg = sum(num_correct_drafts_per_req) / len(num_correct_drafts_per_req)
+        if self.current_steps > 0:
-        self.ema_accept_len = (
+            batch_avg = sum(num_correct_drafts_per_req) / len(
-            1 - self.ema_alpha
+                num_correct_drafts_per_req
-        ) * self.ema_accept_len + self.ema_alpha * batch_avg
+            )
+            self.ema_accept_len = (
+                1 - self.ema_alpha
+            ) * self.ema_accept_len + self.ema_alpha * batch_avg
 
         self._batch_count += 1
         if self._batch_count <= self.warmup_batches:
@@ -190,23 +199,39 @@ class AdaptiveStepSlot:
         """Recompute steps from EMA. Returns True if params changed."""
         old_steps = self.current_steps
         current_idx = self.candidate_steps.index(old_steps)
+        old_idx = current_idx
+
+        # Probe the smallest positive step after a zero-step nospec interval.
+        if old_steps == 0:
+            current_idx = min(current_idx + 1, len(self.candidate_steps) - 1)
+            target = self.candidate_steps[current_idx]
+            if target > 0 and self.ema_accept_len < 0:
+                # A slot initialized at steps=0 has no draft acceptance history;
+                # start the first positive-step probe from that step's neutral EMA.
+                self.ema_accept_len = float(target - 1)
+            return self._apply_target_steps(old_steps, target)
 
         # TODO: Consider limiting step changes to avoid overshooting.
         while current_idx > 0:
             prev_step = self.candidate_steps[current_idx - 1]
-            drop_threshold = prev_step - 0.5 + self.down_hysteresis
+            # A zero-step candidate disables drafting. Treat zero accepted drafts
+            # as low enough to reach it when it is the floor candidate.
+            drop_threshold = 0.5 if prev_step == 0 else prev_step - 0.5
+            drop_threshold += self.down_hysteresis
             if self.ema_accept_len <= drop_threshold:
                 current_idx -= 1
             else:
                 break
 
-        while current_idx < len(self.candidate_steps) - 1:
+        moved_down = current_idx < old_idx
-            current_step = self.candidate_steps[current_idx]
+        if not moved_down:
-            rise_threshold = current_step - 0.5 + self.up_hysteresis
+            while current_idx < len(self.candidate_steps) - 1:
-            if self.ema_accept_len > rise_threshold:
+                current_step = self.candidate_steps[current_idx]
-                current_idx += 1
+                rise_threshold = current_step - 0.5 + self.up_hysteresis
-            else:
+                if self.ema_accept_len > rise_threshold:
-                break
+                    current_idx += 1
+                else:
+                    break
 
         target = self.candidate_steps[current_idx]
         # EMA ceiling: only caps downward — never blocks step-ups, so the
@@ -218,6 +243,9 @@ class AdaptiveStepSlot:
                     current_idx -= 1
                 target = self.candidate_steps[current_idx]
 
+        return self._apply_target_steps(old_steps, target)
+
+    def _apply_target_steps(self, old_steps: int, target: int) -> bool:
         if target != old_steps:
             self.current_steps = target
             log_info_on_rank0(

python/sglang/srt/speculative/draft_utils.py

@@ -37,7 +37,8 @@ class DraftBackendFactory:
         return backend_map[backend_type]()
 
     def create_decode_backend(self):
-        if self.speculative_num_steps == 1:
+        # No multi-step draft backend for steps=0 (nospec) or steps=1.
+        if self.speculative_num_steps <= 1:
             return None
 
         backend_map = {

python/sglang/srt/speculative/eagle_worker_v2.py

@@ -1011,33 +1011,129 @@ class EAGLEWorkerV2(BaseSpecWorker):
                     topk=self.topk,
                     capture_hidden_mode=capture_mode,
                 )
-            with (
+            if self.speculative_num_steps == 0:
-                self.draft_worker.draft_tp_context(
+                # Drafting disabled (high batch size). _draft_extend below still
-                    self.draft_worker.draft_runner.tp_group
+                # runs, keeping draft KV warm for when the batch shrinks.
-                ),
+                verify_input = self._build_trivial_verify_input(batch)
-                speculative_moe_backend_context(),
+            else:
-                speculative_moe_a2a_backend_context(),
+                with (
-                spec_stage_span("draft"),
+                    self.draft_worker.draft_tp_context(
-            ):
+                        self.draft_worker.draft_runner.tp_group
-                verify_input: EagleVerifyInput = self.draft_worker.draft(batch)
+                    ),
+                    speculative_moe_backend_context(),
+                    speculative_moe_a2a_backend_context(),
+                    spec_stage_span("draft"),
+                ):
+                    verify_input: EagleVerifyInput = self.draft_worker.draft(batch)
             assert verify_input.is_verify_input()
             batch.spec_info = verify_input
             batch_output = self.verify(batch)
             # Publish before draft_extend so the fence is at verify-end.
             if on_publish is not None:
                 on_publish(batch_output.new_seq_lens)
-            with (
+            if (
-                self.draft_worker.draft_tp_context(
+                self.speculative_num_steps == 0
-                    self.draft_worker.draft_runner.tp_group
+                and envs.SGLANG_SPEC_SKIP_ZERO_STEP_DRAFT_EXTEND.get()
-                ),
-                speculative_moe_backend_context(),
-                speculative_moe_a2a_backend_context(),
-                spec_stage_span("draft_extend"),
             ):
-                self.draft_worker._draft_extend_for_decode(batch, batch_output)
+                self._stub_skipped_draft_extend(batch, batch_output)
+            else:
+                with (
+                    self.draft_worker.draft_tp_context(
+                        self.draft_worker.draft_runner.tp_group
+                    ),
+                    speculative_moe_backend_context(),
+                    speculative_moe_a2a_backend_context(),
+                    spec_stage_span("draft_extend"),
+                ):
+                    self.draft_worker._draft_extend_for_decode(batch, batch_output)
 
             return batch_output
 
+    def _build_trivial_verify_input(self, batch: ScheduleBatch) -> EagleVerifyInput:
+        """Build a 1-node EagleVerifyInput rooted at the previous bonus token.
+
+        Used when ``speculative_num_steps == 0`` to skip drafting while still
+        routing through the existing TARGET_VERIFY graph captured at
+        ``draft_token_num=1``: the kernel always accepts the root and samples
+        one new bonus token from target logits -- functionally a plain decode.
+        """
+        if batch.forward_mode.is_idle():
+            return EagleVerifyInput.create_idle_input(
+                topk=self.topk, spec_steps=0, num_verify_tokens=1
+            )
+
+        draft_input: EagleDraftInput = batch.spec_info
+        bs = batch.seq_lens.shape[0]
+        device = self.device
+
+        retrieve_index = torch.arange(bs, dtype=torch.long, device=device).unsqueeze(1)
+        retrieve_next_token = torch.full((bs, 1), -1, dtype=torch.long, device=device)
+        retrieve_next_sibling = torch.full((bs, 1), -1, dtype=torch.long, device=device)
+
+        attn_backend = self._target_worker.model_runner.attn_backend
+        mask_buf, position_buf = attn_backend.get_verify_buffers_to_fill_after_draft()
+        if mask_buf is not None:
+            custom_mask = mask_buf
+            custom_mask.fill_(True)
+        else:
+            if batch.seq_lens_sum is not None:
+                seq_lens_sum = batch.seq_lens_sum
+            elif batch.seq_lens_cpu is not None:
+                seq_lens_sum = int(batch.seq_lens_cpu.sum())
+            else:
+                seq_lens_sum = bs * attn_backend.max_context_len
+            custom_mask = torch.ones(seq_lens_sum + bs, dtype=torch.bool, device=device)
+
+        if position_buf is not None:
+            positions = position_buf
+            positions[:bs].copy_(batch.seq_lens)
+        else:
+            positions = batch.seq_lens.to(torch.int64)
+
+        return EagleVerifyInput(
+            draft_token=draft_input.bonus_tokens,
+            custom_mask=custom_mask,
+            positions=positions,
+            retrieve_index=retrieve_index,
+            retrieve_next_token=retrieve_next_token,
+            retrieve_next_sibling=retrieve_next_sibling,
+            retrieve_cum_len=None,
+            spec_steps=0,
+            topk=self.topk,
+            draft_token_num=1,
+            capture_hidden_mode=CaptureHiddenMode.FULL,
+            seq_lens_sum=None,
+            seq_lens_cpu=None,
+        )
+
+    def _stub_skipped_draft_extend(
+        self, batch: ScheduleBatch, batch_output: GenerationBatchResult
+    ) -> None:
+        """Fill shape-valid stubs on next_draft_input when draft_extend is skipped.
+
+        ``verify`` already set ``bonus_tokens`` (the only field the next steps=0
+        verify reads). The overlap FutureMap still stashes topk_p/topk_index/
+        hidden_states, so provide zeroed tensors of the right shape. They are never
+        consumed while at steps=0; an upshift to steps>0 would draft from this stale
+        state (cold recovery), which is the documented cost of this experimental flag.
+        """
+        next_draft_input: EagleDraftInput = batch_output.next_draft_input
+        bs = batch.seq_lens.shape[0]
+        device = self.device
+        next_draft_input.topk_p = torch.zeros(
+            (bs, self.topk), dtype=torch.float32, device=device
+        )
+        next_draft_input.topk_index = torch.zeros(
+            (bs, self.topk), dtype=torch.int64, device=device
+        )
+        hidden_size = EagleDraftInput.hidden_size_for(self.draft_worker)
+        if hidden_size is not None:
+            next_draft_input.hidden_states = torch.zeros(
+                (bs, hidden_size),
+                dtype=EagleDraftInput.dtype_for(self.draft_worker),
+                device=device,
+            )
+
     def on_verify_complete_cpu(
         self, num_correct_drafts_per_req: list[int], batch_size: int = 0
     ) -> None:

test/registered/spec/eagle/test_adaptive_speculative.py

@@ -18,7 +18,7 @@ from sglang.test.test_utils import (
     popen_launch_server,
 )
 
-register_cuda_ci(est_time=76, stage="base-b", runner_config="1-gpu-large")
+register_cuda_ci(est_time=160, stage="base-b", runner_config="1-gpu-large")
 
 HIGH_ACCEPT_PROMPT = (
     "Output exactly 128 new lines. "
@@ -198,5 +198,118 @@ class TestAdaptiveSpeculativeServer(CustomTestCase):
         )
 
 
+class TestAdaptiveZeroStepBatchSizeServer(CustomTestCase):
+    """steps=0 (nospec) fallback triggered by batch size.
+
+    Config routes BS>=8 -> steps=0 (drafting disabled) and BS<8 -> steps=3, so the
+    server cycles steps=3 -> steps=0 -> steps=3 as load rises and falls.
+    """
+
+    model = DEFAULT_TARGET_MODEL_EAGLE
+    draft_model = DEFAULT_DRAFT_MODEL_EAGLE
+    base_url = DEFAULT_URL_FOR_TEST
+
+    COUNT_PROMPT = "Count from 1 to 400, separated by commas. Output only the numbers."
+
+    @classmethod
+    def setUpClass(cls):
+        with tempfile.NamedTemporaryFile("w", suffix=".json", delete=False) as f:
+            json.dump(
+                {
+                    "1": {"candidate_steps": [3], "warmup_batches": 0},
+                    "8": {"candidate_steps": [0], "warmup_batches": 0},
+                },
+                f,
+            )
+            cls.adaptive_config_path = f.name
+
+        try:
+            cls.process = popen_launch_server(
+                cls.model,
+                cls.base_url,
+                timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+                other_args=[
+                    "--trust-remote-code",
+                    "--attention-backend",
+                    "triton",
+                    "--speculative-algorithm",
+                    "EAGLE",
+                    "--speculative-draft-model-path",
+                    cls.draft_model,
+                    "--speculative-num-steps",
+                    "3",
+                    "--speculative-eagle-topk",
+                    "1",
+                    "--speculative-num-draft-tokens",
+                    "4",
+                    "--speculative-adaptive",
+                    "--speculative-adaptive-config",
+                    cls.adaptive_config_path,
+                    "--max-running-requests",
+                    "32",
+                    "--skip-server-warmup",
+                    "--mem-fraction-static",
+                    "0.7",
+                ],
+            )
+        except Exception:
+            os.unlink(cls.adaptive_config_path)
+            raise
+
+    @classmethod
+    def tearDownClass(cls):
+        if hasattr(cls, "process"):
+            kill_process_tree(cls.process.pid)
+        if os.path.exists(cls.adaptive_config_path):
+            os.unlink(cls.adaptive_config_path)
+
+    def _steps(self) -> int:
+        r = requests.get(self.base_url + "/server_info", timeout=30)
+        self.assertEqual(r.status_code, 200, r.text)
+        return r.json()["internal_states"][0]["speculative_num_steps"]
+
+    def test_batch_size_step_cycle(self):
+        """The server cycles steps=3 -> steps=0 -> steps=3 as load rises and falls:
+        a BS=1 request drafts at steps=3; a 14-way batch (BS>=8) routes the worker
+        to nospec steps=0; a following BS=1 request returns to steps=3 with drafting
+        restored (high accept rate again)."""
+        one = {"temperature": 0, "max_new_tokens": 64, "ignore_eos": True}
+
+        def generate_single() -> dict:
+            r = requests.post(
+                self.base_url + "/generate",
+                json={"text": self.COUNT_PROMPT, "sampling_params": one},
+                timeout=600,
+            )
+            self.assertEqual(r.status_code, 200, r.text)
+            return r.json()["meta_info"]
+
+        # Phase 1: BS=1 -> steps=3, drafting active.
+        m1 = generate_single()
+        self.assertEqual(self._steps(), 3, "expected steps=3 at BS=1")
+        self.assertGreater(
+            m1["spec_accept_rate"], 0.8, f"not drafting at steps=3: {m1}"
+        )
+
+        # Phase 2: BS=14 -> the worker switches to nospec steps=0. Equal-length
+        # requests finish together, so the last decode batch (and thus the state)
+        # is at BS=14 -> steps=0.
+        full = {"temperature": 0, "max_new_tokens": 128, "ignore_eos": True}
+        r = requests.post(
+            self.base_url + "/generate",
+            json={"text": [self.COUNT_PROMPT] * 14, "sampling_params": [full] * 14},
+            timeout=600,
+        )
+        self.assertEqual(r.status_code, 200, r.text)
+        self.assertEqual(self._steps(), 0, "BS>=8 did not switch to steps=0")
+
+        # Phase 3: BS=1 -> steps=3 again, drafting restored.
+        m3 = generate_single()
+        self.assertEqual(self._steps(), 3, "did not reopen to steps=3")
+        self.assertGreater(
+            m3["spec_accept_rate"], 0.8, f"drafting not restored after steps=0: {m3}"
+        )
+
+
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/spec/test_adaptive_spec_params.py

@@ -198,6 +198,30 @@ class TestAdaptiveStepSlot(unittest.TestCase):
         self.assertEqual(params.current_steps, 1)
         self.assertEqual(params.ema_accept_len, 0.375)
 
+    def test_zero_step_mixed_slot_drops_probes_and_rechecks(self):
+        params = self._make_params_from_config(
+            3,
+            {
+                "candidate_steps": [0, 3],
+                "ema_alpha": 1.0,
+                "warmup_batches": 0,
+                "update_interval": 1,
+                "down_hysteresis": 0.0,
+            },
+        )
+
+        self.assertTrue(params.update([0, 0]))
+        self.assertEqual(params.current_steps, 0)
+        self.assertEqual(params.ema_accept_len, 0.0)
+
+        self.assertTrue(params.update([3, 3]))
+        self.assertEqual(params.current_steps, 3)
+        self.assertEqual(params.ema_accept_len, 0.0)
+
+        self.assertTrue(params.update([0, 0]))
+        self.assertEqual(params.current_steps, 0)
+        self.assertEqual(params.ema_accept_len, 0.0)
+
     def test_ceiling_coeff_caps_steps(self):
         params = self._make_params_from_config(
             7,
@@ -223,10 +247,11 @@ class TestAdaptiveStepSlot(unittest.TestCase):
 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._bs_list, [1, 8, 32, 64])
         self.assertEqual(params._slots[1].candidate_steps, [1, 3, 7])
-        self.assertEqual(params._slots[8].candidate_steps, [1, 3])
+        self.assertEqual(params._slots[8].candidate_steps, [0, 1, 3])
-        self.assertEqual(params._slots[32].candidate_steps, [1])
+        self.assertEqual(params._slots[32].candidate_steps, [0, 1])
+        self.assertEqual(params._slots[64].candidate_steps, [0])
 
     def test_config_file(self):
         with tempfile.NamedTemporaryFile("w", suffix=".json") as f:
@@ -287,13 +312,6 @@ class TestAdaptiveSpeculativeParams(unittest.TestCase):
             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(
@@ -333,20 +351,20 @@ class TestBatchSizeRouting(unittest.TestCase):
         # 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(8).candidate_steps, [0, 1, 3])
-        self.assertEqual(params._route(31).candidate_steps, [1, 3])
+        self.assertEqual(params._route(31).candidate_steps, [0, 1, 3])
-        self.assertEqual(params._route(32).candidate_steps, [1])
+        self.assertEqual(params._route(32).candidate_steps, [0, 1])
-        self.assertEqual(params._route(1000).candidate_steps, [1])
+        self.assertEqual(params._route(1000).candidate_steps, [0])
 
     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])
+        self.assertEqual(params._route(5).candidate_steps, [0, 1, 3])
         # bs=17 pads up to 32 -> slot bs=32.
-        self.assertEqual(params._route(17).candidate_steps, [1])
+        self.assertEqual(params._route(17).candidate_steps, [0, 1])
         # A batch larger than every captured BS keeps its own value -> top slot.
-        self.assertEqual(params._route(100).candidate_steps, [1])
+        self.assertEqual(params._route(100).candidate_steps, [0])
 
     def test_cuda_graph_bs_for_step_prunes_unreachable_graphs(self):
         params = self._params()