diff --git a/tests/samplers/test_beam_search_online.py b/tests/samplers/test_beam_search_online.py index 22afce2d78..a6e4967620 100644 --- a/tests/samplers/test_beam_search_online.py +++ b/tests/samplers/test_beam_search_online.py @@ -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) diff --git a/vllm/entrypoints/generate/beam_search/offline.py b/vllm/entrypoints/generate/beam_search/offline.py index 0acd0b43a5..441475cafd 100644 --- a/vllm/entrypoints/generate/beam_search/offline.py +++ b/vllm/entrypoints/generate/beam_search/offline.py @@ -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: diff --git a/vllm/entrypoints/generate/beam_search/online.py b/vllm/entrypoints/generate/beam_search/online.py index 5a95f40d31..9cbe4ea258 100644 --- a/vllm/entrypoints/generate/beam_search/online.py +++ b/vllm/entrypoints/generate/beam_search/online.py @@ -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 diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index 071f73f043..102eeb0e26 100644 --- a/vllm/entrypoints/openai/chat_completion/protocol.py +++ b/vllm/entrypoints/openai/chat_completion/protocol.py @@ -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, diff --git a/vllm/entrypoints/openai/completion/protocol.py b/vllm/entrypoints/openai/completion/protocol.py index 9e8ffb496b..5a72d0e605 100644 --- a/vllm/entrypoints/openai/completion/protocol.py +++ b/vllm/entrypoints/openai/completion/protocol.py @@ -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, diff --git a/vllm/sampling_params.py b/vllm/sampling_params.py index 8331a1707b..a0ff636155 100644 --- a/vllm/sampling_params.py +++ b/vllm/sampling_params.py @@ -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