A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 18:30 UTC
[V1] Refactor Structured Output for multiple backends (#14694)
Signed-off-by: Russell Bryant <rbryant@redhat.com>
Russell Bryant committed
Mar 18, 2025 at 15:49 UTC
3a1e6481586ed7f079275b5d5072a6e246af691e
6 files changed
+290
-185
vllm/v1/engine/processor.py
+15
-10
@@ -119,16 +119,21 @@ class Processor:
119
def _validate_structured_output(self, params: SamplingParams) -> None:
120
if not params.guided_decoding or not self.decoding_config:
121
return
122
- if self.decoding_config.guided_decoding_backend != "xgrammar":
123
- raise ValueError(
124
- "Only xgrammar structured output is supported in V1.")
125
- if (params.guided_decoding.backend
126
- and params.guided_decoding.backend != 'xgrammar'):
127
- raise ValueError(
128
- "Only xgrammar structured output is supported in V1.")
129
- if self.vllm_config.speculative_config:
130
- raise ValueError("Structured output is not supported with "
131
- "speculative decoding.")
122
+
123
+ supported_backends = ["xgrammar"]
124
+ engine_level_backend = self.decoding_config.guided_decoding_backend
125
+ if engine_level_backend not in supported_backends:
126
+ raise ValueError(f"Only {supported_backends} structured output is "
127
+ "supported in V1.")
128
+ if params.guided_decoding.backend:
129
+ if params.guided_decoding.backend != engine_level_backend:
130
+ raise ValueError("Request-level structured output backend "
131
+ "must match engine-level backend. "
132
+ f"{params.guided_decoding.backend}"
133
+ f" != {engine_level_backend}")
134
+ else:
135
+ params.guided_decoding.backend = engine_level_backend
136
+
137
if vllm.platforms.current_platform.is_tpu():
138
raise ValueError("Structured output is not supported on TPU.")
139
vllm/v1/structured_output/__init__.py
+31
-92
@@ -7,75 +7,27 @@ from typing import TYPE_CHECKING, Optional
7
8
from vllm.config import VllmConfig
9
from vllm.logger import init_logger
10
-from vllm.transformers_utils.tokenizer_group import init_tokenizer_from_configs
11
-from vllm.transformers_utils.tokenizers.mistral import MistralTokenizer
12
-from vllm.utils import LazyLoader
13
-from vllm.v1.structured_output.grammar import Grammar, StructuredOutputOptions
10
+from vllm.v1.structured_output.backend_types import (StructuredOutputBackend,
11
+ StructuredOutputGrammar)
12
+from vllm.v1.structured_output.backend_xgrammar import XgrammarBackend
13
14
if TYPE_CHECKING:
15
import numpy as np
16
import numpy.typing as npt
18
- import xgrammar as xgr
17
+ import torch
18
19
from vllm.v1.request import Request
21
-else:
22
- xgr = LazyLoader("xgr", globals(), "xgrammar")
20
21
logger = init_logger(__name__)
22
23
24
class StructuredOutputManager:
25
+ """Engine-level manager for structured output requests."""
26
27
def __init__(self, vllm_config: VllmConfig):
28
+ self.backend: Optional[StructuredOutputBackend] = None
29
self.vllm_config = vllm_config
31
- self.init_complete = False
32
-
33
- def _delayed_init(self):
34
- """Initialization delayed until we know it is needed."""
35
- tokenizer_group = init_tokenizer_from_configs(
36
- model_config=self.vllm_config.model_config,
37
- scheduler_config=self.vllm_config.scheduler_config,
38
- parallel_config=self.vllm_config.parallel_config,
39
- lora_config=self.vllm_config.lora_config) # type: ignore[arg-type]
40
- tokenizer_group.ping()
41
-
42
- tokenizer = tokenizer_group.get_lora_tokenizer(None)
43
- self.vocab_size = self.vllm_config.model_config.get_vocab_size()
44
- if isinstance(tokenizer, MistralTokenizer):
45
- # NOTE: ideally, xgrammar should handle this accordingly.
46
- # refer to https://github.com/mlc-ai/xgrammar/blob/d77c0a0173ef14779c918e3be7966ba852f7910f/python/xgrammar/tokenizer_info.py#L98
47
- try:
48
- encoded_vocab = [
49
- token for token, _ in sorted(
50
- tokenizer.get_vocab().items(),
51
- key=lambda x: x[1],
52
- )
53
- ]
54
- stop_token_ids = None
55
- if hasattr(
56
- tokenizer,
57
- "eos_token_id",
58
- ) and tokenizer.eos_token_id is not None:
59
- stop_token_ids = [tokenizer.eos_token_id]
60
- except AttributeError as e:
61
- raise ValueError(
62
- f"Cannot get the vocabulary of the tokenizer "
63
- f"{type(tokenizer)}. The tokenizer should have a "
64
- "get_vocab method.") from e
65
- tokenizer_info = xgr.TokenizerInfo(
66
- encoded_vocab=encoded_vocab,
67
- # NOTE: https://github.com/mlc-ai/xgrammar/blob/5e141f6ff1ca02bc31f9e512e68b61f2a8ae88e5/tests/python/test_tokenizer_info.py#L43 # noqa: E501
68
- vocab_type=xgr.VocabType.BYTE_FALLBACK,
69
- vocab_size=self.vocab_size,
70
- stop_token_ids=stop_token_ids,
71
- add_prefix_space=True,
72
- )
73
- else:
74
- tokenizer_info = xgr.TokenizerInfo.from_huggingface(
75
- tokenizer,
76
- vocab_size=self.vocab_size,
77
- )
78
- self.compiler = xgr.GrammarCompiler(tokenizer_info, max_threads=8)
30
+ self._grammar_bitmask: Optional[torch.Tensor] = None
31
32
# The default max_workers if not specified is the number of CPUs * 5,
33
# which is way too high since these tasks are CPU-bound, not I/O bound.
@@ -83,28 +35,30 @@ class StructuredOutputManager:
35
# compilation, so we set it to half the number of CPUs.
36
max_workers = max(1, (multiprocessing.cpu_count() + 1) // 2)
37
self.executor = ThreadPoolExecutor(max_workers=max_workers)
86
- self._grammar_bitmask = xgr.allocate_token_bitmask(
87
- self.vllm_config.scheduler_config.max_num_seqs,
88
- self.vocab_size,
89
- )
90
-
91
- self.init_complete = True
38
39
def grammar_init(self, request: Request) -> None:
40
if request.structured_output_request is None:
41
return
42
97
- # The first time this is called, we need to finish initialization
98
- # of xgrammar. We defer it to avoid the import of xgrammar and
99
- # initialization cost if it is not going to be used.
100
- if not self.init_complete:
101
- self._delayed_init()
43
+ # Initialize the backend the first time it is needed.
44
+ #
45
+ # NOTE: We only support a single backend. We do NOT support different
46
+ # backends on a per-request basis in V1 (for now, anyway...).
47
+ if self.backend is None:
48
+ backend_name = request.sampling_params.guided_decoding.backend_name
49
+ if backend_name == "xgrammar":
50
+ self.backend = XgrammarBackend(self.vllm_config)
51
+ else:
52
+ raise ValueError(
53
+ f"Unsupported structured output backend: {backend_name}")
54
103
- grammar: Future[Grammar] = self.executor.submit(
104
- self._async_create_grammar, request)
55
+ grammar: Future[StructuredOutputGrammar] = self.executor.submit(
56
+ self._async_create_grammar, request, self.backend)
57
request.structured_output_request.grammar = grammar # type: ignore[assignment]
58
107
- def _async_create_grammar(self, request: Request) -> Grammar:
59
+ def _async_create_grammar(
60
+ self, request: Request,
61
+ backend: StructuredOutputBackend) -> StructuredOutputGrammar:
62
key = request.structured_output_request.structured_output_key # type: ignore[union-attr]
63
64
# Note that the request was validated in the engine core client,
@@ -114,28 +68,8 @@ class StructuredOutputManager:
68
# though it should be unlikely as we test that up front as well.
69
request_type, grammar_spec = key
70
117
- if request_type == StructuredOutputOptions.JSON:
118
- # TODO -- allow any_whitespace to be configurable
119
- # pending merge of https://github.com/vllm-project/vllm/pull/12744
120
- ctx = self.compiler.compile_json_schema(grammar_spec,
121
- any_whitespace=False)
122
- elif request_type == StructuredOutputOptions.JSON_OBJECT:
123
- ctx = self.compiler.compile_builtin_json_grammar()
124
- elif request_type == StructuredOutputOptions.GRAMMAR:
125
- ctx = self.compiler.compile_grammar(grammar_spec)
126
- elif request_type == StructuredOutputOptions.REGEX:
127
- ctx = self.compiler.compile_regex(grammar_spec)
128
- else:
129
- logger.error("Validation should have already occurred. "
130
- "Please file an issue.")
131
- raise ValueError(
132
- f"grammar is not of valid supported types. ({request_type!s})")
133
-
134
- return Grammar(
135
- matcher=xgr.GrammarMatcher(ctx),
136
- vocab_size=self.vocab_size,
137
- ctx=ctx,
138
- )
71
+ assert self.backend is not None
72
+ return self.backend.compile_grammar(request_type, grammar_spec)
73
74
def grammar_bitmask(
75
self,
@@ -147,6 +81,11 @@ class StructuredOutputManager:
81
if not structured_output_request_ids:
82
return None
83
84
+ if self._grammar_bitmask is None:
85
+ assert self.backend is not None
86
+ self._grammar_bitmask = self.backend.allocate_token_bitmask(
87
+ self.vllm_config.scheduler_config.max_num_seqs)
88
+
89
# Fill the bitmask using the index of each request equal to its
90
# position in the batch. Resize the bitmask down to the size of
91
# the batch.
@@ -154,7 +93,7 @@ class StructuredOutputManager:
93
for req_id, batch_index in structured_output_request_ids.items():
94
request = requests[req_id].structured_output_request
95
assert request is not None and request.grammar is not None
157
- if not request.grammar.matcher.is_terminated():
96
+ if not request.grammar.is_terminated():
97
request.grammar.fill_bitmask(bitmask_tensor, batch_index)
98
if batch_len < self._grammar_bitmask.shape[0]:
99
bitmask_tensor = self._grammar_bitmask[:batch_len]
vllm/v1/structured_output/backend_types.py
new
+89
@@ -0,0 +1,89 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+
3
+import enum
4
+from abc import ABC, abstractmethod
5
+
6
+import torch
7
+
8
+
9
+class StructuredOutputOptions(enum.Enum):
10
+ JSON = enum.auto()
11
+ JSON_OBJECT = enum.auto()
12
+ REGEX = enum.auto()
13
+ GRAMMAR = enum.auto()
14
+ CHOICE = enum.auto()
15
+
16
+
17
+StructuredOutputKey = tuple[StructuredOutputOptions, str]
18
+
19
+
20
+class StructuredOutputGrammar(ABC):
21
+ """Request-level backend for structured output requests."""
22
+
23
+ @abstractmethod
24
+ def accept_tokens(self, request_id: str, tokens: list[int]) -> bool:
25
+ """
26
+ Determines whether the provided tokens are accepted for the
27
+ given request.
28
+
29
+ Args:
30
+ request_id (str): The unique identifier for the request.
31
+ tokens (list[int]): A list of token IDs to evaluate.
32
+
33
+ Returns:
34
+ bool: True if the tokens are accepted, False otherwise.
35
+ """
36
+
37
+ @abstractmethod
38
+ def fill_bitmask(self, bitmask: torch.Tensor, batch_index: int) -> None:
39
+ """
40
+ Fills the bitmask for a specific batch index.
41
+
42
+ Args:
43
+ bitmask (torch.Tensor): The bitmask to fill
44
+ batch_index (int): The index in the bitmask to fill
45
+ """
46
+
47
+ @abstractmethod
48
+ def is_terminated(self) -> bool:
49
+ """
50
+ Checks whether the structured output process has terminated.
51
+
52
+ Returns:
53
+ bool: True if the process is terminated, False otherwise.
54
+ """
55
+
56
+ @abstractmethod
57
+ def reset(self):
58
+ """
59
+ Resets the state of the structured output grammar.
60
+ """
61
+
62
+
63
+class StructuredOutputBackend(ABC):
64
+ """Engine-level backend for structured output requests."""
65
+
66
+ @abstractmethod
67
+ def compile_grammar(self, request_type: StructuredOutputOptions,
68
+ grammar_spec: str) -> StructuredOutputGrammar:
69
+ """
70
+ Compiles a grammar specification into a structured output grammar.
71
+
72
+ Args:
73
+ request_type (StructuredOutputOptions): The type of structured
74
+ output request.
75
+ grammar_spec (str): The grammar specification to compile.
76
+
77
+ Returns:
78
+ StructuredOutputGrammar: The compiled structured output grammar.
79
+ """
80
+
81
+ @abstractmethod
82
+ def allocate_token_bitmask(self, max_num_seqs: int):
83
+ """
84
+ Allocates a token bitmask for the specified maximum number of sequences.
85
+
86
+ Args:
87
+ max_num_seqs (int): The maximum number of sequences for which
88
+ to allocate the bitmask.
89
+ """
vllm/v1/structured_output/backend_xgrammar.py
new
+143
@@ -0,0 +1,143 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+
3
+from dataclasses import dataclass, field
4
+from typing import TYPE_CHECKING
5
+
6
+import torch
7
+
8
+from vllm.config import VllmConfig
9
+from vllm.logger import init_logger
10
+from vllm.transformers_utils.tokenizer_group import init_tokenizer_from_configs
11
+from vllm.transformers_utils.tokenizers.mistral import MistralTokenizer
12
+from vllm.utils import LazyLoader
13
+from vllm.v1.structured_output.backend_types import (StructuredOutputBackend,
14
+ StructuredOutputGrammar,
15
+ StructuredOutputOptions)
16
+
17
+if TYPE_CHECKING:
18
+ import xgrammar as xgr
19
+else:
20
+ xgr = LazyLoader("xgr", globals(), "xgrammar")
21
+
22
+logger = init_logger(__name__)
23
+
24
+
25
+class XgrammarBackend(StructuredOutputBackend):
26
+
27
+ def __init__(self, vllm_config: VllmConfig):
28
+ self.vllm_config = vllm_config
29
+ tokenizer_group = init_tokenizer_from_configs(
30
+ model_config=vllm_config.model_config,
31
+ scheduler_config=vllm_config.scheduler_config,
32
+ parallel_config=vllm_config.parallel_config,
33
+ lora_config=vllm_config.lora_config) # type: ignore[arg-type]
34
+ tokenizer_group.ping()
35
+
36
+ tokenizer = tokenizer_group.get_lora_tokenizer(None)
37
+ self.vocab_size = vllm_config.model_config.get_vocab_size()
38
+ if isinstance(tokenizer, MistralTokenizer):
39
+ # NOTE: ideally, xgrammar should handle this accordingly.
40
+ # refer to https://github.com/mlc-ai/xgrammar/blob/d77c0a0173ef14779c918e3be7966ba852f7910f/python/xgrammar/tokenizer_info.py#L98
41
+ try:
42
+ encoded_vocab = [
43
+ token for token, _ in sorted(
44
+ tokenizer.get_vocab().items(),
45
+ key=lambda x: x[1],
46
+ )
47
+ ]
48
+ stop_token_ids = None
49
+ if hasattr(
50
+ tokenizer,
51
+ "eos_token_id",
52
+ ) and tokenizer.eos_token_id is not None:
53
+ stop_token_ids = [tokenizer.eos_token_id]
54
+ except AttributeError as e:
55
+ raise ValueError(
56
+ f"Cannot get the vocabulary of the tokenizer "
57
+ f"{type(tokenizer)}. The tokenizer should have a "
58
+ "get_vocab method.") from e
59
+ tokenizer_info = xgr.TokenizerInfo( # type: ignore
60
+ encoded_vocab=encoded_vocab,
61
+ # NOTE: https://github.com/mlc-ai/xgrammar/blob/5e141f6ff1ca02bc31f9e512e68b61f2a8ae88e5/tests/python/test_tokenizer_info.py#L43 # noqa: E501
62
+ vocab_type=xgr.VocabType.BYTE_FALLBACK,
63
+ vocab_size=self.vocab_size,
64
+ stop_token_ids=stop_token_ids,
65
+ add_prefix_space=True,
66
+ )
67
+ else:
68
+ tokenizer_info = xgr.TokenizerInfo.from_huggingface(
69
+ tokenizer,
70
+ vocab_size=self.vocab_size,
71
+ )
72
+ self.compiler = xgr.GrammarCompiler(tokenizer_info, max_threads=8)
73
+
74
+ def compile_grammar(self, request_type: StructuredOutputOptions,
75
+ grammar_spec: str) -> StructuredOutputGrammar:
76
+ if request_type == StructuredOutputOptions.JSON:
77
+ ctx = self.compiler.compile_json_schema(grammar_spec,
78
+ any_whitespace=False)
79
+ elif request_type == StructuredOutputOptions.JSON_OBJECT:
80
+ ctx = self.compiler.compile_builtin_json_grammar()
81
+ elif request_type == StructuredOutputOptions.GRAMMAR:
82
+ ctx = self.compiler.compile_grammar(grammar_spec)
83
+ elif request_type == StructuredOutputOptions.REGEX:
84
+ ctx = self.compiler.compile_regex(grammar_spec)
85
+ else:
86
+ logger.error(
87
+ "Validation should have already occurred. Please file an issue."
88
+ )
89
+ raise ValueError(
90
+ f"grammar is not of valid supported types. ({request_type!s})")
91
+
92
+ return XgrammarGrammar(
93
+ matcher=xgr.GrammarMatcher(ctx),
94
+ vocab_size=self.vocab_size,
95
+ ctx=ctx,
96
+ )
97
+
98
+ def allocate_token_bitmask(self, max_num_seqs: int):
99
+ return xgr.allocate_token_bitmask(max_num_seqs, self.vocab_size)
100
+
101
+
102
+@dataclass
103
+class XgrammarGrammar(StructuredOutputGrammar):
104
+ # NOTE: This would be a generic-enough class for
105
+ # supporting different backends, in the future.
106
+ # For now, just xgrammar.
107
+ #
108
+ # TODO: support max_rollback_tokens
109
+ # https://xgrammar.mlc.ai/docs/api/python/index.html#xgrammar.GrammarMatcher.find_jump_forward_string
110
+ # for jump-forward decoding
111
+
112
+ vocab_size: int
113
+ matcher: xgr.GrammarMatcher = field(hash=False)
114
+ ctx: xgr.CompiledGrammar = field(hash=False)
115
+ num_processed_tokens: int = field(default_factory=lambda: 0,
116
+ repr=False,
117
+ hash=False,
118
+ init=False)
119
+
120
+ def accept_tokens(self, request_id: str, tokens: list[int]) -> bool:
121
+ """Accepts a list of tokens and advances the FSM.
122
+
123
+ Returns True if the FSM was advanced successfully.
124
+ Returns False if the FSM failed to advance.
125
+ """
126
+ for token in tokens:
127
+ if not self.matcher.accept_token(token):
128
+ logger.error(
129
+ "Failed to advance FSM for request %s "
130
+ "for tokens %s. Please file an issue.", request_id, token)
131
+ return False
132
+ self.num_processed_tokens += 1
133
+ return True
134
+
135
+ def fill_bitmask(self, bitmask: torch.Tensor, idx: int) -> None:
136
+ self.matcher.fill_next_token_bitmask(bitmask, idx)
137
+
138
+ def is_terminated(self) -> bool:
139
+ return self.matcher.is_terminated()
140
+
141
+ def reset(self):
142
+ self.num_processed_tokens = 0
143
+ self.matcher.reset()
vllm/v1/structured_output/grammar.py
deleted
-77
@@ -1,77 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-from __future__ import annotations
3
-
4
-import enum
5
-from dataclasses import dataclass, field
6
-from typing import TYPE_CHECKING
7
-
8
-import torch
9
-
10
-from vllm.logger import init_logger
11
-from vllm.utils import LazyLoader
12
-
13
-if TYPE_CHECKING:
14
- import xgrammar as xgr
15
-else:
16
- xgr = LazyLoader("xgr", globals(), "xgrammar")
17
-
18
-logger = init_logger(__name__)
19
-
20
-
21
-class StructuredOutputOptions(enum.Enum):
22
- JSON = enum.auto()
23
- JSON_OBJECT = enum.auto()
24
- REGEX = enum.auto()
25
- GRAMMAR = enum.auto()
26
- CHOICE = enum.auto()
27
-
28
-
29
-StructuredOutputKey = tuple[StructuredOutputOptions, str]
30
-
31
-
32
-@dataclass
33
-class Grammar:
34
- # NOTE: This would be a generic-enough class for
35
- # supporting different backends, in the future.
36
- # For now, just xgrammar.
37
- #
38
- # TODO: support max_rollback_tokens
39
- # https://xgrammar.mlc.ai/docs/api/python/index.html#xgrammar.GrammarMatcher.find_jump_forward_string
40
- # for jump-forward decoding
41
-
42
- vocab_size: int
43
- matcher: xgr.GrammarMatcher = field(hash=False)
44
- ctx: xgr.CompiledGrammar = field(hash=False)
45
- num_processed_tokens: int = field(default_factory=lambda: 0,
46
- repr=False,
47
- hash=False,
48
- init=False)
49
-
50
- def accept_tokens(self, request_id: str, tokens: list[int]) -> bool:
51
- """Accepts a list of tokens and advances the FSM.
52
-
53
- Returns True if the FSM was advanced successfully.
54
- Returns False if the FSM failed to advance.
55
- """
56
- for token in tokens:
57
- if not self.matcher.accept_token(token):
58
- logger.error(
59
- "Failed to advance FSM for request %s "
60
- "for tokens %s. Please file an issue.", request_id, token)
61
- return False
62
- self.num_processed_tokens += 1
63
- return True
64
-
65
- def fill_bitmask(self, bitmask: torch.Tensor, idx: int) -> bool:
66
- return self.matcher.fill_next_token_bitmask(bitmask, idx)
67
-
68
- def reset(self):
69
- self.num_processed_tokens = 0
70
- self.matcher.reset()
71
-
72
- def __copy__(self):
73
- return Grammar(
74
- matcher=xgr.GrammarMatcher(self.ctx),
75
- vocab_size=self.vocab_size,
76
- ctx=self.ctx,
77
- )
vllm/v1/structured_output/request.py
+12
-6
@@ -9,15 +9,17 @@ from concurrent.futures._base import TimeoutError
9
from typing import Optional, Union, cast
10
11
from vllm.sampling_params import SamplingParams
12
-from vllm.v1.structured_output.grammar import (Grammar, StructuredOutputKey,
13
- StructuredOutputOptions)
12
+from vllm.v1.structured_output.backend_types import (StructuredOutputGrammar,
13
+ StructuredOutputKey,
14
+ StructuredOutputOptions)
15
16
17
@dataclasses.dataclass
18
class StructuredOutputRequest:
19
20
sampling_params: SamplingParams
20
- _grammar: Optional[Union[Future[Grammar], Grammar]] = None
21
+ _grammar: Optional[Union[Future[StructuredOutputGrammar],
22
+ StructuredOutputGrammar]] = None
23
24
def _check_grammar_completion(self) -> bool:
25
# NOTE: We have to lazy import to gate circular imports
@@ -37,12 +39,16 @@ class StructuredOutputRequest:
39
return self._check_grammar_completion()
40
41
@property
40
- def grammar(self) -> Optional[Grammar]:
42
+ def grammar(self) -> Optional[StructuredOutputGrammar]:
43
completed = self._check_grammar_completion()
42
- return cast(Optional[Grammar], self._grammar) if completed else None
44
+ return cast(Optional[StructuredOutputGrammar],
45
+ self._grammar) if completed else None
46
47
@grammar.setter
45
- def grammar(self, grammar: Union[Grammar, Future[Grammar]]) -> None:
48
+ def grammar(
49
+ self, grammar: Union[StructuredOutputGrammar,
50
+ Future[StructuredOutputGrammar]]
51
+ ) -> None:
52
self._grammar = grammar
53
54
@functools.cached_property