[MoE] Fix moe_fused_gate out-of-range expert id on all-NaN rows (fixes eagle_dp_attention crash) (#30079)

2026-07-03 18:48:51 -07:00
parent b28bc1060f
commit c21f6f19cf
1 changed files with 3 additions and 0 deletions
@@ -167,6 +167,9 @@ def _router_triton_kernel(

    biased = tl.where(mask_n[None, :], biased, -float("inf"))  # [BLOCK_M, BLOCK_N]

+    # Map NaN -> a finite floor
+    biased = tl.where(biased == biased, biased, -1e30)  # [BLOCK_M, BLOCK_N]
+
    # Grouped routing (DeepSeek-V3 noaux_tc): per-group score = sum of the top-2
    # biased values; keep TOPK_GROUP groups (lowest group id wins ties); mask the
    # experts of dropped groups to -inf before the top-k below. Weight is still the