diff --git a/tests/v1/spec_decode/test_rejection_sampler_utils.py b/tests/v1/spec_decode/test_rejection_sampler_utils.py index 982582ffae..2663ab1496 100644 --- a/tests/v1/spec_decode/test_rejection_sampler_utils.py +++ b/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. diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py index 0c59a1c65e..0ca3932a61 100644 --- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py +++ b/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