diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index 27703d3b0..9debcb6d7 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/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 diff --git a/python/sglang/srt/runtime_context.py b/python/sglang/srt/runtime_context.py index 898a39ac2..f379eef83 100644 --- a/python/sglang/srt/runtime_context.py +++ b/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.""" diff --git a/python/sglang/srt/speculative/spec_info.py b/python/sglang/srt/speculative/spec_info.py index c2c494ed8..42bf1d3ac 100644 --- a/python/sglang/srt/speculative/spec_info.py +++ b/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: diff --git a/python/sglang/srt/speculative/spec_registry.py b/python/sglang/srt/speculative/spec_registry.py index 20c0902bf..9205ed701 100644 --- a/python/sglang/srt/speculative/spec_registry.py +++ b/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) diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index 6b17789f5..803a62343 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/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) diff --git a/test/registered/unit/test_runtime_context.py b/test/registered/unit/test_runtime_context.py index aa1d9c52c..9bb5a3250 100644 --- a/test/registered/unit/test_runtime_context.py +++ b/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 - 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. - """ + """The adaptive draft-token bound is snapshotted at each publication.""" 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,