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 3, 19:00 UTC
[Model Runner V2] Apply synthetic mode to probabilistic rejection sampler (#41035)
Giancarlo Delfin committed
May 12, 2026 at 15:37 UTC
fe5b4e0fe7c0a6d290be9bb99f3a0f3c2aca9cce
5 files changed
+138
-165
.buildkite/test_areas/model_runner_v2.yaml
+2
-3
@@ -106,13 +106,12 @@ steps:
106
- vllm/v1/worker/gpu/
107
- vllm/v1/worker/gpu_worker.py
108
- tests/v1/spec_decode/test_max_len.py
109
- - tests/v1/spec_decode/test_probabilistic_rejection_sampler_utils.py
110
- - tests/v1/spec_decode/test_synthetic_rejection_sampler_utils.py
109
+ - tests/v1/spec_decode/test_rejection_sampler_utils.py
110
- tests/v1/e2e/spec_decode/test_spec_decode.py
111
commands:
112
- set -x
113
- export VLLM_USE_V2_MODEL_RUNNER=1
114
- pytest -v -s v1/spec_decode/test_max_len.py -k "eagle or mtp"
116
- - pytest -v -s v1/spec_decode/test_probabilistic_rejection_sampler_utils.py
115
+ - pytest -v -s v1/spec_decode/test_rejection_sampler_utils.py
116
- pytest -v -s v1/spec_decode/test_synthetic_rejection_sampler_utils.py
117
- pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "eagle or mtp"
tests/v1/spec_decode/test_rejection_sampler_utils.py
renamed
+71
-4
@@ -6,8 +6,8 @@ import math
6
import pytest
7
import torch
8
9
-from vllm.v1.worker.gpu.spec_decode.probabilistic_rejection_sampler_utils import (
10
- probabilistic_rejection_sample,
9
+from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import (
10
+ rejection_sample,
11
)
12
13
VOCAB_SIZE = 4096
@@ -167,7 +167,7 @@ def test_stochastic_rejection_sample(num_speculative_steps: int, temperature: fl
167
num_trials=num_trials,
168
)
169
170
- sampled, num_sampled = probabilistic_rejection_sample(
170
+ sampled, num_sampled = rejection_sample(
171
**inputs, num_speculative_steps=num_speculative_steps
172
)
173
@@ -201,7 +201,7 @@ def test_greedy_rejection_sample(num_speculative_steps: int):
201
num_trials=num_trials,
202
)
203
204
- sampled, num_sampled = probabilistic_rejection_sample(
204
+ sampled, num_sampled = rejection_sample(
205
**inputs, num_speculative_steps=num_speculative_steps
206
)
207
@@ -213,3 +213,70 @@ def test_greedy_rejection_sample(num_speculative_steps: int):
213
assert (sampled[accepted_mask] == target_argmax).all(), (
214
"Greedy sampling produced tokens that are not the target argmax"
215
)
216
+
217
+
218
+@pytest.mark.parametrize(
219
+ "num_speculative_steps,temperature,unconditional_rates",
220
+ [
221
+ (3, 1.0, [0.9, 0.5, 0.2]),
222
+ (3, 0.0, [0.9, 0.5, 0.2]),
223
+ (3, 1.0, [1.0, 1.0, 1.0]),
224
+ (3, 0.0, [1.0, 1.0, 1.0]),
225
+ (3, 1.0, [0.0, 0.0, 0.0]),
226
+ (3, 0.0, [0.0, 0.0, 0.0]),
227
+ (1, 1.0, [0.7]),
228
+ (1, 0.0, [0.7]),
229
+ ],
230
+)
231
+def test_synthetic_rejection_sample(
232
+ num_speculative_steps: int,
233
+ temperature: float,
234
+ unconditional_rates: list[float],
235
+):
236
+ """
237
+ Verify that synthetic rejection sampling produces the expected
238
+ per-position acceptance rates. The unconditional rate at position i
239
+ is P(all draft steps 0..i accepted) = product(conditional_rates[0:i+1]).
240
+ This is approximately mean(num accepted >= i + 1) over many trials.
241
+ """
242
+ from vllm.v1.spec_decode.utils import unconditional_to_conditional_rates
243
+
244
+ torch.manual_seed(42)
245
+ device = "cuda"
246
+ num_trials = 10 * VOCAB_SIZE
247
+ deviation_tol = 1e-2
248
+
249
+ target_logits_1d = torch.randn(VOCAB_SIZE, device=device, dtype=torch.float32)
250
+ draft_logits_1d = torch.randn(VOCAB_SIZE, device=device, dtype=torch.float32)
251
+
252
+ if temperature > 0:
253
+ target_logits_1d /= temperature
254
+ draft_logits_1d /= temperature
255
+
256
+ inputs = _build_rejection_sample_inputs(
257
+ target_logits_1d,
258
+ draft_logits_1d,
259
+ num_speculative_steps,
260
+ temperature=temperature,
261
+ num_trials=num_trials,
262
+ )
263
+
264
+ conditional_rates = unconditional_to_conditional_rates(unconditional_rates)
265
+ synthetic_conditional_rates = torch.tensor(
266
+ conditional_rates, dtype=torch.float32, device=device
267
+ )
268
+
269
+ _, num_sampled = rejection_sample(
270
+ **inputs,
271
+ num_speculative_steps=num_speculative_steps,
272
+ synthetic_conditional_rates=synthetic_conditional_rates,
273
+ )
274
+
275
+ # num_sampled includes the resampled/bonus token.
276
+ num_accepted = num_sampled - 1
277
+ for i, expected_rate in enumerate(unconditional_rates):
278
+ observed_rate = (num_accepted >= i + 1).float().mean().item()
279
+ assert abs(observed_rate - expected_rate) < deviation_tol, (
280
+ f"Step {i}: observed rate {observed_rate:.4f} deviates from "
281
+ f"expected rate {expected_rate:.4f} by more than {deviation_tol}."
282
+ )
vllm/v1/worker/gpu/spec_decode/rejection_sampler.py
+34
-54
@@ -12,11 +12,8 @@ from vllm.v1.worker.gpu.sample.logprob import compute_topk_logprobs
12
from vllm.v1.worker.gpu.sample.output import SamplerOutput
13
from vllm.v1.worker.gpu.sample.sampler import Sampler
14
from vllm.v1.worker.gpu.sample.states import NO_LOGPROBS
15
-from vllm.v1.worker.gpu.spec_decode.probabilistic_rejection_sampler_utils import (
16
- probabilistic_rejection_sample,
17
-)
18
-from vllm.v1.worker.gpu.spec_decode.synthetic_rejection_sampler_utils import (
19
- synthetic_rejection_sample,
15
+from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import (
16
+ rejection_sample,
17
)
18
19
@@ -101,59 +98,42 @@ class RejectionSampler:
98
input_batch: InputBatch,
99
draft_logits: torch.Tensor | None = None,
100
) -> SamplerOutput:
104
- draft_sampled = input_batch.input_ids[input_batch.logits_indices]
101
# NOTE(woosuk): We intentionally compute num_nans before sampling to make clear
102
# that num_nans is computed before applying penalties and temperature.
103
num_nans = get_num_nans(logits) if self.sampler.compute_nans else None
104
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,
113
- input_batch.expanded_idx_mapping,
114
- input_batch.idx_mapping_np,
115
- pos,
116
- draft_sampled,
117
- input_batch.expanded_local_pos,
118
- )
119
- sampled, num_sampled = probabilistic_rejection_sample(
120
- processed_logits,
121
- draft_logits,
122
- draft_sampled,
123
- input_batch.cu_num_logits,
124
- pos,
125
- input_batch.idx_mapping,
126
- input_batch.expanded_idx_mapping,
127
- input_batch.expanded_local_pos,
128
- self.sampler.sampling_states.temperature.gpu,
129
- self.sampler.sampling_states.seeds.gpu,
130
- self.num_speculative_steps,
131
- )
132
- logprobs_tensors = self._get_logprobs_tensors(
133
- input_batch,
134
- sampled,
135
- num_sampled,
136
- processed_logits
137
- if self.sampler.logprobs_mode == "processed_logprobs"
138
- else logits,
139
- )
140
- elif self.rejection_sample_method == "synthetic":
141
- sampler_output = self.sampler(logits, input_batch)
142
- logprobs_tensors = sampler_output.logprobs_tensors
143
- sampled, num_sampled = synthetic_rejection_sample(
144
- sampler_output.sampled_token_ids.view(-1),
145
- draft_sampled,
146
- input_batch.cu_num_logits,
147
- input_batch.positions[input_batch.logits_indices],
148
- input_batch.idx_mapping,
149
- self.sampler.sampling_states.seeds.gpu,
150
- self.synthetic_conditional_rates,
151
- self.num_speculative_steps,
152
- )
153
- else:
154
- raise ValueError(
155
- f"Unknown rejection sample method: {self.rejection_sample_method}"
156
- )
105
+ draft_sampled = input_batch.input_ids[input_batch.logits_indices]
106
+ pos = input_batch.positions[input_batch.logits_indices]
107
+ processed_logits = self.sampler.apply_sampling_params(
108
+ logits,
109
+ input_batch.expanded_idx_mapping,
110
+ input_batch.idx_mapping_np,
111
+ pos,
112
+ draft_sampled,
113
+ input_batch.expanded_local_pos,
114
+ )
115
+ sampled, num_sampled = rejection_sample(
116
+ processed_logits,
117
+ draft_logits,
118
+ draft_sampled,
119
+ input_batch.cu_num_logits,
120
+ pos,
121
+ input_batch.idx_mapping,
122
+ input_batch.expanded_idx_mapping,
123
+ input_batch.expanded_local_pos,
124
+ self.sampler.sampling_states.temperature.gpu,
125
+ self.sampler.sampling_states.seeds.gpu,
126
+ self.num_speculative_steps,
127
+ self.synthetic_conditional_rates,
128
+ )
129
+ logprobs_tensors = self._get_logprobs_tensors(
130
+ input_batch,
131
+ sampled,
132
+ num_sampled,
133
+ processed_logits
134
+ if self.sampler.logprobs_mode == "processed_logprobs"
135
+ else logits,
136
+ )
137
138
return SamplerOutput(
139
sampled_token_ids=sampled,
vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py
renamed
+31
-9
@@ -154,7 +154,7 @@ def _compute_block_stats_kernel(
154
155
156
@triton.jit
157
-def _probabilistic_rejection_kernel(
157
+def _rejection_kernel(
158
# [num_reqs, num_speculative_steps + 1]
159
sampled_ptr,
160
sampled_stride,
@@ -198,9 +198,12 @@ def _probabilistic_rejection_kernel(
198
seed_ptr,
199
# [num_logits]
200
pos_ptr,
201
+ # [num_speculative_steps]
202
+ synthetic_conditional_rates_ptr,
203
vocab_num_blocks,
204
PADDED_VOCAB_NUM_BLOCKS: tl.constexpr,
205
HAS_DRAFT_LOGITS: tl.constexpr,
206
+ SYNTHETIC_MODE: tl.constexpr,
207
):
208
req_idx = tl.program_id(0)
209
req_state_idx = tl.load(idx_mapping_ptr + req_idx)
@@ -217,7 +220,7 @@ def _probabilistic_rejection_kernel(
220
for i in range(num_tokens - 1):
221
if accepted:
222
logit_idx = start_idx + i
220
- draft_sampled = tl.load(draft_sampled_ptr + logit_idx + 1)
223
+ draft_sampled = tl.load(draft_sampled_ptr + logit_idx + 1).to(tl.int64)
224
if temp == 0.0:
225
# Greedy sampling. Accept IFF draft matches target argmax.
226
# NOTE: Target argmax is stored directly so that resampling
@@ -236,9 +239,19 @@ def _probabilistic_rejection_kernel(
239
target_local_argmax_ptr
240
+ logit_idx * target_local_argmax_stride
241
+ max_target_block_idx
242
+ ).to(tl.int64)
243
+
244
+ if SYNTHETIC_MODE:
245
+ pos = tl.load(pos_ptr + logit_idx)
246
+ u = tl_rand64(seed, pos, includes_zero=False)
247
+ rate = tl.load(synthetic_conditional_rates_ptr + i)
248
+ accepted &= u < rate
249
+ else:
250
+ accepted &= target_argmax == draft_sampled
251
+ tl.store(
252
+ sampled_ptr + req_idx * sampled_stride + i,
253
+ draft_sampled if accepted else target_argmax,
254
)
240
- accepted &= target_argmax == draft_sampled
241
- tl.store(sampled_ptr + req_idx * sampled_stride + i, target_argmax)
255
else:
256
target_logit = tl.load(
257
target_logits_ptr + logit_idx * target_logits_stride + draft_sampled
@@ -275,9 +288,14 @@ def _probabilistic_rejection_kernel(
288
else:
289
# One-hot draft: q(draft_token) = 1, log_q = 0.
290
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
291
+
292
+ if SYNTHETIC_MODE:
293
+ rate = tl.load(synthetic_conditional_rates_ptr + i)
294
+ accepted &= u < rate
295
+ else:
296
+ # Probability ratio test: p(x) > u * q(x)
297
+ # Equivalent log form: log_p(x) > log(u) + log_q(x)
298
+ accepted &= target_log_prob > tl.log(u) + draft_log_prob
299
tl.store(sampled_ptr + req_idx * sampled_stride + i, draft_sampled)
300
rejected_step += accepted
301
tl.store(rejected_steps_ptr + req_idx, rejected_step)
@@ -470,7 +488,7 @@ def _insert_resampled_kernel(
488
)
489
490
473
-def probabilistic_rejection_sample(
491
+def rejection_sample(
492
# [num_logits, V]
493
target_logits: torch.Tensor,
494
# [max_num_reqs, num_speculative_steps, V]
@@ -492,6 +510,8 @@ def probabilistic_rejection_sample(
510
# [max_num_reqs]
511
seed: torch.Tensor,
512
num_speculative_steps: int,
513
+ # [num_speculative_steps]
514
+ synthetic_conditional_rates: torch.Tensor | None = None,
515
) -> tuple[torch.Tensor, torch.Tensor]:
516
num_reqs = cu_num_logits.shape[0] - 1
517
num_logits, vocab_size = target_logits.shape
@@ -557,7 +577,7 @@ def probabilistic_rejection_sample(
577
num_sampled = sampled.new_empty(num_reqs, dtype=torch.int32)
578
target_rejected_logsumexp = target_logits.new_empty(num_reqs, dtype=torch.float32)
579
draft_rejected_logsumexp = target_logits.new_empty(num_reqs, dtype=torch.float32)
560
- _probabilistic_rejection_kernel[(num_reqs,)](
580
+ _rejection_kernel[(num_reqs,)](
581
sampled,
582
sampled.stride(0),
583
num_sampled,
@@ -584,9 +604,11 @@ def probabilistic_rejection_sample(
604
temperature,
605
seed,
606
pos,
607
+ synthetic_conditional_rates,
608
vocab_num_blocks,
609
PADDED_VOCAB_NUM_BLOCKS=padded_vocab_num_blocks,
610
HAS_DRAFT_LOGITS=has_draft_logits,
611
+ SYNTHETIC_MODE=synthetic_conditional_rates is not None,
612
num_warps=1,
613
)
614
vllm/v1/worker/gpu/spec_decode/synthetic_rejection_sampler_utils.py
deleted
-95
@@ -1,95 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-import torch
4
-
5
-from vllm.triton_utils import tl, triton
6
-from vllm.v1.worker.gpu.sample.gumbel import tl_rand64
7
-
8
-
9
-@triton.jit
10
-def _synthetic_rejection_sample_kernel(
11
- # [num_reqs, num_speculative_steps + 1]
12
- sampled_ptr,
13
- sampled_stride,
14
- # [num_reqs]
15
- num_sampled_ptr,
16
- # [num_draft_tokens + num_reqs]
17
- target_sampled_ptr,
18
- # [num_draft_tokens + num_reqs]
19
- input_ids_ptr,
20
- # [num_reqs + 1]
21
- cu_num_logits_ptr,
22
- # [num_logits]
23
- pos_ptr,
24
- # [num_reqs]
25
- idx_mapping_ptr,
26
- # [max_num_reqs]
27
- seeds_ptr,
28
- # [num_speculative_steps]
29
- acceptance_rates_ptr,
30
-):
31
- req_idx = tl.program_id(0)
32
- start_idx = tl.load(cu_num_logits_ptr + req_idx)
33
- end_idx = tl.load(cu_num_logits_ptr + req_idx + 1)
34
- num_tokens = end_idx - start_idx
35
- req_state_idx = tl.load(idx_mapping_ptr + req_idx)
36
- seed = tl.load(seeds_ptr + req_state_idx)
37
-
38
- num_sampled = 0
39
- rejected = False
40
- for i in range(num_tokens - 1):
41
- if not rejected:
42
- logit_idx = start_idx + i
43
- pos = tl.load(pos_ptr + logit_idx)
44
- u = tl_rand64(seed, pos, includes_zero=False)
45
- acceptance_rate = tl.load(acceptance_rates_ptr + i)
46
- if u < acceptance_rate:
47
- sampled = tl.load(input_ids_ptr + logit_idx + 1).to(tl.int64)
48
- else:
49
- sampled = tl.load(target_sampled_ptr + logit_idx)
50
- rejected = True
51
- tl.store(sampled_ptr + req_idx * sampled_stride + i, sampled)
52
- num_sampled += 1
53
- if not rejected:
54
- target_sampled = tl.load(target_sampled_ptr + start_idx + num_tokens - 1)
55
- tl.store(
56
- sampled_ptr + req_idx * sampled_stride + num_tokens - 1, target_sampled
57
- )
58
- num_sampled += 1
59
- tl.store(num_sampled_ptr + req_idx, num_sampled)
60
-
61
-
62
-def synthetic_rejection_sample(
63
- # [num_draft_tokens + num_reqs]
64
- target_sampled: torch.Tensor,
65
- # [num_draft_tokens + num_reqs]
66
- draft_sampled: torch.Tensor,
67
- # [num_reqs + 1]
68
- cu_num_logits: torch.Tensor,
69
- # [num_logits]
70
- pos: torch.Tensor,
71
- # [num_reqs]
72
- idx_mapping: torch.Tensor,
73
- # [max_num_reqs]
74
- seed: torch.Tensor,
75
- # [num_speculative_steps]
76
- acceptance_rates: torch.Tensor,
77
- num_speculative_steps: int,
78
-) -> tuple[torch.Tensor, torch.Tensor]:
79
- num_reqs = cu_num_logits.shape[0] - 1
80
- sampled = target_sampled.new_empty(num_reqs, num_speculative_steps + 1)
81
- num_sampled = target_sampled.new_empty(num_reqs, dtype=torch.int32)
82
- _synthetic_rejection_sample_kernel[(num_reqs,)](
83
- sampled,
84
- sampled.stride(0),
85
- num_sampled,
86
- target_sampled,
87
- draft_sampled,
88
- cu_num_logits,
89
- pos,
90
- idx_mapping,
91
- seed,
92
- acceptance_rates,
93
- num_warps=1,
94
- )
95
- return sampled, num_sampled