[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

@@ -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(

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 = {

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)

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:

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