minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Cheng Wan
2026-07-04 02:21:43 -07:00
committed by GitHub
co-authored by Claude Fable 5
parent 8d8f17e28c
commit 4bf4db09f6
4 changed files with 159 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/arg_groups/overrides.py
+60 -8
View File
@@ -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
+2
View File
@@ -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
+14 -20
View File
@@ -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:
test/registered/unit/test_model_overrides.py
+83 -2
View File
@@ -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):