utils.py 16.4 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
8
9
import hashlib
import importlib.metadata
import os
from typing import TYPE_CHECKING

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

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

if TYPE_CHECKING:
    import outlines_core as oc
    import transformers.file_utils as file_utils
24
    import xgrammar as xgr
25
    from transformers.convert_slow_tokenizer import bytes_to_unicode
26

27
    from vllm.tokenizers import TokenizerLike
28
    from vllm.v1.worker.gpu_input_batch import InputBatch
29
else:
30
    xgr = LazyLoader("xgr", globals(), "xgrammar")
31
32
    oc = LazyLoader("oc", globals(), "outlines_core")
    file_utils = LazyLoader("file_utils", globals(), "transformers.file_utils")
33
34
    bytes_to_unicode = LazyLoader(
        "bytes_to_unicode", globals(), "transformers.convert_slow_tokenizer"
35
36
    )

37
    TokenizerLike = object
38
39
40
    SchedulerOutput = object
    InputBatch = object

41
42
43
44
45
logger = init_logger(__name__)

CACHE = None


46
47
def apply_grammar_bitmask(
    scheduler_output: SchedulerOutput,
48
    grammar_output: GrammarOutput,
49
50
51
52
53
54
55
56
57
58
59
    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.
    """
60
61
62
    # Serialization of np.ndarray is much more efficient than a tensor,
    # so we receive it in that format.
    grammar_bitmask = grammar_output.grammar_bitmask
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78

    # 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
    seq = sorted(input_batch.req_id_to_index.items(), key=lambda x: x[1])
    for req_id, batch_index in seq:
        logit_index = batch_index + cumulative_offset
        cumulative_offset += len(
79
80
            scheduler_output.scheduled_spec_decode_tokens.get(req_id, [])
        )
81
        if req_id in grammar_output.structured_output_request_ids:
82
83
84
85
86
            struct_out_req_batch_indices[req_id] = logit_index

    out_indices = []

    # Reorder the bitmask to match the order of the requests in the batch.
87
88
89
90
91
    sorted_bitmask = np.full(
        shape=(logits.shape[0], grammar_bitmask.shape[1]),
        fill_value=-1,
        dtype=grammar_bitmask.dtype,
    )
92
    cumulative_index = 0
93
    for req_id in grammar_output.structured_output_request_ids:
94
        num_spec_tokens = len(
95
96
            scheduler_output.scheduled_spec_decode_tokens.get(req_id, [])
        )
97
98
99
        if req_id in struct_out_req_batch_indices:
            logit_index = struct_out_req_batch_indices[req_id]
            for i in range(1 + num_spec_tokens):
100
                sorted_bitmask[logit_index + i] = grammar_bitmask[cumulative_index + i]
101
                out_indices.append(logit_index + i)
102
        cumulative_index += 1 + num_spec_tokens
103
104
105
106
107

    # Copy async to device as tensor.
    grammar_bitmask = torch.from_numpy(sorted_bitmask).to(
        logits.device, non_blocking=True
    )
108
109
110
111
112
113

    # 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]

114
115
116
117
118
119
120
121
122
    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)
123

124
    xgr.apply_token_bitmask_inplace(logits, grammar_bitmask, indices=index_tensor)
125
126


127
128
129
130
131
132
133
134
135
136
137
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.
138
        hex_str = hashlib.sha256(vocabulary.__repr__().encode("utf-8")).hexdigest()
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
        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
    elif xdg_cache_home:
        return os.path.join(xdg_cache_home, ".cache", "outlines")
    # If homedir is "/", we may be inside a container, and thus writing to
155
    # root would be problematic, so we fall back to using a tempfile.
156
    # Also validate the path exists, since os.path.expanduser does
157
    # not guarantee existence.
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
    elif os.path.isdir(home_dir) and home_dir != "/":
        # Default Unix fallback: ~/.cache/outlines
        return os.path.join(home_dir, ".cache", "outlines")
    else:
        import tempfile

        # home_dir may be / inside a docker container without existing user
        tempdir = tempfile.gettempdir()
        return os.path.join(tempdir, ".cache", "outlines")


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:
174
175
176
177
178
        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."
        )
179
180
181
        cache = Cache(cache_dir, eviction_policy="none", cull_limit=0)
        outlines_version = importlib.metadata.version("outlines_core")

182
        cached_version = cache.get("__version__", None)
183
184
        if cached_version != outlines_version:
            cache.clear()
185
        cache.set("__version__", outlines_version)
186
187
188
189
190
191
192
193
194
195
        return cache
    else:
        return LRUCache(maxsize=128)


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(
196
    tokenizer: TokenizerLike,
197
198
199
200
201
202
203
204
    eos_token_id: int,
) -> 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
    """

205
    unicode_to_bytes = {v: k for k, v in bytes_to_unicode().items()}
206
207
208
209
210

    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
211
212
213
214
215
        if (
            type(token) is str
            and token.startswith(file_utils.SPIECE_UNDERLINE)
            or token == "<0x20>"
        ):
216
217
218
219
220
221
222
            return " " + string

        return string

    vocabulary: dict[bytes, list[int]] = {}
    empty_token_ids: list[int] = []
    for token, token_idx in tokenizer.get_vocab().items():
223
        if token in tokenizer.all_special_tokens:
224
225
226
227
228
229
230
231
232
233
234
            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]

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

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

    return vocabulary


