[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>

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)

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: