@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 5769a7382c

[ROCm][CI][Bugfix] Fix flaky parallel tool-call streaming (test assertion + Mistral/Granite parsers) (#47550)

Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com> Signed-off-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>

akii96 committed Jul 7, 2026 at 02:06 UTC 5769a7382cb111288a9c2927342ca908943b9793

6 files changed +142 -47

tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py

+23 -11

@@ -183,27 +183,39 @@ async def test_thinking_token_budget_mixed_requests(client: openai.AsyncOpenAI):

183 async def test_thinking_token_budget_limits_reasoning(client: openai.AsyncOpenAI):

184 """Test that thinking_token_budget limits the number of reasoning tokens.

185

186 - Counts non-empty streaming ``delta.reasoning`` chunks (coarse proxy; each

187 - chunk may represent multiple decode tokens — see

188 - ``_count_reasoning_decode_token_ids_between_markers`` and the Qwen3.5 MTP

189 - test for id-based checks).

186 + Counts reasoning decode tokens by id, which is robust to how tokens are

187 + grouped into streamed chunks (a single chunk can carry several tokens under

188 + async scheduling / stream_interval > 1). Counting chunks under-counts.

189 """

190

192 - reasoning_token_count = 0

191 + tokenizer = get_tokenizer(tokenizer_name=MODEL_NAME)

192 + start_ids = list(tokenizer.encode(REASONING_START_STR, add_special_tokens=False))

193 + end_ids = list(tokenizer.encode(REASONING_END_STR, add_special_tokens=False))

194 +

195 + prompt_token_ids: list[int] = []

196 + decode_token_ids: list[int] = []

197 stream = await client.chat.completions.create(

198 model=MODEL_NAME,

199 messages=MESSAGES,

200 max_tokens=100,

201 stream=True,

198 - extra_body={"thinking_token_budget": THINK_BUDGET},

202 + extra_body={"thinking_token_budget": THINK_BUDGET, "return_token_ids": True},

203 )

204 async for chunk in stream:

201 - delta = chunk.choices[0].delta

202 - if getattr(delta, "reasoning", None):

203 - reasoning_token_count += 1

204 -

205 + if not chunk.choices:

206 + continue

207 + if getattr(chunk, "prompt_token_ids", None):

208 + prompt_token_ids = list(chunk.prompt_token_ids)

209 + delta_ids = getattr(chunk.choices[0], "token_ids", None)

210 + if delta_ids:

211 + decode_token_ids.extend(delta_ids)

212 +

213 + reasoning_token_count = _count_reasoning_decode_token_ids_between_markers(

214 + prompt_token_ids + decode_token_ids, start_ids, end_ids

215 + )

216 + assert reasoning_token_count is not None, "missing reasoning start marker in ids"

217 assert reasoning_token_count == THINK_BUDGET, (

206 - f"reasoning tokens ({reasoning_token_count}) exceeded "

218 + f"reasoning tokens ({reasoning_token_count}) != "

219 f"thinking_token_budget ({THINK_BUDGET})"

220 )

221

tests/tool_parsers/test_granite_tool_parser.py

+44 -1

@@ -2,13 +2,21 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

5 +import json

6 +

7 import pytest

9 from tests.tool_parsers.common_tests import (

10 ToolParserTestConfig,

11 ToolParserTests,

12 )

11 -from tests.tool_parsers.utils import run_tool_extraction

13 +from tests.tool_parsers.utils import (

14 + run_tool_extraction,

15 + run_tool_extraction_streaming,

16 + split_string_into_token_deltas,

17 +)

18 +from vllm.tokenizers import get_tokenizer

19 +from vllm.tool_parsers.granite_tool_parser import GraniteToolParser

22 class TestGraniteToolParser(ToolParserTests):

@@ -116,3 +124,38 @@ I'll get that information.""",

124 f"Expected 1 tool call from string format, got {len(tool_calls)}"

125 )

126 assert tool_calls[0].function.name == "get_weather"

127 +

128 +

129 +# granite emits arguments before name and its own tokenizer (not gpt2) is used

130 +# here so the token boundaries match production; get_tokenizer only fetches the

131 +# small tokenizer files, not the model weights.

132 +@pytest.fixture(scope="module")

133 +def granite_tokenizer():

134 + return get_tokenizer(tokenizer_name="ibm-granite/granite-3.1-8b-instruct")

135 +

136 +

137 +@pytest.mark.parametrize("chunk_size", [2, 3, 4, 5])

138 +def test_streaming_parallel_calls_batched_deltas(granite_tokenizer, chunk_size):

139 + """A batched delta (multiple tokens) spanning the boundary between two

140 + parallel calls must not drop the first call's name. granite streams

141 + arguments before name, so the name only completes as the next call appears.

142 + """

143 + parser = GraniteToolParser(granite_tokenizer)

