[Bugfix][Structured Output][Spec Decode] Constrain bitmask and trim grammar advance at the reasoning boundary (#44297)

Signed-off-by: Allen.Yu <[email protected]>
Signed-off-by: yue.yu <[email protected]>
Co-authored-by: Benjamin Chislett <[email protected]>

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

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

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, </think>]) 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

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

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

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.