@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 80e9afb5bc
Fork 0

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, 19:45 UTC

[V1][Core] Support for Structured Outputs (#12388)

Signed-off-by: Aaron Pham <contact@aarnphm.xyz> Signed-off-by: Russell Bryant <rbryant@redhat.com> Co-authored-by: Russell Bryant <rbryant@redhat.com> Co-authored-by: Michael Goin <mgoin64@gmail.com> Co-authored-by: Nick Hill <nhill@redhat.com>

Aaron Pham committed Mar 7, 2025 at 10:19 UTC 80e9afb5bc53a5ca9f2b229e46c3cdead2704f5c
26 files changed +1528 -715
.buildkite/test-pipeline.yaml
+1
@@ -204,6 +204,7 @@ steps:
204 - VLLM_USE_V1=1 pytest -v -s v1/engine
205 - VLLM_USE_V1=1 pytest -v -s v1/sample
206 - VLLM_USE_V1=1 pytest -v -s v1/worker
207 + - VLLM_USE_V1=1 pytest -v -s v1/structured_output
208 - VLLM_USE_V1=1 pytest -v -s v1/test_stats.py
209 - VLLM_USE_V1=1 pytest -v -s v1/test_utils.py
210 # TODO: accuracy does not match, whether setting
.gitignore
+1 -1
@@ -197,7 +197,7 @@ _build/
197 hip_compat.h
198
199 # Benchmark dataset
200 -benchmarks/*.json
200 +benchmarks/**/*.json
201
202 # Linting
203 actionlint
benchmarks/benchmark_guided.py deleted
-507
@@ -1,507 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -"""Benchmark guided decoding throughput."""
3 -import argparse
4 -import dataclasses
5 -import json
6 -import os
7 -import random
8 -import time
9 -
10 -import datasets
11 -import pandas as pd
12 -import uvloop
13 -from transformers import AutoTokenizer, PreTrainedTokenizerBase
14 -
15 -from vllm.engine.arg_utils import AsyncEngineArgs, EngineArgs
16 -from vllm.entrypoints.openai.api_server import (
17 - build_async_engine_client_from_engine_args)
18 -from vllm.sampling_params import GuidedDecodingParams
19 -from vllm.utils import FlexibleArgumentParser, merge_async_iterators
20 -
21 -
22 -@dataclasses.dataclass
23 -class SampleRequest:
24 - """A class representing a single inference request for benchmarking.
25 -
26 - Attributes:
27 - prompt: The input text prompt for the model.
28 - multi_modal_data: Optional dictionary containing multi-modal data (e.g.
29 - images).
30 - prompt_len: The length of the prompt in tokens.
31 - expected_output_len: The expected length of the output in tokens.
32 - """
33 - prompt: str
34 - prompt_len: int
35 - expected_output_len: int
36 - schema: dict
37 - structure_type: str = 'json'
38 - completion: str = None
39 -
40 -
41 -def run_vllm(requests: list[SampleRequest],
42 - engine_args: EngineArgs,
43 - n: int,
44 - guided_decoding_rate: float = 1.0,
45 - warmup: bool = False) -> float:
46 - from vllm import LLM, SamplingParams
47 - llm = LLM(**vars(engine_args))
48 - assert all(
49 - llm.llm_engine.model_config.max_model_len >= (
50 - request.prompt_len + request.expected_output_len)
51 - for request in requests), (
52 - "Please ensure that max_model_len is greater than the sum of"
53 - " prompt_len and expected_output_len for all requests.")
54 -
55 - # Add the requests to the engine.
56 - prompts: list[str] = []
57 - sampling_params: list[SamplingParams] = []
58 - # create a list containing random selected true or false
59 - guided_decoding_req_idx = random.sample(
60 - range(len(requests)), int(len(requests) * guided_decoding_rate))
61 -
62 - if warmup:
63 - print(">>>>> Running warmup prompt, for the first 5")
64 - # We setup the first 5 requests to warmup FSM
65 - # if using xgrammar dataset, we will skip warmup
66 - warmup_requests = requests[:5]
67 - for i, request in enumerate(warmup_requests):
68 - prompts.append(request.prompt)
69 - sampling_params.append(
70 - SamplingParams(
71 - n=n,
72 - temperature=1.0,
73 - top_p=1.0,
74 - ignore_eos=True,
75 - max_tokens=request.expected_output_len,
76 - guided_decoding=GuidedDecodingParams(json=request.schema)
77 - if guided_decoding_rate > 0 else None,
78 - ))
79 - llm.generate(prompts, sampling_params, use_tqdm=False)
80 -
81 - print(">>>>> Benchmark started...")
82 - prompts = []
83 - sampling_params = []
84 - for i, request in enumerate(requests):
85 - prompts.append(request.prompt)
86 - sampling_params.append(
87 - SamplingParams(
88 - n=n,
89 - temperature=1.0,
90 - top_p=1.0,
91 - ignore_eos=True,
92 - max_tokens=request.expected_output_len,
93 - guided_decoding=GuidedDecodingParams(
94 - **{request.structure_type: request.schema})
95 - if i in guided_decoding_req_idx else None,
96 - ))
97 -
98 - start = time.perf_counter()
99 - outputs = llm.generate(prompts, sampling_params, use_tqdm=False)
100 - ret = []
101 - for output, request in zip(outputs, requests):
102 - generated_text = output.outputs[0].text
103 - ret.append({
104 - "generated": generated_text,
105 - "expected": request.completion
106 - })
107 - end = time.perf_counter()
108 - return end - start, ret
109 -
110 -
111 -async def run_vllm_async(
112 - requests: list[SampleRequest],
113 - engine_args: AsyncEngineArgs,
114 - n: int,
115 - guided_decoding_rate: float = 1.0,
116 - warmup: bool = False,
117 - disable_frontend_multiprocessing: bool = False) -> float:
118 - from vllm import SamplingParams
119 -
120 - async with build_async_engine_client_from_engine_args(
121 - engine_args, disable_frontend_multiprocessing) as llm:
122 -
123 - assert all(
124 - llm.model_config.max_model_len >= (request.prompt_len +
125 - request.expected_output_len)
126 - for request in requests), (
127 - "Please ensure that max_model_len is greater than the sum of"
128 - " prompt_len and expected_output_len for all requests.")
129 -
130 - # Add the requests to the engine.
131 - prompts: list[str] = []
132 - sampling_params: list[SamplingParams] = []
133 - guided_decoding_req_idx = random.sample(
134 - range(len(requests)), int(len(requests) * guided_decoding_rate))
135 -
136 - if warmup:
137 - print(">>>>>> Running warmup prompt, for the first 5")
138 - # We setup the first 5 requests to warmup FSM
139 - # if using xgrammar dataset, we will skip warmup
140 - warmup_requests = requests[:5]
141 - for i, request in enumerate(warmup_requests):
142 - prompts.append(request.prompt)
143 - sampling_params.append(
144 - SamplingParams(
145 - n=n,
146 - temperature=1.0,
147 - top_p=1.0,
148 - ignore_eos=True,
149 - max_tokens=request.expected_output_len,
150 - guided_decoding=GuidedDecodingParams(
151 - json=request.schema)
152 - if guided_decoding_rate > 0 else None,
153 - ))
154 - generators = []
155 - for i, (prompt, sp) in enumerate(zip(prompts, sampling_params)):
156 - generator = llm.generate(prompt, sp, request_id=f"test{i}")
157 - generators.append(generator)
158 - all_gens = merge_async_iterators(*generators)
159 - async for i, res in all_gens:
160 - pass
161 -
162 - print(">>>>> Benchmark started...")
163 - prompts = []
164 - sampling_params = []
165 - for i, request in enumerate(requests):
166 - prompts.append(request.prompt)
167 - sampling_params.append(
168 - SamplingParams(
169 - n=n,
170 - temperature=1.0,
171 - top_p=1.0,
172 - ignore_eos=True,
173 - max_tokens=request.expected_output_len,
174 - guided_decoding=GuidedDecodingParams(json=request.schema)
175 - if i in guided_decoding_req_idx else None,
176 - ))
177 -
178 - generators = []
179 - start_time = []
180 - latencies = []
181 - start = time.perf_counter()
182 - for i, (prompt, sp) in enumerate(zip(prompts, sampling_params)):
183 - generator = llm.generate(prompt, sp, request_id=f"test{i}")
184 - generators.append(generator)
185 - start_time.append(time.perf_counter())
186 - latencies.append([])
187 - all_gens = merge_async_iterators(*generators)
188 - generated_texts = [''] * len(requests)
189 - async for i, res in all_gens:
190 - generated_texts[i] = res.outputs[0].text
191 - lat = time.perf_counter() - start_time[i]
192 - latencies[i].append(lat)
193 - ret = [{
194 - 'generated': gt,
195 - 'expected': req.completion
196 - } for gt, req in zip(generated_texts, requests)]
197 - end = time.perf_counter()
198 - first_latency = pd.Series([lat[0] * 1000 for lat in latencies])
199 - next_latency = pd.Series([(lat[-1] - lat[0]) / len(lat[1:]) * 1000
200 - for lat in latencies])
201 - return end - start, ret, (first_latency, next_latency)
202 -
203 -
204 -def sample_requests(tokenizer: PreTrainedTokenizerBase,
205 - args: argparse.Namespace) -> list[SampleRequest]:
206 - if args.dataset == 'json':
207 - if args.json_schema_path is None:
208 - dir_path = os.path.dirname(os.path.realpath(__file__))
209 - args.json_schema_path = os.path.join(dir_path,
210 - "structured_schemas",
211 - "structured_schema_1.json")
212 - with open(args.json_schema_path) as f:
213 - schema = json.load(f)
214 - prompt = f"Generate an example of a user profile given the following schema: {json.dumps(schema)}" # noqa: E501
215 - input_len = len(tokenizer(prompt).input_ids)
216 - print(f"Input length of the prompt: {input_len} tokens")
217 - requests = [
218 - SampleRequest(prompt=prompt,
219 - prompt_len=input_len,
220 - expected_output_len=args.output_len,
221 - schema=schema,
222 - structure_type=args.structure_type)
223 - for _ in range(args.num_prompts)
224 - ]
225 -
226 - elif args.dataset == "grammar":
227 - schema = """
228 - ?start: select_statement
229 -
230 - ?select_statement: "SELECT " column_list " FROM " table_name
231 -
232 - ?column_list: column_name ("," column_name)*
233 -
234 - ?table_name: identifier
235 -
236 - ?column_name: identifier
237 -
238 - ?identifier: /[a-zA-Z_][a-zA-Z0-9_]*/
239 - """
240 - prompt = "Generate an SQL query to show the 'username' \
241 - and 'email' from the 'users' table."
242 -
243 - input_len = len(tokenizer(prompt).input_ids)
244 - print(f"Input length of the prompt: {input_len} tokens")
245 - requests = [
246 - SampleRequest(prompt=prompt,
247 - prompt_len=input_len,
248 - expected_output_len=args.output_len,
249 - schema=schema,
250 - structure_type=args.structure_type)
251 - for _ in range(args.num_prompts)
252 - ]
253 -
254 - elif args.dataset == "regex":
255 - regex = r"\w+@\w+\.com\n"
256 - args.regex = regex
257 - prompt = "Generate an email address for Alan Turing, \
258 - who works in Enigma. End in .com and new line. \
259 - Example result: alan.turing@enigma.com\n"
260 -
261 - input_len = len(tokenizer(prompt).input_ids)
262 - print(f"Input length of the prompt: {input_len} tokens")
263 - requests = [
264 - SampleRequest(prompt=prompt,
265 - prompt_len=input_len,
266 - expected_output_len=args.output_len,
267 - schema=regex,
268 - structure_type=args.structure_type)
269 - for _ in range(args.num_prompts)
270 - ]
271 -
272 - elif args.dataset == "choice":
273 - choice = ["Positive", "Negative"]
274 - args.choice = choice
275 - prompt = "Classify this sentiment: vLLM is wonderful!"
276 - input_len = len(tokenizer(prompt).input_ids)
277 - print(f"Input length of the prompt: {input_len} tokens")
278 - requests = [
279 - SampleRequest(prompt=prompt,
280 - prompt_len=input_len,
281 - expected_output_len=args.output_len,
282 - schema=choice,
283 - structure_type=args.structure_type)
284 - for _ in range(args.num_prompts)
285 - ]
286 -
287 - elif args.dataset == "xgrammar_bench":
288 - args.warmup = False
289 - requests: list[SampleRequest] = []
290 - dataset = datasets.load_dataset("NousResearch/json-mode-eval",
291 - split="train")
292 - print(f"dataset has {len(dataset)} entries")
293 - len_dataset = len(dataset)
294 - for data_point_idx in range(args.num_prompts):
295 - idx = data_point_idx
296 - while idx >= len_dataset:
297 - idx -= len_dataset
298 - schema = dataset["schema"][idx]
299 - prompt = tokenizer.apply_chat_template(dataset["prompt"][idx],
300 - tokenize=False)
301 - input_len = len(tokenizer(prompt).input_ids)
302 - completion = dataset["completion"][idx]
303 -
304 - requests.append(
305 - SampleRequest(prompt=prompt,
306 - prompt_len=input_len,
307 - expected_output_len=args.output_len,
308 - schema=schema,
309 - completion=completion))
310 -
311 - return requests
312 -
313 -
314 -def evaluate(ret, args):
315 -
316 - def _eval_correctness_json(expected, actual):
317 - # extract json string from string using regex
318 - import re
319 - actual = actual.replace('\n', '').replace(' ', '').strip()
320 - try:
321 - actual = re.search(r'\{.*\}', actual).group()
322 - actual = json.loads(actual)
323 - except Exception:
324 - return False
325 -
326 - return True
327 -
328 - def _eval_correctness_choice(expected, actual):
329 - return actual in args.choice
330 -
331 - def _eval_correctness_regex(expected, actual):
332 - import re
333 - return re.match(args.regex, actual) is not None
334 -
335 - def _eval_correctness(expected, actual):
336 - if args.structure_type == 'json':
337 - return _eval_correctness_json(expected, actual)
338 - elif args.structure_type == 'regex':
339 - return _eval_correctness_regex(expected, actual)
340 - elif args.structure_type == 'choice':
341 - return _eval_correctness_choice(expected, actual)
342 - else:
343 - return None
344 -
345 - scores = []
346 - for res in ret:
347 - score = _eval_correctness(res['expected'], res['generated'])
348 - res['correctness'] = score
349 - scores.append(score)
350 -
351 - not_none_scores = [score for score in scores if score is not None]
352 -
353 - return (sum(not_none_scores) / len(not_none_scores) *
354 - 100) if len(not_none_scores) > 0 else None
355 -
356 -
357 -def main(args: argparse.Namespace):
358 - print(args)
359 - random.seed(args.seed)
360 -
361 - # async engine is working for 'regex', 'choice' and 'grammar'
362 - if args.dataset == 'grammar':
363 - args.structure_type = 'grammar'
364 - args.async_engine = False
365 - elif args.dataset == 'regex':
366 - args.structure_type = 'regex'
367 - args.async_engine = False
368 - elif args.dataset == 'choice':
369 - args.structure_type = 'choice'
370 - args.async_engine = False
371 - else:
372 - args.structure_type = 'json'
373 -
374 - if args.no_guided_decoding:
375 - args.guided_decoding_ratio = 0
376 - if args.save_results:
377 - result_file_name = f'{args.guided_decoding_ratio}guided'
378 - result_file_name += f"_{args.model.split('/')[-1]}"
379 - result_file_name += f"_{args.dataset}"
380 - result_file_name += f"_{args.num_prompts}"
381 - result_file_name += f"_out{args.output_len}"
382 - result_file_name += f"_async{args.async_engine}"
383 - result_file_name += f"_warmup{args.warmup}"
384 - result_file_name += f"_chunkedprefill{args.enable_chunked_prefill}"
385 - result_file_name += ".txt"
386 - else:
387 - result_file_name = None
388 -
389 - # Synthesize a prompt with the given input length.
390 - tokenizer = AutoTokenizer.from_pretrained(
391 - args.tokenizer, trust_remote_code=args.trust_remote_code)
392 - requests = sample_requests(tokenizer, args)
393 -
394 - if args.async_engine:
395 - engine_args = AsyncEngineArgs.from_cli_args(args)
396 - elapsed_time, ret, (first_latency, next_latency) = uvloop.run(
397 - run_vllm_async(requests, engine_args, args.n,
398 - args.guided_decoding_ratio, args.warmup,
399 - args.disable_frontend_multiprocessing))
400 - else:
401 - engine_args = EngineArgs.from_cli_args(args)
402 - elapsed_time, ret = run_vllm(requests, engine_args, args.n,
403 - args.guided_decoding_ratio, args.warmup)
404 - first_latency, next_latency = None, None
405 -
406 - score = evaluate(ret, args)
407 - total_num_tokens = sum(request.prompt_len + request.expected_output_len
408 - for request in requests)
409 - total_output_tokens = sum(request.expected_output_len
410 - for request in requests)
411 - if first_latency is not None:
412 - latency_breakdown = "\nFirst token latency(msecs):\n"
413 - latency_breakdown += f"{first_latency.describe()}"
414 - latency_breakdown += "\nNext token latency(msecs):\n"
415 - latency_breakdown += f"{next_latency.describe()}"
416 - print(
417 - f"Throughput: {len(requests) / elapsed_time:.2f} requests/s, "
418 - f"{total_num_tokens / elapsed_time:.2f} total tokens/s, "
419 - f"{total_output_tokens / elapsed_time:.2f} output tokens/s",
420 - f"Correct rate is {score} %",
421 - f"{latency_breakdown if first_latency is not None else ''}")
422 -
423 - # Output JSON results if specified
424 - if args.output_json or result_file_name:
425 - results = {
426 - "elapsed_time": elapsed_time,
427 - "num_requests": len(requests),
428 - "total_num_tokens": total_num_tokens,
429 - "total_output_tokens": total_output_tokens,
430 - "requests_per_second": len(requests) / elapsed_time,
431 - "tokens_per_second": f"{total_num_tokens / elapsed_time:.2f}",
432 - "output_tokens_per_second":
433 - f"{total_output_tokens / elapsed_time:.2f}",
434 - "correct_rate(%)": score
435 - }
436 - results = {"outputs": ret, **results}
437 - if first_latency is not None:
438 - results["first_token_latency(msecs)"] = first_latency.describe(
439 - ).to_dict()
440 - results["next_token_latency(msecs)"] = next_latency.describe(
441 - ).to_dict()
442 - if args.output_json:
443 - with open(args.output_json, "w") as f:
444 - json.dump(results, f, indent=4)
445 - elif result_file_name:
446 - with open(result_file_name, "w") as f:
447 - json.dump(results, f, indent=4)
448 -
449 -
450 -if __name__ == "__main__":
451 - parser = FlexibleArgumentParser(description="Benchmark guided decoding.")
452 - parser = AsyncEngineArgs.add_cli_args(parser)
453 -
454 - parser.add_argument("--output-len",
455 - type=int,
456 - default=512,
457 - help="Output length for each request. Overrides the "
458 - "output length from the dataset.")
459 - parser.add_argument(
460 - "--dataset",
461 - default='json',
462 - choices=['json', 'grammar', 'regex', 'choice', 'xgrammar_bench'])
463 - parser.add_argument("--json_schema_path",
464 - type=str,
465 - default=None,
466 - help="Path to json schema.")
467 - parser.add_argument("--n",
468 - type=int,
469 - default=1,
470 - help="Number of generated sequences per prompt.")
471 - parser.add_argument("--num-prompts",
472 - type=int,
473 - default=10,
474 - help="Number of prompts to process.")
475 - parser.add_argument(
476 - '--output-json',
477 - type=str,
478 - default=None,
479 - help='Path to save the throughput results in JSON format.')
480 - parser.add_argument("--async-engine",
481 - action='store_true',
482 - default=False,
483 - help="Use vLLM async engine rather than LLM class.")
484 - parser.add_argument("--no-guided-decoding",
485 - action='store_true',
486 - default=False,
487 - help="Whether to disable JSON decoding or not.")
488 - parser.add_argument("--guided-decoding-ratio",
489 - type=float,
490 - default=1.0,
491 - help="Ratio of Guided Decoding requests")
492 - parser.add_argument("--disable-frontend-multiprocessing",
493 - action='store_true',
494 - default=False,
495 - help="Disable decoupled async engine frontend.")
496 - parser.add_argument("--warmup",
497 - action="store_true",
498 - default=False,
499 - help="Run warmup prompts before benchmark.")
500 - parser.add_argument("--save-results",
501 - action="store_true",
502 - default=False,
503 - help="save output results.")
504 - args = parser.parse_args()
505 - if args.tokenizer is None:
506 - args.tokenizer = args.model
507 - main(args)
benchmarks/benchmark_serving_structured_output.py renamed
+47 -32
@@ -1,5 +1,5 @@
1 # SPDX-License-Identifier: Apache-2.0
2 -r"""Benchmark online serving throughput with guided decoding.
2 +r"""Benchmark online serving throughput with structured outputs.
3
4 On the server side, run one of the following commands:
5 (vLLM OpenAI API server)
@@ -9,12 +9,12 @@ On the server side, run one of the following commands:
9 ./launch_tgi_server.sh <your_model> <max_batch_total_tokens>
10
11 On the client side, run:
12 - python benchmarks/benchmark_serving_guided.py \
12 + python benchmarks/benchmark_serving_structured_output.py \
13 --backend <backend> \
14 --model <your_model> \
15 --dataset json \
16 - --guided-decoding-ratio 1.0 \
17 - --guided-decoding-backend xgrammar \
16 + --structured-output-ratio 1.0 \
17 + --structured-output-backend xgrammar \
18 --request-rate 10 \
19 --num-prompts 1000
20
@@ -52,6 +52,9 @@ try:
52 except ImportError:
53 from argparse import ArgumentParser as FlexibleArgumentParser
54
55 +from vllm.v1.structured_output.utils import (
56 + has_xgrammar_unsupported_json_features)
57 +
58 MILLISECONDS_TO_SECONDS_CONVERSION = 1000
59
60
@@ -191,7 +194,17 @@ def sample_requests(tokenizer: PreTrainedTokenizerBase,
194 requests: list[SampleRequest] = []
195 dataset = datasets.load_dataset("NousResearch/json-mode-eval",
196 split="train")
194 - print(f"dataset has {len(dataset)} entries")
197 + full_dataset_len = len(dataset)
198 +
199 + def _filter_func(item):
200 + import json
201 + schema = json.loads(item["schema"])
202 + return not has_xgrammar_unsupported_json_features(schema)
203 +
204 + dataset = dataset.filter(_filter_func)
205 + num_filtered_out = full_dataset_len - len(dataset)
206 + print(f"dataset has {len(dataset)} entries after filtering "
207 + f"out {num_filtered_out} entries with unsupported features")
208 len_dataset = len(dataset)
209 for data_point_idx in range(args.num_prompts):
210 idx = data_point_idx
@@ -220,21 +233,21 @@ async def get_request(
233 burstiness: float = 1.0,
234 ) -> AsyncGenerator[tuple[int, SampleRequest], None]:
235 """
223 - Asynchronously generates requests at a specified rate
236 + Asynchronously generates requests at a specified rate
237 with OPTIONAL burstiness.
225 -
238 +
239 Args:
227 - input_requests:
240 + input_requests:
241 A list of input requests, each represented as a tuple.
229 - request_rate:
242 + request_rate:
243 The rate at which requests are generated (requests/s).
231 - burstiness (optional):
232 - The burstiness factor of the request generation.
244 + burstiness (optional):
245 + The burstiness factor of the request generation.
246 Only takes effect when request_rate is not inf.
247 Default value is 1, which follows a Poisson process.
248 Otherwise, the request intervals follow a gamma distribution.
236 - A lower burstiness value (0 < burstiness < 1) results
237 - in more bursty requests, while a higher burstiness value
249 + A lower burstiness value (0 < burstiness < 1) results
250 + in more bursty requests, while a higher burstiness value
251 (burstiness > 1) results in a more uniform arrival of requests.
252 """
253 input_requests = iter(input_requests)
@@ -378,8 +391,8 @@ async def benchmark(
391 selected_percentiles: list[str],
392 ignore_eos: bool,
393 max_concurrency: Optional[int],
381 - guided_decoding_ratio: float,
382 - guided_decoding_backend: str,
394 + structured_output_ratio: float,
395 + structured_output_backend: str,
396 goodput_config_dict: Optional[dict[str, float]] = None,
397 ):
398 if backend in ASYNC_REQUEST_FUNCS:
@@ -391,16 +404,18 @@ async def benchmark(
404 extra_body = {}
405 # Add the schema to the extra_body
406 extra_body[request.structure_type] = request.schema
394 - # Add the specific guided_decoding_backend
395 - extra_body["guided_decoding_backend"] = guided_decoding_backend
407 + # Add the specific structured_output_backend
408 + extra_body["guided_decoding_backend"] = structured_output_backend
409 return extra_body
410
411 print("Starting initial single prompt test run...")
399 - guided_decoding_req_idx = random.sample(
412 + structured_output_req_idx = random.sample(
413 range(len(input_requests)),
401 - int(len(input_requests) * guided_decoding_ratio))
414 + int(len(input_requests) * structured_output_ratio))
415
416 test_request = input_requests[0]
417 + test_req_extra_body = (prepare_extra_body(test_request)
418 + if 0 in structured_output_req_idx else None)
419 test_input = RequestFuncInput(
420 model=model_id,
421 prompt=test_request.prompt,
@@ -408,7 +423,7 @@ async def benchmark(
423 prompt_len=test_request.prompt_len,
424 output_len=test_request.expected_output_len,
425 ignore_eos=ignore_eos,
411 - extra_body=prepare_extra_body(test_request),
426 + extra_body=test_req_extra_body,
427 )
428 test_output = await request_func(request_func_input=test_input)
429 if not test_output.success:
@@ -427,7 +442,7 @@ async def benchmark(
442 prompt_len=test_request.prompt_len,
443 output_len=test_request.expected_output_len,
444 ignore_eos=ignore_eos,
430 - extra_body=prepare_extra_body(test_request),
445 + extra_body=test_req_extra_body,
446 )
447 profile_output = await request_func(request_func_input=profile_input)
448 if profile_output.success:
@@ -465,7 +480,7 @@ async def benchmark(
480 async for i, request in get_request(input_requests, request_rate,
481 burstiness):
482 extra_body = prepare_extra_body(
468 - request) if i in guided_decoding_req_idx else None
483 + request) if i in structured_output_req_idx else None
484 request_func_input = RequestFuncInput(
485 model=model_id,
486 prompt=request.prompt,
@@ -708,10 +723,10 @@ def main(args: argparse.Namespace):
723 else:
724 args.structure_type = 'guided_json'
725
711 - if args.no_guided_decoding:
712 - args.guided_decoding_ratio = 0
726 + if args.no_structured_output:
727 + args.structured_output_ratio = 0
728 if args.save_results:
714 - result_file_name = f'{args.guided_decoding_ratio}guided'
729 + result_file_name = f'{args.structured_output_ratio}guided'
730 result_file_name += f"_{backend}"
731 result_file_name += f"_{args.request_rate}qps"
732 result_file_name += f"_{args.model.split('/')[-1]}"
@@ -744,8 +759,8 @@ def main(args: argparse.Namespace):
759 ],
760 ignore_eos=args.ignore_eos,
761 max_concurrency=args.max_concurrency,
747 - guided_decoding_ratio=args.guided_decoding_ratio,
748 - guided_decoding_backend=args.guided_decoding_backend,
762 + structured_output_ratio=args.structured_output_ratio,
763 + structured_output_backend=args.structured_output_backend,
764 goodput_config_dict=goodput_config_dict,
765 ))
766
@@ -943,19 +958,19 @@ if __name__ == "__main__":
958 "goodput, refer to DistServe paper: https://arxiv.org/pdf/2401.09670 "
959 "and the blog: https://hao-ai-lab.github.io/blogs/distserve")
960
946 - parser.add_argument("--no-guided-decoding",
961 + parser.add_argument("--no-structured-output",
962 action='store_true',
963 default=False,
964 help="Whether to disable JSON decoding or not.")
950 - parser.add_argument("--guided-decoding-ratio",
965 + parser.add_argument("--structured-output-ratio",
966 type=float,
967 default=1.0,
953 - help="Ratio of Guided Decoding requests")
954 - parser.add_argument("--guided-decoding-backend",
968 + help="Ratio of Structured Outputs requests")
969 + parser.add_argument("--structured-output-backend",
970 type=str,
971 choices=["outlines", "lm-format-enforcer", "xgrammar"],
972 default="xgrammar",
958 - help="Backend to use for guided decoding")
973 + help="Backend to use for structured outputs")
974
975 args = parser.parse_args()
976 main(args)
benchmarks/run_structured_output_benchmark.sh new
+64
@@ -0,0 +1,64 @@
1 +#!/bin/bash
2 +
3 +# Define the model to use
4 +MODEL=${1:-"Qwen/Qwen2.5-7B-Instruct"}
5 +
6 +# Define the backend to use
7 +BACKEND=${2:-"vllm"}
8 +
9 +# Define the dataset to use
10 +DATASET=${3:-"xgrammar_bench"}
11 +
12 +# Define the guided decoding backend
13 +GUIDED_BACKEND=${4:-"xgrammar"}
14 +
15 +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
16 +OUTPUT_DIR=${5:-"$SCRIPT_DIR/structured_output_benchmark_results"}
17 +
18 +GUIDED_RATIO=${6:-0.5}
19 +
20 +# Create output directory if it doesn't exist
21 +mkdir -p "$OUTPUT_DIR"
22 +
23 +# Define QPS values to test
24 +QPS_VALUES=(70 60 50 25 20 15 10)
25 +
26 +# Common parameters
27 +COMMON_PARAMS="--backend $BACKEND \
28 + --model $MODEL \
29 + --dataset $DATASET \
30 + --structured-output-backend $GUIDED_BACKEND \
31 + --structured-output-ratio $GUIDED_RATIO \
32 + --save-results \
33 + --result-dir $OUTPUT_DIR"
34 +
35 +echo "Starting structured output benchmark with model: $MODEL"
36 +echo "Backend: $BACKEND"
37 +echo "Dataset: $DATASET"
38 +echo "Structured output backend: $GUIDED_BACKEND"
39 +echo "Results will be saved to: $OUTPUT_DIR"
40 +echo "----------------------------------------"
41 +
42 +# Run benchmarks with different QPS values
43 +for qps in "${QPS_VALUES[@]}"; do
44 + echo "Running benchmark with QPS: $qps"
45 +
46 + # Get git hash and branch for the filename
47 + GIT_HASH=$(git rev-parse --short HEAD 2>/dev/null || echo "unknown")
48 + GIT_BRANCH=$(git rev-parse --abbrev-ref HEAD 2>/dev/null || echo "unknown")
49 +
50 + # Construct filename for this run
51 + FILENAME="${GUIDED_BACKEND}_${BACKEND}_${qps}qps_$(basename $MODEL)_${DATASET}_${GIT_HASH}.json"
52 +
53 + # Run the benchmark
54 + python "$SCRIPT_DIR/benchmark_serving_structured_output.py" $COMMON_PARAMS \
55 + --request-rate $qps \
56 + --result-filename "$FILENAME" \
57 + --port ${PORT:-8000}
58 +
59 + echo "Completed benchmark with QPS: $qps"
60 + echo "----------------------------------------"
61 +done
62 +
63 +echo "All benchmarks completed!"
64 +echo "Results saved to: $OUTPUT_DIR"
benchmarks/structured_schemas/structured_schema_1.json
+22 -110
@@ -1,113 +1,25 @@
1 {
2 - "$schema":
3 - "https://json-schema.org/draft/2020-12/schema",
4 - "title":
5 - "User Profile",
6 - "type":
7 - "object",
2 + "type": "array",
3 + "items": {
4 + "type": "object",
5 "properties": {
9 - "userId": {
10 - "type": "string",
11 - "description": "Unique identifier for the user."
12 - },
13 - "personalInfo": {
14 - "type": "object",
15 - "properties": {
16 - "firstName": {
17 - "type": "string",
18 - "description": "The user's first name."
19 - },
20 - "lastName": {
21 - "type": "string",
22 - "description": "The user's last name."
23 - },
24 - "age": {
25 - "type": "integer",
26 - "minimum": 0,
27 - "description": "The user's age."
28 - },
29 - "phoneNumbers": {
30 - "type":
31 - "array",
32 - "items": {
33 - "type": "object",
34 - "properties": {
35 - "type": {
36 - "type": "string",
37 - "enum": ["home", "work", "mobile"],
38 - "description": "Type of phone number."
39 - },
40 - "number": {
41 - "type": "string",
42 - "pattern": "^\\+?[1-9]\\d{1,14}$",
43 - "description": "Phone number in E.164 format."
44 - }
45 - },
46 - "required": ["type", "number"]
47 - },
48 - "description":
49 - "List of phone numbers associated with the user."
50 - }
51 - },
52 - "required": ["firstName", "lastName"]
53 - },
54 - "address": {
55 - "type": "object",
56 - "properties": {
57 - "street": {
58 - "type": "string",
59 - "description": "Street address."
60 - },
61 - "city": {
62 - "type": "string",
63 - "description": "City name."
64 - },
65 - "state": {
66 - "type": "string",
67 - "description": "State or province."
68 - },
69 - "postalCode": {
70 - "type": "string",
71 - "pattern": "^\\d{5}(-\\d{4})?$",
72 - "description": "Postal code."
73 - },
74 - "country": {
75 - "type": "string",
76 - "description": "Country name."
77 - }
78 - },
79 - "required": ["street", "city", "state", "postalCode", "country"]
80 - },
81 - "preferences": {
82 - "type": "object",
83 - "properties": {
84 - "newsletterSubscribed": {
85 - "type":
86 - "boolean",
87 - "description":
88 - "Indicates if the user is subscribed to the newsletter."
89 - },
90 - "favoriteCategories": {
91 - "type": "array",
92 - "items": {
93 - "type": "string"
94 - },
95 - "description": "List of user's favorite categories."
96 - }
97 - },
98 - "required": ["newsletterSubscribed"]
99 - },
100 - "accountStatus": {
101 - "type": "string",
102 - "enum": ["active", "inactive", "suspended"],
103 - "description": "Current status of the user's account."
104 - },
105 - "registrationDate": {
106 - "type": "string",
107 - "format": "date-time",
108 - "description": "ISO 8601 formatted date-time of user registration."
109 - }
6 + "name": { "type": "string" },
7 + "race": { "type": "string" },
8 + "class": { "type": "string" },
9 + "level": { "type": "integer" },
10 + "background": { "type": "string" },
11 + "alignment": { "type": "string" },
12 + "backstory": { "type": "string" }
13 },
111 - "required":
112 - ["userId", "personalInfo", "address", "accountStatus", "registrationDate"]
113 -}
\ No newline at end of file
14 + "required": [
15 + "name",
16 + "race",
17 + "class",
18 + "level",
19 + "background",
20 + "alignment",
21 + "backstory"
22 + ]
23 + }
24 +}
25 +
tests/v1/core/test_scheduler.py
+28 -11
@@ -1,12 +1,13 @@
1 # SPDX-License-Identifier: Apache-2.0
2 from typing import Optional
3
4 -from vllm.config import CacheConfig, ModelConfig, SchedulerConfig
4 +from vllm.config import CacheConfig, ModelConfig, SchedulerConfig, VllmConfig
5 from vllm.multimodal.inputs import MultiModalKwargs, PlaceholderRange
6 from vllm.sampling_params import SamplingParams
7 from vllm.v1.core.scheduler import Scheduler, SchedulerOutput
8 from vllm.v1.outputs import ModelRunnerOutput
9 from vllm.v1.request import Request, RequestStatus
10 +from vllm.v1.structured_output import StructuredOutputManager
11
12 EOS_TOKEN_ID = 50256
13
@@ -36,13 +37,21 @@ def create_scheduler(
37 swap_space=0,
38 cache_dtype="auto",
39 )
40 + vllm_config = VllmConfig(
41 + scheduler_config=scheduler_config,
42 + model_config=model_config,
43 + cache_config=cache_config,
44 + )
45 cache_config.num_gpu_blocks = 10000
40 - return Scheduler(scheduler_config,
41 - model_config,
42 - cache_config,
43 - speculative_config=None,
44 - lora_config=None,
45 - log_stats=True)
46 + return Scheduler(
47 + scheduler_config,
48 + model_config,
49 + cache_config,
50 + speculative_config=None,
51 + lora_config=None,
52 + log_stats=True,
53 + structured_output_manager=StructuredOutputManager(vllm_config),
54 + )
55
56
57 def create_requests(
@@ -249,7 +258,9 @@ def test_stop_via_update_from_output():
258 },
259 num_common_prefix_blocks=0,
260 finished_req_ids=set(),
252 - free_encoder_input_ids=[])
261 + free_encoder_input_ids=[],
262 + structured_output_request_ids={},
263 + grammar_bitmask=None)
264
265 model_output = ModelRunnerOutput(
266 req_ids=[req.request_id for req in requests],
@@ -299,7 +310,9 @@ def test_stop_via_update_from_output():
310 },
311 num_common_prefix_blocks=0,
312 finished_req_ids=set(),
302 - free_encoder_input_ids=[])
313 + free_encoder_input_ids=[],
314 + structured_output_request_ids={},
315 + grammar_bitmask=None)
316
317 model_output = ModelRunnerOutput(
318 req_ids=[req.request_id for req in requests],
@@ -347,7 +360,9 @@ def test_stop_via_update_from_output():
360 },
361 num_common_prefix_blocks=0,
362 finished_req_ids=set(),
350 - free_encoder_input_ids=[])
363 + free_encoder_input_ids=[],
364 + structured_output_request_ids={},
365 + grammar_bitmask=None)
366
367 model_output = ModelRunnerOutput(
368 req_ids=[req.request_id for req in requests],
@@ -392,7 +407,9 @@ def test_stop_via_update_from_output():
407 },
408 num_common_prefix_blocks=0,
409 finished_req_ids=set(),
395 - free_encoder_input_ids=[])
410 + free_encoder_input_ids=[],
411 + structured_output_request_ids={},
412 + grammar_bitmask=None)
413
414 model_output = ModelRunnerOutput(
415 req_ids=[requests[0].request_id],
tests/v1/entrypoints/conftest.py
+17 -5
@@ -29,6 +29,7 @@ def sample_regex():
29 r"(25[0-5]|(2[0-4]|1\d|[1-9]|)\d)")
30
31
32 +# Note: Ensure this only uses attributes compatible with xgrammar
33 @pytest.fixture
34 def sample_json_schema():
35 return {
@@ -44,9 +45,7 @@ def sample_json_schema():
45 "type": "array",
46 "items": {
47 "type": "string",
47 - "maxLength": 10
48 - },
49 - "minItems": 3
48 + }
49 },
50 "work_history": {
51 "type": "array",
@@ -71,8 +70,9 @@ def sample_json_schema():
70 }
71
72
73 +# A schema unsupported by xgrammar
74 @pytest.fixture
75 -def sample_complex_json_schema():
75 +def unsupported_json_schema():
76 return {
77 "type": "object",
78 "properties": {
@@ -150,7 +150,19 @@ def sample_guided_choice():
150
151
152 @pytest.fixture
153 -def sample_sql_statements():
153 +def sample_sql_ebnf():
154 + return """
155 +root ::= select_statement
156 +select_statement ::= "SELECT" column "from" table "where" condition
157 +column ::= "col_1" | "col_2"
158 +table ::= "table_1" | "table_2"
159 +condition ::= column "=" number
160 +number ::= "1" | "2"
161 +"""
162 +
163 +
164 +@pytest.fixture
165 +def sample_sql_lark():
166 return ("""
167 start: select_statement
168 select_statement: "SELECT" column "from" table "where" condition
tests/v1/entrypoints/llm/__init__.py
tests/v1/entrypoints/llm/test_struct_output_generate.py new
+269
@@ -0,0 +1,269 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +import json
4 +
5 +import jsonschema
6 +import pytest
7 +
8 +from vllm.entrypoints.llm import LLM
9 +from vllm.outputs import RequestOutput
10 +from vllm.sampling_params import GuidedDecodingParams, SamplingParams
11 +
12 +MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct"
13 +GUIDED_DECODING_BACKENDS_V1 = ["xgrammar"]
14 +
15 +
16 +@pytest.mark.skip_global_cleanup
17 +@pytest.mark.parametrize("guided_decoding_backend",
18 + GUIDED_DECODING_BACKENDS_V1)
19 +def test_guided_json_completion(monkeypatch, sample_json_schema,
20 + guided_decoding_backend: str):
21 + monkeypatch.setenv("VLLM_USE_V1", "1")
22 + llm = LLM(model=MODEL_NAME, max_model_len=1024)
23 + sampling_params = SamplingParams(temperature=1.0,
24 + max_tokens=1000,
25 + guided_decoding=GuidedDecodingParams(
26 + json=sample_json_schema,
27 + backend=guided_decoding_backend))
28 + outputs = llm.generate(prompts=[
29 + f"Give an example JSON for an employee profile "
30 + f"that fits this schema: {sample_json_schema}"
31 + ] * 2,
32 + sampling_params=sampling_params,
33 + use_tqdm=True)
34 +
35 + assert outputs is not None
36 +
37 + for output in outputs:
38 + assert output is not None
39 + assert isinstance(output, RequestOutput)
40 + prompt = output.prompt
41 +
42 + generated_text = output.outputs[0].text
43 + assert generated_text is not None
44 + print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
45 + output_json = json.loads(generated_text)
46 + jsonschema.validate(instance=output_json, schema=sample_json_schema)
47 +
48 +
49 +@pytest.mark.skip_global_cleanup
50 +@pytest.mark.parametrize("guided_decoding_backend",
51 + GUIDED_DECODING_BACKENDS_V1)
52 +def test_guided_json_object(monkeypatch, guided_decoding_backend: str):
53 + monkeypatch.setenv("VLLM_USE_V1", "1")
54 + llm = LLM(model=MODEL_NAME, max_model_len=1024)
55 + sampling_params = SamplingParams(temperature=1.0,
56 + max_tokens=100,
57 + n=2,
58 + guided_decoding=GuidedDecodingParams(
59 + json_object=True,
60 + backend=guided_decoding_backend))
61 +
62 + outputs = llm.generate(
63 + prompts=("Generate a JSON object with curly braces for a person with "
64 + "name and age fields for John Smith who is 31 years old."),
65 + sampling_params=sampling_params,
66 + use_tqdm=True)
67 +
68 + assert outputs is not None
69 + for output in outputs:
70 + assert output is not None
71 + assert isinstance(output, RequestOutput)
72 +
73 + for i in range(2):
74 + generated_text = output.outputs[i].text
75 + print(generated_text)
76 + assert generated_text is not None
77 +
78 + # Parse to verify it is valid JSON
79 + parsed_json = json.loads(generated_text)
80 + assert isinstance(parsed_json, dict)
81 +
82 +
83 +@pytest.mark.skip_global_cleanup
84 +@pytest.mark.parametrize("guided_decoding_backend",
85 + GUIDED_DECODING_BACKENDS_V1)
86 +def test_guided_json_unsupported_schema(monkeypatch, unsupported_json_schema,
87 + guided_decoding_backend: str):
88 + monkeypatch.setenv("VLLM_USE_V1", "1")
89 + llm = LLM(model=MODEL_NAME, max_model_len=1024)
90 + sampling_params = SamplingParams(temperature=1.0,
91 + max_tokens=1000,
92 + guided_decoding=GuidedDecodingParams(
93 + json=unsupported_json_schema,
94 + backend=guided_decoding_backend))
95 + with pytest.raises(ValueError,
96 + match="The provided JSON schema contains features "
97 + "not supported by xgrammar."):
98 + llm.generate(prompts=[
99 + f"Give an example JSON for an employee profile "
100 + f"that fits this schema: {unsupported_json_schema}"
101 + ] * 2,
102 + sampling_params=sampling_params,
103 + use_tqdm=True)
104 +
105 +
106 +@pytest.mark.skip_global_cleanup
107 +@pytest.mark.parametrize("guided_decoding_backend",
108 + GUIDED_DECODING_BACKENDS_V1)
109 +def test_guided_grammar_ebnf(monkeypatch, sample_sql_ebnf,
110 + guided_decoding_backend: str):
111 + monkeypatch.setenv("VLLM_USE_V1", "1")
112 + llm = LLM(model=MODEL_NAME, max_model_len=1024)
113 + sampling_params = SamplingParams(temperature=0.8,
114 + top_p=0.95,
115 + max_tokens=1000,
116 + guided_decoding=GuidedDecodingParams(
117 + grammar=sample_sql_ebnf,
118 + backend=guided_decoding_backend))
119 + outputs = llm.generate(
120 + prompts=("Generate a sql statement that selects col_1 from "
121 + "table_1 where it is equal to 1"),
122 + sampling_params=sampling_params,
123 + use_tqdm=True,
124 + )
125 +
126 + assert outputs is not None
127 + for output in outputs:
128 + assert output is not None
129 + assert isinstance(output, RequestOutput)
130 + prompt = output.prompt
131 +
132 + generated_text = output.outputs[0].text
133 + assert generated_text is not None
134 +
135 + # remove spaces for comparison b/c we removed them in the grammar
136 + ground_truth = "SELECT col_1 from table_1 where col_1 = 1".replace(
137 + " ", "")
138 +
139 + assert generated_text.strip() == ground_truth
140 +
141 + print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
142 +
143 +
144 +@pytest.mark.skip_global_cleanup
145 +@pytest.mark.parametrize("guided_decoding_backend",
146 + GUIDED_DECODING_BACKENDS_V1)
147 +def test_guided_grammar_lark(monkeypatch, sample_sql_lark,
148 + guided_decoding_backend: str):
149 + monkeypatch.setenv("VLLM_USE_V1", "1")
150 + llm = LLM(model=MODEL_NAME, max_model_len=1024)
151 + sampling_params = SamplingParams(temperature=0.8,
152 + top_p=0.95,
153 + max_tokens=1000,
154 + guided_decoding=GuidedDecodingParams(
155 + grammar=sample_sql_lark,
156 + backend=guided_decoding_backend))
157 + outputs = llm.generate(
158 + prompts=("Generate a sql statement that selects col_1 from "
159 + "table_1 where it is equal to 1"),
160 + sampling_params=sampling_params,
161 + use_tqdm=True,
162 + )
163 +
164 + assert outputs is not None
165 + for output in outputs:
166 + assert output is not None
167 + assert isinstance(output, RequestOutput)
168 + prompt = output.prompt
169 +
170 + generated_text = output.outputs[0].text
171 + assert generated_text is not None
172 +
173 + # use Lark to parse the output, and make sure it's a valid parse tree
174 + from lark import Lark
175 + parser = Lark(sample_sql_lark)
176 + parser.parse(generated_text)
177 +
178 + # remove spaces for comparison b/c we removed them in the grammar
179 + ground_truth = "SELECT col_1 from table_1 where col_1 = 1".replace(
180 + " ", "")
181 +
182 + assert generated_text.strip() == ground_truth
183 +
184 + print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
185 +
186 +
187 +@pytest.mark.skip_global_cleanup
188 +@pytest.mark.parametrize("guided_decoding_backend",
189 + GUIDED_DECODING_BACKENDS_V1)
190 +def test_guided_grammar_ebnf_invalid(monkeypatch,
191 + guided_decoding_backend: str):
192 + monkeypatch.setenv("VLLM_USE_V1", "1")
193 + llm = LLM(model=MODEL_NAME, max_model_len=1024)
194 + sampling_params = SamplingParams(temperature=0.8,
195 + top_p=0.95,
196 + max_tokens=1000,
197 + guided_decoding=GuidedDecodingParams(
198 + grammar="not a grammar",
199 + backend=guided_decoding_backend))
200 + with pytest.raises(ValueError,
201 + match="Failed to convert the grammar "
202 + "from Lark to EBNF."):
203 + llm.generate(
204 + prompts=("Generate a sql statement that selects col_1 from "
205 + "table_1 where it is equal to 1"),
206 + sampling_params=sampling_params,
207 + use_tqdm=True,
208 + )
209 +
210 +
211 +@pytest.mark.skip_global_cleanup
212 +@pytest.mark.parametrize("guided_decoding_backend",
213 + GUIDED_DECODING_BACKENDS_V1)
214 +def test_guided_regex(monkeypatch, sample_regex, guided_decoding_backend: str):
215 + monkeypatch.setenv("VLLM_USE_V1", "1")
216 + llm = LLM(model=MODEL_NAME, max_model_len=1024)
217 + sampling_params = SamplingParams(temperature=0.8,
218 + top_p=0.95,
219 + guided_decoding=GuidedDecodingParams(
220 + regex=sample_regex,
221 + backend=guided_decoding_backend))
222 + with pytest.raises(ValueError,
223 + match="Regex guided decoding is not supported."):
224 + llm.generate(prompts=[
225 + f"Give an example IPv4 address with this regex: {sample_regex}"
226 + ] * 2,
227 + sampling_params=sampling_params,
228 + use_tqdm=True)
229 +
230 + # Once regex is supported --
231 + #assert outputs is not None
232 + #for output in outputs:
233 + # assert output is not None
234 + # assert isinstance(output, RequestOutput)
235 + # prompt = output.prompt
236 + # generated_text = output.outputs[0].text
237 + # print(generated_text)
238 + # assert generated_text is not None
239 + # assert re.fullmatch(sample_regex, generated_text) is not None
240 + # print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
241 +
242 +
243 +@pytest.mark.skip_global_cleanup
244 +@pytest.mark.parametrize("guided_decoding_backend",
245 + GUIDED_DECODING_BACKENDS_V1)
246 +def test_guided_choice_completion(monkeypatch, sample_guided_choice,
247 + guided_decoding_backend: str):
248 + monkeypatch.setenv("VLLM_USE_V1", "1")
249 + llm = LLM(model=MODEL_NAME, max_model_len=1024)
250 + sampling_params = SamplingParams(temperature=0.8,
251 + top_p=0.95,
252 + guided_decoding=GuidedDecodingParams(
253 + choice=sample_guided_choice,
254 + backend=guided_decoding_backend))
255 + outputs = llm.generate(
256 + prompts="The best language for type-safe systems programming is ",
257 + sampling_params=sampling_params,
258 + use_tqdm=True)
259 +
260 + assert outputs is not None
261 + for output in outputs:
262 + assert output is not None
263 + assert isinstance(output, RequestOutput)
264 + prompt = output.prompt
265 + generated_text = output.outputs[0].text
266 + print(generated_text)
267 + assert generated_text is not None
268 + assert generated_text in sample_guided_choice
269 + print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
tests/v1/structured_output/__init__.py
tests/v1/structured_output/test_utils.py new
+196
@@ -0,0 +1,196 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +import pytest
4 +
5 +from vllm.v1.structured_output.utils import (
6 + has_xgrammar_unsupported_json_features)
7 +
8 +
9 +@pytest.fixture
10 +def unsupported_string_schemas():
11 + return [
12 + {
13 + "type": "string",
14 + "pattern": "^[a-zA-Z]+$"
15 + },
16 + {
17 + "type": "string",
18 + "enum": ["active", "inactive", "pending"]
19 + },
20 + {
21 + "type": "string",
22 + "minLength": 1
23 + },
24 + {
25 + "type": "string",
26 + "maxLength": 100
27 + },
28 + {
29 + "type": "string",
30 + "format": "email"
31 + },
32 + ]
33 +
34 +
35 +@pytest.fixture
36 +def unsupported_integer_schemas():
37 + return [
38 + {
39 + "type": "integer",
40 + "minimum": 0
41 + },
42 + {
43 + "type": "integer",
44 + "maximum": 120
45 + },
46 + {
47 + "type": "integer",
48 + "exclusiveMinimum": 120
49 + },
50 + {
51 + "type": "integer",
52 + "exclusiveMaximum": 120
53 + },
54 + {
55 + "type": "integer",
56 + "multipleOf": 120
57 + },
58 + ]
59 +
60 +
61 +@pytest.fixture
62 +def unsupported_number_schemas():
63 + return [
64 + {
65 + "type": "number",
66 + "minimum": 0
67 + },
68 + {
69 + "type": "number",
70 + "maximum": 120
71 + },
72 + {
73 + "type": "number",
74 + "exclusiveMinimum": 120
75 + },
76 + {
77 + "type": "number",
78 + "exclusiveMaximum": 120
79 + },
80 + {
81 + "type": "number",
82 + "multipleOf": 120
83 + },
84 + ]
85 +
86 +
87 +@pytest.fixture
88 +def unsupported_array_schemas():
89 + return [
90 + {
91 + "type": "array",
92 + "uniqueItems": True
93 + },
94 + {
95 + "type": "array",
96 + "contains": {
97 + "type": "string"
98 + }
99 + },
100 + {
101 + "type": "array",
102 + "minContains": 1
103 + },
104 + {
105 + "type": "array",
106 + "maxContains": 5
107 + },
108 + {
109 + "type": "array",
110 + "minItems": 1
111 + },
112 + {
113 + "type": "array",
114 + "maxItems": 10
115 + },
116 + ]
117 +
118 +
119 +@pytest.fixture
120 +def unsupported_object_schemas():
121 + return [
122 + {
123 + "type": "object",
124 + "minProperties": 1
125 + },
126 + {
127 + "type": "object",
128 + "maxProperties": 5
129 + },
130 + {
131 + "type": "object",
132 + "propertyNames": {
133 + "pattern": "^[a-z]+$"
134 + }
135 + },
136 + {
137 + "type": "object",
138 + "patternProperties": {
139 + "^S": {
140 + "type": "string"
141 + }
142 + }
143 + },
144 + ]
145 +
146 +
147 +@pytest.fixture
148 +def supported_schema():
149 + return {
150 + "type": "object",
151 + "properties": {
152 + "name": {
153 + "type": "string"
154 + },
155 + "age": {
156 + "type": "integer"
157 + },
158 + "status": {
159 + "type": "string"
160 + },
161 + "scores": {
162 + "type": "array",
163 + "items": {
164 + "type": "number"
165 + }
166 + },
167 + "address": {
168 + "type": "object",
169 + "properties": {
170 + "street": {
171 + "type": "string"
172 + },
173 + "city": {
174 + "type": "string"
175 + }
176 + }
177 + }
178 + }
179 + }
180 +
181 +
182 +@pytest.mark.parametrize("schema_type", [
183 + "unsupported_string_schemas", "unsupported_integer_schemas",
184 + "unsupported_number_schemas", "unsupported_array_schemas",
185 + "unsupported_object_schemas"
186 +])
187 +def test_unsupported_json_features_by_type(schema_type, request):
188 + schemas = request.getfixturevalue(schema_type)
189 + for schema in schemas:
190 + assert has_xgrammar_unsupported_json_features(
191 + schema), f"Schema should be unsupported: {schema}"
192 +
193 +
194 +def test_supported_json_features(supported_schema):
195 + assert not has_xgrammar_unsupported_json_features(
196 + supported_schema), "Schema should be supported"
tests/v1/worker/test_gpu_model_runner.py
+12
@@ -72,6 +72,8 @@ def _schedule_new_request(*req_ids: str) -> SchedulerOutput:
72 num_common_prefix_blocks=0,
73 finished_req_ids=set(),
74 free_encoder_input_ids=[],
75 + structured_output_request_ids={},
76 + grammar_bitmask=None,
77 )
78
79
@@ -135,6 +137,8 @@ def test_update_states_request_finished(model_runner):
137 num_common_prefix_blocks=0,
138 finished_req_ids={req_id},
139 free_encoder_input_ids=[],
140 + structured_output_request_ids={},
141 + grammar_bitmask=None,
142 )
143
144 metadata_before = model_runner.input_batch.sampling_metadata
@@ -165,6 +169,8 @@ def test_update_states_request_resumed(model_runner):
169 num_common_prefix_blocks=0,
170 finished_req_ids=set(),
171 free_encoder_input_ids=[],
172 + structured_output_request_ids={},
173 + grammar_bitmask=None,
174 )
175
176 model_runner._update_states(scheduler_output)
@@ -190,6 +196,8 @@ def test_update_states_request_resumed(model_runner):
196 num_common_prefix_blocks=0,
197 finished_req_ids=set(),
198 free_encoder_input_ids=[],
199 + structured_output_request_ids={},
200 + grammar_bitmask=None,
201 )
202
203 metadata_before = model_runner.input_batch.sampling_metadata
@@ -221,6 +229,8 @@ def test_update_states_no_changes(model_runner):
229 num_common_prefix_blocks=0,
230 finished_req_ids=set(),
231 free_encoder_input_ids=[],
232 + structured_output_request_ids={},
233 + grammar_bitmask=None,
234 )
235
236 metadata_before = model_runner.input_batch.sampling_metadata
@@ -256,6 +266,8 @@ def test_update_states_request_unscheduled(model_runner):
266 num_common_prefix_blocks=0,
267 finished_req_ids=set(),
268 free_encoder_input_ids=[],
269 + structured_output_request_ids={},
270 + grammar_bitmask=None,
271 )
272
273 metadata_before = model_runner._update_states(scheduler_output)
vllm/utils.py
+57 -2
@@ -1,5 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2
3 +from __future__ import annotations
4 +
5 import argparse
6 import asyncio
7 import concurrent
@@ -8,6 +10,7 @@ import datetime
10 import enum
11 import gc
12 import getpass
13 +import importlib
14 import importlib.metadata
15 import importlib.util
16 import inspect
@@ -23,6 +26,7 @@ import tempfile
26 import threading
27 import time
28 import traceback
29 +import types
30 import uuid
31 import warnings
32 import weakref
@@ -982,7 +986,7 @@ def current_stream() -> torch.cuda.Stream:
986 return _current_stream
987
988
985 -def enable_trace_function_call_for_thread(vllm_config: "VllmConfig") -> None:
989 +def enable_trace_function_call_for_thread(vllm_config: VllmConfig) -> None:
990 """Set up function tracing for the current thread,
991 if enabled via the VLLM_TRACE_FUNCTION environment variable
992 """
@@ -1977,7 +1981,7 @@ class MemorySnapshot:
1981 self.non_torch_memory = self.cuda_memory - self.torch_memory
1982 self.timestamp = time.time()
1983
1980 - def __sub__(self, other: "MemorySnapshot") -> "MemorySnapshot":
1984 + def __sub__(self, other: MemorySnapshot) -> MemorySnapshot:
1985 return MemorySnapshot(
1986 torch_peak=self.torch_peak - other.torch_peak,
1987 cuda_memory=self.cuda_memory - other.cuda_memory,
@@ -2306,3 +2310,54 @@ def warn_for_unimplemented_methods(cls: type[T]) -> type[T]:
2310
2311 type.__setattr__(cls, '__init__', wrapped_init)
2312 return cls
2313 +
2314 +
2315 +class LazyLoader(types.ModuleType):
2316 + """
2317 + LazyLoader module borrowed from Tensorflow
2318 + https://github.com/tensorflow/tensorflow/blob/main/tensorflow/python/util/lazy_loader.py
2319 + with a addition of "module caching".
2320 +
2321 + Lazily import a module, mainly to avoid pulling in large dependencies.
2322 + Modules such as `xgrammar` might do additional side effects, so we
2323 + only want to use this when it is needed, delaying all eager effects
2324 + """
2325 +
2326 + def __init__(
2327 + self,
2328 + local_name: str,
2329 + parent_module_globals: dict[str, Any],
2330 + name: str,
2331 + ):
2332 + self._local_name = local_name
2333 + self._parent_module_globals = parent_module_globals
2334 + self._module: types.ModuleType | None = None
2335 +
2336 + super().__init__(str(name))
2337 +
2338 + def _load(self) -> types.ModuleType:
2339 + # Import the target module and insert it into the parent's namespace
2340 + try:
2341 + module = importlib.import_module(self.__name__)
2342 + self._parent_module_globals[self._local_name] = module
2343 + # The additional add to sys.modules
2344 + # ensures library is actually loaded.
2345 + sys.modules[self._local_name] = module
2346 + except ModuleNotFoundError as err:
2347 + raise err from None
2348 +
2349 + # Update this object's dict so that if someone keeps a
2350 + # reference to the LazyLoader, lookups are efficient
2351 + # (__getattr__ is only called on lookups that fail).
2352 + self.__dict__.update(module.__dict__)
2353 + return module
2354 +
2355 + def __getattr__(self, item: Any) -> Any:
2356 + if self._module is None:
2357 + self._module = self._load()
2358 + return getattr(self._module, item)
2359 +
2360 + def __dir__(self) -> list[str]:
2361 + if self._module is None:
2362 + self._module = self._load()
2363 + return dir(self._module)
vllm/v1/core/scheduler.py
+61 -4
@@ -1,5 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2
3 +from __future__ import annotations
4 +
5 import time
6 from collections import deque
7 from collections.abc import Iterable
@@ -18,6 +20,7 @@ from vllm.v1.engine import (EngineCoreEvent, EngineCoreEventType,
20 from vllm.v1.metrics.stats import SchedulerStats
21 from vllm.v1.outputs import ModelRunnerOutput
22 from vllm.v1.request import Request, RequestStatus
23 +from vllm.v1.structured_output import StructuredOutputManager
24
25 logger = init_logger(__name__)
26
@@ -32,12 +35,14 @@ class Scheduler:
35 lora_config: Optional[LoRAConfig],
36 speculative_config: Optional[SpeculativeConfig],
37 log_stats: bool,
38 + structured_output_manager: StructuredOutputManager,
39 ) -> None:
40 self.scheduler_config = scheduler_config
41 self.cache_config = cache_config
42 self.lora_config = lora_config
43 self.speculative_config = speculative_config
44 self.log_stats = log_stats
45 + self.structured_output_manager = structured_output_manager
46
47 # Scheduling constraints.
48 self.max_num_running_reqs = self.scheduler_config.max_num_seqs
@@ -97,7 +102,7 @@ class Scheduler:
102 self.encoder_cache_manager = EncoderCacheManager(
103 cache_size=encoder_cache_size)
104
100 - def schedule(self) -> "SchedulerOutput":
105 + def schedule(self) -> SchedulerOutput:
106 # NOTE(woosuk) on the scheduling algorithm:
107 # There's no "decoding phase" nor "prefill phase" in the scheduler.
108 # Each request just has the num_computed_tokens and
@@ -114,6 +119,14 @@ class Scheduler:
119 scheduled_running_reqs: list[Request] = []
120 preempted_reqs: list[Request] = []
121
122 + # NOTE: structured_output_request_ids maps
123 + # a request's (request that uses structured output)
124 + # request_id to the running request index.
125 + # This will helps us determine to slice the grammar bitmask
126 + # and only applies valid mask for requests that
127 + # uses structured decoding.
128 + structured_output_request_ids: dict[str, int] = {}
129 +
130 req_to_new_block_ids: dict[str, list[int]] = {}
131 num_scheduled_tokens: dict[str, int] = {}
132 token_budget = self.max_num_scheduled_tokens
@@ -184,6 +197,12 @@ class Scheduler:
197 # Schedule the request.
198 scheduled_running_reqs.append(request)
199 self.scheduled_req_ids.add(request.request_id)
200 + if request.use_structured_output:
201 + # PERF: in case of chunked prefill,
202 + # request might not include any new tokens.
203 + # Therefore, we might introduce some additional
204 + # cycle to fill in the bitmask, which could be a big no-op.
205 + structured_output_request_ids[request.request_id] = req_index
206 req_to_new_block_ids[request.request_id] = [
207 b.block_id for b in new_blocks
208 ]
@@ -219,6 +238,10 @@ class Scheduler:
238 if req.lora_request and req.lora_request.lora_int_id > 0)
239 assert len(requested_loras) <= self.lora_config.max_loras
240
241 + # Use a temporary deque to collect requests that need to be skipped
242 + # and put back at the head of the waiting queue later
243 + waiting_for_fsm: deque[Request] = deque()
244 +
245 # Next, schedule the WAITING requests.
246 if not preempted_reqs:
247 while self.waiting and token_budget > 0:
@@ -227,6 +250,16 @@ class Scheduler:
250
251 request = self.waiting[0]
252
253 + if request.status == RequestStatus.WAITING_FOR_FSM:
254 + structured_output_req = request.structured_output_request
255 + if structured_output_req and structured_output_req.grammar:
256 + request.status = RequestStatus.WAITING
257 + else:
258 + waiting_structured_output_req = self.waiting.popleft()
259 + waiting_for_fsm.appendleft(
260 + waiting_structured_output_req)
261 + continue
262 +
263 # Check that adding the request still respects the max_loras
264 # constraint.
265 if self.lora_config and request.lora_request:
@@ -281,6 +314,10 @@ class Scheduler:
314 break
315
316 self.waiting.popleft()
317 + if request.use_structured_output:
318 + structured_output_request_ids[
319 + request.request_id] = req_index
320 + req_index += 1
321 self.running.append(request)
322 self.scheduled_req_ids.add(request.request_id)
323 self.request_scheduled(request, scheduled_timestamp)
@@ -311,6 +348,10 @@ class Scheduler:
348 self.encoder_cache_manager.allocate(request, i)
349 encoder_budget = new_encoder_budget
350
351 + # Put back any skipped requests at the head of the waiting queue
352 + if waiting_for_fsm:
353 + self.waiting.extendleft(waiting_for_fsm)
354 +
355 # Check if the scheduling constraints are satisfied.
356 total_num_scheduled_tokens = sum(num_scheduled_tokens.values())
357 assert total_num_scheduled_tokens <= self.max_num_scheduled_tokens
@@ -331,6 +372,11 @@ class Scheduler:
372 self.kv_cache_manager.get_num_common_prefix_blocks(
373 any_request, len(self.running)))
374
375 + grammar_bitmask = self.structured_output_manager.grammar_bitmask(
376 + self.requests,
377 + structured_output_request_ids,
378 + len(self.running),
379 + )
380 # Construct the scheduler output.
381 new_reqs_data = [
382 NewRequestData.from_request(req,
@@ -369,6 +415,8 @@ class Scheduler:
415 # the previous and the current steps.
416 finished_req_ids=self.finished_req_ids,
417 free_encoder_input_ids=self.encoder_cache_manager.get_freed_ids(),
418 + structured_output_request_ids=structured_output_request_ids,
419 + grammar_bitmask=grammar_bitmask,
420 )
421
422 self.finished_req_ids = set()
@@ -381,7 +429,7 @@ class Scheduler:
429 num_scheduled_spec_tokens: int,
430 new_block_ids: list[int],
431 resumed_from_preemption: bool,
384 - ) -> "CachedRequestData":
432 + ) -> CachedRequestData:
433 # OPTIMIZATION: Cache the CachedRequestData objects to avoid creating
434 # them at each scheduling step.
435 num_computed_tokens = request.num_computed_tokens
@@ -474,8 +522,8 @@ class Scheduler:
522
523 def update_from_output(
524 self,
477 - scheduler_output: "SchedulerOutput",
478 - model_runner_output: "ModelRunnerOutput",
525 + scheduler_output: SchedulerOutput,
526 + model_runner_output: ModelRunnerOutput,
527 ) -> EngineCoreOutputs:
528 sampled_token_ids = model_runner_output.sampled_token_ids
529 spec_token_ids = model_runner_output.spec_token_ids
@@ -565,6 +613,15 @@ class Scheduler:
613 # the outer lists can be of length > 1.
614 new_logprobs = logprobs.slice(req_index, req_index + 1)
615
616 + if new_token_ids and request.use_structured_output:
617 + # NOTE: structured_output_request
618 + # should not be None if use_structured_output, we have
619 + # check above, so safe to ignore type warning
620 + request.structured_output_request.grammar.accept_tokens( # type: ignore[union-attr]
621 + request.request_id,
622 + new_token_ids,
623 + )
624 +
625 # Transmit partial if chunked prefill & prompt logprobs is enabled
626 if new_token_ids or prompt_logprobs_tensors is not None:
627 # Add EngineCoreOutput for this Request.
vllm/v1/core/scheduler_output.py
+19 -8
@@ -1,9 +1,14 @@
1 # SPDX-License-Identifier: Apache-2.0
2
3 +from __future__ import annotations
4 +
5 from dataclasses import dataclass
6 from typing import TYPE_CHECKING, Optional
7
8 if TYPE_CHECKING:
9 + import numpy as np
10 + import numpy.typing as npt
11 +
12 from vllm.lora.request import LoRARequest
13 from vllm.multimodal import MultiModalKwargs
14 from vllm.multimodal.base import PlaceholderRange
@@ -17,20 +22,20 @@ class NewRequestData:
22 req_id: str
23 prompt_token_ids: list[int]
24 prompt: Optional[str]
20 - mm_inputs: list["MultiModalKwargs"]
25 + mm_inputs: list[MultiModalKwargs]
26 mm_hashes: list[str]
22 - mm_positions: list["PlaceholderRange"]
23 - sampling_params: "SamplingParams"
27 + mm_positions: list[PlaceholderRange]
28 + sampling_params: SamplingParams
29 block_ids: list[int]
30 num_computed_tokens: int
26 - lora_request: Optional["LoRARequest"]
31 + lora_request: Optional[LoRARequest]
32
33 @classmethod
34 def from_request(
35 cls,
31 - request: "Request",
36 + request: Request,
37 block_ids: list[int],
33 - ) -> "NewRequestData":
38 + ) -> NewRequestData:
39 return cls(
40 req_id=request.request_id,
41 prompt_token_ids=request.prompt_token_ids,
@@ -60,11 +65,11 @@ class CachedRequestData:
65 @classmethod
66 def from_request(
67 cls,
63 - request: "Request",
68 + request: Request,
69 resumed_from_preemption: bool,
70 new_token_ids: list[int],
71 new_block_ids: list[int],
67 - ) -> "CachedRequestData":
72 + ) -> CachedRequestData:
73 return cls(
74 req_id=request.request_id,
75 resumed_from_preemption=resumed_from_preemption,
@@ -111,3 +116,9 @@ class SchedulerOutput:
116 # list of (req_id, encoder_input_index) tuples.
117 # Used to free the encoder cache.
118 free_encoder_input_ids: list[tuple[str, int]]
119 +
120 + # Dict of request ids to their index within the batch
121 + # for filling the next token bitmask
122 + structured_output_request_ids: dict[str, int]
123 + # the bitmask for the whole batch
124 + grammar_bitmask: Optional[npt.NDArray[np.int32]]
vllm/v1/engine/async_llm.py
+3 -5
@@ -72,9 +72,7 @@ class AsyncLLM(EngineClient):
72
73 # Processor (converts Inputs --> EngineCoreRequests).
74 self.processor = Processor(
75 - model_config=vllm_config.model_config,
76 - cache_config=vllm_config.cache_config,
77 - lora_config=vllm_config.lora_config,
75 + vllm_config=vllm_config,
76 tokenizer=self.tokenizer,
77 input_registry=input_registry,
78 )
@@ -194,8 +192,8 @@ class AsyncLLM(EngineClient):
192 * 3) Adding the Request to the Detokenizer.
193 * 4) Adding the Request to the EngineCore (separate process).
194
197 - A separate output_handler loop runs in a background AsyncIO task,
198 - pulling outputs from EngineCore and putting them into the
195 + A separate output_handler loop runs in a background AsyncIO task,
196 + pulling outputs from EngineCore and putting them into the
197 per-request AsyncStream.
198
199 The caller of generate() iterates the returned AsyncGenerator,
vllm/v1/engine/core.py
+21 -1
@@ -29,6 +29,7 @@ from vllm.v1.executor.abstract import Executor
29 from vllm.v1.outputs import ModelRunnerOutput
30 from vllm.v1.request import Request, RequestStatus
31 from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder
32 +from vllm.v1.structured_output import StructuredOutputManager
33 from vllm.version import __version__ as VLLM_VERSION
34
35 logger = init_logger(__name__)
@@ -61,6 +62,8 @@ class EngineCore:
62 vllm_config.cache_config.num_gpu_blocks = num_gpu_blocks
63 vllm_config.cache_config.num_cpu_blocks = num_cpu_blocks
64
65 + self.structured_output_manager = StructuredOutputManager(vllm_config)
66 +
67 # Setup scheduler.
68 self.scheduler = Scheduler(
69 scheduler_config=vllm_config.scheduler_config,
@@ -69,6 +72,7 @@ class EngineCore:
72 lora_config=vllm_config.lora_config,
73 speculative_config=vllm_config.speculative_config,
74 log_stats=self.log_stats,
75 + structured_output_manager=self.structured_output_manager,
76 )
77
78 # Setup MM Input Mapper.
@@ -131,6 +135,9 @@ class EngineCore:
135 request.mm_inputs, request.mm_hashes)
136
137 req = Request.from_engine_core_request(request)
138 + if req.use_structured_output:
139 + # Start grammar compilation asynchronously
140 + self.structured_output_manager.populate_cache(req)
141
142 self.scheduler.add_request(req)
143
@@ -148,11 +155,24 @@ class EngineCore:
155
156 if not self.scheduler.has_unfinished_requests():
157 return EngineCoreOutputs(
151 - outputs=[], scheduler_stats=self.scheduler.make_stats())
158 + outputs=[],
159 + scheduler_stats=self.scheduler.make_stats(),
160 + )
161 scheduler_output = self.scheduler.schedule()
162 +
163 + # This case may occur when the only unfinished requests are
164 + # structured output requests where the grammar has not finished
165 + # compiling yet, so there's nothing to run.
166 + if scheduler_output.total_num_scheduled_tokens == 0:
167 + return EngineCoreOutputs(
168 + outputs=[],
169 + scheduler_stats=self.scheduler.make_stats(),
170 + )
171 +
172 output = self.model_executor.execute_model(scheduler_output)
173 engine_core_outputs = self.scheduler.update_from_output(
174 scheduler_output, output) # type: ignore
175 +
176 return engine_core_outputs
177
178 def step_with_batch_queue(self) -> Optional[EngineCoreOutputs]:
vllm/v1/engine/llm_engine.py
+1 -3
@@ -66,9 +66,7 @@ class LLMEngine:
66 self.tokenizer.ping()
67
68 # Processor (convert Inputs --> EngineCoreRequests)
69 - self.processor = Processor(model_config=vllm_config.model_config,
70 - cache_config=vllm_config.cache_config,
71 - lora_config=vllm_config.lora_config,
69 + self.processor = Processor(vllm_config=vllm_config,
70 tokenizer=self.tokenizer,
71 input_registry=input_registry,
72 mm_registry=mm_registry)
vllm/v1/engine/processor.py
+33 -14
@@ -4,7 +4,7 @@ import time
4 from collections.abc import Mapping
5 from typing import Optional, Union
6
7 -from vllm.config import CacheConfig, LoRAConfig, ModelConfig
7 +from vllm.config import VllmConfig
8 from vllm.inputs import (INPUT_REGISTRY, InputRegistry, ProcessorInputs,
9 PromptType, SingletonInputsAdapter)
10 from vllm.inputs.parse import is_encoder_decoder_inputs
@@ -19,39 +19,41 @@ from vllm.sampling_params import SamplingParams
19 from vllm.transformers_utils.tokenizer_group import BaseTokenizerGroup
20 from vllm.v1.engine import EngineCoreRequest
21 from vllm.v1.engine.mm_input_cache import MMInputCacheClient
22 +from vllm.v1.structured_output.utils import validate_structured_output_request
23
24
25 class Processor:
26
27 def __init__(
28 self,
28 - model_config: ModelConfig,
29 - cache_config: CacheConfig,
30 - lora_config: Optional[LoRAConfig],
29 + vllm_config: VllmConfig,
30 tokenizer: BaseTokenizerGroup,
31 input_registry: InputRegistry = INPUT_REGISTRY,
32 mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY,
33 ):
34
36 - self.model_config = model_config
37 - self.cache_config = cache_config
38 - self.lora_config = lora_config
35 + self.vllm_config = vllm_config
36 + self.model_config = vllm_config.model_config
37 + self.cache_config = vllm_config.cache_config
38 + self.lora_config = vllm_config.lora_config
39 + self.decoding_config = vllm_config.decoding_config
40 self.tokenizer = tokenizer
41
41 - self.generation_config_fields = model_config.try_get_generation_config(
42 - )
43 - self.input_preprocessor = InputPreprocessor(model_config,
42 + self.generation_config_fields = (
43 + self.model_config.try_get_generation_config())
44 + self.input_preprocessor = InputPreprocessor(self.model_config,
45 self.tokenizer,
46 mm_registry)
47 self.input_processor = input_registry.create_input_processor(
47 - model_config)
48 + self.model_config)
49
50 # Multi-modal (huggingface) input mapper
50 - self.mm_input_cache_client = MMInputCacheClient(model_config)
51 + self.mm_input_cache_client = MMInputCacheClient(self.model_config)
52
53 # Multi-modal hasher (for images)
53 - self.use_hash = (not model_config.disable_mm_preprocessor_cache) or \
54 - cache_config.enable_prefix_caching
54 + self.use_hash = (
55 + not self.model_config.disable_mm_preprocessor_cache) or \
56 + self.cache_config.enable_prefix_caching
57
58 def _validate_logprobs(
59 self,
@@ -80,6 +82,8 @@ class Processor:
82 self,
83 params: SamplingParams,
84 ) -> None:
85 + self._validate_structured_output(params)
86 +
87 if params.allowed_token_ids is None:
88 return
89 if not params.allowed_token_ids:
@@ -125,6 +129,21 @@ class Processor:
129 raise ValueError(f"Got lora_request {lora_request} but LoRA is "
130 "not enabled!")
131
132 + def _validate_structured_output(self, params: SamplingParams) -> None:
133 + if not params.guided_decoding or not self.decoding_config:
134 + return
135 + if self.decoding_config.guided_decoding_backend != "xgrammar":
136 + raise ValueError(
137 + "Only xgrammar structured output is supported in V1.")
138 + if (params.guided_decoding.backend
139 + and params.guided_decoding.backend != 'xgrammar'):
140 + raise ValueError(
141 + "Only xgrammar structured output is supported in V1.")
142 + if self.vllm_config.speculative_config:
143 + raise ValueError("Structured output is not supported with "
144 + "speculative decoding.")
145 + validate_structured_output_request(params)
146 +
147 def process_inputs(
148 self,
149 request_id: str,
vllm/v1/request.py
+24 -11
@@ -3,13 +3,15 @@
3 import enum
4 from typing import TYPE_CHECKING, Optional, Union
5
6 -from vllm.lora.request import LoRARequest
6 from vllm.sampling_params import SamplingParams
7 from vllm.v1.engine import (EngineCoreEvent, EngineCoreEventType,
8 EngineCoreRequest, FinishReason)
9 +from vllm.v1.structured_output.request import StructuredOutputRequest
10 from vllm.v1.utils import ConstantList
11
12 if TYPE_CHECKING:
13 +
14 + from vllm.lora.request import LoRARequest
15 from vllm.multimodal import MultiModalKwargs
16 from vllm.multimodal.inputs import PlaceholderRange
17
@@ -27,15 +29,19 @@ class Request:
29 sampling_params: SamplingParams,
30 eos_token_id: Optional[int],
31 arrival_time: float,
30 - lora_request: Optional[LoRARequest] = None,
32 + lora_request: Optional["LoRARequest"] = None,
33 + structured_output_request: Optional["StructuredOutputRequest"] = None,
34 ) -> None:
35 self.request_id = request_id
36 self.sampling_params = sampling_params
37 # Because of LoRA, the eos token id can be different for each request.
38 self.eos_token_id = eos_token_id
39 self.lora_request = lora_request
40 + self.structured_output_request = structured_output_request
41
38 - self.status = RequestStatus.WAITING
42 + self.status = (RequestStatus.WAITING_FOR_FSM
43 + if sampling_params.guided_decoding is not None else
44 + RequestStatus.WAITING)
45 self.events: list[EngineCoreEvent] = []
46 self.stop_reason: Union[int, str, None] = None
47 assert sampling_params.max_tokens is not None
@@ -78,6 +84,8 @@ class Request:
84 eos_token_id=request.eos_token_id,
85 arrival_time=request.arrival_time,
86 lora_request=request.lora_request,
87 + structured_output_request=StructuredOutputRequest(
88 + sampling_params=request.sampling_params),
89 )
90
91 def queued(self, timestamp: Optional[float] = None) -> None:
@@ -134,18 +142,23 @@ class Request:
142 num_tokens = self.mm_positions[input_id]["length"]
143 return num_tokens
144
145 + @property
146 + def use_structured_output(self) -> bool:
147 + return self.sampling_params.guided_decoding is not None
148 +
149
150 class RequestStatus(enum.IntEnum):
151 """Status of a request."""
140 - WAITING = 0
141 - RUNNING = 1
142 - PREEMPTED = 2
143 - # Note: anything after PREEMPTED (2) will be considered
152 + WAITING = enum.auto()
153 + WAITING_FOR_FSM = enum.auto()
154 + RUNNING = enum.auto()
155 + PREEMPTED = enum.auto()
156 + # Note: anything after PREEMPTED will be considered
157 # as a finished status.
145 - FINISHED_STOPPED = 3
146 - FINISHED_LENGTH_CAPPED = 4
147 - FINISHED_ABORTED = 5
148 - FINISHED_IGNORED = 6
158 + FINISHED_STOPPED = enum.auto()
159 + FINISHED_LENGTH_CAPPED = enum.auto()
160 + FINISHED_ABORTED = enum.auto()
161 + FINISHED_IGNORED = enum.auto()
162
163 @staticmethod
164 def is_finished(status: "RequestStatus") -> bool:
vllm/v1/structured_output/__init__.py new
+152
@@ -0,0 +1,152 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +from __future__ import annotations
3 +
4 +import copy
5 +import multiprocessing
6 +from collections import OrderedDict
7 +from concurrent.futures import ThreadPoolExecutor
8 +from typing import TYPE_CHECKING, Optional
9 +
10 +from vllm.config import VllmConfig
11 +from vllm.logger import init_logger
12 +from vllm.transformers_utils.tokenizer_group import init_tokenizer_from_configs
13 +from vllm.utils import LazyLoader
14 +from vllm.v1.structured_output.grammar import (Grammar, StructuredOutputKey,
15 + StructuredOutputOptions)
16 +
17 +if TYPE_CHECKING:
18 + import numpy as np
19 + import numpy.typing as npt
20 + import xgrammar as xgr
21 +
22 + from vllm.v1.request import Request
23 +else:
24 + xgr = LazyLoader("xgr", globals(), "xgrammar")
25 +
26 +logger = init_logger(__name__)
27 +
28 +
29 +class StructuredOutputManager:
30 +
31 + def __init__(self, vllm_config: VllmConfig, max_cache_size: int = 500):
32 + tokenizer_group = init_tokenizer_from_configs(
33 + model_config=vllm_config.model_config,
34 + scheduler_config=vllm_config.scheduler_config,
35 + parallel_config=vllm_config.parallel_config,
36 + lora_config=vllm_config.lora_config) # type: ignore[arg-type]
37 + tokenizer_group.ping()
38 + self.vocab_size = vllm_config.model_config.get_vocab_size()
39 + self.vllm_config = vllm_config
40 +
41 + tokenizer = tokenizer_group.get_lora_tokenizer(None)
42 + tokenizer_info = xgr.TokenizerInfo.from_huggingface(
43 + tokenizer, vocab_size=self.vocab_size)
44 + self.compiler = xgr.GrammarCompiler(tokenizer_info, max_threads=8)
45 +
46 + self.max_cache_size = max_cache_size
47 + self.request_key_to_grammar: OrderedDict[StructuredOutputKey,
48 + Grammar] = OrderedDict()
49 +
50 + # The default max_workers if not specified is the number of CPUs * 5,
51 + # which is way too high since these tasks are CPU-bound, not I/O bound.
52 + # We also know we would never dominate CPU usage with just grammar
53 + # compilation, so we set it to half the number of CPUs.
54 + max_workers = max(1, (multiprocessing.cpu_count() + 1) // 2)
55 + self.executor = ThreadPoolExecutor(max_workers=max_workers)
56 + self._grammar_bitmask = xgr.allocate_token_bitmask(
57 + self.vllm_config.scheduler_config.max_num_seqs, self.vocab_size)
58 +
59 + def __getitem__(self, key: StructuredOutputKey) -> Optional[Grammar]:
60 + # We need to pop and re-insert the grammar here for LRU cache
61 + # of request_key_to_grammar
62 + if key in self.request_key_to_grammar:
63 + # Move accessed item to the end (most recently used)
64 + value = self.request_key_to_grammar.pop(key)
65 + if value is not None:
66 + self.request_key_to_grammar[key] = value
67 + return value
68 + return None
69 +
70 + def populate_cache(self, request: Request) -> None:
71 + if request.structured_output_request is None:
72 + return
73 +
74 + grammar = self.request_key_to_grammar.get(
75 + request.structured_output_request.structured_output_key)
76 + if grammar:
77 + request.structured_output_request.grammar = copy.copy(grammar)
78 + return
79 + request.structured_output_request.grammar = self.cache(request)
80 +
81 + def cache(self, request: Request):
82 + return self.executor.submit(self._executor_loop, request)
83 +
84 + def _executor_loop(self, request: Request) -> Grammar:
85 + # NOTE: The structured_output_request should never be
86 + # None in this case, but mypy can't infer this
87 + # correctly, so we need to ignore the error here.
88 + key = request.structured_output_request.structured_output_key # type: ignore[union-attr]
89 + grammar = self.request_key_to_grammar.get(key)
90 + if grammar is not None:
91 + return copy.copy(grammar)
92 + grammar = self.initialize_grammar(key)
93 + # If cache is full, remove the least recently used item
94 + if len(self.request_key_to_grammar) >= self.max_cache_size:
95 + self.request_key_to_grammar.popitem(last=False)
96 + self.request_key_to_grammar[key] = grammar
97 + return copy.copy(grammar)
98 +
99 + def initialize_grammar(self, key: StructuredOutputKey) -> Grammar:
100 + # Note that the request was validated in the engine core client,
101 + # so at this point we know it is a supported type of request.
102 + #
103 + # TODO: we still need to handle xgrammar compilation failures
104 + request_type, grammar_spec = key
105 +
106 + if request_type == StructuredOutputOptions.JSON:
107 + # TODO -- allow any_whitespace to be configurable
108 + # pending merge of https://github.com/vllm-project/vllm/pull/12744
109 + ctx = self.compiler.compile_json_schema(grammar_spec,
110 + any_whitespace=False)
111 + elif request_type == StructuredOutputOptions.JSON_OBJECT:
112 + ctx = self.compiler.compile_builtin_json_grammar()
113 + elif request_type == StructuredOutputOptions.GRAMMAR:
114 + ctx = self.compiler.compile_grammar(grammar_spec)
115 + else:
116 + logger.error("Validation should have already occurred. "
117 + "Please file an issue.")
118 + raise ValueError(
119 + f"grammar is not of valid supported types. ({request_type!s})")
120 +
121 + return Grammar(
122 + matcher=xgr.GrammarMatcher(ctx),
123 + vocab_size=self.vocab_size,
124 + ctx=ctx,
125 + )
126 +
127 + def grammar_bitmask(
128 + self,
129 + requests: dict[str, Request],
130 + structured_output_request_ids: dict[str, int],
131 + batch_len: int,
132 + ) -> Optional[npt.NDArray[np.int32]]:
133 + # Prepare the structured output bitmask for this batch.
134 + if not structured_output_request_ids:
135 + return None
136 +
137 + # Fill the bitmask using the index of each request equal to its
138 + # position in the batch. Resize the bitmask down to the size of
139 + # the batch.
140 + bitmask_tensor = self._grammar_bitmask
141 + for req_id, batch_index in structured_output_request_ids.items():
142 + request = requests[req_id].structured_output_request
143 + assert request is not None and request.grammar is not None
144 + if not request.grammar.matcher.is_terminated():
145 + request.grammar.fill_bitmask(bitmask_tensor, batch_index)
146 + if batch_len < self._grammar_bitmask.shape[0]:
147 + bitmask_tensor = self._grammar_bitmask[:batch_len]
148 +
149 + # After finishing with the xgrammar operations, we convert to
150 + # np.ndarray, because that is much more efficient for serialization
151 + # and deserialization when sending this to the GPU workers.
152 + return bitmask_tensor.numpy()
vllm/v1/structured_output/grammar.py new
+77
@@ -0,0 +1,77 @@
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 new
+71
@@ -0,0 +1,71 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +from __future__ import annotations
3 +
4 +import dataclasses
5 +import functools
6 +import json
7 +from concurrent.futures import Future
8 +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)
14 +
15 +
16 +@dataclasses.dataclass
17 +class StructuredOutputRequest:
18 +
19 + sampling_params: SamplingParams
20 + _grammar: Optional[Union[Future[Grammar], Grammar]] = None
21 +
22 + def _check_grammar_completion(self) -> bool:
23 + # NOTE: We have to lazy import to gate circular imports
24 + from vllm.v1.request import RequestStatus
25 +
26 + if isinstance(self._grammar, Future):
27 + try:
28 + # We will check whether the future is ready within 100 us
29 + self._grammar = self._grammar.result(timeout=0.0001)
30 + self.status = RequestStatus.WAITING
31 + except TimeoutError:
32 + return False
33 + return True
34 +
35 + @property
36 + def is_grammar_ready(self) -> bool:
37 + return self._check_grammar_completion()
38 +
39 + @property
40 + def grammar(self) -> Optional[Grammar]:
41 + completed = self._check_grammar_completion()
42 + return cast(Optional[Grammar], self._grammar) if completed else None
43 +
44 + @grammar.setter
45 + def grammar(self, grammar: Union[Grammar, Future[Grammar]]) -> None:
46 + self._grammar = grammar
47 +
48 + @functools.cached_property
49 + def structured_output_key(self) -> StructuredOutputKey:
50 + params = self.sampling_params.guided_decoding
51 + assert params is not None, "params can't be None."
52 + if params.json is not None:
53 + if not isinstance(params.json, str):
54 + json_str = json.dumps(params.json)
55 + else:
56 + json_str = params.json
57 + return (StructuredOutputOptions.JSON, json_str)
58 + elif params.json_object:
59 + return (StructuredOutputOptions.JSON_OBJECT, "")
60 + elif params.regex is not None:
61 + return (StructuredOutputOptions.REGEX, params.regex)
62 + elif params.choice is not None:
63 + if not isinstance(params.choice, str):
64 + json_str = json.dumps(params.choice)
65 + else:
66 + json_str = params.choice
67 + return (StructuredOutputOptions.CHOICE, json_str)
68 + elif params.grammar is not None:
69 + return (StructuredOutputOptions.GRAMMAR, params.grammar)
70 + else:
71 + raise ValueError("No valid structured output parameter found")
vllm/v1/structured_output/utils.py new
+295
@@ -0,0 +1,295 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +from __future__ import annotations
4 +
5 +import json
6 +import re
7 +from typing import TYPE_CHECKING, Any
8 +
9 +from vllm.sampling_params import SamplingParams
10 +from vllm.utils import LazyLoader
11 +
12 +if TYPE_CHECKING:
13 + import xgrammar as xgr
14 +else:
15 + xgr = LazyLoader("xgr", globals(), "xgrammar")
16 +
17 +
18 +def has_xgrammar_unsupported_json_features(schema: dict[str, Any]) -> bool:
19 + """Check if JSON schema contains features unsupported by xgrammar."""
20 +
21 + def check_object(obj: dict[str, Any]) -> bool:
22 + if not isinstance(obj, dict):
23 + return False
24 +
25 + # Check for pattern restrictions
26 + if "pattern" in obj:
27 + return True
28 +
29 + # Check for enum restrictions
30 + if "enum" in obj:
31 + return True
32 +
33 + # Check for numeric ranges
34 + if obj.get("type") in ("integer", "number") and any(
35 + key in obj
36 + for key in ("minimum", "maximum", "exclusiveMinimum",
37 + "exclusiveMaximum", "multipleOf")):
38 + return True
39 +
40 + # Check for array unsupported keywords
41 + if obj.get("type") == "array" and any(
42 + key in obj
43 + for key in ("uniqueItems", "contains", "minContains",
44 + "maxContains", "minItems", "maxItems")):
45 + return True
46 +
47 + # Unsupported keywords for strings
48 + if obj.get("type") == "string" and any(
49 + key in obj for key in ("minLength", "maxLength", "format")):
50 + return True
51 +
52 + # Unsupported keywords for objects
53 + if obj.get("type") == "object" and any(
54 + key in obj for key in ("minProperties", "maxProperties",
55 + "propertyNames", "patternProperties")):
56 + return True
57 +
58 + # Recursively check all nested objects and arrays
59 + for value in obj.values():
60 + if isinstance(value, dict):
61 + if check_object(value):
62 + return True
63 + elif isinstance(value, list):
64 + for item in value:
65 + if isinstance(item, dict) and check_object(item):
66 + return True
67 +
68 + return False
69 +
70 + return check_object(schema)
71 +
72 +
73 +def grammar_is_likely_lark(grammar_str: str) -> bool:
74 + """
75 + Check if grammar appears to use Lark syntax.
76 +
77 + Args:
78 + grammar_str: Input grammar string
79 +
80 + Returns:
81 + bool: True if grammar appears to be in Lark format, False otherwise
82 +
83 + Examples:
84 + >>> grammar_is_likely_lark("rule: 'abc'")
85 + True
86 + >>> grammar_is_likely_lark("rule ::= 'abc'")
87 + False
88 + """
89 + if not grammar_str or not isinstance(grammar_str, str):
90 + return False
91 +
92 + for line in grammar_str.split('\n'):
93 + # Remove both comment styles
94 + line = re.sub(r'(#|//).*$', '', line).strip()
95 + if not line:
96 + continue
97 +
98 + # Look for EBNF rule definition
99 + if '::=' in line:
100 + return False
101 +
102 + return True
103 +
104 +
105 +def convert_lark_to_ebnf(grammar_str: str) -> str:
106 + """
107 + Convert a Lark grammar string to EBNF format.
108 +
109 + EBNF reference:
110 + https://github.com/ggerganov/llama.cpp/blob/master/grammars/README.md
111 + Lark grammar reference:
112 + https://lark-parser.readthedocs.io/en/latest/grammar.html
113 +
114 + Args:
115 + grammar_str: Input grammar in Lark format
116 +
117 + Returns:
118 + str: Converted grammar in EBNF format
119 +
120 + Examples:
121 + >>> print(convert_lark_to_ebnf("rule: 'hello'"))
122 + root ::= rule
123 + rule ::= "hello"
124 + """
125 + if not isinstance(grammar_str, str):
126 + raise ValueError(f"Grammar must be a string, got {type(grammar_str)}")
127 + if not grammar_str.strip():
128 + raise ValueError("Grammar string cannot be empty")
129 +
130 + defined_rules = set()
131 + referenced_rules = set()
132 + output_lines = []
133 +
134 + def clean_line(line: str) -> str:
135 + """Remove comments and whitespace from line."""
136 + return re.sub(r'(#|//).*$', '', line).strip()
137 +
138 + def check_quotes(text: str, rule_name: str, line_num: int) -> None:
139 + """Validate quote matching in text."""
140 + if text.count("'") % 2 != 0 or text.count('"') % 2 != 0:
141 + raise ValueError(
142 + f"Mismatched quotes in {rule_name} on line {line_num}")
143 +
144 + def extract_references(text: str) -> set:
145 + """Extract rule references from text."""
146 + # Remove quoted strings and special characters
147 + text = re.sub(r'"[^"]*"', '', text)
148 + text = re.sub(r'[+*?()|\[\]{}]', ' ', text)
149 + return set(re.findall(r'\b[a-zA-Z_][a-zA-Z0-9_]*\b', text))
150 +
151 + # First pass: Find root rule and validate rule definitions
152 + lines = [clean_line(line) for line in grammar_str.split('\n')]
153 + first_rule = None
154 +
155 + for line_num, line in enumerate(lines, 1):
156 + if not line or line.startswith('|'):
157 + continue
158 +
159 + if ':' in line:
160 + try:
161 + name = line.split(':', 1)[0].strip().strip('?')
162 + defined_rules.add(name)
163 + if first_rule is None:
164 + first_rule = name
165 + if name == 'start':
166 + first_rule = 'start'
167 + except IndexError as e:
168 + raise ValueError(f"Invalid rule format on line {line_num}. "
169 + "Expected 'rule_name: definition'") from e
170 +
171 + if not defined_rules:
172 + raise ValueError("No valid rules found in grammar")
173 +
174 + # Add root rule
175 + output_lines.append(f"root ::= {first_rule}")
176 +
177 + # Second pass: Process rule definitions and alternatives
178 + current_rule = None
179 + current_definition = []
180 +
181 + for line_num, line in enumerate(lines, 1):
182 + if not line:
183 + continue
184 +
185 + try:
186 + if ':' in line and not line.startswith('|'):
187 + # Save previous rule if exists
188 + if current_rule:
189 + output_lines.append(
190 + f"{current_rule} ::= {' | '.join(current_definition)}")
191 +
192 + # Process new rule
193 + name, definition = line.split(':', 1)
194 + current_rule = name.strip().strip('?')
195 +
196 + check_quotes(definition, f"rule '{current_rule}'", line_num)
197 + definition = re.sub(r"'([^']*)'", r'"\1"', definition)
198 + referenced_rules.update(extract_references(definition))
199 + current_definition = [definition.strip()]
200 +
201 + elif line.startswith('|'):
202 + if not current_rule:
203 + raise ValueError(f"Alternative '|' on line {line_num} "
204 + "without a preceding rule definition")
205 +
206 + alt_def = line[1:].strip()
207 + check_quotes(alt_def, f"alternative for rule '{current_rule}'",
208 + line_num)
209 + alt_def = re.sub(r"'([^']*)'", r'"\1"', alt_def)
210 + referenced_rules.update(extract_references(alt_def))
211 + current_definition.append(alt_def)
212 +
213 + except ValueError as e:
214 + raise ValueError(f"Error on line {line_num}: {str(e)}") from e
215 +
216 + # Add final rule if exists
217 + if current_rule:
218 + output_lines.append(
219 + f"{current_rule} ::= {' | '.join(current_definition)}")
220 +
221 + # Validate all rules are defined
222 + undefined_rules = referenced_rules - defined_rules - {'root'}
223 + if undefined_rules:
224 + raise ValueError("Referenced rules are not defined: "
225 + f"{', '.join(sorted(undefined_rules))}")
226 +
227 + return '\n'.join(output_lines)
228 +
229 +
230 +def choice_as_grammar(choice: list[str]) -> str:
231 +
232 + def escape_ebnf_string(s: str) -> str:
233 + """Escape special characters in a EBNF string."""
234 + # Escape double quotes and backslashes
235 + return re.sub(r'(["\\])', r'\\\1', s)
236 +
237 + escaped_choices = (escape_ebnf_string(c) for c in choice)
238 + grammar = ('root ::= ' + ' | '.join(f'"{c}"' for c in escaped_choices))
239 + return grammar
240 +
241 +
242 +def validate_structured_output_request(
243 + sampling_params: SamplingParams) -> None:
244 + """Validate that the request is supported by structured output.
245 +
246 + Raises ValueError if the request is not supported.
247 + """
248 + if sampling_params.guided_decoding is None:
249 + return
250 +
251 + gd_params = sampling_params.guided_decoding
252 +
253 + if gd_params.regex:
254 + raise ValueError("Regex structured output is not supported.")
255 +
256 + if gd_params.choice:
257 + choice_grammar = choice_as_grammar(gd_params.choice)
258 + try:
259 + xgr.Grammar.from_ebnf(choice_grammar)
260 + except Exception as err:
261 + raise ValueError("Failed to transform choices into a grammar: "
262 + "{err}") from err
263 + gd_params.choice = None
264 + gd_params.grammar = choice_grammar
265 + return
266 +
267 + if gd_params.json:
268 + if isinstance(gd_params.json, str):
269 + try:
270 + schema = json.loads(gd_params.json)
271 + except json.JSONDecodeError as e:
272 + raise ValueError("Invalid JSON grammar specification.") from e
273 + else:
274 + schema = gd_params.json
275 +
276 + if has_xgrammar_unsupported_json_features(schema):
277 + raise ValueError("The provided JSON schema contains features not "
278 + "supported by xgrammar.")
279 + return
280 +
281 + if gd_params.grammar:
282 + if grammar_is_likely_lark(gd_params.grammar):
283 + # xgrammar supports EBNF grammars only
284 + try:
285 + gd_params.grammar = convert_lark_to_ebnf(gd_params.grammar)
286 + except ValueError as e:
287 + raise ValueError(
288 + "Failed to convert the grammar from Lark to EBNF. ") from e
289 +
290 + # Test parsing EBNF grammar, possibly already converted from Lark
291 + try:
292 + # parse the grammar, but we aren't compiling it.
293 + xgr.Grammar.from_ebnf(gd_params.grammar)
294 + except Exception as e:
295 + raise ValueError("Invalid grammar specification.") from e
vllm/v1/worker/gpu_model_runner.py
+57 -1
@@ -25,7 +25,8 @@ from vllm.multimodal.utils import group_mm_inputs_by_modality
25 from vllm.sampling_params import SamplingType
26 from vllm.sequence import IntermediateTensors
27 from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, DeviceMemoryProfiler,
28 - LayerBlockType, cdiv, is_pin_memory_available)
28 + LayerBlockType, LazyLoader, cdiv,
29 + is_pin_memory_available)
30 from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata
31 from vllm.v1.core.encoder_cache_manager import compute_encoder_budget
32 from vllm.v1.engine.mm_input_cache import MMInputCacheClient
@@ -40,7 +41,11 @@ from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
41 from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin
42
43 if TYPE_CHECKING:
44 + import xgrammar as xgr
45 +
46 from vllm.v1.core.scheduler_output import SchedulerOutput
47 +else:
48 + xgr = LazyLoader("xgr", globals(), "xgrammar")
49
50 logger = init_logger(__name__)
51
@@ -860,6 +865,53 @@ class GPUModelRunner(LoRAModelRunnerMixin):
865 def get_model(self) -> nn.Module:
866 return self.model
867
868 + def apply_grammar_bitmask(
869 + self,
870 + scheduler_output: "SchedulerOutput",
871 + logits: torch.Tensor,
872 + ):
873 + # Serialization of np.ndarray is much more efficient than a tensor,
874 + # so we receive it in that format.
875 + grammar_bitmask = scheduler_output.grammar_bitmask
876 + if grammar_bitmask is None:
877 + return
878 +
879 + # We receive the structured output bitmask from the scheduler, but the
880 + # indices of the requests in the batch may not match the indices of
881 + # the bitmask since the scheduler doesn't know how the gpu runner is
882 + # ordering the requests in the batch. We need to sort the bitmask to
883 + # match the order of the requests used here.
884 + struct_out_req_batch_indices: dict[str, int] = {}
885 + indices_match = True
886 + for req_id in self.input_batch.req_ids:
887 + mask_index = scheduler_output.structured_output_request_ids.get(
888 + req_id)
889 + if mask_index is None:
890 + # not a structured output request
891 + continue
892 + batch_index = self.input_batch.req_id_to_index[req_id]
893 + if batch_index != mask_index:
894 + indices_match = False
895 + struct_out_req_batch_indices[req_id] = batch_index
896 +
897 + if not indices_match:
898 + # Sort the bitmask to match the order of the requests
899 + sorted_bitmask = np.zeros_like(grammar_bitmask)
900 + for req_id, batch_index in struct_out_req_batch_indices.items():
901 + orig_index = scheduler_output.structured_output_request_ids[
902 + req_id]
903 + sorted_bitmask[batch_index] = grammar_bitmask[orig_index]
904 + grammar_bitmask = sorted_bitmask
905 +
906 + grammar_bitmask = torch.from_numpy(grammar_bitmask)
907 +
908 + # TODO: compatibility with spec decode
909 + xgr.apply_token_bitmask_inplace(
910 + logits,
911 + grammar_bitmask.to(self.device, non_blocking=True),
912 + indices=list(struct_out_req_batch_indices.values()),
913 + )
914 +
915 @torch.inference_mode()
916 def execute_model(
917 self,
@@ -945,6 +997,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
997 sample_hidden_states = hidden_states[logits_indices]
998 logits = self.model.compute_logits(sample_hidden_states, None)
999
1000 + # Apply structured output bitmasks if present
1001 + if scheduler_output.grammar_bitmask is not None:
1002 + self.apply_grammar_bitmask(scheduler_output, logits)
1003 +
1004 # Sample the next token and get logprobs if needed.
1005 sampling_metadata = self.input_batch.sampling_metadata
1006 if not self.use_spec_decode: