| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import json |
| 5 | from dataclasses import dataclass, field |
| 6 | from typing import TYPE_CHECKING, Any |
| 7 | |
| 8 | import torch |
| 9 | |
| 10 | import vllm.envs |
| 11 | from vllm.exceptions import VLLMValidationError |
| 12 | from vllm.logger import init_logger |
| 13 | from vllm.sampling_params import SamplingParams |
| 14 | from vllm.utils.import_utils import LazyLoader |
| 15 | from vllm.utils.mistral import is_mistral_tokenizer |
| 16 | from vllm.v1.structured_output.backend_types import ( |
| 17 | StructuredOutputBackend, |
| 18 | StructuredOutputGrammar, |
| 19 | StructuredOutputOptions, |
| 20 | ) |
| 21 | from vllm.v1.structured_output.utils import ( |
| 22 | choice_as_grammar, |
| 23 | compile_regex_with_timeout, |
| 24 | grammar_is_likely_lark, |
| 25 | ) |
| 26 | |
| 27 | if TYPE_CHECKING: |
| 28 | import xgrammar as xgr |
| 29 | else: |
| 30 | xgr = LazyLoader("xgr", globals(), "xgrammar") |
| 31 | |
| 32 | logger = init_logger(__name__) |
| 33 | |
| 34 | |
| 35 | @dataclass |
| 36 | class XgrammarBackend(StructuredOutputBackend): |
| 37 | def __post_init__(self): |
| 38 | self.disable_any_whitespace = ( |
| 39 | self.vllm_config.structured_outputs_config.disable_any_whitespace |
| 40 | ) |
| 41 | model_config = self.vllm_config.model_config |
| 42 | is_plamo3 = ( |
| 43 | model_config is not None and model_config.hf_config.model_type == "plamo3" |
| 44 | ) |
| 45 | |
| 46 | if is_mistral_tokenizer(self.tokenizer): |
| 47 | # NOTE: ideally, xgrammar should handle this accordingly. |
| 48 | # refer to https://github.com/mlc-ai/xgrammar/blob/d77c0a0173ef14779c918e3be7966ba852f7910f/python/xgrammar/tokenizer_info.py#L98 |
| 49 | stop_token_ids = [self.tokenizer.eos_token_id] |
| 50 | |
| 51 | # not self.tokenizer.vocab_size as self.tokenizer.vocab |
| 52 | # collapses all decoded errors into a single token. |
| 53 | self.vocab_size = len(self.tokenizer.vocab) |
| 54 | tokenizer_info = xgr.TokenizerInfo( # type: ignore |
| 55 | encoded_vocab=self.tokenizer.vocab, |
| 56 | # NOTE: https://github.com/mlc-ai/xgrammar/blob/5e141f6ff1ca02bc31f9e512e68b61f2a8ae88e5/tests/python/test_tokenizer_info.py#L43 # noqa: E501 |
| 57 | vocab_type=xgr.VocabType.RAW |
| 58 | if self.tokenizer.is_tekken |
| 59 | else xgr.VocabType.BYTE_FALLBACK, |
| 60 | vocab_size=self.vocab_size, |
| 61 | stop_token_ids=stop_token_ids, |
| 62 | add_prefix_space=True, |
| 63 | ) |
| 64 | elif is_plamo3 and callable( |
| 65 | init_xgrammar := getattr(self.tokenizer, "init_xgrammar", None) |
| 66 | ): |
| 67 | tokenizer_info, _ = init_xgrammar() |
| 68 | else: |
| 69 | tokenizer_info = xgr.TokenizerInfo.from_huggingface( |
| 70 | self.tokenizer, |
| 71 | vocab_size=self.vocab_size, |
| 72 | ) |
| 73 | self.compiler = xgr.GrammarCompiler( |
| 74 | tokenizer_info, |
| 75 | max_threads=8, |
| 76 | cache_enabled=True, |
| 77 | cache_limit_bytes=vllm.envs.VLLM_XGRAMMAR_CACHE_MB * 1024 * 1024, |
| 78 | ) |
| 79 | |
| 80 | self.num_speculative_tokens = 0 |
| 81 | if self.vllm_config.speculative_config is not None: |
| 82 | self.num_speculative_tokens = ( |
| 83 | self.vllm_config.speculative_config.num_speculative_tokens |
| 84 | ) |
| 85 | |
| 86 | def compile_grammar( |
| 87 | self, |
| 88 | request_type: StructuredOutputOptions, |
| 89 | grammar_spec: str, |
| 90 | stop_token_ids: set[int] | None = None, |
| 91 | ) -> StructuredOutputGrammar: |
| 92 | if request_type == StructuredOutputOptions.JSON: |
| 93 | ctx = self.compiler.compile_json_schema( |
| 94 | grammar_spec, any_whitespace=not self.disable_any_whitespace |
| 95 | ) |
| 96 | elif request_type == StructuredOutputOptions.JSON_OBJECT: |
| 97 | ctx = self.compiler.compile_json_schema( |
| 98 | '{"type": "object"}', any_whitespace=not self.disable_any_whitespace |
| 99 | ) |
| 100 | elif request_type == StructuredOutputOptions.GRAMMAR: |
| 101 | if grammar_is_likely_lark(grammar_spec): |
| 102 | ctx = self.compiler.compile_lark(grammar_spec) |
| 103 | else: |
| 104 | ctx = self.compiler.compile_grammar(grammar_spec) |
| 105 | elif request_type == StructuredOutputOptions.REGEX: |
| 106 | ctx = compile_regex_with_timeout( |
| 107 | self.compiler.compile_regex, |
| 108 | grammar_spec, |
| 109 | ) |
| 110 | elif request_type == StructuredOutputOptions.STRUCTURAL_TAG: |
| 111 | s_tag = json.loads(grammar_spec) |
| 112 | if "structures" in s_tag: |
| 113 | # Falling back to deprecated method of compiling structural tag |
| 114 | tags = [ |
| 115 | xgr.StructuralTagItem( |
| 116 | begin=s["begin"], |
| 117 | schema=json.dumps(s["schema"]), |
| 118 | end=s["end"], |
| 119 | ) |
| 120 | for s in s_tag["structures"] |
| 121 | ] |
| 122 | ctx = self.compiler.compile_structural_tag(tags, s_tag["triggers"]) |
| 123 | else: |
| 124 | ctx = self.compiler.compile_structural_tag(grammar_spec) |
| 125 | else: |
| 126 | logger.error( |
| 127 | "Validation should have already occurred. Please file an issue." |
| 128 | ) |
| 129 | raise ValueError( |
| 130 | f"grammar is not of valid supported types. ({request_type!s})" |
| 131 | ) |
| 132 | |
| 133 | return XgrammarGrammar( |
| 134 | matcher=xgr.GrammarMatcher( |
| 135 | ctx, |
| 136 | override_stop_tokens=list(stop_token_ids) if stop_token_ids else None, |
| 137 | max_rollback_tokens=self.num_speculative_tokens, |
| 138 | ), |
| 139 | vocab_size=self.vocab_size, |
| 140 | ctx=ctx, |
| 141 | ) |
| 142 | |
| 143 | def allocate_token_bitmask(self, max_num_seqs: int): |
| 144 | return xgr.allocate_token_bitmask(max_num_seqs, self.vocab_size) |
| 145 | |
| 146 | def destroy(self): |
| 147 | del self.compiler |
| 148 | |
| 149 | |
| 150 | @dataclass |
| 151 | class XgrammarGrammar(StructuredOutputGrammar): |
| 152 | # NOTE: This would be a generic-enough class for |
| 153 | # supporting different backends, in the future. |
| 154 | # For now, just xgrammar. |
| 155 | # |
| 156 | # https://xgrammar.mlc.ai/docs/api/python/index.html#xgrammar.GrammarMatcher.find_jump_forward_string |
| 157 | # for jump-forward decoding |
| 158 | |
| 159 | vocab_size: int |
| 160 | matcher: xgr.GrammarMatcher = field(hash=False) |
| 161 | ctx: xgr.CompiledGrammar = field(hash=False) |
| 162 | num_processed_tokens: int = field( |
| 163 | default_factory=lambda: 0, repr=False, hash=False, init=False |
| 164 | ) |
| 165 | _is_terminated: bool = field(default=False, repr=False, hash=False) |
| 166 | |
| 167 | def accept_tokens(self, request_id: str, tokens: list[int]) -> bool: |
| 168 | """Accepts a list of tokens and advances the FSM. |
| 169 | |
| 170 | Returns True if all grammar-constrained tokens were accepted. |
| 171 | Tokens after termination are ignored. Returns False if the FSM |
| 172 | failed to advance. |
| 173 | """ |
| 174 | if self._is_terminated: |
| 175 | return True |
| 176 | for token in tokens: |
| 177 | if not self.matcher.accept_token(token): |
| 178 | logger.error( |
| 179 | "Failed to advance FSM for request %s " |
| 180 | "for tokens %s. Please file an issue.", |
| 181 | request_id, |
| 182 | token, |
| 183 | ) |
| 184 | return False |
| 185 | self.num_processed_tokens += 1 |
| 186 | self._is_terminated = self.matcher.is_terminated() |
| 187 | if self._is_terminated: |
| 188 | break |
| 189 | return True |
| 190 | |
| 191 | def validate_tokens(self, tokens: list[int]) -> list[int]: |
| 192 | """Checks if the list of tokens are accepted by the FSM in sequence. |
| 193 | Will not advance the FSM. |
| 194 | |
| 195 | Returns the prefix list of tokens that are accepted by the FSM. |
| 196 | """ |
| 197 | if self._is_terminated: |
| 198 | return [] |
| 199 | |
| 200 | accepted_tokens = [] |
| 201 | for token in tokens: |
| 202 | if self.matcher.accept_token(token): |
| 203 | accepted_tokens.append(token) |
| 204 | if self.matcher.is_terminated(): |
| 205 | break |
| 206 | else: |
| 207 | break |
| 208 | if len(accepted_tokens) > 0: |
| 209 | # Rollback the FSM to the initial state |
| 210 | self.matcher.rollback(len(accepted_tokens)) |
| 211 | return accepted_tokens |
| 212 | |
| 213 | def rollback(self, num_tokens: int) -> None: |
| 214 | self.matcher.rollback(num_tokens) |
| 215 | self.num_processed_tokens -= num_tokens |
| 216 | self._is_terminated = self.matcher.is_terminated() |
| 217 | |
| 218 | def fill_bitmask(self, bitmask: torch.Tensor, idx: int) -> None: |
| 219 | self.matcher.fill_next_token_bitmask(bitmask, idx) |
| 220 | |
| 221 | def is_terminated(self) -> bool: |
| 222 | return self._is_terminated |
| 223 | |
| 224 | def reset(self): |
| 225 | self.matcher.reset() |
| 226 | self.num_processed_tokens = 0 |
| 227 | self._is_terminated = False |
| 228 | |
| 229 | |
| 230 | # cf https://github.com/mlc-ai/xgrammar/blob/a32ac892676d2eedc0327416105b9b06edfb94b2/cpp/json_schema_converter.cc |
| 231 | STRING_SUPPORTED_FORMATS = { |
| 232 | "email", |
| 233 | "date", |
| 234 | "time", |
| 235 | "date-time", |
| 236 | "duration", |
| 237 | "ipv4", |
| 238 | "ipv6", |
| 239 | "hostname", |
| 240 | "uuid", |
| 241 | "uri", |
| 242 | "uri-reference", |
| 243 | "uri-template", |
| 244 | "json-pointer", |
| 245 | "relative-json-pointer", |
| 246 | } |
| 247 | |
| 248 | |
| 249 | def _has_pattern_and_length_bounds(schema: dict[str, Any]) -> bool: |
| 250 | return ("pattern" in schema or "format" in schema) and ( |
| 251 | "minLength" in schema or "maxLength" in schema |
| 252 | ) |
| 253 | |
| 254 | |
| 255 | # FIXME(arpera): The approach used here needs to be redesigned because of |
| 256 | # existing bugs: https://github.com/vllm-project/vllm/issues/57550 |
| 257 | def _schema_types(schema: dict[str, Any]) -> set[str]: |
| 258 | """Normalize a scalar or list-valued JSON Schema type.""" |
| 259 | schema_type = schema.get("type") |
| 260 | if isinstance(schema_type, str): |
| 261 | return {schema_type} |
| 262 | if isinstance(schema_type, list): |
| 263 | return {item for item in schema_type if isinstance(item, str)} |
| 264 | return set() |
| 265 | |
| 266 | |
| 267 | def has_xgrammar_unsupported_json_features(schema: dict[str, Any]) -> bool: |
| 268 | """Check if JSON schema contains features unsupported by xgrammar.""" |
| 269 | |
| 270 | def check_object(obj: dict[str, Any]) -> bool: |
| 271 | if not isinstance(obj, dict): |
| 272 | return False |
| 273 | |
| 274 | schema_types = _schema_types(obj) |
| 275 | |
| 276 | # Check for numeric ranges |
| 277 | if (schema_types & {"integer", "number"}) and ("multipleOf" in obj): |
| 278 | return True |
| 279 | |
| 280 | # Check for array unsupported keywords |
| 281 | if "array" in schema_types and any( |
| 282 | key in obj |
| 283 | for key in ("uniqueItems", "contains", "minContains", "maxContains") |
| 284 | ): |
| 285 | return True |
| 286 | |
| 287 | # Unsupported keywords for strings |
| 288 | if ( |
| 289 | "string" in schema_types |
| 290 | and "format" in obj |
| 291 | and obj["format"] not in STRING_SUPPORTED_FORMATS |
| 292 | ): |
| 293 | return True |
| 294 | |
| 295 | # A string mixing a generative constraint (pattern or format) with |
| 296 | # explicit length bounds. xgrammar compiles the pattern/format side |
| 297 | # and silently drops minLength/maxLength from the grammar, so output |
| 298 | # can violate the bound without any error surfacing. Verified against |
| 299 | # the compiled EBNF: pattern/format grammars come out byte-identical |
| 300 | # with and without the length keywords, while maxLength alone lowers |
| 301 | # to {0, N} correctly. |
| 302 | if "string" in schema_types and _has_pattern_and_length_bounds(obj): |
| 303 | return True |
| 304 | |
| 305 | # propertyNames validates names, so it is a string schema even when it |
| 306 | # omits "type", which is the form that escapes the check above. |
| 307 | if ( |
| 308 | "object" in schema_types |
| 309 | and isinstance(obj.get("propertyNames"), dict) |
| 310 | and _has_pattern_and_length_bounds(obj["propertyNames"]) |
| 311 | ): |
| 312 | return True |
| 313 | |
| 314 | # FIXME: propertyNames conflicts with properties/patternProperties/ |
| 315 | # additionalProperties/unevaluatedProperties under xgrammar. |
| 316 | # https://github.com/mlc-ai/xgrammar/issues/826 |
| 317 | if ( |
| 318 | "object" in schema_types |
| 319 | and "propertyNames" in obj |
| 320 | and ( |
| 321 | "properties" in obj |
| 322 | or "patternProperties" in obj |
| 323 | or isinstance(obj.get("additionalProperties"), dict) |
| 324 | or obj.get("unevaluatedProperties", True) is not True |
| 325 | ) |
| 326 | ): |
| 327 | return True |
| 328 | |
| 329 | # FIXME: multiple patternProperties, or patternProperties alongside |
| 330 | # properties, conflict under xgrammar. |
| 331 | if ( |
| 332 | "object" in schema_types |
| 333 | and isinstance(obj.get("patternProperties"), dict) |
| 334 | and ("properties" in obj or len(obj["patternProperties"]) > 1) |
| 335 | ): |
| 336 | return True |
| 337 | |
| 338 | # Recursively check all nested objects and arrays |
| 339 | for value in obj.values(): |
| 340 | if isinstance(value, dict): |
| 341 | if check_object(value): |
| 342 | return True |
| 343 | elif isinstance(value, list): |
| 344 | for item in value: |
| 345 | if isinstance(item, dict) and check_object(item): |
| 346 | return True |
| 347 | |
| 348 | return False |
| 349 | |
| 350 | return check_object(schema) |
| 351 | |
| 352 | |
| 353 | def validate_xgrammar_grammar(sampling_params: SamplingParams) -> None: |
| 354 | """Validate that the request is supported by structured output. |
| 355 | |
| 356 | Raises VLLMValidationError if the request is not supported. |
| 357 | """ |
| 358 | if sampling_params.structured_outputs is None: |
| 359 | return |
| 360 | |
| 361 | so_params = sampling_params.structured_outputs |
| 362 | |
| 363 | if so_params.regex: |
| 364 | # A NUL byte is never meaningful in a regex pattern and is not handled |
| 365 | # by xgrammar's native regex converter. Reject it here, before the |
| 366 | # pattern reaches that native code; the try/except below does not cover |
| 367 | # this case. |
| 368 | if "\x00" in so_params.regex: |
| 369 | raise ValueError( |
| 370 | "structured_outputs.regex must not contain a NUL character ('\\x00')" |
| 371 | ) |
| 372 | try: |
| 373 | compile_regex_with_timeout( |
| 374 | xgr.Grammar.from_regex, |
| 375 | so_params.regex, |
| 376 | ) |
| 377 | except Exception as err: |
| 378 | raise VLLMValidationError( |
| 379 | f"Failed to transform regex into a grammar: {err}" |
| 380 | ) from err |
| 381 | |
| 382 | if so_params.choice: |
| 383 | choice_grammar = choice_as_grammar(so_params.choice) |
| 384 | try: |
| 385 | xgr.Grammar.from_ebnf(choice_grammar) |
| 386 | except Exception as err: |
| 387 | raise VLLMValidationError( |
| 388 | f"Failed to transform choices into a grammar: {err}" |
| 389 | ) from err |
| 390 | so_params.choice = None |
| 391 | so_params.grammar = choice_grammar |
| 392 | return |
| 393 | |
| 394 | if so_params.json: |
| 395 | if isinstance(so_params.json, str): |
| 396 | try: |
| 397 | schema = json.loads(so_params.json) |
| 398 | except json.JSONDecodeError as e: |
| 399 | raise VLLMValidationError("Invalid JSON grammar specification.") from e |
| 400 | else: |
| 401 | schema = so_params.json |
| 402 | |
| 403 | if has_xgrammar_unsupported_json_features(schema): |
| 404 | raise VLLMValidationError( |
| 405 | "The provided JSON schema contains features not supported by xgrammar." |
| 406 | ) |
| 407 | |
| 408 | try: |
| 409 | xgr.Grammar.from_json_schema(schema) |
| 410 | except Exception as err: |
| 411 | raise VLLMValidationError( |
| 412 | f"Failed to transform json schema into a grammar: {err}" |
| 413 | ) from err |
| 414 | return |
| 415 | |
| 416 | if so_params.grammar: |
| 417 | # Parse the grammar with the same syntax `compile_grammar` will use, |
| 418 | # but don't compile it. The grammar is passed on unchanged. |
| 419 | try: |
| 420 | if grammar_is_likely_lark(so_params.grammar): |
| 421 | xgr.Grammar.from_lark(so_params.grammar) |
| 422 | else: |
| 423 | xgr.Grammar.from_ebnf(so_params.grammar) |
| 424 | except Exception as e: |
| 425 | raise VLLMValidationError("Invalid grammar specification.") from e |
| 426 | return |
| 427 | |
| 428 | if so_params.structural_tag: |
| 429 | try: |
| 430 | s_tag = json.loads(so_params.structural_tag) |
| 431 | |
| 432 | # Using the deprecated method of compiling structural tag |
| 433 | if "structures" in s_tag: |
| 434 | tags = [ |
| 435 | xgr.StructuralTagItem( |
| 436 | begin=s["begin"], |
| 437 | schema=json.dumps(s["schema"]), |
| 438 | end=s["end"], |
| 439 | ) |
| 440 | for s in s_tag["structures"] |
| 441 | ] |
| 442 | xgr.Grammar.from_structural_tag(tags, s_tag["triggers"]) |
| 443 | else: |
| 444 | xgr.Grammar.from_structural_tag(so_params.structural_tag) |
| 445 | except Exception as e: |
| 446 | raise VLLMValidationError("Invalid structural tag specification.") from e |