spec: size the speculative buffers from the bags, not the startup record (#35024)

test/registered/unit/disaggregation/test_disaggregation_wire.py

@@ -34,6 +34,7 @@ from sglang.srt.environ import envs
 from sglang.srt.layers.attention.dsa.utils import should_use_dsa_fused_topk
 from sglang.srt.managers.overlap_utils import FutureMap, RelayPayload
 from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
+from sglang.srt.runtime_context import get_context
 from sglang.srt.speculative.eagle_disaggregation import (
     build_eagle_disagg_draft_input,
 )
@@ -302,17 +303,17 @@ class TestEagleDsaSeedTransfer(unittest.TestCase):
             device="cpu",
             enable_overlap=False,
         )
-        server_args = SimpleNamespace(
+        # The draft-input shape comes from the spec bag.
+        override = get_context().override_server_args(
             speculative_eagle_topk=1,
             speculative_num_steps=5,
             enable_multi_layer_eagle=False,
-            disaggregation_mode="null",
         )
+        override.install()
+        self.addCleanup(override.restore)
         last_tokens = torch.tensor([11, 12], dtype=torch.int64)
 
-        draft_input = build_eagle_disagg_draft_input(
-            batch, server_args, last_tokens, None
-        )
+        draft_input = build_eagle_disagg_draft_input(batch, last_tokens, None)
         self.assertTrue(torch.equal(draft_input.dsa_topk_indices, torch.stack(seeds)))
 
         for invalid_seed in (
@@ -320,9 +321,7 @@ class TestEagleDsaSeedTransfer(unittest.TestCase):
             torch.full((3,), -1, dtype=torch.int32),
         ):
             batch.reqs[1].output_dsa_topk_indices = invalid_seed
-            draft_input = build_eagle_disagg_draft_input(
-                batch, server_args, last_tokens, None
-            )
+            draft_input = build_eagle_disagg_draft_input(batch, last_tokens, None)
             self.assertIsNone(draft_input.dsa_topk_indices)
 
     def test_pd_decode_fused_topk_remaps_wire_positions_to_local_slots(self):
@@ -347,24 +346,24 @@ class TestEagleDsaSeedTransfer(unittest.TestCase):
             req_to_token_pool=SimpleNamespace(req_to_token=req_to_token),
             seq_lens=torch.tensor([4, 4], dtype=torch.int32),
         )
