generalize structured output manager and backends

Signed-off-by: william-baker-inflection <william.baker@inflection.ai>

Author: william-baker-inflection

vllm/v1/core/sched/output.py

@@ -7,15 +7,15 @@
 from typing import TYPE_CHECKING, Optional
 
 if TYPE_CHECKING:
-    import numpy as np
-    import numpy.typing as npt
 
     from vllm.distributed.kv_transfer.kv_connector.v1.base import (
         KVConnectorMetadata)
     from vllm.lora.request import LoRARequest
     from vllm.multimodal.inputs import MultiModalKwargs, PlaceholderRange
     from vllm.sampling_params import SamplingParams
     from vllm.v1.request import Request
+    from vllm.v1.structured_output.backend_types import (
+        StructuredOutputBatchMetaData)
 
 
 @dataclass
@@ -144,11 +144,10 @@ class SchedulerOutput:
     # Used to free the encoder cache.
     free_encoder_input_ids: list[tuple[str, int]]
 
-    # Dict of request ids to their index within the batch
-    # for filling the next token bitmask
-    structured_output_request_ids: dict[str, int]
-    # the bitmask for the whole batch
-    grammar_bitmask: Optional[npt.NDArray[np.int32]]
+    # Meta data for structured output batches
+    # By default this holds only the structured_output_request_ids
+    # but backends may extend this to hold more data for the batch
+    structured_output_meta: Optional[StructuredOutputBatchMetaData]
 
     # KV Cache Connector metadata.
     kv_connector_metadata: Optional[KVConnectorMetadata] = None

vllm/v1/core/sched/scheduler.py

@@ -505,11 +505,15 @@ def schedule(self) -> SchedulerOutput:
                 self.kv_cache_manager.get_num_common_prefix_blocks(
                     any_request, len(self.running)))
 
-        grammar_bitmask = self.structured_output_manager.grammar_bitmask(
-            self.requests,
-            structured_output_request_ids,
-            scheduled_spec_decode_tokens,
-        )
+        if self.structured_output_manager.backend is not None:
+            structured_output_meta = self.structured_output_manager.init_batch(
+                self.requests,
+                structured_output_request_ids,
+                scheduled_spec_decode_tokens,
+            )
+        else:
+            structured_output_meta = None
+
         # Construct the scheduler output.
         new_reqs_data = [
             NewRequestData.from_request(req,
@@ -548,9 +552,7 @@ def schedule(self) -> SchedulerOutput:
             # the previous and the current steps.
             finished_req_ids=self.finished_req_ids,
             free_encoder_input_ids=self.encoder_cache_manager.get_freed_ids(),
-            structured_output_request_ids=structured_output_request_ids,
-            grammar_bitmask=grammar_bitmask,
-        )
+            structured_output_meta=structured_output_meta)
 
         # NOTE(Kuntai): this function is designed for multiple purposes:
         # 1. Plan the KV cache store
@@ -784,8 +786,8 @@ def update_from_output(
                 # NOTE: structured_output_request
                 # should not be None if use_structured_output, we have
                 # check above, so safe to ignore type warning
-                request.structured_output_request.grammar.accept_tokens(  # type: ignore[union-attr]
-                    req_id, new_token_ids)
+                self.structured_output_manager.accept_tokens(
+                    request, req_id, new_token_ids)
 
             # Add newly generated spec token ids to the request.
             if spec_token_ids is not None:

vllm/v1/engine/processor.py

@@ -21,10 +21,7 @@
 from vllm.transformers_utils.tokenizer_group import TokenizerGroup
 from vllm.v1.engine import EngineCoreRequest
 from vllm.v1.engine.mm_input_cache import MirroredProcessingCache
-from vllm.v1.structured_output.backend_guidance import (
-    validate_guidance_grammar)
-from vllm.v1.structured_output.backend_xgrammar import (
-    validate_xgrammar_grammar)
+from vllm.v1.structured_output import StructuredOutputManager
 
 
 class Processor:
@@ -81,7 +78,7 @@ def _validate_sampling_params(
         params: SamplingParams,
         lora_request: Optional[LoRARequest],
     ) -> None:
-        self._validate_structured_output(params)
+        StructuredOutputManager.validate_request(params, self.vllm_config)
         self._validate_logit_bias(params)
 
         if params.allowed_token_ids is None:
@@ -148,59 +145,6 @@ def _validate_lora(self, lora_request: Optional[LoRARequest]) -> None:
             raise ValueError(f"Got lora_request {lora_request} but LoRA is "
                              "not enabled!")
 
-    def _validate_structured_output(self, params: SamplingParams) -> None:
-        if not params.guided_decoding or not self.decoding_config:
-            return
-
-        engine_level_backend = self.decoding_config.backend
-        if params.guided_decoding.backend:
-            # Request-level backend selection is not supported in V1.
-            # The values may differ if `params` is reused and was set
-            # to a specific backend based on `auto` behavior in a previous
-            # request. We remember that it was set as a result of `auto`
-            # using the `_auto` option set on the backend in the params.
-            if (params.guided_decoding.backend != engine_level_backend
-                    and not (engine_level_backend == "auto"
-                             and params.guided_decoding.backend_was_auto)):
-                raise ValueError(
-                    "Request-level structured output backend selection is no "
-                    "longer supported. The request specified "
-                    f"'{params.guided_decoding.backend}', but vLLM was "
-                    f"initialised with '{engine_level_backend}'. This error "
-                    "can be resolved by removing backend selection from the "
-                    "request.")
-        else:
-            params.guided_decoding.backend = engine_level_backend
-
-        # Request content validation
-        if engine_level_backend.startswith("xgrammar"):
-            # xgrammar with no fallback
-            validate_xgrammar_grammar(params)
-        elif engine_level_backend.startswith("guidance"):
-            # TODO: ideally we would have the LLTokenizer here as Lark syntax
-            # allows <|special_token|> and similar, see
-            # https://github.com/guidance-ai/llguidance/blob/main/docs/syntax.md#special-tokens
-            # Without tokenizer these are disallowed in grammars.
-            validate_guidance_grammar(params, tokenizer=None)
-        else:
-            # NOTE: engine_level_backend must be "auto" here, because we have
-            # checked supported_backends above.
-            # "auto" is an opt-in to opinionated behavior where we try to
-            # choose a backend based on request contents. This is not the
-            # default as it is less predictable and subject to change
-            # between releases as feature support changes.
-            try:
-                validate_xgrammar_grammar(params)
-                params.guided_decoding.backend = "xgrammar"
-            except ValueError:
-                # The request either failed validation
-                # or includes some jsonschema feature(s) that
-                # are not supported in xgrammar. Fall back to guidance.
-                validate_guidance_grammar(params, tokenizer=None)
-                params.guided_decoding.backend = "guidance"
-            # Remember that this backend was set automatically
-            params.guided_decoding.backend_was_auto = True
-
     def process_inputs(
         self,
         request_id: str,