[ROCm][MoE] Support MiMo-V2.6 MXFP4 on gfx942 (#58262)
Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>
yhcheong (vllmellm) committed
Sep 30, 2026 at 17:23 UTC
16d4ac4ba3a0a7d0b6a5eb11e203cd171a2d64a1
2 files changed
+192
-5
tests/models/quantization/test_mimo_v2_w4a16_routing.py
new
+176
@@ -0,0 +1,176 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Routing support for the AITER Triton MXFP4 W4A16 MoE backend.
4
+
5
+MiMo-V2.6 uses an ungrouped sigmoid router with a per-expert correction bias
6
+(``scoring_func=sigmoid``, ``topk_method=noaux_tc``, ``n_group == 1``), which
7
+``get_routing_method_type`` classifies as ``RoutingMethodType.DeepSeekV3``.
8
+``AiterW4A16ExpertsMonolithic`` used to reject that router ("kernel does not
9
+support routing method ..."), and on gfx942 it is the only MXFP4 kernel whose
10
+device gate accepts the card, so the checkpoint could not be served natively at
11
+all.
12
+
13
+aiter's *flat* top-k handles that router natively (``score_mode="sigmoid"`` plus
14
+a correction bias), so the backend only has to ask for it. This module covers
15
+both halves: the gate accepts the router family, and the routing the wrapper
16
+requests reproduces the model's own math.
17
+"""
18
+
19
+import pytest
20
+import torch
21
+
22
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
23
+from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
24
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
25
+from vllm.model_executor.layers.fused_moe.config import (
26
+ FusedMoEConfig,
27
+ FusedMoEParallelConfig,
28
+ RoutingMethodType,
29
+ get_routing_method_type,
30
+)
31
+from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp4_w4a16_moe import (
32
+ AiterW4A16ExpertsMonolithic,
33
+)
34
+from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp4Static
35
+
36
+MIMO_EXPERTS = 256 # MiMo-V2.6-Flash
37
+MIMO_PRO_EXPERTS = 384 # MiMo-V2.6-Pro
38
+MIMO_TOP_K = 8
39
+
40
+
41
+def mimo_routing_method() -> RoutingMethodType:
42
+ """MiMo-V2.6: sigmoid + correction bias, a single expert group."""
43
+ return get_routing_method_type(
44
+ scoring_func="sigmoid",
45
+ top_k=MIMO_TOP_K,
46
+ renormalize=True,
47
+ num_expert_group=1,
48
+ has_e_score_bias=True,
49
+ )
50
+
51
+
52
+def mimo_moe_config(num_experts: int = MIMO_EXPERTS) -> FusedMoEConfig:
53
+ return FusedMoEConfig(
54
+ num_experts=num_experts,
55
+ experts_per_token=MIMO_TOP_K,
56
+ hidden_dim=4096,
57
+ intermediate_size=2048,
58
+ num_local_experts=num_experts,
59
+ num_logical_experts=num_experts,
60
+ moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
61
+ activation=MoEActivation.SILU,
62
+ in_dtype=torch.bfloat16,
63
+ device=torch.device("cuda"),
64
+ routing_method=mimo_routing_method(),
65
+ router_logits_dtype=torch.float32,
66
+ )
67
+
68
+
69
+def test_ungrouped_sigmoid_router_is_deepseekv3() -> None:
70
+ """Guard the premise of the fix: MiMo's router is DeepSeekV3-classified."""
71
+ assert mimo_routing_method() == RoutingMethodType.DeepSeekV3
72
+
73
+
74
+def test_aiter_w4a16_supports_mimo_routing() -> None:
75
+ """The gate accepts MiMo's router family — the reported construction abort.
76
+
77
+ Only the predicate the fix touches is exercised, so this runs on any host;
78
+ the full config path is covered by the test below.
79
+ """
80
+ assert AiterW4A16ExpertsMonolithic._supports_routing_method(
81
+ RoutingMethodType.DeepSeekV3, kMxfp4Static, None
82
+ )
83
+
84
+
85
+@pytest.mark.skipif(
86
+ not is_aiter_found_and_supported(),
87
+ reason="only runs on ROCm with a supported AITER install",
88
+)
89
+def test_aiter_w4a16_supports_mimo_config(monkeypatch: pytest.MonkeyPatch) -> None:
90
+ """The same gate through the real config path, device check included."""
91
+ monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
92
+ rocm_aiter_ops.refresh_env_variables()
93
+
94
+ supported, reason = AiterW4A16ExpertsMonolithic.is_supported_config(
95
+ AiterW4A16ExpertsMonolithic,
96
+ mimo_moe_config(),
97
+ kMxfp4Static,
98
+ None,
99
+ mk.FusedMoEActivationFormat.Standard,
100
+ )
101
+
102
+ assert supported, reason
103
+
104
+
105
+@pytest.mark.parametrize("num_experts", [MIMO_EXPERTS, MIMO_PRO_EXPERTS])
106
+@pytest.mark.skipif(
107
+ not is_aiter_found_and_supported(),
108
+ reason="only runs on ROCm with a supported AITER install",
109
+)
110
+def test_flat_sigmoid_routing_matches_mimo_reference(num_experts: int) -> None:
111
+ """The routing the wrapper requests must equal MiMo's own router math.
112
+
113
+ MiMo's router is ungrouped: sigmoid scores, selection on
114
+ ``sigmoid(logits) + bias``, weights = the *unbiased* sigmoid scores, then
115
+ renormalized. Both checkpoints use it — 256 experts is Flash, 384 is Pro.
116
+
117
+ The ``sigmoid`` score mode comes from ROCm/aiter#4688 (aiter >= 0.1.20).
118
+ """
119
+ aiter_routing = pytest.importorskip(
120
+ "aiter.ops.triton.moe.moe_routing.routing"
121
+ ).routing
122
+
123
+ torch.manual_seed(0)
124
+ top_k, num_tokens = MIMO_TOP_K, 16
125
+ logits = torch.randn(num_tokens, num_experts, device="cuda") * 2.0
126
+ bias = torch.randn(num_experts, device="cuda") * 0.5
127
+
128
+ # Reference: select on the biased score, return the unbiased score, renorm.
129
+ scores = torch.sigmoid(logits.float())
130
+ ref_ids = (scores + bias).topk(top_k, dim=-1).indices
131
+ ref_weights = scores.gather(1, ref_ids)
132
+ ref_weights = ref_weights / ref_weights.sum(dim=-1, keepdim=True)
133
+
134
+ def routed(score_mode: str) -> tuple[torch.Tensor, torch.Tensor]:
135
+ routing_data, topk_indx, _ = aiter_routing(
136
+ logits,
137
+ top_k,
138
+ score_mode=score_mode,
139
+ bias=bias,
140
+ renorm=True,
141
+ routed_scaling_factor=1.0,
142
+ use_grouped_topk=False,
143
+ )
144
+ # `topk_indx[i]` is the token-slot index (token * top_k + slot) of the
145
+ # i-th entry of the expert-sorted array whose weights are gate_scal.
146
+ hist = routing_data.expt_hist.to(torch.long).cpu()
147
+ expert_of_slot = torch.repeat_interleave(torch.arange(hist.numel()), hist)
148
+ slot = topk_indx.reshape(-1).to(torch.long).cpu()
149
+ weights = routing_data.gate_scal.reshape(-1).to(torch.float32).cpu()
150
+ ids = torch.full((num_tokens, top_k), -1, dtype=torch.long)
151
+ out = torch.zeros(num_tokens, top_k)
152
+ for i in range(slot.numel()):
153
+ token, position = int(slot[i]) // top_k, int(slot[i]) % top_k
154
+ ids[token, position] = int(expert_of_slot[i])
155
+ out[token, position] = float(weights[i])
156
+ return ids, out
157
+
158
+ got_ids, got_weights = routed("sigmoid")
159
+
160
+ # aiter's triton top-k emits each token's selection sorted by *ascending
161
+ # expert id*, while the reference is in descending ``sigmoid(logits)+bias``
162
+ # rank order — align both sides on the expert id so the weight check
163
+ # becomes elementwise and still validates the expert<->weight pairing.
164
+ def align(ids: torch.Tensor, w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
165
+ order = torch.argsort(ids, dim=-1)
166
+ return ids.gather(1, order), w.gather(1, order)
167
+
168
+ got_ids, got_weights = align(got_ids, got_weights)
169
+ ref_ids, ref_weights = align(ref_ids.cpu(), ref_weights.cpu())
170
+ assert torch.equal(got_ids, ref_ids)
171
+ assert torch.allclose(got_weights, ref_weights, atol=1e-3)
172
+
173
+ # Control: the correction bias must steer selection only. Ignoring it would
174
+ # pick from the plain sigmoid scores instead.
175
+ plain_ids = scores.topk(top_k, dim=-1).indices
176
+ 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
+16
-5
@@ -295,6 +295,7 @@ class AiterW4A16ExpertsMonolithic(mk.FusedMoEExpertsMonolithic):
295
RoutingMethodType.Renormalize,
296
RoutingMethodType.RenormalizeNaive,
297
RoutingMethodType.DeepseekV4,
298
+ RoutingMethodType.DeepSeekV3,
299
)
300
301
@staticmethod
@@ -344,6 +345,7 @@ class AiterW4A16ExpertsMonolithic(mk.FusedMoEExpertsMonolithic):
345
RoutingMethodType.Renormalize,
346
RoutingMethodType.RenormalizeNaive,
347
RoutingMethodType.DeepseekV4,
348
+ RoutingMethodType.DeepSeekV3,
349
]
350
351
@staticmethod
@@ -373,11 +375,20 @@ class AiterW4A16ExpertsMonolithic(mk.FusedMoEExpertsMonolithic):
375
routed_scaling_factor: float | None = None,
376
topk_group: int | None = None,
377
) -> torch.Tensor:
376
- score_mode = (
377
- "sqrtsoftplus"
378
- if self.moe_config.routing_method == RoutingMethodType.DeepseekV4
379
- else None
380
- )
378
+ routing_method = self.moe_config.routing_method
379
+ if routing_method == RoutingMethodType.DeepseekV4:
380
+ score_mode = "sqrtsoftplus"
381
+ elif routing_method == RoutingMethodType.DeepSeekV3:
382
+ # MiMo-V2.6: an ungrouped sigmoid router with a per-expert
383
+ # correction bias. aiter's flat top-k implements this natively
384
+ # (select on sigmoid(logits) + bias, return the unbiased sigmoid
385
+ # scores).
386
+ assert e_score_correction_bias is not None, (
387
+ "DeepSeekV3 routing requires e_score_correction_bias"
388
+ )
389
+ score_mode = "sigmoid"
390
+ else:
391
+ score_mode = None
392
return aiter_triton_kernel_w4a16_moe_forward(
393
hidden_states=hidden_states,
394
w1=w1,