144 + model_output = (

145 + '<|tool_call|> [{"arguments": {"city": "Tokyo"}, "name": "get_weather"}, '

146 + '{"arguments": {"timezone": "Asia/Tokyo"}, "name": "get_time"}]'

147 + )

148 + token_deltas = split_string_into_token_deltas(granite_tokenizer, model_output)

149 + batched = [

150 + "".join(token_deltas[i : i + chunk_size])

151 + for i in range(0, len(token_deltas), chunk_size)

152 + ]

153 + reconstructor = run_tool_extraction_streaming(

154 + parser, batched, assert_one_tool_per_delta=False

155 + )

156 + names = [tc.function.name for tc in reconstructor.tool_calls]

157 + assert names == ["get_weather", "get_time"]

158 + # trailing args of the final call are flushed by the serving layer

159 + assert json.loads(reconstructor.tool_calls[0].function.arguments) == {

160 + "city": "Tokyo"

161 + }

tests/tool_parsers/test_mistral_tool_parser.py

+60 -13

@@ -144,6 +144,7 @@ def stream_delta_message_generator(

144 mistral_tokenizer: TokenizerLike,

145 model_output: str | None,

146 tools: list[tuple[str, str]] | None,

147 + chunk_size: int = 1,

148 ) -> Generator[DeltaMessage, None, None]:

