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
            elif "\ufffd" in token_str and not re_replacement_seq.match(token_str):
230
231
232
233
234
235
236
237
238
239
                # 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}`"
240
241
                            f" ({token_idx}) to bytes: {token_str}"
                        )
242
243
244
245
                    # safe to ignore, since if None in byte_vals,
                    # an error is thrown.
                    token_bytes = bytes(byte_vals)  # type: ignore[arg-type]
            else:
246
                token_bytes = token_str.encode("utf-8")
247
248
249
250
251
252
253
254
255

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

    return vocabulary


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

    try:
262
        if hasattr(tokenizer, "eos_token_id") and tokenizer.eos_token_id is not None:
263
264
265
            eos_token_id = tokenizer.eos_token_id
        else:
            raise ValueError(
266
267
268
                "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."
269
270
271
272
            )

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

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


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

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

        # Look for EBNF rule definition
313
        if "::=" in line:
314
315
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
            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."""
350
        return re.sub(r"(#|//).*$", "", line).strip()
351
352
353
354

    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:
355
            raise ValueError(f"Mismatched quotes in {rule_name} on line {line_num}")
356

357
    def extract_references(text: str) -> set[str]:
358
359
        """Extract rule references from text."""
        # Remove quoted strings and special characters
360
361
362
        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))
363
364

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

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

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

    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:
401
            if ":" in line and not line.startswith("|"):
402
403
404
                # Save previous rule if exists
                if current_rule:
                    output_lines.append(
405
406
                        f"{current_rule} ::= {' | '.join(current_definition)}"
                    )
407
408

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

                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()]

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

                alt_def = line[1:].strip()
425
426
427
                check_quotes(
                    alt_def, f"alternative for rule '{current_rule}'", line_num
                )
428
429
430
431
432
433
434
435
436
                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:
437
        output_lines.append(f"{current_rule} ::= {' | '.join(current_definition)}")
438
439

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

446
    return "\n".join(output_lines)
447
448
449
450
451
452


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
453
        return re.sub(r'(["\\])', r"\\\1", s)
454
455

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