iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Bugfix][OpenAI] Respect allowed_token_ids in beam search (#47630)

Browse Source 
Signed-off-by: Ting Sun <suntcrick@gmail.com>
This commit is contained in:
Ting SUN
2026-10-17 08:50:40 +00:00
committed by GitHub
parent b05ede5616
commit eb8de8d1a1
6 changed files with 185 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/samplers/test_beam_search_online.py
+98 -2
View File
@@ -2,13 +2,20 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from collections.abc import Callable, Sequence
from types import SimpleNamespace
from typing import Any, cast
from unittest.mock import MagicMock
import pytest
import torch
from vllm import CompletionOutput, PromptType, RequestOutput, SamplingParams
from vllm import (
CompletionOutput,
PromptType,
RequestOutput,
SamplingParams,
TokensPrompt,
)
from vllm import logger as vllm_logger
from vllm.config import VllmConfig
from vllm.entrypoints.generate.beam_search.offline import BeamSearchOfflineMixin
@@ -17,9 +24,10 @@ from vllm.entrypoints.generate.beam_search.utils import (
BeamSearchInstance,
BeamSearchSequence,
)
from vllm.exceptions import VLLMValidationError
from vllm.inputs import EngineInput, TokensInput
from vllm.logprobs import Logprob, SampleLogprobs
from vllm.sampling_params import BeamSearchParams
from vllm.sampling_params import BeamSearchParams, StructuredOutputsParams
from vllm.v1.structured_output.backend_types import StructuredOutputBackend
@@ -48,8 +56,10 @@ class _Renderer:
class _VllmConfig:
use_v2_model_runner = True
watermark_config = object()
speculative_config = None
model_config = SimpleNamespace(is_diffusion=False)
def _check_supports_watermarking(self, params: BeamSearchParams) -> bool:
return VllmConfig._check_supports_watermarking(cast(VllmConfig, self), params)
@@ -322,3 +332,89 @@ def test_offline_structured_beam_search_disables_internal_watermarking() -> None
assert beam_params.watermarking is False
assert beam_params.allowed_token_ids == [0]
assert base_params.allowed_token_ids is None
@pytest.mark.asyncio
@pytest.mark.parametrize("offline", [False, True])
@pytest.mark.parametrize("exhausted", [False, True])
async def test_explicit_beam_allowlist_uses_exact_scores(offline, exhausted):
prompt: TokensInput = {"type": "token", "prompt_token_ids": [1]}
base = SamplingParams(watermarking=False)
template = [r async for r in _EngineClient().generate(prompt, base)][0]
template.outputs[0].logprobs[0] = dict(
reversed(list(template.outputs[0].logprobs[0].items()))
)
calls = []
if exhausted:
for logprob in template.outputs[0].logprobs[0].values():
logprob.logprob = float("-inf")
def score(params):
assert params.logprob_token_ids == [13, 14]
assert params.logprobs is None and params.allowed_token_ids is None
calls.append(params)
return template
class Client(_EngineClient):
async def generate(self, prompt, params, *args, **kwargs):
yield score(params)
class Offline(_OfflineServing):
_beam_search_step = BeamSearchOfflineMixin._beam_search_step
def _render_and_run_requests(self, prompts, params, output_type, **kwargs):
return [score(p) for p in params]
params = BeamSearchParams(beam_width=2, max_tokens=2, allowed_token_ids=[13, 14])
ids: list[Sequence[int]]
scores: list[float | None]
if offline:
result = (
Offline().beam_search([cast(TokensPrompt, prompt)], params)[0].sequences
)
ids = [b.tokens[1:] for b in result]
scores = [b.cum_logprob for b in result]
step_scores = [step for b in result for step in b.logprobs]
else:
serving = _AsyncServing()
serving.engine_client = Client()
online_result = [r async for r in serving.beam_search(prompt, "test", params)][
0
]
ids = [b.token_ids for b in online_result.outputs]
scores = [b.cumulative_logprob for b in online_result.outputs]
step_scores = [
step for b in online_result.outputs for step in (b.logprobs or [])
]
if exhausted:
assert len(calls) == 1
assert ids == scores == []
return
assert len(calls) == 3
assert ids[0] == [13, 13] and all(set(x) <= {13, 14} for x in ids)
assert scores == [-6.0, -7.0]
assert all(list(step) == [13, 14] for step in step_scores)
assert all(score.rank is None for step in step_scores for score in step.values())
@pytest.mark.parametrize(
"case", ["empty", "grammar", "spec", "diffusion", "v1", "size", "id"]
)
def test_beam_allowlist_rejects_unsupported_requests(case):
config = SimpleNamespace(
use_v2_model_runner=case != "v1",
speculative_config=object() if case == "spec" else None,
model_config=SimpleNamespace(
is_diffusion=case == "diffusion",
max_logprobs=20,
get_vocab_size=lambda: 256,
),
)
ids = {"empty": [], "size": list(range(129)), "id": [256]}.get(case, [13])
params = BeamSearchParams(2, 1, allowed_token_ids=ids)
if case == "grammar":
params.structured_outputs = StructuredOutputsParams(regex="a")
sampling = SamplingParams(logprobs=4)
with pytest.raises((ValueError, VLLMValidationError)):
params.configure_logprobs(sampling, config)
sampling._validate_logprobs(config.model_config)
vllm/entrypoints/generate/beam_search/offline.py
+31 -9
View File
@@ -2,6 +2,7 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import itertools
import math
from collections.abc import Callable, Sequence
import torch
@@ -104,14 +105,6 @@ class BeamSearchOfflineMixin(OfflineInferenceMixin):
structured_output_backend: StructuredOutputBackend | None = None
structured_output_key = None
structured_output_bitmask = None
if params.structured_outputs is not None:
(
structured_output_backend,
structured_output_key,
structured_output_bitmask,
) = self._init_beam_search_structured_output(
params.structured_outputs, tokenizer
)
# generate 2 * beam_width candidates at each step
# following the huggingface transformers implementation
@@ -124,6 +117,18 @@ class BeamSearchOfflineMixin(OfflineInferenceMixin):
detokenize=False,
skip_clone=True, # Internal beam search, safe to skip clone
)
params.configure_logprobs(
base_sampling_params, self.llm_engine.input_processor.vllm_config
)
if params.structured_outputs is not None:
(
structured_output_backend,
structured_output_key,
structured_output_bitmask,
) = self._init_beam_search_structured_output(
params.structured_outputs, tokenizer
)
instances: list[BeamSearchInstance] = []
for lora_req, prompt in zip(lora_requests, engine_inputs):
@@ -287,7 +292,12 @@ class BeamSearchOfflineMixin(OfflineInferenceMixin):
# tokens outside the grammar's allowed set. This filtering is also
# the only grammar enforcement for beams whose allowed set exceeds
# the engine-side allowed_token_ids cap.
allowed_sets: list[set[int] | None] = [None] * len(all_beams)
request_allowed = (
set(base_sampling_params.logprob_token_ids)
if base_sampling_params.logprob_token_ids is not None
else None
)
allowed_sets: list[set[int] | None] = [request_allowed] * len(all_beams)
if structured_output_backend is not None:
for i, entry in enumerate(beam_entries):
if entry is not None:
@@ -316,6 +326,18 @@ class BeamSearchOfflineMixin(OfflineInferenceMixin):
if result.outputs[0].logprobs is not None:
logprobs = result.outputs[0].logprobs[0]
if request_allowed is not None:
logprobs = {
token: score
for token, score in sorted(
logprobs.items(),
key=lambda item: item[1].logprob,
reverse=True,
)
if token in request_allowed and math.isfinite(score.logprob)
}
for score in logprobs.values():
score.rank = None
allowed = allowed_sets[i]
for token_id, logprob_obj in logprobs.items():
if allowed is not None and token_id not in allowed:
vllm/entrypoints/generate/beam_search/online.py
+21 -1
View File
@@ -2,6 +2,7 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import asyncio
import math
from abc import ABC
from collections.abc import AsyncGenerator, Mapping
@@ -16,7 +17,10 @@ from vllm.sampling_params import BeamSearchParams, SamplingParams
from vllm.utils import random_uuid
from vllm.utils.async_utils import collect_from_async_generator
from .utils import BeamSearchSequence, create_sort_beams_key_function
from .utils import (
BeamSearchSequence,
create_sort_beams_key_function,
)
class BeamSearchOnlineMixin(ABC):
@@ -66,6 +70,10 @@ class BeamSearchOnlineMixin(ABC):
watermarking=False,
detokenize=False,
)
params.configure_logprobs(
sampling_params, self.engine_client.input_processor.vllm_config
)
allowed = set(params.allowed_token_ids) if params.allowed_token_ids else None
all_beams = [
BeamSearchSequence(
orig_prompt=prompt,
@@ -135,6 +143,18 @@ class BeamSearchOnlineMixin(ABC):
if result.outputs[0].logprobs is not None:
logprobs = result.outputs[0].logprobs[0]
if allowed is not None:
logprobs = {
token: score
for token, score in sorted(
logprobs.items(),
key=lambda item: item[1].logprob,
reverse=True,
)
if token in allowed and math.isfinite(score.logprob)
}
for score in logprobs.values():
score.rank = None
for token_id, logprob_obj in logprobs.items():
candidate_logprob = (
current_beam.cum_logprob + logprob_obj.logprob
vllm/entrypoints/openai/chat_completion/protocol.py
+1
View File
@@ -710,6 +710,7 @@ class ChatCompletionRequest(OpenAIBaseModel):
return BeamSearchParams(
beam_width=n,
allowed_token_ids=self.allowed_token_ids,
max_tokens=max_tokens,
ignore_eos=self.ignore_eos,
temperature=temperature,
vllm/entrypoints/openai/completion/protocol.py
+1
View File
@@ -334,6 +334,7 @@ class CompletionRequest(OpenAIBaseModel):
return BeamSearchParams(
beam_width=n,
allowed_token_ids=self.allowed_token_ids,
max_tokens=max_tokens,
ignore_eos=self.ignore_eos,
temperature=temperature,
vllm/sampling_params.py
+33 -1
View File
@@ -8,7 +8,7 @@ import math
from dataclasses import field
from enum import Enum, IntEnum
from functools import cached_property
from typing import Annotated, Any
from typing import TYPE_CHECKING, Annotated, Any
import msgspec
import numpy as np
@@ -24,6 +24,9 @@ from vllm.utils.mistral import is_mistral_tokenizer
from vllm.utils.torch_utils import make_ndarray_with_pad
from vllm.v1.serial_utils import PydanticMsgspecMixin
if TYPE_CHECKING:
from vllm.config import VllmConfig
logger = init_logger(__name__)
_SAMPLING_EPS = 1e-5
@@ -1507,6 +1510,35 @@ class BeamSearchParams(
structured_outputs: StructuredOutputsParams | None = None
skip_special_tokens: bool = True
watermarking: bool | None = None
allowed_token_ids: list[int] | None = None
def __post_init__(self) -> None:
_verify_num_sequences(self.beam_width, "beam_width")
def configure_logprobs(
self, sampling_params: SamplingParams, config: "VllmConfig"
) -> None:
ids = self.allowed_token_ids
if ids is None:
return
if not ids:
raise ValueError("Beam search allowed_token_ids must not be empty")
if self.structured_outputs is not None:
# Grammar candidates can exceed the exact-ID limit of 128 tokens.
raise ValueError(
"Beam allowlists with structured outputs are not supported"
)
if config.speculative_config is not None:
# Rejection sampling does not yet fully support exact-ID logprobs.
raise ValueError(
"Beam allowlists with speculative decoding are not supported"
)
if config.model_config.is_diffusion:
# The diffusion sampler does not consume logprob_token_ids.
raise ValueError("Beam allowlists with diffusion models are not supported")
if not config.use_v2_model_runner:
# V1 exact-ID logprobs overwrite other rows in a mixed batch.
raise ValueError("Beam allowlists require the V2 model runner")
# Score the original distribution without applying a sampling mask.
sampling_params.logprobs = None
sampling_params.logprob_token_ids = ids