ultravox.py 26.3 KB
Newer Older
1
2
# SPDX-License-Identifier: Apache-2.0

3
4
# Adapted from https://github.com/fixie-ai/ultravox/blob/ecd58c4041030bae2ad15aa6bcf04ab43199ea02/ultravox/model/ultravox_model.py
"""PyTorch Ultravox model."""
5
from collections.abc import Iterable, Mapping, Sequence
6
from typing import Any, Literal, Optional, TypedDict, Union
7
8
9
10

import torch
from torch import nn
from torch.nn import functional as F
11
from transformers import BatchFeature, ProcessorMixin
12
13
14
from transformers.models.whisper import WhisperFeatureExtractor
from transformers.models.whisper.modeling_whisper import WhisperEncoder

15
from vllm import envs
16
from vllm.config import VllmConfig
17
from vllm.forward_context import get_forward_context
18
from vllm.model_executor.layers.activation import MulAndSilu, get_act_fn
19
from vllm.model_executor.layers.layernorm import RMSNorm
20
from vllm.model_executor.model_loader import DefaultModelLoader
21
from vllm.model_executor.models.module_mapping import MultiModelKeys
22
from vllm.model_executor.sampling_metadata import SamplingMetadata
23
from vllm.multimodal import MULTIMODAL_REGISTRY
24
25
from vllm.multimodal.inputs import (MultiModalDataDict, MultiModalFieldConfig,
                                    MultiModalKwargs, NestedTensors)
26
from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser
27
from vllm.multimodal.processing import (BaseMultiModalProcessor,
28
29
                                        BaseProcessingInfo, PromptReplacement,
                                        PromptUpdate)
30
from vllm.multimodal.profiling import BaseDummyInputsBuilder
31
from vllm.sequence import IntermediateTensors
32
33
from vllm.transformers_utils.configs.ultravox import UltravoxConfig

34
from .interfaces import (MultiModalEmbeddings, SupportsLoRA,
35
                         SupportsMultiModal, SupportsPP)
36
from .utils import (AutoWeightsLoader, WeightsMapper, flatten_bn,
37
                    init_vllm_registered_model, maybe_prefix,
38
                    merge_multimodal_embeddings,
39
                    merge_multimodal_embeddings_from_map)
40

41
42
_AUDIO_PLACEHOLDER_OVERRIDE = "<|reserved_special_token_0|>"
_AUDIO_PLACEHOLDER_TOKEN = 128002
43
_AUDIO_TOKENS_PER_SECOND = 6.25
44
_MAX_ENCODER_BATCH_SIZE = 16
45
46
47
48


class UltravoxAudioFeatureInputs(TypedDict):
    type: Literal["audio_features"]
49
    data: Union[torch.Tensor, list[torch.Tensor], list[list[torch.Tensor]]]
50
    """Shape: `(batch_size, num_chunks, 80, M)`"""
51
    lens: Union[torch.Tensor, list[torch.Tensor]]
52
53
54
55
    """
    Length of the audio frames. Used for attention mask in WhisperEncoder.
    Shape: `(batch_size, num_chunks)`
    """
56
    token_len: Union[torch.Tensor, list[torch.Tensor]]
57
58
59
60
    """
    Length of the audio tokens. Used for flattening the audio features.
    Shape: `(batch_size, num_chunks)`
    """
61
62
63
64


class UltravoxAudioEmbeddingInputs(TypedDict):
    type: Literal["audio_embeds"]
65
    data: NestedTensors
66
    """Shape: `(batch_size, num_audios, audio_feature_size, hidden_size)`"""
67
68
69
70
71
72


UltravoxAudioInputs = Union[UltravoxAudioFeatureInputs,
                            UltravoxAudioEmbeddingInputs]


73
class UltravoxProcessingInfo(BaseProcessingInfo):
74

75
    def get_hf_processor(
76
77
78
79
        self,
        *,
        # Ignored in initialization
        sampling_rate: Optional[int] = None,
80
        **kwargs: object,
81
    ) -> ProcessorMixin:
