@tduncan / Nvidia-Optimization-Stack-VLLM / commits / b908a21f9a

[Model][LoRA] Add tower/connector LoRA support for Ultravox (#48215)

Signed-off-by: arthurgao2003 <arthurgao2003@gmail.com> Signed-off-by: Arthur Gao <arthurgao2003@gmail.com> Signed-off-by: 高杨懿 <15312196+gao-yangyi@user.noreply.gitee.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: 高杨懿 <15312196+gao-yangyi@user.noreply.gitee.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: OpenAI Codex <noreply@openai.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Yangyi Gao committed Aug 13, 2026 at 13:57 UTC b908a21f9a449592ca06352381d9118dbc39c6d2

2 files changed +345 -129

vllm/model_executor/models/ultravox.py

+277 -119

@@ -4,27 +4,26 @@

4 # Adapted from https://github.com/fixie-ai/ultravox/blob/ecd58c4041030bae2ad15aa6bcf04ab43199ea02/ultravox/model/ultravox_model.py

5 """PyTorch Ultravox model."""

7 -import copy

7 +import math

8 from collections.abc import Iterable, Mapping, Sequence

9 -from types import SimpleNamespace

9 from typing import Annotated, Any, Literal, TypeAlias

11 +import numpy as np

12 import torch

13 from torch import nn

14 from torch.nn import functional as F

15 from transformers import BatchFeature, ProcessorMixin

16 -from transformers.modeling_utils import ModuleUtilsMixin

16 from transformers.models.whisper import WhisperFeatureExtractor

18 -from transformers.models.whisper.modeling_whisper import (

19 - WhisperEncoder,

20 - WhisperEncoderLayer,

21 -)

18 from vllm.config import VllmConfig

19 from vllm.config.multimodal import BaseDummyOptions

20 +from vllm.distributed import get_tensor_model_parallel_world_size

21 from vllm.inputs import MultiModalDataDict

22 from vllm.model_executor.layers.activation import MulAndSilu, get_act_fn

23 +from vllm.model_executor.layers.attention import MMEncoderAttention

24 from vllm.model_executor.layers.layernorm import RMSNorm

25 +from vllm.model_executor.layers.linear import ReplicatedLinear

26 +from vllm.model_executor.layers.quantization import QuantizationConfig

27 from vllm.model_executor.model_loader import DefaultModelLoader

28 from vllm.model_executor.models.module_mapping import MultiModelKeys

29 from vllm.multimodal import MULTIMODAL_REGISTRY

@@ -59,6 +58,11 @@ from .utils import (

58 init_vllm_registered_model,

59 maybe_prefix,

60 )

61 +from .whisper import (

62 + WhisperEncoder,

63 + WhisperEncoderLayer,

64 + _create_fake_bias_for_k_proj,

65 +)

67 _AUDIO_PLACEHOLDER_OVERRIDE = "<|audio|>"

68 _MAX_ENCODER_BATCH_SIZE = 16

@@ -304,14 +308,97 @@ class StackAudioFrames(nn.Module):

308 return audio_embeds

309

310

311 +def _build_chunk_attn_metadata(

312 + attn: MMEncoderAttention,

313 + feature_lens: torch.Tensor,

314 + seq_len: int,

315 + hidden_size: int,

316 + device: torch.device,

317 +) -> dict[str, torch.Tensor | None]:

318 + """Segmented varlen attention metadata for a padded batch of audio chunks.

319 +

320 + Each padded row contributes up to two sequences to `cu_seqlens`: its valid

321 + frames and its padding tail. Attention therefore never crosses a

322 + valid/padding boundary (equivalent to the key-padding mask the HF

323 + implementation uses), while every row still flows through the

324 + (potentially LoRA-wrapped) linears, keeping the per-chunk token counts

325 + constant as required by `get_num_mm_encoder_tokens` /

326 + `get_num_mm_connector_tokens`. Queries at padding positions produce

327 + (garbage) outputs, which are trimmed by `audio_token_len` downstream.

328 + """

329 + batch_size = feature_lens.shape[0]

330 + starts = np.arange(batch_size, dtype=np.int64) * seq_len

331 + lens_np = feature_lens.cpu().numpy().astype(np.int64)

332 + bounds = np.stack([starts + lens_np, starts + seq_len], axis=1).reshape(-1)

333 + cu_seqlens_np = np.concatenate(([0], bounds))

334 + # Fully-valid rows produce empty padding segments; drop the duplicates.

335 + cu_seqlens_np = np.unique(cu_seqlens_np).astype(np.int32)

336 +

337 + attn_backend = attn.attn_backend

338 + sequence_lengths = MMEncoderAttention.maybe_compute_seq_lens(

339 + attn_backend, cu_seqlens_np, device

340 + )

341 + max_seqlen = torch.tensor(

342 + MMEncoderAttention.compute_max_seqlen(attn_backend, cu_seqlens_np),

343 + dtype=torch.int32,

344 + )

345 + cu_seqlens = MMEncoderAttention.maybe_recompute_cu_seqlens(

346 + attn_backend,

347 + cu_seqlens_np,

348 + hidden_size,

349 + get_tensor_model_parallel_world_size(),

350 + device,

351 + )

352 + return {

353 + "cu_seqlens": cu_seqlens,

354 + "max_seqlen": max_seqlen,

355 + "sequence_lengths": sequence_lengths,

356 + }

357 +

358 +

359 +# Maps HF whisper encoder layer names to vLLM's fused qkv_proj and nested

360 +# mlp; `_create_fake_bias_for_k_proj` supplies the fused bias' k-slice

361 +# (HF whisper k_proj has no bias).

362 +_WHISPER_ENCODER_MAPPER = WeightsMapper(

363 + orig_to_new_substr={".fc1.": ".mlp.fc1.", ".fc2.": ".mlp.fc2."},

364 + orig_to_new_stacked={

365 + ".self_attn.q_proj": (".self_attn.qkv_proj", "q"),

366 + ".self_attn.k_proj": (".self_attn.qkv_proj", "k"),

367 + ".self_attn.v_proj": (".self_attn.qkv_proj", "v"),

368 + },

369 +)

370 +

371 +

372 +def _load_whisper_layer_weights(

373 + module: nn.Module, weights: Iterable[tuple[str, torch.Tensor]]

374 +) -> set[str]:

375 + weights = _create_fake_bias_for_k_proj(weights, ".self_attn.k_proj.weight")

376 + loader = AutoWeightsLoader(module)

377 + return loader.load_weights(weights, mapper=_WHISPER_ENCODER_MAPPER)

378 +

379 +

380 class UltravoxFeedForwardProjector(nn.Module):

308 - def __init__(self, config: UltravoxConfig):

381 + def __init__(

382 + self,

383 + config: UltravoxConfig,

384 + *,

385 + quant_config: QuantizationConfig | None = None,

386 + prefix: str = "",

387 + ):

388 super().__init__()

389 self.hidden_dim = config.hidden_size

390 self._pad_and_stack = StackAudioFrames(config.stack_factor)

391 dim_in = config.audio_config.hidden_size * config.stack_factor

392 self.ln_pre = RMSNorm(dim_in)

314 - self.linear_1 = nn.Linear(dim_in, self.hidden_dim, bias=False)

393 + # ReplicatedLinear (a vLLM-native linear) so the connector can be

394 + # wrapped for LoRA; a plain nn.Linear is left unwrapped by from_layer.

395 + self.linear_1 = ReplicatedLinear(

396 + dim_in,

397 + self.hidden_dim,

398 + bias=False,

399 + quant_config=quant_config,

400 + prefix=f"{prefix}.linear_1",

401 + )

402 dim_mid = self.hidden_dim

403

404 if config.projector_act == "swiglu":

@@ -321,7 +408,13 @@ class UltravoxFeedForwardProjector(nn.Module):

408 self.act = get_act_fn(config.projector_act)

409

410 dim_out = config.text_config.hidden_size

324 - self.linear_2 = nn.Linear(dim_mid, dim_out, bias=False)

411 + self.linear_2 = ReplicatedLinear(

412 + dim_mid,

413 + dim_out,

414 + bias=False,

415 + quant_config=quant_config,

416 + prefix=f"{prefix}.linear_2",

417 + )

418

419 # Ultravox v0.4.1 and below use layer_norm after the second linear layer

420 # while v0.5.0 and above uses layer_norm after the first linear layer.

@@ -337,42 +430,64 @@ class UltravoxFeedForwardProjector(nn.Module):

430 ) -> torch.Tensor:

