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 prompt logprobs (#23868)
Signed-off-by: Xingyu Liu <charlotteliu12x@gmail.com> Co-authored-by: 22quinn <33176974+22quinn@users.noreply.github.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
Xingyu Liu committed
Sep 7, 2025 at 19:34 UTC
b3d7e3c845d3f4905e9a95776c7d605c743f0ea1
4 files changed
+38
-18
tests/v1/sample/test_logprobs.py
+10
-2
@@ -430,7 +430,7 @@ def test_zero_logprobs(vllm_model, example_prompts,
430
431
432
def test_all_logprobs(example_prompts, monkeypatch: pytest.MonkeyPatch):
433
- """Engine should return all vocabulary logprobs
433
+ """Engine should return all vocabulary logprobs and prompt logprobs
434
435
Args:
436
example_prompts: list of example prompts (test fixture)
@@ -444,16 +444,24 @@ def test_all_logprobs(example_prompts, monkeypatch: pytest.MonkeyPatch):
444
# 2 other llms alive during whole session
445
gpu_memory_utilization=0.15,
446
max_model_len=256)
447
+
448
sampling_params_logprobs_all = SamplingParams(max_tokens=5,
448
- logprobs=-1)
449
+ logprobs=-1,
450
+ prompt_logprobs=-1)
451
results_logprobs_all = runner.llm.generate(
452
example_prompts, sampling_params=sampling_params_logprobs_all)
453
vocab_size = runner.llm.llm_engine.get_model_config().get_vocab_size()
454
+
455
for i in range(len(results_logprobs_all)):
456
logprobs = results_logprobs_all[i].outputs[0].logprobs
457
+ prompt_logprobs = results_logprobs_all[i].prompt_logprobs
458
assert logprobs is not None
459
for logprob in logprobs:
460
assert len(logprob) == vocab_size
461
+ assert prompt_logprobs is not None
462
+ assert prompt_logprobs[0] is None
463
+ for prompt_logprob in prompt_logprobs[1:]:
464
+ assert len(prompt_logprob) == vocab_size
465
466
467
@pytest.mark.parametrize("logprobs_mode", list(LogprobsMode))
vllm/sampling_params.py
+7
-4
@@ -165,7 +165,8 @@ class SamplingParams(
165
the sampled token, so there may be up to `logprobs+1` elements in the
166
response. When set to -1, return all `vocab_size` log probabilities."""
167
prompt_logprobs: Optional[int] = None
168
- """Number of log probabilities to return per prompt token."""
168
+ """Number of log probabilities to return per prompt token.
169
+ When set to -1, return all `vocab_size` log probabilities."""
170
# NOTE: This parameter is only exposed at the engine level for now.
171
# It is not exposed in the OpenAI API server, as the OpenAI API does
172
# not support returning only a list of token IDs.
@@ -409,9 +410,11 @@ class SamplingParams(
410
and self.logprobs < 0):
411
raise ValueError(
412
f"logprobs must be non-negative or -1, got {self.logprobs}.")
412
- if self.prompt_logprobs is not None and self.prompt_logprobs < 0:
413
- raise ValueError(f"prompt_logprobs must be non-negative, got "
414
- f"{self.prompt_logprobs}.")
413
+ if (self.prompt_logprobs is not None and self.prompt_logprobs != -1
414
+ and self.prompt_logprobs < 0):
415
+ raise ValueError(
416
+ f"prompt_logprobs must be non-negative or -1, got "
417
+ f"{self.prompt_logprobs}.")
418
if (self.truncate_prompt_tokens is not None
419
and (self.truncate_prompt_tokens == 0
420
or self.truncate_prompt_tokens < -1)):
vllm/v1/engine/processor.py
+18
-10
@@ -65,19 +65,27 @@ class Processor:
65
) -> None:
66
max_logprobs = self.model_config.max_logprobs
67
if max_logprobs == -1:
68
- return
68
+ max_logprobs = self.model_config.get_vocab_size()
69
+
70
# Validate sample logprobs.
70
- if params.logprobs and (params.logprobs == -1
71
- or params.logprobs > max_logprobs):
72
- raise ValueError(
73
- f"Requested sample logprobs of {params.logprobs}, "
74
- f"which is greater than max allowed: {max_logprobs}")
71
+ if params.logprobs:
72
+ num_logprobs = params.logprobs
73
+ if num_logprobs == -1:
74
+ num_logprobs = self.model_config.get_vocab_size()
75
+ if num_logprobs > max_logprobs:
76
+ raise ValueError(
77
+ f"Requested sample logprobs of {num_logprobs}, "
78
+ f"which is is greater than max allowed: {max_logprobs}")
79
80
# Validate prompt logprobs.
77
- if params.prompt_logprobs and params.prompt_logprobs > max_logprobs:
78
- raise ValueError(
79
- f"Requested prompt logprobs of {params.prompt_logprobs}, "
80
- f"which is greater than max allowed: {max_logprobs}")
81
+ if params.prompt_logprobs:
82
+ num_prompt_logprobs = params.prompt_logprobs
83
+ if num_prompt_logprobs == -1:
84
+ num_prompt_logprobs = self.model_config.get_vocab_size()
85
+ if num_prompt_logprobs > max_logprobs:
86
+ raise ValueError(
87
+ f"Requested prompt logprobs of {num_prompt_logprobs}, "
88
+ f"which is is greater than max allowed: {max_logprobs}")
89
90
def _validate_sampling_params(
91
self,
vllm/v1/worker/gpu_input_batch.py
+3
-2
@@ -360,8 +360,9 @@ class InputBatch:
360
if sampling_params.logprobs == -1
361
else sampling_params.logprobs)
362
if sampling_params.prompt_logprobs is not None:
363
- self.num_prompt_logprobs[
364
- req_id] = sampling_params.prompt_logprobs
363
+ self.num_prompt_logprobs[req_id] = (
364
+ self.vocab_size if sampling_params.prompt_logprobs == -1
365
+ else sampling_params.prompt_logprobs)
366
367
if sampling_params.allowed_token_ids:
368
self.has_allowed_token_ids.add(req_id)