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 8, 02:30 UTC
[Frontend] Fix reasoning_tokens for text-based parsers in Responses API (#33513)
Signed-off-by: Jaeyeon Kim <anencore94@gmail.com>
Jaeyeon Kim(김재연) committed
Feb 19, 2026 at 08:16 UTC
9681068cf995af3cf651f1150fbabd8604537660
7 files changed
+208
-3
tests/entrypoints/openai/responses/test_simple.py
+47
@@ -134,6 +134,53 @@ async def test_streaming_output_consistency(client: OpenAI, model_name: str):
134
)
135
136
137
+@pytest.mark.asyncio
138
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
139
+async def test_streaming_reasoning_tokens_e2e(client: OpenAI, model_name: str):
140
+ """Verify final usage includes reasoning_tokens in streaming mode."""
141
+ response = await client.responses.create(
142
+ model=model_name,
143
+ input="Compute 17 * 19 and explain briefly.",
144
+ reasoning={"effort": "low"},
145
+ temperature=0.0,
146
+ stream=True,
147
+ )
148
+
149
+ completed_event = None
150
+ async for event in response:
151
+ if event.type == "response.completed":
152
+ completed_event = event
153
+
154
+ assert completed_event is not None
155
+ assert completed_event.response.status == "completed"
156
+ assert completed_event.response.usage is not None
157
+ assert completed_event.response.usage.output_tokens_details is not None
158
+ assert completed_event.response.usage.output_tokens_details.reasoning_tokens > 0, (
159
+ "Expected reasoning_tokens > 0 for streamed Qwen3 response."
160
+ )
161
+
162
+
163
+@pytest.mark.asyncio
164
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
165
+async def test_non_streaming_reasoning_tokens_e2e(client: OpenAI, model_name: str):
166
+ """Verify usage includes reasoning_tokens in non-streaming mode."""
167
+ response = await client.responses.create(
168
+ model=model_name,
169
+ input="Compute 23 * 17 and explain briefly.",
170
+ reasoning={"effort": "low"},
171
+ temperature=0.0,
172
+ stream=False,
173
+ )
174
+
175
+ assert response is not None
176
+ assert response.status == "completed"
177
+ assert response.usage is not None
178
+ assert response.usage.output_tokens_details is not None
179
+ assert response.usage.output_tokens_details.reasoning_tokens > 0, (
180
+ "Expected reasoning_tokens > 0 for non-streamed Qwen3 response."
181
+ )
182
+
183
+
184
@pytest.mark.asyncio
185
@pytest.mark.parametrize("model_name", [MODEL_NAME])
186
async def test_max_tokens(client: OpenAI, model_name: str):
tests/entrypoints/openai/test_serving_responses.py
+89
-2
@@ -13,9 +13,13 @@ from openai.types.responses.tool import (
13
Tool,
14
)
15
16
+import vllm.envs as envs
17
from vllm.entrypoints.mcp.tool_server import ToolServer
17
-from vllm.entrypoints.openai.engine.protocol import ErrorResponse
18
-from vllm.entrypoints.openai.responses.context import ConversationContext
18
+from vllm.entrypoints.openai.engine.protocol import (
19
+ ErrorResponse,
20
+ RequestResponseMetadata,
21
+)
22
+from vllm.entrypoints.openai.responses.context import ConversationContext, SimpleContext
23
from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
24
from vllm.entrypoints.openai.responses.serving import (
25
OpenAIServingResponses,
@@ -23,6 +27,8 @@ from vllm.entrypoints.openai.responses.serving import (
27
extract_tool_types,
28
)
29
from vllm.inputs.data import TokensPrompt
30
+from vllm.outputs import CompletionOutput, RequestOutput
31
+from vllm.sampling_params import SamplingParams
32
33
34
class MockConversationContext(ConversationContext):
@@ -259,6 +265,87 @@ class TestValidateGeneratorInput:
265
assert isinstance(result, ErrorResponse)
266
267
268
+@pytest.mark.asyncio
269
+async def test_reasoning_tokens_counted_for_text_reasoning_model(monkeypatch):
270
+ """Ensure reasoning_tokens usage is derived from thinking token spans."""
271
+
272
+ class FakeTokenizer:
273
+ def __init__(self):
274
+ self._vocab = {"<think>": 1, "</think>": 2, "reason": 3, "final": 4}
275
+
276
+ def get_vocab(self):
277
+ return self._vocab
278
+
279
+ # Force non-harmony, SimpleContext path
280
+ monkeypatch.setattr(envs, "VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT", False)
281
+
282
+ engine_client = MagicMock()
283
+ model_config = MagicMock()
284
+ model_config.hf_config.model_type = "test"
285
+ model_config.hf_text_config = MagicMock()
286
+ model_config.get_diff_sampling_param.return_value = {}
287
+ engine_client.model_config = model_config
288
+ engine_client.input_processor = MagicMock()
289
+ engine_client.io_processor = MagicMock()
290
+ engine_client.renderer = MagicMock()
291
+
292
+ tokenizer = FakeTokenizer()
293
+ engine_client.renderer.get_tokenizer.return_value = tokenizer
294
+
295
+ models = MagicMock()
296
+
297
+ serving = OpenAIServingResponses(
298
+ engine_client=engine_client,
299
+ models=models,
300
+ request_logger=None,
301
+ chat_template=None,
302
+ chat_template_content_format="auto",
303
+ reasoning_parser="qwen3",
304
+ )
305
+
306
+ # Build a SimpleContext with thinking tokens in the output.
307
+ context = SimpleContext()
308
+ token_ids = [1, 10, 2, 20] # <think> 10 </think> 20 -> reasoning token count = 1
309
+ completion = CompletionOutput(
310
+ index=0,
311
+ text="<think>reason</think>final",
312
+ token_ids=token_ids,
313
+ cumulative_logprob=0.0,
314
+ logprobs=None,
315
+ finish_reason="stop",
316
+ stop_reason=None,
317
+ )
318
+ req_output = RequestOutput(
319
+ request_id="req",
320
+ prompt="hi",
321
+ prompt_token_ids=[7, 8],
322
+ prompt_logprobs=None,
323
+ outputs=[completion],
324
+ finished=True,
325
+ num_cached_tokens=0,
326
+ )
327
+ context.append_output(req_output)
328
+
329
+ async def dummy_result_generator():
330
+ yield None
331
+
332
+ request = ResponsesRequest(input="hi", tools=[], stream=False)
333
+ sampling_params = SamplingParams(max_tokens=16)
334
+ metadata = RequestResponseMetadata(request_id="req")
335
+
336
+ response = await serving.responses_full_generator(
337
+ request=request,
338
+ sampling_params=sampling_params,
339
+ result_generator=dummy_result_generator(),
340
+ context=context,
341
+ model_name="test-model",
342
+ tokenizer=tokenizer,
343
+ request_metadata=metadata,
344
+ )
345
+
346
+ assert response.usage.output_tokens_details.reasoning_tokens == 1
347
+
348
+
349
class TestExtractAllowedToolsFromMcpRequests:
350
"""Test class for _extract_allowed_tools_from_mcp_requests function"""
351
tests/reasoning/test_base_thinking_reasoning_parser.py
+17
@@ -167,6 +167,23 @@ class TestBaseThinkingReasoningParserMethods:
167
is False
168
)
169
170
+ def test_count_reasoning_tokens(self, test_tokenizer):
171
+ """Count tokens between start/end markers."""
172
+ parser = TestThinkingReasoningParser(test_tokenizer)
173
+ start = parser.start_token_id
174
+ end = parser.end_token_id
175
+ token_ids = [0, start, 11, 12, end, 99]
176
+ assert parser.count_reasoning_tokens(token_ids) == 2
177
+
178
+ def test_count_reasoning_tokens_nested(self, test_tokenizer):
179
+ """Ensure nested thinking spans count all inner tokens safely."""
180
+ parser = TestThinkingReasoningParser(test_tokenizer)
181
+ s = parser.start_token_id
182
+ e = parser.end_token_id
183
+ token_ids = [s, 1, s, 2, e, 3, e]
184
+ # Tokens 1,2,3 are inside reasoning (depth>0) => 3 tokens
185
+ assert parser.count_reasoning_tokens(token_ids) == 3
186
+
187
def test_extract_content_ids(self, test_tokenizer):
188
"""Test the extract_content_ids method."""
189
parser = TestThinkingReasoningParser(test_tokenizer)
vllm/entrypoints/openai/responses/context.py
+3
-1
@@ -280,7 +280,6 @@ class ParsableContext(ConversationContext):
280
self.num_prompt_tokens = 0
281
self.num_output_tokens = 0
282
self.num_cached_tokens = 0
283
- # TODO: num_reasoning_tokens is not implemented yet.
283
self.num_reasoning_tokens = 0
284
# not implemented yet for ParsableContext
285
self.all_turn_metrics: list[TurnMetrics] = []
@@ -308,12 +307,15 @@ class ParsableContext(ConversationContext):
307
308
self.input_messages: list[ResponseRawMessageAndToken] = []
309
self.output_messages: list[ResponseRawMessageAndToken] = []
310
+ self._accumulated_token_ids: list[int] = []
311
312
def append_output(self, output: RequestOutput) -> None:
313
self.num_prompt_tokens = len(output.prompt_token_ids or [])
314
self.num_cached_tokens = output.num_cached_tokens or 0
315
self.num_output_tokens += len(output.outputs[0].token_ids or [])
316
self.parser.process(output.outputs[0])
317
+ output_token_ids = output.outputs[0].token_ids or []
318
+ self._accumulated_token_ids.extend(output_token_ids)
319
320
# only store if enable_response_messages is True, save memory
321
if self.request.enable_response_messages:
vllm/entrypoints/openai/responses/serving.py
+13
@@ -759,6 +759,19 @@ class OpenAIServingResponses(OpenAIServing):
759
num_generated_tokens = context.num_output_tokens
760
num_cached_tokens = context.num_cached_tokens
761
num_reasoning_tokens = context.num_reasoning_tokens
762
+ # For text-based reasoning parsers (e.g., <think>...</think>),
763
+ # HarmonyContext already counts reasoning tokens via channels.
764
+ # For Simple/Parsable contexts, derive reasoning_tokens from
765
+ # accumulated output token IDs using the parser if not already set.
766
+ if (
767
+ num_reasoning_tokens == 0
768
+ and self.parser is not None
769
+ and self.parser.reasoning_parser_cls is not None
770
+ and isinstance(context, (SimpleContext, ParsableContext))
771
+ ):
772
+ reasoning_parser = self.parser.reasoning_parser_cls(tokenizer)
773
+ accumulated = getattr(context, "_accumulated_token_ids", []) or []
774
+ num_reasoning_tokens = reasoning_parser.count_reasoning_tokens(accumulated)
775
776
usage = ResponseUsage(
777
input_tokens=num_prompt_tokens,
vllm/reasoning/abs_reasoning_parsers.py
+19
@@ -104,6 +104,25 @@ class ReasoningParser:
104
The extracted content from the input_ids.
105
"""
106
107
+ def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int:
108
+ """Count the number of reasoning tokens in a sequence.
109
+
110
+ Text-based reasoning models typically wrap their chain-of-thought
111
+ between special start/end tokens (e.g., ``<think> ... </think>``).
112
+ Implementations that support reasoning token counting should override
113
+ this method. The default implementation returns ``0`` so existing
114
+ parsers remain unchanged unless they explicitly opt in.
115
+
116
+ Args:
117
+ token_ids: Sequence of generated token ids (excluding prompt).
118
+
119
+ Returns:
120
+ int: Number of tokens that belong to reasoning content.
121
+ """
122
+
123
+ # By default, assume the parser cannot detect reasoning spans.
124
+ return 0
125
+
126
@abstractmethod
127
def extract_reasoning(
128
self,
vllm/reasoning/basic_parsers.py
+20
@@ -175,3 +175,23 @@ class BaseThinkingReasoningParser(ReasoningParser):
175
# If generation stops right after end-of-think, return null content
176
final_content = content or None
177
return reasoning, final_content
178
+
179
+ def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int:
180
+ """Count tokens that fall within start/end thinking markers.
181
+
182
+ Uses a depth counter so nested spans are handled safely and stray end
183
+ tokens do not drive the counter negative.
184
+ """
185
+ count = 0
186
+ depth = 0
187
+ for token_id in token_ids:
188
+ if token_id == self.start_token_id:
189
+ depth += 1
190
+ continue
191
+ if token_id == self.end_token_id:
192
+ if depth > 0:
193
+ depth -= 1
194
+ continue
195
+ if depth > 0:
196
+ count += 1
197
+ return count