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 4, 15:00 UTC
[Model Runner V2] Fix rejection sampling acceptance rate gap vs MRV1 (#40651)
Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai>
Giancarlo Delfin committed
Apr 26, 2026 at 19:12 UTC
f5f987851493e6f09ab2ddeb3f33ae878eda0353
4 files changed
+126
-148
vllm/config/speculative.py
+15
-6
@@ -63,7 +63,8 @@ SpeculativeMethod = Literal[
63
EagleModelTypes,
64
NgramGPUTypes,
65
]
66
-RejectionSampleMethod = Literal["strict", "probabilistic", "synthetic"]
66
+RejectionSampleMethod = Literal["standard", "synthetic"]
67
+DraftSampleMethod = Literal["greedy", "gumbel"]
68
69
70
@config
@@ -183,11 +184,11 @@ class SpeculativeConfig:
184
"""Load config for the draft model. If not specified, will use the load
185
config from the target model."""
186
186
- rejection_sample_method: RejectionSampleMethod = "strict"
187
- """Whether to use strict (target and draft sampled tokens match exactly)
188
- or probabilistic rejection sampling. Both respect the target model
189
- distribution, but the latter yields a higher acceptance rate at the cost
190
- of more memory to cache draft logits."""
187
+ rejection_sample_method: RejectionSampleMethod = "standard"
188
+ """The rejection sampling method to use. 'standard' uses probabilistic
189
+ rejection sampling (with or without cached draft logits, controlled by
190
+ draft_sample_method). 'synthetic' accepts draft tokens with a decaying
191
+ probability calibrated to synthetic_acceptance_rate."""
192
193
synthetic_acceptance_rates: list[float] | None = None
194
"""Per-position *unconditional* acceptance rates for synthetic rejection
@@ -248,6 +249,14 @@ class SpeculativeConfig:
249
)
250
return SpeculativeConfig._acceptance_length_to_rates(length, n)
251
252
+ draft_sample_method: DraftSampleMethod = "greedy"
253
+ """How the draft model samples tokens. 'greedy' always picks the argmax
254
+ token, and the draft probabilities are treated as one-hot during rejection
255
+ sampling. 'gumbel' adds Gumbel noise for stochastic sampling, and the full
256
+ draft logits are used for the probability ratio test during rejection
257
+ sampling. This comes at the cost of additional GPU memory usage. This
258
+ parameter currently only applies to Model Runner V2."""
259
+
260
def compute_hash(self) -> str:
261
"""
262
WARNING: Whenever a new field is added to this config,
vllm/v1/worker/gpu/spec_decode/eagle/speculator.py
+29
-21
@@ -99,7 +99,7 @@ class EagleSpeculator:
99
)
100
101
self.draft_logits: torch.Tensor | None = None
102
- if self.speculative_config.rejection_sample_method == "probabilistic":
102
+ if self.speculative_config.draft_sample_method == "gumbel":
103
self.draft_logits = torch.zeros(
104
self.max_num_reqs,
105
self.num_speculative_steps,
@@ -220,6 +220,28 @@ class EagleSpeculator:
220
last_hidden_states, hidden_states = ret_hidden_states
221
return last_hidden_states, hidden_states
222
223
+ def _sample_draft(
224
+ self,
225
+ logits: torch.Tensor,
226
+ idx_mapping: torch.Tensor,
227
+ pos: torch.Tensor,
228
+ step: int,
229
+ ) -> torch.Tensor:
230
+ if self.draft_logits is not None:
231
+ # NOTE(woosuk): We must add 1 to the positions to match the Gumbel noise
232
+ # used for draft and target sampling.
233
+ return gumbel_sample(
234
+ logits,
235
+ idx_mapping,
236
+ self.temperature,
237
+ self.seeds,
238
+ pos + 1,
239
+ apply_temperature=True,
240
+ processed_logits_out=self.draft_logits[:, step],
241
+ )
242
+ else:
243
+ return logits.argmax(dim=-1)
244
+
245
def prefill(
246
self,
247
num_reqs: int,
@@ -245,18 +267,11 @@ class EagleSpeculator:
267
sample_hidden_states = last_hidden_states[last_token_indices]
268
logits = self.model.compute_logits(sample_hidden_states)
269
248
- # NOTE(woosuk): We must add 1 to the positions to match the Gumbel noise
249
- # used for draft and target sampling.
250
- self.draft_tokens[:num_reqs, 0] = gumbel_sample(
270
+ self.draft_tokens[:num_reqs, 0] = self._sample_draft(
271
logits,
272
idx_mapping,
253
- self.temperature,
254
- self.seeds,
255
- pos + 1,
256
- apply_temperature=True,
257
- processed_logits_out=self.draft_logits[:, 0]
258
- if self.draft_logits is not None
259
- else None,
273
+ pos,
274
+ step=0,
275
)
276
self.hidden_states[:num_reqs] = hidden_states[last_token_indices]
277
self.input_buffers.positions[:num_reqs] = pos
@@ -286,18 +301,11 @@ class EagleSpeculator:
301
hidden_states = hidden_states[:num_reqs]
302
logits = self.model.compute_logits(last_hidden_states)
303
289
- # NOTE(woosuk): We must add 1 to the positions to match the Gumbel noise
290
- # used for draft and target sampling.
291
- draft_tokens = gumbel_sample(
304
+ draft_tokens = self._sample_draft(
305
logits,
306
idx_mapping,
294
- self.temperature,
295
- self.seeds,
296
- pos + 1,
297
- apply_temperature=True,
298
- processed_logits_out=self.draft_logits[:, step]
299
- if self.draft_logits is not None
300
- else None,
307
+ pos,
308
+ step=step,
309
)
310
self.draft_tokens[:num_reqs, step] = draft_tokens
311
vllm/v1/worker/gpu/spec_decode/probabilistic_rejection_sampler_utils.py
+81
-53
@@ -45,7 +45,7 @@ def _compute_global_lse(
45
46
47
@triton.jit
48
-def _compute_block_max_and_sumexp_kernel(
48
+def _compute_block_stats_kernel(
49
# [num_logits, num_blocks]
50
target_local_argmax_ptr,
51
target_local_argmax_stride,
@@ -77,6 +77,7 @@ def _compute_block_max_and_sumexp_kernel(
77
vocab_size,
78
num_speculative_steps,
79
BLOCK_SIZE: tl.constexpr,
80
+ HAS_DRAFT_LOGITS: tl.constexpr,
81
):
82
logit_idx = tl.program_id(0)
83
draft_step_idx = tl.load(expanded_local_pos_ptr + logit_idx)
@@ -112,24 +113,6 @@ def _compute_block_max_and_sumexp_kernel(
113
value,
114
)
115
else:
115
- # Get local draft max and summed exponentials.
116
- draft_logits = tl.load(
117
- draft_logits_ptr
118
- + req_state_idx * draft_logits_stride_0
119
- + draft_step_idx * draft_logits_stride_1
120
- + block_offsets,
121
- mask=mask,
122
- other=float("-inf"),
123
- ).to(tl.float32)
124
- draft_max, draft_sumexp = _compute_block_max_and_sumexp(draft_logits)
125
- tl.store(
126
- draft_local_max_ptr + logit_idx * draft_local_max_stride + block_idx,
127
- draft_max,
128
- )
129
- tl.store(
130
- draft_local_sumexp_ptr + logit_idx * draft_local_sumexp_stride + block_idx,
131
- draft_sumexp,
132
- )
116
# Get local target max and summed exponentials.
117
target_logits = tl.load(
118
target_logits_ptr + logit_idx * target_logits_stride + block_offsets,
@@ -147,6 +130,27 @@ def _compute_block_max_and_sumexp_kernel(
130
+ block_idx,
131
target_sumexp,
132
)
133
+ if HAS_DRAFT_LOGITS:
134
+ # Get local draft max and summed exponentials.
135
+ draft_logits = tl.load(
136
+ draft_logits_ptr
137
+ + req_state_idx * draft_logits_stride_0
138
+ + draft_step_idx * draft_logits_stride_1
139
+ + block_offsets,
140
+ mask=mask,
141
+ other=float("-inf"),
142
+ ).to(tl.float32)
143
+ draft_max, draft_sumexp = _compute_block_max_and_sumexp(draft_logits)
144
+ tl.store(
145
+ draft_local_max_ptr + logit_idx * draft_local_max_stride + block_idx,
146
+ draft_max,
147
+ )
148
+ tl.store(
149
+ draft_local_sumexp_ptr
150
+ + logit_idx * draft_local_sumexp_stride
151
+ + block_idx,
152
+ draft_sumexp,
153
+ )
154
155
156
@triton.jit
@@ -196,6 +200,7 @@ def _probabilistic_rejection_kernel(
200
pos_ptr,
201
vocab_num_blocks,
202
PADDED_VOCAB_NUM_BLOCKS: tl.constexpr,
203
+ HAS_DRAFT_LOGITS: tl.constexpr,
204
):
205
req_idx = tl.program_id(0)
206
req_state_idx = tl.load(idx_mapping_ptr + req_idx)
@@ -238,12 +243,6 @@ def _probabilistic_rejection_kernel(
243
target_logit = tl.load(
244
target_logits_ptr + logit_idx * target_logits_stride + draft_sampled
245
).to(tl.float32)
241
- draft_logit = tl.load(
242
- draft_logits_ptr
243
- + req_state_idx * draft_logits_stride_0
244
- + i * draft_logits_stride_1
245
- + draft_sampled
246
- ).to(tl.float32)
246
target_lse = _compute_global_lse(
247
target_local_max_ptr,
248
target_local_max_stride,
@@ -253,19 +252,29 @@ def _probabilistic_rejection_kernel(
252
vocab_num_blocks,
253
PADDED_VOCAB_NUM_BLOCKS,
254
)
256
- draft_lse = _compute_global_lse(
257
- draft_local_max_ptr,
258
- draft_local_max_stride,
259
- draft_local_sumexp_ptr,
260
- draft_local_sumexp_stride,
261
- logit_idx,
262
- vocab_num_blocks,
263
- PADDED_VOCAB_NUM_BLOCKS,
264
- )
255
target_log_prob = target_logit - target_lse
266
- draft_log_prob = draft_logit - draft_lse
256
pos = tl.load(pos_ptr + logit_idx)
257
u = tl_rand64(seed, pos, includes_zero=False)
258
+ if HAS_DRAFT_LOGITS:
259
+ draft_logit = tl.load(
260
+ draft_logits_ptr
261
+ + req_state_idx * draft_logits_stride_0
262
+ + i * draft_logits_stride_1
263
+ + draft_sampled
264
+ ).to(tl.float32)
265
+ draft_lse = _compute_global_lse(
266
+ draft_local_max_ptr,
267
+ draft_local_max_stride,
268
+ draft_local_sumexp_ptr,
269
+ draft_local_sumexp_stride,
270
+ logit_idx,
271
+ vocab_num_blocks,
272
+ PADDED_VOCAB_NUM_BLOCKS,
273
+ )
274
+ draft_log_prob = draft_logit - draft_lse
275
+ else:
276
+ # One-hot draft: q(draft_token) = 1, log_q = 0.
277
+ draft_log_prob = 0
278
# Probability ratio test: p(x) > u * q(x)
279
# Equivalent log form: log_p(x) > log(u) + log_q(x)
280
accepted &= target_log_prob > tl.log(u) + draft_log_prob
@@ -301,6 +310,8 @@ def _resample_kernel(
310
cu_num_logits_ptr,
311
# [num_logits]
312
expanded_idx_mapping_ptr,
313
+ # [num_logits]
314
+ draft_sampled_ptr,
315
# [max_num_reqs]
316
temp_ptr,
317
# [max_num_reqs]
@@ -309,6 +320,7 @@ def _resample_kernel(
320
pos_ptr,
321
vocab_size,
322
BLOCK_SIZE: tl.constexpr,
323
+ HAS_DRAFT_LOGITS: tl.constexpr,
324
):
325
req_idx = tl.program_id(0)
326
resample_idx = tl.load(rejected_step_ptr + req_idx)
@@ -327,22 +339,17 @@ def _resample_kernel(
339
block_idx = tl.program_id(1)
340
block = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
341
mask = block < vocab_size
342
+ target_logits = tl.load(
343
+ target_logits_ptr + resample_token_idx * target_logits_stride + block,
344
+ mask=mask,
345
+ other=float("-inf"),
346
+ ).to(tl.float32)
347
331
- # Compute the residual logits to resample the rejected token
332
- # from. In the case of no rejections (bonus token), we directly
333
- # use the target logits.
348
+ # Compute the residual logits to resample the rejected token from.
349
if is_bonus:
335
- residual_logits = tl.load(
336
- target_logits_ptr + resample_token_idx * target_logits_stride + block,
337
- mask=mask,
338
- other=float("-inf"),
339
- ).to(tl.float32)
340
- else:
341
- target_logits = tl.load(
342
- target_logits_ptr + resample_token_idx * target_logits_stride + block,
343
- mask=mask,
344
- other=float("-inf"),
345
- ).to(tl.float32)
350
+ # Bonus token (no rejections). Directly use the target logits.
351
+ residual_logits = target_logits
352
+ elif HAS_DRAFT_LOGITS:
353
draft_logits = tl.load(
354
draft_logits_ptr
355
+ req_state_idx * draft_logits_stride_0
@@ -365,6 +372,15 @@ def _resample_kernel(
372
target_log_probs + tl.log(1 - ratio),
373
float("-inf"),
374
).to(tl.float32)
375
+ else:
376
+ # One-hot draft. The residual is just the target distribution with
377
+ # the rejected draft token probability zeroed out.
378
+ rejected_draft_token = tl.load(draft_sampled_ptr + resample_token_idx + 1)
379
+ residual_logits = tl.where(
380
+ block != rejected_draft_token,
381
+ target_logits,
382
+ float("-inf"),
383
+ ).to(tl.float32)
384
385
# Resample the rejected/bonus token.
386
value, idx = gumbel_block_argmax(
@@ -456,7 +472,7 @@ def probabilistic_rejection_sample(
472
# [num_logits, V]
473
target_logits: torch.Tensor,
474
# [max_num_reqs, num_speculative_steps, V]
459
- draft_logits: torch.Tensor,
475
+ draft_logits: torch.Tensor | None,
476
# [num_logits]
477
draft_sampled: torch.Tensor,
478
# [num_reqs + 1]
@@ -477,9 +493,17 @@ def probabilistic_rejection_sample(
493
) -> tuple[torch.Tensor, torch.Tensor]:
494
num_reqs = cu_num_logits.shape[0] - 1
495
num_logits, vocab_size = target_logits.shape
496
+ has_draft_logits = draft_logits is not None
497
481
- # Gather draft logits, compute target argmax for greedy, and
482
- # compute per-block LSE and max for non-greedy requests.
498
+ if draft_logits is None:
499
+ # When draft_logits is None, create a dummy tensor so that Triton
500
+ # kernel signatures receive valid pointers/strides. The kernels
501
+ # will never read from it when HAS_DRAFT_LOGITS=False.
502
+ draft_logits = target_logits.new_empty(1, 1, 1)
503
+
504
+ # Compute the block-level logits stats, such as target argmax
505
+ # (for greedy requests), and target max + softmax exponential
506
+ # (for non-greedy requests).
507
VOCAB_BLOCK_SIZE = 8192
508
vocab_num_blocks = triton.cdiv(vocab_size, VOCAB_BLOCK_SIZE)
509
padded_vocab_num_blocks = triton.next_power_of_2(vocab_num_blocks)
@@ -498,7 +522,7 @@ def probabilistic_rejection_sample(
522
draft_local_sumexp = target_logits.new_empty(
523
num_logits, vocab_num_blocks, dtype=torch.float32
524
)
501
- _compute_block_max_and_sumexp_kernel[(num_logits, vocab_num_blocks)](
525
+ _compute_block_stats_kernel[(num_logits, vocab_num_blocks)](
526
target_local_argmax,
527
target_local_argmax.stride(0),
528
target_local_max,
@@ -520,6 +544,7 @@ def probabilistic_rejection_sample(
544
vocab_size,
545
num_speculative_steps,
546
BLOCK_SIZE=VOCAB_BLOCK_SIZE,
547
+ HAS_DRAFT_LOGITS=has_draft_logits,
548
)
549
550
# Sample up until the first rejected/bonus token, and store
@@ -559,6 +584,7 @@ def probabilistic_rejection_sample(
584
pos,
585
vocab_num_blocks,
586
PADDED_VOCAB_NUM_BLOCKS=padded_vocab_num_blocks,
587
+ HAS_DRAFT_LOGITS=has_draft_logits,
588
num_warps=1,
589
)
590
@@ -587,11 +613,13 @@ def probabilistic_rejection_sample(
613
num_sampled,
614
cu_num_logits,
615
expanded_idx_mapping,
616
+ draft_sampled,
617
temperature,
618
seed,
619
pos,
620
vocab_size,
621
BLOCK_SIZE=RESAMPLE_BLOCK_SIZE,
622
+ HAS_DRAFT_LOGITS=has_draft_logits,
623
)
624
625
# Insert the resampled tokens into the output sampled.
vllm/v1/worker/gpu/spec_decode/rejection_sampler.py
+1
-68
@@ -20,63 +20,6 @@ from vllm.v1.worker.gpu.spec_decode.synthetic_rejection_sampler_utils import (
20
)
21
22
23
-@triton.jit
24
-def _strict_rejection_sample_kernel(
25
- sampled_ptr, # [num_reqs, num_speculative_steps + 1]
26
- sampled_stride,
27
- num_sampled_ptr, # [num_reqs]
28
- target_sampled_ptr, # [num_draft_tokens + num_reqs]
29
- input_ids_ptr, # [num_draft_tokens + num_reqs]
30
- cu_num_logits_ptr, # [num_reqs + 1]
31
-):
32
- req_idx = tl.program_id(0)
33
- start_idx = tl.load(cu_num_logits_ptr + req_idx)
34
- end_idx = tl.load(cu_num_logits_ptr + req_idx + 1)
35
- num_tokens = end_idx - start_idx
36
-
37
- num_sampled = 0
38
- rejected = False
39
- for i in range(num_tokens - 1):
40
- if not rejected:
41
- target_sampled = tl.load(target_sampled_ptr + start_idx + i)
42
- draft_sampled = tl.load(input_ids_ptr + start_idx + i + 1)
43
- tl.store(sampled_ptr + req_idx * sampled_stride + i, target_sampled)
44
- num_sampled += 1
45
- if target_sampled != draft_sampled:
46
- rejected = True
47
- if not rejected:
48
- target_sampled = tl.load(target_sampled_ptr + start_idx + num_tokens - 1)
49
- tl.store(
50
- sampled_ptr + req_idx * sampled_stride + num_tokens - 1, target_sampled
51
- )
52
- num_sampled += 1
53
- tl.store(num_sampled_ptr + req_idx, num_sampled)
54
-
55
-
56
-def strict_rejection_sample(
57
- # [num_draft_tokens + num_reqs]
58
- target_sampled: torch.Tensor,
59
- # [num_draft_tokens + num_reqs]
60
- draft_sampled: torch.Tensor,
61
- # [num_reqs + 1]
62
- cu_num_logits: torch.Tensor,
63
- num_speculative_steps,
64
-) -> tuple[torch.Tensor, torch.Tensor]:
65
- num_reqs = cu_num_logits.shape[0] - 1
66
- sampled = target_sampled.new_empty(num_reqs, num_speculative_steps + 1)
67
- num_sampled = target_sampled.new_empty(num_reqs, dtype=torch.int32)
68
- _strict_rejection_sample_kernel[(num_reqs,)](
69
- sampled,
70
- sampled.stride(0),
71
- num_sampled,
72
- target_sampled,
73
- draft_sampled,
74
- cu_num_logits,
75
- num_warps=1,
76
- )
77
- return sampled, num_sampled
78
-
79
-
23
@triton.jit
24
def _flatten_sampled_kernel(
25
# [num_logits]
@@ -163,17 +106,7 @@ class RejectionSampler:
106
# that num_nans is computed before applying penalties and temperature.
107
num_nans = get_num_nans(logits) if self.sampler.compute_nans else None
108
166
- if self.rejection_sample_method == "strict":
167
- sampler_output = self.sampler(logits, input_batch)
168
- logprobs_tensors = sampler_output.logprobs_tensors
169
- sampled, num_sampled = strict_rejection_sample(
170
- sampler_output.sampled_token_ids.view(-1),
171
- draft_sampled,
172
- input_batch.cu_num_logits,
173
- self.num_speculative_steps,
174
- )
175
- elif self.rejection_sample_method == "probabilistic":
176
- assert draft_logits is not None
109
+ if self.rejection_sample_method == "standard":
110
pos = input_batch.positions[input_batch.logits_indices]
111
processed_logits = self.sampler.apply_sampling_params(
112
logits,