24bfbd5d4d · tduncan/Nvidia-Optimization-Stack-VLLM

tests/parser/engine/test_gemma4_streaming_reasoning.py

+33

@@ -1616,3 +1616,36 @@ class TestCommaInStringValueRegression:

1616 assert result.tools_called is True

1617 args = json.loads(result.tool_calls[0].function.arguments)

1618 assert args["destination"] == "456 Oakwood Avenue, Rivermist, 83214"

1619 +

1620 +

1621 +class TestGemma4IsReasoningEnd:

1622 + @staticmethod

1623 + def _parser(thinking: bool):

1624 + return Gemma4Parser(

1625 + _make_tokenizer(_PLAIN_ANSWER_TOKENS),

1626 + chat_template_kwargs={"enable_thinking": thinking},

1627 + )

1628 +

1629 + @pytest.mark.parametrize(

1630 + ("ids", "thinking", "ended"),

1631 + [

1632 + ([CHANNEL_START_ID, 3000, CHANNEL_END_ID], True, True),

1633 + ([CHANNEL_START_ID, 3000, TOOL_CALL_START_ID], True, True),

1634 + ([CHANNEL_START_ID, 3000], True, False),

1635 + ([NEW_TURN_ID, 9100], True, False),

1636 + ([CHANNEL_END_ID, NEW_TURN_ID, 9100], True, False),

1637 + ([NEW_TURN_ID, 9100], False, True),

1638 + # The bundled template's thinking-off prompt: closed channel.

1639 + ([NEW_TURN_ID, 9100, CHANNEL_START_ID, 3000, CHANNEL_END_ID], False, True),

1640 + # A marker still wins with thinking off: the grammar honours it.

1641 + ([CHANNEL_START_ID, 3000], False, False),

1642 + # A tool-call opener after the channel closed: still ended.

1643 + (

1644 + [CHANNEL_START_ID, 3000, CHANNEL_END_ID, TOOL_CALL_START_ID, 9200],

1645 + True,

1646 + True,

1647 + ),

1648 + ],

1649 + )

1650 + def test_is_reasoning_end(self, ids, thinking, ended):

1651 + assert self._parser(thinking).is_reasoning_end(ids) is ended

tests/parser/engine/test_qwen3_reasoning.py

+39 -3

@@ -180,9 +180,8 @@ class TestIsReasoningEnd:

180 """Tool examples before the generation <think> must not end reasoning."""

181 assert not parser.is_reasoning_end([_TOOL_CALL_ID, _TEXT_ID, _THINK_START_ID])

182

183 - def test_paired_tool_call_not_end(self, parser):

184 - """Paired <tool_call>...</tool_call> (from template) is NOT end."""