82
        hf_processor = self.ctx.get_hf_processor(**kwargs)
83
84
85
86
87
88

        # NOTE: Ultravox processing definition uses '<|eot_id|>' as the
        # placeholder that will cause confusion with the actual end of turn
        # token, thus we override placeholder with a reserved special
        # token.
        hf_processor.audio_token_replacement = _AUDIO_PLACEHOLDER_OVERRIDE
89
        hf_processor.audio_replacement_token_id = _AUDIO_PLACEHOLDER_TOKEN
90
        return hf_processor
91

92
    def get_feature_extractor(
93
94
95
96
97
        self,
        *,
        # Ignored in initialization
        sampling_rate: Optional[int] = None,
    ) -> WhisperFeatureExtractor:
98
        hf_processor = self.get_hf_processor(sampling_rate=sampling_rate)
99
100
101
102
103
        audio_processor = hf_processor.audio_processor  # type: ignore
        feature_extractor = audio_processor.feature_extractor  # type: ignore
        assert isinstance(feature_extractor, WhisperFeatureExtractor)
        return feature_extractor

104
105
    def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:
        return {"audio": None}
106

107
108
109
110

class UltravoxDummyInputsBuilder(BaseDummyInputsBuilder[UltravoxProcessingInfo]
                                 ):

111
112
113
114
115
116
    def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
        num_audios = mm_counts.get("audio", 0)

        return "<|audio|>" * num_audios

    def get_dummy_mm_data(
117
        self,
118
119
        seq_len: int,
        mm_counts: Mapping[str, int],
120
    ) -> MultiModalDataDict:
121
        feature_extractor = self.info.get_feature_extractor()
122
123

        sampling_rate = feature_extractor.sampling_rate
124
125
        audio_len = (feature_extractor.chunk_length * sampling_rate *
                     _MAX_ENCODER_BATCH_SIZE)
126
127
        num_audios = mm_counts.get("audio", 0)

128
        return {
129
130
131
132
133
            "audio":
            self._get_dummy_audios(length=audio_len, num_audios=num_audios)
        }


134
135
class UltravoxMultiModalProcessor(
        BaseMultiModalProcessor[UltravoxProcessingInfo]):
136

137
    def _get_data_parser(self) -> MultiModalDataParser:
138
        feature_extractor = self.info.get_feature_extractor()
139
        return MultiModalDataParser(target_sr=feature_extractor.sampling_rate)
140
141

    def _call_hf_processor(
142
143
        self,
        prompt: str,
144
145
        mm_data: Mapping[str, object],
        mm_kwargs: Mapping[str, object],
146
    ) -> BatchFeature:
147
        # Text-only input not supported in composite processor
148
        if not mm_data.get("audios", []):
149
150
            prompt_ids = self.info.get_tokenizer().encode(
                prompt, add_special_tokens=False)
151
            prompt_ids = self._apply_hf_processor_tokens_only(prompt_ids)
152
153
154
155
            return BatchFeature(dict(input_ids=[prompt_ids]), tensor_type="pt")

        mm_data = dict(mm_data)
        audios = mm_data.pop("audios", [])
156
        assert isinstance(audios, list)
157

158
        feature_extractor = self.info.get_feature_extractor()
159
160
        mm_kwargs = dict(
            **mm_kwargs,
161
            sampling_rate=feature_extractor.sampling_rate,
162
            include_audio_num_chunks=True,
163
164
        )

165
        item_processor_data = dict(**mm_data, audios=audios)
166

167
168
169
170
        output = super()._call_hf_processor(
            prompt=prompt,
            mm_data=item_processor_data,
            mm_kwargs=mm_kwargs,
171
        )
172
173
174
        output['audio_features'] = output.pop('audio_values')

        return output
175

176
177
178
179
180
    def _get_mm_fields_config(
        self,
        hf_inputs: BatchFeature,
        hf_processor_mm_kwargs: Mapping[str, object],
    ) -> Mapping[str, MultiModalFieldConfig]:
