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
[Sampler] Support returning final logprobs (#22387)
Signed-off-by: 22quinn <33176974+22quinn@users.noreply.github.com> Co-authored-by: Nick Hill <nhill@redhat.com> Co-authored-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
22quinn committed
Aug 20, 2025 at 21:28 UTC
f571ff8eb6d9117c6a418f7f925921968dff8ac8
7 files changed
+126
-70
docs/usage/v1_guide.md
+5
-2
@@ -154,12 +154,15 @@ differences compared to V0:
154
155
##### Logprobs Calculation
156
157
-Logprobs in V1 are now returned immediately once computed from the model’s raw output (i.e.
157
+By default, logprobs in V1 are now returned immediately once computed from the model’s raw output (i.e.
158
before applying any logits post-processing such as temperature scaling or penalty
159
adjustments). As a result, the returned logprobs do not reflect the final adjusted
160
probabilities used during sampling.
161
162
-Support for logprobs with post-sampling adjustments is in progress and will be added in future updates.
162
+You can adjust this behavior by setting the `--logprobs-mode` flag.
163
+Four modes are supported: `raw_logprobs` (default), `processed_logprobs`, `raw_logits`, `processed_logits`.
164
+Raw means the values before applying any logit processors, like bad words.
165
+Processed means the values after applying all processors, including temperature and top_k/top_p.
166
167
##### Prompt Logprobs with Prefix Caching
168
tests/v1/sample/test_logprobs.py
+5
-5
@@ -456,9 +456,7 @@ def test_all_logprobs(example_prompts, monkeypatch: pytest.MonkeyPatch):
456
assert len(logprob) == vocab_size
457
458
459
-@pytest.mark.parametrize(
460
- "logprobs_mode",
461
- ["raw_logprobs", "raw_logits", "processed_logprobs", "processed_logits"])
459
+@pytest.mark.parametrize("logprobs_mode", list(LogprobsMode))
460
def test_logprobs_mode(logprobs_mode: LogprobsMode,
461
monkeypatch: pytest.MonkeyPatch):
462
"""Test with LLM engine with different logprobs_mode.
@@ -487,12 +485,14 @@ def test_logprobs_mode(logprobs_mode: LogprobsMode,
485
for logprobs in output.logprobs:
486
for token_id in logprobs:
487
logprob = logprobs[token_id]
490
- if "logprobs" in logprobs_mode:
488
+ if logprobs_mode in (LogprobsMode.RAW_LOGPROBS,
489
+ LogprobsMode.PROCESSED_LOGPROBS):
490
assert logprob.logprob <= 0
491
if logprob.logprob > 0:
492
positive_values = positive_values + 1
493
total_token_with_logprobs = total_token_with_logprobs + 1
494
assert total_token_with_logprobs >= len(results[0].outputs)
496
- if "logits" in logprobs_mode:
495
+ if logprobs_mode in (LogprobsMode.RAW_LOGITS,
496
+ LogprobsMode.PROCESSED_LOGITS):
497
assert positive_values > 0
498
del llm
vllm/config/__init__.py
+18
-12
@@ -257,11 +257,16 @@ def is_init_field(cls: ConfigType, name: str) -> bool:
257
258
TokenizerMode = Literal["auto", "slow", "mistral", "custom"]
259
ModelDType = Literal["auto", "half", "float16", "bfloat16", "float", "float32"]
260
-LogprobsMode = Literal["raw_logprobs", "raw_logits", "processed_logprobs",
261
- "processed_logits"]
260
MMEncoderTPMode = Literal["weights", "data"]
261
262
263
+class LogprobsMode(enum.Enum):
264
+ RAW_LOGITS = "raw_logits"
265
+ RAW_LOGPROBS = "raw_logprobs"
266
+ PROCESSED_LOGITS = "processed_logits"
267
+ PROCESSED_LOGPROBS = "processed_logprobs"
268
+
269
+
270
@config
271
@dataclass(config=ConfigDict(arbitrary_types_allowed=True))
272
class ModelConfig:
@@ -363,12 +368,13 @@ class ModelConfig:
368
specified in `SamplingParams`. The default value comes the default for the
369
OpenAI Chat Completions API. -1 means no cap, i.e. all (output_length *
370
vocab_size) logprobs are allowed to be returned and it may cause OOM."""
366
- logprobs_mode: LogprobsMode = "raw_logprobs"
371
+ logprobs_mode: LogprobsMode = LogprobsMode.RAW_LOGPROBS
372
"""Indicates the content returned in the logprobs and prompt_logprobs.
373
Supported mode:
374
1) raw_logprobs, 2) processed_logprobs, 3) raw_logits, 4) processed_logits.
370
- Raw means the values before applying logit processors, like bad words.
371
- Processed means the values after applying such processors.
375
+ Raw means the values before applying any logit processors, like bad words.
376
+ Processed means the values after applying all processors, including
377
+ temperature and top_k/top_p.
378
"""
379
disable_sliding_window: bool = False
380
"""Whether to disable sliding window. If True, we will disable the sliding
@@ -2586,7 +2592,7 @@ class MultiModalConfig:
2592
2593
skip_mm_profiling: bool = False
2594
"""
2589
- When enabled, skips multimodal memory profiling and only profiles with
2595
+ When enabled, skips multimodal memory profiling and only profiles with
2596
language backbone model during engine initialization.
2597
2598
This reduces engine startup time but shifts the responsibility to users for
@@ -2649,24 +2655,24 @@ class PoolerConfig:
2655
## for embeddings models
2656
normalize: Optional[bool] = None
2657
"""
2652
- Whether to normalize the embeddings outputs.
2658
+ Whether to normalize the embeddings outputs.
2659
"""
2660
dimensions: Optional[int] = None
2661
"""
2656
- Reduce the dimensions of embeddings if model
2662
+ Reduce the dimensions of embeddings if model
2663
support matryoshka representation.
2664
"""
2665
2666
## for classification models
2667
activation: Optional[bool] = None
2668
"""
2663
- Whether to apply activation function to the classification outputs.
2669
+ Whether to apply activation function to the classification outputs.
2670
"""
2671
2672
## for reward models
2673
softmax: Optional[bool] = None
2674
"""
2669
- Whether to apply softmax to the reward outputs.
2675
+ Whether to apply softmax to the reward outputs.
2676
"""
2677
step_tag_id: Optional[int] = None
2678
"""
@@ -2692,9 +2698,9 @@ class PoolerConfig:
2698
2699
max_embed_len: Optional[int] = None
2700
"""
2695
- Maximum input length allowed for embedding generation. When set, allows
2701
+ Maximum input length allowed for embedding generation. When set, allows
2702
inputs longer than max_embed_len to be accepted for embedding models.
2697
- This parameter enables accepting long inputs without requiring
2703
+ This parameter enables accepting long inputs without requiring
2704
VLLM_ALLOW_LONG_MAX_MODEL_LEN environment variable. When an input exceeds
2705
max_embed_len, it will be handled according to the original max_model_len
2706
validation logic. Defaults to None (i.e. set to max_model_len).
vllm/engine/arg_utils.py
+1
@@ -516,6 +516,7 @@ class EngineArgs:
516
model_group.add_argument("--max-logprobs",
517
**model_kwargs["max_logprobs"])
518
model_group.add_argument("--logprobs-mode",
519
+ choices=[f.value for f in LogprobsMode],
520
**model_kwargs["logprobs_mode"])
521
model_group.add_argument("--disable-sliding-window",
522
**model_kwargs["disable_sliding_window"])
vllm/v1/sample/ops/topk_topp_sampler.py
+35
-30
@@ -8,6 +8,7 @@ import torch.nn as nn
8
from packaging import version
9
10
from vllm import envs
11
+from vllm.config import LogprobsMode
12
from vllm.logger import init_logger
13
from vllm.platforms import current_platform
14
@@ -28,9 +29,16 @@ class TopKTopPSampler(nn.Module):
29
Implementations may update the logits tensor in-place.
30
"""
31
31
- def __init__(self):
32
+ def __init__(
33
+ self,
34
+ logprobs_mode: LogprobsMode = LogprobsMode.RAW_LOGPROBS) -> None:
35
super().__init__()
33
- if current_platform.is_cuda():
36
+ self.logprobs_mode = logprobs_mode
37
+ # flashinfer optimization does not apply if intermediate
38
+ # logprobs/logits after top_k/top_p need to be returned
39
+ if logprobs_mode not in (LogprobsMode.PROCESSED_LOGITS,
40
+ LogprobsMode.PROCESSED_LOGPROBS
41
+ ) and current_platform.is_cuda():
42
if is_flashinfer_available:
43
flashinfer_version = flashinfer.__version__
44
if version.parse(flashinfer_version) < version.parse("0.2.3"):
@@ -63,10 +71,12 @@ class TopKTopPSampler(nn.Module):
71
"native implementation of top-p & top-k sampling. For the "
72
"best performance, please install FlashInfer.")
73
self.forward = self.forward_native
66
- elif current_platform.is_tpu():
67
- self.forward = self.forward_tpu
74
else:
75
self.forward = self.forward_native
76
+ if current_platform.is_tpu():
77
+ self.apply_top_k_top_p = apply_top_k_top_p_tpu
78
+ else:
79
+ self.apply_top_k_top_p = apply_top_k_top_p
80
81
def forward_native(
82
self,
@@ -74,15 +84,20 @@ class TopKTopPSampler(nn.Module):
84
generators: dict[int, torch.Generator],
85
k: Optional[torch.Tensor],
86
p: Optional[torch.Tensor],
77
- ) -> torch.Tensor:
87
+ ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
88
"""
89
PyTorch-native implementation of top-k and top-p sampling.
90
91
The logits tensor may be updated in-place.
92
"""
83
- logits = apply_top_k_top_p(logits, k, p)
93
+ logits = self.apply_top_k_top_p(logits, k, p)
94
+ logits_to_return = None
95
+ if self.logprobs_mode == LogprobsMode.PROCESSED_LOGITS:
96
+ logits_to_return = logits
97
+ elif self.logprobs_mode == LogprobsMode.PROCESSED_LOGPROBS:
98
+ logits_to_return = logits.log_softmax(dim=-1, dtype=torch.float32)
99
probs = logits.softmax(dim=-1, dtype=torch.float32)
85
- return random_sample(probs, generators)
100
+ return random_sample(probs, generators), logits_to_return
101
102
def forward_cuda(
103
self,
@@ -90,34 +105,24 @@ class TopKTopPSampler(nn.Module):
105
generators: dict[int, torch.Generator],
106
k: Optional[torch.Tensor],
107
p: Optional[torch.Tensor],
93
- ) -> torch.Tensor:
108
+ ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
109
"""More optimized implementation for top-k and top-p sampling."""
95
- if k is None and p is None:
96
- # We prefer `random_sample` over `flashinfer_sample` when sorting is
97
- # not needed. This is because `random_sample` does not require
98
- # CPU-GPU synchronization while `flashinfer_sample` does.
99
- probs = logits.softmax(dim=-1, dtype=torch.float32)
100
- return random_sample(probs, generators)
101
- if generators:
102
- logger.warning_once("FlashInfer 0.2.3+ does not support "
103
- "per-request generators. Falling back to "
104
- "PyTorch-native implementation.")
110
+ # We prefer `random_sample` over `flashinfer_sample` when sorting is
111
+ # not needed. This is because `random_sample` does not require
112
+ # CPU-GPU synchronization while `flashinfer_sample` does.
113
+ if (k is None and p is None) or generators:
114
+ if generators:
115
+ logger.warning_once("FlashInfer 0.2.3+ does not support "
116
+ "per-request generators. Falling back to "
117
+ "PyTorch-native implementation.")
118
return self.forward_native(logits, generators, k, p)
119
+ assert self.logprobs_mode not in (
120
+ LogprobsMode.PROCESSED_LOGITS, LogprobsMode.PROCESSED_LOGPROBS
121
+ ), "FlashInfer does not support returning logits/logprobs"
122
# flashinfer sampling functions expect contiguous logits.
123
# In flex_attn/triton_attn fp32 inference, logits can be non-contiguous
124
# because of slicing operation in logits_processor.
109
- return flashinfer_sample(logits.contiguous(), k, p, generators)
110
-
111
- def forward_tpu(
112
- self,
113
- logits: torch.Tensor,
114
- generators: dict[int, torch.Generator],
115
- k: Optional[torch.Tensor],
116
- p: Optional[torch.Tensor],
117
- ) -> torch.Tensor:
118
- logits = apply_top_k_top_p_tpu(logits, k, p)
119
- probs = logits.softmax(dim=-1, dtype=torch.float32)
120
- return random_sample(probs, generators)
125
+ return flashinfer_sample(logits.contiguous(), k, p, generators), None
126
127
128
def apply_top_k_top_p_tpu(
vllm/v1/sample/sampler.py
+61
-20
@@ -2,6 +2,8 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""A layer that samples the next tokens from the model's outputs."""
4
5
+from typing import Optional
6
+
7
import torch
8
import torch.nn as nn
9
@@ -18,10 +20,50 @@ _SAMPLING_EPS = 1e-5
20
21
22
class Sampler(nn.Module):
21
-
22
- def __init__(self, logprobs_mode: LogprobsMode = "raw_logprobs"):
23
+ """
24
+ A layer that samples the next tokens from the model's outputs
25
+ with the following steps in order:
26
+
27
+ 1. If logprobs are requested:
28
+ a) If `logprobs_mode` is `raw_logprobs`, compute logprobs
29
+ as the final logprobs to return.
30
+ b) If `logprobs_mode` is `raw_logits`, clone the logits
31
+ as the final logprobs to return.
32
+ 2. Convert logits to float32.
33
+ 3. Apply allowed token ids whitelist.
34
+ 4. Apply bad words exclusion.
35
+ 5. Apply logit processors which are not argmax-invariant,
36
+ i.e. that can impact greedy sampling.
37
+ a) Min tokens processor
38
+ b) Logit bias processor
39
+ 6. Apply penalties
40
+ a) Repetition penalty
41
+ b) Frequency penalty
42
+ c) Presence penalty
43
+ 7. Sample the next tokens. `sample` method performs the following steps:
44
+ a) If not `all_random`, perform greedy sampling. If `all_greedy`,
45
+ return the greedily sampled tokens and final logprobs if requested.
46
+ b) Apply temperature.
47
+ c) Apply logit processors which are argmax-invariant, by default
48
+ the min_p processor.
49
+ d) Apply top_k and/or top_p.
50
+ e) Sample the next tokens with the probability distribution.
51
+ f) If `all_random` or temperature >= epsilon (1e-5), return the
52
+ randomly sampled tokens and final logprobs if requested. Else,
53
+ return the greedily sampled tokens and logprobs if requested.
54
+ 8. Gather the logprobs of the top `max_num_logprobs` and sampled token
55
+ (if requested). Note that if the sampled token is within the top
56
+ `max_num_logprobs`, the logprob will be eventually merged in
57
+ `LogprobsProcessor` during output processing. Therefore, the
58
+ final output may contain either `max_num_logprobs + 1` or
59
+ `max_num_logprobs` logprobs.
60
+ 9. Return the final `SamplerOutput`.
61
+ """
62
+
63
+ def __init__(self,
64
+ logprobs_mode: LogprobsMode = LogprobsMode.RAW_LOGPROBS):
65
super().__init__()
24
- self.topk_topp_sampler = TopKTopPSampler()
66
+ self.topk_topp_sampler = TopKTopPSampler(logprobs_mode)
67
self.pin_memory = is_pin_memory_available()
68
self.logprobs_mode = logprobs_mode
69
@@ -34,13 +76,11 @@ class Sampler(nn.Module):
76
# temperature scaling) for the top-k logprobs.
77
# This is different from the V0 sampler, which uses the logits that
78
# is used for sampling (after penalties and temperature scaling).
37
- # TODO(rob): provide option for logprobs post sampling.
38
- # See https://vllm-dev.slack.com/archives/C07UUL8E61Z/p1735907856007919 # noqa: E501
79
num_logprobs = sampling_metadata.max_num_logprobs
80
if num_logprobs is not None:
41
- if self.logprobs_mode == "raw_logprobs":
81
+ if self.logprobs_mode == LogprobsMode.RAW_LOGPROBS:
82
raw_logprobs = self.compute_logprobs(logits)
43
- elif self.logprobs_mode == "raw_logits":
83
+ elif self.logprobs_mode == LogprobsMode.RAW_LOGITS:
84
raw_logprobs = logits.clone()
85
86
# Use float32 for the logits.
@@ -57,15 +97,10 @@ class Sampler(nn.Module):
97
# Apply penalties (e.g., min_tokens, freq_penalties).
98
logits = self.apply_penalties(logits, sampling_metadata)
99
60
- # Get the process logprobs or logits.
61
- if num_logprobs is not None:
62
- if self.logprobs_mode == "processed_logprobs":
63
- raw_logprobs = self.compute_logprobs(logits)
64
- elif self.logprobs_mode == "processed_logits":
65
- raw_logprobs = logits.clone()
66
-
100
# Sample the next token.
68
- sampled = self.sample(logits, sampling_metadata)
101
+ sampled, processed_logprobs = self.sample(logits, sampling_metadata)
102
+ if processed_logprobs is not None:
103
+ raw_logprobs = processed_logprobs
104
# Convert sampled token ids to int64 (long) type to ensure compatibility
105
# with subsequent operations that may use these values as indices.
106
# This conversion is necessary because FlashInfer sampling operations
@@ -105,7 +140,7 @@ class Sampler(nn.Module):
140
self,
141
logits: torch.Tensor,
142
sampling_metadata: SamplingMetadata,
108
- ) -> torch.Tensor:
143
+ ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
144
"""Sample logits based on sampling metadata.
145
146
The various logits processing functions called in this method
@@ -119,7 +154,13 @@ class Sampler(nn.Module):
154
else:
155
greedy_sampled = self.greedy_sample(logits)
156
if sampling_metadata.all_greedy:
122
- return greedy_sampled
157
+ processed_logprobs = None
158
+ if sampling_metadata.max_num_logprobs is not None:
159
+ if self.logprobs_mode == LogprobsMode.PROCESSED_LOGITS:
160
+ processed_logprobs = logits
161
+ elif self.logprobs_mode == LogprobsMode.PROCESSED_LOGPROBS:
162
+ processed_logprobs = self.compute_logprobs(logits)
163
+ return greedy_sampled, processed_logprobs
164
165
assert sampling_metadata.temperature is not None
166
@@ -132,7 +173,7 @@ class Sampler(nn.Module):
173
logits = processor.apply(logits)
174
175
# Apply top_k and/or top_p.
135
- random_sampled = self.topk_topp_sampler(
176
+ random_sampled, processed_logprobs = self.topk_topp_sampler(
177
logits,
178
sampling_metadata.generators,
179
sampling_metadata.top_k,
@@ -140,7 +181,7 @@ class Sampler(nn.Module):
181
)
182
183
if greedy_sampled is None:
143
- return random_sampled
184
+ return random_sampled, processed_logprobs
185
186
sampled = torch.where(
187
sampling_metadata.temperature < _SAMPLING_EPS,
@@ -148,7 +189,7 @@ class Sampler(nn.Module):
189
random_sampled,
190
out=greedy_sampled, # Reuse tensor
191
)
151
- return sampled
192
+ return sampled, processed_logprobs
193
194
def compute_logprobs(self, logits: torch.Tensor) -> torch.Tensor:
195
return logits.log_softmax(dim=-1, dtype=torch.float32)
vllm/v1/sample/tpu/sampler.py
+1
-1
@@ -65,7 +65,7 @@ class Sampler(nn.Module):
65
logits = self.apply_min_p(logits, sampling_metadata.min_p)
66
67
# Apply top_k and/or top_p.
68
- random_sampled = self.topk_topp_sampler(
68
+ random_sampled, _ = self.topk_topp_sampler(
69
logits,
70
sampling_metadata.generators,
71
sampling_metadata.top_k,