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

A high-throughput and memory-efficient inference and serving engine for LLMs

[Core] Support structured outputs for beam search (#35022)

Signed-off-by: Guan-Ming (Wesley) Chiu <guanmingchiu@gmail.com> Signed-off-by: Guan-Ming (Wesley) Chiu <105915352+guan404ming@users.noreply.github.com>

Guan-Ming (Wesley) Chiu committed Jun 12, 2026 at 21:56 UTC c7aa3d263049ac9eefd0f59a10f5ecc6a78927df
3 files changed +405 -70
tests/samplers/test_beam_search.py
+63
@@ -5,11 +5,16 @@
5 Run `pytest tests/samplers/test_beam_search.py`.
6 """
7
8 +import json
9 +
10 +import jsonschema
11 import pytest
12 from transformers import AutoModelForSeq2SeqLM
13
14 from vllm.assets.audio import AudioAsset
15 +from vllm.entrypoints.llm import LLM
16 from vllm.platforms import current_platform
17 +from vllm.sampling_params import BeamSearchParams, StructuredOutputsParams
18
19 # Extra engine kwargs needed for numerically deterministic beam search.
20 # On ROCm, floating-point reductions in attention and GEMM kernels are
@@ -223,3 +228,61 @@ def test_beam_search_passes_multimodal_data(
228
229 # NOTE: encoder/decoder tests are currently located under
230 # tests/models/multimodal/generation/test_whisper.py
231 +
232 +
233 +@pytest.mark.parametrize("model", MODELS)
234 +@pytest.mark.parametrize("dtype", ["half"])
235 +@pytest.mark.parametrize("beam_width", BEAM_WIDTHS)
236 +def test_beam_search_structured_output(
237 + model: str,
238 + dtype: str,
239 + beam_width: int,
240 +) -> None:
241 + """Ensure beam search with structured output produces valid JSON."""
242 + json_schema = {
243 + "type": "object",
244 + "properties": {
245 + "name": {"type": "string"},
246 + "age": {"type": "integer"},
247 + },
248 + "required": ["name", "age"],
249 + "additionalProperties": False,
250 + }
251 +
252 + llm = LLM(
253 + model=model,
254 + dtype=dtype,
255 + max_model_len=512,
256 + structured_outputs_config=dict(
257 + backend="xgrammar",
258 + disable_any_whitespace=True,
259 + ),
260 + **(dict(enforce_eager=True) | EXTRA_ENGINE_KWARGS),
261 + )
262 +
263 + params = BeamSearchParams(
264 + beam_width=beam_width,
265 + max_tokens=64,
266 + structured_outputs=StructuredOutputsParams(json=json_schema),
267 + )
268 +
269 + prompts = [
270 + "Generate a JSON object for a person with name and age:",
271 + ]
272 +
273 + outputs = llm.beam_search(prompts, params)
274 +
275 + assert len(outputs) == len(prompts)
276 + for output in outputs:
277 + assert len(output.sequences) > 0
278 + for seq in output.sequences:
279 + assert seq.text is not None
280 + print(f"Full text: {seq.text!r}")
281 + # seq.text includes the prompt, extract generated JSON.
282 + gen_start = seq.text.find("{")
283 + assert gen_start != -1, f"No JSON found in output: {seq.text!r}"
284 + generated = seq.text[gen_start:]
285 + generated = generated.replace("</s>", "").strip()
286 + print(f"Generated JSON: {generated!r}")
287 + parsed = json.loads(generated)
288 + jsonschema.validate(instance=parsed, schema=json_schema)
vllm/entrypoints/generate/beam_search/offline.py
+341 -70
@@ -2,14 +2,24 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import itertools
5 +from collections.abc import Callable, Sequence
6
7 +import torch
8 from tqdm import tqdm
9
10 from vllm import RequestOutput, TextPrompt, TokensPrompt
11 from vllm.entrypoints.offline_utils import OfflineInferenceMixin
12 from vllm.logger import init_logger
13 from vllm.lora.request import LoRARequest
12 -from vllm.sampling_params import BeamSearchParams, SamplingParams
14 +from vllm.pooling_params import PoolingParams
15 +from vllm.sampling_params import (
16 + BeamSearchParams,
17 + SamplingParams,
18 + StructuredOutputsParams,
19 +)
20 +from vllm.tokenizers import TokenizerLike
21 +from vllm.v1.structured_output.backend_types import StructuredOutputBackend
22 +from vllm.v1.structured_output.request import get_structured_output_key
23
24 from .utils import (
25 BeamSearchInstance,
@@ -20,6 +30,27 @@ from .utils import (
30
31 logger = init_logger(__name__)
32
33 +# Engine-side cap on `SamplingParams.allowed_token_ids`; keep in sync with
34 +# MAX_NUM_ALLOWED_TOKEN_IDS in vllm/v1/worker/gpu/sample/logit_bias.py.
35 +_MAX_NUM_ALLOWED_TOKEN_IDS = 1024
36 +
37 +
38 +_bitmask_cache: dict[int, tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = {}
39 +
40 +
41 +def _bitmask_to_token_ids(bitmask_row: torch.Tensor, vocab_size: int) -> list[int]:
42 + """Convert a packed int32 bitmask row to a list of allowed token IDs."""
43 + if vocab_size not in _bitmask_cache:
44 + indices = torch.arange(vocab_size)
45 + _bitmask_cache[vocab_size] = (
46 + indices,
47 + indices >> 5, # i // 32
48 + indices & 31, # i % 32
49 + )
50 + indices, word_indices, bit_indices = _bitmask_cache[vocab_size]
51 + mask = ((bitmask_row[word_indices] >> bit_indices) & 1).bool()
52 + return indices[mask].tolist()
53 +
54
55 class BeamSearchOfflineMixin(OfflineInferenceMixin):
56 """Offline inference for beam search"""
@@ -69,10 +100,22 @@ class BeamSearchOfflineMixin(OfflineInferenceMixin):
100 if concurrency_limit is None:
101 concurrency_limit = len(engine_inputs)
102
103 + structured_output_backend: StructuredOutputBackend | None = None
104 + structured_output_key = None
105 + structured_output_bitmask = None
106 + if params.structured_outputs is not None:
107 + (
108 + structured_output_backend,
109 + structured_output_key,
110 + structured_output_bitmask,
111 + ) = self._init_beam_search_structured_output(
112 + params.structured_outputs, tokenizer
113 + )
114 +
115 # generate 2 * beam_width candidates at each step
116 # following the huggingface transformers implementation
117 # at https://github.com/huggingface/transformers/blob/e15687fffe5c9d20598a19aeab721ae0a7580f8a/src/transformers/generation/beam_search.py#L534 # noqa
75 - sampling_params = SamplingParams(
118 + base_sampling_params = SamplingParams(
119 logprobs=2 * beam_width,
120 max_tokens=1,
121 temperature=temperature,
@@ -94,77 +137,43 @@ class BeamSearchOfflineMixin(OfflineInferenceMixin):
137 ),
138 )
139
97 - for prompt_start in range(0, len(instances), concurrency_limit):
98 - instances_batch = instances[prompt_start : prompt_start + concurrency_limit]
140 + try:
141 + for prompt_start in range(0, len(instances), concurrency_limit):
142 + instances_batch = instances[
143 + prompt_start : prompt_start + concurrency_limit
144 + ]
145
100 - token_iter = range(max_tokens)
101 - if use_tqdm:
102 - token_iter = tqdm(
103 - token_iter, desc="Beam search", unit="token", unit_scale=False
104 - )
105 - logger.warning(
106 - "The progress bar shows the upper bound on token steps and "
107 - "may finish early due to stopping conditions. It does not "
108 - "reflect instance-level progress."
109 - )
110 - for _ in token_iter:
111 - all_beams: list[BeamSearchSequence] = list(
112 - sum((instance.beams for instance in instances_batch), [])
113 - )
114 - pos = [0] + list(
115 - itertools.accumulate(
116 - len(instance.beams) for instance in instances_batch
146 + token_iter = range(max_tokens)
147 + if use_tqdm:
148 + token_iter = tqdm(
149 + token_iter,
150 + desc="Beam search",
151 + unit="token",
152 + unit_scale=False,
153 )
118 - )
119 - instance_start_and_end: list[tuple[int, int]] = list(
120 - zip(pos[:-1], pos[1:])
121 - )
122 -
123 - if len(all_beams) == 0:
124 - break
125 -
126 - # only runs for one step
127 - # we don't need to use tqdm here
128 - output = self._render_and_run_requests(
129 - prompts=(beam.get_prompt() for beam in all_beams),
130 - params=self._params_to_seq(sampling_params, len(all_beams)),
131 - output_type=RequestOutput,
132 - lora_requests=[beam.lora_request for beam in all_beams],
133 - use_tqdm=False,
134 - )
135 -
136 - for (start, end), instance in zip(
137 - instance_start_and_end, instances_batch
138 - ):
139 - instance_new_beams = []
140 - for i in range(start, end):
141 - current_beam = all_beams[i]
142 - result = output[i]
143 -
144 - if result.outputs[0].logprobs is not None:
145 - # if `result.outputs[0].logprobs` is None, it means
146 - # the sequence is completed because of the
147 - # max-model-len or abortion. we don't need to add
148 - # it to the new beams.
149 - logprobs = result.outputs[0].logprobs[0]
150 - for token_id, logprob_obj in logprobs.items():
151 - new_beam = BeamSearchSequence(
152 - current_beam.orig_prompt,
153 - tokens=current_beam.tokens + [token_id],
154 - logprobs=current_beam.logprobs + [logprobs],
155 - lora_request=current_beam.lora_request,
156 - cum_logprob=current_beam.cum_logprob
157 - + logprob_obj.logprob,
158 - )
159 -
160 - if token_id == eos_token_id and not ignore_eos:
161 - instance.completed.append(new_beam)
162 - else:
163 - instance_new_beams.append(new_beam)
164 - sorted_beams = sorted(
165 - instance_new_beams, key=sort_beams_key, reverse=True
154 + logger.warning(
155 + "The progress bar shows the upper bound on token "
156 + "steps and may finish early due to stopping "
157 + "conditions. It does not reflect instance-level "
158 + "progress."
159 )
167 - instance.beams = sorted_beams[:beam_width]
160 + for _ in token_iter:
161 + should_stop = self._beam_search_step(
162 + instances_batch=instances_batch,
163 + base_sampling_params=base_sampling_params,
164 + eos_token_id=eos_token_id,
165 + ignore_eos=ignore_eos,
166 + beam_width=beam_width,
167 + sort_beams_key=sort_beams_key,
168 + structured_output_backend=structured_output_backend,
169 + structured_output_key=structured_output_key,
170 + structured_output_bitmask=structured_output_bitmask,
171 + )
172 + if should_stop:
173 + break
174 + finally:
175 + if structured_output_backend is not None:
176 + structured_output_backend.destroy()
177
178 outputs = []
179 for instance in instances:
@@ -180,3 +189,265 @@ class BeamSearchOfflineMixin(OfflineInferenceMixin):
189 outputs.append(BeamSearchOutput(sequences=best_beams))
190
191 return outputs
192 +
193 + def _beam_search_step(
194 + self,
195 + instances_batch: list[BeamSearchInstance],
196 + base_sampling_params: SamplingParams,
197 + eos_token_id: int | None,
198 + ignore_eos: bool,
199 + beam_width: int,
200 + sort_beams_key: Callable,
201 + structured_output_backend: StructuredOutputBackend | None,
202 + structured_output_key: tuple | None,
203 + structured_output_bitmask: torch.Tensor | None,
204 + ) -> bool:
205 + """Run one token step of beam search across a batch of instances.
206 +
207 + Returns True if all beams are exhausted and search should stop.
208 + """
209 + all_beams: list[BeamSearchSequence] = list(
210 + sum((instance.beams for instance in instances_batch), [])
211 + )
212 + pos = [0] + list(
213 + itertools.accumulate(len(instance.beams) for instance in instances_batch)
214 + )
215 + instance_start_and_end: list[tuple[int, int]] = list(zip(pos[:-1], pos[1:]))
216 +
217 + if len(all_beams) == 0:
218 + return True
219 +
220 + if structured_output_backend is not None:
221 + assert (
222 + structured_output_key is not None
223 + and structured_output_bitmask is not None
224 + )
225 + beam_entries = self._build_beam_sampling_params(
226 + all_beams,
227 + base_sampling_params,
228 + structured_output_backend,
229 + structured_output_key,
230 + structured_output_bitmask,
231 + )
232 + active_indices = [
233 + i for i, entry in enumerate(beam_entries) if entry is not None
234 + ]
235 + for i, entry in enumerate(beam_entries):
236 + if entry is None:
237 + beam = all_beams[i]
238 + assert beam.orig_prompt["type"] != "enc_dec"
239 + prompt_len = len(beam.orig_prompt["prompt_token_ids"])
240 + if len(beam.tokens) > prompt_len:
241 + for (s, e), inst in zip(
242 + instance_start_and_end,
243 + instances_batch,
244 + ):
245 + if s <= i < e:
246 + inst.completed.append(beam)
247 + break
248 +
249 + if not active_indices:
250 + return True
251 +
252 + active_beams = [all_beams[i] for i in active_indices]
253 + active_params: Sequence[SamplingParams | PoolingParams] = [
254 + beam_entries[i][0] # type: ignore[index]
255 + for i in active_indices
256 + ]
257 + else:
258 + active_indices = list(range(len(all_beams)))
259 + active_beams = all_beams
260 + active_params = self._params_to_seq( # type: ignore[assignment]
261 + base_sampling_params, len(all_beams)
262 + )
263 +
264 + # only runs for one step
265 + # we don't need to use tqdm here
266 + active_output = self._render_and_run_requests(
267 + prompts=(beam.get_prompt() for beam in active_beams),
268 + params=active_params,
269 + output_type=RequestOutput,
270 + lora_requests=[beam.lora_request for beam in active_beams],
271 + use_tqdm=False,
272 + )
273 +
274 + output: list[RequestOutput | None] = [None] * len(all_beams)
275 + for idx, active_idx in enumerate(active_indices):
276 + output[active_idx] = active_output[idx]
277 +
278 + # Logprobs are computed from raw logits before
279 + # allowed_token_ids masking, so they may contain
280 + # tokens outside the grammar's allowed set. This filtering is also
281 + # the only grammar enforcement for beams whose allowed set exceeds
282 + # the engine-side allowed_token_ids cap.
283 + allowed_sets: list[set[int] | None] = [None] * len(all_beams)
284 + if structured_output_backend is not None:
285 + for i, entry in enumerate(beam_entries):
286 + if entry is not None:
287 + allowed_sets[i] = set(entry[1])
288 +
289 + for (start, end), instance in zip(instance_start_and_end, instances_batch):
290 + instance_new_beams = []
291 + for i in range(start, end):
292 + current_beam = all_beams[i]
293 + result = output[i]
294 +
295 + if result is None:
296 + continue
297 +
298 + if result.outputs[0].logprobs is not None:
299 + # if logprobs is None, the sequence completed
300 + # due to max-model-len or abortion.
301 + logprobs = result.outputs[0].logprobs[0]
302 + allowed = allowed_sets[i]
303 + for token_id, logprob_obj in logprobs.items():
304 + if allowed is not None and token_id not in allowed:
305 + continue
306 + new_beam = BeamSearchSequence(
307 + current_beam.orig_prompt,
308 + tokens=current_beam.tokens + [token_id],
309 + logprobs=current_beam.logprobs + [logprobs],
310 + lora_request=current_beam.lora_request,
311 + cum_logprob=current_beam.cum_logprob + logprob_obj.logprob,
312 + )
313 +
314 + if token_id == eos_token_id and not ignore_eos:
315 + instance.completed.append(new_beam)
316 + else:
317 + instance_new_beams.append(new_beam)
318 + sorted_beams = sorted(
319 + instance_new_beams,
320 + key=sort_beams_key,
321 + reverse=True,
322 + )
323 + instance.beams = sorted_beams[:beam_width]
324 +
325 + return False
326 +
327 + def _init_beam_search_structured_output(
328 + self,
329 + structured_outputs: StructuredOutputsParams,
330 + tokenizer: TokenizerLike,
331 + ) -> tuple[StructuredOutputBackend, tuple, torch.Tensor]:
332 + """Initialize the structured output backend for beam search."""
333 + vllm_config = self.llm_engine.vllm_config
334 + so_config = vllm_config.structured_outputs_config
335 + if so_config is None:
336 + raise ValueError(
337 + "structured_outputs_config is required for beam search "
338 + "with structured outputs"
339 + )
340 +
341 + # Resolve the backend name from engine config if not already set.
342 + if not structured_outputs._backend:
343 + structured_outputs._backend = so_config.backend
344 +
345 + backend_name = structured_outputs._backend
346 + vocab_size = self.model_config.get_vocab_size()
347 +
348 + backend: StructuredOutputBackend
349 + if backend_name == "xgrammar":
350 + from vllm.v1.structured_output.backend_xgrammar import (
351 + XgrammarBackend,
352 + )
353 +
354 + backend = XgrammarBackend(
355 + vllm_config=vllm_config,
356 + tokenizer=tokenizer,
357 + vocab_size=vocab_size,
358 + )
359 + elif backend_name == "guidance":
360 + from vllm.v1.structured_output.backend_guidance import (
361 + GuidanceBackend,
362 + )
363 +
364 + backend = GuidanceBackend(
365 + vllm_config=vllm_config,
366 + tokenizer=tokenizer,
367 + vocab_size=vocab_size,
368 + )
369 + elif backend_name == "outlines":
370 + from vllm.v1.structured_output.backend_outlines import (
371 + OutlinesBackend,
372 + )
373 +
374 + backend = OutlinesBackend(
375 + vllm_config=vllm_config,
376 + tokenizer=tokenizer,
377 + vocab_size=vocab_size,
378 + )
379 + elif backend_name == "lm-format-enforcer":
380 + from vllm.v1.structured_output.backend_lm_format_enforcer import (
381 + LMFormatEnforcerBackend,
382 + )
383 +
384 + backend = LMFormatEnforcerBackend(
385 + vllm_config=vllm_config,
386 + tokenizer=tokenizer,
387 + vocab_size=vocab_size,
388 + )
389 + else:
390 + raise ValueError(f"Unsupported structured output backend: {backend_name}")
391 +
392 + structured_output_key = get_structured_output_key(structured_outputs)
393 + bitmask = backend.allocate_token_bitmask(1)
394 +
395 + return backend, structured_output_key, bitmask
396 +
397 + def _build_beam_sampling_params(
398 + self,
399 + beams: list[BeamSearchSequence],
400 + base_params: SamplingParams,
401 + backend: StructuredOutputBackend,
402 + structured_output_key: tuple,
403 + bitmask: torch.Tensor,
404 + ) -> list[tuple[SamplingParams, list[int]] | None]:
405 + """Build per-beam SamplingParams and allowed token IDs from grammar.
406 +
407 + Returns None for beams where the grammar has terminated.
408 + """
409 + vocab_size = self.model_config.get_vocab_size()
410 + request_type, grammar_spec = structured_output_key
411 + result: list[tuple[SamplingParams, list[int]] | None] = []
412 +
413 + for beam in beams:
414 + # Fresh grammar per beam, replaying generated tokens.
415 + # Backends don't support cloning grammar state, so
416 + # replay is needed to reconstruct the FSM position.
417 + grammar = backend.compile_grammar(request_type, grammar_spec)
418 + assert beam.orig_prompt["type"] != "enc_dec"
419 + prompt_len = len(beam.orig_prompt["prompt_token_ids"])
420 + generated_tokens = beam.tokens[prompt_len:]
421 +
422 + if generated_tokens:
423 + grammar.accept_tokens("beam", generated_tokens)
424 +
425 + if grammar.is_terminated():
426 + result.append(None)
427 + continue
428 +
429 + grammar.fill_bitmask(bitmask, 0)
430 + allowed_ids = _bitmask_to_token_ids(bitmask[0], vocab_size)
431 +
432 + if not allowed_ids:
433 + result.append(None)
434 + continue
435 +
436 + # The engine caps the size of allowed_token_ids. While the
437 + # grammar still allows more tokens than the cap (e.g. inside
438 + # free-form strings), skip the engine-side constraint and rely
439 + # on the logprobs filtering in _beam_search_step instead.
440 + beam_params = SamplingParams(
441 + logprobs=base_params.logprobs,
442 + max_tokens=1,
443 + temperature=base_params.temperature,
444 + allowed_token_ids=(
445 + allowed_ids
446 + if len(allowed_ids) <= _MAX_NUM_ALLOWED_TOKEN_IDS
447 + else None
448 + ),
449 + skip_clone=True,
450 + )
451 + result.append((beam_params, allowed_ids))
452 +
453 + return result
vllm/sampling_params.py
+1
@@ -1048,3 +1048,4 @@ class BeamSearchParams(
1048 temperature: float = 0.0
1049 length_penalty: float = 1.0
1050 include_stop_str_in_output: bool = False
1051 + structured_outputs: StructuredOutputsParams | None = None