-        server_args = SimpleNamespace(
+        override = get_context().override_server_args(
             speculative_eagle_topk=1,
             speculative_num_steps=5,
             enable_multi_layer_eagle=False,
             disaggregation_mode="decode",
             enable_hisparse=False,
         )
+        override.install()
+        self.addCleanup(override.restore)
 
         with envs.SGLANG_DSA_FUSE_TOPK.override(True), patch(
             "sglang.srt.layers.attention.dsa.utils.is_cuda", return_value=True
         ):
             self.assertTrue(
-                should_use_dsa_fused_topk(
-                    server_args, seed_dsa_topk_from_draft_extend=True
-                )
+                should_use_dsa_fused_topk(seed_dsa_topk_from_draft_extend=True)
             )
             draft_input = build_eagle_disagg_draft_input(
-                batch, server_args, torch.tensor([11, 12], dtype=torch.int64), None
+                batch, torch.tensor([11, 12], dtype=torch.int64), None
             )
 
         self.assertEqual(

test/registered/unit/managers/test_priority_scheduling_disaggregation.py

@@ -432,8 +432,6 @@ class TestDecodePrebuiltPriority(unittest.TestCase):
         scheduler.enable_overlap = False
         scheduler.spec_algorithm = MagicMock()
         scheduler.max_running_requests = 1
-        # Passed whole into the (mocked) batch's process_prebuilt; never read.
-        scheduler.server_args = SimpleNamespace()
         scheduler.future_map = MagicMock()
         scheduler.policy = MagicMock()
         scheduler.policy.calc_priority.side_effect = lambda waiting_queue, _: (

test/registered/unit/spec/test_plugin_hook_signatures.py

@@ -0,0 +1,109 @@
+"""The plugin hook takes the call the dispatch makes.
+
+`CustomSpecAlgo` is the out-of-tree extension point: a registered algorithm's
+method is called through the same dispatch as the built-in ones, and nothing in
+the tree implements it, so a drift between the two sides only ever surfaces in
+somebody's plugin. It has drifted twice, both times on the disaggregation
+draft-input builder: the built-in dropped a parameter the hook kept, so every
+plugin call would have hit a TypeError.
+
+What is pinned here:
+
+  * every method the dispatch may call on either type takes the same arguments
+    on both -- the set is intersected out of the two types, never listed;
+  * the call the dispatch actually writes binds on both types.
+"""
+
+import ast
+import inspect
+import unittest
+from pathlib import Path
+
+from sglang.srt.disaggregation import decode_schedule_batch_mixin
+from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
+from sglang.srt.speculative.spec_registry import CustomSpecAlgo
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+
+def _dispatched_methods():
+    """Methods carried by both types. The dispatch calls them on whichever it
+    holds without knowing which, so their argument lists must agree."""
+    enum_methods = {
+        name
+        for name, value in vars(SpeculativeAlgorithm).items()
+        if inspect.isfunction(value)
+    }
+    hook_methods = {
+        name
+        for name, value in vars(CustomSpecAlgo).items()
+        if inspect.isfunction(value)
+    }
+    return sorted(enum_methods & hook_methods)
+
+
+def _dispatch_calls():
+    """Every call the decode dispatch makes on the algorithm object, read from
+    its source: `(method, positional count, keyword names)`."""
+    source = Path(decode_schedule_batch_mixin.__file__).read_text(encoding="utf-8")
+    calls = []
+    for node in ast.walk(ast.parse(source)):
+        if not isinstance(node, ast.Call):
+            continue
+        func = node.func
+        if not (
+            isinstance(func, ast.Attribute)
+            and isinstance(func.value, ast.Attribute)
+            and func.value.attr == "spec_algorithm"
+        ):
+            continue
+        calls.append((func.attr, len(node.args), tuple(kw.arg for kw in node.keywords)))
+    return calls
+
+
+def _parameters(function):
+    """Parameter names, without any catch-alls."""
+    return [
+        name
+        for name, parameter in inspect.signature(function).parameters.items()
+        if parameter.kind
+        not in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD)
+    ]
+
+
+class TestDispatchedSignatures(CustomTestCase):
+    def test_the_hook_and_the_built_in_agree(self):
+        methods = _dispatched_methods()
+        self.assertNotEqual(methods, [], "the dispatched set derived to nothing")
+        mismatches = []
+        for name in methods:
+            hook = _parameters(getattr(CustomSpecAlgo, name))
+            builtin = _parameters(getattr(SpeculativeAlgorithm, name))
+            if hook != builtin:
+                mismatches.append(
+                    f"{name}: CustomSpecAlgo{tuple(hook)} vs "
+                    f"SpeculativeAlgorithm{tuple(builtin)}"
+                )
+        self.assertEqual(
+            mismatches,
+            [],
+            "a plugin implementing the hook would be called with the "
+            "dispatch's arguments:\n  " + "\n  ".join(mismatches),
+        )
+
+    def test_the_dispatch_call_binds_on_both_types(self):
+        calls = _dispatch_calls()
+        self.assertNotEqual(calls, [], "no dispatch call found to bind against")
+        for method, positional, keywords in calls:
+            self.assertIn(method, _dispatched_methods())
+            for owner in (CustomSpecAlgo, SpeculativeAlgorithm):
+                arguments = [None] * (1 + positional)
+                inspect.signature(getattr(owner, method)).bind(
+                    *arguments, **{name: None for name in keywords}
+                )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/test_publish_precedes_bag_reads.py

@@ -111,10 +111,6 @@ _UNREAD_ENTRIES: dict = {
     ("srt/managers/detokenizer_manager.py", "run_detokenizer_process"): (
         "DetokenizerManager reads the handed instance at this revision"
     ),
-    ("srt/managers/tokenizer_manager.py", "__init__"): (
-        "the constructor and the init_* helpers it calls read the handed "
-        "instance at this revision"
-    ),
 }
 
 # `publish` itself and its named wrappers live here; a call inside them is the

test/registered/unit/test_supplied_instance_exposure_ratchet.py

@@ -242,7 +242,6 @@ _EXPOSED = {
     ("kv_canary/capacities.py", "cuda_graph_config"),
     ("kv_canary/capacities.py", "speculative_num_draft_tokens"),
     ("kv_canary/token_oracle/install.py", "sampling_backend"),
-    ("layers/attention/dsa/utils.py", "disaggregation_mode"),
     ("layers/cp/base.py", "attn_cp_size"),
     ("layers/cp/base.py", "cp_strategy"),
     ("layers/cp/base.py", "enable_prefill_cp"),
@@ -334,13 +333,6 @@ _EXPOSED = {
     ("managers/tp_worker.py", "random_seed"),
     ("managers/tp_worker.py", "speculative_algorithm"),
     ("managers/tp_worker.py", "tokenizer_path"),
-    ("managers/utils.py", "speculative_algorithm"),
-    ("managers/utils.py", "speculative_eagle_topk"),
-    ("managers/utils.py", "speculative_num_steps"),
-    ("mem_cache/allocation_sizing.py", "page_size"),
-    ("mem_cache/allocation_sizing.py", "speculative_algorithm"),
-    ("mem_cache/allocation_sizing.py", "speculative_eagle_topk"),
-    ("mem_cache/allocation_sizing.py", "speculative_num_steps"),
     ("mem_cache/hiradix_cache.py", "hicache_io_backend"),
     ("mem_cache/hiradix_cache.py", "hicache_mem_layout"),
     ("mem_cache/hiradix_cache.py", "served_model_name"),
@@ -431,9 +423,6 @@ _EXPOSED = {
         "speculative/dspark_components/dspark_worker_v2.py",
         "speculative_num_draft_tokens",
     ),
-    ("speculative/eagle_disaggregation.py", "enable_multi_layer_eagle"),
-    ("speculative/eagle_disaggregation.py", "speculative_eagle_topk"),
-    ("speculative/eagle_disaggregation.py", "speculative_num_steps"),
     ("speculative/eagle_worker_v2.py", "device"),
     ("speculative/eagle_worker_v2.py", "enable_dp_attention"),
     ("speculative/eagle_worker_v2.py", "speculative_adaptive"),
@@ -460,12 +449,7 @@ _EXPOSED = {
     ("speculative/ngram_worker.py", "speculative_num_draft_tokens"),
     ("speculative/ngram_worker.py", "speculative_num_steps"),
     ("speculative/spec_info.py", "enable_multi_layer_eagle"),
-    ("speculative/spec_info.py", "speculative_eagle_topk"),
-    ("speculative/spec_info.py", "speculative_num_draft_tokens"),
-    ("speculative/spec_info.py", "speculative_num_steps"),
     ("speculative/spec_registry.py", "disable_overlap_schedule"),
-    ("speculative/spec_utils.py", "speculative_eagle_topk"),
-    ("speculative/spec_utils.py", "speculative_num_draft_tokens"),
     ("speculative/standalone_worker_v2.py", "device"),
     ("speculative/standalone_worker_v2.py", "enable_dp_attention"),
     ("speculative/standalone_worker_v2.py", "speculative_algorithm"),
@@ -533,8 +517,6 @@ _OVERRIDDEN_AND_READ = {
     ("managers/tokenizer_manager.py", "model_path"),
     ("managers/tokenizer_manager.py", "speculative_num_draft_tokens"),
     ("managers/tp_worker.py", "model_path"),
-    ("managers/utils.py", "speculative_num_steps"),
-    ("mem_cache/allocation_sizing.py", "speculative_num_steps"),
     ("mem_cache/hiradix_cache.py", "hicache_storage_backend"),
     ("mem_cache/hiradix_cache.py", "hicache_storage_backend_extra_config"),
     ("mem_cache/hiradix_cache.py", "hicache_storage_prefetch_policy"),
@@ -561,7 +543,6 @@ _OVERRIDDEN_AND_READ = {
         "speculative/dspark_components/dspark_worker_v2.py",
         "speculative_num_draft_tokens",
     ),
-    ("speculative/eagle_disaggregation.py", "speculative_num_steps"),
     ("speculative/eagle_worker_v2.py", "speculative_num_draft_tokens"),
     ("speculative/eagle_worker_v2.py", "speculative_num_steps"),
     ("speculative/frozen_kv_mtp_worker_v2.py", "speculative_num_draft_tokens"),
@@ -570,9 +551,6 @@ _OVERRIDDEN_AND_READ = {
     ("speculative/multi_layer_eagle_worker_v2.py", "speculative_num_steps"),
     ("speculative/ngram_worker.py", "speculative_num_draft_tokens"),
     ("speculative/ngram_worker.py", "speculative_num_steps"),
-    ("speculative/spec_info.py", "speculative_num_draft_tokens"),
-    ("speculative/spec_info.py", "speculative_num_steps"),
-    ("speculative/spec_utils.py", "speculative_num_draft_tokens"),
     ("speculative/standalone_worker_v2.py", "speculative_num_draft_tokens"),
     ("speculative/standalone_worker_v2.py", "speculative_num_steps"),
     ("utils/common.py", "speculative_num_draft_tokens"),