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

[model] support FireRedLID (#39290)

Signed-off-by: PatchouliTaisa <patchychen@tencent.com> Co-authored-by: PatchouliTaisa <patchychen@tencent.com>

PatchyTIS committed Apr 10, 2026 at 16:43 UTC 967146e7bdfb0cc3cb16fb5cc547bff9667ae0a2
14 files changed +1803 -347
docs/models/supported_models.md
+1
@@ -661,6 +661,7 @@ Speech2Text models trained specifically for Automatic Speech Recognition.
661 | ------------ | ------ | ----------------- | -------------------- | ------------------------- |
662 | `CohereAsrForConditionalGeneration` | Cohere-Transcribe | `CohereLabs/cohere-transcribe-03-2026` | | |
663 | `FireRedASR2ForConditionalGeneration` | FireRedASR2 | `allendou/FireRedASR2-LLM-vllm`, etc. | | |
664 +| `FireRedLIDForConditionalGeneration` | FireRedLID | `PatchyTisa/FireRedLID-vllm`, etc. | | |
665 | `FunASRForConditionalGeneration` | FunASR | `allendou/Fun-ASR-Nano-2512-vllm`, etc. | | |
666 | `Gemma3nForConditionalGeneration` | Gemma3n | `google/gemma-3n-E2B-it`, `google/gemma-3n-E4B-it`, etc. | | |
667 | `GlmAsrForConditionalGeneration` | GLM-ASR | `zai-org/GLM-ASR-Nano-2512` | ✅︎ | ✅︎ |
examples/offline_inference/audio_language.py
+21
@@ -537,9 +537,30 @@ def run_whisper(question: str, audio_count: int) -> ModelRequestData:
537 )
538
539
540 +# FireRedLID
541 +def run_fireredlid(question: str, audio_count: int) -> ModelRequestData:
542 + assert audio_count == 1, "FireRedLID only supports single audio input per prompt"
543 + model_name = "PatchyTisa/FireRedLID-vllm"
544 +
545 + prompt = "<sos>"
546 +
547 + engine_args = EngineArgs(
548 + model=model_name,
549 + max_model_len=8,
550 + max_num_seqs=5,
551 + limit_mm_per_prompt={"audio": audio_count},
552 + )
553 +
554 + return ModelRequestData(
555 + engine_args=engine_args,
556 + prompt=prompt,
557 + )
558 +
559 +
560 model_example_map = {
561 "audioflamingo3": run_audioflamingo3,
562 "cohere_asr": run_cohere_asr,
563 + "fireredlid": run_fireredlid,
564 "funaudiochat": run_funaudiochat,
565 "gemma3n": run_gemma3n,
566 "glmasr": run_glmasr,
examples/offline_inference/encoder_decoder_multimodal.py
+84
@@ -55,7 +55,91 @@ def run_whisper():
55 )
56
57
58 +def run_fireredasr2():
59 + """
60 + FireRedASR2 – Automatic Speech Recognition model.
61 +
62 + This model uses a Conformer encoder + Qwen2 LLM decoder architecture
63 + for speech-to-text transcription. Audio is passed via the implicit
64 + prompt format with the ``<|AUDIO|>`` placeholder token.
65 + """
66 + engine_args = EngineArgs(
67 + model="allendou/FireRedASR2-LLM-vllm",
68 + max_model_len=448,
69 + max_num_seqs=16,
70 + limit_mm_per_prompt={"audio": 1},
71 + )
72 +
73 + prompt_str = (
74 + "<|im_start|>user\n<|AUDIO|>请转写音频为文字<|im_end|>\n<|im_start|>assistant\n"
75 + )
76 +
77 + prompts = [
78 + { # Implicit prompt with audio
79 + "prompt": prompt_str,
80 + "multi_modal_data": {
81 + "audio": AudioAsset("mary_had_lamb").audio_and_sample_rate,
82 + },
83 + },
84 + { # Another audio sample
85 + "prompt": prompt_str,
86 + "multi_modal_data": {
87 + "audio": AudioAsset("winning_call").audio_and_sample_rate,
88 + },
89 + },
90 + ]
91 +
92 + return ModelRequestData(
93 + engine_args=engine_args,
94 + prompts=prompts,
95 + )
96 +
97 +
98 +def run_fireredlid():
99 + """
100 + FireRedLID – Language Identification model.
101 +
102 + This encoder-decoder model identifies the spoken language of an audio
103 + clip. It outputs at most 2 tokens representing the detected language
104 + (e.g. "en", "zh mandarin").
105 + """
106 + engine_args = EngineArgs(
107 + model="PatchyTisa/FireRedLID-vllm",
108 + max_model_len=8,
109 + max_num_seqs=16,
110 + limit_mm_per_prompt={"audio": 1},
111 + )
112 +
113 + prompts = [
114 + { # Test explicit encoder/decoder prompt
115 + "encoder_prompt": {
116 + "prompt": "",
117 + "multi_modal_data": {
118 + "audio": AudioAsset("mary_had_lamb").audio_and_sample_rate,
119 + },
120 + },
121 + "decoder_prompt": "<sos>",
122 + },
123 + { # Another audio sample
124 + "encoder_prompt": {
125 + "prompt": "",
126 + "multi_modal_data": {
127 + "audio": AudioAsset("winning_call").audio_and_sample_rate,
128 + },
129 + },
130 + "decoder_prompt": "<sos>",
131 + },
132 + ]
133 +
134 + return ModelRequestData(
135 + engine_args=engine_args,
136 + prompts=prompts,
137 + )
138 +
139 +
140 model_example_map = {
141 + "fireredasr2": run_fireredasr2,
142 + "fireredlid": run_fireredlid,
143 "whisper": run_whisper,
144 }
145
examples/online_serving/openai_lid_client.py new
+193
@@ -0,0 +1,193 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Language Identification (LID) demo using the FireRedLID model on vLLM.
5 +
6 +FireRedLID is an audio encoder-decoder model that identifies the spoken
7 +language of an audio clip. Unlike ASR models that output full transcriptions,
8 +FireRedLID outputs at most 2 tokens representing the detected language
9 +(e.g. "en", "zh mandarin").
10 +
11 +Start the vLLM server:
12 +
13 + vllm serve PatchyTisa/FireRedLID-vllm
14 +
15 +Then run this script:
16 +
17 + # Use the built-in sample audio
18 + python examples/online_serving/openai_lid_client.py
19 +
20 + # Use your own audio file(s)
21 + python examples/online_serving/openai_lid_client.py \
22 + --audio_paths audio_en.wav audio_zh.wav audio_fr.wav
23 +
24 + # Batch-identify multiple files in one run
25 + python examples/online_serving/openai_lid_client.py \
26 + --audio_paths /path/to/dir/*.wav
27 +
28 +Requirements:
29 +- vLLM with audio support
30 +- openai Python SDK
31 +- kaldi_native_fbank (pulled in by the model)
32 +"""
33 +
34 +import argparse
35 +import json
36 +import os
37 +
38 +from openai import OpenAI
39 +
40 +from vllm.assets.audio import AudioAsset
41 +
42 +# ──────────────────────────────────────────────────────────────────────
43 +# Helpers
44 +# ──────────────────────────────────────────────────────────────────────
45 +
46 +
47 +def identify_language(
48 + audio_path: str,
49 + client: OpenAI,
50 + model: str,
51 +) -> str:
52 + """
53 + Send a single audio file to the vLLM transcription endpoint and return
54 + the detected language tag.
55 +
56 + FireRedLID re-uses the OpenAI-compatible ``/v1/audio/transcriptions``
57 + endpoint. The "transcription" it returns is actually the language label
58 + (e.g. ``"en"`` or ``"zh mandarin"``).
59 + """
60 + with open(audio_path, "rb") as f:
61 + result = client.audio.transcriptions.create(
62 + file=f,
63 + model=model,
64 + response_format="json",
65 + temperature=0.0,
66 + )
67 + return result.text.strip()
68 +
69 +
70 +def identify_language_raw(
71 + audio_path: str,
72 + model: str,
73 + api_base: str,
74 +) -> str:
75 + """
76 + Same as :func:`identify_language` but uses raw HTTP so that the demo
77 + works without the ``openai`` SDK (useful for quick debugging).
78 + """
79 + import requests
80 +
81 + url = f"{api_base}/audio/transcriptions"
82 + with open(audio_path, "rb") as f:
83 + files = {"file": (os.path.basename(audio_path), f)}
84 + data = {
85 + "model": model,
86 + "response_format": "json",
87 + }
88 + resp = requests.post(url, files=files, data=data)
89 + resp.raise_for_status()
90 + return resp.json()["text"].strip()
91 +
92 +
93 +def identify_language_streaming(
94 + audio_path: str,
95 + model: str,
96 + api_base: str,
97 +) -> str:
98 + """
99 + Streaming variant – demonstrates the streaming transcription endpoint.
100 + For a 1-2 token output the stream finishes almost instantly, but this
101 + shows that the API path works end-to-end.
102 + """
103 + import requests
104 +
105 + url = f"{api_base}/audio/transcriptions"
106 + with open(audio_path, "rb") as f:
107 + files = {"file": (os.path.basename(audio_path), f)}
108 + data = {
109 + "stream": "true",
110 + "model": model,
111 + "response_format": "json",
112 + }
113 + response = requests.post(url, files=files, data=data, stream=True)
114 + response.raise_for_status()
115 +
116 + tokens: list[str] = []
117 + for chunk in response.iter_lines(
118 + chunk_size=8192, decode_unicode=False, delimiter=b"\n"
119 + ):
120 + if not chunk:
121 + continue
122 + payload = json.loads(chunk[len("data: ") :].decode("utf-8"))
123 + choice = payload["choices"][0]
124 + delta = choice.get("delta", {}).get("content", "")
125 + if delta:
126 + tokens.append(delta)
127 + if choice.get("finish_reason") is not None:
128 + break
129 +
130 + return "".join(tokens).strip()
131 +
132 +
133 +# ──────────────────────────────────────────────────────────────────────
134 +# Main
135 +# ──────────────────────────────────────────────────────────────────────
136 +
137 +
138 +def main(args: argparse.Namespace) -> None:
139 + api_base = args.api_base.rstrip("/")
140 + client = OpenAI(api_key="EMPTY", base_url=api_base)
141 + model = client.models.list().data[0].id
142 + print(f"Model : {model}")
143 + print(f"Server: {api_base}\n")
144 +
145 + # Resolve audio paths ------------------------------------------------
146 + if args.audio_paths:
147 + audio_paths = args.audio_paths
148 + else:
149 + # Fall back to the built-in vLLM sample audios (both are English).
150 + audio_paths = [
151 + str(AudioAsset("mary_had_lamb").get_local_path()),
152 + str(AudioAsset("winning_call").get_local_path()),
153 + ]
154 +
155 + # Run LID for each file ----------------------------------------------
156 + print(f"{'Audio File':<50} {'Language (sync)':<20} {'Language (stream)'}")
157 + print("-" * 90)
158 +
159 + for path in audio_paths:
160 + basename = os.path.basename(path)
161 +
162 + # 1) Synchronous via OpenAI SDK
163 + lang_sync = identify_language(path, client, model)
164 +
165 + # 2) Streaming via raw HTTP
166 + lang_stream = identify_language_streaming(path, model, api_base)
167 +
168 + print(f"{basename:<50} {lang_sync:<20} {lang_stream}")
169 +
170 + print()
171 +
172 +
173 +if __name__ == "__main__":
174 + parser = argparse.ArgumentParser(
175 + description="FireRedLID – Language Identification demo via vLLM",
176 + )
177 + parser.add_argument(
178 + "--audio_paths",
179 + nargs="+",
180 + default=None,
181 + help=(
182 + "One or more audio files to identify. "
183 + "If omitted, uses vLLM's built-in sample audios."
184 + ),
185 + )
186 + parser.add_argument(
187 + "--api_base",
188 + type=str,
189 + default="http://localhost:8000/v1",
190 + help="vLLM API base URL (default: http://localhost:8000/v1)",
191 + )
192 + args = parser.parse_args()
193 + main(args)
tests/models/registry.py
+3
@@ -820,6 +820,9 @@ _MULTIMODAL_EXAMPLE_MODELS = {
820 "FireRedASR2ForConditionalGeneration": _HfExamplesInfo(
821 "allendou/FireRedASR2-LLM-vllm",
822 ),
823 + "FireRedLIDForConditionalGeneration": _HfExamplesInfo(
824 + "PatchyTisa/FireRedLID-vllm",
825 + ),
826 "FunASRForConditionalGeneration": _HfExamplesInfo(
827 "allendou/Fun-ASR-Nano-2512-vllm",
828 ),
vllm/model_executor/models/conformer_encoder.py new
+350
@@ -0,0 +1,350 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Shared Conformer encoder components for FireRedASR2 and FireRedLID.
5 +
6 +Both models use the same Conformer-based audio encoder architecture
7 +(Conv2dSubsampling → RelPositionalEncoding → N × RelPosEmbConformerBlock).
8 +This module factors out the common building blocks to avoid duplication.
9 +"""
10 +
11 +import torch
12 +import torch.nn.functional as F
13 +from torch import nn
14 +
15 +from vllm.model_executor.layers.linear import ReplicatedLinear
16 +
17 +
18 +class Conv2dSubsampling(nn.Module):
19 + def __init__(self, idim: int, d_model: int, out_channels: int = 32):
20 + super().__init__()
21 + self.conv = nn.Sequential(
22 + nn.Conv2d(1, out_channels, 3, 2),
23 + nn.ReLU(),
24 + nn.Conv2d(out_channels, out_channels, 3, 2),
25 + nn.ReLU(),
26 + )
27 + subsample_idim = ((idim - 1) // 2 - 1) // 2
28 + self.out = ReplicatedLinear(
29 + input_size=out_channels * subsample_idim,
30 + output_size=d_model,
31 + bias=True,
32 + )
33 +
34 + self.subsampling = 4
35 + left_context = right_context = 3 # both exclude current frame
36 + self.context = left_context + 1 + right_context # 7
37 +
38 + def forward(
39 + self, x: torch.Tensor, x_mask: torch.Tensor
40 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
41 + x = x.unsqueeze(1)
42 + x = self.conv(x)
43 + N, C, T, D = x.size()
44 + x, _ = self.out(x.transpose(1, 2).contiguous().view(N, T, C * D))
45 + mask = x_mask[:, :, :-2:2][:, :, :-2:2]
46 + input_lengths = mask[:, -1, :].sum(dim=-1)
47 + return x, input_lengths, mask
48 +
49 +
50 +class Swish(nn.Module):
51 + def forward(self, x: torch.Tensor) -> torch.Tensor:
52 + return x * torch.sigmoid(x)
53 +
54 +
55 +class RelPositionalEncoding(nn.Module):
56 + def __init__(self, d_model: int, max_len: int = 5000):
57 + super().__init__()
58 + pe_positive = torch.zeros(max_len, d_model, requires_grad=False)
59 + pe_negative = torch.zeros(max_len, d_model, requires_grad=False)
60 + position = torch.arange(0, max_len).unsqueeze(1).float()
61 + div_term = torch.exp(
62 + torch.arange(0, d_model, 2).float()
63 + * -(torch.log(torch.tensor(10000.0)).item() / d_model)
64 + )
65 + pe_positive[:, 0::2] = torch.sin(position * div_term)
66 + pe_positive[:, 1::2] = torch.cos(position * div_term)
67 + pe_negative[:, 0::2] = torch.sin(-1 * position * div_term)
68 + pe_negative[:, 1::2] = torch.cos(-1 * position * div_term)
69 +
70 + pe_positive = torch.flip(pe_positive, [0]).unsqueeze(0)
71 + pe_negative = pe_negative[1:].unsqueeze(0)
72 + self.pe = torch.cat([pe_positive, pe_negative], dim=1)
73 +
74 + def forward(self, x: torch.Tensor) -> torch.Tensor:
75 + # Tmax = 2 * max_len - 1
76 + Tmax, T = self.pe.size(1), x.size(1)
77 + pos_emb = self.pe[:, Tmax // 2 - T + 1 : Tmax // 2 + T].clone().detach()
78 + return pos_emb
79 +
80 +
81 +class ConformerFeedForward(nn.Module):
82 + def __init__(self, d_model: int):
83 + super().__init__()
84 + self.pre_layer_norm = nn.LayerNorm(d_model)
85 + self.linear_expand = ReplicatedLinear(
86 + input_size=d_model,
87 + output_size=d_model * 4,
88 + bias=True,
89 + )
90 + self.nonlinear = Swish()
91 + self.linear_project = ReplicatedLinear(
92 + input_size=d_model * 4,
93 + output_size=d_model,
94 + bias=True,
95 + )
96 +
97 + def forward(self, x: torch.Tensor) -> torch.Tensor:
98 + residual = x
99 + x = self.pre_layer_norm(x)
100 + x, _ = self.linear_expand(x)
101 + x = self.nonlinear(x)
102 + x, _ = self.linear_project(x)
103 + return x + residual
104 +
105 +
106 +class EncoderMultiHeadAttention(nn.Module):
107 + def __init__(self, n_head: int, d_model: int):
108 + super().__init__()
109 + assert d_model % n_head == 0
110 + self.n_head = n_head
111 + self.d_k = d_model // n_head
112 + self.d_v = self.d_k
113 +
114 + self.w_qs = ReplicatedLinear(d_model, n_head * self.d_k, bias=False)
115 + self.w_ks = ReplicatedLinear(d_model, n_head * self.d_k, bias=False)
116 + self.w_vs = ReplicatedLinear(d_model, n_head * self.d_v, bias=False)
117 +
118 + self.layer_norm_q = nn.LayerNorm(d_model)
119 + self.layer_norm_k = nn.LayerNorm(d_model)
120 + self.layer_norm_v = nn.LayerNorm(d_model)
121 +
122 + self.fc = ReplicatedLinear(n_head * self.d_v, d_model, bias=False)
123 +
124 + def forward_qkv(
125 + self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor
126 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
127 + d_k, d_v, n_head = self.d_k, self.d_v, self.n_head
128 + sz_b, len_q, len_k, len_v = q.size(0), q.size(1), k.size(1), v.size(1)
129 +
130 + q = self.layer_norm_q(q)
131 + k = self.layer_norm_k(k)
132 + v = self.layer_norm_v(v)
133 +
134 + q = self.w_qs(q)[0].view(sz_b, len_q, n_head, d_k)
135 + k = self.w_ks(k)[0].view(sz_b, len_k, n_head, d_k)
136 + v = self.w_vs(v)[0].view(sz_b, len_v, n_head, d_v)
137 + q = q.transpose(1, 2)
138 + k = k.transpose(1, 2)
139 + v = v.transpose(1, 2)
140 + return q, k, v
141 +
142 + def forward_output(
143 + self,
144 + output: torch.Tensor,
145 + residual: torch.Tensor,
146 + sz_b: int,
147 + len_q: int,
148 + ) -> torch.Tensor:
149 + output = output.transpose(1, 2).contiguous().view(sz_b, len_q, -1)
150 + fc_out, _ = self.fc(output)
151 + return fc_out + residual
152 +
153 + def forward_attention(
154 + self,
155 + attn: torch.Tensor,
156 + v: torch.Tensor,
157 + mask: torch.Tensor | None = None,
158 + ) -> tuple[torch.Tensor, torch.Tensor]:
159 + if mask is not None:
160 + mask = mask.unsqueeze(1)
161 + mask = mask.eq(0)
162 + attn = attn.masked_fill(mask, -float("inf"))
163 + attn = torch.softmax(attn, dim=-1).masked_fill(mask, 0.0)
164 + else:
165 + attn = torch.softmax(attn, dim=-1)
166 + output = torch.matmul(attn, v)
167 + return output, attn
168 +
169 +
170 +class RelPosMultiHeadAttention(EncoderMultiHeadAttention):
171 + def __init__(self, n_head: int, d_model: int):
172 + super().__init__(n_head, d_model)
173 + d_k = d_model // n_head
174 + self.scale = 1.0 / (d_k**0.5)
175 + self.linear_pos = ReplicatedLinear(d_model, n_head * d_k, bias=False)
176 + self.pos_bias_u = nn.Parameter(torch.empty([n_head, d_k]))
177 + self.pos_bias_v = nn.Parameter(torch.empty([n_head, d_k]))
178 +
179 + def _rel_shift(self, x):
180 + N, H, T1, T2 = x.size()
181 + zero_pad = torch.zeros((N, H, T1, 1), device=x.device, dtype=x.dtype)
182 + x_padded = torch.cat([zero_pad, x], dim=-1)
183 + x_padded = x_padded.view(N, H, T2 + 1, T1)
184 + x = x_padded[:, :, 1:].view_as(x)
185 + x = x[:, :, :, : x.size(-1) // 2 + 1]
186 + return x
187 +
188 + def forward(
189 + self,
190 + q: torch.Tensor,
191 + k: torch.Tensor,
192 + v: torch.Tensor,
193 + pos_emb: torch.Tensor,
194 + mask: torch.Tensor | None = None,
195 + ) -> tuple[torch.Tensor, torch.Tensor]:
196 + sz_b, len_q = q.size(0), q.size(1)
197 + residual = q
198 + q, k, v = self.forward_qkv(q, k, v)
199 +
200 + q = q.transpose(1, 2)
201 + n_batch_pos = pos_emb.size(0)
202 + p = self.linear_pos(pos_emb)[0].view(n_batch_pos, -1, self.n_head, self.d_k)
203 + p = p.transpose(1, 2)
204 +
205 + q_with_bias_u = (q + self.pos_bias_u).transpose(1, 2)
206 + q_with_bias_v = (q + self.pos_bias_v).transpose(1, 2)
207 +
208 + matrix_ac = torch.matmul(q_with_bias_u, k.transpose(-2, -1))
209 + matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1))
210 + matrix_bd = self._rel_shift(matrix_bd)
211 +
212 + attn_scores = matrix_ac + matrix_bd
213 + attn_scores.mul_(self.scale)
214 +
215 + output, attn = self.forward_attention(attn_scores, v, mask=mask)
216 + output = self.forward_output(output, residual, sz_b, len_q)
217 + return output, attn
218 +
219 +
220 +class ConformerConvolution(nn.Module):
221 + def __init__(self, d_model: int, kernel_size: int = 33):
222 + super().__init__()
223 + assert kernel_size % 2 == 1
224 + self.pre_layer_norm = nn.LayerNorm(d_model)
225 + self.pointwise_conv1 = nn.Conv1d(
226 + d_model, d_model * 4, kernel_size=1, bias=False
227 + )
228 + self.padding = (kernel_size - 1) // 2
229 + self.depthwise_conv = nn.Conv1d(
230 + d_model * 2,
231 + d_model * 2,
232 + kernel_size,
233 + stride=1,
234 + padding=self.padding,
235 + groups=d_model * 2,
236 + bias=False,
237 + )
238 + self.batch_norm = nn.LayerNorm(d_model * 2)
239 + self.swish = Swish()
240 + self.pointwise_conv2 = nn.Conv1d(
241 + d_model * 2, d_model, kernel_size=1, bias=False
242 + )
243 +
244 + def forward(
245 + self, x: torch.Tensor, mask: torch.Tensor | None = None
246 + ) -> torch.Tensor:
247 + residual = x
248 + out = self.pre_layer_norm(x)
249 + out = out.transpose(1, 2)
250 + if mask is not None:
251 + out.masked_fill_(mask.ne(1), 0.0)
252 + out = self.pointwise_conv1(out)
253 + out = F.glu(out, dim=1)
254 + out = self.depthwise_conv(out)
255 + out = out.transpose(1, 2)
256 + out = self.swish(self.batch_norm(out))
257 + out = out.transpose(1, 2)
258 + out = self.pointwise_conv2(out)
259 + if mask is not None:
260 + out.masked_fill_(mask.ne(1), 0.0)
261 + out = out.transpose(1, 2)
262 + return out + residual
263 +
264 +
265 +class RelPosEmbConformerBlock(nn.Module):
266 + def __init__(self, d_model: int, n_head: int, kernel_size: int = 33):
267 + super().__init__()
268 + self.ffn1 = ConformerFeedForward(d_model)
269 + self.mhsa = RelPosMultiHeadAttention(n_head, d_model)
270 + self.conv = ConformerConvolution(d_model, kernel_size)
271 + self.ffn2 = ConformerFeedForward(d_model)
272 + self.layer_norm = nn.LayerNorm(d_model)
273 +
274 + def forward(
275 + self,
276 + x: torch.Tensor,
277 + pos_emb: torch.Tensor,
278 + slf_attn_mask: torch.Tensor | None = None,
279 + pad_mask: torch.Tensor | None = None,
280 + ) -> torch.Tensor:
281 + out = 0.5 * x + 0.5 * self.ffn1(x)
282 + out = self.mhsa(out, out, out, pos_emb, mask=slf_attn_mask)[0]
283 + out = self.conv(out, pad_mask)
284 + out = 0.5 * out + 0.5 * self.ffn2(out)
285 + out = self.layer_norm(out)
286 + return out
287 +
288 +
289 +class ConformerEncoder(nn.Module):
290 + """
291 + Conformer encoder shared by FireRedASR2 and FireRedLID.
292 + """
293 +
294 + def __init__(
295 + self,
296 + idim: int,
297 + n_layers_enc: int,
298 + n_head: int,
299 + d_model: int,
300 + kernel_size: int = 33,
301 + pe_maxlen: int = 5000,
302 + ):
303 + super().__init__()
304 + self.odim = d_model
305 +
306 + self.input_preprocessor = Conv2dSubsampling(idim, d_model)
307 + self.positional_encoding = RelPositionalEncoding(d_model, max_len=pe_maxlen)
308 +
309 + self.layer_stack = nn.ModuleList()
310 + for _ in range(n_layers_enc):
311 + block = RelPosEmbConformerBlock(d_model, n_head, kernel_size)
312 + self.layer_stack.append(block)
313 +
314 + def forward(
315 + self,
316 + padded_input: torch.Tensor,
317 + input_lengths: torch.Tensor,
318 + pad: bool = True,
319 + ):
320 + if pad:
321 + padded_input = F.pad(
322 + padded_input,
323 + (0, 0, 0, self.input_preprocessor.context - 1),
324 + "constant",
325 + 0.0,
326 + )
327 + src_mask = self.padding_position_is_0(padded_input, input_lengths)
328 +
329 + embed_output, input_lengths, src_mask = self.input_preprocessor(
330 + padded_input, src_mask
331 + )
332 + enc_output = embed_output
333 +
334 + pos_emb = self.positional_encoding(embed_output)
335 +
336 + for enc_layer in self.layer_stack:
337 + enc_output = enc_layer(
338 + enc_output, pos_emb, slf_attn_mask=src_mask, pad_mask=src_mask
339 + )
340 +
341 + return enc_output, input_lengths, src_mask
342 +
343 + def padding_position_is_0(
344 + self, padded_input: torch.Tensor, input_lengths: torch.Tensor
345 + ) -> torch.Tensor:
346 + N, T = padded_input.size()[:2]
347 + # Use broadcasting instead of a Python loop for efficiency.
348 + positions = torch.arange(T, device=padded_input.device).unsqueeze(0)
349 + mask = (positions < input_lengths.unsqueeze(1)).to(torch.uint8)
350 + return mask.unsqueeze(1)
vllm/model_executor/models/fireredasr2.py
+1 -347
@@ -6,7 +6,6 @@ from typing import Annotated, Literal, cast
6
7 import numpy as np
8 import torch
9 -import torch.nn.functional as F
9 from torch import nn
10 from transformers import (
11 BatchFeature,
@@ -45,6 +44,7 @@ from vllm.transformers_utils.processors.fireredasr2 import (
44 )
45 from vllm.utils.tensor_schema import TensorSchema, TensorShape
46
47 +from .conformer_encoder import ConformerEncoder
48 from .interfaces import (
49 MultiModalEmbeddings,
50 SupportsMultiModal,
@@ -84,352 +84,6 @@ class FireRedASR2AudioInputs(TensorSchema):
84 ]
85
86
87 -class Swish(nn.Module):
88 - def forward(self, x: torch.Tensor) -> torch.Tensor:
89 - return x * torch.sigmoid(x)
90 -
91 -
92 -class Conv2dSubsampling(nn.Module):
93 - def __init__(self, idim: int, d_model: int, out_channels: int = 32):
94 - super().__init__()
95 - self.conv = nn.Sequential(
96 - nn.Conv2d(1, out_channels, 3, 2),
97 - nn.ReLU(),
98 - nn.Conv2d(out_channels, out_channels, 3, 2),
99 - nn.ReLU(),
100 - )
101 - subsample_idim = ((idim - 1) // 2 - 1) // 2
102 - self.out = ReplicatedLinear(
103 - input_size=out_channels * subsample_idim,
104 - output_size=d_model,
105 - bias=True,
106 - )
107 -
108 - self.subsampling = 4
109 - left_context = right_context = 3 # both exclude current frame
110 - self.context = left_context + 1 + right_context # 7
111 -
112 - def forward(
113 - self, x: torch.Tensor, x_mask: torch.Tensor
114 - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
115 - x = x.unsqueeze(1)
116 - x = self.conv(x)
117 - N, C, T, D = x.size()
118 - x, _ = self.out(x.transpose(1, 2).contiguous().view(N, T, C * D))
119 - mask = x_mask[:, :, :-2:2][:, :, :-2:2]
120 - input_lengths = mask[:, -1, :].sum(dim=-1)
121 - return x, input_lengths, mask
122 -
123 -
124 -class RelPositionalEncoding(nn.Module):
125 - def __init__(self, d_model: int, max_len: int = 5000):
126 - super().__init__()
127 - pe_positive = torch.zeros(max_len, d_model, requires_grad=False)
128 - pe_negative = torch.zeros(max_len, d_model, requires_grad=False)
129 - position = torch.arange(0, max_len).unsqueeze(1).float()
130 - div_term = torch.exp(
131 - torch.arange(0, d_model, 2).float()
132 - * -(torch.log(torch.tensor(10000.0)).item() / d_model)
133 - )
134 - pe_positive[:, 0::2] = torch.sin(position * div_term)
135 - pe_positive[:, 1::2] = torch.cos(position * div_term)
136 - pe_negative[:, 0::2] = torch.sin(-1 * position * div_term)
137 - pe_negative[:, 1::2] = torch.cos(-1 * position * div_term)
138 -
139 - pe_positive = torch.flip(pe_positive, [0]).unsqueeze(0)
140 - pe_negative = pe_negative[1:].unsqueeze(0)
141 - self.pe = torch.cat([pe_positive, pe_negative], dim=1)
142 -
143 - def forward(self, x: torch.Tensor) -> torch.Tensor:
144 - # Tmax = 2 * max_len - 1
145 - Tmax, T = self.pe.size(1), x.size(1)
146 - pos_emb = self.pe[:, Tmax // 2 - T + 1 : Tmax // 2 + T].clone().detach()
147 - return pos_emb
148 -
149 -
150 -class ConformerFeedForward(nn.Module):
151 - def __init__(self, d_model: int):
152 - super().__init__()
153 - self.pre_layer_norm = nn.LayerNorm(d_model)
154 - self.linear_expand = ReplicatedLinear(
155 - input_size=d_model,
156 - output_size=d_model * 4,
157 - bias=True,
158 - )
159 - self.nonlinear = Swish()
160 - self.linear_project = ReplicatedLinear(
161 - input_size=d_model * 4,
162 - output_size=d_model,
163 - bias=True,
164 - )
165 -
166 - def forward(self, x: torch.Tensor) -> torch.Tensor:
167 - residual = x
168 - x = self.pre_layer_norm(x)
169 - x, _ = self.linear_expand(x)
170 - x = self.nonlinear(x)
171 - x, _ = self.linear_project(x)
172 - output = x + residual
173 - return output
174 -
175 -
176 -class EncoderMultiHeadAttention(nn.Module):
177 - def __init__(self, n_head: int, d_model: int):
178 - super().__init__()
179 - assert d_model % n_head == 0
180 - self.n_head = n_head
181 - self.d_k = d_model // n_head
182 - self.d_v = self.d_k
183 -
184 - self.w_qs = ReplicatedLinear(
185 - input_size=d_model, output_size=n_head * self.d_k, bias=False
186 - )
187 - self.w_ks = ReplicatedLinear(
188 - input_size=d_model, output_size=n_head * self.d_k, bias=False
189 - )
190 - self.w_vs = ReplicatedLinear(
191 - input_size=d_model, output_size=n_head * self.d_v, bias=False
192 - )
193 -
194 - self.layer_norm_q = nn.LayerNorm(d_model)
195 - self.layer_norm_k = nn.LayerNorm(d_model)
196 - self.layer_norm_v = nn.LayerNorm(d_model)
197 -
198 - self.fc = ReplicatedLinear(
199 - input_size=n_head * self.d_v, output_size=d_model, bias=False
200 - )
201 -
202 - def forward_qkv(
203 - self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor
204 - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
205 - d_k, d_v, n_head = self.d_k, self.d_v, self.n_head
206 - sz_b, len_q, len_k, len_v = q.size(0), q.size(1), k.size(1), v.size(1)
207 -
208 - q = self.layer_norm_q(q)
209 - k = self.layer_norm_k(k)
210 - v = self.layer_norm_v(v)
211 -
212 - q = self.w_qs(q)[0].view(sz_b, len_q, n_head, d_k)
213 - k = self.w_ks(k)[0].view(sz_b, len_k, n_head, d_k)
214 - v = self.w_vs(v)[0].view(sz_b, len_v, n_head, d_v)
215 - q = q.transpose(1, 2)
216 - k = k.transpose(1, 2)
217 - v = v.transpose(1, 2)
218 - return q, k, v
219 -
220 - def forward_output(
221 - self, output: torch.Tensor, residual: torch.Tensor, sz_b: int, len_q: int
222 - ) -> torch.Tensor:
223 - output = output.transpose(1, 2).contiguous().view(sz_b, len_q, -1)
224 - fc_out, _ = self.fc(output)
225 - output = fc_out
226 - output = output + residual
227 - return output
228 -
229 - def forward_attention(
230 - self, attn: torch.Tensor, v: torch.Tensor, mask: torch.Tensor | None = None
231 - ) -> tuple[torch.Tensor, torch.Tensor]:
232 - if mask is not None:
233 - mask = mask.unsqueeze(1)
234 - mask = mask.eq(0)
235 - attn = attn.masked_fill(mask, -float("inf"))
236 - attn = torch.softmax(attn, dim=-1).masked_fill(mask, 0.0)
237 - else:
238 - attn = torch.softmax(attn, dim=-1)
239 -
240 - d_attn = attn
241 - output = torch.matmul(d_attn, v)
242 -
243 - return output, attn
244 -
245 -
246 -class RelPosMultiHeadAttention(EncoderMultiHeadAttention):
247 - def __init__(self, n_head: int, d_model: int):
248 - super().__init__(n_head, d_model)
249 - d_k = d_model // n_head
250 - self.scale = 1.0 / (d_k**0.5)
251 - self.linear_pos = ReplicatedLinear(
252 - input_size=d_model, output_size=n_head * d_k, bias=False
253 - )
254 - self.pos_bias_u = nn.Parameter(torch.empty([n_head, d_k]))
255 - self.pos_bias_v = nn.Parameter(torch.empty([n_head, d_k]))
256 -
257 - def _rel_shift(self, x):
258 - N, H, T1, T2 = x.size()
259 - zero_pad = torch.zeros((N, H, T1, 1), device=x.device, dtype=x.dtype)
260 - x_padded = torch.cat([zero_pad, x], dim=-1)
261 -
262 - x_padded = x_padded.view(N, H, T2 + 1, T1)
263 - x = x_padded[:, :, 1:].view_as(x)
264 - x = x[:, :, :, : x.size(-1) // 2 + 1]
265 - return x
266 -
267 - def forward(
268 - self,
269 - q: torch.Tensor,
270 - k: torch.Tensor,
271 - v: torch.Tensor,
272 - pos_emb: torch.Tensor,
273 - mask: torch.Tensor | None = None,
274 - ) -> tuple[torch.Tensor, torch.Tensor]:
275 - sz_b, len_q = q.size(0), q.size(1)
276 -
277 - residual = q
278 - q, k, v = self.forward_qkv(q, k, v)
279 -
280 - q = q.transpose(1, 2)
281 - n_batch_pos = pos_emb.size(0)
282 - p = self.linear_pos(pos_emb)[0].view(n_batch_pos, -1, self.n_head, self.d_k)
283 - p = p.transpose(1, 2)
284 -
285 - q_with_bias_u = (q + self.pos_bias_u).transpose(1, 2)
286 - q_with_bias_v = (q + self.pos_bias_v).transpose(1, 2)
287 -
288 - matrix_ac = torch.matmul(q_with_bias_u, k.transpose(-2, -1))
289 -
290 - matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1))
291 - matrix_bd = self._rel_shift(matrix_bd)
292 -
293 - attn_scores = matrix_ac + matrix_bd
294 - attn_scores.mul_(self.scale)
295 -
296 - output, attn = self.forward_attention(attn_scores, v, mask=mask)
297 -
298 - output = self.forward_output(output, residual, sz_b, len_q)
299 - return output, attn
300 -
301 -
302 -class ConformerConvolution(nn.Module):
303 - def __init__(self, d_model: int, kernel_size: int = 33):
304 - super().__init__()
305 - assert kernel_size % 2 == 1
306 - self.pre_layer_norm = nn.LayerNorm(d_model)
307 - self.pointwise_conv1 = nn.Conv1d(
308 - d_model, d_model * 4, kernel_size=1, bias=False
309 - )
310 - self.padding = (kernel_size - 1) // 2
311 - self.depthwise_conv = nn.Conv1d(
312 - d_model * 2,
313 - d_model * 2,
314 - kernel_size,
315 - stride=1,
316 - padding=self.padding,
317 - groups=d_model * 2,
318 - bias=False,
319 - )
320 - self.batch_norm = nn.LayerNorm(d_model * 2)
321 - self.swish = Swish()
322 - self.pointwise_conv2 = nn.Conv1d(
323 - d_model * 2, d_model, kernel_size=1, bias=False
324 - )
325 -
326 - def forward(
327 - self, x: torch.Tensor, mask: torch.Tensor | None = None
328 - ) -> torch.Tensor:
329 - residual = x
330 - out = self.pre_layer_norm(x)
331 - out = out.transpose(1, 2)
332 - if mask is not None:
333 - out.masked_fill_(mask.ne(1), 0.0)
334 - out = self.pointwise_conv1(out)
335 - out = F.glu(out, dim=1)
336 - out = self.depthwise_conv(out)
337 -
338 - out = out.transpose(1, 2)
339 - out = self.swish(self.batch_norm(out))
340 - out = out.transpose(1, 2)
341 -
342 - out = self.pointwise_conv2(out)
343 - if mask is not None:
344 - out.masked_fill_(mask.ne(1), 0.0)
345 - out = out.transpose(1, 2)
346 - return out + residual
347 -
348 -
349 -class RelPosEmbConformerBlock(nn.Module):
350 - def __init__(self, d_model, n_head, kernel_size=33):
351 - super().__init__()
352 - self.ffn1 = ConformerFeedForward(d_model)
353 - self.mhsa = RelPosMultiHeadAttention(n_head, d_model)
354 - self.conv = ConformerConvolution(d_model, kernel_size)
355 - self.ffn2 = ConformerFeedForward(d_model)
356 - self.layer_norm = nn.LayerNorm(d_model)
357 -
358 - def forward(
359 - self,
360 - x: torch.Tensor,
361 - pos_emb: torch.Tensor,
362 - slf_attn_mask: torch.Tensor | None = None,
363 - pad_mask: torch.Tensor | None = None,
364 - ) -> torch.Tensor:
365 - out = 0.5 * x + 0.5 * self.ffn1(x)
366 - out = self.mhsa(out, out, out, pos_emb, mask=slf_attn_mask)[0]
367 - out = self.conv(out, pad_mask)
368 - out = 0.5 * out + 0.5 * self.ffn2(out)
369 - out = self.layer_norm(out)
370 - return out
371 -
372 -
373 -class ConformerEncoder(nn.Module):
374 - def __init__(
375 - self,
376 - idim: int,
377 - n_layers_enc: int,
378 - n_head: int,
379 - d_model: int,
380 - kernel_size: int = 33,
381 - pe_maxlen: int = 5000,
382 - ):
383 - super().__init__()
384 - self.odim = d_model
385 -
386 - self.input_preprocessor = Conv2dSubsampling(idim, d_model)
387 - self.positional_encoding = RelPositionalEncoding(d_model)
388 -
389 - self.layer_stack = nn.ModuleList()
390 - for _ in range(n_layers_enc):
391 - block = RelPosEmbConformerBlock(d_model, n_head, kernel_size)
392 - self.layer_stack.append(block)
393 -
394 - def forward(
395 - self, padded_input: torch.Tensor, input_lengths: torch.Tensor, pad: bool = True
396 - ):
397 - if pad:
398 - padded_input = F.pad(
399 - padded_input,
400 - (0, 0, 0, self.input_preprocessor.context - 1),
401 - "constant",
402 - 0.0,
403 - )
404 - src_mask = self.padding_position_is_0(padded_input, input_lengths)
405 -
406 - embed_output, input_lengths, src_mask = self.input_preprocessor(
407 - padded_input, src_mask
408 - )
409 - enc_output = embed_output
410 -
411 - pos_emb = self.positional_encoding(embed_output)
412 -
413 - enc_outputs = []
414 - for enc_layer in self.layer_stack:
415 - enc_output = enc_layer(
416 - enc_output, pos_emb, slf_attn_mask=src_mask, pad_mask=src_mask
417 - )
418 - enc_outputs.append(enc_output)
419 -
420 - return enc_output, input_lengths, src_mask
421 -
422 - def padding_position_is_0(
423 - self, padded_input: torch.Tensor, input_lengths: torch.Tensor
424 - ) -> torch.Tensor:
425 - N, T = padded_input.size()[:2]
426 - mask = torch.ones((N, T)).to(padded_input.device)
427 - for i in range(N):
428 - mask[i, input_lengths[i] :] = 0
429 - mask = mask.unsqueeze(dim=1)
430 - return mask.to(torch.uint8)
431 -
432 -
87 class FireRedASR2Adapter(nn.Module):
88 def __init__(self, encoder_dim: int, llm_dim: int, downsample_rate: int = 2):
89 super().__init__()
vllm/model_executor/models/fireredlid.py new
+792
@@ -0,0 +1,792 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +FireRedLID – Language Identification model adapted for vLLM.
5 +
6 +Architecture: ConformerEncoder + TransformerDecoder (6-layer cross-attn)
7 +Vocabulary: 120 LID tokens (dict.txt)
8 +Output: Up to 2 tokens (e.g. "en", "zh mandarin")
9 +
10 +This implementation follows the Whisper-style encoder-decoder pattern:
11 + • Encoder processes audio features (Fbank + CMVN via FeatureExtractor)
12 + • Decoder performs single-step autoregressive forward
13 + • vLLM's generation loop handles beam search / sampling
14 +"""
15 +
16 +from __future__ import annotations
17 +
18 +from collections.abc import Iterable, Mapping, Sequence
19 +from typing import Annotated, Literal
20 +
21 +import numpy as np
22 +import torch
23 +from torch import nn
24 +from transformers import BatchFeature
25 +
26 +from vllm.config import ModelConfig, VllmConfig
27 +from vllm.config.multimodal import BaseDummyOptions
28 +from vllm.config.speech_to_text import SpeechToTextConfig
29 +from vllm.distributed import get_tensor_model_parallel_world_size
30 +from vllm.inputs import MultiModalDataDict, PromptType
31 +from vllm.logger import init_logger
32 +from vllm.model_executor.layers.attention import Attention, CrossAttention
33 +from vllm.model_executor.layers.linear import (
34 + ColumnParallelLinear,
35 + ReplicatedLinear,
36 + RowParallelLinear,
37 +)
38 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
39 +from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
40 +from vllm.multimodal import MULTIMODAL_REGISTRY
41 +from vllm.multimodal.inputs import (
42 + MultiModalFieldConfig,
43 + MultiModalKwargsItems,
44 +)
45 +from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser
46 +from vllm.multimodal.processing import (
47 + BaseDummyInputsBuilder,
48 + BaseProcessingInfo,
49 + EncDecMultiModalProcessor,
50 + PromptReplacement,
51 + PromptUpdate,
52 +)
53 +from vllm.transformers_utils.processor import cached_processor_from_config
54 +from vllm.utils.tensor_schema import TensorSchema, TensorShape
55 +
56 +from .conformer_encoder import ConformerEncoder
57 +from .interfaces import (
58 + MultiModalEmbeddings,
59 + SupportsMultiModal,
60 + SupportsTranscription,
61 +)
62 +from .utils import (
63 + AutoWeightsLoader,
64 + WeightsMapper,
65 + maybe_prefix,
66 +)
67 +from .whisper_utils import ISO639_1_SUPPORTED_LANGS
68 +
69 +logger = init_logger(__name__)
70 +
71 +
72 +class FireRedLIDAudioInputs(TensorSchema):
73 + """
74 + Dimensions:
75 + - b: Batch size
76 + - t: Time frames (variable across utterances)
77 + - nmb: Number of mel bins (80)
78 + """
79 +
80 + input_features: Annotated[
81 + list[torch.Tensor] | None,
82 + TensorShape("b", "t", "nmb", dynamic_dims={"t"}),
83 + ]
84 + speech_lengths: Annotated[
85 + list[torch.Tensor] | None,
86 + TensorShape("b"),
87 + ]
88 + fake_token_lengths: Annotated[
89 + list[torch.Tensor] | None,
90 + TensorShape("b"),
91 + ]
92 +
93 +
94 +FireRedLIDEncoder = ConformerEncoder
95 +
96 +
97 +class FireRedLIDPositionalEmbedding(nn.Module):
98 + """Absolute sinusoidal positional embedding indexed by `positions`."""
99 +
100 + def __init__(self, d_model: int, max_len: int = 5000):
101 + super().__init__()
102 + assert d_model % 2 == 0
103 + pe = torch.zeros(max_len, d_model, requires_grad=False)
104 + position = torch.arange(0, max_len).unsqueeze(1).float()
105 + div_term = torch.exp(
106 + torch.arange(0, d_model, 2).float()
107 + * -(torch.log(torch.tensor(10000.0)).item() / d_model)
108 + )
109 + pe[:, 0::2] = torch.sin(position * div_term)
110 + pe[:, 1::2] = torch.cos(position * div_term)
111 + self.register_buffer("pe", pe, persistent=False)
112 +
113 + def forward(self, position_ids: torch.Tensor) -> torch.Tensor:
114 + return self.pe[position_ids]
115 +
116 +
117 +class FireRedLIDAttention(nn.Module):
118 + """Base attention with shared QKV/FC projections for the LID decoder."""
119 +
120 + def __init__(
121 + self,
122 + d_model: int,
123 + n_head: int,
124 + *,
125 + vllm_config: VllmConfig,
126 + prefix: str = "",
127 + ):
128 + super().__init__()
129 + tp_size = get_tensor_model_parallel_world_size()
130 + assert n_head % tp_size == 0
131 + self.total_num_heads = n_head
132 + self.num_heads = n_head // tp_size
133 + self.num_kv_heads = max(1, n_head // tp_size)
134 + self.head_dim = d_model // n_head
135 + self.scaling = self.head_dim**-0.5
136 +
137 + cache_config = vllm_config.cache_config
138 + quant_config = vllm_config.quant_config
139 +
140 + self.w_qs = ColumnParallelLinear(
141 + d_model,
142 + d_model,
143 + bias=True,
144 + quant_config=quant_config,
145 + prefix=f"{prefix}.w_qs",
146 + )
147 + self.w_ks = ColumnParallelLinear(
148 + d_model,
149 + d_model,
150 + bias=False,
151 + quant_config=quant_config,
152 + prefix=f"{prefix}.w_ks",
153 + )
154 + self.w_vs = ColumnParallelLinear(
155 + d_model,
156 + d_model,
157 + bias=True,
158 + quant_config=quant_config,
159 + prefix=f"{prefix}.w_vs",
160 + )
161 + self.fc = RowParallelLinear(
162 + d_model,
163 + d_model,
164 + bias=True,
165 + quant_config=quant_config,
166 + prefix=f"{prefix}.fc",
167 + )
168 + self._init_attn(cache_config, quant_config, prefix)
169 +
170 + def _init_attn(self, cache_config, quant_config, prefix: str) -> None:
171 + raise NotImplementedError
172 +
173 +
174 +class FireRedLIDSelfAttention(FireRedLIDAttention):
175 + def _init_attn(self, cache_config, quant_config, prefix: str) -> None:
176 + self.attn = Attention(
177 + self.num_heads,
178 + self.head_dim,
179 + self.scaling,
180 + num_kv_heads=self.num_kv_heads,
181 + cache_config=cache_config,
182 + quant_config=quant_config,
183 + prefix=f"{prefix}.attn",
184 + )
185 +
186 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
187 + q, _ = self.w_qs(hidden_states)
188 + k, _ = self.w_ks(hidden_states)
189 + v, _ = self.w_vs(hidden_states)
190 + attn_output = self.attn(q, k, v)
191 + output, _ = self.fc(attn_output)
192 + return output
193 +
194 +
195 +class FireRedLIDCrossAttention(FireRedLIDAttention):
196 + def _init_attn(self, cache_config, quant_config, prefix: str) -> None:
197 + self.attn = CrossAttention(
198 + self.num_heads,
199 + self.head_dim,
200 + self.scaling,
201 + num_kv_heads=self.num_kv_heads,
202 + cache_config=cache_config,
203 + quant_config=quant_config,
204 + prefix=f"{prefix}.attn",
205 + )
206 +
207 + def forward(
208 + self,
209 + hidden_states: torch.Tensor,
210 + encoder_hidden_states: torch.Tensor | None,
211 + ) -> torch.Tensor:
212 + q, _ = self.w_qs(hidden_states)
213 + if encoder_hidden_states is not None:
214 + k, _ = self.w_ks(encoder_hidden_states)
215 + v, _ = self.w_vs(encoder_hidden_states)
216 + else:
217 + k = v = None
218 +
219 + attn_output = self.attn(q, k, v)
220 + output, _ = self.fc(attn_output)
221 + return output
222 +
223 +
224 +class FireRedLIDFFN(nn.Module):
225 + def __init__(self, d_model: int, d_ff: int):
226 + super().__init__()
227 + self.w_1 = ReplicatedLinear(d_model, d_ff, bias=True)
228 + self.act = nn.GELU()
229 + self.w_2 = ReplicatedLinear(d_ff, d_model, bias=True)
230 +
231 + def forward(self, x: torch.Tensor) -> torch.Tensor:
232 + x, _ = self.w_1(x)
233 + x = self.act(x)
234 + x, _ = self.w_2(x)
235 + return x
236 +
237 +
238 +class FireRedLIDDecoderLayer(nn.Module):
239 + """vLLM-native decoder layer while preserving FireRedLID parameter names."""
240 +
241 + def __init__(
242 + self,
243 + d_model: int,
244 + n_head: int,
245 + *,
246 + vllm_config: VllmConfig,
247 + prefix: str = "",
248 + ):
249 + super().__init__()
250 + self.self_attn_norm = nn.LayerNorm(d_model)
251 + self.self_attn = FireRedLIDSelfAttention(
252 + d_model,
253 + n_head,
254 + vllm_config=vllm_config,
255 + prefix=f"{prefix}.self_attn",
256 + )
257 +
258 + self.cross_attn_norm = nn.LayerNorm(d_model)
259 + self.cross_attn = FireRedLIDCrossAttention(
260 + d_model,
261 + n_head,
262 + vllm_config=vllm_config,
263 + prefix=f"{prefix}.cross_attn",
264 + )
265 +
266 + self.mlp_norm = nn.LayerNorm(d_model)
267 + self.mlp = FireRedLIDFFN(d_model, d_model * 4)
268 +
269 + def forward(
270 + self,
271 + hidden_states: torch.Tensor,
272 + encoder_hidden_states: torch.Tensor | None,
273 + ) -> torch.Tensor:
274 + residual = hidden_states
275 + hidden_states = self.self_attn_norm(hidden_states)
276 + hidden_states = self.self_attn(hidden_states)
277 + hidden_states = residual + hidden_states
278 +
279 + residual = hidden_states
280 + hidden_states = self.cross_attn_norm(hidden_states)
281 + hidden_states = self.cross_attn(hidden_states, encoder_hidden_states)
282 + hidden_states = residual + hidden_states
283 +
284 + residual = hidden_states
285 + hidden_states = self.mlp_norm(hidden_states)
286 + hidden_states = residual + self.mlp(hidden_states)
287 +
288 + return hidden_states
289 +
290 +
291 +class FireRedLIDDecoder(nn.Module):
292 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
293 + super().__init__()
294 + config = vllm_config.model_config.hf_config
295 + self.pad_id = getattr(config, "pad_token_id", 2)
296 + self.n_layers = getattr(config, "n_layers_lid_dec", 6)
297 + self.d_model = getattr(config, "d_model", 1280)
298 + self.scale = self.d_model**0.5
299 +
300 + self.tgt_word_emb = nn.Embedding(
301 + getattr(config, "vocab_size", 120),
302 + self.d_model,
303 + padding_idx=self.pad_id,
304 + )
305 + self.positional_encoding = FireRedLIDPositionalEmbedding(
306 + self.d_model,
307 + max_len=getattr(config, "pe_maxlen", 5000),
308 + )
309 +
310 + self.layer_stack = nn.ModuleList(
311 + [
312 + FireRedLIDDecoderLayer(
313 + self.d_model,
314 + getattr(config, "n_head", 20),
315 + vllm_config=vllm_config,
316 + prefix=f"{prefix}.layer_stack.{idx}",
317 + )
318 + for idx in range(self.n_layers)
319 + ]
320 + )
321 + self.layer_norm_out = nn.LayerNorm(self.d_model)
322 +
323 + def forward(
324 + self,
325 + input_ids: torch.Tensor,
326 + positions: torch.Tensor,
327 + encoder_hidden_states: torch.Tensor | None,
328 + ) -> torch.Tensor:
329 + hidden_states = self.tgt_word_emb(input_ids) * self.scale
330 + hidden_states = hidden_states + self.positional_encoding(positions)
331 +
332 + for layer in self.layer_stack:
333 + hidden_states = layer(hidden_states, encoder_hidden_states)
334 +
335 + hidden_states = self.layer_norm_out(hidden_states)
336 + return hidden_states
337 +
338 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
339 + return self.tgt_word_emb(input_ids)
340 +
341 +
342 +class FireRedLIDModel(nn.Module):
343 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
344 + super().__init__()
345 + config = vllm_config.model_config.hf_config
346 +
347 + self.encoder = FireRedLIDEncoder(
348 + idim=getattr(config, "idim", 80),
349 + n_layers_enc=getattr(config, "n_layers_enc", 16),
350 + n_head=getattr(config, "n_head", 20),
351 + d_model=getattr(config, "d_model", 1280),
352 + kernel_size=getattr(config, "kernel_size", 33),
353 + pe_maxlen=getattr(config, "pe_maxlen", 5000),
354 + )
355 +
356 + self.decoder = FireRedLIDDecoder(
357 + vllm_config=vllm_config,
358 + prefix=maybe_prefix(prefix, "decoder"),
359 + )
360 +
361 + def forward(
362 + self,
363 + input_ids: torch.Tensor,
364 + positions: torch.Tensor,
365 + encoder_outputs: list[torch.Tensor] | None = None,
366 + ) -> torch.Tensor:
367 + enc_states = (
368 + torch.cat(encoder_outputs, dim=0)
369 + if encoder_outputs and len(encoder_outputs) > 0
370 + else None
371 + )
372 + decoder_outputs = self.decoder(
373 + input_ids=input_ids,
374 + positions=positions,
375 + encoder_hidden_states=enc_states,
376 + )
377 + return decoder_outputs
378 +
379 + def get_encoder_outputs(
380 + self,
381 + speech: torch.Tensor | list[torch.Tensor],
382 + speech_lengths: torch.Tensor | list[torch.Tensor],
383 + ) -> tuple[torch.Tensor, torch.Tensor]:
384 + """Run the encoder and return padded outputs plus true sequence lengths."""
385 + enc_output, enc_lengths, _ = self.encoder(speech, speech_lengths)
386 + return enc_output, enc_lengths
387 +
388 +
389 +class FireRedLIDProcessingInfo(BaseProcessingInfo):
390 + def get_hf_config(self):
391 + return self.ctx.get_hf_config()
392 +
393 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:
394 + return {"audio": 1}
395 +
396 + def get_feature_extractor(self, **kwargs):
397 + hf_processor = self.get_hf_processor(**kwargs)
398 + feature_extractor = hf_processor.feature_extractor
399 + return feature_extractor
400 +
401 + def get_data_parser(self) -> MultiModalDataParser:
402 + feature_extractor = self.get_feature_extractor()
403 + return MultiModalDataParser(
404 + target_sr=feature_extractor.sampling_rate,
405 + target_channels=1,
406 + )
407 +
408 + @property
409 + def skip_prompt_length_check(self) -> bool:
410 + return True
411 +
412 + def get_num_audio_tokens(self) -> int:
413 + # For encoder profiling – return a reasonable dummy length.
414 + # This doesn't affect actual inference since encoder processes
415 + # variable-length features.
416 + return 1
417 +
418 +
419 +class FireRedLIDDummyInputsBuilder(BaseDummyInputsBuilder[FireRedLIDProcessingInfo]):
420 + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
421 + return "<sos>"
422 +
423 + def get_dummy_mm_data(
424 + self,
425 + seq_len: int,
426 + mm_counts: Mapping[str, int],
427 + mm_options: Mapping[str, BaseDummyOptions],
428 + ) -> MultiModalDataDict:
429 + feature_extractor = self.info.get_feature_extractor()
430 + sampling_rate = feature_extractor.sampling_rate
431 + audio_len = feature_extractor.chunk_length * sampling_rate
432 + num_audios = mm_counts.get("audio", 0)
433 + audio_overrides = mm_options.get("audio")
434 + return {
435 + "audio": self._get_dummy_audios(
436 + length=audio_len,
437 + num_audios=num_audios,
438 + overrides=audio_overrides,
439 + )
440 + }
441 +
442 +
443 +class FireRedLIDMultiModalProcessor(
444 + EncDecMultiModalProcessor[FireRedLIDProcessingInfo]
445 +):
446 + def create_encoder_prompt(
447 + self,
448 + prompt: str | list[int],
449 + mm_items: MultiModalDataItems,
450 + ) -> str | list[int]:
451 + # Dummy encoder prompt for profiling (encoder only processes audio).
452 + return [0]
453 +
454 + def _call_hf_processor(
455 + self,
456 + prompt: str,
457 + mm_data: Mapping[str, object],
458 + mm_kwargs: Mapping[str, object],
459 + tok_kwargs: Mapping[str, object],
460 + ) -> BatchFeature:
461 + if mm_data:
462 + feature_extractor = self.info.get_feature_extractor(**mm_kwargs)
463 + mm_data = dict(audio=mm_data.pop("audios"))
464 + mm_kwargs = dict(
465 + **mm_kwargs,
466 + sampling_rate=feature_extractor.sampling_rate,
467 + )
468 + processed_outputs = super()._call_hf_processor(
469 + prompt=prompt,
470 + mm_data=mm_data,
471 + mm_kwargs=mm_kwargs,
472 + tok_kwargs=tok_kwargs,
473 + )
474 + if "labels" in processed_outputs:
475 + processed_outputs["input_ids"] = processed_outputs.pop("labels")
476 + return processed_outputs
477 +
478 + def _get_mm_fields_config(
479 + self,
480 + hf_inputs: BatchFeature,
481 + hf_processor_mm_kwargs: Mapping[str, object],
482 + ) -> Mapping[str, MultiModalFieldConfig]:
483 + return dict(
484 + input_features=MultiModalFieldConfig.batched("audio"),
485 + speech_lengths=MultiModalFieldConfig.batched("audio"),
486 + fake_token_lengths=MultiModalFieldConfig.batched("audio"),
487 + )
488 +
489 + def _get_prompt_updates(
490 + self,
491 + mm_items: MultiModalDataItems,
492 + hf_processor_mm_kwargs: Mapping[str, object],
493 + out_mm_kwargs: MultiModalKwargsItems,
494 + ) -> Sequence[PromptUpdate]:
495 + out_mm_data = out_mm_kwargs.get_data()
496 + fake_token_lengths = out_mm_data.get("fake_token_lengths")
497 +
498 + if fake_token_lengths is None:
499 + # Fallback to max encoder output length if not available
500 + audio_output_lengths = []
501 + else:
502 + assert isinstance(fake_token_lengths, torch.Tensor)
503 + audio_output_lengths = fake_token_lengths.tolist()
504 +
505 + def get_replacement(item_idx: int):
506 + if audio_output_lengths:
507 + num_tokens = int(audio_output_lengths[item_idx])
508 + else:
509 + num_tokens = self.info.get_num_audio_tokens()
510 + return [0] * num_tokens
511 +
512 + return [
513 + PromptReplacement(
514 + modality="audio",
515 + target=[0],
516 + replacement=get_replacement,
517 + )
518 + ]
519 +
520 +
521 +# FireRedLID supports a wider set of languages than Whisper's shared list.
522 +# Only ISO 639-1 codes are listed; FireRedLID's dialect tokens (mandarin,
523 +# xinan, wu, …) are output tokens but not valid language *request* codes.
524 +_FIREREDLID_SUPPORTED_LANGUAGES: Mapping[str, str] = {
525 + **ISO639_1_SUPPORTED_LANGS,
526 + "am": "Amharic",
527 + "as": "Assamese",
528 + "ba": "Bashkir",
529 + "bn": "Bengali",
530 + "bo": "Tibetan",
531 + "br": "Breton",
532 + "eu": "Basque",
533 + "fo": "Faroese",
534 + "gu": "Gujarati",
535 + "ha": "Hausa",
536 + "haw": "Hawaiian",
537 + "ht": "Haitian Creole",
538 + "jw": "Javanese",
539 + "ka": "Georgian",
540 + "km": "Khmer",
541 + "la": "Latin",
542 + "lb": "Luxembourgish",
543 + "ln": "Lingala",
544 + "lo": "Lao",
545 + "mg": "Malagasy",
546 + "ml": "Malayalam",
547 + "mn": "Mongolian",
548 + "mt": "Maltese",
549 + "my": "Myanmar",
550 + "nn": "Nynorsk",
551 + "oc": "Occitan",
552 + "pa": "Panjabi",
553 + "ps": "Pashto",
554 + "sa": "Sanskrit",
555 + "sd": "Sindhi",
556 + "si": "Sinhala",
557 + "sn": "Shona",
558 + "so": "Somali",
559 + "sq": "Albanian",
560 + "su": "Sundanese",
561 + "te": "Telugu",
562 + "tg": "Tajik",
563 + "tk": "Turkmen",
564 + "tt": "Tatar",
565 + "uz": "Uzbek",
566 + "yi": "Yiddish",
567 + "yo": "Yoruba",
568 + "yue": "Cantonese",
569 +}
570 +
571 +
572 +@MULTIMODAL_REGISTRY.register_processor(
573 + FireRedLIDMultiModalProcessor,
574 + info=FireRedLIDProcessingInfo,
575 + dummy_inputs=FireRedLIDDummyInputsBuilder,
576 +)
577 +class FireRedLIDForConditionalGeneration(
578 + nn.Module, SupportsTranscription, SupportsMultiModal
579 +):
580 + # -- SupportsTranscription protocol attributes --
581 + supports_transcription_only = True
582 + supported_languages = _FIREREDLID_SUPPORTED_LANGUAGES
583 +
584 + hf_to_vllm_mapper = WeightsMapper(
585 + orig_to_new_substr={
586 + "encoder.": "model.encoder.",
587 + "lid_decoder.": "model.decoder.",
588 + # Encoder FFN: nn.Sequential indices → named children
589 + "net.0": "pre_layer_norm",
590 + "net.1": "linear_expand",
591 + "net.4": "linear_project",
592 + }
593 + )
594 +
595 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
596 + super().__init__()
597 + config = vllm_config.model_config.hf_config
598 + self.config = config
599 + self.dtype = vllm_config.model_config.dtype
600 +
601 + with self._mark_composite_model(
602 + vllm_config,
603 + language_targets=FireRedLIDDecoder,
604 + tower_targets={"audio": FireRedLIDEncoder},
605 + ):
606 + self.model = FireRedLIDModel(
607 + vllm_config=vllm_config,
608 + prefix=maybe_prefix(prefix, "model"),
609 + )
610 +
611 + self.proj_out = ParallelLMHead(
612 + getattr(config, "vocab_size", 120),
613 + getattr(config, "d_model", 1280),
614 + quant_config=vllm_config.quant_config,
615 + prefix=maybe_prefix(prefix, "proj_out"),
616 + )
617 + self.proj_out = self.proj_out.tie_weights(self.model.decoder.tgt_word_emb)
618 +
619 + logit_scale = getattr(config, "logit_scale", 1.0)
620 + self.logits_processor = LogitsProcessor(
621 + getattr(config, "vocab_size", 120),
622 + scale=logit_scale,
623 + )
624 +
625 + def forward(
626 + self,
627 + input_ids: torch.Tensor,
628 + positions: torch.Tensor,
629 + encoder_outputs: list[torch.Tensor] | None = None,
630 + **kwargs,
631 + ) -> torch.Tensor:
632 + if encoder_outputs is None:
633 + encoder_outputs = []
634 + decoder_outputs = self.model(
635 + input_ids=input_ids,
636 + positions=positions,
637 + encoder_outputs=encoder_outputs,
638 + )
639 + return decoder_outputs
640 +
641 + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
642 + """Run encoder on audio features and return per-item embeddings."""
643 + audio_input = self._parse_and_validate_audio_input(**kwargs)
644 +
645 + speech = audio_input["input_features"]
646 + speech_lengths = audio_input["speech_lengths"]
647 + if speech is None or speech_lengths is None:
648 + return []
649 +
650 + # When audio items have different time lengths, vLLM's
651 + # MultiModalBatchedField._reduce_data returns a plain
652 + # list[Tensor] instead of a stacked Tensor. The encoder
653 + # expects a padded [B, Tmax, feat_dim] Tensor, so we
654 + # normalise both speech and speech_lengths here.
655 + if isinstance(speech, (list, tuple)):
656 + # Each element: [Ti, feat_dim] (or [1, Ti, feat_dim])
657 + tensors = [
658 + s.squeeze(0) if s.dim() == 3 and s.size(0) == 1 else s for s in speech
659 + ]
660 + device = tensors[0].device
661 + dtype = tensors[0].dtype
662 + feat_dim = tensors[0].shape[-1]
663 + lengths = torch.tensor(
664 + [t.size(0) for t in tensors],
665 + device=device,
666 + dtype=torch.int32,
667 + )
668 + t_max = int(lengths.max().item())
669 + # Pre-allocate zero-padded batch tensor
670 + speech = torch.zeros(
671 + (len(tensors), t_max, feat_dim),
672 + device=device,
673 + dtype=dtype,
674 + )
675 + for i, t in enumerate(tensors):
676 + speech[i, : t.size(0)] = t
677 + speech_lengths = lengths
678 + else:
679 + # Already a batched Tensor [B, T, feat_dim]
680 + if speech.dim() == 2:
681 + speech = speech.unsqueeze(0)
682 +
683 + speech_lengths = torch.as_tensor(
684 + speech_lengths, dtype=torch.int32, device=speech.device
685 + )
686 +
687 + enc_output, enc_lengths = self.model.get_encoder_outputs(
688 + speech=speech,
689 + speech_lengths=speech_lengths,
690 + )
691 +
692 + # vLLM expects one 2D tensor per multimodal item. Slice each batch entry
693 + # by the true encoder length so cross-attention never sees padded frames.
694 + return tuple(
695 + enc_output[i, : max(0, int(enc_lengths[i].item()))]
696 + for i in range(enc_output.size(0))
697 + )
698 +
699 + def embed_input_ids(
700 + self,
701 + input_ids: torch.Tensor,
702 + multimodal_embeddings: MultiModalEmbeddings | None = None,
703 + *,
704 + is_multimodal: torch.Tensor | None = None,
705 + ) -> torch.Tensor:
706 + return self.model.decoder.embed_input_ids(input_ids)
707 +
708 + def _parse_and_validate_audio_input(
709 + self, **kwargs: object
710 + ) -> FireRedLIDAudioInputs:
711 + input_features = kwargs.pop("input_features", None)
712 + speech_lengths = kwargs.pop("speech_lengths", None)
713 + fake_token_lengths = kwargs.pop("fake_token_lengths", None)
714 + return FireRedLIDAudioInputs(
715 + input_features=input_features,
716 + speech_lengths=speech_lengths,
717 + fake_token_lengths=fake_token_lengths,
718 + )
719 +
720 + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
721 + logits = self.logits_processor(self.proj_out, hidden_states)
722 + return logits
723 +
724 + @classmethod
725 + def validate_language(cls, language: str | None) -> str | None:
726 + # FireRedLID is a language *identification* model – the caller does
727 + # not need to specify a language up-front. Accept None silently.
728 + if language is None:
729 + return None
730 + return super().validate_language(language)
731 +
732 + @classmethod
733 + def get_generation_prompt(
734 + cls,
735 + audio: np.ndarray,
736 + stt_config: SpeechToTextConfig,
737 + model_config: ModelConfig,
738 + language: str | None,
739 + task_type: Literal["transcribe", "translate"],
740 + request_prompt: str,
741 + to_language: str | None,
742 + ) -> PromptType:
743 + """Build the prompt for the FireRedLID encoder-decoder model.
744 +
745 + The decoder receives a single <sos> token; the encoder processes
746 + the raw audio waveform via the multimodal pipeline.
747 + """
748 + prompt: PromptType = {
749 + "encoder_prompt": {
750 + "prompt": "",
751 + "multi_modal_data": {
752 + "audio": (audio, int(stt_config.sample_rate)),
753 + },
754 + },
755 + "decoder_prompt": {
756 + "prompt": "<sos>",
757 + },
758 + }
759 + return prompt
760 +
761 + @classmethod
762 + def get_speech_to_text_config(
763 + cls,
764 + model_config: ModelConfig,
765 + task_type: Literal["transcribe", "translate"],
766 + ) -> SpeechToTextConfig:
767 + processor = cached_processor_from_config(model_config)
768 + return SpeechToTextConfig(
769 + max_audio_clip_s=processor.feature_extractor.chunk_length,
770 + sample_rate=processor.feature_extractor.sampling_rate,
771 + # LID output is at most 2 tokens – no chunking needed.
772 + min_energy_split_window_size=None,
773 + )
774 +
775 + @classmethod
776 + def post_process_output(cls, text: str) -> str:
777 + # Strip any leading/trailing whitespace from the raw LID output.
778 + return text.strip()
779 +
780 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
781 + loader = AutoWeightsLoader(
782 + self,
783 + skip_prefixes=[
784 + # Position encoding buffers are rebuilt at init
785 + "model.encoder.positional_encoding.pe",
786 + "model.decoder.positional_encoding.pe",
787 + # Tied output projection (shared with embedding)
788 + "model.decoder.tgt_word_prj.weight",
789 + "proj_out.",
790 + ],
791 + )
792 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/registry.py
+4
@@ -380,6 +380,10 @@ _MULTIMODAL_MODELS = {
380 "FireRedASR2ForConditionalGeneration",
381 ),
382 "FunASRForConditionalGeneration": ("funasr", "FunASRForConditionalGeneration"),
383 + "FireRedLIDForConditionalGeneration": (
384 + "fireredlid",
385 + "FireRedLIDForConditionalGeneration",
386 + ),
387 "FunAudioChatForConditionalGeneration": (
388 "funaudiochat",
389 "FunAudioChatForConditionalGeneration",
vllm/transformers_utils/config.py
+1
@@ -90,6 +90,7 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(
90 deepseek_vl_v2="DeepseekVLV2Config",
91 deepseek_v32="DeepseekV3Config",
92 flex_olmo="FlexOlmoConfig",
93 + fireredlid="FireRedLIDConfig",
94 funaudiochat="FunAudioChatConfig",
95 hunyuan_vl="HunYuanVLConfig",
96 isaac="IsaacConfig",
vllm/transformers_utils/configs/__init__.py
+2
@@ -28,6 +28,7 @@ _CLASS_TO_MODULE: dict[str, str] = {
28 "DeepseekVLV2Config": "vllm.transformers_utils.configs.deepseek_vl2",
29 "DotsOCRConfig": "vllm.transformers_utils.configs.dotsocr",
30 "EAGLEConfig": "vllm.transformers_utils.configs.eagle",
31 + "FireRedLIDConfig": "vllm.transformers_utils.configs.fireredlid",
32 "FlexOlmoConfig": "vllm.transformers_utils.configs.flex_olmo",
33 "FunAudioChatConfig": "vllm.transformers_utils.configs.funaudiochat",
34 "FunAudioChatAudioEncoderConfig": "vllm.transformers_utils.configs.funaudiochat",
@@ -88,6 +89,7 @@ __all__ = [
89 "DotsOCRConfig",
90 "EAGLEConfig",
91 "FlexOlmoConfig",
92 + "FireRedLIDConfig",
93 "FunAudioChatConfig",
94 "FunAudioChatAudioEncoderConfig",
95 "HunYuanVLConfig",
vllm/transformers_utils/configs/fireredlid.py new
+73
@@ -0,0 +1,73 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from __future__ import annotations
5 +
6 +import contextlib
7 +
8 +from transformers import AutoConfig
9 +from transformers.configuration_utils import PretrainedConfig
10 +
11 +
12 +class FireRedLIDConfig(PretrainedConfig):
13 + """Minimal config class for native vLLM FireRedLID support."""
14 +
15 + model_type = "fireredlid"
16 + keys_to_ignore_at_inference = ["past_key_values"]
17 +
18 + def __init__(
19 + self,
20 + vocab_size: int = 120,
21 + lid_odim: int = 120,
22 + idim: int = 80,
23 + d_model: int = 1280,
24 + n_head: int = 20,
25 + n_layers_enc: int = 16,
26 + n_layers_lid_dec: int = 6,
27 + kernel_size: int = 33,
28 + residual_dropout: float = 0.05,
29 + dropout_rate: float = 0.05,
30 + pe_maxlen: int = 5000,
31 + pad_token_id: int = 2,
32 + bos_token_id: int = 3,
33 + eos_token_id: int = 4,
34 + decoder_start_token_id: int = 3,
35 + tie_word_embeddings: bool = True,
36 + is_encoder_decoder: bool = True,
37 + architectures: list[str] | None = None,
38 + **kwargs,
39 + ):
40 + self.vocab_size = vocab_size
41 + self.lid_odim = lid_odim
42 + self.idim = idim
43 + self.d_model = d_model
44 + self.hidden_size = d_model
45 + self.n_head = n_head
46 + self.num_attention_heads = n_head
47 + self.n_layers_enc = n_layers_enc
48 + self.encoder_layers = n_layers_enc
49 + self.n_layers_lid_dec = n_layers_lid_dec
50 + self.decoder_layers = n_layers_lid_dec
51 + self.num_hidden_layers = n_layers_lid_dec
52 + self.kernel_size = kernel_size
53 + self.residual_dropout = residual_dropout
54 + self.dropout_rate = dropout_rate
55 + self.pe_maxlen = pe_maxlen
56 + self.tie_word_embeddings = tie_word_embeddings
57 + self.is_encoder_decoder = is_encoder_decoder
58 + self.architectures = architectures or ["FireRedLIDForConditionalGeneration"]
59 +
60 + super().__init__(
61 + pad_token_id=pad_token_id,
62 + bos_token_id=bos_token_id,
63 + eos_token_id=eos_token_id,
64 + decoder_start_token_id=decoder_start_token_id,
65 + tie_word_embeddings=tie_word_embeddings,
66 + is_encoder_decoder=is_encoder_decoder,
67 + architectures=self.architectures,
68 + **kwargs,
69 + )
70 +
71 +
72 +with contextlib.suppress(ValueError):
73 + AutoConfig.register(FireRedLIDConfig.model_type, FireRedLIDConfig)
vllm/transformers_utils/processors/__init__.py
+2
@@ -16,6 +16,7 @@ __all__ = [
16 "CohereASRProcessor",
17 "DeepseekVLV2Processor",
18 "FireRedASR2Processor",
19 + "FireRedLIDProcessor",
20 "FunASRProcessor",
21 "GLM4VProcessor",
22 "H2OVLProcessor",
@@ -44,6 +45,7 @@ _CLASS_TO_MODULE: dict[str, str] = {
45 "CohereASRProcessor": "vllm.transformers_utils.processors.cohere_asr",
46 "DeepseekVLV2Processor": "vllm.transformers_utils.processors.deepseek_vl2",
47 "FireRedASR2Processor": "vllm.transformers_utils.processors.fireredasr2",
48 + "FireRedLIDProcessor": "vllm.transformers_utils.processors.fireredlid",
49 "FunASRProcessor": "vllm.transformers_utils.processors.funasr",
50 "GLM4VProcessor": "vllm.transformers_utils.processors.glm4v",
51 "H2OVLProcessor": "vllm.transformers_utils.processors.h2ovl",
vllm/transformers_utils/processors/fireredlid.py new
+276
@@ -0,0 +1,276 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +FireRedLID feature extractor and processor.
5 +
6 +The FeatureExtractor handles:
7 + - Raw waveform → 80-dim log-mel filterbank (via kaldi_native_fbank)
8 + - CMVN normalization (means / inverse_std_variences from preprocessor_config)
9 + - Padding + length tracking
10 +
11 +The Processor wraps the FeatureExtractor and a tokenizer.
12 +"""
13 +
14 +from __future__ import annotations
15 +
16 +from typing import TYPE_CHECKING
17 +
18 +import numpy as np
19 +import torch
20 +import torch.nn.functional as F
21 +from transformers import (
22 + AutoFeatureExtractor,
23 + BatchFeature,
24 +)
25 +from transformers.feature_extraction_sequence_utils import SequenceFeatureExtractor
26 +from transformers.processing_utils import ProcessorMixin
27 +from transformers.utils import TensorType
28 +
29 +from vllm.logger import init_logger
30 +from vllm.utils.import_utils import LazyLoader
31 +
32 +if TYPE_CHECKING:
33 + import kaldi_native_fbank as knf
34 +else:
35 + knf = LazyLoader("knf", globals(), "kaldi_native_fbank")
36 +
37 +
38 +logger = init_logger(__name__)
39 +
40 +
41 +# ---------------------------------------------------------------------------
42 +# Helpers (shared with FireRedASR2 processor)
43 +# ---------------------------------------------------------------------------
44 +
45 +
46 +class CMVN:
47 + def __init__(self, dim, means, inverse_std_variences):
48 + self.dim = dim
49 + self.means = np.array(means)
50 + self.inverse_std_variences = np.array(inverse_std_variences)
51 +
52 + def __call__(self, x):
53 + assert x.shape[-1] == self.dim, "CMVN dim mismatch"
54 + out = x - self.means
55 + out = out * self.inverse_std_variences
56 + return out
57 +
58 +
59 +class KaldifeatFbank:
60 + def __init__(
61 + self,
62 + num_mel_bins: int = 80,
63 + frame_length: int = 25,
64 + frame_shift: int = 10,
65 + dither: float = 0.0,
66 + ):
67 + self.dither = dither
68 + opts = knf.FbankOptions()
69 + opts.frame_opts.dither = dither
70 + opts.mel_opts.num_bins = num_mel_bins
71 + opts.frame_opts.snip_edges = True
72 + opts.mel_opts.debug_mel = False
73 + self.opts = opts
74 +
75 + def __call__(self, sample_rate, wav_np, is_train=False):
76 + dither = self.dither if is_train else 0.0
77 + self.opts.frame_opts.dither = dither
78 + fbank = knf.OnlineFbank(self.opts)
79 + fbank.accept_waveform(sample_rate, wav_np.tolist())
80 + feat = []
81 + for i in range(fbank.num_frames_ready):
82 + feat.append(fbank.get_frame(i))
83 + if len(feat) == 0:
84 + return np.zeros((0, self.opts.mel_opts.num_bins))
85 + return np.vstack(feat)
86 +
87 +
88 +# ---------------------------------------------------------------------------
89 +# Feature Extractor
90 +# ---------------------------------------------------------------------------
91 +
92 +
93 +class FireRedLIDFeatureExtractor(SequenceFeatureExtractor):
94 + """
95 + Extracts 80-dim log-mel filterbank features from raw waveforms,
96 + applies CMVN, and returns padded feature tensors with lengths.
97 +
98 + Also computes ``fake_token_lengths`` — the actual encoder output
99 + length for each audio — so that vLLM can allocate the correct
100 + number of cross-attention KV cache slots.
101 + """
102 +
103 + model_input_names = ["input_features"]
104 +
105 + def __init__(
106 + self,
107 + feature_size=80,
108 + sampling_rate=16000,
109 + chunk_length=30,
110 + padding_value=0.0,
111 + return_attention_mask=False,
112 + dim=80,
113 + means=None,
114 + inverse_std_variences=None,
115 + num_mel_bins=80,
116 + frame_length=25,
117 + frame_shift=10,
118 + dither=0.0,
119 + left_context=3,
120 + right_context=3,
121 + **kwargs,
122 + ):
123 + super().__init__(
124 + feature_size=feature_size,
125 + sampling_rate=sampling_rate,
126 + padding_value=padding_value,
127 + return_attention_mask=return_attention_mask,
128 + **kwargs,
129 + )
130 + self.chunk_length = chunk_length
131 + self.dim = dim
132 + self.means = means
133 + self.inverse_std_variences = inverse_std_variences
134 + self.num_mel_bins = num_mel_bins
135 + self.frame_length = frame_length
136 + self.frame_shift = frame_shift
137 + self.dither = dither
138 + self.sampling_rate = sampling_rate
139 + self.context = left_context + 1 + right_context
140 +
141 + def __call__(
142 + self,
143 + raw_speech: np.ndarray | list[float] | list[np.ndarray] | list[list[float]],
144 + truncation: bool = True,
145 + pad_to_multiple_of: int | None = None,
146 + return_tensors: str | TensorType | None = None,
147 + return_attention_mask: bool | None = None,
148 + padding: str | None = "max_length",
149 + max_length: int | None = None,
150 + sampling_rate: int | None = None,
151 + do_normalize: bool | None = None,
152 + **kwargs,
153 + ) -> BatchFeature:
154 + if sampling_rate is not None and sampling_rate != self.sampling_rate:
155 + raise ValueError(
156 + f"FireRedLIDFeatureExtractor expects sampling_rate="
157 + f"{self.sampling_rate}, got {sampling_rate}."
158 + )
159 +
160 + # Initialize helpers
161 + cmvn = CMVN(self.dim, self.means, self.inverse_std_variences)
162 + fbank = KaldifeatFbank(
163 + num_mel_bins=self.num_mel_bins,
164 + frame_length=self.frame_length,
165 + frame_shift=self.frame_shift,
166 + dither=self.dither,
167 + )
168 +
169 + def padding_position_is_0(padded_input, input_lengths):
170 + N, T = padded_input.size()[:2]
171 + mask = torch.ones((N, T)).to(padded_input.device)
172 + for i in range(N):
173 + mask[i, input_lengths[i] :] = 0
174 + mask = mask.unsqueeze(dim=1)
175 + return mask.to(torch.uint8)
176 +
177 + feats = []
178 + speech_lengths = []
179 + fake_token_lengths = []
180 +
181 + for speech in raw_speech:
182 + # vLLM loads audio via librosa (float32 in [-1,1]),
183 + # but kaldi_native_fbank expects int16-scale values.
184 + speech_scaled = speech * 32768
185 + feat = fbank(self.sampling_rate, speech_scaled)
186 + feat = cmvn(feat)
187 + feat = torch.from_numpy(feat).float()
188 + length = feat.size(0)
189 + feats.append(feat)
190 + speech_lengths.append(length)
191 +
192 + # Compute the actual Conv2dSubsampling output length.
193 + # This mirrors the mask logic in Conv2dSubsampling.forward:
194 + # pad context frames, then mask[:, :, :-2:2][:, :, :-2:2].sum()
195 + padded_input = F.pad(feat, (0, 0, 0, self.context - 1), "constant", 0.0)
196 + src_mask = padding_position_is_0(
197 + padded_input[None, :, :],
198 + torch.tensor([length], dtype=torch.int32),
199 + )
200 + mask = src_mask[:, :, :-2:2][:, :, :-2:2]
201 + enc_len = mask[:, -1, :].sum(dim=-1)
202 + fake_token_len = torch.clamp(enc_len, min=1)
203 + fake_token_lengths.append(fake_token_len)
204 +
205 + if len(feats) == 0:
206 + return BatchFeature()
207 +
208 + # Pad to uniform length
209 + max_feat_len = max(f.size(0) for f in feats)
210 + padded = feats[0].new_zeros(len(feats), max_feat_len, feats[0].size(1))
211 + for i, feat in enumerate(feats):
212 + padded[i, : feat.size(0)] = feat
213 +
214 + result = BatchFeature({"input_features": padded})
215 +
216 + if return_tensors is not None:
217 + result = result.convert_to_tensors(return_tensors)
218 +
219 + result["speech_lengths"] = torch.tensor(speech_lengths, dtype=torch.long)
220 + result["fake_token_lengths"] = torch.concat(fake_token_lengths)
221 + return result
222 +
223 +
224 +# ---------------------------------------------------------------------------
225 +# Processor
226 +# ---------------------------------------------------------------------------
227 +
228 +
229 +class FireRedLIDProcessor(ProcessorMixin):
230 + """
231 + Wraps FireRedLIDFeatureExtractor + a tokenizer.
232 + """
233 +
234 + feature_extractor_class = "FireRedLIDFeatureExtractor"
235 + tokenizer_class = ("PreTrainedTokenizer", "PreTrainedTokenizerFast")
236 +
237 + def __init__(self, feature_extractor, tokenizer):
238 + super().__init__(feature_extractor, tokenizer)
239 + self.current_processor = self.feature_extractor
240 + self._in_target_context_manager = False
241 +
242 + def __call__(self, *args, **kwargs):
243 + if self._in_target_context_manager:
244 + return self.current_processor(*args, **kwargs)
245 +
246 + audio = kwargs.pop("audio", None)
247 + sampling_rate = kwargs.pop("sampling_rate", None)
248 + text = kwargs.pop("text", None)
249 + if len(args) > 0:
250 + audio = args[0]
251 + args = args[1:]
252 +
253 + if audio is not None:
254 + inputs = self.feature_extractor(
255 + audio, *args, sampling_rate=sampling_rate, **kwargs
256 + )
257 + else:
258 + inputs = BatchFeature()
259 +
260 + if text is not None:
261 + if isinstance(text, str):
262 + text = [text]
263 + encodings = self.tokenizer(text, **kwargs)
264 + if audio is not None:
265 + inputs["labels"] = encodings["input_ids"]
266 + else:
267 + return encodings
268 +
269 + return inputs
270 +
271 +
272 +# ---------------------------------------------------------------------------
273 +# Registration
274 +# ---------------------------------------------------------------------------
275 +
276 +AutoFeatureExtractor.register("FireRedLIDFeatureExtractor", FireRedLIDFeatureExtractor)