utils.py 16.3 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from __future__ import annotations

5
6
7
import hashlib
import importlib.metadata
import os
8
import tempfile
9
10
from typing import TYPE_CHECKING

11
import numpy as np
12
import regex as re
13
import torch
14
15
16
17
18
from cachetools import LRUCache
from diskcache import Cache

import vllm.envs as envs
from vllm.logger import init_logger
19
from vllm.utils.import_utils import LazyLoader
20
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
21
22
23

if TYPE_CHECKING:
    import outlines_core as oc
Harry Mellor's avatar
Harry Mellor committed
24
    import transformers.convert_slow_tokenizer as convert_slow_tokenizer
25
    import transformers.file_utils as file_utils
26
    import xgrammar as xgr
27

28
    from vllm.tokenizers import TokenizerLike
29
    from vllm.v1.worker.gpu_input_batch import InputBatch
30
else:
31
    xgr = LazyLoader("xgr", globals(), "xgrammar")
32
33
    oc = LazyLoader("oc", globals(), "outlines_core")
    file_utils = LazyLoader("file_utils", globals(), "transformers.file_utils")
Harry Mellor's avatar
Harry Mellor committed
34
35
    convert_slow_tokenizer = LazyLoader(
        "convert_slow_tokenizer", globals(), "transformers.convert_slow_tokenizer"
36
37
    )

38

39
40
41
42
43
logger = init_logger(__name__)

CACHE = None


44
45
def apply_grammar_bitmask(
    scheduler_output: SchedulerOutput,
46
    grammar_output: GrammarOutput,
47
48
49
50
51
52
53
54
55
56
57
    input_batch: InputBatch,
    logits: torch.Tensor,
) -> None:
    """
    Apply grammar bitmask to output logits of the model with xgrammar function.

    Args:
        scheduler_output (SchedulerOutput): The result of engine scheduling.
        input_batch (InputBatch): The input of model runner.
        logits (torch.Tensor): The output logits of model forward.
    """
58
59
60
    # Serialization of np.ndarray is much more efficient than a tensor,
    # so we receive it in that format.
    grammar_bitmask = grammar_output.grammar_bitmask
61
62
63
64
65
66
67
68
69
70
71
72

    # We receive the structured output bitmask from the scheduler,
    # compacted to contain bitmasks only for structured output requests.
    # The order of the requests in the bitmask is not guaranteed to be the
    # same as the order of the requests in the gpu runner's batch. We need
    # to sort the bitmask to match the order of the requests used here.

    # Get the batch indices of the structured output requests.
    # Keep track of the number of speculative tokens scheduled for every
    # request in the batch, as the logit indices are offset by this amount.
    struct_out_req_batch_indices: dict[str, int] = {}
    cumulative_offset = 0
73
74
75
    spec_tokens = scheduler_output.scheduled_spec_decode_tokens
    struct_out_req_ids = set(grammar_output.structured_output_request_ids)
    for batch_index, req_id in enumerate(input_batch.req_ids):
76
        logit_index = batch_index + cumulative_offset
77
78
        cumulative_offset += len(spec_tokens.get(req_id, ()))
        if req_id in struct_out_req_ids:
79
80
81
82
83
            struct_out_req_batch_indices[req_id] = logit_index

    out_indices = []

    # Reorder the bitmask to match the order of the requests in the batch.
84
85
86
87
88
    sorted_bitmask = np.full(
        shape=(logits.shape[0], grammar_bitmask.shape[1]),
        fill_value=-1,
        dtype=grammar_bitmask.dtype,
    )
89
    cumulative_index = 0
90
    for req_id in grammar_output.structured_output_request_ids:
91
92
        num_spec_tokens = len(spec_tokens.get(req_id, ()))
        if (logit_idx := struct_out_req_batch_indices.get(req_id)) is not None:
93
            for i in range(1 + num_spec_tokens):
94
95
96
                bitmask_index = logit_idx + i
                sorted_bitmask[bitmask_index] = grammar_bitmask[cumulative_index + i]
                out_indices.append(bitmask_index)
97
        cumulative_index += 1 + num_spec_tokens
98
99
100
101
102

    # Copy async to device as tensor.
    grammar_bitmask = torch.from_numpy(sorted_bitmask).to(
        logits.device, non_blocking=True
    )
103
104
105
106
107
108

    # If the length of out indices and the logits have the same shape
    # we don't need to pass indices to the kernel,
    # since the bitmask is already aligned with the logits.
    skip_out_indices = len(out_indices) == logits.shape[0]

