@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 208 lines 6.82 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """End-to-end tests for the `logprob_token_ids` field on the OpenAI-compat
4 chat-completion endpoint.
5
6 `logprob_token_ids` lets a caller pin the set of vocab ids whose logprobs
7 should appear in the response, independent of where those ids would rank in
8 the natural top-k distribution. This is the primitive that multilabel
9 scoring postprocessors use to gather logprobs at a fixed small label
10 vocabulary (e.g. PII detection where each label corresponds to a known
11 digit-token vocab id).
12 """
13
14 import math
15
16 import pytest
17
18 from tests.utils import RemoteOpenAIServer
19 from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
20 from vllm.entrypoints.openai.completion.protocol import CompletionRequest
21 from vllm.exceptions import VLLMValidationError
22
23 MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct"
24
25
26 @pytest.fixture(scope="module")
27 def server():
28 args = [
29 "--max-model-len",
30 "1024",
31 "--max-num-seqs",
32 "8",
33 "--enforce-eager",
34 ]
35 with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
36 yield remote_server
37
38
39 def _logprob_entries(resp) -> list:
40 """Return the top_logprobs list from the first generated position."""
41 return resp.choices[0].logprobs.content[0].top_logprobs
42
43
44 def _token_id(token: str) -> int:
45 assert token.startswith("token_id:"), (
46 "expected return_tokens_as_token_ids=True to yield the "
47 f"`token_id:<int>` form, got {token!r}"
48 )
49 return int(token.removeprefix("token_id:"))
50
51
52 def _top_logprob_token_ids(resp) -> list[int]:
53 return [_token_id(entry.token) for entry in _logprob_entries(resp)]
54
55
56 def test_chat_request_decouples_top_k_from_explicit_token_ids():
57 request = ChatCompletionRequest(
58 model=MODEL_NAME,
59 messages=[{"role": "user", "content": "Hello"}],
60 logprobs=True,
61 top_logprobs=5,
62 logprob_token_ids=[5000],
63 )
64
65 sampling_params = request.to_sampling_params(
66 max_tokens=1, default_sampling_params={}
67 )
68
69 assert sampling_params.logprobs is None
70 assert sampling_params.logprob_token_ids == [5000]
71 assert sampling_params.num_logprobs == 1
72
73
74 def test_completion_request_decouples_top_k_from_explicit_token_ids():
75 request = CompletionRequest(
76 model=MODEL_NAME,
77 prompt="Hello",
78 logprobs=5,
79 logprob_token_ids=[5000],
80 )
81
82 sampling_params = request.to_sampling_params(max_tokens=1)
83
84 assert sampling_params.logprobs is None
85 assert sampling_params.logprob_token_ids == [5000]
86 assert sampling_params.num_logprobs == 1
87
88
89 def test_completion_rejects_explicit_token_ids_without_generated_tokens():
90 with pytest.raises(VLLMValidationError, match="no output tokens are generated"):
91 CompletionRequest(
92 model=MODEL_NAME,
93 prompt="Hello",
94 echo=True,
95 max_tokens=0,
96 logprobs=5,
97 logprob_token_ids=[5000],
98 )
99
100
101 def test_requests_reject_explicit_token_ids_with_beam_search():
102 with pytest.raises(VLLMValidationError, match="not supported with beam search"):
103 ChatCompletionRequest(
104 model=MODEL_NAME,
105 messages=[{"role": "user", "content": "Hello"}],
106 logprobs=True,
107 logprob_token_ids=[5000],
108 use_beam_search=True,
109 )
110
111 with pytest.raises(VLLMValidationError, match="not supported with beam search"):
112 CompletionRequest(
113 model=MODEL_NAME,
114 prompt="Hello",
115 logprobs=5,
116 logprob_token_ids=[5000],
117 use_beam_search=True,
118 )
119
120
121 @pytest.mark.asyncio
122 @pytest.mark.parametrize(
123 ("requested_ids", "top_logprobs", "sampled_id"),
124 [
125 pytest.param([5000], 5, 42, id="original-issue-shape"),
126 pytest.param([100, 1000, 5000], 3, 42, id="sampled-outside-set"),
127 pytest.param([100, 1000, 5000], 3, 5000, id="sampled-inside-set"),
128 ],
129 )
130 async def test_logprob_token_ids_returns_requested_ids(
131 server, requested_ids, top_logprobs, sampled_id
132 ):
133 async with server.get_async_client() as client:
134 resp = await client.chat.completions.create(
135 model=MODEL_NAME,
136 messages=[{"role": "user", "content": "Hello"}],
137 max_tokens=1,
138 temperature=0,
139 logprobs=True,
140 top_logprobs=top_logprobs,
141 extra_body={
142 "logprob_token_ids": requested_ids,
143 "allowed_token_ids": [sampled_id],
144 "return_tokens_as_token_ids": True,
145 },
146 )
147 entries = _logprob_entries(resp)
148 returned_token_ids = _top_logprob_token_ids(resp)
149
150 assert _token_id(resp.choices[0].logprobs.content[0].token) == sampled_id
151 assert set(returned_token_ids) == {*requested_ids, sampled_id}
152 assert len(returned_token_ids) == len(set(returned_token_ids))
153
154 for e in entries:
155 assert isinstance(e.logprob, float)
156 assert not math.isnan(e.logprob)
157 assert not math.isinf(e.logprob)
158
159
160 @pytest.mark.asyncio
161 async def test_logprob_token_ids_stream_with_no_top_k(server):
162 async with server.get_async_client() as client:
163 stream = await client.chat.completions.create(
164 model=MODEL_NAME,
165 messages=[{"role": "user", "content": "Hello"}],
166 max_tokens=1,
167 temperature=0,
168 logprobs=True,
169 top_logprobs=None,
170 stream=True,
171 extra_body={
172 "logprob_token_ids": [5000],
173 "allowed_token_ids": [42],
174 "return_tokens_as_token_ids": True,
175 },
176 )
177
178 returned_token_ids: list[int] = []
179 async for chunk in stream:
180 if not chunk.choices or chunk.choices[0].logprobs is None:
181 continue
182 for content in chunk.choices[0].logprobs.content:
183 returned_token_ids.extend(
184 _token_id(entry.token) for entry in content.top_logprobs
185 )
186
187 assert set(returned_token_ids) == {42, 5000}
188
189
190 @pytest.mark.asyncio
191 async def test_logprob_token_ids_default_behavior_unchanged(server):
192 """Without `logprob_token_ids`, the response carries the natural top-k
193 most-likely tokens. This guards against the new field accidentally
194 changing the default-path output."""
195 async with server.get_async_client() as client:
196 resp = await client.chat.completions.create(
197 model=MODEL_NAME,
198 messages=[{"role": "user", "content": "Hello"}],
199 max_tokens=1,
200 temperature=0,
201 logprobs=True,
202 top_logprobs=5,
203 )
204 entries = _logprob_entries(resp)
205 assert len(entries) == 5
206 # Default mode emits log_softmax, so all values are <= 0.
207 for e in entries:
208 assert e.logprob <= 0.0