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

A high-throughput and memory-efficient inference and serving engine for LLMs

[Bugfix][Gemma4] Pre-initialise streaming reasoning state when prompt ends inside an open `<|channel>` (fixes #45834) (#45852)

Signed-off-by: nikhilesh-csa <nchhetri@csa1.com>

Nikhilesh Chhetri committed Jun 17, 2026 at 05:16 UTC 3c6084bb0d5168be132df395d1859098807f0fec
6 files changed +270
tests/parser/engine/test_gemma4_streaming_reasoning.py
+208
@@ -25,12 +25,14 @@ CHANNEL_END_ID = 51 # <channel|>
25 TOOL_CALL_START_ID = 48 # <|tool_call>
26 TOOL_CALL_END_ID = 49 # <tool_call|>
27 QUOTED_ID = 52 # <|"|>
28 +NEW_TURN_ID = 53 # <|turn>
29 SPECIAL_TOKEN_MAP = {
30 CHANNEL_START_ID: "<|channel>",
31 CHANNEL_END_ID: "<channel|>",
32 TOOL_CALL_START_ID: "<|tool_call>",
33 TOOL_CALL_END_ID: "<tool_call|>",
34 QUOTED_ID: '<|"|>',
35 + NEW_TURN_ID: "<|turn>",
36 }
37
38 SPECIAL_TEXT_TO_ID = {v: k for k, v in SPECIAL_TOKEN_MAP.items()}
@@ -253,6 +255,212 @@ class TestGemma4StreamingReasoningThenToolCall:
255 )
256
257
258 +# ── Prompt ends inside an open <|channel>thought\n block ─────────────
259 +
260 +_OPEN_REASONING_GEN_SEQUENCE: list[tuple[int, str]] = [
261 + (7001, "Sure"),
262 + (7002, ","),
263 + (7003, " the"),
264 + (7004, " answer"),
265 + (7005, " is"),
266 + (7006, " 42"),
267 + (CHANNEL_END_ID, "<channel|>"),
268 + (7007, "Hello"),
269 + (7008, " world"),
270 +]
271 +
272 +
273 +class TestGemma4PromptOpenReasoning:
274 + """When ``add_generation_prompt=True`` after a final tool response with
275 + ``enable_thinking=True``, the Gemma4 chat template leaves the prompt
276 + ending with ``<|channel>thought\\n`` — i.e. inside an open reasoning
277 + channel. Tokens generated before ``<channel|>`` must be classified as
278 + ``reasoning``, not visible ``content``.
279 +
280 + Regression test for vllm-project/vllm#45834.
281 + """
282 +
283 + @pytest.fixture
284 + def open_reasoning_tokenizer(self):
285 + return _make_tokenizer(_OPEN_REASONING_GEN_SEQUENCE)
286 +
287 + @pytest.fixture
288 + def open_reasoning_parser(self, open_reasoning_tokenizer):
289 + return Gemma4Parser(open_reasoning_tokenizer)
290 +
291 + @staticmethod
292 + def _prompt_ids_open_channel() -> list[int]:
293 + # Mimics a prompt that ends with ``...<|channel>thought\n``. The
294 + # specific token ids for ``thought`` and ``\n`` are arbitrary — only
295 + # the trailing ``<|channel>`` start token matters for detection.
296 + return [CHANNEL_START_ID, 3000, 3001]
297 +
298 + def test_reasoning_not_leaked_into_content(
299 + self, open_reasoning_parser, open_reasoning_tokenizer, request_obj
300 + ):
301 + results = _stream_tokens_batched(
302 + open_reasoning_parser,
303 + open_reasoning_tokenizer,
304 + request_obj,
305 + batch_size=1,
306 + prompt_token_ids=self._prompt_ids_open_channel(),
307 + )
308 +
309 + reasoning, content, _ = _collect_fields(results)
310 +
311 + assert "Sure, the answer is 42" in reasoning, (
312 + f"Expected pre-<channel|> tokens in reasoning, got "
313 + f"reasoning={reasoning!r} content={content!r}"
314 + )
315 + for leaked in ("Sure", "answer", "42"):
316 + assert leaked not in content, (
317 + f"Reasoning text leaked into content: {content!r}"
318 + )
319 +
320 + def test_post_reasoning_text_in_content(
321 + self, open_reasoning_parser, open_reasoning_tokenizer, request_obj
322 + ):
323 + results = _stream_tokens_batched(
324 + open_reasoning_parser,
325 + open_reasoning_tokenizer,
326 + request_obj,
327 + batch_size=1,
328 + prompt_token_ids=self._prompt_ids_open_channel(),
329 + )
330 +
331 + _, content, _ = _collect_fields(results)
332 +
333 + assert "Hello world" in content, (
334 + f"Post-<channel|> text missing from content: {content!r}"
335 + )
336 +
337 + def test_new_turn_prompt_unchanged(self, parser, mock_tokenizer, request_obj):
338 + """When the prompt does NOT end in an open reasoning channel (e.g. a
339 + new turn that ends with ``<|turn>model\\n``), behaviour must match
340 + the existing flow — the model itself opens ``<|channel>``.
341 + """
342 + results = _stream_tokens_batched(
343 + parser,
344 + mock_tokenizer,
345 + request_obj,
346 + batch_size=10,
347 + # No <|channel> in the prompt tail.
348 + prompt_token_ids=[9000, 9001],
349 + )
350 +
351 + reasoning, content, tool_calls = _collect_fields(results)
352 +
353 + assert "weather" in reasoning.lower(), (
354 + f"Expected reasoning about weather, got: {reasoning[:100]!r}"
355 + )
356 + assert len(tool_calls) > 0, f"Tool calls missing — content={content!r}"
357 +
358 +
359 +# ── Engine pre-initialised to REASONING + model still emits channel open ──
360 +
361 +_PRE_INIT_THOUGHT_GEN_SEQUENCE: list[tuple[int, str]] = [
362 + # Model naively emits the full reasoning opener even though the engine
363 + # was pre-initialised to REASONING from the prompt.
364 + (CHANNEL_START_ID, "<|channel>"),
365 + (8000, "thought"),
366 + (8001, "\n"),
367 + (8002, "Reason"),
368 + (8003, "ing"),
369 + (8004, " body"),
370 + (CHANNEL_END_ID, "<channel|>"),
371 + (8005, "Final"),
372 + (8006, " content"),
373 +]
374 +
375 +
376 +class TestGemma4PreInitReasoningRobustness:
377 + """Tests for the ``(REASONING, THINK_START)`` no-op transition and
378 + cooperating ``thought\\n`` prefix stripping when the engine has been
379 + pre-initialised to ``REASONING`` from the prompt.
380 +
381 + These cover the case the reviewer raised: prompt ends with
382 + ``<|turn>model\\n`` (``is_reasoning_end`` returns ``False`` because
383 + thinking is enabled, so the engine is pre-initialised), but the model
384 + still emits its own ``<|channel>thought\\n…<channel|>content``. The
385 + ``thought\\n`` prefix must be stripped, the ``<|channel>`` must not
386 + leak as text, and the post-``<channel|>`` text must appear as content.
387 + """
388 +
389 + @pytest.fixture
390 + def pre_init_tokenizer(self):
391 + return _make_tokenizer(_PRE_INIT_THOUGHT_GEN_SEQUENCE)
392 +
393 + @pytest.fixture
394 + def pre_init_parser(self, pre_init_tokenizer):
395 + return Gemma4Parser(pre_init_tokenizer)
396 +
397 + def test_redundant_channel_open_swallowed_after_new_turn(
398 + self, pre_init_parser, pre_init_tokenizer, request_obj
399 + ):
400 + # Prompt ends with ``<|turn>model\n``-style sentinel. With
401 + # ``enable_thinking=True`` (the default), ``is_reasoning_end``
402 + # returns ``False`` for a ``<|turn>`` tail, so the engine is
403 + # pre-initialised to ``REASONING``.
404 + results = _stream_tokens_batched(
405 + pre_init_parser,
406 + pre_init_tokenizer,
407 + request_obj,
408 + batch_size=1,
409 + prompt_token_ids=[NEW_TURN_ID, 9100, 9101],
410 + )
411 +
412 + reasoning, content, _ = _collect_fields(results)
413 +
414 + # ``thought\n`` prefix must be stripped from reasoning even though
415 + # the engine was pre-initialised to REASONING.
416 + assert reasoning.startswith("Reason"), (
417 + f"thought\\n prefix leaked into reasoning: {reasoning!r}"
418 + )
419 + assert "thought\n" not in reasoning, (
420 + f"thought\\n prefix leaked into reasoning: {reasoning!r}"
421 + )
422 + assert "Reasoning body" in reasoning, f"Reasoning body missing: {reasoning!r}"
423 +
424 + # The redundant ``<|channel>`` opener must not appear as text.
425 + assert "<|channel>" not in content, (
426 + f"<|channel> leaked into content: {content!r}"
427 + )
428 + assert "<|channel>" not in reasoning, (
429 + f"<|channel> leaked into reasoning: {reasoning!r}"
430 + )
431 +
432 + # Post-``<channel|>`` text must appear as content.
433 + assert "Final content" in content, (
434 + f"Post-<channel|> text missing from content: {content!r}"
435 + )
436 +
437 + def test_redundant_channel_open_swallowed_after_open_channel_prompt(
438 + self, pre_init_parser, pre_init_tokenizer, request_obj
439 + ):
440 + # Prompt already ends inside an open ``<|channel>`` block. Engine
441 + # is pre-initialised to ``REASONING`` via the start-token check.
442 + # Even if the model redundantly re-emits ``<|channel>thought\n``,
443 + # the no-op transition + prefix stripping must keep output clean.
444 + results = _stream_tokens_batched(
445 + pre_init_parser,
446 + pre_init_tokenizer,
447 + request_obj,
448 + batch_size=1,
449 + prompt_token_ids=[CHANNEL_START_ID, 3000, 3001],
450 + )
451 +
452 + reasoning, content, _ = _collect_fields(results)
453 +
454 + assert "<|channel>" not in content, (
455 + f"<|channel> leaked into content: {content!r}"
456 + )
457 + assert "thought\n" not in reasoning, (
458 + f"thought\\n prefix leaked into reasoning: {reasoning!r}"
459 + )
460 + assert "Reasoning body" in reasoning
461 + assert "Final content" in content
462 +
463 +
464 # ── Second model output: two tool calls with holdback ────────────────
465
466 REASONING_TEXT_2 = (
vllm/parser/abstract_parser.py
+7
@@ -758,6 +758,13 @@ class DelegatingParser(Parser):
758 prompt_token_ids
759 ):
760 state.reasoning_ended = True
761 + else:
762 + # Reasoning is still open at the end of the prompt; let the
763 + # reasoning parser adjust its initial parsing state so the
764 + # first generated tokens are classified correctly.
765 + self._reasoning_parser.adjust_initial_state_from_prompt(
766 + prompt_token_ids
767 + )
768
769 current_text, current_token_ids = state.advance(delta_text, delta_token_ids)
770 delta_message: DeltaMessage | None = None
vllm/parser/engine/adapters.py
+3
@@ -60,6 +60,9 @@ class ParserEngineReasoningAdapter(ReasoningParser):
60 def is_reasoning_end(self, input_ids: Sequence[int]) -> bool:
61 return self._parser_engine.is_reasoning_end(list(input_ids))
62
63 + def adjust_initial_state_from_prompt(self, prompt_token_ids: Sequence[int]) -> None:
64 + self._parser_engine.adjust_initial_state_from_prompt(prompt_token_ids)
65 +
66 def extract_content_ids(self, input_ids: list[int]) -> list[int]:
67 return self._parser_engine.extract_content_ids(input_ids)
68
vllm/parser/engine/parser_engine.py
+12
@@ -101,6 +101,7 @@ class ParserEngine(Parser):
101
102 self._reasoning_ended: bool = False
103 self._streaming_initialized: bool = False
104 + self._prompt_streaming_prepared: bool = False
105
106 self._tool_slots: list[ToolCallSlot] = []
107 self._deferred_content: str = ""
@@ -165,6 +166,10 @@ class ParserEngine(Parser):
166 self._streaming_initialized = True
167 self._reset(initial_state=initial_state)
168
169 + def adjust_initial_state_from_prompt(self, prompt_token_ids: Sequence[int]) -> None:
170 + """See :meth:`ReasoningParser.adjust_initial_state_from_prompt`."""
171 + return
172 +
173 def finish_streaming(self) -> DeltaMessage | None:
174 events = self._engine.finish()
175 return self._events_to_delta(events) if events else None
@@ -176,6 +181,7 @@ class ParserEngine(Parser):
181 self._deferred_content = ""
182 self._deferred_reasoning = ""
183 self._content_has_nonws = False
184 + self._prompt_streaming_prepared = False
185
186 def adjust_request(
187 self, request: ChatCompletionRequest | ResponsesRequest
@@ -364,6 +370,12 @@ class ParserEngine(Parser):
370 *,
371 finished: bool,
372 ) -> DeltaMessage | None:
373 + if not self._prompt_streaming_prepared and prompt_token_ids is not None:
374 + # NOTE: call the hook BEFORE setting the flag, because the hook
375 + # may invoke ``_reset`` (e.g. via ``initialize_streaming``) which
376 + # clears ``_prompt_streaming_prepared``.
377 + self.adjust_initial_state_from_prompt(prompt_token_ids)
378 + self._prompt_streaming_prepared = True
379 self._check_skip_tool_parsing(request)
380 events = self._feed(delta_text, delta_token_ids)
381 if finished:
vllm/parser/gemma4.py
+28
@@ -353,6 +353,14 @@ def gemma4_config() -> ParserEngineConfig:
353 ParserState.REASONING,
354 (EventType.REASONING_START,),
355 ),
356 + # No-op: if we pre-initialised the engine to REASONING from the
357 + # prompt (see ``adjust_initial_state_from_prompt``) but the model
358 + # still emits its own ``<|channel>`` opener, swallow it instead
359 + # of leaking it as TEXT_CHUNK.
360 + (ParserState.REASONING, "THINK_START"): Transition(
361 + ParserState.REASONING,
362 + (),
363 + ),
364 (ParserState.REASONING, "THINK_END"): Transition(
365 ParserState.CONTENT,
366 (EventType.REASONING_END,),
@@ -524,6 +532,26 @@ class Gemma4Parser(ParserEngine):
532 return True
533 return True
534
535 + def adjust_initial_state_from_prompt(self, prompt_token_ids: Sequence[int]) -> None:
536 + """Pre-initialise the engine to ``REASONING`` when the prompt does
537 + not already end with reasoning concluded.
538 +
539 + This covers the post-tool-response continuation case where the chat
540 + template leaves the prompt ending inside an open ``<|channel>``
541 + block (issue #45834). It is also safe in the common new-turn case
542 + where the model itself emits ``<|channel>`` first: the no-op
543 + ``(REASONING, THINK_START)`` transition swallows it, and the
544 + ``thought\n`` prefix in the first reasoning chunk is stripped by
545 + ``_events_to_delta`` as it already is in the default flow.
546 + """
547 + if self.is_reasoning_end(list(prompt_token_ids)):
548 + return
549 + self._engine.reset(initial_state=ParserState.REASONING)
550 + # Prevent a later default ``initialize_streaming()`` (e.g. from
551 + # ``ParserEngineReasoningAdapter.extract_reasoning_streaming``) from
552 + # clobbering this with ``CONTENT``.
553 + self._streaming_initialized = True
554 +
555 def _events_to_delta(
556 self,
557 events: list[SemanticEvent],
vllm/reasoning/abs_reasoning_parsers.py
+12
@@ -187,6 +187,18 @@ class ReasoningParser:
187 """Adjust request parameters; override in subclasses as needed."""
188 return request
189
190 + def adjust_initial_state_from_prompt(self, prompt_token_ids: Sequence[int]) -> None:
191 + """Hook called once at the start of streaming with the prompt tokens.
192 +
193 + Gives parsers a chance to adjust their initial parsing state based on
194 + the prompt — for example, when the chat template leaves the prompt
195 + inside an open reasoning channel and the engine's default initial
196 + state would otherwise misclassify the first generated tokens.
197 +
198 + Default is a no-op; override in subclasses as needed.
199 + """
200 + return
201 +
202 def prepare_structured_tag(
203 self,
204 original_tag: str | None,