backend_xgrammar.py 12.2 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

4
5
from __future__ import annotations

6
import json
7
from dataclasses import dataclass, field
8
from typing import TYPE_CHECKING, Any
9
10
11

import torch

12
import vllm.envs
13
from vllm.logger import init_logger
14
from vllm.sampling_params import SamplingParams
15
16
17
18
19
from vllm.transformers_utils.tokenizers.mistral import MistralTokenizer
from vllm.utils import LazyLoader
from vllm.v1.structured_output.backend_types import (StructuredOutputBackend,
                                                     StructuredOutputGrammar,
                                                     StructuredOutputOptions)
20
21
22
from vllm.v1.structured_output.utils import (choice_as_grammar,
                                             convert_lark_to_ebnf,
                                             grammar_is_likely_lark)
23
24
25
26
27
28
29
30
31

if TYPE_CHECKING:
    import xgrammar as xgr
else:
    xgr = LazyLoader("xgr", globals(), "xgrammar")

logger = init_logger(__name__)


32
@dataclass
33
34
class XgrammarBackend(StructuredOutputBackend):

35
    def __post_init__(self):
36
        self.disable_any_whitespace = \
37
            self.vllm_config.structured_outputs_config.disable_any_whitespace
38

39
        if isinstance(self.tokenizer, MistralTokenizer):
40
41
42
            # NOTE: ideally, xgrammar should handle this accordingly.
            # refer to https://github.com/mlc-ai/xgrammar/blob/d77c0a0173ef14779c918e3be7966ba852f7910f/python/xgrammar/tokenizer_info.py#L98
            try:
43
44
                if self.tokenizer.is_tekken:
                    encoded_vocab = self.tokenizer._vocab
45
46
47
                else:
                    encoded_vocab = [
                        token for token, _ in sorted(
48
                            self.tokenizer.get_vocab().items(),
49
50
51
                            key=lambda x: x[1],
                        )
                    ]
52
                stop_token_ids = None
53
54
                if (hasattr(
                        self.tokenizer,
55
                        "eos_token_id",
56
57
                ) and self.tokenizer.eos_token_id is not None):
                    stop_token_ids = [self.tokenizer.eos_token_id]
58
59
60
            except AttributeError as e:
                raise ValueError(
                    f"Cannot get the vocabulary of the tokenizer "
61
                    f"{type(self.tokenizer)}. The tokenizer should have a "
62
63
64
65
                    "get_vocab method.") from e
            tokenizer_info = xgr.TokenizerInfo(  # type: ignore
                encoded_vocab=encoded_vocab,
                # NOTE: https://github.com/mlc-ai/xgrammar/blob/5e141f6ff1ca02bc31f9e512e68b61f2a8ae88e5/tests/python/test_tokenizer_info.py#L43 # noqa: E501
66
                vocab_type=xgr.VocabType.RAW
67
                if self.tokenizer.is_tekken else xgr.VocabType.BYTE_FALLBACK,
68
69
70
71
72
73
                vocab_size=self.vocab_size,
                stop_token_ids=stop_token_ids,
                add_prefix_space=True,
            )
        else:
            tokenizer_info = xgr.TokenizerInfo.from_huggingface(
74
                self.tokenizer,
75
76
                vocab_size=self.vocab_size,
            )
77
78
79
80
81
82
        self.compiler = xgr.GrammarCompiler(
            tokenizer_info,
            max_threads=8,
            cache_enabled=True,
            cache_limit_bytes=vllm.envs.VLLM_XGRAMMAR_CACHE_MB * 1024 * 1024,
        )
83

84
85
86
87
88
        self.num_speculative_tokens = 0
        if self.vllm_config.speculative_config is not None:
            self.num_speculative_tokens = \
                self.vllm_config.speculative_config.num_speculative_tokens

89
90
91
    def compile_grammar(self, request_type: StructuredOutputOptions,
                        grammar_spec: str) -> StructuredOutputGrammar:
        if request_type == StructuredOutputOptions.JSON:
92
93
            ctx = self.compiler.compile_json_schema(
                grammar_spec, any_whitespace=not self.disable_any_whitespace)