109
110
111
112
113
114
115
116
117
    index_tensor = None
    if not skip_out_indices:
        # xgrammar expects a python list of indices but it will actually work with
        # a tensor. If we copy the tensor ourselves here we can do it in a non_blocking
        # manner and there should be no cpu sync within xgrammar.
        index_tensor = torch.tensor(
            out_indices, dtype=torch.int32, device="cpu", pin_memory=True
        )
        index_tensor = index_tensor.to(logits.device, non_blocking=True)
118

119
    xgr.apply_token_bitmask_inplace(logits, grammar_bitmask, indices=index_tensor)
120
121


122
123
124
125
126
127
128
129
130
131
132
class OutlinesVocabulary:
    """
    Wrapper class for `outlines_core.Vocabulary`,
    which allows us to store a hash with the vocabulary
    """

    def __init__(self, vocabulary: oc.Vocabulary) -> None:
        # Actual vocabulary object
        self.inner = vocabulary
        # Have to do abs(hash()) because python hashes can
        # be negative, and we are using hash as a cache key.
133
        hex_str = hashlib.sha256(vocabulary.__repr__().encode("utf-8")).hexdigest()
134
135
136
137
138
139
140
141
142
143
144
145
146
        hash_int = int(hex_str, 16)
        self._hash = hash_int


def get_outlines_cache_path() -> str:
    """Get the context object that contains previously-computed return values"""
    outlines_cache_dir = os.getenv("OUTLINES_CACHE_DIR")
    xdg_cache_home = os.getenv("XDG_CACHE_HOME")
    home_dir = os.path.expanduser("~")

    if outlines_cache_dir:
        # OUTLINES_CACHE_DIR takes precedence
        return outlines_cache_dir
147
    if xdg_cache_home:
148
149
        return os.path.join(xdg_cache_home, ".cache", "outlines")
    # If homedir is "/", we may be inside a container, and thus writing to
150
    # root would be problematic, so we fall back to using a tempfile.
151
    # Also validate the path exists, since os.path.expanduser does
152
    # not guarantee existence.
153
    if os.path.isdir(home_dir) and home_dir != "/":
154
155
156
        # Default Unix fallback: ~/.cache/outlines
        return os.path.join(home_dir, ".cache", "outlines")

157
158
159
    # home_dir may be / inside a docker container without existing user
    tempdir = tempfile.gettempdir()
    return os.path.join(tempdir, ".cache", "outlines")
160
161
162
163
164
165
166


def get_outlines_cache():
    """Get the Cache instance to be used for index caching"""

    cache_dir = get_outlines_cache_path()
    if envs.VLLM_V1_USE_OUTLINES_CACHE:
167
168
169
170
171
        logger.warning(
            "Enabling outlines cache. This is an unbounded on-disk "
            "cache. It may consume a lot of disk space and should "
            "not be used with untrusted clients."
        )
172
173
174
        cache = Cache(cache_dir, eviction_policy="none", cull_limit=0)
        outlines_version = importlib.metadata.version("outlines_core")

175
        cached_version = cache.get("__version__", None)
176
177
        if cached_version != outlines_version:
            cache.clear()
178
        cache.set("__version__", outlines_version)
179
        return cache
180
181

    return LRUCache(maxsize=128)
182
183
184
185
186
187
188


re_llama_byte_token = re.compile(r"^<0x[0-9A-F]{2}>$")
re_replacement_seq = re.compile(r"^.{0,6}�+.{0,6}$")


def _reduced_vocabulary(
189
    tokenizer: TokenizerLike, eos_token_id: int
190
191
192
193
194
195
196
) -> dict[bytes, list[int]]:
    """Create a map from vocabulary tokens to lists of equivalent token ids.

    Returns:
        A Dict of token string -> equivalent token ids
    """

Harry Mellor's avatar
Harry Mellor committed
197
198
199
    unicode_to_bytes = {
        v: k for k, v in convert_slow_tokenizer.bytes_to_unicode().items()
    }
200
201
202
203
204

    def convert_token_to_string(token: str) -> str:
        string = tokenizer.convert_tokens_to_string([token])

        # A hack to handle missing spaces to HF's Llama tokenizers
205
206
207
208
209
        if (
            type(token) is str
            and token.startswith(file_utils.SPIECE_UNDERLINE)
            or token == "<0x20>"
        ):
210
211
212
213
214
215
216
            return " " + string

        return string

    vocabulary: dict[bytes, list[int]] = {}
    empty_token_ids: list[int] = []
    for token, token_idx in tokenizer.get_vocab().items():
