diff --git a/python/sglang/srt/arg_groups/model_overrides/mimo_v2.py b/python/sglang/srt/arg_groups/model_overrides/mimo_v2.py index bef9a36304ae..504f0dc34b8b 100644 --- a/python/sglang/srt/arg_groups/model_overrides/mimo_v2.py +++ b/python/sglang/srt/arg_groups/model_overrides/mimo_v2.py @@ -8,6 +8,7 @@ from sglang.srt.arg_groups.model_override_base import ( _register_for, + model_config_of, resolving_view, ) from sglang.srt.runtime_context import get_platform @@ -34,4 +35,16 @@ def _mimo_v2_overrides(server_args: Any, hf_config: Any) -> dict: ): overrides["moe_runner_backend"] = "flashinfer_trtllm" logger.info("MiMoV2 FP8 on SM100: moe_runner_backend=flashinfer_trtllm.") + + if ( + get_platform().is_sm120 + and cfg.device == "cuda" + and cfg.moe_runner_backend == "auto" + and cfg.moe_a2a_backend == "none" + ): + model_config = model_config_of(server_args) + if model_config.quantization == "fp8" and model_config.is_fp4_experts: + # The default FP8 Triton runner cannot consume packed MXFP4 experts. + overrides["moe_runner_backend"] = "flashinfer_mxfp4" + logger.info("MiMoV2 MXFP4 on SM120: moe_runner_backend=flashinfer_mxfp4.") return overrides diff --git a/test/registered/unit/test_model_overrides.py b/test/registered/unit/test_model_overrides.py index 594096c16921..ffbfac2675ca 100644 --- a/test/registered/unit/test_model_overrides.py +++ b/test/registered/unit/test_model_overrides.py @@ -908,6 +908,69 @@ def _args(**kw): # pinned to flashinfer_trtllm. self.assertEqual(_mimo_v2_overrides(_args(), _hf("mxfp4")), {}) + def test_mimo_v2_sm120_mxfp4_selects_native_runner(self): + """Packed MXFP4 experts must not fall through to the FP8 Triton runner.""" + with override_platform(is_sm100=False, is_sm120=True, is_hip=False): + for architecture in ("MiMoV2ForCausalLM", "MiMoV2FlashForCausalLM"): + with self.subTest(architecture=architecture): + args = SimpleNamespace( + speculative_algorithm=None, + device="cuda", + moe_a2a_backend="none", + moe_runner_backend="auto", + _model_config=SimpleNamespace( + quantization="fp8", is_fp4_experts=True + ), + ) + self.assertEqual( + collect_model_override_declarations( + architecture, + args, + SimpleNamespace(architectures=[architecture]), + ), + [ + ( + "_mimo_v2_overrides", + {"moe_runner_backend": "flashinfer_mxfp4"}, + ) + ], + ) + + def test_mimo_v2_sm120_mxfp4_preserves_other_configurations(self): + from sglang.srt.arg_groups.model_overrides.mimo_v2 import _mimo_v2_overrides + + with override_platform(is_sm100=False, is_sm120=True, is_hip=False): + for override in ( + {"moe_runner_backend": "marlin"}, + {"moe_a2a_backend": "deepep"}, + {"device": "cpu"}, + { + "_model_config": SimpleNamespace( + quantization="fp8", is_fp4_experts=False + ) + }, + { + "_model_config": SimpleNamespace( + quantization=None, is_fp4_experts=False + ) + }, + ): + with self.subTest(override=override): + values = dict( + speculative_algorithm="EAGLE", + device="cuda", + moe_a2a_backend="none", + moe_runner_backend="auto", + _model_config=SimpleNamespace( + quantization="fp8", is_fp4_experts=True + ), + ) + values.update(override) + self.assertEqual( + _mimo_v2_overrides(SimpleNamespace(**values), _hf("fp8")), + {"enable_multi_layer_eagle": True}, + ) + def test_mimo_v2_family_is_registered(self): with override_platform(is_sm100=False): self.assertEqual(