minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: Dash <dash@minimaxi.com>
Co-authored-by: 赵晨阳 <zhaochen20@outlook.com>
This commit is contained in:
Da Chen
2025-11-26 10:01:23 -08:00
committed by GitHub
co-authored by Dash 赵晨阳
parent 5102d00901
commit e0e8a99630
1 changed files with 3 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/models/minimax_m2.py
+3 -7
View File
@@ -222,7 +222,7 @@ class MiniMaxM2MoE(nn.Module):
if hidden_states.shape[0] > 0: if hidden_states.shape[0] > 0:
# router_logits: (num_tokens, n_experts) # router_logits: (num_tokens, n_experts)
router_logits, _ = self.gate(hidden_states.to(torch.float32)) router_logits, _ = self.gate(hidden_states.to(torch.float32))
topk_weights, topk_idx, _ = self.topk( topk_output = self.topk(
hidden_states, hidden_states,
router_logits, router_logits,
num_token_non_padded=forward_batch.num_token_non_padded, num_token_non_padded=forward_batch.num_token_non_padded,
@@ -231,14 +231,10 @@ class MiniMaxM2MoE(nn.Module):
), ),
) )
else: else:
topk_weights, topk_idx, _ = self.topk.empty_topk_output( topk_output = self.topk.empty_topk_output(device=hidden_states.device)
hidden_states.shape[0], self.top_k
)
final_hidden_states = self.experts( final_hidden_states = self.experts(
hidden_states=hidden_states, hidden_states=hidden_states,
topk_idx=topk_idx, topk_output=topk_output,
topk_weights=topk_weights,
forward_batch=forward_batch,
) )
return final_hidden_states return final_hidden_states