[refactor] Add predicate-keyed registration; migrate the Step3p family (stack 8/15) (#30070)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

python/sglang/srt/arg_groups/overrides.py

@@ -48,6 +48,11 @@ MODEL_OVERRIDES: Dict[str, Dict[str, Any]] = {
 # Derived per-architecture override providers, in registration order.
 _MODEL_OVERRIDE_FNS: Dict[str, List[Callable[..., dict]]] = {}
 
+# Predicate-keyed providers, in registration order — for legacy branches
+# matched by substring/predicate on the architecture string rather than an
+# exact name (e.g. '"Step3p5ForCausalLM" in model_arch').
+_PREDICATE_OVERRIDE_FNS: List[Tuple[Callable[[str], bool], Callable[..., dict]]] = []
+
 
 def register_model_override(architecture: str):
     """Register a derived-override provider for ``architecture``.
@@ -66,26 +71,50 @@ def register_model_override(architecture: str):
     return decorator
 
 
-def collect_model_override_declarations(
+def register_model_override_predicate(predicate: Callable[[str], bool]):
-    architecture: str, server_args: Any, hf_config: Any
+    """Register a derived-override provider keyed by an architecture
-) -> List[Tuple[str, Dict[str, Any]]]:
+    predicate. Same callable contract as ``register_model_override``."""
-    """Collect ``(source, declaration)`` pairs for one architecture.
 
-    Application order (last writer wins downstream in the gate): the constant
+    def decorator(fn: Callable[..., dict]) -> Callable[..., dict]:
-    ``MODEL_OVERRIDES`` entry first, then registered callables in registration
+        _PREDICATE_OVERRIDE_FNS.append((predicate, fn))
-    order. Empty declarations are dropped.
+        return fn
-    """
+
-    declarations: List[Tuple[str, Dict[str, Any]]] = []
+    return decorator
-    const = MODEL_OVERRIDES.get(architecture)
+
-    if const:
+
-        declarations.append((f"MODEL_OVERRIDES[{architecture!r}]", dict(const)))
+def _invoke_provider(
-    for fn in _MODEL_OVERRIDE_FNS.get(architecture, ()):
+    fn: Callable[..., dict], server_args: Any, hf_config: Any
+) -> Dict[str, Any]:
     declared = fn(server_args, hf_config)
     if not isinstance(declared, dict):
         raise TypeError(
             f"model override provider {fn.__qualname__} must return a dict, "
             f"got {type(declared).__name__}"
         )
+    return declared
+
+
+def collect_model_override_declarations(
+    architecture: str, server_args: Any, hf_config: Any
+) -> List[Tuple[str, Dict[str, Any]]]:
+    """Collect ``(source, declaration)`` pairs for one architecture.
+
+    Application order (last writer wins downstream in the gate): the constant
+    ``MODEL_OVERRIDES`` entry first, then exact-keyed callables in
+    registration order, then matching predicate-keyed callables in
+    registration order. Empty declarations are dropped.
+    """
+    declarations: List[Tuple[str, Dict[str, Any]]] = []
+    const = MODEL_OVERRIDES.get(architecture)
+    if const:
+        declarations.append((f"MODEL_OVERRIDES[{architecture!r}]", dict(const)))
+    for fn in _MODEL_OVERRIDE_FNS.get(architecture, ()):
+        declared = _invoke_provider(fn, server_args, hf_config)
+        if declared:
+            declarations.append((fn.__qualname__, dict(declared)))
+    for predicate, fn in _PREDICATE_OVERRIDE_FNS:
+        if predicate(architecture):
+            declared = _invoke_provider(fn, server_args, hf_config)
             if declared:
                 declarations.append((fn.__qualname__, dict(declared)))
     return declarations
@@ -126,6 +155,29 @@ def _minimax_m2_overrides(server_args: Any, hf_config: Any) -> dict:
     return {"enable_tf32_matmul": True}
 
 
+@register_model_override_predicate(
+    lambda arch: "Step3p5ForCausalLM" in arch
+    or "Step3p7ForConditionalGeneration" in arch
+)
+def _step3p_overrides(server_args: Any, hf_config: Any) -> dict:
+    overrides: Dict[str, Any] = {}
+    if server_args.speculative_algorithm == "EAGLE":
+        logger.info(
+            "Enable multi-layer EAGLE speculative decoding for Step3p5ForCausalLM model."
+        )
+        overrides["enable_multi_layer_eagle"] = True
+    if server_args.enable_hierarchical_cache:
+        logger.warning(
+            "Reset swa_full_tokens_ratio to 1.0 for Step3p5ForCausalLM model with hierarchical cache"
+        )
+        overrides["swa_full_tokens_ratio"] = 1.0
+        logger.warning(
+            "Disable hybrid SWA memory for Step3p5ForCausalLM model with hierarchical cache"
+        )
+        overrides["disable_hybrid_swa_memory"] = True
+    return overrides
+
+
 @dataclasses.dataclass(frozen=True)
 class OverrideRecord:
     """Provenance of one resolved write: ``base`` is the value before this

python/sglang/srt/runtime_context.py

@@ -312,6 +312,8 @@ class Flags(_StaticFlags):
     dtype: str = "auto"
     enable_tf32_matmul: bool = False
     enable_multi_layer_eagle: bool = False
+    swa_full_tokens_ratio: float = 0.8
+    disable_hybrid_swa_memory: bool = False
 
     def freeze(self) -> None:
         for field in dataclasses.fields(self):

python/sglang/srt/server_args.py

@@ -754,14 +754,20 @@ class ServerArgs:
     page_size: A[Optional[int], "The number of tokens in a page."] = None
     swa_full_tokens_ratio: A[
         float,
-        (
+        Arg(
+            help=(
                 "The ratio of SWA layer KV tokens / full layer KV tokens, regardless "
                 "of the number of swa:full layers. It should be between 0 and 1. "
                 "E.g. 0.5 means if each swa layer has 50 tokens, then each full "
                 "layer has 100 tokens."
             ),
+            model_overridable=True,
+        ),
     ] = 0.8
-    disable_hybrid_swa_memory: A[bool, "Disable the hybrid SWA memory pool."] = False
+    disable_hybrid_swa_memory: A[
+        bool,
+        Arg(help="Disable the hybrid SWA memory pool.", model_overridable=True),
+    ] = False
     radix_eviction_policy: A[
         str,
         Arg(
@@ -4235,20 +4241,8 @@ class ServerArgs:
                     logger.info(
                         "Auto-select fa3 attention backend for Step3p7 on Hopper."
                     )
-            if self.speculative_algorithm == "EAGLE":
+            # EAGLE multi-layer + hierarchical-cache SWA writes moved to the
-                self.enable_multi_layer_eagle = True
+            # override registry (arg_groups/overrides.py: _step3p_overrides).
-                logger.info(
-                    "Enable multi-layer EAGLE speculative decoding for Step3p5ForCausalLM model."
-                )
-            if self.enable_hierarchical_cache:
-                self.swa_full_tokens_ratio = 1.0
-                logger.warning(
-                    "Reset swa_full_tokens_ratio to 1.0 for Step3p5ForCausalLM model with hierarchical cache"
-                )
-                self.disable_hybrid_swa_memory = True
-                logger.warning(
-                    "Disable hybrid SWA memory for Step3p5ForCausalLM model with hierarchical cache"
-                )
         elif model_arch in LLAMA4_MODEL_ARCHS and self.device != "cpu":
             # Auto-select attention backend for Llama4 if not specified
             if self.attention_backend is None:

test/registered/unit/test_model_overrides.py

@@ -60,7 +60,15 @@ class TestModelOverridableWhitelist(CustomTestCase):
 
         self.assertEqual(
             model_overridable_fields(ServerArgs),
-            frozenset({"dtype", "enable_tf32_matmul", "enable_multi_layer_eagle"}),
+            frozenset(
+                {
+                    "dtype",
+                    "enable_tf32_matmul",
+                    "enable_multi_layer_eagle",
+                    "swa_full_tokens_ratio",
+                    "disable_hybrid_swa_memory",
+                }
+            ),
         )
 
     def test_non_dataclass_yields_empty_whitelist(self):
@@ -75,6 +83,7 @@ class _IsolatedRegistry(CustomTestCase):
         self._patches = [
             patch.dict(overrides_module.MODEL_OVERRIDES, clear=True),
             patch.dict(overrides_module._MODEL_OVERRIDE_FNS, clear=True),
+            patch.object(overrides_module, "_PREDICATE_OVERRIDE_FNS", []),
         ]
         for p in self._patches:
             p.start()
@@ -131,6 +140,27 @@ class TestModelOverrideRegistry(_IsolatedRegistry):
         with self.assertRaises(TypeError):
             collect_model_override_declarations("FakeForCausalLM", None, None)
 
+    def test_predicate_keyed_provider(self):
+        from sglang.srt.arg_groups.overrides import register_model_override_predicate
+
+        @register_model_override("FakeStep9ForCausalLM")
+        def _exact(server_args, hf_config):
+            return {"a": 1}
+
+        @register_model_override_predicate(lambda arch: "Step9" in arch)
+        def _by_predicate(server_args, hf_config):
+            return {"b": 2}
+
+        # matching arch: exact-keyed first, then predicate-keyed
+        self.assertEqual(
+            collect_model_override_declarations("FakeStep9ForCausalLM", None, None),
+            [(_exact.__qualname__, {"a": 1}), (_by_predicate.__qualname__, {"b": 2})],
+        )
+        # non-matching arch: predicate does not fire
+        self.assertEqual(
+            collect_model_override_declarations("OtherForCausalLM", None, None), []
+        )
+
 
 @dataclasses.dataclass
 class _FakeAttnGroup(_StaticFlags):
@@ -292,13 +322,14 @@ class TestGoldenModelOverrides(_IsolatedPublish):
         "v_head_dim": 16,
     }
 
-    def _construct(self, arch, model_type, **server_kwargs):
+    def _construct(self, arch, model_type, config_extra=None, **server_kwargs):
         from sglang.srt.server_args import ServerArgs
 
         # Golden resolution must be host-independent: accelerator-less CI
         # runners resolve only the base platform, where get_device() raises.
         server_kwargs.setdefault("device", "cuda")
         config = dict(self._MINI_CONFIG, architectures=[arch], model_type=model_type)
+        config.update(config_extra or {})
         config_dir = tempfile.mkdtemp(prefix="golden_override_")
         self.addCleanup(shutil.rmtree, config_dir, ignore_errors=True)
         with open(os.path.join(config_dir, "config.json"), "w") as f:
@@ -377,6 +408,56 @@ class TestGoldenModelOverrides(_IsolatedPublish):
             [("_mimo_v2_overrides", {"enable_multi_layer_eagle": True})],
         )
 
+    def test_step3p_hierarchical_cache_golden(self):
+        # SWA-hybrid arch: the mini config needs layer_types/sliding_window.
+        config_extra = {
+            "layer_types": ["sliding_attention", "full_attention"],
+            "sliding_window": 64,
+        }
+        sa = self._construct(
+            "Step3p5ForCausalLM",
+            "llama",
+            config_extra=config_extra,
+            enable_hierarchical_cache=True,
+        )
+        # dual-apply == legacy writes
+        self.assertEqual(sa.swa_full_tokens_ratio, 1.0)
+        self.assertTrue(sa.disable_hybrid_swa_memory)
+        flags = self._publish(sa)
+        self.assertEqual(flags.swa_full_tokens_ratio, 1.0)
+        self.assertTrue(flags.disable_hybrid_swa_memory)
+
+    def test_step3p_declarations_at_callable_level(self):
+        from sglang.srt.arg_groups.overrides import _step3p_overrides
+
+        self.assertEqual(
+            _step3p_overrides(
+                SimpleNamespace(
+                    speculative_algorithm="EAGLE", enable_hierarchical_cache=False
+                ),
+                None,
+            ),
+            {"enable_multi_layer_eagle": True},
+        )
+        self.assertEqual(
+            _step3p_overrides(
+                SimpleNamespace(
+                    speculative_algorithm=None, enable_hierarchical_cache=True
+                ),
+                None,
+            ),
+            {"swa_full_tokens_ratio": 1.0, "disable_hybrid_swa_memory": True},
+        )
+        self.assertEqual(
+            _step3p_overrides(
+                SimpleNamespace(
+                    speculative_algorithm=None, enable_hierarchical_cache=False
+                ),
+                None,
+            ),
+            {},
+        )
+
 
 class TestDualApplyParity(CustomTestCase):
     def test_dual_apply_replays_and_parity_holds(self):