[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>
Elvir Crnčević committed
Sep 26, 2026 at 01:29 UTC
9bf44c40134c024956df9134036b1d5b3d4eef7b
2 files changed
+91
-4
tests/reasoning/test_kimi_k3_reasoning_parser.py
+51
-4
@@ -19,6 +19,9 @@ SEP = "<|sep|>"
19
THINK_OPEN = f"{OPEN}think{SEP}"
20
THINK_CLOSE = f"{CLOSE}think{SEP}"
21
RESPONSE_OPEN = f"{OPEN}response{SEP}"
22
+OPEN_IDS = [1, 2, 3]
23
+CLOSE_IDS = [4, 2, 3]
24
+RESPONSE_OPEN_IDS = [ord(ch) for ch in RESPONSE_OPEN]
25
26
27
class DummyTokenizer:
@@ -119,6 +122,54 @@ def test_is_reasoning_end_ignores_stale_close_from_prior_turn():
122
assert not parser.is_reasoning_end([*new_open])
123
124
125
+@pytest.mark.parametrize(
126
+ ("token_ids", "expected"),
127
+ [
128
+ pytest.param(
129
+ [*OPEN_IDS, 9, 10, *CLOSE_IDS, *RESPONSE_OPEN_IDS, 11],
130
+ 2,
131
+ id="open_and_close_markers",
132
+ ),
133
+ pytest.param(
134
+ [9, 10, *CLOSE_IDS, *RESPONSE_OPEN_IDS, 11],
135
+ 2,
136
+ id="open_marker_consumed_as_generation_prefix",
137
+ ),
138
+ pytest.param([*RESPONSE_OPEN_IDS, 11, 12], 0, id="response_only"),
139
+ pytest.param([*OPEN_IDS, 9, 10], 2, id="unterminated_after_open"),
140
+ pytest.param([9, 10, 11], 3, id="unterminated_without_markers"),
141
+ pytest.param([], 0, id="empty"),
142
+ pytest.param([9, 4, 2, 10], 4, id="partial_close_marker_is_reasoning"),
143
+ pytest.param(
144
+ [*OPEN_IDS, 9, *RESPONSE_OPEN_IDS, 11],
145
+ 2 + len(RESPONSE_OPEN_IDS),
146
+ id="response_open_inside_open_think_is_reasoning",
147
+ ),
148
+ ],
149
+)
150
+def test_count_reasoning_tokens_matches_think_channel(token_ids, expected):
151
+ """reasoning_tokens must cover exactly what extract_reasoning labels as
152
+ reasoning: marker tokens are excluded and a consumed generation prefix
153
+ means the output starts inside the think channel."""
154
+ parser = KimiK3ReasoningParser(DummyTokenizer())
155
+
156
+ assert parser.count_reasoning_tokens(token_ids) == expected
157
+
158
+
159
+def test_count_reasoning_tokens_is_zero_when_thinking_disabled():
160
+ parser = KimiK3ReasoningParser(
161
+ DummyTokenizer(), chat_template_kwargs={"thinking": False}
162
+ )
163
+
164
+ assert parser.count_reasoning_tokens([*OPEN_IDS, 9, 10, *CLOSE_IDS]) == 0
165
+
166
+
167
+def test_count_reasoning_tokens_through_delegating_parser():
168
+ parser = ReasoningOnlyParser(DummyTokenizer())
169
+
170
+ assert parser.count_reasoning_tokens([*OPEN_IDS, 9, *CLOSE_IDS, 11]) == 1
171
+
172
+
173
def test_streaming_split_open_marker_is_held_back():
174
parser = KimiK3ReasoningParser(DummyTokenizer())
175
@@ -250,10 +301,6 @@ def test_adjust_request_keeps_xtml_markers_contiguous():
301
assert adjusted.spaces_between_special_tokens is False
302
303
253
-OPEN_IDS = [1, 2, 3]
254
-CLOSE_IDS = [4, 2, 3]
255
-
256
-
304
def _reference_is_reasoning_end(input_ids: list[int]) -> bool:
305
"""Full-sequence reference for the streaming check to be measured against.
306
vllm/reasoning/kimi_k3_reasoning_parser.py
+40
@@ -134,6 +134,9 @@ class KimiK3ReasoningParser(ReasoningParser):
134
self._think_close_ids = tokenizer.encode(
135
self._think_close, add_special_tokens=False
136
)
137
+ self._response_open_ids = tokenizer.encode(
138
+ self._response_open, add_special_tokens=False
139
+ )
140
self._last_streaming_delta_token_ids: tuple[int, ...] | None = None
141
self._last_streaming_content_token_ids: list[int] | None = None
142
@@ -211,6 +214,43 @@ class KimiK3ReasoningParser(ReasoningParser):
214
return cached_content_ids
215
return self._extract_content_ids(input_ids)
216
217
+ def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int:
218
+ if not self._thinking_enabled:
219
+ return 0
220
+ think_open = self._think_open_ids
221
+ think_close = self._think_close_ids
222
+ response_open = self._response_open_ids
223
+ first_ids = {m[0] for m in (think_open, think_close, response_open) if m}
224
+ n = len(token_ids)
225
+ count = 0
226
+ in_reasoning = True
227
+ seen_open = False
228
+ i = 0
229
+ while i < n:
230
+ head = token_ids[i]
231
+ if head in first_ids:
232
+ if i + len(think_open) <= n and _match_at(token_ids, i, think_open):
233
+ in_reasoning = True
234
+ seen_open = True
235
+ i += len(think_open)
236
+ continue
237
+ if i + len(think_close) <= n and _match_at(token_ids, i, think_close):
238
+ in_reasoning = False
239
+ i += len(think_close)
240
+ continue
241
+ if (
242
+ not seen_open
243
+ and i + len(response_open) <= n
244
+ and _match_at(token_ids, i, response_open)
245
+ ):
246
+ in_reasoning = False
247
+ i += len(response_open)
248
+ continue
249
+ if in_reasoning:
250
+ count += 1
251
+ i += 1
252
+ return count
253
+
254
def _strip_content_wrapper(self, text: str) -> str:
255
"""Strip ``<|open|>response<|sep|>…<|close|>response<|sep|>`` wrapper and
256
``<|close|>message<|sep|>`` from *text*.