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 13, 00:45 UTC
[Core] Switch Flat logprob control from environment variable to SamplingParams (#28914)
Signed-off-by: Jialin Ouyang <Jialin.Ouyang@gmail.com> Co-authored-by: 22quinn <33176974+22quinn@users.noreply.github.com>
Jialin Ouyang committed
Nov 18, 2025 at 18:10 UTC
40b6b38f2c8f8df1dbc145b48df99575f191014f
6 files changed
+33
-41
tests/samplers/test_logprobs.py
+1
-2
@@ -24,9 +24,7 @@ def test_ranks(
24
greedy,
25
flat_logprobs,
26
example_prompts,
27
- monkeypatch: pytest.MonkeyPatch,
27
):
29
- monkeypatch.setenv("VLLM_FLAT_LOGPROBS", "1" if flat_logprobs else "0")
28
with vllm_runner(model, dtype=dtype, max_logprobs=MAX_LOGPROBS) as vllm_model:
29
tokenizer = vllm_model.llm.get_tokenizer()
30
example_prompt_tokens = [tokenizer.encode(prompt) for prompt in example_prompts]
@@ -36,6 +34,7 @@ def test_ranks(
34
max_tokens=MAX_TOKENS,
35
logprobs=NUM_TOP_LOGPROBS,
36
prompt_logprobs=NUM_PROMPT_LOGPROBS,
37
+ flat_logprobs=flat_logprobs,
38
)
39
results = vllm_model.generate_w_logprobs(example_prompts, sampling_params)
40
tests/test_logprobs.py
+10
-22
@@ -2,8 +2,6 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
5
-import pytest
6
-
5
from vllm.logprobs import (
6
FlatLogprobs,
7
Logprob,
@@ -14,24 +12,20 @@ from vllm.logprobs import (
12
)
13
14
17
-def test_create_logprobs_non_flat(monkeypatch: pytest.MonkeyPatch) -> None:
18
- monkeypatch.setenv("VLLM_FLAT_LOGPROBS", "0")
19
-
20
- prompt_logprobs = create_prompt_logprobs()
15
+def test_create_logprobs_non_flat() -> None:
16
+ prompt_logprobs = create_prompt_logprobs(flat_logprobs=False)
17
assert isinstance(prompt_logprobs, list)
18
# Ensure first prompt position logprobs is None
19
assert len(prompt_logprobs) == 1
20
assert prompt_logprobs[0] is None
21
26
- sample_logprobs = create_sample_logprobs()
22
+ sample_logprobs = create_sample_logprobs(flat_logprobs=False)
23
assert isinstance(sample_logprobs, list)
24
assert len(sample_logprobs) == 0
25
26
31
-def test_create_logprobs_flat(monkeypatch: pytest.MonkeyPatch) -> None:
32
- monkeypatch.setenv("VLLM_FLAT_LOGPROBS", "1")
33
-
34
- prompt_logprobs = create_prompt_logprobs()
27
+def test_create_logprobs_flat() -> None:
28
+ prompt_logprobs = create_prompt_logprobs(flat_logprobs=True)
29
assert isinstance(prompt_logprobs, FlatLogprobs)
30
assert prompt_logprobs.start_indices == [0]
31
assert prompt_logprobs.end_indices == [0]
@@ -43,7 +37,7 @@ def test_create_logprobs_flat(monkeypatch: pytest.MonkeyPatch) -> None:
37
assert len(prompt_logprobs) == 1
38
assert prompt_logprobs[0] == dict()
39
46
- sample_logprobs = create_sample_logprobs()
40
+ sample_logprobs = create_sample_logprobs(flat_logprobs=True)
41
assert isinstance(sample_logprobs, FlatLogprobs)
42
assert len(sample_logprobs.start_indices) == 0
43
assert len(sample_logprobs.end_indices) == 0
@@ -54,11 +48,8 @@ def test_create_logprobs_flat(monkeypatch: pytest.MonkeyPatch) -> None:
48
assert len(sample_logprobs) == 0
49
50
57
-def test_append_logprobs_for_next_position_none_flat(
58
- monkeypatch: pytest.MonkeyPatch,
59
-) -> None:
60
- monkeypatch.setenv("VLLM_FLAT_LOGPROBS", "0")
61
- logprobs = create_sample_logprobs()
51
+def test_append_logprobs_for_next_position_none_flat() -> None:
52
+ logprobs = create_sample_logprobs(flat_logprobs=False)
53
append_logprobs_for_next_position(
54
logprobs,
55
token_ids=[1],
@@ -85,11 +76,8 @@ def test_append_logprobs_for_next_position_none_flat(
76
]
77
78
88
-def test_append_logprobs_for_next_position_flat(
89
- monkeypatch: pytest.MonkeyPatch,
90
-) -> None:
91
- monkeypatch.setenv("VLLM_FLAT_LOGPROBS", "1")
92
- logprobs = create_sample_logprobs()
79
+def test_append_logprobs_for_next_position_flat() -> None:
80
+ logprobs = create_sample_logprobs(flat_logprobs=True)
81
append_logprobs_for_next_position(
82
logprobs,
83
token_ids=[1],
vllm/envs.py
-6
@@ -225,7 +225,6 @@ if TYPE_CHECKING:
225
VLLM_DISABLE_SHARED_EXPERTS_STREAM: bool = False
226
VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD: int = 256
227
VLLM_COMPILE_CACHE_SAVE_FORMAT: Literal["binary", "unpacked"] = "binary"
228
- VLLM_FLAT_LOGPROBS: bool = False
228
229
230
def get_default_cache_root():
@@ -1499,11 +1498,6 @@ environment_variables: dict[str, Callable[[], Any]] = {
1498
"VLLM_COMPILE_CACHE_SAVE_FORMAT": env_with_choices(
1499
"VLLM_COMPILE_CACHE_SAVE_FORMAT", "binary", ["binary", "unpacked"]
1500
),
1502
- # Flag to enable FlatLogprobs whose GC overhead is significantly smaller than
1503
- # the original list[dict[int, Logprob]] approach.
1504
- # After enabled, PromptLogprobs and SampleLogprobs would populated as
1505
- # FlatLogprobs.
1506
- "VLLM_FLAT_LOGPROBS": lambda: bool(int(os.getenv("VLLM_FLAT_LOGPROBS", "0"))),
1501
}
1502
1503
# --8<-- [end:env-vars-definition]
vllm/logprobs.py
+4
-6
@@ -5,8 +5,6 @@ from collections.abc import Iterable, Iterator, MutableSequence
5
from dataclasses import dataclass, field
6
from typing import overload
7
8
-import vllm.envs as envs
9
-
8
9
# We use dataclass for now because it is used for
10
# openai server output, and msgspec is not serializable.
@@ -161,17 +159,17 @@ PromptLogprobs = FlatLogprobs | list[LogprobsOnePosition | None]
159
SampleLogprobs = FlatLogprobs | list[LogprobsOnePosition]
160
161
164
-def create_prompt_logprobs() -> PromptLogprobs:
162
+def create_prompt_logprobs(flat_logprobs: bool) -> PromptLogprobs:
163
"""Creates a container to store prompt logprobs for a request"""
166
- logprobs = FlatLogprobs() if envs.VLLM_FLAT_LOGPROBS else []
164
+ logprobs = FlatLogprobs() if flat_logprobs else []
165
# NOTE: logprob of first prompt token is None.
166
logprobs.append(None)
167
return logprobs
168
169
172
-def create_sample_logprobs() -> SampleLogprobs:
170
+def create_sample_logprobs(flat_logprobs: bool) -> SampleLogprobs:
171
"""Creates a container to store decode logprobs for a request"""
174
- return FlatLogprobs() if envs.VLLM_FLAT_LOGPROBS else []
172
+ return FlatLogprobs() if flat_logprobs else []
173
174
175
def append_logprobs_for_next_position(
vllm/sampling_params.py
+6
@@ -204,6 +204,12 @@ class SamplingParams(
204
prompt_logprobs: int | None = None
205
"""Number of log probabilities to return per prompt token.
206
When set to -1, return all `vocab_size` log probabilities."""
207
+ flat_logprobs: bool = False
208
+ """Whether to return logprobs in flatten format (i.e. FlatLogprob)
209
+ for better performance.
210
+ NOTE: GC costs of FlatLogprobs is significantly smaller than
211
+ list[dict[int, Logprob]]. After enabled, PromptLogprobs and
212
+ SampleLogprobs would populated as FlatLogprobs."""
213
# NOTE: This parameter is only exposed at the engine level for now.
214
# It is not exposed in the OpenAI API server, as the OpenAI API does
215
# not support returning only a list of token IDs.
vllm/v1/engine/logprobs.py
+12
-5
@@ -43,15 +43,22 @@ class LogprobsProcessor:
43
tokenizer: AnyTokenizer | None,
44
request: EngineCoreRequest,
45
) -> "LogprobsProcessor":
46
- assert request.sampling_params is not None
47
- num_logprobs = request.sampling_params.logprobs
48
- num_prompt_logprobs = request.sampling_params.prompt_logprobs
46
+ sampling_params = request.sampling_params
47
+ assert sampling_params is not None
48
+ num_logprobs = sampling_params.logprobs
49
+ num_prompt_logprobs = sampling_params.prompt_logprobs
50
return cls(
51
tokenizer=tokenizer,
52
cumulative_logprob=(None if num_logprobs is None else 0.0),
52
- logprobs=(None if num_logprobs is None else create_sample_logprobs()),
53
+ logprobs=(
54
+ None
55
+ if num_logprobs is None
56
+ else create_sample_logprobs(sampling_params.flat_logprobs)
57
+ ),
58
prompt_logprobs=(
54
- None if num_prompt_logprobs is None else create_prompt_logprobs()
59
+ None
60
+ if num_prompt_logprobs is None
61
+ else create_prompt_logprobs(sampling_params.flat_logprobs)
62
),
63
num_prompt_logprobs=num_prompt_logprobs,
64
num_logprobs=num_logprobs,