[Bugfix][Model Runner V2][Spec Decode] Fix int32 offset overflow in block verification kernels (#47383)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Woosuk Kwon committed
Jul 2, 2026 at 07:32 UTC
a47f38f82569c236d7d23b7ad0c8792ac6d62247
2 files changed
+148
-4
tests/v1/worker/test_gpu_rejection_sampler_i64.py
new
+144
@@ -0,0 +1,144 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""int64 indexing in the block-verification rejection sampler kernels.
4
+
5
+With a GLM-scale vocab (~155k), logit_idx * vocab_stride exceeds int32 once
6
+logit_idx >= ~13.8k (and req_state_idx * draft_stride_0 similarly), so the
7
+block-verification kernels must promote indices to int64 before the stride
8
+multiply. These tests place one request at a high logit/request-state index
9
+and check its outputs match the same request run alone at index 0.
10
+
11
+Each case allocates ~5 GiB of GPU memory.
12
+"""
13
+
14
+import pytest
15
+import torch
16
+
17
+from vllm.platforms import current_platform
18
+from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import rejection_sample
19
+
20
+VOCAB_SIZE = 155264
21
+NUM_SPECULATIVE_STEPS = 2
22
+LOGITS_PER_REQ = NUM_SPECULATIVE_STEPS + 1
23
+
24
+
25
+def _run(
26
+ target_logits: torch.Tensor, # [LOGITS_PER_REQ, V], the request under test
27
+ draft_logits: torch.Tensor, # [NUM_SPECULATIVE_STEPS, V]
28
+ draft_tokens: list[int],
29
+ num_reqs: int,
30
+ max_num_reqs: int,
31
+ roi_req: int,
32
+ req_state_idx: int,
33
+ device: torch.device,
34
+) -> tuple[torch.Tensor, torch.Tensor]:
35
+ """Run rejection_sample on a batch of uniform k=2 requests, with the
36
+ request under test at request index `roi_req` and draft-logits row
37
+ `req_state_idx`. Returns that request's (sampled, num_sampled)."""
38
+ num_logits = num_reqs * LOGITS_PER_REQ
39
+ roi_row = roi_req * LOGITS_PER_REQ
40
+
41
+ full_target = torch.randn(
42
+ num_logits, VOCAB_SIZE, dtype=torch.bfloat16, device=device
43
+ )
44
+ full_target[roi_row : roi_row + LOGITS_PER_REQ] = target_logits
45
+ full_draft = torch.randn(
46
+ max_num_reqs,
47
+ NUM_SPECULATIVE_STEPS,
48
+ VOCAB_SIZE,
49
+ dtype=torch.bfloat16,
50
+ device=device,
51
+ )
52
+ full_draft[req_state_idx] = draft_logits
53
+
54
+ # Draft token for step i of request r lives at row r*LOGITS_PER_REQ + i + 1.
55
+ draft_sampled = torch.zeros(num_logits, dtype=torch.int64, device=device)
56
+ for i, tok in enumerate(draft_tokens):
57
+ draft_sampled[roi_row + i + 1] = tok
58
+
59
+ cu_num_logits = torch.arange(
60
+ 0, num_logits + 1, LOGITS_PER_REQ, dtype=torch.int32, device=device
61
+ )
62
+ # Filler requests all map to draft-logits row 0.
63
+ idx_mapping = torch.zeros(num_reqs, dtype=torch.int32, device=device)
64
+ idx_mapping[roi_req] = req_state_idx
65
+ expanded_idx_mapping = idx_mapping.repeat_interleave(LOGITS_PER_REQ)
66
+ expanded_local_pos = (
67
+ torch.arange(LOGITS_PER_REQ, dtype=torch.int32, device=device)
68
+ .repeat(num_reqs)
69
+ .contiguous()
70
+ )
71
+
72
+ # Positions only need to match between the reference and big-batch runs
73
+ # for the request under test (they key the Gumbel noise and uniforms).
74
+ pos = torch.arange(num_logits, dtype=torch.int64, device=device)
75
+ pos[roi_row : roi_row + LOGITS_PER_REQ] = torch.arange(
76
+ 1000, 1000 + LOGITS_PER_REQ, dtype=torch.int64, device=device
77
+ )
78
+ temperature = torch.ones(max_num_reqs, dtype=torch.float32, device=device)
79
+ seeds = torch.full((max_num_reqs,), 42, dtype=torch.int64, device=device)
80
+
81
+ sampled, num_sampled = rejection_sample(
82
+ target_logits=full_target,
83
+ draft_logits=full_draft,
84
+ draft_sampled=draft_sampled,
85
+ cu_num_logits=cu_num_logits,
86
+ pos=pos,
87
+ idx_mapping=idx_mapping,
88
+ expanded_idx_mapping=expanded_idx_mapping,
89
+ expanded_local_pos=expanded_local_pos,
90
+ temperature=temperature,
91
+ seed=seeds,
92
+ num_speculative_steps=NUM_SPECULATIVE_STEPS,
93
+ use_block_verification=True,
94
+ )
95
+ return sampled[roi_req].clone(), num_sampled[roi_req].clone()
96
+
97
+
98
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA")
99
+@pytest.mark.parametrize(
100
+ # Overflow triggers when roi_req * LOGITS_PER_REQ * VOCAB_SIZE (target
101
+ # logits) or req_state_idx * NUM_SPECULATIVE_STEPS * VOCAB_SIZE (draft
102
+ # logits) exceeds 2**31.
103
+ ("num_reqs", "max_num_reqs", "roi_req", "req_state_idx"),
104
+ [
105
+ (5500, 4, 5000, 2), # target-side: logit_idx 15000 * V > 2**31
106
+ (22, 8192, 11, 8000), # draft-side: req_state_idx 8000 * 2V > 2**31
107
+ ],
108
+)
109
+def test_block_verification_i64_indexing(
110
+ num_reqs: int, max_num_reqs: int, roi_req: int, req_state_idx: int
111
+):
112
+ device = torch.device("cuda:0")
113
+ torch.manual_seed(0)
114
+ target_logits = torch.randn(
115
+ LOGITS_PER_REQ, VOCAB_SIZE, dtype=torch.bfloat16, device=device
116
+ )
117
+ draft_logits = torch.randn(
118
+ NUM_SPECULATIVE_STEPS, VOCAB_SIZE, dtype=torch.bfloat16, device=device
119
+ )
120
+ draft_tokens = [123, 45678]
121
+
122
+ sampled_ref, num_sampled_ref = _run(
123
+ target_logits,
124
+ draft_logits,
125
+ draft_tokens,
126
+ num_reqs=1,
127
+ max_num_reqs=1,
128
+ roi_req=0,
129
+ req_state_idx=0,
130
+ device=device,
131
+ )
132
+ sampled, num_sampled = _run(
133
+ target_logits,
134
+ draft_logits,
135
+ draft_tokens,
136
+ num_reqs=num_reqs,
137
+ max_num_reqs=max_num_reqs,
138
+ roi_req=roi_req,
139
+ req_state_idx=req_state_idx,
140
+ device=device,
141
+ )
142
+ assert num_sampled.item() == num_sampled_ref.item()
143
+ n = num_sampled_ref.item()
144
+ assert torch.equal(sampled[:n], sampled_ref[:n])
vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py
+4
-4
@@ -328,8 +328,8 @@ def _compute_cumulative_log_p_kernel(
328
HAS_DRAFT_LOGITS: tl.constexpr,
329
):
330
req_idx = tl.program_id(0)
331
- req_state_idx = tl.load(idx_mapping_ptr + req_idx)
332
- start_idx = tl.load(cu_num_logits_ptr + req_idx)
331
+ req_state_idx = tl.load(idx_mapping_ptr + req_idx).to(tl.int64)
332
+ start_idx = tl.load(cu_num_logits_ptr + req_idx).to(tl.int64)
333
end_idx = tl.load(cu_num_logits_ptr + req_idx + 1)
334
num_draft_tokens = end_idx - start_idx - 1
335
temp = tl.load(temp_ptr + req_state_idx).to(tl.float32)
@@ -405,7 +405,7 @@ def _compute_local_residual_mass_kernel(
405
BLOCK_SIZE: tl.constexpr,
406
PADDED_VOCAB_NUM_BLOCKS: tl.constexpr,
407
):
408
- logit_idx = tl.program_id(0)
408
+ logit_idx = tl.program_id(0).to(tl.int64)
409
draft_step_idx = tl.load(expanded_local_pos_ptr + logit_idx)
410
if draft_step_idx == 0 or draft_step_idx >= num_speculative_steps:
411
# The acceptance threshold, h, looks one position ahead and sums
@@ -413,7 +413,7 @@ def _compute_local_residual_mass_kernel(
413
# first and last (bonus) positions aren't needed for this computation.
414
return
415
416
- req_state_idx = tl.load(expanded_idx_mapping_ptr + logit_idx)
416
+ req_state_idx = tl.load(expanded_idx_mapping_ptr + logit_idx).to(tl.int64)
417
temp = tl.load(temp_ptr + req_state_idx).to(tl.float32)
418
if temp == 0.0:
419
return