[ROCm][MoE] Support MiMo-V2.6 MXFP4 on gfx942 (#58262)

Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>

tests/models/quantization/test_mimo_v2_w4a16_routing.py

@@ -0,0 +1,176 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Routing support for the AITER Triton MXFP4 W4A16 MoE backend.
+
+MiMo-V2.6 uses an ungrouped sigmoid router with a per-expert correction bias
+(``scoring_func=sigmoid``, ``topk_method=noaux_tc``, ``n_group == 1``), which
+``get_routing_method_type`` classifies as ``RoutingMethodType.DeepSeekV3``.
+``AiterW4A16ExpertsMonolithic`` used to reject that router ("kernel does not
+support routing method ..."), and on gfx942 it is the only MXFP4 kernel whose
+device gate accepts the card, so the checkpoint could not be served natively at
+all.
+
+aiter's *flat* top-k handles that router natively (``score_mode="sigmoid"`` plus
+a correction bias), so the backend only has to ask for it. This module covers
+both halves: the gate accepts the router family, and the routing the wrapper
+requests reproduces the model's own math.
+"""
+
+import pytest
+import torch
+
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
+from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+from vllm.model_executor.layers.fused_moe.config import (
+    FusedMoEConfig,
+    FusedMoEParallelConfig,
+    RoutingMethodType,
+    get_routing_method_type,
+)
+from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp4_w4a16_moe import (
+    AiterW4A16ExpertsMonolithic,
+)
+from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp4Static
+
+MIMO_EXPERTS = 256  # MiMo-V2.6-Flash
+MIMO_PRO_EXPERTS = 384  # MiMo-V2.6-Pro
+MIMO_TOP_K = 8
+
+
+def mimo_routing_method() -> RoutingMethodType:
+    """MiMo-V2.6: sigmoid + correction bias, a single expert group."""
+    return get_routing_method_type(
+        scoring_func="sigmoid",
+        top_k=MIMO_TOP_K,
+        renormalize=True,
+        num_expert_group=1,
+        has_e_score_bias=True,
+    )
+
+
+def mimo_moe_config(num_experts: int = MIMO_EXPERTS) -> FusedMoEConfig:
+    return FusedMoEConfig(
+        num_experts=num_experts,
+        experts_per_token=MIMO_TOP_K,
+        hidden_dim=4096,
+        intermediate_size=2048,
+        num_local_experts=num_experts,
+        num_logical_experts=num_experts,
+        moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
+        activation=MoEActivation.SILU,
+        in_dtype=torch.bfloat16,
+        device=torch.device("cuda"),
+        routing_method=mimo_routing_method(),
+        router_logits_dtype=torch.float32,
+    )
+
+
+def test_ungrouped_sigmoid_router_is_deepseekv3() -> None:
+    """Guard the premise of the fix: MiMo's router is DeepSeekV3-classified."""
+    assert mimo_routing_method() == RoutingMethodType.DeepSeekV3
+
+
+def test_aiter_w4a16_supports_mimo_routing() -> None:
+    """The gate accepts MiMo's router family — the reported construction abort.
+
+    Only the predicate the fix touches is exercised, so this runs on any host;
+    the full config path is covered by the test below.
+    """
+    assert AiterW4A16ExpertsMonolithic._supports_routing_method(
+        RoutingMethodType.DeepSeekV3, kMxfp4Static, None
+    )
+
+
+@pytest.mark.skipif(
+    not is_aiter_found_and_supported(),
+    reason="only runs on ROCm with a supported AITER install",
+)
+def test_aiter_w4a16_supports_mimo_config(monkeypatch: pytest.MonkeyPatch) -> None:
+    """The same gate through the real config path, device check included."""
+    monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
+    rocm_aiter_ops.refresh_env_variables()
+
+    supported, reason = AiterW4A16ExpertsMonolithic.is_supported_config(
+        AiterW4A16ExpertsMonolithic,
+        mimo_moe_config(),
+        kMxfp4Static,
+        None,
+        mk.FusedMoEActivationFormat.Standard,
+    )
+
+    assert supported, reason
+
+
+@pytest.mark.parametrize("num_experts", [MIMO_EXPERTS, MIMO_PRO_EXPERTS])
+@pytest.mark.skipif(
+    not is_aiter_found_and_supported(),
+    reason="only runs on ROCm with a supported AITER install",
+)
+def test_flat_sigmoid_routing_matches_mimo_reference(num_experts: int) -> None:
+    """The routing the wrapper requests must equal MiMo's own router math.
+
+    MiMo's router is ungrouped: sigmoid scores, selection on
+    ``sigmoid(logits) + bias``, weights = the *unbiased* sigmoid scores, then
+    renormalized. Both checkpoints use it — 256 experts is Flash, 384 is Pro.
+
+    The ``sigmoid`` score mode comes from ROCm/aiter#4688 (aiter >= 0.1.20).
+    """
+    aiter_routing = pytest.importorskip(
+        "aiter.ops.triton.moe.moe_routing.routing"
+    ).routing
+
+    torch.manual_seed(0)
+    top_k, num_tokens = MIMO_TOP_K, 16
+    logits = torch.randn(num_tokens, num_experts, device="cuda") * 2.0
+    bias = torch.randn(num_experts, device="cuda") * 0.5
+
+    # Reference: select on the biased score, return the unbiased score, renorm.
+    scores = torch.sigmoid(logits.float())
+    ref_ids = (scores + bias).topk(top_k, dim=-1).indices
+    ref_weights = scores.gather(1, ref_ids)
+    ref_weights = ref_weights / ref_weights.sum(dim=-1, keepdim=True)
+
+    def routed(score_mode: str) -> tuple[torch.Tensor, torch.Tensor]:
+        routing_data, topk_indx, _ = aiter_routing(
+            logits,
+            top_k,
+            score_mode=score_mode,
+            bias=bias,
+            renorm=True,
+            routed_scaling_factor=1.0,
+            use_grouped_topk=False,
+        )
+        # `topk_indx[i]` is the token-slot index (token * top_k + slot) of the
+        # i-th entry of the expert-sorted array whose weights are gate_scal.
+        hist = routing_data.expt_hist.to(torch.long).cpu()
+        expert_of_slot = torch.repeat_interleave(torch.arange(hist.numel()), hist)
+        slot = topk_indx.reshape(-1).to(torch.long).cpu()
+        weights = routing_data.gate_scal.reshape(-1).to(torch.float32).cpu()
+        ids = torch.full((num_tokens, top_k), -1, dtype=torch.long)
+        out = torch.zeros(num_tokens, top_k)
+        for i in range(slot.numel()):
+            token, position = int(slot[i]) // top_k, int(slot[i]) % top_k
+            ids[token, position] = int(expert_of_slot[i])
+            out[token, position] = float(weights[i])
+        return ids, out
+
+    got_ids, got_weights = routed("sigmoid")
+
+    # aiter's triton top-k emits each token's selection sorted by *ascending
+    # expert id*, while the reference is in descending ``sigmoid(logits)+bias``
+    # rank order — align both sides on the expert id so the weight check
+    # becomes elementwise and still validates the expert<->weight pairing.
+    def align(ids: torch.Tensor, w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        order = torch.argsort(ids, dim=-1)
+        return ids.gather(1, order), w.gather(1, order)
+
+    got_ids, got_weights = align(got_ids, got_weights)
+    ref_ids, ref_weights = align(ref_ids.cpu(), ref_weights.cpu())
+    assert torch.equal(got_ids, ref_ids)
+    assert torch.allclose(got_weights, ref_weights, atol=1e-3)
+
+    # Control: the correction bias must steer selection only. Ignoring it would
+    # pick from the plain sigmoid scores instead.
+    plain_ids = scores.topk(top_k, dim=-1).indices
+    assert not torch.equal(torch.sort(plain_ids.cpu(), dim=-1).values, ref_ids)

vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a16_moe.py

@@ -295,6 +295,7 @@ class AiterW4A16ExpertsMonolithic(mk.FusedMoEExpertsMonolithic):
             RoutingMethodType.Renormalize,
             RoutingMethodType.RenormalizeNaive,
             RoutingMethodType.DeepseekV4,
+            RoutingMethodType.DeepSeekV3,
         )
 
     @staticmethod
@@ -344,6 +345,7 @@ class AiterW4A16ExpertsMonolithic(mk.FusedMoEExpertsMonolithic):
             RoutingMethodType.Renormalize,
             RoutingMethodType.RenormalizeNaive,
             RoutingMethodType.DeepseekV4,
+            RoutingMethodType.DeepSeekV3,
         ]
 
     @staticmethod
@@ -373,11 +375,20 @@ class AiterW4A16ExpertsMonolithic(mk.FusedMoEExpertsMonolithic):
         routed_scaling_factor: float | None = None,
         topk_group: int | None = None,
     ) -> torch.Tensor:
-        score_mode = (
-            "sqrtsoftplus"
-            if self.moe_config.routing_method == RoutingMethodType.DeepseekV4
-            else None
-        )
+        routing_method = self.moe_config.routing_method
+        if routing_method == RoutingMethodType.DeepseekV4:
+            score_mode = "sqrtsoftplus"
+        elif routing_method == RoutingMethodType.DeepSeekV3:
+            # MiMo-V2.6: an ungrouped sigmoid router with a per-expert
+            # correction bias. aiter's flat top-k implements this natively
+            # (select on sigmoid(logits) + bias, return the unbiased sigmoid
+            # scores).
+            assert e_score_correction_bias is not None, (
+                "DeepSeekV3 routing requires e_score_correction_bias"
+            )
+            score_mode = "sigmoid"
+        else:
+            score_mode = None
         return aiter_triton_kernel_w4a16_moe_forward(
             hidden_states=hidden_states,
             w1=w1,