[Bugfix][Structured Output] Mask grammar-invalid implicit reasoning ends (#61574)

Signed-off-by: Maroon Ayoub <mayoub@redhat.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>

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(

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):

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"),

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,

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

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:

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(