From 29f46ac475322cfc14ced890825b9a149ada3608 Mon Sep 17 00:00:00 2001 From: Flora Feng <4florafeng@gmail.com> Date: Tue, 8 Sep 2026 17:07:16 -0400 Subject: [PATCH] [Perf] Eliminate full-history reasoning scans for structured outputs (#55223) Signed-off-by: sfeng33 <4florafeng@gmail.com> Signed-off-by: Flora Feng <4florafeng@gmail.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> --- tests/parser/engine/test_deepseek_v4.py | 11 ++ tests/parser/engine/test_parser_engine.py | 63 ++++++++ .../spec_decode/test_mtp_structured_output.py | 117 ++++++++++++++ .../test_reasoning_structured_output.py | 152 ++++++++++++++++++ vllm/parser/abstract_parser.py | 15 ++ vllm/parser/engine/adapters.py | 7 + vllm/parser/engine/parser_engine.py | 39 +++++ vllm/v1/structured_output/__init__.py | 113 ++++++++----- 8 files changed, 475 insertions(+), 42 deletions(-) diff --git a/tests/parser/engine/test_deepseek_v4.py b/tests/parser/engine/test_deepseek_v4.py index 873a1e6bd2..0936dccfea 100644 --- a/tests/parser/engine/test_deepseek_v4.py +++ b/tests/parser/engine/test_deepseek_v4.py @@ -416,6 +416,17 @@ class TestImplicitReasoningEnd: args = json.loads(tool_calls[0].arguments) assert args == {"location": "NYC"} + def test_reasoning_end_token_ids_drop_multi_token_tool_opener( + self, thinking_parser + ): + # Like the real DeepSeek-V4 tokenizer, the mock vocab has no single + # token for the DSML tool opener, so only is tracked on the + # token-ID fast path; the implicit tool-call exit still goes through + # the text lexer. + assert thinking_parser.reasoning_end_token_ids == {_THINK_END_ID} + assert thinking_parser.find_reasoning_end_offset([7, _THINK_END_ID, 8]) == 1 + assert thinking_parser.find_reasoning_end_offset([7, 8]) is None + def test_streaming_reasoning_implicit_end(self, thinking_parser): chunks = [ "Let me look up the weather.\n\n", diff --git a/tests/parser/engine/test_parser_engine.py b/tests/parser/engine/test_parser_engine.py index e759853cbc..f27ce52cb3 100644 --- a/tests/parser/engine/test_parser_engine.py +++ b/tests/parser/engine/test_parser_engine.py @@ -7,6 +7,7 @@ DeltaMessage / ExtractedToolCallInformation protocol. from __future__ import annotations +import dataclasses import json from types import SimpleNamespace from unittest.mock import MagicMock @@ -132,6 +133,68 @@ def _make_engine( ) +# ── TestReasoningEndTokenIds ───────────────────────────────────────── + + +def _with_reasoning_exits( + *exits: tuple[str, ParserState, tuple[EventType, ...]], +) -> ParserEngineConfig: + """Combined config plus extra transitions out of REASONING.""" + base = _combined_config() + transitions = dict(base.transitions) + for terminal, next_state, events in exits: + transitions[(ParserState.REASONING, terminal)] = Transition(next_state, events) + return dataclasses.replace(base, transitions=transitions) + + +class TestReasoningEndTokenIds: + """ParserEngine derives the reasoning-end token set from its config.""" + + def test_every_reasoning_end_exit_contributes(self): + cfg = _with_reasoning_exits( + ( + "TOOL_START", + ParserState.TOOL_ARGS, + (EventType.REASONING_END, EventType.TOOL_CALL_START), + ), + ) + assert _make_engine(cfg).reasoning_end_token_ids == {201, 202} + + def test_transition_staying_in_reasoning_is_ignored(self): + cfg = _with_reasoning_exits(("THINK_START", ParserState.REASONING, ())) + assert _make_engine(cfg).reasoning_end_token_ids == {201} + + def test_unreported_exit_fails_closed(self): + cfg = _with_reasoning_exits( + ("TOOL_START", ParserState.TOOL_ARGS, (EventType.TOOL_CALL_START,)), + ) + assert _make_engine(cfg).reasoning_end_token_ids == frozenset() + + def test_unresolved_think_end_fails_closed(self): + vocab = {k: v for k, v in _VOCAB.items() if k != ""} + assert _make_engine(vocab=vocab).reasoning_end_token_ids == frozenset() + + text_only = dataclasses.replace( + _combined_config(), token_id_terminals={"THINK_START": ""} + ) + assert _make_engine(text_only).reasoning_end_token_ids == frozenset() + + def test_config_without_reasoning_has_empty_set(self): + assert _make_engine(_hermes_config()).reasoning_end_token_ids == frozenset() + + def test_find_reasoning_end_offset_returns_first_match(self): + engine = _make_engine() + assert engine.find_reasoning_end_offset([5, 201, 6, 201]) == 1 + assert engine.find_reasoning_end_offset([5, 6]) is None + assert engine.find_reasoning_end_offset([]) is None + # Rejected-draft placeholders never match. + assert engine.find_reasoning_end_offset([-1, 201]) == 1 + + def test_find_reasoning_end_offset_with_empty_set_returns_none(self): + engine = _make_engine(_hermes_config()) + assert engine.find_reasoning_end_offset([201]) is None + + # ── TestEventsToDelta ──────────────────────────────────────────────── diff --git a/tests/v1/spec_decode/test_mtp_structured_output.py b/tests/v1/spec_decode/test_mtp_structured_output.py index 1fce9037d4..330e10438f 100644 --- a/tests/v1/spec_decode/test_mtp_structured_output.py +++ b/tests/v1/spec_decode/test_mtp_structured_output.py @@ -8,6 +8,7 @@ from transformers import AutoTokenizer from vllm.config import StructuredOutputsConfig, VllmConfig from vllm.config.model import ModelConfig from vllm.config.speculative import SpeculativeConfig +from vllm.parser.engine.adapters import ParserEngineReasoningAdapter from vllm.sampling_params import SamplingParams, StructuredOutputsParams from vllm.v1.request import Request from vllm.v1.structured_output import StructuredOutputManager @@ -390,3 +391,119 @@ def test_trim_reasoning_for_advance(): next_step = [post, post] request.append_output_token_ids(next_step) assert manager.trim_reasoning_for_advance(request, next_step) == next_step + + +class _EngineReasonerStub(ParserEngineReasoningAdapter): + """Adapter-typed reasoner with a fixed end-token set and no real engine.""" + + def __init__(self, end_token_ids): + self._end_token_ids = frozenset(end_token_ids) + self.windows: list[list[int]] = [] + + @property + def reasoning_end_token_ids(self): + return self._end_token_ids + + def find_reasoning_end_offset(self, token_ids): + self.windows.append(list(token_ids)) + for offset, token in enumerate(token_ids): + if token in self._end_token_ids: + return offset + return None + + def is_reasoning_end(self, input_ids): + return any(token in self._end_token_ids for token in input_ids) + + def is_reasoning_end_streaming(self, input_ids, delta_ids): + raise AssertionError("engine path must not rescan draft prefixes") + + +@pytest.mark.parametrize("backend", ["xgrammar", "guidance"]) +def test_bitmask_engine_reasoner_ends_midwindow_with_padding(backend): + """Engine reasoners see the draft window once, without -1 padding.""" + tokenizer, manager, request, prompt = _make_manager_and_request(backend) + grammar = request.structured_output_request.grammar + + assert grammar.accept_tokens(request.request_id, prompt) + + marker = tokenizer.encode("\n")[0] + reasoner = _EngineReasonerStub({marker}) + manager.reasoner_cls = _EngineReasonerStub + request.structured_output_request.reasoner = reasoner + request.structured_output_request.reasoning_ended = False + + pre = tokenizer.encode(" ")[0] + post = tokenizer.encode(",")[0] + drafts = [pre, marker, post, -1] + + bitmask = manager.grammar_bitmask( + requests={request.request_id: request}, + structured_output_request_ids=[request.request_id], + scheduled_spec_decode_tokens={request.request_id: drafts}, + ) + + assert bitmask is not None + assert bitmask.shape[0] == len(drafts) + 1 + assert (bitmask[0] == -1).all() + assert (bitmask[1] == -1).all() + assert not (bitmask[2] == -1).all() + assert reasoner.windows == [[pre, marker, post]] + assert not grammar.is_terminated() + + +@pytest.mark.parametrize("backend", ["xgrammar", "guidance"]) +def test_bitmask_legacy_order_sensitive_predicate_probed_per_token(backend): + """Legacy predicates are probed one draft at a time, never on the whole + window. A predicate that reports which marker is *newest* (KimiK3) would + otherwise answer False for a window that closes and then reopens + reasoning, leaving the post-marker drafts and the bonus row + unconstrained. + """ + tokenizer, manager, request, prompt = _make_manager_and_request(backend) + grammar = request.structured_output_request.grammar + + assert grammar.accept_tokens(request.request_id, prompt) + + close = tokenizer.encode("\n")[0] + reopen = tokenizer.encode("#")[0] + assert close != reopen + + class NewestMarkerReasoner: + def __init__(self, *_, **__): + pass + + def is_reasoning_end(self, input_ids): + return self.is_reasoning_end_streaming(input_ids, input_ids) + + def is_reasoning_end_streaming(self, input_ids, delta_ids): + for token in reversed(list(delta_ids)): + if token == close: + return True + if token == reopen: + return False + return False + + manager.reasoner_cls = NewestMarkerReasoner + request.structured_output_request.reasoner = NewestMarkerReasoner() + request.structured_output_request.reasoning_ended = False + + pre = tokenizer.encode(" ")[0] + post = tokenizer.encode(",")[0] + drafts = [pre, close, post, reopen] + # Sanity: the whole window says "reopened", so only per-token probing + # can find the close marker. + assert not NewestMarkerReasoner().is_reasoning_end_streaming([], drafts) + + bitmask = manager.grammar_bitmask( + requests={request.request_id: request}, + structured_output_request_ids=[request.request_id], + scheduled_spec_decode_tokens={request.request_id: drafts}, + ) + + assert bitmask is not None + assert bitmask.shape[0] == len(drafts) + 1 + assert (bitmask[0] == -1).all() + assert (bitmask[1] == -1).all() + assert not (bitmask[2] == -1).all() + assert not (bitmask[-1] == -1).all() + assert not grammar.is_terminated() diff --git a/tests/v1/structured_output/test_reasoning_structured_output.py b/tests/v1/structured_output/test_reasoning_structured_output.py index 00a6cd87bc..a57e5da72b 100644 --- a/tests/v1/structured_output/test_reasoning_structured_output.py +++ b/tests/v1/structured_output/test_reasoning_structured_output.py @@ -3,6 +3,7 @@ """Unit tests for reasoning-aware structured output functionality (PR #25515).""" +from collections.abc import Sequence from unittest.mock import Mock import pytest @@ -13,6 +14,33 @@ from vllm.v1.structured_output import StructuredOutputManager from vllm.v1.structured_output.backend_types import StructuredOutputOptions +class UntouchableTokenIds(Sequence[int]): + """A token history that fails the test if anything reads it.""" + + def __init__(self, length: int) -> None: + self._length = length + + def __len__(self) -> int: + return self._length + + def __getitem__(self, item): + raise AssertionError("token history must not be read") + + def __iter__(self): + raise AssertionError("token history must not be iterated") + + +def make_qwen3_reasoner(): + """A real engine-based reasoning adapter over a tiny mock vocabulary.""" + from tests.parser.engine.conftest import make_mock_tokenizer + from vllm.parser.engine.registered_adapters import Qwen3ParserReasoningAdapter + + tokenizer = make_mock_tokenizer( + {"": 200, "": 201, "": 202, "": 203} + ) + return Qwen3ParserReasoningAdapter(tokenizer) + + class MockReasoner: def __init__(self, tokenizer): self.is_reasoning_end = Mock(return_value=False) @@ -364,6 +392,31 @@ class TestReasoningStructuredOutput: assert list(called_delta) == [4, 5] assert result is False + def test_should_advance_whole_delta_fallback_pins_last_sequence_index( + self, + manager_with_reasoner, + mock_request_with_structured_output, + ): + """When only the whole-delta predicate fires, the boundary is the + last index of the sequence, even if the placeholder window + overshoots it and the delta is empty. + """ + structured_req = mock_request_with_structured_output.structured_output_request + structured_req.reasoning_ended = False + reasoner = MockReasoner(tokenizer=Mock()) + # Base-default shape: rescans the whole sequence, ignores the delta. + reasoner.is_reasoning_end_streaming.return_value = True + structured_req.reasoner = reasoner + + mock_request_with_structured_output.all_token_ids = [1, 2, 3, 4, 5] + # placeholder window: start = 9 - 1 = 8, past the end -> delta = [] + mock_request_with_structured_output.num_computed_tokens = 9 + mock_request_with_structured_output.num_output_placeholders = 1 + + assert manager_with_reasoner.should_advance(mock_request_with_structured_output) + assert structured_req.reasoning_ended is True + assert structured_req.reasoning_end_token_index == 4 + def test_should_advance_trims_reasoning_prefix_for_json( self, manager_with_reasoner, @@ -403,3 +456,102 @@ class TestReasoningStructuredOutput: assert manager_with_reasoner.trim_reasoning_for_advance( mock_request_with_structured_output, new_token_ids ) == [271, 5005] + + def test_should_advance_engine_adapter_scans_only_the_delta( + self, + manager_with_reasoner, + mock_request_with_structured_output, + ): + """Engine-based parsers locate the boundary from the delta alone.""" + structured_req = mock_request_with_structured_output.structured_output_request + structured_req.reasoning_ended = False + reasoner = make_qwen3_reasoner() + assert reasoner.reasoning_end_token_ids == {201, 202} + reasoner.is_reasoning_end_streaming = Mock( + side_effect=AssertionError("engine path must not rescan the prefix") + ) + structured_req.reasoner = reasoner + + new_token_ids = [9, 198, 201, 271] + mock_request_with_structured_output.all_token_ids = [1, 2, 3] + new_token_ids + + assert manager_with_reasoner.should_advance( + mock_request_with_structured_output, new_token_ids=new_token_ids + ) + assert structured_req.reasoning_ended is True + assert structured_req.reasoning_end_token_index == 5 + assert manager_with_reasoner.trim_reasoning_for_advance( + mock_request_with_structured_output, new_token_ids + ) == [271] + + def test_should_advance_engine_adapter_tool_start_ends_reasoning( + self, + manager_with_reasoner, + mock_request_with_structured_output, + ): + structured_req = mock_request_with_structured_output.structured_output_request + structured_req.reasoning_ended = False + structured_req.reasoner = make_qwen3_reasoner() + + new_token_ids = [9, 202, 271] + mock_request_with_structured_output.all_token_ids = [1, 2, 3] + new_token_ids + + assert manager_with_reasoner.should_advance( + mock_request_with_structured_output, new_token_ids=new_token_ids + ) + assert structured_req.reasoning_end_token_index == 4 + + def test_should_advance_engine_adapter_no_boundary( + self, + manager_with_reasoner, + mock_request_with_structured_output, + ): + structured_req = mock_request_with_structured_output.structured_output_request + structured_req.reasoning_ended = False + structured_req.reasoner = make_qwen3_reasoner() + + new_token_ids = [9, 198, 271] + mock_request_with_structured_output.all_token_ids = [1, 2, 3] + new_token_ids + + assert not manager_with_reasoner.should_advance( + mock_request_with_structured_output, new_token_ids=new_token_ids + ) + assert structured_req.reasoning_ended is False + + def test_should_advance_engine_adapter_never_reads_history( + self, + manager_with_reasoner, + mock_request_with_structured_output, + ): + """Per-step cost must not depend on how much has been generated.""" + structured_req = mock_request_with_structured_output.structured_output_request + structured_req.reasoning_ended = False + structured_req.reasoner = make_qwen3_reasoner() + + new_token_ids = [9, 201] + total = (1 << 20) + len(new_token_ids) + mock_request_with_structured_output.all_token_ids = UntouchableTokenIds(total) + + assert manager_with_reasoner.should_advance( + mock_request_with_structured_output, new_token_ids=new_token_ids + ) + assert structured_req.reasoning_end_token_index == total - 1 + + def test_should_advance_engine_adapter_without_new_tokens_persists( + self, + manager_with_reasoner, + mock_request_with_structured_output, + ): + """Draft-validation callers persist the boundary like the legacy path.""" + structured_req = mock_request_with_structured_output.structured_output_request + structured_req.reasoning_ended = False + structured_req.reasoner = make_qwen3_reasoner() + + mock_request_with_structured_output.all_token_ids = [1, 2, 3, 4, 5, 201, 7] + # placeholder window: start = 7 - 2 = 5, delta = [201, 7] + mock_request_with_structured_output.num_computed_tokens = 7 + mock_request_with_structured_output.num_output_placeholders = 2 + + assert manager_with_reasoner.should_advance(mock_request_with_structured_output) + assert structured_req.reasoning_ended is True + assert structured_req.reasoning_end_token_index == 5 diff --git a/vllm/parser/abstract_parser.py b/vllm/parser/abstract_parser.py index c9b7796e45..cf83bb996e 100644 --- a/vllm/parser/abstract_parser.py +++ b/vllm/parser/abstract_parser.py @@ -219,6 +219,21 @@ class Parser: """ return self.is_reasoning_end(input_ids) + def find_reasoning_end_offset(self, token_ids: Sequence[int]) -> int | None: + """ + Locate the token that ends reasoning within one window of tokens. + + Args: + token_ids: The tokens to examine. + + Returns: + The offset within ``token_ids`` of the last reasoning token, or + ``None`` when reasoning does not end in this window. Parsers that + cannot answer from a window alone always return ``None``; callers + fall back to :meth:`is_reasoning_end_streaming`. + """ + return None + @abstractmethod def extract_content_ids(self, input_ids: list[int]) -> list[int]: """ diff --git a/vllm/parser/engine/adapters.py b/vllm/parser/engine/adapters.py index cd723f5f8b..4f4aa9f5a1 100644 --- a/vllm/parser/engine/adapters.py +++ b/vllm/parser/engine/adapters.py @@ -65,6 +65,9 @@ class ParserEngineReasoningAdapter(ReasoningParser): def is_reasoning_end(self, input_ids: Sequence[int]) -> bool: return self._parser_engine.is_reasoning_end(list(input_ids)) + def find_reasoning_end_offset(self, token_ids: Sequence[int]) -> int | None: + return self._parser_engine.find_reasoning_end_offset(token_ids) + def adjust_initial_state_from_prompt(self, prompt_token_ids: Sequence[int]) -> None: self._parser_engine.adjust_initial_state_from_prompt(prompt_token_ids) @@ -108,6 +111,10 @@ class ParserEngineReasoningAdapter(ReasoningParser): def reasoning_end_str(self) -> str | None: return self._parser_engine.reasoning_end_str + @property + def reasoning_end_token_ids(self) -> frozenset[int]: + return self._parser_engine.reasoning_end_token_ids + def adjust_request( self, request: ChatCompletionRequest | ResponsesRequest, diff --git a/vllm/parser/engine/parser_engine.py b/vllm/parser/engine/parser_engine.py index 2e6bd92344..36bb4d0cb5 100644 --- a/vllm/parser/engine/parser_engine.py +++ b/vllm/parser/engine/parser_engine.py @@ -152,6 +152,9 @@ class ParserEngine(Parser): for token in parser_engine_config.turn_boundary_tokens if (token_id := vocab.get(token)) is not None ) + self._reasoning_end_token_ids: frozenset[int] = ( + self._derive_reasoning_end_token_ids(parser_engine_config, vocab) + ) @property def reasoning_start_str(self) -> str | None: @@ -608,6 +611,42 @@ class ParserEngine(Parser): # ── Reasoning state queries ─────────────────────────────────────── + @staticmethod + def _derive_reasoning_end_token_ids( + config: ParserEngineConfig, vocab: dict[str, int] + ) -> frozenset[int]: + end_terminals: dict[str, ParserState] = {} + for (state, terminal), transition in config.transitions.items(): + if state != ParserState.REASONING: + continue + if EventType.REASONING_END in transition.events: + end_terminals.setdefault(terminal, transition.next_state) + elif transition.next_state != ParserState.REASONING: + return frozenset() + + token_ids: set[int] = set() + for terminal, next_state in end_terminals.items(): + text = config.token_id_terminals.get(terminal) + token_id = vocab.get(text) if text is not None else None + if token_id is not None: + token_ids.add(token_id) + elif next_state == ParserState.CONTENT: + return frozenset() + return frozenset(token_ids) + + @property + def reasoning_end_token_ids(self) -> frozenset[int]: + return self._reasoning_end_token_ids + + def find_reasoning_end_offset(self, token_ids: Sequence[int]) -> int | None: + end_ids = self._reasoning_end_token_ids + if not end_ids: + return None + for offset, token_id in enumerate(token_ids): + if token_id in end_ids: + return offset + return None + def is_reasoning_end(self, input_ids: list[int]) -> bool: end_id = self._reasoning_end_token_id start_id = self._reasoning_start_token_id diff --git a/vllm/v1/structured_output/__init__.py b/vllm/v1/structured_output/__init__.py index d0ec080d84..fd48c99d1c 100644 --- a/vllm/v1/structured_output/__init__.py +++ b/vllm/v1/structured_output/__init__.py @@ -8,6 +8,7 @@ from typing import TYPE_CHECKING from vllm.config import VllmConfig from vllm.logger import init_logger +from vllm.parser.engine.adapters import ParserEngineReasoningAdapter from vllm.reasoning import ReasoningParserManager from vllm.tokenizers import cached_tokenizer_from_config from vllm.utils.import_utils import LazyLoader @@ -291,40 +292,43 @@ class StructuredOutputManager: apply_bitmask = self.should_fill_bitmask(request) reasoner = None if apply_bitmask else self._get_reasoner(request) - detect_reasoning_end = reasoner is not None - simulated_buf: list[int] | None = None - history_len = 0 state_advancements = 0 post_reasoning_end_in_window = False req_tokens = scheduled_spec_decode_tokens.get(req_id, ()) + # Locate the reasoning-end marker in the draft window once. + # Rejected-draft padding (-1) is trailing, so the window is + # cut there and the reasoner never sees a placeholder. + window_start = len(request.all_token_ids) + reasoning_end_index: int | None = None + if reasoner is not None and req_tokens: + window = req_tokens + if -1 in window: + window = window[: window.index(-1)] + reasoning_end_index = self._find_reasoning_end_index( + reasoner, + request.all_token_ids, + window_start, + window, + delta_appended=False, + ) for i, token in enumerate(req_tokens): self._fill_bitmasks(((grammar, cumulative_index, apply_bitmask),)) advance_grammar = apply_bitmask if token == -1: apply_bitmask = False advance_grammar = False - elif ( - detect_reasoning_end - and reasoner is not None - and not apply_bitmask - ): - if simulated_buf is None: - history = list(request.all_token_ids) - history_len = len(history) - simulated_buf = history + list(req_tokens) - simulated = simulated_buf[: history_len + i + 1] - if reasoner.is_reasoning_end_streaming(simulated, [token]): - # Reasoning ended mid-window. Constrain the rest - # of the window via bitmask. Skip grammar advance - # through the marker (it is reasoning content); - # try to advance through subsequent drafts so the - # next bitmask row reflects the post-advance state, - # but tolerate rejection since those drafts predate - # the bitmask and are not guaranteed valid. - apply_bitmask = True - advance_grammar = False - post_reasoning_end_in_window = True + elif not apply_bitmask and window_start + i == reasoning_end_index: + # Reasoning ended mid-window. Constrain the rest + # of the window via bitmask. Skip grammar advance + # through the marker (it is reasoning content); + # try to advance through subsequent drafts so the + # next bitmask row reflects the post-advance state, + # but tolerate rejection since those drafts predate + # the bitmask and are not guaranteed valid. + apply_bitmask = True + advance_grammar = False + post_reasoning_end_in_window = True if advance_grammar and not grammar.is_terminated(): if post_reasoning_end_in_window: accepted = bool(grammar.validate_tokens([token])) @@ -431,21 +435,24 @@ class StructuredOutputManager: # The tokens were already appended this step, so the step window # starts exactly len(new_token_ids) from the end. start = len(all_token_ids) - len(new_token_ids) - delta_ids: Iterable[int] = new_token_ids + delta_ids: Sequence[int] = new_token_ids else: + # Draft-validation callers arrive here after update_from_output + # has already persisted any boundary found this step, so this + # window only needs to cover tokens not examined yet. delta_from = request.num_computed_tokens - request.num_output_placeholders start = ( delta_from if delta_from >= 0 else max(len(all_token_ids) + delta_from, 0) ) - delta_ids = itertools.islice(all_token_ids, start, None) - if reasoner.is_reasoning_end_streaming(all_token_ids, delta_ids): + delta_ids = all_token_ids[start:] + end_index = self._find_reasoning_end_index( + reasoner, all_token_ids, start, delta_ids, delta_appended=True + ) + if end_index is not None: structured_req.reasoning_ended = True - # Record the boundary so the scheduler can exclude reasoning tokens. - end_index = self._find_reasoning_end_index(reasoner, all_token_ids, start) - structured_req.reasoning_end_token_index = end_index return True @@ -453,24 +460,46 @@ class StructuredOutputManager: @staticmethod def _find_reasoning_end_index( - reasoner: "ReasoningParser", all_token_ids: Sequence[int], start: int - ) -> int: - """Locates the last reasoning token within ``all_token_ids[start:]``. + reasoner: "ReasoningParser", + all_token_ids: Sequence[int], + start: int, + delta_ids: Sequence[int], + *, + delta_appended: bool, + ) -> int | None: + """Locates the token that ends reasoning within ``delta_ids``. + + ``delta_ids`` occupy positions ``start..``; ``delta_appended`` says + whether they are already in ``all_token_ids`` (accepted tokens) or + not (speculative drafts). + + Legacy parsers are probed one token at a time. Accepted tokens get a + whole-delta check as well, and a marker seen only that way reports + the final token. Drafts skip it: some predicates are order-sensitive + over a multi-token window and would miss the marker. Returns: - The absolute index of the token at which - ``is_reasoning_end_streaming`` first fires. Falls back to the - final index when no single token triggers the detection (e.g. - a multi-token marker only recognized on the full delta), which - conservatively treats the whole step as reasoning content. + Absolute index of the last reasoning token, or ``None``. """ + if ( + isinstance(reasoner, ParserEngineReasoningAdapter) + and reasoner.reasoning_end_token_ids + ): + offset = reasoner.find_reasoning_end_offset(delta_ids) + return None if offset is None else start + offset + + if delta_appended and not reasoner.is_reasoning_end_streaming( + all_token_ids, delta_ids + ): + return None prefix = list(itertools.islice(all_token_ids, start)) - for idx in range(start, len(all_token_ids)): - token = all_token_ids[idx] + for offset, token in enumerate(delta_ids): prefix.append(token) if reasoner.is_reasoning_end_streaming(prefix, [token]): - return idx - return len(all_token_ids) - 1 + return start + offset + if delta_appended: + return len(all_token_ids) - 1 + return None def trim_reasoning_for_advance( self, request: "Request", new_token_ids: list[int]