94
        elif request_type == StructuredOutputOptions.JSON_OBJECT:
95
96
97
            ctx = self.compiler.compile_json_schema(
                '{"type": "object"}',
                any_whitespace=not self.disable_any_whitespace)
98
99
100
101
        elif request_type == StructuredOutputOptions.GRAMMAR:
            ctx = self.compiler.compile_grammar(grammar_spec)
        elif request_type == StructuredOutputOptions.REGEX:
            ctx = self.compiler.compile_regex(grammar_spec)
102
103
104
105
106
107
108
109
110
111
        elif request_type == StructuredOutputOptions.STRUCTURAL_TAG:
            s_tag = json.loads(grammar_spec)
            tags = [
                xgr.StructuralTagItem(
                    begin=s["begin"],
                    schema=json.dumps(s["schema"]),
                    end=s["end"],
                ) for s in s_tag["structures"]
            ]
            ctx = self.compiler.compile_structural_tag(tags, s_tag["triggers"])
112
113
114
115
116
117
118
119
        else:
            logger.error(
                "Validation should have already occurred. Please file an issue."
            )
            raise ValueError(
                f"grammar is not of valid supported types. ({request_type!s})")

        return XgrammarGrammar(
120
121
122
123
            matcher=xgr.GrammarMatcher(
                ctx,
                max_rollback_tokens=self.num_speculative_tokens,
            ),
124
125
126
127
128
129
130
            vocab_size=self.vocab_size,
            ctx=ctx,
        )

    def allocate_token_bitmask(self, max_num_seqs: int):
        return xgr.allocate_token_bitmask(max_num_seqs, self.vocab_size)

131
132
133
    def destroy(self):
        del self.compiler

134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150

@dataclass
class XgrammarGrammar(StructuredOutputGrammar):
    # NOTE: This would be a generic-enough class for
    # supporting different backends, in the future.
    # For now, just xgrammar.
    #
    # https://xgrammar.mlc.ai/docs/api/python/index.html#xgrammar.GrammarMatcher.find_jump_forward_string
    # for jump-forward decoding

    vocab_size: int
    matcher: xgr.GrammarMatcher = field(hash=False)
    ctx: xgr.CompiledGrammar = field(hash=False)
    num_processed_tokens: int = field(default_factory=lambda: 0,
                                      repr=False,
                                      hash=False,
                                      init=False)
151
    _is_terminated: bool = field(default=False, repr=False, hash=False)
152
153
154
155
156
157
158

    def accept_tokens(self, request_id: str, tokens: list[int]) -> bool:
        """Accepts a list of tokens and advances the FSM.

        Returns True if the FSM was advanced successfully.
        Returns False if the FSM failed to advance.
        """
159
160
        if self._is_terminated:
            return False
161
162
163
164
165
166
167
        for token in tokens:
            if not self.matcher.accept_token(token):
                logger.error(
                    "Failed to advance FSM for request %s "
                    "for tokens %s. Please file an issue.", request_id, token)
                return False
            self.num_processed_tokens += 1
168
        self._is_terminated = self.matcher.is_terminated()
169
170
        return True

171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
    def validate_tokens(self, tokens: list[int]) -> list[int]:
        """Checks if the list of tokens are accepted by the FSM in sequence.
        Will not advance the FSM.

        Returns the prefix list of tokens that are accepted by the FSM.
        """
        accepted_tokens = []
        for token in tokens:
            if self.matcher.accept_token(token):
                accepted_tokens.append(token)
            else:
                break
        if len(accepted_tokens) > 0:
            # Rollback the FSM to the initial state
            self.matcher.rollback(len(accepted_tokens))
        return accepted_tokens

    def rollback(self, num_tokens: int) -> None:
        self.matcher.rollback(num_tokens)
        self.num_processed_tokens -= num_tokens
191
        self._is_terminated = self.matcher.is_terminated()
192

193
194
195
196
    def fill_bitmask(self, bitmask: torch.Tensor, idx: int) -> None:
        self.matcher.fill_next_token_bitmask(bitmask, idx)

    def is_terminated(self) -> bool:
197
        return self._is_terminated
