fix(moe): cast filtered-activation expert_ids to int32 for torch.compile (#38085)

Co-authored-by: BBuf <bbuf@users.noreply.github.com>
2026-09-05 17:13:07 +08:00
parent a74470e904
commit d49180019b
2 changed files with 35 additions and 0 deletions
@@ -134,6 +134,11 @@ def run_activation(
    if expert_ids is None:
        _run_activation_inplace(op_name, input, out)
    else:
+        # The JIT kernel indexes expert ids as int32. Routing ids may arrive as
+        # int64 (e.g. from torch.topk) and torch.compile realizes them at their
+        # true dtype, so normalize here instead of asserting downstream.
+        if expert_ids.dtype != torch.int32:
+            expert_ids = expert_ids.to(torch.int32)
        _run_activation_filtered_inplace(op_name, input, out, expert_ids, expert_step)
    return out