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 13, 19:30 UTC
[Model, Core] Support Granite Speech & LoRA for STT (#24455)
Alex Brooks committed
Nov 5, 2025 at 00:33 UTC
b7cbc254169128a4203d111f3b87edaa17839a32
5 files changed
+169
-8
docs/models/supported_models.md
+1
@@ -761,6 +761,7 @@ Speech2Text models trained specifically for Automatic Speech Recognition.
761
| `WhisperForConditionalGeneration` | Whisper | `openai/whisper-small`, `openai/whisper-large-v3-turbo`, etc. | | |
762
| `VoxtralForConditionalGeneration` | Voxtral (Mistral format) | `mistralai/Voxtral-Mini-3B-2507`, `mistralai/Voxtral-Small-24B-2507`, etc. | ✅︎ | ✅︎ |
763
| `Gemma3nForConditionalGeneration` | Gemma3n | `google/gemma-3n-E2B-it`, `google/gemma-3n-E4B-it`, etc. | | |
764
+| `GraniteSpeechForConditionalGeneration` | Granite Speech | `ibm-granite/granite-speech-3.3-2b`, `ibm-granite/granite-speech-3.3-8b`, etc. | ✅︎ | ✅︎ |
765
766
### Pooling Models
767
tests/entrypoints/openai/test_transcription_validation.py
+35
@@ -65,6 +65,41 @@ async def test_basic_audio(mary_had_lamb, model_name):
65
assert out_usage["seconds"] == 16, out_usage["seconds"]
66
67
68
+@pytest.mark.asyncio
69
+async def test_basic_audio_with_lora(mary_had_lamb):
70
+ """Ensure STT (transcribe) requests can pass LoRA through to generate."""
71
+ model_name = "ibm-granite/granite-speech-3.3-2b"
72
+ lora_model_name = "speech"
73
+ server_args = [
74
+ "--enforce-eager",
75
+ "--enable-lora",
76
+ "--max-lora-rank",
77
+ "64",
78
+ "--lora-modules",
79
+ f"{lora_model_name}={model_name}",
80
+ "--max-model-len",
81
+ "2048",
82
+ "--max-num-seqs",
83
+ "1",
84
+ ]
85
+
86
+ # Based on https://github.com/openai/openai-cookbook/blob/main/examples/Whisper_prompting_guide.ipynb.
87
+ with RemoteOpenAIServer(model_name, server_args) as remote_server:
88
+ client = remote_server.get_async_client()
89
+ transcription = await client.audio.transcriptions.create(
90
+ model=lora_model_name,
91
+ file=mary_had_lamb,
92
+ language="en",
93
+ response_format="text",
94
+ temperature=0.0,
95
+ )
96
+ out = json.loads(transcription)
97
+ out_text = out["text"]
98
+ out_usage = out["usage"]
99
+ assert "mary had a little lamb" in out_text
100
+ assert out_usage["seconds"] == 16, out_usage["seconds"]
101
+
102
+
103
@pytest.mark.asyncio
104
async def test_basic_audio_gemma(foscolo):
105
# Gemma accuracy on some of the audio samples we use is particularly bad,
tests/entrypoints/openai/test_translation_validation.py
+34
@@ -48,6 +48,40 @@ async def test_non_asr_model(foscolo):
48
assert err["message"] == "The model does not support Translations API"
49
50
51
+@pytest.mark.asyncio
52
+async def test_basic_audio_with_lora(mary_had_lamb):
53
+ """Ensure STT (translate) requests can pass LoRA through to generate."""
54
+ # NOTE - careful to call this test before the module scoped server
55
+ # fixture, otherwise it'll OOMkill the CI
56
+ model_name = "ibm-granite/granite-speech-3.3-2b"
57
+ lora_model_name = "speech"
58
+ server_args = [
59
+ "--enforce-eager",
60
+ "--enable-lora",
61
+ "--max-lora-rank",
62
+ "64",
63
+ "--lora-modules",
64
+ f"{lora_model_name}={model_name}",
65
+ "--max-model-len",
66
+ "2048",
67
+ "--max-num-seqs",
68
+ "1",
69
+ ]
70
+
71
+ # Based on https://github.com/openai/openai-cookbook/blob/main/examples/Whisper_prompting_guide.ipynb.
72
+ with RemoteOpenAIServer(model_name, server_args) as remote_server:
73
+ client = remote_server.get_async_client()
74
+ translation = await client.audio.translations.create(
75
+ model=lora_model_name,
76
+ file=mary_had_lamb,
77
+ extra_body=dict(language="en", to_language="es"),
78
+ response_format="text",
79
+ temperature=0.0,
80
+ )
81
+ out = json.loads(translation)["text"].strip().lower()
82
+ assert "mary tenía un pequeño cordero" in out
83
+
84
+
85
# NOTE: (NickLucche) the large-v3-turbo model was not trained on translation!
86
@pytest.mark.asyncio
87
async def test_basic_audio(foscolo, client_and_model):
vllm/entrypoints/openai/speech_to_text.py
+2
-6
@@ -170,11 +170,6 @@ class OpenAISpeechToText(OpenAIServing):
170
try:
171
lora_request = self._maybe_get_adapters(request)
172
173
- if lora_request:
174
- return self.create_error_response(
175
- f"Currently do not support LoRA for {self.task_type.title()}."
176
- )
177
-
173
prompts, duration_s = await self._preprocess_speech_to_text(
174
request=request,
175
audio_data=audio_data,
@@ -199,7 +194,7 @@ class OpenAISpeechToText(OpenAIServing):
194
# It will not display special tokens like <|startoftranscript|>
195
request.prompt,
196
params=sampling_params,
202
- lora_request=None,
197
+ lora_request=lora_request,
198
)
199
200
list_result_generator = [
@@ -207,6 +202,7 @@ class OpenAISpeechToText(OpenAIServing):
202
prompt,
203
sampling_params,
204
request_id,
205
+ lora_request=lora_request,
206
)
207
for prompt in prompts
208
]
vllm/model_executor/models/granite_speech.py
+97
-2
@@ -26,15 +26,17 @@
26
27
import math
28
from collections.abc import Iterable, Mapping
29
-from typing import Annotated
29
+from typing import Annotated, Literal, cast
30
31
+import numpy as np
32
import torch
33
import torch.nn.functional as F
34
from torch import nn
35
from transformers import BatchFeature, PretrainedConfig
36
36
-from vllm.config import CacheConfig, VllmConfig
37
+from vllm.config import CacheConfig, ModelConfig, SpeechToTextConfig, VllmConfig
38
from vllm.config.multimodal import BaseDummyOptions
39
+from vllm.inputs.data import PromptType
40
from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear
41
from vllm.model_executor.layers.quantization import QuantizationConfig
42
from vllm.model_executor.models.module_mapping import MultiModelKeys
@@ -57,6 +59,8 @@ from vllm.multimodal.processing import (
59
)
60
from vllm.multimodal.profiling import BaseDummyInputsBuilder
61
from vllm.sequence import IntermediateTensors
62
+from vllm.transformers_utils.processor import cached_get_processor
63
+from vllm.transformers_utils.tokenizer import cached_get_tokenizer
64
from vllm.utils.tensor_schema import TensorSchema, TensorShape
65
66
from .blip2 import Blip2QFormerModel
@@ -65,9 +69,22 @@ from .interfaces import (
69
SupportsLoRA,
70
SupportsMultiModal,
71
SupportsPP,
72
+ SupportsTranscription,
73
)
74
from .utils import AutoWeightsLoader, init_vllm_registered_model, maybe_prefix
75
76
+# NOTE lang support is based on what is written here:
77
+# https://huggingface.co/ibm-granite/granite-speech-3.3-2b
78
+# Though this may vary from model to model, and also many langs
79
+# work pretty well with zero shot.
80
+ISO639_1_SUPPORTED_LANGS = {
81
+ "en": "English",
82
+ "fr": "French",
83
+ "de": "German",
84
+ "pt": "Portuguese",
85
+ "es": "Spanish",
86
+}
87
+
88
89
### Audio Input
90
class GraniteSpeechAudioInputs(TensorSchema):
@@ -545,8 +562,10 @@ class GraniteSpeechForConditionalGeneration(
562
SupportsMultiModal,
563
SupportsPP,
564
SupportsLoRA,
565
+ SupportsTranscription,
566
):
567
merge_by_field_config = True
568
+ supported_languages = ISO639_1_SUPPORTED_LANGS
569
570
packed_modules_mapping = {
571
"qkv_proj": [
@@ -816,3 +835,79 @@ class GraniteSpeechForConditionalGeneration(
835
connector="projector",
836
tower_model="encoder",
837
)
838
+
839
+ ### Support for speech-to-text Transcription
840
+ @classmethod
841
+ def get_generation_prompt(
842
+ cls,
843
+ audio: np.ndarray,
844
+ model_config: ModelConfig,
845
+ stt_config: SpeechToTextConfig,
846
+ language: str | None,
847
+ task_type: Literal["transcribe", "translate"],
848
+ request_prompt: str,
849
+ to_language: str | None,
850
+ ) -> PromptType:
851
+ """Get the generation prompt to be used for transcription requests."""
852
+ # Audio placeholders don't use an index, so value doesn't matter
853
+ audio_tok = cls.get_placeholder_str("audio", 0)
854
+
855
+ if task_type == "translate":
856
+ full_lang_name_to = cls.supported_languages.get(to_language, to_language)
857
+ user_prompt = f"{audio_tok}translate the speech to {full_lang_name_to}" # noqa: E501
858
+ elif task_type == "transcribe":
859
+ user_prompt = (
860
+ f"{audio_tok}can you transcribe the speech into a written format?" # noqa: E501
861
+ )
862
+ else:
863
+ raise ValueError(f"Unsupported task type {task_type}")
864
+
865
+ tokenizer = cached_get_tokenizer(model_config.model)
866
+ chat = [dict(role="user", content=user_prompt)]
867
+ prompt = tokenizer.apply_chat_template(
868
+ chat,
869
+ tokenize=False,
870
+ add_generation_prompt=True,
871
+ )
872
+
873
+ prompt_token_ids = tokenizer.encode(prompt)
874
+ prompt = {
875
+ "prompt_token_ids": prompt_token_ids,
876
+ "multi_modal_data": {"audio": audio},
877
+ }
878
+ return cast(PromptType, prompt)
879
+
880
+ # Adapted from https://github.com/huggingface/transformers/blob/v4.56.0/src/transformers/models/granite_speech/feature_extraction_granite_speech.py#L122 # noqa: E501
881
+ @classmethod
882
+ def get_num_audio_tokens(
883
+ cls,
884
+ audio_duration_s: float,
885
+ stt_config: SpeechToTextConfig,
886
+ model_config: ModelConfig,
887
+ ) -> int | None:
888
+ """Get the number of audio tokens for an audio duration in sec."""
889
+ processor = cached_get_processor(model_config.model)
890
+ hop_length = processor.audio_processor.melspec_kwargs["hop_length"]
891
+ proj_win_size = processor.audio_processor.projector_window_size
892
+ ds_rate = processor.audio_processor.projector_downsample_rate
893
+ effective_window_size = proj_win_size // ds_rate
894
+
895
+ raw_length = audio_duration_s * stt_config.sample_rate
896
+
897
+ # mel sequence length computation
898
+ mel_length = raw_length // hop_length + 1
899
+ # encoder frame takes two mel features
900
+ encoder_length = mel_length // 2
901
+ nblocks = math.ceil(encoder_length / proj_win_size)
902
+ # projector output length
903
+ return nblocks * effective_window_size
904
+
905
+ @classmethod
906
+ def get_speech_to_text_config(
907
+ cls, model_config: ModelConfig, task_type: str
908
+ ) -> SpeechToTextConfig:
909
+ """Get the stt config for this model."""
910
+ # Default settings are reasonable for this model and we don't currently
911
+ # expose this information in the model configs, but this may change in
912
+ # the future
913
+ return SpeechToTextConfig()