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 all logprobs or logits (#21792)
Signed-off-by: 22quinn <33176974+22quinn@users.noreply.github.com>
22quinn committed
Aug 4, 2025 at 03:04 UTC
54de71d0dfbb6340fdbc620f4ebeb4236d165a37
6 files changed
+45
-9
tests/v1/sample/test_logprobs.py
+27
@@ -429,6 +429,33 @@ def test_zero_logprobs(vllm_model, example_prompts,
429
assert len(prompt_token_ids) == len(prompt_logprobs)
430
431
432
+def test_all_logprobs(example_prompts, monkeypatch: pytest.MonkeyPatch):
433
+ """Engine should return all vocabulary logprobs
434
+
435
+ Args:
436
+ example_prompts: list of example prompts (test fixture)
437
+ """
438
+ with monkeypatch.context() as m:
439
+ m.setenv("VLLM_USE_V1", "1")
440
+ runner = VllmRunner(
441
+ "facebook/opt-125m",
442
+ max_logprobs=-1,
443
+ enable_prefix_caching=False,
444
+ # 2 other llms alive during whole session
445
+ gpu_memory_utilization=0.15,
446
+ max_model_len=256)
447
+ sampling_params_logprobs_all = SamplingParams(max_tokens=5,
448
+ logprobs=-1)
449
+ results_logprobs_all = runner.llm.generate(
450
+ example_prompts, sampling_params=sampling_params_logprobs_all)
451
+ vocab_size = runner.llm.llm_engine.get_model_config().get_vocab_size()
452
+ for i in range(len(results_logprobs_all)):
453
+ logprobs = results_logprobs_all[i].outputs[0].logprobs
454
+ assert logprobs is not None
455
+ for logprob in logprobs:
456
+ assert len(logprob) == vocab_size
457
+
458
+
459
@pytest.mark.parametrize(
460
"logprobs_mode",
461
["raw_logprobs", "raw_logits", "processed_logprobs", "processed_logits"])
vllm/config.py
+4
-3
@@ -377,7 +377,8 @@ class ModelConfig:
377
max_logprobs: int = 20
378
"""Maximum number of log probabilities to return when `logprobs` is
379
specified in `SamplingParams`. The default value comes the default for the
380
- OpenAI Chat Completions API."""
380
+ OpenAI Chat Completions API. -1 means no cap, i.e. all (output_length *
381
+ vocab_size) logprobs are allowed to be returned and it may cause OOM."""
382
logprobs_mode: LogprobsMode = "raw_logprobs"
383
"""Indicates the content returned in the logprobs and prompt_logprobs.
384
Supported mode:
@@ -1585,7 +1586,7 @@ class ModelConfig:
1586
"""
1587
This method attempts to retrieve the non-default values of the
1588
generation config for this model.
1588
-
1589
+
1590
The generation config can contain information about special tokens, as
1591
well as sampling parameters. Which is why this method exists separately
1592
to `get_diff_sampling_param`.
@@ -2066,7 +2067,7 @@ class ParallelConfig:
2067
and when data_parallel_size > 0. Enables running an AsyncLLM
2068
and API server on a "per-node" basis where vLLM load balances
2069
between local data parallel ranks, but an external LB balances
2069
- between vLLM nodes/replicas. Set explicitly in conjunction with
2070
+ between vLLM nodes/replicas. Set explicitly in conjunction with
2071
--data-parallel-start-rank."""
2072
enable_expert_parallel: bool = False
2073
"""Use expert parallelism instead of tensor parallelism for MoE layers."""
vllm/sampling_params.py
+4
-2
@@ -156,6 +156,7 @@ class SamplingParams(
156
Note that the implementation follows the OpenAI API: The API will
157
always return the log probability of the sampled token, so there
158
may be up to `logprobs+1` elements in the response.
159
+ When set to -1, return all `vocab_size` log probabilities.
160
prompt_logprobs: Number of log probabilities to return per prompt token.
161
detokenize: Whether to detokenize the output. Defaults to True.
162
skip_special_tokens: Whether to skip special tokens in the output.
@@ -414,9 +415,10 @@ class SamplingParams(
415
raise ValueError(
416
f"min_tokens must be less than or equal to "
417
f"max_tokens={self.max_tokens}, got {self.min_tokens}.")
417
- if self.logprobs is not None and self.logprobs < 0:
418
+ if (self.logprobs is not None and self.logprobs != -1
419
+ and self.logprobs < 0):
420
raise ValueError(
419
- f"logprobs must be non-negative, got {self.logprobs}.")
421
+ f"logprobs must be non-negative or -1, got {self.logprobs}.")
422
if self.prompt_logprobs is not None and self.prompt_logprobs < 0:
423
raise ValueError(f"prompt_logprobs must be non-negative, got "
424
f"{self.prompt_logprobs}.")
vllm/v1/engine/logprobs.py
+3
-2
@@ -138,7 +138,7 @@ class LogprobsProcessor:
138
139
def pop_prompt_logprobs(self) -> Optional[PromptLogprobs]:
140
"""Pop and return all request prompt logprobs
141
-
141
+
142
The logprobs processor aggregates prompt chunk logprobs
143
over one or more prefill chunks. This method returns
144
all prompt logprobs at once and then forgets them.
@@ -176,7 +176,8 @@ class LogprobsProcessor:
176
Returns:
177
dict[token id, Logprob]
178
"""
179
-
179
+ if num_logprobs == -1:
180
+ num_logprobs = len(logprobs)
181
# We do not need a special case for the sampled token
182
# being in the topk, since inserting duplicated data
183
# into a dictionary twice is the same as doing it once.
vllm/v1/engine/processor.py
+4
-1
@@ -65,8 +65,11 @@ class Processor:
65
params: SamplingParams,
66
) -> None:
67
max_logprobs = self.model_config.max_logprobs
68
+ if max_logprobs == -1:
69
+ return
70
# Validate sample logprobs.
69
- if params.logprobs and params.logprobs > max_logprobs:
71
+ if params.logprobs and (params.logprobs == -1
72
+ or params.logprobs > max_logprobs):
73
raise ValueError(
74
f"Requested sample logprobs of {params.logprobs}, "
75
f"which is greater than max allowed: {max_logprobs}")
vllm/v1/worker/gpu_input_batch.py
+3
-1
@@ -337,7 +337,9 @@ class InputBatch:
337
self.generators[req_index] = request.generator
338
339
if sampling_params.logprobs is not None:
340
- self.num_logprobs[req_id] = sampling_params.logprobs
340
+ self.num_logprobs[req_id] = (self.vocab_size
341
+ if sampling_params.logprobs == -1
342
+ else sampling_params.logprobs)
343
if sampling_params.prompt_logprobs is not None:
344
self.num_prompt_logprobs[
345
req_id] = sampling_params.prompt_logprobs