431 audio_features = self._pad_and_stack(audio_features)

432 audio_features = self.ln_pre(audio_features)

340 - hidden_states = self.linear_1(audio_features)

433 + hidden_states, _ = self.linear_1(audio_features)

434 hidden_states = self.act(hidden_states)

435 hidden_states = self.ln_mid(hidden_states)

343 - hidden_states = self.linear_2(hidden_states)

436 + hidden_states, _ = self.linear_2(hidden_states)

437 hidden_states = self.ln_post(hidden_states)

438 return hidden_states

439

440

348 -class UltravoxTransformerProjector(nn.Module, ModuleUtilsMixin):

349 - def __init__(self, config: UltravoxConfig):

441 +class UltravoxTransformerProjector(nn.Module):

442 + def __init__(

443 + self,

444 + vllm_config: VllmConfig,

445 + *,

446 + prefix: str = "",

447 + ):

448 super().__init__()

351 - self.config = SimpleNamespace(is_decoder=False)

449 + config: UltravoxConfig = vllm_config.model_config.hf_config

450 + quant_config = vllm_config.quant_config

451 + audio_config = config.audio_config

452

453 self._pad_and_stack = StackAudioFrames(config.stack_factor)

354 - dim_in = config.audio_config.hidden_size * config.stack_factor

