[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,28 +71,52 @@ def register_model_override(architecture: str):
     return decorator
 
 
+def register_model_override_predicate(predicate: Callable[[str], bool]):
+    """Register a derived-override provider keyed by an architecture
+    predicate. Same callable contract as ``register_model_override``."""
+
+    def decorator(fn: Callable[..., dict]) -> Callable[..., dict]:
+        _PREDICATE_OVERRIDE_FNS.append((predicate, fn))
+        return fn
+
+    return decorator
+
+
+def _invoke_provider(
+    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 registered callables in registration
-    order. Empty declarations are dropped.
+    ``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 = 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__}"
-            )
+        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,
-        (
-            "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."
+        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":
-                self.enable_multi_layer_eagle = True
-                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"
-                )
+            # EAGLE multi-layer + hierarchical-cache SWA writes moved to the
+            # override registry (arg_groups/overrides.py: _step3p_overrides).
         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: