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, 17:30 UTC
[Feat] Unified Synthetic Acceptance Rate for V1 and V2 (#40662)
Signed-off-by: Benjamin Chislett <bchislett@nvidia.com> Signed-off-by: Benjamin Chislett <chislett.ben@gmail.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Benjamin Chislett committed
Apr 23, 2026 at 20:48 UTC
626daa2076b00068572f570c2b1567786eeab141
10 files changed
+340
-137
tests/v1/e2e/spec_decode/test_spec_decode.py
+48
@@ -1310,6 +1310,54 @@ def test_dflash_acceptance_rates(dflash_config):
1310
cleanup_dist_env_and_memory()
1311
1312
1313
+@single_gpu_only
1314
+def test_synthetic_acceptance_rate():
1315
+ """Verify that synthetic rejection sampling produces an acceptance
1316
+ length close to the requested mean acceptance length."""
1317
+ num_spec_tokens = 3
1318
+ expected_acceptance_len = 1.875
1319
+ tolerance = 0.15
1320
+
1321
+ spec_llm = LLM(
1322
+ model="meta-llama/Llama-3.2-1B-Instruct",
1323
+ trust_remote_code=True,
1324
+ speculative_config={
1325
+ "method": "eagle3",
1326
+ "model": "nm-testing/Llama3_2_1B_speculator.eagle3",
1327
+ "num_speculative_tokens": num_spec_tokens,
1328
+ "max_model_len": 2048,
1329
+ "rejection_sample_method": "synthetic",
1330
+ "synthetic_acceptance_length": expected_acceptance_len,
1331
+ },
1332
+ max_model_len=2048,
1333
+ enforce_eager=True,
1334
+ disable_log_stats=False,
1335
+ )
1336
+
1337
+ test_prompts = get_test_prompts(mm_enabled=False, num_prompts=50)
1338
+ spec_llm.chat(
1339
+ test_prompts,
1340
+ SamplingParams(temperature=0, max_tokens=64, ignore_eos=True),
1341
+ )
1342
+
1343
+ metrics = spec_llm.get_metrics()
1344
+ acceptance_len = compute_acceptance_len(metrics)
1345
+
1346
+ print(
1347
+ f"Synthetic acceptance length: {acceptance_len:.3f}"
1348
+ f" (expected={expected_acceptance_len:.3f},"
1349
+ f" tolerance=±{tolerance})"
1350
+ )
1351
+ assert abs(acceptance_len - expected_acceptance_len) <= tolerance, (
1352
+ f"Synthetic acceptance length {acceptance_len:.3f} is not within"
1353
+ f" ±{tolerance} of expected {expected_acceptance_len:.3f}"
1354
+ )
1355
+
1356
+ del spec_llm
1357
+ torch.accelerator.empty_cache()
1358
+ cleanup_dist_env_and_memory()
1359
+
1360
+
1361
def test_dflash_correctness(dflash_config):
1362
"""
1363
E2E test for DFlash (block diffusion) speculative decoding.
tests/v1/sample/test_rejection_sampler.py
+61
@@ -933,3 +933,64 @@ def test_sample_recovered_tokens(
933
device=DEVICE_TYPE,
934
)
935
assert torch.equal(recovered_token_ids, ref_recovered_token_ids)
936
+
937
+
938
+########################### Tests for Synthetic Rejection Sampling #########
939
+
940
+
941
+def _make_synthetic_sampler(rates: list[float]) -> RejectionSampler:
942
+ mock_sampler = Mock(spec=Sampler)
943
+ mock_sampler.logprobs_mode = "raw_logprobs"
944
+ spec_config = Mock()
945
+ spec_config.rejection_sample_method = "synthetic"
946
+ spec_config.synthetic_acceptance_rates = rates
947
+ return RejectionSampler(mock_sampler, spec_config, torch.device(DEVICE_TYPE))
948
+
949
+
950
+def _make_sampling_metadata(all_greedy: bool) -> SamplingMetadata:
951
+ temperature = None if all_greedy else torch.tensor([1.0, 1.0], device=DEVICE_TYPE)
952
+ return create_sampling_metadata(all_greedy=all_greedy, temperature=temperature)
953
+
954
+
955
+@pytest.mark.parametrize("all_greedy", [True, False])
956
+def test_synthetic_all_accepted(all_greedy: bool):
957
+ """With all rates=1.0, every draft token is accepted."""
958
+ sampler = _make_synthetic_sampler([1.0, 1.0])
959
+ spec_tokens = [[1, 2], [3]]
960
+ output_tokens = [[10, 20, 50], [30, 40]]
961
+
962
+ metadata = _make_sampling_metadata(all_greedy)
963
+ logits = create_logits_tensor(output_tokens)
964
+ bonus = torch.tensor([50, 40], device=DEVICE_TYPE)
965
+ spec_decode_metadata = create_spec_decode_metadata(spec_tokens, logits)
966
+
967
+ mock_sampler_output(sampler, bonus)
968
+ output = sampler(spec_decode_metadata, None, logits, metadata)
969
+ expected = torch.tensor(
970
+ [[1, 2, 50], [3, 40, PLACEHOLDER_TOKEN_ID]],
971
+ dtype=torch.int,
972
+ device=DEVICE_TYPE,
973
+ )
974
+ assert torch.equal(output.sampled_token_ids, expected)
975
+
976
+
977
+@pytest.mark.parametrize("all_greedy", [True, False])
978
+def test_synthetic_all_rejected(all_greedy: bool):
979
+ """With all rates=0.0, the first token is always rejected."""
980
+ sampler = _make_synthetic_sampler([0.0, 0.0])
981
+ spec_tokens = [[1, 2], [3]]
982
+ output_tokens = [[10, 20, 50], [30, 40]]
983
+
984
+ metadata = _make_sampling_metadata(all_greedy)
985
+ logits = create_logits_tensor(output_tokens)
986
+ bonus = torch.tensor([50, 40], device=DEVICE_TYPE)
987
+ spec_decode_metadata = create_spec_decode_metadata(spec_tokens, logits)
988
+
989
+ mock_sampler_output(sampler, bonus)
990
+ output = sampler(spec_decode_metadata, None, logits, metadata)
991
+ result = output.sampled_token_ids
992
+ # Exactly one token emitted per sequence (the rejection fallback),
993
+ # followed by placeholders.
994
+ for row in result:
995
+ assert row[0] != PLACEHOLDER_TOKEN_ID
996
+ assert (row[1:] == PLACEHOLDER_TOKEN_ID).all()
tests/v1/spec_decode/test_synthetic_rejection_sampler_utils.py
+43
-28
@@ -2,33 +2,48 @@
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,
5
+from vllm.config.speculative import SpeculativeConfig
6
+from vllm.v1.spec_decode.utils import unconditional_to_conditional_rates
7
+
8
+
9
+def test_unconditional_to_conditional_rates_basic():
10
+ # c_0 = p_0; c_i = p_i / p_{i-1}
11
+ assert unconditional_to_conditional_rates([0.9, 0.5, 0.2]) == pytest.approx(
12
+ [0.9, 0.5 / 0.9, 0.2 / 0.5]
13
+ )
14
+
15
+
16
+def test_unconditional_to_conditional_rates_handles_zero():
17
+ # After a zero, subsequent conditional rates are clamped to 0 (the chain
18
+ # has already terminated in the kernel, so these values are unused).
19
+ assert unconditional_to_conditional_rates([1.0, 0.6, 0.0, 0.0]) == pytest.approx(
20
+ [1.0, 0.6, 0.0, 0.0]
21
+ )
22
+
23
+
24
+def test_unconditional_to_conditional_rates_all_ones():
25
+ assert unconditional_to_conditional_rates([1.0, 1.0, 1.0]) == pytest.approx(
26
+ [1.0, 1.0, 1.0]
27
+ )
28
+
29
+
30
+@pytest.mark.parametrize(
31
+ "length,n,expected",
32
+ [
33
+ (2.6, 3, [1.0, 0.6, 0.0]),
34
+ (1.0, 3, [0.0, 0.0, 0.0]),
35
+ (4.0, 3, [1.0, 1.0, 1.0]),
36
+ (2.0, 3, [1.0, 0.0, 0.0]),
37
+ (3.5, 4, [1.0, 1.0, 0.5, 0.0]),
38
+ ],
39
)
40
+def test_acceptance_length_to_rates(length, n, expected):
41
+ assert SpeculativeConfig._acceptance_length_to_rates(length, n) == pytest.approx(
42
+ expected
43
+ )
44
+
45
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
46
+def test_resolve_length_produces_minvariance_schedule():
47
+ assert SpeculativeConfig._resolve_synthetic_acceptance_rates(
48
+ 3, None, 2.6
49
+ ) == pytest.approx([1.0, 0.6, 0.0])
vllm/config/speculative.py
+75
-6
@@ -189,12 +189,64 @@ class SpeculativeConfig:
189
distribution, but the latter yields a higher acceptance rate at the cost
190
of more memory to cache draft logits."""
191
192
- synthetic_acceptance_rate: float | None = None
193
- """Average acceptance rate for synthetic rejection sampling. Draft
194
- tokens are accepted with a position-dependent probability that decays
195
- geometrically, calibrated so that the mean rate across all speculative
196
- positions equals this value. Only used when rejection_sample_method
197
- is 'synthetic'. Must be in [0, 1]."""
192
+ synthetic_acceptance_rates: list[float] | None = None
193
+ """Per-position *unconditional* acceptance rates for synthetic rejection
194
+ sampling. Position i's entry is the marginal probability that the first
195
+ i+1 draft tokens are all accepted; the list must have length
196
+ num_speculative_tokens, each entry in [0, 1], and be monotonically
197
+ non-increasing. Only valid when rejection_sample_method is 'synthetic'.
198
+ Mutually exclusive with synthetic_acceptance_length."""
199
+
200
+ synthetic_acceptance_length: float | None = None
201
+ """Target mean acceptance length for synthetic rejection sampling, in
202
+ [1, num_speculative_tokens + 1]. Resolved internally to
203
+ synthetic_acceptance_rates. Only valid when rejection_sample_method is 'synthetic'.
204
+ Mutually exclusive with synthetic_acceptance_rates."""
205
+
206
+ @staticmethod
207
+ def _acceptance_length_to_rates(length: float, n: int) -> list[float]:
208
+ """Mean acceptance length to unconditional per-position rates, using
209
+ the minimum-variance schedule."""
210
+ num_drafts = length - 1 # expected number of accepted draft tokens
211
+ num_full = int(num_drafts)
212
+ return (
213
+ [1.0] * num_full + [num_drafts - num_full] + [0.0] * (n - num_full - 1)
214
+ )[:n]
215
+
216
+ @staticmethod
217
+ def _resolve_synthetic_acceptance_rates(
218
+ n: int,
219
+ rates: list[float] | None,
220
+ length: float | None,
221
+ ) -> list[float]:
222
+ """Return per-position unconditional acceptance rates from exactly one
223
+ of `rates` or `length` (validates range, length, and monotonicity)."""
224
+ if (rates is None) == (length is None):
225
+ raise ValueError(
226
+ "rejection_sample_method='synthetic' requires exactly one of "
227
+ "synthetic_acceptance_rates or synthetic_acceptance_length."
228
+ )
229
+ if rates is not None:
230
+ if len(rates) != n:
231
+ raise ValueError(
232
+ f"synthetic_acceptance_rates must have length {n}, got {rates}."
233
+ )
234
+ if not all(0.0 <= r <= 1.0 for r in rates):
235
+ raise ValueError(
236
+ f"synthetic_acceptance_rates entries must be in [0, 1], "
237
+ f"got {rates}."
238
+ )
239
+ if any(rates[i] > rates[i - 1] for i in range(1, n)):
240
+ raise ValueError(
241
+ f"synthetic_acceptance_rates must be non-increasing, got {rates}."
242
+ )
243
+ return list(rates)
244
+ assert length is not None
245
+ if not 1.0 <= length <= float(n + 1):
246
+ raise ValueError(
247
+ f"synthetic_acceptance_length must be in [1, {n + 1}], got {length}."
248
+ )
249
+ return SpeculativeConfig._acceptance_length_to_rates(length, n)
250
251
def compute_hash(self) -> str:
252
"""
@@ -818,6 +870,23 @@ class SpeculativeConfig:
870
f"than zero ({self.num_speculative_tokens})."
871
)
872
873
+ if self.rejection_sample_method == "synthetic":
874
+ # Consolidate to per-position rates
875
+ self.synthetic_acceptance_rates = self._resolve_synthetic_acceptance_rates(
876
+ self.num_speculative_tokens,
877
+ self.synthetic_acceptance_rates,
878
+ self.synthetic_acceptance_length,
879
+ )
880
+ self.synthetic_acceptance_length = None
881
+ elif (
882
+ self.synthetic_acceptance_rates is not None
883
+ or self.synthetic_acceptance_length is not None
884
+ ):
885
+ raise ValueError(
886
+ "synthetic_acceptance_rates / synthetic_acceptance_length "
887
+ "are only valid with rejection_sample_method='synthetic'."
888
+ )
889
+
890
if self.draft_model_config:
891
self.draft_model_config.verify_with_parallel_config(
892
self.draft_parallel_config
vllm/v1/sample/rejection_sampler.py
+86
-28
@@ -1,8 +1,11 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+from __future__ import annotations
5
+
6
from collections.abc import Sequence
7
from dataclasses import replace
8
+from typing import TYPE_CHECKING
9
10
import torch
11
import torch.nn as nn
@@ -17,6 +20,10 @@ from vllm.v1.sample.ops.penalties import apply_all_penalties
20
from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p
21
from vllm.v1.sample.sampler import Sampler
22
from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
23
+from vllm.v1.spec_decode.utils import unconditional_to_conditional_rates
24
+
25
+if TYPE_CHECKING:
26
+ from vllm.config.speculative import SpeculativeConfig
27
28
logger = init_logger(__name__)
29
@@ -50,13 +57,33 @@ class RejectionSampler(nn.Module):
57
output tokens = accepted tokens + recovered tokens + bonus tokens
58
"""
59
53
- def __init__(self, sampler: Sampler):
60
+ def __init__(
61
+ self,
62
+ sampler: Sampler,
63
+ spec_config: SpeculativeConfig | None = None,
64
+ device: torch.device | None = None,
65
+ ):
66
super().__init__()
67
self.sampler = sampler
68
logprobs_mode = self.sampler.logprobs_mode
69
self.is_processed_logprobs_mode = logprobs_mode.startswith("processed")
70
self.is_logits_logprobs_mode = logprobs_mode.endswith("logits")
71
72
+ self.synthetic_conditional_rates: torch.Tensor | None = None
73
+ if (
74
+ spec_config is not None
75
+ and spec_config.rejection_sample_method == "synthetic"
76
+ ):
77
+ assert spec_config.synthetic_acceptance_rates is not None
78
+ self.synthetic_conditional_rates = torch.tensor(
79
+ unconditional_to_conditional_rates(
80
+ spec_config.synthetic_acceptance_rates
81
+ ),
82
+ dtype=torch.float32,
83
+ device=device,
84
+ )
85
+ self.synthetic_mode = self.synthetic_conditional_rates is not None
86
+
87
def forward(
88
self,
89
metadata: SpecDecodeMetadata,
@@ -147,6 +174,8 @@ class RejectionSampler(nn.Module):
174
target_logits,
175
bonus_token_ids,
176
sampling_metadata,
177
+ synthetic_mode=self.synthetic_mode,
178
+ synthetic_conditional_rates=self.synthetic_conditional_rates,
179
)
180
181
logprobs_tensors = None
@@ -362,6 +391,8 @@ def rejection_sample(
391
# [batch_size, 1]
392
bonus_token_ids: torch.Tensor,
393
sampling_metadata: SamplingMetadata,
394
+ synthetic_mode: bool = False,
395
+ synthetic_conditional_rates: torch.Tensor | None = None,
396
) -> torch.Tensor:
397
assert draft_token_ids.ndim == 1
398
assert draft_probs is None or draft_probs.ndim == 2
@@ -389,6 +420,20 @@ def rejection_sample(
420
is_greedy = None
421
else:
422
is_greedy = sampling_metadata.temperature == GREEDY_TEMPERATURE
423
+
424
+ # Generate uniform probabilities before either kernel because synthetic
425
+ # mode needs them in the greedy kernel too. Skip only when all requests
426
+ # are greedy *and* synthetic mode is off (the standard fast-path).
427
+ # [num_tokens]
428
+ uniform_probs: torch.Tensor | None = None
429
+ if synthetic_mode or not sampling_metadata.all_greedy:
430
+ uniform_probs = generate_uniform_probs(
431
+ num_tokens,
432
+ num_draft_tokens,
433
+ sampling_metadata.generators,
434
+ device,
435
+ )
436
+
437
if not sampling_metadata.all_random:
438
# Rejection sampling for greedy sampling requests.
439
target_argmax = target_logits.argmax(dim=-1)
@@ -400,6 +445,9 @@ def rejection_sample(
445
bonus_token_ids,
446
is_greedy,
447
max_spec_len,
448
+ uniform_probs,
449
+ synthetic_conditional_rates,
450
+ SYNTHETIC_MODE=synthetic_mode,
451
)
452
if sampling_metadata.all_greedy:
453
return output_token_ids
@@ -408,15 +456,6 @@ def rejection_sample(
456
target_probs = target_logits.softmax(dim=-1, dtype=torch.float32)
457
assert target_probs.is_contiguous()
458
411
- # Generate uniform probabilities for rejection sampling.
412
- # [num_tokens]
413
- uniform_probs = generate_uniform_probs(
414
- num_tokens,
415
- num_draft_tokens,
416
- sampling_metadata.generators,
417
- device,
418
- )
419
-
459
# Sample recovered tokens for each position.
460
# [num_tokens]
461
recovered_token_ids = sample_recovered_tokens(
@@ -431,6 +470,7 @@ def rejection_sample(
470
)
471
472
# Rejection sampling for random sampling requests.
473
+ assert uniform_probs is not None
474
rejection_random_sample_kernel[(batch_size,)](
475
output_token_ids,
476
cu_num_draft_tokens,
@@ -443,7 +483,9 @@ def rejection_sample(
483
is_greedy,
484
max_spec_len,
485
vocab_size,
486
+ synthetic_conditional_rates,
487
NO_DRAFT_PROBS=draft_probs is None,
488
+ SYNTHETIC_MODE=synthetic_mode,
489
)
490
return output_token_ids
491
@@ -658,6 +700,9 @@ def rejection_greedy_sample_kernel(
700
bonus_token_ids_ptr, # [batch_size]
701
is_greedy_ptr, # [batch_size] or None
702
max_spec_len,
703
+ uniform_probs_ptr, # [num_tokens] or None (synthetic mode only)
704
+ synthetic_conditional_rates_ptr, # [num_speculative_tokens] or None
705
+ SYNTHETIC_MODE: tl.constexpr,
706
):
707
req_idx = tl.program_id(0)
708
# FIXME(woosuk): Because is_greedy_ptr is not None at profiling run,
@@ -675,14 +720,20 @@ def rejection_greedy_sample_kernel(
720
for pos in range(num_draft_tokens):
721
if not rejected:
722
draft_token_id = tl.load(draft_token_ids_ptr + start_idx + pos)
678
- target_argmax_id = tl.load(target_argmax_ptr + start_idx + pos)
723
+ target_argmax_id = tl.load(target_argmax_ptr + start_idx + pos).to(tl.int32)
724
+ if SYNTHETIC_MODE:
725
+ uniform_prob = tl.load(uniform_probs_ptr + start_idx + pos)
726
+ rate = tl.load(synthetic_conditional_rates_ptr + pos)
727
+ accepted = uniform_prob < rate
728
+ token_id = draft_token_id if accepted else target_argmax_id
729
+ rejected = not accepted
730
+ else:
731
+ token_id = target_argmax_id
732
+ rejected = draft_token_id != target_argmax_id
733
tl.store(
734
output_token_ids_ptr + req_idx * (max_spec_len + 1) + pos,
681
- target_argmax_id,
735
+ token_id,
736
)
683
- if draft_token_id != target_argmax_id:
684
- # Reject.
685
- rejected = True
737
738
if not rejected:
739
# If all tokens are accepted, append the bonus token.
@@ -707,7 +758,9 @@ def rejection_random_sample_kernel(
758
is_greedy_ptr, # [batch_size]
759
max_spec_len,
760
vocab_size,
761
+ synthetic_conditional_rates_ptr, # [num_speculative_tokens] or None
762
NO_DRAFT_PROBS: tl.constexpr,
763
+ SYNTHETIC_MODE: tl.constexpr,
764
):
765
req_idx = tl.program_id(0)
766
is_greedy = tl.load(is_greedy_ptr + req_idx)
@@ -723,23 +776,28 @@ def rejection_random_sample_kernel(
776
for pos in range(num_draft_tokens):
777
if not rejected:
778
draft_token_id = tl.load(draft_token_ids_ptr + start_idx + pos)
726
- if NO_DRAFT_PROBS:
727
- draft_prob = 1
779
+ uniform_prob = tl.load(uniform_probs_ptr + start_idx + pos)
780
+ if SYNTHETIC_MODE:
781
+ rate = tl.load(synthetic_conditional_rates_ptr + pos)
782
+ accepted = uniform_prob < rate
783
else:
729
- draft_prob = tl.load(
730
- draft_probs_ptr + (start_idx + pos) * vocab_size + draft_token_id
784
+ if NO_DRAFT_PROBS:
785
+ draft_prob = 1
786
+ else:
787
+ draft_prob = tl.load(
788
+ draft_probs_ptr
789
+ + (start_idx + pos) * vocab_size
790
+ + draft_token_id
791
+ )
792
+ target_prob = tl.load(
793
+ target_probs_ptr + (start_idx + pos) * vocab_size + draft_token_id
794
)
732
- target_prob = tl.load(
733
- target_probs_ptr + (start_idx + pos) * vocab_size + draft_token_id
734
- )
735
- uniform_prob = tl.load(uniform_probs_ptr + start_idx + pos)
736
- # NOTE(woosuk): While the draft probability should never be 0,
737
- # we check it to avoid NaNs. If it happens to be 0, we reject.
738
- if draft_prob > 0 and target_prob / draft_prob >= uniform_prob:
739
- # Accept.
795
+ # NOTE(woosuk): While the draft probability should never be 0,
796
+ # we check it to avoid NaNs. If it happens to be 0, we reject.
797
+ accepted = draft_prob > 0 and target_prob / draft_prob >= uniform_prob
798
+ if accepted:
799
token_id = draft_token_id
800
else:
742
- # Reject. Use recovered token.
801
rejected = True
802
token_id = tl.load(recovered_token_ids_ptr + start_idx + pos)
803
tl.store(
vllm/v1/spec_decode/utils.py
+6
@@ -594,3 +594,9 @@ def update_num_computed_tokens_for_batch_change(
594
num_accepted_tokens.copy_(
595
torch.where(participating, valid_counts, num_accepted_tokens)
596
)
597
+
598
+
599
+def unconditional_to_conditional_rates(rates: list[float]) -> list[float]:
600
+ """Convert per-position unconditional rates to per-position conditional
601
+ rates for the early-terminating rejection loop (c_i = p_i / p_{i-1})."""
602
+ return [p / q if q > 0.0 else 0.0 for p, q in zip(rates, [1.0, *rates[:-1]])]
vllm/v1/worker/gpu/model_runner.py
+1
@@ -220,6 +220,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
220
self.rejection_sampler = RejectionSampler(
221
self.sampler,
222
self.speculative_config,
223
+ self.device,
224
)
225
self.prompt_logprobs_worker = PromptLogprobsWorker(self.max_num_reqs)
226
self.structured_outputs_worker = StructuredOutputsWorker(
vllm/v1/worker/gpu/spec_decode/rejection_sampler.py
+11
-16
@@ -5,6 +5,7 @@ import torch
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.spec_decode.utils import unconditional_to_conditional_rates
9
from vllm.v1.worker.gpu.input_batch import InputBatch
10
from vllm.v1.worker.gpu.metrics.logits import get_num_nans
11
from vllm.v1.worker.gpu.sample.logprob import compute_topk_logprobs
@@ -15,7 +16,6 @@ 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 (
18
- compute_synthetic_rejection_sampler_params,
19
synthetic_rejection_sample,
20
)
21
@@ -102,24 +102,20 @@ class RejectionSampler:
102
self,
103
sampler: Sampler,
104
spec_config: SpeculativeConfig,
105
+ device: torch.device,
106
):
107
self.sampler = sampler
108
self.num_speculative_steps = spec_config.num_speculative_tokens
109
self.rejection_sample_method = spec_config.rejection_sample_method
110
+ self.synthetic_conditional_rates: torch.Tensor | None = None
111
if self.rejection_sample_method == "synthetic":
110
- synthetic_acceptance_rate = spec_config.synthetic_acceptance_rate
111
- if (
112
- synthetic_acceptance_rate is None
113
- or not 0.0 <= synthetic_acceptance_rate <= 1.0
114
- ):
115
- raise ValueError(
116
- f"synthetic_acceptance_rate must be in [0, 1], "
117
- f"but got {synthetic_acceptance_rate}"
118
- )
119
- self.base_acceptance_rate, self.decay_factor = (
120
- compute_synthetic_rejection_sampler_params(
121
- synthetic_acceptance_rate, self.num_speculative_steps
122
- )
112
+ assert spec_config.synthetic_acceptance_rates is not None
113
+ self.synthetic_conditional_rates = torch.tensor(
114
+ unconditional_to_conditional_rates(
115
+ spec_config.synthetic_acceptance_rates
116
+ ),
117
+ dtype=torch.float32,
118
+ device=device,
119
)
120
121
def _get_logprobs_tensors(
@@ -218,8 +214,7 @@ class RejectionSampler:
214
input_batch.positions[input_batch.logits_indices],
215
input_batch.idx_mapping,
216
self.sampler.sampling_states.seeds.gpu,
221
- self.base_acceptance_rate,
222
- self.decay_factor,
217
+ self.synthetic_conditional_rates,
218
self.num_speculative_steps,
219
)
220
else:
vllm/v1/worker/gpu/spec_decode/synthetic_rejection_sampler_utils.py
+6
-58
@@ -5,8 +5,6 @@ import torch
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
-
8
9
@triton.jit
10
def _synthetic_rejection_sample_kernel(
@@ -27,8 +25,8 @@ def _synthetic_rejection_sample_kernel(
25
idx_mapping_ptr,
26
# [max_num_reqs]
27
seeds_ptr,
30
- base_acceptance_rate,
31
- decay_factor,
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)
@@ -38,13 +36,13 @@ def _synthetic_rejection_sample_kernel(
36
seed = tl.load(seeds_ptr + req_state_idx)
37
38
num_sampled = 0
41
- acceptance_rate = base_acceptance_rate
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:
@@ -52,7 +50,6 @@ def _synthetic_rejection_sample_kernel(
50
rejected = True
51
tl.store(sampled_ptr + req_idx * sampled_stride + i, sampled)
52
num_sampled += 1
55
- acceptance_rate *= decay_factor
53
if not rejected:
54
target_sampled = tl.load(target_sampled_ptr + start_idx + num_tokens - 1)
55
tl.store(
@@ -75,8 +72,8 @@ def synthetic_rejection_sample(
72
idx_mapping: torch.Tensor,
73
# [max_num_reqs]
74
seed: torch.Tensor,
78
- base_acceptance_rate: float,
79
- decay_factor: float,
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
@@ -92,56 +89,7 @@ def synthetic_rejection_sample(
89
pos,
90
idx_mapping,
91
seed,
95
- base_acceptance_rate,
96
- decay_factor,
92
+ acceptance_rates,
93
num_warps=1,
94
)
95
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
vllm/v1/worker/gpu_model_runner.py
+3
-1
@@ -577,7 +577,9 @@ class GPUModelRunner(
577
"Unknown speculative decoding method: "
578
f"{self.speculative_config.method}"
579
)
580
- self.rejection_sampler = RejectionSampler(self.sampler)
580
+ self.rejection_sampler = RejectionSampler(
581
+ self.sampler, self.speculative_config, self.device
582
+ )
583
584
self.num_spec_tokens = 0
585
self.valid_sampled_token_count_gpu: torch.Tensor | None = None