262
def get_outlines_vocabulary(tokenizer: TokenizerLike) -> oc.Vocabulary:
263
    """Get the `Vocabulary` object for a given tokenizer."""
264
265
266
267
    if hasattr(tokenizer, "_outlines_vocabulary"):
        return tokenizer._outlines_vocabulary  # type: ignore

    try:
268
269
        if (
            hasattr(
270
271
                tokenizer,
                "eos_token_id",
272
273
274
            )
            and tokenizer.eos_token_id is not None
        ):
275
276
277
278
279
280
281
282
            eos_token_id = tokenizer.eos_token_id
        else:
            raise ValueError(
                f"Error during structured outputs setup for outlines: Tokenizer ({type(tokenizer)}) has no `eos_token_id` property, but `eos_token_id` is required for structured outputs to work properly."  # noqa: E501
            )

        reduced_vocab = _reduced_vocabulary(
            tokenizer,
283
            eos_token_id,  # type: ignore
284
        )
285
        vocabulary = OutlinesVocabulary(oc.Vocabulary(eos_token_id, reduced_vocab))
286
287
288
289
        tokenizer._outlines_vocabulary = vocabulary  # type: ignore

        return vocabulary
    except AttributeError as e:
290
291
292
293
294
        raise ValueError(
            f"Cannot get the vocabulary of the tokenizer "
            f"({type(tokenizer)}). The tokenizer should have a "
            "get_vocab method."
        ) from e
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315


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

316
    for line in grammar_str.split("\n"):
317
        # Remove both comment styles
318
        line = re.sub(r"(#|//).*$", "", line).strip()
319
320
321
322
        if not line:
            continue

        # Look for EBNF rule definition
323
        if "::=" in line:
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
352
353
354
355
356
357
358
359
            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."""
360
        return re.sub(r"(#|//).*$", "", line).strip()
361
362
363
364

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

367
    def extract_references(text: str) -> set[str]:
368
369
        """Extract rule references from text."""
        # Remove quoted strings and special characters
370
371
372
        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))
373
374

    # First pass: Find root rule and validate rule definitions
375
    lines = [clean_line(line) for line in grammar_str.split("\n")]
376
377
378
    first_rule = None

    for line_num, line in enumerate(lines, 1):
379
        if not line or line.startswith("|"):
380
381
            continue

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

    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:
411
            if ":" in line and not line.startswith("|"):
412
413
414
                # Save previous rule if exists
                if current_rule:
                    output_lines.append(
415
416
                        f"{current_rule} ::= {' | '.join(current_definition)}"
                    )
417
418

                # Process new rule
419
420
                name, definition = line.split(":", 1)
                current_rule = name.strip().strip("?")
421
422
423
424
425
426

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

427
            elif line.startswith("|"):
428
                if not current_rule:
429
430
431
432
                    raise ValueError(
                        f"Alternative '|' on line {line_num} "
                        "without a preceding rule definition"
                    )
433
434

                alt_def = line[1:].strip()
435
436
437
                check_quotes(
                    alt_def, f"alternative for rule '{current_rule}'", line_num
                )
438
439
440
441
442
443
444
445
446
                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:
447
        output_lines.append(f"{current_rule} ::= {' | '.join(current_definition)}")
448
449

    # Validate all rules are defined
450
    undefined_rules = referenced_rules - defined_rules - {"root"}
451
    if undefined_rules:
452
453
454
        raise ValueError(
            f"Referenced rules are not defined: {', '.join(sorted(undefined_rules))}"
        )
455

456
    return "\n".join(output_lines)
457
458
459
460
461
462


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
463
        return re.sub(r'(["\\])', r"\\\1", s)
464
465

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