[Structured Output] Keep invalid structured-output requests from stopping the engine (#51450)

Co-authored-by: Lucas Bourtoule <35483370+dhalf@users.noreply.github.com>
Co-authored-by: Vadim Gimpelson <156319763+vadiklyutiy@users.noreply.github.com>

tests/v1/core/test_scheduler.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import dataclasses
 from concurrent.futures import Future
+from types import SimpleNamespace
 from unittest.mock import Mock
 
 import numpy as np
@@ -6425,3 +6426,92 @@ def test_encoder_input_skipped_when_connector_already_has_the_item(ec_role: str)
 
     assert output.num_scheduled_tokens[req_id] > 0
     assert not output.scheduled_encoder_inputs.get(req_id)
+
+
+class _RecordingGrammar:
+    """Records what reaches ``validate_tokens`` and raises on a negative id.
+
+    llguidance raises ``OverflowError`` on a negative token id, so this mirrors
+    the real failure: the tests below fail on the unpatched scheduler (the -1
+    padding reaches ``validate_tokens``) and pass once it is stripped first.
+    """
+
+    def __init__(self) -> None:
+        self.seen: list[list[int]] = []
+
+    def validate_tokens(self, tokens: list[int]) -> list[int]:
+        if any(t < 0 for t in tokens):
+            raise OverflowError("llguidance rejects negative token ids")
+        self.seen.append(list(tokens))
+        return tokens
+
+
+def _decode_ready_request(scheduler):
+    """Add one request and advance it out of prefill into decode."""
+    request = create_requests(num_requests=1, num_tokens=1)[0]
+    scheduler.add_request(request)
+    output = scheduler.schedule()
+    scheduler.update_from_output(
+        output,
+        ModelRunnerOutput(
+            req_ids=[request.request_id],
+            req_id_to_index={request.request_id: 0},
+            sampled_token_ids=[[0]],
+            logprobs=None,
+            prompt_logprobs_dict={},
+            pooler_output=[],
+        ),
+    )
+    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
+    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
+    )
+
+    scheduler.update_draft_token_ids(
+        DraftTokenIds([request.request_id], [[10, 11, -1, -1]])
+    )
+
+    assert grammar.seen == [[10, 11]]
+    assert request.spec_token_ids == [10, 11]
+
+
+def test_update_draft_token_ids_in_output_strips_padding(monkeypatch):
+    """Same guard on the output path; the -1 pad-back for the rejected count
+    is preserved (only the input to validate_tokens 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
+    )
+
+    scheduler_output = SimpleNamespace(
+        scheduled_spec_decode_tokens={request.request_id: [0, 0, 0, 0]}
+    )
+    scheduler.update_draft_token_ids_in_output(
+        DraftTokenIds([request.request_id], [[10, 11, -1, -1]]),
+        scheduler_output,
+    )
+
+    # The grammar only saw the stripped prefix, never a -1.
+    assert grammar.seen == [[10, 11]]
+    # Two drafts were rejected (4 scheduled - 2 valid), padded back with -1.
+    assert scheduler_output.scheduled_spec_decode_tokens[request.request_id] == [
+        10,
+        11,
+        -1,
+        -1,
+    ]
+    assert scheduler_output.num_invalid_spec_tokens == {request.request_id: 2}

tests/v1/structured_output/test_guidance_negative_draft_tokens.py

@@ -0,0 +1,61 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Guidance backend must not choke on speculative -1 padding.
+
+Speculative drafters (e.g. ngram_gpu) pad unfilled draft slots with a -1
+sentinel. llguidance raises ``OverflowError`` on a negative id, so
+``GuidanceGrammar.validate_tokens`` must stop at the first sentinel. The mock
+matcher below emulates that raise, so each test fails on the unpatched backend
+and passes with the guard.
+"""
+
+import pytest
+
+from vllm.v1.structured_output.backend_guidance import GuidanceGrammar
+
+pytestmark = pytest.mark.skip_global_cleanup
+
+
+class _Matcher:
+    """Stand-in for ``llguidance.LLMatcher`` that rejects negative ids."""
+
+    def __init__(self) -> None:
+        self.tokens: list[int] | None = None
+
+    def is_stopped(self) -> bool:
+        return False
+
+    def validate_tokens(self, tokens: list[int]) -> int:
+        if any(t < 0 for t in tokens):
+            raise OverflowError("llguidance rejects negative token ids")
+        self.tokens = tokens
+        return len(tokens)
+
+    def get_error(self):
+        return None
+
+
+class _Tokenizer:
+    eos_token = 0
+
+
+def _grammar(matcher: _Matcher) -> GuidanceGrammar:
+    return GuidanceGrammar(ll_matcher=matcher, ll_tokenizer=_Tokenizer(), vocab_size=16)
+
+
+def test_validate_tokens_stops_at_trailing_padding() -> None:
+    matcher = _Matcher()
+    assert _grammar(matcher).validate_tokens([3, 4, -1, -1]) == [3, 4]
+    assert matcher.tokens == [3, 4]
+
+
+def test_validate_tokens_single_negative_does_not_raise() -> None:
+    matcher = _Matcher()
+    assert _grammar(matcher).validate_tokens([-1]) == []
+    assert matcher.tokens is None
+
+
+def test_validate_tokens_all_padding_returns_empty() -> None:
+    matcher = _Matcher()
+    assert _grammar(matcher).validate_tokens([-1, -1]) == []
+    assert matcher.tokens is None

