[Refactor] StructuredOutputManager x Speculative Decoding Refactor (#48200)
This commit is contained in:
yzong-rh
@@ -574,7 +574,7 @@ class TestImplicitReasoningEnd:
|
||||
# the text lexer.
|
||||
assert thinking_parser.reasoning_end_token_ids == {_THINK_END_ID}
|
||||
assert thinking_parser.find_reasoning_end_offset([7, _THINK_END_ID, 8]) == 1
|
||||
assert thinking_parser.find_reasoning_end_offset([7, 8]) is None
|
||||
assert thinking_parser.find_reasoning_end_offset([7, 8]) == 2
|
||||
|
||||
def test_streaming_reasoning_implicit_end(self, thinking_parser):
|
||||
chunks = [
|
||||
|
||||
@@ -185,8 +185,8 @@ class TestReasoningEndTokenIds:
|
||||
def test_find_reasoning_end_offset_returns_first_match(self):
|
||||
engine = _make_engine()
|
||||
assert engine.find_reasoning_end_offset([5, 201, 6, 201]) == 1
|
||||
assert engine.find_reasoning_end_offset([5, 6]) is None
|
||||
assert engine.find_reasoning_end_offset([]) is None
|
||||
assert engine.find_reasoning_end_offset([5, 6]) == 2
|
||||
assert engine.find_reasoning_end_offset([]) == 0
|
||||
# Rejected-draft placeholders never match.
|
||||
assert engine.find_reasoning_end_offset([-1, 201]) == 1
|
||||
|
||||
|
||||
@@ -334,7 +334,6 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
|
||||
request = create_requests(num_requests=1, num_tokens=1)[0]
|
||||
request.structured_output_request = Mock()
|
||||
request.structured_output_request.grammar = Mock(spec=StructuredOutputGrammar)
|
||||
request.structured_output_request.grammar.accept_tokens.return_value = False
|
||||
request.status = RequestStatus.RUNNING
|
||||
request.num_computed_tokens = request.num_tokens
|
||||
request.num_output_placeholders = 1
|
||||
@@ -343,10 +342,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
|
||||
scheduler.connector = None
|
||||
scheduler.ec_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.structured_output_manager.accept_tokens.return_value = False
|
||||
scheduler.requests = {request.request_id: request}
|
||||
scheduler.running = [request]
|
||||
scheduler.waiting = Mock()
|
||||
|
||||
@@ -3674,7 +3674,6 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
|
||||
)
|
||||
request.structured_output_request = Mock()
|
||||
request.structured_output_request.grammar = Mock(spec=StructuredOutputGrammar)
|
||||
request.structured_output_request.grammar.accept_tokens.return_value = False
|
||||
request.status = RequestStatus.RUNNING
|
||||
request.num_computed_tokens = request.num_tokens
|
||||
|
||||
@@ -3682,10 +3681,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
|
||||
scheduler.connector = None
|
||||
scheduler.ec_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.structured_output_manager.accept_tokens.return_value = False
|
||||
scheduler.requests = {request.request_id: request}
|
||||
scheduler.running = [request]
|
||||
scheduler.waiting = Mock()
|
||||
@@ -3734,8 +3730,8 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
|
||||
)
|
||||
engine_core_outputs = scheduler.update_from_output(output, model_runner_output)
|
||||
|
||||
request.structured_output_request.grammar.accept_tokens.assert_called_once_with(
|
||||
request.request_id, [123]
|
||||
scheduler.structured_output_manager.accept_tokens.assert_called_once_with(
|
||||
request, [123]
|
||||
)
|
||||
assert request.resumable is False
|
||||
assert request.status == RequestStatus.FINISHED_ERROR
|
||||
@@ -6573,16 +6569,15 @@ def _decode_ready_request(scheduler):
|
||||
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
|
||||
def test_update_draft_token_ids_strips_ngram_padding():
|
||||
"""ngram_gpu pads unfilled draft slots with -1; manager validation 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
|
||||
request.structured_output_request = SimpleNamespace(
|
||||
grammar=grammar, reasoning_ended=True
|
||||
)
|
||||
|
||||
scheduler.update_draft_token_ids(
|
||||
@@ -6593,16 +6588,15 @@ def test_update_draft_token_ids_strips_ngram_padding(monkeypatch):
|
||||
assert request.spec_token_ids == [10, 11]
|
||||
|
||||
|
||||
def test_update_draft_token_ids_in_output_strips_padding(monkeypatch):
|
||||
def test_update_draft_token_ids_in_output_strips_padding():
|
||||
"""Same guard on the output path; the -1 pad-back for the rejected count
|
||||
is preserved (only the input to validate_tokens is stripped)."""
|
||||
is preserved (only the input to manager validation 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
|
||||
request.structured_output_request = SimpleNamespace(
|
||||
grammar=grammar, reasoning_ended=True
|
||||
)
|
||||
|
||||
scheduler_output = SimpleNamespace(
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""grammar_bitmask under spec-decode draft padding (#44006)."""
|
||||
"""Backend and engine-reasoner checks that are not manager-flow tests."""
|
||||
|
||||
import pytest
|
||||
from transformers import AutoTokenizer
|
||||
@@ -49,221 +49,6 @@ def _make_manager_and_request(backend: str, prompt_str: str = '{"a": "b"}'):
|
||||
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, caplog):
|
||||
"""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()
|
||||
assert "Failed to advance FSM" not in caplog.text
|
||||
|
||||
|
||||
@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()
|
||||
|
||||
|
||||
def test_xgrammar_accept_tokens_stops_at_termination(capfd):
|
||||
"""Tokens after a terminating EOS do not reach the matcher."""
|
||||
tokenizer, _, request, prompt = _make_manager_and_request("xgrammar")
|
||||
@@ -302,7 +87,6 @@ def test_xgrammar_validate_tokens_stops_at_termination(capfd):
|
||||
|
||||
assert grammar.validate_tokens([eos, trailing]) == [eos]
|
||||
assert "trying to accept new token" not in capfd.readouterr().err
|
||||
# Check matcher state directly to verify validation rolled it back.
|
||||
assert not grammar.matcher.is_terminated()
|
||||
|
||||
assert grammar.accept_tokens(request.request_id, [eos])
|
||||
@@ -312,87 +96,6 @@ def test_xgrammar_validate_tokens_stops_at_termination(capfd):
|
||||
assert "trying to accept new token" not in capfd.readouterr().err
|
||||
|
||||
|
||||
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
|
||||
|
||||
|
||||
class _EngineReasonerStub(ParserEngineReasoningAdapter):
|
||||
"""Adapter-typed reasoner with a fixed end-token set and no real engine."""
|
||||
|
||||
@@ -409,7 +112,7 @@ class _EngineReasonerStub(ParserEngineReasoningAdapter):
|
||||
for offset, token in enumerate(token_ids):
|
||||
if token in self._end_token_ids:
|
||||
return offset
|
||||
return None
|
||||
return len(token_ids)
|
||||
|
||||
def is_reasoning_end(self, input_ids):
|
||||
return any(token in self._end_token_ids for token in input_ids)
|
||||
@@ -449,61 +152,3 @@ def test_bitmask_engine_reasoner_ends_midwindow_with_padding(backend):
|
||||
assert not (bitmask[2] == -1).all()
|
||||
assert reasoner.windows == [[pre, marker, post]]
|
||||
assert not grammar.is_terminated()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", ["xgrammar", "guidance"])
|
||||
def test_bitmask_legacy_order_sensitive_predicate_probed_per_token(backend):
|
||||
"""Legacy predicates are probed one draft at a time, never on the whole
|
||||
window. A predicate that reports which marker is *newest* (KimiK3) would
|
||||
otherwise answer False for a window that closes and then reopens
|
||||
reasoning, leaving the post-marker drafts and the bonus row
|
||||
unconstrained.
|
||||
"""
|
||||
tokenizer, manager, request, prompt = _make_manager_and_request(backend)
|
||||
grammar = request.structured_output_request.grammar
|
||||
|
||||
assert grammar.accept_tokens(request.request_id, prompt)
|
||||
|
||||
close = tokenizer.encode("\n")[0]
|
||||
reopen = tokenizer.encode("#")[0]
|
||||
assert close != reopen
|
||||
|
||||
class NewestMarkerReasoner:
|
||||
def __init__(self, *_, **__):
|
||||
pass
|
||||
|
||||
def is_reasoning_end(self, input_ids):
|
||||
return self.is_reasoning_end_streaming(input_ids, input_ids)
|
||||
|
||||
def is_reasoning_end_streaming(self, input_ids, delta_ids):
|
||||
for token in reversed(list(delta_ids)):
|
||||
if token == close:
|
||||
return True
|
||||
if token == reopen:
|
||||
return False
|
||||
return False
|
||||
|
||||
manager.reasoner_cls = NewestMarkerReasoner
|
||||
request.structured_output_request.reasoner = NewestMarkerReasoner()
|
||||
request.structured_output_request.reasoning_ended = False
|
||||
|
||||
pre = tokenizer.encode(" ")[0]
|
||||
post = tokenizer.encode(",")[0]
|
||||
drafts = [pre, close, post, reopen]
|
||||
# Sanity: the whole window says "reopened", so only per-token probing
|
||||
# can find the close marker.
|
||||
assert not NewestMarkerReasoner().is_reasoning_end_streaming([], 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 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()
|
||||
|
||||
@@ -1,557 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""Unit tests for reasoning-aware structured output functionality (PR #25515)."""
|
||||
|
||||
from collections.abc import Sequence
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.config import ModelConfig, SchedulerConfig, VllmConfig
|
||||
from vllm.v1.request import Request
|
||||
from vllm.v1.structured_output import StructuredOutputManager
|
||||
from vllm.v1.structured_output.backend_types import StructuredOutputOptions
|
||||
|
||||
|
||||
class UntouchableTokenIds(Sequence[int]):
|
||||
"""A token history that fails the test if anything reads it."""
|
||||
|
||||
def __init__(self, length: int) -> None:
|
||||
self._length = length
|
||||
|
||||
def __len__(self) -> int:
|
||||
return self._length
|
||||
|
||||
def __getitem__(self, item):
|
||||
raise AssertionError("token history must not be read")
|
||||
|
||||
def __iter__(self):
|
||||
raise AssertionError("token history must not be iterated")
|
||||
|
||||
|
||||
def make_qwen3_reasoner():
|
||||
"""A real engine-based reasoning adapter over a tiny mock vocabulary."""
|
||||
from tests.parser.engine.conftest import make_mock_tokenizer
|
||||
from vllm.parser.engine.registered_adapters import Qwen3ParserReasoningAdapter
|
||||
|
||||
tokenizer = make_mock_tokenizer(
|
||||
{"<think>": 200, "</think>": 201, "<tool_call>": 202, "</tool_call>": 203}
|
||||
)
|
||||
return Qwen3ParserReasoningAdapter(tokenizer)
|
||||
|
||||
|
||||
class MockReasoner:
|
||||
def __init__(self, tokenizer):
|
||||
self.is_reasoning_end = Mock(return_value=False)
|
||||
self.is_reasoning_end_streaming = Mock(return_value=False)
|
||||
|
||||
|
||||
class TestReasoningStructuredOutput:
|
||||
"""Test reasoning-aware structured output functionality."""
|
||||
|
||||
@pytest.fixture
|
||||
def mock_model_config(self):
|
||||
"""Create a mock ModelConfig."""
|
||||
config = Mock(spec=ModelConfig)
|
||||
config.skip_tokenizer_init = True # Skip tokenizer init to avoid network calls
|
||||
config.get_vocab_size = Mock(return_value=50000)
|
||||
# Add missing runner_type attribute that tokenizer initialization expects
|
||||
config.runner_type = "generate"
|
||||
# Add other attributes that tokenizer initialization might need
|
||||
config.tokenizer = "test-tokenizer"
|
||||
config.tokenizer_mode = "auto"
|
||||
config.trust_remote_code = False
|
||||
config.tokenizer_revision = None
|
||||
return config
|
||||
|
||||
@pytest.fixture
|
||||
def mock_scheduler_config(self):
|
||||
"""Create a mock SchedulerConfig."""
|
||||
config = Mock(spec=SchedulerConfig)
|
||||
config.max_num_seqs = 128
|
||||
return config
|
||||
|
||||
@pytest.fixture
|
||||
def mock_vllm_config(self, mock_model_config, mock_scheduler_config):
|
||||
"""Create a mock VllmConfig."""
|
||||
config = Mock(spec=VllmConfig)
|
||||
config.model_config = mock_model_config
|
||||
config.scheduler_config = mock_scheduler_config
|
||||
config.structured_outputs_config = Mock()
|
||||
config.structured_outputs_config.reasoning_parser = None
|
||||
config.structured_outputs_config.enable_in_reasoning = False
|
||||
config.speculative_config = None
|
||||
return config
|
||||
|
||||
@pytest.fixture
|
||||
def mock_request_with_structured_output(self):
|
||||
"""Create a mock request with structured output."""
|
||||
request = Mock(spec=Request)
|
||||
request.structured_output_request = Mock()
|
||||
request.structured_output_request.reasoning_ended = None
|
||||
request.structured_output_request.grammar = Mock()
|
||||
request.structured_output_request.reasoning_parser_kwargs = None
|
||||
request.structured_output_request.reasoner = None
|
||||
request.structured_output_request.grammar.is_terminated = Mock(
|
||||
return_value=False
|
||||
)
|
||||
request.use_structured_output = True
|
||||
request.prompt_token_ids = [1, 2, 3, 4, 5]
|
||||
request.all_token_ids = [1, 2, 3, 4, 5, 6, 7, 8]
|
||||
request.num_computed_tokens = 5
|
||||
request.num_output_placeholders = 0
|
||||
request.request_id = "mock_req"
|
||||
return request
|
||||
|
||||
@pytest.fixture
|
||||
def manager_with_reasoner(self, mock_vllm_config):
|
||||
manager = StructuredOutputManager(mock_vllm_config)
|
||||
manager.reasoner_cls = MockReasoner
|
||||
manager.tokenizer = Mock()
|
||||
return manager
|
||||
|
||||
def test_should_fill_bitmask_with_enable_in_reasoning(
|
||||
self, manager_with_reasoner, mock_request_with_structured_output
|
||||
):
|
||||
"""Test should_fill_bitmask when enable_in_reasoning is True."""
|
||||
manager_with_reasoner.enable_in_reasoning = True
|
||||
|
||||
# Should always return True when enable_in_reasoning is enabled
|
||||
result = manager_with_reasoner.should_fill_bitmask(
|
||||
mock_request_with_structured_output
|
||||
)
|
||||
assert result is True
|
||||
assert (
|
||||
mock_request_with_structured_output.structured_output_request.reasoner
|
||||
is None
|
||||
)
|
||||
|
||||
def test_should_fill_bitmask_reasoning_already_ended(
|
||||
self,
|
||||
manager_with_reasoner,
|
||||
mock_request_with_structured_output,
|
||||
):
|
||||
"""An active grammar does not need a request-local reasoner."""
|
||||
structured_req = mock_request_with_structured_output.structured_output_request
|
||||
structured_req.reasoning_ended = True
|
||||
|
||||
result = manager_with_reasoner.should_fill_bitmask(
|
||||
mock_request_with_structured_output
|
||||
)
|
||||
|
||||
assert result is True
|
||||
assert structured_req.reasoner is None
|
||||
|
||||
def test_grammar_bitmask_skips_reasoner_when_already_active(
|
||||
self,
|
||||
manager_with_reasoner,
|
||||
mock_request_with_structured_output,
|
||||
):
|
||||
"""Bitmask generation skips reasoning-boundary detection."""
|
||||
manager_with_reasoner.vllm_config.num_speculative_tokens = 1
|
||||
manager_with_reasoner._grammar_bitmask = Mock()
|
||||
manager_with_reasoner._grammar_bitmask.shape = (1, 1)
|
||||
expected_bitmask = Mock()
|
||||
manager_with_reasoner._grammar_bitmask.numpy.return_value = expected_bitmask
|
||||
manager_with_reasoner._fill_bitmasks = Mock()
|
||||
|
||||
structured_req = mock_request_with_structured_output.structured_output_request
|
||||
structured_req.reasoning_ended = True
|
||||
request_id = mock_request_with_structured_output.request_id
|
||||
|
||||
bitmask = manager_with_reasoner.grammar_bitmask(
|
||||
requests={request_id: mock_request_with_structured_output},
|
||||
structured_output_request_ids=[request_id],
|
||||
scheduled_spec_decode_tokens={},
|
||||
)
|
||||
|
||||
assert bitmask is expected_bitmask
|
||||
assert structured_req.reasoner is None
|
||||
|
||||
def test_should_fill_bitmask_without_enable_in_reasoning(
|
||||
self,
|
||||
manager_with_reasoner,
|
||||
mock_request_with_structured_output,
|
||||
):
|
||||
"""Test should_fill_bitmask when enable_in_reasoning is False."""
|
||||
# Keep enable_in_reasoning as False (default)
|
||||
config = manager_with_reasoner.vllm_config.structured_outputs_config
|
||||
assert config.enable_in_reasoning is False
|
||||
|
||||
result = manager_with_reasoner.should_fill_bitmask(
|
||||
mock_request_with_structured_output
|
||||
)
|
||||
|
||||
# Should set reasoning_ended and return its value
|
||||
assert (
|
||||
mock_request_with_structured_output.structured_output_request.reasoning_ended
|
||||
is False
|
||||
)
|
||||
assert result is False
|
||||
|
||||
def test_should_fill_bitmask_no_reasoner(
|
||||
self, mock_vllm_config, mock_request_with_structured_output
|
||||
):
|
||||
"""Test should_fill_bitmask when no reasoner is configured."""
|
||||
manager = StructuredOutputManager(mock_vllm_config)
|
||||
|
||||
result = manager.should_fill_bitmask(mock_request_with_structured_output)
|
||||
|
||||
# Should default to True when no reasoner
|
||||
assert result is True
|
||||
|
||||
def test_should_fill_bitmask_uses_request_reasoning_parser_kwargs(
|
||||
self, mock_vllm_config, mock_request_with_structured_output
|
||||
):
|
||||
"""Test request-level parser kwargs override the default reasoner."""
|
||||
|
||||
class KwargReasoner:
|
||||
def __init__(self, tokenizer, chat_template_kwargs=None):
|
||||
self.chat_template_kwargs = chat_template_kwargs or {}
|
||||
|
||||
def is_reasoning_end(self, input_ids):
|
||||
return not self.chat_template_kwargs.get("enable_thinking", False)
|
||||
|
||||
manager = StructuredOutputManager(mock_vllm_config)
|
||||
manager.reasoner_cls = KwargReasoner
|
||||
manager.tokenizer = Mock()
|
||||
|
||||
structured_req = mock_request_with_structured_output.structured_output_request
|
||||
structured_req.reasoning_parser_kwargs = {
|
||||
"chat_template_kwargs": {"enable_thinking": True}
|
||||
}
|
||||
|
||||
result = manager.should_fill_bitmask(mock_request_with_structured_output)
|
||||
|
||||
assert result is False
|
||||
assert (
|
||||
mock_request_with_structured_output.structured_output_request.reasoner
|
||||
is not None
|
||||
)
|
||||
|
||||
def test_should_advance_with_enable_in_reasoning(
|
||||
self,
|
||||
manager_with_reasoner,
|
||||
mock_request_with_structured_output,
|
||||
):
|
||||
"""Test should_advance when enable_in_reasoning is True."""
|
||||
# Enable enable_in_reasoning
|
||||
manager_with_reasoner.enable_in_reasoning = True
|
||||
|
||||
# Should always return True when enable_in_reasoning is enabled
|
||||
result = manager_with_reasoner.should_advance(
|
||||
mock_request_with_structured_output
|
||||
)
|
||||
assert result is True
|
||||
assert (
|
||||
mock_request_with_structured_output.structured_output_request.reasoner
|
||||
is None
|
||||
)
|
||||
|
||||
def test_should_advance_reasoning_not_ended(
|
||||
self,
|
||||
manager_with_reasoner,
|
||||
mock_request_with_structured_output,
|
||||
):
|
||||
"""Test should_advance when reasoning has not ended."""
|
||||
# Set reasoning as not ended
|
||||
(
|
||||
mock_request_with_structured_output.structured_output_request
|
||||
).reasoning_ended = False
|
||||
|
||||
result = manager_with_reasoner.should_advance(
|
||||
mock_request_with_structured_output
|
||||
)
|
||||
|
||||
# Should return False since reasoning hasn't ended
|
||||
assert result is False
|
||||
|
||||
def test_should_advance_reasoning_just_ended(
|
||||
self,
|
||||
manager_with_reasoner,
|
||||
mock_request_with_structured_output,
|
||||
):
|
||||
"""Test should_advance when reasoning ends in current step."""
|
||||
# Set reasoning as not ended initially, but ends in this step
|
||||
(
|
||||
mock_request_with_structured_output.structured_output_request
|
||||
).reasoning_ended = False
|
||||
reasoner = MockReasoner(tokenizer=Mock())
|
||||
reasoner.is_reasoning_end_streaming.return_value = True
|
||||
structured_req = mock_request_with_structured_output.structured_output_request
|
||||
structured_req.reasoner = reasoner
|
||||
|
||||
result = manager_with_reasoner.should_advance(
|
||||
mock_request_with_structured_output
|
||||
)
|
||||
|
||||
# The scheduler trims the reasoning prefix before advancing the grammar.
|
||||
assert (
|
||||
mock_request_with_structured_output.structured_output_request.reasoning_ended
|
||||
is True
|
||||
)
|
||||
assert result is True
|
||||
|
||||
def test_should_advance_reasoning_already_ended(
|
||||
self,
|
||||
manager_with_reasoner,
|
||||
mock_request_with_structured_output,
|
||||
):
|
||||
"""Test should_advance when reasoning has already ended."""
|
||||
# Set reasoning as already ended
|
||||
(
|
||||
mock_request_with_structured_output.structured_output_request
|
||||
).reasoning_ended = True
|
||||
|
||||
result = manager_with_reasoner.should_advance(
|
||||
mock_request_with_structured_output
|
||||
)
|
||||
|
||||
# Should return True since reasoning has ended
|
||||
assert result is True
|
||||
assert (
|
||||
mock_request_with_structured_output.structured_output_request.reasoner
|
||||
is None
|
||||
)
|
||||
|
||||
def test_should_advance_uses_new_token_ids_when_provided(
|
||||
self,
|
||||
manager_with_reasoner,
|
||||
mock_request_with_structured_output,
|
||||
):
|
||||
"""Regression for #43388: when caller passes new_token_ids, the
|
||||
reasoner sees the exact multi-token delta rather than the
|
||||
placeholder-derived window.
|
||||
"""
|
||||
structured_req = mock_request_with_structured_output.structured_output_request
|
||||
structured_req.reasoning_ended = False
|
||||
|
||||
end_token_id = 248069
|
||||
|
||||
reasoner = MockReasoner(tokenizer=Mock())
|
||||
# Detection mirrors the real Qwen3 parser: end token in the delta.
|
||||
reasoner.is_reasoning_end_streaming = Mock(
|
||||
side_effect=lambda input_ids, delta_ids: end_token_id in list(delta_ids)
|
||||
)
|
||||
structured_req.reasoner = reasoner
|
||||
|
||||
# Scenario from #43388: async + spec decode K=4, 4 tokens accepted
|
||||
# but only 1 placeholder remains (some drafts were rejected).
|
||||
# The placeholder math would yield delta=[271] and miss </think>.
|
||||
# Passing new_token_ids must override that.
|
||||
new_token_ids = [9, 198, end_token_id, 271]
|
||||
mock_request_with_structured_output.all_token_ids = [
|
||||
1,
|
||||
2,
|
||||
3,
|
||||
4,
|
||||
5,
|
||||
] + new_token_ids
|
||||
mock_request_with_structured_output.num_computed_tokens = 9
|
||||
mock_request_with_structured_output.num_output_placeholders = 1
|
||||
|
||||
result = manager_with_reasoner.should_advance(
|
||||
mock_request_with_structured_output,
|
||||
new_token_ids=new_token_ids,
|
||||
)
|
||||
|
||||
# First call to is_reasoning_end_streaming was with the full
|
||||
# new_token_ids (not the truncated placeholder window).
|
||||
first_call = reasoner.is_reasoning_end_streaming.call_args_list[0]
|
||||
_, called_delta = first_call.args
|
||||
assert list(called_delta) == new_token_ids
|
||||
|
||||
assert structured_req.reasoning_ended is True
|
||||
assert result is True
|
||||
|
||||
def test_should_advance_without_new_token_ids_falls_back(
|
||||
self,
|
||||
manager_with_reasoner,
|
||||
mock_request_with_structured_output,
|
||||
):
|
||||
"""Backward compat: callers that don't pass new_token_ids keep
|
||||
the original placeholder-derived delta window.
|
||||
"""
|
||||
structured_req = mock_request_with_structured_output.structured_output_request
|
||||
structured_req.reasoning_ended = False
|
||||
reasoner = MockReasoner(tokenizer=Mock())
|
||||
reasoner.is_reasoning_end_streaming.return_value = False
|
||||
structured_req.reasoner = reasoner
|
||||
|
||||
mock_request_with_structured_output.all_token_ids = [1, 2, 3, 4, 5]
|
||||
mock_request_with_structured_output.num_computed_tokens = 5
|
||||
mock_request_with_structured_output.num_output_placeholders = 2
|
||||
|
||||
result = manager_with_reasoner.should_advance(
|
||||
mock_request_with_structured_output
|
||||
)
|
||||
|
||||
# placeholder window: start = 5 - 2 = 3, delta = [4, 5]
|
||||
_, called_delta = reasoner.is_reasoning_end_streaming.call_args[0]
|
||||
assert list(called_delta) == [4, 5]
|
||||
assert result is False
|
||||
|
||||
def test_should_advance_whole_delta_fallback_pins_last_sequence_index(
|
||||
self,
|
||||
manager_with_reasoner,
|
||||
mock_request_with_structured_output,
|
||||
):
|
||||
"""When only the whole-delta predicate fires, the boundary is the
|
||||
last index of the sequence, even if the placeholder window
|
||||
overshoots it and the delta is empty.
|
||||
"""
|
||||
structured_req = mock_request_with_structured_output.structured_output_request
|
||||
structured_req.reasoning_ended = False
|
||||
reasoner = MockReasoner(tokenizer=Mock())
|
||||
# Base-default shape: rescans the whole sequence, ignores the delta.
|
||||
reasoner.is_reasoning_end_streaming.return_value = True
|
||||
structured_req.reasoner = reasoner
|
||||
|
||||
mock_request_with_structured_output.all_token_ids = [1, 2, 3, 4, 5]
|
||||
# placeholder window: start = 9 - 1 = 8, past the end -> delta = []
|
||||
mock_request_with_structured_output.num_computed_tokens = 9
|
||||
mock_request_with_structured_output.num_output_placeholders = 1
|
||||
|
||||
assert manager_with_reasoner.should_advance(mock_request_with_structured_output)
|
||||
assert structured_req.reasoning_ended is True
|
||||
assert structured_req.reasoning_end_token_index == 4
|
||||
|
||||
def test_should_advance_trims_reasoning_prefix_for_json(
|
||||
self,
|
||||
manager_with_reasoner,
|
||||
mock_request_with_structured_output,
|
||||
):
|
||||
"""JSON uses the common trim-then-advance path at the boundary."""
|
||||
structured_req = mock_request_with_structured_output.structured_output_request
|
||||
structured_req.reasoning_ended = False
|
||||
structured_req.structured_output_key = (
|
||||
StructuredOutputOptions.JSON_OBJECT,
|
||||
"{}",
|
||||
)
|
||||
|
||||
marker = 248069
|
||||
|
||||
class MarkerReasoner:
|
||||
def __init__(self, *_, **__):
|
||||
pass
|
||||
|
||||
def is_reasoning_end_streaming(self, input_ids, delta_ids):
|
||||
return marker in list(delta_ids)
|
||||
|
||||
structured_req.reasoner = MarkerReasoner()
|
||||
|
||||
new_token_ids = [9, 198, marker, 271, 5005]
|
||||
mock_request_with_structured_output.all_token_ids = [1, 2, 3] + new_token_ids
|
||||
|
||||
result = manager_with_reasoner.should_advance(
|
||||
mock_request_with_structured_output,
|
||||
new_token_ids=new_token_ids,
|
||||
)
|
||||
|
||||
structured_req.grammar.accept_tokens.assert_not_called()
|
||||
assert structured_req.reasoning_ended is True
|
||||
assert result is True
|
||||
assert structured_req.reasoning_end_token_index == 5
|
||||
assert manager_with_reasoner.trim_reasoning_for_advance(
|
||||
mock_request_with_structured_output, new_token_ids
|
||||
) == [271, 5005]
|
||||
|
||||
def test_should_advance_engine_adapter_scans_only_the_delta(
|
||||
self,
|
||||
manager_with_reasoner,
|
||||
mock_request_with_structured_output,
|
||||
):
|
||||
"""Engine-based parsers locate the boundary from the delta alone."""
|
||||
structured_req = mock_request_with_structured_output.structured_output_request
|
||||
structured_req.reasoning_ended = False
|
||||
reasoner = make_qwen3_reasoner()
|
||||
assert reasoner.reasoning_end_token_ids == {201, 202}
|
||||
reasoner.is_reasoning_end_streaming = Mock(
|
||||
side_effect=AssertionError("engine path must not rescan the prefix")
|
||||
)
|
||||
structured_req.reasoner = reasoner
|
||||
|
||||
new_token_ids = [9, 198, 201, 271]
|
||||
mock_request_with_structured_output.all_token_ids = [1, 2, 3] + new_token_ids
|
||||
|
||||
assert manager_with_reasoner.should_advance(
|
||||
mock_request_with_structured_output, new_token_ids=new_token_ids
|
||||
)
|
||||
assert structured_req.reasoning_ended is True
|
||||
assert structured_req.reasoning_end_token_index == 5
|
||||
assert manager_with_reasoner.trim_reasoning_for_advance(
|
||||
mock_request_with_structured_output, new_token_ids
|
||||
) == [271]
|
||||
|
||||
def test_should_advance_engine_adapter_tool_start_ends_reasoning(
|
||||
self,
|
||||
manager_with_reasoner,
|
||||
mock_request_with_structured_output,
|
||||
):
|
||||
structured_req = mock_request_with_structured_output.structured_output_request
|
||||
structured_req.reasoning_ended = False
|
||||
structured_req.reasoner = make_qwen3_reasoner()
|
||||
|
||||
new_token_ids = [9, 202, 271]
|
||||
mock_request_with_structured_output.all_token_ids = [1, 2, 3] + new_token_ids
|
||||
|
||||
assert manager_with_reasoner.should_advance(
|
||||
mock_request_with_structured_output, new_token_ids=new_token_ids
|
||||
)
|
||||
assert structured_req.reasoning_end_token_index == 4
|
||||
|
||||
def test_should_advance_engine_adapter_no_boundary(
|
||||
self,
|
||||
manager_with_reasoner,
|
||||
mock_request_with_structured_output,
|
||||
):
|
||||
structured_req = mock_request_with_structured_output.structured_output_request
|
||||
structured_req.reasoning_ended = False
|
||||
structured_req.reasoner = make_qwen3_reasoner()
|
||||
|
||||
new_token_ids = [9, 198, 271]
|
||||
mock_request_with_structured_output.all_token_ids = [1, 2, 3] + new_token_ids
|
||||
|
||||
assert not manager_with_reasoner.should_advance(
|
||||
mock_request_with_structured_output, new_token_ids=new_token_ids
|
||||
)
|
||||
assert structured_req.reasoning_ended is False
|
||||
|
||||
def test_should_advance_engine_adapter_never_reads_history(
|
||||
self,
|
||||
manager_with_reasoner,
|
||||
mock_request_with_structured_output,
|
||||
):
|
||||
"""Per-step cost must not depend on how much has been generated."""
|
||||
structured_req = mock_request_with_structured_output.structured_output_request
|
||||
structured_req.reasoning_ended = False
|
||||
structured_req.reasoner = make_qwen3_reasoner()
|
||||
|
||||
new_token_ids = [9, 201]
|
||||
total = (1 << 20) + len(new_token_ids)
|
||||
mock_request_with_structured_output.all_token_ids = UntouchableTokenIds(total)
|
||||
|
||||
assert manager_with_reasoner.should_advance(
|
||||
mock_request_with_structured_output, new_token_ids=new_token_ids
|
||||
)
|
||||
assert structured_req.reasoning_end_token_index == total - 1
|
||||
|
||||
def test_should_advance_engine_adapter_without_new_tokens_persists(
|
||||
self,
|
||||
manager_with_reasoner,
|
||||
mock_request_with_structured_output,
|
||||
):
|
||||
"""Draft-validation callers persist the boundary like the legacy path."""
|
||||
structured_req = mock_request_with_structured_output.structured_output_request
|
||||
structured_req.reasoning_ended = False
|
||||
structured_req.reasoner = make_qwen3_reasoner()
|
||||
|
||||
mock_request_with_structured_output.all_token_ids = [1, 2, 3, 4, 5, 201, 7]
|
||||
# placeholder window: start = 7 - 2 = 5, delta = [201, 7]
|
||||
mock_request_with_structured_output.num_computed_tokens = 7
|
||||
mock_request_with_structured_output.num_output_placeholders = 2
|
||||
|
||||
assert manager_with_reasoner.should_advance(mock_request_with_structured_output)
|
||||
assert structured_req.reasoning_ended is True
|
||||
assert structured_req.reasoning_end_token_index == 5
|
||||
@@ -0,0 +1,414 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
|
||||
import pytest
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
from vllm.config import DeviceConfig, 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"
|
||||
THINK_END = "\n" # reasoning-end marker (single GPT-2 token)
|
||||
EOS = "<|eos|>" # resolved to tokenizer.eos_token_id
|
||||
JSON_SCHEMA = '{"type": "object"}'
|
||||
BACKENDS = ("xgrammar", "guidance")
|
||||
MAX_WAIT_SECONDS = 5
|
||||
NUM_SPEC_TOKENS = 8
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def tokenizer():
|
||||
return AutoTokenizer.from_pretrained(TOKENIZER)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class FlowCase:
|
||||
"""One validate -> bitmask -> accept -> post-accept bitmask scenario."""
|
||||
|
||||
raw_drafts: tuple[str, ...]
|
||||
expected_row_pattern: str
|
||||
expected_reasoning: bool | None
|
||||
expect_terminated: bool
|
||||
prefix: str = ""
|
||||
expected_validated: tuple[str, ...] | None = None
|
||||
reasoning_ended: bool | None = False
|
||||
xfail_guidance: str | None = None
|
||||
|
||||
|
||||
class MockReasoner:
|
||||
def __init__(self, tokenizer, marker: int | None = None):
|
||||
self.marker = marker
|
||||
|
||||
def is_reasoning_end(self, input_ids):
|
||||
if self.marker is None:
|
||||
return True
|
||||
return self.marker in list(input_ids)
|
||||
|
||||
def is_reasoning_end_streaming(self, input_ids, delta_ids):
|
||||
return self.is_reasoning_end(delta_ids)
|
||||
|
||||
|
||||
def _single_token(tokenizer, text: str) -> int:
|
||||
token_ids = tokenizer.encode(text)
|
||||
assert len(token_ids) == 1, (text, token_ids)
|
||||
return token_ids[0]
|
||||
|
||||
|
||||
def _to_token_ids(tokenizer, texts: tuple[str, ...]) -> list[int]:
|
||||
"""Convert literal text (or EOS sentinel) to token IDs."""
|
||||
return [
|
||||
tokenizer.eos_token_id if t == EOS else _single_token(tokenizer, t)
|
||||
for t in texts
|
||||
]
|
||||
|
||||
|
||||
def _wait_for_grammar(request: Request) -> None:
|
||||
structured_req = request.structured_output_request
|
||||
assert structured_req is not None
|
||||
|
||||
deadline = time.time() + MAX_WAIT_SECONDS
|
||||
while not structured_req._check_grammar_completion():
|
||||
if time.time() > deadline:
|
||||
pytest.fail("Grammar compilation timed out")
|
||||
time.sleep(0.01)
|
||||
|
||||
|
||||
def _build_harness(
|
||||
tokenizer,
|
||||
backend: str,
|
||||
prefix: str = "",
|
||||
use_reasoner: bool = True,
|
||||
reasoning_ended: bool | None = None,
|
||||
enable_in_reasoning: bool = False,
|
||||
reasoning_parser_kwargs: dict | None = None,
|
||||
) -> tuple[StructuredOutputManager, Request]:
|
||||
vllm_config = VllmConfig(
|
||||
model_config=ModelConfig(tokenizer=TOKENIZER),
|
||||
device_config=DeviceConfig(device="cpu"),
|
||||
structured_outputs_config=StructuredOutputsConfig(
|
||||
backend=backend,
|
||||
enable_in_reasoning=enable_in_reasoning,
|
||||
),
|
||||
speculative_config=SpeculativeConfig(
|
||||
model="[ngram]",
|
||||
num_speculative_tokens=NUM_SPEC_TOKENS,
|
||||
),
|
||||
)
|
||||
manager = StructuredOutputManager(vllm_config)
|
||||
if use_reasoner:
|
||||
manager.reasoner_cls = MockReasoner
|
||||
|
||||
sampling_params = SamplingParams(
|
||||
structured_outputs=StructuredOutputsParams(json=JSON_SCHEMA)
|
||||
)
|
||||
sampling_params.structured_outputs._backend = backend
|
||||
sampling_params.update_from_generation_config({}, tokenizer.eos_token_id)
|
||||
|
||||
prompt_ids = tokenizer.encode(prefix) if prefix else []
|
||||
request = Request(
|
||||
request_id=f"{backend}-flow",
|
||||
prompt_token_ids=prompt_ids,
|
||||
sampling_params=sampling_params,
|
||||
pooling_params=None,
|
||||
reasoning_ended=reasoning_ended,
|
||||
reasoning_parser_kwargs=reasoning_parser_kwargs,
|
||||
)
|
||||
manager.grammar_init(request)
|
||||
_wait_for_grammar(request)
|
||||
|
||||
structured_req = request.structured_output_request
|
||||
assert request.prompt_token_ids is not None
|
||||
assert structured_req is not None
|
||||
assert structured_req.grammar is not None
|
||||
if prompt_ids:
|
||||
assert structured_req.grammar.accept_tokens(
|
||||
request.request_id,
|
||||
prompt_ids,
|
||||
)
|
||||
return manager, request
|
||||
|
||||
|
||||
def _row_pattern(bitmask) -> str:
|
||||
return "".join("U" if (row == -1).all() else "C" for row in bitmask)
|
||||
|
||||
|
||||
def _run_real_flow(
|
||||
manager: StructuredOutputManager,
|
||||
request: Request,
|
||||
raw_drafts: list[int],
|
||||
expected_validated: list[int],
|
||||
expected_row_pattern: str,
|
||||
expected_reasoning: bool | None,
|
||||
expect_terminated: bool,
|
||||
) -> None:
|
||||
assert len(expected_validated) <= len(raw_drafts)
|
||||
|
||||
validated = manager.validate_tokens(request, list(raw_drafts))
|
||||
assert validated == expected_validated
|
||||
|
||||
padded = validated + [-1] * (len(raw_drafts) - len(validated))
|
||||
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(raw_drafts) + 1
|
||||
assert _row_pattern(bitmask) == expected_row_pattern
|
||||
|
||||
# grammar_bitmask() must rollback any speculative state it advanced.
|
||||
assert manager.validate_tokens(request, list(raw_drafts)) == expected_validated
|
||||
|
||||
structured_req = request.structured_output_request
|
||||
assert structured_req is not None
|
||||
grammar = structured_req.grammar
|
||||
assert grammar is not None
|
||||
assert not grammar.is_terminated()
|
||||
# `bitmask[i]` is the grammar state after the first `i` scheduled tokens.
|
||||
# These tests commit the validated prefix, so the matching post-accept
|
||||
# state is already present in the scheduled bitmask.
|
||||
expected_post_accept_row = bitmask[len(expected_validated)].copy()
|
||||
|
||||
# Mirror the scheduler flow: sampled tokens are appended before
|
||||
# StructuredOutputManager.accept_tokens() trims them to the
|
||||
# grammar-constrained suffix.
|
||||
request.append_output_token_ids(list(expected_validated))
|
||||
assert manager.accept_tokens(request, list(expected_validated))
|
||||
|
||||
post_accept_bitmask = manager.grammar_bitmask(
|
||||
requests={request.request_id: request},
|
||||
structured_output_request_ids=[request.request_id],
|
||||
scheduled_spec_decode_tokens={},
|
||||
)
|
||||
assert post_accept_bitmask is not None
|
||||
assert post_accept_bitmask.shape[0] == 1
|
||||
assert (post_accept_bitmask[0] == expected_post_accept_row).all()
|
||||
assert structured_req.reasoning_ended == expected_reasoning
|
||||
assert grammar.is_terminated() is expect_terminated
|
||||
|
||||
|
||||
FLOW_CASES = [
|
||||
pytest.param(
|
||||
FlowCase(
|
||||
raw_drafts=(" ", "z", " "),
|
||||
expected_row_pattern="UUUU",
|
||||
expected_reasoning=False,
|
||||
expect_terminated=False,
|
||||
),
|
||||
id="inactive",
|
||||
),
|
||||
pytest.param(
|
||||
FlowCase(
|
||||
prefix='{"a"',
|
||||
raw_drafts=(":", ' "', "b", '"}'),
|
||||
expected_row_pattern="CCCCC",
|
||||
expected_reasoning=True,
|
||||
expect_terminated=False,
|
||||
reasoning_ended=True,
|
||||
),
|
||||
id="active_all_valid",
|
||||
),
|
||||
pytest.param(
|
||||
FlowCase(
|
||||
prefix='{"a"',
|
||||
raw_drafts=("z", ":", ' "', "b"),
|
||||
expected_validated=(),
|
||||
expected_row_pattern="CUUUU",
|
||||
expected_reasoning=True,
|
||||
expect_terminated=False,
|
||||
reasoning_ended=True,
|
||||
),
|
||||
id="active_first_token_invalid",
|
||||
),
|
||||
pytest.param(
|
||||
FlowCase(
|
||||
prefix='{"a"',
|
||||
raw_drafts=(":", "z", ' "', "b"),
|
||||
expected_validated=(":",),
|
||||
expected_row_pattern="CCUUU",
|
||||
expected_reasoning=True,
|
||||
expect_terminated=False,
|
||||
reasoning_ended=True,
|
||||
),
|
||||
id="active_later_token_invalid",
|
||||
),
|
||||
pytest.param(
|
||||
FlowCase(
|
||||
prefix='{"a": "b"}',
|
||||
raw_drafts=(EOS, " ", " "),
|
||||
expected_validated=(EOS,),
|
||||
expected_row_pattern="CUUU",
|
||||
expected_reasoning=True,
|
||||
expect_terminated=True,
|
||||
reasoning_ended=True,
|
||||
xfail_guidance=(
|
||||
"guidance validate_tokens rejects EOS on an already-complete "
|
||||
"object even though accept_tokens would accept it"
|
||||
),
|
||||
),
|
||||
id="active_first_token_terminates",
|
||||
),
|
||||
pytest.param(
|
||||
FlowCase(
|
||||
prefix='{"a"',
|
||||
raw_drafts=(":", ' "', "b", '"}', EOS, " "),
|
||||
expected_validated=(":", ' "', "b", '"}', EOS),
|
||||
expected_row_pattern="CCCCCUU",
|
||||
expected_reasoning=True,
|
||||
expect_terminated=True,
|
||||
reasoning_ended=True,
|
||||
),
|
||||
id="active_later_token_terminates",
|
||||
),
|
||||
pytest.param(
|
||||
FlowCase(
|
||||
raw_drafts=(THINK_END, "{", ' "', "b"),
|
||||
expected_row_pattern="UCCCC",
|
||||
expected_reasoning=True,
|
||||
expect_terminated=False,
|
||||
),
|
||||
id="becomes_active_on_first_token",
|
||||
),
|
||||
pytest.param(
|
||||
FlowCase(
|
||||
raw_drafts=(" ", THINK_END, "{", ' "'),
|
||||
expected_row_pattern="UUCCC",
|
||||
expected_reasoning=True,
|
||||
expect_terminated=False,
|
||||
),
|
||||
id="becomes_active_on_middle_token",
|
||||
),
|
||||
pytest.param(
|
||||
FlowCase(
|
||||
raw_drafts=(" ", " ", THINK_END),
|
||||
expected_row_pattern="UUUC",
|
||||
expected_reasoning=True,
|
||||
expect_terminated=False,
|
||||
),
|
||||
id="becomes_active_on_last_token",
|
||||
),
|
||||
pytest.param(
|
||||
FlowCase(
|
||||
raw_drafts=(" ", THINK_END, "z", "{"),
|
||||
expected_validated=(" ", THINK_END),
|
||||
expected_row_pattern="UUCUU",
|
||||
expected_reasoning=True,
|
||||
expect_terminated=False,
|
||||
),
|
||||
id="becomes_active_first_token_invalid",
|
||||
),
|
||||
pytest.param(
|
||||
FlowCase(
|
||||
raw_drafts=(" ", THINK_END, "{", "z"),
|
||||
expected_validated=(" ", THINK_END, "{"),
|
||||
expected_row_pattern="UUCCU",
|
||||
expected_reasoning=True,
|
||||
expect_terminated=False,
|
||||
),
|
||||
id="becomes_active_later_token_invalid",
|
||||
),
|
||||
pytest.param(
|
||||
FlowCase(
|
||||
raw_drafts=(" ", THINK_END, "{", "}", EOS, " "),
|
||||
expected_validated=(" ", THINK_END, "{", "}", EOS),
|
||||
expected_row_pattern="UUCCCUU",
|
||||
expected_reasoning=True,
|
||||
expect_terminated=True,
|
||||
),
|
||||
id="becomes_active_terminates",
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", BACKENDS)
|
||||
@pytest.mark.parametrize("case", FLOW_CASES)
|
||||
def test_real_flow(
|
||||
tokenizer,
|
||||
backend: str,
|
||||
case: FlowCase,
|
||||
):
|
||||
if backend == "guidance" and case.xfail_guidance:
|
||||
pytest.xfail(case.xfail_guidance)
|
||||
|
||||
reasoner_kwargs = {"marker": _single_token(tokenizer, THINK_END)}
|
||||
manager, request = _build_harness(
|
||||
tokenizer,
|
||||
backend,
|
||||
prefix=case.prefix,
|
||||
reasoning_ended=case.reasoning_ended,
|
||||
reasoning_parser_kwargs=reasoner_kwargs,
|
||||
)
|
||||
|
||||
raw_drafts = _to_token_ids(tokenizer, case.raw_drafts)
|
||||
expected_texts = (
|
||||
case.raw_drafts if case.expected_validated is None else case.expected_validated
|
||||
)
|
||||
expected_validated = _to_token_ids(tokenizer, expected_texts)
|
||||
|
||||
_run_real_flow(
|
||||
manager,
|
||||
request,
|
||||
raw_drafts=raw_drafts,
|
||||
expected_validated=expected_validated,
|
||||
expected_row_pattern=case.expected_row_pattern,
|
||||
expected_reasoning=case.expected_reasoning,
|
||||
expect_terminated=case.expect_terminated,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", BACKENDS)
|
||||
@pytest.mark.parametrize(
|
||||
("use_reasoner", "reasoning_ended", "enable_in_reasoning"),
|
||||
[
|
||||
pytest.param(True, False, True, id="enable_in_reasoning"),
|
||||
pytest.param(True, True, False, id="reasoning_ended"),
|
||||
pytest.param(False, False, False, id="no_reasoner"),
|
||||
pytest.param(True, None, False, id="inferred"),
|
||||
],
|
||||
)
|
||||
def test_initial_constraint_activation(
|
||||
tokenizer,
|
||||
backend: str,
|
||||
use_reasoner: bool,
|
||||
reasoning_ended: bool | None,
|
||||
enable_in_reasoning: bool,
|
||||
):
|
||||
manager, request = _build_harness(
|
||||
tokenizer,
|
||||
backend,
|
||||
use_reasoner=use_reasoner,
|
||||
reasoning_ended=reasoning_ended,
|
||||
enable_in_reasoning=enable_in_reasoning,
|
||||
reasoning_parser_kwargs={
|
||||
"marker": _single_token(tokenizer, THINK_END)
|
||||
if reasoning_ended is not None
|
||||
else None
|
||||
},
|
||||
)
|
||||
|
||||
# "{" is valid JSON start; "z" is not. Truncation proves grammar is active.
|
||||
open_brace = _single_token(tokenizer, "{")
|
||||
z = _single_token(tokenizer, "z")
|
||||
assert manager.validate_tokens(request, [open_brace, z]) == [open_brace]
|
||||
|
||||
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
|
||||
request.append_output_token_ids([open_brace])
|
||||
assert manager.accept_tokens(request, [open_brace])
|
||||
|
||||
# Short-circuit paths must not construct a request-local reasoner (#52573).
|
||||
if enable_in_reasoning or reasoning_ended is True:
|
||||
structured_req = request.structured_output_request
|
||||
assert structured_req is not None
|
||||
assert structured_req.reasoner is None
|
||||
@@ -197,6 +197,8 @@ class Parser:
|
||||
Used by structured engines like `xgrammar` to check if the
|
||||
reasoning content ends in the model output.
|
||||
|
||||
Must be a pure function of the input_ids.
|
||||
|
||||
Args:
|
||||
input_ids: The token IDs of the model output.
|
||||
|
||||
@@ -210,6 +212,8 @@ class Parser:
|
||||
"""
|
||||
Check if the reasoning content ends during a decode step.
|
||||
|
||||
Must be a pure function of the input_ids and delta_ids.
|
||||
|
||||
Args:
|
||||
input_ids: The entire model output token IDs.
|
||||
delta_ids: The last few computed tokens at the current decode step.
|
||||
@@ -228,8 +232,8 @@ class Parser:
|
||||
|
||||
Returns:
|
||||
The offset within ``token_ids`` of the last reasoning token, or
|
||||
``None`` when reasoning does not end in this window. Parsers that
|
||||
cannot answer from a window alone always return ``None``; callers
|
||||
``len(token_ids)`` when reasoning does not end in this window. Parsers
|
||||
that cannot answer from a window alone always return ``None``; callers
|
||||
fall back to :meth:`is_reasoning_end_streaming`.
|
||||
"""
|
||||
return None
|
||||
|
||||
@@ -645,7 +645,7 @@ class ParserEngine(Parser):
|
||||
for offset, token_id in enumerate(token_ids):
|
||||
if token_id in end_ids:
|
||||
return offset
|
||||
return None
|
||||
return len(token_ids)
|
||||
|
||||
def is_reasoning_end(self, input_ids: list[int]) -> bool:
|
||||
config = self.parser_engine_config
|
||||
|
||||
@@ -70,8 +70,7 @@ from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput
|
||||
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.structured_output import StructuredOutputManager
|
||||
from vllm.v1.utils import record_function_or_nullcontext
|
||||
|
||||
logger = init_logger(__name__)
|
||||
@@ -2052,33 +2051,18 @@ class Scheduler(SchedulerInterface):
|
||||
request.status = RequestStatus.FINISHED_STOPPED
|
||||
stopped = True
|
||||
|
||||
if new_token_ids and self.structured_output_manager.should_advance(
|
||||
request, new_token_ids=new_token_ids
|
||||
if new_token_ids and not self.structured_output_manager.accept_tokens(
|
||||
request, new_token_ids
|
||||
):
|
||||
struct_output_request = request.structured_output_request
|
||||
assert struct_output_request is not None
|
||||
grammar = struct_output_request.grammar
|
||||
assert isinstance(grammar, StructuredOutputGrammar)
|
||||
# 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
|
||||
)
|
||||
logger.error(
|
||||
"Unexpected: grammar rejected tokens %s for request %s. "
|
||||
"Terminating request.",
|
||||
new_token_ids,
|
||||
req_id,
|
||||
)
|
||||
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.",
|
||||
advance_token_ids,
|
||||
req_id,
|
||||
)
|
||||
request.status = RequestStatus.FINISHED_ERROR
|
||||
request.resumable = False
|
||||
stopped = True
|
||||
request.status = RequestStatus.FINISHED_ERROR
|
||||
request.resumable = False
|
||||
stopped = True
|
||||
|
||||
routed_experts = None
|
||||
if (
|
||||
@@ -2434,11 +2418,9 @@ class Scheduler(SchedulerInterface):
|
||||
continue
|
||||
|
||||
# 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
|
||||
request.spec_token_ids = self.structured_output_manager.validate_tokens(
|
||||
request, spec_token_ids
|
||||
)
|
||||
|
||||
def update_draft_token_ids_in_output(
|
||||
self, draft_token_ids: DraftTokenIds, scheduler_output: SchedulerOutput
|
||||
@@ -2464,10 +2446,9 @@ class Scheduler(SchedulerInterface):
|
||||
# (needed for chunked prefill case for example).
|
||||
del spec_token_ids[orig_num_spec_tokens:]
|
||||
# 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]
|
||||
spec_token_ids = self.structured_output_manager.validate_tokens(
|
||||
request, spec_token_ids
|
||||
)
|
||||
# Pad to original number of spec tokens.
|
||||
num_invalid_tokens = orig_num_spec_tokens - len(spec_token_ids)
|
||||
if num_invalid_tokens:
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import itertools
|
||||
import multiprocessing
|
||||
from collections.abc import Iterable, Sequence
|
||||
from concurrent.futures import Future, ThreadPoolExecutor
|
||||
@@ -18,6 +17,7 @@ from vllm.v1.structured_output.backend_types import (
|
||||
StructuredOutputGrammar,
|
||||
)
|
||||
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
|
||||
@@ -217,6 +217,100 @@ class StructuredOutputManager:
|
||||
) -> Future:
|
||||
return self.executor_for_fillmask.submit(self._fill_bitmasks, batch)
|
||||
|
||||
def _get_constraint_start(
|
||||
self,
|
||||
request: "Request",
|
||||
spec_tokens: Sequence[int],
|
||||
spec_tokens_committed: bool = False,
|
||||
) -> int:
|
||||
"""Return the index into `spec_tokens` where tokens should start being
|
||||
constrained by the grammar.
|
||||
Assumes `spec_tokens` are stripped of -1 padding tokens.
|
||||
Returns `len(spec_tokens) + 1` if no token should be constrained after
|
||||
accepting all `spec_tokens`.
|
||||
|
||||
`spec_tokens_committed` is True if spec_tokens is already in
|
||||
`request.all_token_ids`.
|
||||
"""
|
||||
if self.enable_in_reasoning:
|
||||
return 0
|
||||
|
||||
structured_req = request.structured_output_request
|
||||
assert structured_req is not None
|
||||
if structured_req.reasoning_ended:
|
||||
return 0
|
||||
|
||||
reasoner = self._get_reasoner(request)
|
||||
if reasoner is None:
|
||||
return 0
|
||||
|
||||
if structured_req.reasoning_ended is None:
|
||||
# This should be removed here, but since `openai_gptoss`
|
||||
# is an independent code path, it is kept for now.
|
||||
# After unifying the `openai_gptoss` and non-`openai_gptoss` styles,
|
||||
# it can be removed.
|
||||
if reasoner.is_reasoning_end(request.prompt_token_ids or []):
|
||||
return 0
|
||||
structured_req.reasoning_ended = False
|
||||
|
||||
num_spec_tokens = len(spec_tokens)
|
||||
if num_spec_tokens <= 0:
|
||||
return num_spec_tokens + 1
|
||||
|
||||
# Use `find_reasoning_end_offset` to find constraint start if supported
|
||||
if (
|
||||
isinstance(reasoner, ParserEngineReasoningAdapter)
|
||||
and reasoner.reasoning_end_token_ids
|
||||
):
|
||||
offset = reasoner.find_reasoning_end_offset(spec_tokens)
|
||||
if offset is not None:
|
||||
return offset + 1
|
||||
|
||||
# Fallback to `find_reasoning_end_offset`
|
||||
# TODO: Build a read-only Sequence view over all_token_ids
|
||||
# instead of copying the entire all_token_ids into input_ids.
|
||||
if spec_tokens_committed:
|
||||
if not spec_tokens or not reasoner.is_reasoning_end_streaming(
|
||||
request.all_token_ids, spec_tokens
|
||||
):
|
||||
return num_spec_tokens + 1
|
||||
|
||||
input_ids = request.all_token_ids.copy()
|
||||
delta_ids = list(spec_tokens)
|
||||
else:
|
||||
input_ids = request.all_token_ids.copy()
|
||||
input_ids.extend(spec_tokens)
|
||||
if not reasoner.is_reasoning_end_streaming(input_ids, spec_tokens):
|
||||
return num_spec_tokens + 1
|
||||
delta_ids = list(spec_tokens)
|
||||
|
||||
for i in range(num_spec_tokens - 1, 0, -1):
|
||||
input_ids.pop()
|
||||
delta_ids.pop()
|
||||
if not reasoner.is_reasoning_end_streaming(input_ids, delta_ids):
|
||||
return i + 1
|
||||
return 1
|
||||
|
||||
def validate_tokens(self, request: "Request", spec_tokens: list[int]) -> list[int]:
|
||||
"""Return the longest unconstrained or grammar-valid prefix of `spec_tokens`."""
|
||||
if not request.use_structured_output:
|
||||
return spec_tokens
|
||||
|
||||
spec_tokens = strip_speculative_padding(spec_tokens)
|
||||
constraint_start = self._get_constraint_start(request, spec_tokens)
|
||||
if constraint_start >= len(spec_tokens):
|
||||
return spec_tokens
|
||||
|
||||
structured_req = request.structured_output_request
|
||||
if TYPE_CHECKING:
|
||||
assert structured_req is not None
|
||||
grammar = structured_req.grammar
|
||||
if TYPE_CHECKING:
|
||||
assert isinstance(grammar, StructuredOutputGrammar)
|
||||
prefix = spec_tokens[:constraint_start]
|
||||
validated = grammar.validate_tokens(spec_tokens[constraint_start:])
|
||||
return prefix + validated
|
||||
|
||||
def grammar_bitmask(
|
||||
self,
|
||||
requests: dict[str, "Request"],
|
||||
@@ -264,7 +358,7 @@ class StructuredOutputManager:
|
||||
if TYPE_CHECKING:
|
||||
assert isinstance(grammar, StructuredOutputGrammar)
|
||||
|
||||
apply_bitmask = self.should_fill_bitmask(request)
|
||||
apply_bitmask = self._get_constraint_start(request, ()) == 0
|
||||
batch.append((grammar, cumulative_index, apply_bitmask))
|
||||
if len(batch) == self.fill_bitmask_parallel_batch_size:
|
||||
promises.append(self._async_submit_fill_bitmask(batch))
|
||||
@@ -288,76 +382,47 @@ class StructuredOutputManager:
|
||||
grammar = structured_output_request.grammar
|
||||
if TYPE_CHECKING:
|
||||
assert isinstance(grammar, StructuredOutputGrammar)
|
||||
apply_bitmask = self.should_fill_bitmask(request)
|
||||
|
||||
reasoner = None if apply_bitmask else self._get_reasoner(request)
|
||||
|
||||
req_tokens = scheduled_spec_decode_tokens.get(req_id, list())
|
||||
constraint_start = self._get_constraint_start(
|
||||
request, strip_speculative_padding(req_tokens)
|
||||
)
|
||||
state_advancements = 0
|
||||
post_reasoning_end_in_window = False
|
||||
req_tokens = scheduled_spec_decode_tokens.get(req_id, ())
|
||||
# Locate the reasoning-end marker in the draft window once.
|
||||
# Rejected-draft padding (-1) is trailing, so the window is
|
||||
# cut there and the reasoner never sees a placeholder.
|
||||
window_start = len(request.all_token_ids)
|
||||
reasoning_end_index: int | None = None
|
||||
if reasoner is not None and req_tokens:
|
||||
window = req_tokens
|
||||
if -1 in window:
|
||||
window = window[: window.index(-1)]
|
||||
reasoning_end_index = self._find_reasoning_end_index(
|
||||
reasoner,
|
||||
request.all_token_ids,
|
||||
window_start,
|
||||
window,
|
||||
delta_appended=False,
|
||||
)
|
||||
seen_padding = False
|
||||
failed = False
|
||||
for i, token in enumerate(req_tokens):
|
||||
apply_bitmask = (
|
||||
not failed and not seen_padding and i >= constraint_start
|
||||
)
|
||||
self._fill_bitmasks(((grammar, cumulative_index, apply_bitmask),))
|
||||
advance_grammar = apply_bitmask
|
||||
if token == -1:
|
||||
apply_bitmask = False
|
||||
advance_grammar = False
|
||||
elif not apply_bitmask and window_start + i == reasoning_end_index:
|
||||
# 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():
|
||||
if post_reasoning_end_in_window:
|
||||
accepted = bool(grammar.validate_tokens([token]))
|
||||
if accepted:
|
||||
accepted = grammar.accept_tokens(req_id, [token])
|
||||
else:
|
||||
accepted = grammar.accept_tokens(req_id, [token])
|
||||
if accepted:
|
||||
seen_padding = True
|
||||
elif apply_bitmask:
|
||||
if not grammar.is_terminated() and grammar.accept_tokens(
|
||||
req_id, [token]
|
||||
):
|
||||
state_advancements += 1
|
||||
elif not post_reasoning_end_in_window:
|
||||
raise AssertionError(
|
||||
(token, req_id, scheduled_spec_decode_tokens)
|
||||
else:
|
||||
failed = True
|
||||
logger.error(
|
||||
"Unexpected: grammar terminated or rejected draft "
|
||||
"token %s for request %s during bitmask fill.",
|
||||
token,
|
||||
req_id,
|
||||
)
|
||||
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
|
||||
bonus_apply = (
|
||||
not failed
|
||||
and not seen_padding
|
||||
and constraint_start <= len(req_tokens)
|
||||
)
|
||||
self._fill_bitmasks(((grammar, cumulative_index, bonus_apply),))
|
||||
cumulative_index += 1
|
||||
|
||||
if state_advancements > 0:
|
||||
grammar.rollback(state_advancements)
|
||||
|
||||
@@ -370,161 +435,29 @@ class StructuredOutputManager:
|
||||
# and deserialization when sending this to the GPU workers.
|
||||
return bitmask_tensor.numpy()
|
||||
|
||||
def should_fill_bitmask(self, request: "Request") -> bool:
|
||||
# NOTE (Hanchen) if enable_in_reasoning is True, it means that
|
||||
# the model needs to be constrained in reasoning. So we should always
|
||||
# enable the bitmask filling.
|
||||
structured_req = request.structured_output_request
|
||||
if self.enable_in_reasoning or (
|
||||
structured_req is not None and structured_req.reasoning_ended is True
|
||||
):
|
||||
def accept_tokens(self, request: "Request", new_token_ids: list[int]) -> bool:
|
||||
"""Advance grammar with accepted tokens. Returns False on rejection."""
|
||||
if not request.use_structured_output or not new_token_ids:
|
||||
return True
|
||||
|
||||
reasoner = self._get_reasoner(request)
|
||||
if reasoner is not None:
|
||||
assert structured_req is not None
|
||||
if structured_req.reasoning_ended is None:
|
||||
# This should be removed here, but since `openai_gptoss`
|
||||
# is an independent code path, it is kept for now.
|
||||
# After unifying the `openai_gptoss` and non-`openai_gptoss` styles,
|
||||
# it can be removed.
|
||||
structured_req.reasoning_ended = reasoner.is_reasoning_end(
|
||||
request.prompt_token_ids or []
|
||||
)
|
||||
return structured_req.reasoning_ended
|
||||
return True
|
||||
|
||||
def should_advance(
|
||||
self,
|
||||
request: "Request",
|
||||
new_token_ids: list[int] | None = None,
|
||||
) -> bool:
|
||||
if not request.use_structured_output:
|
||||
return False
|
||||
|
||||
# To determine whether we can advance the FSM.
|
||||
# Supports thinking usage where we skip the reasoning components.
|
||||
structured_req = request.structured_output_request
|
||||
if TYPE_CHECKING:
|
||||
assert request.structured_output_request is not None
|
||||
assert request.structured_output_request.grammar is not None
|
||||
structured_req = request.structured_output_request
|
||||
if self.enable_in_reasoning or (
|
||||
structured_req is not None and structured_req.reasoning_ended is True
|
||||
):
|
||||
return True
|
||||
assert structured_req is not None
|
||||
grammar = structured_req.grammar
|
||||
if TYPE_CHECKING:
|
||||
assert isinstance(grammar, StructuredOutputGrammar)
|
||||
|
||||
# by default, we should always advance
|
||||
# for cases that don't use thinking mode.
|
||||
reasoner = self._get_reasoner(request)
|
||||
if reasoner is None:
|
||||
return True
|
||||
|
||||
assert structured_req is not None
|
||||
|
||||
# Check if reasoning ends in *this* step.
|
||||
# When the caller passes new_token_ids (the tokens that were just
|
||||
# appended this step), use it directly as the delta window. The
|
||||
# placeholder-derived fallback assumes num_output_placeholders ==
|
||||
# len(new_token_ids), which breaks under async scheduling + spec
|
||||
# decode when some drafts are rejected (#43388): the placeholder
|
||||
# count remains > 0 after the step and the computed delta window
|
||||
# starts past the reasoning-end marker.
|
||||
all_token_ids = request.all_token_ids
|
||||
if new_token_ids:
|
||||
# The tokens were already appended this step, so the step window
|
||||
# starts exactly len(new_token_ids) from the end.
|
||||
start = len(all_token_ids) - len(new_token_ids)
|
||||
delta_ids: Sequence[int] = new_token_ids
|
||||
else:
|
||||
# Draft-validation callers arrive here after update_from_output
|
||||
# has already persisted any boundary found this step, so this
|
||||
# window only needs to cover tokens not examined yet.
|
||||
delta_from = request.num_computed_tokens - request.num_output_placeholders
|
||||
start = (
|
||||
delta_from
|
||||
if delta_from >= 0
|
||||
else max(len(all_token_ids) + delta_from, 0)
|
||||
)
|
||||
delta_ids = all_token_ids[start:]
|
||||
end_index = self._find_reasoning_end_index(
|
||||
reasoner, all_token_ids, start, delta_ids, delta_appended=True
|
||||
constraint_start = self._get_constraint_start(
|
||||
request, new_token_ids, spec_tokens_committed=True
|
||||
)
|
||||
if end_index is not None:
|
||||
structured_req.reasoning_ended = True
|
||||
# Record the boundary so the scheduler can exclude reasoning tokens.
|
||||
structured_req.reasoning_end_token_index = end_index
|
||||
# Early return only when the constraint hasn't started.
|
||||
# Otherwise, latch structured_req.reasoning_ended.
|
||||
if constraint_start > len(new_token_ids):
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def _find_reasoning_end_index(
|
||||
reasoner: "ReasoningParser",
|
||||
all_token_ids: Sequence[int],
|
||||
start: int,
|
||||
delta_ids: Sequence[int],
|
||||
*,
|
||||
delta_appended: bool,
|
||||
) -> int | None:
|
||||
"""Locates the token that ends reasoning within ``delta_ids``.
|
||||
|
||||
``delta_ids`` occupy positions ``start..``; ``delta_appended`` says
|
||||
whether they are already in ``all_token_ids`` (accepted tokens) or
|
||||
not (speculative drafts).
|
||||
|
||||
Legacy parsers are probed one token at a time. Accepted tokens get a
|
||||
whole-delta check as well, and a marker seen only that way reports
|
||||
the final token. Drafts skip it: some predicates are order-sensitive
|
||||
over a multi-token window and would miss the marker.
|
||||
|
||||
Returns:
|
||||
Absolute index of the last reasoning token, or ``None``.
|
||||
"""
|
||||
if (
|
||||
isinstance(reasoner, ParserEngineReasoningAdapter)
|
||||
and reasoner.reasoning_end_token_ids
|
||||
):
|
||||
offset = reasoner.find_reasoning_end_offset(delta_ids)
|
||||
return None if offset is None else start + offset
|
||||
|
||||
if delta_appended and not reasoner.is_reasoning_end_streaming(
|
||||
all_token_ids, delta_ids
|
||||
):
|
||||
return None
|
||||
prefix = list(itertools.islice(all_token_ids, start))
|
||||
for offset, token in enumerate(delta_ids):
|
||||
prefix.append(token)
|
||||
if reasoner.is_reasoning_end_streaming(prefix, [token]):
|
||||
return start + offset
|
||||
if delta_appended:
|
||||
return len(all_token_ids) - 1
|
||||
return None
|
||||
|
||||
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:]
|
||||
structured_req.reasoning_ended = True
|
||||
return grammar.accept_tokens(
|
||||
request.request_id, new_token_ids[constraint_start:]
|
||||
)
|
||||
|
||||
def clear_backend(self) -> None:
|
||||
if self.backend is not None:
|
||||
|
||||
@@ -25,12 +25,6 @@ class StructuredOutputRequest:
|
||||
Future[StructuredOutputGrammar] | StructuredOutputGrammar | Exception | 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.
|
||||
|
||||
Reference in New Issue
Block a user