@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 3c43237233

[Bugfix][Model Runner V2][Spec Decode] Fix int32 offset overflow in sampler kernels (#46560)

Signed-off-by: xiaojun.wei <jessiewei747@gmail.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>

JessieWei committed Jun 25, 2026 at 02:00 UTC 3c43237233a8f753a0ed24d95c0978b09f009d29

10 files changed +90 -16

tests/v1/sample/test_logprobs.py

+34

@@ -1263,3 +1263,37 @@ def test_prompt_logprobs_with_chunking_and_preemption():

1263 assert preemptions > 0, "Test did not trigger any preemptions"

1264

1265 print(f"Test passed with {preemptions} preemptions")

1266 +

1267 +

1268 +@large_gpu_mark(min_gb=24)

1269 +def test_token_logprobs_large_batch_int64_row_offset():

1270 + """Regression: logprob kernel row offset (row * vocab_size) must use int64.

1271 +

1272 + The rejection-sampler logprobs path runs the logprob kernels over the

1273 + spec-expanded logits batch, so batch_size * vocab_size can exceed 2**31

1274 + (e.g. DFlash drafts K tokens per request). With int32 offset arithmetic the

1275 + per-row pointer wraps to a negative address and the kernel hits a CUDA

1276 + illegal memory access. Run over a batch where batch_size * vocab_size > 2**31

1277 + and check the highest-offset row matches a reference log-softmax.

1278 + """

1279 + if not current_platform.is_cuda():

1280 + pytest.skip("int32 row-offset overflow is a CUDA kernel issue")

1281 + from vllm.v1.worker.gpu.sample.logprob import compute_token_logprobs

1282 +

1283 + device = torch.device("cuda")

1284 + vocab_size = 131072

1285 + batch_size = 2**31 // vocab_size + 64 # batch_size * vocab_size > 2**31

1286 + # logits (the large input) plus small logprob/rank outputs; ~1 GB headroom.

1287 + required_bytes = batch_size * vocab_size * 4 + (1 << 30)

1288 + if torch.cuda.mem_get_info()[0] < required_bytes:

1289 + pytest.skip(f"needs ~{required_bytes / 1e9:.0f} GB of free GPU memory")

1290 +

1291 + logits = torch.randn(batch_size, vocab_size, device=device, dtype=torch.float32)

1292 + token_ids = torch.full((batch_size, 1), 7, device=device, dtype=torch.int64)

1293 + logprobs = compute_token_logprobs(logits, token_ids)

1294 + torch.accelerator.synchronize() # surface any async illegal memory access

1295 + last = batch_size - 1

1296 + ref = torch.log_softmax(logits[last].float(), dim=-1)[7]

1297 + assert torch.allclose(logprobs[last, 0], ref, atol=1e-2), (

1298 + f"logprob {logprobs[last, 0].item()} != ref {ref.item()}"

1299 + )

tests/v1/sample/test_topk_topp_sampler.py

+40

@@ -4,6 +4,7 @@ import pytest

4 import torch

5 from torch import Generator

7 +from tests.utils import large_gpu_mark

8 from vllm.platforms import current_platform

9 from vllm.triton_utils import HAS_TRITON

10 from vllm.utils.torch_utils import set_random_seed

@@ -404,6 +405,45 @@ class TestTritonTopkTopp:

405

406 self._compare_results(logits, k, p)

407

408 + @large_gpu_mark(min_gb=24)

409 + def test_large_batch_int64_row_offset(self):

410 + """Regression: per-row offset (row * vocab_size) must not overflow int32.

411 +

412 + Speculative decoding expands the logits batch (e.g. DFlash drafts K

413 + tokens per request), so batch_size * vocab_size can exceed 2**31. With

414 + int32 offset arithmetic the per-row pointer wraps to a negative address

415 + and the kernel hits a CUDA illegal memory access. Use a batch where

416 + batch_size * vocab_size > 2**31 and give the highest-offset row the same

417 + logits as row 0: an overflow there would read a different row and change

418 + the kept set.

419 + """

420 + from vllm.v1.sample.ops.topk_topp_triton import apply_top_k_top_p_triton

421 +

422 + if not current_platform.is_cuda():

423 + pytest.skip("int32 row-offset overflow is a CUDA kernel issue")

424 + vocab_size = 131072

425 + batch_size = 2**31 // vocab_size + 64 # batch_size * vocab_size > 2**31

426 + # logits is modified in place; the only extra device memory is the

427 + # per-SM scratch buffer (~num_sm * vocab), so allow ~1 GB of headroom.

428 + required_bytes = batch_size * vocab_size * 4 + (1 << 30)

429 + if torch.cuda.mem_get_info()[0] < required_bytes:

430 + pytest.skip(f"needs ~{required_bytes / 1e9:.0f} GB of free GPU memory")

431 +

432 + logits = torch.randn(

433 + batch_size, vocab_size, generator=self.generator, dtype=torch.float32

434 + )

435 + logits[batch_size - 1] = logits[0]

436 + k = torch.full((batch_size,), 5, dtype=torch.int32)

437 + result = apply_top_k_top_p_triton(logits, k, None)

438 + torch.accelerator.synchronize() # surface any async illegal memory access

439 + kept_first = (result[0] > float("-inf")).nonzero(as_tuple=True)[0]

440 + kept_last = (result[batch_size - 1] > float("-inf")).nonzero(as_tuple=True)[0]

441 + assert kept_first.numel() == 5, f"row 0 kept {kept_first.numel()}, expected 5"

442 + assert torch.equal(kept_first, kept_last), (

443 + "highest-offset row produced a different top-k mask than the "

444 + "identical row 0 (int32 row-offset overflow)"

445 + )

446 +

447 @pytest.mark.parametrize(

448 "mode",

449 ["topk_only", "topp_only", "topk_and_topp"],

vllm/v1/sample/ops/topk_topp_triton.py

+1 -1

@@ -111,7 +111,7 @@ def _topk_topp_kernel(

111 pid = tl.program_id(0)

112 num_programs = tl.num_programs(0)

113 for row_id in tl.range(pid, BATCH_SIZE, num_programs):

114 - LOGITS_ROW = LOGITS + row_id * LOGITS_STRIDE_0

114 + LOGITS_ROW = LOGITS + row_id.to(tl.int64) * LOGITS_STRIDE_0

115 BUFFER_ROW = BUFFER + pid * VOCAB_SIZE

116

117 final_pivot = -float("inf")

vllm/v1/worker/gpu/sample/bad_words.py

+1 -1

@@ -114,7 +114,7 @@ def _bad_words_kernel(

114 input_ids_ptr,

115 expanded_local_pos_ptr,

116 ):

117 - token_idx = tl.program_id(0)

117 + token_idx = tl.program_id(0).to(tl.int64)

118 bw_idx = tl.program_id(1)

119

120 req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx)

vllm/v1/worker/gpu/sample/gumbel.py

+3 -3

@@ -23,7 +23,7 @@ def _temperature_kernel(

23 vocab_size,

24 BLOCK_SIZE: tl.constexpr,

25 ):

26 - token_idx = tl.program_id(0)

26 + token_idx = tl.program_id(0).to(tl.int64)

27 req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx)

28 temperature = tl.load(temperature_ptr + req_state_idx).to(tl.float32)

29 if temperature == 0.0 or temperature == 1.0:

@@ -91,7 +91,7 @@ def gumbel_block_argmax(

91 USE_FP64: tl.constexpr,

92 PER_TOKEN_COL: tl.constexpr = False,

93 ):

94 - req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx)

