diff --git a/tests/parser/engine/test_parser_engine.py b/tests/parser/engine/test_parser_engine.py index 811f9cb7b8..ac0203bc0b 100644 --- a/tests/parser/engine/test_parser_engine.py +++ b/tests/parser/engine/test_parser_engine.py @@ -161,12 +161,15 @@ class TestReasoningEndTokenIds: ) engine = _make_engine(cfg) assert engine.reasoning_end_token_ids == {201, 202} + assert engine.implicit_reasoning_end_token_ids == {202} assert engine.find_reasoning_end([5, 201]) == ReasoningEnd(1, False) assert engine.find_reasoning_end([5, 202]) == ReasoningEnd(1, True) def test_transition_staying_in_reasoning_is_ignored(self): cfg = _with_reasoning_exits(("THINK_START", ParserState.REASONING, ())) - assert _make_engine(cfg).reasoning_end_token_ids == {201} + engine = _make_engine(cfg) + assert engine.reasoning_end_token_ids == {201} + assert engine.implicit_reasoning_end_token_ids == frozenset() def test_unreported_exit_fails_closed(self): cfg = _with_reasoning_exits( diff --git a/tests/v1/spec_decode/test_mtp_structured_output.py b/tests/v1/spec_decode/test_mtp_structured_output.py index 8a72765db1..72fa3e7ad7 100644 --- a/tests/v1/spec_decode/test_mtp_structured_output.py +++ b/tests/v1/spec_decode/test_mtp_structured_output.py @@ -114,6 +114,10 @@ class _EngineReasonerStub(ParserEngineReasoningAdapter): def reasoning_end_token_ids(self): return self._end_token_ids + @property + def implicit_reasoning_end_token_ids(self): + return frozenset() + def find_reasoning_end(self, token_ids): self.windows.append(list(token_ids)) for offset, token in enumerate(token_ids): diff --git a/tests/v1/structured_output/test_structured_output_manager.py b/tests/v1/structured_output/test_structured_output_manager.py index 1d1e7b78df..0cb051f0fd 100644 --- a/tests/v1/structured_output/test_structured_output_manager.py +++ b/tests/v1/structured_output/test_structured_output_manager.py @@ -6,6 +6,7 @@ from concurrent.futures import Future, ThreadPoolExecutor from dataclasses import dataclass from threading import Event +import numpy as np import pytest from transformers import AutoTokenizer @@ -72,20 +73,26 @@ class MockReasoner: class MockEngineReasoner(ParserEngineReasoningAdapter): - """Ends reasoning implicitly on ``marker``, which is also content.""" + """Ends reasoning implicitly on ``marker``, which is also content, and + explicitly on ``end_marker`` when one is given.""" - def __init__(self, tokenizer, marker: int): + def __init__(self, tokenizer, marker: int, end_marker: int | None = None): self.marker = marker + self.end_marker = end_marker @property def reasoning_end_token_ids(self): + return frozenset({self.marker, self.end_marker} - {None}) + + @property + def implicit_reasoning_end_token_ids(self): return frozenset({self.marker}) def find_reasoning_end(self, token_ids): - ids = list(token_ids) - if self.marker in ids: - return ReasoningEnd(ids.index(self.marker), True) - return ReasoningEnd(len(ids), False) + for offset, token in enumerate(token_ids): + if token in self.reasoning_end_token_ids: + return ReasoningEnd(offset, token == self.marker) + return ReasoningEnd(len(token_ids), False) def _single_token(tokenizer, text: str) -> int: @@ -122,6 +129,7 @@ def _build_harness( enable_in_reasoning: bool = False, reasoning_parser_kwargs: dict | None = None, structured_outputs: StructuredOutputsParams | None = None, + num_spec_tokens: int = NUM_SPEC_TOKENS, ) -> tuple[StructuredOutputManager, Request]: vllm_config = VllmConfig( model_config=ModelConfig(tokenizer=TOKENIZER), @@ -132,8 +140,10 @@ def _build_harness( ), speculative_config=SpeculativeConfig( model="[ngram]", - num_speculative_tokens=NUM_SPEC_TOKENS, - ), + num_speculative_tokens=num_spec_tokens, + ) + if num_spec_tokens + else None, ) manager = StructuredOutputManager(vllm_config) if use_reasoner: @@ -171,8 +181,20 @@ def _build_harness( return manager, request -def _row_pattern(bitmask) -> str: - return "".join("U" if (row == -1).all() else "C" for row in bitmask) +def _row_pattern(bitmask, masked_ids=()) -> str: + """U: unconstrained; R: all but `masked_ids` allowed; C: grammar rows.""" + reasoning_row = np.full(bitmask.shape[1], -1, dtype=np.int32) + for token_id in masked_ids: + reasoning_row.view(np.uint32)[token_id >> 5] &= ~np.uint32(1 << (token_id & 31)) + + def kind(row) -> str: + if (row == -1).all(): + return "U" + if masked_ids and (row == reasoning_row).all(): + return "R" + return "C" + + return "".join(kind(row) for row in bitmask) def _run_real_flow( @@ -183,6 +205,7 @@ def _run_real_flow( expected_row_pattern: str, expected_reasoning: bool | None, expect_terminated: bool, + masked_ids: list[int] | None = None, ) -> None: assert len(expected_validated) <= len(raw_drafts) @@ -197,7 +220,7 @@ def _run_real_flow( ) assert bitmask is not None assert bitmask.shape[0] == len(raw_drafts) + 1 - assert _row_pattern(bitmask) == expected_row_pattern + assert _row_pattern(bitmask, masked_ids or ()) == expected_row_pattern # grammar_bitmask() must rollback any speculative state it advanced. assert manager.validate_tokens(request, list(raw_drafts)) == expected_validated @@ -375,7 +398,8 @@ FLOW_CASES = [ FlowCase( raw_drafts=(" ", IMPLICIT_THINK_END, "{", "}"), expected_validated=(" ",), - expected_row_pattern="UUUUU", + # The grammar cannot start with the marker, so reasoning rows mask it. + expected_row_pattern="RRUUU", expected_reasoning=False, expect_terminated=False, engine_reasoner=True, @@ -421,9 +445,93 @@ def test_real_flow( expected_row_pattern=case.expected_row_pattern, expected_reasoning=case.expected_reasoning, expect_terminated=case.expect_terminated, + masked_ids=[reasoner_kwargs["marker"]] if case.engine_reasoner else None, ) +@pytest.mark.parametrize("backend", BACKENDS) +@pytest.mark.parametrize( + ("structured_outputs", "raw_drafts", "expected_row_pattern"), + [ + # JSON cannot start with the marker: reasoning rows mask it and keep + # every other token, including the explicit end, allowed. + pytest.param( + StructuredOutputsParams(json=JSON_SCHEMA), (" ", " "), "RRR", id="json" + ), + # The grammar can start with the marker: it stays allowed and, once + # sampled, is the first token the grammar consumes. + pytest.param( + StructuredOutputsParams(regex="z[0-9]+"), + (" ", IMPLICIT_THINK_END, "1"), + "UUCC", + id="regex_starts_with_marker", + ), + ], +) +def test_reasoning_rows_mask_grammar_invalid_implicit_end( + tokenizer, + backend: StructuredOutputsBackend, + structured_outputs: StructuredOutputsParams, + raw_drafts: tuple[str, ...], + expected_row_pattern: str, +): + marker = _single_token(tokenizer, IMPLICIT_THINK_END) + manager, request = _build_harness( + tokenizer, + backend, + reasoning_ended=False, + reasoning_parser_kwargs={ + "marker": marker, + "end_marker": _single_token(tokenizer, THINK_END), + }, + structured_outputs=structured_outputs, + ) + manager.reasoner_cls = MockEngineReasoner + drafts = _to_token_ids(tokenizer, raw_drafts) + assert manager.validate_tokens(request, drafts) == 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 _row_pattern(bitmask, [marker]) == expected_row_pattern + + +@pytest.mark.parametrize("backend", BACKENDS) +def test_parallel_fill_masks_grammar_invalid_implicit_end( + tokenizer, backend: StructuredOutputsBackend +): + marker, end = _to_token_ids(tokenizer, (IMPLICIT_THINK_END, THINK_END)) + manager, request = _build_harness( + tokenizer, + backend, + reasoning_ended=False, + reasoning_parser_kwargs={"marker": marker, "end_marker": end}, + num_spec_tokens=0, + ) + manager.reasoner_cls = MockEngineReasoner + + def row_pattern() -> str: + 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 + return _row_pattern(bitmask, [marker]) + + manager.fill_bitmask_parallel_threshold = 0 + manager.fill_bitmask_parallel_batch_size = 16 + with ThreadPoolExecutor(max_workers=1) as executor: + manager.executor_for_fillmask = executor + assert row_pattern() == "R" + request.append_output_token_ids([end]) + assert manager.accept_tokens(request, [end]) + assert row_pattern() == "C" + + @pytest.mark.parametrize("backend", BACKENDS) @pytest.mark.parametrize( ("use_reasoner", "reasoning_ended", "enable_in_reasoning"), diff --git a/vllm/parser/engine/adapters.py b/vllm/parser/engine/adapters.py index cd7aaa7297..a4a307a43a 100644 --- a/vllm/parser/engine/adapters.py +++ b/vllm/parser/engine/adapters.py @@ -115,6 +115,10 @@ class ParserEngineReasoningAdapter(ReasoningParser): def reasoning_end_token_ids(self) -> frozenset[int]: return self._parser_engine.reasoning_end_token_ids + @property + def implicit_reasoning_end_token_ids(self) -> frozenset[int]: + return self._parser_engine.implicit_reasoning_end_token_ids + def adjust_request( self, request: ChatCompletionRequest | ResponsesRequest, diff --git a/vllm/parser/engine/parser_engine.py b/vllm/parser/engine/parser_engine.py index 95df03eab5..bbb3605ed7 100644 --- a/vllm/parser/engine/parser_engine.py +++ b/vllm/parser/engine/parser_engine.py @@ -678,6 +678,11 @@ class ParserEngine(Parser): def reasoning_end_token_ids(self) -> frozenset[int]: return self._reasoning_end_token_ids + @property + def implicit_reasoning_end_token_ids(self) -> frozenset[int]: + """Reasoning-end tokens that are also content, e.g. a tool-call start.""" + return self._reasoning_end_token_ids - {self._reasoning_end_token_id} + def find_reasoning_end(self, token_ids: Sequence[int]) -> ReasoningEnd | None: """First reasoning end in `token_ids`, and if its token is content.""" end_ids = self._reasoning_end_token_ids diff --git a/vllm/v1/structured_output/__init__.py b/vllm/v1/structured_output/__init__.py index 3f461ed47f..7de9dc2216 100644 --- a/vllm/v1/structured_output/__init__.py +++ b/vllm/v1/structured_output/__init__.py @@ -5,6 +5,8 @@ from collections.abc import Iterable, Sequence from concurrent.futures import Future, ThreadPoolExecutor from typing import TYPE_CHECKING, NamedTuple +import numpy as np + from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.parser.engine.adapters import ParserEngineReasoningAdapter @@ -20,7 +22,6 @@ 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 import numpy.typing as npt import torch @@ -211,10 +212,16 @@ class StructuredOutputManager: raise def _fill_bitmasks( - self, batch: Iterable[tuple[StructuredOutputGrammar, int, bool]] + self, + batch: Iterable[tuple[StructuredOutputGrammar, int, bool, tuple[int, ...]]], ) -> None: + """Fill one bitmask row per entry. + + A row the grammar does not constrain allows every token except + `masked_ids`, the implicit reasoning ends masked while reasoning. + """ assert self._grammar_bitmask is not None - for grammar, index, apply_bitmask in batch: + for grammar, index, apply_bitmask, masked_ids in batch: if apply_bitmask and not grammar.is_terminated(): grammar.fill_bitmask(self._grammar_bitmask, index) else: @@ -222,12 +229,41 @@ class StructuredOutputManager: # reset the relevant part of the bitmask for consequent # requests here. self._grammar_bitmask[index].fill_(self._full_mask) + if masked_ids: + row = self._grammar_bitmask[index].numpy().view(np.uint32) + for token_id in masked_ids: + row[token_id >> 5] &= ~np.uint32(1 << (token_id & 31)) def _async_submit_fill_bitmask( - self, batch: list[tuple[StructuredOutputGrammar, int, bool]] + self, batch: list[tuple[StructuredOutputGrammar, int, bool, tuple[int, ...]]] ) -> Future: return self.executor_for_fillmask.submit(self._fill_bitmasks, batch) + def _masked_reasoning_ends(self, request: "Request") -> tuple[int, ...]: + """Implicit reasoning-end tokens the grammar rejects as first content. + + Such a token (e.g. a tool-call start) ends reasoning and is also the + first content token, so sampling one where the grammar cannot start + with it would make the output invalid. These stay masked while the + request reasons. The grammar does not advance before reasoning ends, + so the set is fixed per request. + """ + structured_req = request.structured_output_request + assert structured_req is not None + if structured_req.masked_reasoning_end_ids is None: + masked: tuple[int, ...] = () + reasoner = self._get_reasoner(request) + if isinstance(reasoner, ParserEngineReasoningAdapter): + grammar = structured_req.grammar + assert isinstance(grammar, StructuredOutputGrammar) + masked = tuple( + token_id + for token_id in sorted(reasoner.implicit_reasoning_end_token_ids) + if not grammar.validate_tokens([token_id]) + ) + structured_req.masked_reasoning_end_ids = masked + return structured_req.masked_reasoning_end_ids + def _get_constraint_bounds( self, request: "Request", @@ -377,7 +413,10 @@ class StructuredOutputManager: apply_bitmask = ( self._get_constraint_bounds(request, ()).constraint_start == 0 ) - batch.append((grammar, cumulative_index, apply_bitmask)) + masked_ids = ( + () if apply_bitmask else self._masked_reasoning_ends(request) + ) + batch.append((grammar, cumulative_index, apply_bitmask, masked_ids)) if len(batch) == self.fill_bitmask_parallel_batch_size: promises.append(self._async_submit_fill_bitmask(batch)) batch = [] @@ -405,6 +444,10 @@ class StructuredOutputManager: grammar_start, constraint_start = self._get_constraint_bounds( request, strip_speculative_padding(req_tokens) ) + # Rows before constraint_start are still reasoning. + masked_ids = ( + self._masked_reasoning_ends(request) if constraint_start else () + ) state_advancements = 0 seen_padding = False # Row filled from the last valid grammar state before a draft @@ -417,7 +460,11 @@ class StructuredOutputManager: cumulative_index += 1 continue apply_bitmask = not seen_padding and i >= constraint_start - self._fill_bitmasks(((grammar, cumulative_index, apply_bitmask),)) + reasoning = not seen_padding and i < constraint_start + row_masked = masked_ids if reasoning else () + self._fill_bitmasks( + ((grammar, cumulative_index, apply_bitmask, row_masked),) + ) if token == -1: seen_padding = True elif not seen_padding and i >= grammar_start: @@ -444,7 +491,11 @@ class StructuredOutputManager: bonus_apply = not seen_padding and constraint_start <= len( req_tokens ) - self._fill_bitmasks(((grammar, cumulative_index, bonus_apply),)) + bonus_reasoning = not seen_padding and not bonus_apply + bonus_masked = masked_ids if bonus_reasoning else () + self._fill_bitmasks( + ((grammar, cumulative_index, bonus_apply, bonus_masked),) + ) cumulative_index += 1 if state_advancements > 0: diff --git a/vllm/v1/structured_output/request.py b/vllm/v1/structured_output/request.py index cd044f5466..2091cc12f9 100644 --- a/vllm/v1/structured_output/request.py +++ b/vllm/v1/structured_output/request.py @@ -28,6 +28,9 @@ class StructuredOutputRequest: # Cached per request; do not share reasoning parsers across requests because # their behavior can depend on reasoning_parser_kwargs. reasoner: "ReasoningParser | None" = None + # Implicit reasoning-end tokens the grammar rejects as the first content + # token; masked while the request reasons. Computed on first use. + masked_reasoning_end_ids: tuple[int, ...] | None = None @staticmethod def from_sampling_params(