217
        if token in tokenizer.all_special_tokens:
218
219
220
221
222
223
224
225
226
227
228
            continue

        token_str = convert_token_to_string(token)
        if token_str:
            if isinstance(token, (bytes, bytearray)):
                # For BPE tokenizers where tokens are stored as bytes.

                # safe to ignore since token_str is of type (bytearray, bytes)
                # by this point.
                token_bytes = bytes(token_str)  # type: ignore[arg-type]

229
230
231
            elif (token_str == "\ufffd" and token != "\ufffd") or (
                "\ufffd" in token_str and not re_replacement_seq.match(token_str)
            ):
232
233
234
235
236
237
238
239
240
241
                # Handle tokens with invalid UTF-8 sequences.
                if re_llama_byte_token.match(token):
                    # Llama-like tokenizers use <0xXX> for incomplete sequences.
                    token_bytes = bytes([int(token[3:5], 16)])
                else:
                    # GPT2 tokenizers: map each byte back using unicode_to_bytes
                    byte_vals = [unicode_to_bytes.get(c) for c in token]
                    if None in byte_vals:
                        raise RuntimeError(
                            f"Cannot convert token `{token}`"
242
243
                            f" ({token_idx}) to bytes: {token_str}"
                        )
244
245
246
247
                    # safe to ignore, since if None in byte_vals,
                    # an error is thrown.
                    token_bytes = bytes(byte_vals)  # type: ignore[arg-type]
            else:
248
                token_bytes = token_str.encode("utf-8")
249
250
251
252
253
254
255
256
257

            if token_idx != eos_token_id:
                vocabulary.setdefault(token_bytes, []).append(token_idx)
        else:
            empty_token_ids.append(token_idx)

    return vocabulary


258
def get_outlines_vocabulary(tokenizer: TokenizerLike) -> oc.Vocabulary:
259
    """Get the `Vocabulary` object for a given tokenizer."""
260
261
262
263
    if hasattr(tokenizer, "_outlines_vocabulary"):
        return tokenizer._outlines_vocabulary  # type: ignore

    try:
264
        if hasattr(tokenizer, "eos_token_id") and tokenizer.eos_token_id is not None:
265
266
267
            eos_token_id = tokenizer.eos_token_id
        else:
            raise ValueError(
268
269
270
                "Error during structured outputs setup for outlines: Tokenizer "
                f"({type(tokenizer)}) has no `eos_token_id` property, but "
                "`eos_token_id` is required for structured outputs to work properly."
271
272
273
274
            )

        reduced_vocab = _reduced_vocabulary(
            tokenizer,
275
            eos_token_id,  # type: ignore
276
        )
277
        vocabulary = OutlinesVocabulary(oc.Vocabulary(eos_token_id, reduced_vocab))
278
279
280
281
        tokenizer._outlines_vocabulary = vocabulary  # type: ignore

        return vocabulary
    except AttributeError as e:
282
        raise ValueError(
283
            "Cannot get the vocabulary of the tokenizer "
284
285
286
            f"({type(tokenizer)}). The tokenizer should have a "
            "get_vocab method."
        ) from e
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307


def grammar_is_likely_lark(grammar_str: str) -> bool:
    """
    Check if grammar appears to use Lark syntax.

    Args:
        grammar_str: Input grammar string

    Returns:
        bool: True if grammar appears to be in Lark format, False otherwise

    Examples:
        >>> grammar_is_likely_lark("rule: 'abc'")
        True
        >>> grammar_is_likely_lark("rule ::= 'abc'")
        False
    """
    if not grammar_str or not isinstance(grammar_str, str):
        return False

308
    for line in grammar_str.split("\n"):
309
        # Remove both comment styles
310
        line = re.sub(r"(#|//).*$", "", line).strip()
311
312
313
314
        if not line:
            continue

        # Look for EBNF rule definition
315
        if "::=" in line:
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
            return False

    return True


def convert_lark_to_ebnf(grammar_str: str) -> str:
    """
    Convert a Lark grammar string to EBNF format.

    EBNF reference:
    https://github.com/ggerganov/llama.cpp/blob/master/grammars/README.md
    Lark grammar reference:
    https://lark-parser.readthedocs.io/en/latest/grammar.html

    Args:
        grammar_str: Input grammar in Lark format

    Returns:
        str: Converted grammar in EBNF format

    Examples:
        >>> print(convert_lark_to_ebnf("rule: 'hello'"))
        root ::= rule
        rule ::= "hello"
    """
    if not isinstance(grammar_str, str):
        raise ValueError(f"Grammar must be a string, got {type(grammar_str)}")
    if not grammar_str.strip():
        raise ValueError("Grammar string cannot be empty")

    defined_rules = set()
    referenced_rules = set()
    output_lines = []

    def clean_line(line: str) -> str:
        """Remove comments and whitespace from line."""
