@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 5f7e8a916a

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

[Model] Define merge_by_field_config MM interface (U-Z) (#26261)

Signed-off-by: Ayush Satyam <ayushsatyam146@gmail.com> Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk> Co-authored-by: DarkLight1337 <tlleungac@connect.ust.hk>

Ayush Satyam committed Oct 7, 2025 at 12:15 UTC 5f7e8a916a18577c36cdea14d0a5666e0e1280b6

4 files changed +32 -25

vllm/model_executor/models/ultravox.py

+18 -15

@@ -69,18 +69,16 @@ class UltravoxAudioFeatureInputs(TensorSchema):

69 type: Literal["audio_features"]

70 data: Annotated[

71 Union[torch.Tensor, list[torch.Tensor], list[list[torch.Tensor]]],

72 - TensorShape("b", "n", "nmb", "t", dynamic_dims={"n"}),

72 + TensorShape("bn", "nmb", "t"),

73 ]

74 - lens: Annotated[

75 - Union[torch.Tensor, list[torch.Tensor]],

76 - TensorShape("b", "n", dynamic_dims={"n"}),

77 - ]

78 - """Length of the audio frames. Used for attention mask in WhisperEncoder."""

79 - token_len: Annotated[

80 - Union[torch.Tensor, list[torch.Tensor]],

81 - TensorShape("b", "n", dynamic_dims={"n"}),

82 - ]

83 - """Length of the audio tokens. Used for flattening the audio features."""

74 + lens: Annotated[torch.Tensor, TensorShape("bn")]

75 + """

76 + Length of the audio frames per chunk. Used for attention mask in WhisperEncoder.

77 + """

78 + token_len: Annotated[torch.Tensor, TensorShape("bn")]

79 + """Length of the audio tokens per chunk. Used for flattening the audio features."""

80 + num_chunks: Annotated[torch.Tensor, TensorShape("n")]

81 + """Number of chunks per audio. Used for flattening the audio features."""

84 class UltravoxAudioEmbeddingInputs(TensorSchema):

@@ -421,6 +419,8 @@ class ModifiedWhisperEncoder(WhisperEncoder):

419 dummy_inputs=UltravoxDummyInputsBuilder,

420 )

421 class UltravoxModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):

422 + merge_by_field_config = True

423 +

424 packed_modules_mapping = {

425 "qkv_proj": ["q_proj", "k_proj", "v_proj"],

426 "gate_up_proj": ["gate_proj", "up_proj"],

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

519 audio_embeds = kwargs.pop("audio_embeds", None)

520 audio_lens = kwargs.pop("audio_lens", None)

521 audio_token_len = kwargs.pop("audio_token_len", None)

522 + audio_num_chunks = kwargs.pop("audio_num_chunks", None)

523

524 if audio_features is None and audio_embeds is None:

525 return None

@@ -529,6 +530,7 @@ class UltravoxModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):

530 data=audio_features,

531 lens=audio_lens,

532 token_len=audio_token_len,

533 + num_chunks=audio_num_chunks,

534 )

535

536 if audio_embeds is not None:

@@ -547,9 +549,8 @@ class UltravoxModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):

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

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

551

550 - # [B1, B2] -> [B1+B2]

551 - audio_lens = flatten_bn(audio_input["lens"], concat=True)

552 - audio_token_len = flatten_bn(audio_input["token_len"], concat=True)

552 + audio_lens = audio_input["lens"]

553 + audio_token_len = audio_input["token_len"]

554

555 embeddings = self._audio_features_to_embeddings(audio_features, audio_lens)

556

@@ -568,7 +569,8 @@ class UltravoxModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):

569

570 # Return one tensor per input audio

571 embed_lens = [

571 - token_len_item.sum().item() for token_len_item in audio_input["token_len"]

572 + chunk_lens.sum().item()

573 + for chunk_lens in audio_token_len.split(audio_input["num_chunks"].tolist())

574 ]

575 return flattened_embeddings.split(embed_lens)

576

