[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>
Samyabrata Maji committed
Sep 25, 2026 at 05:42 UTC
064bc747ab42f8fc25cff5cec321f486f6864938
6 files changed
+136
-3
tests/parser/test_harmony.py
+61
@@ -873,6 +873,67 @@ class TestProcessChunk:
873
]
874
875
876
+class TestCountReasoningTokens:
877
+ def test_matches_process_chunk(self, harmony_parser, gpt_oss_tokenizer):
878
+ """Chat usage must report the same count as the Responses API path."""
879
+ token_ids = get_model_output_tokens(
880
+ [
881
+ assistant("Let me check the weather.", "analysis"),
882
+ tool_call("functions.get_weather", '{"location": "SF"}'),
883
+ ]
884
+ )
885
+ count = harmony_parser.count_reasoning_tokens(token_ids)
886
+
887
+ fresh = HarmonyParser(gpt_oss_tokenizer)
888
+ assert count == fresh.process_chunk(token_ids).reasoning_token_count
889
+ assert count > 0
890
+
891
+ def test_excludes_harmony_control_tokens(self, harmony_parser):
892
+ text = "Let me think about this."
893
+ token_ids = get_model_output_tokens(
894
+ [assistant(text, "analysis"), assistant("Hi", "final")]
895
+ )
896
+ expected = len(get_encoding().encode(text))
897
+ assert harmony_parser.count_reasoning_tokens(token_ids) == expected
898
+
899
+ def test_final_only_has_no_reasoning(self, harmony_parser):
900
+ token_ids = get_model_output_tokens([assistant("Hello", "final")])
901
+ assert harmony_parser.count_reasoning_tokens(token_ids) == 0
902
+
903
+ def test_streaming_reuses_running_total(
904
+ self, harmony_parser, gpt_oss_tokenizer, chat_request, monkeypatch
905
+ ):
906
+ """Called per streamed chunk with all tokens so far; must not replay."""
907
+ token_ids = get_model_output_tokens(
908
+ [assistant("Think", "analysis"), assistant("Answer", "final")]
909
+ )
910
+ expected = HarmonyParser(gpt_oss_tokenizer).count_reasoning_tokens(token_ids)
911
+ harmony_parser.parse_delta("", token_ids[:3], chat_request, finished=False)
912
+
913
+ def fail():
914
+ raise AssertionError("count_reasoning_tokens replayed the output")
915
+
916
+ monkeypatch.setattr(
917
+ "vllm.parser.harmony.get_streamable_parser_for_assistant", fail
918
+ )
919
+ for end in range(3, len(token_ids)):
920
+ harmony_parser.process_chunk(token_ids[end : end + 1])
921
+ harmony_parser.count_reasoning_tokens(token_ids[: end + 1])
922
+ assert harmony_parser.count_reasoning_tokens(token_ids) == expected
923
+
924
+ def test_does_not_disturb_stream_state(self, harmony_parser, chat_request):
925
+ """Streaming counts after each chunk, so it must not consume parser state."""
926
+ token_ids = get_model_output_tokens(
927
+ [assistant("Think", "analysis"), assistant("Answer", "final")]
928
+ )
929
+ harmony_parser.count_reasoning_tokens(token_ids)
930
+
931
+ reasoning, content, _ = harmony_parser.parse(
932
+ "", chat_request, model_output_token_ids=token_ids
933
+ )
934
+ assert (reasoning, content) == ("Think", "Answer")
935
+
936
+
937
class TestAdjustRequest:
938
REQUEST_TEXT = "Hello"
939
TOOL_TYPE = "function"
tests/reasoning/test_deepseekv3_reasoning_parser.py
+11
@@ -87,3 +87,14 @@ def test_identity_reasoning_parser_basic(tokenizer):
87
delta_token_ids=[],
88
)
89
assert result_none is None
90
+
91
+
92
+@pytest.mark.parametrize("thinking,expected", [(True, 3), (False, 0)])
93
+def test_count_reasoning_tokens_delegates(tokenizer, thinking, expected):
94
+ parser = DeepSeekV3ReasoningParser(
95
+ tokenizer, chat_template_kwargs={"thinking": thinking}
96
+ )
97
+ think = tokenizer.convert_tokens_to_ids(["<think>", "</think>"])
98
+ token_ids = [think[0], 11, 12, 13, think[1], 14]
99
+
100
+ assert parser.count_reasoning_tokens(token_ids) == expected
tests/reasoning/test_step3_reasoning_parser.py
new
+27
@@ -0,0 +1,27 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+
6
+from vllm.reasoning import ReasoningParserManager
7
+from vllm.tokenizers import get_tokenizer
8
+
9
+pytestmark = pytest.mark.skip_global_cleanup
10
+
11
+REASONING_MODEL_NAME = "stepfun-ai/step3"
12
+
13
+
14
+@pytest.fixture(scope="module")
15
+def step3_parser():
16
+ tokenizer = get_tokenizer(tokenizer_name=REASONING_MODEL_NAME)
17
+ return ReasoningParserManager.get_reasoning_parser("step3")(tokenizer)
18
+
19
+
20
+def test_count_reasoning_tokens_stops_at_think_end(step3_parser):
21
+ end = step3_parser.think_end_token_id
22
+ assert step3_parser.count_reasoning_tokens([11, 12, end, 13]) == 2
23
+
24
+
25
+def test_count_reasoning_tokens_without_think_end(step3_parser):
26
+ """Step3 treats output as reasoning until </think> appears."""
27
+ assert step3_parser.count_reasoning_tokens([11, 12, 13]) == 3
vllm/parser/harmony.py
+29
-3
@@ -38,6 +38,7 @@ from vllm.entrypoints.openai.chat_completion.protocol import (
38
)
39
from vllm.entrypoints.openai.parser.harmony_utils import (
40
extract_function_from_recipient,
41
+ get_encoding,
42
get_streamable_parser_for_assistant,
43
is_function_recipient,
44
)
@@ -114,6 +115,9 @@ class HarmonyParser(DelegatingParser):
115
self._next_tool_call_index = 0
116
self._num_processed_messages = 0
117
118
+ self._num_counted_tokens = 0
119
+ self._num_reasoning_tokens = 0
120
+
121
# For error recovery
122
self._current_message_tokens: list[int] = []
123
@@ -348,9 +352,7 @@ class HarmonyParser(DelegatingParser):
352
else:
353
self._current_message_tokens.append(token_id)
354
351
- if channel == "analysis" or (
352
- channel == "commentary" and recipient is not None
353
- ):
355
+ if self._is_reasoning_token(token_id, channel, recipient):
356
reasoning_token_count += 1
357
358
segments.append(
@@ -364,11 +366,35 @@ class HarmonyParser(DelegatingParser):
366
367
# TODO: Optionally merge and suppress empty Segments
368
369
+ self._num_counted_tokens += len(token_ids)
370
+ self._num_reasoning_tokens += reasoning_token_count
371
return ChunkResult(
372
segments=segments,
373
reasoning_token_count=reasoning_token_count,
374
)
375
376
+ def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int:
377
+ if len(token_ids) == self._num_counted_tokens:
378
+ return self._num_reasoning_tokens
379
+
380
+ parser = get_streamable_parser_for_assistant()
381
+ count = 0
382
+ for token_id in token_ids:
383
+ parser.process(token_id)
384
+ recipient = self._normalize_recipient(parser.current_recipient)
385
+ if self._is_reasoning_token(token_id, parser.current_channel, recipient):
386
+ count += 1
387
+ return count
388
+
389
+ @staticmethod
390
+ def _is_reasoning_token(
391
+ token_id: int, channel: str | None, recipient: str | None
392
+ ) -> bool:
393
+ is_reasoning_channel = channel == "analysis" or (
394
+ channel == "commentary" and recipient is not None
395
+ )
396
+ return is_reasoning_channel and not get_encoding().is_special_token(token_id)
397
+
398
def adjust_request(
399
self, request: ChatCompletionRequest | ResponsesRequest
400
) -> ChatCompletionRequest | ResponsesRequest:
vllm/reasoning/deepseek_v3_reasoning_parser.py
+3
@@ -55,6 +55,9 @@ class DeepSeekV3ReasoningParser(ReasoningParser):
55
def extract_content_ids(self, input_ids: list[int]) -> list[int]:
56
return self._parser.extract_content_ids(input_ids)
57
58
+ def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int:
59
+ return self._parser.count_reasoning_tokens(token_ids)
60
+
61
def extract_reasoning(
62
self, model_output: str, request: "ChatCompletionRequest | ResponsesRequest"
63
) -> tuple[str | None, str | None]:
vllm/reasoning/step3_reasoning_parser.py
+5
@@ -121,3 +121,8 @@ class Step3ReasoningParser(ReasoningParser):
121
return []
122
else:
123
return input_ids[input_ids.index(self.think_end_token_id) + 1 :]
124
+
125
+ def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int:
126
+ if self.think_end_token_id not in token_ids:
127
+ return len(token_ids)
128
+ return list(token_ids).index(self.think_end_token_id)