198
199
200
201

    def reset(self):
        self.num_processed_tokens = 0
        self.matcher.reset()
202
203
204
205
206
207
208
209
210
211


def has_xgrammar_unsupported_json_features(schema: dict[str, Any]) -> bool:
    """Check if JSON schema contains features unsupported by xgrammar."""

    def check_object(obj: dict[str, Any]) -> bool:
        if not isinstance(obj, dict):
            return False

        # Check for numeric ranges
212
        if obj.get("type") in ("integer", "number") and ("multipleOf" in obj):
213
214
215
216
            return True

        # Check for array unsupported keywords
        if obj.get("type") == "array" and any(
217
218
                key in obj for key in ("uniqueItems", "contains",
                                       "minContains", "maxContains")):
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
            return True

        # Unsupported keywords for strings
        if obj.get("type") == "string" and "format" in obj:
            return True

        # Unsupported keywords for objects
        if obj.get("type") == "object" and any(
                key in obj for key in ("minProperties", "maxProperties",
                                       "propertyNames", "patternProperties")):
            return True

        # Recursively check all nested objects and arrays
        for value in obj.values():
            if isinstance(value, dict):
                if check_object(value):
                    return True
            elif isinstance(value, list):
                for item in value:
                    if isinstance(item, dict) and check_object(item):
                        return True

        return False

    return check_object(schema)


def validate_xgrammar_grammar(sampling_params: SamplingParams) -> None:
    """Validate that the request is supported by structured output.

    Raises ValueError if the request is not supported.
    """
251
    if sampling_params.structured_outputs is None:
252
253
        return

254
    so_params = sampling_params.structured_outputs
255

256
    if so_params.regex:
257
        try:
258
            xgr.Grammar.from_regex(so_params.regex)
259
260
261
262
        except Exception as err:
            raise ValueError("Failed to transform regex into a grammar: "
                             f"{err}") from err

263
264
    if so_params.choice:
        choice_grammar = choice_as_grammar(so_params.choice)
265
266
267
268
269
        try:
            xgr.Grammar.from_ebnf(choice_grammar)
        except Exception as err:
            raise ValueError("Failed to transform choices into a grammar: "
                             "{err}") from err
270
271
        so_params.choice = None
        so_params.grammar = choice_grammar
272
273
        return

274
275
    if so_params.json:
        if isinstance(so_params.json, str):
276
            try:
277
                schema = json.loads(so_params.json)
278
279
280
            except json.JSONDecodeError as e:
                raise ValueError("Invalid JSON grammar specification.") from e
        else:
281
            schema = so_params.json
282

283
284
285
286
287
288
        try:
            xgr.Grammar.from_json_schema(schema)
        except Exception as err:
            raise ValueError("Failed to transform json schema into a grammar: "
                             f"{err}") from err

289
290
291
292
293
        if has_xgrammar_unsupported_json_features(schema):
            raise ValueError("The provided JSON schema contains features not "
                             "supported by xgrammar.")
        return

294
295
    if so_params.grammar:
        if grammar_is_likely_lark(so_params.grammar):
296
297
            # xgrammar supports EBNF grammars only
            try:
298
                so_params.grammar = convert_lark_to_ebnf(so_params.grammar)
299
300
301
302
303
304
305
            except ValueError as e:
                raise ValueError(
                    "Failed to convert the grammar from Lark to EBNF. ") from e

        # Test parsing EBNF grammar, possibly already converted from Lark
        try:
            # parse the grammar, but we aren't compiling it.
306
            xgr.Grammar.from_ebnf(so_params.grammar)
307
308
        except Exception as e:
            raise ValueError("Invalid grammar specification.") from e
309
310
        return

311
    if so_params.structural_tag:
312
        try:
313
            s_tag = json.loads(so_params.structural_tag)
314
315
316
317
318
319
320
321
322
323
            tags = [
                xgr.StructuralTagItem(
                    begin=s["begin"],
                    schema=json.dumps(s["schema"]),
                    end=s["end"],
                ) for s in s_tag["structures"]
            ]
            xgr.Grammar.from_structural_tag(tags, s_tag["triggers"])
        except Exception as e:
            raise ValueError("Invalid structural tag specification.") from e