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 14, 05:30 UTC
[V1] Logit processors for rejection sampler (#19482)
Signed-off-by: southfreebird <yvorott@gmail.com> Signed-off-by: Sergei Skvortsov <sergeyskv@nebius.com> Signed-off-by: Sergei Skvortsov <yvorott@gmail.com> Co-authored-by: Sergei Skvortsov <sergeyskv@nebius.com> Co-authored-by: Nick Hill <nhill@redhat.com>
Sergei Skvortsov committed
Oct 7, 2025 at 21:02 UTC
6ebaf43ee4a6fbbeba685315d605536db1c0c471
12 files changed
+468
-89
tests/v1/logits_processors/test_custom_offline.py
+38
-21
@@ -2,7 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import random
4
import sys
5
-from typing import Union
5
+from typing import Any, Union
6
7
import pytest
8
@@ -25,6 +25,7 @@ from tests.v1.logits_processors.utils import entry_points as fake_entry_points
25
from vllm import LLM, SamplingParams
26
from vllm.v1.sample.logits_processor import (
27
STR_POOLING_REJECTS_LOGITSPROCS,
28
+ STR_SPEC_DEC_REJECTS_LOGITSPROCS,
29
LogitsProcessor,
30
)
31
@@ -205,6 +206,7 @@ def test_custom_logitsprocs_req(monkeypatch):
206
207
208
@create_new_process_for_each_test()
209
+@pytest.mark.parametrize("model_scenario", ["pooling", "spec_dec"])
210
@pytest.mark.parametrize(
211
"logitproc_source",
212
[
@@ -213,11 +215,12 @@ def test_custom_logitsprocs_req(monkeypatch):
215
CustomLogitprocSource.LOGITPROC_SOURCE_CLASS,
216
],
217
)
216
-def test_pooling_rejects_custom_logitsprocs(
217
- monkeypatch, logitproc_source: CustomLogitprocSource
218
+def test_rejects_custom_logitsprocs(
219
+ monkeypatch, model_scenario: str, logitproc_source: CustomLogitprocSource
220
):
221
"""Validate that vLLM engine initialization properly rejects custom
220
- logitsprocs when the model is a pooling model.
222
+ logitsprocs when the model is a pooling model or speculative decoding
223
+ enabled.
224
225
Use `LLM` entrypoint. We expect `LLM` initialization to fail before the
226
logitproc is actually loaded.
@@ -241,8 +244,32 @@ def test_pooling_rejects_custom_logitsprocs(
244
monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
245
random.seed(40)
246
247
+ test_params: dict[str, dict[str, Any]] = {
248
+ "pooling": {
249
+ "runner": "pooling",
250
+ "model": POOLING_MODEL_NAME,
251
+ "error_message": STR_POOLING_REJECTS_LOGITSPROCS,
252
+ "speculative_config": None,
253
+ },
254
+ "spec_dec": {
255
+ "runner": "auto",
256
+ "model": MODEL_NAME,
257
+ "error_message": STR_SPEC_DEC_REJECTS_LOGITSPROCS,
258
+ "speculative_config": {"model": "ngram", "num_speculative_tokens": 1},
259
+ },
260
+ }
261
+
262
+ config = test_params[model_scenario]
263
+
264
+ llm_kwargs: dict[str, Any] = {
265
+ "runner": config["runner"],
266
+ "model": config["model"],
267
+ "gpu_memory_utilization": 0.1,
268
+ "speculative_config": config["speculative_config"],
269
+ }
270
+
271
if logitproc_source == CustomLogitprocSource.LOGITPROC_SOURCE_ENTRYPOINT:
245
- # Scenario: vLLM loads a pooling model and ignores a logitproc that is
272
+ # Scenario: vLLM loads a model and ignores a logitproc that is
273
# available at a preconfigured entrypoint
274
275
# Patch in dummy logitproc entrypoint
@@ -254,30 +281,20 @@ def test_pooling_rejects_custom_logitsprocs(
281
# although they should ignore the entrypoint patch anyway
282
monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "fork")
283
257
- llm = LLM(
258
- runner="pooling",
259
- model=POOLING_MODEL_NAME,
260
- gpu_memory_utilization=0.1,
261
- )
284
+ llm = LLM(**llm_kwargs)
285
# Require that no logitsprocs have been loaded
286
worker = llm.llm_engine.model_executor.driver_worker.worker
287
assert sum([1 for _ in worker.model_runner.input_batch.logitsprocs.all]) == 0
288
return
289
267
- kwargs: dict[str, list[Union[str, type[LogitsProcessor]]]] = {}
290
if logitproc_source == CustomLogitprocSource.LOGITPROC_SOURCE_FQCN:
291
# Scenario: load logitproc based on fully-qualified class name (FQCN)
270
- kwargs["logits_processors"] = [DUMMY_LOGITPROC_FQCN]
292
+ llm_kwargs["logits_processors"] = [DUMMY_LOGITPROC_FQCN]
293
elif logitproc_source == CustomLogitprocSource.LOGITPROC_SOURCE_CLASS:
294
# Scenario: load logitproc from provided class object
273
- kwargs["logits_processors"] = [DummyLogitsProcessor]
295
+ llm_kwargs["logits_processors"] = [DummyLogitsProcessor]
296
275
- with pytest.raises(ValueError, match=STR_POOLING_REJECTS_LOGITSPROCS):
276
- # Require that loading a pooling model alongside the logitproc raises
297
+ with pytest.raises(ValueError, match=config["error_message"]):
298
+ # Require that loading a model alongside the logitproc raises
299
# the appropriate exception.
278
- LLM(
279
- runner="pooling",
280
- model=POOLING_MODEL_NAME,
281
- gpu_memory_utilization=0.1,
282
- **kwargs,
283
- )
300
+ LLM(**llm_kwargs)
tests/v1/sample/test_rejection_sampler.py
+171
-9
@@ -6,6 +6,7 @@ import pytest
6
import torch
7
import torch.nn.functional as F
8
9
+from tests.v1.sample.utils import create_allowed_token_ids
10
from vllm.platforms import current_platform
11
from vllm.v1.sample.logits_processor import LogitsProcessors
12
from vllm.v1.sample.metadata import SamplingMetadata
@@ -21,7 +22,9 @@ def rejection_sampler():
22
23
24
def create_logits_tensor(
24
- output_token_ids: list[list[int]], vocab_size: int = 100
25
+ output_token_ids: list[list[int]],
26
+ vocab_size: int = 100,
27
+ token_idx_to_override: Optional[int] = None,
28
) -> torch.Tensor:
29
"""Helper function to create logits tensor that
30
will produce desired token ids on argmax"""
@@ -33,15 +36,25 @@ def create_logits_tensor(
36
for j, token_id in enumerate(tokens):
37
logits[start_loc + j, token_id] = 100.0
38
start_loc += len(tokens)
39
+ if token_idx_to_override:
40
+ logits[:, token_idx_to_override] = 99.0
41
return logits
42
43
44
def create_sampling_metadata(
45
all_greedy: bool,
46
+ output_token_ids: Optional[list[list[int]]] = None,
47
+ prompt_token_ids: Optional[torch.Tensor] = None,
48
+ spec_token_ids: Optional[torch.Tensor] = None,
49
temperature: Optional[torch.Tensor] = None,
50
top_k: Optional[torch.Tensor] = None,
51
top_p: Optional[torch.Tensor] = None,
52
generators: Optional[dict[int, Any]] = None,
53
+ frequency_penalties: Optional[list[float]] = None,
54
+ presence_penalties: Optional[list[float]] = None,
55
+ repetition_penalties: Optional[list[float]] = None,
56
+ bad_words_token_ids: Optional[dict[int, list[list[int]]]] = None,
57
+ allowed_token_ids_mask: Optional[torch.Tensor] = None,
58
) -> SamplingMetadata:
59
"""Create a v1 sampling metadata object with all_greedy set
60
to the given value. Either all greedy or all random sampling
@@ -53,6 +66,21 @@ def create_sampling_metadata(
66
else:
67
assert temperature is not None
68
69
+ if any([frequency_penalties, presence_penalties, repetition_penalties]):
70
+ no_penalties = False
71
+
72
+ assert output_token_ids
73
+ assert len(output_token_ids) > 0
74
+
75
+ frequency_penalties = torch.tensor(frequency_penalties, device=DEVICE)
76
+ presence_penalties = torch.tensor(presence_penalties, device=DEVICE)
77
+ repetition_penalties = torch.tensor(repetition_penalties, device=DEVICE)
78
+ else:
79
+ no_penalties = True
80
+ frequency_penalties = torch.tensor([])
81
+ presence_penalties = torch.tensor([])
82
+ repetition_penalties = torch.tensor([])
83
+
84
return SamplingMetadata(
85
temperature=temperature,
86
all_greedy=all_greedy,
@@ -61,14 +89,15 @@ def create_sampling_metadata(
89
top_k=top_k,
90
generators=generators,
91
max_num_logprobs=0,
64
- no_penalties=False,
65
- prompt_token_ids=None,
66
- frequency_penalties=torch.tensor([]),
67
- presence_penalties=torch.tensor([]),
68
- repetition_penalties=torch.tensor([]),
69
- output_token_ids=[],
70
- allowed_token_ids_mask=None,
71
- bad_words_token_ids={},
92
+ no_penalties=no_penalties,
93
+ prompt_token_ids=prompt_token_ids,
94
+ frequency_penalties=frequency_penalties,
95
+ presence_penalties=presence_penalties,
96
+ repetition_penalties=repetition_penalties,
97
+ output_token_ids=[] if output_token_ids is None else output_token_ids,
98
+ spec_token_ids=[] if spec_token_ids is None else spec_token_ids,
99
+ allowed_token_ids_mask=allowed_token_ids_mask,
100
+ bad_words_token_ids={} if bad_words_token_ids is None else bad_words_token_ids,
101
logitsprocs=LogitsProcessors(),
102
)
103
@@ -611,3 +640,136 @@ def test_top_p(rejection_sampler, top_p):
640
unmasked_indices=top_p_indices,
641
sampling_metadata=sampling_metadata,
642
)
643
+
644
+
645
+########################### Tests for Logit Processors ###################
646
+def test_frequency_penalties(rejection_sampler):
647
+ """Test rejection sampling with frequency penalties"""
648
+ spec_tokens = [[1, 1, 1], [], [1, 1, 1]]
649
+ output_tokens = [[1, 1, 1, 1], [7], [1, 1, 1, 1]] # 1, 7 and 1 are the bonus tokens
650
+
651
+ num_requsts = len(spec_tokens)
652
+ logits = create_logits_tensor(output_tokens, token_idx_to_override=15)
653
+ metadata = create_sampling_metadata(
654
+ all_greedy=True,
655
+ output_token_ids=[[2], [3], [4]],
656
+ spec_token_ids=spec_tokens,
657
+ prompt_token_ids=torch.tensor([[5, 6, 7], [6, 7, 8], [7, 8, 9]], device=DEVICE),
658
+ frequency_penalties=[1.5, 1.5, 0.7],
659
+ presence_penalties=[0.0] * num_requsts,
660
+ repetition_penalties=[1.0] * num_requsts,
661
+ )
662
+ bonus_token_tensor = torch.tensor(
663
+ [output_tokens[i][-1] for i in range(len(output_tokens))], device=logits.device
664
+ )
665
+ spec_decode_metadata = SpecDecodeMetadata.make_dummy(
666
+ spec_tokens, device=logits.device
667
+ )
668
+ output = rejection_sampler(
669
+ spec_decode_metadata,
670
+ draft_probs=None,
671
+ target_logits=logits,
672
+ bonus_token_ids=bonus_token_tensor,
673
+ sampling_metadata=metadata,
674
+ )
675
+ expected = torch.tensor(
676
+ [[1, 15, -1, -1], [7, -1, -1, -1], [1, 1, 15, -1]],
677
+ dtype=torch.int,
678
+ device=logits.device,
679
+ )
680
+ assert torch.equal(output, expected)
681
+
682
+
683
+def test_bad_words(rejection_sampler):
684
+ """Test rejection sampling with bad words constraints"""
685
+ spec_tokens = [[1, 2, 3], [1, 15, 3], [1, 2, 3]]
686
+ output_tokens = [[1, 2, 3, 4], [1, 2, 3, 4], [1, 2, 3, 4]]
687
+
688
+ logits = create_logits_tensor(output_tokens, token_idx_to_override=15)
689
+ metadata = create_sampling_metadata(
690
+ all_greedy=True,
691
+ output_token_ids=[[2], [3], [4]],
692
+ spec_token_ids=spec_tokens,
693
+ bad_words_token_ids={
694
+ 0: [
695
+ [
696
+ 2,
697
+ ]
698
+ ],
699
+ 1: [
700
+ [
701
+ 2,
702
+ ]
703
+ ],
704
+ # Do not apply bad words to the last request
705
+ },
706
+ )
707
+ bonus_token_tensor = torch.tensor(
708
+ [output_tokens[i][-1] for i in range(len(output_tokens))], device=logits.device
709
+ )
710
+ spec_decode_metadata = SpecDecodeMetadata.make_dummy(
711
+ spec_tokens, device=logits.device
712
+ )
713
+ output = rejection_sampler(
714
+ spec_decode_metadata,
715
+ draft_probs=None,
716
+ target_logits=logits,
717
+ bonus_token_ids=bonus_token_tensor,
718
+ sampling_metadata=metadata,
719
+ )
720
+
721
+ expected = torch.tensor(
722
+ [[1, 15, -1, -1], [1, 15, 3, 4], [1, 2, 3, 4]],
723
+ dtype=torch.int,
724
+ device=logits.device,
725
+ )
726
+ assert torch.equal(output, expected)
727
+
728
+
729
+def test_allowed_token_ids(rejection_sampler):
730
+ """Test rejection sampling with allowed token ids"""
731
+ spec_tokens = [[1, 2, 10], [10, 5, 3], [7, 10, 12]]
732
+ output_tokens = [[1, 2, 10, 5], [10, 5, 10, 5], [7, 10, 12, 5]]
733
+ # Not allowed tokens:
734
+ # 0: 0-4
735
+ # 1: 1-5
736
+ # 2: 2-6
737
+ num_allowed_token_ids = 5
738
+
739
+ # Use the token 15 as the sampler choose if a token rejected
740
+ logits = create_logits_tensor(output_tokens, token_idx_to_override=15)
741
+
742
+ batch_size = len(output_tokens)
743
+ _, vocab_size = logits.size()
744
+ mask = create_allowed_token_ids(
745
+ batch_size=batch_size,
746
+ vocab_size=vocab_size,
747
+ num_allowed_token_ids=num_allowed_token_ids,
748
+ device=logits.device,
749
+ )
750
+ metadata = create_sampling_metadata(
751
+ all_greedy=True,
752
+ output_token_ids=[[], [], []],
753
+ spec_token_ids=spec_tokens,
754
+ allowed_token_ids_mask=mask,
755
+ )
756
+ bonus_token_tensor = torch.tensor(
757
+ [output_tokens[i][-1] for i in range(len(output_tokens))], device=logits.device
758
+ )
759
+ spec_decode_metadata = SpecDecodeMetadata.make_dummy(
760
+ spec_tokens, device=logits.device
761
+ )
762
+ output = rejection_sampler(
763
+ spec_decode_metadata,
764
+ draft_probs=None,
765
+ target_logits=logits,
766
+ bonus_token_ids=bonus_token_tensor,
767
+ sampling_metadata=metadata,
768
+ )
769
+
770
+ expected = torch.tensor(
771
+ [[15, -1, -1, -1], [10, 5, 10, -1], [7, 10, 12, 5]],
772
+ dtype=torch.int,
773
+ device=logits.device,
774
+ )
775
+ assert torch.equal(output, expected)
tests/v1/sample/test_sampler.py
+18
-28
@@ -1,12 +1,11 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from typing import Optional
5
-
4
import numpy as np
5
import pytest
6
import torch
7
8
+from tests.v1.sample.utils import create_allowed_token_ids
9
from vllm.platforms import current_platform
10
from vllm.utils import is_pin_memory_available, make_tensor_with_pad
11
from vllm.v1.sample.logits_processor import LogitsProcessors
@@ -51,26 +50,6 @@ def _create_prompt_tokens_tensor(
50
)
51
52
54
-def _create_allowed_token_ids(
55
- batch_size: int,
56
- vocab_size: int,
57
- num_allowed_token_ids: int,
58
- device: torch.device,
59
-) -> Optional[torch.Tensor]:
60
- mask: Optional[torch.Tensor] = None
61
- for i in range(batch_size):
62
- if i % 2 == 1:
63
- continue
64
- if mask is None:
65
- mask = torch.zeros(
66
- (batch_size, vocab_size), dtype=torch.bool, device=device
67
- )
68
- start = min(i, vocab_size - 1)
69
- end = min(i + num_allowed_token_ids, vocab_size - 1)
70
- mask[i, start:end] = True
71
- return mask
72
-
73
-
53
def _create_bad_words_token_ids(
54
batch_size: int,
55
vocab_size: int,
@@ -173,6 +152,7 @@ def _create_default_sampling_metadata(
152
prompt_token_ids, vocab_size, device
153
),
154
output_token_ids=output_token_ids,
155
+ spec_token_ids=[[] for _ in range(batch_size)],
156
frequency_penalties=_create_penalty_tensor(batch_size, 0.0, device),
157
presence_penalties=_create_penalty_tensor(batch_size, 0.0, device),
158
repetition_penalties=_create_penalty_tensor(batch_size, 1.0, device),
@@ -241,7 +221,9 @@ def test_sampler_presence_penalty(
221
)
222
sampling_metadata.no_penalties = False
223
sampler = Sampler()
244
- logits = sampler.apply_penalties(fake_logits, sampling_metadata)
224
+ logits = sampler.apply_penalties(
225
+ fake_logits, sampling_metadata, sampling_metadata.output_token_ids
226
+ )
227
logits = logits.cpu()
228
for batch_idx in range(batch_size):
229
# Since all tokens initially have the same logits, the non-penalized
@@ -293,7 +275,9 @@ def test_sampler_frequency_penalty(
275
sampling_metadata.output_token_ids = output_token_ids
276
sampling_metadata.no_penalties = False
277
sampler = Sampler()
296
- logits = sampler.apply_penalties(fake_logits, sampling_metadata)
278
+ logits = sampler.apply_penalties(
279
+ fake_logits, sampling_metadata, sampling_metadata.output_token_ids
280
+ )
281
logits = logits.cpu()
282
for batch_idx in range(batch_size):
283
non_penalized_token_id = logits[batch_idx].argmax().item()
@@ -343,7 +327,9 @@ def test_sampler_repetition_penalty(
327
)
328
sampling_metadata.no_penalties = False
329
sampler = Sampler()
346
- logits = sampler.apply_penalties(fake_logits, sampling_metadata)
330
+ logits = sampler.apply_penalties(
331
+ fake_logits, sampling_metadata, sampling_metadata.output_token_ids
332
+ )
333
logits = logits.cpu()
334
for batch_idx in range(batch_size):
335
non_penalized_token_id = logits[batch_idx].argmax().item()
@@ -394,7 +380,7 @@ def test_sampler_allowed_token_ids(
380
sampling_metadata = _create_default_sampling_metadata(
381
NUM_OUTPUT_TOKENS, batch_size, VOCAB_SIZE, torch.device(device)
382
)
397
- mask = _create_allowed_token_ids(
383
+ mask = create_allowed_token_ids(
384
batch_size=batch_size,
385
vocab_size=VOCAB_SIZE,
386
num_allowed_token_ids=num_allowed_token_ids,
@@ -402,7 +388,9 @@ def test_sampler_allowed_token_ids(
388
)
389
sampling_metadata.allowed_token_ids_mask = mask
390
sampler = Sampler()
405
- logits = sampler.apply_allowed_token_ids(fake_logits, sampling_metadata)
391
+ logits = sampler.apply_logits_processors(
392
+ fake_logits, sampling_metadata, predict_bonus_token=False
393
+ )
394
logits = logits.cpu()
395
for batch_idx in range(batch_size):
396
logits_for_req = logits[batch_idx]
@@ -444,7 +432,9 @@ def test_sampler_bad_words(
432
sampling_metadata, VOCAB_SIZE
433
)
434
sampler = Sampler()
447
- logits = sampler.apply_bad_words(fake_logits, sampling_metadata)
435
+ logits = sampler.apply_logits_processors(
436
+ fake_logits, sampling_metadata, predict_bonus_token=False
437
+ )
438
logits = logits.cpu()
439
for batch_idx in range(batch_size):
440
logits_for_req = logits[batch_idx]
tests/v1/sample/utils.py
+20
@@ -215,3 +215,23 @@ def fake_apply_logitsprocs(
215
for processor in test_fakes.get_logitsprocs():
216
logits = processor.apply(logits)
217
return logits
218
+
219
+
220
+def create_allowed_token_ids(
221
+ batch_size: int,
222
+ vocab_size: int,
223
+ num_allowed_token_ids: int,
224
+ device: torch.device,
225
+) -> Optional[torch.Tensor]:
226
+ mask: Optional[torch.Tensor] = None
227
+ for i in range(batch_size):
228
+ if i % 2 == 1:
229
+ continue
230
+ if mask is None:
231
+ mask = torch.zeros(
232
+ (batch_size, vocab_size), dtype=torch.bool, device=device
233
+ )
234
+ start = min(i, vocab_size - 1)
235
+ end = min(i + num_allowed_token_ids, vocab_size - 1)
236
+ mask[i, start:end] = True
237
+ return mask
tests/v1/worker/test_gpu_input_batch.py
+1
@@ -170,6 +170,7 @@ def _construct_expected_sampling_metadata(
170
repetition_penalties, dtype=torch.float, device=device
171
),
172
output_token_ids=output_token_ids,
173
+ spec_token_ids=[[] for _ in range(len(output_token_ids))],
174
no_penalties=(
175
all(x == 0 for x in presence_penalties)
176
and all(x == 0 for x in frequency_penalties)
vllm/v1/sample/logits_processor/__init__.py
+17
@@ -37,6 +37,12 @@ STR_POOLING_REJECTS_LOGITSPROCS = (
37
"Pooling models do not support custom logits processors."
38
)
39
40
+# Error message when the user tries to initialize vLLM with a speculative
41
+# decoding enabled and custom logitsproces
42
+STR_SPEC_DEC_REJECTS_LOGITSPROCS = (
43
+ "Custom logits processors are not supportedwhen speculative decoding is enabled."
44
+)
45
+
46
LOGITSPROCS_GROUP = "vllm.logits_processors"
47
48
BUILTIN_LOGITS_PROCESSORS: list[type[LogitsProcessor]] = [
@@ -185,6 +191,17 @@ def build_logitsprocs(
191
" do not support logits processors."
192
)
193
return LogitsProcessors()
194
+
195
+ # Check if speculative decoding is enabled.
196
+ if vllm_config.speculative_config:
197
+ if custom_logitsprocs:
198
+ raise ValueError(STR_SPEC_DEC_REJECTS_LOGITSPROCS)
199
+ logger.warning(
200
+ "min_p, logit_bias, and min_tokens parameters won't currently work "
201
+ "with speculative decoding enabled."
202
+ )
203
+ return LogitsProcessors()
204
+
205
custom_logitsprocs_classes = _load_custom_logitsprocs(custom_logitsprocs)
206
return LogitsProcessors(
207
ctor(vllm_config, device, is_pin_memory)
vllm/v1/sample/metadata.py
+3
@@ -40,3 +40,6 @@ class SamplingMetadata:
40
41
# Loaded logits processors
42
logitsprocs: LogitsProcessors
43
+
44
+ # Speculative token ids
45
+ spec_token_ids: Optional[list[list[int]]] = None
vllm/v1/sample/ops/bad_words.py
+17
@@ -33,3 +33,20 @@ def apply_bad_words(
33
) -> None:
34
for i, bad_words_ids in bad_words_token_ids.items():
35
_apply_bad_words_single_batch(logits[i], bad_words_ids, past_tokens_ids[i])
36
+
37
+
38
+def apply_bad_words_with_drafts(
39
+ logits: torch.Tensor,
40
+ bad_words_token_ids: dict[int, list[list[int]]],
41
+ past_tokens_ids: list[list[int]],
42
+ num_draft_tokens: list[int],
43
+) -> None:
44
+ start_idx = 0
45
+ for i, bad_words_ids in bad_words_token_ids.items():
46
+ for draft_idx in range(num_draft_tokens[i]):
47
+ _apply_bad_words_single_batch(
48
+ logits[start_idx + draft_idx],
49
+ bad_words_ids,
50
+ past_tokens_ids[start_idx + draft_idx],
51
+ )
52
+ start_idx += num_draft_tokens[i]
vllm/v1/sample/rejection_sampler.py
+104
@@ -8,6 +8,8 @@ import torch.nn as nn
8
from vllm.logger import init_logger
9
from vllm.triton_utils import tl, triton
10
from vllm.v1.sample.metadata import SamplingMetadata
11
+from vllm.v1.sample.ops.bad_words import apply_bad_words_with_drafts
12
+from vllm.v1.sample.ops.penalties import apply_all_penalties
13
from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p
14
from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
15
@@ -83,6 +85,14 @@ class RejectionSampler(nn.Module):
85
A tensor containing the final output token IDs.
86
"""
87
assert metadata.max_spec_len <= MAX_SPEC_LEN
88
+
89
+ # Use float32 for the target_logits.
90
+ target_logits = target_logits.to(torch.float32)
91
+
92
+ target_logits = self.apply_logits_processors(
93
+ target_logits, sampling_metadata, metadata
94
+ )
95
+
96
# [num_tokens, vocab_size]
97
# NOTE(woosuk): `target_logits` can be updated in place inside the
98
# `compute_probs` function.
@@ -131,6 +141,100 @@ class RejectionSampler(nn.Module):
141
]
142
return outputs
143
144
+ def apply_logits_processors(
145
+ self,
146
+ logits: torch.Tensor,
147
+ sampling_metadata: SamplingMetadata,
148
+ metadata: SpecDecodeMetadata,
149
+ ) -> torch.Tensor:
150
+ any_penalties_or_bad_words = (
151
+ sampling_metadata.bad_words_token_ids or not sampling_metadata.no_penalties
152
+ )
153
+
154
+ output_token_ids = sampling_metadata.output_token_ids
155
+ if any_penalties_or_bad_words:
156
+ output_token_ids = self._combine_outputs_with_spec_tokens(
157
+ sampling_metadata.output_token_ids,
158
+ sampling_metadata.spec_token_ids,
159
+ )
160
+
161
+ # Calculate indices of target logits.
162
+ if (
163
+ sampling_metadata.allowed_token_ids_mask is not None
164
+ or not sampling_metadata.no_penalties
165
+ ):
166
+ num_requests = len(sampling_metadata.output_token_ids)
167
+ num_draft_tokens = torch.tensor(metadata.num_draft_tokens, device="cpu")
168
+ original_indices = torch.arange(num_requests, device="cpu")
169
+ repeat_indices_cpu = original_indices.repeat_interleave(num_draft_tokens)
170
+ repeat_indices = repeat_indices_cpu.to(
171
+ device=logits.device, non_blocking=True
172
+ )
173
+ logits = self.apply_penalties(
174
+ logits, sampling_metadata, metadata, repeat_indices, output_token_ids
175
+ )
176
+
177
+ # Apply allowed token ids.
178
+ if sampling_metadata.allowed_token_ids_mask is not None:
179
+ token_mask = sampling_metadata.allowed_token_ids_mask[repeat_indices]
180
+ logits.masked_fill_(token_mask, float("-inf"))
181
+
182
+ # Apply bad words exclusion.
183
+ if sampling_metadata.bad_words_token_ids:
184
+ apply_bad_words_with_drafts(
185
+ logits,
186
+ sampling_metadata.bad_words_token_ids,
187
+ output_token_ids,
188
+ metadata.num_draft_tokens,
189
+ )
190
+
191
+ return logits
192
+
193
+ def apply_penalties(
194
+ self,
195
+ logits: torch.Tensor,
196
+ sampling_metadata: SamplingMetadata,
197
+ metadata: SpecDecodeMetadata,
198
+ repeat_indices: torch.Tensor,
199
+ output_token_ids: list[list[int]],
200
+ ) -> torch.Tensor:
201
+ if sampling_metadata.no_penalties:
202
+ return logits
203
+
204
+ assert sampling_metadata.prompt_token_ids is not None
205
+
206
+ prompt_token_ids = sampling_metadata.prompt_token_ids[repeat_indices]
207
+ presence_penalties = sampling_metadata.presence_penalties[repeat_indices]
208
+ frequency_penalties = sampling_metadata.frequency_penalties[repeat_indices]
209
+ repetition_penalties = sampling_metadata.repetition_penalties[repeat_indices]
210
+
211
+ logits = apply_all_penalties(
212
+ logits,
213
+ prompt_token_ids,
214
+ presence_penalties,
215
+ frequency_penalties,
216
+ repetition_penalties,
217
+ output_token_ids,
218
+ )
219
+ return logits
220
+
221
+ def _combine_outputs_with_spec_tokens(
222
+ self,
223
+ output_token_ids: list[list[int]],
224
+ spec_token_ids: Optional[list[list[int]]] = None,
225
+ ) -> list[list[int]]:
226
+ if spec_token_ids is None:
227
+ return output_token_ids
228
+
229
+ result = []
230
+ for out, spec in zip(output_token_ids, spec_token_ids):
231
+ if len(spec) == 0:
232
+ continue
233
+ result.append(out)
234
+ for i in range(len(spec) - 1):
235
+ result.append([*result[-1], spec[i]])
236
+ return result
237
+
238
239
def rejection_sample(
240
# [num_tokens]
vllm/v1/sample/sampler.py
+59
-31
@@ -70,6 +70,7 @@ class Sampler(nn.Module):
70
self,
71
logits: torch.Tensor,
72
sampling_metadata: SamplingMetadata,
73
+ predict_bonus_token: bool = False,
74
) -> SamplerOutput:
75
# NOTE(woosuk): Use the original logits (before any penalties or
76
# temperature scaling) for the top-k logprobs.
@@ -84,18 +85,10 @@ class Sampler(nn.Module):
85
86
# Use float32 for the logits.
87
logits = logits.to(torch.float32)
87
- # Apply allowed token ids.
88
- logits = self.apply_allowed_token_ids(logits, sampling_metadata)
89
- # Apply bad words exclusion.
90
- logits = self.apply_bad_words(logits, sampling_metadata)
91
-
92
- # Apply logits processors which can impact greedy sampling
93
- for processor in sampling_metadata.logitsprocs.non_argmax_invariant:
94
- logits = processor.apply(logits)
95
-
96
- # Apply penalties (e.g., min_tokens, freq_penalties).
97
- logits = self.apply_penalties(logits, sampling_metadata)
88
89
+ logits = self.apply_logits_processors(
90
+ logits, sampling_metadata, predict_bonus_token
91
+ )
92
# Sample the next token.
93
sampled, processed_logprobs = self.sample(logits, sampling_metadata)
94
if processed_logprobs is not None:
@@ -245,41 +238,76 @@ class Sampler(nn.Module):
238
239
return LogprobsTensors(indices, logprobs, token_ranks)
240
248
- def apply_penalties(
241
+ def _combine_outputs_with_spec_tokens(
242
+ self,
243
+ output_token_ids: list[list[int]],
244
+ spec_token_ids: Optional[list[list[int]]] = None,
245
+ ) -> list[list[int]]:
246
+ if spec_token_ids is None:
247
+ return output_token_ids
248
+
249
+ return [
250
+ [*out, *spec] if spec else out
251
+ for out, spec in zip(output_token_ids, spec_token_ids)
252
+ ]
253
+
254
+ def apply_logits_processors(
255
self,
256
logits: torch.Tensor,
257
sampling_metadata: SamplingMetadata,
258
+ predict_bonus_token: bool,
259
) -> torch.Tensor:
253
- if not sampling_metadata.no_penalties:
254
- assert sampling_metadata.prompt_token_ids is not None
255
- logits = apply_all_penalties(
256
- logits,
257
- sampling_metadata.prompt_token_ids,
258
- sampling_metadata.presence_penalties,
259
- sampling_metadata.frequency_penalties,
260
- sampling_metadata.repetition_penalties,
260
+ any_penalties_or_bad_words = (
261
+ sampling_metadata.bad_words_token_ids or not sampling_metadata.no_penalties
262
+ )
263
+
264
+ output_token_ids = sampling_metadata.output_token_ids
265
+ if predict_bonus_token and any_penalties_or_bad_words:
266
+ # Combine base outputs with spec tokens when speculative decoding
267
+ # is enabled.
268
+ output_token_ids = self._combine_outputs_with_spec_tokens(
269
sampling_metadata.output_token_ids,
270
+ sampling_metadata.spec_token_ids,
271
)
263
- return logits
272
265
- def apply_allowed_token_ids(
266
- self,
267
- logits: torch.Tensor,
268
- sampling_metadata: SamplingMetadata,
269
- ) -> torch.Tensor:
273
+ # Apply allowed token ids.
274
if sampling_metadata.allowed_token_ids_mask is not None:
275
logits.masked_fill_(sampling_metadata.allowed_token_ids_mask, float("-inf"))
276
+
277
+ # Apply bad words exclusion.
278
+ if sampling_metadata.bad_words_token_ids:
279
+ apply_bad_words(
280
+ logits,
281
+ sampling_metadata.bad_words_token_ids,
282
+ output_token_ids
283
+ if output_token_ids is not None
284
+ else sampling_metadata.output_token_ids,
285
+ )
286
+
287
+ # Apply logits processors which can impact greedy sampling.
288
+ for processor in sampling_metadata.logitsprocs.non_argmax_invariant:
289
+ logits = processor.apply(logits)
290
+
291
+ # Apply penalties (e.g., freq_penalties).
292
+ logits = self.apply_penalties(logits, sampling_metadata, output_token_ids)
293
return logits
294
274
- def apply_bad_words(
295
+ def apply_penalties(
296
self,
297
logits: torch.Tensor,
298
sampling_metadata: SamplingMetadata,
299
+ output_token_ids: Optional[list[list[int]]] = None,
300
) -> torch.Tensor:
279
- if sampling_metadata.bad_words_token_ids:
280
- apply_bad_words(
301
+ if not sampling_metadata.no_penalties:
302
+ assert sampling_metadata.prompt_token_ids is not None
303
+ logits = apply_all_penalties(
304
logits,
282
- sampling_metadata.bad_words_token_ids,
283
- sampling_metadata.output_token_ids,
305
+ sampling_metadata.prompt_token_ids,
306
+ sampling_metadata.presence_penalties,
307
+ sampling_metadata.frequency_penalties,
308
+ sampling_metadata.repetition_penalties,
309
+ output_token_ids
310
+ if output_token_ids is not None
311
+ else sampling_metadata.output_token_ids,
312
)
313
return logits
vllm/v1/worker/gpu_input_batch.py
+17
@@ -240,6 +240,9 @@ class InputBatch:
240
# data structure
241
self.logitsprocs = logitsprocs or LogitsProcessors()
242
243
+ # Store last speculative tokens for sampler.
244
+ self.spec_token_ids: list[Optional[list[int]]] = []
245
+
246
# This is updated each time the batch constituents change.
247
self.sampling_metadata = self._make_sampling_metadata()
248
@@ -292,9 +295,11 @@ class InputBatch:
295
if req_index == len(self._req_ids):
296
self._req_ids.append(req_id)
297
self.req_output_token_ids.append(request.output_token_ids)
298
+ self.spec_token_ids.append([])
299
else:
300
self._req_ids[req_index] = req_id
301
self.req_output_token_ids[req_index] = request.output_token_ids
302
+ self.spec_token_ids[req_index] = []
303
304
self.req_id_to_index[req_id] = req_index
305
@@ -443,6 +448,7 @@ class InputBatch:
448
self.batch_update_builder.removed_append(req_index)
449
self._req_ids[req_index] = None
450
self.req_output_token_ids[req_index] = None
451
+ self.spec_token_ids[req_index] = None
452
453
# LoRA
454
lora_id = self.request_lora_mapping[req_index]
@@ -486,6 +492,10 @@ class InputBatch:
492
self.req_output_token_ids[i2],
493
self.req_output_token_ids[i1],
494
)
495
+ self.spec_token_ids[i1], self.spec_token_ids[i2] = (
496
+ self.spec_token_ids[i2],
497
+ self.spec_token_ids[i1],
498
+ )
499
assert old_id_i1 is not None and old_id_i2 is not None
500
self.req_id_to_index[old_id_i1], self.req_id_to_index[old_id_i2] = (
501
self.req_id_to_index[old_id_i2],
@@ -601,6 +611,7 @@ class InputBatch:
611
# The batched states are empty.
612
self._req_ids.clear()
613
self.req_output_token_ids.clear()
614
+ self.spec_token_ids.clear()
615
return
616
617
# NOTE(woosuk): This function assumes that the empty_req_indices
@@ -629,6 +640,10 @@ class InputBatch:
640
self.req_output_token_ids[last_req_index] = None
641
self.req_id_to_index[req_id] = empty_index
642
643
+ spec_token_ids = self.spec_token_ids[last_req_index]
644
+ self.spec_token_ids[empty_index] = spec_token_ids
645
+ self.spec_token_ids[last_req_index] = None
646
+
647
num_tokens = self.num_tokens[last_req_index]
648
self.token_ids_cpu[empty_index, :num_tokens] = self.token_ids_cpu[
649
last_req_index, :num_tokens
@@ -700,6 +715,7 @@ class InputBatch:
715
# Trim lists to the batch size.
716
del self._req_ids[num_reqs:]
717
del self.req_output_token_ids[num_reqs:]
718
+ del self.spec_token_ids[num_reqs:]
719
720
def refresh_metadata(self):
721
"""Apply any batch updates to sampling metadata."""
@@ -784,6 +800,7 @@ class InputBatch:
800
presence_penalties=self.presence_penalties[:num_reqs],
801
repetition_penalties=self.repetition_penalties[:num_reqs],
802
output_token_ids=cast(list[list[int]], self.req_output_token_ids),
803
+ spec_token_ids=cast(list[list[int]], self.spec_token_ids),
804
no_penalties=self.no_penalties,
805
allowed_token_ids_mask=allowed_token_ids_mask,
806
bad_words_token_ids=self.bad_words_token_ids,
vllm/v1/worker/gpu_model_runner.py
+3
@@ -783,6 +783,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
783
] = spec_token_ids
784
# NOTE(woosuk): `num_tokens` here may include spec tokens.
785
self.input_batch.num_tokens[req_index] += num_spec_tokens
786
+ self.input_batch.spec_token_ids[req_index] = spec_token_ids
787
788
# Add the new or resumed requests to the persistent batch.
789
# The smaller empty indices are filled first.
@@ -2197,6 +2198,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2198
sampler_output = self.sampler(
2199
logits=bonus_logits,
2200
sampling_metadata=sampling_metadata,
2201
+ predict_bonus_token=True,
2202
)
2203
bonus_token_ids = sampler_output.sampled_token_ids
2204
@@ -3491,6 +3493,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
3493
presence_penalties=dummy_tensors(0.1),
3494
repetition_penalties=dummy_tensors(0.1),
3495
output_token_ids=[[] for _ in range(num_reqs)],
3496
+ spec_token_ids=[[] for _ in range(num_reqs)],
3497
allowed_token_ids_mask=None,
3498
bad_words_token_ids={},
3499
logitsprocs=LogitsProcessors(),