class StructuredOutputManager:
"""Engine-level manager for structured output requests."""
def __init__(self, vllm_config: VllmConfig):
self.backend: StructuredOutputBackend | None = None
# We only store the class of the reasoner in the manager.
# The parser instance is request-scoped because some reasoning parsers
# depend on per-request chat-template kwargs.
self.reasoner_cls: type[ReasoningParser] | None = None
self.vllm_config = vllm_config
# When in external_launcher mode, async grammar compilation causes deadlocks
# due to external_launcher mode having a scheduler for each TP rank.
# Async grammar compilation causes the
# WAITING_FOR_STRUCTURED_OUTPUT_GRAMMAR → WAITING transition to
# happen at different times on different TP ranks,
# breaking the determinism assumption that external_launcher relies on.
self._use_async_grammar_compilation = (
vllm_config.parallel_config.distributed_executor_backend
!= "external_launcher"
)
self._grammar_bitmask: torch.Tensor | None = None
self._full_mask = torch.tensor(-1, dtype=torch.int32)
max_batch_size = self.vllm_config.scheduler_config.max_num_seqs
self.fill_bitmask_parallel_threshold = 128
if self.fill_bitmask_parallel_threshold < max_batch_size:
self.fill_bitmask_parallel_batch_size = 16
# Use:
# - at least 1 CPU
# - at most half the number of CPUs or 8, whichever is less
max_workers = max(1, min(multiprocessing.cpu_count() // 2, 8))
self.executor_for_fillmask = ThreadPoolExecutor(max_workers=max_workers)
if not self.vllm_config.model_config.skip_tokenizer_init:
# The default max_workers if not specified is the number of
# CPUs * 5, which is way too high since these tasks are CPU-bound,
# not I/O bound. We also know we would never dominate CPU usage
# with just grammar compilation, so we set it to half the number
# of CPUs.
max_workers = max(1, (multiprocessing.cpu_count() + 1) // 2)
self.executor = ThreadPoolExecutor(max_workers=max_workers)
self.tokenizer = cached_tokenizer_from_config(
model_config=self.vllm_config.model_config
)
reasoning_parser_plugin = (
self.vllm_config.structured_outputs_config.reasoning_parser_plugin
)
if reasoning_parser_plugin and len(reasoning_parser_plugin) > 3:
ReasoningParserManager.import_reasoning_parser(reasoning_parser_plugin)
reasoning_parser = (
self.vllm_config.structured_outputs_config.reasoning_parser
)
if reasoning_parser:
self.reasoner_cls = ReasoningParserManager.get_reasoning_parser(
reasoning_parser
)
self.enable_in_reasoning = (
self.vllm_config.structured_outputs_config.enable_in_reasoning
)
def _get_reasoner(self, request: "Request") -> "ReasoningParser | None":
structured_req = request.structured_output_request
if structured_req is None or self.reasoner_cls is None:
return None
if structured_req.reasoner is None:
# Lazily build the request-local parser so the structured-output
# gate observes the same template kwargs used by the frontend.
parser_kwargs = structured_req.reasoning_parser_kwargs or {}
structured_req.reasoner = self.reasoner_cls(
tokenizer=self.tokenizer,
**parser_kwargs,
)
return structured_req.reasoner
def grammar_init(self, request: "Request") -> None:
if request.structured_output_request is None:
return
if TYPE_CHECKING:
assert (
request.sampling_params is not None
and request.sampling_params.structured_outputs is not None
)
# Initialize the backend the first time it is needed.
#
# NOTE: We only support a single backend. We do NOT support different
# backends on a per-request basis in V1 (for now, anyway...).
# _backend is set in Processor._validate_structured_output
if self.backend is None:
assert request.sampling_params is not None
backend = request.sampling_params.structured_outputs._backend
vocab_size = self.vllm_config.model_config.get_vocab_size()
if backend == "xgrammar":
self.backend = XgrammarBackend(
self.vllm_config,
tokenizer=self.tokenizer,
vocab_size=vocab_size,
)
elif backend == "guidance":
self.backend = GuidanceBackend(
self.vllm_config,
tokenizer=self.tokenizer,
vocab_size=vocab_size,
)
elif backend == "outlines":
from vllm.v1.structured_output.backend_outlines import OutlinesBackend
self.backend = OutlinesBackend(
self.vllm_config,
tokenizer=self.tokenizer,
vocab_size=vocab_size,
)
elif backend == "lm-format-enforcer":
from vllm.v1.structured_output.backend_lm_format_enforcer import ( # noqa: E501
LMFormatEnforcerBackend,
)
self.backend = LMFormatEnforcerBackend(
self.vllm_config,
tokenizer=self.tokenizer,
vocab_size=vocab_size,
)
else:
raise ValueError(f"Unsupported structured output backend: {backend}")
grammar: Future[StructuredOutputGrammar] | StructuredOutputGrammar
if self._use_async_grammar_compilation:
grammar = self.executor.submit(self._create_grammar, request)
else:
try:
grammar = self._create_grammar(request)
except Exception as e:
grammar = Future()
grammar.set_exception(e)
request.structured_output_request.grammar = grammar
def _create_grammar(self, request: "Request") -> StructuredOutputGrammar:
struct_request = request.structured_output_request
assert struct_request is not None
# Note that the request was validated in the engine core client,
# so at this point we know it is a supported type of request. Grammar
# compilation may still fail; the Future carries that error to the
# scheduler so it can fail only this request.
try:
request_type, grammar_spec = struct_request.structured_output_key
assert self.backend is not None
stop_token_ids = (
request.sampling_params.all_stop_token_ids
if request.sampling_params is not None
else None
)
return self.backend.compile_grammar(
request_type, grammar_spec, stop_token_ids=stop_token_ids
)
except Exception:
logger.exception(
"Failed to compile grammar for request %s", request.request_id
)
raise
def _fill_bitmasks(
self, batch: Iterable[tuple[StructuredOutputGrammar, int, bool]]
) -> None:
assert self._grammar_bitmask is not None
for grammar, index, apply_bitmask in batch:
if apply_bitmask and not grammar.is_terminated():
grammar.fill_bitmask(self._grammar_bitmask, index)
else:
# Note that for thinking support, we will need to
# reset the relevant part of the bitmask for consequent
# requests here.
self._grammar_bitmask[index].fill_(self._full_mask)
def _async_submit_fill_bitmask(
self, batch: list[tuple[StructuredOutputGrammar, int, bool]]
) -> 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"],
structured_output_request_ids: list[str],
scheduled_spec_decode_tokens: dict[str, list[int]],
) -> "npt.NDArray[np.int32] | None":
# Prepare the structured output bitmask for this batch.
if not structured_output_request_ids:
return None
# Covers both speculative decoding and diffusion LLMs (canvas_length).
max_num_spec_tokens = self.vllm_config.num_speculative_tokens
if self._grammar_bitmask is None:
assert self.backend is not None
max_batch_size = self.vllm_config.scheduler_config.max_num_seqs
# Allocate a bitmask for each token needing to be checked:
# one for each speculative position, and one more for the
# bonus token / non-speculative token.
self._grammar_bitmask = self.backend.allocate_token_bitmask(
max_batch_size * (1 + max_num_spec_tokens)
)
# Generate a batched bitmask for all structured output requests.
# When speculative decoding is enabled, we need to include multiple
# masks for each request, one for each possible bonus token position.
# These are stored inline in the tensor and unpacked by the gpu runner.
cumulative_index = 0
# Optimized parallel filling of bitmasks for
# non-spec, large-batch-size cases
if (
len(structured_output_request_ids) > self.fill_bitmask_parallel_threshold
and max_num_spec_tokens == 0
):
promises = []
batch = []
for req_id in structured_output_request_ids:
request = requests[req_id]
structured_output_request = request.structured_output_request
if TYPE_CHECKING:
assert structured_output_request is not None
grammar = structured_output_request.grammar
if TYPE_CHECKING:
assert isinstance(grammar, StructuredOutputGrammar)
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))
batch = []
cumulative_index += 1
if batch:
promises.append(self._async_submit_fill_bitmask(batch))
# Wait for all bitmask filling tasks to complete.
for promise in promises:
promise.result()
else:
# Fallback to serial filling of bitmasks for small-batch-size cases
for req_id in structured_output_request_ids:
request = requests[req_id]
structured_output_request = request.structured_output_request
if TYPE_CHECKING:
assert structured_output_request is not None
grammar = structured_output_request.grammar
if TYPE_CHECKING:
assert isinstance(grammar, StructuredOutputGrammar)
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
seen_padding = False
# Row filled from the last valid grammar state before a draft
# was rejected; later rows reuse it rather than unconstraining.
failed_index: int | None = None
bitmask = self._grammar_bitmask
for i, token in enumerate(req_tokens):
if failed_index is not None:
bitmask[cumulative_index].copy_(bitmask[failed_index])
cumulative_index += 1
continue
apply_bitmask = not seen_padding and i >= constraint_start
self._fill_bitmasks(((grammar, cumulative_index, apply_bitmask),))
if token == -1:
seen_padding = True
elif apply_bitmask:
if not grammar.is_terminated() and grammar.accept_tokens(
req_id, [token]
):
state_advancements += 1
else:
failed_index = cumulative_index
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):
if failed_index is not None:
bitmask[cumulative_index].copy_(bitmask[failed_index])
else:
bonus_apply = 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)
bitmask_tensor = self._grammar_bitmask
if cumulative_index < bitmask_tensor.shape[0]:
bitmask_tensor = bitmask_tensor[:cumulative_index]
# After finishing with the xgrammar operations, we convert to
# np.ndarray, because that is much more efficient for serialization
# and deserialization when sending this to the GPU workers.
return bitmask_tensor.numpy()
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
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)
constraint_start = self._get_constraint_start(
request, new_token_ids, spec_tokens_committed=True
)
# Early return only when the constraint hasn't started.
# Otherwise, latch structured_req.reasoning_ended.
if constraint_start > len(new_token_ids):
return True
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:
self.backend.destroy()