[Bugfix] Fix Step-3.5 reasoning parser for structured outputs (#53174)
Signed-off-by: Yifan Zong <yzong@redhat.com>
yzong-rh committed
Sep 9, 2026 at 17:10 UTC
26fec6d183c3957ee4dc0022452c430e745dd20f
2 files changed
+26
-91
tests/reasoning/test_step3p5_reasoning_parser.py
+20
-23
@@ -25,12 +25,11 @@ SIMPLE_REASONING = {
25
"content": "This is the rest",
26
"is_reasoning_end": True,
27
}
28
-# need to get into parser again to remove newline after </think>
28
COMPLETE_REASONING = {
29
"output": "This is a reasoning section</think>",
30
"reasoning": "This is a reasoning section",
31
"content": None,
33
- "is_reasoning_end": False,
32
+ "is_reasoning_end": True,
33
}
34
NO_CONTENT = {
35
"output": "This is content",
@@ -72,7 +71,7 @@ COMPLETE_REASONING_WITH_THINK = {
71
"output": "<think>This is a reasoning section</think>",
72
"reasoning": "This is a reasoning section",
73
"content": None,
75
- "is_reasoning_end": False,
74
+ "is_reasoning_end": True,
75
}
76
MULTIPLE_LINES_WITH_THINK = {
77
"output": "<think>This\nThat</think>This is the rest\nThat",
@@ -113,21 +112,21 @@ EMPTY_STREAMING = {
112
NEW_LINE = {
113
"output": "\n<think>This is a reasoning section</think>\nThis is the rest",
114
"reasoning": "This is a reasoning section",
116
- "content": "This is the rest",
115
+ "content": "\nThis is the rest",
116
"is_reasoning_end": True,
117
}
118
119
NEW_LINE_STREAMING = {
120
"output": "\n<think>This is a reasoning section\n</think>\nThis is the rest",
121
"reasoning": "\nThis is a reasoning section",
123
- "content": "This is the rest",
122
+ "content": "\nThis is the rest",
123
"is_reasoning_end": True,
124
}
125
126
NEW_LINE_STREAMING_COMPLEX_CONTENT = {
127
"output": "\n This is a \n reasoning section\n\n\n</think>\n\nThis is the rest",
128
"reasoning": "\n This is a \n reasoning section\n\n",
130
- "content": "\nThis is the rest",
129
+ "content": "\n\nThis is the rest",
130
"is_reasoning_end": True,
131
}
132
@@ -301,9 +300,20 @@ def test_reasoning(
300
301
# Test is_reasoning_end
302
output_ids = step3p5_tokenizer.convert_tokens_to_ids(output)
304
- if streaming:
305
- is_reasoning_end = parser.is_reasoning_end(output_ids)
306
- assert is_reasoning_end == param_dict["is_reasoning_end"]
303
+ is_reasoning_end = parser.is_reasoning_end(output_ids)
304
+ assert is_reasoning_end == param_dict["is_reasoning_end"]
305
+
306
+ # Test is_reasoning_end_streaming
307
+ current_ids: list[int] = []
308
+ for token_id in output_ids:
309
+ current_ids.append(token_id)
310
+ if token_id == parser.end_token_id:
311
+ assert parser.is_reasoning_end_streaming(current_ids, [token_id])
312
+ assert parser.is_reasoning_end_streaming(output_ids, output_ids)
313
+ break
314
+ assert not parser.is_reasoning_end_streaming(current_ids, [token_id])
315
+ else:
316
+ assert not parser.is_reasoning_end_streaming(output_ids, output_ids)
317
318
# Test extract_content
319
if param_dict["content"] is not None:
@@ -324,18 +334,5 @@ def test_reasoning(
334
)
335
assert content == expected_content_ids
336
else:
327
- content = parser.extract_content_ids(output)
337
+ content = parser.extract_content_ids(output_ids)
338
assert content == []
329
-
330
-
331
-def test_step3p5_streaming_drops_leading_newline(step3p5_tokenizer):
332
- parser_cls = ReasoningParserManager.get_reasoning_parser("step3p5")
333
- parser = parser_cls(step3p5_tokenizer)
334
- output = "<think>calc</think>\nAnswer"
335
- tokens = step3p5_tokenizer.tokenize(output)
336
- output_tokens = [
337
- step3p5_tokenizer.convert_tokens_to_string([token]) for token in tokens
338
- ]
339
-
340
- _, content = run_reasoning_extraction(parser, output_tokens, streaming=True)
341
- assert content == "Answer"
vllm/reasoning/step3p5_reasoning_parser.py
+6
-68
@@ -1,7 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from collections.abc import Iterable, Sequence
4
+from collections.abc import Sequence
5
from typing import TYPE_CHECKING
6
7
from vllm.entrypoints.generate.base.protocol import DeltaMessage
@@ -18,9 +18,11 @@ class Step3p5ReasoningParser(BaseThinkingReasoningParser):
18
Reasoning parser for Step3p5 model.
19
20
Step3p5 uses the <think>...</think> format, but it tends to emit an extra
21
- newline immediately before and/or after the </think> token. This parser trims:
22
- - the newline right before </think>
23
- - the newline right after </think>
21
+ newline immediately before and/or after the </think> token. This parser
22
+ trims the newline before </think> in extract_reasoning*.
23
+
24
+ Reasoning end detection doesn't wait for an extra token after </think> so
25
+ structured output can constrain the first post-think token.
26
"""
27
28
@property
@@ -38,60 +40,6 @@ class Step3p5ReasoningParser(BaseThinkingReasoningParser):
40
# whether it is immediately before </think>.
41
self._pending_reasoning_newline = False
42
41
- # Tracks whether we've seen </think> but are still waiting for one more
42
- # token to confirm the end.
43
- self._end_token_pending = False
44
-
45
- def is_reasoning_end(self, input_ids: Sequence[int]) -> bool:
46
- return self._is_reasoning_end_from_ids(input_ids)
47
-
48
- def is_reasoning_end_streaming(
49
- self, input_ids: Sequence[int], delta_ids: Iterable[int]
50
- ) -> bool:
51
- # Only examine newly generated tokens; they may contain multiple ids.
52
- return self._is_reasoning_end_from_ids(tuple(delta_ids))
53
-
54
- def _is_reasoning_end_from_ids(self, input_ids: Sequence[int]) -> bool:
55
- # Scan backwards to find the last special token, <think> or </think>.
56
- last_special = None
57
- last_idx = -1
58
- for i in range(len(input_ids) - 1, -1, -1):
59
- token_id = input_ids[i]
60
- if token_id == self.start_token_id:
61
- last_special = "start"
62
- last_idx = i
63
- break
64
- if token_id == self.end_token_id:
65
- last_special = "end"
66
- last_idx = i
67
- break
68
-
69
- if last_special == "start":
70
- # If we're already waiting for one token after </think>, do not
71
- # clear the pending state just because the prompt contains <think>.
72
- # Streaming deltas should not include <think> for this model.
73
- if self._end_token_pending:
74
- return False
75
- # A start token after any end token means reasoning is ongoing.
76
- self._end_token_pending = False
77
- return False
78
-
79
- if last_special == "end":
80
- # Require at least one token after </think> before ending.
81
- if last_idx < len(input_ids) - 1:
82
- self._end_token_pending = False
83
- return True
84
- self._end_token_pending = True
85
- return False
86
-
87
- # No special tokens in this input. If we were waiting for one token
88
- # after </think>, any new token completes the end.
89
- if self._end_token_pending and input_ids:
90
- self._end_token_pending = False
91
- return True
92
-
93
- return False
94
-
43
def extract_reasoning(
44
self,
45
model_output: str,
@@ -100,8 +48,6 @@ class Step3p5ReasoningParser(BaseThinkingReasoningParser):
48
reasoning, content = super().extract_reasoning(model_output, request)
49
if reasoning is not None:
50
reasoning = reasoning.removesuffix("\n")
103
- if content is not None:
104
- content = content.removeprefix("\n")
51
return reasoning or None, content or None
52
53
def extract_reasoning_streaming(
@@ -113,14 +59,6 @@ class Step3p5ReasoningParser(BaseThinkingReasoningParser):
59
current_token_ids: Sequence[int],
60
delta_token_ids: Sequence[int],
61
) -> DeltaMessage | None:
116
- # Drop the immediate newline that models often emit after </think>.
117
- if previous_text.endswith(self.end_token) and delta_text:
118
- if delta_text == "\n":
119
- return None
120
- elif delta_text.startswith("\n"):
121
- remaining = delta_text.removeprefix("\n")
122
- return DeltaMessage(content=remaining) if remaining else None
123
-
62
ret = super().extract_reasoning_streaming(
63
previous_text,
64
current_text,