[Feature] support bf16 MoE router and mxfp4 MoE for MiMo V2 (#57784)
Signed-off-by: Zyann7 <62597503+Zyann7@users.noreply.github.com> Co-authored-by: Abatom <182586866+Abatom@users.noreply.github.com>
Zyann committed
Sep 20, 2026 at 21:25 UTC
9b2f34cad446f73b1699e8236ec0b611a65f48af
3 files changed
+21
-6
vllm/model_executor/models/mimo_v2.py
+8
-5
@@ -27,6 +27,7 @@ from vllm.model_executor.layers.fused_moe import (
27
FusedMoEFactory,
28
fused_moe_make_expert_params_mapping,
29
)
30
+from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear
31
from vllm.model_executor.layers.layernorm import RMSNorm
32
from vllm.model_executor.layers.linear import (
33
MergedColumnParallelLinear,
@@ -156,14 +157,16 @@ class MiMoV2MoE(nn.Module):
157
158
dtype = getattr(config, "moe_router_dtype", "float32")
159
self.gate_dtype = str_dtype_to_torch_dtype(dtype)
159
- self.gate = nn.Linear(
160
+ self.gate = GateLinear(
161
config.hidden_size,
162
config.n_routed_experts,
163
bias=False,
163
- dtype=self.gate_dtype,
164
+ params_dtype=self.gate_dtype,
165
+ out_dtype=torch.float32,
166
+ prefix=f"{prefix}.gate",
167
)
168
self.gate.e_score_correction_bias = nn.Parameter(
166
- torch.empty(config.n_routed_experts, dtype=self.gate_dtype)
169
+ torch.empty(config.n_routed_experts, dtype=self.gate.out_dtype)
170
)
171
172
self.experts = FusedMoEFactory(
@@ -182,7 +185,7 @@ class MiMoV2MoE(nn.Module):
185
num_expert_group=config.n_group,
186
topk_group=config.topk_group,
187
scoring_func="sigmoid",
185
- router_logits_dtype=self.gate_dtype,
188
+ router_logits_dtype=self.gate.out_dtype,
189
)
190
191
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
@@ -198,7 +201,7 @@ class MiMoV2MoE(nn.Module):
201
gate_input = hidden_states.to(self.gate_dtype)
202
else:
203
gate_input = hidden_states
201
- router_logits = self.gate(gate_input)
204
+ router_logits, _ = self.gate(gate_input)
205
final_hidden_states = self.experts(
206
hidden_states=hidden_states, router_logits=router_logits
207
)
vllm/model_executor/models/mimo_v2_omni.py
+4
-1
@@ -52,6 +52,7 @@ from vllm.transformers_utils.processors.mimo_v2_omni import (
52
53
from .interfaces import (
54
MultiModalEmbeddings,
55
+ SupportsEagle3,
56
SupportsMultiModal,
57
SupportsPP,
58
SupportsQuant,
@@ -1209,7 +1210,9 @@ class MiMoV2OmniDummyInputsBuilder(BaseDummyInputsBuilder[MiMoV2OmniProcessingIn
1210
info=MiMoV2OmniProcessingInfo,
1211
dummy_inputs=MiMoV2OmniDummyInputsBuilder,
1212
)
1212
-class MiMoV2OmniForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsQuant):
1213
+class MiMoV2OmniForCausalLM(
1214
+ nn.Module, SupportsMultiModal, SupportsPP, SupportsQuant, SupportsEagle3
1215
+):
1216
# To ensure correct weight loading and mapping.
1217
hf_to_vllm_mapper = WeightsMapper(
1218
orig_to_new_prefix={
vllm/model_executor/models/qwen3_dflash.py
+9
@@ -168,6 +168,7 @@ class DFlashQwen3Attention(nn.Module):
168
rms_norm_eps: float = 1e-06,
169
attention_bias: bool = False,
170
add_swa_attention_sink_bias: bool = False,
171
+ v_scale: float | None = None,
172
sliding_window: int | None = None,
173
causal: bool = False,
174
is_neox_style: bool = True,
@@ -193,6 +194,7 @@ class DFlashQwen3Attention(nn.Module):
194
self.q_size = self.num_heads * self.head_dim
195
self.kv_size = self.num_kv_heads * self.head_dim
196
self.scaling = self.head_dim**-0.5
197
+ self.v_scale = v_scale
198
199
self.qkv_proj = QKVParallelLinear(
200
hidden_size,
@@ -264,6 +266,8 @@ class DFlashQwen3Attention(nn.Module):
266
267
q, k = self.rotary_emb(positions, q, k)
268
269
+ if self.v_scale is not None:
270
+ v = v * self.v_scale
271
attn_output = self.attn(q, k, v)
272
output, _ = self.o_proj(attn_output)
273
return output
@@ -312,6 +316,7 @@ class DFlashQwen3DecoderLayer(nn.Module):
316
rms_norm_eps=config.rms_norm_eps,
317
attention_bias=getattr(config, "attention_bias", False),
318
add_swa_attention_sink_bias=add_swa_attention_sink_bias,
319
+ v_scale=dflash_config.get("attention_value_scale"),
320
sliding_window=sliding_window,
321
causal=causal,
322
is_neox_style=is_neox_style,
@@ -633,6 +638,10 @@ class DFlashQwen3Model(nn.Module):
638
if context_slot_mapping is None:
639
return
640
641
+ v_scale = getattr(self.layers[0].self_attn, "v_scale", None)
642
+ if v_scale is not None:
643
+ all_v.mul_(v_scale)
644
+
645
# --- Per-layer cache insert ---
646
all_k_final = all_k_flat.view(L, num_ctx, nkv, hd)
647
per_layer = isinstance(context_slot_mapping, (list, tuple))