94 + req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx).to(tl.int64)

95 temp = tl.load(temp_ptr + req_state_idx).to(tl.float32)

96 if temp != 0.0 and APPLY_TEMPERATURE:

97 # Apply temperature.

@@ -169,7 +169,7 @@ def _gumbel_sample_kernel(

169 USE_FP64: tl.constexpr,

170 PER_TOKEN_COL: tl.constexpr,

171 ):

172 - token_idx = tl.program_id(0)

172 + token_idx = tl.program_id(0).to(tl.int64)

173 block_idx = tl.program_id(1)

174 block = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)

175 mask = block < vocab_size

vllm/v1/worker/gpu/sample/logit_bias.py

+1 -1

@@ -169,7 +169,7 @@ def _bias_kernel(

169 BLOCK_SIZE: tl.constexpr,

170 LOGITS_BLOCK_SIZE: tl.constexpr,

171 ):

172 - token_idx = tl.program_id(0)

172 + token_idx = tl.program_id(0).to(tl.int64)

173 req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx)

174

175 block = tl.arange(0, BLOCK_SIZE)

vllm/v1/worker/gpu/sample/logprob.py

+2 -2

@@ -21,7 +21,7 @@ def _topk_log_softmax_kernel(

21 BLOCK_SIZE: tl.constexpr,

22 PADDED_TOPK: tl.constexpr,

23 ):

