@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 341 lines 11.4 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 """E2E tests for ``thinking_token_budget`` with reasoning models.
5
6 Covers Qwen3-0.6B and Qwen3.5 FP8 + MTP.
7 """
8
9 import asyncio
10 import json
11 from typing import Literal
12
13 import openai
14 import pytest
15 import pytest_asyncio
16
17 from tests.utils import RemoteOpenAIServer, multi_gpu_only, requires_fp8
18 from vllm.platforms import current_platform
19 from vllm.tokenizers import get_tokenizer
20
21 MODEL_NAME = "Qwen/Qwen3-0.6B"
22 QWEN35_FP8_MTP_MODEL = "Qwen/Qwen3.5-35B-A3B-FP8"
23 MESSAGES = [{"role": "user", "content": "What is 1+1? Be concise."}]
24 THINK_BUDGET = 5
25
26 REASONING_START_STR = "<think>"
27 REASONING_END_STR = "</think>"
28
29
30 def _count_reasoning_decode_token_ids_between_markers(
31 full_token_ids: list[int],
32 reasoning_start_ids: list[int],
33 reasoning_end_ids: list[int],
34 ) -> int | None:
35 """Count decode tokens in the thinking span (after last start, before first end)."""
36 if not reasoning_start_ids or not reasoning_end_ids:
37 raise ValueError("reasoning marker token id lists must be non-empty")
38
39 def _last_subseq_index(haystack: list[int], needle: list[int]) -> int:
40 n = len(needle)
41 if n > len(haystack):
42 return -1
43 for i in range(len(haystack) - n, -1, -1):
44 if haystack[i : i + n] == needle:
45 return i
46 return -1
47
48 last_start = _last_subseq_index(full_token_ids, reasoning_start_ids)
49 if last_start < 0:
50 return None
51
52 pos_after_start = last_start + len(reasoning_start_ids)
53 end_n = len(reasoning_end_ids)
54 for j in range(pos_after_start, len(full_token_ids) - end_n + 1):
55 if full_token_ids[j : j + end_n] == reasoning_end_ids:
56 return j - pos_after_start
57 return len(full_token_ids) - pos_after_start
58
59
60 @pytest.fixture(scope="module")
61 def server():
62 args = [
63 "--reasoning-parser",
64 "qwen3",
65 "--reasoning-config",
66 '{"reasoning_start_str": "<think>", "reasoning_end_str": "</think>"}',
67 "--max-model-len",
68 "2048",
69 "--enforce-eager",
70 "--gpu-memory-utilization",
71 "0.4",
72 "--no-async-scheduling",
73 ]
74 with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
75 yield remote_server
76
77
78 @pytest.fixture(scope="module")
79 def server_with_auto_reasoning_config():
80 args = [
81 "--reasoning-parser",
82 "qwen3",
83 "--max-model-len",
84 "2048",
85 "--enforce-eager",
86 "--gpu-memory-utilization",
87 "0.4",
88 "--no-async-scheduling",
89 ]
90 with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
91 yield remote_server
92
93
94 @pytest.fixture(scope="module")
95 def server_qwen35_fp8_mtp_tp2():
96 """Qwen3.5-35B FP8 with MTP speculative decoding and tensor parallel size 2."""
97 if current_platform.device_count() < 2:
98 pytest.skip("Need at least 2 GPUs for --tensor-parallel-size 2")
99 if not current_platform.supports_fp8():
100 pytest.skip("FP8 is not supported on this platform")
101
102 spec_cfg = {
103 "method": "mtp",
104 "num_speculative_tokens": 2,
105 "max_model_len": 32768,
106 }
107 args = [
108 "--tensor-parallel-size",
109 "2",
110 "--max-model-len",
111 "32768",
112 "--speculative-config",
113 json.dumps(spec_cfg),
114 "--reasoning-parser",
115 "qwen3",
116 "--reasoning-config",
117 json.dumps(
118 {
119 "reasoning_start_str": REASONING_START_STR,
120 "reasoning_end_str": REASONING_END_STR,
121 }
122 ),
123 ]
124 env_dict: dict[str, str] = {}
125 # With 4+ GPUs, run TP=2 on physical devices 2,3 so module-scoped 0.6B servers
126 # on 0,1 do not exhaust memory on the same devices as this worker.
127 if current_platform.device_count() >= 4:
128 env_dict["CUDA_VISIBLE_DEVICES"] = "2,3"
129
130 with RemoteOpenAIServer(
131 QWEN35_FP8_MTP_MODEL,
132 args,
133 max_wait_seconds=3000,
134 env_dict=env_dict,
135 ) as remote_server:
136 yield remote_server
137
138
139 @pytest_asyncio.fixture
140 async def client(request, server, server_with_auto_reasoning_config):
141 server_map = {
142 "default": server,
143 "auto_config": server_with_auto_reasoning_config,
144 }
145 target_server = server_map[request.param]
146 async with target_server.get_async_client() as async_client:
147 yield async_client
148
149
150 @pytest.mark.asyncio
151 @pytest.mark.parametrize("client", ["default", "auto_config"], indirect=True)
152 async def test_thinking_token_budget_mixed_requests(client: openai.AsyncOpenAI):
153 """Test that mixed requests (some with thinking_token_budget, some without)
154 complete successfully without errors."""
155 response_with_budget = await client.chat.completions.create(
156 model=MODEL_NAME,
157 messages=MESSAGES,
158 max_tokens=100,
159 extra_body={"thinking_token_budget": THINK_BUDGET},
160 )
161 response_without_budget = await client.chat.completions.create(
162 model=MODEL_NAME,
163 messages=MESSAGES,
164 max_tokens=100,
165 )
166
167 msg_with = response_with_budget.choices[0].message
168 msg_without = response_without_budget.choices[0].message
169
170 assert msg_with.content or getattr(msg_with, "reasoning", None)
171 assert msg_without.content or getattr(msg_without, "reasoning", None)
172
173
174 @pytest.mark.asyncio
175 @pytest.mark.parametrize("client", ["default", "auto_config"], indirect=True)
176 async def test_thinking_token_budget_limits_reasoning(client: openai.AsyncOpenAI):
177 """Test that thinking_token_budget limits the number of reasoning tokens.
178
179 Counts reasoning decode tokens by id, which is robust to how tokens are
180 grouped into streamed chunks (a single chunk can carry several tokens under
181 async scheduling / stream_interval > 1). Counting chunks under-counts.
182 """
183 tokenizer = get_tokenizer(tokenizer_name=MODEL_NAME)
184 start_ids = list(tokenizer.encode(REASONING_START_STR, add_special_tokens=False))
185 end_ids = list(tokenizer.encode(REASONING_END_STR, add_special_tokens=False))
186
187 prompt_token_ids: list[int] = []
188 decode_token_ids: list[int] = []
189 stream = await client.chat.completions.create(
190 model=MODEL_NAME,
191 messages=MESSAGES,
192 max_tokens=100,
193 stream=True,
194 extra_body={"thinking_token_budget": THINK_BUDGET, "return_token_ids": True},
195 )
196 async for chunk in stream:
197 if not chunk.choices:
198 continue
199 if getattr(chunk, "prompt_token_ids", None):
200 prompt_token_ids = list(chunk.prompt_token_ids)
201 delta_ids = getattr(chunk.choices[0], "token_ids", None)
202 if delta_ids:
203 decode_token_ids.extend(delta_ids)
204
205 reasoning_token_count = _count_reasoning_decode_token_ids_between_markers(
206 prompt_token_ids + decode_token_ids, start_ids, end_ids
207 )
208 assert reasoning_token_count is not None, "missing reasoning start marker in ids"
209 assert reasoning_token_count == THINK_BUDGET, (
210 f"reasoning tokens ({reasoning_token_count}) != "
211 f"thinking_token_budget ({THINK_BUDGET})"
212 )
213
214
215 @pytest.mark.asyncio
216 @multi_gpu_only(num_gpus=2)
217 @requires_fp8
218 async def test_thinking_token_budget_qwen35_fp8_mtp_concurrent_mixed_budget_and_plain(
219 server_qwen35_fp8_mtp_tp2,
220 ):
221 """Concurrent chat requests: some with ``thinking_token_budget``, some without.
222
223 Exercises the scheduler / input processor under a mixed batch on the same
224 Qwen3.5 FP8 + MTP (TP=2) server. Budgeted calls are checked with
225 ``_count_reasoning_decode_token_ids_between_markers`` on full token ids.
226 """
227 _batch_spec: list[tuple[Literal["budget"], int] | tuple[Literal["plain"], None]] = [
228 ("budget", 1),
229 ("budget", 12),
230 ("plain", None),
231 ("budget", 20),
232 ("budget", 14),
233 ("plain", None),
234 ("plain", None),
235 ("budget", 12),
236 ("plain", None),
237 ]
238
239 tokenizer = get_tokenizer(tokenizer_name=QWEN35_FP8_MTP_MODEL)
240 start_ids = list(tokenizer.encode(REASONING_START_STR, add_special_tokens=False))
241 end_ids = list(tokenizer.encode(REASONING_END_STR, add_special_tokens=False))
242
243 async with server_qwen35_fp8_mtp_tp2.get_async_client() as client:
244
245 async def budgeted_call(expected_budget: int):
246 return await client.chat.completions.create(
247 model=QWEN35_FP8_MTP_MODEL,
248 messages=MESSAGES,
249 max_tokens=256,
250 stream=False,
251 extra_body={
252 "thinking_token_budget": expected_budget,
253 "return_token_ids": True,
254 },
255 )
256
257 async def plain_call():
258 return await client.chat.completions.create(
259 model=QWEN35_FP8_MTP_MODEL,
260 messages=MESSAGES,
261 max_tokens=256,
262 stream=False,
263 )
264
265 coros = []
266 for row in _batch_spec:
267 if row[0] == "budget":
268 b = row[1]
269 assert isinstance(b, int)
270 coros.append(budgeted_call(b))
271 else:
272 coros.append(plain_call())
273 results = await asyncio.gather(*coros)
274
275 for i, (response, (kind, expected_budget)) in enumerate(
276 zip(results, _batch_spec, strict=True)
277 ):
278 msg = response.choices[0].message
279 assert msg.content or getattr(msg, "reasoning", None), (
280 f"index {i} ({kind}): empty message"
281 )
282
283 if kind == "budget":
284 assert expected_budget is not None
285 assert response.prompt_token_ids is not None
286 assert response.choices[0].token_ids is not None
287 full_ids = list(response.prompt_token_ids) + list(
288 response.choices[0].token_ids
289 )
290 n_reason = _count_reasoning_decode_token_ids_between_markers(
291 full_ids, start_ids, end_ids
292 )
293 assert n_reason is not None, f"index {i}: missing reasoning start in ids"
294 assert n_reason == expected_budget, (
295 f"index {i}: reasoning decode token ids ({n_reason}) != "
296 f"thinking_token_budget ({expected_budget})"
297 )
298
299
300 @pytest.mark.asyncio
301 @pytest.mark.parametrize("client", ["default", "auto_config"], indirect=True)
302 async def test_streaming_with_thinking_disabled_stays_in_content(
303 client: openai.AsyncOpenAI,
304 ):
305 request_kwargs = {
306 "model": MODEL_NAME,
307 "messages": [
308 {
309 "role": "user",
310 "content": "Which is larger, 4 or 12?"
311 " Output exactly one token: 4 or 12.",
312 }
313 ],
314 "max_tokens": 16,
315 "temperature": 0.0,
316 "extra_body": {"chat_template_kwargs": {"enable_thinking": False}},
317 }
318
319 response = await client.chat.completions.create(**request_kwargs)
320 message = response.choices[0].message
321 assert message.content is not None and message.content.strip() != ""
322 assert getattr(message, "reasoning", None) in (None, "")
323
324 stream = await client.chat.completions.create(
325 **request_kwargs,
326 stream=True,
327 )
328
329 content_chunks = []
330 reasoning_chunks = []
331 async for chunk in stream:
332 if not chunk.choices:
333 continue
334 delta = chunk.choices[0].delta
335 if getattr(delta, "content", None):
336 content_chunks.append(delta.content)
337 if getattr(delta, "reasoning", None):
338 reasoning_chunks.append(delta.reasoning)
339
340 assert "".join(content_chunks).strip() != ""
341 assert reasoning_chunks == []