fix: correct usage of minimax-m2 deepep moe forward (#13892)

Co-authored-by: Dash <dash@minimaxi.com>
Co-authored-by: 赵晨阳 <zhaochen20@outlook.com>
This commit is contained in:
Da Chen
2025-11-27 02:01:23 +08:00
committed by GitHub
parent 5102d00901
commit e0e8a99630

View File

@@ -222,7 +222,7 @@ class MiniMaxM2MoE(nn.Module):
if hidden_states.shape[0] > 0:
# router_logits: (num_tokens, n_experts)
router_logits, _ = self.gate(hidden_states.to(torch.float32))
topk_weights, topk_idx, _ = self.topk(
topk_output = self.topk(
hidden_states,
router_logits,
num_token_non_padded=forward_batch.num_token_non_padded,
@@ -231,14 +231,10 @@ class MiniMaxM2MoE(nn.Module):
),
)
else:
topk_weights, topk_idx, _ = self.topk.empty_topk_output(
hidden_states.shape[0], self.top_k
)
topk_output = self.topk.empty_topk_output(device=hidden_states.device)
final_hidden_states = self.experts(
hidden_states=hidden_states,
topk_idx=topk_idx,
topk_weights=topk_weights,
forward_batch=forward_batch,
topk_output=topk_output,
)
return final_hidden_states