24 - req_idx = tl.program_id(0)

24 + req_idx = tl.program_id(0).to(tl.int64)

25 row_ptr = logits_ptr + req_idx * logits_stride

27 max_val = float("-inf")

@@ -61,7 +61,7 @@ def _ranks_kernel(

61 vocab_size,

62 BLOCK_SIZE: tl.constexpr,

63 ):

64 - req_idx = tl.program_id(0)

64 + req_idx = tl.program_id(0).to(tl.int64)

65 row_ptr = logits_ptr + req_idx * logits_stride

67 token_id = tl.load(token_ids_ptr + req_idx)

vllm/v1/worker/gpu/sample/min_p.py

+1 -1

@@ -14,7 +14,7 @@ def _min_p_kernel(

14 vocab_size,

15 BLOCK_SIZE: tl.constexpr,

16 ):

17 - token_idx = tl.program_id(0)

17 + token_idx = tl.program_id(0).to(tl.int64)

18 req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx)

19 min_p = tl.load(min_p_ptr + req_state_idx).to(tl.float32)

20 if min_p == 0.0:

vllm/v1/worker/gpu/sample/penalties.py

+1 -1

@@ -120,7 +120,7 @@ def _penalties_kernel(

120 vocab_size,

121 BLOCK_SIZE: tl.constexpr,

122 ):

123 - token_idx = tl.program_id(0)

123 + token_idx = tl.program_id(0).to(tl.int64)

124 req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx)

125 rep_penalty = tl.load(repetition_penalty_ptr + req_state_idx)

126 freq_penalty = tl.load(frequency_penalty_ptr + req_state_idx)

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

+6 -6

@@ -79,14 +79,14 @@ def _compute_block_stats_kernel(

79 BLOCK_SIZE: tl.constexpr,

80 HAS_DRAFT_LOGITS: tl.constexpr,

81 ):

82 - logit_idx = tl.program_id(0)

82 + logit_idx = tl.program_id(0).to(tl.int64)

83 draft_step_idx = tl.load(expanded_local_pos_ptr + logit_idx)

85 if draft_step_idx >= num_speculative_steps:

86 # Bonus token. Max/argmax and summed exponentials are not needed.

87 return

89 - req_state_idx = tl.load(expanded_idx_mapping_ptr + logit_idx)

89 + req_state_idx = tl.load(expanded_idx_mapping_ptr + logit_idx).to(tl.int64)

90 temp = tl.load(temp_ptr + req_state_idx).to(tl.float32)

92 block_idx = tl.program_id(1)

@@ -206,8 +206,8 @@ def _rejection_kernel(

206 SYNTHETIC_MODE: tl.constexpr,

207 ):

208 req_idx = tl.program_id(0)

209 - req_state_idx = tl.load(idx_mapping_ptr + req_idx)

210 - start_idx = tl.load(cu_num_logits_ptr + req_idx)

209 + req_state_idx = tl.load(idx_mapping_ptr + req_idx).to(tl.int64)

210 + start_idx = tl.load(cu_num_logits_ptr + req_idx).to(tl.int64)

211 end_idx = tl.load(cu_num_logits_ptr + req_idx + 1)

212 num_tokens = end_idx - start_idx

213 seed = tl.load(seed_ptr + req_state_idx)

@@ -349,10 +349,10 @@ def _resample_kernel(

349 ):

350 req_idx = tl.program_id(0)

351 resample_idx = tl.load(rejected_step_ptr + req_idx)

352 - start_idx = tl.load(cu_num_logits_ptr + req_idx)

352 + start_idx = tl.load(cu_num_logits_ptr + req_idx).to(tl.int64)

353 end_idx = tl.load(cu_num_logits_ptr + req_idx + 1)

354 resample_token_idx = start_idx + resample_idx

355 - req_state_idx = tl.load(expanded_idx_mapping_ptr + resample_token_idx)

355 + req_state_idx = tl.load(expanded_idx_mapping_ptr + resample_token_idx).to(tl.int64)

356

357 temp = tl.load(temp_ptr + req_state_idx).to(tl.float32)

358 is_bonus = resample_token_idx == end_idx - 1