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:
-        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,
-        )
+    from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
 
-        candidate_steps = resolve_candidate_steps_from_config(
-            cfg_path=cfg.speculative_adaptive_config,
+    result = SpeculativeAlgorithm.from_string(
+        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
-    follows a post-publish override; ``overrides.max_speculative_num_draft_tokens``
-    is the pre-publish equivalent. Adaptive spec resolves the count from its
-    candidate-step table instead of the flat field.
+    Adaptive algorithms may switch to a longer state after the scheduler
+    reserves KV for the next decode batch, so include the capacity captured
+    when the resolved configuration was published.
     """
     spec = get_spec()
-    if spec.speculative_num_draft_tokens is None:
-        return None
-    if not spec.speculative_adaptive:
-        return spec.speculative_num_draft_tokens
-    # The adaptive branch parses a JSON config, and this is called per decode
-    # batch (`spec_prepare_for_decode`), so memoize on the inputs -- keyed, not
-    # cached once, so a post-publish override still recomputes.
-    return _adaptive_draft_token_bound(spec.speculative_adaptive_config)
-
-
-@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,
+    return max(
+        (
+            bound
+            for bound in (
+                spec.speculative_num_draft_tokens,
+                spec.max_speculative_num_draft_tokens,
+            )
+            if bound is not None
+        ),
+        default=None,
     )
 
-    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)