diff --git a/tests/models/quantization/test_mimo_v2_w4a16_routing.py b/tests/models/quantization/test_mimo_v2_w4a16_routing.py new file mode 100644 index 0000000000..9dc22a77ca --- /dev/null +++ b/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) diff --git a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a16_moe.py b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a16_moe.py index 13c34698c8..4776a908fb 100644 --- a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a16_moe.py +++ b/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,