355 -

356 - projector_audio_config = copy.deepcopy(config.audio_config)

454 + dim_in = audio_config.hidden_size * config.stack_factor

455

456 self.ln_pre = RMSNorm(dim_in)

359 - self.linear_in = nn.Linear(dim_in, projector_audio_config.d_model)

457 + self.linear_in = ReplicatedLinear(

458 + dim_in,

459 + audio_config.d_model,

460 + bias=True,

461 + quant_config=quant_config,

462 + prefix=f"{prefix}.linear_in",

463 + )

464

465 self.embed_positions = nn.Embedding(

362 - projector_audio_config.max_source_positions,

363 - projector_audio_config.d_model,

466 + audio_config.max_source_positions,

467 + audio_config.d_model,

468 )

469

470 + audio_vllm_config = vllm_config.with_hf_config(

471 + audio_config, architectures=["UltravoxModel"]

472 + )

473 self.layers = nn.ModuleList(

474 [

368 - WhisperEncoderLayer(projector_audio_config)

369 - for _ in range(config.num_projector_layers)

475 + WhisperEncoderLayer(

476 + vllm_config=audio_vllm_config,

477 + prefix=f"{prefix}.layers.{idx}",

478 + )

479 + for idx in range(config.num_projector_layers)

480 ]

481 )

372 -

373 - self.ln_post = RMSNorm(projector_audio_config.d_model)

