[refactor] Sweep disable_hybrid_swa_memory writers; close the dtype family (stack 9/15) (#30071)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

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

python/sglang/srt/server_args.py

@@ -4073,17 +4073,8 @@ class ServerArgs:
                 else:
                     self.attention_backend = "triton"
 
-            if is_xpu():
+            # XPU dtype validation moved to the override registry
-                # Check for bf16 dtype on Intel XPU
+            # (arg_groups/overrides.py: _gpt_oss_overrides).
-                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."
-                    )
 
             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:
+            # The mxfp4 dtype override moved to the override registry
-                # use bf16 for mxfp4 triton kernels
+            # (arg_groups/overrides.py: _gpt_oss_overrides).
-                self.dtype = "bfloat16"
 
             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 [
+        # Gemma2/Gemma3 (disable_hybrid_swa_memory) moved to the override registry
-            "Gemma2ForCausalLM",
+        # (arg_groups/overrides.py: _gemma2_gemma3_overrides).
-            "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
         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(
+                # disable_hybrid_swa_memory moved to the override registry
-                    f"Disabling hybrid SWA memory for {model_arch} as it is not yet supported."
+                # (arg_groups/overrides.py: _exaone_overrides).
-                )
-                self.disable_hybrid_swa_memory = True
                 # 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.
+            # disable_hybrid_swa_memory moved to the override registry
-            logger.warning(
+            # (arg_groups/overrides.py: _olmo2_overrides).
-                f"Disabling hybrid SWA memory for {model_arch} as it is not yet supported."
-            )
-            self.disable_hybrid_swa_memory = True
-
             if self.attention_backend is None:
                 if is_cuda() and is_sm100_supported():
                     self.attention_backend = "trtllm_mha"

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