@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / fa93cd9f60
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 15, 12:15 UTC

[Model] Add Granite Speech Support (#16246)

Signed-off-by: Alex-Brooks <Alex.brooks@ibm.com> Signed-off-by: Alex-Brooks <Alex.Brooks@ibm.com>

Alex Brooks committed Apr 28, 2025 at 04:05 UTC fa93cd9f608a43db8ad55685114c38f3e1ee1289
11 files changed +1025 -28
docs/source/models/supported_models.md
+7
@@ -895,6 +895,13 @@ See [this page](#generative-models) for more information on how to use generativ
895 * ✅︎
896 * ✅︎
897 * ✅︎
898 +- * `GraniteSpeechForConditionalGeneration`
899 + * Granite Speech
900 + * T + A
901 + * `ibm-granite/granite-speech-3.3-8b`
902 + * ✅︎
903 + * ✅︎
904 + * ✅︎
905 - * `H2OVLChatModel`
906 * H2OVL
907 * T + I<sup>E+</sup>
examples/offline_inference/audio_language.py
+32
@@ -38,6 +38,37 @@ class ModelRequestData(NamedTuple):
38 # Unless specified, these settings have been tested to work on a single L4.
39
40
41 +# Granite Speech
42 +def run_granite_speech(question: str, audio_count: int) -> ModelRequestData:
43 + # NOTE - the setting in this example are somehat different than what is
44 + # optimal for granite speech, and it is generally recommended to use beam
45 + # search. Check the model README for suggested settings.
46 + # https://huggingface.co/ibm-granite/granite-speech-3.3-8b
47 + model_name = "ibm-granite/granite-speech-3.3-8b"
48 +
49 + engine_args = EngineArgs(
50 + model=model_name,
51 + trust_remote_code=True,
52 + max_model_len=2048,
53 + max_num_seqs=2,
54 + enable_lora=True,
55 + max_lora_rank=64,
56 + limit_mm_per_prompt={"audio": audio_count},
57 + )
58 +
59 + # The model has an audio-specific lora directly in its model dir;
60 + # it should be enabled whenever you pass audio inputs to the model.
61 + speech_lora_path = model_name
62 + audio_placeholder = "<|audio|>" * audio_count
63 + prompts = f"<|start_of_role|>system<|end_of_role|>Knowledge Cutoff Date: April 2024.\nToday's Date: December 19, 2024.\nYou are Granite, developed by IBM. You are a helpful AI assistant<|end_of_text|>\n<|start_of_role|>user<|end_of_role|>{audio_placeholder}{question}<|end_of_text|>\n<|start_of_role|>assistant<|end_of_role|>" # noqa: E501
64 +
65 + return ModelRequestData(
66 + engine_args=engine_args,
67 + prompt=prompts,
68 + lora_requests=[LoRARequest("speech", 1, speech_lora_path)],
69 + )
70 +
71 +
72 # MiniCPM-O
73 def run_minicpmo(question: str, audio_count: int) -> ModelRequestData:
74 model_name = "openbmb/MiniCPM-o-2_6"
@@ -209,6 +240,7 @@ def run_whisper(question: str, audio_count: int) -> ModelRequestData:
240
241
242 model_example_map = {
243 + "granite_speech": run_granite_speech,
244 "minicpmo": run_minicpmo,
245 "phi4_mm": run_phi4mm,
246 "qwen2_audio": run_qwen2_audio,
tests/conftest.py
+30 -4
@@ -21,6 +21,7 @@ from transformers.models.auto.auto_factory import _BaseAutoModelClass
21 from tests.models.utils import (TokensTextLogprobs,
22 TokensTextLogprobsPromptLogprobs)
23 from vllm import LLM, SamplingParams
24 +from vllm.assets.audio import AudioAsset
25 from vllm.assets.image import ImageAsset
26 from vllm.assets.video import VideoAsset
27 from vllm.config import TaskOption, _get_and_verify_dtype
@@ -103,10 +104,25 @@ class _VideoAssets(_VideoAssetsBase):
104 return [prompts["sample_demo_1"]]
105
106
107 +class _AudioAssetsBase(UserList[AudioAsset]):
108 + pass
109 +
110 +
111 +class _AudioAssets(_AudioAssetsBase):
112 +
113 + def __init__(self) -> None:
114 + super().__init__([
115 + AudioAsset("mary_had_lamb"),
116 + AudioAsset("winning_call"),
117 + ])
118 +
119 +
120 IMAGE_ASSETS = _ImageAssets()
121 """Singleton instance of :class:`_ImageAssets`."""
122 VIDEO_ASSETS = _VideoAssets()
123 """Singleton instance of :class:`_VideoAssets`."""
124 +AUDIO_ASSETS = _AudioAssets()
125 +"""Singleton instance of :class:`_AudioAssets`."""
126
127
128 @pytest.fixture(scope="function", autouse=True)
@@ -263,6 +279,11 @@ def video_assets() -> _VideoAssets:
279 return VIDEO_ASSETS
280
281
282 +@pytest.fixture(scope="session")
283 +def audio_assets() -> _AudioAssets:
284 + return AUDIO_ASSETS
285 +
286 +
287 _T = TypeVar("_T", nn.Module, torch.Tensor, BatchEncoding, BatchFeature, dict)
288 _R = TypeVar("_R")
289
@@ -390,10 +411,15 @@ class HfRunner:
411 processor_kwargs["images"] = image
412 if videos is not None and (video := videos[i]) is not None:
413 processor_kwargs["videos"] = video
393 - if audios is not None and (audio_tuple := audios[i]) is not None:
394 - audio, sr = audio_tuple
395 - processor_kwargs["audio"] = audio
396 - processor_kwargs["sampling_rate"] = sr
414 + if audios is not None and (audio_inputs := audios[i]) is not None:
415 + # HACK - not all processors take sampling_rate; we should
416 + # clean this up in the future.
417 + if len(audio_inputs) == 2:
418 + audio, sr = audio_inputs
419 + processor_kwargs["audio"] = audio
420 + processor_kwargs["sampling_rate"] = sr
421 + else:
422 + processor_kwargs["audio"] = audio_inputs
423
424 inputs = self.processor(**processor_kwargs)
425 if isinstance(inputs, BatchFeature):
tests/models/decoder_only/audio_language/test_granite_speech.py new
+143
@@ -0,0 +1,143 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +from collections.abc import Sequence
4 +from typing import Optional
5 +
6 +import pytest
7 +from transformers import AutoModelForSpeechSeq2Seq
8 +
9 +from vllm.lora.request import LoRARequest
10 +from vllm.sequence import SampleLogprobs
11 +
12 +from ....conftest import HfRunner, PromptAudioInput, VllmRunner, _AudioAssets
13 +from ...registry import HF_EXAMPLE_MODELS
14 +from ...utils import check_logprobs_close
15 +
16 +HF_AUDIO_PROMPT = "<|start_of_role|>system<|end_of_role|>Knowledge Cutoff Date: April 2024.\nToday's Date: December 19, 2024.\nYou are Granite, developed by IBM. You are a helpful AI assistant<|end_of_text|>\n<|start_of_role|>user<|end_of_role|><|audio|>can you transcribe the speech into a written format?<|end_of_text|>\n<|start_of_role|>assistant<|end_of_role|>" # noqa: E501
17 +
18 +
19 +def vllm_to_hf_output(
20 + vllm_output: tuple[list[int], str, Optional[SampleLogprobs]],
21 +) -> tuple[list[int], str, Optional[SampleLogprobs]]:
22 + """Sanitize hf output to be comparable with vllm output."""
23 + output_ids, output_str, out_logprobs = vllm_output
24 +
25 + hf_output_str = output_str + "<|end_of_text|>"
26 +
27 + return output_ids, hf_output_str, out_logprobs
28 +
29 +
30 +MODEL_NAME = "ibm-granite/granite-speech-3.3-8b"
31 +# Audio lora co-exists directly in the model directory, but
32 +# currently still needs to be passed directly to vLLM.
33 +audio_lora_path = MODEL_NAME
34 +models = [MODEL_NAME]
35 +
36 +
37 +def run_test(
38 + hf_runner: type[HfRunner],
39 + vllm_runner: type[VllmRunner],
40 + inputs: Sequence[tuple[list[str], PromptAudioInput]],
41 + model: str,
42 + *,
43 + max_model_len: int,
44 + dtype: str,
45 + max_tokens: int,
46 + num_logprobs: int,
47 + tensor_parallel_size: int,
48 + distributed_executor_backend: Optional[str] = None,
49 +):
50 + """Inference result should be the same between hf and vllm.
51 +
52 + All the audio fixtures for the test are from AUDIO_ASSETS.
53 + For huggingface runner, we provide the audio as input.
54 + For vllm runner, we provide MultiModalDataDict objects
55 + and corresponding MultiModalConfig as input.
56 + Note, the text input is also adjusted to abide by vllm contract.
57 + The text output is sanitized to be able to compare with hf.
58 + """
59 + # NOTE: take care of the order. run vLLM first, and then run HF.
60 + # vLLM needs a fresh new process without cuda initialization.
61 + # if we run HF first, the cuda initialization will be done and it
62 + # will hurt multiprocessing backend with fork method (the default method).
63 + # max_model_len should be greater than image_feature_size
64 + with vllm_runner(
65 + model,
66 + task="generate",
67 + max_model_len=max_model_len,
68 + max_num_seqs=1,
69 + dtype=dtype,
70 + limit_mm_per_prompt={"audio": 1},
71 + tensor_parallel_size=tensor_parallel_size,
72 + distributed_executor_backend=distributed_executor_backend,
73 + enable_lora=True,
74 + max_lora_rank=64,
75 + enforce_eager=True,
76 + ) as vllm_model:
77 + lora_request = LoRARequest("audio", 1, audio_lora_path)
78 + vllm_outputs_per_case = [
79 + vllm_model.generate_greedy_logprobs(prompts,
80 + max_tokens,
81 + num_logprobs=num_logprobs,
82 + audios=audios,
83 + lora_request=lora_request)
84 + for prompts, audios in inputs
85 + ]
86 +
87 + with hf_runner(model, dtype=dtype,
88 + auto_cls=AutoModelForSpeechSeq2Seq) as hf_model:
89 +
90 + hf_processor = hf_model.processor
91 + eos_token_id = hf_processor.tokenizer.eos_token_id
92 +
93 + hf_outputs_per_case = [
94 + hf_model.generate_greedy_logprobs_limit(prompts,
95 + max_tokens,
96 + num_logprobs=num_logprobs,
97 + audios=[audios],
98 + eos_token_id=eos_token_id)
99 + for prompts, audios in inputs
100 + ]
101 +
102 + for hf_outputs, vllm_outputs in zip(hf_outputs_per_case,
103 + vllm_outputs_per_case):
104 + check_logprobs_close(
105 + outputs_0_lst=hf_outputs,
106 + outputs_1_lst=[
107 + vllm_to_hf_output(output) for output in vllm_outputs
108 + ],
109 + name_0="hf",
110 + name_1="vllm",
111 + )
112 +
113 +
114 +@pytest.mark.parametrize("model", models)
115 +@pytest.mark.parametrize("dtype", ["bfloat16"])
116 +@pytest.mark.parametrize("max_model_len", [2048])
117 +@pytest.mark.parametrize("max_tokens", [128])
118 +@pytest.mark.parametrize("num_logprobs", [10])
119 +def test_models(hf_runner, vllm_runner, model: str, audio_assets: _AudioAssets,
120 + dtype: str, max_model_len: int, max_tokens: int,
121 + num_logprobs: int) -> None:
122 + model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
123 + model_info.check_available_online(on_fail="skip")
124 + model_info.check_transformers_version(on_fail="skip")
125 +
126 + audio, sr = audio_assets[0].audio_and_sample_rate
127 + # This model expects 16k sample rate, which our test audio
128 + # already is; if this changes, it may break this test,
129 + # so we check it directly
130 + assert sr == 16000
131 + run_test(
132 + hf_runner,
133 + vllm_runner,
134 + [
135 + ([HF_AUDIO_PROMPT], [audio]),
136 + ],
137 + model,
138 + dtype=dtype,
139 + max_model_len=max_model_len,
140 + max_tokens=max_tokens,
141 + num_logprobs=num_logprobs,
142 + tensor_parallel_size=1,
143 + )
tests/models/decoder_only/audio_language/test_ultravox.py
+6 -11
@@ -11,7 +11,7 @@ from transformers import AutoModel, AutoTokenizer
11 from vllm.multimodal.audio import resample_audio_librosa
12 from vllm.sequence import SampleLogprobs
13
14 -from ....conftest import HfRunner, VllmRunner
14 +from ....conftest import HfRunner, VllmRunner, _AudioAssets
15 from ....utils import RemoteOpenAIServer
16 from ...registry import HF_EXAMPLE_MODELS
17 from ...utils import check_logprobs_close
@@ -31,12 +31,6 @@ CHUNKED_PREFILL_KWARGS = {
31 }
32
33
34 -@pytest.fixture(scope="session")
35 -def audio_assets():
36 - from vllm.assets.audio import AudioAsset
37 - return [AudioAsset("mary_had_lamb"), AudioAsset("winning_call")]
38 -
39 -
34 @pytest.fixture(scope="module", params=("mary_had_lamb", "winning_call"))
35 def audio(request):
36 from vllm.assets.audio import AudioAsset
@@ -59,7 +53,7 @@ def params_kwargs_to_cli_args(params_kwargs: dict[str, Any]) -> list[str]:
53 pytest.param({}, marks=pytest.mark.cpu_model),
54 pytest.param(CHUNKED_PREFILL_KWARGS),
55 ])
62 -def server(request, audio_assets):
56 +def server(request, audio_assets: _AudioAssets):
57 args = [
58 "--dtype", "bfloat16", "--max-model-len", "4096", "--enforce-eager",
59 "--limit-mm-per-prompt",
@@ -230,8 +224,9 @@ def test_models(hf_runner, vllm_runner, audio, dtype: str, max_tokens: int,
224 pytest.param({}, marks=pytest.mark.cpu_model),
225 pytest.param(CHUNKED_PREFILL_KWARGS),
226 ])
233 -def test_models_with_multiple_audios(vllm_runner, audio_assets, dtype: str,
234 - max_tokens: int, num_logprobs: int,
227 +def test_models_with_multiple_audios(vllm_runner, audio_assets: _AudioAssets,
228 + dtype: str, max_tokens: int,
229 + num_logprobs: int,
230 vllm_kwargs: dict) -> None:
231
232 vllm_prompt = _get_prompt(len(audio_assets),
@@ -250,7 +245,7 @@ def test_models_with_multiple_audios(vllm_runner, audio_assets, dtype: str,
245
246
247 @pytest.mark.asyncio
253 -async def test_online_serving(client, audio_assets):
248 +async def test_online_serving(client, audio_assets: _AudioAssets):
249 """Exercises online serving with/without chunked prefill enabled."""
250
251 messages = [{
tests/models/multimodal/processing/test_common.py
+1
@@ -254,6 +254,7 @@ def _test_processing_correctness_mistral(
254 "adept/fuyu-8b",
255 "google/gemma-3-4b-it",
256 "THUDM/glm-4v-9b",
257 + "ibm-granite/granite-speech-3.3-8b",
258 "h2oai/h2ovl-mississippi-800m",
259 "OpenGVLab/InternVL2-1B",
260 "HuggingFaceM4/Idefics3-8B-Llama3",
tests/models/registry.py
+3 -1
@@ -298,9 +298,11 @@ _MULTIMODAL_EXAMPLE_MODELS = {
298 extras={"fork": "Isotr0py/deepseek-vl2-tiny"}, # noqa: E501
299 max_transformers_version="4.48", # noqa: E501
300 transformers_version_reason="HF model is not compatible.", # noqa: E501
301 - hf_overrides={"architectures": ["DeepseekVLV2ForCausalLM"]}), # noqa: E501
301 + hf_overrides={"architectures": ["DeepseekVLV2ForCausalLM"]}), # noqa: E501
302 "FuyuForCausalLM": _HfExamplesInfo("adept/fuyu-8b"),
303 "Gemma3ForConditionalGeneration": _HfExamplesInfo("google/gemma-3-4b-it"),
304 + "GraniteSpeechForConditionalGeneration": _HfExamplesInfo("ibm-granite/granite-speech-3.3-8b", # noqa: E501
305 + min_transformers_version="4.52.0"), # noqa: E501
306 "GLM4VForCausalLM": _HfExamplesInfo("THUDM/glm-4v-9b",
307 trust_remote_code=True,
308 hf_overrides={"architectures": ["GLM4VForCausalLM"]}), # noqa: E501
vllm/entrypoints/chat_utils.py
+1 -1
@@ -517,7 +517,7 @@ class BaseMultiModalItemTracker(ABC, Generic[_T]):
517
518 raise TypeError(f"Unknown {modality} model type: {model_type}")
519 elif modality == "audio":
520 - if model_type == "ultravox":
520 + if model_type in ("ultravox", "granite_speech"):
521 return "<|audio|>"
522 if model_type == "phi4mm":
523 return f"<|audio_{current_count}|>"
vllm/model_executor/models/blip2.py
+24 -11
@@ -60,6 +60,7 @@ class Blip2QFormerMultiHeadAttention(nn.Module):
60 quant_config: Optional[QuantizationConfig],
61 cache_config: Optional[CacheConfig],
62 is_cross_attention: bool = False,
63 + prefix: str = "",
64 ) -> None:
65 super().__init__()
66
@@ -139,7 +140,7 @@ class Blip2QFormerMultiHeadAttention(nn.Module):
140
141 class Blip2QFormerSelfOutput(nn.Module):
142
142 - def __init__(self, config: Blip2QFormerConfig) -> None:
143 + def __init__(self, config: Blip2QFormerConfig, prefix: str = "") -> None:
144 super().__init__()
145
146 self.dense = nn.Linear(config.hidden_size, config.hidden_size)
@@ -167,6 +168,7 @@ class Blip2QFormerAttention(nn.Module):
168 quant_config: Optional[QuantizationConfig],
169 cache_config: Optional[CacheConfig],
170 is_cross_attention: bool = False,
171 + prefix: str = "",
172 ) -> None:
173 super().__init__()
174
@@ -175,9 +177,10 @@ class Blip2QFormerAttention(nn.Module):
177 quant_config=quant_config,
178 cache_config=cache_config,
179 is_cross_attention=is_cross_attention,
180 + prefix=f"{prefix}.attention",
181 )
182
180 - self.output = Blip2QFormerSelfOutput(config)
183 + self.output = Blip2QFormerSelfOutput(config, prefix=f"{prefix}.output")
184
185 def forward(
186 self,
@@ -195,7 +198,7 @@ class Blip2QFormerAttention(nn.Module):
198
199 class Blip2QFormerIntermediate(nn.Module):
200
198 - def __init__(self, config: Blip2QFormerConfig) -> None:
201 + def __init__(self, config: Blip2QFormerConfig, prefix: str = "") -> None:
202 super().__init__()
203
204 self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
@@ -209,7 +212,7 @@ class Blip2QFormerIntermediate(nn.Module):
212
213 class Blip2QFormerOutput(nn.Module):
214
212 - def __init__(self, config: Blip2QFormerConfig) -> None:
215 + def __init__(self, config: Blip2QFormerConfig, prefix: str = "") -> None:
216 super().__init__()
217
218 self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
@@ -237,6 +240,7 @@ class Blip2QFormerLayer(nn.Module):
240 quant_config: Optional[QuantizationConfig],
241 cache_config: Optional[CacheConfig],
242 layer_idx: int,
243 + prefix: str = "",
244 ) -> None:
245 super().__init__()
246
@@ -244,7 +248,8 @@ class Blip2QFormerLayer(nn.Module):
248 self.seq_len_dim = 1
249 self.attention = Blip2QFormerAttention(config,
250 quant_config=quant_config,
247 - cache_config=cache_config)
251 + cache_config=cache_config,
252 + prefix=f"{prefix}.attention")
253
254 self.layer_idx = layer_idx
255
@@ -253,13 +258,16 @@ class Blip2QFormerLayer(nn.Module):
258 config,
259 quant_config=quant_config,
260 cache_config=cache_config,
256 - is_cross_attention=True)
261 + is_cross_attention=True,
262 + prefix=f"{prefix}.crossattention")
263 self.has_cross_attention = True
264 else:
265 self.has_cross_attention = False
266
261 - self.intermediate_query = Blip2QFormerIntermediate(config)
262 - self.output_query = Blip2QFormerOutput(config)
267 + self.intermediate_query = Blip2QFormerIntermediate(
268 + config, prefix=f"{prefix}.intermediate_query")
269 + self.output_query = Blip2QFormerOutput(config,
270 + prefix=f"{prefix}.output_query")
271
272 def forward(
273 self,
@@ -325,6 +333,7 @@ class Blip2QFormerEncoder(nn.Module):
333 *,
334 quant_config: Optional[QuantizationConfig],
335 cache_config: Optional[CacheConfig],
336 + prefix: str = "",
337 ) -> None:
338 super().__init__()
339
@@ -334,7 +343,8 @@ class Blip2QFormerEncoder(nn.Module):
343 Blip2QFormerLayer(config,
344 quant_config=quant_config,
345 cache_config=cache_config,
337 - layer_idx=layer_idx)
346 + layer_idx=layer_idx,
347 + prefix=f"{prefix}.layer.{layer_idx}")
348 for layer_idx in range(config.num_hidden_layers)
349 ])
350
@@ -365,6 +375,7 @@ class Blip2QFormerModel(nn.Module):
375 *,
376 quant_config: Optional[QuantizationConfig],
377 cache_config: Optional[CacheConfig],
378 + prefix: str = "",
379 ) -> None:
380 super().__init__()
381
@@ -376,7 +387,8 @@ class Blip2QFormerModel(nn.Module):
387
388 self.encoder = Blip2QFormerEncoder(config,
389 quant_config=quant_config,
379 - cache_config=cache_config)
390 + cache_config=cache_config,
391 + prefix=f"{prefix}.encoder")
392
393 def forward(
394 self,
@@ -511,7 +523,8 @@ class Blip2ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP,
523
524 self.qformer = Blip2QFormerModel(config.qformer_config,
525 cache_config=cache_config,
514 - quant_config=quant_config)
526 + quant_config=quant_config,
527 + prefix=f"{prefix}.qformer")
528
529 self.language_projection = nn.Linear(
530 config.qformer_config.hidden_size,
vllm/model_executor/models/granite_speech.py new
+777
@@ -0,0 +1,777 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +# Adapted from
4 +# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/llama/modeling_llama.py
5 +# Copyright 2025 The vLLM team.
6 +# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
7 +#
8 +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
9 +# and OPT implementations in this library. It has been modified from its
10 +# original forms to accommodate minor architectural differences compared
11 +# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
12 +#
13 +# Licensed under the Apache License, Version 2.0 (the "License");
14 +# you may not use this file except in compliance with the License.
15 +# You may obtain a copy of the License at
16 +#
17 +# http://www.apache.org/licenses/LICENSE-2.0
18 +#
19 +# Unless required by applicable law or agreed to in writing, software
20 +# distributed under the License is distributed on an "AS IS" BASIS,
21 +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
22 +# See the License for the specific language governing permissions and
23 +# limitations under the License.
24 +"""Inference-only IBM Granite speeech model."""
25 +import math
26 +from typing import Iterable, Mapping, Optional, Set, Tuple, TypedDict, Union
27 +
28 +import torch
29 +import torch.nn.functional as F
30 +from torch import nn
31 +from transformers import BatchFeature, PretrainedConfig
32 +
33 +from vllm.config import CacheConfig, VllmConfig
34 +from vllm.model_executor.layers.linear import (ColumnParallelLinear,
35 + RowParallelLinear)
36 +from vllm.model_executor.layers.quantization import QuantizationConfig
37 +from vllm.model_executor.layers.sampler import get_sampler
38 +from vllm.model_executor.models.module_mapping import MultiModelKeys
39 +from vllm.model_executor.sampling_metadata import SamplingMetadata
40 +from vllm.multimodal import MULTIMODAL_REGISTRY
41 +from vllm.multimodal.inputs import (MultiModalDataDict, MultiModalFieldConfig,
42 + MultiModalKwargs)
43 +from vllm.multimodal.parse import (AudioProcessorItems, MultiModalDataItems,
44 + MultiModalDataParser)
45 +from vllm.multimodal.processing import (BaseMultiModalProcessor,
46 + BaseProcessingInfo, PromptReplacement,
47 + PromptUpdate)
48 +from vllm.multimodal.profiling import BaseDummyInputsBuilder
49 +from vllm.sequence import IntermediateTensors
50 +
51 +from .blip2 import Blip2QFormerModel
52 +from .interfaces import (MultiModalEmbeddings, SupportsLoRA,
53 + SupportsMultiModal, SupportsPP)
54 +from .utils import (AutoWeightsLoader, embed_multimodal,
55 + init_vllm_registered_model, maybe_prefix)
56 +
57 +
58 +### Audio Input
59 +class GraniteSpeechAudioInputs(TypedDict):
60 +
61 + input_features: torch.Tensor
62 + """Shape: `(bsz, num_features, 160)`"""
63 +
64 + input_features_mask: torch.Tensor
65 + """Shape: `(bsz, num_features)`"""
66 +
67 + audio_embed_sizes: list[int]
68 + """List of length `bsz`"""
69 +
70 +
71 +class GraniteSpeechMultiModalProcessingInfo(BaseProcessingInfo):
72 +
73 + def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:
74 + return {"audio": 1}
75 +
76 + # There is no limit to the maximum number of audio tokens that can be
77 + # encoded as features; we pick ~5000 as a number that is probably higher
78 + # than we would expect to encounter. The sequence of length
79 + # get_max_audio_len() produces get_max_audio_tokens().
80 + def get_max_audio_tokens(self):
81 + return 5001
82 +
83 + def get_max_audio_len(self):
84 + return 8000000
85 +
86 +
87 +### Input Processing & Multimodal utils
88 +class GraniteSpeechMultiModalProcessor(
89 + BaseMultiModalProcessor[GraniteSpeechMultiModalProcessingInfo]):
90 +
91 + def _get_data_parser(self) -> MultiModalDataParser:
92 + feature_extractor = self.info.get_hf_processor().audio_processor
93 + sampling_rate = feature_extractor.melspec_kwargs["sample_rate"]
94 + return MultiModalDataParser(target_sr=sampling_rate)
95 +
96 + def _get_mm_fields_config(
97 + self,
98 + hf_inputs: BatchFeature,
99 + hf_processor_mm_kwargs: Mapping[str, object],
100 + ) -> Mapping[str, MultiModalFieldConfig]:
101 + return dict(
102 + input_features=MultiModalFieldConfig.batched("audio"),
103 + audio_embed_sizes=MultiModalFieldConfig.batched("audio"),
104 + )
105 +
106 + def _get_prompt_updates(
107 + self,
108 + mm_items: MultiModalDataItems,
109 + hf_processor_mm_kwargs: Mapping[str, object],
110 + out_mm_kwargs: MultiModalKwargs,
111 + ) -> list[PromptUpdate]:
112 + processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
113 + tokenizer = self.info.get_tokenizer()
114 + feature_extractor = processor.audio_processor
115 + vocab = tokenizer.get_vocab()
116 +
117 + # Use getattr with default to be compatible with transformers<4.48
118 + audio_token = getattr(processor, "audio_token", "<|audio|>")
119 + audio_token_id = vocab[audio_token]
120 +
121 + def get_replacement(item_idx: int):
122 + audios = mm_items.get_items("audio", AudioProcessorItems)
123 + audio = audios.get(item_idx)
124 + audio_length = audio.shape[-1]
125 + num_projector_features = feature_extractor._get_num_audio_features(
126 + [audio_length])[0]
127 + return [audio_token_id] * num_projector_features
128 +
129 + return [
130 + PromptReplacement(
131 + modality="audio",
132 + target=[audio_token_id],
133 + replacement=get_replacement,
134 + )
135 + ]
136 +
137 + def _call_hf_processor(
138 + self,
139 + prompt: str,
140 + mm_data: Mapping[str, object],
141 + mm_kwargs: Mapping[str, object],
142 + ) -> BatchFeature:
143 + mm_data = dict(mm_data)
144 + audios = mm_data.pop("audios", [])
145 +
146 + if audios:
147 + # GraniteSpeechFeatureExtractor accepts "audio"
148 + mm_data["audio"] = audios
149 +
150 + processed_outputs = super()._call_hf_processor(
151 + prompt=prompt,
152 + mm_data=mm_data,
153 + mm_kwargs=mm_kwargs,
154 + )
155 +
156 + if "audio" in mm_data:
157 + # Calculate the number of audio tokens per entry in the batch;
158 + # This is used to split the batch back out after padding.
159 + audio_token_index = self.info.get_hf_config().audio_token_index
160 + processed_outputs["audio_embed_sizes"] = [
161 + torch.sum(indices == audio_token_index).item()
162 + for indices in processed_outputs["input_ids"]
163 + ]
164 +
165 + return processed_outputs
166 +
167 +
168 +class GraniteSpeechDummyInputsBuilder(
169 + BaseDummyInputsBuilder[GraniteSpeechMultiModalProcessingInfo]):
170 +
171 + def get_dummy_mm_data(
172 + self,
173 + seq_len: int,
174 + mm_counts: Mapping[str, int],
175 + ) -> MultiModalDataDict:
176 + num_audios = mm_counts.get("audio", 0)
177 + return {
178 + "audio":
179 + self._get_dummy_audios(
180 + length=self.info.get_max_audio_len(),
181 + num_audios=num_audios,
182 + )
183 + }
184 +
185 + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
186 + num_audios = mm_counts.get("audio", 0)
187 + hf_processor = self.info.get_hf_processor()
188 + audio_token = getattr(hf_processor, "audio_token", "<|audio|>")
189 + return audio_token * num_audios
190 +
191 +
192 +### QFormer Projector
193 +class GraniteSpeechEncoderProjector(nn.Module):
194 +
195 + def __init__(
196 + self,
197 + config: PretrainedConfig,
198 + cache_config: CacheConfig,
199 + quant_config: Optional[QuantizationConfig] = None,
200 + prefix: str = "",
201 + ):
202 + super().__init__()
203 + self.hidden_size = config.projector_config.hidden_size
204 + self.downsample_rate = config.downsample_rate
205 + self.window_size = config.window_size
206 + self.num_queries = config.window_size // config.downsample_rate
207 +
208 + self.query = nn.Parameter(
209 + torch.zeros(1, self.num_queries,
210 + config.projector_config.hidden_size))
211 +
212 + # NOTE - this is implemented generically in transformers,
213 + # but for now we create the QFormer model directly since
214 + # all existing models use this for the projector.
215 + self.qformer = Blip2QFormerModel(
216 + config.projector_config,
217 + quant_config=quant_config,
218 + cache_config=cache_config,
219 + prefix=f"{prefix}.qformer",
220 + )
221 + self.linear = nn.Linear(config.projector_config.hidden_size,
222 + config.text_config.hidden_size)
223 +
224 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
225 + batch_size, seq_len, dim = hidden_states.size()
226 + nblocks = math.ceil(seq_len / self.window_size)
227 + pad = nblocks * self.window_size - seq_len
228 + hidden_states = nn.functional.pad(hidden_states, (0, 0, 0, pad),
229 + "constant", 0)
230 + hidden_states = hidden_states.view(batch_size * nblocks,
231 + self.window_size, dim)
232 +
233 + last_hidden_state = self.qformer(
234 + query_embeds=self.query.data,
235 + encoder_hidden_states=hidden_states,
236 + )
237 +
238 + query_proj = self.linear(
239 + last_hidden_state.view(
240 + batch_size,
241 + nblocks * self.window_size // self.downsample_rate,
242 + -1,
243 + ))
244 + return query_proj
245 +
246 +
247 +# Encoder - conformer is adapted from: https://github.com/lucidrains/conformer.git
248 +# NOTE - it would be nice to see if we can align this with other models using
249 +# conformer in vLLM, e.g., phi4mm audio.
250 +class GraniteSpeechConformerFeedForward(nn.Module):
251 + """Feedforward module for conformer encoder blocks."""
252 +
253 + def __init__(self,
254 + config: PretrainedConfig,
255 + quant_config: Optional[QuantizationConfig] = None,
256 + prefix: str = ""):
257 + super().__init__()
258 + self.pre_norm = nn.LayerNorm(config.hidden_dim)
259 +
260 + self.up_proj = ColumnParallelLinear(
261 + input_size=config.hidden_dim,
262 + output_size=config.hidden_dim * config.feedforward_mult,
263 + quant_config=quant_config,
264 + prefix=f"{prefix}.up_proj",
265 + )
266 + self.silu = nn.SiLU()
267 +
268 + self.down_proj = RowParallelLinear(
269 + input_size=config.hidden_dim * config.feedforward_mult,
270 + output_size=config.hidden_dim,
271 + quant_config=quant_config,
272 + prefix=f"{prefix}.down_proj",
273 + )
274 +
275 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
276 + hidden_states = self.pre_norm(hidden_states)
277 + hidden_states, _ = self.up_proj(hidden_states)
278 + hidden_states = self.silu(hidden_states)
279 + hidden_states, _ = self.down_proj(hidden_states)
280 + return hidden_states
281 +
282 +
283 +class GraniteSpeechConformerAttention(nn.Module):
284 + """Attention for conformer blocks using Shaw's relative positional
285 + embeddings. See the following [paper](https://arxiv.org/pdf/1803.02155)
286 + for more details.
287 + """
288 +
289 + def __init__(self, config: PretrainedConfig, prefix: str = ""):
290 + super().__init__()
291 +
292 + inner_dim = config.dim_head * config.num_heads
293 + self.max_pos_emb = config.max_pos_emb
294 + self.context_size = config.context_size
295 + self.num_heads = config.num_heads
296 + self.dim_head = config.dim_head
297 + self.scale = self.dim_head**-0.5
298 + self.pre_norm = nn.LayerNorm(config.hidden_dim)
299 + self.to_q = nn.Linear(config.hidden_dim, inner_dim, bias=False)
300 + self.to_kv = nn.Linear(config.hidden_dim, inner_dim * 2, bias=False)
301 + self.to_out = nn.Linear(inner_dim, config.hidden_dim)
302 + self.rel_pos_emb = nn.Embedding(2 * self.max_pos_emb + 1,
303 + self.dim_head)
304 +
305 + if self.context_size <= 0 or self.context_size > self.max_pos_emb:
306 + raise ValueError(
307 + "Context size is either less than 0 or exceeds the max_pos_emb"
308 + )
309 +
310 + def forward(self, hidden_states: torch.Tensor,
311 + attention_dists: torch.Tensor) -> torch.Tensor:
312 + hidden_states = self.pre_norm(hidden_states)
313 + bsz, num_features, _ = hidden_states.shape
314 +
315 + num_blocks = math.ceil(num_features / self.context_size)
316 + remainder = num_features % self.context_size
317 + if remainder > 0:
318 + # right padding to reach block size
319 + hidden_states = torch.nn.functional.pad(
320 + hidden_states, (0, 0, 0, self.context_size - remainder))
321 +
322 + # NOTE: would be nice to try to use qkvparallellinear
323 + # here for this block attention implementation if possible
324 + query_states = self.to_q(hidden_states)
325 + key_states, value_states = self.to_kv(hidden_states).chunk(2, dim=-1)
326 +
327 + query_states = query_states.reshape(bsz, num_blocks, self.context_size,
328 + self.num_heads,
329 + -1).transpose(2, 3)
330 + key_states = key_states.reshape(bsz, num_blocks, self.context_size,
331 + self.num_heads, -1).transpose(2, 3)
332 + value_states = value_states.reshape(bsz, num_blocks, self.context_size,
333 + self.num_heads,
334 + -1).transpose(2, 3)
335 +
336 + # shaw's relative positional embedding
337 + dist = attention_dists.to(hidden_states.device)
338 + rel_pos_emb = self.rel_pos_emb(dist)
339 + rel_pos_emb_expanded = rel_pos_emb.view([1, 1, 1] +
340 + list(rel_pos_emb.shape))
341 + pos_attn = torch.sum(query_states.unsqueeze(-2) * rel_pos_emb_expanded,
342 + dim=-1) * self.scale
343 +
344 + if remainder > 0:
345 + # masked attention in the extended block
346 + mask = torch.ones(self.context_size,
347 + self.context_size,
348 + dtype=bool,
349 + device=hidden_states.device)
350 + mask[:remainder, :remainder] = 0
351 + mask_value = -torch.finfo(pos_attn.dtype).max
352 + pos_attn[:, -1, :].masked_fill_(mask, mask_value)
353 +
354 + with torch.nn.attention.sdpa_kernel(
355 + torch.nn.attention.SDPBackend.MATH):
356 + out = F.scaled_dot_product_attention(query_states,
357 + key_states,
358 + value_states,
359 + attn_mask=pos_attn,
360 + scale=self.scale)
361 + out = out.transpose(2, 3).reshape(bsz, hidden_states.shape[1], -1)
362 + return self.to_out(out[:, :num_features, :])
363 +
364 +
365 +class GraniteSpeechConformerDepthWiseConv1d(nn.Module):
366 + """Wrapper for padded 1D pointwise convolution."""
367 +
368 + def __init__(self,
369 + chan_in: int,
370 + chan_out: int,
371 + kernel_size: int,
372 + prefix: str = ""):
373 + super().__init__()
374 + # Padding for the 1D conv is symmetric or close (i.e., offset by one).
375 + pad = kernel_size // 2
376 + pad_offset = (kernel_size + 1) % 2
377 + self.padding = (pad, pad - pad_offset)
378 +
379 + self.conv = nn.Conv1d(chan_in,
380 + chan_out,
381 + kernel_size,
382 + groups=chan_in,
383 + bias=False)
384 +
385 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
386 + hidden_states = F.pad(hidden_states, self.padding)
387 + return self.conv(hidden_states)
388 +
389 +
390 +class GraniteSpeechConformerConvModule(nn.Module):
391 + """Conformer conv module consisting of several 1D/depthwise 1D
392 + convolutional layers.
393 + """
394 +
395 + def __init__(self, config: PretrainedConfig, prefix: str = ""):
396 + super().__init__()
397 + inner_dim = config.hidden_dim * config.conv_expansion_factor
398 +
399 + self.norm = nn.LayerNorm(config.hidden_dim)
400 + self.up_conv = nn.Conv1d(config.hidden_dim, inner_dim * 2, 1)
401 + self.glu = nn.GLU(dim=1)
402 + self.depth_conv = GraniteSpeechConformerDepthWiseConv1d(
403 + inner_dim,
404 + inner_dim,
405 + kernel_size=config.conv_kernel_size,
406 + prefix=f"{prefix}.depth_conv",
407 + )
408 + self.silu = nn.SiLU()
409 + self.batch_norm = nn.BatchNorm1d(inner_dim)
410 + self.down_conv = nn.Conv1d(inner_dim, config.hidden_dim, 1)
411 +
412 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
413 + hidden_states = self.norm(hidden_states)
414 + hidden_states = self.up_conv(hidden_states.permute(0, 2, 1))
415 + hidden_states = self.glu(hidden_states)
416 + hidden_states = self.depth_conv(hidden_states)
417 + hidden_states = self.silu(self.batch_norm(hidden_states))
418 + hidden_states = self.down_conv(hidden_states).permute(0, 2, 1)
419 + return hidden_states
420 +
421 +
422 +class GraniteSpeechConformerBlock(nn.Module):
423 + """Conformer block, consisting largely of linear layers,
424 + attention, and convolutional layers."""
425 +
426 + def __init__(self, config: PretrainedConfig, prefix: str = ""):
427 + super().__init__()
428 + self.ff1 = GraniteSpeechConformerFeedForward(config,
429 + prefix=f"{prefix}.ff1")
430 + self.attn = GraniteSpeechConformerAttention(config,
431 + prefix=f"{prefix}.attn")
432 + self.conv = GraniteSpeechConformerConvModule(config,
433 + prefix=f"{prefix}.conv")
434 + self.ff2 = GraniteSpeechConformerFeedForward(config,
435 + prefix=f"{prefix}.ff2")
436 + self.post_norm = nn.LayerNorm(config.hidden_dim)
437 +
438 + def forward(self, hidden_states: torch.Tensor,
439 + attention_dists: torch.Tensor) -> torch.Tensor:
440 + hidden_states = 0.5 * self.ff1(hidden_states) + hidden_states
441 + hidden_states = self.attn(
442 + hidden_states, attention_dists=attention_dists) + hidden_states
443 + hidden_states = self.conv(hidden_states) + hidden_states
444 + hidden_states = 0.5 * self.ff2(hidden_states) + hidden_states
445 + hidden_states = self.post_norm(hidden_states)
446 + return hidden_states
447 +
448 +
449 +class GraniteSpeechCTCEncoder(nn.Module):
450 + """CTC Encoder comprising conformer blocks and additional linear layers."""
451 +
452 + def __init__(self,
453 + config: PretrainedConfig,
454 + prefix: str,
455 + quant_config: Optional[QuantizationConfig] = None):
456 + super().__init__()
457 + self.config = config
458 +
459 + # Precompute clamped relative positional encoding distances
460 + seq = torch.arange(config.context_size)
461 + relpos_dist = seq.view(-1, 1) - seq.view(1, -1)
462 + self.attention_dists = torch.clamp(
463 + relpos_dist, -config.context_size,
464 + config.context_size) + config.max_pos_emb
465 +
466 + self.input_linear = nn.Linear(config.input_dim,
467 + config.hidden_dim,
468 + bias=True)
469 + self.layers = nn.ModuleList([
470 + GraniteSpeechConformerBlock(
471 + config,
472 + prefix=f"{prefix}.layers.{idx}",
473 + ) for idx in range(config.num_layers)
474 + ])
475 +
476 + self.out = ColumnParallelLinear(
477 + input_size=config.hidden_dim,
478 + output_size=config.output_dim,
479 + bias=True,
480 + quant_config=quant_config,
481 + prefix=f"{prefix}.out",
482 + )
483 +
484 + self.out_mid = RowParallelLinear(
485 + input_size=config.output_dim,
486 + output_size=config.hidden_dim,
487 + bias=True,
488 + quant_config=quant_config,
489 + prefix=f"{prefix}.out_mid",
490 + )
491 + self.softmax = nn.Softmax(dim=-1)
492 + self.num_layers = config.num_layers
493 +
494 + def forward(self, hidden_states: torch.Tensor):
495 + hidden_states = self.input_linear(hidden_states)
496 + for idx, layer in enumerate(self.layers, start=1):
497 + hidden_states = layer(hidden_states,
498 + attention_dists=self.attention_dists)
499 +
500 + if idx == self.num_layers // 2:
501 + hidden_states_mid = hidden_states.clone()
502 + hidden_states_mid, _ = self.out(hidden_states_mid)
503 + hidden_states_mid = self.softmax(hidden_states_mid)
504 + hidden_states_mid, _ = self.out_mid(hidden_states_mid)
505 + hidden_states += hidden_states_mid
506 + return hidden_states
507 +
508 +
509 +@MULTIMODAL_REGISTRY.register_processor(
510 + GraniteSpeechMultiModalProcessor,
511 + info=GraniteSpeechMultiModalProcessingInfo,
512 + dummy_inputs=GraniteSpeechDummyInputsBuilder)
513 +class GraniteSpeechForConditionalGeneration(
514 + nn.Module,
515 + SupportsMultiModal,
516 + SupportsPP,
517 + SupportsLoRA,
518 +):
519 +
520 + packed_modules_mapping = {
521 + "qkv_proj": [
522 + "q_proj",
523 + "k_proj",
524 + "v_proj",
525 + ],
526 + "gate_up_proj": [
527 + "gate_proj",
528 + "up_proj",
529 + ],
530 + }
531 +
532 + def __init__(self, *, vllm_config: VllmConfig, prefix: str):
533 + super().__init__()
534 + config = vllm_config.model_config.hf_config
535 + quant_config = vllm_config.quant_config
536 + cache_config = vllm_config.cache_config
537 +
538 + self.config = config
539 + self.quant_config = quant_config
540 + self.cache_config = cache_config
541 + self.sampler = get_sampler()
542 +
543 + # The language model is typically a Granite LLM
544 + self.language_model = init_vllm_registered_model(
545 + vllm_config=vllm_config,
546 + hf_config=config.text_config,
547 + prefix=maybe_prefix(prefix, "language_model"),
548 + )
549 +
550 + # Conformer encoder
551 + self.encoder = GraniteSpeechCTCEncoder(
552 + config=config.encoder_config,
553 + quant_config=quant_config,
554 + prefix=f"{prefix}.encoder",
555 + )
556 +
557 + # Blip2 QFormer
558 + self.projector = GraniteSpeechEncoderProjector(
559 + config=config,
560 + quant_config=quant_config,
561 + cache_config=cache_config,
562 + prefix=f"{prefix}.projector",
563 + )
564 +
565 + self.make_empty_intermediate_tensors = (
566 + self.language_model.make_empty_intermediate_tensors)
567 +
568 + def _parse_and_validate_audio_input(
569 + self,
570 + **kwargs: object,
571 + ) -> Optional[GraniteSpeechAudioInputs]:
572 + input_features = kwargs.pop("input_features", None)
573 + input_features_mask = kwargs.pop("input_features_mask", None)
574 + audio_embed_sizes = kwargs.pop("audio_embed_sizes", None)
575 + if input_features is None:
576 + return None
577 +
578 + # If we have a batch of variable feature length audio clips, we need
579 + # to mask the features; usually we would get an input_features_mask
580 + # from the processor, but we handle rebuilding it here since
581 + # vLLM generally processes everything independently + batches.
582 + if input_features_mask is None:
583 + input_features_mask = self._build_input_features_mask(
584 + audio_embed_sizes)
585 +
586 + if not isinstance(input_features, (torch.Tensor, list)):
587 + raise ValueError("Incorrect type of audio input features. "
588 + f"Got type: {type(input_features)}")
589 +
590 + if input_features_mask is not None and not isinstance(
591 + input_features_mask, torch.Tensor):
592 + raise ValueError("Incorrect type of audio input features mask. "
593 + f"Got type: {type(input_features_mask)}")
594 +
595 + if isinstance(input_features, torch.Tensor):
596 + # Granite speech currently only allows one audio token per instance
597 + # and features are already unsqueezed in the processor, so one
598 + # instance will have shape [1, {num_features}, 160]. As such,
599 + # input features will usually be of shape
600 + # [bsz, 1, num_features, 160], which we squeeze to be 3D here.
601 + if len(input_features.shape) == 4:
602 + input_features = input_features.squeeze(1)
603 + if len(input_features.shape) != 3:
604 + raise ValueError(
605 + "Squeezed input features should be 3D but are of shape "
606 + f"{input_features.shape}")
607 + input_features = input_features.to(
608 + self.encoder.input_linear.weight.dtype)
609 +
610 + else:
611 + # Otherwise we have a list of tensors, which are almost certainly
612 + # differing in their respective numbers of audio features;
613 + # stack them into a 3D tensor of size [bsz, most_num_features, 160].
614 + input_features = self._pad_and_stack_input_features(
615 + input_features, ).to(self.encoder.input_linear.weight.dtype)
616 +
617 + return GraniteSpeechAudioInputs(
618 + input_features=input_features,
619 + input_features_mask=input_features_mask,
620 + audio_embed_sizes=audio_embed_sizes.flatten().tolist(),
621 + )
622 +
623 + def _build_input_features_mask(
624 + self,
625 + audio_embed_sizes: torch.Tensor,
626 + ) -> torch.Tensor:
627 + """Calculate the input features mask, which will generally be used
628 + to mask the the padded features for all entries in the batch except
629 + for those with the most audio features.
630 +
631 + Args:
632 + audio_embed_sizes: torch.Tensor
633 + Tensor of num features in each seq in the batch.
634 + Returns:
635 + torch.Tensor: Mask of shape (bsz, num_features) to be applied to
636 + the audio features prior to splitting the audio embeddings.
637 + """
638 + most_audio_features = torch.max(audio_embed_sizes).item()
639 + mask_indices = torch.arange(
640 + most_audio_features,
641 + device=audio_embed_sizes.device,
642 + ).view(1, -1)
643 + input_features_mask = mask_indices < audio_embed_sizes.view(-1, 1)
644 + return input_features_mask
645 +
646 + def _pad_and_stack_input_features(
647 + self,
648 + input_features: list[torch.Tensor],
649 + ) -> torch.Tensor:
650 + """Given a list of input features of varying length, pad them to the
651 + same length and stack them into a torch.Tensor.
652 +
653 + NOTE: Usually, padding is done in the input processor/feature extractor
654 + and zero padded prior to the computation of the Mel features; the
655 + resulting values are only constant within a batch and generally nonzero
656 + (i.e., slightly negative nums); we should validate that this is okay
657 + since we don't use a feature attention mask, but the more important
658 + thing is that we apply the input_features_mask with variable len
659 + batches.
660 +
661 + Args:
662 + input_features: list[torch.Tensor]
663 + Input features to be coerced into a tensor.
664 + Returns:
665 + torch.Tensor: Tensor of shape [bsz, num_features, 160], where
666 + num_features is the max number of features of any entry in the
667 + batch.
668 + """
669 + # Input features are of shape [bsz, num_features, 160]
670 + feat_lens = [feats.shape[1] for feats in input_features]
671 + padding = [max(feat_lens) - length for length in feat_lens]
672 + # TODO (Alex) - Validate that it's okay to zero pad like this;
673 + # in transformers we zero pad prior to calculating the speech features,
674 + # so the value is not zero and is dependent on the batched features.
675 + padded = [
676 + torch.nn.functional.pad(feats, (0, 0, 0, pad, 0, 0))
677 + for feats, pad in zip(input_features, padding)
678 + ]
679 + stacked_features = torch.cat(padded, dim=0).to(input_features[0])
680 + return stacked_features
681 +
682 + def _process_audio_input(
683 + self,
684 + audio_input: GraniteSpeechAudioInputs,
685 + ) -> tuple[torch.Tensor]:
686 + """Compute the audio features to be merged into the LLM embeddings.
687 +
688 + Args:
689 + audio_input: GraniteSpeechAudioInputs
690 + Audio inputs object containing Mel features, an input features
691 + mask, and the (flattened) number of audio tokens per instance.
692 + Returns:
693 + tuple[torch.Tensor]: List of length bsz.
694 + """
695 + # TODO (Alex) - support embedding inputs
696 + encoder_embeds = self.encoder(audio_input["input_features"])
697 + # [bsz, <max feature size>, 4096]
698 + projected_embeds = self.projector(encoder_embeds)
699 + # Apply mask on variable length audio features
700 + masked_embeds = projected_embeds[audio_input["input_features_mask"]]
701 + # Split variable length features into a tuple
702 + return torch.split(masked_embeds, audio_input["audio_embed_sizes"])
703 +
704 + def get_multimodal_embeddings(
705 + self,
706 + **kwargs: object,
707 + ) -> Optional[MultiModalEmbeddings]:
708 + """Compute the audio embeddings if audio inputs are present."""
709 + audio_input = self._parse_and_validate_audio_input(**kwargs)
710 + if audio_input is None:
711 + return None
712 + audio_features = self._process_audio_input(audio_input)
713 + return audio_features
714 +
715 + def get_input_embeddings(
716 + self,
717 + input_ids: torch.Tensor,
718 + multimodal_embeddings: Optional[MultiModalEmbeddings] = None,
719 + ) -> torch.Tensor:
720 + """Compute the merged LLM / audio embeddings."""
721 + if multimodal_embeddings is None:
722 + return self.language_model.get_input_embeddings(input_ids)
723 +
724 + inputs_embeds = embed_multimodal(
725 + input_ids,
726 + self.config.audio_token_index,
727 + self.language_model.model.get_input_embeddings,
728 + multimodal_embeddings,
729 + )
730 + return inputs_embeds
731 +
732 + def forward(
733 + self,
734 + input_ids: torch.Tensor,
735 + positions: torch.Tensor,
736 + intermediate_tensors: Optional[IntermediateTensors] = None,
737 + inputs_embeds: Optional[torch.Tensor] = None,
738 + **kwargs: object,
739 + ) -> Union[torch.Tensor, IntermediateTensors]:
740 + if intermediate_tensors is not None:
741 + inputs_embeds = None
742 +
743 + # NOTE: In v1, inputs_embeds is always generated at model runner, this
744 + # condition is for v0 compatibility.
745 + elif inputs_embeds is None:
746 + audio_embeds = self.get_multimodal_embeddings(**kwargs)
747 + inputs_embeds = self.get_input_embeddings(input_ids, audio_embeds)
748 + input_ids = None
749 +
750 + model_output = self.language_model(input_ids, positions,
751 + intermediate_tensors, inputs_embeds)
752 + return model_output
753 +
754 + def compute_logits(
755 + self,
756 + hidden_states: torch.Tensor,
757 + sampling_metadata: SamplingMetadata,
758 + ) -> Optional[torch.Tensor]:
759 + return self.language_model.compute_logits(
760 + hidden_states,
761 + sampling_metadata,
762 + )
763 +
764 + def load_weights(
765 + self,
766 + weights: Iterable[Tuple[str, torch.Tensor]],
767 + ) -> Set[str]:
768 + loader = AutoWeightsLoader(self)
769 + return loader.load_weights(weights)
770 +
771 + def get_mm_mapping(self) -> MultiModelKeys:
772 + """Get the module prefix in multimodal models."""
773 + return MultiModelKeys.from_string_field(
774 + language_model="language_model",
775 + connector="projector",
776 + tower_model="encoder",
777 + )
vllm/model_executor/models/registry.py
+1
@@ -178,6 +178,7 @@ _MULTIMODAL_MODELS = {
178 "FuyuForCausalLM": ("fuyu", "FuyuForCausalLM"),
179 "Gemma3ForConditionalGeneration": ("gemma3_mm", "Gemma3ForConditionalGeneration"), # noqa: E501
180 "GLM4VForCausalLM": ("glm4v", "GLM4VForCausalLM"),
181 + "GraniteSpeechForConditionalGeneration": ("granite_speech", "GraniteSpeechForConditionalGeneration"), # noqa: E501
182 "H2OVLChatModel": ("h2ovl", "H2OVLChatModel"),
183 "InternVLChatModel": ("internvl", "InternVLChatModel"),
184 "Idefics3ForConditionalGeneration":("idefics3","Idefics3ForConditionalGeneration"),