374 - self.linear_out = nn.Linear(

375 - projector_audio_config.d_model, config.text_config.hidden_size

482 + self._d_model = audio_config.d_model

483 +

484 + self.ln_post = RMSNorm(audio_config.d_model)

485 + self.linear_out = ReplicatedLinear(

486 + audio_config.d_model,

487 + config.text_config.hidden_size,

488 + bias=True,

489 + quant_config=quant_config,

490 + prefix=f"{prefix}.linear_out",

491 )

492

493 def forward(

@@ -380,100 +495,65 @@ class UltravoxTransformerProjector(nn.Module, ModuleUtilsMixin):

495 ) -> torch.Tensor:

496 audio_features = self._pad_and_stack(audio_features)

497

383 - max_len_stacked = audio_features.shape[1]

384 - attention_mask = torch.arange(max_len_stacked, device=audio_features.device)[

385 - None, :

386 - ].lt(audio_token_len[:, None])

387 - extended_attention_mask = self.get_extended_attention_mask(

388 - attention_mask, attention_mask.shape, audio_features.dtype

389 - )

390 -

498 hidden_states = self.ln_pre(audio_features)

392 - hidden_states = self.linear_in(hidden_states)

499 + hidden_states, _ = self.linear_in(hidden_states)

500

501 positions = self.embed_positions(

502 torch.arange(hidden_states.size(1), device=hidden_states.device)

503 )

504 hidden_states = hidden_states + positions

505

506 + batch_size, seq_len, d_model = hidden_states.shape

507 + attn_metadata = _build_chunk_attn_metadata(

508 + self.layers[0].self_attn.attn,

509 + audio_token_len,

510 + seq_len,

511 + d_model,

512 + hidden_states.device,

513 + )

514 + hidden_states = hidden_states.reshape(batch_size * seq_len, d_model)

515 for layer in self.layers:

400 - hidden_states = layer(

401 - hidden_states,

402 - attention_mask=extended_attention_mask,

403 - )

404 - # BC version that allows for the old tupled output

405 - if isinstance(hidden_states, tuple):

406 - hidden_states = hidden_states[0]

516 + hidden_states = layer(hidden_states, **attn_metadata)

517 + hidden_states = hidden_states.reshape(batch_size, seq_len, d_model)

518

519 hidden_states = self.ln_post(hidden_states)

409 - hidden_states = self.linear_out(hidden_states)

520 + hidden_states, _ = self.linear_out(hidden_states)

521 return hidden_states

522

523 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

524 + return _load_whisper_layer_weights(self, weights)

525

413 -class ModifiedWhisperEncoder(WhisperEncoder):

414 - """

415 - Encoder portion of OpenAI's Whisper model.

416 -

417 - This implementation is a slightly modified version of HF Transformers'

418 - Whisper Encoder, with only a few fixes:

419 - 1. base_model_prefix updated to allow for doing `.from_pretrained`

420 - directly on the encoder

421 - 2. allow less than 30 second of audio padding to be passed in:

422 - - relaxed ValueError check for `input_features` length to be less

423 - than or equal to `expected_seq_length` instead of strictly equal

424 - - embed_pos is now sliced to match the length of `inputs_embeds`

425 -

426 - Original: https://github.com/huggingface/transformers/blob/main/src/transformers/models/whisper/modeling_whisper.py

427 - See commentary: https://github.com/huggingface/transformers/issues/25744

428 - """

526

430 - base_model_prefix = "model.encoder"

527 +class UltravoxWhisperEncoder(WhisperEncoder):

528 + """Ultravox's `ModifiedWhisperEncoder` on top of vLLM's `WhisperEncoder`.

529

432 - def __init__(self, *args, **kwargs):

433 - super().__init__(*args, **kwargs)

434 - self.config.is_decoder = False

530 + Like the original (a modified HF whisper encoder,

531 + see https://github.com/huggingface/transformers/issues/25744), it accepts

532 + mel inputs shorter than 30s (positions are sliced to the input length) and

533 + confines attention to each chunk's valid frames based on `audio_lens`

534 + (via segmented `cu_seqlens` instead of a dense key-padding mask), so its

535 + outputs match the HF implementation for valid positions. The linears are

536 + vLLM-native so the tower can be wrapped for LoRA.

537 + """

538

539 @property

437 - def max_context_length(self):

438 - return (

439 - self.config.max_source_positions

440 - * self.conv1.stride[0]

441 - * self.conv2.stride[0]

442 - )

540 + def dtype(self) -> torch.dtype:

541 + return self.conv1.weight.dtype

542

444 - def get_attention_mask_by_audio_len(

445 - self, audio_lens: torch.Tensor | None, hidden_states: torch.Tensor

446 - ):

447 - """

448 - Create attention mask based on audio lengths to mask out padding tokens

449 - For each sample in batch:

450 - - Convert raw audio length to feature length after convolutions

451 - - Create bool mask: True for valid positions and False for padding

452 - - Convert to attention mask format expected by transformer layers

453 - (1.0 for positions to attend to, large negative for positions to ignore)

454 - This masking ensures consistent behavior between training and inference

455 - by preventing the model from attending to padding tokens in both cases

456 - """

457 - if audio_lens is None:

458 - return None

543 + @property

544 + def max_context_length(self) -> int:

545 + return self.max_source_positions * self.total_stride

546

460 - audio_feature_len = self._get_feat_extract_output_lengths(audio_lens)

461 - max_seq_len = hidden_states.shape[1]

462 - attention_mask = torch.arange(max_seq_len, device=hidden_states.device)[

463 - None, :

464 - ].lt(audio_feature_len.view(-1, 1))

465 - attention_mask = self.get_extended_attention_mask(

466 - attention_mask,

467 - None,

468 - dtype=hidden_states.dtype,

469 - )

470 - return attention_mask

547 + def _get_feat_extract_output_lengths(

548 + self, input_lengths: torch.Tensor

549 + ) -> torch.Tensor:

550 + return (input_lengths - 1) // 2 + 1

551

552 def forward(

553 self,

554 input_features: torch.Tensor,

475 - audio_lens: torch.Tensor | None = None,

476 - ):

555 + audio_lens: torch.Tensor,

556 + ) -> torch.Tensor:

