From ff01a6b14f329ccd4162fce309913542e2abcd5b Mon Sep 17 00:00:00 2001 From: alp Date: Sun, 18 Oct 2026 09:49:42 -0700 Subject: [PATCH] [Bugfix][Frontend] Stop beam search on generation config EOS tokens (#62072) Signed-off-by: mottopanikeiku <176798723+mottopanikeiku@users.noreply.github.com> Co-authored-by: mottopanikeiku <176798723+mottopanikeiku@users.noreply.github.com> --- tests/samplers/test_beam_search.py | 1 + tests/samplers/test_beam_search_online.py | 71 ++++++++++++++++++- .../generate/beam_search/offline.py | 14 ++-- .../generate/beam_search/online.py | 16 +++-- .../entrypoints/generate/beam_search/utils.py | 24 +++++-- 5 files changed, 110 insertions(+), 16 deletions(-) diff --git a/tests/samplers/test_beam_search.py b/tests/samplers/test_beam_search.py index 823202ba32..4e00c41559 100644 --- a/tests/samplers/test_beam_search.py +++ b/tests/samplers/test_beam_search.py @@ -99,6 +99,7 @@ def test_beam_search_abort_returns_partial_outputs_and_continues_other_prompts( llm = LLM.__new__(LLM) llm.llm_engine = Mock() + llm.llm_engine.input_processor.generation_config_fields = {} tokenizer = SimpleNamespace( eos_token_id=0, decode=lambda tokens, skip_special_tokens=False: " ".join( diff --git a/tests/samplers/test_beam_search_online.py b/tests/samplers/test_beam_search_online.py index a6e4967620..b8134e6c8d 100644 --- a/tests/samplers/test_beam_search_online.py +++ b/tests/samplers/test_beam_search_online.py @@ -67,6 +67,7 @@ class _VllmConfig: class _InputProcessor: vllm_config = _VllmConfig() + generation_config_fields: dict[str, Any] = {} def resolve_watermarking(self, params): return self.vllm_config._check_supports_watermarking(params) @@ -128,7 +129,7 @@ class _OfflineServing(BeamSearchOfflineMixin): self, instances_batch: list[BeamSearchInstance], base_sampling_params: SamplingParams, - eos_token_id: int | None, + eos_token_ids: set[int], ignore_eos: bool, beam_width: int, sort_beams_key: Callable[..., Any], @@ -140,6 +141,74 @@ class _OfflineServing(BeamSearchOfflineMixin): return True +@pytest.mark.asyncio +@pytest.mark.parametrize("offline", [False, True]) +@pytest.mark.parametrize("config_eos", [None, 11, [0, 11]]) +@pytest.mark.parametrize("ignore_eos", [False, True]) +async def test_beam_search_respects_generation_config_eos( + monkeypatch, offline: bool, config_eos: int | list[int] | None, ignore_eos: bool +) -> None: + """Model-specific end-of-turn tokens must finish beams like tokenizer EOS.""" + monkeypatch.setattr( + _InputProcessor, "generation_config_fields", {"eos_token_id": config_eos} + ) + prompt: TokensInput = { + "type": "token", + "prompt_token_ids": [1], + } + params = BeamSearchParams( + beam_width=1, + max_tokens=3, + ignore_eos=ignore_eos, + include_stop_str_in_output=True, + watermarking=False, + ) + result = await anext( + _EngineClient().generate(prompt, SamplingParams(watermarking=False)) + ) + result.outputs[0].logprobs = [{11: Logprob(-0.1)}] + should_stop = config_eos is not None and not ignore_eos + expected_tokens = [11] if should_stop else [11, 11, 11] + engine_params: list[SamplingParams] = [] + + if offline: + offline_serving = _OfflineServing() + monkeypatch.setattr( + _OfflineServing, + "_beam_search_step", + BeamSearchOfflineMixin._beam_search_step, + ) + + def render_and_run_requests(*, params, **kwargs): + engine_params.extend(params) + return [result] + + monkeypatch.setattr( + offline_serving, "_render_and_run_requests", render_and_run_requests + ) + beam = offline_serving.beam_search(cast(list, [prompt]), params)[0].sequences[0] + assert beam.tokens == [1, *expected_tokens] + else: + + async def generate(engine_prompt, sampling_params, *args, **kwargs): + engine_params.append(sampling_params) + yield result + + online_serving = _AsyncServing() + monkeypatch.setattr(online_serving.engine_client, "generate", generate) + output = ( + await anext(online_serving.beam_search(prompt, "request", params)) + ).outputs[0] + assert output.token_ids == expected_tokens + assert output.finish_reason == ("stop" if should_stop else "length") + assert output.stop_reason == (11 if should_stop else None) + + # The engine validates stop_token_ids against the vocab before it applies + # the generation config, so the config EOS ids must not be copied into them. + assert engine_params + assert all(not p.stop_token_ids for p in engine_params) + + @pytest.mark.asyncio async def test_beam_search_handles_extra_logprob_candidates() -> None: prompt: TokensInput = { diff --git a/vllm/entrypoints/generate/beam_search/offline.py b/vllm/entrypoints/generate/beam_search/offline.py index 441475cafd..53297a4736 100644 --- a/vllm/entrypoints/generate/beam_search/offline.py +++ b/vllm/entrypoints/generate/beam_search/offline.py @@ -27,6 +27,7 @@ from .utils import ( BeamSearchOutput, BeamSearchSequence, create_sort_beams_key_function, + get_beam_search_eos_token_ids, ) logger = init_logger(__name__) @@ -86,8 +87,6 @@ class BeamSearchOfflineMixin(OfflineInferenceMixin): self.llm_engine.input_processor.resolve_watermarking(params) tokenizer = self.renderer.get_tokenizer() - eos_token_id = tokenizer.eos_token_id - sort_beams_key = create_sort_beams_key_function(eos_token_id, length_penalty) engine_inputs = self._preprocess_cmpl(prompts) lora_requests = self._lora_request_to_seq(lora_request, len(engine_inputs)) @@ -129,6 +128,11 @@ class BeamSearchOfflineMixin(OfflineInferenceMixin): params.structured_outputs, tokenizer ) + eos_token_ids = get_beam_search_eos_token_ids( + self.llm_engine.input_processor.generation_config_fields, + tokenizer.eos_token_id, + ) + sort_beams_key = create_sort_beams_key_function(eos_token_ids, length_penalty) instances: list[BeamSearchInstance] = [] for lora_req, prompt in zip(lora_requests, engine_inputs): @@ -169,7 +173,7 @@ class BeamSearchOfflineMixin(OfflineInferenceMixin): should_stop = self._beam_search_step( instances_batch=instances_batch, base_sampling_params=base_sampling_params, - eos_token_id=eos_token_id, + eos_token_ids=eos_token_ids, ignore_eos=ignore_eos, beam_width=beam_width, sort_beams_key=sort_beams_key, @@ -204,7 +208,7 @@ class BeamSearchOfflineMixin(OfflineInferenceMixin): self, instances_batch: list[BeamSearchInstance], base_sampling_params: SamplingParams, - eos_token_id: int | None, + eos_token_ids: set[int], ignore_eos: bool, beam_width: int, sort_beams_key: Callable, @@ -350,7 +354,7 @@ class BeamSearchOfflineMixin(OfflineInferenceMixin): cum_logprob=current_beam.cum_logprob + logprob_obj.logprob, ) - if token_id == eos_token_id and not ignore_eos: + if token_id in eos_token_ids and not ignore_eos: instance.completed.append(new_beam) else: instance_new_beams.append(new_beam) diff --git a/vllm/entrypoints/generate/beam_search/online.py b/vllm/entrypoints/generate/beam_search/online.py index 9cbe4ea258..1df1be0145 100644 --- a/vllm/entrypoints/generate/beam_search/online.py +++ b/vllm/entrypoints/generate/beam_search/online.py @@ -20,6 +20,7 @@ from vllm.utils.async_utils import collect_from_async_generator from .utils import ( BeamSearchSequence, create_sort_beams_key_function, + get_beam_search_eos_token_ids, ) @@ -47,8 +48,6 @@ class BeamSearchOnlineMixin(ABC): self.engine_client.input_processor.resolve_watermarking(params) tokenizer = self.renderer.get_tokenizer() - eos_token_id = tokenizer.eos_token_id - sort_beams_key = create_sort_beams_key_function(eos_token_id, length_penalty) if prompt["type"] == "embeds": raise NotImplementedError("Embedding prompt not supported for beam search") @@ -74,6 +73,11 @@ class BeamSearchOnlineMixin(ABC): sampling_params, self.engine_client.input_processor.vllm_config ) allowed = set(params.allowed_token_ids) if params.allowed_token_ids else None + eos_token_ids = get_beam_search_eos_token_ids( + self.engine_client.input_processor.generation_config_fields, + tokenizer.eos_token_id, + ) + sort_beams_key = create_sort_beams_key_function(eos_token_ids, length_penalty) all_beams = [ BeamSearchSequence( orig_prompt=prompt, @@ -159,17 +163,17 @@ class BeamSearchOnlineMixin(ABC): candidate_logprob = ( current_beam.cum_logprob + logprob_obj.logprob ) - if token_id == eos_token_id and not ignore_eos: + if token_id in eos_token_ids and not ignore_eos: completed.append( BeamSearchSequence( orig_prompt=prompt, - tokens=current_beam.tokens + [eos_token_id] + tokens=current_beam.tokens + [token_id] if include_stop_str_in_output else current_beam.tokens, logprobs=current_beam.logprobs + [logprobs], cum_logprob=candidate_logprob, finish_reason="stop", - stop_reason=eos_token_id, + stop_reason=token_id, ) ) else: @@ -219,7 +223,7 @@ class BeamSearchOnlineMixin(ABC): best_beams = sorted_completed[:beam_width] for beam in best_beams: - if beam.tokens[-1] == eos_token_id and not ignore_eos: + if beam.tokens[-1] in eos_token_ids and not ignore_eos: # Skip the eos token in the text. tokens = beam.tokens[tokenized_length:-1] else: diff --git a/vllm/entrypoints/generate/beam_search/utils.py b/vllm/entrypoints/generate/beam_search/utils.py index 61f97963b6..b0312b75f4 100644 --- a/vllm/entrypoints/generate/beam_search/utils.py +++ b/vllm/entrypoints/generate/beam_search/utils.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from dataclasses import dataclass +from typing import Any from vllm.inputs import ( DecoderOnlyEngineInput, @@ -13,6 +14,7 @@ from vllm.inputs import ( ) from vllm.logprobs import Logprob from vllm.lora.request import LoRARequest +from vllm.sampling_params import SamplingParams @dataclass @@ -137,7 +139,7 @@ class BeamSearchInstance: def get_beam_search_score( tokens: list[int], cumulative_logprob: float, - eos_token_id: int, + eos_token_ids: set[int], length_penalty: float = 1.0, ) -> float: """Calculate the beam search score with length penalty. @@ -147,17 +149,31 @@ def get_beam_search_score( https://github.com/huggingface/transformers/blob/ccb92be23def445f2afdea94c31286f84b89eb5b/src/transformers/generation/beam_search.py#L938 """ seq_len = len(tokens) - if tokens[-1] == eos_token_id: + if tokens[-1] in eos_token_ids: seq_len -= 1 # An aborted beam may contain only an EOS prompt token. return cumulative_logprob / (max(seq_len, 1) ** length_penalty) -def create_sort_beams_key_function(eos_token_id: int, length_penalty: float): +def create_sort_beams_key_function(eos_token_ids: set[int], length_penalty: float): def sort_beams_key(x: BeamSearchSequence) -> float: return get_beam_search_score( - x.tokens, x.cum_logprob, eos_token_id, length_penalty + x.tokens, x.cum_logprob, eos_token_ids, length_penalty ) return sort_beams_key + + +def get_beam_search_eos_token_ids( + generation_config: dict[str, Any], eos_token_id: int | None +) -> set[int]: + """Return the token ids that end generation: the tokenizer EOS plus the + generation config's ``eos_token_id``. + + This uses a separate SamplingParams so the params sent to the engine keep + their own ``stop_token_ids`` and are validated as before. + """ + params = SamplingParams() + params.update_from_generation_config(generation_config, eos_token_id) + return params.all_stop_token_ids