| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103 |
- 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(
|
