diff --git a/tests/v1/core/test_async_scheduler.py b/tests/v1/core/test_async_scheduler.py index 5e9c9280dbe..8253a8422e2 100644 --- a/tests/v1/core/test_async_scheduler.py +++ b/tests/v1/core/test_async_scheduler.py @@ -272,6 +272,9 @@ def test_abort_request_when_structured_output_fsm_cannot_advance(): scheduler.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.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 dad345c643a..3e2b7dc5832 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -2911,6 +2911,9 @@ def test_abort_request_when_structured_output_fsm_cannot_advance(): scheduler.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.requests = {request.request_id: request} scheduler.running = [request] scheduler.waiting = Mock() diff --git a/tests/v1/spec_decode/test_mtp_structured_output.py b/tests/v1/spec_decode/test_mtp_structured_output.py new file mode 100644 index 00000000000..619f3ad6fde --- /dev/null +++ b/tests/v1/spec_decode/test_mtp_structured_output.py @@ -0,0 +1,343 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""grammar_bitmask under spec-decode draft padding (#44006).""" + +import pytest +from transformers import AutoTokenizer + +from vllm.config import 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" +NUM_SPEC_TOKENS = 4 + + +def _make_manager_and_request(backend: str, prompt_str: str = '{"a": "b"}'): + tokenizer = AutoTokenizer.from_pretrained(TOKENIZER) + prompt = tokenizer.encode(prompt_str) + + vllm_config = VllmConfig( + model_config=ModelConfig(tokenizer=TOKENIZER), + structured_outputs_config=StructuredOutputsConfig(backend=backend), + speculative_config=SpeculativeConfig( + model="[ngram]", num_speculative_tokens=NUM_SPEC_TOKENS + ), + ) + manager = StructuredOutputManager(vllm_config) + + sampling_params = SamplingParams( + structured_outputs=StructuredOutputsParams(json='{"type": "object"}'), + ) + sampling_params.structured_outputs._backend = backend + sampling_params.update_from_generation_config({}, tokenizer.eos_token_id) + + request = Request( + "mtp_req", + prompt_token_ids=prompt, + sampling_params=sampling_params, + pooling_params=None, + ) + manager.grammar_init(request) + while not request.structured_output_request._check_grammar_completion(): + continue + + 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): + """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() + + +@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() + + +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 diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index ea657d7e949..9ce1d94ef3c 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -1636,14 +1636,23 @@ class Scheduler(SchedulerInterface): if new_token_ids and self.structured_output_manager.should_advance(request): struct_output_request = request.structured_output_request assert struct_output_request is not None - assert struct_output_request.grammar is not None - if not struct_output_request.grammar.accept_tokens( # type: ignore[union-attr] - req_id, new_token_ids + grammar = struct_output_request.grammar + assert grammar is not None + # 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 + ) + ) + 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.", - new_token_ids, + advance_token_ids, req_id, ) request.status = RequestStatus.FINISHED_ERROR diff --git a/vllm/v1/structured_output/__init__.py b/vllm/v1/structured_output/__init__.py index 30921f3d74a..34f775257be 100644 --- a/vllm/v1/structured_output/__init__.py +++ b/vllm/v1/structured_output/__init__.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import itertools import multiprocessing -from collections.abc import Iterable +from collections.abc import Iterable, Sequence from concurrent.futures import Future, ThreadPoolExecutor from typing import TYPE_CHECKING @@ -272,23 +272,69 @@ class StructuredOutputManager: grammar = structured_output_request.grammar apply_bitmask = self.should_fill_bitmask(request) + reasoner = self._get_reasoner(request) + detect_reasoning_end = ( + not apply_bitmask + and reasoner is not None + and not self.enable_in_reasoning + ) + simulated_buf: list[int] | None = None + history_len = 0 + state_advancements = 0 + post_reasoning_end_in_window = False req_tokens = scheduled_spec_decode_tokens.get(req_id, ()) - if self.vllm_config.model_config.is_diffusion and req_tokens: - # Diffusion LLMs don't sample a bonus token after the - # scheduled positions, so don't append the -1 placeholder. - token_iter: Iterable[int] = req_tokens - else: - token_iter = itertools.chain(req_tokens, (-1,)) - for token in token_iter: + for i, token in enumerate(req_tokens): self._fill_bitmasks(((grammar, cumulative_index, apply_bitmask),)) + advance_grammar = apply_bitmask if token == -1: - # Stop advancing the grammar once we hit a padding token. apply_bitmask = False - if apply_bitmask and not grammar.is_terminated(): + advance_grammar = False + elif ( + detect_reasoning_end + and reasoner is not None + and not apply_bitmask + ): + if simulated_buf is None: + history = list(request.all_token_ids) + history_len = len(history) + simulated_buf = history + list(req_tokens) + simulated = simulated_buf[: history_len + i + 1] + if reasoner.is_reasoning_end_streaming(simulated, [token]): + # Reasoning ended mid-window. Constrain the rest + # of the window via bitmask. Skip grammar advance + # through the marker (it is reasoning content); + # try to advance through subsequent drafts so the + # next bitmask row reflects the post-advance state, + # but tolerate rejection since those drafts predate + # the bitmask and are not guaranteed valid. + apply_bitmask = True + advance_grammar = False + post_reasoning_end_in_window = True + if advance_grammar and not grammar.is_terminated(): accepted = grammar.accept_tokens(req_id, [token]) - assert accepted, (token, req_id, scheduled_spec_decode_tokens) - state_advancements += 1 + if accepted: + state_advancements += 1 + elif not post_reasoning_end_in_window: + raise AssertionError( + (token, req_id, scheduled_spec_decode_tokens) + ) + 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 + self._fill_bitmasks(((grammar, cumulative_index, bonus_apply),)) cumulative_index += 1 if state_advancements > 0: grammar.rollback(state_advancements) @@ -368,10 +414,64 @@ class StructuredOutputManager: and structured_req.structured_output_key[0] == StructuredOutputOptions.STRUCTURAL_TAG ): + # The scheduler will advance the grammar with this step's + # tokens right away, but the step still contains reasoning + # content up to and including the end marker. Record where + # it ends so trim_reasoning_for_advance() can drop it. + structured_req.reasoning_end_token_index = ( + self._find_reasoning_end_index(reasoner, all_token_ids, start) + ) return True return False + @staticmethod + def _find_reasoning_end_index( + reasoner: "ReasoningParser", all_token_ids: Sequence[int], start: int + ) -> int: + """Locates the last reasoning token within ``all_token_ids[start:]``. + + Returns: + The absolute index of the token at which + ``is_reasoning_end_streaming`` first fires. Falls back to the + final index when no single token triggers the detection (e.g. + a multi-token marker only recognized on the full delta), which + conservatively treats the whole step as reasoning content. + """ + prefix = list(itertools.islice(all_token_ids, start)) + for idx in range(start, len(all_token_ids)): + token = all_token_ids[idx] + prefix.append(token) + if reasoner.is_reasoning_end_streaming(prefix, [token]): + return idx + return len(all_token_ids) - 1 + + 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:] + def clear_backend(self) -> None: if self.backend is not None: self.backend.destroy() diff --git a/vllm/v1/structured_output/request.py b/vllm/v1/structured_output/request.py index dfa8c7efcae..f9ab54a0471 100644 --- a/vllm/v1/structured_output/request.py +++ b/vllm/v1/structured_output/request.py @@ -23,6 +23,12 @@ class StructuredOutputRequest: params: StructuredOutputsParams _grammar: Future[StructuredOutputGrammar] | StructuredOutputGrammar | 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.