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 7, 19:00 UTC
[Bugfix] Fix step3p5 reasoning with interleaved thinking (#34211)
Signed-off-by: mariohong <mariohong128@gmail.com> Co-authored-by: Chauncey <chaunceyjiang@gmail.com>
Mario Hong committed
Feb 25, 2026 at 23:13 UTC
af5e6afa0af28b75409104906fe47ac9e8c03cf1
2 files changed
+387
-14
tests/reasoning/test_step3p5_reasoning_parser.py
new
+341
@@ -0,0 +1,341 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+from transformers import AutoTokenizer
6
+
7
+from tests.reasoning.utils import run_reasoning_extraction
8
+from vllm.reasoning import ReasoningParser, ReasoningParserManager
9
+
10
+parser_name = "step3p5"
11
+start_token = "<think>"
12
+end_token = "</think>"
13
+
14
+REASONING_MODEL_NAME = "stepfun-ai/Step-3.5-Flash"
15
+
16
+
17
+@pytest.fixture(scope="module")
18
+def step3p5_tokenizer():
19
+ return AutoTokenizer.from_pretrained(REASONING_MODEL_NAME)
20
+
21
+
22
+SIMPLE_REASONING = {
23
+ "output": "This is a reasoning section</think>This is the rest",
24
+ "reasoning_content": "This is a reasoning section",
25
+ "content": "This is the rest",
26
+ "is_reasoning_end": True,
27
+}
28
+# need to get into parser again to remove newline after </think>
29
+COMPLETE_REASONING = {
30
+ "output": "This is a reasoning section</think>",
31
+ "reasoning_content": "This is a reasoning section",
32
+ "content": None,
33
+ "is_reasoning_end": False,
34
+}
35
+NO_CONTENT = {
36
+ "output": "This is content",
37
+ "reasoning_content": "This is content",
38
+ "content": None,
39
+ "is_reasoning_end": False,
40
+}
41
+NO_REASONING_STREAMING = {
42
+ "output": "This is a reasoning section",
43
+ "reasoning_content": "This is a reasoning section",
44
+ "content": None,
45
+ "is_reasoning_end": False,
46
+}
47
+MULTIPLE_LINES = {
48
+ "output": "This\nThat</think>This is the rest\nThat",
49
+ "reasoning_content": "This\nThat",
50
+ "content": "This is the rest\nThat",
51
+ "is_reasoning_end": True,
52
+}
53
+SHORTEST_REASONING_NO_STREAMING = {
54
+ "output": "</think>This is the rest",
55
+ "reasoning_content": None,
56
+ "content": "This is the rest",
57
+ "is_reasoning_end": True,
58
+}
59
+SHORTEST_REASONING = {
60
+ "output": "</think>This is the rest",
61
+ "reasoning_content": None,
62
+ "content": "This is the rest",
63
+ "is_reasoning_end": True,
64
+}
65
+REASONING_WITH_THINK = {
66
+ "output": "<think>This is a reasoning section</think>This is the rest",
67
+ "reasoning_content": "This is a reasoning section",
68
+ "content": "This is the rest",
69
+ "is_reasoning_end": True,
70
+}
71
+COMPLETE_REASONING_WITH_THINK = {
72
+ "output": "<think>This is a reasoning section</think>",
73
+ "reasoning_content": "This is a reasoning section",
74
+ "content": None,
75
+ "is_reasoning_end": False,
76
+}
77
+MULTIPLE_LINES_WITH_THINK = {
78
+ "output": "<think>This\nThat</think>This is the rest\nThat",
79
+ "reasoning_content": "This\nThat",
80
+ "content": "This is the rest\nThat",
81
+ "is_reasoning_end": True,
82
+}
83
+SHORTEST_REASONING_NO_STREAMING_WITH_THINK = {
84
+ "output": "</think>This is the rest",
85
+ "reasoning_content": None,
86
+ "content": "This is the rest",
87
+ "is_reasoning_end": True,
88
+}
89
+SHORTEST_REASONING_WITH_THINK = {
90
+ "output": "</think>This is the rest",
91
+ "reasoning_content": None,
92
+ "content": "This is the rest",
93
+ "is_reasoning_end": True,
94
+}
95
+THINK_NO_END = {
96
+ "output": "<think>This is a reasoning section",
97
+ "reasoning_content": "This is a reasoning section",
98
+ "content": None,
99
+ "is_reasoning_end": False,
100
+}
101
+EMPTY = {
102
+ "output": "",
103
+ "reasoning_content": None,
104
+ "content": None,
105
+ "is_reasoning_end": False,
106
+}
107
+EMPTY_STREAMING = {
108
+ "output": "",
109
+ "reasoning_content": None,
110
+ "content": None,
111
+ "is_reasoning_end": False,
112
+}
113
+NEW_LINE = {
114
+ "output": "\n<think>This is a reasoning section</think>\nThis is the rest",
115
+ "reasoning_content": "This is a reasoning section",
116
+ "content": "This is the rest",
117
+ "is_reasoning_end": True,
118
+}
119
+
120
+NEW_LINE_STREAMING = {
121
+ "output": "\n<think>This is a reasoning section\n</think>\nThis is the rest",
122
+ "reasoning_content": "\nThis is a reasoning section",
123
+ "content": "This is the rest",
124
+ "is_reasoning_end": True,
125
+}
126
+
127
+NEW_LINE_STREAMING_COMPLEX_CONTENT = {
128
+ "output": "\n This is a \n reasoning section\n\n\n</think>\n\nThis is the rest",
129
+ "reasoning_content": "\n This is a \n reasoning section\n\n",
130
+ "content": "\nThis is the rest",
131
+ "is_reasoning_end": True,
132
+}
133
+
134
+MULTI_TURN_PROMPT_CONTENT = {
135
+ "output": "<think> This is last turn's reasoning section </think> hello <think>",
136
+ "reasoning_content": "",
137
+ "content": "",
138
+ "is_reasoning_end": False,
139
+}
140
+
141
+TEST_CASES = [
142
+ pytest.param(
143
+ False,
144
+ SIMPLE_REASONING,
145
+ id="simple_reasoning",
146
+ ),
147
+ pytest.param(
148
+ True,
149
+ SIMPLE_REASONING,
150
+ id="simple_reasoning_streaming",
151
+ ),
152
+ pytest.param(
153
+ False,
154
+ COMPLETE_REASONING,
155
+ id="complete_reasoning",
156
+ ),
157
+ pytest.param(
158
+ True,
159
+ COMPLETE_REASONING,
160
+ id="complete_reasoning_streaming",
161
+ ),
162
+ pytest.param(
163
+ False,
164
+ NO_CONTENT,
165
+ id="no_content_token",
166
+ ),
167
+ pytest.param(
168
+ True,
169
+ NO_REASONING_STREAMING,
170
+ id="no_reasoning_token_streaming",
171
+ ),
172
+ pytest.param(
173
+ False,
174
+ MULTIPLE_LINES,
175
+ id="multiple_lines",
176
+ ),
177
+ pytest.param(
178
+ True,
179
+ MULTIPLE_LINES,
180
+ id="multiple_lines_streaming",
181
+ ),
182
+ pytest.param(
183
+ True,
184
+ SHORTEST_REASONING,
185
+ id="shortest",
186
+ ),
187
+ pytest.param(
188
+ False,
189
+ SHORTEST_REASONING_NO_STREAMING,
190
+ id="shortest_streaming",
191
+ ),
192
+ pytest.param(
193
+ False,
194
+ REASONING_WITH_THINK,
195
+ id="reasoning_with_think",
196
+ ),
197
+ pytest.param(
198
+ True,
199
+ REASONING_WITH_THINK,
200
+ id="reasoning_with_think_streaming",
201
+ ),
202
+ pytest.param(
203
+ False,
204
+ COMPLETE_REASONING_WITH_THINK,
205
+ id="complete_reasoning_with_think",
206
+ ),
207
+ pytest.param(
208
+ True,
209
+ COMPLETE_REASONING_WITH_THINK,
210
+ id="complete_reasoning_with_think_streaming",
211
+ ),
212
+ pytest.param(
213
+ False,
214
+ MULTIPLE_LINES_WITH_THINK,
215
+ id="multiple_lines_with_think",
216
+ ),
217
+ pytest.param(
218
+ True,
219
+ MULTIPLE_LINES_WITH_THINK,
220
+ id="multiple_lines_with_think_streaming",
221
+ ),
222
+ pytest.param(
223
+ False,
224
+ SHORTEST_REASONING_NO_STREAMING_WITH_THINK,
225
+ id="shortest_with_think",
226
+ ),
227
+ pytest.param(
228
+ True,
229
+ SHORTEST_REASONING_WITH_THINK,
230
+ id="shortest_with_think_streaming",
231
+ ),
232
+ pytest.param(
233
+ False,
234
+ THINK_NO_END,
235
+ id="think_no_end",
236
+ ),
237
+ pytest.param(
238
+ True,
239
+ THINK_NO_END,
240
+ id="think_no_end_streaming",
241
+ ),
242
+ pytest.param(
243
+ False,
244
+ EMPTY,
245
+ id="empty",
246
+ ),
247
+ pytest.param(
248
+ True,
249
+ EMPTY_STREAMING,
250
+ id="empty_streaming",
251
+ ),
252
+ pytest.param(
253
+ False,
254
+ NEW_LINE,
255
+ id="new_line",
256
+ ),
257
+ pytest.param(
258
+ True,
259
+ NEW_LINE_STREAMING,
260
+ id="new_line_streaming",
261
+ ),
262
+ pytest.param(
263
+ True,
264
+ NEW_LINE_STREAMING_COMPLEX_CONTENT,
265
+ id="new_line_streaming_complex_content",
266
+ ),
267
+ pytest.param(
268
+ True,
269
+ MULTI_TURN_PROMPT_CONTENT,
270
+ id="multi_turn_prompt_content",
271
+ ),
272
+]
273
+
274
+
275
+@pytest.mark.parametrize("streaming, param_dict", TEST_CASES)
276
+def test_reasoning(
277
+ streaming: bool,
278
+ param_dict: dict,
279
+ step3p5_tokenizer,
280
+ request,
281
+):
282
+ output = step3p5_tokenizer.tokenize(param_dict["output"])
283
+ # decode everything to tokens
284
+ output_tokens: list[str] = [
285
+ step3p5_tokenizer.convert_tokens_to_string([token]) for token in output
286
+ ]
287
+ parser: ReasoningParser = ReasoningParserManager.get_reasoning_parser(parser_name)(
288
+ step3p5_tokenizer
289
+ )
290
+
291
+ reasoning, content = run_reasoning_extraction(
292
+ parser, output_tokens, streaming=streaming
293
+ )
294
+
295
+ print(f"reasoning: {reasoning}")
296
+ print(f"content: {content}")
297
+ test_id = request.node.callspec.id if hasattr(request.node, "callspec") else None
298
+ if request.node.callspec.id != "multi_turn_prompt_content":
299
+ assert reasoning == param_dict["reasoning_content"]
300
+ assert content == param_dict["content"]
301
+
302
+ # Test is_reasoning_end
303
+ 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"]
307
+
308
+ # Test extract_content
309
+ if param_dict["content"] is not None:
310
+ content = parser.extract_content_ids(output_ids)
311
+ # Fixed expected token ids for specific test cases
312
+ test_id = (
313
+ request.node.callspec.id if hasattr(request.node, "callspec") else None
314
+ )
315
+ # Match most specific first
316
+ if test_id not in [
317
+ "new_line_streaming_complex_content",
318
+ "new_line_streaming",
319
+ "new_line",
320
+ "multi_turn_prompt_content",
321
+ ]:
322
+ expected_content_ids = step3p5_tokenizer.convert_tokens_to_ids(
323
+ step3p5_tokenizer.tokenize(param_dict["content"])
324
+ )
325
+ assert content == expected_content_ids
326
+ else:
327
+ content = parser.extract_content_ids(output)
328
+ 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
+46
-14
@@ -39,24 +39,59 @@ class Step3p5ReasoningParser(BaseThinkingReasoningParser):
39
# whether it is immediately before </think>.
40
self._pending_reasoning_newline = False
41
42
- # Used to delay the reasoning end detection.
43
- # This is necessary to remove the newline appears immediately after </think>,
44
- # which may cause the end detection to be delayed by one round.
45
- self.end_offset = 1
42
+ # Tracks whether we've seen </think> but are still waiting for one more
43
+ # token to confirm the end.
44
+ self._end_token_pending = False
45
46
def is_reasoning_end(self, input_ids: Sequence[int]) -> bool:
48
- if self.end_token_id in input_ids and self.end_offset > 0:
49
- self.end_offset -= 1
50
- return False
51
- return self.end_offset < 1
47
+ return self._is_reasoning_end_from_ids(input_ids)
48
49
def is_reasoning_end_streaming(
50
self, input_ids: Sequence[int], delta_ids: Iterable[int]
51
) -> bool:
56
- if self.end_token_id in input_ids and self.end_offset > 0:
57
- self.end_offset -= 1
52
+ # Only examine newly generated tokens; they may contain multiple ids.
53
+ return self._is_reasoning_end_from_ids(delta_ids)
54
+
55
+ def _is_reasoning_end_from_ids(self, input_ids: Sequence[int]) -> bool:
56
+ # Scan backwards to find the last special token, <think> or </think>.
57
+ last_special = None
58
+ last_idx = -1
59
+ for i in range(len(input_ids) - 1, -1, -1):
60
+ token_id = input_ids[i]
61
+ if token_id == self.start_token_id:
62
+ last_special = "start"
63
+ last_idx = i
64
+ break
65
+ if token_id == self.end_token_id:
66
+ last_special = "end"
67
+ last_idx = i
68
+ break
69
+
70
+ if last_special == "start":
71
+ # If we're already waiting for one token after </think>, do not
72
+ # clear the pending state just because the prompt contains <think>.
73
+ # Streaming deltas should not include <think> for this model.
74
+ if self._end_token_pending:
75
+ return False
76
+ # A start token after any end token means reasoning is ongoing.
77
+ self._end_token_pending = False
78
+ return False
79
+
80
+ if last_special == "end":
81
+ # Require at least one token after </think> before ending.
82
+ if last_idx < len(input_ids) - 1:
83
+ self._end_token_pending = False
84
+ return True
85
+ self._end_token_pending = True
86
return False
59
- return self.end_offset < 1
87
+
88
+ # No special tokens in this input. If we were waiting for one token
89
+ # after </think>, any new token completes the end.
90
+ if self._end_token_pending and input_ids:
91
+ self._end_token_pending = False
92
+ return True
93
+
94
+ return False
95
96
def extract_reasoning(
97
self,
@@ -136,9 +171,6 @@ class Step3p5ReasoningParser(BaseThinkingReasoningParser):
171
172
# Content: handle the newline immediately after </think>.
173
if content_to_output is not None:
139
- # No need to get into parser again to remove newline after </think>.
140
- self.end_offset -= 1
141
-
174
# If we have content, reasoning must have ended.
175
self._pending_reasoning_newline = False
176