@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 65a4da1504
Fork 0

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, 01:30 UTC

[Frontend] Add Support for MM Encoder/Decoder Beam Search (Online Transcriptions) (#36160)

Signed-off-by: Alex Brooks <albrooks@redhat.com>

Alex Brooks committed Mar 8, 2026 at 23:46 UTC 65a4da15043f11e86ffcc036f9eb9ad549f0ad17
6 files changed +216 -25
docs/serving/openai_compatible_server.md
+2
@@ -439,6 +439,8 @@ you can use the [official OpenAI Python client](https://github.com/openai/openai
439
440 Code example: [examples/online_serving/openai_transcription_client.py](../../examples/online_serving/openai_transcription_client.py)
441
442 +NOTE: beam search is currently supported in the transcriptions endpoint for encoder-decoder multimodal models, e.g., whisper, but highly inefficient as work for handling the encoder/decoder cache is actively ongoing. This is an active point of ongoing optimization and will be handled properly in the very near future.
443 +
444 #### API Enforced Limits
445
446 Set the maximum audio file size (in MB) that VLLM will accept, via the
tests/entrypoints/openai/test_transcription_validation_whisper.py
+69
@@ -317,3 +317,72 @@ async def test_language_auto_detect(
317 assert any(word.lower() in text_lower for word in expected_text), (
318 f"Expected {expected_lang} text but got: {transcription.text}"
319 )
320 +
321 +
322 +@pytest.mark.asyncio
323 +async def test_whisper_beam_search_single_beam(mary_had_lamb, whisper_client):
324 + """Test beam search with encoder-decoder model (Whisper) on transcriptions with
325 + one beam aligns with greedy decoding.
326 + """
327 + beam_transcription = await whisper_client.audio.transcriptions.create(
328 + model=MODEL_NAME,
329 + file=mary_had_lamb,
330 + language="en",
331 + response_format="text",
332 + temperature=0.0,
333 + extra_body=dict(
334 + use_beam_search=True,
335 + n=1,
336 + ),
337 + )
338 +
339 + greedy_transcription = await whisper_client.audio.transcriptions.create(
340 + model=MODEL_NAME,
341 + file=mary_had_lamb,
342 + response_format="text",
343 + temperature=0.0,
344 + )
345 +
346 + greedy_res = json.loads(greedy_transcription)["text"]
347 + beam_res = json.loads(beam_transcription)["text"]
348 + assert greedy_res == beam_res
349 +
350 +
351 +@pytest.mark.asyncio
352 +async def test_whisper_beam_search_multibeam(mary_had_lamb, whisper_client):
353 + """Test n>1 for beam search returns one transcription (best beam)."""
354 + transcription = await whisper_client.audio.transcriptions.create(
355 + model=MODEL_NAME,
356 + file=mary_had_lamb,
357 + language="en",
358 + response_format="text",
359 + temperature=0.0,
360 + extra_body=dict(
361 + use_beam_search=True,
362 + n=2,
363 + ),
364 + )
365 +
366 + result = json.loads(transcription)
367 +
368 + text = result["text"]
369 +
370 + assert text is not None
371 + assert len(text) > 0
372 + assert "mary had a little lamb" in text.lower()
373 +
374 +
375 +@pytest.mark.asyncio
376 +async def test_stream_with_beams_raises(winning_call, whisper_client):
377 + """Test that stream=True + beam search raises bad request for now."""
378 + with pytest.raises(openai.BadRequestError):
379 + await whisper_client.audio.transcriptions.create(
380 + model=MODEL_NAME,
381 + file=winning_call,
382 + language="en",
383 + stream=True,
384 + extra_body=dict(
385 + use_beam_search=True,
386 + n=2,
387 + ),
388 + )
vllm/entrypoints/openai/completion/serving.py
+6 -4
@@ -129,6 +129,11 @@ class OpenAIServingCompletion(OpenAIServing):
129 - suffix (the language models we currently support do not support
130 suffix)
131 """
132 + if request.stream and request.use_beam_search:
133 + return self.create_error_response(
134 + "Streaming is not currently supported with beam search"
135 + )
136 +
137 result = await self.render_completion_request(request)
138 if isinstance(result, ErrorResponse):
139 return result
@@ -211,13 +216,10 @@ class OpenAIServingCompletion(OpenAIServing):
216 model_name = self.models.model_name(lora_request)
217 num_prompts = len(engine_prompts)
218
214 - # We do not stream the results when using beam search.
215 - stream = request.stream and not request.use_beam_search
216 -
219 # Streaming response
220 tokenizer = self.renderer.tokenizer
221
220 - if stream:
222 + if request.stream:
223 return self.completion_stream_generator(
224 request,
225 engine_prompts,
vllm/entrypoints/openai/engine/serving.py
+7 -6
@@ -237,13 +237,14 @@ class OpenAIServing:
237
238 if prompt["type"] == "embeds":
239 raise NotImplementedError("Embedding prompt not supported for beam search")
240 - if prompt["type"] == "enc_dec":
241 - raise NotImplementedError(
242 - "Encoder-decoder prompt not supported for beam search"
243 - )
240
245 - prompt_text = prompt.get("prompt")
246 - prompt_token_ids = prompt["prompt_token_ids"]
241 + # Extract prompt tokens and text based on model type
242 + decoder_prompt = (
243 + prompt if prompt["type"] != "enc_dec" else prompt["decoder_prompt"]
244 + )
245 + prompt_text = decoder_prompt.get("prompt")
246 + prompt_token_ids = decoder_prompt["prompt_token_ids"]
247 +
248 tokenized_length = len(prompt_token_ids)
249
250 logprobs_num = 2 * beam_width
vllm/entrypoints/openai/speech_to_text/protocol.py
+71
@@ -20,6 +20,7 @@ from vllm.entrypoints.openai.engine.protocol import (
20 from vllm.exceptions import VLLMValidationError
21 from vllm.logger import init_logger
22 from vllm.sampling_params import (
23 + BeamSearchParams,
24 RequestOutputKind,
25 SamplingParams,
26 )
@@ -123,6 +124,18 @@ class TranscriptionRequest(OpenAIBaseModel):
124 """
125
126 # --8<-- [start:transcription-sampling-params]
127 + use_beam_search: bool = False
128 + """Whether or not beam search should be used."""
129 +
130 + n: int = 1
131 + """The number of beams to be used in beam search."""
132 +
133 + length_penalty: float = 1.0
134 + """Length penalty to be used for beam search."""
135 +
136 + include_stop_str_in_output: bool = False
137 + """Whether to include the stop strings in output text."""
138 +
139 temperature: float = Field(default=0.0)
140 """The sampling temperature, between 0 and 1.
141
@@ -170,6 +183,29 @@ class TranscriptionRequest(OpenAIBaseModel):
183 "min_p": 0.0,
184 }
185
186 + def to_beam_search_params(
187 + self,
188 + default_max_tokens: int,
189 + default_sampling_params: dict | None = None,
190 + ) -> BeamSearchParams:
191 + if default_sampling_params is None:
192 + default_sampling_params = {}
193 +
194 + max_tokens = default_max_tokens
195 + n = self.n if self.n is not None else 1
196 +
197 + # NOTE: Temp 0 is a different fallback than completions
198 + if (temperature := self.temperature) is None:
199 + temperature = default_sampling_params.get("temperature", 0)
200 +
201 + return BeamSearchParams(
202 + beam_width=n,
203 + max_tokens=max_tokens,
204 + temperature=temperature,
205 + length_penalty=self.length_penalty,
206 + include_stop_str_in_output=self.include_stop_str_in_output,
207 + )
208 +
209 def to_sampling_params(
210 self, default_max_tokens: int, default_sampling_params: dict | None = None
211 ) -> SamplingParams:
@@ -376,6 +412,18 @@ class TranslationRequest(OpenAIBaseModel):
412
413 # TODO support additional sampling parameters
414 # --8<-- [start:translation-sampling-params]
415 + use_beam_search: bool = False
416 + """Whether or not beam search should be used."""
417 +
418 + n: int = 1
419 + """The number of beams to be used in beam search."""
420 +
421 + length_penalty: float = 1.0
422 + """Length penalty to be used for beam search."""
423 +
424 + include_stop_str_in_output: bool = False
425 + """Whether to include the stop strings in output text."""
426 +
427 seed: int | None = Field(None, ge=_LONG_INFO.min, le=_LONG_INFO.max)
428 """The seed to use for sampling."""
429
@@ -424,6 +472,29 @@ class TranslationRequest(OpenAIBaseModel):
472 "temperature": 0,
473 }
474
475 + def to_beam_search_params(
476 + self,
477 + default_max_tokens: int,
478 + default_sampling_params: dict | None = None,
479 + ) -> BeamSearchParams:
480 + if default_sampling_params is None:
481 + default_sampling_params = {}
482 +
483 + max_tokens = default_max_tokens
484 + n = self.n if self.n is not None else 1
485 +
486 + # NOTE: Temp 0 is a different fallback than completions
487 + if (temperature := self.temperature) is None:
488 + temperature = default_sampling_params.get("temperature", 0)
489 +
490 + return BeamSearchParams(
491 + beam_width=n,
492 + max_tokens=max_tokens,
493 + temperature=temperature,
494 + length_penalty=self.length_penalty,
495 + include_stop_str_in_output=self.include_stop_str_in_output,
496 + )
497 +
498 def to_sampling_params(
499 self, default_max_tokens: int, default_sampling_params: dict | None = None
500 ) -> SamplingParams:
vllm/entrypoints/openai/speech_to_text/speech_to_text.py
+61 -15
@@ -39,7 +39,7 @@ from vllm.entrypoints.openai.speech_to_text.protocol import (
39 )
40 from vllm.entrypoints.utils import get_max_tokens
41 from vllm.exceptions import VLLMValidationError
42 -from vllm.inputs import ProcessorInputs
42 +from vllm.inputs import EncoderDecoderInputs, ProcessorInputs
43 from vllm.logger import init_logger
44 from vllm.logprobs import FlatLogprobs, Logprob
45 from vllm.model_executor.models import (
@@ -50,6 +50,7 @@ from vllm.multimodal.audio import split_audio
50 from vllm.outputs import RequestOutput
51 from vllm.renderers.inputs import DictPrompt, EncoderDecoderDictPrompt
52 from vllm.renderers.inputs.preprocess import parse_enc_dec_prompt, parse_model_prompt
53 +from vllm.sampling_params import BeamSearchParams, SamplingParams
54 from vllm.tokenizers import get_tokenizer
55 from vllm.utils.import_utils import PlaceholderModule
56
@@ -264,8 +265,6 @@ class OpenAISpeechToText(OpenAIServing):
265 via ``get_language_detection_prompt`` and
266 ``parse_language_detection_output``.
267 """
267 - from vllm.sampling_params import SamplingParams
268 -
268 prompt = self.model_cls.get_language_detection_prompt(
269 audio_chunk,
270 self.asr_config,
@@ -403,6 +402,26 @@ class OpenAISpeechToText(OpenAIServing):
402
403 return prompt
404
405 + @staticmethod
406 + def _get_decoder_prompt_len(engine_prompts: list[ProcessorInputs]) -> int:
407 + """Get the length of the decoder prompt. Currently we need to offset
408 + by the decoder prompt length when running beam search because the mm
409 + encoder is not currently cached and runs on decode calls; because of
410 + this, we need to make sure the redundant encoder calls won't exceed
411 + the context :(
412 +
413 + FIXME (Alex) - this will be removed in the very near future once the
414 + encoder/decoder caching is implemented.
415 + """
416 + input_len = 0
417 + assert len(engine_prompts) > 0
418 + first_eng_prompt = engine_prompts[0]
419 +
420 + if first_eng_prompt.get("type") == "enc_dec":
421 + first_eng_prompt = cast(EncoderDecoderInputs, first_eng_prompt)
422 + input_len = len(first_eng_prompt["decoder_prompt"]["prompt_token_ids"])
423 + return input_len
424 +
425 def _get_verbose_segments(
426 self,
427 tokens: tuple,
@@ -481,6 +500,11 @@ class OpenAISpeechToText(OpenAIServing):
500 ) -> T | V | AsyncGenerator[str, None] | ErrorResponse:
501 """Base method for speech-to-text operations like transcription and
502 translation."""
503 + if request.stream and request.use_beam_search:
504 + return self.create_error_response(
505 + "Streaming is not currently supported with beam search"
506 + )
507 +
508 error_check_ret = await self._check_model(request)
509 if error_check_ret is not None:
510 return error_check_ret
@@ -526,6 +550,13 @@ class OpenAISpeechToText(OpenAIServing):
550 # Schedule the request and get the result generator.
551 max_model_len = self.model_config.max_model_len
552 list_result_generator: list[AsyncGenerator[RequestOutput, None]] | None = None
553 +
554 + input_len = (
555 + OpenAISpeechToText._get_decoder_prompt_len(engine_prompts)
556 + if request.use_beam_search
557 + else 0
558 + )
559 +
560 # Unlike most decoder-only models, whisper generation length is not
561 # constrained by the size of the input audio, which is mapped to a
562 # fixed-size log-mel-spectogram. Still, allow for fewer tokens to be
@@ -533,14 +564,20 @@ class OpenAISpeechToText(OpenAIServing):
564 max_tokens = get_max_tokens(
565 max_model_len,
566 request.max_completion_tokens,
536 - 0,
567 + input_len,
568 self.default_sampling_params,
569 )
570
540 - sampling_params = request.to_sampling_params(
541 - max_tokens,
542 - self.default_sampling_params,
543 - )
571 + if request.use_beam_search:
572 + sampling_params = request.to_beam_search_params(
573 + max_tokens, self.default_sampling_params
574 + )
575 + else:
576 + sampling_params = request.to_sampling_params(
577 + max_tokens,
578 + self.default_sampling_params,
579 + )
580 +
581 if request.response_format == "verbose_json":
582 sampling_params.logprobs = 1
583
@@ -561,13 +598,22 @@ class OpenAISpeechToText(OpenAIServing):
598 else await self._get_trace_headers(raw_request.headers)
599 )
600
564 - generator = self.engine_client.generate(
565 - engine_prompt,
566 - sampling_params,
567 - request_id_item,
568 - lora_request=lora_request,
569 - trace_headers=trace_headers,
570 - )
601 + if isinstance(sampling_params, BeamSearchParams):
602 + generator = self.beam_search(
603 + prompt=engine_prompt,
604 + params=sampling_params,
605 + request_id=request_id_item,
606 + lora_request=lora_request,
607 + trace_headers=trace_headers,
608 + )
609 + else:
610 + generator = self.engine_client.generate(
611 + engine_prompt,
612 + sampling_params,
613 + request_id_item,
614 + lora_request=lora_request,
615 + trace_headers=trace_headers,
616 + )
617
618 list_result_generator.append(generator)
619