@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 0efd9f867c

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 11, 15:00 UTC

[Core] Whisper Enable Encoder Batching (#29421)

Signed-off-by: NickLucche <nlucches@redhat.com>

Nicolò Lucchesi committed Dec 11, 2025 at 22:06 UTC 0efd9f867c6a7617fbcb8a335677bb8295f1bcb8

5 files changed +87 -25

vllm/config/model.py

@@ -539,6 +539,11 @@ class ModelConfig:

539

540 self.original_max_model_len = self.max_model_len

541 self.max_model_len = self.get_and_verify_max_len(self.max_model_len)

542 +

543 + if self.is_encoder_decoder:

544 + self.mm_processor_cache_gb = 0

545 + logger.info("Encoder-decoder model detected, disabling mm processor cache.")

546 +

547 # Init multimodal config if needed

548 if self._model_info.supports_multimodal:

549 if (

vllm/config/vllm.py

+10 -20

@@ -750,27 +750,17 @@ class VllmConfig:

750 # TODO: Move after https://github.com/vllm-project/vllm/pull/26847 lands

751 self._set_compile_ranges()

752

753 - if self.model_config and self.model_config.is_encoder_decoder:

754 - from vllm.multimodal import MULTIMODAL_REGISTRY

755 -

756 - self.scheduler_config.max_num_encoder_input_tokens = (

757 - MULTIMODAL_REGISTRY.get_encdec_max_encoder_len(self.model_config)

758 - )

759 - logger.debug(

760 - "Encoder-decoder model detected: setting "

761 - "`max_num_encoder_input_tokens` to encoder length (%s)",

762 - self.scheduler_config.max_num_encoder_input_tokens,

753 + if (

754 + self.model_config

755 + and self.model_config.architecture == "WhisperForConditionalGeneration"

756 + and os.environ.get("VLLM_WORKER_MULTIPROC_METHOD") != "spawn"

757 + ):

758 + logger.warning(

759 + "Whisper is known to have issues with "

760 + "forked workers. If startup is hanging, "

761 + "try setting 'VLLM_WORKER_MULTIPROC_METHOD' "

762 + "to 'spawn'."

763 )

764 - if (

765 - self.model_config.architecture == "WhisperForConditionalGeneration"

766 - and os.environ.get("VLLM_WORKER_MULTIPROC_METHOD") != "spawn"

767 - ):

768 - logger.warning(

769 - "Whisper is known to have issues with "

770 - "forked workers. If startup is hanging, "

771 - "try setting 'VLLM_WORKER_MULTIPROC_METHOD' "

772 - "to 'spawn'."

773 - )

764

765 if (

766 self.kv_events_config is not None

vllm/model_executor/models/whisper.py

+13 -4

@@ -522,6 +522,7 @@ class WhisperEncoder(nn.Module):

522

523 def forward(self, input_features: torch.Tensor | list[torch.Tensor]):

524 hidden_states = []

525 + input_is_batched = False

526 for features in input_features:

527 embeds = nn.functional.gelu(self.conv1(features))

528 embeds = nn.functional.gelu(self.conv2(embeds))

@@ -530,7 +531,13 @@ class WhisperEncoder(nn.Module):

531 embeds.dtype

532 )

533 hidden_states.append(embeds)

533 - hidden_states = torch.cat(hidden_states)

534 + input_is_batched = embeds.ndim > 2

535 + # Input to MHA must be B x T x D

536 + if input_is_batched:

537 + # Models using WhisperEncoder may handle batching internally.

538 + hidden_states = torch.cat(hidden_states)

539 + else:

540 + hidden_states = torch.stack(hidden_states, dim=0)

541

542 for encoder_layer in self.layers:

543 hidden_states = encoder_layer(hidden_states)

@@ -603,8 +610,7 @@ class WhisperModel(nn.Module):

610 positions: torch.Tensor,

611 encoder_outputs: list[torch.Tensor],

612 ) -> torch.Tensor:

606 - assert len(encoder_outputs) in (0, 1)

607 - enc_states = encoder_outputs[0] if len(encoder_outputs) == 1 else None

613 + enc_states = torch.cat(encoder_outputs, dim=0) if len(encoder_outputs) else None

614 decoder_outputs = self.decoder(

615 input_ids=input_ids,

616 positions=positions,

@@ -913,7 +919,10 @@ class WhisperForConditionalGeneration(

919 def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:

920 # Required as part of SupportsMultiModal interface.

921 audio_input = self._parse_and_validate_audio_input(**kwargs)

916 - return [self.model.get_encoder_outputs(audio_input["input_features"])]

922 + # Split concatenated encoder outputs into one tensor per audio input

923 + enc_output = self.model.get_encoder_outputs(audio_input["input_features"])

924 + # The assumption is we can only process whole mm items (audios)

925 + return enc_output.unbind(dim=0)

926

927 def embed_input_ids(

928 self,

vllm/v1/core/encoder_cache_manager.py

+53

@@ -341,3 +341,56 @@ def compute_mm_encoder_budget(

341 )

342

343 return encoder_compute_budget, encoder_cache_size

344 +

345 +

346 +# NOTE (NickLucche): Temporary implementation for encoder-decoder models that only

347 +# use the manager for scheduling purposes. Encoder-decoder models will eventually

348 +# utilize the cache and this class will fold into EncoderCacheManager, as

349 +# differences with MM models shrink.

350 +class EncoderDecoderCacheManager(EncoderCacheManager):

351 + def __init__(self, cache_size: int):

352 + self.cache_size = cache_size

353 + self.num_free_slots = cache_size

354 + self.freed: list[str] = []

355 +

356 + def check_and_update_cache(self, request: Request, input_id: int) -> bool:

357 + return False

358 +

359 + def can_allocate(

360 + self,

361 + request: Request,

362 + input_id: int,

363 + encoder_compute_budget: int,

364 + num_tokens_to_schedule: int,

365 + ) -> bool:

366 + num_tokens = request.get_num_encoder_tokens(input_id)

367 + # Not enough compute budget

368 + if num_tokens > encoder_compute_budget:

369 + return False

370 +

371 + num_tokens += num_tokens_to_schedule

372 + # Enough free slots

373 + return num_tokens <= self.num_free_slots

374 +

375 + def allocate(self, request: Request, input_id: int) -> None:

376 + num_encoder_tokens = request.get_num_encoder_tokens(input_id)

377 + self.num_free_slots -= num_encoder_tokens

378 +

379 + mm_hash = request.mm_features[input_id].identifier

380 + self.freed.append(mm_hash)

381 +

382 + def free(self, request: Request) -> None:

383 + for input_id in range(len(request.mm_features)):

384 + self.free_encoder_input(request, input_id)

385 +

386 + def get_cached_input_ids(self, request: Request) -> set[int]:

387 + return set(range(len(request.mm_features)))

388 +

389 + def get_freed_mm_hashes(self) -> list[str]:

390 + freed = self.freed

391 + self.freed = []

392 + return freed

393 +

394 + def free_encoder_input(self, request: Request, input_id: int) -> None:

395 + num_tokens = request.get_num_encoder_tokens(input_id)

396 + self.num_free_slots += num_tokens

vllm/v1/core/sched/scheduler.py

+6 -1

@@ -27,6 +27,7 @@ from vllm.logger import init_logger

27 from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry

28 from vllm.v1.core.encoder_cache_manager import (

29 EncoderCacheManager,

30 + EncoderDecoderCacheManager,

31 compute_encoder_budget,

32 )

33 from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager

@@ -181,7 +182,11 @@ class Scheduler(SchedulerInterface):

182 # NOTE: For the models without encoder (e.g., text-only models),

183 # the encoder cache will not be initialized because cache size is 0

184 # for these models.

184 - self.encoder_cache_manager = EncoderCacheManager(cache_size=encoder_cache_size)

185 + self.encoder_cache_manager = (

186 + EncoderDecoderCacheManager(cache_size=encoder_cache_size)

187 + if self.is_encoder_decoder

188 + else EncoderCacheManager(cache_size=encoder_cache_size)

189 + )

190

191 speculative_config = vllm_config.speculative_config

192 self.use_eagle = False