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, 15:45 UTC
[Perf] Introduce FlattenLogprobs to store logprobs results to reduce GC overhead (#28171)
Signed-off-by: Jialin Ouyang <Jialin.Ouyang@gmail.com>
Jialin Ouyang committed
Nov 7, 2025 at 00:27 UTC
ccd98b59c15b50478bb214d6824889570c6f4b8c
6 files changed
+534
-125
tests/samplers/test_logprobs.py
new
+96
@@ -0,0 +1,96 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+
6
+from vllm import SamplingParams
7
+from vllm.logprobs import FlattenLogprobs
8
+
9
+MODELS = ["distilbert/distilgpt2"]
10
+MAX_TOKENS = 5
11
+NUM_TOP_LOGPROBS = 5
12
+NUM_PROMPT_LOGPROBS = 7
13
+MAX_LOGPROBS = max(NUM_TOP_LOGPROBS, NUM_PROMPT_LOGPROBS)
14
+
15
+
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])
20
+def test_ranks(
21
+ vllm_runner,
22
+ model,
23
+ dtype,
24
+ greedy,
25
+ flatten_logprobs,
26
+ example_prompts,
27
+ monkeypatch: pytest.MonkeyPatch,
28
+):
29
+ monkeypatch.setenv("VLLM_FLATTEN_LOGPROBS", "1" if flatten_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]
33
+ sampling_params = SamplingParams(
34
+ temperature=0.0 if greedy else 1.0,
35
+ top_p=1.0,
36
+ max_tokens=MAX_TOKENS,
37
+ logprobs=NUM_TOP_LOGPROBS,
38
+ prompt_logprobs=NUM_PROMPT_LOGPROBS,
39
+ )
40
+ results = vllm_model.generate_w_logprobs(example_prompts, sampling_params)
41
+
42
+ assert len(results) == len(example_prompt_tokens)
43
+ for i, (result, prompt_tokens) in enumerate(zip(results, example_prompt_tokens)):
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
+ )
53
+
54
+ ########################
55
+ # Check prompt logprobs
56
+ ########################
57
+ assert len(prompt_tokens) == len(prompt_logprobs)
58
+ # No logprob for first prompt token
59
+ assert not prompt_logprobs[0]
60
+ for position, (token, logprobs) in enumerate(
61
+ zip(prompt_tokens[1:], prompt_logprobs[1:]), start=1
62
+ ):
63
+ # Ensure logprobs of prompt token is always returned
64
+ logprob = logprobs.get(token)
65
+ assert logprob is not None
66
+ assert logprob.rank >= 1
67
+ # Ensure # of returned logprobs should be
68
+ # either NUM_PROMPT_LOGPROBS or NUM_PROMPT_LOGPROBS+1
69
+ assert NUM_PROMPT_LOGPROBS <= len(logprobs) <= NUM_PROMPT_LOGPROBS + 1
70
+ # Ensure top NUM_PROMPT_LOGPROBS is always extracted
71
+ assert set(range(1, NUM_PROMPT_LOGPROBS + 1)).issubset(
72
+ {logprob.rank for logprob in logprobs.values()}
73
+ )
74
+
75
+ ########################
76
+ # Check sample logprobs
77
+ ########################
78
+ assert len(decode_tokens) == len(decode_logprobs)
79
+ for position, (token, logprobs) in enumerate(
80
+ zip(decode_tokens, decode_logprobs)
81
+ ):
82
+ # Ensure logprobs of chosen token is always returned
83
+ logprob = logprobs.get(token)
84
+ assert logprob is not None
85
+ if greedy:
86
+ # For greedy sampling, all chosen logprob should be top ranked
87
+ assert logprob.rank == 1
88
+ else:
89
+ assert logprob.rank >= 1
90
+ # Ensure # of returned logprobs should be
91
+ # either NUM_TOP_LOGPROBS or NUM_TOP_LOGPROBS+1
92
+ assert NUM_TOP_LOGPROBS <= len(logprobs) <= NUM_TOP_LOGPROBS + 1
93
+ # Ensure top NUM_TOP_LOGPROBS logprobs is always extracted
94
+ assert set(range(1, NUM_TOP_LOGPROBS + 1)).issubset(
95
+ {logprob.rank for logprob in logprobs.values()}
96
+ )
tests/samplers/test_ranks.py
deleted
-59
@@ -1,59 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import pytest
5
-
6
-from vllm import SamplingParams
7
-
8
-MODELS = ["distilbert/distilgpt2"]
9
-
10
-
11
-@pytest.mark.parametrize("model", MODELS)
12
-@pytest.mark.parametrize("dtype", ["half"])
13
-def test_ranks(
14
- vllm_runner,
15
- model,
16
- dtype,
17
- example_prompts,
18
-):
19
- max_tokens = 5
20
- num_top_logprobs = 5
21
- num_prompt_logprobs = 5
22
-
23
- with vllm_runner(model, dtype=dtype, max_logprobs=num_top_logprobs) as vllm_model:
24
- ## Test greedy logprobs ranks
25
- vllm_sampling_params = SamplingParams(
26
- temperature=0.0,
27
- top_p=1.0,
28
- max_tokens=max_tokens,
29
- logprobs=num_top_logprobs,
30
- prompt_logprobs=num_prompt_logprobs,
31
- )
32
- vllm_results = vllm_model.generate_w_logprobs(
33
- example_prompts, vllm_sampling_params
34
- )
35
-
36
- ## Test non-greedy logprobs ranks
37
- sampling_params = SamplingParams(
38
- temperature=1.0,
39
- top_p=1.0,
40
- max_tokens=max_tokens,
41
- logprobs=num_top_logprobs,
42
- prompt_logprobs=num_prompt_logprobs,
43
- )
44
- res = vllm_model.generate_w_logprobs(example_prompts, sampling_params)
45
-
46
- for result in vllm_results:
47
- assert result[2] is not None
48
- assert len(result[2]) == len(result[0])
49
- # check whether all chosen tokens have ranks = 1
50
- for token, logprobs in zip(result[0], result[2]):
51
- assert token in logprobs
52
- assert logprobs[token].rank == 1
53
-
54
- for result in res:
55
- assert result[2] is not None
56
- assert len(result[2]) == len(result[0])
57
- # check whether all chosen tokens have ranks
58
- for token, logprobs in zip(result[0], result[2]):
59
- assert logprobs[token].rank >= 1
tests/test_logprobs.py
new
+222
@@ -0,0 +1,222 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+
5
+import pytest
6
+
7
+from vllm.logprobs import (
8
+ FlattenLogprobs,
9
+ Logprob,
10
+ LogprobsOnePosition,
11
+ append_logprobs_for_next_position,
12
+ create_prompt_logprobs,
13
+ create_sample_logprobs,
14
+)
15
+
16
+
17
+def test_create_logprobs_non_flatten(monkeypatch: pytest.MonkeyPatch) -> None:
18
+ monkeypatch.setenv("VLLM_FLATTEN_LOGPROBS", "0")
19
+
20
+ prompt_logprobs = create_prompt_logprobs()
21
+ assert isinstance(prompt_logprobs, list)
22
+ # Ensure first prompt position logprobs is None
23
+ assert len(prompt_logprobs) == 1
24
+ assert prompt_logprobs[0] is None
25
+
26
+ sample_logprobs = create_sample_logprobs()
27
+ assert isinstance(sample_logprobs, list)
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")
33
+
34
+ prompt_logprobs = create_prompt_logprobs()
35
+ assert isinstance(prompt_logprobs, FlattenLogprobs)
36
+ assert prompt_logprobs.start_indices == [0]
37
+ assert prompt_logprobs.end_indices == [0]
38
+ assert len(prompt_logprobs.token_ids) == 0
39
+ assert len(prompt_logprobs.logprobs) == 0
40
+ assert len(prompt_logprobs.ranks) == 0
41
+ assert len(prompt_logprobs.decoded_tokens) == 0
42
+ # Ensure first prompt position logprobs is empty
43
+ assert len(prompt_logprobs) == 1
44
+ assert prompt_logprobs[0] == dict()
45
+
46
+ sample_logprobs = create_sample_logprobs()
47
+ assert isinstance(sample_logprobs, FlattenLogprobs)
48
+ assert len(sample_logprobs.start_indices) == 0
49
+ assert len(sample_logprobs.end_indices) == 0
50
+ assert len(sample_logprobs.token_ids) == 0
51
+ assert len(sample_logprobs.logprobs) == 0
52
+ assert len(sample_logprobs.ranks) == 0
53
+ assert len(sample_logprobs.decoded_tokens) == 0
54
+ assert len(sample_logprobs) == 0
55
+
56
+
57
+def test_append_logprobs_for_next_position_none_flatten(
58
+ monkeypatch: pytest.MonkeyPatch,
59
+) -> None:
60
+ monkeypatch.setenv("VLLM_FLATTEN_LOGPROBS", "0")
61
+ logprobs = create_sample_logprobs()
62
+ append_logprobs_for_next_position(
63
+ logprobs,
64
+ token_ids=[1],
65
+ logprobs=[0.1],
66
+ decoded_tokens=["1"],
67
+ rank=10,
68
+ num_logprobs=-1,
69
+ )
70
+ append_logprobs_for_next_position(
71
+ logprobs,
72
+ token_ids=[2, 3],
73
+ logprobs=[0.2, 0.3],
74
+ decoded_tokens=["2", "3"],
75
+ rank=11,
76
+ num_logprobs=-1,
77
+ )
78
+ assert isinstance(logprobs, list)
79
+ assert logprobs == [
80
+ {1: Logprob(logprob=0.1, rank=10, decoded_token="1")},
81
+ {
82
+ 2: Logprob(logprob=0.2, rank=11, decoded_token="2"),
83
+ 3: Logprob(logprob=0.3, rank=1, decoded_token="3"),
84
+ },
85
+ ]
86
+
87
+
88
+def test_append_logprobs_for_next_position_flatten(
89
+ monkeypatch: pytest.MonkeyPatch,
90
+) -> None:
91
+ monkeypatch.setenv("VLLM_FLATTEN_LOGPROBS", "1")
92
+ logprobs = create_sample_logprobs()
93
+ append_logprobs_for_next_position(
94
+ logprobs,
95
+ token_ids=[1],
96
+ logprobs=[0.1],
97
+ decoded_tokens=["1"],
98
+ rank=10,
99
+ num_logprobs=-1,
100
+ )
101
+ append_logprobs_for_next_position(
102
+ logprobs,
103
+ token_ids=[2, 3],
104
+ logprobs=[0.2, 0.3],
105
+ decoded_tokens=["2", "3"],
106
+ rank=11,
107
+ num_logprobs=-1,
108
+ )
109
+ assert isinstance(logprobs, FlattenLogprobs)
110
+ assert logprobs.start_indices == [0, 1]
111
+ assert logprobs.end_indices == [1, 3]
112
+ assert logprobs.token_ids == [1, 2, 3]
113
+ assert logprobs.logprobs == [0.1, 0.2, 0.3]
114
+ assert logprobs.ranks == [10, 11, 1]
115
+ assert logprobs.decoded_tokens == ["1", "2", "3"]
116
+
117
+
118
+LOGPROBS_ONE_POSITION_0: LogprobsOnePosition = {
119
+ 1: Logprob(logprob=0.1, rank=10, decoded_token="10")
120
+}
121
+LOGPROBS_ONE_POSITION_1: LogprobsOnePosition = {
122
+ 2: Logprob(logprob=0.2, rank=20, decoded_token="20"),
123
+ 3: Logprob(logprob=0.3, rank=30, decoded_token="30"),
124
+}
125
+LOGPROBS_ONE_POSITION_2: LogprobsOnePosition = {
126
+ 4: Logprob(logprob=0.4, rank=40, decoded_token="40"),
127
+ 5: Logprob(logprob=0.5, rank=50, decoded_token="50"),
128
+ 6: Logprob(logprob=0.6, rank=60, decoded_token="60"),
129
+}
130
+
131
+
132
+def test_flatten_logprobs_append() -> None:
133
+ logprobs = FlattenLogprobs()
134
+ logprobs.append(LOGPROBS_ONE_POSITION_0)
135
+ logprobs.append(LOGPROBS_ONE_POSITION_1)
136
+ assert logprobs.start_indices == [0, 1]
137
+ assert logprobs.end_indices == [1, 3]
138
+ assert logprobs.token_ids == [1, 2, 3]
139
+ assert logprobs.logprobs == [0.1, 0.2, 0.3]
140
+ assert logprobs.ranks == [10, 20, 30]
141
+ assert logprobs.decoded_tokens == ["10", "20", "30"]
142
+
143
+ logprobs.append(LOGPROBS_ONE_POSITION_2)
144
+ assert logprobs.start_indices == [0, 1, 3]
145
+ assert logprobs.end_indices == [1, 3, 6]
146
+ assert logprobs.token_ids == [1, 2, 3, 4, 5, 6]
147
+ assert logprobs.logprobs == [0.1, 0.2, 0.3, 0.4, 0.5, 0.6]
148
+ assert logprobs.ranks == [10, 20, 30, 40, 50, 60]
149
+ assert logprobs.decoded_tokens == ["10", "20", "30", "40", "50", "60"]
150
+
151
+
152
+def test_flatten_logprobs_extend() -> None:
153
+ logprobs = FlattenLogprobs()
154
+ # Extend with list[LogprobsOnePosition]
155
+ logprobs.extend([LOGPROBS_ONE_POSITION_2, LOGPROBS_ONE_POSITION_0])
156
+ assert logprobs.start_indices == [0, 3]
157
+ assert logprobs.end_indices == [3, 4]
158
+ assert logprobs.token_ids == [4, 5, 6, 1]
159
+ assert logprobs.logprobs == [0.4, 0.5, 0.6, 0.1]
160
+ assert logprobs.ranks == [40, 50, 60, 10]
161
+ assert logprobs.decoded_tokens == ["40", "50", "60", "10"]
162
+
163
+ other_logprobs = FlattenLogprobs()
164
+ other_logprobs.extend([LOGPROBS_ONE_POSITION_1, LOGPROBS_ONE_POSITION_0])
165
+ # Extend with another FlattenLogprobs
166
+ logprobs.extend(other_logprobs)
167
+ assert logprobs.start_indices == [0, 3, 4, 6]
168
+ assert logprobs.end_indices == [3, 4, 6, 7]
169
+ assert logprobs.token_ids == [4, 5, 6, 1, 2, 3, 1]
170
+ assert logprobs.logprobs == [0.4, 0.5, 0.6, 0.1, 0.2, 0.3, 0.1]
171
+ assert logprobs.ranks == [40, 50, 60, 10, 20, 30, 10]
172
+ assert logprobs.decoded_tokens == ["40", "50", "60", "10", "20", "30", "10"]
173
+
174
+
175
+def test_flatten_logprobs_access() -> None:
176
+ logprobs = FlattenLogprobs()
177
+ logprobs.extend(
178
+ [LOGPROBS_ONE_POSITION_1, LOGPROBS_ONE_POSITION_2, LOGPROBS_ONE_POSITION_0]
179
+ )
180
+ assert logprobs.start_indices == [0, 2, 5]
181
+ assert logprobs.end_indices == [2, 5, 6]
182
+ assert logprobs.token_ids == [2, 3, 4, 5, 6, 1]
183
+ assert logprobs.logprobs == [0.2, 0.3, 0.4, 0.5, 0.6, 0.1]
184
+ assert logprobs.ranks == [20, 30, 40, 50, 60, 10]
185
+ assert logprobs.decoded_tokens == ["20", "30", "40", "50", "60", "10"]
186
+
187
+ # Test __len__
188
+ assert len(logprobs) == 3
189
+
190
+ # Test __iter__
191
+ for actual_logprobs, expected_logprobs in zip(
192
+ logprobs,
193
+ [LOGPROBS_ONE_POSITION_1, LOGPROBS_ONE_POSITION_2, LOGPROBS_ONE_POSITION_0],
194
+ ):
195
+ assert actual_logprobs == expected_logprobs
196
+
197
+ # Test __getitem__ : single item
198
+ assert logprobs[0] == LOGPROBS_ONE_POSITION_1
199
+ assert logprobs[1] == LOGPROBS_ONE_POSITION_2
200
+ assert logprobs[2] == LOGPROBS_ONE_POSITION_0
201
+
202
+ # Test __getitem__ : slice
203
+ logprobs02 = logprobs[:2]
204
+ assert len(logprobs02) == 2
205
+ assert logprobs02[0] == LOGPROBS_ONE_POSITION_1
206
+ assert logprobs02[1] == LOGPROBS_ONE_POSITION_2
207
+ assert logprobs02.start_indices == [0, 2]
208
+ assert logprobs02.end_indices == [2, 5]
209
+ assert logprobs02.token_ids == [2, 3, 4, 5, 6]
210
+ assert logprobs02.logprobs == [0.2, 0.3, 0.4, 0.5, 0.6]
211
+ assert logprobs02.ranks == [20, 30, 40, 50, 60]
212
+ assert logprobs02.decoded_tokens == ["20", "30", "40", "50", "60"]
213
+ logprobs_last2 = logprobs[-2:]
214
+ assert len(logprobs_last2) == 2
215
+ assert logprobs_last2[0] == LOGPROBS_ONE_POSITION_2
216
+ assert logprobs_last2[1] == LOGPROBS_ONE_POSITION_0
217
+ assert logprobs_last2.start_indices == [0, 3]
218
+ assert logprobs_last2.end_indices == [3, 4]
219
+ assert logprobs_last2.token_ids == [4, 5, 6, 1]
220
+ assert logprobs_last2.logprobs == [0.4, 0.5, 0.6, 0.1]
221
+ assert logprobs_last2.ranks == [40, 50, 60, 10]
222
+ assert logprobs_last2.decoded_tokens == ["40", "50", "60", "10"]
vllm/envs.py
+6
@@ -220,6 +220,7 @@ if TYPE_CHECKING:
220
VLLM_GC_DEBUG: str = ""
221
VLLM_DISABLE_SHARED_EXPERTS_STREAM: bool = False
222
VLLM_COMPILE_CACHE_SAVE_FORMAT: Literal["binary", "unpacked"] = "binary"
223
+ VLLM_FLATTEN_LOGPROBS: bool = False
224
225
226
def get_default_cache_root():
@@ -1463,6 +1464,11 @@ environment_variables: dict[str, Callable[[], Any]] = {
1464
"VLLM_COMPILE_CACHE_SAVE_FORMAT": env_with_choices(
1465
"VLLM_COMPILE_CACHE_SAVE_FORMAT", "binary", ["binary", "unpacked"]
1466
),
1467
+ # Flag to enable FlattenLogprobs whose GC overhead is significantly smaller than
1468
+ # the original list[dict[int, Logprob]] approach.
1469
+ # After enabled, PromptLogprobs and SampleLogprobs would populated as
1470
+ # FlattenLogprobs.
1471
+ "VLLM_FLATTEN_LOGPROBS": lambda: bool(int(os.getenv("VLLM_FLATTEN_LOGPROBS", "0"))),
1472
}
1473
1474
# --8<-- [end:env-vars-definition]
vllm/logprobs.py
+183
-3
@@ -1,6 +1,11 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from dataclasses import dataclass
3
+import itertools
4
+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
10
11
# We use dataclass for now because it is used for
@@ -21,8 +26,183 @@ class Logprob:
26
decoded_token: str | None = None
27
28
29
+LogprobsOnePosition = dict[int, Logprob]
30
+
31
+
32
+@dataclass
33
+class FlattenLogprobs(MutableSequence[LogprobsOnePosition]):
34
+ """
35
+ Flatten 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.
43
+
44
+ As each position might contains different amount of ranks,
45
+ start_indices_per_position would be used to access the logprob ranges
46
+ for different positions.
47
+
48
+ NOTE: To reduce the migration overhead and improve backward compatibility,
49
+ we support the key Sequence APIs of list, so it could act as
50
+ list[LogprobsOnePosition]
51
+ """
52
+
53
+ # Start / end indices to indicate the range of logprobs for each position.
54
+ start_indices: list[int] = field(default_factory=list)
55
+ end_indices: list[int] = field(default_factory=list)
56
+
57
+ # Flatten Logprob information for (each position, rank).
58
+ # For position <i>, the logprobs are ranged
59
+ # from self.start_indices[i] to self.end_indices[i] (exclusive).
60
+ token_ids: list[int] = field(default_factory=list)
61
+ logprobs: list[float] = field(default_factory=list)
62
+ ranks: list[int | None] = field(default_factory=list)
63
+ decoded_tokens: list[str | None] = field(default_factory=list)
64
+
65
+ def append(self, logprobs_one_position: LogprobsOnePosition | None) -> None:
66
+ """Appends the container with logprobs for the next position"""
67
+ self.start_indices.append(len(self.logprobs))
68
+ if logprobs_one_position:
69
+ for token_id, logprob in logprobs_one_position.items():
70
+ self.token_ids.append(token_id)
71
+ self.logprobs.append(logprob.logprob)
72
+ self.ranks.append(logprob.rank)
73
+ self.decoded_tokens.append(logprob.decoded_token)
74
+ self.end_indices.append(len(self.logprobs))
75
+
76
+ def append_fast(
77
+ self,
78
+ token_ids: list[int],
79
+ logprobs: list[float],
80
+ ranks: itertools.chain[int],
81
+ decoded_tokens: Iterable[str | None],
82
+ ) -> None:
83
+ """
84
+ Appends logprobs for the next position without creating
85
+ the intermediate logprob dictionary.
86
+ """
87
+ self.start_indices.append(len(self.logprobs))
88
+ for token_id, logprob, rank, decoded_token in zip(
89
+ token_ids, logprobs, ranks, decoded_tokens
90
+ ):
91
+ self.token_ids.append(token_id)
92
+ self.logprobs.append(logprob)
93
+ self.ranks.append(rank)
94
+ self.decoded_tokens.append(decoded_token)
95
+ self.end_indices.append(len(self.logprobs))
96
+
97
+ def extend(self, logprobs_multi_positions) -> None:
98
+ """Extends the container with logprobs for the next multiple positions"""
99
+ for logprobs_one_position in logprobs_multi_positions:
100
+ self.append(logprobs_one_position)
101
+
102
+ def __len__(self) -> int:
103
+ """Gets number of positions stored in the container"""
104
+ return len(self.start_indices)
105
+
106
+ @overload
107
+ def __getitem__(self, position: int) -> LogprobsOnePosition: ...
108
+
109
+ @overload
110
+ def __getitem__(self, s: slice, /) -> "FlattenLogprobs": ...
111
+
112
+ def __getitem__(self, index: int | slice):
113
+ """Extracts logprobs of a given position or slice"""
114
+ if isinstance(index, int):
115
+ return {
116
+ self.token_ids[i]: Logprob(
117
+ logprob=self.logprobs[i],
118
+ rank=self.ranks[i],
119
+ decoded_token=self.decoded_tokens[i],
120
+ )
121
+ for i in range(self.start_indices[index], self.end_indices[index])
122
+ }
123
+ elif isinstance(index, slice):
124
+ min_index = self.start_indices[index][0]
125
+ max_index = self.end_indices[index][-1]
126
+ return FlattenLogprobs(
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]],
130
+ end_indices=[i - min_index for i in self.end_indices[index]],
131
+ token_ids=self.token_ids[min_index:max_index],
132
+ logprobs=self.logprobs[min_index:max_index],
133
+ ranks=self.ranks[min_index:max_index],
134
+ decoded_tokens=self.decoded_tokens[min_index:max_index],
135
+ )
136
+ else:
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")
141
+
142
+ def __delitem__(self, item) -> None:
143
+ raise TypeError("Cannot delete logprobs from FlattenLogprobs")
144
+
145
+ def insert(self, item) -> None:
146
+ raise TypeError("Cannot insert logprobs to FlattenLogprobs")
147
+
148
+ def __iter__(self) -> Iterator[LogprobsOnePosition]:
149
+ """
150
+ Iterates the container and yields LogprobsOnePosition for
151
+ each position.
152
+ """
153
+ for i in range(0, len(self.start_indices)):
154
+ yield self.__getitem__(i)
155
+
156
+
157
# {token_id -> logprob} per each sequence group. None if the corresponding
158
# sequence group doesn't require prompt logprob.
26
-PromptLogprobs = list[dict[int, Logprob] | None]
159
+PromptLogprobs = FlattenLogprobs | list[LogprobsOnePosition | None]
160
# {token_id -> logprob} for each sequence group.
28
-SampleLogprobs = list[dict[int, Logprob]]
161
+SampleLogprobs = FlattenLogprobs | 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 []
167
+ # NOTE: logprob of first prompt token is None.
168
+ logprobs.append(None)
169
+ return logprobs
170
+
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 []
175
+
176
+
177
+def append_logprobs_for_next_position(
178
+ request_logprobs: PromptLogprobs | SampleLogprobs,
179
+ token_ids: list[int],
180
+ logprobs: list[float],
181
+ decoded_tokens: Iterable[str | None],
182
+ rank: int,
183
+ num_logprobs: int,
184
+) -> None:
185
+ """Appends logprobs for the next position"""
186
+ if num_logprobs == -1:
187
+ num_logprobs = len(logprobs)
188
+ # We do not need a special case for the sampled token
189
+ # being in the topk, since inserting duplicated data
190
+ # into a dictionary twice is the same as doing it once.
191
+ topk_ranks = range(1, num_logprobs + 1)
192
+ ranks = itertools.chain((rank,), topk_ranks)
193
+
194
+ if isinstance(request_logprobs, FlattenLogprobs):
195
+ request_logprobs.append_fast(token_ids, logprobs, ranks, decoded_tokens)
196
+ else:
197
+ request_logprobs.append(
198
+ {
199
+ token_id: Logprob(
200
+ logprob=logprob,
201
+ rank=rank,
202
+ decoded_token=token,
203
+ )
204
+ for token_id, logprob, rank, token in zip(
205
+ token_ids, logprobs, ranks, decoded_tokens
206
+ )
207
+ }
208
+ )
vllm/v1/engine/logprobs.py
+27
-63
@@ -2,11 +2,16 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import itertools
5
-from collections.abc import Iterable
5
from dataclasses import dataclass
6
7
from vllm.logger import init_logger
9
-from vllm.logprobs import Logprob, PromptLogprobs, SampleLogprobs
8
+from vllm.logprobs import (
9
+ PromptLogprobs,
10
+ SampleLogprobs,
11
+ append_logprobs_for_next_position,
12
+ create_prompt_logprobs,
13
+ create_sample_logprobs,
14
+)
15
from vllm.transformers_utils.detokenizer_utils import (
16
AnyTokenizer,
17
convert_ids_list_to_tokens,
@@ -44,9 +49,10 @@ class LogprobsProcessor:
49
return cls(
50
tokenizer=tokenizer,
51
cumulative_logprob=(None if num_logprobs is None else 0.0),
47
- logprobs=(None if num_logprobs is None else []),
48
- # NOTE: logprob of first prompt token is None.
49
- prompt_logprobs=(None if num_prompt_logprobs is None else [None]),
52
+ logprobs=(None if num_logprobs is None else create_sample_logprobs()),
53
+ prompt_logprobs=(
54
+ None if num_prompt_logprobs is None else create_prompt_logprobs()
55
+ ),
56
num_prompt_logprobs=num_prompt_logprobs,
57
num_logprobs=num_logprobs,
58
)
@@ -80,15 +86,14 @@ class LogprobsProcessor:
86
sampled_token_logprob = logprobs[0]
87
self.cumulative_logprob += sampled_token_logprob
88
83
- # Update with the Logprob dictionary for this pos.
84
- self.logprobs.append(
85
- self._make_logprob_dict(
86
- logprobs,
87
- token_ids,
88
- decoded_tokens,
89
- rank,
90
- self.num_logprobs,
91
- )
89
+ # Update with the Logprob container for this pos.
90
+ append_logprobs_for_next_position(
91
+ self.logprobs,
92
+ token_ids,
93
+ logprobs,
94
+ decoded_tokens,
95
+ rank,
96
+ self.num_logprobs,
97
)
98
99
def _update_prompt_logprobs(
@@ -136,15 +141,14 @@ class LogprobsProcessor:
141
NONES if decoded_tokens is None else decoded_tokens[offset:offset_end]
142
)
143
139
- # Update with the Logprob dictionary for this pos.
140
- self.prompt_logprobs.append(
141
- self._make_logprob_dict(
142
- prompt_logprobs[pos],
143
- token_ids[pos],
144
- decoded_tokens_for_pos,
145
- prompt_token_ranks[pos],
146
- self.num_prompt_logprobs,
147
- )
144
+ # Update with the Logprob container for this pos.
145
+ append_logprobs_for_next_position(
146
+ self.prompt_logprobs,
147
+ token_ids[pos],
148
+ prompt_logprobs[pos],
149
+ decoded_tokens_for_pos,
150
+ prompt_token_ranks[pos],
151
+ self.num_prompt_logprobs,
152
)
153
154
def pop_prompt_logprobs(self) -> PromptLogprobs | None:
@@ -166,46 +170,6 @@ class LogprobsProcessor:
170
self.prompt_logprobs = []
171
return plp
172
169
- @staticmethod
170
- def _make_logprob_dict(
171
- logprobs: list[float],
172
- logprob_token_ids: list[int],
173
- decoded_tokens: Iterable[str | None],
174
- rank: int,
175
- num_logprobs: int,
176
- ) -> dict[int, Logprob]:
177
- """Make a Logprob dictionary for a position.
178
-
179
- Args:
180
- logprobs: list of log probabilities
181
- logprob_token_ids: list of top token ids
182
- decoded_tokens: list of decoded top tokens
183
- rank: rank of the sampled token
184
- num_logprobs: number of logprobs requested
185
- by the user (in addition to sampled logprob)
186
-
187
- Returns:
188
- dict[token id, Logprob]
189
- """
190
- if num_logprobs == -1:
191
- num_logprobs = len(logprobs)
192
- # We do not need a special case for the sampled token
193
- # being in the topk, since inserting duplicated data
194
- # into a dictionary twice is the same as doing it once.
195
- topk_ranks = range(1, num_logprobs + 1)
196
- ranks = itertools.chain((rank,), topk_ranks)
197
-
198
- return {
199
- token_id: Logprob(
200
- logprob=logprob,
201
- rank=rank,
202
- decoded_token=token,
203
- )
204
- for token_id, logprob, rank, token in zip(
205
- logprob_token_ids, logprobs, ranks, decoded_tokens
206
- )
207
- }
208
-
173
def update_from_output(self, output: EngineCoreOutput) -> None:
174
if output.new_logprobs is not None:
175
self._update_sample_logprobs(output.new_logprobs)