From e0e8a996304383b857adae8357149606b48d08c9 Mon Sep 17 00:00:00 2001 From: Da Chen Date: Thu, 27 Nov 2025 02:01:23 +0800 Subject: [PATCH] fix: correct usage of minimax-m2 deepep moe forward (#13892) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Dash Co-authored-by: 赵晨阳 --- python/sglang/srt/models/minimax_m2.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/python/sglang/srt/models/minimax_m2.py b/python/sglang/srt/models/minimax_m2.py index 011e45d8d..3ac16102f 100644 --- a/python/sglang/srt/models/minimax_m2.py +++ b/python/sglang/srt/models/minimax_m2.py @@ -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