[Bugfix] Define router_logits_dtype for remaining MoE models (#33737)

Signed-off-by: mgoin <mgoin64@gmail.com>
2026-02-04 00:24:14 -05:00
parent 2647163674
commit eb5ed20743
6 changed files with 9 additions and 4 deletions
--- a/vllm/model_executor/models/mimo_v2_flash.py
+++ b/vllm/model_executor/models/mimo_v2_flash.py
@@ -174,6 +174,7 @@ class MiMoV2MoE(nn.Module):
            num_expert_group=config.n_group,
            topk_group=config.topk_group,
            scoring_func="sigmoid",
+            router_logits_dtype=self.gate_dtype,
        )

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: