diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index 8833fe741..ccf01daef 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/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 diff --git a/python/sglang/srt/runtime_context.py b/python/sglang/srt/runtime_context.py index 0380062e4..246f68fed 100644 --- a/python/sglang/srt/runtime_context.py +++ b/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): diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index bd00ef016..4cbf17ff4 100644 --- a/python/sglang/srt/server_args.py +++ b/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: diff --git a/test/registered/unit/test_model_overrides.py b/test/registered/unit/test_model_overrides.py index 168ef2db7..9868d9405 100644 --- a/test/registered/unit/test_model_overrides.py +++ b/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):