@tduncan / Nvidia-Optimization-Stack-VLLM / commits / e6960af33b

[Refactor] StructuredOutputManager x Speculative Decoding Refactor (#48200)

yzong-rh committed Sep 15, 2026 at 09:48 UTC e6960af33b379d502f409e3e2241bbf2b2c2f68d
12 files changed +595 -1191
tests/parser/engine/test_deepseek_v4.py
+1 -1
@@ -574,7 +574,7 @@ class TestImplicitReasoningEnd:
574 # the text lexer.
575 assert thinking_parser.reasoning_end_token_ids == {_THINK_END_ID}
576 assert thinking_parser.find_reasoning_end_offset([7, _THINK_END_ID, 8]) == 1
577 - assert thinking_parser.find_reasoning_end_offset([7, 8]) is None
577 + assert thinking_parser.find_reasoning_end_offset([7, 8]) == 2
578
579 def test_streaming_reasoning_implicit_end(self, thinking_parser):
580 chunks = [
tests/parser/engine/test_parser_engine.py
+2 -2
@@ -185,8 +185,8 @@ class TestReasoningEndTokenIds:
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
188 + assert engine.find_reasoning_end_offset([5, 6]) == 2
189 + assert engine.find_reasoning_end_offset([]) == 0
190 # Rejected-draft placeholders never match.
191 assert engine.find_reasoning_end_offset([-1, 201]) == 1
192
tests/v1/core/test_async_scheduler.py
+1 -5
@@ -334,7 +334,6 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
334 request = create_requests(num_requests=1, num_tokens=1)[0]
335 request.structured_output_request = Mock()
336 request.structured_output_request.grammar = Mock(spec=StructuredOutputGrammar)
337 - request.structured_output_request.grammar.accept_tokens.return_value = False
337 request.status = RequestStatus.RUNNING
338 request.num_computed_tokens = request.num_tokens
339 request.num_output_placeholders = 1
@@ -343,10 +342,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
342 scheduler.connector = None
343 scheduler.ec_connector = None
344 scheduler.structured_output_manager = Mock()
346 - scheduler.structured_output_manager.should_advance.return_value = True
347 - scheduler.structured_output_manager.trim_reasoning_for_advance.side_effect = (
348 - lambda request, new_token_ids: new_token_ids
349 - )
345 + scheduler.structured_output_manager.accept_tokens.return_value = False
346 scheduler.requests = {request.request_id: request}
347 scheduler.running = [request]
348 scheduler.waiting = Mock()
tests/v1/core/test_scheduler.py
+11 -17
@@ -3674,7 +3674,6 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
3674 )
3675 request.structured_output_request = Mock()
3676 request.structured_output_request.grammar = Mock(spec=StructuredOutputGrammar)
3677 - request.structured_output_request.grammar.accept_tokens.return_value = False
3677 request.status = RequestStatus.RUNNING
3678 request.num_computed_tokens = request.num_tokens
3679
@@ -3682,10 +3681,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
3681 scheduler.connector = None
3682 scheduler.ec_connector = None
3683 scheduler.structured_output_manager = Mock()
3685 - scheduler.structured_output_manager.should_advance.return_value = True
3686 - scheduler.structured_output_manager.trim_reasoning_for_advance.side_effect = (
3687 - lambda request, new_token_ids: new_token_ids
3688 - )
3684 + scheduler.structured_output_manager.accept_tokens.return_value = False
3685 scheduler.requests = {request.request_id: request}
3686 scheduler.running = [request]
3687 scheduler.waiting = Mock()
@@ -3734,8 +3730,8 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
3730 )
3731 engine_core_outputs = scheduler.update_from_output(output, model_runner_output)
3732
3737 - request.structured_output_request.grammar.accept_tokens.assert_called_once_with(
3738 - request.request_id, [123]
3733 + scheduler.structured_output_manager.accept_tokens.assert_called_once_with(
3734 + request, [123]
3735 )
3736 assert request.resumable is False
3737 assert request.status == RequestStatus.FINISHED_ERROR
@@ -6573,16 +6569,15 @@ def _decode_ready_request(scheduler):
6569 return request
6570
6571
6576 -def test_update_draft_token_ids_strips_ngram_padding(monkeypatch):
6577 - """ngram_gpu pads unfilled draft slots with -1; update_draft_token_ids must
6572 +def test_update_draft_token_ids_strips_ngram_padding():
6573 + """ngram_gpu pads unfilled draft slots with -1; manager validation must
6574 strip them before grammar.validate_tokens (which otherwise raises)."""
6575 scheduler = create_scheduler(num_speculative_tokens=4)
6576 request = _decode_ready_request(scheduler)
6577
6578 grammar = _RecordingGrammar()
6583 - request.structured_output_request = SimpleNamespace(grammar=grammar)
6584 - monkeypatch.setattr(
6585 - scheduler.structured_output_manager, "should_advance", lambda req: True
6579 + request.structured_output_request = SimpleNamespace(
6580 + grammar=grammar, reasoning_ended=True
6581 )
6582
6583 scheduler.update_draft_token_ids(
@@ -6593,16 +6588,15 @@ def test_update_draft_token_ids_strips_ngram_padding(monkeypatch):
6588 assert request.spec_token_ids == [10, 11]
6589
6590
6596 -def test_update_draft_token_ids_in_output_strips_padding(monkeypatch):
6591 +def test_update_draft_token_ids_in_output_strips_padding():
6592 """Same guard on the output path; the -1 pad-back for the rejected count
6598 - is preserved (only the input to validate_tokens is stripped)."""
6593 + is preserved (only the input to manager validation is stripped)."""
6594 scheduler = create_scheduler(num_speculative_tokens=4)
6595 request = _decode_ready_request(scheduler)
6596
6597 grammar = _RecordingGrammar()
6603 - request.structured_output_request = SimpleNamespace(grammar=grammar)
6604 - monkeypatch.setattr(
6605 - scheduler.structured_output_manager, "should_advance", lambda req: True
6598 + request.structured_output_request = SimpleNamespace(
6599 + grammar=grammar, reasoning_ended=True
6600 )
6601
6602 scheduler_output = SimpleNamespace(
tests/v1/spec_decode/test_mtp_structured_output.py
+2 -357
@@ -1,6 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""grammar_bitmask under spec-decode draft padding (#44006)."""
3 +"""Backend and engine-reasoner checks that are not manager-flow tests."""
4
5 import pytest
6 from transformers import AutoTokenizer
@@ -49,221 +49,6 @@ def _make_manager_and_request(backend: str, prompt_str: str = '{"a": "b"}'):
49 return tokenizer, manager, request, prompt
50
51
52 -@pytest.mark.parametrize("backend", ["xgrammar", "guidance"])
53 -def test_bitmask_with_padded_invalid_drafts(backend):
54 - """Bitmask handles -1 padded drafts and returns N+1 rows."""
55 - tokenizer, manager, request, prompt = _make_manager_and_request(
56 - backend, prompt_str='{"a"'
57 - )
58 - grammar = request.structured_output_request.grammar
59 -
60 - assert grammar.accept_tokens(request.request_id, prompt)
61 -
62 - valid_drafts = [tokenizer.encode(":")[0], tokenizer.encode(' "')[0]]
63 - padded = valid_drafts + [-1, -1]
64 -
65 - bitmask = manager.grammar_bitmask(
66 - requests={request.request_id: request},
67 - structured_output_request_ids=[request.request_id],
68 - scheduled_spec_decode_tokens={request.request_id: padded},
69 - )
70 -
71 - assert bitmask is not None
72 - assert bitmask.shape[0] == len(padded) + 1
73 - assert not grammar.is_terminated()
74 -
75 -
76 -@pytest.mark.parametrize("backend", ["xgrammar", "guidance"])
77 -def test_bitmask_when_grammar_terminates_mid_window(backend):
78 - """Drafts following an EOS that terminates the grammar are a no-op."""
79 - tokenizer, manager, request, prompt = _make_manager_and_request(backend)
80 - grammar = request.structured_output_request.grammar
81 -
82 - assert grammar.accept_tokens(request.request_id, prompt)
83 - eos = tokenizer.eos_token_id
84 - drafts = [eos] + [tokenizer.encode(" ")[0]] * (NUM_SPEC_TOKENS - 1)
85 -
86 - bitmask = manager.grammar_bitmask(
87 - requests={request.request_id: request},
88 - structured_output_request_ids=[request.request_id],
89 - scheduled_spec_decode_tokens={request.request_id: drafts},
90 - )
91 -
92 - assert bitmask is not None
93 - assert bitmask.shape[0] == NUM_SPEC_TOKENS + 1
94 - assert not grammar.is_terminated()
95 -
96 -
97 -@pytest.mark.parametrize("backend", ["xgrammar", "guidance"])
98 -def test_bitmask_idempotent_across_calls(backend):
99 - """Repeated calls with the same input return the same bitmask."""
100 - tokenizer, manager, request, prompt = _make_manager_and_request(
101 - backend, prompt_str='{"a"'
102 - )
103 - grammar = request.structured_output_request.grammar
104 -
105 - assert grammar.accept_tokens(request.request_id, prompt)
106 -
107 - drafts = [tokenizer.encode(":")[0], -1, -1, -1]
108 -
109 - first = manager.grammar_bitmask(
110 - requests={request.request_id: request},
111 - structured_output_request_ids=[request.request_id],
112 - scheduled_spec_decode_tokens={request.request_id: drafts},
113 - )
114 - second = manager.grammar_bitmask(
115 - requests={request.request_id: request},
116 - structured_output_request_ids=[request.request_id],
117 - scheduled_spec_decode_tokens={request.request_id: drafts},
118 - )
119 -
120 - assert first is not None and second is not None
121 - assert (first == second).all()
122 - assert not grammar.is_terminated()
123 -
124 -
125 -@pytest.mark.parametrize("backend", ["xgrammar", "guidance"])
126 -def test_bonus_position_constrained_after_invalid_drafts(backend):
127 - """Regression for #44006: bonus row stays constrained after -1 padding."""
128 - tokenizer, manager, request, prompt = _make_manager_and_request(
129 - backend, prompt_str='{"a"'
130 - )
131 - grammar = request.structured_output_request.grammar
132 -
133 - assert grammar.accept_tokens(request.request_id, prompt)
134 -
135 - valid = tokenizer.encode(":")[0]
136 - drafts = [valid, -1, -1, -1]
137 - bitmask = manager.grammar_bitmask(
138 - requests={request.request_id: request},
139 - structured_output_request_ids=[request.request_id],
140 - scheduled_spec_decode_tokens={request.request_id: drafts},
141 - )
142 - assert bitmask is not None
143 - assert bitmask.shape[0] == len(drafts) + 1
144 -
145 - assert not (bitmask[-1] == -1).all()
146 - assert not grammar.is_terminated()
147 -
148 -
149 -@pytest.mark.parametrize("backend", ["xgrammar", "guidance"])
150 -def test_bitmask_constrained_when_reasoning_ends_midwindow(backend):
151 - """Drafts after a mid-window reasoning-end marker stay constrained."""
152 - tokenizer, manager, request, prompt = _make_manager_and_request(backend)
153 - grammar = request.structured_output_request.grammar
154 -
155 - assert grammar.accept_tokens(request.request_id, prompt)
156 -
157 - marker = tokenizer.encode("\n")[0]
158 -
159 - class StubReasoner:
160 - def __init__(self, *_, **__):
161 - self.end_token_id = marker
162 -
163 - def is_reasoning_end(self, input_ids):
164 - return marker in list(input_ids)
165 -
166 - def is_reasoning_end_streaming(self, input_ids, delta_ids):
167 - return marker in list(delta_ids)
168 -
169 - manager.reasoner_cls = StubReasoner
170 - request.structured_output_request.reasoner = StubReasoner()
171 - request.structured_output_request.reasoning_ended = False
172 -
173 - pre = tokenizer.encode(" ")[0]
174 - post = tokenizer.encode(",")[0]
175 - drafts = [pre, marker, post]
176 -
177 - bitmask = manager.grammar_bitmask(
178 - requests={request.request_id: request},
179 - structured_output_request_ids=[request.request_id],
180 - scheduled_spec_decode_tokens={request.request_id: drafts},
181 - )
182 -
183 - assert bitmask is not None
184 - assert bitmask.shape[0] == len(drafts) + 1
185 - assert (bitmask[0] == -1).all()
186 - assert (bitmask[1] == -1).all()
187 - assert not (bitmask[2] == -1).all()
188 - assert not (bitmask[-1] == -1).all()
189 - assert not grammar.is_terminated()
190 -
191 -
192 -@pytest.mark.parametrize("backend", ["xgrammar", "guidance"])
193 -def test_bitmask_post_reasoning_end_drafts_skip_grammar_advance(backend, caplog):
194 - """Post-marker drafts predate the bitmask and may be grammar-invalid;
195 - grammar_bitmask must skip the grammar advance instead of asserting.
196 - """
197 - tokenizer, manager, request, prompt = _make_manager_and_request(
198 - backend, prompt_str="{"
199 - )
200 - grammar = request.structured_output_request.grammar
201 -
202 - assert grammar.accept_tokens(request.request_id, prompt)
203 - assert not grammar.is_terminated()
204 -
205 - marker = tokenizer.encode("\n")[0]
206 -
207 - class StubReasoner:
208 - def __init__(self, *_, **__):
209 - self.end_token_id = marker
210 -
211 - def is_reasoning_end(self, input_ids):
212 - return marker in list(input_ids)
213 -
214 - def is_reasoning_end_streaming(self, input_ids, delta_ids):
215 - return marker in list(delta_ids)
216 -
217 - manager.reasoner_cls = StubReasoner
218 - request.structured_output_request.reasoner = StubReasoner()
219 - request.structured_output_request.reasoning_ended = False
220 -
221 - pre = tokenizer.encode(" ")[0]
222 - # A token that the JSON grammar would reject as the first post-marker
223 - # token; without the fix grammar.accept_tokens fires the assertion.
224 - invalid_post = tokenizer.encode("z")[0]
225 - drafts = [pre, marker, invalid_post]
226 -
227 - bitmask = manager.grammar_bitmask(
228 - requests={request.request_id: request},
229 - structured_output_request_ids=[request.request_id],
230 - scheduled_spec_decode_tokens={request.request_id: drafts},
231 - )
232 -
233 - assert bitmask is not None
234 - assert bitmask.shape[0] == len(drafts) + 1
235 - # Post-marker position is still bitmask-constrained.
236 - assert not (bitmask[2] == -1).all()
237 - # Grammar must not have advanced through the unvalidated draft.
238 - assert not grammar.is_terminated()
239 - assert "Failed to advance FSM" not in caplog.text
240 -
241 -
242 -@pytest.mark.parametrize("backend", ["xgrammar", "guidance"])
243 -def test_validate_tokens_then_bitmask_round_trip(backend):
244 - """validate_tokens -> pad with -1 -> grammar_bitmask must not assert."""
245 - tokenizer, manager, request, prompt = _make_manager_and_request(backend)
246 - grammar = request.structured_output_request.grammar
247 -
248 - assert grammar.accept_tokens(request.request_id, prompt)
249 -
250 - raw_drafts = [tokenizer.encode(",")[0], 99999, 12345, 67890]
251 - accepted = grammar.validate_tokens(raw_drafts)
252 - assert len(accepted) <= len(raw_drafts)
253 -
254 - padded = accepted + [-1] * (len(raw_drafts) - len(accepted))
255 - assert len(padded) == len(raw_drafts)
256 -
257 - bitmask = manager.grammar_bitmask(
258 - requests={request.request_id: request},
259 - structured_output_request_ids=[request.request_id],
260 - scheduled_spec_decode_tokens={request.request_id: padded},
261 - )
262 - assert bitmask is not None
263 - assert bitmask.shape[0] == len(padded) + 1
264 - assert not grammar.is_terminated()
265 -
266 -
52 def test_xgrammar_accept_tokens_stops_at_termination(capfd):
53 """Tokens after a terminating EOS do not reach the matcher."""
54 tokenizer, _, request, prompt = _make_manager_and_request("xgrammar")
@@ -302,7 +87,6 @@ def test_xgrammar_validate_tokens_stops_at_termination(capfd):
87
88 assert grammar.validate_tokens([eos, trailing]) == [eos]
89 assert "trying to accept new token" not in capfd.readouterr().err
305 - # Check matcher state directly to verify validation rolled it back.
90 assert not grammar.matcher.is_terminated()
91
92 assert grammar.accept_tokens(request.request_id, [eos])
@@ -312,87 +96,6 @@ def test_xgrammar_validate_tokens_stops_at_termination(capfd):
96 assert "trying to accept new token" not in capfd.readouterr().err
97
98
315 -class _MarkerReasoner:
316 - """Stub reasoner whose reasoning-end marker is a single fixed token."""
317 -
318 - def __init__(self, marker: int):
319 - self.marker = marker
320 -
321 - def is_reasoning_end(self, input_ids):
322 - return self.marker in list(input_ids)
323 -
324 - def is_reasoning_end_streaming(self, input_ids, delta_ids):
325 - return self.marker in list(delta_ids)
326 -
327 -
328 -def _setup_boundary_request(backend: str):
329 - """Request with a structural-tag key and reasoning not yet ended."""
330 - from vllm.v1.structured_output.backend_types import StructuredOutputOptions
331 -
332 - tokenizer, manager, request, prompt = _make_manager_and_request(backend)
333 - marker = tokenizer.encode("\n")[0]
334 - structured_req = request.structured_output_request
335 - # The grammar itself is JSON (cheap to build); only the key kind matters
336 - # for the should_advance structural-tag branch, so pre-seed the cached
337 - # property.
338 - structured_req.__dict__["structured_output_key"] = (
339 - StructuredOutputOptions.STRUCTURAL_TAG,
340 - "",
341 - )
342 - manager.reasoner_cls = _MarkerReasoner
343 - structured_req.reasoner = _MarkerReasoner(marker)
344 - structured_req.reasoning_ended = False
345 - return tokenizer, manager, request, prompt, marker
346 -
347 -
348 -def test_should_advance_records_reasoning_end_index():
349 - """Regression for #44006 on post-#42452 main: the boundary step must
350 - record where reasoning ends so the scheduler can trim before advancing.
351 - """
352 - tokenizer, manager, request, prompt, marker = _setup_boundary_request("xgrammar")
353 - structured_req = request.structured_output_request
354 -
355 - pre = tokenizer.encode(" ")[0]
356 - post = tokenizer.encode("{")[0]
357 - request.append_output_token_ids([pre, marker, post])
358 -
359 - assert manager.should_advance(request)
360 - assert structured_req.reasoning_ended
361 - # Marker sits at absolute index len(prompt) + 1.
362 - assert structured_req.reasoning_end_token_index == len(prompt) + 1
363 -
364 -
365 -def test_trim_reasoning_for_advance():
366 - """trim drops the marker and everything before it; later steps and
367 - requests without a recorded boundary pass through unchanged.
368 - """
369 - tokenizer, manager, request, prompt, marker = _setup_boundary_request("xgrammar")
370 - structured_req = request.structured_output_request
371 -
372 - pre = tokenizer.encode(" ")[0]
373 - post = tokenizer.encode("{")[0]
374 -
375 - # No boundary recorded yet: pass-through.
376 - assert manager.trim_reasoning_for_advance(request, [pre]) == [pre]
377 -
378 - # Boundary step: marker mid-step keeps only the suffix.
379 - step_tokens = [pre, marker, post]
380 - request.append_output_token_ids(step_tokens)
381 - assert manager.should_advance(request)
382 - assert manager.trim_reasoning_for_advance(request, step_tokens) == [post]
383 -
384 - # Boundary step variant: marker last (the #44006 crash shape
385 - # [198, </think>]) trims to empty -> scheduler skips accept_tokens.
386 - structured_req.reasoning_end_token_index = len(request.all_token_ids) - 1
387 - assert manager.trim_reasoning_for_advance(request, step_tokens) == []
388 -
389 - # Later steps: tokens are past the boundary, returned unchanged.
390 - structured_req.reasoning_end_token_index = len(prompt) + 1
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 -
99 class _EngineReasonerStub(ParserEngineReasoningAdapter):
100 """Adapter-typed reasoner with a fixed end-token set and no real engine."""
101
@@ -409,7 +112,7 @@ class _EngineReasonerStub(ParserEngineReasoningAdapter):
112 for offset, token in enumerate(token_ids):
113 if token in self._end_token_ids:
114 return offset
412 - return None
115 + return len(token_ids)
116
117 def is_reasoning_end(self, input_ids):
118 return any(token in self._end_token_ids for token in input_ids)
@@ -449,61 +152,3 @@ def test_bitmask_engine_reasoner_ends_midwindow_with_padding(backend):
152 assert not (bitmask[2] == -1).all()
153 assert reasoner.windows == [[pre, marker, post]]
154 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 deleted
-557
@@ -1,557 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
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
10 -
11 -from vllm.config import ModelConfig, SchedulerConfig, VllmConfig
12 -from vllm.v1.request import Request
13 -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)
47 - self.is_reasoning_end_streaming = Mock(return_value=False)
48 -
49 -
50 -class TestReasoningStructuredOutput:
51 - """Test reasoning-aware structured output functionality."""
52 -
53 - @pytest.fixture
54 - def mock_model_config(self):
55 - """Create a mock ModelConfig."""
56 - config = Mock(spec=ModelConfig)
57 - config.skip_tokenizer_init = True # Skip tokenizer init to avoid network calls
58 - config.get_vocab_size = Mock(return_value=50000)
59 - # Add missing runner_type attribute that tokenizer initialization expects
60 - config.runner_type = "generate"
61 - # Add other attributes that tokenizer initialization might need
62 - config.tokenizer = "test-tokenizer"
63 - config.tokenizer_mode = "auto"
64 - config.trust_remote_code = False
65 - config.tokenizer_revision = None
66 - return config
67 -
68 - @pytest.fixture
69 - def mock_scheduler_config(self):
70 - """Create a mock SchedulerConfig."""
71 - config = Mock(spec=SchedulerConfig)
72 - config.max_num_seqs = 128
73 - return config
74 -
75 - @pytest.fixture
76 - def mock_vllm_config(self, mock_model_config, mock_scheduler_config):
77 - """Create a mock VllmConfig."""
78 - config = Mock(spec=VllmConfig)
79 - config.model_config = mock_model_config
80 - config.scheduler_config = mock_scheduler_config
81 - config.structured_outputs_config = Mock()
82 - config.structured_outputs_config.reasoning_parser = None
83 - config.structured_outputs_config.enable_in_reasoning = False
84 - config.speculative_config = None
85 - return config
86 -
87 - @pytest.fixture
88 - def mock_request_with_structured_output(self):
89 - """Create a mock request with structured output."""
90 - request = Mock(spec=Request)
91 - request.structured_output_request = Mock()
92 - request.structured_output_request.reasoning_ended = None
93 - request.structured_output_request.grammar = Mock()
94 - request.structured_output_request.reasoning_parser_kwargs = None
95 - request.structured_output_request.reasoner = None
96 - request.structured_output_request.grammar.is_terminated = Mock(
97 - return_value=False
98 - )
99 - request.use_structured_output = True
100 - request.prompt_token_ids = [1, 2, 3, 4, 5]
101 - request.all_token_ids = [1, 2, 3, 4, 5, 6, 7, 8]
102 - request.num_computed_tokens = 5
103 - request.num_output_placeholders = 0
104 - request.request_id = "mock_req"
105 - return request
106 -
107 - @pytest.fixture
108 - def manager_with_reasoner(self, mock_vllm_config):
109 - manager = StructuredOutputManager(mock_vllm_config)
110 - manager.reasoner_cls = MockReasoner
111 - manager.tokenizer = Mock()
112 - return manager
113 -
114 - def test_should_fill_bitmask_with_enable_in_reasoning(
115 - self, manager_with_reasoner, mock_request_with_structured_output
116 - ):
117 - """Test should_fill_bitmask when enable_in_reasoning is True."""
118 - manager_with_reasoner.enable_in_reasoning = True
119 -
120 - # Should always return True when enable_in_reasoning is enabled
121 - result = manager_with_reasoner.should_fill_bitmask(
122 - mock_request_with_structured_output
123 - )
124 - assert result is True
125 - assert (
126 - mock_request_with_structured_output.structured_output_request.reasoner
127 - is None
128 - )
129 -
130 - def test_should_fill_bitmask_reasoning_already_ended(
131 - self,
132 - manager_with_reasoner,
133 - mock_request_with_structured_output,
134 - ):
135 - """An active grammar does not need a request-local reasoner."""
136 - structured_req = mock_request_with_structured_output.structured_output_request
137 - structured_req.reasoning_ended = True
138 -
139 - result = manager_with_reasoner.should_fill_bitmask(
140 - mock_request_with_structured_output
141 - )
142 -
143 - assert result is True
144 - assert structured_req.reasoner is None
145 -
146 - def test_grammar_bitmask_skips_reasoner_when_already_active(
147 - self,
148 - manager_with_reasoner,
149 - mock_request_with_structured_output,
150 - ):
151 - """Bitmask generation skips reasoning-boundary detection."""
152 - manager_with_reasoner.vllm_config.num_speculative_tokens = 1
153 - manager_with_reasoner._grammar_bitmask = Mock()
154 - manager_with_reasoner._grammar_bitmask.shape = (1, 1)
155 - expected_bitmask = Mock()
156 - manager_with_reasoner._grammar_bitmask.numpy.return_value = expected_bitmask
157 - manager_with_reasoner._fill_bitmasks = Mock()
158 -
159 - structured_req = mock_request_with_structured_output.structured_output_request
160 - structured_req.reasoning_ended = True
161 - request_id = mock_request_with_structured_output.request_id
162 -
163 - bitmask = manager_with_reasoner.grammar_bitmask(
164 - requests={request_id: mock_request_with_structured_output},
165 - structured_output_request_ids=[request_id],
166 - scheduled_spec_decode_tokens={},
167 - )
168 -
169 - assert bitmask is expected_bitmask
170 - assert structured_req.reasoner is None
171 -
172 - def test_should_fill_bitmask_without_enable_in_reasoning(
173 - self,
174 - manager_with_reasoner,
175 - mock_request_with_structured_output,
176 - ):
177 - """Test should_fill_bitmask when enable_in_reasoning is False."""
178 - # Keep enable_in_reasoning as False (default)
179 - config = manager_with_reasoner.vllm_config.structured_outputs_config
180 - assert config.enable_in_reasoning is False
181 -
182 - result = manager_with_reasoner.should_fill_bitmask(
183 - mock_request_with_structured_output
184 - )
185 -
186 - # Should set reasoning_ended and return its value
187 - assert (
188 - mock_request_with_structured_output.structured_output_request.reasoning_ended
189 - is False
190 - )
191 - assert result is False
192 -
193 - def test_should_fill_bitmask_no_reasoner(
194 - self, mock_vllm_config, mock_request_with_structured_output
195 - ):
196 - """Test should_fill_bitmask when no reasoner is configured."""
197 - manager = StructuredOutputManager(mock_vllm_config)
198 -
199 - result = manager.should_fill_bitmask(mock_request_with_structured_output)
200 -
201 - # Should default to True when no reasoner
202 - assert result is True
203 -
204 - def test_should_fill_bitmask_uses_request_reasoning_parser_kwargs(
205 - self, mock_vllm_config, mock_request_with_structured_output
206 - ):
207 - """Test request-level parser kwargs override the default reasoner."""
208 -
209 - class KwargReasoner:
210 - def __init__(self, tokenizer, chat_template_kwargs=None):
211 - self.chat_template_kwargs = chat_template_kwargs or {}
212 -
213 - def is_reasoning_end(self, input_ids):
214 - return not self.chat_template_kwargs.get("enable_thinking", False)
215 -
216 - manager = StructuredOutputManager(mock_vllm_config)
217 - manager.reasoner_cls = KwargReasoner
218 - manager.tokenizer = Mock()
219 -
220 - structured_req = mock_request_with_structured_output.structured_output_request
221 - structured_req.reasoning_parser_kwargs = {
222 - "chat_template_kwargs": {"enable_thinking": True}
223 - }
224 -
225 - result = manager.should_fill_bitmask(mock_request_with_structured_output)
226 -
227 - assert result is False
228 - assert (
229 - mock_request_with_structured_output.structured_output_request.reasoner
230 - is not None
231 - )
232 -
233 - def test_should_advance_with_enable_in_reasoning(
234 - self,
235 - manager_with_reasoner,
236 - mock_request_with_structured_output,
237 - ):
238 - """Test should_advance when enable_in_reasoning is True."""
239 - # Enable enable_in_reasoning
240 - manager_with_reasoner.enable_in_reasoning = True
241 -
242 - # Should always return True when enable_in_reasoning is enabled
243 - result = manager_with_reasoner.should_advance(
244 - mock_request_with_structured_output
245 - )
246 - assert result is True
247 - assert (
248 - mock_request_with_structured_output.structured_output_request.reasoner
249 - is None
250 - )
251 -
252 - def test_should_advance_reasoning_not_ended(
253 - self,
254 - manager_with_reasoner,
255 - mock_request_with_structured_output,
256 - ):
257 - """Test should_advance when reasoning has not ended."""
258 - # Set reasoning as not ended
259 - (
260 - mock_request_with_structured_output.structured_output_request
261 - ).reasoning_ended = False
262 -
263 - result = manager_with_reasoner.should_advance(
264 - mock_request_with_structured_output
265 - )
266 -
267 - # Should return False since reasoning hasn't ended
268 - assert result is False
269 -
270 - def test_should_advance_reasoning_just_ended(
271 - self,
272 - manager_with_reasoner,
273 - mock_request_with_structured_output,
274 - ):
275 - """Test should_advance when reasoning ends in current step."""
276 - # Set reasoning as not ended initially, but ends in this step
277 - (
278 - mock_request_with_structured_output.structured_output_request
279 - ).reasoning_ended = False
280 - reasoner = MockReasoner(tokenizer=Mock())
281 - reasoner.is_reasoning_end_streaming.return_value = True
282 - structured_req = mock_request_with_structured_output.structured_output_request
283 - structured_req.reasoner = reasoner
284 -
285 - result = manager_with_reasoner.should_advance(
286 - mock_request_with_structured_output
287 - )
288 -
289 - # The scheduler trims the reasoning prefix before advancing the grammar.
290 - assert (
291 - mock_request_with_structured_output.structured_output_request.reasoning_ended
292 - is True
293 - )
294 - assert result is True
295 -
296 - def test_should_advance_reasoning_already_ended(
297 - self,
298 - manager_with_reasoner,
299 - mock_request_with_structured_output,
300 - ):
301 - """Test should_advance when reasoning has already ended."""
302 - # Set reasoning as already ended
303 - (
304 - mock_request_with_structured_output.structured_output_request
305 - ).reasoning_ended = True
306 -
307 - result = manager_with_reasoner.should_advance(
308 - mock_request_with_structured_output
309 - )
310 -
311 - # Should return True since reasoning has ended
312 - assert result is True
313 - assert (
314 - mock_request_with_structured_output.structured_output_request.reasoner
315 - is None
316 - )
317 -
318 - def test_should_advance_uses_new_token_ids_when_provided(
319 - self,
320 - manager_with_reasoner,
321 - mock_request_with_structured_output,
322 - ):
323 - """Regression for #43388: when caller passes new_token_ids, the
324 - reasoner sees the exact multi-token delta rather than the
325 - placeholder-derived window.
326 - """
327 - structured_req = mock_request_with_structured_output.structured_output_request
328 - structured_req.reasoning_ended = False
329 -
330 - end_token_id = 248069
331 -
332 - reasoner = MockReasoner(tokenizer=Mock())
333 - # Detection mirrors the real Qwen3 parser: end token in the delta.
334 - reasoner.is_reasoning_end_streaming = Mock(
335 - side_effect=lambda input_ids, delta_ids: end_token_id in list(delta_ids)
336 - )
337 - structured_req.reasoner = reasoner
338 -
339 - # Scenario from #43388: async + spec decode K=4, 4 tokens accepted
340 - # but only 1 placeholder remains (some drafts were rejected).
341 - # The placeholder math would yield delta=[271] and miss </think>.
342 - # Passing new_token_ids must override that.
343 - new_token_ids = [9, 198, end_token_id, 271]
344 - mock_request_with_structured_output.all_token_ids = [
345 - 1,
346 - 2,
347 - 3,
348 - 4,
349 - 5,
350 - ] + new_token_ids
351 - mock_request_with_structured_output.num_computed_tokens = 9
352 - mock_request_with_structured_output.num_output_placeholders = 1
353 -
354 - result = manager_with_reasoner.should_advance(
355 - mock_request_with_structured_output,
356 - new_token_ids=new_token_ids,
357 - )
358 -
359 - # First call to is_reasoning_end_streaming was with the full
360 - # new_token_ids (not the truncated placeholder window).
361 - first_call = reasoner.is_reasoning_end_streaming.call_args_list[0]
362 - _, called_delta = first_call.args
363 - assert list(called_delta) == new_token_ids
364 -
365 - assert structured_req.reasoning_ended is True
366 - assert result is True
367 -
368 - def test_should_advance_without_new_token_ids_falls_back(
369 - self,
370 - manager_with_reasoner,
371 - mock_request_with_structured_output,
372 - ):
373 - """Backward compat: callers that don't pass new_token_ids keep
374 - the original placeholder-derived delta window.
375 - """
376 - structured_req = mock_request_with_structured_output.structured_output_request
377 - structured_req.reasoning_ended = False
378 - reasoner = MockReasoner(tokenizer=Mock())
379 - reasoner.is_reasoning_end_streaming.return_value = False
380 - structured_req.reasoner = reasoner
381 -
382 - mock_request_with_structured_output.all_token_ids = [1, 2, 3, 4, 5]
383 - mock_request_with_structured_output.num_computed_tokens = 5
384 - mock_request_with_structured_output.num_output_placeholders = 2
385 -
386 - result = manager_with_reasoner.should_advance(
387 - mock_request_with_structured_output
388 - )
389 -
390 - # placeholder window: start = 5 - 2 = 3, delta = [4, 5]
391 - _, called_delta = reasoner.is_reasoning_end_streaming.call_args[0]
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,
423 - mock_request_with_structured_output,
424 - ):
425 - """JSON uses the common trim-then-advance path at the boundary."""
426 - structured_req = mock_request_with_structured_output.structured_output_request
427 - structured_req.reasoning_ended = False
428 - structured_req.structured_output_key = (
429 - StructuredOutputOptions.JSON_OBJECT,
430 - "{}",
431 - )
432 -
433 - marker = 248069
434 -
435 - class MarkerReasoner:
436 - def __init__(self, *_, **__):
437 - pass
438 -
439 - def is_reasoning_end_streaming(self, input_ids, delta_ids):
440 - return marker in list(delta_ids)
441 -
442 - structured_req.reasoner = MarkerReasoner()
443 -
444 - new_token_ids = [9, 198, marker, 271, 5005]
445 - mock_request_with_structured_output.all_token_ids = [1, 2, 3] + new_token_ids
446 -
447 - result = manager_with_reasoner.should_advance(
448 - mock_request_with_structured_output,
449 - new_token_ids=new_token_ids,
450 - )
451 -
452 - structured_req.grammar.accept_tokens.assert_not_called()
453 - assert structured_req.reasoning_ended is True
454 - assert result is True
455 - assert structured_req.reasoning_end_token_index == 5
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
tests/v1/structured_output/test_structured_output_manager.py new
+414
@@ -0,0 +1,414 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import time
5 +from dataclasses import dataclass
6 +
7 +import pytest
8 +from transformers import AutoTokenizer
9 +
10 +from vllm.config import DeviceConfig, StructuredOutputsConfig, VllmConfig
11 +from vllm.config.model import ModelConfig
12 +from vllm.config.speculative import SpeculativeConfig
13 +from vllm.sampling_params import SamplingParams, StructuredOutputsParams
14 +from vllm.v1.request import Request
15 +from vllm.v1.structured_output import StructuredOutputManager
16 +
17 +TOKENIZER = "gpt2"
18 +THINK_END = "\n" # reasoning-end marker (single GPT-2 token)
19 +EOS = "<|eos|>" # resolved to tokenizer.eos_token_id
20 +JSON_SCHEMA = '{"type": "object"}'
21 +BACKENDS = ("xgrammar", "guidance")
22 +MAX_WAIT_SECONDS = 5
23 +NUM_SPEC_TOKENS = 8
24 +
25 +
26 +@pytest.fixture(scope="module")
27 +def tokenizer():
28 + return AutoTokenizer.from_pretrained(TOKENIZER)
29 +
30 +
31 +@dataclass(frozen=True)
32 +class FlowCase:
33 + """One validate -> bitmask -> accept -> post-accept bitmask scenario."""
34 +
35 + raw_drafts: tuple[str, ...]
36 + expected_row_pattern: str
37 + expected_reasoning: bool | None
38 + expect_terminated: bool
39 + prefix: str = ""
40 + expected_validated: tuple[str, ...] | None = None
41 + reasoning_ended: bool | None = False
42 + xfail_guidance: str | None = None
43 +
44 +
45 +class MockReasoner:
46 + def __init__(self, tokenizer, marker: int | None = None):
47 + self.marker = marker
48 +
49 + def is_reasoning_end(self, input_ids):
50 + if self.marker is None:
51 + return True
52 + return self.marker in list(input_ids)
53 +
54 + def is_reasoning_end_streaming(self, input_ids, delta_ids):
55 + return self.is_reasoning_end(delta_ids)
56 +
57 +
58 +def _single_token(tokenizer, text: str) -> int:
59 + token_ids = tokenizer.encode(text)
60 + assert len(token_ids) == 1, (text, token_ids)
61 + return token_ids[0]
62 +
63 +
64 +def _to_token_ids(tokenizer, texts: tuple[str, ...]) -> list[int]:
65 + """Convert literal text (or EOS sentinel) to token IDs."""
66 + return [
67 + tokenizer.eos_token_id if t == EOS else _single_token(tokenizer, t)
68 + for t in texts
69 + ]
70 +
71 +
72 +def _wait_for_grammar(request: Request) -> None:
73 + structured_req = request.structured_output_request
74 + assert structured_req is not None
75 +
76 + deadline = time.time() + MAX_WAIT_SECONDS
77 + while not structured_req._check_grammar_completion():
78 + if time.time() > deadline:
79 + pytest.fail("Grammar compilation timed out")
80 + time.sleep(0.01)
81 +
82 +
83 +def _build_harness(
84 + tokenizer,
85 + backend: str,
86 + prefix: str = "",
87 + use_reasoner: bool = True,
88 + reasoning_ended: bool | None = None,
89 + enable_in_reasoning: bool = False,
90 + reasoning_parser_kwargs: dict | None = None,
91 +) -> tuple[StructuredOutputManager, Request]:
92 + vllm_config = VllmConfig(
93 + model_config=ModelConfig(tokenizer=TOKENIZER),
94 + device_config=DeviceConfig(device="cpu"),
95 + structured_outputs_config=StructuredOutputsConfig(
96 + backend=backend,
97 + enable_in_reasoning=enable_in_reasoning,
98 + ),
99 + speculative_config=SpeculativeConfig(
100 + model="[ngram]",
101 + num_speculative_tokens=NUM_SPEC_TOKENS,
102 + ),
103 + )
104 + manager = StructuredOutputManager(vllm_config)
105 + if use_reasoner:
106 + manager.reasoner_cls = MockReasoner
107 +
108 + sampling_params = SamplingParams(
109 + structured_outputs=StructuredOutputsParams(json=JSON_SCHEMA)
110 + )
111 + sampling_params.structured_outputs._backend = backend
112 + sampling_params.update_from_generation_config({}, tokenizer.eos_token_id)
113 +
114 + prompt_ids = tokenizer.encode(prefix) if prefix else []
115 + request = Request(
116 + request_id=f"{backend}-flow",
117 + prompt_token_ids=prompt_ids,
118 + sampling_params=sampling_params,
119 + pooling_params=None,
120 + reasoning_ended=reasoning_ended,
121 + reasoning_parser_kwargs=reasoning_parser_kwargs,
122 + )
123 + manager.grammar_init(request)
124 + _wait_for_grammar(request)
125 +
126 + structured_req = request.structured_output_request
127 + assert request.prompt_token_ids is not None
128 + assert structured_req is not None
129 + assert structured_req.grammar is not None
130 + if prompt_ids:
131 + assert structured_req.grammar.accept_tokens(
132 + request.request_id,
133 + prompt_ids,
134 + )
135 + return manager, request
136 +
137 +
138 +def _row_pattern(bitmask) -> str:
139 + return "".join("U" if (row == -1).all() else "C" for row in bitmask)
140 +
141 +
142 +def _run_real_flow(
143 + manager: StructuredOutputManager,
144 + request: Request,
145 + raw_drafts: list[int],
146 + expected_validated: list[int],
147 + expected_row_pattern: str,
148 + expected_reasoning: bool | None,
149 + expect_terminated: bool,
150 +) -> None:
151 + assert len(expected_validated) <= len(raw_drafts)
152 +
153 + validated = manager.validate_tokens(request, list(raw_drafts))
154 + assert validated == expected_validated
155 +
156 + padded = validated + [-1] * (len(raw_drafts) - len(validated))
157 + bitmask = manager.grammar_bitmask(
158 + requests={request.request_id: request},
159 + structured_output_request_ids=[request.request_id],
160 + scheduled_spec_decode_tokens={request.request_id: padded},
161 + )
162 + assert bitmask is not None
163 + assert bitmask.shape[0] == len(raw_drafts) + 1
164 + assert _row_pattern(bitmask) == expected_row_pattern
165 +
166 + # grammar_bitmask() must rollback any speculative state it advanced.
167 + assert manager.validate_tokens(request, list(raw_drafts)) == expected_validated
168 +
169 + structured_req = request.structured_output_request
170 + assert structured_req is not None
171 + grammar = structured_req.grammar
172 + assert grammar is not None
173 + assert not grammar.is_terminated()
174 + # `bitmask[i]` is the grammar state after the first `i` scheduled tokens.
175 + # These tests commit the validated prefix, so the matching post-accept
176 + # state is already present in the scheduled bitmask.
177 + expected_post_accept_row = bitmask[len(expected_validated)].copy()
178 +
179 + # Mirror the scheduler flow: sampled tokens are appended before
180 + # StructuredOutputManager.accept_tokens() trims them to the
181 + # grammar-constrained suffix.
182 + request.append_output_token_ids(list(expected_validated))
183 + assert manager.accept_tokens(request, list(expected_validated))
184 +
185 + post_accept_bitmask = manager.grammar_bitmask(
186 + requests={request.request_id: request},
187 + structured_output_request_ids=[request.request_id],
188 + scheduled_spec_decode_tokens={},
189 + )
190 + assert post_accept_bitmask is not None
191 + assert post_accept_bitmask.shape[0] == 1
192 + assert (post_accept_bitmask[0] == expected_post_accept_row).all()
193 + assert structured_req.reasoning_ended == expected_reasoning
194 + assert grammar.is_terminated() is expect_terminated
195 +
196 +
197 +FLOW_CASES = [
198 + pytest.param(
199 + FlowCase(
200 + raw_drafts=(" ", "z", " "),
201 + expected_row_pattern="UUUU",
202 + expected_reasoning=False,
203 + expect_terminated=False,
204 + ),
205 + id="inactive",
206 + ),
207 + pytest.param(
208 + FlowCase(
209 + prefix='{"a"',
210 + raw_drafts=(":", ' "', "b", '"}'),
211 + expected_row_pattern="CCCCC",
212 + expected_reasoning=True,
213 + expect_terminated=False,
214 + reasoning_ended=True,
215 + ),
216 + id="active_all_valid",
217 + ),
218 + pytest.param(
219 + FlowCase(
220 + prefix='{"a"',
221 + raw_drafts=("z", ":", ' "', "b"),
222 + expected_validated=(),
223 + expected_row_pattern="CUUUU",
224 + expected_reasoning=True,
225 + expect_terminated=False,
226 + reasoning_ended=True,
227 + ),
228 + id="active_first_token_invalid",
229 + ),
230 + pytest.param(
231 + FlowCase(
232 + prefix='{"a"',
233 + raw_drafts=(":", "z", ' "', "b"),
234 + expected_validated=(":",),
235 + expected_row_pattern="CCUUU",
236 + expected_reasoning=True,
237 + expect_terminated=False,
238 + reasoning_ended=True,
239 + ),
240 + id="active_later_token_invalid",
241 + ),
242 + pytest.param(
243 + FlowCase(
244 + prefix='{"a": "b"}',
245 + raw_drafts=(EOS, " ", " "),
246 + expected_validated=(EOS,),
247 + expected_row_pattern="CUUU",
248 + expected_reasoning=True,
249 + expect_terminated=True,
250 + reasoning_ended=True,
251 + xfail_guidance=(
252 + "guidance validate_tokens rejects EOS on an already-complete "
253 + "object even though accept_tokens would accept it"
254 + ),
255 + ),
256 + id="active_first_token_terminates",
257 + ),
258 + pytest.param(
259 + FlowCase(
260 + prefix='{"a"',
261 + raw_drafts=(":", ' "', "b", '"}', EOS, " "),
262 + expected_validated=(":", ' "', "b", '"}', EOS),
263 + expected_row_pattern="CCCCCUU",
264 + expected_reasoning=True,
265 + expect_terminated=True,
266 + reasoning_ended=True,
267 + ),
268 + id="active_later_token_terminates",
269 + ),
270 + pytest.param(
271 + FlowCase(
272 + raw_drafts=(THINK_END, "{", ' "', "b"),
273 + expected_row_pattern="UCCCC",
274 + expected_reasoning=True,
275 + expect_terminated=False,
276 + ),
277 + id="becomes_active_on_first_token",
278 + ),
279 + pytest.param(
280 + FlowCase(
281 + raw_drafts=(" ", THINK_END, "{", ' "'),
282 + expected_row_pattern="UUCCC",
283 + expected_reasoning=True,
284 + expect_terminated=False,
285 + ),
286 + id="becomes_active_on_middle_token",
287 + ),
288 + pytest.param(
289 + FlowCase(
290 + raw_drafts=(" ", " ", THINK_END),
291 + expected_row_pattern="UUUC",
292 + expected_reasoning=True,
293 + expect_terminated=False,
294 + ),
295 + id="becomes_active_on_last_token",
296 + ),
297 + pytest.param(
298 + FlowCase(
299 + raw_drafts=(" ", THINK_END, "z", "{"),
300 + expected_validated=(" ", THINK_END),
301 + expected_row_pattern="UUCUU",
302 + expected_reasoning=True,
303 + expect_terminated=False,
304 + ),
305 + id="becomes_active_first_token_invalid",
306 + ),
307 + pytest.param(
308 + FlowCase(
309 + raw_drafts=(" ", THINK_END, "{", "z"),
310 + expected_validated=(" ", THINK_END, "{"),
311 + expected_row_pattern="UUCCU",
312 + expected_reasoning=True,
313 + expect_terminated=False,
314 + ),
315 + id="becomes_active_later_token_invalid",
316 + ),
317 + pytest.param(
318 + FlowCase(
319 + raw_drafts=(" ", THINK_END, "{", "}", EOS, " "),
320 + expected_validated=(" ", THINK_END, "{", "}", EOS),
321 + expected_row_pattern="UUCCCUU",
322 + expected_reasoning=True,
323 + expect_terminated=True,
324 + ),
325 + id="becomes_active_terminates",
326 + ),
327 +]
328 +
329 +
330 +@pytest.mark.parametrize("backend", BACKENDS)
331 +@pytest.mark.parametrize("case", FLOW_CASES)
332 +def test_real_flow(
333 + tokenizer,
334 + backend: str,
335 + case: FlowCase,
336 +):
337 + if backend == "guidance" and case.xfail_guidance:
338 + pytest.xfail(case.xfail_guidance)
339 +
340 + reasoner_kwargs = {"marker": _single_token(tokenizer, THINK_END)}
341 + manager, request = _build_harness(
342 + tokenizer,
343 + backend,
344 + prefix=case.prefix,
345 + reasoning_ended=case.reasoning_ended,
346 + reasoning_parser_kwargs=reasoner_kwargs,
347 + )
348 +
349 + raw_drafts = _to_token_ids(tokenizer, case.raw_drafts)
350 + expected_texts = (
351 + case.raw_drafts if case.expected_validated is None else case.expected_validated
352 + )
353 + expected_validated = _to_token_ids(tokenizer, expected_texts)
354 +
355 + _run_real_flow(
356 + manager,
357 + request,
358 + raw_drafts=raw_drafts,
359 + expected_validated=expected_validated,
360 + expected_row_pattern=case.expected_row_pattern,
361 + expected_reasoning=case.expected_reasoning,
362 + expect_terminated=case.expect_terminated,
363 + )
364 +
365 +
366 +@pytest.mark.parametrize("backend", BACKENDS)
367 +@pytest.mark.parametrize(
368 + ("use_reasoner", "reasoning_ended", "enable_in_reasoning"),
369 + [
370 + pytest.param(True, False, True, id="enable_in_reasoning"),
371 + pytest.param(True, True, False, id="reasoning_ended"),
372 + pytest.param(False, False, False, id="no_reasoner"),
373 + pytest.param(True, None, False, id="inferred"),
374 + ],
375 +)
376 +def test_initial_constraint_activation(
377 + tokenizer,
378 + backend: str,
379 + use_reasoner: bool,
380 + reasoning_ended: bool | None,
381 + enable_in_reasoning: bool,
382 +):
383 + manager, request = _build_harness(
384 + tokenizer,
385 + backend,
386 + use_reasoner=use_reasoner,
387 + reasoning_ended=reasoning_ended,
388 + enable_in_reasoning=enable_in_reasoning,
389 + reasoning_parser_kwargs={
390 + "marker": _single_token(tokenizer, THINK_END)
391 + if reasoning_ended is not None
392 + else None
393 + },
394 + )
395 +
396 + # "{" is valid JSON start; "z" is not. Truncation proves grammar is active.
397 + open_brace = _single_token(tokenizer, "{")
398 + z = _single_token(tokenizer, "z")
399 + assert manager.validate_tokens(request, [open_brace, z]) == [open_brace]
400 +
401 + bitmask = manager.grammar_bitmask(
402 + requests={request.request_id: request},
403 + structured_output_request_ids=[request.request_id],
404 + scheduled_spec_decode_tokens={},
405 + )
406 + assert bitmask is not None
407 + request.append_output_token_ids([open_brace])
408 + assert manager.accept_tokens(request, [open_brace])
409 +
410 + # Short-circuit paths must not construct a request-local reasoner (#52573).
411 + if enable_in_reasoning or reasoning_ended is True:
412 + structured_req = request.structured_output_request
413 + assert structured_req is not None
414 + assert structured_req.reasoner is None
vllm/parser/abstract_parser.py
+6 -2
@@ -197,6 +197,8 @@ class Parser:
197 Used by structured engines like `xgrammar` to check if the
198 reasoning content ends in the model output.
199
200 + Must be a pure function of the input_ids.
201 +
202 Args:
203 input_ids: The token IDs of the model output.
204
@@ -210,6 +212,8 @@ class Parser:
212 """
213 Check if the reasoning content ends during a decode step.
214
215 + Must be a pure function of the input_ids and delta_ids.
216 +
217 Args:
218 input_ids: The entire model output token IDs.
219 delta_ids: The last few computed tokens at the current decode step.
@@ -228,8 +232,8 @@ class Parser:
232
233 Returns:
234 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
235 + ``len(token_ids)`` when reasoning does not end in this window. Parsers
236 + that cannot answer from a window alone always return ``None``; callers
237 fall back to :meth:`is_reasoning_end_streaming`.
238 """
239 return None
vllm/parser/engine/parser_engine.py
+1 -1
@@ -645,7 +645,7 @@ class ParserEngine(Parser):
645 for offset, token_id in enumerate(token_ids):
646 if token_id in end_ids:
647 return offset
648 - return None
648 + return len(token_ids)
649
650 def is_reasoning_end(self, input_ids: list[int]) -> bool:
651 config = self.parser_engine_config
vllm/v1/core/sched/scheduler.py
+17 -36
@@ -70,8 +70,7 @@ from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput
70 from vllm.v1.request import Request, RequestStatus, StreamingUpdate
71 from vllm.v1.spec_decode.dynamic.utils import build_dynamic_sd_schedule_lookup
72 from vllm.v1.spec_decode.metrics import SpecDecodingStats
73 -from vllm.v1.structured_output import StructuredOutputGrammar, StructuredOutputManager
74 -from vllm.v1.structured_output.utils import strip_speculative_padding
73 +from vllm.v1.structured_output import StructuredOutputManager
74 from vllm.v1.utils import record_function_or_nullcontext
75
76 logger = init_logger(__name__)
@@ -2052,33 +2051,18 @@ class Scheduler(SchedulerInterface):
2051 request.status = RequestStatus.FINISHED_STOPPED
2052 stopped = True
2053
2055 - if new_token_ids and self.structured_output_manager.should_advance(
2056 - request, new_token_ids=new_token_ids
2054 + if new_token_ids and not self.structured_output_manager.accept_tokens(
2055 + request, new_token_ids
2056 ):
2058 - struct_output_request = request.structured_output_request
2059 - assert struct_output_request is not None
2060 - grammar = struct_output_request.grammar
2061 - assert isinstance(grammar, StructuredOutputGrammar)
2062 - # new_token_ids can be a mixed block of reasoning content, then
2063 - # the reasoning end marker, then the start of the grammar content.
2064 - # Trim the reasoning content so the grammar only sees grammar content.
2065 - advance_token_ids = (
2066 - self.structured_output_manager.trim_reasoning_for_advance(
2067 - request, new_token_ids
2068 - )
2057 + logger.error(
2058 + "Unexpected: grammar rejected tokens %s for request %s. "
2059 + "Terminating request.",
2060 + new_token_ids,
2061 + req_id,
2062 )
2070 - if advance_token_ids and not grammar.accept_tokens(
2071 - req_id, advance_token_ids
2072 - ):
2073 - logger.error(
2074 - "Unexpected: grammar rejected tokens %s for request %s. "
2075 - "Terminating request.",
2076 - advance_token_ids,
2077 - req_id,
2078 - )
2079 - request.status = RequestStatus.FINISHED_ERROR
2080 - request.resumable = False
2081 - stopped = True
2063 + request.status = RequestStatus.FINISHED_ERROR
2064 + request.resumable = False
2065 + stopped = True
2066
2067 routed_experts = None
2068 if (
@@ -2434,11 +2418,9 @@ class Scheduler(SchedulerInterface):
2418 continue
2419
2420 # Add newly generated spec token ids to the request.
2437 - if self.structured_output_manager.should_advance(request):
2438 - metadata = request.structured_output_request
2439 - spec_token_ids = strip_speculative_padding(spec_token_ids)
2440 - spec_token_ids = metadata.grammar.validate_tokens(spec_token_ids) # type: ignore[union-attr]
2441 - request.spec_token_ids = spec_token_ids
2421 + request.spec_token_ids = self.structured_output_manager.validate_tokens(
2422 + request, spec_token_ids
2423 + )
2424
2425 def update_draft_token_ids_in_output(
2426 self, draft_token_ids: DraftTokenIds, scheduler_output: SchedulerOutput
@@ -2464,10 +2446,9 @@ class Scheduler(SchedulerInterface):
2446 # (needed for chunked prefill case for example).
2447 del spec_token_ids[orig_num_spec_tokens:]
2448 # Filter out spec tokens which do not adhere to the grammar.
2467 - if self.structured_output_manager.should_advance(request):
2468 - metadata = request.structured_output_request
2469 - spec_token_ids = strip_speculative_padding(spec_token_ids)
2470 - spec_token_ids = metadata.grammar.validate_tokens(spec_token_ids) # type: ignore[union-attr]
2449 + spec_token_ids = self.structured_output_manager.validate_tokens(
2450 + request, spec_token_ids
2451 + )
2452 # Pad to original number of spec tokens.
2453 num_invalid_tokens = orig_num_spec_tokens - len(spec_token_ids)
2454 if num_invalid_tokens:
vllm/v1/structured_output/__init__.py
+140 -207
@@ -1,6 +1,5 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -import itertools
3 import multiprocessing
4 from collections.abc import Iterable, Sequence
5 from concurrent.futures import Future, ThreadPoolExecutor
@@ -18,6 +17,7 @@ from vllm.v1.structured_output.backend_types import (
17 StructuredOutputGrammar,
18 )
19 from vllm.v1.structured_output.backend_xgrammar import XgrammarBackend
20 +from vllm.v1.structured_output.utils import strip_speculative_padding
21
22 if TYPE_CHECKING:
23 import numpy as np
@@ -217,6 +217,100 @@ class StructuredOutputManager:
217 ) -> Future:
218 return self.executor_for_fillmask.submit(self._fill_bitmasks, batch)
219
220 + def _get_constraint_start(
221 + self,
222 + request: "Request",
223 + spec_tokens: Sequence[int],
224 + spec_tokens_committed: bool = False,
225 + ) -> int:
226 + """Return the index into `spec_tokens` where tokens should start being
227 + constrained by the grammar.
228 + Assumes `spec_tokens` are stripped of -1 padding tokens.
229 + Returns `len(spec_tokens) + 1` if no token should be constrained after
230 + accepting all `spec_tokens`.
231 +
232 + `spec_tokens_committed` is True if spec_tokens is already in
233 + `request.all_token_ids`.
234 + """
235 + if self.enable_in_reasoning:
236 + return 0
237 +
238 + structured_req = request.structured_output_request
239 + assert structured_req is not None
240 + if structured_req.reasoning_ended:
241 + return 0
242 +
243 + reasoner = self._get_reasoner(request)
244 + if reasoner is None:
245 + return 0
246 +
247 + if structured_req.reasoning_ended is None:
248 + # This should be removed here, but since `openai_gptoss`
249 + # is an independent code path, it is kept for now.
250 + # After unifying the `openai_gptoss` and non-`openai_gptoss` styles,
251 + # it can be removed.
252 + if reasoner.is_reasoning_end(request.prompt_token_ids or []):
253 + return 0
254 + structured_req.reasoning_ended = False
255 +
256 + num_spec_tokens = len(spec_tokens)
257 + if num_spec_tokens <= 0:
258 + return num_spec_tokens + 1
259 +
260 + # Use `find_reasoning_end_offset` to find constraint start if supported
261 + if (
262 + isinstance(reasoner, ParserEngineReasoningAdapter)
263 + and reasoner.reasoning_end_token_ids
264 + ):
265 + offset = reasoner.find_reasoning_end_offset(spec_tokens)
266 + if offset is not None:
267 + return offset + 1
268 +
269 + # Fallback to `find_reasoning_end_offset`
270 + # TODO: Build a read-only Sequence view over all_token_ids
271 + # instead of copying the entire all_token_ids into input_ids.
272 + if spec_tokens_committed:
273 + if not spec_tokens or not reasoner.is_reasoning_end_streaming(
274 + request.all_token_ids, spec_tokens
275 + ):
276 + return num_spec_tokens + 1
277 +
278 + input_ids = request.all_token_ids.copy()
279 + delta_ids = list(spec_tokens)
280 + else:
281 + input_ids = request.all_token_ids.copy()
282 + input_ids.extend(spec_tokens)
283 + if not reasoner.is_reasoning_end_streaming(input_ids, spec_tokens):
284 + return num_spec_tokens + 1
285 + delta_ids = list(spec_tokens)
286 +
287 + for i in range(num_spec_tokens - 1, 0, -1):
288 + input_ids.pop()
289 + delta_ids.pop()
290 + if not reasoner.is_reasoning_end_streaming(input_ids, delta_ids):
291 + return i + 1
292 + return 1
293 +
294 + def validate_tokens(self, request: "Request", spec_tokens: list[int]) -> list[int]:
295 + """Return the longest unconstrained or grammar-valid prefix of `spec_tokens`."""
296 + if not request.use_structured_output:
297 + return spec_tokens
298 +
299 + spec_tokens = strip_speculative_padding(spec_tokens)
300 + constraint_start = self._get_constraint_start(request, spec_tokens)
301 + if constraint_start >= len(spec_tokens):
302 + return spec_tokens
303 +
304 + structured_req = request.structured_output_request
305 + if TYPE_CHECKING:
306 + assert structured_req is not None
307 + grammar = structured_req.grammar
308 + if TYPE_CHECKING:
309 + assert isinstance(grammar, StructuredOutputGrammar)
310 + prefix = spec_tokens[:constraint_start]
311 + validated = grammar.validate_tokens(spec_tokens[constraint_start:])
312 + return prefix + validated
313 +
314 def grammar_bitmask(
315 self,
316 requests: dict[str, "Request"],
@@ -264,7 +358,7 @@ class StructuredOutputManager:
358 if TYPE_CHECKING:
359 assert isinstance(grammar, StructuredOutputGrammar)
360
267 - apply_bitmask = self.should_fill_bitmask(request)
361 + apply_bitmask = self._get_constraint_start(request, ()) == 0
362 batch.append((grammar, cumulative_index, apply_bitmask))
363 if len(batch) == self.fill_bitmask_parallel_batch_size:
364 promises.append(self._async_submit_fill_bitmask(batch))
@@ -288,76 +382,47 @@ class StructuredOutputManager:
382 grammar = structured_output_request.grammar
383 if TYPE_CHECKING:
384 assert isinstance(grammar, StructuredOutputGrammar)
291 - apply_bitmask = self.should_fill_bitmask(request)
292 -
293 - reasoner = None if apply_bitmask else self._get_reasoner(request)
385
386 + req_tokens = scheduled_spec_decode_tokens.get(req_id, list())
387 + constraint_start = self._get_constraint_start(
388 + request, strip_speculative_padding(req_tokens)
389 + )
390 state_advancements = 0
296 - post_reasoning_end_in_window = False
297 - req_tokens = scheduled_spec_decode_tokens.get(req_id, ())
298 - # Locate the reasoning-end marker in the draft window once.
299 - # Rejected-draft padding (-1) is trailing, so the window is
300 - # cut there and the reasoner never sees a placeholder.
301 - window_start = len(request.all_token_ids)
302 - reasoning_end_index: int | None = None
303 - if reasoner is not None and req_tokens:
304 - window = req_tokens
305 - if -1 in window:
306 - window = window[: window.index(-1)]
307 - reasoning_end_index = self._find_reasoning_end_index(
308 - reasoner,
309 - request.all_token_ids,
310 - window_start,
311 - window,
312 - delta_appended=False,
313 - )
391 + seen_padding = False
392 + failed = False
393 for i, token in enumerate(req_tokens):
394 + apply_bitmask = (
395 + not failed and not seen_padding and i >= constraint_start
396 + )
397 self._fill_bitmasks(((grammar, cumulative_index, apply_bitmask),))
316 - advance_grammar = apply_bitmask
398 if token == -1:
318 - apply_bitmask = False
319 - advance_grammar = False
320 - elif not apply_bitmask and window_start + i == reasoning_end_index:
321 - # Reasoning ended mid-window. Constrain the rest
322 - # of the window via bitmask. Skip grammar advance
323 - # through the marker (it is reasoning content);
324 - # try to advance through subsequent drafts so the
325 - # next bitmask row reflects the post-advance state,
326 - # but tolerate rejection since those drafts predate
327 - # the bitmask and are not guaranteed valid.
328 - apply_bitmask = True
329 - advance_grammar = False
330 - post_reasoning_end_in_window = True
331 - if advance_grammar and not grammar.is_terminated():
332 - if post_reasoning_end_in_window:
333 - accepted = bool(grammar.validate_tokens([token]))
334 - if accepted:
335 - accepted = grammar.accept_tokens(req_id, [token])
336 - else:
337 - accepted = grammar.accept_tokens(req_id, [token])
338 - if accepted:
399 + seen_padding = True
400 + elif apply_bitmask:
401 + if not grammar.is_terminated() and grammar.accept_tokens(
402 + req_id, [token]
403 + ):
404 state_advancements += 1
340 - elif not post_reasoning_end_in_window:
341 - raise AssertionError(
342 - (token, req_id, scheduled_spec_decode_tokens)
405 + else:
406 + failed = True
407 + logger.error(
408 + "Unexpected: grammar terminated or rejected draft "
409 + "token %s for request %s during bitmask fill.",
410 + token,
411 + req_id,
412 )
413 cumulative_index += 1
414 +
415 # Diffusion LLMs don't sample a bonus token after the
416 # scheduled positions, so skip its bitmask in that case.
417 if not (self.vllm_config.model_config.is_diffusion and req_tokens):
348 - # bonus_apply must be True when the bonus-row position
349 - # should be grammar-constrained. Two triggers:
350 - # - should_fill_bitmask(request): reasoning was already
351 - # over at step start (or no reasoner /
352 - # enable_in_reasoning).
353 - # - apply_bitmask: reasoning ended mid-window in this
354 - # call and was flipped True after the marker;
355 - # should_fill_bitmask still returns False here because
356 - # reasoning_ended is only persisted later by
357 - # should_advance.
358 - bonus_apply = self.should_fill_bitmask(request) or apply_bitmask
418 + bonus_apply = (
419 + not failed
420 + and not seen_padding
421 + and constraint_start <= len(req_tokens)
422 + )
423 self._fill_bitmasks(((grammar, cumulative_index, bonus_apply),))
424 cumulative_index += 1
425 +
426 if state_advancements > 0:
427 grammar.rollback(state_advancements)
428
@@ -370,161 +435,29 @@ class StructuredOutputManager:
435 # and deserialization when sending this to the GPU workers.
436 return bitmask_tensor.numpy()
437
373 - def should_fill_bitmask(self, request: "Request") -> bool:
374 - # NOTE (Hanchen) if enable_in_reasoning is True, it means that
375 - # the model needs to be constrained in reasoning. So we should always
376 - # enable the bitmask filling.
377 - structured_req = request.structured_output_request
378 - if self.enable_in_reasoning or (
379 - structured_req is not None and structured_req.reasoning_ended is True
380 - ):
438 + def accept_tokens(self, request: "Request", new_token_ids: list[int]) -> bool:
439 + """Advance grammar with accepted tokens. Returns False on rejection."""
440 + if not request.use_structured_output or not new_token_ids:
441 return True
442
383 - reasoner = self._get_reasoner(request)
384 - if reasoner is not None:
443 + structured_req = request.structured_output_request
444 + if TYPE_CHECKING:
445 assert structured_req is not None
386 - if structured_req.reasoning_ended is None:
387 - # This should be removed here, but since `openai_gptoss`
388 - # is an independent code path, it is kept for now.
389 - # After unifying the `openai_gptoss` and non-`openai_gptoss` styles,
390 - # it can be removed.
391 - structured_req.reasoning_ended = reasoner.is_reasoning_end(
392 - request.prompt_token_ids or []
393 - )
394 - return structured_req.reasoning_ended
395 - return True
396 -
397 - def should_advance(
398 - self,
399 - request: "Request",
400 - new_token_ids: list[int] | None = None,
401 - ) -> bool:
402 - if not request.use_structured_output:
403 - return False
404 -
405 - # To determine whether we can advance the FSM.
406 - # Supports thinking usage where we skip the reasoning components.
446 + grammar = structured_req.grammar
447 if TYPE_CHECKING:
408 - assert request.structured_output_request is not None
409 - assert request.structured_output_request.grammar is not None
410 - structured_req = request.structured_output_request
411 - if self.enable_in_reasoning or (
412 - structured_req is not None and structured_req.reasoning_ended is True
413 - ):
414 - return True
415 -
416 - # by default, we should always advance
417 - # for cases that don't use thinking mode.
418 - reasoner = self._get_reasoner(request)
419 - if reasoner is None:
420 - return True
448 + assert isinstance(grammar, StructuredOutputGrammar)
449
422 - assert structured_req is not None
423 -
424 - # Check if reasoning ends in *this* step.
425 - # When the caller passes new_token_ids (the tokens that were just
426 - # appended this step), use it directly as the delta window. The
427 - # placeholder-derived fallback assumes num_output_placeholders ==
428 - # len(new_token_ids), which breaks under async scheduling + spec
429 - # decode when some drafts are rejected (#43388): the placeholder
430 - # count remains > 0 after the step and the computed delta window
431 - # starts past the reasoning-end marker.
432 - all_token_ids = request.all_token_ids
433 - if new_token_ids:
434 - # The tokens were already appended this step, so the step window
435 - # starts exactly len(new_token_ids) from the end.
436 - start = len(all_token_ids) - len(new_token_ids)
437 - delta_ids: Sequence[int] = new_token_ids
438 - else:
439 - # Draft-validation callers arrive here after update_from_output
440 - # has already persisted any boundary found this step, so this
441 - # window only needs to cover tokens not examined yet.
442 - delta_from = request.num_computed_tokens - request.num_output_placeholders
443 - start = (
444 - delta_from
445 - if delta_from >= 0
446 - else max(len(all_token_ids) + delta_from, 0)
447 - )
448 - delta_ids = all_token_ids[start:]
449 - end_index = self._find_reasoning_end_index(
450 - reasoner, all_token_ids, start, delta_ids, delta_appended=True
450 + constraint_start = self._get_constraint_start(
451 + request, new_token_ids, spec_tokens_committed=True
452 )
452 - if end_index is not None:
453 - structured_req.reasoning_ended = True
454 - # Record the boundary so the scheduler can exclude reasoning tokens.
455 - structured_req.reasoning_end_token_index = end_index
453 + # Early return only when the constraint hasn't started.
454 + # Otherwise, latch structured_req.reasoning_ended.
455 + if constraint_start > len(new_token_ids):
456 return True
457 -
458 - return False
459 -
460 - @staticmethod
461 - def _find_reasoning_end_index(
462 - reasoner: "ReasoningParser",
463 - all_token_ids: Sequence[int],
464 - start: int,
465 - delta_ids: Sequence[int],
466 - *,
467 - delta_appended: bool,
468 - ) -> int | None:
469 - """Locates the token that ends reasoning within ``delta_ids``.
470 -
471 - ``delta_ids`` occupy positions ``start..``; ``delta_appended`` says
472 - whether they are already in ``all_token_ids`` (accepted tokens) or
473 - not (speculative drafts).
474 -
475 - Legacy parsers are probed one token at a time. Accepted tokens get a
476 - whole-delta check as well, and a marker seen only that way reports
477 - the final token. Drafts skip it: some predicates are order-sensitive
478 - over a multi-token window and would miss the marker.
479 -
480 - Returns:
481 - Absolute index of the last reasoning token, or ``None``.
482 - """
483 - if (
484 - isinstance(reasoner, ParserEngineReasoningAdapter)
485 - and reasoner.reasoning_end_token_ids
486 - ):
487 - offset = reasoner.find_reasoning_end_offset(delta_ids)
488 - return None if offset is None else start + offset
489 -
490 - if delta_appended and not reasoner.is_reasoning_end_streaming(
491 - all_token_ids, delta_ids
492 - ):
493 - return None
494 - prefix = list(itertools.islice(all_token_ids, start))
495 - for offset, token in enumerate(delta_ids):
496 - prefix.append(token)
497 - if reasoner.is_reasoning_end_streaming(prefix, [token]):
498 - return start + offset
499 - if delta_appended:
500 - return len(all_token_ids) - 1
501 - return None
502 -
503 - def trim_reasoning_for_advance(
504 - self, request: "Request", new_token_ids: list[int]
505 - ) -> list[int]:
506 - """Drops reasoning content from tokens about to advance the grammar.
507 -
508 - When reasoning ends mid-step (see should_advance), the step's output
509 - still contains reasoning tokens up to and including the end marker.
510 - Those are not grammar content: feeding them to accept_tokens makes
511 - the grammar reject the marker and kills the request (#44006).
512 -
513 - Returns:
514 - The suffix of ``new_token_ids`` that follows the reasoning-end
515 - marker. Steps fully after the boundary are returned unchanged.
516 - """
517 - structured_req = request.structured_output_request
518 - if structured_req is None:
519 - return new_token_ids
520 - end_idx = structured_req.reasoning_end_token_index
521 - if end_idx is None:
522 - return new_token_ids
523 - first_idx = len(request.all_token_ids) - len(new_token_ids)
524 - num_reasoning = end_idx + 1 - first_idx
525 - if num_reasoning <= 0:
526 - return new_token_ids
527 - return new_token_ids[num_reasoning:]
457 + structured_req.reasoning_ended = True
458 + return grammar.accept_tokens(
459 + request.request_id, new_token_ids[constraint_start:]
460 + )
461
462 def clear_backend(self) -> None:
463 if self.backend is not None:
vllm/v1/structured_output/request.py
-6
@@ -25,12 +25,6 @@ class StructuredOutputRequest:
25 Future[StructuredOutputGrammar] | StructuredOutputGrammar | Exception | None
26 ) = None
27 reasoning_ended: bool | None = None
28 - # Absolute index into the request's all_token_ids of the last reasoning
29 - # token (the reasoning-end marker). Tokens at or before this index are
30 - # reasoning content and must never be fed to the grammar. Only set when
31 - # reasoning ends in a step whose tokens the scheduler advances immediately
32 - # (structural tags + speculative decoding, see #42452).
33 - reasoning_end_token_index: int | None = None
28 reasoning_parser_kwargs: dict[str, Any] | None = None
29 # Cached per request; do not share reasoning parsers across requests because
30 # their behavior can depend on reasoning_parser_kwargs.