From bd47ec97ff7a2881f9bb0316a4a657000a50c020 Mon Sep 17 00:00:00 2001 From: weireweire Date: Wed, 29 Jul 2026 15:07:42 +0800 Subject: [PATCH] [EAGLE] Handle NaNs in fused top-k=1 (#32396) Co-authored-by: weireweire <20922698+weireweire@users.noreply.github.com> --- python/sglang/kernels/ops/speculative/topk1.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/python/sglang/kernels/ops/speculative/topk1.py b/python/sglang/kernels/ops/speculative/topk1.py index a63a25617..93f658f2a 100644 --- a/python/sglang/kernels/ops/speculative/topk1.py +++ b/python/sglang/kernels/ops/speculative/topk1.py @@ -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)