[Bugfix][Reasoning] Count Kimi K3 reasoning tokens (#58372)

Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>
Signed-off-by: Flora Feng <4florafeng@gmail.com>
Co-authored-by: Flora Feng <4florafeng@gmail.com>

tests/reasoning/test_kimi_k3_reasoning_parser.py

@@ -19,6 +19,9 @@ SEP = "<|sep|>"
 THINK_OPEN = f"{OPEN}think{SEP}"
 THINK_CLOSE = f"{CLOSE}think{SEP}"
 RESPONSE_OPEN = f"{OPEN}response{SEP}"
+OPEN_IDS = [1, 2, 3]
+CLOSE_IDS = [4, 2, 3]
+RESPONSE_OPEN_IDS = [ord(ch) for ch in RESPONSE_OPEN]
 
 
 class DummyTokenizer:
@@ -119,6 +122,54 @@ def test_is_reasoning_end_ignores_stale_close_from_prior_turn():
     assert not parser.is_reasoning_end([*new_open])
 
 
+@pytest.mark.parametrize(
+    ("token_ids", "expected"),
+    [
+        pytest.param(
+            [*OPEN_IDS, 9, 10, *CLOSE_IDS, *RESPONSE_OPEN_IDS, 11],
+            2,
+            id="open_and_close_markers",
+        ),
+        pytest.param(
+            [9, 10, *CLOSE_IDS, *RESPONSE_OPEN_IDS, 11],
+            2,
+            id="open_marker_consumed_as_generation_prefix",
+        ),
+        pytest.param([*RESPONSE_OPEN_IDS, 11, 12], 0, id="response_only"),
+        pytest.param([*OPEN_IDS, 9, 10], 2, id="unterminated_after_open"),
+        pytest.param([9, 10, 11], 3, id="unterminated_without_markers"),
+        pytest.param([], 0, id="empty"),
+        pytest.param([9, 4, 2, 10], 4, id="partial_close_marker_is_reasoning"),
+        pytest.param(
+            [*OPEN_IDS, 9, *RESPONSE_OPEN_IDS, 11],
+            2 + len(RESPONSE_OPEN_IDS),
+            id="response_open_inside_open_think_is_reasoning",
+        ),
+    ],
+)
+def test_count_reasoning_tokens_matches_think_channel(token_ids, expected):
+    """reasoning_tokens must cover exactly what extract_reasoning labels as
+    reasoning: marker tokens are excluded and a consumed generation prefix
+    means the output starts inside the think channel."""
+    parser = KimiK3ReasoningParser(DummyTokenizer())
+
+    assert parser.count_reasoning_tokens(token_ids) == expected
+
+
+def test_count_reasoning_tokens_is_zero_when_thinking_disabled():
+    parser = KimiK3ReasoningParser(
+        DummyTokenizer(), chat_template_kwargs={"thinking": False}
+    )
+
+    assert parser.count_reasoning_tokens([*OPEN_IDS, 9, 10, *CLOSE_IDS]) == 0
+
+
+def test_count_reasoning_tokens_through_delegating_parser():
+    parser = ReasoningOnlyParser(DummyTokenizer())
+
+    assert parser.count_reasoning_tokens([*OPEN_IDS, 9, *CLOSE_IDS, 11]) == 1
+
+
 def test_streaming_split_open_marker_is_held_back():
     parser = KimiK3ReasoningParser(DummyTokenizer())
 
@@ -250,10 +301,6 @@ def test_adjust_request_keeps_xtml_markers_contiguous():
         assert adjusted.spaces_between_special_tokens is False
 
 
-OPEN_IDS = [1, 2, 3]
-CLOSE_IDS = [4, 2, 3]
-
-
 def _reference_is_reasoning_end(input_ids: list[int]) -> bool:
     """Full-sequence reference for the streaming check to be measured against.
 

vllm/reasoning/kimi_k3_reasoning_parser.py

@@ -134,6 +134,9 @@ class KimiK3ReasoningParser(ReasoningParser):
         self._think_close_ids = tokenizer.encode(
             self._think_close, add_special_tokens=False
         )
+        self._response_open_ids = tokenizer.encode(
+            self._response_open, add_special_tokens=False
+        )
         self._last_streaming_delta_token_ids: tuple[int, ...] | None = None
         self._last_streaming_content_token_ids: list[int] | None = None
 
@@ -211,6 +214,43 @@ class KimiK3ReasoningParser(ReasoningParser):
             return cached_content_ids
         return self._extract_content_ids(input_ids)
 
+    def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int:
+        if not self._thinking_enabled:
+            return 0
+        think_open = self._think_open_ids
+        think_close = self._think_close_ids
+        response_open = self._response_open_ids
+        first_ids = {m[0] for m in (think_open, think_close, response_open) if m}
+        n = len(token_ids)
+        count = 0
+        in_reasoning = True
+        seen_open = False
+        i = 0
+        while i < n:
+            head = token_ids[i]
+            if head in first_ids:
+                if i + len(think_open) <= n and _match_at(token_ids, i, think_open):
+                    in_reasoning = True
+                    seen_open = True
+                    i += len(think_open)
+                    continue
+                if i + len(think_close) <= n and _match_at(token_ids, i, think_close):
+                    in_reasoning = False
+                    i += len(think_close)
+                    continue
+                if (
+                    not seen_open
+                    and i + len(response_open) <= n
+                    and _match_at(token_ids, i, response_open)
+                ):
+                    in_reasoning = False
+                    i += len(response_open)
+                    continue
+            if in_reasoning:
+                count += 1
+            i += 1
+        return count
+
     def _strip_content_wrapper(self, text: str) -> str:
         """Strip ``<|open|>response<|sep|>…<|close|>response<|sep|>`` wrapper and
         ``<|close|>message<|sep|>`` from *text*.