[Bugfix][Spec Decode] Fix NaN handling in rejection sampler tl.argmax (#50183)

Signed-off-by: Gabriel Peracio <gabriel@sigsonic.com>
Co-authored-by: Kimi K3 <noreply@moonshot.ai>

tests/v1/spec_decode/test_rejection_sampler_utils.py

@@ -282,6 +282,47 @@ def test_synthetic_rejection_sample(
         )
 
 
+@pytest.mark.parametrize("temperature", [0.0, 1.0])
+def test_all_nan_target_logits_in_range(temperature: float):
+    """Regression test for NaN breaking tl.argmax index bounds.
+
+    An all-NaN target logits row makes every per-block local max NaN.
+    tl.argmax over such a block returns an index into the padded region
+    (>= num_blocks), causing an OOB read of the local-argmax tensors in
+    _compute_global_target_argmax (greedy) and _insert_resampled_kernel
+    (stochastic). The vocab size below gives a non-power-of-2 block count
+    (3 blocks of 8192, padded to 4) so the padded region exists. Post-fix,
+    NaN is mapped to -inf, so the kernels must complete without error and
+    emit in-range token ids.
+    """
+    torch.manual_seed(0)
+    device = "cuda"
+    num_trials = 4
+    K = 1
+    vocab_size = 20000  # 3 vocab blocks of 8192, padded to 4
+
+    target_logits_1d = torch.full(
+        (vocab_size,), float("nan"), device=device, dtype=torch.float32
+    )
+    draft_logits_1d = torch.randn(vocab_size, device=device, dtype=torch.float32)
+
+    inputs = _build_rejection_sample_inputs(
+        target_logits_1d,
+        draft_logits_1d,
+        K,
+        temperature=temperature,
+        num_trials=num_trials,
+    )
+
+    sampled, num_sampled = rejection_sample(**inputs, num_speculative_steps=K)
+
+    assert ((num_sampled >= 1) & (num_sampled <= K + 1)).all()
+    steps = torch.arange(K + 1, device=device).unsqueeze(0)
+    valid = steps < num_sampled.unsqueeze(1)
+    assert (sampled[valid] >= 0).all()
+    assert (sampled[valid] < vocab_size).all()
+
+
 def test_placeholder_draft_token_rejected():
     """A placeholder draft id (-1) must be rejected without reading the logit
     tensors out of bounds, for any sampling method.

vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py

@@ -108,6 +108,8 @@ def _compute_global_target_argmax(
         mask=blocks_mask,
         other=float("-inf"),
     )
+    # See _insert_resampled_kernel: NaN breaks tl.argmax index bounds.
+    local_max = tl.where(local_max != local_max, float("-inf"), local_max)
     max_block_idx = tl.argmax(local_max, axis=0)
     return tl.load(
         target_local_argmax_ptr + logit_idx * target_local_argmax_stride + max_block_idx
@@ -849,6 +851,14 @@ def _insert_resampled_kernel(
         mask=mask,
         other=float("-inf"),
     )
+    # NaN max values (from NaN target logits) make tl.argmax return an
+    # out-of-range block index (into the padded region), causing an OOB read
+    # of resampled_local_argmax. Map NaN to -inf so argmax stays in range.
+    resampled_local_max = tl.where(
+        resampled_local_max != resampled_local_max,
+        float("-inf"),
+        resampled_local_max,
+    )
     resampled_max_block_idx = tl.argmax(resampled_local_max, axis=0)
     resampled = tl.load(
         resampled_local_argmax_ptr