181
        num_chunks = hf_inputs.get('audio_num_chunks', torch.zeros(0))
182
        return dict(
183
184
185
186
187
188
189
190
191
192
193
            # to handle longer than 30s audio, each audio might be split
            # into multiple chunks as such, their batch dimension can be
            # higher than the number of audio samples
            audio_features=MultiModalFieldConfig.flat_from_sizes(
                "audio", num_chunks),
            audio_token_len=MultiModalFieldConfig.flat_from_sizes(
                "audio", num_chunks),
            audio_lens=MultiModalFieldConfig.flat_from_sizes(
                "audio", num_chunks),
            # num_chunks can convert audio_chunked to audio batch dimension
            audio_num_chunks=MultiModalFieldConfig.batched("audio"),
194
195
196
            audio_embeds=MultiModalFieldConfig.batched("audio"),
        )

197
    def _get_prompt_updates(
198
199
        self,
        mm_items: MultiModalDataItems,
200
        hf_processor_mm_kwargs: Mapping[str, Any],
201
        out_mm_kwargs: MultiModalKwargs,
202
    ) -> Sequence[PromptUpdate]:
203
        hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
204

205
206
207
208
209
210
211
212
213
214
215
        replacement_id = hf_processor.audio_replacement_token_id  # type: ignore

        # Each audio can be split into multiple chunks.
        # chunks_start_idx[i] indicates the start index of the chunks
        # belonging to the i-th audio.
        num_chunks = out_mm_kwargs.get("audio_num_chunks", torch.zeros(0))
        chunks_start_idx: torch.Tensor = torch.cumsum(num_chunks,
                                                      dim=0,
                                                      dtype=torch.int32)
        chunks_start_idx = torch.cat(
            [torch.tensor([0], dtype=torch.int32), chunks_start_idx])
216
217

        def get_replacement_ultravox(item_idx: int):
218
219
220
            start = chunks_start_idx[item_idx]
            end = chunks_start_idx[item_idx + 1]
            audio_token_len = out_mm_kwargs["audio_token_len"][start:end].sum()
221
            return [replacement_id] * int(audio_token_len)  # type: ignore
222
223
224
225

        return [
            PromptReplacement(
                modality="audio",
226
                target="<|audio|>",
227
228
229
                replacement=get_replacement_ultravox,
            )
        ]
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258


class StackAudioFrames(nn.Module):
    """
    Stack the audio embedding frames to reduce the sequence length by a factor
    of `stack_factor`.
    """

    def __init__(self, stack_factor: int = 8):
        super().__init__()
        self.stack_factor = stack_factor

    def forward(self, audio_embeds: torch.Tensor) -> torch.Tensor:
        B, T, C = audio_embeds.shape
        T_pad = (T + self.stack_factor -
                 1) // self.stack_factor * self.stack_factor
        audio_embeds = F.pad(audio_embeds, (0, 0, 0, T_pad - T))
        B, T, C = audio_embeds.shape
        audio_embeds = audio_embeds.view(B, T // self.stack_factor,
                                         C * self.stack_factor)
        return audio_embeds


class UltravoxProjector(nn.Module):

    def __init__(self, config: UltravoxConfig):
        super().__init__()
        self.hidden_dim = config.hidden_size
        self._pad_and_stack = StackAudioFrames(config.stack_factor)
259
260
261
262
        dim_in = config.audio_config.hidden_size * config.stack_factor
        self.ln_pre = RMSNorm(dim_in)
        self.linear_1 = nn.Linear(dim_in, self.hidden_dim, bias=False)
        dim_mid = self.hidden_dim
263
264

        if config.projector_act == "swiglu":
265
            self.act = MulAndSilu()
266
            dim_mid = dim_mid // 2
267
268
269
        else:
            self.act = get_act_fn(config.projector_act)

