From 064bc747ab42f8fc25cff5cec321f486f6864938 Mon Sep 17 00:00:00 2001 From: Samyabrata Maji Date: Fri, 25 Sep 2026 05:42:29 +0530 Subject: [PATCH] [Bugfix][Frontend] Count reasoning tokens for Harmony, DeepSeek-V3 and Step3 parsers (#58626) Signed-off-by: Samyabrata Maji <116789799+sammaji@users.noreply.github.com> Co-authored-by: Flora Feng <4florafeng@gmail.com> --- tests/parser/test_harmony.py | 61 +++++++++++++++++++ .../test_deepseekv3_reasoning_parser.py | 11 ++++ .../reasoning/test_step3_reasoning_parser.py | 27 ++++++++ vllm/parser/harmony.py | 32 +++++++++- .../reasoning/deepseek_v3_reasoning_parser.py | 3 + vllm/reasoning/step3_reasoning_parser.py | 5 ++ 6 files changed, 136 insertions(+), 3 deletions(-) create mode 100644 tests/reasoning/test_step3_reasoning_parser.py diff --git a/tests/parser/test_harmony.py b/tests/parser/test_harmony.py index 86f157c768..3d6fff04bf 100644 --- a/tests/parser/test_harmony.py +++ b/tests/parser/test_harmony.py @@ -873,6 +873,67 @@ class TestProcessChunk: ] +class TestCountReasoningTokens: + def test_matches_process_chunk(self, harmony_parser, gpt_oss_tokenizer): + """Chat usage must report the same count as the Responses API path.""" + token_ids = get_model_output_tokens( + [ + assistant("Let me check the weather.", "analysis"), + tool_call("functions.get_weather", '{"location": "SF"}'), + ] + ) + count = harmony_parser.count_reasoning_tokens(token_ids) + + fresh = HarmonyParser(gpt_oss_tokenizer) + assert count == fresh.process_chunk(token_ids).reasoning_token_count + assert count > 0 + + def test_excludes_harmony_control_tokens(self, harmony_parser): + text = "Let me think about this." + token_ids = get_model_output_tokens( + [assistant(text, "analysis"), assistant("Hi", "final")] + ) + expected = len(get_encoding().encode(text)) + assert harmony_parser.count_reasoning_tokens(token_ids) == expected + + def test_final_only_has_no_reasoning(self, harmony_parser): + token_ids = get_model_output_tokens([assistant("Hello", "final")]) + assert harmony_parser.count_reasoning_tokens(token_ids) == 0 + + def test_streaming_reuses_running_total( + self, harmony_parser, gpt_oss_tokenizer, chat_request, monkeypatch + ): + """Called per streamed chunk with all tokens so far; must not replay.""" + token_ids = get_model_output_tokens( + [assistant("Think", "analysis"), assistant("Answer", "final")] + ) + expected = HarmonyParser(gpt_oss_tokenizer).count_reasoning_tokens(token_ids) + harmony_parser.parse_delta("", token_ids[:3], chat_request, finished=False) + + def fail(): + raise AssertionError("count_reasoning_tokens replayed the output") + + monkeypatch.setattr( + "vllm.parser.harmony.get_streamable_parser_for_assistant", fail + ) + for end in range(3, len(token_ids)): + harmony_parser.process_chunk(token_ids[end : end + 1]) + harmony_parser.count_reasoning_tokens(token_ids[: end + 1]) + assert harmony_parser.count_reasoning_tokens(token_ids) == expected + + def test_does_not_disturb_stream_state(self, harmony_parser, chat_request): + """Streaming counts after each chunk, so it must not consume parser state.""" + token_ids = get_model_output_tokens( + [assistant("Think", "analysis"), assistant("Answer", "final")] + ) + harmony_parser.count_reasoning_tokens(token_ids) + + reasoning, content, _ = harmony_parser.parse( + "", chat_request, model_output_token_ids=token_ids + ) + assert (reasoning, content) == ("Think", "Answer") + + class TestAdjustRequest: REQUEST_TEXT = "Hello" TOOL_TYPE = "function" diff --git a/tests/reasoning/test_deepseekv3_reasoning_parser.py b/tests/reasoning/test_deepseekv3_reasoning_parser.py index 03867fb87b..d6bbea2f91 100644 --- a/tests/reasoning/test_deepseekv3_reasoning_parser.py +++ b/tests/reasoning/test_deepseekv3_reasoning_parser.py @@ -87,3 +87,14 @@ def test_identity_reasoning_parser_basic(tokenizer): delta_token_ids=[], ) assert result_none is None + + +@pytest.mark.parametrize("thinking,expected", [(True, 3), (False, 0)]) +def test_count_reasoning_tokens_delegates(tokenizer, thinking, expected): + parser = DeepSeekV3ReasoningParser( + tokenizer, chat_template_kwargs={"thinking": thinking} + ) + think = tokenizer.convert_tokens_to_ids(["", ""]) + token_ids = [think[0], 11, 12, 13, think[1], 14] + + assert parser.count_reasoning_tokens(token_ids) == expected diff --git a/tests/reasoning/test_step3_reasoning_parser.py b/tests/reasoning/test_step3_reasoning_parser.py new file mode 100644 index 0000000000..999d1d2645 --- /dev/null +++ b/tests/reasoning/test_step3_reasoning_parser.py @@ -0,0 +1,27 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest + +from vllm.reasoning import ReasoningParserManager +from vllm.tokenizers import get_tokenizer + +pytestmark = pytest.mark.skip_global_cleanup + +REASONING_MODEL_NAME = "stepfun-ai/step3" + + +@pytest.fixture(scope="module") +def step3_parser(): + tokenizer = get_tokenizer(tokenizer_name=REASONING_MODEL_NAME) + return ReasoningParserManager.get_reasoning_parser("step3")(tokenizer) + + +def test_count_reasoning_tokens_stops_at_think_end(step3_parser): + end = step3_parser.think_end_token_id + assert step3_parser.count_reasoning_tokens([11, 12, end, 13]) == 2 + + +def test_count_reasoning_tokens_without_think_end(step3_parser): + """Step3 treats output as reasoning until appears.""" + assert step3_parser.count_reasoning_tokens([11, 12, 13]) == 3 diff --git a/vllm/parser/harmony.py b/vllm/parser/harmony.py index 8c944733e3..5748bd6653 100644 --- a/vllm/parser/harmony.py +++ b/vllm/parser/harmony.py @@ -38,6 +38,7 @@ from vllm.entrypoints.openai.chat_completion.protocol import ( ) from vllm.entrypoints.openai.parser.harmony_utils import ( extract_function_from_recipient, + get_encoding, get_streamable_parser_for_assistant, is_function_recipient, ) @@ -114,6 +115,9 @@ class HarmonyParser(DelegatingParser): self._next_tool_call_index = 0 self._num_processed_messages = 0 + self._num_counted_tokens = 0 + self._num_reasoning_tokens = 0 + # For error recovery self._current_message_tokens: list[int] = [] @@ -348,9 +352,7 @@ class HarmonyParser(DelegatingParser): else: self._current_message_tokens.append(token_id) - if channel == "analysis" or ( - channel == "commentary" and recipient is not None - ): + if self._is_reasoning_token(token_id, channel, recipient): reasoning_token_count += 1 segments.append( @@ -364,11 +366,35 @@ class HarmonyParser(DelegatingParser): # TODO: Optionally merge and suppress empty Segments + self._num_counted_tokens += len(token_ids) + self._num_reasoning_tokens += reasoning_token_count return ChunkResult( segments=segments, reasoning_token_count=reasoning_token_count, ) + def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int: + if len(token_ids) == self._num_counted_tokens: + return self._num_reasoning_tokens + + parser = get_streamable_parser_for_assistant() + count = 0 + for token_id in token_ids: + parser.process(token_id) + recipient = self._normalize_recipient(parser.current_recipient) + if self._is_reasoning_token(token_id, parser.current_channel, recipient): + count += 1 + return count + + @staticmethod + def _is_reasoning_token( + token_id: int, channel: str | None, recipient: str | None + ) -> bool: + is_reasoning_channel = channel == "analysis" or ( + channel == "commentary" and recipient is not None + ) + return is_reasoning_channel and not get_encoding().is_special_token(token_id) + def adjust_request( self, request: ChatCompletionRequest | ResponsesRequest ) -> ChatCompletionRequest | ResponsesRequest: diff --git a/vllm/reasoning/deepseek_v3_reasoning_parser.py b/vllm/reasoning/deepseek_v3_reasoning_parser.py index 072048af91..a7613112f4 100644 --- a/vllm/reasoning/deepseek_v3_reasoning_parser.py +++ b/vllm/reasoning/deepseek_v3_reasoning_parser.py @@ -55,6 +55,9 @@ class DeepSeekV3ReasoningParser(ReasoningParser): def extract_content_ids(self, input_ids: list[int]) -> list[int]: return self._parser.extract_content_ids(input_ids) + def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int: + return self._parser.count_reasoning_tokens(token_ids) + def extract_reasoning( self, model_output: str, request: "ChatCompletionRequest | ResponsesRequest" ) -> tuple[str | None, str | None]: diff --git a/vllm/reasoning/step3_reasoning_parser.py b/vllm/reasoning/step3_reasoning_parser.py index 16ab1cf9f4..e0ada498c2 100644 --- a/vllm/reasoning/step3_reasoning_parser.py +++ b/vllm/reasoning/step3_reasoning_parser.py @@ -121,3 +121,8 @@ class Step3ReasoningParser(ReasoningParser): return [] else: return input_ids[input_ids.index(self.think_end_token_id) + 1 :] + + def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int: + if self.think_end_token_id not in token_ids: + return len(token_ids) + return list(token_ids).index(self.think_end_token_id)