[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>
Flora Feng committed
Sep 8, 2026 at 17:07 UTC
29f46ac475322cfc14ced890825b9a149ada3608
8 files changed
+475
-42
tests/parser/engine/test_deepseek_v4.py
+11
@@ -416,6 +416,17 @@ class TestImplicitReasoningEnd:
416
args = json.loads(tool_calls[0].arguments)
417
assert args == {"location": "NYC"}
418
419
+ def test_reasoning_end_token_ids_drop_multi_token_tool_opener(
420
+ self, thinking_parser
421
+ ):
422
+ # Like the real DeepSeek-V4 tokenizer, the mock vocab has no single
423
+ # token for the DSML tool opener, so only </think> is tracked on the
424
+ # token-ID fast path; the implicit tool-call exit still goes through
425
+ # the text lexer.
426
+ assert thinking_parser.reasoning_end_token_ids == {_THINK_END_ID}
427
+ assert thinking_parser.find_reasoning_end_offset([7, _THINK_END_ID, 8]) == 1
428
+ assert thinking_parser.find_reasoning_end_offset([7, 8]) is None
429
+
430
def test_streaming_reasoning_implicit_end(self, thinking_parser):
431
chunks = [
432
"Let me look up the weather.\n\n",
tests/parser/engine/test_parser_engine.py
+63
@@ -7,6 +7,7 @@ DeltaMessage / ExtractedToolCallInformation protocol.
7
8
from __future__ import annotations
9
10
+import dataclasses
11
import json
12
from types import SimpleNamespace
13
from unittest.mock import MagicMock
@@ -132,6 +133,68 @@ def _make_engine(
133
)
134
135
136
+# ── TestReasoningEndTokenIds ─────────────────────────────────────────
137
+
138
+
139
+def _with_reasoning_exits(
140
+ *exits: tuple[str, ParserState, tuple[EventType, ...]],
141
+) -> ParserEngineConfig:
142
+ """Combined config plus extra transitions out of REASONING."""
143
+ base = _combined_config()
144
+ transitions = dict(base.transitions)
145
+ for terminal, next_state, events in exits:
146
+ transitions[(ParserState.REASONING, terminal)] = Transition(next_state, events)
147
+ return dataclasses.replace(base, transitions=transitions)
148
+
149
+
150
+class TestReasoningEndTokenIds:
151
+ """ParserEngine derives the reasoning-end token set from its config."""
152
+
153
+ def test_every_reasoning_end_exit_contributes(self):
154
+ cfg = _with_reasoning_exits(
155
+ (
156
+ "TOOL_START",
157
+ ParserState.TOOL_ARGS,
158
+ (EventType.REASONING_END, EventType.TOOL_CALL_START),
159
+ ),
160
+ )
161
+ assert _make_engine(cfg).reasoning_end_token_ids == {201, 202}
162
+
163
+ def test_transition_staying_in_reasoning_is_ignored(self):
164
+ cfg = _with_reasoning_exits(("THINK_START", ParserState.REASONING, ()))
165
+ assert _make_engine(cfg).reasoning_end_token_ids == {201}
166
+
167
+ def test_unreported_exit_fails_closed(self):
168
+ cfg = _with_reasoning_exits(
169
+ ("TOOL_START", ParserState.TOOL_ARGS, (EventType.TOOL_CALL_START,)),
170
+ )
171
+ assert _make_engine(cfg).reasoning_end_token_ids == frozenset()
172
+
173
+ def test_unresolved_think_end_fails_closed(self):
174
+ vocab = {k: v for k, v in _VOCAB.items() if k != "</think>"}
175
+ assert _make_engine(vocab=vocab).reasoning_end_token_ids == frozenset()
176
+
177
+ text_only = dataclasses.replace(
178
+ _combined_config(), token_id_terminals={"THINK_START": "<think>"}
179
+ )
180
+ assert _make_engine(text_only).reasoning_end_token_ids == frozenset()
181
+
182
+ def test_config_without_reasoning_has_empty_set(self):
183
+ assert _make_engine(_hermes_config()).reasoning_end_token_ids == frozenset()
184
+
185
+ def test_find_reasoning_end_offset_returns_first_match(self):
186
+ engine = _make_engine()
187
+ assert engine.find_reasoning_end_offset([5, 201, 6, 201]) == 1
188
+ assert engine.find_reasoning_end_offset([5, 6]) is None
189
+ assert engine.find_reasoning_end_offset([]) is None
190
+ # Rejected-draft placeholders never match.
191
+ assert engine.find_reasoning_end_offset([-1, 201]) == 1
192
+
193
+ def test_find_reasoning_end_offset_with_empty_set_returns_none(self):
194
+ engine = _make_engine(_hermes_config())
195
+ assert engine.find_reasoning_end_offset([201]) is None
196
+
197
+
198
# ── TestEventsToDelta ────────────────────────────────────────────────
199
200
tests/v1/spec_decode/test_mtp_structured_output.py
+117
@@ -8,6 +8,7 @@ from transformers import AutoTokenizer
8
from vllm.config import StructuredOutputsConfig, VllmConfig
9
from vllm.config.model import ModelConfig
10
from vllm.config.speculative import SpeculativeConfig
11
+from vllm.parser.engine.adapters import ParserEngineReasoningAdapter
12
from vllm.sampling_params import SamplingParams, StructuredOutputsParams
13
from vllm.v1.request import Request
14
from vllm.v1.structured_output import StructuredOutputManager
@@ -390,3 +391,119 @@ def test_trim_reasoning_for_advance():
391
next_step = [post, post]
392
request.append_output_token_ids(next_step)
393
assert manager.trim_reasoning_for_advance(request, next_step) == next_step
394
+
395
+
396
+class _EngineReasonerStub(ParserEngineReasoningAdapter):
397
+ """Adapter-typed reasoner with a fixed end-token set and no real engine."""
398
+
399
+ def __init__(self, end_token_ids):
400
+ self._end_token_ids = frozenset(end_token_ids)
401
+ self.windows: list[list[int]] = []
402
+
403
+ @property
404
+ def reasoning_end_token_ids(self):
405
+ return self._end_token_ids
406
+
407
+ def find_reasoning_end_offset(self, token_ids):
408
+ self.windows.append(list(token_ids))
409
+ for offset, token in enumerate(token_ids):
410
+ if token in self._end_token_ids:
411
+ return offset
412
+ return None
413
+
414
+ def is_reasoning_end(self, input_ids):
415
+ return any(token in self._end_token_ids for token in input_ids)
416
+
417
+ def is_reasoning_end_streaming(self, input_ids, delta_ids):
418
+ raise AssertionError("engine path must not rescan draft prefixes")
419
+
420
+
421
+@pytest.mark.parametrize("backend", ["xgrammar", "guidance"])
422
+def test_bitmask_engine_reasoner_ends_midwindow_with_padding(backend):
423
+ """Engine reasoners see the draft window once, without -1 padding."""
424
+ tokenizer, manager, request, prompt = _make_manager_and_request(backend)
425
+ grammar = request.structured_output_request.grammar
426
+
427
+ assert grammar.accept_tokens(request.request_id, prompt)
428
+
429
+ marker = tokenizer.encode("\n")[0]
430
+ reasoner = _EngineReasonerStub({marker})
431
+ manager.reasoner_cls = _EngineReasonerStub
432
+ request.structured_output_request.reasoner = reasoner
433
+ request.structured_output_request.reasoning_ended = False
434
+
435
+ pre = tokenizer.encode(" ")[0]
436
+ post = tokenizer.encode(",")[0]
437
+ drafts = [pre, marker, post, -1]
438
+
439
+ bitmask = manager.grammar_bitmask(
440
+ requests={request.request_id: request},
441
+ structured_output_request_ids=[request.request_id],
442
+ scheduled_spec_decode_tokens={request.request_id: drafts},
443
+ )
444
+
445
+ assert bitmask is not None
446
+ assert bitmask.shape[0] == len(drafts) + 1
447
+ assert (bitmask[0] == -1).all()
448
+ assert (bitmask[1] == -1).all()
449
+ assert not (bitmask[2] == -1).all()
450
+ assert reasoner.windows == [[pre, marker, post]]
451
+ assert not grammar.is_terminated()
452
+
453
+
454
+@pytest.mark.parametrize("backend", ["xgrammar", "guidance"])
455
+def test_bitmask_legacy_order_sensitive_predicate_probed_per_token(backend):
456
+ """Legacy predicates are probed one draft at a time, never on the whole
457
+ window. A predicate that reports which marker is *newest* (KimiK3) would
458
+ otherwise answer False for a window that closes and then reopens
459
+ reasoning, leaving the post-marker drafts and the bonus row
460
+ unconstrained.
461
+ """
462
+ tokenizer, manager, request, prompt = _make_manager_and_request(backend)
463
+ grammar = request.structured_output_request.grammar
464
+
465
+ assert grammar.accept_tokens(request.request_id, prompt)
466
+
467
+ close = tokenizer.encode("\n")[0]
468
+ reopen = tokenizer.encode("#")[0]
469
+ assert close != reopen
470
+
471
+ class NewestMarkerReasoner:
472
+ def __init__(self, *_, **__):
473
+ pass
474
+
475
+ def is_reasoning_end(self, input_ids):
476
+ return self.is_reasoning_end_streaming(input_ids, input_ids)
477
+
478
+ def is_reasoning_end_streaming(self, input_ids, delta_ids):
479
+ for token in reversed(list(delta_ids)):
480
+ if token == close:
481
+ return True
482
+ if token == reopen:
483
+ return False
484
+ return False
485
+
486
+ manager.reasoner_cls = NewestMarkerReasoner
487
+ request.structured_output_request.reasoner = NewestMarkerReasoner()
488
+ request.structured_output_request.reasoning_ended = False
489
+
490
+ pre = tokenizer.encode(" ")[0]
491
+ post = tokenizer.encode(",")[0]
492
+ drafts = [pre, close, post, reopen]
493
+ # Sanity: the whole window says "reopened", so only per-token probing
494
+ # can find the close marker.
495
+ assert not NewestMarkerReasoner().is_reasoning_end_streaming([], drafts)
496
+
497
+ bitmask = manager.grammar_bitmask(
498
+ requests={request.request_id: request},
499
+ structured_output_request_ids=[request.request_id],
500
+ scheduled_spec_decode_tokens={request.request_id: drafts},
501
+ )
502
+
503
+ assert bitmask is not None
504
+ assert bitmask.shape[0] == len(drafts) + 1
505
+ assert (bitmask[0] == -1).all()
506
+ assert (bitmask[1] == -1).all()
507
+ assert not (bitmask[2] == -1).all()
508
+ assert not (bitmask[-1] == -1).all()
509
+ assert not grammar.is_terminated()
tests/v1/structured_output/test_reasoning_structured_output.py
+152
@@ -3,6 +3,7 @@
3
4
"""Unit tests for reasoning-aware structured output functionality (PR #25515)."""
5
6
+from collections.abc import Sequence
7
from unittest.mock import Mock
8
9
import pytest
@@ -13,6 +14,33 @@ from vllm.v1.structured_output import StructuredOutputManager
14
from vllm.v1.structured_output.backend_types import StructuredOutputOptions
15
16
17
+class UntouchableTokenIds(Sequence[int]):
18
+ """A token history that fails the test if anything reads it."""
19
+
20
+ def __init__(self, length: int) -> None:
21
+ self._length = length
22
+
23
+ def __len__(self) -> int:
24
+ return self._length
25
+
26
+ def __getitem__(self, item):
27
+ raise AssertionError("token history must not be read")
28
+
29
+ def __iter__(self):
30
+ raise AssertionError("token history must not be iterated")
31
+
32
+
33
+def make_qwen3_reasoner():
34
+ """A real engine-based reasoning adapter over a tiny mock vocabulary."""
35
+ from tests.parser.engine.conftest import make_mock_tokenizer
36
+ from vllm.parser.engine.registered_adapters import Qwen3ParserReasoningAdapter
37
+
38
+ tokenizer = make_mock_tokenizer(
39
+ {"<think>": 200, "</think>": 201, "<tool_call>": 202, "</tool_call>": 203}
40
+ )
41
+ return Qwen3ParserReasoningAdapter(tokenizer)
42
+
43
+
44
class MockReasoner:
45
def __init__(self, tokenizer):
46
self.is_reasoning_end = Mock(return_value=False)
@@ -364,6 +392,31 @@ class TestReasoningStructuredOutput:
392
assert list(called_delta) == [4, 5]
393
assert result is False
394
395
+ def test_should_advance_whole_delta_fallback_pins_last_sequence_index(
396
+ self,
397
+ manager_with_reasoner,
398
+ mock_request_with_structured_output,
399
+ ):
400
+ """When only the whole-delta predicate fires, the boundary is the
401
+ last index of the sequence, even if the placeholder window
402
+ overshoots it and the delta is empty.
403
+ """
404
+ structured_req = mock_request_with_structured_output.structured_output_request
405
+ structured_req.reasoning_ended = False
406
+ reasoner = MockReasoner(tokenizer=Mock())
407
+ # Base-default shape: rescans the whole sequence, ignores the delta.
408
+ reasoner.is_reasoning_end_streaming.return_value = True
409
+ structured_req.reasoner = reasoner
410
+
411
+ mock_request_with_structured_output.all_token_ids = [1, 2, 3, 4, 5]
412
+ # placeholder window: start = 9 - 1 = 8, past the end -> delta = []
413
+ mock_request_with_structured_output.num_computed_tokens = 9
414
+ mock_request_with_structured_output.num_output_placeholders = 1
415
+
416
+ assert manager_with_reasoner.should_advance(mock_request_with_structured_output)
417
+ assert structured_req.reasoning_ended is True
418
+ assert structured_req.reasoning_end_token_index == 4
419
+
420
def test_should_advance_trims_reasoning_prefix_for_json(
421
self,
422
manager_with_reasoner,
@@ -403,3 +456,102 @@ class TestReasoningStructuredOutput:
456
assert manager_with_reasoner.trim_reasoning_for_advance(
457
mock_request_with_structured_output, new_token_ids
458
) == [271, 5005]
459
+
460
+ def test_should_advance_engine_adapter_scans_only_the_delta(
461
+ self,
462
+ manager_with_reasoner,
463
+ mock_request_with_structured_output,
464
+ ):
465
+ """Engine-based parsers locate the boundary from the delta alone."""
466
+ structured_req = mock_request_with_structured_output.structured_output_request
467
+ structured_req.reasoning_ended = False
468
+ reasoner = make_qwen3_reasoner()
469
+ assert reasoner.reasoning_end_token_ids == {201, 202}
470
+ reasoner.is_reasoning_end_streaming = Mock(
471
+ side_effect=AssertionError("engine path must not rescan the prefix")
472
+ )
473
+ structured_req.reasoner = reasoner
474
+
475
+ new_token_ids = [9, 198, 201, 271]
476
+ mock_request_with_structured_output.all_token_ids = [1, 2, 3] + new_token_ids
477
+
478
+ assert manager_with_reasoner.should_advance(
479
+ mock_request_with_structured_output, new_token_ids=new_token_ids
480
+ )
481
+ assert structured_req.reasoning_ended is True
482
+ assert structured_req.reasoning_end_token_index == 5
483
+ assert manager_with_reasoner.trim_reasoning_for_advance(
484
+ mock_request_with_structured_output, new_token_ids
485
+ ) == [271]
486
+
487
+ def test_should_advance_engine_adapter_tool_start_ends_reasoning(
488
+ self,
489
+ manager_with_reasoner,
490
+ mock_request_with_structured_output,
491
+ ):
492
+ structured_req = mock_request_with_structured_output.structured_output_request
493
+ structured_req.reasoning_ended = False
494
+ structured_req.reasoner = make_qwen3_reasoner()
495
+
496
+ new_token_ids = [9, 202, 271]
497
+ mock_request_with_structured_output.all_token_ids = [1, 2, 3] + new_token_ids
498
+
499
+ assert manager_with_reasoner.should_advance(
500
+ mock_request_with_structured_output, new_token_ids=new_token_ids
501
+ )
502
+ assert structured_req.reasoning_end_token_index == 4
503
+
504
+ def test_should_advance_engine_adapter_no_boundary(
505
+ self,
506
+ manager_with_reasoner,
507
+ mock_request_with_structured_output,
508
+ ):
509
+ structured_req = mock_request_with_structured_output.structured_output_request
510
+ structured_req.reasoning_ended = False
511
+ structured_req.reasoner = make_qwen3_reasoner()
512
+
513
+ new_token_ids = [9, 198, 271]
514
+ mock_request_with_structured_output.all_token_ids = [1, 2, 3] + new_token_ids
515
+
516
+ assert not manager_with_reasoner.should_advance(
517
+ mock_request_with_structured_output, new_token_ids=new_token_ids
518
+ )
519
+ assert structured_req.reasoning_ended is False
520
+
521
+ def test_should_advance_engine_adapter_never_reads_history(
522
+ self,
523
+ manager_with_reasoner,
524
+ mock_request_with_structured_output,
525
+ ):
526
+ """Per-step cost must not depend on how much has been generated."""
527
+ structured_req = mock_request_with_structured_output.structured_output_request
528
+ structured_req.reasoning_ended = False
529
+ structured_req.reasoner = make_qwen3_reasoner()
530
+
531
+ new_token_ids = [9, 201]
532
+ total = (1 << 20) + len(new_token_ids)
533
+ mock_request_with_structured_output.all_token_ids = UntouchableTokenIds(total)
534
+
535
+ assert manager_with_reasoner.should_advance(
536
+ mock_request_with_structured_output, new_token_ids=new_token_ids
537
+ )
538
+ assert structured_req.reasoning_end_token_index == total - 1
539
+
540
+ def test_should_advance_engine_adapter_without_new_tokens_persists(
541
+ self,
542
+ manager_with_reasoner,
543
+ mock_request_with_structured_output,
544
+ ):
545
+ """Draft-validation callers persist the boundary like the legacy path."""
546
+ structured_req = mock_request_with_structured_output.structured_output_request
547
+ structured_req.reasoning_ended = False
548
+ structured_req.reasoner = make_qwen3_reasoner()
549
+
550
+ mock_request_with_structured_output.all_token_ids = [1, 2, 3, 4, 5, 201, 7]
551
+ # placeholder window: start = 7 - 2 = 5, delta = [201, 7]
552
+ mock_request_with_structured_output.num_computed_tokens = 7
553
+ mock_request_with_structured_output.num_output_placeholders = 2
554
+
555
+ assert manager_with_reasoner.should_advance(mock_request_with_structured_output)
556
+ assert structured_req.reasoning_ended is True
557
+ assert structured_req.reasoning_end_token_index == 5
vllm/parser/abstract_parser.py
+15
@@ -219,6 +219,21 @@ class Parser:
219
"""
220
return self.is_reasoning_end(input_ids)
221
222
+ def find_reasoning_end_offset(self, token_ids: Sequence[int]) -> int | None:
223
+ """
224
+ Locate the token that ends reasoning within one window of tokens.
225
+
226
+ Args:
227
+ token_ids: The tokens to examine.
228
+
229
+ Returns:
230
+ The offset within ``token_ids`` of the last reasoning token, or
231
+ ``None`` when reasoning does not end in this window. Parsers that
232
+ cannot answer from a window alone always return ``None``; callers
233
+ fall back to :meth:`is_reasoning_end_streaming`.
234
+ """
235
+ return None
236
+
237
@abstractmethod
238
def extract_content_ids(self, input_ids: list[int]) -> list[int]:
239
"""
vllm/parser/engine/adapters.py
+7
@@ -65,6 +65,9 @@ class ParserEngineReasoningAdapter(ReasoningParser):
65
def is_reasoning_end(self, input_ids: Sequence[int]) -> bool:
66
return self._parser_engine.is_reasoning_end(list(input_ids))
67
68
+ def find_reasoning_end_offset(self, token_ids: Sequence[int]) -> int | None:
69
+ return self._parser_engine.find_reasoning_end_offset(token_ids)
70
+
71
def adjust_initial_state_from_prompt(self, prompt_token_ids: Sequence[int]) -> None:
72
self._parser_engine.adjust_initial_state_from_prompt(prompt_token_ids)
73
@@ -108,6 +111,10 @@ class ParserEngineReasoningAdapter(ReasoningParser):
111
def reasoning_end_str(self) -> str | None:
112
return self._parser_engine.reasoning_end_str
113
114
+ @property
115
+ def reasoning_end_token_ids(self) -> frozenset[int]:
116
+ return self._parser_engine.reasoning_end_token_ids
117
+
118
def adjust_request(
119
self,
120
request: ChatCompletionRequest | ResponsesRequest,
vllm/parser/engine/parser_engine.py
+39
@@ -152,6 +152,9 @@ class ParserEngine(Parser):
152
for token in parser_engine_config.turn_boundary_tokens
153
if (token_id := vocab.get(token)) is not None
154
)
155
+ self._reasoning_end_token_ids: frozenset[int] = (
156
+ self._derive_reasoning_end_token_ids(parser_engine_config, vocab)
157
+ )
158
159
@property
160
def reasoning_start_str(self) -> str | None:
@@ -608,6 +611,42 @@ class ParserEngine(Parser):
611
612
# ── Reasoning state queries ───────────────────────────────────────
613
614
+ @staticmethod
615
+ def _derive_reasoning_end_token_ids(
616
+ config: ParserEngineConfig, vocab: dict[str, int]
617
+ ) -> frozenset[int]:
618
+ end_terminals: dict[str, ParserState] = {}
619
+ for (state, terminal), transition in config.transitions.items():
620
+ if state != ParserState.REASONING:
621
+ continue
622
+ if EventType.REASONING_END in transition.events:
623
+ end_terminals.setdefault(terminal, transition.next_state)
624
+ elif transition.next_state != ParserState.REASONING:
625
+ return frozenset()
626
+
627
+ token_ids: set[int] = set()
628
+ for terminal, next_state in end_terminals.items():
629
+ text = config.token_id_terminals.get(terminal)
630
+ token_id = vocab.get(text) if text is not None else None
631
+ if token_id is not None:
632
+ token_ids.add(token_id)
633
+ elif next_state == ParserState.CONTENT:
634
+ return frozenset()
635
+ return frozenset(token_ids)
636
+
637
+ @property
638
+ def reasoning_end_token_ids(self) -> frozenset[int]:
639
+ return self._reasoning_end_token_ids
640
+
641
+ def find_reasoning_end_offset(self, token_ids: Sequence[int]) -> int | None:
642
+ end_ids = self._reasoning_end_token_ids
643
+ if not end_ids:
644
+ return None
645
+ for offset, token_id in enumerate(token_ids):
646
+ if token_id in end_ids:
647
+ return offset
648
+ return None
649
+
650
def is_reasoning_end(self, input_ids: list[int]) -> bool:
651
end_id = self._reasoning_end_token_id
652
start_id = self._reasoning_start_token_id
vllm/v1/structured_output/__init__.py
+71
-42
@@ -8,6 +8,7 @@ from typing import TYPE_CHECKING
8
9
from vllm.config import VllmConfig
10
from vllm.logger import init_logger
11
+from vllm.parser.engine.adapters import ParserEngineReasoningAdapter
12
from vllm.reasoning import ReasoningParserManager
13
from vllm.tokenizers import cached_tokenizer_from_config
14
from vllm.utils.import_utils import LazyLoader
@@ -291,40 +292,43 @@ class StructuredOutputManager:
292
apply_bitmask = self.should_fill_bitmask(request)
293
294
reasoner = None if apply_bitmask else self._get_reasoner(request)
294
- detect_reasoning_end = reasoner is not None
295
- simulated_buf: list[int] | None = None
296
- history_len = 0
295
296
state_advancements = 0
297
post_reasoning_end_in_window = False
298
req_tokens = scheduled_spec_decode_tokens.get(req_id, ())
299
+ # Locate the reasoning-end marker in the draft window once.
300
+ # Rejected-draft padding (-1) is trailing, so the window is
301
+ # cut there and the reasoner never sees a placeholder.
302
+ window_start = len(request.all_token_ids)
303
+ reasoning_end_index: int | None = None
304
+ if reasoner is not None and req_tokens:
305
+ window = req_tokens
306
+ if -1 in window:
307
+ window = window[: window.index(-1)]
308
+ reasoning_end_index = self._find_reasoning_end_index(
309
+ reasoner,
310
+ request.all_token_ids,
311
+ window_start,
312
+ window,
313
+ delta_appended=False,
314
+ )
315
for i, token in enumerate(req_tokens):
316
self._fill_bitmasks(((grammar, cumulative_index, apply_bitmask),))
317
advance_grammar = apply_bitmask
318
if token == -1:
319
apply_bitmask = False
320
advance_grammar = False
307
- elif (
308
- detect_reasoning_end
309
- and reasoner is not None
310
- and not apply_bitmask
311
- ):
312
- if simulated_buf is None:
313
- history = list(request.all_token_ids)
314
- history_len = len(history)
315
- simulated_buf = history + list(req_tokens)
316
- simulated = simulated_buf[: history_len + i + 1]
317
- if reasoner.is_reasoning_end_streaming(simulated, [token]):
318
- # Reasoning ended mid-window. Constrain the rest
319
- # of the window via bitmask. Skip grammar advance
320
- # through the marker (it is reasoning content);
321
- # try to advance through subsequent drafts so the
322
- # next bitmask row reflects the post-advance state,
323
- # but tolerate rejection since those drafts predate
324
- # the bitmask and are not guaranteed valid.
325
- apply_bitmask = True
326
- advance_grammar = False
327
- post_reasoning_end_in_window = True
321
+ elif not apply_bitmask and window_start + i == reasoning_end_index:
322
+ # Reasoning ended mid-window. Constrain the rest
323
+ # of the window via bitmask. Skip grammar advance
324
+ # through the marker (it is reasoning content);
325
+ # try to advance through subsequent drafts so the
326
+ # next bitmask row reflects the post-advance state,
327
+ # but tolerate rejection since those drafts predate
328
+ # the bitmask and are not guaranteed valid.
329
+ apply_bitmask = True
330
+ advance_grammar = False
331
+ post_reasoning_end_in_window = True
332
if advance_grammar and not grammar.is_terminated():
333
if post_reasoning_end_in_window:
334
accepted = bool(grammar.validate_tokens([token]))
@@ -431,21 +435,24 @@ class StructuredOutputManager:
435
# The tokens were already appended this step, so the step window
436
# starts exactly len(new_token_ids) from the end.
437
start = len(all_token_ids) - len(new_token_ids)
434
- delta_ids: Iterable[int] = new_token_ids
438
+ delta_ids: Sequence[int] = new_token_ids
439
else:
440
+ # Draft-validation callers arrive here after update_from_output
441
+ # has already persisted any boundary found this step, so this
442
+ # window only needs to cover tokens not examined yet.
443
delta_from = request.num_computed_tokens - request.num_output_placeholders
444
start = (
445
delta_from
446
if delta_from >= 0
447
else max(len(all_token_ids) + delta_from, 0)
448
)
442
- delta_ids = itertools.islice(all_token_ids, start, None)
443
- if reasoner.is_reasoning_end_streaming(all_token_ids, delta_ids):
449
+ delta_ids = all_token_ids[start:]
450
+ end_index = self._find_reasoning_end_index(
451
+ reasoner, all_token_ids, start, delta_ids, delta_appended=True
452
+ )
453
+ if end_index is not None:
454
structured_req.reasoning_ended = True
445
-
455
# Record the boundary so the scheduler can exclude reasoning tokens.
447
- end_index = self._find_reasoning_end_index(reasoner, all_token_ids, start)
448
-
456
structured_req.reasoning_end_token_index = end_index
457
return True
458
@@ -453,24 +460,46 @@ class StructuredOutputManager:
460
461
@staticmethod
462
def _find_reasoning_end_index(
456
- reasoner: "ReasoningParser", all_token_ids: Sequence[int], start: int
457
- ) -> int:
458
- """Locates the last reasoning token within ``all_token_ids[start:]``.
463
+ reasoner: "ReasoningParser",
464
+ all_token_ids: Sequence[int],
465
+ start: int,
466
+ delta_ids: Sequence[int],
467
+ *,
468
+ delta_appended: bool,
469
+ ) -> int | None:
470
+ """Locates the token that ends reasoning within ``delta_ids``.
471
+
472
+ ``delta_ids`` occupy positions ``start..``; ``delta_appended`` says
473
+ whether they are already in ``all_token_ids`` (accepted tokens) or
474
+ not (speculative drafts).
475
+
476
+ Legacy parsers are probed one token at a time. Accepted tokens get a
477
+ whole-delta check as well, and a marker seen only that way reports
478
+ the final token. Drafts skip it: some predicates are order-sensitive
479
+ over a multi-token window and would miss the marker.
480
481
Returns:
461
- The absolute index of the token at which
462
- ``is_reasoning_end_streaming`` first fires. Falls back to the
463
- final index when no single token triggers the detection (e.g.
464
- a multi-token marker only recognized on the full delta), which
465
- conservatively treats the whole step as reasoning content.
482
+ Absolute index of the last reasoning token, or ``None``.
483
"""
484
+ if (
485
+ isinstance(reasoner, ParserEngineReasoningAdapter)
486
+ and reasoner.reasoning_end_token_ids
487
+ ):
488
+ offset = reasoner.find_reasoning_end_offset(delta_ids)
489
+ return None if offset is None else start + offset
490
+
491
+ if delta_appended and not reasoner.is_reasoning_end_streaming(
492
+ all_token_ids, delta_ids
493
+ ):
494
+ return None
495
prefix = list(itertools.islice(all_token_ids, start))
468
- for idx in range(start, len(all_token_ids)):
469
- token = all_token_ids[idx]
496
+ for offset, token in enumerate(delta_ids):
497
prefix.append(token)
498
if reasoner.is_reasoning_end_streaming(prefix, [token]):
472
- return idx
473
- return len(all_token_ids) - 1
499
+ return start + offset
500
+ if delta_appended:
501
+ return len(all_token_ids) - 1
502
+ return None
503
504
def trim_reasoning_for_advance(
505
self, request: "Request", new_token_ids: list[int]