A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 8, 22:45 UTC
[Bugfix] Define router_logits_dtype for remaining MoE models (#33737)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Feb 4, 2026 at 00:24 UTC
eb5ed207437d31daf09cd4ae51e993fc1efa6928
6 files changed
+9
-4
vllm/model_executor/models/afmoe.py
+1
@@ -142,6 +142,7 @@ class AfmoeMoE(nn.Module):
142
e_score_correction_bias=self.expert_bias,
143
enable_eplb=self.enable_eplb,
144
num_redundant_experts=self.n_redundant_experts,
145
+ router_logits_dtype=torch.float32,
146
)
147
148
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
vllm/model_executor/models/bailing_moe.py
+1
@@ -300,6 +300,7 @@ class BailingMoE(nn.Module):
300
num_expert_group=self.n_group,
301
topk_group=self.topk_group,
302
use_grouped_topk=self.use_grouped_topk,
303
+ router_logits_dtype=self.router_dtype,
304
)
305
306
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
vllm/model_executor/models/flex_olmo.py
+1
-1
@@ -71,7 +71,6 @@ class FlexOlmoMoE(nn.Module):
71
prefix=f"{prefix}.gate",
72
)
73
74
- # Gate always runs at half / full precision for now.
74
self.experts = FusedMoE(
75
num_experts=hf_config.num_experts,
76
top_k=hf_config.num_experts_per_tok,
@@ -82,6 +81,7 @@ class FlexOlmoMoE(nn.Module):
81
quant_config=None,
82
tp_size=tp_size,
83
prefix=f"{prefix}.experts",
84
+ router_logits_dtype=torch.float32,
85
)
86
87
self.top_k = hf_config.num_experts_per_tok
vllm/model_executor/models/longcat_flash.py
+4
-3
@@ -236,9 +236,9 @@ class FlashMLP(nn.Module):
236
class LongcatRouter(nn.Module):
237
def __init__(
238
self,
239
- config,
240
- zero_expert_num=0,
241
- rounter_params_dtype=torch.bfloat16,
239
+ config: FlashConfig,
240
+ zero_expert_num: int,
241
+ rounter_params_dtype: torch.dtype,
242
prefix: str = "",
243
):
244
super().__init__()
@@ -309,6 +309,7 @@ class LongcatMoe(nn.Module):
309
prefix=f"{prefix}.experts",
310
enable_eplb=enable_eplb,
311
routed_scaling_factor=config.routed_scaling_factor,
312
+ router_logits_dtype=self.rounter_params_dtype,
313
)
314
315
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
vllm/model_executor/models/mimo_v2_flash.py
+1
@@ -174,6 +174,7 @@ class MiMoV2MoE(nn.Module):
174
num_expert_group=config.n_group,
175
topk_group=config.topk_group,
176
scoring_func="sigmoid",
177
+ router_logits_dtype=self.gate_dtype,
178
)
179
180
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
vllm/model_executor/models/step3p5.py
+1
@@ -388,6 +388,7 @@ class FusedMoEBlock(nn.Module):
388
routed_scaling_factor=config.moe_router_scaling_factor,
389
enable_eplb=self.enable_eplb,
390
num_redundant_experts=self.n_redundant_experts,
391
+ router_logits_dtype=torch.float32,
392
)
393
394
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: