Structured outputs: JSON schemas, grammars and token bitmasks in vLLM

StructuredOutputsBackend = Literal[
    "auto", "xgrammar", "guidance", "outlines", "lm-format-enforcer"
]


@config
class StructuredOutputsConfig:
    """Dataclass which contains structured outputs config for the engine."""

    backend: StructuredOutputsBackend = "auto"
    """Which engine will be used for structured outputs (e.g. JSON schema,
    regex, etc) by default. With "auto", we will make opinionated choices
    based on request contents and what the backend libraries currently support,
    so the behavior is subject to change in each release."""
    disable_any_whitespace: bool = False
    """If `True`, json output will always be compact without any whitespace.
    If `False`, the model may generate whitespace between JSON fields,
    which is still valid JSON. This is only supported for xgrammar
    and guidance backends."""
    disable_additional_properties: bool = False
    """If `True`, the `guidance` backend will not use `additionalProperties`
    in the JSON schema. This is only supported for the `guidance` backend and
    is used to better align its behaviour with `outlines` and `xgrammar`."""
    reasoning_parser: str = ""
    """Select the reasoning parser depending on the model that you're using.
    This is used to parse the reasoning content into OpenAI API format."""
    reasoning_parser_plugin: str = ""
    """Path to a dynamically reasoning parser plugin that can be dynamically
    loaded and registered."""
    enable_in_reasoning: bool = False
    """Whether to use structured input for reasoning."""

        else:
            # NOTE: backend must be "auto" here, because we have
            # checked supported_backends above.
            # In this mode, we set opinionated defaults based on what we think
            # will satisfy the most use cases without having to worry about
            # this setting. We include fallback behavior here, but not with any
            # other setting where a specific backend was specified.
            try:
                validate_xgrammar_grammar(self)
                self.structured_outputs._backend = "xgrammar"
            except VLLMValidationError:
                # The request either failed validation
                # or includes some jsonschema feature(s) that
                # are not supported in xgrammar.

                skip_guidance = _is_non_tekken_mistral(tokenizer)

                # Check if schema has features unsupported by guidance
                so_params = self.structured_outputs
                if not skip_guidance and so_params.json:
                    if isinstance(so_params.json, str):
                        try:
                            schema = json_mod.loads(so_params.json)
                        except json_mod.JSONDecodeError as e:
                            raise VLLMValidationError(
                                "Invalid JSON grammar specification."
                            ) from e
                    else:
                        schema = so_params.json
                    skip_guidance = has_guidance_unsupported_json_features(schema)

                if skip_guidance:
                    # Fall back to outlines for non-Tekken Mistral tokenizers or
                    # schemas containing features unsupported by guidance.
                    validate_structured_output_request_outlines(self)
                    self.structured_outputs._backend = "outlines"
                else:
                    # Fall back to guidance by default.
                    if not GuidanceBackend.is_tokenizer_supported(tokenizer):
                        raise VLLMValidationError(
                            "No compatible structured output backend was found for "
                            "this request. XGrammar rejected the constraint, and "
                            "Guidance does not support this tokenizer. Configure a "
                            "fast tokenizer or select a compatible backend explicitly."
                        ) from None
                    validate_guidance_grammar(
                        self,
                        tokenizer=_get_llg_tokenizer(tokenizer),
                    )
                    self.structured_outputs._backend = "guidance"
            # Remember that this backend was set automatically
            self.structured_outputs._backend_was_auto = True

        request_backend = structured_outputs._backend
        assert request_backend is not None
        if self.backend_name is not None and request_backend != self.backend_name:
            raise VLLMValidationError(
                "Structured output processing only supports one backend per engine. "
                f"The engine is already using '{self.backend_name}', but this "
                f"request resolved to '{request_backend}'. Configure "
                "`structured_outputs_config.backend` explicitly or use schemas "
                "supported by the initialized backend."
            )

        # 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:
            vocab_size = self.vllm_config.model_config.get_vocab_size()
            if request_backend == "xgrammar":
                self.backend = XgrammarBackend(
                    self.vllm_config,
                    tokenizer=self.tokenizer,
                    vocab_size=vocab_size,
                )
            elif request_backend == "guidance":
                self.backend = GuidanceBackend(
                    self.vllm_config,
                    tokenizer=self.tokenizer,
                    vocab_size=vocab_size,
                )
            elif request_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 request_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: {request_backend}"
                )
            self.backend_name = request_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 get_grammar_bitmask(
        self, scheduler_output: SchedulerOutput
    ) -> GrammarOutput | None:
        # Collect list of scheduled request ids that use structured output.
        # The corresponding rows of the bitmask will be in this order.
        if not scheduler_output.has_structured_output_requests:
            return None

        structured_output_request_ids = [
            req_id
            for req_id in scheduler_output.num_scheduled_tokens
            if (req := self.requests.get(req_id))
            and (req.use_structured_output and not req.is_prefill_chunk)
        ]
        if not structured_output_request_ids:
            return None

        bitmask = self.structured_output_manager.grammar_bitmask(
            self.requests,
            structured_output_request_ids,
            scheduler_output.scheduled_spec_decode_tokens,
        )
        spec_tokens = scheduler_output.scheduled_spec_decode_tokens
        num_acceptable_drafts = [
            len(strip_speculative_padding(spec_tokens.get(req_id, [])))
            for req_id in structured_output_request_ids
        ]
        return GrammarOutput(
            structured_output_request_ids, bitmask, num_acceptable_drafts
        )

            sample_hidden_states = hidden_states[input_batch.logits_indices]
            logits = self.model.compute_logits(sample_hidden_states)

        invalid_drafts = None
        # A diffusion prefill has no logit rows even when a bitmask row
        # arrived for it.
        if grammar_output is not None and logits.shape[0] > 0:
            # Apply grammar bitmask to the logits in-place.
            assert self.structured_outputs_worker is not None
            self.structured_outputs_worker.apply_grammar_bitmask(
                logits,
                input_batch,
                grammar_output.structured_output_request_ids,
                grammar_output.grammar_bitmask,
            )
            invalid_drafts = grammar_invalid_drafts(
                input_batch,
                grammar_output.structured_output_request_ids,
                grammar_output.num_acceptable_drafts,
            )

        sampler_output: SamplerOutput | None
        if input_batch.num_reqs == 0:
            # This rank owns no requests this step. It contributes an
            # all-padding block to the gather below.
            sampler_output = None
        elif input_batch.num_draft_tokens == 0 or self.rejection_sampler is None:
            assert self.sampler is not None
            sampler_output = self.sampler(logits, input_batch)

    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)

        grammar_start, _ = self._get_constraint_bounds(
            request, new_token_ids, spec_tokens_committed=True
        )
        # Early return only when no reasoning-end token is present.
        # Otherwise, latch structured_req.reasoning_ended.
        if grammar_start > len(new_token_ids):
            return True
        structured_req.reasoning_ended = True
        return grammar.accept_tokens(request.request_id, new_token_ids[grammar_start:])