185 - assert not parser.is_reasoning_end(

183 + def test_completed_tool_call_after_open_think_is_end(self, parser):

184 + assert parser.is_reasoning_end(

185 [_THINK_START_ID, 1, _TOOL_CALL_ID, 2, _TOOL_CALL_END_ID]

186 )

187

@@ -196,6 +195,43 @@ class TestIsReasoningEnd:

195 assert not parser.is_reasoning_end([])

196

197

198 +class TestIsReasoningEndDerivedFromGrammar:

199 + @pytest.mark.parametrize(

200 + ("ids", "thinking", "ended"),

201 + [

202 + # Explicit close, and the implicit close on <tool_call>.

203 + ([_THINK_START_ID, 1, _THINK_END_ID], True, True),

204 + ([_THINK_START_ID, 1, _TOOL_CALL_ID], True, True),

205 + # Still inside the block.

206 + ([_THINK_START_ID, 1], True, False),

207 + ([_THINK_END_ID, _THINK_START_ID, 1], True, False),

208 + # Uncommitted: a bare turn, an empty sequence. Thinking on starts

209 + # in REASONING, so reasoning is expected.

210 + ([_IM_END_ID, _IM_START_ID, 1], True, False),

211 + ([], True, False),

212 + # Thinking off starts in CONTENT with no way into REASONING, so

213 + # a <think> the template left behind is inert, and a silent tail

214 + # is ended because no block is expected.

215 + ([_THINK_START_ID, 1], False, True),

216 + ([_THINK_END_ID, _THINK_START_ID, 1], False, True),

217 + ([_IM_END_ID, _IM_START_ID, 1], False, True),

218 + ([], False, True),

219 + ],

220 + )

221 + def test_is_reasoning_end(self, mock_tokenizer, ids, thinking, ended):

222 + parser = Qwen3Parser(

223 + mock_tokenizer, chat_template_kwargs={"enable_thinking": thinking}

224 + )

225 + assert parser.is_reasoning_end(ids) is ended

226 +

227 + @pytest.mark.parametrize("thinking", [True, False])

228 + def test_wait_for_reasoning_follows_thinking(self, mock_tokenizer, thinking):

229 + parser = Qwen3Parser(

230 + mock_tokenizer, chat_template_kwargs={"enable_thinking": thinking}

231 + )

232 + assert parser.parser_engine_config.wait_for_reasoning is thinking

233 +

234 +

235 class TestIsReasoningEndTurnBoundaries:

236 """Reasoning markers from earlier turns must not mark the new turn's

237 reasoning as finished (think-marker injection / interleaved-thinking

vllm/parser/deepseek_v32.py

+1

@@ -53,6 +53,7 @@ def deepseek_v32_config() -> ParserEngineConfig:

53 return ParserEngineConfig(

54 name="deepseek_v32",

55 initial_state=ParserState.CONTENT,

56 + wait_for_reasoning=False,

57 terminals={

58 "TOOL_START": DSML_FUNC_START,

59 "TOOL_END": DSML_FUNC_END,

vllm/parser/deepseek_v4.py

+1

@@ -140,6 +140,7 @@ def deepseek_v4_config(thinking: bool = False) -> ParserEngineConfig:

140 return ParserEngineConfig(

141 name="deepseek_v4",

142 initial_state=ParserState.REASONING if thinking else ParserState.CONTENT,

143 + wait_for_reasoning=thinking,

144 terminals={

145 "THINK_START": DSML_THINK_START,

146 "THINK_END": DSML_THINK_END,

vllm/parser/engine/parser_engine.py

+20 -16

@@ -648,24 +648,28 @@ class ParserEngine(Parser):

648 return None

649

650 def is_reasoning_end(self, input_ids: list[int]) -> bool:

651 - end_id = self._reasoning_end_token_id

651 + config = self.parser_engine_config

652 + wait_for_reasoning = config.wait_for_reasoning

653 + if wait_for_reasoning is None:

654 + wait_for_reasoning = config.initial_state is ParserState.REASONING

655 + start_transition = config.transitions.get((config.initial_state, "THINK_START"))

656 + start_opens_reasoning = (

657 + start_transition is not None

658 + and start_transition.next_state is ParserState.REASONING

659 + )

660 + end_ids = self._reasoning_end_token_ids

661 start_id = self._reasoning_start_token_id

653 - if end_id is not None:

654 - if not input_ids:

655 - return self.parser_engine_config.initial_state != ParserState.REASONING

656 - boundary_ids = self._turn_boundary_token_ids

657 - for i in range(len(input_ids) - 1, -1, -1):

658 - token_id = input_ids[i]

659 - if token_id == end_id:

660 - return True

661 - if start_id is not None and token_id == start_id:

662 + boundary_ids = self._turn_boundary_token_ids

663 + for token_id in reversed(input_ids):

664 + if token_id in end_ids:

665 + return True

666 + if token_id == start_id:

667 + if start_opens_reasoning:

668 return False

663 - if token_id in boundary_ids:

664 - return (

665 - self.parser_engine_config.initial_state != ParserState.REASONING

666 - )

667 - return False

668 - return self._reasoning_ended

669 + break

670 + if token_id in boundary_ids:

671 + break

672 + return not wait_for_reasoning

673

674 def extract_content_ids(self, input_ids: list[int]) -> list[int]:

675 end_id = self._reasoning_end_token_id

vllm/parser/engine/parser_engine_config.py

+2

@@ -74,6 +74,8 @@ class ParserEngineConfig:

74

75 initial_state: ParserState = ParserState.CONTENT

76

77 + wait_for_reasoning: bool | None = None

78 +

79 arg_converter: Callable[[str, bool], str] | None = None

80

81 stream_arg_deltas: bool = True

vllm/parser/gemma4.py

+7 -27

@@ -293,10 +293,14 @@ def _gemma4_arg_converter(raw_args: str, partial: bool) -> str:

293

294

295 @functools.cache

296 -def gemma4_config() -> ParserEngineConfig:

296 +def gemma4_config(thinking: bool = False) -> ParserEngineConfig:

297 return ParserEngineConfig(

298 name="gemma4",

299 initial_state=ParserState.CONTENT,

300 + # Decides a silent tail: the E-series HF templates write no marker

301 + # with thinking off, and no template does after a tool response.

302 + wait_for_reasoning=thinking,

303 + turn_boundary_tokens=frozenset({"<|turn>", "<|tool_response>"}),

304 terminals={

305 "THINK_START": CHANNEL_START,

306 "THINK_END": CHANNEL_END,

@@ -417,11 +421,11 @@ class Gemma4Parser(ParserEngine):

421 **kwargs,

422 ) -> None:

423 chat_kwargs = kwargs.get("chat_template_kwargs", {}) or {}

420 - self._thinking_enabled = chat_kwargs.get("enable_thinking", False)

424 + thinking = chat_kwargs.get("enable_thinking", False)

425 super().__init__(

426 tokenizer,

427 tools,

424 - parser_engine_config=gemma4_config(),

428 + parser_engine_config=gemma4_config(thinking=thinking),

429 **kwargs,

430 )

431 vocab = self.vocab

@@ -472,30 +476,6 @@ class Gemma4Parser(ParserEngine):

476

477 return delta_text, delta_token_ids

478

475 - def is_reasoning_end(self, input_ids: list[int]) -> bool:

476 - end_id = self._reasoning_end_token_id

477 - start_id = self._reasoning_start_token_id

478 - tool_call_id = self._tool_call_token_id

479 - new_turn_id = self._new_turn_token_id

480 - tool_response_id = self._tool_response_token_id

481 -

482 - if end_id is not None and not input_ids:

483 - return self.parser_engine_config.initial_state != ParserState.REASONING

484 -

485 - for i in range(len(input_ids) - 1, -1, -1):

486 - tid = input_ids[i]

487 - if start_id is not None and tid == start_id:

488 - return False

489 - if tool_call_id is not None and tid == tool_call_id:

490 - return True

491 - if new_turn_id is not None and tid == new_turn_id:

492 - return not self._thinking_enabled

493 - if tool_response_id is not None and tid == tool_response_id:

494 - return not self._thinking_enabled

495 - if end_id is not None and tid == end_id:

496 - return True

497 - return True

498 -

479 def _prompt_ends_in_open_reasoning(self, prompt_token_ids: Sequence[int]) -> bool:

480 """Whether the prompt tail is inside an open ``<|channel>`` block.

481

vllm/parser/glm47_moe.py

+1 -5

@@ -123,6 +123,7 @@ def glm47_moe_config(thinking: bool = True) -> ParserEngineConfig:

123 return ParserEngineConfig(

124 name="glm47_moe",

125 initial_state=ParserState.REASONING if thinking else ParserState.CONTENT,

126 + wait_for_reasoning=thinking,

127 terminals={

128 **reasoning_terminals,

129 "TOOL_START": TOOL_CALL_START,

@@ -206,11 +207,6 @@ class Glm47MoeParser(ParserEngine):

207 self._tool_slots[idx].name = self._tool_slots[idx].name.strip()

208 super()._handle_tool_end(event, deltas)

209

209 - def is_reasoning_end(self, input_ids: list[int]) -> bool:

210 - if not self.thinking_enabled:

211 - return True

212 - return super().is_reasoning_end(input_ids)

213 -

210 def extract_content_ids(self, input_ids: list[int]) -> list[int]:

211 if not self.thinking_enabled:

212 return input_ids

vllm/parser/inkling.py

+1

@@ -282,6 +282,7 @@ def inkling_config() -> ParserEngineConfig:

282 # `<|message_model|>`. Non-streaming parsing receives only the generated

283 # suffix, so begin in the corresponding message-header state as well.

284 initial_state=ParserState.MESSAGE_HEADER,

285 + wait_for_reasoning=True,

286 terminals=terminals,

287 # Inkling content-kind markers are the grammar. When the engine is

288 # used through DelegatingParser, the reasoning pass can hand the tool

vllm/parser/kimi_k2.py

+1 -18

@@ -81,6 +81,7 @@ def kimi_k2_config(thinking: bool = True) -> ParserEngineConfig:

81 return ParserEngineConfig(

82 name="kimi_k2",

83 initial_state=ParserState.REASONING if thinking else ParserState.CONTENT,

84 + wait_for_reasoning=thinking,

85 terminals={

86 **reasoning_terminals,

87 "TOOL_SECTION_START": TOOL_SECTION_START,

@@ -236,24 +237,6 @@ class KimiK2Parser(ParserEngine):

237 def _extract_args_json(self, raw_args: str, func_name: str) -> str:

238 return raw_args.strip() or "{}"

239

239 - def is_reasoning_end(self, input_ids: list[int]) -> bool:

240 - if not self.thinking_enabled:

241 - return True

242 -

243 - start_id = self._start_token_id

244 - end_id = self._end_token_id

245 - tool_section_id = self._tool_section_start_token_id

246 -

247 - for i in range(len(input_ids) - 1, -1, -1):

248 - token_id = input_ids[i]

249 - if start_id is not None and token_id == start_id:

250 - return False

251 - if end_id is not None and token_id == end_id:

252 - return True

253 - if tool_section_id is not None and token_id == tool_section_id:

254 - return True

255 - return False

256 -

240 def extract_content_ids(self, input_ids: list[int]) -> list[int]:

241 if not self.thinking_enabled:

242 return input_ids

vllm/parser/minimax_m2.py

+1

@@ -95,6 +95,7 @@ def minimax_m2_config() -> ParserEngineConfig:

95 return ParserEngineConfig(

96 name="minimax_m2",

97 initial_state=ParserState.REASONING,

98 + wait_for_reasoning=True,

99 terminals={

100 "THINK_START": THINK_START,

101 "THINK_END": THINK_END,

vllm/parser/mistral.py

+1 -18

@@ -193,6 +193,7 @@ def mistral_config(

193 return ParserEngineConfig(

194 name=name,

195 initial_state=ParserState.CONTENT,

196 + wait_for_reasoning=reasoning_encoding != "none",

197 terminals={

198 **reasoning_terminals,

199 "TOOL_CALLS": _TOOL_CALLS,

@@ -891,24 +892,6 @@ class MistralParser(ParserEngine):

892

893 return (delta_to_be_parsed, "")

894

894 - def is_reasoning_end(self, input_ids: list[int]) -> bool:

895 - if self._reasoning_encoding == "none":

896 - return True

897 - if super().is_reasoning_end(input_ids):

898 - return True

899 - # [TOOL_CALLS] acts as an implicit reasoning-end marker

900 - if self.bot_token_id is not None:

901 - reasoning_start_id = self._reasoning_start_token_id

902 - for i in range(len(input_ids) - 1, -1, -1):

903 - if (

904 - reasoning_start_id is not None

905 - and input_ids[i] == reasoning_start_id

906 - ):

907 - return False

908 - if input_ids[i] == self.bot_token_id:

909 - return True

910 - return False

911 -

895 def extract_reasoning(

896 self,

897 model_output: str,

vllm/parser/qwen3.py

+3 -30

@@ -100,6 +100,7 @@ def qwen3_config(

100 return ParserEngineConfig(

101 name=name,

102 initial_state=ParserState.REASONING if thinking else ParserState.CONTENT,

103 + wait_for_reasoning=thinking,

104 turn_boundary_tokens=turn_boundary_tokens,

105 terminals={

106 # Reasoning terminals

@@ -205,8 +206,8 @@ class Qwen3Parser(ParserEngine):

206 """Qwen3 parser: ``<think>``/``</think>`` reasoning +

207 ``<tool_call>`` XML tool calls in a single engine.

208

208 - - ``<tool_call>`` as implicit reasoning end

209 - - Unpaired ``<tool_call>`` token ID detection for ``is_reasoning_end``

209 + - ``<tool_call>`` as implicit reasoning end (a grammar transition, so it

210 + also feeds ``is_reasoning_end`` and the structured-output gate)

211

212 Subclasses that share the grammar but differ only in the four wrapper

213 token strings (reasoning + tool-call) override the class attributes

@@ -245,9 +246,6 @@ class Qwen3Parser(ParserEngine):

246 tools,

247 **kwargs,

248 )

248 - vocab = self.vocab

249 - self._tool_call_token_id: int | None = vocab.get(self.TOOL_START)

250 - self._tool_call_end_token_id: int | None = vocab.get(self.TOOL_END)

249

250 def extract_reasoning(

251 self,

@@ -257,28 +255,3 @@ class Qwen3Parser(ParserEngine):

255 if not self.thinking_enabled:

256 return None, model_output

257 return super().extract_reasoning(model_output, request)

260 -

261 - def is_reasoning_end(self, input_ids: list[int]) -> bool:

262 - if super().is_reasoning_end(input_ids):

263 - return True

264 - tool_call_id = self._tool_call_token_id

265 - tool_call_end_id = self._tool_call_end_token_id

266 - reasoning_start_id = self._reasoning_start_token_id

267 - boundary_ids = self._turn_boundary_token_ids

268 - if tool_call_id is not None:

269 - for i in range(len(input_ids) - 1, -1, -1):

270 - if (

271 - reasoning_start_id is not None

272 - and input_ids[i] == reasoning_start_id

273 - ):

274 - return False

275 - if input_ids[i] in boundary_ids:

276 - return False

277 - if input_ids[i] == tool_call_id:

278 - if tool_call_end_id is not None and any(

279 - input_ids[j] == tool_call_end_id

280 - for j in range(i + 1, len(input_ids))

281 - ):

282 - continue

283 - return True

284 - return False

[Refactor] Derive is_reasoning_end from the engine grammar (#56200)