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 7, 16:30 UTC
[Core] Support `min_tokens` with speculative decoding (#32642)
Signed-off-by: qianlihuang <yiliu.dong@qq.com> Co-authored-by: qianlihuang <yiliu.dong@qq.com>
Yiliu Dong committed
Feb 27, 2026 at 01:31 UTC
d940607629b03602f34ba4dd75c747162b01aedd
7 files changed
+75
-11
tests/v1/e2e/test_async_scheduling.py
+1
-2
@@ -32,8 +32,7 @@ example_prompts = [first_prompt, "In one word, the capital of France is "] + [
32
default_params = dict(
33
temperature=0.0, # greedy
34
max_tokens=30,
35
- # spec decoding currently doesn't support min_tokens
36
- # min_tokens=28,
35
+ min_tokens=28,
36
)
37
38
tests/v1/logits_processors/test_custom_offline.py
+5
-2
@@ -276,9 +276,12 @@ def test_rejects_custom_logitsprocs(
276
monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "fork")
277
278
llm = LLM(**llm_kwargs)
279
- # Require that no logitsprocs have been loaded
279
+ # Require that no custom logitsprocs have been loaded
280
+ # (built-in processors may exist: MinTokensLogitsProcessor,
281
+ # LogitBiasLogitsProcessor, MinPLogitsProcessor)
282
worker = llm.llm_engine.model_executor.driver_worker.worker
281
- assert sum([1 for _ in worker.model_runner.input_batch.logitsprocs.all]) == 0
283
+ for proc in worker.model_runner.input_batch.logitsprocs.all:
284
+ assert not isinstance(proc, DummyLogitsProcessor)
285
return
286
287
if logitproc_source == CustomLogitprocSource.LOGITPROC_SOURCE_FQCN:
vllm/sampling_params.py
+2
-2
@@ -678,9 +678,9 @@ class SamplingParams(
678
return
679
680
# Some sampling parameters are not yet compatible with spec decoding.
681
- if self.min_tokens > 1 or self.min_p > _SAMPLING_EPS or self.logit_bias:
681
+ if self.min_p > _SAMPLING_EPS or self.logit_bias:
682
raise ValueError(
683
- "The min_tokens, min_p, and logit_bias sampling parameters "
683
+ "The min_p and logit_bias sampling parameters "
684
"are not yet supported with speculative decoding."
685
)
686
vllm/v1/sample/logits_processor/__init__.py
+4
-3
@@ -202,10 +202,11 @@ def build_logitsprocs(
202
if custom_logitsprocs:
203
raise ValueError(STR_SPEC_DEC_REJECTS_LOGITSPROCS)
204
logger.warning(
205
- "min_p, logit_bias, and min_tokens parameters won't currently work "
206
- "with speculative decoding enabled."
205
+ "min_p and logit_bias parameters won't work with speculative decoding."
206
+ )
207
+ return LogitsProcessors(
208
+ [MinTokensLogitsProcessor(vllm_config, device, is_pin_memory)]
209
)
208
- return LogitsProcessors()
210
211
custom_logitsprocs_classes = _load_custom_logitsprocs(custom_logitsprocs)
212
return LogitsProcessors(
vllm/v1/sample/logits_processor/builtin.py
+54
@@ -3,6 +3,7 @@
3
from collections.abc import Callable, Sequence
4
from typing import TYPE_CHECKING, TypeVar
5
6
+import numpy as np
7
import torch
8
9
from vllm import SamplingParams
@@ -236,6 +237,59 @@ class MinTokensLogitsProcessor(LogitsProcessor):
237
logits.index_put_(self.logits_slice, self.neg_inf_tensor)
238
return logits
239
240
+ def apply_with_spec_decode(
241
+ self,
242
+ logits: torch.Tensor,
243
+ num_draft_tokens: list[int],
244
+ ) -> torch.Tensor:
245
+ """Spec-decode version of apply().
246
+ Priority: ``min_tokens`` > ``stop_token_ids`` / EOS.
247
+ Example: ``num_draft_tokens = [2, 3, 1]``
248
+ → ``logits`` shape ``[6, V]``, ``cumsum = [0, 2, 5, 6]``
249
+ → request 0 owns rows 0‑1, request 1 rows 2‑4, request 2 row 5.
250
+ """
251
+ if not self.min_toks:
252
+ return logits
253
+
254
+ num_draft_arr = np.array(num_draft_tokens, dtype=np.int64)
255
+ cumsum = np.concatenate([[0], np.cumsum(num_draft_arr)])
256
+
257
+ entries = [
258
+ (req_idx, min_tok, len(out_tok_ids), list(stop_tok_ids))
259
+ for req_idx, (min_tok, out_tok_ids, stop_tok_ids) in self.min_toks.items()
260
+ if stop_tok_ids
261
+ ]
262
+
263
+ if not entries:
264
+ return logits
265
+
266
+ all_rows: list[np.ndarray] = [] # row indices to mask
267
+ all_toks: list[np.ndarray] = [] # stop-token ids at those rows
268
+
269
+ for req_idx, min_tok, current_len, stop_toks in entries:
270
+ remaining = min_tok - current_len
271
+ # How many leading draft positions still need stop-token masking.
272
+ n_mask = int(min(max(remaining, 0), num_draft_arr[req_idx]))
273
+
274
+ if n_mask > 0:
275
+ offset = cumsum[req_idx]
276
+ row_indices = np.arange(offset, offset + n_mask, dtype=np.int64)
277
+ n_stop = len(stop_toks)
278
+ all_rows.append(np.repeat(row_indices, n_stop))
279
+ all_toks.append(np.tile(stop_toks, n_mask))
280
+
281
+ if all_rows:
282
+ rows_arr = np.concatenate(all_rows)
283
+ toks_arr = np.concatenate(all_toks)
284
+ # (row_indices, token_indices) for index_put_ to set -inf.
285
+ logits_slice = (
286
+ torch.from_numpy(rows_arr).to(self.device, non_blocking=True),
287
+ torch.from_numpy(toks_arr).to(self.device, non_blocking=True),
288
+ )
289
+ logits.index_put_(logits_slice, self.neg_inf_tensor)
290
+
291
+ return logits
292
+
293
294
def process_dict_updates(
295
req_entries: dict[int, T],
vllm/v1/sample/logits_processor/state.py
+2
-2
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from collections.abc import Iterator
3
+from collections.abc import Iterable, Iterator
4
from itertools import chain
5
from typing import TYPE_CHECKING
6
@@ -148,7 +148,7 @@ class BatchUpdateBuilder:
148
class LogitsProcessors:
149
"""Encapsulates initialized logitsproc objects."""
150
151
- def __init__(self, logitsprocs: Iterator["LogitsProcessor"] | None = None) -> None:
151
+ def __init__(self, logitsprocs: Iterable["LogitsProcessor"] | None = None) -> None:
152
self.argmax_invariant: list[LogitsProcessor] = []
153
self.non_argmax_invariant: list[LogitsProcessor] = []
154
if logitsprocs:
vllm/v1/sample/rejection_sampler.py
+7
@@ -10,6 +10,7 @@ import torch.nn as nn
10
from vllm.logger import init_logger
11
from vllm.triton_utils import tl, triton
12
from vllm.v1.outputs import LogprobsLists, LogprobsTensors, SamplerOutput
13
+from vllm.v1.sample.logits_processor.builtin import MinTokensLogitsProcessor
14
from vllm.v1.sample.metadata import SamplingMetadata
15
from vllm.v1.sample.ops.bad_words import apply_bad_words_with_drafts
16
from vllm.v1.sample.ops.penalties import apply_all_penalties
@@ -292,6 +293,12 @@ class RejectionSampler(nn.Module):
293
logits, bad_words_token_ids, output_token_ids, metadata.num_draft_tokens
294
)
295
296
+ for processor in sampling_metadata.logitsprocs.non_argmax_invariant:
297
+ if isinstance(processor, MinTokensLogitsProcessor):
298
+ logits = processor.apply_with_spec_decode(
299
+ logits, metadata.num_draft_tokens
300
+ )
301
+
302
return logits
303
304
@staticmethod