557 expected_seq_length = self.max_context_length

558 if input_features.shape[-1] > expected_seq_length:

559 raise ValueError(

@@ -489,25 +569,27 @@ class ModifiedWhisperEncoder(WhisperEncoder):

569 inputs_embeds = inputs_embeds.permute(0, 2, 1)

570 embed_pos = self.embed_positions.weight[: inputs_embeds.size(-2)]

571

492 - hidden_states = inputs_embeds + embed_pos

493 - hidden_states = nn.functional.dropout(

494 - hidden_states, p=self.dropout, training=self.training

495 - )

496 -

497 - attention_mask = self.get_attention_mask_by_audio_len(audio_lens, hidden_states)

572 + hidden_states = (inputs_embeds + embed_pos).to(inputs_embeds.dtype)

573

574 + batch_size, seq_len, d_model = hidden_states.shape

575 + attn_metadata = _build_chunk_attn_metadata(

576 + self.layers[0].self_attn.attn,

577 + self._get_feat_extract_output_lengths(audio_lens),

578 + seq_len,

579 + d_model,

580 + hidden_states.device,

581 + )

582 + hidden_states = hidden_states.reshape(batch_size * seq_len, d_model)

583 for encoder_layer in self.layers:

500 - hidden_states = encoder_layer(

501 - hidden_states,

502 - attention_mask,

503 - )

504 - # BC version that allows for the old tupled output

505 - if isinstance(hidden_states, tuple):

506 - hidden_states = hidden_states[0]

584 + hidden_states = encoder_layer(hidden_states, **attn_metadata)

585 + hidden_states = hidden_states.reshape(batch_size, seq_len, d_model)

586

587 hidden_states = self.layer_norm(hidden_states)

588 return hidden_states

589

590 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

591 + return _load_whisper_layer_weights(self, weights)

592 +

593

594 @MULTIMODAL_REGISTRY.register_processor(

595 UltravoxMultiModalProcessor,

@@ -521,7 +603,13 @@ class UltravoxModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):

603 }

604

605 hf_to_vllm_mapper = WeightsMapper(

524 - orig_to_new_prefix={"audio_tower.model.encoder.": "audio_tower."}

606 + orig_to_new_prefix={

607 + "audio_tower.model.encoder.": "audio_tower.",

608 + # A whisper checkpoint loaded via `audio_model_id` also carries

609 + # decoder and LM-head weights the tower does not use.

610 + "audio_tower.model.": None,

611 + "audio_tower.proj_out.": None,

612 + }

613 )

614

615 @classmethod

@@ -535,10 +623,20 @@ class UltravoxModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):