352
        return re.sub(r"(#|//).*$", "", line).strip()
353
354
355
356

    def check_quotes(text: str, rule_name: str, line_num: int) -> None:
        """Validate quote matching in text."""
        if text.count("'") % 2 != 0 or text.count('"') % 2 != 0:
357
            raise ValueError(f"Mismatched quotes in {rule_name} on line {line_num}")
358

359
    def extract_references(text: str) -> set[str]:
360
361
        """Extract rule references from text."""
        # Remove quoted strings and special characters
362
363
364
        text = re.sub(r'"[^"]*"', "", text)
        text = re.sub(r"[+*?()|\[\]{}]", " ", text)
        return set(re.findall(r"\b[a-zA-Z_][a-zA-Z0-9_]*\b", text))
365
366

    # First pass: Find root rule and validate rule definitions
367
    lines = [clean_line(line) for line in grammar_str.split("\n")]
368
369
370
    first_rule = None

    for line_num, line in enumerate(lines, 1):
371
        if not line or line.startswith("|"):
372
373
            continue

374
        if ":" in line:
375
            try:
376
                name = line.split(":", 1)[0].strip().strip("?")
377
378
379
                defined_rules.add(name)
                if first_rule is None:
                    first_rule = name
380
381
                if name == "start":
                    first_rule = "start"
382
            except IndexError as e:
383
384
385
386
                raise ValueError(
                    f"Invalid rule format on line {line_num}. "
                    "Expected 'rule_name: definition'"
                ) from e
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402

    if not defined_rules:
        raise ValueError("No valid rules found in grammar")

    # Add root rule
    output_lines.append(f"root ::= {first_rule}")

    # Second pass: Process rule definitions and alternatives
    current_rule = None
    current_definition = []

    for line_num, line in enumerate(lines, 1):
        if not line:
            continue

        try:
403
            if ":" in line and not line.startswith("|"):
404
405
406
                # Save previous rule if exists
                if current_rule:
                    output_lines.append(
407
408
                        f"{current_rule} ::= {' | '.join(current_definition)}"
                    )
409
410

                # Process new rule
411
412
                name, definition = line.split(":", 1)
                current_rule = name.strip().strip("?")
413
414
415
416
417
418

                check_quotes(definition, f"rule '{current_rule}'", line_num)
                definition = re.sub(r"'([^']*)'", r'"\1"', definition)
                referenced_rules.update(extract_references(definition))
                current_definition = [definition.strip()]

419
            elif line.startswith("|"):
420
                if not current_rule:
421
422
423
424
                    raise ValueError(
                        f"Alternative '|' on line {line_num} "
                        "without a preceding rule definition"
                    )
425
426

                alt_def = line[1:].strip()
427
428
429
                check_quotes(
                    alt_def, f"alternative for rule '{current_rule}'", line_num
                )
430
431
432
433
434
435
436
437
438
                alt_def = re.sub(r"'([^']*)'", r'"\1"', alt_def)
                referenced_rules.update(extract_references(alt_def))
                current_definition.append(alt_def)

        except ValueError as e:
            raise ValueError(f"Error on line {line_num}: {str(e)}") from e

    # Add final rule if exists
    if current_rule:
439
        output_lines.append(f"{current_rule} ::= {' | '.join(current_definition)}")
440
441

    # Validate all rules are defined
442
    undefined_rules = referenced_rules - defined_rules - {"root"}
443
    if undefined_rules:
444
445
446
        raise ValueError(
            f"Referenced rules are not defined: {', '.join(sorted(undefined_rules))}"
        )
447

448
    return "\n".join(output_lines)
449
450
451
452
453
454


def choice_as_grammar(choice: list[str]) -> str:
    def escape_ebnf_string(s: str) -> str:
        """Escape special characters in a EBNF string."""
        # Escape double quotes and backslashes
455
        return re.sub(r'(["\\])', r"\\\1", s)
456
457

    escaped_choices = (escape_ebnf_string(c) for c in choice)
458
    grammar = "root ::= " + " | ".join(f'"{c}"' for c in escaped_choices)
459
    return grammar