tests/v1/structured_output/test_scheduler_speculative_padding.py

@@ -0,0 +1,20 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm.v1.structured_output.utils import strip_speculative_padding
+
+
+def test_strips_trailing_padding():
+    assert strip_speculative_padding([3, 4, -1, -1]) == [3, 4]
+
+
+def test_empty_when_leading_padding():
+    assert strip_speculative_padding([-1, -1]) == []
+
+
+def test_passthrough_without_padding():
+    assert strip_speculative_padding([3, 4]) == [3, 4]
+
+
+def test_truncates_at_first_sentinel():
+    assert strip_speculative_padding([3, -1, 4]) == [3]

vllm/v1/core/sched/scheduler.py

@@ -71,6 +71,7 @@ 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.utils import record_function_or_nullcontext
 
 logger = init_logger(__name__)
@@ -2424,6 +2425,7 @@ class Scheduler(SchedulerInterface):
             # 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
 
@@ -2453,6 +2455,7 @@ class Scheduler(SchedulerInterface):
             # 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]
             # Pad to original number of spec tokens.
             num_invalid_tokens = orig_num_spec_tokens - len(spec_token_ids)

vllm/v1/structured_output/backend_guidance.py

@@ -21,6 +21,7 @@ from vllm.v1.structured_output.backend_types import (
     StructuredOutputOptions,
 )
 from vllm.v1.structured_output.request import get_structured_output_key
+from vllm.v1.structured_output.utils import strip_speculative_padding
 
 if TYPE_CHECKING:
     import llguidance
@@ -198,6 +199,10 @@ class GuidanceGrammar(StructuredOutputGrammar):
         if self.ll_matcher.is_stopped():
             return []
 
+        tokens = strip_speculative_padding(tokens)
+        if len(tokens) == 0:
+            return []
+
         num_tokens = self.ll_matcher.validate_tokens(tokens)
 
         self.check_error()

vllm/v1/structured_output/utils.py

@@ -46,6 +46,20 @@ _T = TypeVar("_T")
 CACHE = None
 
 
+def strip_speculative_padding(token_ids: list[int]) -> list[int]:
+    """Drop speculative-decoding padding from a token block.
+
+    ngram and other speculative backends pad rejected draft positions with a
+    -1 sentinel. Structured-output grammars treat every entry as a real token
+    id, so the sentinels (and everything after the first one) are removed here,
+    before tokens reach any backend, rather than inside a single backend.
+    """
+    for i, token_id in enumerate(token_ids):
+        if token_id < 0:
+            return token_ids[:i]
+    return token_ids
+
+
 def compile_regex_with_timeout(fn: Callable[[str], _T], pattern: str) -> _T:
     """Run a regex compilation callable with a timeout.