[EAGLE] Handle NaNs in fused top-k=1 (#32396)

Co-authored-by: weireweire <20922698+weireweire@users.noreply.github.com>
2026-07-29 00:07:42 -07:00
parent 7c248dde7f
commit bd47ec97ff
1 changed files with 2 additions and 0 deletions
@@ -27,6 +27,8 @@ def _draft_topk1_partial_argmax_kernel(
        mask=mask,
        other=-float("inf"),
    ).to(tl.float32)
+    # Keep NaNs on valid lanes from selecting the masked tail.
+    vals = tl.where(vals == vals, vals, -1e30)

    max_val = tl.max(vals, axis=0)
    local_index = tl.argmax(vals, axis=0)