From c55e15a44ec4127832d4a86928a356fdd9e68dbd Mon Sep 17 00:00:00 2001 From: Clinton Thomas <1033162+KernelClint@users.noreply.github.com> Date: Wed, 9 Sep 2026 11:32:01 -0400 Subject: [PATCH] [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 | 90 +++++++++++++++++++ .../test_guidance_negative_draft_tokens.py | 61 +++++++++++++ .../test_scheduler_speculative_padding.py | 20 +++++ vllm/v1/core/sched/scheduler.py | 3 + vllm/v1/structured_output/backend_guidance.py | 5 ++ vllm/v1/structured_output/utils.py | 14 +++ 6 files changed, 193 insertions(+) create mode 100644 tests/v1/structured_output/test_guidance_negative_draft_tokens.py create mode 100644 tests/v1/structured_output/test_scheduler_speculative_padding.py diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 7a6d9f7887..1c56eaa6a5 100644 --- a/tests/v1/core/test_scheduler.py +++ b/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} diff --git a/tests/v1/structured_output/test_guidance_negative_draft_tokens.py b/tests/v1/structured_output/test_guidance_negative_draft_tokens.py new file mode 100644 index 0000000000..b354a4b63f --- /dev/null +++ b/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 diff --git a/tests/v1/structured_output/test_scheduler_speculative_padding.py b/tests/v1/structured_output/test_scheduler_speculative_padding.py new file mode 100644 index 0000000000..a9375d32f5 --- /dev/null +++ b/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] diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 275dc77b6a..838358a364 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/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) diff --git a/vllm/v1/structured_output/backend_guidance.py b/vllm/v1/structured_output/backend_guidance.py index 310576cd99..16e1f54be6 100644 --- a/vllm/v1/structured_output/backend_guidance.py +++ b/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() diff --git a/vllm/v1/structured_output/utils.py b/vllm/v1/structured_output/utils.py index 0633c55fe6..aa5d59c9c5 100644 --- a/vllm/v1/structured_output/utils.py +++ b/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.