A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 6, 10:45 UTC
[Model Runner V2] Spec decode rejection sampler greedy support (#37238)
Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai>
Giancarlo Delfin committed
Mar 18, 2026 at 15:59 UTC
04244fd0e1d082134d22ae1021a9bad993db4f59
2 files changed
+207
-71
vllm/v1/worker/gpu/model_runner.py
+1
-3
@@ -821,9 +821,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
821
logits,
822
input_batch,
823
# Draft logits are needed for probabilistic rejection sampling.
824
- self.req_states.draft_logits[input_batch.idx_mapping]
825
- if self.req_states.draft_logits is not None
826
- else None,
824
+ self.req_states.draft_logits,
825
)
826
827
# Get the number of sampled and rejected tokens.
vllm/v1/worker/gpu/spec_decode/rejection_sampler.py
+206
-68
@@ -68,55 +68,158 @@ def strict_rejection_sample(
68
69
70
@triton.jit
71
-def _probabilistic_rejection_sample_kernel(
71
+def _gather_draft_logits_and_target_argmax_kernel(
72
+ local_target_argmax_ptr,
73
+ local_target_argmax_stride,
74
+ local_target_max_ptr,
75
+ local_target_max_stride,
76
+ # [num_logits, V]
77
+ out_draft_logits_ptr,
78
+ out_draft_logits_stride,
79
+ # [num_logits, V]
80
+ target_logits_ptr,
81
+ target_logits_stride,
82
+ # [max_num_reqs, num_speculative_steps, V]
83
+ draft_logits_ptr,
84
+ draft_logits_stride_0,
85
+ draft_logits_stride_1,
86
+ # [num_logits]
87
+ expanded_idx_mapping_ptr,
88
+ # [num_logits]
89
+ expanded_local_pos_ptr,
90
+ # [max_num_reqs]
91
+ temp_ptr,
92
+ vocab_size,
93
+ num_speculative_steps,
94
+ BLOCK_SIZE: tl.constexpr,
95
+):
96
+ logit_idx = tl.program_id(0)
97
+ req_state_idx = tl.load(expanded_idx_mapping_ptr + logit_idx)
98
+ draft_step_idx = tl.load(expanded_local_pos_ptr + logit_idx)
99
+
100
+ block_idx = tl.program_id(1)
101
+ block_offsets = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
102
+ mask = block_offsets < vocab_size
103
+ temp = tl.load(temp_ptr + req_state_idx).to(tl.float32)
104
+
105
+ if temp == 0.0:
106
+ # Greedy sampling. Get the target logits argmax.
107
+ target_logits = tl.load(
108
+ target_logits_ptr + logit_idx * target_logits_stride + block_offsets,
109
+ mask=mask,
110
+ other=float("-inf"),
111
+ ).to(tl.float32)
112
+ value, idx = tl.max(target_logits, axis=0, return_indices=True)
113
+ token_id = block_idx * BLOCK_SIZE + idx
114
+ tl.store(
115
+ local_target_argmax_ptr
116
+ + logit_idx * local_target_argmax_stride
117
+ + block_idx,
118
+ token_id,
119
+ )
120
+ tl.store(
121
+ local_target_max_ptr + logit_idx * local_target_max_stride + block_idx,
122
+ value,
123
+ )
124
+ elif draft_step_idx < num_speculative_steps:
125
+ draft_logits = tl.load(
126
+ draft_logits_ptr
127
+ + req_state_idx * draft_logits_stride_0
128
+ + draft_step_idx * draft_logits_stride_1
129
+ + block_offsets,
130
+ mask=mask,
131
+ other=float("-inf"),
132
+ ).to(tl.float32)
133
+ tl.store(
134
+ out_draft_logits_ptr + logit_idx * out_draft_logits_stride + block_offsets,
135
+ draft_logits,
136
+ mask=mask,
137
+ )
138
+
139
+
140
+@triton.jit
141
+def _probabilistic_rejection_kernel(
142
# [num_reqs, num_speculative_steps + 1]
143
sampled_ptr,
144
sampled_stride,
145
# [num_reqs]
146
rejected_steps_ptr,
147
+ # [num_reqs]
148
+ rejected_pos_ptr,
149
# [num_logits]
150
draft_sampled_ptr,
151
# [num_logits, V]
152
target_probs_ptr,
153
target_probs_stride,
82
- # [num_reqs, num_speculative_steps, V]
154
+ # [num_logits, V]
155
draft_probs_ptr,
84
- draft_probs_stride_0,
85
- draft_probs_stride_1,
156
+ draft_probs_stride,
157
+ # [num_logits, num_blocks]
158
+ local_target_argmax_ptr,
159
+ local_target_argmax_stride,
160
+ # [num_logits, num_blocks]
161
+ local_target_max_ptr,
162
+ local_target_max_stride,
163
# [num_reqs + 1]
164
cu_num_logits_ptr,
165
# [num_logits]
166
pos_ptr,
167
# [num_reqs]
168
idx_mapping_ptr,
92
- # [num_reqs]
169
+ # [max_num_reqs]
170
+ temp_ptr,
171
+ # [max_num_reqs]
172
seeds_ptr,
173
+ NUM_BLOCKS: tl.constexpr,
174
+ PADDED_NUM_BLOCKS: tl.constexpr,
175
):
176
req_idx = tl.program_id(0)
177
start_idx = tl.load(cu_num_logits_ptr + req_idx)
178
num_tokens = tl.load(cu_num_logits_ptr + req_idx + 1) - start_idx
98
- seed = tl.load(seeds_ptr + tl.load(idx_mapping_ptr + req_idx))
179
+ req_state_idx = tl.load(idx_mapping_ptr + req_idx)
180
+ seed = tl.load(seeds_ptr + req_state_idx)
181
+ temp = tl.load(temp_ptr + req_state_idx).to(tl.float32)
182
183
rejected_step = 0
184
accepted = True
185
for i in range(num_tokens - 1):
186
if accepted:
104
- draft_sampled = tl.load(draft_sampled_ptr + start_idx + i + 1)
105
- target_prob = tl.load(
106
- target_probs_ptr + (start_idx + i) * target_probs_stride + draft_sampled
107
- )
108
- draft_prob = tl.load(
109
- draft_probs_ptr
110
- + req_idx * draft_probs_stride_0
111
- + i * draft_probs_stride_1
112
- + draft_sampled
113
- )
114
- pos = tl.load(pos_ptr + start_idx + i)
115
- u = tl.sum(tl.rand(seed, pos + tl.arange(0, 1)))
116
- accepted &= target_prob > u * draft_prob
187
+ logit_idx = start_idx + i
188
+ draft_sampled = tl.load(draft_sampled_ptr + logit_idx + 1)
189
+ if temp == 0.0:
190
+ # Greedy sampling. Only accept the sampled draft token if
191
+ # it exactly matches the target argmax.
192
+ block_offsets = tl.arange(0, PADDED_NUM_BLOCKS)
193
+ block_mask = block_offsets < NUM_BLOCKS
194
+ local_max = tl.load(
195
+ local_target_max_ptr
196
+ + logit_idx * local_target_max_stride
197
+ + block_offsets,
198
+ mask=block_mask,
199
+ other=float("-inf"),
200
+ )
201
+ max_block = tl.argmax(local_max, axis=0)
202
+ target_argmax = tl.load(
203
+ local_target_argmax_ptr
204
+ + logit_idx * local_target_argmax_stride
205
+ + max_block
206
+ )
207
+ accepted &= target_argmax == draft_sampled
208
+ else:
209
+ target_prob = tl.load(
210
+ target_probs_ptr + logit_idx * target_probs_stride + draft_sampled
211
+ )
212
+ draft_prob = tl.load(
213
+ draft_probs_ptr + logit_idx * draft_probs_stride + draft_sampled
214
+ )
215
+ pos = tl.load(pos_ptr + logit_idx)
216
+ u = tl.sum(tl.rand(seed, pos + tl.arange(0, 1)))
217
+ accepted &= target_prob > u * draft_prob
218
tl.store(sampled_ptr + req_idx * sampled_stride + i, draft_sampled)
219
rejected_step += accepted
220
tl.store(rejected_steps_ptr + req_idx, rejected_step)
221
+ pos_val = tl.load(pos_ptr + start_idx + rejected_step)
222
+ tl.store(rejected_pos_ptr + req_idx, pos_val)
223
224
225
@triton.jit
@@ -124,63 +227,60 @@ def _compute_residual_logits_kernel(
227
# [num_reqs, V]
228
residual_logits_ptr,
229
residual_logits_stride,
127
- # [num_reqs]
128
- residual_pos_ptr,
129
- # [num_logits, V]
130
- target_logits_ptr,
131
- target_logits_stride,
230
# [num_logits, V]
231
target_probs_ptr,
232
target_probs_stride,
135
- # [num_reqs, num_speculative_steps, V]
233
+ # [num_logits, V]
234
draft_probs_ptr,
137
- draft_probs_stride_0,
138
- draft_probs_stride_1,
235
+ draft_probs_stride,
236
+ # [num_logits, V]
237
+ target_logits_ptr,
238
+ target_logits_stride,
239
# [num_reqs]
240
rejected_step_ptr,
241
# [num_reqs + 1]
242
cu_num_logits_ptr,
143
- # [num_logits]
144
- pos_ptr,
243
+ # [num_reqs]
244
+ idx_mapping_ptr,
245
+ # [max_num_reqs]
246
+ temp_ptr,
247
vocab_size,
248
BLOCK_SIZE: tl.constexpr,
249
):
250
req_idx = tl.program_id(0)
251
block_idx = tl.program_id(1)
252
253
+ req_state_idx = tl.load(idx_mapping_ptr + req_idx)
254
start_idx = tl.load(cu_num_logits_ptr + req_idx)
255
end_idx = tl.load(cu_num_logits_ptr + req_idx + 1)
153
- rejected_draft_step = tl.load(rejected_step_ptr + req_idx)
154
- rejected_logit_idx = start_idx + rejected_draft_step
155
-
256
+ rejected_logit_idx = start_idx + tl.load(rejected_step_ptr + req_idx)
257
+ temp = tl.load(temp_ptr + req_state_idx).to(tl.float32)
258
block_offsets = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
259
mask = block_offsets < vocab_size
260
159
- if rejected_logit_idx < end_idx - 1:
261
+ if temp == 0.0 or (rejected_logit_idx == end_idx - 1):
262
+ # Greedy sampling / bonus token. In either case, use the
263
+ # target logits directly to reduce numerical error.
264
+ residual_logits = tl.load(
265
+ target_logits_ptr
266
+ + rejected_logit_idx * target_logits_stride
267
+ + block_offsets,
268
+ mask=mask,
269
+ other=float("-inf"),
270
+ )
271
+ else:
272
target_probs = tl.load(
273
target_probs_ptr + rejected_logit_idx * target_probs_stride + block_offsets,
274
mask=mask,
275
other=0.0,
276
)
277
draft_probs = tl.load(
166
- draft_probs_ptr
167
- + req_idx * draft_probs_stride_0
168
- + rejected_draft_step * draft_probs_stride_1
169
- + block_offsets,
278
+ draft_probs_ptr + rejected_logit_idx * draft_probs_stride + block_offsets,
279
mask=mask,
280
other=0.0,
281
)
282
residual_probs = tl.maximum(target_probs - draft_probs, 0.0)
283
residual_logits = tl.log(residual_probs)
175
- else:
176
- # This is a bonus token. Directly return the target logits.
177
- residual_logits = tl.load(
178
- target_logits_ptr
179
- + rejected_logit_idx * target_logits_stride
180
- + block_offsets,
181
- mask=mask,
182
- other=0.0,
183
- )
284
285
tl.store(
286
residual_logits_ptr + req_idx * residual_logits_stride + block_offsets,
@@ -188,18 +288,13 @@ def _compute_residual_logits_kernel(
288
mask=mask,
289
)
290
191
- # First block computes the residual logit positions.
192
- if block_idx == 0:
193
- pos_val = tl.load(pos_ptr + rejected_logit_idx)
194
- tl.store(residual_pos_ptr + req_idx, pos_val)
195
-
291
292
def probabilistic_rejection_sample(
198
- # [num_draft_tokens + num_reqs, V]
293
+ # [num_logits, V]
294
target_logits: torch.Tensor,
200
- # [num_reqs, num_speculative_steps, V]
295
+ # [max_num_reqs, num_speculative_steps, V]
296
draft_logits: torch.Tensor,
202
- # [num_draft_tokens + num_reqs]
297
+ # [num_logits]
298
draft_sampled: torch.Tensor,
299
# [num_reqs + 1]
300
cu_num_logits: torch.Tensor,
@@ -207,16 +302,53 @@ def probabilistic_rejection_sample(
302
pos: torch.Tensor,
303
# [num_reqs]
304
idx_mapping: torch.Tensor,
305
+ # [num_logits]
306
+ expanded_idx_mapping: torch.Tensor,
307
+ # [num_logits]
308
+ expanded_local_pos: torch.Tensor,
309
+ # [max_num_reqs]
310
temperature: torch.Tensor,
311
+ # [max_num_reqs]
312
seed: torch.Tensor,
313
num_speculative_steps: int,
314
) -> tuple[torch.Tensor, torch.Tensor]:
315
num_reqs = cu_num_logits.shape[0] - 1
215
- vocab_size = target_logits.shape[-1]
316
+ num_logits, vocab_size = target_logits.shape
317
+
318
+ BLOCK_SIZE = 1024
319
+ num_blocks = triton.cdiv(vocab_size, BLOCK_SIZE)
320
+
321
+ # Gather draft logits and target argmax for greedy sampling.
322
+ gathered_draft_logits = target_logits.new_empty(target_logits.shape)
323
+ local_target_argmax = target_logits.new_empty(
324
+ num_logits, num_blocks, dtype=torch.int64
325
+ )
326
+ local_target_max = target_logits.new_empty(
327
+ num_logits, num_blocks, dtype=torch.float32
328
+ )
329
+ _gather_draft_logits_and_target_argmax_kernel[(num_logits, num_blocks)](
330
+ local_target_argmax,
331
+ local_target_argmax.stride(0),
332
+ local_target_max,
333
+ local_target_max.stride(0),
334
+ gathered_draft_logits,
335
+ gathered_draft_logits.stride(0),
336
+ target_logits,
337
+ target_logits.stride(0),
338
+ draft_logits,
339
+ draft_logits.stride(0),
340
+ draft_logits.stride(1),
341
+ expanded_idx_mapping,
342
+ expanded_local_pos,
343
+ temperature,
344
+ vocab_size,
345
+ num_speculative_steps,
346
+ BLOCK_SIZE=BLOCK_SIZE,
347
+ )
348
349
# Compute target and draft probs.
350
target_probs = torch.softmax(target_logits, dim=-1)
219
- draft_probs = torch.softmax(draft_logits, dim=-1)
351
+ draft_probs = torch.softmax(gathered_draft_logits, dim=-1)
352
353
# Rejection sample.
354
# [num_reqs, num_speculative_steps + 1]
@@ -225,45 +357,49 @@ def probabilistic_rejection_sample(
357
)
358
# [num_reqs]
359
rejected_steps = sampled.new_empty(num_reqs)
228
- _probabilistic_rejection_sample_kernel[(num_reqs,)](
360
+ # [num_reqs]
361
+ rejected_pos = pos.new_empty(num_reqs)
362
+ _probabilistic_rejection_kernel[(num_reqs,)](
363
sampled,
364
sampled.stride(0),
365
rejected_steps,
366
+ rejected_pos,
367
draft_sampled,
368
target_probs,
369
target_probs.stride(0),
370
draft_probs,
371
draft_probs.stride(0),
237
- draft_probs.stride(1),
372
+ local_target_argmax,
373
+ local_target_argmax.stride(0),
374
+ local_target_max,
375
+ local_target_max.stride(0),
376
cu_num_logits,
377
pos,
378
idx_mapping,
379
+ temperature,
380
seed,
381
num_warps=1,
382
+ NUM_BLOCKS=num_blocks,
383
+ PADDED_NUM_BLOCKS=triton.next_power_of_2(num_blocks),
384
)
385
386
# Compute the logits and positions to resample the rejected/bonus
387
# tokens from.
388
# [num_reqs, vocab_size]
389
residual_logits = target_logits.new_empty(num_reqs, vocab_size)
249
- # [num_reqs]
250
- residual_pos = pos.new_empty(num_reqs)
251
- BLOCK_SIZE = 1024
252
- num_blocks = triton.cdiv(vocab_size, BLOCK_SIZE)
390
_compute_residual_logits_kernel[(num_reqs, num_blocks)](
391
residual_logits,
392
residual_logits.stride(0),
256
- residual_pos,
257
- target_logits,
258
- target_logits.stride(0),
393
target_probs,
394
target_probs.stride(0),
395
draft_probs,
396
draft_probs.stride(0),
263
- draft_probs.stride(1),
397
+ target_logits,
398
+ target_logits.stride(0),
399
rejected_steps,
400
cu_num_logits,
266
- pos,
401
+ idx_mapping,
402
+ temperature,
403
vocab_size,
404
BLOCK_SIZE=BLOCK_SIZE,
405
)
@@ -274,7 +410,7 @@ def probabilistic_rejection_sample(
410
idx_mapping,
411
temperature,
412
seed,
277
- residual_pos,
413
+ rejected_pos,
414
apply_temperature=False,
415
)
416
sampled.scatter_(1, rejected_steps.unsqueeze(1), resampled.unsqueeze(1))
@@ -333,6 +469,8 @@ class RejectionSampler:
469
input_batch.cu_num_logits,
470
pos,
471
input_batch.idx_mapping,
472
+ input_batch.expanded_idx_mapping,
473
+ input_batch.expanded_local_pos,
474
self.sampler.sampling_states.temperature.gpu,
475
self.sampler.sampling_states.seeds.gpu,
476
self.num_speculative_steps,