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 5, 23:30 UTC
[Model Runner V2] Enable forcing a specific acceptance rate during rejection sampling (#38045)
Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai>
Giancarlo Delfin committed
Mar 26, 2026 at 13:38 UTC
c32e97602d54f34fb60ff956dae09f9a557ec50a
7 files changed
+240
-18
.buildkite/test_areas/model_runner_v2.yaml
+2
@@ -101,9 +101,11 @@ steps:
101
- vllm/v1/worker/gpu/
102
- vllm/v1/worker/gpu_worker.py
103
- tests/v1/spec_decode/test_max_len.py
104
+ - tests/v1/spec_decode/test_synthetic_rejection_sampler_utils.py
105
- tests/v1/e2e/spec_decode/test_spec_decode.py
106
commands:
107
- set -x
108
- export VLLM_USE_V2_MODEL_RUNNER=1
109
- pytest -v -s v1/spec_decode/test_max_len.py -k "eagle or mtp"
110
+ - pytest -v -s v1/spec_decode/test_synthetic_rejection_sampler_utils.py
111
- pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "eagle or mtp"
tests/v1/spec_decode/test_synthetic_rejection_sampler_utils.py
new
+34
@@ -0,0 +1,34 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import pytest
4
+
5
+from vllm.v1.worker.gpu.spec_decode.synthetic_rejection_sampler_utils import (
6
+ compute_synthetic_rejection_sampler_params,
7
+)
8
+
9
+NUM_SPECULATIVE_STEPS = [1, 2, 3, 4, 5, 7, 10]
10
+ACCEPTANCE_RATES = [i / 100 for i in range(0, 100)]
11
+
12
+
13
+@pytest.mark.parametrize("num_speculative_steps", NUM_SPECULATIVE_STEPS)
14
+def test_compute_synthetic_rejection_sampler_params(num_speculative_steps: int):
15
+ """Test that the base acceptance rate and decay factor generated for
16
+ synthetic rejection sampling have a mean joint acceptance probability
17
+ that matches the desired acceptance rate."""
18
+ tol = 1e-9
19
+ for desired_acceptance_rate in ACCEPTANCE_RATES:
20
+ base_rate, decay_factor = compute_synthetic_rejection_sampler_params(
21
+ desired_acceptance_rate, num_speculative_steps, tol=tol
22
+ )
23
+
24
+ # Compute the mean of joint acceptance probabilities across
25
+ # all speculative positions.
26
+ joint_prob = 1.0
27
+ mean_joint = 0.0
28
+ for i in range(num_speculative_steps):
29
+ joint_prob *= base_rate * decay_factor**i
30
+ mean_joint += joint_prob
31
+ mean_joint /= num_speculative_steps
32
+
33
+ assert abs(desired_acceptance_rate - mean_joint) < 10 * tol
34
+ assert base_rate <= 1.0
vllm/config/speculative.py
+8
-1
@@ -58,7 +58,7 @@ SpeculativeMethod = Literal[
58
EagleModelTypes,
59
NgramGPUTypes,
60
]
61
-RejectionSampleMethod = Literal["strict", "probabilistic"]
61
+RejectionSampleMethod = Literal["strict", "probabilistic", "synthetic"]
62
63
64
@config
@@ -184,6 +184,13 @@ class SpeculativeConfig:
184
distribution, but the latter yields a higher acceptance rate at the cost
185
of more memory to cache draft logits."""
186
187
+ synthetic_acceptance_rate: float | None = None
188
+ """Average acceptance rate for synthetic rejection sampling. Draft
189
+ tokens are accepted with a position-dependent probability that decays
190
+ geometrically, calibrated so that the mean rate across all speculative
191
+ positions equals this value. Only used when rejection_sample_method
192
+ is 'synthetic'. Must be in [0, 1]."""
193
+
194
def compute_hash(self) -> str:
195
"""
196
WARNING: Whenever a new field is added to this config,
vllm/v1/worker/gpu/model_runner.py
+5
-9
@@ -163,12 +163,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
163
self.speculator = None
164
self.num_speculative_steps = 0
165
self.use_aux_hidden_state_outputs = False
166
- use_strict_rejection_sampling = False
166
if self.speculative_config is not None:
167
self.num_speculative_steps = self.speculative_config.num_speculative_tokens
169
- use_strict_rejection_sampling = (
170
- self.speculative_config.rejection_sample_method == "strict"
171
- )
168
169
if self.is_last_pp_rank:
170
self.speculator = init_speculator(self.vllm_config, self.device)
@@ -217,11 +213,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):
213
logprobs_mode=self.model_config.logprobs_mode,
214
num_speculative_tokens=self.num_speculative_steps + 1,
215
)
220
- self.rejection_sampler = RejectionSampler(
221
- self.sampler,
222
- num_speculative_steps=self.num_speculative_steps,
223
- use_strict_rejection_sampling=use_strict_rejection_sampling,
224
- )
216
+ if self.speculative_config is not None:
217
+ self.rejection_sampler = RejectionSampler(
218
+ self.sampler,
219
+ self.speculative_config,
220
+ )
221
self.prompt_logprobs_worker = PromptLogprobsWorker(self.max_num_reqs)
222
self.structured_outputs_worker = StructuredOutputsWorker(
223
max_num_logits=self.max_num_reqs * (self.num_speculative_steps + 1),
vllm/v1/worker/gpu/spec_decode/eagle/speculator.py
+1
-2
@@ -85,9 +85,8 @@ class EagleSpeculator:
85
self.max_num_tokens, self.hidden_size, dtype=self.dtype, device=device
86
)
87
88
- cache_draft_logits = self.speculative_config.rejection_sample_method != "strict"
88
self.draft_logits: torch.Tensor | None = None
90
- if cache_draft_logits:
89
+ if self.speculative_config.rejection_sample_method == "probabilistic":
90
self.draft_logits = torch.zeros(
91
self.max_num_reqs,
92
self.num_speculative_steps,
vllm/v1/worker/gpu/spec_decode/rejection_sampler.py
+43
-6
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import torch
4
5
+from vllm.config import SpeculativeConfig
6
from vllm.triton_utils import tl, triton
7
from vllm.v1.outputs import LogprobsTensors
8
from vllm.v1.worker.gpu.input_batch import InputBatch
@@ -11,6 +12,10 @@ 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.synthetic_rejection_sampler_utils import (
16
+ compute_synthetic_rejection_sampler_params,
17
+ synthetic_rejection_sample,
18
+)
19
20
21
@triton.jit
@@ -445,12 +450,26 @@ class RejectionSampler:
450
def __init__(
451
self,
452
sampler: Sampler,
448
- num_speculative_steps,
449
- use_strict_rejection_sampling: bool = True,
453
+ spec_config: SpeculativeConfig,
454
):
455
self.sampler = sampler
452
- self.num_speculative_steps = num_speculative_steps
453
- self.use_strict_rejection_sampling = use_strict_rejection_sampling
456
+ self.num_speculative_steps = spec_config.num_speculative_tokens
457
+ self.rejection_sample_method = spec_config.rejection_sample_method
458
+ if self.rejection_sample_method == "synthetic":
459
+ synthetic_acceptance_rate = spec_config.synthetic_acceptance_rate
460
+ if (
461
+ synthetic_acceptance_rate is None
462
+ or not 0.0 <= synthetic_acceptance_rate <= 1.0
463
+ ):
464
+ raise ValueError(
465
+ f"synthetic_acceptance_rate must be in [0, 1], "
466
+ f"but got {synthetic_acceptance_rate}"
467
+ )
468
+ self.base_acceptance_rate, self.decay_factor = (
469
+ compute_synthetic_rejection_sampler_params(
470
+ synthetic_acceptance_rate, self.num_speculative_steps
471
+ )
472
+ )
473
474
def _get_logprobs_tensors(
475
self,
@@ -497,7 +516,7 @@ class RejectionSampler:
516
# that num_nans is computed before applying penalties and temperature.
517
num_nans = get_num_nans(logits) if self.sampler.compute_nans else None
518
500
- if self.use_strict_rejection_sampling:
519
+ if self.rejection_sample_method == "strict":
520
sampler_output = self.sampler(logits, input_batch)
521
logprobs_tensors = sampler_output.logprobs_tensors
522
sampled, num_sampled = strict_rejection_sample(
@@ -506,7 +525,7 @@ class RejectionSampler:
525
input_batch.cu_num_logits,
526
self.num_speculative_steps,
527
)
509
- else:
528
+ elif self.rejection_sample_method == "probabilistic":
529
assert draft_logits is not None
530
pos = input_batch.positions[input_batch.logits_indices]
531
processed_logits = self.sampler.apply_sampling_params(
@@ -538,6 +557,24 @@ class RejectionSampler:
557
if self.sampler.logprobs_mode == "processed_logprobs"
558
else logits,
559
)
560
+ elif self.rejection_sample_method == "synthetic":
561
+ sampler_output = self.sampler(logits, input_batch)
562
+ logprobs_tensors = sampler_output.logprobs_tensors
563
+ sampled, num_sampled = synthetic_rejection_sample(
564
+ sampler_output.sampled_token_ids.view(-1),
565
+ draft_sampled,
566
+ input_batch.cu_num_logits,
567
+ input_batch.positions[input_batch.logits_indices],
568
+ input_batch.idx_mapping,
569
+ self.sampler.sampling_states.seeds.gpu,
570
+ self.base_acceptance_rate,
571
+ self.decay_factor,
572
+ self.num_speculative_steps,
573
+ )
574
+ else:
575
+ raise ValueError(
576
+ f"Unknown rejection sample method: {self.rejection_sample_method}"
577
+ )
578
579
return SamplerOutput(
580
sampled_token_ids=sampled,
vllm/v1/worker/gpu/spec_decode/synthetic_rejection_sampler_utils.py
new
+147
@@ -0,0 +1,147 @@
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
+MIN_ACCEPTANCE_DECAY_FACTOR = 0.85
9
+
10
+
11
+@triton.jit
12
+def _synthetic_rejection_sample_kernel(
13
+ # [num_reqs, num_speculative_steps + 1]
14
+ sampled_ptr,
15
+ sampled_stride,
16
+ # [num_reqs]
17
+ num_sampled_ptr,
18
+ # [num_draft_tokens + num_reqs]
19
+ target_sampled_ptr,
20
+ # [num_draft_tokens + num_reqs]
21
+ input_ids_ptr,
22
+ # [num_reqs + 1]
23
+ cu_num_logits_ptr,
24
+ # [num_logits]
25
+ pos_ptr,
26
+ # [num_reqs]
27
+ idx_mapping_ptr,
28
+ # [max_num_reqs]
29
+ seeds_ptr,
30
+ base_acceptance_rate,
31
+ decay_factor,
32
+):
33
+ req_idx = tl.program_id(0)
34
+ start_idx = tl.load(cu_num_logits_ptr + req_idx)
35
+ end_idx = tl.load(cu_num_logits_ptr + req_idx + 1)
36
+ num_tokens = end_idx - start_idx
37
+ req_state_idx = tl.load(idx_mapping_ptr + req_idx)
38
+ seed = tl.load(seeds_ptr + req_state_idx)
39
+
40
+ num_sampled = 0
41
+ acceptance_rate = base_acceptance_rate
42
+ rejected = False
43
+ for i in range(num_tokens - 1):
44
+ if not rejected:
45
+ logit_idx = start_idx + i
46
+ pos = tl.load(pos_ptr + logit_idx)
47
+ u = tl_rand64(seed, pos, includes_zero=False)
48
+ if u < acceptance_rate:
49
+ sampled = tl.load(input_ids_ptr + logit_idx + 1).to(tl.int64)
50
+ else:
51
+ sampled = tl.load(target_sampled_ptr + logit_idx)
52
+ rejected = True
53
+ tl.store(sampled_ptr + req_idx * sampled_stride + i, sampled)
54
+ num_sampled += 1
55
+ acceptance_rate *= decay_factor
56
+ if not rejected:
57
+ target_sampled = tl.load(target_sampled_ptr + start_idx + num_tokens - 1)
58
+ tl.store(
59
+ sampled_ptr + req_idx * sampled_stride + num_tokens - 1, target_sampled
60
+ )
61
+ num_sampled += 1
62
+ tl.store(num_sampled_ptr + req_idx, num_sampled)
63
+
64
+
65
+def synthetic_rejection_sample(
66
+ # [num_draft_tokens + num_reqs]
67
+ target_sampled: torch.Tensor,
68
+ # [num_draft_tokens + num_reqs]
69
+ draft_sampled: torch.Tensor,
70
+ # [num_reqs + 1]
71
+ cu_num_logits: torch.Tensor,
72
+ # [num_logits]
73
+ pos: torch.Tensor,
74
+ # [num_reqs]
75
+ idx_mapping: torch.Tensor,
76
+ # [max_num_reqs]
77
+ seed: torch.Tensor,
78
+ base_acceptance_rate: float,
79
+ decay_factor: float,
80
+ num_speculative_steps: int,
81
+) -> tuple[torch.Tensor, torch.Tensor]:
82
+ num_reqs = cu_num_logits.shape[0] - 1
83
+ sampled = target_sampled.new_empty(num_reqs, num_speculative_steps + 1)
84
+ num_sampled = target_sampled.new_empty(num_reqs, dtype=torch.int32)
85
+ _synthetic_rejection_sample_kernel[(num_reqs,)](
86
+ sampled,
87
+ sampled.stride(0),
88
+ num_sampled,
89
+ target_sampled,
90
+ draft_sampled,
91
+ cu_num_logits,
92
+ pos,
93
+ idx_mapping,
94
+ seed,
95
+ base_acceptance_rate,
96
+ decay_factor,
97
+ num_warps=1,
98
+ )
99
+ return sampled, num_sampled
100
+
101
+
102
+def compute_synthetic_rejection_sampler_params(
103
+ p_avg: float, n: int, tol: float = 1e-9
104
+) -> tuple[float, float]:
105
+ def mean_joint_prob(a_0: float, gamma: float, n: int):
106
+ total = 0.0
107
+ for i in range(n):
108
+ total += a_0 ** (i + 1) * gamma ** (i * (i + 1) // 2)
109
+ return total / n
110
+
111
+ def min_valid_decay_factor(p: float, n: int, tol: float = 1e-9) -> float:
112
+ low, high = MIN_ACCEPTANCE_DECAY_FACTOR, 1.0
113
+ if mean_joint_prob(1, low, n) >= p:
114
+ return low
115
+
116
+ # Sweep for a gamma decay factor that is guaranteed
117
+ # to yield a base acceptance rate <= 1.
118
+ while (high - low) > tol:
119
+ mid = (low + high) / 2
120
+ if mean_joint_prob(1, mid, n) >= p:
121
+ high = mid
122
+ else:
123
+ low = mid
124
+ return high
125
+
126
+ def compute_base_acceptance_rate(
127
+ p_avg: float, gamma: float, n: int, tol: float = 1e-9
128
+ ) -> float:
129
+ if p_avg <= 0.0:
130
+ return 0.0
131
+ if p_avg >= 1.0:
132
+ return 1.0
133
+
134
+ # Sweep for a base acceptance rate that yields
135
+ # the desired mean joint probability.
136
+ low, high = 0.0, 1.0
137
+ while (high - low) > tol:
138
+ mid = (low + high) / 2
139
+ if mean_joint_prob(mid, gamma, n) >= p_avg:
140
+ high = mid
141
+ else:
142
+ low = mid
143
+ return high
144
+
145
+ decay_factor = min_valid_decay_factor(p_avg, n)
146
+ base_rate = compute_base_acceptance_rate(p_avg, decay_factor, n)
147
+ return base_rate, decay_factor