Fix native MoE handling of noncontiguous top-k IDs (#36407)

Co-authored-by: BBuf <1182563586@qq.com>
2026-09-02 14:16:19 -07:00
parent 2d799c28f4
commit acea43079f
2 changed files with 56 additions and 1 deletions
@@ -76,7 +76,7 @@ def moe_forward_native(
    cnts = topk_ids.new_zeros((topk_ids.shape[0], len_experts))
    cnts.scatter_(1, topk_ids.to(torch.int64), 1)
    tokens_per_expert = cnts.sum(dim=0)
-    idxs = topk_ids.view(-1).argsort()
+    idxs = topk_ids.reshape(-1).argsort()

    sorted_tokens = x[idxs // topk_ids.shape[1]]
    tokens_per_expert = tokens_per_expert.cpu().numpy()