@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / d084e9fca7
Fork 0

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 10, 04:00 UTC

[MODEL] Fix handling of multiple channels for gpt-oss with speculative decoding (#26291)

Signed-off-by: Aleksandr Samarin <astrlrd@nebius.com> Signed-off-by: southfreebird <yvorott@gmail.com> Co-authored-by: southfreebird <yvorott@gmail.com>

Aleksandr Samarin committed Jan 14, 2026 at 21:20 UTC d084e9fca7d5d40cbb62eb5fe8ab5cbc6c769cf0
4 files changed +650 -361
tests/entrypoints/openai/test_serving_chat.py
+335 -275
@@ -35,6 +35,7 @@ from .utils import (
35 )
36
37 GPT_OSS_MODEL_NAME = "openai/gpt-oss-20b"
38 +GPT_OSS_SPECULATOR_NAME = "RedHatAI/gpt-oss-20b-speculator.eagle3"
39
40
41 @pytest.fixture(scope="module")
@@ -66,7 +67,8 @@ def exclude_tools_when_tool_choice_none(request) -> bool:
67
68 @pytest.fixture(scope="module")
69 def default_server_args(
69 - with_tool_parser: bool, exclude_tools_when_tool_choice_none: bool
70 + with_tool_parser: bool,
71 + exclude_tools_when_tool_choice_none: bool,
72 ):
73 args = [
74 # use half precision for speed and memory savings in CI environment
@@ -76,7 +78,7 @@ def default_server_args(
78 "--reasoning-parser",
79 "openai_gptoss",
80 "--gpu-memory-utilization",
79 - "0.8",
81 + "0.85",
82 ]
83 if with_tool_parser:
84 args.extend(
@@ -91,327 +93,385 @@ def default_server_args(
93 return args
94
95
94 -@pytest.fixture(scope="module")
96 +@pytest.fixture(scope="class")
97 def gptoss_server(default_server_args: list[str]):
98 server_args = default_server_args + ["--attention-backend=TRITON_ATTN"]
99 with RemoteOpenAIServer(GPT_OSS_MODEL_NAME, server_args) as remote_server:
100 yield remote_server
101
102
103 +@pytest.fixture(scope="class")
104 +def gptoss_speculative_server(default_server_args: list[str]):
105 + server_args = default_server_args + [
106 + "--speculative-config",
107 + f'{{"model": "{GPT_OSS_SPECULATOR_NAME}", '
108 + f'"method": "eagle3", "num_speculative_tokens": 3}}',
109 + "--attention-backend=TRITON_ATTN",
110 + ]
111 + with RemoteOpenAIServer(GPT_OSS_MODEL_NAME, server_args) as remote_server:
112 + yield remote_server
113 +
114 +
115 @pytest_asyncio.fixture
116 async def gptoss_client(gptoss_server):
117 async with gptoss_server.get_async_client() as async_client:
118 yield async_client
119
120
107 -@pytest.mark.asyncio
108 -async def test_gpt_oss_chat_tool_call_streaming(
109 - gptoss_client: OpenAI, with_tool_parser: bool
110 -):
111 - tools = [
112 - {
113 - "type": "function",
114 - "function": {
115 - "name": "get_current_weather",
116 - "description": "Get the current weather in a given location",
117 - "parameters": {
118 - "type": "object",
119 - "properties": {
120 - "city": {"type": "string"},
121 - "state": {"type": "string"},
122 - "unit": {
123 - "type": "string",
124 - "enum": ["celsius", "fahrenheit"],
121 +@pytest_asyncio.fixture
122 +async def gptoss_speculative_client(gptoss_speculative_server):
123 + async with gptoss_speculative_server.get_async_client() as async_client:
124 + yield async_client
125 +
126 +
127 +class TestGPTOSSChat:
128 + @pytest.mark.asyncio
129 + async def test_gpt_oss_chat_tool_call_streaming(
130 + self, gptoss_client: OpenAI, with_tool_parser: bool
131 + ):
132 + tools = [
133 + {
134 + "type": "function",
135 + "function": {
136 + "name": "get_current_weather",
137 + "description": "Get the current weather in a given location",
138 + "parameters": {
139 + "type": "object",
140 + "properties": {
141 + "city": {"type": "string"},
142 + "state": {"type": "string"},
143 + "unit": {
144 + "type": "string",
145 + "enum": ["celsius", "fahrenheit"],
146 + },
147 },
148 + "required": ["city", "state", "unit"],
149 },
127 - "required": ["city", "state", "unit"],
150 },
129 - },
130 - }
131 - ]
132 -
133 - messages = [
134 - {"role": "user", "content": "What is the weather in Dallas, TX?"},
135 - ]
151 + }
152 + ]
153
137 - stream = await gptoss_client.chat.completions.create(
138 - model=GPT_OSS_MODEL_NAME,
139 - messages=messages,
140 - tools=tools if with_tool_parser else None,
141 - stream=True,
142 - )
154 + messages = [
155 + {"role": "user", "content": "What is the weather in Dallas, TX?"},
156 + ]
157
144 - name = None
145 - args_buf = ""
146 - content_buf = ""
147 - async for chunk in stream:
148 - delta = chunk.choices[0].delta
149 - if delta.tool_calls:
150 - tc = delta.tool_calls[0]
151 - if tc.function and tc.function.name:
152 - name = tc.function.name
153 - if tc.function and tc.function.arguments:
154 - args_buf += tc.function.arguments
155 - if getattr(delta, "content", None):
156 - content_buf += delta.content
157 - if with_tool_parser:
158 - assert name is not None
159 - assert len(args_buf) > 0
160 - else:
161 - assert name is None
162 - assert len(args_buf) == 0
163 - assert len(content_buf) > 0
158 + stream = await gptoss_client.chat.completions.create(
159 + model=GPT_OSS_MODEL_NAME,
160 + messages=messages,
161 + tools=tools if with_tool_parser else None,
162 + stream=True,
163 + )
164
165 + name = None
166 + args_buf = ""
167 + content_buf = ""
168 + async for chunk in stream:
169 + delta = chunk.choices[0].delta
170 + if delta.tool_calls:
171 + tc = delta.tool_calls[0]
172 + if tc.function and tc.function.name:
173 + name = tc.function.name
174 + if tc.function and tc.function.arguments:
175 + args_buf += tc.function.arguments
176 + if getattr(delta, "content", None):
177 + content_buf += delta.content
178 + if with_tool_parser:
179 + assert name is not None
180 + assert len(args_buf) > 0
181 + else:
182 + assert name is None
183 + assert len(args_buf) == 0
184 + assert len(content_buf) > 0
185
166 -@pytest.mark.asyncio
167 -async def test_gpt_oss_multi_turn_chat(gptoss_client: OpenAI, with_tool_parser: bool):
168 - if not with_tool_parser:
169 - pytest.skip("skip non-tool for multi-turn tests")
170 - tools = [
171 - {
172 - "type": "function",
173 - "function": {
174 - "name": "get_current_weather",
175 - "description": "Get the current weather in a given location",
176 - "parameters": {
177 - "type": "object",
178 - "properties": {
179 - "city": {"type": "string"},
180 - "state": {"type": "string"},
181 - "unit": {
182 - "type": "string",
183 - "enum": ["celsius", "fahrenheit"],
186 + @pytest.mark.asyncio
187 + async def test_gpt_oss_multi_turn_chat(
188 + self, gptoss_client: OpenAI, with_tool_parser: bool
189 + ):
190 + if not with_tool_parser:
191 + pytest.skip("skip non-tool for multi-turn tests")
192 + tools = [
193 + {
194 + "type": "function",
195 + "function": {
196 + "name": "get_current_weather",
197 + "description": "Get the current weather in a given location",
198 + "parameters": {
199 + "type": "object",
200 + "properties": {
201 + "city": {"type": "string"},
202 + "state": {"type": "string"},
203 + "unit": {
204 + "type": "string",
205 + "enum": ["celsius", "fahrenheit"],
206 + },
207 },
208 + "required": ["city", "state", "unit"],
209 },
186 - "required": ["city", "state", "unit"],
210 },
188 - },
189 - }
190 - ]
191 -
192 - messages = [
193 - {"role": "system", "content": "you are a helpful assistant"},
194 - {"role": "user", "content": "What is the weather in Dallas, TX with celsius?"},
195 - ]
211 + }
212 + ]
213
197 - first = await gptoss_client.chat.completions.create(
198 - model=GPT_OSS_MODEL_NAME,
199 - messages=messages,
200 - tools=tools,
201 - temperature=0.0,
202 - )
203 - first_msg = first.choices[0].message
204 - assert first_msg.tool_calls is not None and len(first_msg.tool_calls) > 0
205 - tc = first_msg.tool_calls[0]
206 - assert tc.function is not None and tc.function.name == "get_current_weather"
207 - args1 = tc.function.arguments
208 - assert args1 is not None and len(args1) > 0
209 - assert not first_msg.content
210 -
211 - messages.append({"role": "assistant", "content": args1})
212 - messages.append(
213 - {"role": "user", "content": "Now convert to celsius and return JSON only"}
214 - )
214 + messages = [
215 + {"role": "system", "content": "you are a helpful assistant"},
216 + {
217 + "role": "user",
218 + "content": "What is the weather in Dallas, TX with celsius?",
219 + },
220 + ]
221
216 - second = await gptoss_client.chat.completions.create(
217 - model=GPT_OSS_MODEL_NAME,
218 - messages=messages,
219 - tools=tools,
220 - temperature=0.0,
221 - )
222 - second_msg = second.choices[0].message
223 - assert (second_msg.content is not None and len(second_msg.content) > 0) or (
224 - second_msg.tool_calls is not None and len(second_msg.tool_calls) > 0
225 - )
222 + first = await gptoss_client.chat.completions.create(
223 + model=GPT_OSS_MODEL_NAME,
224 + messages=messages,
225 + tools=tools,
226 + temperature=0.0,
227 + )
228 + first_msg = first.choices[0].message
229 + assert first_msg.tool_calls is not None and len(first_msg.tool_calls) > 0
230 + tc = first_msg.tool_calls[0]
231 + assert tc.function is not None and tc.function.name == "get_current_weather"
232 + args1 = tc.function.arguments
233 + assert args1 is not None and len(args1) > 0
234 + assert not first_msg.content
235 +
236 + messages.append({"role": "assistant", "content": args1})
237 + messages.append(
238 + {"role": "user", "content": "Now convert to celsius and return JSON only"}
239 + )
240
241 + second = await gptoss_client.chat.completions.create(
242 + model=GPT_OSS_MODEL_NAME,
243 + messages=messages,
244 + tools=tools,
245 + temperature=0.0,
246 + )
247 + second_msg = second.choices[0].message
248 + assert (second_msg.content is not None and len(second_msg.content) > 0) or (
249 + second_msg.tool_calls is not None and len(second_msg.tool_calls) > 0
250 + )
251
228 -@pytest.mark.asyncio
229 -async def test_gpt_oss_tool_message_array_content(
230 - gptoss_client: OpenAI, with_tool_parser: bool
231 -):
232 - """Test that tool messages support both string and array content formats."""
233 - if not with_tool_parser:
234 - pytest.skip("skip non-tool for array content tests")
252 + @pytest.mark.asyncio
253 + async def test_gpt_oss_tool_message_array_content(
254 + self, gptoss_client: OpenAI, with_tool_parser: bool
255 + ):
256 + """Test that tool messages support both string and array content formats."""
257 + if not with_tool_parser:
258 + pytest.skip("skip non-tool for array content tests")
259
236 - tools = [
237 - {
238 - "type": "function",
239 - "function": {
240 - "name": "get_weather",
241 - "description": "Get the current weather in a given location",
242 - "parameters": {
243 - "type": "object",
244 - "properties": {
245 - "city": {"type": "string"},
246 - "state": {"type": "string"},
260 + tools = [
261 + {
262 + "type": "function",
263 + "function": {
264 + "name": "get_weather",
265 + "description": "Get the current weather in a given location",
266 + "parameters": {
267 + "type": "object",
268 + "properties": {
269 + "city": {"type": "string"},
270 + "state": {"type": "string"},
271 + },
272 + "required": ["city", "state"],
273 },
248 - "required": ["city", "state"],
274 },
250 - },
251 - }
252 - ]
253 -
254 - # Test 1: Tool message with string content
255 - messages_string = [
256 - {"role": "user", "content": "What's the weather in Paris?"},
257 - {
258 - "role": "assistant",
259 - "tool_calls": [
260 - {
261 - "id": "call_123",
262 - "type": "function",
263 - "function": {
264 - "name": "get_weather",
265 - "arguments": '{"city": "Paris", "state": "TX"}',
266 - },
267 - }
268 - ],
269 - },
270 - {"role": "tool", "content": "The weather in Paris, TX is sunny, 22°C"},
271 - ]
275 + }
276 + ]
277
273 - response_string = await gptoss_client.chat.completions.create(
274 - model=GPT_OSS_MODEL_NAME,
275 - messages=messages_string,
276 - tools=tools,
277 - temperature=0.0,
278 - )
278 + # Test 1: Tool message with string content
279 + messages_string = [
280 + {"role": "user", "content": "What's the weather in Paris?"},
281 + {
282 + "role": "assistant",
283 + "tool_calls": [
284 + {
285 + "id": "call_123",
286 + "type": "function",
287 + "function": {
288 + "name": "get_weather",
289 + "arguments": '{"city": "Paris", "state": "TX"}',
290 + },
291 + }
292 + ],
293 + },
294 + {"role": "tool", "content": "The weather in Paris, TX is sunny, 22°C"},
295 + ]
296
280 - assert response_string is not None
281 - assert response_string.choices[0].message is not None
297 + response_string = await gptoss_client.chat.completions.create(
298 + model=GPT_OSS_MODEL_NAME,
299 + messages=messages_string,
300 + tools=tools,
301 + temperature=0.0,
302 + )
303
283 - # Test 2: Tool message with array content
284 - messages_array = [
285 - {"role": "user", "content": "What's the weather in Dallas?"},
286 - {
287 - "role": "assistant",
288 - "tool_calls": [
289 - {
290 - "id": "call_456",
291 - "type": "function",
292 - "function": {
293 - "name": "get_weather",
294 - "arguments": '{"city": "Dallas", "state": "TX"}',
295 - },
296 - }
297 - ],
298 - },
299 - {
300 - "role": "tool",
301 - "content": [
302 - {"type": "text", "text": "f2e897a7-2705-4337-8193-2a8f57b81618"}
303 - ],
304 - },
305 - ]
304 + assert response_string is not None
305 + assert response_string.choices[0].message is not None
306
307 - response_array = await gptoss_client.chat.completions.create(
308 - model=GPT_OSS_MODEL_NAME,
309 - messages=messages_array,
310 - tools=tools,
311 - temperature=0.0,
312 - )
307 + # Test 2: Tool message with array content
308 + messages_array = [
309 + {"role": "user", "content": "What's the weather in Dallas?"},
310 + {
311 + "role": "assistant",
312 + "tool_calls": [
313 + {
314 + "id": "call_456",
315 + "type": "function",
316 + "function": {
317 + "name": "get_weather",
318 + "arguments": '{"city": "Dallas", "state": "TX"}',
319 + },
320 + }
321 + ],
322 + },
323 + {
324 + "role": "tool",
325 + "content": [
326 + {"type": "text", "text": "f2e897a7-2705-4337-8193-2a8f57b81618"}
327 + ],
328 + },
329 + ]
330
314 - assert response_array is not None
315 - assert response_array.choices[0].message is not None
331 + response_array = await gptoss_client.chat.completions.create(
332 + model=GPT_OSS_MODEL_NAME,
333 + messages=messages_array,
334 + tools=tools,
335 + temperature=0.0,
336 + )
337
317 - # Test 3: Tool message with multiple array content items
318 - messages_multi_array = [
319 - {"role": "user", "content": "Search for information"},
320 - {
321 - "role": "assistant",
322 - "tool_calls": [
323 - {
324 - "id": "call_789",
325 - "type": "function",
326 - "function": {
327 - "name": "get_weather",
328 - "arguments": '{"city": "Austin", "state": "TX"}',
329 - },
330 - }
331 - ],
332 - },
333 - {
334 - "role": "tool",
335 - "content": [
336 - {"type": "text", "text": "Weather data: "},
337 - {"type": "text", "text": "Austin, TX - Partly cloudy, 25°C"},
338 - {"type": "text", "text": " with 60% humidity"},
339 - ],
340 - },
341 - ]
338 + assert response_array is not None
339 + assert response_array.choices[0].message is not None
340
343 - response_multi_array = await gptoss_client.chat.completions.create(
344 - model=GPT_OSS_MODEL_NAME,
345 - messages=messages_multi_array,
346 - tools=tools,
347 - temperature=0.0,
348 - )
341 + # Test 3: Tool message with multiple array content items
342 + messages_multi_array = [
343 + {"role": "user", "content": "Search for information"},
344 + {
345 + "role": "assistant",
346 + "tool_calls": [
347 + {
348 + "id": "call_789",
349 + "type": "function",
350 + "function": {
351 + "name": "get_weather",
352 + "arguments": '{"city": "Austin", "state": "TX"}',
353 + },
354 + }
355 + ],
356 + },
357 + {
358 + "role": "tool",
359 + "content": [
360 + {"type": "text", "text": "Weather data: "},
361 + {"type": "text", "text": "Austin, TX - Partly cloudy, 25°C"},
362 + {"type": "text", "text": " with 60% humidity"},
363 + ],
364 + },
365 + ]
366
350 - assert response_multi_array is not None
351 - assert response_multi_array.choices[0].message is not None
367 + response_multi_array = await gptoss_client.chat.completions.create(
368 + model=GPT_OSS_MODEL_NAME,
369 + messages=messages_multi_array,
370 + tools=tools,
371 + temperature=0.0,
372 + )
373
374 + assert response_multi_array is not None
375 + assert response_multi_array.choices[0].message is not None
376
354 -@pytest.mark.asyncio
355 -async def test_gpt_oss_tool_choice_none(
356 - gptoss_client: OpenAI,
357 - with_tool_parser: bool,
358 - exclude_tools_when_tool_choice_none: bool,
359 -):
360 - if not (with_tool_parser and exclude_tools_when_tool_choice_none):
361 - pytest.skip(
362 - "skip tool_choice tests when non-tool or "
363 - "--exclude-tools-when-tool-choice-none not set"
364 - )
377 + @pytest.mark.asyncio
378 + async def test_gpt_oss_tool_choice_none(
379 + self,
380 + gptoss_client: OpenAI,
381 + with_tool_parser: bool,
382 + exclude_tools_when_tool_choice_none: bool,
383 + ):
384 + if not (with_tool_parser and exclude_tools_when_tool_choice_none):
385 + pytest.skip(
386 + "skip tool_choice tests when non-tool or "
387 + "--exclude-tools-when-tool-choice-none not set"
388 + )
389
366 - tools = [
367 - {
368 - "type": "function",
369 - "function": {
370 - "name": "get_current_weather",
371 - "description": "Get the current weather in a given location",
372 - "parameters": {
373 - "type": "object",
374 - "properties": {
375 - "city": {"type": "string"},
376 - "state": {"type": "string"},
377 - "unit": {
378 - "type": "string",
379 - "enum": ["celsius", "fahrenheit"],
390 + tools = [
391 + {
392 + "type": "function",
393 + "function": {
394 + "name": "get_current_weather",
395 + "description": "Get the current weather in a given location",
396 + "parameters": {
397 + "type": "object",
398 + "properties": {
399 + "city": {"type": "string"},
400 + "state": {"type": "string"},
401 + "unit": {
402 + "type": "string",
403 + "enum": ["celsius", "fahrenheit"],
404 + },
405 },
406 + "required": ["city", "state", "unit"],
407 },
382 - "required": ["city", "state", "unit"],
408 },
409 + }
410 + ]
411 +
412 + messages = [
413 + {
414 + "role": "user",
415 + "content": "What's the temperature(in degrees Celsius) in Dallas?",
416 },
385 - }
386 - ]
417 + ]
418
388 - messages = [
389 - {
390 - "role": "user",
391 - "content": "What's the temperature(in degrees Celsius) in Dallas?",
392 - },
393 - ]
419 + tool_choice_auto = await gptoss_client.chat.completions.create(
420 + model=GPT_OSS_MODEL_NAME,
421 + messages=messages,
422 + tools=tools,
423 + tool_choice="auto",
424 + temperature=0.0,
425 + )
426 + msg = tool_choice_auto.choices[0].message
427 + assert len(msg.tool_calls) == 1
428 +
429 + tool_choice_none = await gptoss_client.chat.completions.create(
430 + model=GPT_OSS_MODEL_NAME,
431 + messages=messages,
432 + tools=tools,
433 + tool_choice="none",
434 + temperature=0.0,
435 + )
436
395 - tool_choice_auto = await gptoss_client.chat.completions.create(
396 - model=GPT_OSS_MODEL_NAME,
397 - messages=messages,
398 - tools=tools,
399 - tool_choice="auto",
400 - temperature=0.0,
401 - )
402 - msg = tool_choice_auto.choices[0].message
403 - assert len(msg.tool_calls) == 1
437 + msg = tool_choice_none.choices[0].message
438 + assert len(msg.tool_calls) == 0
439
405 - tool_choice_none = await gptoss_client.chat.completions.create(
406 - model=GPT_OSS_MODEL_NAME,
407 - messages=messages,
408 - tools=tools,
409 - tool_choice="none",
410 - temperature=0.0,
411 - )
440
413 - msg = tool_choice_none.choices[0].message
414 - assert len(msg.tool_calls) == 0
441 +class TestGPTOSSSpeculativeChat:
442 + @pytest.mark.asyncio
443 + async def test_gpt_oss_speculative_reasoning_leakage(
444 + self,
445 + gptoss_speculative_client: OpenAI,
446 + with_tool_parser: bool,
447 + ):
448 + if not with_tool_parser:
449 + pytest.skip("skip non-tool for array content tests")
450 +
451 + messages = [
452 + {"role": "user", "content": "Calculate 2+2. Return the answer 4 only."},
453 + ]
454 +
455 + stream = await gptoss_speculative_client.chat.completions.create(
456 + model=GPT_OSS_MODEL_NAME,
457 + messages=messages,
458 + stream=True,
459 + temperature=0.0,
460 + )
461 +
462 + content = ""
463 + reasoning_content = ""
464 + async for chunk in stream:
465 + delta = chunk.choices[0].delta
466 + if delta.content:
467 + content += delta.content
468 +
469 + chunk_reasoning = getattr(delta, "reasoning", None)
470 + if chunk_reasoning:
471 + reasoning_content += delta.reasoning
472 +
473 + assert len(reasoning_content) > 0, "No reasoning was generated."
474 + assert content.strip() == "4"
475
476
477 MODEL_NAME = "openai-community/gpt2"
tests/entrypoints/openai/test_serving_chat_stream_harmony.py
+171 -27
@@ -10,6 +10,7 @@ from unittest.mock import patch
10 import pytest
11
12 from vllm.entrypoints.openai.chat_completion.stream_harmony import (
13 + TokenState,
14 extract_harmony_streaming_delta,
15 )
16
@@ -42,12 +43,14 @@ class TestExtractHarmonyStreamingDelta:
43 def test_final_channel_returns_content_delta(self, delta_text, expected_content):
44 """Test that final channel returns a DeltaMessage with content."""
45 parser = MockStreamableParser()
46 +
47 + # Updated to use TokenState list
48 + token_states = [TokenState(channel="final", recipient=None, text=delta_text)]
49 +
50 delta_message, tools_streamed = extract_harmony_streaming_delta(
51 harmony_parser=parser,
47 - cur_channel="final",
48 - cur_recipient=None,
52 + token_states=token_states,
53 prev_recipient=None,
50 - delta_text=delta_text,
54 include_reasoning=False,
55 )
56
@@ -65,18 +68,19 @@ class TestExtractHarmonyStreamingDelta:
68 def test_analysis_channel_reasoning(self, include_reasoning, expected_has_message):
69 """Test analysis channel respects include_reasoning flag."""
70 parser = MockStreamableParser()
71 + text = "Let me think..."
72 + token_states = [TokenState(channel="analysis", recipient=None, text=text)]
73 +
74 delta_message, tools_streamed = extract_harmony_streaming_delta(
75 harmony_parser=parser,
70 - cur_channel="analysis",
71 - cur_recipient=None,
76 + token_states=token_states,
77 prev_recipient=None,
73 - delta_text="Let me think...",
78 include_reasoning=include_reasoning,
79 )
80
81 if expected_has_message:
82 assert delta_message is not None
79 - assert delta_message.reasoning == "Let me think..."
83 + assert delta_message.reasoning == text
84 else:
85 assert delta_message is None
86 assert tools_streamed is False
@@ -88,12 +92,14 @@ class TestExtractHarmonyStreamingDelta:
92 mock_make_tool_call_id.return_value = "call_test123"
93 parser = MockStreamableParser()
94
95 + token_states = [
96 + TokenState(channel=channel, recipient="functions.get_weather", text="")
97 + ]
98 +
99 delta_message, tools_streamed = extract_harmony_streaming_delta(
100 harmony_parser=parser,
93 - cur_channel=channel,
94 - cur_recipient="functions.get_weather",
101 + token_states=token_states,
102 prev_recipient=None,
96 - delta_text="",
103 include_reasoning=False,
104 )
105
@@ -111,20 +117,25 @@ class TestExtractHarmonyStreamingDelta:
117 def test_tool_call_argument_streaming(self, channel):
118 """Test streaming tool call arguments (same recipient)."""
119 parser = MockStreamableParser()
120 + args_text = '{"location": "Paris"}'
121 +
122 + token_states = [
123 + TokenState(
124 + channel=channel, recipient="functions.get_weather", text=args_text
125 + )
126 + ]
127
128 delta_message, tools_streamed = extract_harmony_streaming_delta(
129 harmony_parser=parser,
117 - cur_channel=channel,
118 - cur_recipient="functions.get_weather",
130 + token_states=token_states,
131 prev_recipient="functions.get_weather",
120 - delta_text='{"location": "Paris"}',
132 include_reasoning=False,
133 )
134
135 assert delta_message is not None
136 tool_call = delta_message.tool_calls[0]
137 assert tool_call.id is None
127 - assert tool_call.function.arguments == '{"location": "Paris"}'
138 + assert tool_call.function.arguments == args_text
139 assert tool_call.index == 0
140 assert tools_streamed is True
141
@@ -133,12 +144,14 @@ class TestExtractHarmonyStreamingDelta:
144 """Test empty delta_text with same recipient returns None."""
145 parser = MockStreamableParser()
146
147 + token_states = [
148 + TokenState(channel=channel, recipient="functions.get_weather", text="")
149 + ]
150 +
151 delta_message, tools_streamed = extract_harmony_streaming_delta(
152 harmony_parser=parser,
138 - cur_channel=channel,
139 - cur_recipient="functions.get_weather",
153 + token_states=token_states,
154 prev_recipient="functions.get_weather",
141 - delta_text="",
155 include_reasoning=False,
156 )
157
@@ -154,12 +167,14 @@ class TestExtractHarmonyStreamingDelta:
167 ]
168 parser = MockStreamableParser(messages=messages)
169
170 + token_states = [
171 + TokenState(channel="commentary", recipient="functions.tool2", text="args")
172 + ]
173 +
174 delta_message, _ = extract_harmony_streaming_delta(
175 harmony_parser=parser,
159 - cur_channel="commentary",
160 - cur_recipient="functions.tool2",
176 + token_states=token_states,
177 prev_recipient="functions.tool2",
162 - delta_text="args",
178 include_reasoning=False,
179 )
180
@@ -173,15 +188,18 @@ class TestExtractHarmonyStreamingDelta:
188 ],
189 )
190 def test_returns_tool_call_preambles(self, channel, recipient):
176 - """Test that invalid channel/recipient combinations return None."""
191 + """Test that invalid tool recipient on commentary is treated as content."""
192 parser = MockStreamableParser()
193 delta_text = "some text"
194 +
195 + token_states = [
196 + TokenState(channel=channel, recipient=recipient, text=delta_text)
197 + ]
198 +
199 delta_message, tools_streamed = extract_harmony_streaming_delta(
200 harmony_parser=parser,
181 - cur_channel=channel,
182 - cur_recipient=recipient,
201 + token_states=token_states,
202 prev_recipient=None,
184 - delta_text=delta_text,
203 include_reasoning=True,
204 )
205
@@ -199,14 +217,140 @@ class TestExtractHarmonyStreamingDelta:
217 """Test that invalid channel/recipient combinations return None."""
218 parser = MockStreamableParser()
219
220 + token_states = [
221 + TokenState(channel=channel, recipient=recipient, text="some text")
222 + ]
223 +
224 delta_message, tools_streamed = extract_harmony_streaming_delta(
225 harmony_parser=parser,
204 - cur_channel=channel,
205 - cur_recipient=recipient,
226 + token_states=token_states,
227 prev_recipient=None,
207 - delta_text="some text",
228 include_reasoning=True,
229 )
230
231 assert delta_message is None
232 assert tools_streamed is False
233 +
234 + def test_consecutive_token_grouping(self):
235 + """
236 + Test that consecutive tokens with the same channel/recipient
237 + are merged into a single processing group.
238 + """
239 + parser = MockStreamableParser()
240 + token_states = [
241 + TokenState("final", None, "H"),
242 + TokenState("final", None, "el"),
243 + TokenState("final", None, "lo"),
244 + TokenState("final", None, ","),
245 + TokenState("final", None, " World"),
246 + ]
247 +
248 + delta_message, _ = extract_harmony_streaming_delta(
249 + harmony_parser=parser,
250 + token_states=token_states,
251 + prev_recipient=None,
252 + include_reasoning=False,
253 + )
254 +
255 + assert delta_message is not None
256 + assert delta_message.content == "Hello, World"
257 +
258 + @patch("vllm.entrypoints.openai.chat_completion.stream_harmony.make_tool_call_id")
259 + def test_complex_batch_permutation(self, mock_make_id):
260 + """
261 + Test a complex permutation: Reasoning -> Tool Call -> Content.
262 + This verifies that multiple distinct actions in one batch
263 + are all captured in the single DeltaMessage.
264 + """
265 + mock_make_id.return_value = "call_batch_test"
266 + parser = MockStreamableParser()
267 +
268 + token_states = [
269 + # 1. Reasoning
270 + TokenState("analysis", None, "Reasoning about query..."),
271 + # 2. Tool Calling
272 + TokenState("commentary", "functions.search", '{"query":'),
273 + TokenState("commentary", "functions.search", ' "vllm"}'),
274 + # 3. Final Content
275 + TokenState("final", None, "."),
276 + ]
277 +
278 + delta_message, tools_streamed = extract_harmony_streaming_delta(
279 + harmony_parser=parser,
280 + token_states=token_states,
281 + prev_recipient=None,
282 + include_reasoning=True,
283 + )
284 +
285 + assert delta_message is not None
286 +
287 + assert delta_message.reasoning == "Reasoning about query..."
288 +
289 + # We expect 2 objects for 1 logical tool call:
290 + # 1. The definition (id, name, type)
291 + # 2. The arguments payload
292 + assert len(delta_message.tool_calls) == 2
293 +
294 + header = delta_message.tool_calls[0]
295 + payload = delta_message.tool_calls[1]
296 +
297 + assert header.function.name == "search"
298 + assert header.id == "call_batch_test"
299 + assert header.index == 0
300 +
301 + assert payload.index == 0
302 + assert payload.function.arguments == '{"query": "vllm"}'
303 +
304 + assert delta_message.content == "."
305 + assert tools_streamed is True
306 +
307 + @patch("vllm.entrypoints.openai.chat_completion.stream_harmony.make_tool_call_id")
308 + def test_tool_call_index_consistency_with_ongoing_call(self, mock_make_id):
309 + """
310 + Test that an ongoing tool call continuation and subsequent new calls
311 + maintain correct indexing when interleaved with content.
312 + """
313 + mock_make_id.side_effect = ["id_b", "id_c"]
314 +
315 + messages = [
316 + MockMessage(channel="commentary", recipient="functions.previous_tool")
317 + ]
318 + parser = MockStreamableParser(messages=messages)
319 +
320 + token_states = [
321 + TokenState("commentary", "functions.tool_a", '{"key_a": "val_a"}'),
322 + TokenState("final", None, "Thinking..."),
323 + TokenState("commentary", "functions.tool_b", '{"key_b": "val_b"}'),
324 + TokenState("final", None, " Thinking again..."),
325 + TokenState("commentary", "functions.tool_c", '{"key_c": "val_c"}'),
326 + ]
327 +
328 + delta_message, _ = extract_harmony_streaming_delta(
329 + harmony_parser=parser,
330 + token_states=token_states,
331 + prev_recipient="functions.tool_a",
332 + include_reasoning=False,
333 + )
334 +
335 + assert delta_message is not None
336 +
337 + tool_a_deltas = [t for t in delta_message.tool_calls if t.index == 1]
338 + assert len(tool_a_deltas) > 0
339 + assert tool_a_deltas[0].id is None
340 + assert tool_a_deltas[0].function.arguments == '{"key_a": "val_a"}'
341 +
342 + tool_b_header = next(t for t in delta_message.tool_calls if t.id == "id_b")
343 + assert tool_b_header.index == 2
344 + tool_b_args = next(
345 + t for t in delta_message.tool_calls if t.index == 2 and t.id is None
346 + )
347 + assert tool_b_args.function.arguments == '{"key_b": "val_b"}'
348 +
349 + tool_c_start = next(t for t in delta_message.tool_calls if t.id == "id_c")
350 + assert tool_c_start.index == 3
351 + tool_c_args = next(
352 + t for t in delta_message.tool_calls if t.index == 3 and t.id is None
353 + )
354 + assert tool_c_args.function.arguments == '{"key_c": "val_c"}'
355 +
356 + assert delta_message.content == "Thinking... Thinking again..."
vllm/entrypoints/openai/chat_completion/serving.py
+26 -11
@@ -36,6 +36,7 @@ from vllm.entrypoints.openai.chat_completion.protocol import (
36 ChatMessage,
37 )
38 from vllm.entrypoints.openai.chat_completion.stream_harmony import (
39 + TokenState,
40 extract_harmony_streaming_delta,
41 )
42 from vllm.entrypoints.openai.engine.protocol import (
@@ -826,12 +827,22 @@ class OpenAIServingChat(OpenAIServing):
827 if self.use_harmony:
828 harmony_parser = harmony_parsers[i]
829 prev_recipient = harmony_parser.current_recipient
829 - delta_text = ""
830 +
831 + # Track accumulated content per token with their state
832 + token_states: list[TokenState] = []
833 for token_id in output.token_ids:
834 harmony_parser.process(token_id)
832 - delta_text += harmony_parser.last_content_delta or ""
835 + token_delta = harmony_parser.last_content_delta or ""
836 + token_states.append(
837 + TokenState(
838 + harmony_parser.current_channel,
839 + harmony_parser.current_recipient,
840 + token_delta,
841 + )
842 + )
843 + delta_text = "".join(delta for _, _, delta in token_states)
844 cur_channel = harmony_parser.current_channel
834 - cur_recipient = harmony_parser.current_recipient
845 +
846 # handle the case where several tokens where generated at once
847 # including the final token, leading to a delta in the text
848 # but the current channel to be empty (start state)
@@ -869,10 +880,8 @@ class OpenAIServingChat(OpenAIServing):
880 delta_message, tools_streamed_flag = (
881 extract_harmony_streaming_delta(
882 harmony_parser=harmony_parser,
872 - cur_channel=cur_channel,
873 - cur_recipient=cur_recipient,
883 + token_states=token_states,
884 prev_recipient=prev_recipient,
875 - delta_text=delta_text,
885 include_reasoning=request.include_reasoning,
886 )
887 )
@@ -1139,17 +1148,23 @@ class OpenAIServingChat(OpenAIServing):
1148
1149 # Log streaming delta if output logging is enabled
1150 if self.enable_log_outputs and self.request_logger:
1142 - delta_content = ""
1151 + delta_content_parts = []
1152 if delta_message.content:
1144 - delta_content = delta_message.content
1145 - elif delta_message.tool_calls:
1146 - delta_content = "".join(
1153 + delta_content_parts.append(delta_message.content)
1154 + if delta_message.reasoning_content:
1155 + reasoning = delta_message.reasoning_content
1156 + delta_content_parts.append(f"[reasoning: {reasoning}]")
1157 + if delta_message.tool_calls:
1158 + tool_args = "".join(
1159 tc.function.arguments
1160 for tc in delta_message.tool_calls
1161 if tc.function and tc.function.arguments
1162 )
1163 + if tool_args:
1164 + delta_content_parts.append(f"[tool_calls: {tool_args}]")
1165
1152 - if delta_content and self.enable_log_deltas:
1166 + if delta_content_parts and self.enable_log_deltas:
1167 + delta_content = " ".join(delta_content_parts)
1168 self.request_logger.log_outputs(
1169 request_id=request_id,
1170 outputs=delta_content,
vllm/entrypoints/openai/chat_completion/stream_harmony.py
+118 -48
@@ -7,6 +7,8 @@ This module handles the extraction of DeltaMessage objects from
7 harmony parser state during streaming chat completions.
8 """
9
10 +from typing import NamedTuple
11 +
12 from openai_harmony import StreamableParser
13
14 from vllm.entrypoints.chat_utils import make_tool_call_id
@@ -17,12 +19,16 @@ from vllm.entrypoints.openai.engine.protocol import (
19 )
20
21
22 +class TokenState(NamedTuple):
23 + channel: str | None
24 + recipient: str | None
25 + text: str
26 +
27 +
28 def extract_harmony_streaming_delta(
29 harmony_parser: StreamableParser,
22 - cur_channel: str | None,
23 - cur_recipient: str | None,
30 + token_states: list[TokenState],
31 prev_recipient: str | None,
25 - delta_text: str,
32 include_reasoning: bool,
33 ) -> tuple[DeltaMessage | None, bool]:
34 """
@@ -30,38 +36,81 @@ def extract_harmony_streaming_delta(
36
37 Args:
38 harmony_parser: The StreamableParser instance tracking parse state
33 - cur_channel: Current channel ("final", "analysis", "commentary", etc.)
34 - cur_recipient: Current recipient (e.g., "functions.my_func")
39 + token_states: List of TokenState tuples for each token
40 prev_recipient: Previous recipient for detecting tool call transitions
36 - delta_text: The text delta to include in the message
41 include_reasoning: Whether to include reasoning content
42
43 Returns:
44 A tuple of (DeltaMessage or None, tools_streamed_flag)
45 """
46 +
47 + if not token_states:
48 + return None, False
49 +
50 tools_streamed = False
51
44 - if cur_channel == "final":
45 - delta_message = DeltaMessage(content=delta_text)
46 - elif (
47 - (cur_channel == "commentary" or cur_channel == "analysis")
48 - and cur_recipient
49 - and cur_recipient.startswith("functions.")
50 - ):
51 - # Count completed tool calls to determine index
52 - base_index = 0
53 - for msg in harmony_parser.messages:
54 - if (
55 - (msg.channel == "commentary" or msg.channel == "analysis")
56 - and msg.recipient
57 - and msg.recipient.startswith("functions.")
58 - ):
59 - base_index += 1
60 -
61 - if prev_recipient != cur_recipient:
62 - tool_name = cur_recipient.split("functions.", 1)[1]
63 - delta_message = DeltaMessage(
64 - tool_calls=[
52 + # Group consecutive tokens with same channel/recipient
53 + groups: list[TokenState] = []
54 +
55 + current_channel = token_states[0].channel
56 + current_recipient = token_states[0].recipient
57 + current_text = token_states[0].text
58 +
59 + for i in range(1, len(token_states)):
60 + state = token_states[i]
61 + if state.channel == current_channel and state.recipient == current_recipient:
62 + current_text += state.text
63 + else:
64 + groups.append(TokenState(current_channel, current_recipient, current_text))
65 + current_channel = state.channel
66 + current_recipient = state.recipient
67 + current_text = state.text
68 +
69 + groups.append(TokenState(current_channel, current_recipient, current_text))
70 +
71 + # Process each group and create delta messages
72 + delta_message = None
73 + combined_content = ""
74 + combined_reasoning = ""
75 + tool_messages = []
76 + content_encountered = False
77 +
78 + # Calculate base_index once before the loop
79 + # This counts completed tool calls in messages
80 + base_index = 0
81 + for msg in harmony_parser.messages:
82 + if (
83 + (msg.channel == "commentary" or msg.channel == "analysis")
84 + and msg.recipient
85 + and msg.recipient.startswith("functions.")
86 + ):
87 + base_index += 1
88 +
89 + # If there's an ongoing tool call from previous chunk,
90 + # the next new tool call starts at base_index + 1
91 + if prev_recipient and prev_recipient.startswith("functions."):
92 + next_tool_index = base_index + 1
93 + # Ongoing call is at base_index
94 + ongoing_tool_index = base_index
95 + else:
96 + # No ongoing call, next new call is at base_index
97 + next_tool_index = base_index
98 + ongoing_tool_index = None
99 +
100 + for group in groups:
101 + if group.channel == "final":
102 + combined_content += group.text
103 + content_encountered = True
104 + elif (
105 + (group.channel == "commentary" or group.channel == "analysis")
106 + and group.recipient
107 + and group.recipient.startswith("functions.")
108 + ):
109 + opened_new_call = False
110 + if prev_recipient != group.recipient:
111 + # New tool call - emit the opening message
112 + tool_name = group.recipient.split("functions.", 1)[1]
113 + tool_messages.append(
114 DeltaToolCall(
115 id=make_tool_call_id(),
116 type="function",
@@ -69,32 +118,53 @@ def extract_harmony_streaming_delta(
118 name=tool_name,
119 arguments="",
120 ),
72 - index=base_index,
121 + index=next_tool_index,
122 )
74 - ]
75 - )
76 - elif delta_text:
77 - delta_message = DeltaMessage(
78 - tool_calls=[
123 + )
124 + opened_new_call = True
125 + prev_recipient = group.recipient
126 + # Increment for subsequent new tool calls
127 + next_tool_index += 1
128 +
129 + if group.text:
130 + # Stream arguments for the ongoing tool call
131 + if opened_new_call:
132 + # Just opened in this group
133 + tool_call_index = next_tool_index - 1
134 + else:
135 + # Continuing from previous chunk
136 + # If ongoing_tool_index is None here, it means
137 + # we're continuing a call but prev_recipient
138 + # wasn't a function. Use base_index.
139 + tool_call_index = (
140 + ongoing_tool_index
141 + if ongoing_tool_index is not None
142 + else base_index
143 + )
144 + tool_messages.append(
145 DeltaToolCall(
80 - index=base_index,
81 - function=DeltaFunctionCall(arguments=delta_text),
146 + index=tool_call_index,
147 + function=DeltaFunctionCall(arguments=group.text),
148 )
83 - ]
84 - )
85 - else:
86 - delta_message = None
149 + )
150 + elif group.channel == "commentary":
151 + # Tool call preambles meant to be shown to the user
152 + combined_content += group.text
153 + content_encountered = True
154 + elif group.channel == "analysis" and include_reasoning:
155 + combined_reasoning += group.text
156
88 - if delta_message is not None:
157 + # Combine all non-empty fields into a single message
158 + if content_encountered or combined_reasoning or tool_messages:
159 + delta_kwargs: dict[str, str | list[DeltaToolCall]] = {}
160 + if content_encountered:
161 + delta_kwargs["content"] = combined_content
162 + if combined_reasoning:
163 + delta_kwargs["reasoning"] = combined_reasoning
164 + if tool_messages:
165 + delta_kwargs["tool_calls"] = tool_messages
166 tools_streamed = True
90 - elif cur_channel == "commentary":
91 - # Tool call preambles meant to be shown to the user
92 - delta_message = DeltaMessage(content=delta_text)
93 - elif cur_channel == "analysis":
94 - if include_reasoning:
95 - delta_message = DeltaMessage(reasoning=delta_text)
96 - else:
97 - delta_message = None
167 + delta_message = DeltaMessage(**delta_kwargs)
168 else:
169 delta_message = None
170