diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index ccf01daef..b08cef78a 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -34,6 +34,7 @@ from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Tupl from sglang.srt.arg_groups.arg_utils import model_overridable_fields from sglang.srt.runtime_context import resolve_flag_leaf +from sglang.srt.utils.common import is_xpu logger = logging.getLogger(__name__) @@ -155,6 +156,66 @@ def _minimax_m2_overrides(server_args: Any, hf_config: Any) -> dict: return {"enable_tf32_matmul": True} +@_register_for( + "Gemma2ForCausalLM", + "Gemma3ForCausalLM", + "Gemma3ForConditionalGeneration", + "Gemma3nForCausalLM", + "Gemma3nForConditionalGeneration", +) +def _gemma2_gemma3_overrides(server_args: Any, hf_config: Any) -> dict: + # FIXME: https://github.com/sgl-project/sglang/pull/7367 is not compatible with gemma2 model. + # It failed at this test: https://github.com/sgl-project/sglang/actions/runs/16255155597/job/45890331952#step:4:736 + logger.warning( + f"Disable hybrid SWA memory for {hf_config.architectures[0]} as it is not yet supported." + ) + return {"disable_hybrid_swa_memory": True} + + +@_register_for("Exaone4ForCausalLM", "ExaoneMoEForCausalLM") +def _exaone_overrides(server_args: Any, hf_config: Any) -> dict: + if hf_config.sliding_window_pattern is not None: + logger.warning( + f"Disabling hybrid SWA memory for {hf_config.architectures[0]} as it is not yet supported." + ) + return {"disable_hybrid_swa_memory": True} + return {} + + +@_register_for("GptOssForCausalLM") +def _gpt_oss_overrides(server_args: Any, hf_config: Any) -> dict: + if is_xpu(): + # Check for bf16 dtype on Intel XPU. Reads the pristine dtype request, + # which equals the legacy mid-branch read: dtype had no earlier writer + # for this arch. + if server_args.dtype == "auto": + logger.warning( + "GptOssForCausalLM on Intel XPU currently supports bfloat16 dtype only" + ) + elif server_args.dtype not in ["bfloat16"]: + raise NotImplementedError( + f"GptOssForCausalLM on Intel XPU only supports bfloat16 dtype, " + f"but got '{server_args.dtype}'. Please use --dtype bfloat16 or remove --dtype to use auto." + ) + quantization_config = getattr(hf_config, "quantization_config", None) + if ( + quantization_config is not None + and quantization_config.get("quant_method") == "mxfp4" + ): + # use bf16 for mxfp4 triton kernels + return {"dtype": "bfloat16"} + return {} + + +@_register_for("Olmo2ForCausalLM") +def _olmo2_overrides(server_args: Any, hf_config: Any) -> dict: + # FIXME: https://github.com/sgl-project/sglang/pull/7367 is not compatible with Olmo3 model. + logger.warning( + f"Disabling hybrid SWA memory for {hf_config.architectures[0]} as it is not yet supported." + ) + return {"disable_hybrid_swa_memory": True} + + @register_model_override_predicate( lambda arch: "Step3p5ForCausalLM" in arch or "Step3p7ForConditionalGeneration" in arch diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 4cbf17ff4..f6131faef 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -4073,17 +4073,8 @@ class ServerArgs: else: self.attention_backend = "triton" - if is_xpu(): - # Check for bf16 dtype on Intel XPU - if self.dtype == "auto": - logger.warning( - "GptOssForCausalLM on Intel XPU currently supports bfloat16 dtype only" - ) - elif self.dtype not in ["bfloat16"]: - raise NotImplementedError( - f"GptOssForCausalLM on Intel XPU only supports bfloat16 dtype, " - f"but got '{self.dtype}'. Please use --dtype bfloat16 or remove --dtype to use auto." - ) + # XPU dtype validation moved to the override registry + # (arg_groups/overrides.py: _gpt_oss_overrides). supported_backends = [ "triton", @@ -4116,9 +4107,8 @@ class ServerArgs: quantization_config is not None and quantization_config.get("quant_method") == "mxfp4" ) - if is_mxfp4_quant_format: - # use bf16 for mxfp4 triton kernels - self.dtype = "bfloat16" + # The mxfp4 dtype override moved to the override registry + # (arg_groups/overrides.py: _gpt_oss_overrides). if self.moe_runner_backend == "auto": if is_sm100_supported() and is_mxfp4_quant_format: @@ -4273,19 +4263,8 @@ class ServerArgs: logger.info( "Use flashinfer_trtllm as MoE runner backend on SM100 for Llama4" ) - elif model_arch in [ - "Gemma2ForCausalLM", - "Gemma3ForCausalLM", - "Gemma3ForConditionalGeneration", - "Gemma3nForCausalLM", - "Gemma3nForConditionalGeneration", - ]: - # FIXME: https://github.com/sgl-project/sglang/pull/7367 is not compatible with gemma2 model. - # It failed at this test: https://github.com/sgl-project/sglang/actions/runs/16255155597/job/45890331952#step:4:736 - logger.warning( - f"Disable hybrid SWA memory for {model_arch} as it is not yet supported." - ) - self.disable_hybrid_swa_memory = True + # Gemma2/Gemma3 (disable_hybrid_swa_memory) moved to the override registry + # (arg_groups/overrides.py: _gemma2_gemma3_overrides). elif model_arch in ( "Gemma4ForConditionalGeneration", "Gemma4ForCausalLM", @@ -4337,22 +4316,16 @@ class ServerArgs: ) elif model_arch in ["Exaone4ForCausalLM", "ExaoneMoEForCausalLM"]: if hf_config.sliding_window_pattern is not None: - logger.warning( - f"Disabling hybrid SWA memory for {model_arch} as it is not yet supported." - ) - self.disable_hybrid_swa_memory = True + # disable_hybrid_swa_memory moved to the override registry + # (arg_groups/overrides.py: _exaone_overrides). # https://docs.sglang.ai/advanced_features/attention_backend.html accepted_backends = ["fa3", "triton", "trtllm_mha"] assert ( self.attention_backend in accepted_backends ), f"One of the attention backends in {accepted_backends} is required for {model_arch}, but got {self.attention_backend}" elif model_arch in ["Olmo2ForCausalLM"]: - # FIXME: https://github.com/sgl-project/sglang/pull/7367 is not compatible with Olmo3 model. - logger.warning( - f"Disabling hybrid SWA memory for {model_arch} as it is not yet supported." - ) - self.disable_hybrid_swa_memory = True - + # disable_hybrid_swa_memory moved to the override registry + # (arg_groups/overrides.py: _olmo2_overrides). if self.attention_backend is None: if is_cuda() and is_sm100_supported(): self.attention_backend = "trtllm_mha" diff --git a/test/registered/unit/test_model_overrides.py b/test/registered/unit/test_model_overrides.py index 9868d9405..a4ae88219 100644 --- a/test/registered/unit/test_model_overrides.py +++ b/test/registered/unit/test_model_overrides.py @@ -427,6 +427,69 @@ class TestGoldenModelOverrides(_IsolatedPublish): self.assertEqual(flags.swa_full_tokens_ratio, 1.0) self.assertTrue(flags.disable_hybrid_swa_memory) + def test_gemma2_disables_hybrid_swa_memory(self): + sa = self._construct("Gemma2ForCausalLM", "llama") + self.assertTrue(sa.disable_hybrid_swa_memory) # dual-apply == legacy + self.assertEqual( + sa._resolved_overrides, + [("_gemma2_gemma3_overrides", {"disable_hybrid_swa_memory": True})], + ) + self.assertTrue(self._publish(sa).disable_hybrid_swa_memory) + + def test_olmo2_disables_hybrid_swa_memory(self): + sa = self._construct("Olmo2ForCausalLM", "llama") + self.assertTrue(sa.disable_hybrid_swa_memory) + self.assertTrue(self._publish(sa).disable_hybrid_swa_memory) + + def test_exaone_conditional_on_sliding_window_pattern(self): + # With the pattern the branch also asserts an explicit backend. + sa = self._construct( + "Exaone4ForCausalLM", + "llama", + config_extra={"sliding_window_pattern": "LLLG"}, + attention_backend="fa3", + ) + self.assertTrue(sa.disable_hybrid_swa_memory) + self.assertTrue(self._publish(sa).disable_hybrid_swa_memory) + + def test_exaone_without_pattern_declares_nothing(self): + from sglang.srt.arg_groups.overrides import _exaone_overrides + + self.assertEqual( + _exaone_overrides(None, SimpleNamespace(sliding_window_pattern=None)), + {}, + ) + + def test_gpt_oss_mxfp4_forces_bfloat16(self): + from sglang.srt.layers.quantization import QUANTIZATION_METHODS + + if "mxfp4" not in QUANTIZATION_METHODS: + # Registration is platform-gated (CUDA / CPU engine / MXFP-HIP); + # plain CPU CI runners cannot construct an mxfp4 ModelConfig. + self.skipTest("mxfp4 quantization is not registered on this platform") + sa = self._construct( + "GptOssForCausalLM", + "llama", + config_extra={"quantization_config": {"quant_method": "mxfp4"}}, + ) + self.assertEqual(sa.dtype, "bfloat16") # dual-apply == legacy + self.assertEqual(self._publish(sa).dtype, "bfloat16") + + def test_gpt_oss_without_mxfp4_keeps_pristine_dtype(self): + sa = self._construct("GptOssForCausalLM", "llama") + self.assertEqual(sa.dtype, "auto") + self.assertEqual(self._publish(sa).dtype, "auto") + + def test_gpt_oss_xpu_dtype_validation_reads_pristine(self): + from sglang.srt.arg_groups.overrides import _gpt_oss_overrides + + with patch.object(overrides_module, "is_xpu", return_value=True): + with self.assertRaises(NotImplementedError): + _gpt_oss_overrides( + SimpleNamespace(dtype="float16"), + SimpleNamespace(architectures=["GptOssForCausalLM"]), + ) + def test_step3p_declarations_at_callable_level(self): from sglang.srt.arg_groups.overrides import _step3p_overrides