@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 176 lines 7.01 KB
Raw
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)