270
271
272
273
274
275
276
277
278
279
280
        dim_out = config.text_config.hidden_size
        self.linear_2 = nn.Linear(dim_mid, dim_out, bias=False)

        # Ultravox v0.4.1 and below use layer_norm after the second linear layer
        # while v0.5.0 and above uses layer_norm after the first linear layer.
        if config.projector_ln_mid:
            self.ln_mid: nn.Module = RMSNorm(dim_mid)
            self.ln_post = nn.Identity()
        else:
            self.ln_mid = nn.Identity()
            self.ln_post = RMSNorm(dim_out)
281
282
283
284
285
286

    def forward(self, audio_features: torch.Tensor) -> torch.Tensor:
        audio_features = self._pad_and_stack(audio_features)
        audio_features = self.ln_pre(audio_features)
        hidden_states = self.linear_1(audio_features)
        hidden_states = self.act(hidden_states)
287
        hidden_states = self.ln_mid(hidden_states)
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
        hidden_states = self.linear_2(hidden_states)
        hidden_states = self.ln_post(hidden_states)
        return hidden_states


class ModifiedWhisperEncoder(WhisperEncoder):
    """
    Encoder portion of OpenAI's Whisper model.

    This implementation is a slightly modified version of HF Transformers'
    Whisper Encoder, with only a few fixes:
    1. base_model_prefix updated to allow for doing `.from_pretrained`
       directly on the encoder
    2. allow less than 30 second of audio padding to be passed in:
        - relaxed ValueError check for `input_features` length to be less
           than or equal to `expected_seq_length` instead of strictly equal
        - embed_pos is now sliced to match the length of `inputs_embeds`

    Original: https://github.com/huggingface/transformers/blob/main/src/transformers/models/whisper/modeling_whisper.py
    See commentary: https://github.com/huggingface/transformers/issues/25744
    """

    base_model_prefix = "model.encoder"

312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.config.is_decoder = False

    @property
    def max_context_length(self):
        return (self.config.max_source_positions * self.conv1.stride[0] *
                self.conv2.stride[0])

    def get_attention_mask_by_audio_len(self,
                                        audio_lens: Optional[torch.Tensor],
                                        hidden_states: torch.Tensor):
        """
        Create attention mask based on audio lengths to mask out padding tokens
        For each sample in batch:
        - Convert raw audio length to feature length after convolutions
        - Create bool mask: True for valid positions and False for padding
        - Convert to attention mask format expected by transformer layers
        (1.0 for positions to attend to, large negative for positions to ignore)
        This masking ensures consistent behavior between training and inference
        by preventing the model from attending to padding tokens in both cases
        """
        if audio_lens is None:
            return None

        audio_feature_len = self._get_feat_extract_output_lengths(audio_lens)
        max_seq_len = hidden_states.shape[1]
        attention_mask = torch.arange(max_seq_len,
                                      device=hidden_states.device)[None, :].lt(
                                          audio_feature_len.view(-1, 1))
        attention_mask = self.get_extended_attention_mask(
            attention_mask,
            None,
            dtype=hidden_states.dtype,
        )
        return attention_mask

349
350
    def forward(
        self,
351
352
        input_features: torch.Tensor,
        audio_lens: Optional[torch.Tensor] = None,
353
    ):
354
        expected_seq_length = self.max_context_length
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
        if input_features.shape[-1] > expected_seq_length:
            raise ValueError(
                f"Whisper expects the mel input features to be of length "
                f"{expected_seq_length} or less, but found "
                f"{input_features.shape[-1]}. Make sure to pad the input mel "
                f"features to {expected_seq_length}.")

        inputs_embeds = nn.functional.gelu(self.conv1(input_features))
        inputs_embeds = nn.functional.gelu(self.conv2(inputs_embeds))

        inputs_embeds = inputs_embeds.permute(0, 2, 1)
        embed_pos = self.embed_positions.weight[:inputs_embeds.size(-2)]

        hidden_states = inputs_embeds + embed_pos
        hidden_states = nn.functional.dropout(hidden_states,
                                              p=self.dropout,
                                              training=self.training)

373
374
375
        attention_mask = self.get_attention_mask_by_audio_len(
            audio_lens, hidden_states)

376
377
378
        for encoder_layer in self.layers:
            layer_outputs = encoder_layer(
                hidden_states,
379
                attention_mask,
380
381
382
383
384
385
386
387
388
                layer_head_mask=None,
            )

            hidden_states = layer_outputs[0]

        hidden_states = self.layer_norm(hidden_states)
        return hidden_states


389
390
391
392
@MULTIMODAL_REGISTRY.register_processor(
    UltravoxMultiModalProcessor,
    info=UltravoxProcessingInfo,
    dummy_inputs=UltravoxDummyInputsBuilder)
393
class UltravoxModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):
394
395
396
397
398
399

    packed_modules_mapping = {
        "qkv_proj": ["q_proj", "k_proj", "v_proj"],
        "gate_up_proj": ["gate_proj", "up_proj"]
    }

400
401
402
    hf_to_vllm_mapper = WeightsMapper(
        orig_to_new_prefix={"audio_tower.model.encoder.": "audio_tower."})

403
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
404
        super().__init__()
405
406
        config = vllm_config.model_config.hf_config
        multimodal_config = vllm_config.model_config.multimodal_config
407
408
409
410
        self.config = config
        self.multi_modal_config = multimodal_config
        assert self.multi_modal_config

411
412
        self.secondary_weights = []
        self.audio_tower = ModifiedWhisperEncoder(config.audio_config)
413
        if config.audio_model_id is not None:
414
415
            # this prefix is not for initialization, but for loading weights
            # note the trailing dot
416
417
418
419
420
421
            self.secondary_weights.append(
                DefaultModelLoader.Source(
                    model_or_path=config.audio_model_id,
                    revision=None,
                    prefix="audio_tower.",
                ))
422
423
        self.multi_modal_projector = UltravoxProjector(config)
        self.language_model = init_vllm_registered_model(
424
            vllm_config=vllm_config,
425
426
427
            hf_config=config.text_config,
            prefix=maybe_prefix(prefix, "language_model"),
        )
428
        if config.text_model_id is not None:
429
430
            # this prefix is not for initialization, but for loading weights
            # note the trailing dot
431
432
433
434
            self.secondary_weights.append(
                DefaultModelLoader.Source(model_or_path=config.text_model_id,
                                          revision=None,
                                          prefix="language_model."))
435

436
437
438
        self.make_empty_intermediate_tensors = (
            self.language_model.make_empty_intermediate_tensors)

439
440
441
442
443
444
445
446
447
448
    def get_mm_mapping(self) -> MultiModelKeys:
        """
        Get the module prefix in multimodal models
        """
        return MultiModelKeys.from_string_field(
            language_model="language_model.",
            connector="multi_modal_projector.",
            tower_model="audio_tower.",
        )

449
    def _audio_features_to_embeddings(
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
            self, input_features: torch.Tensor,
            audio_lens: torch.Tensor) -> torch.Tensor:
        audio_features = input_features.to(self.audio_tower.dtype)
        batch_size = audio_features.size(0)
        audio_embeddings = []

        # Process audio features in batches to keep memory usage predictable
        for start in range(0, batch_size, _MAX_ENCODER_BATCH_SIZE):
            end = min(start + _MAX_ENCODER_BATCH_SIZE, batch_size)
            # Process through audio tower
            batch_features = self.audio_tower(audio_features[start:end],
                                              audio_lens[start:end])
            batch_features = batch_features.to(self.audio_tower.dtype)

            # Process through projector
            batch_embeddings = self.multi_modal_projector(batch_features)
            audio_embeddings.append(batch_embeddings)

        # Concatenate results
        audio_embeddings = torch.cat(audio_embeddings, dim=0)
470
471
472
473
474
475
        return audio_embeddings

    def _parse_and_validate_audio_input(
            self, **kwargs: object) -> Optional[UltravoxAudioInputs]:
        audio_features = kwargs.pop("audio_features", None)
        audio_embeds = kwargs.pop("audio_embeds", None)
