[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>
gabriel-peracio committed
Aug 5, 2026 at 22:01 UTC
47a4e410ba3e027a69ec34ca03b3dd313b4350d2
2 files changed
+51
tests/v1/spec_decode/test_rejection_sampler_utils.py
+41
@@ -282,6 +282,47 @@ def test_synthetic_rejection_sample(
282
)
283
284
285
+@pytest.mark.parametrize("temperature", [0.0, 1.0])
286
+def test_all_nan_target_logits_in_range(temperature: float):
287
+ """Regression test for NaN breaking tl.argmax index bounds.
288
+
289
+ An all-NaN target logits row makes every per-block local max NaN.
290
+ tl.argmax over such a block returns an index into the padded region
291
+ (>= num_blocks), causing an OOB read of the local-argmax tensors in
292
+ _compute_global_target_argmax (greedy) and _insert_resampled_kernel
293
+ (stochastic). The vocab size below gives a non-power-of-2 block count
294
+ (3 blocks of 8192, padded to 4) so the padded region exists. Post-fix,
295
+ NaN is mapped to -inf, so the kernels must complete without error and
296
+ emit in-range token ids.
297
+ """
298
+ torch.manual_seed(0)
299
+ device = "cuda"
300
+ num_trials = 4
301
+ K = 1
302
+ vocab_size = 20000 # 3 vocab blocks of 8192, padded to 4
303
+
304
+ target_logits_1d = torch.full(
305
+ (vocab_size,), float("nan"), device=device, dtype=torch.float32
306
+ )
307
+ draft_logits_1d = torch.randn(vocab_size, device=device, dtype=torch.float32)
308
+
309
+ inputs = _build_rejection_sample_inputs(
310
+ target_logits_1d,
311
+ draft_logits_1d,
312
+ K,
313
+ temperature=temperature,
314
+ num_trials=num_trials,
315
+ )
316
+
317
+ sampled, num_sampled = rejection_sample(**inputs, num_speculative_steps=K)
318
+
319
+ assert ((num_sampled >= 1) & (num_sampled <= K + 1)).all()
320
+ steps = torch.arange(K + 1, device=device).unsqueeze(0)
321
+ valid = steps < num_sampled.unsqueeze(1)
322
+ assert (sampled[valid] >= 0).all()
323
+ assert (sampled[valid] < vocab_size).all()
324
+
325
+
326
def test_placeholder_draft_token_rejected():
327
"""A placeholder draft id (-1) must be rejected without reading the logit
328
tensors out of bounds, for any sampling method.
vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py
+10
@@ -108,6 +108,8 @@ def _compute_global_target_argmax(
108
mask=blocks_mask,
109
other=float("-inf"),
110
)
111
+ # See _insert_resampled_kernel: NaN breaks tl.argmax index bounds.
112
+ local_max = tl.where(local_max != local_max, float("-inf"), local_max)
113
max_block_idx = tl.argmax(local_max, axis=0)
114
return tl.load(
115
target_local_argmax_ptr + logit_idx * target_local_argmax_stride + max_block_idx
@@ -849,6 +851,14 @@ def _insert_resampled_kernel(
851
mask=mask,
852
other=float("-inf"),
853
)
854
+ # NaN max values (from NaN target logits) make tl.argmax return an
855
+ # out-of-range block index (into the padded region), causing an OOB read
856
+ # of resampled_local_argmax. Map NaN to -inf so argmax stays in range.
857
+ resampled_local_max = tl.where(
858
+ resampled_local_max != resampled_local_max,
859
+ float("-inf"),
860
+ resampled_local_max,
861
+ )
862
resampled_max_block_idx = tl.argmax(resampled_local_max, axis=0)
863
resampled = tl.load(
864
resampled_local_argmax_ptr