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, 13:15 UTC
[Misc] FlattenLogprobs -> FlatLogprobs (#28335)
Zhuohan Li committed
Nov 10, 2025 at 19:41 UTC
8d706cca903a008169e7ac8f1dc1f65c8ffd85c0
4 files changed
+43
-47
tests/samplers/test_logprobs.py
+6
-10
@@ -4,7 +4,7 @@
4
import pytest
5
6
from vllm import SamplingParams
7
-from vllm.logprobs import FlattenLogprobs
7
+from vllm.logprobs import FlatLogprobs
8
9
MODELS = ["distilbert/distilgpt2"]
10
MAX_TOKENS = 5
@@ -16,17 +16,17 @@ MAX_LOGPROBS = max(NUM_TOP_LOGPROBS, NUM_PROMPT_LOGPROBS)
16
@pytest.mark.parametrize("model", MODELS)
17
@pytest.mark.parametrize("dtype", ["half"])
18
@pytest.mark.parametrize("greedy", [True, False])
19
-@pytest.mark.parametrize("flatten_logprobs", [True, False])
19
+@pytest.mark.parametrize("flat_logprobs", [True, False])
20
def test_ranks(
21
vllm_runner,
22
model,
23
dtype,
24
greedy,
25
- flatten_logprobs,
25
+ flat_logprobs,
26
example_prompts,
27
monkeypatch: pytest.MonkeyPatch,
28
):
29
- monkeypatch.setenv("VLLM_FLATTEN_LOGPROBS", "1" if flatten_logprobs else "0")
29
+ monkeypatch.setenv("VLLM_FLAT_LOGPROBS", "1" if flat_logprobs else "0")
30
with vllm_runner(model, dtype=dtype, max_logprobs=MAX_LOGPROBS) as vllm_model:
31
tokenizer = vllm_model.llm.get_tokenizer()
32
example_prompt_tokens = [tokenizer.encode(prompt) for prompt in example_prompts]
@@ -44,12 +44,8 @@ def test_ranks(
44
decode_tokens, _, decode_logprobs, prompt_logprobs = result
45
46
# Ensure the return type of logprobs is accurate
47
- assert isinstance(
48
- prompt_logprobs, FlattenLogprobs if flatten_logprobs else list
49
- )
50
- assert isinstance(
51
- decode_logprobs, FlattenLogprobs if flatten_logprobs else list
52
- )
47
+ assert isinstance(prompt_logprobs, FlatLogprobs if flat_logprobs else list)
48
+ assert isinstance(decode_logprobs, FlatLogprobs if flat_logprobs else list)
49
50
########################
51
# Check prompt logprobs
tests/test_logprobs.py
+20
-20
@@ -5,7 +5,7 @@
5
import pytest
6
7
from vllm.logprobs import (
8
- FlattenLogprobs,
8
+ FlatLogprobs,
9
Logprob,
10
LogprobsOnePosition,
11
append_logprobs_for_next_position,
@@ -14,8 +14,8 @@ from vllm.logprobs import (
14
)
15
16
17
-def test_create_logprobs_non_flatten(monkeypatch: pytest.MonkeyPatch) -> None:
18
- monkeypatch.setenv("VLLM_FLATTEN_LOGPROBS", "0")
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()
21
assert isinstance(prompt_logprobs, list)
@@ -28,11 +28,11 @@ def test_create_logprobs_non_flatten(monkeypatch: pytest.MonkeyPatch) -> None:
28
assert len(sample_logprobs) == 0
29
30
31
-def test_create_logprobs_flatten(monkeypatch: pytest.MonkeyPatch) -> None:
32
- monkeypatch.setenv("VLLM_FLATTEN_LOGPROBS", "1")
31
+def test_create_logprobs_flat(monkeypatch: pytest.MonkeyPatch) -> None:
32
+ monkeypatch.setenv("VLLM_FLAT_LOGPROBS", "1")
33
34
prompt_logprobs = create_prompt_logprobs()
35
- assert isinstance(prompt_logprobs, FlattenLogprobs)
35
+ assert isinstance(prompt_logprobs, FlatLogprobs)
36
assert prompt_logprobs.start_indices == [0]
37
assert prompt_logprobs.end_indices == [0]
38
assert len(prompt_logprobs.token_ids) == 0
@@ -44,7 +44,7 @@ def test_create_logprobs_flatten(monkeypatch: pytest.MonkeyPatch) -> None:
44
assert prompt_logprobs[0] == dict()
45
46
sample_logprobs = create_sample_logprobs()
47
- assert isinstance(sample_logprobs, FlattenLogprobs)
47
+ assert isinstance(sample_logprobs, FlatLogprobs)
48
assert len(sample_logprobs.start_indices) == 0
49
assert len(sample_logprobs.end_indices) == 0
50
assert len(sample_logprobs.token_ids) == 0
@@ -54,10 +54,10 @@ def test_create_logprobs_flatten(monkeypatch: pytest.MonkeyPatch) -> None:
54
assert len(sample_logprobs) == 0
55
56
57
-def test_append_logprobs_for_next_position_none_flatten(
57
+def test_append_logprobs_for_next_position_none_flat(
58
monkeypatch: pytest.MonkeyPatch,
59
) -> None:
60
- monkeypatch.setenv("VLLM_FLATTEN_LOGPROBS", "0")
60
+ monkeypatch.setenv("VLLM_FLAT_LOGPROBS", "0")
61
logprobs = create_sample_logprobs()
62
append_logprobs_for_next_position(
63
logprobs,
@@ -85,10 +85,10 @@ def test_append_logprobs_for_next_position_none_flatten(
85
]
86
87
88
-def test_append_logprobs_for_next_position_flatten(
88
+def test_append_logprobs_for_next_position_flat(
89
monkeypatch: pytest.MonkeyPatch,
90
) -> None:
91
- monkeypatch.setenv("VLLM_FLATTEN_LOGPROBS", "1")
91
+ monkeypatch.setenv("VLLM_FLAT_LOGPROBS", "1")
92
logprobs = create_sample_logprobs()
93
append_logprobs_for_next_position(
94
logprobs,
@@ -106,7 +106,7 @@ def test_append_logprobs_for_next_position_flatten(
106
rank=11,
107
num_logprobs=-1,
108
)
109
- assert isinstance(logprobs, FlattenLogprobs)
109
+ assert isinstance(logprobs, FlatLogprobs)
110
assert logprobs.start_indices == [0, 1]
111
assert logprobs.end_indices == [1, 3]
112
assert logprobs.token_ids == [1, 2, 3]
@@ -129,8 +129,8 @@ LOGPROBS_ONE_POSITION_2: LogprobsOnePosition = {
129
}
130
131
132
-def test_flatten_logprobs_append() -> None:
133
- logprobs = FlattenLogprobs()
132
+def test_flat_logprobs_append() -> None:
133
+ logprobs = FlatLogprobs()
134
logprobs.append(LOGPROBS_ONE_POSITION_0)
135
logprobs.append(LOGPROBS_ONE_POSITION_1)
136
assert logprobs.start_indices == [0, 1]
@@ -149,8 +149,8 @@ def test_flatten_logprobs_append() -> None:
149
assert logprobs.decoded_tokens == ["10", "20", "30", "40", "50", "60"]
150
151
152
-def test_flatten_logprobs_extend() -> None:
153
- logprobs = FlattenLogprobs()
152
+def test_flat_logprobs_extend() -> None:
153
+ logprobs = FlatLogprobs()
154
# Extend with list[LogprobsOnePosition]
155
logprobs.extend([LOGPROBS_ONE_POSITION_2, LOGPROBS_ONE_POSITION_0])
156
assert logprobs.start_indices == [0, 3]
@@ -160,9 +160,9 @@ def test_flatten_logprobs_extend() -> None:
160
assert logprobs.ranks == [40, 50, 60, 10]
161
assert logprobs.decoded_tokens == ["40", "50", "60", "10"]
162
163
- other_logprobs = FlattenLogprobs()
163
+ other_logprobs = FlatLogprobs()
164
other_logprobs.extend([LOGPROBS_ONE_POSITION_1, LOGPROBS_ONE_POSITION_0])
165
- # Extend with another FlattenLogprobs
165
+ # Extend with another FlatLogprobs
166
logprobs.extend(other_logprobs)
167
assert logprobs.start_indices == [0, 3, 4, 6]
168
assert logprobs.end_indices == [3, 4, 6, 7]
@@ -172,8 +172,8 @@ def test_flatten_logprobs_extend() -> None:
172
assert logprobs.decoded_tokens == ["40", "50", "60", "10", "20", "30", "10"]
173
174
175
-def test_flatten_logprobs_access() -> None:
176
- logprobs = FlattenLogprobs()
175
+def test_flat_logprobs_access() -> None:
176
+ logprobs = FlatLogprobs()
177
logprobs.extend(
178
[LOGPROBS_ONE_POSITION_1, LOGPROBS_ONE_POSITION_2, LOGPROBS_ONE_POSITION_0]
179
)
vllm/envs.py
+4
-4
@@ -223,7 +223,7 @@ if TYPE_CHECKING:
223
VLLM_GC_DEBUG: str = ""
224
VLLM_DISABLE_SHARED_EXPERTS_STREAM: bool = False
225
VLLM_COMPILE_CACHE_SAVE_FORMAT: Literal["binary", "unpacked"] = "binary"
226
- VLLM_FLATTEN_LOGPROBS: bool = False
226
+ VLLM_FLAT_LOGPROBS: bool = False
227
228
229
def get_default_cache_root():
@@ -1481,11 +1481,11 @@ environment_variables: dict[str, Callable[[], Any]] = {
1481
"VLLM_COMPILE_CACHE_SAVE_FORMAT": env_with_choices(
1482
"VLLM_COMPILE_CACHE_SAVE_FORMAT", "binary", ["binary", "unpacked"]
1483
),
1484
- # Flag to enable FlattenLogprobs whose GC overhead is significantly smaller than
1484
+ # Flag to enable FlatLogprobs whose GC overhead is significantly smaller than
1485
# the original list[dict[int, Logprob]] approach.
1486
# After enabled, PromptLogprobs and SampleLogprobs would populated as
1487
- # FlattenLogprobs.
1488
- "VLLM_FLATTEN_LOGPROBS": lambda: bool(int(os.getenv("VLLM_FLATTEN_LOGPROBS", "0"))),
1487
+ # FlatLogprobs.
1488
+ "VLLM_FLAT_LOGPROBS": lambda: bool(int(os.getenv("VLLM_FLAT_LOGPROBS", "0"))),
1489
}
1490
1491
# --8<-- [end:env-vars-definition]
vllm/logprobs.py
+13
-13
@@ -30,16 +30,16 @@ LogprobsOnePosition = dict[int, Logprob]
30
31
32
@dataclass
33
-class FlattenLogprobs(MutableSequence[LogprobsOnePosition]):
33
+class FlatLogprobs(MutableSequence[LogprobsOnePosition]):
34
"""
35
- Flatten logprobs of a request into multiple primitive type lists.
35
+ Flat logprobs of a request into multiple primitive type lists.
36
37
Compared to list[dict[int, Logprob]], this data structure reduced GC
38
overhead significantly. As it flattened logprob information for
39
all positions and ranks in to multiple primitive type lists (i.e.
40
logprobs, token_ids, ranks per token_ids, decoded_tokens).
41
So regardless of the sequence length and top_logprobs setup,
42
- FlattenLogprobs would only introduce a constant amount of objects.
42
+ FlatLogprobs would only introduce a constant amount of objects.
43
44
As each position might contains different amount of ranks,
45
start_indices_per_position would be used to access the logprob ranges
@@ -107,7 +107,7 @@ class FlattenLogprobs(MutableSequence[LogprobsOnePosition]):
107
def __getitem__(self, position: int) -> LogprobsOnePosition: ...
108
109
@overload
110
- def __getitem__(self, s: slice, /) -> "FlattenLogprobs": ...
110
+ def __getitem__(self, s: slice, /) -> "FlatLogprobs": ...
111
112
def __getitem__(self, index: int | slice):
113
"""Extracts logprobs of a given position or slice"""
@@ -123,7 +123,7 @@ class FlattenLogprobs(MutableSequence[LogprobsOnePosition]):
123
elif isinstance(index, slice):
124
min_index = self.start_indices[index][0]
125
max_index = self.end_indices[index][-1]
126
- return FlattenLogprobs(
126
+ return FlatLogprobs(
127
# Shift updated start_indices and end_indices to
128
# be 0-indexed
129
start_indices=[i - min_index for i in self.start_indices[index]],
@@ -137,13 +137,13 @@ class FlattenLogprobs(MutableSequence[LogprobsOnePosition]):
137
raise TypeError(f"Invalid index type: {type(index)}")
138
139
def __setitem__(self, item, value) -> None:
140
- raise TypeError("Cannot set logprobs in FlattenLogprobs")
140
+ raise TypeError("Cannot set logprobs in FlatLogprobs")
141
142
def __delitem__(self, item) -> None:
143
- raise TypeError("Cannot delete logprobs from FlattenLogprobs")
143
+ raise TypeError("Cannot delete logprobs from FlatLogprobs")
144
145
def insert(self, item) -> None:
146
- raise TypeError("Cannot insert logprobs to FlattenLogprobs")
146
+ raise TypeError("Cannot insert logprobs to FlatLogprobs")
147
148
def __iter__(self) -> Iterator[LogprobsOnePosition]:
149
"""
@@ -156,14 +156,14 @@ class FlattenLogprobs(MutableSequence[LogprobsOnePosition]):
156
157
# {token_id -> logprob} per each sequence group. None if the corresponding
158
# sequence group doesn't require prompt logprob.
159
-PromptLogprobs = FlattenLogprobs | list[LogprobsOnePosition | None]
159
+PromptLogprobs = FlatLogprobs | list[LogprobsOnePosition | None]
160
# {token_id -> logprob} for each sequence group.
161
-SampleLogprobs = FlattenLogprobs | list[LogprobsOnePosition]
161
+SampleLogprobs = FlatLogprobs | list[LogprobsOnePosition]
162
163
164
def create_prompt_logprobs() -> PromptLogprobs:
165
"""Creates a container to store prompt logprobs for a request"""
166
- logprobs = FlattenLogprobs() if envs.VLLM_FLATTEN_LOGPROBS else []
166
+ logprobs = FlatLogprobs() if envs.VLLM_FLAT_LOGPROBS else []
167
# NOTE: logprob of first prompt token is None.
168
logprobs.append(None)
169
return logprobs
@@ -171,7 +171,7 @@ def create_prompt_logprobs() -> PromptLogprobs:
171
172
def create_sample_logprobs() -> SampleLogprobs:
173
"""Creates a container to store decode logprobs for a request"""
174
- return FlattenLogprobs() if envs.VLLM_FLATTEN_LOGPROBS else []
174
+ return FlatLogprobs() if envs.VLLM_FLAT_LOGPROBS else []
175
176
177
def append_logprobs_for_next_position(
@@ -191,7 +191,7 @@ def append_logprobs_for_next_position(
191
topk_ranks = range(1, num_logprobs + 1)
192
ranks = itertools.chain((rank,), topk_ranks)
193
194
- if isinstance(request_logprobs, FlattenLogprobs):
194
+ if isinstance(request_logprobs, FlatLogprobs):
195
request_logprobs.append_fast(token_ids, logprobs, ranks, decoded_tokens)
196
else:
197
request_logprobs.append(