A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 04:45 UTC
[Core][Bugfix] Fix Online MM Beam Search (#19688)
Signed-off-by: Alex-Brooks <Alex.Brooks@ibm.com>
Alex Brooks committed
Jun 19, 2025 at 11:18 UTC
ead2110297a65d8df8a28da3953eb449a0b6449c
3 files changed
+45
-12
tests/entrypoints/openai/test_vision.py
+27
-4
@@ -25,6 +25,25 @@ TEST_IMAGE_URLS = [
25
"https://upload.wikimedia.org/wikipedia/commons/0/0b/RGBA_comp.png",
26
]
27
28
+EXPECTED_MM_BEAM_SEARCH_RES = [
29
+ [
30
+ "The image shows a wooden boardwalk leading through a",
31
+ "The image shows a wooden boardwalk extending into a",
32
+ ],
33
+ [
34
+ "The image shows two parrots perched on",
35
+ "The image shows two birds perched on a cur",
36
+ ],
37
+ [
38
+ "The image shows a Venn diagram with three over",
39
+ "This image shows a Venn diagram with three over",
40
+ ],
41
+ [
42
+ "This image displays a gradient of colors ranging from",
43
+ "This image displays a gradient of colors transitioning from",
44
+ ],
45
+]
46
+
47
48
@pytest.fixture(scope="module")
49
def server():
@@ -270,10 +289,13 @@ async def test_single_chat_session_image_base64encoded(
289
290
@pytest.mark.asyncio
291
@pytest.mark.parametrize("model_name", [MODEL_NAME])
273
-@pytest.mark.parametrize("image_url", TEST_IMAGE_URLS)
292
+@pytest.mark.parametrize("image_idx", list(range(len(TEST_IMAGE_URLS))))
293
async def test_single_chat_session_image_base64encoded_beamsearch(
275
- client: openai.AsyncOpenAI, model_name: str, image_url: str,
294
+ client: openai.AsyncOpenAI, model_name: str, image_idx: int,
295
base64_encoded_image: dict[str, str]):
296
+ # NOTE: This test also validates that we pass MM data through beam search
297
+ image_url = TEST_IMAGE_URLS[image_idx]
298
+ expected_res = EXPECTED_MM_BEAM_SEARCH_RES[image_idx]
299
300
messages = [{
301
"role":
@@ -297,10 +319,11 @@ async def test_single_chat_session_image_base64encoded_beamsearch(
319
messages=messages,
320
n=2,
321
max_completion_tokens=10,
322
+ temperature=0.0,
323
extra_body=dict(use_beam_search=True))
324
assert len(chat_completion.choices) == 2
302
- assert chat_completion.choices[
303
- 0].message.content != chat_completion.choices[1].message.content
325
+ for actual, expected_str in zip(chat_completion.choices, expected_res):
326
+ assert actual.message.content == expected_str
327
328
329
@pytest.mark.asyncio
vllm/engine/protocol.py
+11
-2
@@ -88,9 +88,18 @@ class EngineClient(ABC):
88
if processed_inputs["type"] == "embeds":
89
raise NotImplementedError
90
91
- prompt_token_ids = processed_inputs["prompt_token_ids"]
91
+ # This is a workaround to fix multimodal beam search; this is a
92
+ # bandaid fix for 2 small problems:
93
+ # 1. Multi_modal_data on the processed_inputs currently resolves to
94
+ # `None`.
95
+ # 2. preprocessing above expands the multimodal placeholders. However,
96
+ # this happens again in generation, so the double expansion causes
97
+ # a mismatch.
98
+ # TODO - would be ideal to handle this more gracefully.
99
+ prompt_token_ids = prompt.get("prompt_token_ids")
100
+ multi_modal_data = prompt.get("multi_modal_data")
101
+
102
prompt_text = processed_inputs.get("prompt")
93
- multi_modal_data = processed_inputs.get("multi_modal_data")
103
mm_processor_kwargs = processed_inputs.get("mm_processor_kwargs")
104
105
tokenized_length = len(prompt_token_ids)
vllm/entrypoints/llm.py
+7
-6
@@ -15,7 +15,8 @@ from tqdm.auto import tqdm
15
from typing_extensions import TypeVar, deprecated
16
17
from vllm.beam_search import (BeamSearchInstance, BeamSearchOutput,
18
- BeamSearchSequence, get_beam_search_score)
18
+ BeamSearchSequence,
19
+ create_sort_beams_key_function)
20
from vllm.config import (CompilationConfig, ModelDType, TokenizerMode,
21
is_init_field)
22
from vllm.engine.arg_utils import (EngineArgs, HfOverrides, PoolerConfig,
@@ -575,10 +576,11 @@ class LLM:
576
lora_requests = self._get_beam_search_lora_requests(
577
lora_request, prompts)
578
578
- def sort_beams_key(x: BeamSearchSequence) -> float:
579
- return get_beam_search_score(x.tokens, x.cum_logprob,
580
- tokenizer.eos_token_id,
581
- length_penalty)
579
+ tokenizer = self.get_tokenizer()
580
+ sort_beams_key = create_sort_beams_key_function(
581
+ tokenizer.eos_token_id,
582
+ length_penalty,
583
+ )
584
585
def create_tokens_prompt_from_beam(
586
beam: BeamSearchSequence) -> TokensPrompt:
@@ -593,7 +595,6 @@ class LLM:
595
"mm_processor_kwargs"] = beam.mm_processor_kwargs
596
return TokensPrompt(**token_prompt_kwargs)
597
596
- tokenizer = self.get_tokenizer()
598
# generate 2 * beam_width candidates at each step
599
# following the huggingface transformers implementation
600
# at https://github.com/huggingface/transformers/blob/e15687fffe5c9d20598a19aeab721ae0a7580f8a/src/transformers/generation/beam_search.py#L534 # noqa