[Bugfix][Reasoning] kimi_k3: O(delta) reasoning-end check on the decode path (#50886)
Signed-off-by: abmfy <abmfy@icloud.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Bowen Wang committed
Aug 4, 2026 at 10:37 UTC
adbf08d977fb3fa26c4f19826745a02abd6dd7ca
2 files changed
+157
-10
tests/reasoning/test_kimi_k3_reasoning_parser.py
+95
@@ -248,3 +248,98 @@ def test_adjust_request_keeps_xtml_markers_contiguous():
248
assert adjusted.skip_special_tokens is False
249
if hasattr(adjusted, "spaces_between_special_tokens"):
250
assert adjusted.spaces_between_special_tokens is False
251
+
252
+
253
+OPEN_IDS = [1, 2, 3]
254
+CLOSE_IDS = [4, 2, 3]
255
+
256
+
257
+def _reference_is_reasoning_end(input_ids: list[int]) -> bool:
258
+ """Full-sequence reference for the streaming check to be measured against.
259
+
260
+ Two independent last-occurrence scans, i.e. the straightforward reading of
261
+ "reasoning ended iff the newest think marker is a close marker".
262
+ """
263
+
264
+ def last(needle: list[int]) -> int:
265
+ for i in range(len(input_ids) - len(needle), -1, -1):
266
+ if input_ids[i : i + len(needle)] == needle:
267
+ return i
268
+ return -1
269
+
270
+ last_close, last_open = last(CLOSE_IDS), last(OPEN_IDS)
271
+ if last_open == -1:
272
+ return last_close != -1
273
+ return last_close > last_open
274
+
275
+
276
+def test_is_reasoning_end_streaming_only_scans_the_step_window():
277
+ """The decode-step check must not re-derive the answer from the whole
278
+ sequence: an already-closed think block earlier in the sequence is the
279
+ engine's business (it latches `reasoning_ended`), not this call's."""
280
+ parser = KimiK3ReasoningParser(DummyTokenizer())
281
+ prompt = [*CLOSE_IDS, *OPEN_IDS, 9, 9, 9]
282
+
283
+ assert not parser.is_reasoning_end_streaming([*prompt, 7], [7])
284
+ assert parser.is_reasoning_end_streaming([*prompt, *CLOSE_IDS], CLOSE_IDS)
285
+
286
+
287
+def test_is_reasoning_end_streaming_detects_marker_across_step_boundary():
288
+ """A 3-token marker can be split over decode steps; the check carries the
289
+ preceding len(marker)-1 tokens so the final token still completes it."""
290
+ parser = KimiK3ReasoningParser(DummyTokenizer())
291
+ history = [*OPEN_IDS, 5, *CLOSE_IDS[:2]]
292
+
293
+ assert not parser.is_reasoning_end_streaming(history, [CLOSE_IDS[1]])
294
+ assert parser.is_reasoning_end_streaming([*history, CLOSE_IDS[2]], [CLOSE_IDS[2]])
295
+
296
+
297
+def test_is_reasoning_end_streaming_reopened_block_is_not_ended():
298
+ """Kimi K3 may close and immediately reopen the think channel inside one
299
+ window (speculative decoding lands several tokens per step). The newest
300
+ marker wins, matching is_reasoning_end."""
301
+ parser = KimiK3ReasoningParser(DummyTokenizer())
302
+ window = [*CLOSE_IDS, *OPEN_IDS]
303
+
304
+ assert not parser.is_reasoning_end_streaming([*OPEN_IDS, 5, *window], window)
305
+
306
+
307
+def test_is_reasoning_end_streaming_accepts_an_iterator_delta():
308
+ """`should_advance` may hand over an islice rather than a list."""
309
+ parser = KimiK3ReasoningParser(DummyTokenizer())
310
+ full = [*OPEN_IDS, 5, *CLOSE_IDS]
311
+
312
+ assert parser.is_reasoning_end_streaming(full, iter(CLOSE_IDS))
313
+
314
+
315
+def test_is_reasoning_end_streaming_thinking_disabled():
316
+ parser = KimiK3ReasoningParser(
317
+ DummyTokenizer(), chat_template_kwargs={"thinking": False}
318
+ )
319
+
320
+ assert parser.is_reasoning_end_streaming([1], [1])
321
+
322
+
323
+@pytest.mark.parametrize("seed", range(24))
324
+def test_reasoning_end_matches_reference_over_marker_dense_sequences(seed):
325
+ """Both the full-sequence check and the per-step check must agree with the
326
+ two-scan reference, on sequences where the markers (which share the suffix
327
+ [2, 3]) collide and overlap constantly."""
328
+ import random
329
+
330
+ rnd = random.Random(seed)
331
+ parser = KimiK3ReasoningParser(DummyTokenizer())
332
+
333
+ for _ in range(200):
334
+ head = [rnd.choice([1, 2, 3, 4]) for _ in range(rnd.randrange(0, 14))]
335
+ delta = [rnd.choice([1, 2, 3, 4]) for _ in range(rnd.randrange(1, 5))]
336
+ full = [*head, *delta]
337
+
338
+ assert parser.is_reasoning_end(full) == _reference_is_reasoning_end(full)
339
+
340
+ # The engine only calls the streaming check while reasoning is still
341
+ # open, so that is the only case it has to agree on.
342
+ if not _reference_is_reasoning_end(head):
343
+ assert parser.is_reasoning_end_streaming(
344
+ full, delta
345
+ ) == _reference_is_reasoning_end(full), (head, delta)
vllm/reasoning/kimi_k3_reasoning_parser.py
+62
-10
@@ -22,7 +22,7 @@ When thinking is disabled (``chat_template_kwargs={"thinking": False}`` or
22
delta as normal content; there is simply no think channel to extract.
23
"""
24
25
-from collections.abc import Sequence
25
+from collections.abc import Iterable, Sequence
26
from typing import TYPE_CHECKING
27
28
import regex as re
@@ -36,17 +36,47 @@ if TYPE_CHECKING:
36
from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
37
38
39
+def _match_at(haystack: Sequence[int], i: int, needle: Sequence[int]) -> bool:
40
+ """Whether *needle* occurs in *haystack* starting at *i*.
41
+
42
+ Compares element by element rather than slicing: ``haystack`` is usually a
43
+ ``ConstantList``, whose slices cost a Python ``__getitem__`` plus a list
44
+ allocation on every probe.
45
+ """
46
+ return all(haystack[i + k] == needle[k] for k in range(len(needle)))
47
+
48
+
49
def _subseq_index(haystack: Sequence[int], needle: Sequence[int]) -> int:
50
"""Return start index of the last occurrence of needle in haystack, or -1."""
51
n = len(needle)
52
if n == 0:
53
return -1
54
+ first = needle[0]
55
for i in range(len(haystack) - n, -1, -1):
45
- if list(haystack[i : i + n]) == list(needle):
56
+ if haystack[i] == first and _match_at(haystack, i, needle):
57
return i
58
return -1
59
60
61
+def _newest_marker(haystack: Sequence[int], a: Sequence[int], b: Sequence[int]) -> int:
62
+ """Report which of *a* / *b* occurs last in *haystack*.
63
+
64
+ Returns 0 if *a* is the newest marker, 1 if *b* is, -1 if neither occurs.
65
+ Equivalent to comparing two ``_subseq_index`` results, but a single backward
66
+ pass that stops at the first hit instead of walking to index 0 twice.
67
+ """
68
+ if not a or not b:
69
+ return -1
70
+ a0, b0 = a[0], b[0]
71
+ for i in range(len(haystack) - min(len(a), len(b)), -1, -1):
72
+ head = haystack[i]
73
+ if head == a0 and i + len(a) <= len(haystack) and _match_at(haystack, i, a):
74
+ return 0
75
+ if head == b0 and i + len(b) <= len(haystack) and _match_at(haystack, i, b):
76
+ return 1
77
+ return -1
78
+
79
+
80
class KimiK3ReasoningParser(ReasoningParser):
81
"""Reasoning parser for the Kimi K3 (XTML) think channel."""
82
@@ -133,14 +163,36 @@ class KimiK3ReasoningParser(ReasoningParser):
163
# agent continuations, where the chat template keeps a prior turn's
164
# think channel (with its <|close|>think<|sep|>) in the prompt while the
165
# current turn is still reasoning (its <|open|>think<|sep|> is the newest
136
- # marker). See _subseq_index (returns the last occurrence).
137
- last_close = _subseq_index(input_ids, self._think_close_ids)
138
- last_open = _subseq_index(input_ids, self._think_open_ids)
139
- if last_open == -1:
140
- # No open marker in scope (e.g. the open was consumed as the
141
- # generation prefix): a close marker means reasoning has ended.
142
- return last_close != -1
143
- return last_close > last_open
166
+ # marker). A missing open marker (e.g. it was consumed as the generation
167
+ # prefix) means a close marker alone ends reasoning, which is what
168
+ # "close is the newest marker" already encodes.
169
+ return (
170
+ _newest_marker(input_ids, self._think_close_ids, self._think_open_ids) == 0
171
+ )
172
+
173
+ def is_reasoning_end_streaming(
174
+ self, input_ids: Sequence[int], delta_ids: Iterable[int]
175
+ ) -> bool:
176
+ """Reasoning-end check for a single decode step.
177
+
178
+ The engine calls this once per structured-output request per decode step
179
+ while the request is still inside the think channel, so it only has to
180
+ look at the tokens generated this step, plus ``len(marker) - 1`` tokens
181
+ of context in case a marker straddles the step boundary.
182
+
183
+ The inherited default re-runs the full-sequence ``is_reasoning_end``,
184
+ which makes the scheduler O(context) per request per step and starves the
185
+ GPU on long agentic contexts.
186
+ """
187
+ if not self._thinking_enabled:
188
+ return True
189
+ delta = list(delta_ids)
190
+ if not delta:
191
+ return False
192
+ carry = max(len(self._think_close_ids), len(self._think_open_ids)) - 1
193
+ head = len(input_ids) - len(delta)
194
+ window = list(input_ids[max(0, head - carry) : head]) + delta
195
+ return _newest_marker(window, self._think_close_ids, self._think_open_ids) == 0
196
197
def _extract_content_ids(self, input_ids: list[int]) -> list[int]:
198
if not self._thinking_enabled: