From 8ae4a39b50ffcdb44175cd2fc5594032740ba888 Mon Sep 17 00:00:00 2001 From: Sam Shleifer Date: Wed, 16 Sep 2026 23:39:12 -0700 Subject: [PATCH] Gate mamba extra-buffer predicates on uses_mamba_radix_cache (#37474) Co-authored-by: Claude Fable 5 Co-authored-by: Ke Bao Co-authored-by: Sam Shleifer --- python/sglang/srt/arg_groups/model_hook.py | 7 +++++++ .../srt/arg_groups/model_overrides/inkling.py | 4 ++++ python/sglang/srt/arg_groups/overrides.py | 4 ++++ test/registered/unit/test_model_overrides.py | 15 +++++++++++++++ 4 files changed, 30 insertions(+) diff --git a/python/sglang/srt/arg_groups/model_hook.py b/python/sglang/srt/arg_groups/model_hook.py index e49392131..86e3ff97c 100644 --- a/python/sglang/srt/arg_groups/model_hook.py +++ b/python/sglang/srt/arg_groups/model_hook.py @@ -921,6 +921,13 @@ def handle_mamba_radix_cache(server_args: Any, model_arch: str): run_post_process_pass(server_args, _mamba_radix_cache_resolution) view = resolved_view(server_args) if not view.uses_mamba_radix_cache: + # auto is arch-gated, so only an explicit strategy reaches a non-mamba + # arch here, where it would arm the mamba paths and crash at prefill. + if mamba_extra_buffer_of(view): + raise ValueError( + f"--mamba-radix-cache-strategy {view.mamba_radix_cache_strategy} " + f"needs mamba state, got {model_arch}." + ) return if mamba_extra_buffer_of(view): diff --git a/python/sglang/srt/arg_groups/model_overrides/inkling.py b/python/sglang/srt/arg_groups/model_overrides/inkling.py index 8cf128289..f13de63f6 100644 --- a/python/sglang/srt/arg_groups/model_overrides/inkling.py +++ b/python/sglang/srt/arg_groups/model_overrides/inkling.py @@ -58,6 +58,10 @@ def _inkling_overrides(server_args: Any, hf_config: Any) -> dict: # spec today, but giving it one would silently stop this pin from firing. if cfg.mamba_radix_cache_strategy == "auto": overrides["mamba_radix_cache_strategy"] = "extra_buffer" + # The generic resolution never sets the arch-derived leaf for Inkling, and + # handle_mamba_radix_cache rejects extra_buffer on a model without it. + if not cfg.disable_radix_cache: + overrides["uses_mamba_radix_cache"] = True # Inkling attention runs only on the fa4 (Blackwell) or triton backends -- # models/inkling_common/attn.py asserts attention_backend in {fa4, triton}. # The generic resolver would otherwise pick trtllm_mha (SM100) / fa3 diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index e01cc31ac..de96b45e6 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -488,6 +488,10 @@ _MAMBA_EXTRA_BUFFER_ARCHS = frozenset( # KDA backend's track-snapshot writes (decode + extend) so donated # slots hold real states for prefix-cache restores. "KimiK3ForConditionalGeneration", + # Inkling asserts enable_mamba_extra_buffer and _inkling_overrides pins it, + # so validate_mamba_extra_buffer runs for these archs and must accept them. + "InklingForConditionalGeneration", + "InklingForConditionalGenerationMTP", } ) diff --git a/test/registered/unit/test_model_overrides.py b/test/registered/unit/test_model_overrides.py index 0dfeb162c..5fc30bcaf 100644 --- a/test/registered/unit/test_model_overrides.py +++ b/test/registered/unit/test_model_overrides.py @@ -591,6 +591,21 @@ class TestGoldenModelOverrides(_IsolatedPublish): set_global_server_args_for_scheduler(server_args) return get_server_args() + def test_explicit_extra_buffer_without_mamba_state_fails_fast(self): + with self.assertRaisesRegex(ValueError, "needs mamba state"): + self._construct( + "LlamaForCausalLM", "llama", mamba_radix_cache_strategy="extra_buffer" + ) + + def test_explicit_extra_buffer_is_harmless_with_radix_cache_disabled(self): + sa = self._construct( + "LlamaForCausalLM", + "llama", + mamba_radix_cache_strategy="extra_buffer", + disable_radix_cache=True, + ) + self.assertFalse(self._resolved(sa, "uses_mamba_radix_cache")) + def test_mistral_large3_forces_bfloat16(self): sa = self._construct("MistralLarge3ForCausalLM", "mistral") self.assertEqual(