@@ -663,6 +665,7 @@ def pad_and_concat_to_dim3(

665 if features.ndim > 3:

666 # Flatten [B, N, 80, M] -> [B * N, 80, M]

667 features = flatten_bn(features)

668 +

669 return features

670

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

vllm/model_executor/models/voxtral.py

+3 -2

@@ -61,7 +61,7 @@ from vllm.transformers_utils.tokenizer import (

61 )

63 from .interfaces import SupportsLoRA, SupportsMultiModal, SupportsTranscription

64 -from .utils import flatten_bn, init_vllm_registered_model, maybe_prefix

64 +from .utils import init_vllm_registered_model, maybe_prefix

66 logger = init_logger(__name__)

@@ -337,6 +337,8 @@ class VoxtralMultiModalProcessor(BaseMultiModalProcessor[VoxtralProcessingInfo])

337 class VoxtralForConditionalGeneration(

338 nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA, SupportsTranscription

339 ):

340 + merge_by_field_config = True

341 +

342 supported_languages = ISO639_1_SUPPORTED_LANGS

343

344 packed_modules_mapping = {

@@ -445,7 +447,6 @@ class VoxtralForConditionalGeneration(

447 f"Incorrect type of audio_arrays. Got type: {type(audio_arrays)}"

448 )

449

448 - audio_arrays = flatten_bn(audio_arrays)

450 if isinstance(audio_arrays, torch.Tensor):

451 audio_arrays = list(audio_arrays.unbind(0))

452 return audio_arrays

vllm/model_executor/models/whisper.py

+8 -8

@@ -36,7 +36,7 @@ from vllm.model_executor.layers.quantization import QuantizationConfig

36 from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead

37 from vllm.model_executor.model_loader.utils import set_default_torch_dtype

38 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

39 -from vllm.multimodal import MULTIMODAL_REGISTRY, NestedTensors

39 +from vllm.multimodal import MULTIMODAL_REGISTRY

40 from vllm.multimodal.inputs import (

41 MultiModalDataDict,

42 MultiModalFieldConfig,

@@ -51,6 +51,7 @@ from vllm.multimodal.processing import (

51 )

52 from vllm.multimodal.profiling import BaseDummyInputsBuilder

53 from vllm.transformers_utils.processor import cached_get_processor

54 +from vllm.utils.jsontree import json_map_leaves

55 from vllm.utils.tensor_schema import TensorSchema, TensorShape

57 from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsTranscription

@@ -135,7 +136,10 @@ class WhisperAudioInputs(TensorSchema):

136 - t: Time frames (M)

137 """

138

138 - input_features: Annotated[Optional[NestedTensors], TensorShape("b", "nmb", "t")]

139 + input_features: Annotated[

140 + Optional[list[torch.Tensor]],

141 + TensorShape("b", "nmb", "t"),

142 + ]

143

144

145 class WhisperEncoderAttention(MultiHeadAttention):

@@ -781,6 +785,7 @@ class WhisperMultiModalProcessor(EncDecMultiModalProcessor[WhisperProcessingInfo

785 class WhisperForConditionalGeneration(

786 nn.Module, SupportsTranscription, SupportsMultiModal

787 ):

788 + merge_by_field_config = True

789 packed_modules_mapping = {

790 "self_attn.qkv_proj": [

791 "self_attn.q_proj",

@@ -936,12 +941,7 @@ class WhisperForConditionalGeneration(

941 input_features = kwargs.pop("input_features", None)

942

943 if input_features is not None:

939 - if not isinstance(input_features, (torch.Tensor, list)):

940 - raise ValueError(

941 - "Incorrect type of audio features. "

942 - f"Got type: {type(input_features)}"

943 - )

944 - input_features = torch.cat([feat.to(self.dtype) for feat in input_features])

944 + input_features = json_map_leaves(lambda x: x.to(self.dtype), input_features)

945

946 return WhisperAudioInputs(input_features=input_features)

947

vllm/multimodal/inputs.py

@@ -677,6 +677,9 @@ class MultiModalFieldConfig:

677 self.field = field

678 self.modality = modality

679

680 + def __repr__(self) -> str:

681 + return f"MultiModalFieldConfig(field={self.field}, modality={self.modality})"

682 +

683 def build_elems(

684 self,

685 key: str,