623 super().__init__()

624 config: UltravoxConfig = vllm_config.model_config.hf_config

625 multimodal_config = vllm_config.model_config.multimodal_config

626 + quant_config = vllm_config.quant_config

627 + lora_config = vllm_config.lora_config

628 self.config = config

629 self.multi_modal_config = multimodal_config

630 assert self.multi_modal_config

631

632 + # LoRA on the tower/connector requires per-item token counts that are

633 + # predictable from the placeholder count alone (see

634 + # `get_num_mm_encoder_tokens`), so pad every audio chunk's mel input

635 + # to the tower's full context instead of the batch's max length.

636 + self.pad_audio_to_max_context = bool(

637 + lora_config is not None and lora_config.enable_tower_connector_lora

638 + )

639 +

640 self.configure_mm_token_handling(

641 self.config.vocab_size,

642 [self.config.audio_token_index],

@@ -569,11 +667,24 @@ class UltravoxModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):

667 )

668

669 with self._mark_tower_model(vllm_config, "audio"):

572 - self.audio_tower = ModifiedWhisperEncoder(config.audio_config)

670 + self.audio_tower = UltravoxWhisperEncoder(

671 + vllm_config=vllm_config.with_hf_config(

672 + config.audio_config, architectures=["UltravoxModel"]

673 + ),

674 + prefix=maybe_prefix(prefix, "audio_tower"),

675 + enable_pp=False,

676 + )

677 if config.num_projector_layers > 0:

574 - self.multi_modal_projector = UltravoxTransformerProjector(config)

678 + self.multi_modal_projector = UltravoxTransformerProjector(

679 + vllm_config,

680 + prefix=maybe_prefix(prefix, "multi_modal_projector"),

681 + )

682 else:

576 - self.multi_modal_projector = UltravoxFeedForwardProjector(config)

683 + self.multi_modal_projector = UltravoxFeedForwardProjector(

684 + config,

685 + quant_config=quant_config,

686 + prefix=maybe_prefix(prefix, "multi_modal_projector"),

687 + )

688

689 with self._mark_language_model(vllm_config):

690 self.language_model = init_vllm_registered_model(

@@ -596,6 +707,32 @@ class UltravoxModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):

707 tower_model="audio_tower.",

708 )

709

710 + def _get_max_tokens_per_chunk(self) -> int:

711 + # Audio is chunked at the tower's full context (30s -> 1500 encoder

712 + # frames); the projector stacks `stack_factor` frames per LM token, so

713 + # a full chunk yields ceil(max_source_positions / stack_factor) tokens

714 + # and only an audio's last chunk can yield fewer.

715 + return math.ceil(

716 + self.config.audio_config.max_source_positions / self.config.stack_factor

717 + )

718 +

719 + def get_num_mm_encoder_tokens(self, num_audio_tokens: int) -> int:

720 + # With `pad_audio_to_max_context` (enforced when tower/connector LoRA

721 + # is enabled), the tower's LoRA-wrapped linears always run on

722 + # `max_source_positions` conv-downsampled tokens per chunk, regardless

723 + # of the valid frame count.

724 + num_chunks = math.ceil(num_audio_tokens / self._get_max_tokens_per_chunk())

725 + return num_chunks * self.config.audio_config.max_source_positions

726 +

727 + def get_num_mm_connector_tokens(self, num_encoder_tokens: int) -> int:

728 + # The connector runs on the frame-stacked tower output. Stacking pads

729 + # each chunk to a multiple of `stack_factor`, so this is

730 + # ceil(max_source_positions / stack_factor) tokens per chunk (188 for

731 + # whisper's 1500), not `num_encoder_tokens // stack_factor` (187).

732 + max_source_positions = self.config.audio_config.max_source_positions

733 + num_chunks = num_encoder_tokens // max_source_positions

734 + return num_chunks * self._get_max_tokens_per_chunk()

735 +

736 def _audio_features_to_embeddings(

737 self,

738 input_features: torch.Tensor,

@@ -606,9 +743,19 @@ class UltravoxModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):

