@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 217 lines 7.8 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
5 from vllm.logprobs import (
6 FlatLogprobs,
7 Logprob,
8 LogprobsOnePosition,
9 append_logprobs_for_next_position,
10 create_prompt_logprobs,
11 create_sample_logprobs,
12 )
13
14
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
22 sample_logprobs = create_sample_logprobs(flat_logprobs=False)
23 assert isinstance(sample_logprobs, list)
24 assert len(sample_logprobs) == 0
25
26
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]
32 assert len(prompt_logprobs.token_ids) == 0
33 assert len(prompt_logprobs.logprobs) == 0
34 assert len(prompt_logprobs.ranks) == 0
35 assert len(prompt_logprobs.decoded_tokens) == 0
36 # Ensure first prompt position logprobs is empty
37 assert len(prompt_logprobs) == 1
38 assert prompt_logprobs[0] == dict()
39
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
44 assert len(sample_logprobs.token_ids) == 0
45 assert len(sample_logprobs.logprobs) == 0
46 assert len(sample_logprobs.ranks) == 0
47 assert len(sample_logprobs.decoded_tokens) == 0
48 assert len(sample_logprobs) == 0
49
50
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],
56 logprobs=[0.1],
57 decoded_tokens=["1"],
58 rank=10,
59 num_logprobs=-1,
60 )
61 append_logprobs_for_next_position(
62 logprobs,
63 token_ids=[2, 3],
64 logprobs=[0.2, 0.3],
65 decoded_tokens=["2", "3"],
66 rank=11,
67 num_logprobs=-1,
68 )
69 assert isinstance(logprobs, list)
70 assert logprobs == [
71 {1: Logprob(logprob=0.1, rank=10, decoded_token="1")},
72 {
73 2: Logprob(logprob=0.2, rank=11, decoded_token="2"),
74 3: Logprob(logprob=0.3, rank=1, decoded_token="3"),
75 },
76 ]
77
78
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],
84 logprobs=[0.1],
85 decoded_tokens=["1"],
86 rank=10,
87 num_logprobs=-1,
88 )
89 append_logprobs_for_next_position(
90 logprobs,
91 token_ids=[2, 3],
92 logprobs=[0.2, 0.3],
93 decoded_tokens=["2", "3"],
94 rank=11,
95 num_logprobs=-1,
96 )
97 assert isinstance(logprobs, FlatLogprobs)
98 assert logprobs.start_indices == [0, 1]
99 assert logprobs.end_indices == [1, 3]
100 assert logprobs.token_ids == [1, 2, 3]
101 assert logprobs.logprobs == [0.1, 0.2, 0.3]
102 assert logprobs.ranks == [10, 11, 1]
103 assert logprobs.decoded_tokens == ["1", "2", "3"]
104
105
106 LOGPROBS_ONE_POSITION_0: LogprobsOnePosition = {
107 1: Logprob(logprob=0.1, rank=10, decoded_token="10")
108 }
109 LOGPROBS_ONE_POSITION_1: LogprobsOnePosition = {
110 2: Logprob(logprob=0.2, rank=20, decoded_token="20"),
111 3: Logprob(logprob=0.3, rank=30, decoded_token="30"),
112 }
113 LOGPROBS_ONE_POSITION_2: LogprobsOnePosition = {
114 4: Logprob(logprob=0.4, rank=40, decoded_token="40"),
115 5: Logprob(logprob=0.5, rank=50, decoded_token="50"),
116 6: Logprob(logprob=0.6, rank=60, decoded_token="60"),
117 }
118
119
120 def test_flat_logprobs_append() -> None:
121 logprobs = FlatLogprobs()
122 logprobs.append(LOGPROBS_ONE_POSITION_0)
123 logprobs.append(LOGPROBS_ONE_POSITION_1)
124 assert logprobs.start_indices == [0, 1]
125 assert logprobs.end_indices == [1, 3]
126 assert logprobs.token_ids == [1, 2, 3]
127 assert logprobs.logprobs == [0.1, 0.2, 0.3]
128 assert logprobs.ranks == [10, 20, 30]
129 assert logprobs.decoded_tokens == ["10", "20", "30"]
130
131 logprobs.append(LOGPROBS_ONE_POSITION_2)
132 assert logprobs.start_indices == [0, 1, 3]
133 assert logprobs.end_indices == [1, 3, 6]
134 assert logprobs.token_ids == [1, 2, 3, 4, 5, 6]
135 assert logprobs.logprobs == [0.1, 0.2, 0.3, 0.4, 0.5, 0.6]
136 assert logprobs.ranks == [10, 20, 30, 40, 50, 60]
137 assert logprobs.decoded_tokens == ["10", "20", "30", "40", "50", "60"]
138
139
140 def test_flat_logprobs_extend() -> None:
141 logprobs = FlatLogprobs()
142 # Extend with list[LogprobsOnePosition]
143 logprobs.extend([LOGPROBS_ONE_POSITION_2, LOGPROBS_ONE_POSITION_0])
144 assert logprobs.start_indices == [0, 3]
145 assert logprobs.end_indices == [3, 4]
146 assert logprobs.token_ids == [4, 5, 6, 1]
147 assert logprobs.logprobs == [0.4, 0.5, 0.6, 0.1]
148 assert logprobs.ranks == [40, 50, 60, 10]
149 assert logprobs.decoded_tokens == ["40", "50", "60", "10"]
150
151 other_logprobs = FlatLogprobs()
152 other_logprobs.extend([LOGPROBS_ONE_POSITION_1, LOGPROBS_ONE_POSITION_0])
153 # Extend with another FlatLogprobs
154 logprobs.extend(other_logprobs)
155 assert logprobs.start_indices == [0, 3, 4, 6]
156 assert logprobs.end_indices == [3, 4, 6, 7]
157 assert logprobs.token_ids == [4, 5, 6, 1, 2, 3, 1]
158 assert logprobs.logprobs == [0.4, 0.5, 0.6, 0.1, 0.2, 0.3, 0.1]
159 assert logprobs.ranks == [40, 50, 60, 10, 20, 30, 10]
160 assert logprobs.decoded_tokens == ["40", "50", "60", "10", "20", "30", "10"]
161
162
163 def test_flat_logprobs_access() -> None:
164 logprobs = FlatLogprobs()
165 logprobs.extend(
166 [LOGPROBS_ONE_POSITION_1, LOGPROBS_ONE_POSITION_2, LOGPROBS_ONE_POSITION_0]
167 )
168 assert logprobs.start_indices == [0, 2, 5]
169 assert logprobs.end_indices == [2, 5, 6]
170 assert logprobs.token_ids == [2, 3, 4, 5, 6, 1]
171 assert logprobs.logprobs == [0.2, 0.3, 0.4, 0.5, 0.6, 0.1]
172 assert logprobs.ranks == [20, 30, 40, 50, 60, 10]
173 assert logprobs.decoded_tokens == ["20", "30", "40", "50", "60", "10"]
174
175 # Test __len__
176 assert len(logprobs) == 3
177
178 # Test __iter__
179 for actual_logprobs, expected_logprobs in zip(
180 logprobs,
181 [LOGPROBS_ONE_POSITION_1, LOGPROBS_ONE_POSITION_2, LOGPROBS_ONE_POSITION_0],
182 ):
183 assert actual_logprobs == expected_logprobs
184
185 # Test __getitem__ : single item
186 assert logprobs[0] == LOGPROBS_ONE_POSITION_1
187 assert logprobs[1] == LOGPROBS_ONE_POSITION_2
188 assert logprobs[2] == LOGPROBS_ONE_POSITION_0
189
190 # Test __getitem__ : slice
191 logprobs02 = logprobs[:2]
192 assert len(logprobs02) == 2
193 assert logprobs02[0] == LOGPROBS_ONE_POSITION_1
194 assert logprobs02[1] == LOGPROBS_ONE_POSITION_2
195 assert logprobs02.start_indices == [0, 2]
196 assert logprobs02.end_indices == [2, 5]
197 assert logprobs02.token_ids == [2, 3, 4, 5, 6]
198 assert logprobs02.logprobs == [0.2, 0.3, 0.4, 0.5, 0.6]
199 assert logprobs02.ranks == [20, 30, 40, 50, 60]
200 assert logprobs02.decoded_tokens == ["20", "30", "40", "50", "60"]
201 logprobs_last2 = logprobs[-2:]
202 assert len(logprobs_last2) == 2
203 assert logprobs_last2[0] == LOGPROBS_ONE_POSITION_2
204 assert logprobs_last2[1] == LOGPROBS_ONE_POSITION_0
205 assert logprobs_last2.start_indices == [0, 3]
206 assert logprobs_last2.end_indices == [3, 4]
207 assert logprobs_last2.token_ids == [4, 5, 6, 1]
208 assert logprobs_last2.logprobs == [0.4, 0.5, 0.6, 0.1]
209 assert logprobs_last2.ranks == [40, 50, 60, 10]
210 assert logprobs_last2.decoded_tokens == ["40", "50", "60", "10"]
211
212 for empty_slice in (slice(0, 0), slice(3, 3), slice(10, 10)):
213 empty = logprobs[empty_slice]
214 assert isinstance(empty, FlatLogprobs)
215 assert len(empty) == 0
216 assert empty.start_indices == []
217 assert empty.end_indices == []