476
477
        audio_lens = kwargs.pop("audio_lens", None)
        audio_token_len = kwargs.pop("audio_token_len", None)
478
479
480
481
482
483
484
485

        if audio_features is None and audio_embeds is None:
            return None

        if audio_features is not None:
            if not isinstance(audio_features, (torch.Tensor, list)):
                raise ValueError("Incorrect type of audio features. "
                                 f"Got type: {type(audio_features)}")
486
487
488
489
490
491
            if not isinstance(audio_lens, (torch.Tensor, list)):
                raise ValueError("Incorrect type of audio_lens. "
                                 f"Got type: {type(audio_features)}")
            if not isinstance(audio_token_len, (torch.Tensor, list)):
                raise ValueError("Incorrect type of audio_token_len. "
                                 f"Got type: {type(audio_features)}")
492
493

            return UltravoxAudioFeatureInputs(type="audio_features",
494
495
496
                                              data=audio_features,
                                              lens=audio_lens,
                                              token_len=audio_token_len)
497
498

        if audio_embeds is not None:
499
            if not isinstance(audio_embeds, (torch.Tensor, list)):
500
501
502
503
504
505
506
507
508
                raise ValueError("Incorrect type of audio embeds. "
                                 f"Got type: {type(audio_embeds)}")

            return UltravoxAudioEmbeddingInputs(type="audio_embeds",
                                                data=audio_embeds)

        raise AssertionError("This line should be unreachable.")

    def _process_audio_input(
509
510
511
        self,
        audio_input: UltravoxAudioInputs,
    ) -> Union[NestedTensors, tuple[torch.Tensor, ...]]:
512
513
514
        if audio_input["type"] == "audio_embeds":
            return audio_input["data"]

515
516
517
518
        # Pad and concatenate audio features
        # [[B1, 80, M1], [B2, 80, M2]] -> [B1+B2, 80, max(M1, M2)]
        audio_features = pad_and_concat_to_dim3(audio_input["data"])

519
520
521
        # [B1, B2] -> [B1+B2]
        audio_lens = flatten_bn(audio_input['lens'], concat=True)
        audio_token_len = flatten_bn(audio_input['token_len'], concat=True)
522
523
524
525
526
527
528

        embeddings = self._audio_features_to_embeddings(
            audio_features, audio_lens)

        # We should flatten and concatenate embeddings based on token lengths
        # For example, with token_len = [4, 2, 3], flattened_embeddings will be
        # concat(embeddings[0][:4], embeddings[1][:2], embeddings[2][:3])
529

530
531
532
533
534
535
536
537
        # Create a mask of valid indices based on token lengths
        max_len = embeddings.shape[1]
        indices = torch.arange(max_len, device=embeddings.device).expand(
            embeddings.shape[0], -1)
        mask = indices < audio_token_len[:, None]
        # Apply mask and flatten
        flattened_embeddings = embeddings[mask]

538
539
540
541
542
543
        # Return one tensor per input audio
        embed_lens = [
            token_len_item.sum().item()
            for token_len_item in audio_input['token_len']
        ]
        return flattened_embeddings.split(embed_lens)
544

545
546
547
    def get_language_model(self) -> torch.nn.Module:
        return self.language_model

548
    def get_multimodal_embeddings(
549
            self, **kwargs: object) -> Optional[MultiModalEmbeddings]:
550
551
552
553
554
555
556
557
558
        audio_input = self._parse_and_validate_audio_input(**kwargs)
        if audio_input is None:
            return None
        audio_embeddings = self._process_audio_input(audio_input)
        return audio_embeddings

    def get_input_embeddings(
        self,
        input_ids: torch.Tensor,
559
        multimodal_embeddings: Optional[MultiModalEmbeddings] = None,
560
561
562
563
    ) -> torch.Tensor:
        inputs_embeds = self.language_model.get_input_embeddings(input_ids)
        if multimodal_embeddings is not None:

564
565
            # TODO(ywang96): remove this block after v0 is deprecated.
            if not envs.VLLM_USE_V1:
566
                attn_metadata = get_forward_context().attn_metadata
567
568
569
570
571
572
573
                merge_multimodal_embeddings_from_map(
                    inputs_embeds, multimodal_embeddings,
                    attn_metadata.multi_modal_placeholder_index_maps["audio"])
            else:
                inputs_embeds = merge_multimodal_embeddings(
                    input_ids, inputs_embeds, multimodal_embeddings,
                    _AUDIO_PLACEHOLDER_TOKEN)
574
575
576
577
578
579
580
        return inputs_embeds

    def forward(self,
                input_ids: torch.Tensor,
                positions: torch.Tensor,
                intermediate_tensors: Optional[torch.Tensor] = None,
                inputs_embeds: Optional[torch.Tensor] = None,
581
                **kwargs) -> Union[torch.Tensor, IntermediateTensors]:
582
583
584
585
586
587
588
589
590
591
        """Run forward pass for Ultravox

        One key thing to understand is the `input_ids` already accounts for the
        positions of the to-be-inserted audio embeddings. The to-be-inserted
        audio has a size that is essentially 6.25 tokens per second of audio.

        This way, the `positions` and `attn_metadata` are consistent
        with the `input_ids`.

        Args:
592
593
594
595
596
            audio_features: A batch of audio input chunks [B, N, 80, M].
            audio_lens: Length of audio frames for each audio chunk [B].
            audio_token_len: Length of audio tokens for each audio chunk [B'].
                Note: batch dim is different from batch dim in audio chunks.

597
        """
598

599
        if intermediate_tensors is not None:
600
            inputs_embeds = None
601
602
603
604
605
606
607

        # NOTE: In v1, inputs_embeds is always generated at model runner, this
        # condition is for v0 compatibility.
        elif inputs_embeds is None:
            multimodal_embeddings = self.get_multimodal_embeddings(**kwargs)

            inputs_embeds = self.get_input_embeddings(input_ids,
608
                                                      multimodal_embeddings)
609
610
611
612
613
614
            input_ids = None

        hidden_states = self.language_model.model(input_ids,
                                                  positions,
                                                  intermediate_tensors,
                                                  inputs_embeds=inputs_embeds)
615
616
617
618
619
620
621
        return hidden_states

    def compute_logits(self, hidden_states: torch.Tensor,
                       sampling_metadata: SamplingMetadata) -> torch.Tensor:
        return self.language_model.compute_logits(hidden_states,
                                                  sampling_metadata)

622
623
    def load_weights(self, weights: Iterable[tuple[str,
                                                   torch.Tensor]]) -> set[str]:
624
625
626

        loader = AutoWeightsLoader(self,
                                   ignore_unexpected_prefixes=["audio_tower."])
627
        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
628
629
630


def pad_and_concat_to_dim3(
631
    features: Union[torch.Tensor, list[torch.Tensor], list[list[torch.Tensor]]]
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
) -> torch.Tensor:
    """
    Pad and concatenate a list of tensors.

    output:
        Tensor of shape [B, C, M] where M is the maximum length of the input
        tensors, B is the sum of the batch sizes of the input tensors.
        C must be the same for all input tensors.
    """
    if isinstance(features, torch.Tensor):
        if features.ndim > 3:
            # Flatten [B, N, 80, M] -> [B * N, 80, M]
            features = flatten_bn(features)
        return features

    features = [pad_and_concat_to_dim3(f) for f in features]

    max_len = max(f.shape[-1] for f in features)
    # Ensure all features have dim=3
    features = [f.view(-1, *f.shape[-2:]) for f in features]
    # Pad and oncatenate:
    # [[B1, 80, M1], [B2, 80, M2]] -> [B1+B2, 80, max(M1, M2)]
    features = [F.pad(f, (0, max_len - f.shape[-1])) for f in features]
    return torch.cat(features)