From e6960af33b379d502f409e3e2241bbf2b2c2f68d Mon Sep 17 00:00:00 2001 From: yzong-rh Date: Tue, 15 Sep 2026 09:48:49 -0400 Subject: [PATCH] [Refactor] StructuredOutputManager x Speculative Decoding Refactor (#48200) --- tests/parser/engine/test_deepseek_v4.py | 2 +- tests/parser/engine/test_parser_engine.py | 4 +- tests/v1/core/test_async_scheduler.py | 6 +- tests/v1/core/test_scheduler.py | 28 +- .../spec_decode/test_mtp_structured_output.py | 359 +---------- .../test_reasoning_structured_output.py | 557 ------------------ .../test_structured_output_manager.py | 414 +++++++++++++ vllm/parser/abstract_parser.py | 8 +- vllm/parser/engine/parser_engine.py | 2 +- vllm/v1/core/sched/scheduler.py | 53 +- vllm/v1/structured_output/__init__.py | 349 +++++------ vllm/v1/structured_output/request.py | 6 - 12 files changed, 596 insertions(+), 1192 deletions(-) delete mode 100644 tests/v1/structured_output/test_reasoning_structured_output.py create mode 100644 tests/v1/structured_output/test_structured_output_manager.py diff --git a/tests/parser/engine/test_deepseek_v4.py b/tests/parser/engine/test_deepseek_v4.py index 1c1029b51a..7cbf9791a7 100644 --- a/tests/parser/engine/test_deepseek_v4.py +++ b/tests/parser/engine/test_deepseek_v4.py @@ -574,7 +574,7 @@ class TestImplicitReasoningEnd: # the text lexer. assert thinking_parser.reasoning_end_token_ids == {_THINK_END_ID} assert thinking_parser.find_reasoning_end_offset([7, _THINK_END_ID, 8]) == 1 - assert thinking_parser.find_reasoning_end_offset([7, 8]) is None + assert thinking_parser.find_reasoning_end_offset([7, 8]) == 2 def test_streaming_reasoning_implicit_end(self, thinking_parser): chunks = [ diff --git a/tests/parser/engine/test_parser_engine.py b/tests/parser/engine/test_parser_engine.py index f27ce52cb3..8fd7c4d262 100644 --- a/tests/parser/engine/test_parser_engine.py +++ b/tests/parser/engine/test_parser_engine.py @@ -185,8 +185,8 @@ class TestReasoningEndTokenIds: def test_find_reasoning_end_offset_returns_first_match(self): engine = _make_engine() assert engine.find_reasoning_end_offset([5, 201, 6, 201]) == 1 - assert engine.find_reasoning_end_offset([5, 6]) is None - assert engine.find_reasoning_end_offset([]) is None + assert engine.find_reasoning_end_offset([5, 6]) == 2 + assert engine.find_reasoning_end_offset([]) == 0 # Rejected-draft placeholders never match. assert engine.find_reasoning_end_offset([-1, 201]) == 1 diff --git a/tests/v1/core/test_async_scheduler.py b/tests/v1/core/test_async_scheduler.py index 7ff63d019c..7e27a4517e 100644 --- a/tests/v1/core/test_async_scheduler.py +++ b/tests/v1/core/test_async_scheduler.py @@ -334,7 +334,6 @@ def test_abort_request_when_structured_output_fsm_cannot_advance(): request = create_requests(num_requests=1, num_tokens=1)[0] request.structured_output_request = Mock() request.structured_output_request.grammar = Mock(spec=StructuredOutputGrammar) - request.structured_output_request.grammar.accept_tokens.return_value = False request.status = RequestStatus.RUNNING request.num_computed_tokens = request.num_tokens request.num_output_placeholders = 1 @@ -343,10 +342,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance(): scheduler.connector = None scheduler.ec_connector = None scheduler.structured_output_manager = Mock() - scheduler.structured_output_manager.should_advance.return_value = True - scheduler.structured_output_manager.trim_reasoning_for_advance.side_effect = ( - lambda request, new_token_ids: new_token_ids - ) + scheduler.structured_output_manager.accept_tokens.return_value = False scheduler.requests = {request.request_id: request} scheduler.running = [request] scheduler.waiting = Mock() diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 6ca6a8cf87..061c825fc8 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -3674,7 +3674,6 @@ def test_abort_request_when_structured_output_fsm_cannot_advance(): ) request.structured_output_request = Mock() request.structured_output_request.grammar = Mock(spec=StructuredOutputGrammar) - request.structured_output_request.grammar.accept_tokens.return_value = False request.status = RequestStatus.RUNNING request.num_computed_tokens = request.num_tokens @@ -3682,10 +3681,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance(): scheduler.connector = None scheduler.ec_connector = None scheduler.structured_output_manager = Mock() - scheduler.structured_output_manager.should_advance.return_value = True - scheduler.structured_output_manager.trim_reasoning_for_advance.side_effect = ( - lambda request, new_token_ids: new_token_ids - ) + scheduler.structured_output_manager.accept_tokens.return_value = False scheduler.requests = {request.request_id: request} scheduler.running = [request] scheduler.waiting = Mock() @@ -3734,8 +3730,8 @@ def test_abort_request_when_structured_output_fsm_cannot_advance(): ) engine_core_outputs = scheduler.update_from_output(output, model_runner_output) - request.structured_output_request.grammar.accept_tokens.assert_called_once_with( - request.request_id, [123] + scheduler.structured_output_manager.accept_tokens.assert_called_once_with( + request, [123] ) assert request.resumable is False assert request.status == RequestStatus.FINISHED_ERROR @@ -6573,16 +6569,15 @@ def _decode_ready_request(scheduler): return request -def test_update_draft_token_ids_strips_ngram_padding(monkeypatch): - """ngram_gpu pads unfilled draft slots with -1; update_draft_token_ids must +def test_update_draft_token_ids_strips_ngram_padding(): + """ngram_gpu pads unfilled draft slots with -1; manager validation must strip them before grammar.validate_tokens (which otherwise raises).""" scheduler = create_scheduler(num_speculative_tokens=4) request = _decode_ready_request(scheduler) grammar = _RecordingGrammar() - request.structured_output_request = SimpleNamespace(grammar=grammar) - monkeypatch.setattr( - scheduler.structured_output_manager, "should_advance", lambda req: True + request.structured_output_request = SimpleNamespace( + grammar=grammar, reasoning_ended=True ) scheduler.update_draft_token_ids( @@ -6593,16 +6588,15 @@ def test_update_draft_token_ids_strips_ngram_padding(monkeypatch): assert request.spec_token_ids == [10, 11] -def test_update_draft_token_ids_in_output_strips_padding(monkeypatch): +def test_update_draft_token_ids_in_output_strips_padding(): """Same guard on the output path; the -1 pad-back for the rejected count - is preserved (only the input to validate_tokens is stripped).""" + is preserved (only the input to manager validation is stripped).""" scheduler = create_scheduler(num_speculative_tokens=4) request = _decode_ready_request(scheduler) grammar = _RecordingGrammar() - request.structured_output_request = SimpleNamespace(grammar=grammar) - monkeypatch.setattr( - scheduler.structured_output_manager, "should_advance", lambda req: True + request.structured_output_request = SimpleNamespace( + grammar=grammar, reasoning_ended=True ) scheduler_output = SimpleNamespace( diff --git a/tests/v1/spec_decode/test_mtp_structured_output.py b/tests/v1/spec_decode/test_mtp_structured_output.py index 330e10438f..4e12a6f49d 100644 --- a/tests/v1/spec_decode/test_mtp_structured_output.py +++ b/tests/v1/spec_decode/test_mtp_structured_output.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""grammar_bitmask under spec-decode draft padding (#44006).""" +"""Backend and engine-reasoner checks that are not manager-flow tests.""" import pytest from transformers import AutoTokenizer @@ -49,221 +49,6 @@ def _make_manager_and_request(backend: str, prompt_str: str = '{"a": "b"}'): return tokenizer, manager, request, prompt -@pytest.mark.parametrize("backend", ["xgrammar", "guidance"]) -def test_bitmask_with_padded_invalid_drafts(backend): - """Bitmask handles -1 padded drafts and returns N+1 rows.""" - tokenizer, manager, request, prompt = _make_manager_and_request( - backend, prompt_str='{"a"' - ) - grammar = request.structured_output_request.grammar - - assert grammar.accept_tokens(request.request_id, prompt) - - valid_drafts = [tokenizer.encode(":")[0], tokenizer.encode(' "')[0]] - padded = valid_drafts + [-1, -1] - - bitmask = manager.grammar_bitmask( - requests={request.request_id: request}, - structured_output_request_ids=[request.request_id], - scheduled_spec_decode_tokens={request.request_id: padded}, - ) - - assert bitmask is not None - assert bitmask.shape[0] == len(padded) + 1 - assert not grammar.is_terminated() - - -@pytest.mark.parametrize("backend", ["xgrammar", "guidance"]) -def test_bitmask_when_grammar_terminates_mid_window(backend): - """Drafts following an EOS that terminates the grammar are a no-op.""" - tokenizer, manager, request, prompt = _make_manager_and_request(backend) - grammar = request.structured_output_request.grammar - - assert grammar.accept_tokens(request.request_id, prompt) - eos = tokenizer.eos_token_id - drafts = [eos] + [tokenizer.encode(" ")[0]] * (NUM_SPEC_TOKENS - 1) - - bitmask = manager.grammar_bitmask( - requests={request.request_id: request}, - structured_output_request_ids=[request.request_id], - scheduled_spec_decode_tokens={request.request_id: drafts}, - ) - - assert bitmask is not None - assert bitmask.shape[0] == NUM_SPEC_TOKENS + 1 - assert not grammar.is_terminated() - - -@pytest.mark.parametrize("backend", ["xgrammar", "guidance"]) -def test_bitmask_idempotent_across_calls(backend): - """Repeated calls with the same input return the same bitmask.""" - tokenizer, manager, request, prompt = _make_manager_and_request( - backend, prompt_str='{"a"' - ) - grammar = request.structured_output_request.grammar - - assert grammar.accept_tokens(request.request_id, prompt) - - drafts = [tokenizer.encode(":")[0], -1, -1, -1] - - first = manager.grammar_bitmask( - requests={request.request_id: request}, - structured_output_request_ids=[request.request_id], - scheduled_spec_decode_tokens={request.request_id: drafts}, - ) - second = manager.grammar_bitmask( - requests={request.request_id: request}, - structured_output_request_ids=[request.request_id], - scheduled_spec_decode_tokens={request.request_id: drafts}, - ) - - assert first is not None and second is not None - assert (first == second).all() - assert not grammar.is_terminated() - - -@pytest.mark.parametrize("backend", ["xgrammar", "guidance"]) -def test_bonus_position_constrained_after_invalid_drafts(backend): - """Regression for #44006: bonus row stays constrained after -1 padding.""" - tokenizer, manager, request, prompt = _make_manager_and_request( - backend, prompt_str='{"a"' - ) - grammar = request.structured_output_request.grammar - - assert grammar.accept_tokens(request.request_id, prompt) - - valid = tokenizer.encode(":")[0] - drafts = [valid, -1, -1, -1] - bitmask = manager.grammar_bitmask( - requests={request.request_id: request}, - structured_output_request_ids=[request.request_id], - scheduled_spec_decode_tokens={request.request_id: drafts}, - ) - assert bitmask is not None - assert bitmask.shape[0] == len(drafts) + 1 - - assert not (bitmask[-1] == -1).all() - assert not grammar.is_terminated() - - -@pytest.mark.parametrize("backend", ["xgrammar", "guidance"]) -def test_bitmask_constrained_when_reasoning_ends_midwindow(backend): - """Drafts after a mid-window reasoning-end marker stay constrained.""" - tokenizer, manager, request, prompt = _make_manager_and_request(backend) - grammar = request.structured_output_request.grammar - - assert grammar.accept_tokens(request.request_id, prompt) - - marker = tokenizer.encode("\n")[0] - - class StubReasoner: - def __init__(self, *_, **__): - self.end_token_id = marker - - def is_reasoning_end(self, input_ids): - return marker in list(input_ids) - - def is_reasoning_end_streaming(self, input_ids, delta_ids): - return marker in list(delta_ids) - - manager.reasoner_cls = StubReasoner - request.structured_output_request.reasoner = StubReasoner() - request.structured_output_request.reasoning_ended = False - - pre = tokenizer.encode(" ")[0] - post = tokenizer.encode(",")[0] - drafts = [pre, marker, post] - - bitmask = manager.grammar_bitmask( - requests={request.request_id: request}, - structured_output_request_ids=[request.request_id], - scheduled_spec_decode_tokens={request.request_id: drafts}, - ) - - assert bitmask is not None - assert bitmask.shape[0] == len(drafts) + 1 - assert (bitmask[0] == -1).all() - assert (bitmask[1] == -1).all() - assert not (bitmask[2] == -1).all() - assert not (bitmask[-1] == -1).all() - assert not grammar.is_terminated() - - -@pytest.mark.parametrize("backend", ["xgrammar", "guidance"]) -def test_bitmask_post_reasoning_end_drafts_skip_grammar_advance(backend, caplog): - """Post-marker drafts predate the bitmask and may be grammar-invalid; - grammar_bitmask must skip the grammar advance instead of asserting. - """ - tokenizer, manager, request, prompt = _make_manager_and_request( - backend, prompt_str="{" - ) - grammar = request.structured_output_request.grammar - - assert grammar.accept_tokens(request.request_id, prompt) - assert not grammar.is_terminated() - - marker = tokenizer.encode("\n")[0] - - class StubReasoner: - def __init__(self, *_, **__): - self.end_token_id = marker - - def is_reasoning_end(self, input_ids): - return marker in list(input_ids) - - def is_reasoning_end_streaming(self, input_ids, delta_ids): - return marker in list(delta_ids) - - manager.reasoner_cls = StubReasoner - request.structured_output_request.reasoner = StubReasoner() - request.structured_output_request.reasoning_ended = False - - pre = tokenizer.encode(" ")[0] - # A token that the JSON grammar would reject as the first post-marker - # token; without the fix grammar.accept_tokens fires the assertion. - invalid_post = tokenizer.encode("z")[0] - drafts = [pre, marker, invalid_post] - - bitmask = manager.grammar_bitmask( - requests={request.request_id: request}, - structured_output_request_ids=[request.request_id], - scheduled_spec_decode_tokens={request.request_id: drafts}, - ) - - assert bitmask is not None - assert bitmask.shape[0] == len(drafts) + 1 - # Post-marker position is still bitmask-constrained. - assert not (bitmask[2] == -1).all() - # Grammar must not have advanced through the unvalidated draft. - assert not grammar.is_terminated() - assert "Failed to advance FSM" not in caplog.text - - -@pytest.mark.parametrize("backend", ["xgrammar", "guidance"]) -def test_validate_tokens_then_bitmask_round_trip(backend): - """validate_tokens -> pad with -1 -> grammar_bitmask must not assert.""" - tokenizer, manager, request, prompt = _make_manager_and_request(backend) - grammar = request.structured_output_request.grammar - - assert grammar.accept_tokens(request.request_id, prompt) - - raw_drafts = [tokenizer.encode(",")[0], 99999, 12345, 67890] - accepted = grammar.validate_tokens(raw_drafts) - assert len(accepted) <= len(raw_drafts) - - padded = accepted + [-1] * (len(raw_drafts) - len(accepted)) - assert len(padded) == len(raw_drafts) - - bitmask = manager.grammar_bitmask( - requests={request.request_id: request}, - structured_output_request_ids=[request.request_id], - scheduled_spec_decode_tokens={request.request_id: padded}, - ) - assert bitmask is not None - assert bitmask.shape[0] == len(padded) + 1 - assert not grammar.is_terminated() - - def test_xgrammar_accept_tokens_stops_at_termination(capfd): """Tokens after a terminating EOS do not reach the matcher.""" tokenizer, _, request, prompt = _make_manager_and_request("xgrammar") @@ -302,7 +87,6 @@ def test_xgrammar_validate_tokens_stops_at_termination(capfd): assert grammar.validate_tokens([eos, trailing]) == [eos] assert "trying to accept new token" not in capfd.readouterr().err - # Check matcher state directly to verify validation rolled it back. assert not grammar.matcher.is_terminated() assert grammar.accept_tokens(request.request_id, [eos]) @@ -312,87 +96,6 @@ def test_xgrammar_validate_tokens_stops_at_termination(capfd): assert "trying to accept new token" not in capfd.readouterr().err -class _MarkerReasoner: - """Stub reasoner whose reasoning-end marker is a single fixed token.""" - - def __init__(self, marker: int): - self.marker = marker - - def is_reasoning_end(self, input_ids): - return self.marker in list(input_ids) - - def is_reasoning_end_streaming(self, input_ids, delta_ids): - return self.marker in list(delta_ids) - - -def _setup_boundary_request(backend: str): - """Request with a structural-tag key and reasoning not yet ended.""" - from vllm.v1.structured_output.backend_types import StructuredOutputOptions - - tokenizer, manager, request, prompt = _make_manager_and_request(backend) - marker = tokenizer.encode("\n")[0] - structured_req = request.structured_output_request - # The grammar itself is JSON (cheap to build); only the key kind matters - # for the should_advance structural-tag branch, so pre-seed the cached - # property. - structured_req.__dict__["structured_output_key"] = ( - StructuredOutputOptions.STRUCTURAL_TAG, - "", - ) - manager.reasoner_cls = _MarkerReasoner - structured_req.reasoner = _MarkerReasoner(marker) - structured_req.reasoning_ended = False - return tokenizer, manager, request, prompt, marker - - -def test_should_advance_records_reasoning_end_index(): - """Regression for #44006 on post-#42452 main: the boundary step must - record where reasoning ends so the scheduler can trim before advancing. - """ - tokenizer, manager, request, prompt, marker = _setup_boundary_request("xgrammar") - structured_req = request.structured_output_request - - pre = tokenizer.encode(" ")[0] - post = tokenizer.encode("{")[0] - request.append_output_token_ids([pre, marker, post]) - - assert manager.should_advance(request) - assert structured_req.reasoning_ended - # Marker sits at absolute index len(prompt) + 1. - assert structured_req.reasoning_end_token_index == len(prompt) + 1 - - -def test_trim_reasoning_for_advance(): - """trim drops the marker and everything before it; later steps and - requests without a recorded boundary pass through unchanged. - """ - tokenizer, manager, request, prompt, marker = _setup_boundary_request("xgrammar") - structured_req = request.structured_output_request - - pre = tokenizer.encode(" ")[0] - post = tokenizer.encode("{")[0] - - # No boundary recorded yet: pass-through. - assert manager.trim_reasoning_for_advance(request, [pre]) == [pre] - - # Boundary step: marker mid-step keeps only the suffix. - step_tokens = [pre, marker, post] - request.append_output_token_ids(step_tokens) - assert manager.should_advance(request) - assert manager.trim_reasoning_for_advance(request, step_tokens) == [post] - - # Boundary step variant: marker last (the #44006 crash shape - # [198, ]) trims to empty -> scheduler skips accept_tokens. - structured_req.reasoning_end_token_index = len(request.all_token_ids) - 1 - assert manager.trim_reasoning_for_advance(request, step_tokens) == [] - - # Later steps: tokens are past the boundary, returned unchanged. - structured_req.reasoning_end_token_index = len(prompt) + 1 - next_step = [post, post] - request.append_output_token_ids(next_step) - assert manager.trim_reasoning_for_advance(request, next_step) == next_step - - class _EngineReasonerStub(ParserEngineReasoningAdapter): """Adapter-typed reasoner with a fixed end-token set and no real engine.""" @@ -409,7 +112,7 @@ class _EngineReasonerStub(ParserEngineReasoningAdapter): for offset, token in enumerate(token_ids): if token in self._end_token_ids: return offset - return None + return len(token_ids) def is_reasoning_end(self, input_ids): 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): assert not (bitmask[2] == -1).all() assert reasoner.windows == [[pre, marker, post]] assert not grammar.is_terminated() - - -@pytest.mark.parametrize("backend", ["xgrammar", "guidance"]) -def test_bitmask_legacy_order_sensitive_predicate_probed_per_token(backend): - """Legacy predicates are probed one draft at a time, never on the whole - window. A predicate that reports which marker is *newest* (KimiK3) would - otherwise answer False for a window that closes and then reopens - reasoning, leaving the post-marker drafts and the bonus row - unconstrained. - """ - tokenizer, manager, request, prompt = _make_manager_and_request(backend) - grammar = request.structured_output_request.grammar - - assert grammar.accept_tokens(request.request_id, prompt) - - close = tokenizer.encode("\n")[0] - reopen = tokenizer.encode("#")[0] - assert close != reopen - - class NewestMarkerReasoner: - def __init__(self, *_, **__): - pass - - def is_reasoning_end(self, input_ids): - return self.is_reasoning_end_streaming(input_ids, input_ids) - - def is_reasoning_end_streaming(self, input_ids, delta_ids): - for token in reversed(list(delta_ids)): - if token == close: - return True - if token == reopen: - return False - return False - - manager.reasoner_cls = NewestMarkerReasoner - request.structured_output_request.reasoner = NewestMarkerReasoner() - request.structured_output_request.reasoning_ended = False - - pre = tokenizer.encode(" ")[0] - post = tokenizer.encode(",")[0] - drafts = [pre, close, post, reopen] - # Sanity: the whole window says "reopened", so only per-token probing - # can find the close marker. - assert not NewestMarkerReasoner().is_reasoning_end_streaming([], drafts) - - bitmask = manager.grammar_bitmask( - requests={request.request_id: request}, - structured_output_request_ids=[request.request_id], - scheduled_spec_decode_tokens={request.request_id: drafts}, - ) - - assert bitmask is not None - assert bitmask.shape[0] == len(drafts) + 1 - assert (bitmask[0] == -1).all() - assert (bitmask[1] == -1).all() - assert not (bitmask[2] == -1).all() - assert not (bitmask[-1] == -1).all() - assert not grammar.is_terminated() diff --git a/tests/v1/structured_output/test_reasoning_structured_output.py b/tests/v1/structured_output/test_reasoning_structured_output.py deleted file mode 100644 index a57e5da72b..0000000000 --- a/tests/v1/structured_output/test_reasoning_structured_output.py +++ /dev/null @@ -1,557 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -"""Unit tests for reasoning-aware structured output functionality (PR #25515).""" - -from collections.abc import Sequence -from unittest.mock import Mock - -import pytest - -from vllm.config import ModelConfig, SchedulerConfig, VllmConfig -from vllm.v1.request import Request -from vllm.v1.structured_output import StructuredOutputManager -from vllm.v1.structured_output.backend_types import StructuredOutputOptions - - -class UntouchableTokenIds(Sequence[int]): - """A token history that fails the test if anything reads it.""" - - def __init__(self, length: int) -> None: - self._length = length - - def __len__(self) -> int: - return self._length - - def __getitem__(self, item): - raise AssertionError("token history must not be read") - - def __iter__(self): - raise AssertionError("token history must not be iterated") - - -def make_qwen3_reasoner(): - """A real engine-based reasoning adapter over a tiny mock vocabulary.""" - from tests.parser.engine.conftest import make_mock_tokenizer - from vllm.parser.engine.registered_adapters import Qwen3ParserReasoningAdapter - - tokenizer = make_mock_tokenizer( - {"": 200, "": 201, "": 202, "": 203} - ) - return Qwen3ParserReasoningAdapter(tokenizer) - - -class MockReasoner: - def __init__(self, tokenizer): - self.is_reasoning_end = Mock(return_value=False) - self.is_reasoning_end_streaming = Mock(return_value=False) - - -class TestReasoningStructuredOutput: - """Test reasoning-aware structured output functionality.""" - - @pytest.fixture - def mock_model_config(self): - """Create a mock ModelConfig.""" - config = Mock(spec=ModelConfig) - config.skip_tokenizer_init = True # Skip tokenizer init to avoid network calls - config.get_vocab_size = Mock(return_value=50000) - # Add missing runner_type attribute that tokenizer initialization expects - config.runner_type = "generate" - # Add other attributes that tokenizer initialization might need - config.tokenizer = "test-tokenizer" - config.tokenizer_mode = "auto" - config.trust_remote_code = False - config.tokenizer_revision = None - return config - - @pytest.fixture - def mock_scheduler_config(self): - """Create a mock SchedulerConfig.""" - config = Mock(spec=SchedulerConfig) - config.max_num_seqs = 128 - return config - - @pytest.fixture - def mock_vllm_config(self, mock_model_config, mock_scheduler_config): - """Create a mock VllmConfig.""" - config = Mock(spec=VllmConfig) - config.model_config = mock_model_config - config.scheduler_config = mock_scheduler_config - config.structured_outputs_config = Mock() - config.structured_outputs_config.reasoning_parser = None - config.structured_outputs_config.enable_in_reasoning = False - config.speculative_config = None - return config - - @pytest.fixture - def mock_request_with_structured_output(self): - """Create a mock request with structured output.""" - request = Mock(spec=Request) - request.structured_output_request = Mock() - request.structured_output_request.reasoning_ended = None - request.structured_output_request.grammar = Mock() - request.structured_output_request.reasoning_parser_kwargs = None - request.structured_output_request.reasoner = None - request.structured_output_request.grammar.is_terminated = Mock( - return_value=False - ) - request.use_structured_output = True - request.prompt_token_ids = [1, 2, 3, 4, 5] - request.all_token_ids = [1, 2, 3, 4, 5, 6, 7, 8] - request.num_computed_tokens = 5 - request.num_output_placeholders = 0 - request.request_id = "mock_req" - return request - - @pytest.fixture - def manager_with_reasoner(self, mock_vllm_config): - manager = StructuredOutputManager(mock_vllm_config) - manager.reasoner_cls = MockReasoner - manager.tokenizer = Mock() - return manager - - def test_should_fill_bitmask_with_enable_in_reasoning( - self, manager_with_reasoner, mock_request_with_structured_output - ): - """Test should_fill_bitmask when enable_in_reasoning is True.""" - manager_with_reasoner.enable_in_reasoning = True - - # Should always return True when enable_in_reasoning is enabled - result = manager_with_reasoner.should_fill_bitmask( - mock_request_with_structured_output - ) - assert result is True - assert ( - mock_request_with_structured_output.structured_output_request.reasoner - is None - ) - - def test_should_fill_bitmask_reasoning_already_ended( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """An active grammar does not need a request-local reasoner.""" - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoning_ended = True - - result = manager_with_reasoner.should_fill_bitmask( - mock_request_with_structured_output - ) - - assert result is True - assert structured_req.reasoner is None - - def test_grammar_bitmask_skips_reasoner_when_already_active( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """Bitmask generation skips reasoning-boundary detection.""" - manager_with_reasoner.vllm_config.num_speculative_tokens = 1 - manager_with_reasoner._grammar_bitmask = Mock() - manager_with_reasoner._grammar_bitmask.shape = (1, 1) - expected_bitmask = Mock() - manager_with_reasoner._grammar_bitmask.numpy.return_value = expected_bitmask - manager_with_reasoner._fill_bitmasks = Mock() - - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoning_ended = True - request_id = mock_request_with_structured_output.request_id - - bitmask = manager_with_reasoner.grammar_bitmask( - requests={request_id: mock_request_with_structured_output}, - structured_output_request_ids=[request_id], - scheduled_spec_decode_tokens={}, - ) - - assert bitmask is expected_bitmask - assert structured_req.reasoner is None - - def test_should_fill_bitmask_without_enable_in_reasoning( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """Test should_fill_bitmask when enable_in_reasoning is False.""" - # Keep enable_in_reasoning as False (default) - config = manager_with_reasoner.vllm_config.structured_outputs_config - assert config.enable_in_reasoning is False - - result = manager_with_reasoner.should_fill_bitmask( - mock_request_with_structured_output - ) - - # Should set reasoning_ended and return its value - assert ( - mock_request_with_structured_output.structured_output_request.reasoning_ended - is False - ) - assert result is False - - def test_should_fill_bitmask_no_reasoner( - self, mock_vllm_config, mock_request_with_structured_output - ): - """Test should_fill_bitmask when no reasoner is configured.""" - manager = StructuredOutputManager(mock_vllm_config) - - result = manager.should_fill_bitmask(mock_request_with_structured_output) - - # Should default to True when no reasoner - assert result is True - - def test_should_fill_bitmask_uses_request_reasoning_parser_kwargs( - self, mock_vllm_config, mock_request_with_structured_output - ): - """Test request-level parser kwargs override the default reasoner.""" - - class KwargReasoner: - def __init__(self, tokenizer, chat_template_kwargs=None): - self.chat_template_kwargs = chat_template_kwargs or {} - - def is_reasoning_end(self, input_ids): - return not self.chat_template_kwargs.get("enable_thinking", False) - - manager = StructuredOutputManager(mock_vllm_config) - manager.reasoner_cls = KwargReasoner - manager.tokenizer = Mock() - - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoning_parser_kwargs = { - "chat_template_kwargs": {"enable_thinking": True} - } - - result = manager.should_fill_bitmask(mock_request_with_structured_output) - - assert result is False - assert ( - mock_request_with_structured_output.structured_output_request.reasoner - is not None - ) - - def test_should_advance_with_enable_in_reasoning( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """Test should_advance when enable_in_reasoning is True.""" - # Enable enable_in_reasoning - manager_with_reasoner.enable_in_reasoning = True - - # Should always return True when enable_in_reasoning is enabled - result = manager_with_reasoner.should_advance( - mock_request_with_structured_output - ) - assert result is True - assert ( - mock_request_with_structured_output.structured_output_request.reasoner - is None - ) - - def test_should_advance_reasoning_not_ended( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """Test should_advance when reasoning has not ended.""" - # Set reasoning as not ended - ( - mock_request_with_structured_output.structured_output_request - ).reasoning_ended = False - - result = manager_with_reasoner.should_advance( - mock_request_with_structured_output - ) - - # Should return False since reasoning hasn't ended - assert result is False - - def test_should_advance_reasoning_just_ended( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """Test should_advance when reasoning ends in current step.""" - # Set reasoning as not ended initially, but ends in this step - ( - mock_request_with_structured_output.structured_output_request - ).reasoning_ended = False - reasoner = MockReasoner(tokenizer=Mock()) - reasoner.is_reasoning_end_streaming.return_value = True - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoner = reasoner - - result = manager_with_reasoner.should_advance( - mock_request_with_structured_output - ) - - # The scheduler trims the reasoning prefix before advancing the grammar. - assert ( - mock_request_with_structured_output.structured_output_request.reasoning_ended - is True - ) - assert result is True - - def test_should_advance_reasoning_already_ended( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """Test should_advance when reasoning has already ended.""" - # Set reasoning as already ended - ( - mock_request_with_structured_output.structured_output_request - ).reasoning_ended = True - - result = manager_with_reasoner.should_advance( - mock_request_with_structured_output - ) - - # Should return True since reasoning has ended - assert result is True - assert ( - mock_request_with_structured_output.structured_output_request.reasoner - is None - ) - - def test_should_advance_uses_new_token_ids_when_provided( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """Regression for #43388: when caller passes new_token_ids, the - reasoner sees the exact multi-token delta rather than the - placeholder-derived window. - """ - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoning_ended = False - - end_token_id = 248069 - - reasoner = MockReasoner(tokenizer=Mock()) - # Detection mirrors the real Qwen3 parser: end token in the delta. - reasoner.is_reasoning_end_streaming = Mock( - side_effect=lambda input_ids, delta_ids: end_token_id in list(delta_ids) - ) - structured_req.reasoner = reasoner - - # Scenario from #43388: async + spec decode K=4, 4 tokens accepted - # but only 1 placeholder remains (some drafts were rejected). - # The placeholder math would yield delta=[271] and miss . - # Passing new_token_ids must override that. - new_token_ids = [9, 198, end_token_id, 271] - mock_request_with_structured_output.all_token_ids = [ - 1, - 2, - 3, - 4, - 5, - ] + new_token_ids - mock_request_with_structured_output.num_computed_tokens = 9 - mock_request_with_structured_output.num_output_placeholders = 1 - - result = manager_with_reasoner.should_advance( - mock_request_with_structured_output, - new_token_ids=new_token_ids, - ) - - # First call to is_reasoning_end_streaming was with the full - # new_token_ids (not the truncated placeholder window). - first_call = reasoner.is_reasoning_end_streaming.call_args_list[0] - _, called_delta = first_call.args - assert list(called_delta) == new_token_ids - - assert structured_req.reasoning_ended is True - assert result is True - - def test_should_advance_without_new_token_ids_falls_back( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """Backward compat: callers that don't pass new_token_ids keep - the original placeholder-derived delta window. - """ - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoning_ended = False - reasoner = MockReasoner(tokenizer=Mock()) - reasoner.is_reasoning_end_streaming.return_value = False - structured_req.reasoner = reasoner - - mock_request_with_structured_output.all_token_ids = [1, 2, 3, 4, 5] - mock_request_with_structured_output.num_computed_tokens = 5 - mock_request_with_structured_output.num_output_placeholders = 2 - - result = manager_with_reasoner.should_advance( - mock_request_with_structured_output - ) - - # placeholder window: start = 5 - 2 = 3, delta = [4, 5] - _, called_delta = reasoner.is_reasoning_end_streaming.call_args[0] - assert list(called_delta) == [4, 5] - assert result is False - - def test_should_advance_whole_delta_fallback_pins_last_sequence_index( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """When only the whole-delta predicate fires, the boundary is the - last index of the sequence, even if the placeholder window - overshoots it and the delta is empty. - """ - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoning_ended = False - reasoner = MockReasoner(tokenizer=Mock()) - # Base-default shape: rescans the whole sequence, ignores the delta. - reasoner.is_reasoning_end_streaming.return_value = True - structured_req.reasoner = reasoner - - mock_request_with_structured_output.all_token_ids = [1, 2, 3, 4, 5] - # placeholder window: start = 9 - 1 = 8, past the end -> delta = [] - mock_request_with_structured_output.num_computed_tokens = 9 - mock_request_with_structured_output.num_output_placeholders = 1 - - assert manager_with_reasoner.should_advance(mock_request_with_structured_output) - assert structured_req.reasoning_ended is True - assert structured_req.reasoning_end_token_index == 4 - - def test_should_advance_trims_reasoning_prefix_for_json( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """JSON uses the common trim-then-advance path at the boundary.""" - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoning_ended = False - structured_req.structured_output_key = ( - StructuredOutputOptions.JSON_OBJECT, - "{}", - ) - - marker = 248069 - - class MarkerReasoner: - def __init__(self, *_, **__): - pass - - def is_reasoning_end_streaming(self, input_ids, delta_ids): - return marker in list(delta_ids) - - structured_req.reasoner = MarkerReasoner() - - new_token_ids = [9, 198, marker, 271, 5005] - mock_request_with_structured_output.all_token_ids = [1, 2, 3] + new_token_ids - - result = manager_with_reasoner.should_advance( - mock_request_with_structured_output, - new_token_ids=new_token_ids, - ) - - structured_req.grammar.accept_tokens.assert_not_called() - assert structured_req.reasoning_ended is True - assert result is True - assert structured_req.reasoning_end_token_index == 5 - assert manager_with_reasoner.trim_reasoning_for_advance( - mock_request_with_structured_output, new_token_ids - ) == [271, 5005] - - def test_should_advance_engine_adapter_scans_only_the_delta( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """Engine-based parsers locate the boundary from the delta alone.""" - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoning_ended = False - reasoner = make_qwen3_reasoner() - assert reasoner.reasoning_end_token_ids == {201, 202} - reasoner.is_reasoning_end_streaming = Mock( - side_effect=AssertionError("engine path must not rescan the prefix") - ) - structured_req.reasoner = reasoner - - new_token_ids = [9, 198, 201, 271] - mock_request_with_structured_output.all_token_ids = [1, 2, 3] + new_token_ids - - assert manager_with_reasoner.should_advance( - mock_request_with_structured_output, new_token_ids=new_token_ids - ) - assert structured_req.reasoning_ended is True - assert structured_req.reasoning_end_token_index == 5 - assert manager_with_reasoner.trim_reasoning_for_advance( - mock_request_with_structured_output, new_token_ids - ) == [271] - - def test_should_advance_engine_adapter_tool_start_ends_reasoning( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoning_ended = False - structured_req.reasoner = make_qwen3_reasoner() - - new_token_ids = [9, 202, 271] - mock_request_with_structured_output.all_token_ids = [1, 2, 3] + new_token_ids - - assert manager_with_reasoner.should_advance( - mock_request_with_structured_output, new_token_ids=new_token_ids - ) - assert structured_req.reasoning_end_token_index == 4 - - def test_should_advance_engine_adapter_no_boundary( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoning_ended = False - structured_req.reasoner = make_qwen3_reasoner() - - new_token_ids = [9, 198, 271] - mock_request_with_structured_output.all_token_ids = [1, 2, 3] + new_token_ids - - assert not manager_with_reasoner.should_advance( - mock_request_with_structured_output, new_token_ids=new_token_ids - ) - assert structured_req.reasoning_ended is False - - def test_should_advance_engine_adapter_never_reads_history( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """Per-step cost must not depend on how much has been generated.""" - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoning_ended = False - structured_req.reasoner = make_qwen3_reasoner() - - new_token_ids = [9, 201] - total = (1 << 20) + len(new_token_ids) - mock_request_with_structured_output.all_token_ids = UntouchableTokenIds(total) - - assert manager_with_reasoner.should_advance( - mock_request_with_structured_output, new_token_ids=new_token_ids - ) - assert structured_req.reasoning_end_token_index == total - 1 - - def test_should_advance_engine_adapter_without_new_tokens_persists( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """Draft-validation callers persist the boundary like the legacy path.""" - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoning_ended = False - structured_req.reasoner = make_qwen3_reasoner() - - mock_request_with_structured_output.all_token_ids = [1, 2, 3, 4, 5, 201, 7] - # placeholder window: start = 7 - 2 = 5, delta = [201, 7] - mock_request_with_structured_output.num_computed_tokens = 7 - mock_request_with_structured_output.num_output_placeholders = 2 - - assert manager_with_reasoner.should_advance(mock_request_with_structured_output) - assert structured_req.reasoning_ended is True - assert structured_req.reasoning_end_token_index == 5 diff --git a/tests/v1/structured_output/test_structured_output_manager.py b/tests/v1/structured_output/test_structured_output_manager.py new file mode 100644 index 0000000000..17bda09b3e --- /dev/null +++ b/tests/v1/structured_output/test_structured_output_manager.py @@ -0,0 +1,414 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import time +from dataclasses import dataclass + +import pytest +from transformers import AutoTokenizer + +from vllm.config import DeviceConfig, StructuredOutputsConfig, VllmConfig +from vllm.config.model import ModelConfig +from vllm.config.speculative import SpeculativeConfig +from vllm.sampling_params import SamplingParams, StructuredOutputsParams +from vllm.v1.request import Request +from vllm.v1.structured_output import StructuredOutputManager + +TOKENIZER = "gpt2" +THINK_END = "\n" # reasoning-end marker (single GPT-2 token) +EOS = "<|eos|>" # resolved to tokenizer.eos_token_id +JSON_SCHEMA = '{"type": "object"}' +BACKENDS = ("xgrammar", "guidance") +MAX_WAIT_SECONDS = 5 +NUM_SPEC_TOKENS = 8 + + +@pytest.fixture(scope="module") +def tokenizer(): + return AutoTokenizer.from_pretrained(TOKENIZER) + + +@dataclass(frozen=True) +class FlowCase: + """One validate -> bitmask -> accept -> post-accept bitmask scenario.""" + + raw_drafts: tuple[str, ...] + expected_row_pattern: str + expected_reasoning: bool | None + expect_terminated: bool + prefix: str = "" + expected_validated: tuple[str, ...] | None = None + reasoning_ended: bool | None = False + xfail_guidance: str | None = None + + +class MockReasoner: + def __init__(self, tokenizer, marker: int | None = None): + self.marker = marker + + def is_reasoning_end(self, input_ids): + if self.marker is None: + return True + return self.marker in list(input_ids) + + def is_reasoning_end_streaming(self, input_ids, delta_ids): + return self.is_reasoning_end(delta_ids) + + +def _single_token(tokenizer, text: str) -> int: + token_ids = tokenizer.encode(text) + assert len(token_ids) == 1, (text, token_ids) + return token_ids[0] + + +def _to_token_ids(tokenizer, texts: tuple[str, ...]) -> list[int]: + """Convert literal text (or EOS sentinel) to token IDs.""" + return [ + tokenizer.eos_token_id if t == EOS else _single_token(tokenizer, t) + for t in texts + ] + + +def _wait_for_grammar(request: Request) -> None: + structured_req = request.structured_output_request + assert structured_req is not None + + deadline = time.time() + MAX_WAIT_SECONDS + while not structured_req._check_grammar_completion(): + if time.time() > deadline: + pytest.fail("Grammar compilation timed out") + time.sleep(0.01) + + +def _build_harness( + tokenizer, + backend: str, + prefix: str = "", + use_reasoner: bool = True, + reasoning_ended: bool | None = None, + enable_in_reasoning: bool = False, + reasoning_parser_kwargs: dict | None = None, +) -> tuple[StructuredOutputManager, Request]: + vllm_config = VllmConfig( + model_config=ModelConfig(tokenizer=TOKENIZER), + device_config=DeviceConfig(device="cpu"), + structured_outputs_config=StructuredOutputsConfig( + backend=backend, + enable_in_reasoning=enable_in_reasoning, + ), + speculative_config=SpeculativeConfig( + model="[ngram]", + num_speculative_tokens=NUM_SPEC_TOKENS, + ), + ) + manager = StructuredOutputManager(vllm_config) + if use_reasoner: + manager.reasoner_cls = MockReasoner + + sampling_params = SamplingParams( + structured_outputs=StructuredOutputsParams(json=JSON_SCHEMA) + ) + sampling_params.structured_outputs._backend = backend + sampling_params.update_from_generation_config({}, tokenizer.eos_token_id) + + prompt_ids = tokenizer.encode(prefix) if prefix else [] + request = Request( + request_id=f"{backend}-flow", + prompt_token_ids=prompt_ids, + sampling_params=sampling_params, + pooling_params=None, + reasoning_ended=reasoning_ended, + reasoning_parser_kwargs=reasoning_parser_kwargs, + ) + manager.grammar_init(request) + _wait_for_grammar(request) + + structured_req = request.structured_output_request + assert request.prompt_token_ids is not None + assert structured_req is not None + assert structured_req.grammar is not None + if prompt_ids: + assert structured_req.grammar.accept_tokens( + request.request_id, + prompt_ids, + ) + return manager, request + + +def _row_pattern(bitmask) -> str: + return "".join("U" if (row == -1).all() else "C" for row in bitmask) + + +def _run_real_flow( + manager: StructuredOutputManager, + request: Request, + raw_drafts: list[int], + expected_validated: list[int], + expected_row_pattern: str, + expected_reasoning: bool | None, + expect_terminated: bool, +) -> None: + assert len(expected_validated) <= len(raw_drafts) + + validated = manager.validate_tokens(request, list(raw_drafts)) + assert validated == expected_validated + + padded = validated + [-1] * (len(raw_drafts) - len(validated)) + bitmask = manager.grammar_bitmask( + requests={request.request_id: request}, + structured_output_request_ids=[request.request_id], + scheduled_spec_decode_tokens={request.request_id: padded}, + ) + assert bitmask is not None + assert bitmask.shape[0] == len(raw_drafts) + 1 + assert _row_pattern(bitmask) == expected_row_pattern + + # grammar_bitmask() must rollback any speculative state it advanced. + assert manager.validate_tokens(request, list(raw_drafts)) == expected_validated + + structured_req = request.structured_output_request + assert structured_req is not None + grammar = structured_req.grammar + assert grammar is not None + assert not grammar.is_terminated() + # `bitmask[i]` is the grammar state after the first `i` scheduled tokens. + # These tests commit the validated prefix, so the matching post-accept + # state is already present in the scheduled bitmask. + expected_post_accept_row = bitmask[len(expected_validated)].copy() + + # Mirror the scheduler flow: sampled tokens are appended before + # StructuredOutputManager.accept_tokens() trims them to the + # grammar-constrained suffix. + request.append_output_token_ids(list(expected_validated)) + assert manager.accept_tokens(request, list(expected_validated)) + + post_accept_bitmask = manager.grammar_bitmask( + requests={request.request_id: request}, + structured_output_request_ids=[request.request_id], + scheduled_spec_decode_tokens={}, + ) + assert post_accept_bitmask is not None + assert post_accept_bitmask.shape[0] == 1 + assert (post_accept_bitmask[0] == expected_post_accept_row).all() + assert structured_req.reasoning_ended == expected_reasoning + assert grammar.is_terminated() is expect_terminated + + +FLOW_CASES = [ + pytest.param( + FlowCase( + raw_drafts=(" ", "z", " "), + expected_row_pattern="UUUU", + expected_reasoning=False, + expect_terminated=False, + ), + id="inactive", + ), + pytest.param( + FlowCase( + prefix='{"a"', + raw_drafts=(":", ' "', "b", '"}'), + expected_row_pattern="CCCCC", + expected_reasoning=True, + expect_terminated=False, + reasoning_ended=True, + ), + id="active_all_valid", + ), + pytest.param( + FlowCase( + prefix='{"a"', + raw_drafts=("z", ":", ' "', "b"), + expected_validated=(), + expected_row_pattern="CUUUU", + expected_reasoning=True, + expect_terminated=False, + reasoning_ended=True, + ), + id="active_first_token_invalid", + ), + pytest.param( + FlowCase( + prefix='{"a"', + raw_drafts=(":", "z", ' "', "b"), + expected_validated=(":",), + expected_row_pattern="CCUUU", + expected_reasoning=True, + expect_terminated=False, + reasoning_ended=True, + ), + id="active_later_token_invalid", + ), + pytest.param( + FlowCase( + prefix='{"a": "b"}', + raw_drafts=(EOS, " ", " "), + expected_validated=(EOS,), + expected_row_pattern="CUUU", + expected_reasoning=True, + expect_terminated=True, + reasoning_ended=True, + xfail_guidance=( + "guidance validate_tokens rejects EOS on an already-complete " + "object even though accept_tokens would accept it" + ), + ), + id="active_first_token_terminates", + ), + pytest.param( + FlowCase( + prefix='{"a"', + raw_drafts=(":", ' "', "b", '"}', EOS, " "), + expected_validated=(":", ' "', "b", '"}', EOS), + expected_row_pattern="CCCCCUU", + expected_reasoning=True, + expect_terminated=True, + reasoning_ended=True, + ), + id="active_later_token_terminates", + ), + pytest.param( + FlowCase( + raw_drafts=(THINK_END, "{", ' "', "b"), + expected_row_pattern="UCCCC", + expected_reasoning=True, + expect_terminated=False, + ), + id="becomes_active_on_first_token", + ), + pytest.param( + FlowCase( + raw_drafts=(" ", THINK_END, "{", ' "'), + expected_row_pattern="UUCCC", + expected_reasoning=True, + expect_terminated=False, + ), + id="becomes_active_on_middle_token", + ), + pytest.param( + FlowCase( + raw_drafts=(" ", " ", THINK_END), + expected_row_pattern="UUUC", + expected_reasoning=True, + expect_terminated=False, + ), + id="becomes_active_on_last_token", + ), + pytest.param( + FlowCase( + raw_drafts=(" ", THINK_END, "z", "{"), + expected_validated=(" ", THINK_END), + expected_row_pattern="UUCUU", + expected_reasoning=True, + expect_terminated=False, + ), + id="becomes_active_first_token_invalid", + ), + pytest.param( + FlowCase( + raw_drafts=(" ", THINK_END, "{", "z"), + expected_validated=(" ", THINK_END, "{"), + expected_row_pattern="UUCCU", + expected_reasoning=True, + expect_terminated=False, + ), + id="becomes_active_later_token_invalid", + ), + pytest.param( + FlowCase( + raw_drafts=(" ", THINK_END, "{", "}", EOS, " "), + expected_validated=(" ", THINK_END, "{", "}", EOS), + expected_row_pattern="UUCCCUU", + expected_reasoning=True, + expect_terminated=True, + ), + id="becomes_active_terminates", + ), +] + + +@pytest.mark.parametrize("backend", BACKENDS) +@pytest.mark.parametrize("case", FLOW_CASES) +def test_real_flow( + tokenizer, + backend: str, + case: FlowCase, +): + if backend == "guidance" and case.xfail_guidance: + pytest.xfail(case.xfail_guidance) + + reasoner_kwargs = {"marker": _single_token(tokenizer, THINK_END)} + manager, request = _build_harness( + tokenizer, + backend, + prefix=case.prefix, + reasoning_ended=case.reasoning_ended, + reasoning_parser_kwargs=reasoner_kwargs, + ) + + raw_drafts = _to_token_ids(tokenizer, case.raw_drafts) + expected_texts = ( + case.raw_drafts if case.expected_validated is None else case.expected_validated + ) + expected_validated = _to_token_ids(tokenizer, expected_texts) + + _run_real_flow( + manager, + request, + raw_drafts=raw_drafts, + expected_validated=expected_validated, + expected_row_pattern=case.expected_row_pattern, + expected_reasoning=case.expected_reasoning, + expect_terminated=case.expect_terminated, + ) + + +@pytest.mark.parametrize("backend", BACKENDS) +@pytest.mark.parametrize( + ("use_reasoner", "reasoning_ended", "enable_in_reasoning"), + [ + pytest.param(True, False, True, id="enable_in_reasoning"), + pytest.param(True, True, False, id="reasoning_ended"), + pytest.param(False, False, False, id="no_reasoner"), + pytest.param(True, None, False, id="inferred"), + ], +) +def test_initial_constraint_activation( + tokenizer, + backend: str, + use_reasoner: bool, + reasoning_ended: bool | None, + enable_in_reasoning: bool, +): + manager, request = _build_harness( + tokenizer, + backend, + use_reasoner=use_reasoner, + reasoning_ended=reasoning_ended, + enable_in_reasoning=enable_in_reasoning, + reasoning_parser_kwargs={ + "marker": _single_token(tokenizer, THINK_END) + if reasoning_ended is not None + else None + }, + ) + + # "{" is valid JSON start; "z" is not. Truncation proves grammar is active. + open_brace = _single_token(tokenizer, "{") + z = _single_token(tokenizer, "z") + assert manager.validate_tokens(request, [open_brace, z]) == [open_brace] + + bitmask = manager.grammar_bitmask( + requests={request.request_id: request}, + structured_output_request_ids=[request.request_id], + scheduled_spec_decode_tokens={}, + ) + assert bitmask is not None + request.append_output_token_ids([open_brace]) + assert manager.accept_tokens(request, [open_brace]) + + # Short-circuit paths must not construct a request-local reasoner (#52573). + if enable_in_reasoning or reasoning_ended is True: + structured_req = request.structured_output_request + assert structured_req is not None + assert structured_req.reasoner is None diff --git a/vllm/parser/abstract_parser.py b/vllm/parser/abstract_parser.py index cf83bb996e..f099d7752a 100644 --- a/vllm/parser/abstract_parser.py +++ b/vllm/parser/abstract_parser.py @@ -197,6 +197,8 @@ class Parser: Used by structured engines like `xgrammar` to check if the reasoning content ends in the model output. + Must be a pure function of the input_ids. + Args: input_ids: The token IDs of the model output. @@ -210,6 +212,8 @@ class Parser: """ Check if the reasoning content ends during a decode step. + Must be a pure function of the input_ids and delta_ids. + Args: input_ids: The entire model output token IDs. delta_ids: The last few computed tokens at the current decode step. @@ -228,8 +232,8 @@ class Parser: Returns: The offset within ``token_ids`` of the last reasoning token, or - ``None`` when reasoning does not end in this window. Parsers that - cannot answer from a window alone always return ``None``; callers + ``len(token_ids)`` when reasoning does not end in this window. Parsers + that cannot answer from a window alone always return ``None``; callers fall back to :meth:`is_reasoning_end_streaming`. """ return None diff --git a/vllm/parser/engine/parser_engine.py b/vllm/parser/engine/parser_engine.py index f873e36342..fa306a5da4 100644 --- a/vllm/parser/engine/parser_engine.py +++ b/vllm/parser/engine/parser_engine.py @@ -645,7 +645,7 @@ class ParserEngine(Parser): for offset, token_id in enumerate(token_ids): if token_id in end_ids: return offset - return None + return len(token_ids) def is_reasoning_end(self, input_ids: list[int]) -> bool: config = self.parser_engine_config diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 3af527124d..daf2b316d0 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -70,8 +70,7 @@ from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput from vllm.v1.request import Request, RequestStatus, StreamingUpdate from vllm.v1.spec_decode.dynamic.utils import build_dynamic_sd_schedule_lookup from vllm.v1.spec_decode.metrics import SpecDecodingStats -from vllm.v1.structured_output import StructuredOutputGrammar, StructuredOutputManager -from vllm.v1.structured_output.utils import strip_speculative_padding +from vllm.v1.structured_output import StructuredOutputManager from vllm.v1.utils import record_function_or_nullcontext logger = init_logger(__name__) @@ -2052,33 +2051,18 @@ class Scheduler(SchedulerInterface): request.status = RequestStatus.FINISHED_STOPPED stopped = True - if new_token_ids and self.structured_output_manager.should_advance( - request, new_token_ids=new_token_ids + if new_token_ids and not self.structured_output_manager.accept_tokens( + request, new_token_ids ): - struct_output_request = request.structured_output_request - assert struct_output_request is not None - grammar = struct_output_request.grammar - assert isinstance(grammar, StructuredOutputGrammar) - # new_token_ids can be a mixed block of reasoning content, then - # the reasoning end marker, then the start of the grammar content. - # Trim the reasoning content so the grammar only sees grammar content. - advance_token_ids = ( - self.structured_output_manager.trim_reasoning_for_advance( - request, new_token_ids - ) + logger.error( + "Unexpected: grammar rejected tokens %s for request %s. " + "Terminating request.", + new_token_ids, + req_id, ) - if advance_token_ids and not grammar.accept_tokens( - req_id, advance_token_ids - ): - logger.error( - "Unexpected: grammar rejected tokens %s for request %s. " - "Terminating request.", - advance_token_ids, - req_id, - ) - request.status = RequestStatus.FINISHED_ERROR - request.resumable = False - stopped = True + request.status = RequestStatus.FINISHED_ERROR + request.resumable = False + stopped = True routed_experts = None if ( @@ -2434,11 +2418,9 @@ class Scheduler(SchedulerInterface): continue # Add newly generated spec token ids to the request. - if self.structured_output_manager.should_advance(request): - metadata = request.structured_output_request - spec_token_ids = strip_speculative_padding(spec_token_ids) - spec_token_ids = metadata.grammar.validate_tokens(spec_token_ids) # type: ignore[union-attr] - request.spec_token_ids = spec_token_ids + request.spec_token_ids = self.structured_output_manager.validate_tokens( + request, spec_token_ids + ) def update_draft_token_ids_in_output( self, draft_token_ids: DraftTokenIds, scheduler_output: SchedulerOutput @@ -2464,10 +2446,9 @@ class Scheduler(SchedulerInterface): # (needed for chunked prefill case for example). del spec_token_ids[orig_num_spec_tokens:] # Filter out spec tokens which do not adhere to the grammar. - if self.structured_output_manager.should_advance(request): - metadata = request.structured_output_request - spec_token_ids = strip_speculative_padding(spec_token_ids) - spec_token_ids = metadata.grammar.validate_tokens(spec_token_ids) # type: ignore[union-attr] + spec_token_ids = self.structured_output_manager.validate_tokens( + request, spec_token_ids + ) # Pad to original number of spec tokens. num_invalid_tokens = orig_num_spec_tokens - len(spec_token_ids) if num_invalid_tokens: diff --git a/vllm/v1/structured_output/__init__.py b/vllm/v1/structured_output/__init__.py index ca6adc736d..7fd311c873 100644 --- a/vllm/v1/structured_output/__init__.py +++ b/vllm/v1/structured_output/__init__.py @@ -1,6 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import itertools import multiprocessing from collections.abc import Iterable, Sequence from concurrent.futures import Future, ThreadPoolExecutor @@ -18,6 +17,7 @@ from vllm.v1.structured_output.backend_types import ( StructuredOutputGrammar, ) from vllm.v1.structured_output.backend_xgrammar import XgrammarBackend +from vllm.v1.structured_output.utils import strip_speculative_padding if TYPE_CHECKING: import numpy as np @@ -217,6 +217,100 @@ class StructuredOutputManager: ) -> Future: return self.executor_for_fillmask.submit(self._fill_bitmasks, batch) + def _get_constraint_start( + self, + request: "Request", + spec_tokens: Sequence[int], + spec_tokens_committed: bool = False, + ) -> int: + """Return the index into `spec_tokens` where tokens should start being + constrained by the grammar. + Assumes `spec_tokens` are stripped of -1 padding tokens. + Returns `len(spec_tokens) + 1` if no token should be constrained after + accepting all `spec_tokens`. + + `spec_tokens_committed` is True if spec_tokens is already in + `request.all_token_ids`. + """ + if self.enable_in_reasoning: + return 0 + + structured_req = request.structured_output_request + assert structured_req is not None + if structured_req.reasoning_ended: + return 0 + + reasoner = self._get_reasoner(request) + if reasoner is None: + return 0 + + if structured_req.reasoning_ended is None: + # This should be removed here, but since `openai_gptoss` + # is an independent code path, it is kept for now. + # After unifying the `openai_gptoss` and non-`openai_gptoss` styles, + # it can be removed. + if reasoner.is_reasoning_end(request.prompt_token_ids or []): + return 0 + structured_req.reasoning_ended = False + + num_spec_tokens = len(spec_tokens) + if num_spec_tokens <= 0: + return num_spec_tokens + 1 + + # Use `find_reasoning_end_offset` to find constraint start if supported + if ( + isinstance(reasoner, ParserEngineReasoningAdapter) + and reasoner.reasoning_end_token_ids + ): + offset = reasoner.find_reasoning_end_offset(spec_tokens) + if offset is not None: + return offset + 1 + + # Fallback to `find_reasoning_end_offset` + # TODO: Build a read-only Sequence view over all_token_ids + # instead of copying the entire all_token_ids into input_ids. + if spec_tokens_committed: + if not spec_tokens or not reasoner.is_reasoning_end_streaming( + request.all_token_ids, spec_tokens + ): + return num_spec_tokens + 1 + + input_ids = request.all_token_ids.copy() + delta_ids = list(spec_tokens) + else: + input_ids = request.all_token_ids.copy() + input_ids.extend(spec_tokens) + if not reasoner.is_reasoning_end_streaming(input_ids, spec_tokens): + return num_spec_tokens + 1 + delta_ids = list(spec_tokens) + + for i in range(num_spec_tokens - 1, 0, -1): + input_ids.pop() + delta_ids.pop() + if not reasoner.is_reasoning_end_streaming(input_ids, delta_ids): + return i + 1 + return 1 + + def validate_tokens(self, request: "Request", spec_tokens: list[int]) -> list[int]: + """Return the longest unconstrained or grammar-valid prefix of `spec_tokens`.""" + if not request.use_structured_output: + return spec_tokens + + spec_tokens = strip_speculative_padding(spec_tokens) + constraint_start = self._get_constraint_start(request, spec_tokens) + if constraint_start >= len(spec_tokens): + return spec_tokens + + structured_req = request.structured_output_request + if TYPE_CHECKING: + assert structured_req is not None + grammar = structured_req.grammar + if TYPE_CHECKING: + assert isinstance(grammar, StructuredOutputGrammar) + prefix = spec_tokens[:constraint_start] + validated = grammar.validate_tokens(spec_tokens[constraint_start:]) + return prefix + validated + def grammar_bitmask( self, requests: dict[str, "Request"], @@ -264,7 +358,7 @@ class StructuredOutputManager: if TYPE_CHECKING: assert isinstance(grammar, StructuredOutputGrammar) - apply_bitmask = self.should_fill_bitmask(request) + apply_bitmask = self._get_constraint_start(request, ()) == 0 batch.append((grammar, cumulative_index, apply_bitmask)) if len(batch) == self.fill_bitmask_parallel_batch_size: promises.append(self._async_submit_fill_bitmask(batch)) @@ -288,76 +382,47 @@ class StructuredOutputManager: grammar = structured_output_request.grammar if TYPE_CHECKING: assert isinstance(grammar, StructuredOutputGrammar) - apply_bitmask = self.should_fill_bitmask(request) - - reasoner = None if apply_bitmask else self._get_reasoner(request) + req_tokens = scheduled_spec_decode_tokens.get(req_id, list()) + constraint_start = self._get_constraint_start( + request, strip_speculative_padding(req_tokens) + ) state_advancements = 0 - post_reasoning_end_in_window = False - req_tokens = scheduled_spec_decode_tokens.get(req_id, ()) - # Locate the reasoning-end marker in the draft window once. - # Rejected-draft padding (-1) is trailing, so the window is - # cut there and the reasoner never sees a placeholder. - window_start = len(request.all_token_ids) - reasoning_end_index: int | None = None - if reasoner is not None and req_tokens: - window = req_tokens - if -1 in window: - window = window[: window.index(-1)] - reasoning_end_index = self._find_reasoning_end_index( - reasoner, - request.all_token_ids, - window_start, - window, - delta_appended=False, - ) + seen_padding = False + failed = False for i, token in enumerate(req_tokens): + apply_bitmask = ( + not failed and not seen_padding and i >= constraint_start + ) self._fill_bitmasks(((grammar, cumulative_index, apply_bitmask),)) - advance_grammar = apply_bitmask if token == -1: - apply_bitmask = False - advance_grammar = False - elif not apply_bitmask and window_start + i == reasoning_end_index: - # Reasoning ended mid-window. Constrain the rest - # of the window via bitmask. Skip grammar advance - # through the marker (it is reasoning content); - # try to advance through subsequent drafts so the - # next bitmask row reflects the post-advance state, - # but tolerate rejection since those drafts predate - # the bitmask and are not guaranteed valid. - apply_bitmask = True - advance_grammar = False - post_reasoning_end_in_window = True - if advance_grammar and not grammar.is_terminated(): - if post_reasoning_end_in_window: - accepted = bool(grammar.validate_tokens([token])) - if accepted: - accepted = grammar.accept_tokens(req_id, [token]) - else: - accepted = grammar.accept_tokens(req_id, [token]) - if accepted: + seen_padding = True + elif apply_bitmask: + if not grammar.is_terminated() and grammar.accept_tokens( + req_id, [token] + ): state_advancements += 1 - elif not post_reasoning_end_in_window: - raise AssertionError( - (token, req_id, scheduled_spec_decode_tokens) + else: + failed = True + logger.error( + "Unexpected: grammar terminated or rejected draft " + "token %s for request %s during bitmask fill.", + token, + req_id, ) cumulative_index += 1 + # Diffusion LLMs don't sample a bonus token after the # scheduled positions, so skip its bitmask in that case. if not (self.vllm_config.model_config.is_diffusion and req_tokens): - # bonus_apply must be True when the bonus-row position - # should be grammar-constrained. Two triggers: - # - should_fill_bitmask(request): reasoning was already - # over at step start (or no reasoner / - # enable_in_reasoning). - # - apply_bitmask: reasoning ended mid-window in this - # call and was flipped True after the marker; - # should_fill_bitmask still returns False here because - # reasoning_ended is only persisted later by - # should_advance. - bonus_apply = self.should_fill_bitmask(request) or apply_bitmask + bonus_apply = ( + not failed + and not seen_padding + and constraint_start <= len(req_tokens) + ) self._fill_bitmasks(((grammar, cumulative_index, bonus_apply),)) cumulative_index += 1 + if state_advancements > 0: grammar.rollback(state_advancements) @@ -370,161 +435,29 @@ class StructuredOutputManager: # and deserialization when sending this to the GPU workers. return bitmask_tensor.numpy() - def should_fill_bitmask(self, request: "Request") -> bool: - # NOTE (Hanchen) if enable_in_reasoning is True, it means that - # the model needs to be constrained in reasoning. So we should always - # enable the bitmask filling. - structured_req = request.structured_output_request - if self.enable_in_reasoning or ( - structured_req is not None and structured_req.reasoning_ended is True - ): + def accept_tokens(self, request: "Request", new_token_ids: list[int]) -> bool: + """Advance grammar with accepted tokens. Returns False on rejection.""" + if not request.use_structured_output or not new_token_ids: return True - reasoner = self._get_reasoner(request) - if reasoner is not None: - assert structured_req is not None - if structured_req.reasoning_ended is None: - # This should be removed here, but since `openai_gptoss` - # is an independent code path, it is kept for now. - # After unifying the `openai_gptoss` and non-`openai_gptoss` styles, - # it can be removed. - structured_req.reasoning_ended = reasoner.is_reasoning_end( - request.prompt_token_ids or [] - ) - return structured_req.reasoning_ended - return True - - def should_advance( - self, - request: "Request", - new_token_ids: list[int] | None = None, - ) -> bool: - if not request.use_structured_output: - return False - - # To determine whether we can advance the FSM. - # Supports thinking usage where we skip the reasoning components. + structured_req = request.structured_output_request if TYPE_CHECKING: - assert request.structured_output_request is not None - assert request.structured_output_request.grammar is not None - structured_req = request.structured_output_request - if self.enable_in_reasoning or ( - structured_req is not None and structured_req.reasoning_ended is True - ): - return True + assert structured_req is not None + grammar = structured_req.grammar + if TYPE_CHECKING: + assert isinstance(grammar, StructuredOutputGrammar) - # by default, we should always advance - # for cases that don't use thinking mode. - reasoner = self._get_reasoner(request) - if reasoner is None: - return True - - assert structured_req is not None - - # Check if reasoning ends in *this* step. - # When the caller passes new_token_ids (the tokens that were just - # appended this step), use it directly as the delta window. The - # placeholder-derived fallback assumes num_output_placeholders == - # len(new_token_ids), which breaks under async scheduling + spec - # decode when some drafts are rejected (#43388): the placeholder - # count remains > 0 after the step and the computed delta window - # starts past the reasoning-end marker. - all_token_ids = request.all_token_ids - if new_token_ids: - # The tokens were already appended this step, so the step window - # starts exactly len(new_token_ids) from the end. - start = len(all_token_ids) - len(new_token_ids) - delta_ids: Sequence[int] = new_token_ids - else: - # Draft-validation callers arrive here after update_from_output - # has already persisted any boundary found this step, so this - # window only needs to cover tokens not examined yet. - delta_from = request.num_computed_tokens - request.num_output_placeholders - start = ( - delta_from - if delta_from >= 0 - else max(len(all_token_ids) + delta_from, 0) - ) - delta_ids = all_token_ids[start:] - end_index = self._find_reasoning_end_index( - reasoner, all_token_ids, start, delta_ids, delta_appended=True + constraint_start = self._get_constraint_start( + request, new_token_ids, spec_tokens_committed=True ) - if end_index is not None: - structured_req.reasoning_ended = True - # Record the boundary so the scheduler can exclude reasoning tokens. - structured_req.reasoning_end_token_index = end_index + # Early return only when the constraint hasn't started. + # Otherwise, latch structured_req.reasoning_ended. + if constraint_start > len(new_token_ids): return True - - return False - - @staticmethod - def _find_reasoning_end_index( - reasoner: "ReasoningParser", - all_token_ids: Sequence[int], - start: int, - delta_ids: Sequence[int], - *, - delta_appended: bool, - ) -> int | None: - """Locates the token that ends reasoning within ``delta_ids``. - - ``delta_ids`` occupy positions ``start..``; ``delta_appended`` says - whether they are already in ``all_token_ids`` (accepted tokens) or - not (speculative drafts). - - Legacy parsers are probed one token at a time. Accepted tokens get a - whole-delta check as well, and a marker seen only that way reports - the final token. Drafts skip it: some predicates are order-sensitive - over a multi-token window and would miss the marker. - - Returns: - Absolute index of the last reasoning token, or ``None``. - """ - if ( - isinstance(reasoner, ParserEngineReasoningAdapter) - and reasoner.reasoning_end_token_ids - ): - offset = reasoner.find_reasoning_end_offset(delta_ids) - return None if offset is None else start + offset - - if delta_appended and not reasoner.is_reasoning_end_streaming( - all_token_ids, delta_ids - ): - return None - prefix = list(itertools.islice(all_token_ids, start)) - for offset, token in enumerate(delta_ids): - prefix.append(token) - if reasoner.is_reasoning_end_streaming(prefix, [token]): - return start + offset - if delta_appended: - return len(all_token_ids) - 1 - return None - - def trim_reasoning_for_advance( - self, request: "Request", new_token_ids: list[int] - ) -> list[int]: - """Drops reasoning content from tokens about to advance the grammar. - - When reasoning ends mid-step (see should_advance), the step's output - still contains reasoning tokens up to and including the end marker. - Those are not grammar content: feeding them to accept_tokens makes - the grammar reject the marker and kills the request (#44006). - - Returns: - The suffix of ``new_token_ids`` that follows the reasoning-end - marker. Steps fully after the boundary are returned unchanged. - """ - structured_req = request.structured_output_request - if structured_req is None: - return new_token_ids - end_idx = structured_req.reasoning_end_token_index - if end_idx is None: - return new_token_ids - first_idx = len(request.all_token_ids) - len(new_token_ids) - num_reasoning = end_idx + 1 - first_idx - if num_reasoning <= 0: - return new_token_ids - return new_token_ids[num_reasoning:] + structured_req.reasoning_ended = True + return grammar.accept_tokens( + request.request_id, new_token_ids[constraint_start:] + ) def clear_backend(self) -> None: if self.backend is not None: diff --git a/vllm/v1/structured_output/request.py b/vllm/v1/structured_output/request.py index b0f4f3c3a8..dce0038782 100644 --- a/vllm/v1/structured_output/request.py +++ b/vllm/v1/structured_output/request.py @@ -25,12 +25,6 @@ class StructuredOutputRequest: Future[StructuredOutputGrammar] | StructuredOutputGrammar | Exception | None ) = None reasoning_ended: bool | None = None - # Absolute index into the request's all_token_ids of the last reasoning - # token (the reasoning-end marker). Tokens at or before this index are - # reasoning content and must never be fed to the grammar. Only set when - # reasoning ends in a step whose tokens the scheduler advances immediately - # (structural tags + speculative decoding, see #42452). - reasoning_end_token_index: int | None = None reasoning_parser_kwargs: dict[str, Any] | None = None # Cached per request; do not share reasoning parsers across requests because # their behavior can depend on reasoning_parser_kwargs.