743 batch_size = audio_features.size(0)

744 audio_embeddings = []

745

609 - # Process audio features in batches to keep memory usage predictable

610 - for start in range(0, batch_size, _MAX_ENCODER_BATCH_SIZE):

611 - end = min(start + _MAX_ENCODER_BATCH_SIZE, batch_size)

746 + # Process audio features in batches to keep memory usage predictable.

747 + # With tower/connector LoRA, the punica kernels always map the first

748 + # `x.size(0)` entries of the token-LoRA mapping (set once for the whole

749 + # scheduled encoder batch) to the rows of each linear's input, so

750 + # splitting the batch would apply the wrong LoRA ids to every chunk

751 + # after the first sub-batch; process all chunks in a single pass.

752 + encoder_batch_size = (

753 + max(batch_size, 1)

754 + if self.pad_audio_to_max_context

755 + else _MAX_ENCODER_BATCH_SIZE

756 + )

757 + for start in range(0, batch_size, encoder_batch_size):

758 + end = min(start + encoder_batch_size, batch_size)

759 # Process through audio tower

760 batch_features = self.audio_tower(

761 audio_features[start:end], audio_lens[start:end]

@@ -662,6 +809,17 @@ class UltravoxModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):

809 # [[B1, 80, M1], [B2, 80, M2]] -> [B1+B2, 80, max(M1, M2)]

810 audio_features = pad_and_concat_to_dim3(audio_input["data"])

811

812 + if self.pad_audio_to_max_context:

813 + # Pad every chunk to the tower's full context so the number of

814 + # tokens processed by the tower/connector per chunk is constant,

815 + # keeping the LoRA token mappings computed by

816 + # `get_num_mm_encoder_tokens`/`get_num_mm_connector_tokens` exact.

817 + # Attention to the extra padding is masked via `audio_lens`.

818 + # Over-long inputs are not trimmed here; the tower raises on them.

819 + pad_len = self.audio_tower.max_context_length - audio_features.shape[-1]

820 + if pad_len > 0:

821 + audio_features = F.pad(audio_features, (0, pad_len))

822 +

823 audio_lens = audio_input["lens"]

824 audio_token_len = audio_input["token_len"]

825

@@ -755,7 +913,7 @@ class UltravoxModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):

913 return self.language_model.compute_logits(hidden_states)

914

915 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

758 - loader = AutoWeightsLoader(self, ignore_unexpected_prefixes=["audio_tower."])

916 + loader = AutoWeightsLoader(self)

917 return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

918

919

vllm/model_executor/models/whisper.py

+68 -10

@@ -47,6 +47,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead

47 from vllm.model_executor.models.whisper_utils import (

48 ISO639_1_SUPPORTED_LANGS,

49 )

50 +from vllm.model_executor.offloader import get_offloader

51 from vllm.multimodal import MULTIMODAL_REGISTRY

52 from vllm.multimodal.inputs import (

53 MultiModalFieldConfig,

@@ -114,6 +115,10 @@ class WhisperEncoderAttention(MMEncoderAttention):

115 query: torch.Tensor,

116 key: torch.Tensor,

117 value: torch.Tensor,

118 + cu_seqlens: torch.Tensor | None = None,

119 + max_seqlen: torch.Tensor | None = None, # Only used for Flash Attention

120 + # Only used for FlashInfer CuDNN backend.

121 + sequence_lengths: torch.Tensor | None = None,

122 ) -> torch.Tensor:

