fix moe gate dtype, fix tbo, fix fake dispatch (#7825)

2025-07-19 22:13:46 -07:00
parent f62d75b6a1
commit a589a07167
3 changed files with 3 additions and 3 deletions
@@ -499,7 +499,7 @@ def biased_grouped_topk_gpu(
        and is_power_of_two(correction_bias.shape[0])
    ):
        topk_weights, topk_ids = moe_fused_gate(
-            gating_output,
+            gating_output.to(dtype=torch.float32),
            correction_bias,
            num_expert_group,
            topk_group,