| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from __future__ import annotations |
| 4 | |
| 5 | import hashlib |
| 6 | import importlib.metadata |
| 7 | import os |
| 8 | import sqlite3 |
| 9 | import tempfile |
| 10 | from collections.abc import Callable |
| 11 | from concurrent.futures import ThreadPoolExecutor, TimeoutError |
| 12 | from typing import TYPE_CHECKING, TypeVar |
| 13 | |
| 14 | import regex as re |
| 15 | import torch |
| 16 | from cachetools import LRUCache |
| 17 | |
| 18 | import vllm.envs as envs |
| 19 | from vllm.logger import init_logger |
| 20 | from vllm.utils.import_utils import LazyLoader |
| 21 | from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d |
| 22 | from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput |
| 23 | |
| 24 | if TYPE_CHECKING: |
| 25 | import outlines_core as oc |
| 26 | import transformers.convert_slow_tokenizer as convert_slow_tokenizer |
| 27 | import transformers.file_utils as file_utils |
| 28 | import xgrammar as xgr |
| 29 | |
| 30 | from vllm.tokenizers import TokenizerLike |
| 31 | from vllm.v1.worker.gpu_input_batch import InputBatch |
| 32 | else: |
| 33 | xgr = LazyLoader("xgr", globals(), "xgrammar") |
| 34 | oc = LazyLoader("oc", globals(), "outlines_core") |
| 35 | file_utils = LazyLoader("file_utils", globals(), "transformers.file_utils") |
| 36 | convert_slow_tokenizer = LazyLoader( |
| 37 | "convert_slow_tokenizer", globals(), "transformers.convert_slow_tokenizer" |
| 38 | ) |
| 39 | |
| 40 | |
| 41 | logger = init_logger(__name__) |
| 42 | |
| 43 | _T = TypeVar("_T") |
| 44 | |
| 45 | CACHE = None |
| 46 | |
| 47 | |
| 48 | def strip_speculative_padding(token_ids: list[int]) -> list[int]: |
| 49 | """Drop speculative-decoding padding from a token block. |
| 50 | |
| 51 | ngram and other speculative backends pad rejected draft positions with a |
| 52 | -1 sentinel. Structured-output grammars treat every entry as a real token |
| 53 | id, so the sentinels (and everything after the first one) are removed here, |
| 54 | before tokens reach any backend, rather than inside a single backend. |
| 55 | """ |
| 56 | for i, token_id in enumerate(token_ids): |
| 57 | if token_id < 0: |
| 58 | return token_ids[:i] |
| 59 | return token_ids |
| 60 | |
| 61 | |
| 62 | def compile_regex_with_timeout(fn: Callable[[str], _T], pattern: str) -> _T: |
| 63 | """Run a regex compilation callable with a timeout. |
| 64 | |
| 65 | Prevents ReDoS attacks where adversarial regex patterns (e.g. nested |
| 66 | quantifiers like ``(a+)+b``) cause exponential DFA state-space explosion, |
| 67 | hanging the inference worker indefinitely. |
| 68 | |
| 69 | Args: |
| 70 | fn: Single-argument callable that takes the pattern and performs |
| 71 | the regex compilation. |
| 72 | pattern: The regex pattern string, passed to *fn* and included in |
| 73 | timeout error messages. |
| 74 | |
| 75 | Raises: |
| 76 | ValueError: If compilation exceeds the configured timeout. |
| 77 | |
| 78 | """ |
| 79 | timeout = envs.VLLM_REGEX_COMPILATION_TIMEOUT_S |
| 80 | if timeout <= 0: |
| 81 | return fn(pattern) |
| 82 | |
| 83 | executor = ThreadPoolExecutor(max_workers=1) |
| 84 | future = executor.submit(fn, pattern) |
| 85 | try: |
| 86 | result = future.result(timeout=timeout) |
| 87 | except TimeoutError: |
| 88 | future.cancel() |
| 89 | executor.shutdown(wait=False, cancel_futures=True) |
| 90 | raise ValueError( |
| 91 | f"Regex compilation timed out after {timeout}s. " |
| 92 | "The pattern may be too complex or contain constructs that " |
| 93 | "cause exponential state-space explosion (e.g. nested " |
| 94 | f"quantifiers). Pattern: {pattern[:200]}" |
| 95 | ) from None |
| 96 | else: |
| 97 | executor.shutdown(wait=False) |
| 98 | return result |
| 99 | |
| 100 | |
| 101 | def apply_grammar_bitmask( |
| 102 | scheduler_output: SchedulerOutput, |
| 103 | grammar_output: GrammarOutput, |
| 104 | input_batch: InputBatch, |
| 105 | logits: torch.Tensor, |
| 106 | ) -> None: |
| 107 | """Apply grammar bitmask to output logits of the model with xgrammar function. |
| 108 | |
| 109 | Args: |
| 110 | scheduler_output (SchedulerOutput): The result of engine scheduling. |
| 111 | grammar_output (GrammarOutput): The grammar bitmask to apply. |
| 112 | input_batch (InputBatch): The input of model runner. |
| 113 | logits (torch.Tensor): The output logits of model forward. |
| 114 | |
| 115 | """ |
| 116 | # Serialization of np.ndarray is much more efficient than a tensor, |
| 117 | # so we receive it in that format. |
| 118 | grammar_bitmask = grammar_output.grammar_bitmask |
| 119 | |
| 120 | # We receive the structured output bitmask from the scheduler, |
| 121 | # compacted to contain bitmasks only for structured output requests. |
| 122 | # The order of the requests in the bitmask is not guaranteed to be the |
| 123 | # same as the order of the requests in the gpu runner's batch. We need |
| 124 | # to sort the bitmask to match the order of the requests used here. |
| 125 | |
| 126 | # Get the batch indices of the structured output requests. |
| 127 | # Keep track of the number of speculative tokens scheduled for every |
| 128 | # request in the batch, as the logit indices are offset by this amount. |
| 129 | struct_out_req_batch_indices: dict[str, int] = {} |
| 130 | cumulative_offset = 0 |
| 131 | spec_tokens = scheduler_output.scheduled_spec_decode_tokens |
| 132 | struct_out_req_ids = set(grammar_output.structured_output_request_ids) |
| 133 | for batch_index, req_id in enumerate(input_batch.req_ids): |
| 134 | logit_index = batch_index + cumulative_offset |
| 135 | cumulative_offset += len(spec_tokens.get(req_id, ())) |
| 136 | if req_id in struct_out_req_ids: |
| 137 | struct_out_req_batch_indices[req_id] = logit_index |
| 138 | |
| 139 | out_indices = [] |
| 140 | |
| 141 | # Reorder the bitmask to match the order of the requests in the batch. |
| 142 | sorted_bitmask_tensor = torch.full( |
| 143 | (logits.shape[0], grammar_bitmask.shape[1]), |
| 144 | -1, |
| 145 | dtype=torch.from_numpy(grammar_bitmask[:0]).dtype, |
| 146 | pin_memory=PIN_MEMORY, |
| 147 | ) |
| 148 | sorted_bitmask = sorted_bitmask_tensor.numpy() |
| 149 | cumulative_index = 0 |
| 150 | for req_id in grammar_output.structured_output_request_ids: |
| 151 | num_spec_tokens = len(spec_tokens.get(req_id, ())) |
| 152 | if (logit_idx := struct_out_req_batch_indices.get(req_id)) is not None: |
| 153 | for i in range(1 + num_spec_tokens): |
| 154 | bitmask_index = logit_idx + i |
| 155 | sorted_bitmask[bitmask_index] = grammar_bitmask[cumulative_index + i] |
| 156 | out_indices.append(bitmask_index) |
| 157 | cumulative_index += 1 + num_spec_tokens |
| 158 | |
| 159 | # Copy async to device. |
| 160 | grammar_bitmask = sorted_bitmask_tensor.to(logits.device, non_blocking=True) |
| 161 | |
| 162 | # If the length of out indices and the logits have the same shape |
| 163 | # we don't need to pass indices to the kernel, |
| 164 | # since the bitmask is already aligned with the logits. |
| 165 | skip_out_indices = len(out_indices) == logits.shape[0] |
| 166 | |
| 167 | if not logits.is_cpu: |
| 168 | index_tensor = None |
| 169 | if not skip_out_indices: |
| 170 | # xgrammar expects a python list of indices but it will actually work with |
| 171 | # a tensor. If we copy the tensor ourselves here we can do it in a |
| 172 | # non_blocking manner and there should be no cpu sync within xgrammar. |
| 173 | index_tensor = async_tensor_h2d( |
| 174 | out_indices, dtype=torch.int32, device=logits.device |
| 175 | ) |
| 176 | |
| 177 | xgr.apply_token_bitmask_inplace(logits, grammar_bitmask, indices=index_tensor) |
| 178 | return |
| 179 | |
| 180 | # CPU case, use list for indices. |
| 181 | indices = None if skip_out_indices else out_indices |
| 182 | # Handle dtype conversion for CPU (older xgrammar CPU kernels require float32) |
| 183 | # See: https://github.com/vllm-project/vllm/issues/31901 |
| 184 | if logits.dtype != torch.float32: |
| 185 | # Convert to float32, apply bitmask, then convert back |
| 186 | logits_fp32 = logits.to(torch.float32) |
| 187 | xgr.apply_token_bitmask_inplace(logits_fp32, grammar_bitmask, indices=indices) |
| 188 | # Copy the modified values back to the original tensor |
| 189 | logits.copy_(logits_fp32.to(logits.dtype)) |
| 190 | else: |
| 191 | xgr.apply_token_bitmask_inplace(logits, grammar_bitmask, indices=indices) |
| 192 | |
| 193 | |
| 194 | class OutlinesVocabulary: |
| 195 | """Wrapper class for `outlines_core.Vocabulary`, |
| 196 | which allows us to store a hash with the vocabulary |
| 197 | """ |
| 198 | |
| 199 | def __init__(self, vocabulary: oc.Vocabulary) -> None: |
| 200 | # Actual vocabulary object |
| 201 | self.inner = vocabulary |
| 202 | # Have to do abs(hash()) because python hashes can |
| 203 | # be negative, and we are using hash as a cache key. |
| 204 | hex_str = hashlib.sha256(vocabulary.__repr__().encode("utf-8")).hexdigest() |
| 205 | hash_int = int(hex_str, 16) |
| 206 | self._hash = hash_int |
| 207 | |
| 208 | |
| 209 | def get_outlines_cache_path() -> str: |
| 210 | """Get the context object that contains previously-computed return values.""" |
| 211 | outlines_cache_dir = os.getenv("OUTLINES_CACHE_DIR") |
| 212 | xdg_cache_home = os.getenv("XDG_CACHE_HOME") |
| 213 | home_dir = os.path.expanduser("~") |
| 214 | |
| 215 | if outlines_cache_dir: |
| 216 | # OUTLINES_CACHE_DIR takes precedence |
| 217 | return outlines_cache_dir |
| 218 | if xdg_cache_home: |
| 219 | return os.path.join(xdg_cache_home, ".cache", "outlines") |
| 220 | # If homedir is "/", we may be inside a container, and thus writing to |
| 221 | # root would be problematic, so we fall back to using a tempfile. |
| 222 | # Also validate the path exists, since os.path.expanduser does |
| 223 | # not guarantee existence. |
| 224 | if os.path.isdir(home_dir) and home_dir != "/": |
| 225 | # Default Unix fallback: ~/.cache/outlines |
| 226 | return os.path.join(home_dir, ".cache", "outlines") |
| 227 | |
| 228 | # home_dir may be / inside a docker container without existing user |
| 229 | tempdir = tempfile.gettempdir() |
| 230 | return os.path.join(tempdir, ".cache", "outlines") |
| 231 | |
| 232 | |
| 233 | class OutlinesDiskCache: |
| 234 | """SQLite-backed cache for outlines_core.Index objects. |
| 235 | |
| 236 | Uses outlines_core's native binary serialization (via Rust serde) |
| 237 | instead of pickle, eliminating arbitrary code execution risk on |
| 238 | deserialization. |
| 239 | """ |
| 240 | |
| 241 | _TYPE_INDEX = "I" |
| 242 | _TYPE_STRING = "S" |
| 243 | |
| 244 | def __init__(self, path: str): |
| 245 | os.makedirs(path, exist_ok=True) |
| 246 | db_path = os.path.join(path, "outlines_cache.db") |
| 247 | self._db = sqlite3.connect(db_path, check_same_thread=False) |
| 248 | self._db.execute("PRAGMA journal_mode=WAL") |
| 249 | self._db.execute( |
| 250 | "CREATE TABLE IF NOT EXISTS cache " |
| 251 | "(key TEXT PRIMARY KEY, type_tag TEXT NOT NULL, value BLOB NOT NULL)" |
| 252 | ) |
| 253 | self._db.commit() |
| 254 | |
| 255 | def __contains__(self, key: str) -> bool: |
| 256 | row = self._db.execute("SELECT 1 FROM cache WHERE key=?", (key,)).fetchone() |
| 257 | return row is not None |
| 258 | |
| 259 | def __getitem__(self, key: str): |
| 260 | row = self._db.execute( |
| 261 | "SELECT type_tag, value FROM cache WHERE key=?", (key,) |
| 262 | ).fetchone() |
| 263 | if row is None: |
| 264 | raise KeyError(key) |
| 265 | type_tag, data = row |
| 266 | if type_tag == self._TYPE_STRING: |
| 267 | return data.decode("utf-8") |
| 268 | return oc.Index.from_binary(data) |
| 269 | |
| 270 | def __setitem__(self, key: str, value): |
| 271 | if isinstance(value, str): |
| 272 | type_tag = self._TYPE_STRING |
| 273 | data = value.encode("utf-8") |
| 274 | else: |
| 275 | type_tag = self._TYPE_INDEX |
| 276 | data = value.__reduce__()[1][0] |
| 277 | self._db.execute( |
| 278 | "INSERT OR REPLACE INTO cache (key, type_tag, value) VALUES (?, ?, ?)", |
| 279 | (key, type_tag, data), |
| 280 | ) |
| 281 | self._db.commit() |
| 282 | |
| 283 | def get(self, key: str, default=None): |
| 284 | try: |
| 285 | return self[key] |
| 286 | except KeyError: |
| 287 | return default |
| 288 | |
| 289 | def set(self, key: str, value): |
| 290 | self[key] = value |
| 291 | |
| 292 | def clear(self): |
| 293 | self._db.execute("DELETE FROM cache") |
| 294 | self._db.commit() |
| 295 | |
| 296 | |
| 297 | def get_outlines_cache(): |
| 298 | """Get the Cache instance to be used for index caching.""" |
| 299 | cache_dir = get_outlines_cache_path() |
| 300 | if envs.VLLM_V1_USE_OUTLINES_CACHE: |
| 301 | logger.warning( |
| 302 | "Enabling outlines cache. This is an unbounded on-disk " |
| 303 | "cache. It may consume a lot of disk space and should " |
| 304 | "not be used with untrusted clients." |
| 305 | ) |
| 306 | cache = OutlinesDiskCache(cache_dir) |
| 307 | outlines_version = importlib.metadata.version("outlines_core") |
| 308 | |
| 309 | cached_version = cache.get("__version__", None) |
| 310 | if cached_version != outlines_version: |
| 311 | cache.clear() |
| 312 | cache.set("__version__", outlines_version) |
| 313 | return cache |
| 314 | |
| 315 | return LRUCache(maxsize=128) |
| 316 | |
| 317 | |
| 318 | re_llama_byte_token = re.compile(r"^<0x[0-9A-F]{2}>$") |
| 319 | re_replacement_seq = re.compile(r"^.{0,6}�+.{0,6}$") |
| 320 | |
| 321 | |
| 322 | def _reduced_vocabulary(tokenizer: TokenizerLike) -> dict[bytes, list[int]]: |
| 323 | """Create a map from vocabulary tokens to lists of equivalent token ids. |
| 324 | |
| 325 | Returns: |
| 326 | A Dict of token string -> equivalent token ids |
| 327 | |
| 328 | """ |
| 329 | eos_token_id = tokenizer.eos_token_id |
| 330 | |
| 331 | unicode_to_bytes = { |
| 332 | v: k for k, v in convert_slow_tokenizer.bytes_to_unicode().items() |
| 333 | } |
| 334 | |
| 335 | def convert_token_to_string(token: str) -> str: |
| 336 | string = tokenizer.convert_tokens_to_string([token]) |
| 337 | |
| 338 | # A hack to handle missing spaces to HF's Llama tokenizers |
| 339 | if ( |
| 340 | type(token) is str |
| 341 | and token.startswith(file_utils.SPIECE_UNDERLINE) |
| 342 | or token == "<0x20>" |
| 343 | ): |
| 344 | return " " + string |
| 345 | |
| 346 | return string |
| 347 | |
| 348 | vocabulary: dict[bytes, list[int]] = {} |
| 349 | empty_token_ids: list[int] = [] |
| 350 | for token, token_idx in tokenizer.get_vocab().items(): |
| 351 | if token in tokenizer.all_special_tokens: |
| 352 | continue |
| 353 | |
| 354 | token_str = convert_token_to_string(token) |
| 355 | if token_str: |
| 356 | if isinstance(token, (bytes, bytearray)): |
| 357 | # For BPE tokenizers where tokens are stored as bytes. |
| 358 | |
| 359 | # safe to ignore since token_str is of type (bytearray, bytes) |
| 360 | # by this point. |
| 361 | token_bytes = bytes(token_str) # type: ignore[arg-type] |
| 362 | |
| 363 | elif (token_str == "\ufffd" and token != "\ufffd") or ( |
| 364 | "\ufffd" in token_str and not re_replacement_seq.match(token_str) |
| 365 | ): |
| 366 | # Handle tokens with invalid UTF-8 sequences. |
| 367 | if re_llama_byte_token.match(token): |
| 368 | # Llama-like tokenizers use <0xXX> for incomplete sequences. |
| 369 | token_bytes = bytes([int(token[3:5], 16)]) |
| 370 | else: |
| 371 | # GPT2 tokenizers: map each byte back using unicode_to_bytes |
| 372 | byte_vals = [unicode_to_bytes.get(c) for c in token] |
| 373 | if None in byte_vals: |
| 374 | raise RuntimeError( |
| 375 | f"Cannot convert token `{token}`" |
| 376 | f" ({token_idx}) to bytes: {token_str}" |
| 377 | ) |
| 378 | # safe to ignore, since if None in byte_vals, |
| 379 | # an error is thrown. |
| 380 | token_bytes = bytes(byte_vals) # type: ignore[arg-type] |
| 381 | else: |
| 382 | token_bytes = token_str.encode("utf-8") |
| 383 | |
| 384 | if token_idx != eos_token_id: |
| 385 | vocabulary.setdefault(token_bytes, []).append(token_idx) |
| 386 | else: |
| 387 | empty_token_ids.append(token_idx) |
| 388 | |
| 389 | return vocabulary |
| 390 | |
| 391 | |
| 392 | def get_outlines_vocabulary(tokenizer: TokenizerLike) -> oc.Vocabulary: |
| 393 | """Get the `Vocabulary` object for a given tokenizer.""" |
| 394 | if hasattr(tokenizer, "_outlines_vocabulary"): |
| 395 | return tokenizer._outlines_vocabulary # type: ignore |
| 396 | |
| 397 | reduced_vocab = _reduced_vocabulary(tokenizer) |
| 398 | vocabulary = OutlinesVocabulary( |
| 399 | oc.Vocabulary(tokenizer.eos_token_id, reduced_vocab) |
| 400 | ) |
| 401 | tokenizer._outlines_vocabulary = vocabulary # type: ignore |
| 402 | |
| 403 | return vocabulary |
| 404 | |
| 405 | |
| 406 | def grammar_is_likely_lark(grammar_str: str) -> bool: |
| 407 | """Check if grammar appears to use Lark syntax. |
| 408 | |
| 409 | Args: |
| 410 | grammar_str: Input grammar string |
| 411 | |
| 412 | Returns: |
| 413 | bool: True if grammar appears to be in Lark format, False otherwise |
| 414 | |
| 415 | Examples: |
| 416 | >>> grammar_is_likely_lark("rule: 'abc'") |
| 417 | True |
| 418 | >>> grammar_is_likely_lark("rule ::= 'abc'") |
| 419 | False |
| 420 | |
| 421 | """ |
| 422 | if not grammar_str or not isinstance(grammar_str, str): |
| 423 | return False |
| 424 | |
| 425 | for line in grammar_str.split("\n"): |
| 426 | # Remove both comment styles |
| 427 | line = re.sub(r"(#|//).*$", "", line).strip() |
| 428 | if not line: |
| 429 | continue |
| 430 | |
| 431 | # Look for EBNF rule definition |
| 432 | if "::=" in line: |
| 433 | return False |
| 434 | |
| 435 | return True |
| 436 | |
| 437 | |
| 438 | def choice_as_grammar(choice: list[str]) -> str: |
| 439 | def escape_ebnf_string(s: str) -> str: |
| 440 | """Escape EBNF literals, including raw LF, CR, and NUL terminators.""" |
| 441 | escapes = {"\\": r"\\", '"': r"\"", "\n": r"\n", "\r": r"\r", "\t": r"\t"} |
| 442 | |
| 443 | def escape_char(ch: str) -> str: |
| 444 | if ch in escapes: |
| 445 | return escapes[ch] |
| 446 | # Escape remaining C0 controls (U+0000-U+001F) and DEL (U+007F). |
| 447 | if ord(ch) < 0x20 or ord(ch) == 0x7F: |
| 448 | return f"\\u{ord(ch):04x}" |
| 449 | return ch |
| 450 | |
| 451 | return "".join(escape_char(ch) for ch in s) |
| 452 | |
| 453 | escaped_choices = (escape_ebnf_string(c) for c in choice) |
| 454 | grammar = "root ::= " + " | ".join(f'"{c}"' for c in escaped_choices) |
| 455 | return grammar |