From adbf08d977fb3fa26c4f19826745a02abd6dd7ca Mon Sep 17 00:00:00 2001 From: Bowen Wang Date: Tue, 4 Aug 2026 10:37:03 +0800 Subject: [PATCH] [Bugfix][Reasoning] kimi_k3: O(delta) reasoning-end check on the decode path (#50886) Signed-off-by: abmfy Co-authored-by: Claude Opus 5 (1M context) --- .../test_kimi_k3_reasoning_parser.py | 95 +++++++++++++++++++ vllm/reasoning/kimi_k3_reasoning_parser.py | 72 ++++++++++++-- 2 files changed, 157 insertions(+), 10 deletions(-) diff --git a/tests/reasoning/test_kimi_k3_reasoning_parser.py b/tests/reasoning/test_kimi_k3_reasoning_parser.py index d050c8e354..4f831e851e 100644 --- a/tests/reasoning/test_kimi_k3_reasoning_parser.py +++ b/tests/reasoning/test_kimi_k3_reasoning_parser.py @@ -248,3 +248,98 @@ def test_adjust_request_keeps_xtml_markers_contiguous(): assert adjusted.skip_special_tokens is False if hasattr(adjusted, "spaces_between_special_tokens"): 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. + + Two independent last-occurrence scans, i.e. the straightforward reading of + "reasoning ended iff the newest think marker is a close marker". + """ + + def last(needle: list[int]) -> int: + for i in range(len(input_ids) - len(needle), -1, -1): + if input_ids[i : i + len(needle)] == needle: + return i + return -1 + + last_close, last_open = last(CLOSE_IDS), last(OPEN_IDS) + if last_open == -1: + return last_close != -1 + return last_close > last_open + + +def test_is_reasoning_end_streaming_only_scans_the_step_window(): + """The decode-step check must not re-derive the answer from the whole + sequence: an already-closed think block earlier in the sequence is the + engine's business (it latches `reasoning_ended`), not this call's.""" + parser = KimiK3ReasoningParser(DummyTokenizer()) + prompt = [*CLOSE_IDS, *OPEN_IDS, 9, 9, 9] + + assert not parser.is_reasoning_end_streaming([*prompt, 7], [7]) + assert parser.is_reasoning_end_streaming([*prompt, *CLOSE_IDS], CLOSE_IDS) + + +def test_is_reasoning_end_streaming_detects_marker_across_step_boundary(): + """A 3-token marker can be split over decode steps; the check carries the + preceding len(marker)-1 tokens so the final token still completes it.""" + parser = KimiK3ReasoningParser(DummyTokenizer()) + history = [*OPEN_IDS, 5, *CLOSE_IDS[:2]] + + assert not parser.is_reasoning_end_streaming(history, [CLOSE_IDS[1]]) + assert parser.is_reasoning_end_streaming([*history, CLOSE_IDS[2]], [CLOSE_IDS[2]]) + + +def test_is_reasoning_end_streaming_reopened_block_is_not_ended(): + """Kimi K3 may close and immediately reopen the think channel inside one + window (speculative decoding lands several tokens per step). The newest + marker wins, matching is_reasoning_end.""" + parser = KimiK3ReasoningParser(DummyTokenizer()) + window = [*CLOSE_IDS, *OPEN_IDS] + + assert not parser.is_reasoning_end_streaming([*OPEN_IDS, 5, *window], window) + + +def test_is_reasoning_end_streaming_accepts_an_iterator_delta(): + """`should_advance` may hand over an islice rather than a list.""" + parser = KimiK3ReasoningParser(DummyTokenizer()) + full = [*OPEN_IDS, 5, *CLOSE_IDS] + + assert parser.is_reasoning_end_streaming(full, iter(CLOSE_IDS)) + + +def test_is_reasoning_end_streaming_thinking_disabled(): + parser = KimiK3ReasoningParser( + DummyTokenizer(), chat_template_kwargs={"thinking": False} + ) + + assert parser.is_reasoning_end_streaming([1], [1]) + + +@pytest.mark.parametrize("seed", range(24)) +def test_reasoning_end_matches_reference_over_marker_dense_sequences(seed): + """Both the full-sequence check and the per-step check must agree with the + two-scan reference, on sequences where the markers (which share the suffix + [2, 3]) collide and overlap constantly.""" + import random + + rnd = random.Random(seed) + parser = KimiK3ReasoningParser(DummyTokenizer()) + + for _ in range(200): + head = [rnd.choice([1, 2, 3, 4]) for _ in range(rnd.randrange(0, 14))] + delta = [rnd.choice([1, 2, 3, 4]) for _ in range(rnd.randrange(1, 5))] + full = [*head, *delta] + + assert parser.is_reasoning_end(full) == _reference_is_reasoning_end(full) + + # The engine only calls the streaming check while reasoning is still + # open, so that is the only case it has to agree on. + if not _reference_is_reasoning_end(head): + assert parser.is_reasoning_end_streaming( + full, delta + ) == _reference_is_reasoning_end(full), (head, delta) diff --git a/vllm/reasoning/kimi_k3_reasoning_parser.py b/vllm/reasoning/kimi_k3_reasoning_parser.py index 43e4338f32..f05d5c0448 100644 --- a/vllm/reasoning/kimi_k3_reasoning_parser.py +++ b/vllm/reasoning/kimi_k3_reasoning_parser.py @@ -22,7 +22,7 @@ When thinking is disabled (``chat_template_kwargs={"thinking": False}`` or delta as normal content; there is simply no think channel to extract. """ -from collections.abc import Sequence +from collections.abc import Iterable, Sequence from typing import TYPE_CHECKING import regex as re @@ -36,17 +36,47 @@ if TYPE_CHECKING: from vllm.entrypoints.openai.responses.protocol import ResponsesRequest +def _match_at(haystack: Sequence[int], i: int, needle: Sequence[int]) -> bool: + """Whether *needle* occurs in *haystack* starting at *i*. + + Compares element by element rather than slicing: ``haystack`` is usually a + ``ConstantList``, whose slices cost a Python ``__getitem__`` plus a list + allocation on every probe. + """ + return all(haystack[i + k] == needle[k] for k in range(len(needle))) + + def _subseq_index(haystack: Sequence[int], needle: Sequence[int]) -> int: """Return start index of the last occurrence of needle in haystack, or -1.""" n = len(needle) if n == 0: return -1 + first = needle[0] for i in range(len(haystack) - n, -1, -1): - if list(haystack[i : i + n]) == list(needle): + if haystack[i] == first and _match_at(haystack, i, needle): return i return -1 +def _newest_marker(haystack: Sequence[int], a: Sequence[int], b: Sequence[int]) -> int: + """Report which of *a* / *b* occurs last in *haystack*. + + Returns 0 if *a* is the newest marker, 1 if *b* is, -1 if neither occurs. + Equivalent to comparing two ``_subseq_index`` results, but a single backward + pass that stops at the first hit instead of walking to index 0 twice. + """ + if not a or not b: + return -1 + a0, b0 = a[0], b[0] + for i in range(len(haystack) - min(len(a), len(b)), -1, -1): + head = haystack[i] + if head == a0 and i + len(a) <= len(haystack) and _match_at(haystack, i, a): + return 0 + if head == b0 and i + len(b) <= len(haystack) and _match_at(haystack, i, b): + return 1 + return -1 + + class KimiK3ReasoningParser(ReasoningParser): """Reasoning parser for the Kimi K3 (XTML) think channel.""" @@ -133,14 +163,36 @@ class KimiK3ReasoningParser(ReasoningParser): # agent continuations, where the chat template keeps a prior turn's # think channel (with its <|close|>think<|sep|>) in the prompt while the # current turn is still reasoning (its <|open|>think<|sep|> is the newest - # marker). See _subseq_index (returns the last occurrence). - last_close = _subseq_index(input_ids, self._think_close_ids) - last_open = _subseq_index(input_ids, self._think_open_ids) - if last_open == -1: - # No open marker in scope (e.g. the open was consumed as the - # generation prefix): a close marker means reasoning has ended. - return last_close != -1 - return last_close > last_open + # marker). A missing open marker (e.g. it was consumed as the generation + # prefix) means a close marker alone ends reasoning, which is what + # "close is the newest marker" already encodes. + return ( + _newest_marker(input_ids, self._think_close_ids, self._think_open_ids) == 0 + ) + + def is_reasoning_end_streaming( + self, input_ids: Sequence[int], delta_ids: Iterable[int] + ) -> bool: + """Reasoning-end check for a single decode step. + + The engine calls this once per structured-output request per decode step + while the request is still inside the think channel, so it only has to + look at the tokens generated this step, plus ``len(marker) - 1`` tokens + of context in case a marker straddles the step boundary. + + The inherited default re-runs the full-sequence ``is_reasoning_end``, + which makes the scheduler O(context) per request per step and starves the + GPU on long agentic contexts. + """ + if not self._thinking_enabled: + return True + delta = list(delta_ids) + if not delta: + return False + carry = max(len(self._think_close_ids), len(self._think_open_ids)) - 1 + head = len(input_ids) - len(delta) + window = list(input_ids[max(0, head - carry) : head]) + delta + return _newest_marker(window, self._think_close_ids, self._think_open_ids) == 0 def _extract_content_ids(self, input_ids: list[int]) -> list[int]: if not self._thinking_enabled: