Decouple speculative draft capacity from runtime state (#36897)

Co-authored-by: Ming Yang <minos.future@gmail.com>

python/sglang/srt/arg_groups/overrides.py

@@ -1879,22 +1879,22 @@ def max_speculative_num_draft_tokens(server_args: Any) -> Optional[int]:
     memo = server_args.__dict__.get("_max_speculative_num_draft_tokens")
     if memo is not None:
         return memo
-    if cfg.speculative_num_draft_tokens is None:
+    from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
-        result = None
-    elif not cfg.speculative_adaptive:
-        result = cfg.speculative_num_draft_tokens
-    else:
-        from sglang.srt.speculative.adaptive_spec_params import (
-            resolve_candidate_steps_from_config,
-        )
 
-        candidate_steps = resolve_candidate_steps_from_config(
+    result = SpeculativeAlgorithm.from_string(
-            cfg_path=cfg.speculative_adaptive_config,
+        cfg.speculative_algorithm
+    ).resolve_max_speculative_num_draft_tokens(server_args)
+    if (
+        result is not None
+        and cfg.speculative_num_draft_tokens is not None
+        and result < cfg.speculative_num_draft_tokens
+    ):
+        raise ValueError(
+            "The speculative algorithm declared "
+            f"max_speculative_num_draft_tokens={result}, below the configured "
+            "speculative_num_draft_tokens="
+            f"{cfg.speculative_num_draft_tokens}."
         )
-        # TODO: adaptive spec currently requires topk=1, so each runtime
-        # state needs steps + 1 draft-token slots. Revisit this if topk>1
-        # is supported.
-        result = max(candidate_steps) + 1
     if getattr(server_args, "_resolution_finished", False):
         server_args._max_speculative_num_draft_tokens = result
     return result

python/sglang/srt/runtime_context.py

@@ -939,14 +939,22 @@ class RuntimeContext:
             _resolved_or_field(server_args, "enable_torch_compile", False)
         )
         self._server_args = server_args
-        # The adaptive draft-token bound memoizes on the config *path*, so a new
-        # publication that reuses the path must not inherit the bound computed
-        # from the file's previous contents.
-        _adaptive_draft_token_bound.cache_clear()
         # Snapshot resolved config into the namespace bags (the single source of
         # truth for config reads). Driven by NS(...) metadata; a mock/partial
         # config with no NS markers yields an empty tree (no bags projected).
         self._config_bags = _build_config_bags(server_args)
+        spec = self._config_bags.get("spec")
+        if spec is not None:
+            from sglang.srt.arg_groups.overrides import (
+                max_speculative_num_draft_tokens as max_draft_tokens_of,
+            )
+
+            # Keep the launch-time capacity stable while adaptive algorithms
+            # change the active width in this bag.
+            spec._set(
+                "max_speculative_num_draft_tokens",
+                max_draft_tokens_of(server_args),
+            )
         # Wire the published `parallel` bag onto the live wrapper: it is the slot
         # the `config` property reads, which is how config-only leaves like
         # pp_max_micro_batch_size are spelled.
@@ -1625,7 +1633,6 @@ def restore_context(state: dict[str, Any]) -> None:
             setattr(_CONTEXT, name, value)
     for name, value in state["__parallel__"].items():
         setattr(_CONTEXT.parallel, name, value)
-    _adaptive_draft_token_bound.cache_clear()
     set_global_dwdp_manager(state["__dwdp__"])
 
 
@@ -1639,7 +1646,6 @@ def reset_context() -> None:
     """
     _CONTEXT._server_args = None
     _CONTEXT._config_bags = None
-    _adaptive_draft_token_bound.cache_clear()
     _CONTEXT._overrides_log = []
     _CONTEXT._publish_role = None
     _CONTEXT.parallel._config = None
@@ -1917,33 +1923,23 @@ def mamba_track_grid(tree_page: int) -> int:
 def max_speculative_num_draft_tokens() -> int | None:
     """The largest draft-token count speculative decoding may use.
 
-    All three inputs are ``spec`` leaves, so this derives from the bags and
+    Adaptive algorithms may switch to a longer state after the scheduler
-    follows a post-publish override; ``overrides.max_speculative_num_draft_tokens``
+    reserves KV for the next decode batch, so include the capacity captured
-    is the pre-publish equivalent. Adaptive spec resolves the count from its
+    when the resolved configuration was published.
-    candidate-step table instead of the flat field.
     """
     spec = get_spec()
-    if spec.speculative_num_draft_tokens is None:
+    return max(
-        return None
+        (
-    if not spec.speculative_adaptive:
+            bound
-        return spec.speculative_num_draft_tokens
+            for bound in (
-    # The adaptive branch parses a JSON config, and this is called per decode
+                spec.speculative_num_draft_tokens,
-    # batch (`spec_prepare_for_decode`), so memoize on the inputs -- keyed, not
+                spec.max_speculative_num_draft_tokens,
-    # cached once, so a post-publish override still recomputes.
+            )
-    return _adaptive_draft_token_bound(spec.speculative_adaptive_config)
+            if bound is not None
-
+        ),
-
+        default=None,
-@functools.lru_cache(maxsize=8)
-def _adaptive_draft_token_bound(cfg_path: str | None) -> int:
-    from sglang.srt.speculative.adaptive_spec_params import (
-        resolve_candidate_steps_from_config,
     )
 
-    candidate_steps = resolve_candidate_steps_from_config(cfg_path=cfg_path)
-    # Adaptive spec requires topk=1 today, so each runtime state needs
-    # steps + 1 draft-token slots (mirrors the ServerArgs member).
-    return max(candidate_steps) + 1
-
 
 def uses_mla_backend() -> bool:
     """Whether this process's model runs the MLA attention path."""

python/sglang/srt/speculative/spec_info.py

@@ -226,6 +226,29 @@ class SpeculativeAlgorithm(Enum):
         elif self.is_ngram():
             _handle_ngram(server_args)
 
+    def resolve_max_speculative_num_draft_tokens(
+        self, server_args: ServerArgs
+    ) -> Optional[int]:
+        """Return the largest draft-token width this algorithm may use."""
+        from sglang.srt.arg_groups.overrides import resolving_view
+
+        cfg = resolving_view(server_args)
+        if cfg.speculative_num_draft_tokens is None:
+            return None
+        if not cfg.speculative_adaptive:
+            return cfg.speculative_num_draft_tokens
+
+        from sglang.srt.speculative.adaptive_spec_params import (
+            resolve_candidate_steps_from_config,
+        )
+
+        candidate_steps = resolve_candidate_steps_from_config(
+            cfg_path=cfg.speculative_adaptive_config,
+        )
+        # Adaptive spec requires topk=1 today, so each runtime state needs
+        # steps + 1 draft-token slots. Revisit this if topk>1 is supported.
+        return max(candidate_steps) + 1
+
     def get_num_tokens_per_req_for_target_verify(
         self, num_draft_tokens: int, is_draft_worker: bool
     ) -> int:

python/sglang/srt/speculative/spec_registry.py

@@ -109,6 +109,19 @@ class CustomSpecAlgo:
     def handle_server_args(self, server_args: ServerArgs) -> None:
         pass
 
+    def resolve_max_speculative_num_draft_tokens(
+        self, server_args: ServerArgs
+    ) -> Optional[int]:
+        """Return the largest draft-token width this algorithm may use.
+
+        The default covers static algorithms and adaptive algorithms whose
+        runtime states never exceed their startup width. Overrides must not
+        return less than ``server_args.speculative_num_draft_tokens``.
+        """
+        from sglang.srt.arg_groups.overrides import resolving_view
+
+        return resolving_view(server_args).speculative_num_draft_tokens
+
     def create_worker(self, server_args: ServerArgs) -> Type:
 
         cfg = resolving_view(server_args)

test/registered/unit/server_args/test_server_args.py

@@ -41,6 +41,7 @@ from sglang.srt.arg_groups.moe_hook import (
 )
 from sglang.srt.arg_groups.overrides import (
     cutedsl_moe_max_num_tokens,
+    max_speculative_num_draft_tokens,
     resolution_result,
 )
 from sglang.srt.arg_groups.parallel_hook import (
@@ -1681,6 +1682,7 @@ class TestAdaptiveSpecArgs(CustomTestCase):
             )
 
             handle_speculative_decoding(args)
+            self.assertEqual(max_speculative_num_draft_tokens(args), 6)
 
         self.assertTrue(resolution_result(args, "speculative_adaptive"))
         self.assertEqual(resolution_result(args, "speculative_eagle_topk"), 1)

test/registered/unit/test_runtime_context.py

@@ -530,6 +530,7 @@ class _FakeResolvedArgs:
     decode_attention_backend: A[str | None, Arg(help="dab"), NS("exec.kernel")] = None
     disable_radix_cache: A[bool, Arg(help="drc"), NS("memory")] = False
     mamba_radix_cache_strategy: A[str, Arg(help="mrcs"), NS("exec.mamba")] = "auto"
+    speculative_algorithm: A[str | None, Arg(help="sa"), NS("spec")] = None
     speculative_num_draft_tokens: A[int | None, Arg(help="d"), NS("spec")] = None
     speculative_adaptive: A[bool, Arg(help="a"), NS("spec")] = False
     speculative_adaptive_config: A[str | None, Arg(help="c"), NS("spec")] = None
@@ -1289,13 +1290,7 @@ class TestDerivedPredicatesAgreeAcrossTiers(_IsolatedServerArgs):
 
 
 class TestAdaptiveDraftBoundLifecycle(_IsolatedServerArgs):
-    """The adaptive draft-token bound is memoized on the config path, so the
+    """The adaptive draft-token bound is snapshotted at each publication."""
-    memo has to end with the publication it was computed under.
-
-    Without that, a process that republishes with the same adaptive-config path
-    -- the file having been rewritten in between -- keeps the previous bound and
-    under-allocates the draft-token buffers sized from it.
-    """
 
     def _write_config(self, steps):
         path = os.path.join(tempfile.mkdtemp(prefix="adaptive_cfg_"), "adaptive.json")
@@ -1317,7 +1312,7 @@ class TestAdaptiveDraftBoundLifecycle(_IsolatedServerArgs):
 
         with open(path, "w") as handle:
             json.dump({"1": {"candidate_steps": [4]}}, handle)
-        # Same path, new contents: the memo must not survive the republish.
+        # The new publication must not retain the previous capacity.
         get_context().set_server_args(
             _FakeResolvedArgs(
                 speculative_num_draft_tokens=3,