149 if (

150 isinstance(mistral_tokenizer, MistralTokenizer)

@@ -182,15 +183,13 @@ def stream_delta_message_generator(

183 previous_tokens = None

184 prefix_offset = 0

185 read_offset = 0

186 + pending_text = ""

187 + pending_token_ids: list[int] = []

188 for i, delta_token in enumerate(all_token_ids):

186 - delta_token_ids = [delta_token]

187 - previous_token_ids = all_token_ids[:i]

188 - current_token_ids = all_token_ids[: i + 1]

189 -

189 (new_tokens, delta_text, new_prefix_offset, new_read_offset) = (

190 detokenize_incrementally(

191 tokenizer=mistral_tokenizer,

193 - all_input_ids=current_token_ids,

192 + all_input_ids=all_token_ids[: i + 1],

193 prev_tokens=previous_tokens,

194 prefix_offset=prefix_offset,

195 read_offset=read_offset,

@@ -198,27 +197,39 @@ def stream_delta_message_generator(

197 spaces_between_special_tokens=True,

198 )

199 )

200 + previous_tokens = (

201 + previous_tokens + new_tokens if previous_tokens else new_tokens

202 + )

203 + prefix_offset = new_prefix_offset

204 + read_offset = new_read_offset

205

202 - current_text = previous_text + delta_text

206 + # Buffer tokens so each streamed delta can carry ``chunk_size`` tokens,

207 + # reproducing the multi-token deltas produced by async scheduling /

208 + # stream_interval > 1.

209 + pending_text += delta_text

210 + pending_token_ids.append(delta_token)

211 + if len(pending_token_ids) < chunk_size and i != len(all_token_ids) - 1:

212 + continue

213 +

214 + previous_token_ids = all_token_ids[: i + 1 - len(pending_token_ids)]

215 + current_token_ids = all_token_ids[: i + 1]

216 + current_text = previous_text + pending_text

217

218 delta_message = mistral_tool_parser.extract_tool_calls_streaming(

219 previous_text,

220 current_text,

207 - delta_text,

221 + pending_text,

222 previous_token_ids,

223 current_token_ids,

210 - delta_token_ids,

224 + pending_token_ids,

225 request=_DUMMY_REQUEST,

226 )

227 if delta_message:

228 yield delta_message

229

230 previous_text = current_text

217 - previous_tokens = (

218 - previous_tokens + new_tokens if previous_tokens else new_tokens

219 - )

220 - prefix_offset = new_prefix_offset

221 - read_offset = new_read_offset

231 + pending_text = ""

232 + pending_token_ids = []

233

234

235 @pytest.mark.parametrize(

@@ -1572,3 +1583,39 @@ def test_grammar_from_tool_parser_set_by_adjust_request(

1583 request = _make_request()

1584 result = mistral_tool_parser.adjust_request(request)

1585 assert result._grammar_from_tool_parser is True

1586 +

1587 +

1588 +@pytest.mark.parametrize("chunk_size", [2, 3, 4, 5])

1589 +def test_streaming_pre_v11_parallel_calls_batched_deltas(

1590 + mistral_pre_v11_tool_parser, mistral_pre_v11_tokenizer, chunk_size

1591 +):

1592 + """A batched delta spanning the boundary between two parallel calls must

1593 + keep them on distinct indices (the bug collapsed both onto index 0)."""

1594 + model_output = (

1595 + '[TOOL_CALLS] [{"name": "add", "arguments": {"a": 3.5, "b": 4}}, '

1596 + '{"name": "get_current_weather", "arguments": '

1597 + '{"city": "San Francisco", "state": "CA", "unit": "celsius"}}]'

1598 + )

1599 + names: list[str] = []

1600 + args: list[str] = []

1601 + idx = -1

1602 + for delta_message in stream_delta_message_generator(

1603 + mistral_pre_v11_tool_parser,

1604 + mistral_pre_v11_tokenizer,

1605 + model_output,

1606 + tools=None,

1607 + chunk_size=chunk_size,

1608 + ):

1609 + for tool_call in delta_message.tool_calls or []:

1610 + if tool_call.index != idx:

1611 + idx = tool_call.index

1612 + args.append("")

1613 + if tool_call.function and tool_call.function.name:

1614 + names.append(tool_call.function.name)

1615 + if tool_call.function and tool_call.function.arguments:

1616 + args[tool_call.index] += tool_call.function.arguments

1617 +

1618 + assert names == ["add", "get_current_weather"]

1619 + assert len(args) == 2

1620 + # trailing args of the final call are flushed by the serving layer

1621 + assert json.loads(args[0]) == {"a": 3.5, "b": 4}

tests/tool_use/test_parallel_tool_calls.py

+4 -6

@@ -115,14 +115,12 @@ async def test_parallel_tool_calls(

115 assert not role_name or role_name == "assistant"

116 role_name = "assistant"

117

118 - # if a tool call is streamed make sure there's exactly one

119 - # (based on the request parameters

118 + # a chunk may carry >1 tool-call delta at a parallel-call boundary

119 streamed_tool_calls = chunk.choices[0].delta.tool_calls

120

122 - if streamed_tool_calls and len(streamed_tool_calls) > 0:

123 - # make sure only one diff is present - correct even for parallel

124 - assert len(streamed_tool_calls) == 1

125 - tool_call = streamed_tool_calls[0]

121 + for tool_call in streamed_tool_calls or []:

122 + # deltas arrive in non-decreasing index order

123 + assert tool_call.index >= tool_call_idx

124

125 # if a new tool is being called, set up empty arguments

126 if tool_call.index != tool_call_idx:

vllm/tool_parsers/granite_tool_parser.py

+6 -4

@@ -154,9 +154,11 @@ class GraniteToolParser(ToolParser):

154 current_tool_call: dict = tool_call_arr[self.current_tool_id]

155

156 delta = None

157 - # case: we are starting a new tool in the array

158 - # -> array has > 0 length AND length has moved past cursor

159 - if len(tool_call_arr) > self.current_tool_id + 1:

157 + # Only advance once the current tool name is streamed; granite

158 + # emits arguments before name, so advancing early would drop it.

159 + if len(tool_call_arr) > self.current_tool_id + 1 and (

160 + self.current_tool_id < 0 or self.current_tool_name_sent

161 + ):

162 # if we're moving on to a new call, first make sure we

163 # haven't missed anything in the previous one that was

164 # auto-generated due to JSON completions, but wasn't

@@ -184,7 +186,7 @@ class GraniteToolParser(ToolParser):

186 )

187

188 # re-set stuff pertaining to progress in the current tool

187 - self.current_tool_id = len(tool_call_arr) - 1

189 + self.current_tool_id += 1

190 self.current_tool_name_sent = False

191 self.streamed_args_for_tool.append("")

192 logger.debug("starting on new tool %d", self.current_tool_id)

vllm/tool_parsers/mistral_tool_parser.py

+5 -12

@@ -533,6 +533,7 @@ class MistralToolParser(ToolParser):

533

534 if prefix == "item" and event == "start_map":

535 self.streaming_state = StreamingState.WAITING_FOR_TOOL_KEY

536 + self.starting_new_tool = True

537 if prefix == "item" and event == "map_key" and value == "name":

538 self.streaming_state = StreamingState.PARSING_NAME

539 if prefix == "item.name" and event == "string":

@@ -640,18 +641,10 @@ class MistralToolParser(ToolParser):

641

642 # Given the parsed text and the possible streaming state change,

643 # let's add to the tool delta

643 - if (

644 - (streaming_state_before_parse != self.streaming_state)

645 - and streaming_state_before_parse

646 - in [StreamingState.WAITING_FOR_TOOL_START, StreamingState.TOOL_COMPLETE]

647 - and self.streaming_state

648 - not in [

649 - StreamingState.ALL_TOOLS_COMPLETE,

650 - StreamingState.TOOL_COMPLETE,

651 - StreamingState.WAITING_FOR_TOOL_START,

652 - ]

653 - ):

654 - # starting a new tool call

644 + # start_map is the authoritative new-tool signal and survives

645 + # batched deltas, unlike comparing pre/post streaming states

646 + if self.starting_new_tool:

647 + self.starting_new_tool = False

648 if current_tool_call_modified:

649 if self.current_tool_mistral_id is not None:

650 current_tool_call.id = self.current_tool_mistral_id