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 7, 02:45 UTC
[Frontend] Add Support for MM Encoder/Decoder Beam Search (Offline) (#36153)
Signed-off-by: Alex Brooks <albrooks@redhat.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Alex Brooks committed
Mar 6, 2026 at 02:16 UTC
10f4db4dbecaafc8c0af8b36e9e0bc2f186deb2d
5 files changed
+192
-30
tests/conftest.py
+27
-18
@@ -843,7 +843,10 @@ class VllmRunner:
843
844
def get_inputs(
845
self,
846
- prompts: list[str] | list[torch.Tensor] | list[list[int]],
846
+ prompts: list[str]
847
+ | list[torch.Tensor]
848
+ | list[list[int]]
849
+ | list[dict[str, Any]],
850
images: PromptImageInput | None = None,
851
videos: PromptVideoInput | None = None,
852
audios: PromptAudioInput | None = None,
@@ -857,26 +860,32 @@ class VllmRunner:
860
861
inputs = list[dict[str, Any]]()
862
for i, prompt in enumerate(prompts):
860
- prompt_dict = dict[str, Any]()
861
- if isinstance(prompt, str):
862
- prompt_dict["prompt"] = prompt
863
- elif isinstance(prompt, list):
864
- prompt_dict["prompt_token_ids"] = prompt
863
+ # If we're passing an encoder/decoder prompt, we assume it
864
+ # already contains the multimodal data in the prompt
865
+ if isinstance(prompt, dict):
866
+ assert images is None and audios is None and videos is None
867
+ inputs.append(prompt.copy())
868
else:
866
- prompt_dict["prompt_embeds"] = prompt
867
-
868
- multi_modal_data = dict[str, Any]()
869
- if images is not None and (image := images[i]) is not None:
870
- multi_modal_data["image"] = image
871
- if videos is not None and (video := videos[i]) is not None:
872
- multi_modal_data["video"] = video
873
- if audios is not None and (audio := audios[i]) is not None:
874
- multi_modal_data["audio"] = audio
869
+ prompt_dict = dict[str, Any]()
870
+ if isinstance(prompt, str):
871
+ prompt_dict["prompt"] = prompt
872
+ elif isinstance(prompt, list):
873
+ prompt_dict["prompt_token_ids"] = prompt
874
+ else:
875
+ prompt_dict["prompt_embeds"] = prompt
876
+
877
+ multi_modal_data = dict[str, Any]()
878
+ if images is not None and (image := images[i]) is not None:
879
+ multi_modal_data["image"] = image
880
+ if videos is not None and (video := videos[i]) is not None:
881
+ multi_modal_data["video"] = video
882
+ if audios is not None and (audio := audios[i]) is not None:
883
+ multi_modal_data["audio"] = audio
884
876
- if multi_modal_data:
877
- prompt_dict["multi_modal_data"] = multi_modal_data
885
+ if multi_modal_data:
886
+ prompt_dict["multi_modal_data"] = multi_modal_data
887
879
- inputs.append(prompt_dict)
888
+ inputs.append(prompt_dict)
889
890
return inputs
891
tests/models/multimodal/generation/test_whisper.py
+108
-3
@@ -90,9 +90,9 @@ def run_test(
90
91
92
@pytest.fixture
93
-def input_audios() -> list[tuple[list[str], list[str], list[tuple[Any, int]]]]:
93
+def resampled_assets() -> list[tuple[Any, int]]:
94
audio_assets = [AudioAsset("mary_had_lamb"), AudioAsset("winning_call")]
95
- inputs = []
95
+ sampled_assets = []
96
for asset in audio_assets:
97
audio, orig_sr = asset.audio_and_sample_rate
98
# Resample to Whisper's expected sample rate (16kHz)
@@ -100,8 +100,21 @@ def input_audios() -> list[tuple[list[str], list[str], list[tuple[Any, int]]]]:
100
audio = librosa.resample(
101
audio, orig_sr=orig_sr, target_sr=WHISPER_SAMPLE_RATE
102
)
103
+ sampled_assets.append(
104
+ (audio, WHISPER_SAMPLE_RATE),
105
+ )
106
+ return sampled_assets
107
+
108
+
109
+@pytest.fixture
110
+def input_audios(
111
+ resampled_assets,
112
+) -> list[tuple[list[str], list[str], list[tuple[Any, int]]]]:
113
+ inputs = []
114
+ # audio assets are resampled to WHISPER_SAMPLE_RATE
115
+ for audio_info in resampled_assets:
116
# vLLM prompts, HF prompts, audio inputs
104
- inputs.append(([VLLM_PROMPT], [HF_PROMPT], [(audio, WHISPER_SAMPLE_RATE)]))
117
+ inputs.append(([VLLM_PROMPT], [HF_PROMPT], [audio_info]))
118
return inputs
119
120
@@ -111,6 +124,98 @@ def check_model_available(model: str) -> None:
124
model_info.check_transformers_version(on_fail="skip")
125
126
127
+@pytest.mark.parametrize("dtype", ["half"])
128
+@pytest.mark.parametrize("max_tokens", [64])
129
+@pytest.mark.parametrize("beam_width", [1, 2])
130
+def test_beam_search_encoder_decoder(
131
+ monkeypatch,
132
+ hf_runner,
133
+ vllm_runner,
134
+ dtype: str,
135
+ max_tokens: int,
136
+ beam_width: int,
137
+ resampled_assets,
138
+) -> None:
139
+ """Test beam search with encoder-decoder models (Whisper)."""
140
+ if current_platform.is_rocm():
141
+ monkeypatch.setenv("VLLM_ROCM_USE_SKINNY_GEMM", "0")
142
+
143
+ model = "openai/whisper-large-v3-turbo"
144
+ check_model_available(model)
145
+
146
+ hf_prompts = [
147
+ "<|startoftranscript|>",
148
+ "<|startoftranscript|>",
149
+ ]
150
+
151
+ with hf_runner(model, dtype=dtype, auto_cls=AutoModelForSpeechSeq2Seq) as hf_model:
152
+ hf_outputs = hf_model.generate_beam_search(
153
+ hf_prompts,
154
+ beam_width=beam_width,
155
+ max_tokens=max_tokens,
156
+ audios=resampled_assets,
157
+ )
158
+
159
+ # Test both explicit encoder/decoder prompts
160
+ vllm_prompts = [
161
+ # Implicit encoder/decoder prompt
162
+ {
163
+ "prompt": "<|startoftranscript|>",
164
+ "multi_modal_data": {"audio": resampled_assets[0]},
165
+ },
166
+ # Explicit encoder/decover prompt
167
+ {
168
+ "encoder_prompt": {
169
+ "prompt": "",
170
+ "multi_modal_data": {"audio": resampled_assets[1]},
171
+ },
172
+ "decoder_prompt": "<|startoftranscript|>",
173
+ },
174
+ ]
175
+
176
+ with vllm_runner(
177
+ model,
178
+ dtype="half",
179
+ max_model_len=448,
180
+ tensor_parallel_size=1,
181
+ max_num_seqs=4,
182
+ limit_mm_per_prompt={"audio": 2},
183
+ enforce_eager=True,
184
+ ) as vllm_model:
185
+ vllm_outputs = vllm_model.generate_beam_search(
186
+ vllm_prompts,
187
+ beam_width=beam_width,
188
+ max_tokens=max_tokens,
189
+ )
190
+
191
+ for i in range(len(vllm_prompts)):
192
+ hf_output_ids, hf_output_texts = hf_outputs[i]
193
+ vllm_output_ids, vllm_output_texts = vllm_outputs[i]
194
+
195
+ for j, (hf_text, vllm_text) in enumerate(
196
+ zip(hf_output_texts, vllm_output_texts)
197
+ ):
198
+ print(f">>>{j}-th hf output [NOTE: special tokens are filtered]:")
199
+ print(hf_text)
200
+ print(f">>>{j}-th vllm output:")
201
+ print(vllm_text)
202
+
203
+ # Check that we got the same number of beams
204
+ assert len(hf_output_ids) == len(vllm_output_ids)
205
+
206
+ # For encoder-decoder models, we primarily want to verify that:
207
+ # 1. Beam search completes without errors
208
+ # 2. We get the expected number of beams
209
+ # 3. Outputs are reasonable (non-empty, diverse beams)
210
+ for j in range(len(vllm_output_ids)):
211
+ # Check that outputs are not empty
212
+ assert len(vllm_output_ids[j]) > 0, f"Prompt {i}, beam {j}: empty output"
213
+ # Check that decoded text is not empty
214
+ assert len(vllm_output_texts[j].strip()) > 0, (
215
+ f"Prompt {i}, beam {j}: empty text output"
216
+ )
217
+
218
+
219
def test_parse_language_detection_output():
220
"""Unit test for WhisperForConditionalGeneration.parse_language_detection_output.
221
tests/samplers/test_beam_search.py
+4
@@ -219,3 +219,7 @@ def test_beam_search_passes_multimodal_data(
219
filtered_hf_output_ids = filtered_hf_output_ids[:-1]
220
221
assert filtered_hf_output_ids == filtered_vllm_output_ids
222
+
223
+
224
+# NOTE: encoder/decoder tests are currently located under
225
+# tests/models/multimodal/generation/test_whisper.py
vllm/beam_search.py
+53
-5
@@ -3,7 +3,8 @@
3
4
from dataclasses import dataclass
5
6
-from vllm.inputs import TokenInputs, token_inputs
6
+from vllm.inputs import EncoderDecoderInputs, TokenInputs, token_inputs
7
+from vllm.inputs.data import DecoderInputs
8
from vllm.logprobs import Logprob
9
from vllm.lora.request import LoRARequest
10
from vllm.multimodal.inputs import MultiModalInputs, mm_inputs
@@ -17,9 +18,9 @@ class BeamSearchSequence:
18
about to be returned to the user.
19
"""
20
20
- orig_prompt: TokenInputs | MultiModalInputs
21
+ orig_prompt: TokenInputs | MultiModalInputs | EncoderDecoderInputs
22
22
- # The tokens include the prompt.
23
+ # NOTE: Tokens represents decoder tokens in the encoder / decoder case
24
tokens: list[int]
25
logprobs: list[dict[int, Logprob]]
26
lora_request: LoRARequest | None = None
@@ -31,6 +32,10 @@ class BeamSearchSequence:
32
def get_prompt(self):
33
prompt = self.orig_prompt
34
35
+ if prompt["type"] == "enc_dec":
36
+ return self._build_encoder_decoder_inputs(prompt)
37
+
38
+ # Handle decoder-only inputs
39
prompt_text = prompt.get("prompt")
40
cache_salt = prompt.get("cache_salt")
41
@@ -50,6 +55,44 @@ class BeamSearchSequence:
55
cache_salt=cache_salt,
56
)
57
58
+ def _build_encoder_decoder_inputs(
59
+ self, prompt: EncoderDecoderInputs
60
+ ) -> EncoderDecoderInputs:
61
+ """Rebuild the encoder-decoder inputs with the current beam search
62
+ sequence's tokens.
63
+
64
+ FIXME (alex) - the encoder multimodal cache is not properly wired up
65
+ yet, which means that currently we are running the encoder on every
66
+ new beam because num_computed_tokens is 0 on each new request. This
67
+ will be fixed once the cache is correctly implemented.
68
+ """
69
+ dec_prompt = prompt["decoder_prompt"]
70
+
71
+ # Rebuild decoder prompt with updated tokens,
72
+ # but keep everything else the same.
73
+ new_dec_prompt: DecoderInputs
74
+ if dec_prompt["type"] == "multimodal":
75
+ new_dec_prompt = mm_inputs(
76
+ self.tokens,
77
+ mm_kwargs=dec_prompt["mm_kwargs"],
78
+ mm_hashes=dec_prompt["mm_hashes"],
79
+ mm_placeholders=dec_prompt["mm_placeholders"],
80
+ prompt=dec_prompt.get("prompt"),
81
+ cache_salt=dec_prompt.get("cache_salt"),
82
+ )
83
+ else:
84
+ new_dec_prompt = token_inputs(
85
+ self.tokens,
86
+ prompt=dec_prompt.get("prompt"),
87
+ cache_salt=dec_prompt.get("cache_salt"),
88
+ )
89
+
90
+ return EncoderDecoderInputs(
91
+ type="enc_dec",
92
+ encoder_prompt=prompt["encoder_prompt"],
93
+ decoder_prompt=new_dec_prompt,
94
+ )
95
+
96
97
@dataclass
98
class BeamSearchOutput:
@@ -64,15 +107,20 @@ class BeamSearchOutput:
107
class BeamSearchInstance:
108
def __init__(
109
self,
67
- prompt: TokenInputs | MultiModalInputs,
110
+ prompt: TokenInputs | MultiModalInputs | EncoderDecoderInputs,
111
lora_request: LoRARequest | None = None,
112
logprobs: list[dict[int, Logprob]] | None = None,
113
**kwargs,
114
):
115
+ decoder_prompt = (
116
+ prompt if prompt["type"] != "enc_dec" else prompt["decoder_prompt"]
117
+ )
118
+ initial_tokens = decoder_prompt["prompt_token_ids"]
119
+
120
self.beams: list[BeamSearchSequence] = [
121
BeamSearchSequence(
122
orig_prompt=prompt,
75
- tokens=prompt["prompt_token_ids"],
123
+ tokens=initial_tokens,
124
logprobs=[] if logprobs is None else list(logprobs),
125
lora_request=lora_request,
126
**kwargs,
vllm/entrypoints/llm.py
-4
@@ -734,10 +734,6 @@ class LLM:
734
raise NotImplementedError(
735
"Embedding prompt not supported for beam search"
736
)
737
- if prompt["type"] == "enc_dec":
738
- raise NotImplementedError(
739
- "Encoder-decoder prompt not supported for beam search"
740
- )
737
738
instances.append(
739
BeamSearchInstance(