Decouple speculative draft capacity from runtime state (#36897)

Co-authored-by: Ming Yang <minos.future@gmail.com>
This commit is contained in:
Lianmin Zheng
2026-08-30 23:31:25 -07:00
committed by GitHub
co-authored by Ming Yang
parent 2cb3f32b03
commit 3ed3326631
6 changed files with 80 additions and 51 deletions
+3 -8
View File
@@ -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,