123 """

124 Input shape: batch_size x seq_len x hidden_size

@@ -126,7 +131,14 @@ class WhisperEncoderAttention(MMEncoderAttention):

131 value = value.unsqueeze(0)

132

133 # Call the parent forward method

129 - out = super().forward(query, key, value)

134 + out = super().forward(

135 + query,

136 + key,

137 + value,

138 + cu_seqlens=cu_seqlens,

139 + max_seqlen=max_seqlen,

140 + sequence_lengths=sequence_lengths,

141 + )

142

143 if is_2d:

144 out = out.squeeze(0)

@@ -240,11 +252,29 @@ class WhisperAttention(nn.Module):

252 def forward(

253 self,

254 hidden_states: torch.Tensor,

255 + cu_seqlens: torch.Tensor | None = None,

256 + max_seqlen: torch.Tensor | None = None, # Only used for Flash Attention

257 + # Only used for FlashInfer CuDNN backend.

258 + sequence_lengths: torch.Tensor | None = None,

259 ):

260 qkv, _ = self.qkv_proj(hidden_states)

261 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

262

247 - attn_output = self.attn(q, k, v)

263 + if cu_seqlens is None:

264 + attn_output = self.attn(q, k, v)

265 + else:

266 + assert self.attn_type == AttentionType.ENCODER, (

267 + "Variable-length attention metadata is only supported for "

268 + "encoder self-attention."

269 + )

270 + attn_output = self.attn(

271 + q,

272 + k,

273 + v,

274 + cu_seqlens=cu_seqlens,

275 + max_seqlen=max_seqlen,

276 + sequence_lengths=sequence_lengths,

277 + )

278

279 output, _ = self.out_proj(attn_output)

280

@@ -381,10 +411,19 @@ class WhisperEncoderLayer(nn.Module):

411 def forward(

412 self,

413 hidden_states: torch.Tensor,

414 + cu_seqlens: torch.Tensor | None = None,

415 + max_seqlen: torch.Tensor | None = None, # Only used for Flash Attention

416 + # Only used for FlashInfer CuDNN backend.

417 + sequence_lengths: torch.Tensor | None = None,

418 ):

419 residual = hidden_states

420 hidden_states = self.self_attn_layer_norm(hidden_states)

387 - hidden_states = self.self_attn(hidden_states=hidden_states)

421 + hidden_states = self.self_attn(

422 + hidden_states=hidden_states,

423 + cu_seqlens=cu_seqlens,

424 + max_seqlen=max_seqlen,

425 + sequence_lengths=sequence_lengths,

426 + )

427 hidden_states = residual + hidden_states

428 residual = hidden_states

429 hidden_states = self.final_layer_norm(hidden_states)

@@ -457,7 +496,12 @@ class WhisperDecoderLayer(nn.Module):

496

497 class WhisperEncoder(nn.Module):

498 def __init__(

460 - self, *, vllm_config: VllmConfig, prefix: str = "", init_in_fp32: bool = False

499 + self,

500 + *,

501 + vllm_config: VllmConfig,

502 + prefix: str = "",

503 + init_in_fp32: bool = False,

504 + enable_pp: bool = True,

505 ):

506 super().__init__()

507 config = vllm_config.model_config.hf_config

@@ -474,13 +518,27 @@ class WhisperEncoder(nn.Module):

518 self.conv2 = nn.Conv1d(embed_dim, embed_dim, stride=2, kernel_size=3, padding=1)

519

520 self.total_stride = self.conv1.stride[0] * self.conv2.stride[0]

477 - self.start_layer, self.end_layer, self.layers = make_layers(

478 - config.encoder_layers,

479 - lambda prefix: WhisperEncoderLayer(

521 +

522 + def create_layer(prefix: str) -> nn.Module:

523 + return WhisperEncoderLayer(

524 vllm_config=vllm_config, prefix=f"{prefix}.layers"

481 - ),

482 - prefix=f"{prefix}.layers",

483 - )

525 + )

526 +

527 + if enable_pp:

528 + self.start_layer, self.end_layer, self.layers = make_layers(

529 + config.encoder_layers,

530 + create_layer,

531 + prefix=f"{prefix}.layers",

532 + )

533 + else:

534 + self.start_layer = 0

535 + self.end_layer = config.encoder_layers

536 + self.layers = nn.ModuleList(

537 + get_offloader().wrap_modules(

538 + create_layer(prefix=f"{prefix}.layers.{idx}")

539 + for idx in range(config.encoder_layers)

540 + )

541 + )

542 self.layer_norm = nn.LayerNorm(config.d_model)

543

544 if self.pos_embed_type not in (