# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """ FireRedLID – Language Identification model adapted for vLLM. Architecture: ConformerEncoder + TransformerDecoder (6-layer cross-attn) Vocabulary: 120 LID tokens (dict.txt) Output: Up to 2 tokens (e.g. "en", "zh mandarin") This implementation follows the Whisper-style encoder-decoder pattern: • Encoder processes audio features (Fbank + CMVN via FeatureExtractor) • Decoder performs single-step autoregressive forward • vLLM's generation loop handles beam search / sampling """ from __future__ import annotations from collections.abc import Iterable, Mapping, Sequence from typing import Annotated, Literal import numpy as np import torch from torch import nn from transformers import BatchFeature from vllm.config import ModelConfig, VllmConfig from vllm.config.multimodal import BaseDummyOptions from vllm.config.speech_to_text import SpeechToTextConfig from vllm.distributed import get_tensor_model_parallel_world_size from vllm.inputs import MultiModalDataDict, PromptType from vllm.logger import init_logger from vllm.model_executor.layers.attention import Attention, CrossAttention from vllm.model_executor.layers.linear import ( ColumnParallelLinear, ReplicatedLinear, RowParallelLinear, ) from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import ( MultiModalFieldConfig, MultiModalKwargsItems, ) from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser from vllm.multimodal.processing import ( BaseDummyInputsBuilder, BaseProcessingInfo, EncDecMultiModalProcessor, PromptReplacement, PromptUpdate, ) from vllm.transformers_utils.processor import cached_processor_from_config from vllm.utils.tensor_schema import TensorSchema, TensorShape from .conformer_encoder import ConformerEncoder from .interfaces import ( MultiModalEmbeddings, SupportsMultiModal, SupportsTranscription, ) from .utils import ( AutoWeightsLoader, WeightsMapper, maybe_prefix, ) from .whisper_utils import ISO639_1_SUPPORTED_LANGS logger = init_logger(__name__) class FireRedLIDAudioInputs(TensorSchema): """ Dimensions: - b: Batch size - t: Time frames (variable across utterances) - nmb: Number of mel bins (80) """ input_features: Annotated[ list[torch.Tensor] | None, TensorShape("b", "t", "nmb", dynamic_dims={"t"}), ] speech_lengths: Annotated[ list[torch.Tensor] | None, TensorShape("b"), ] fake_token_lengths: Annotated[ list[torch.Tensor] | None, TensorShape("b"), ] FireRedLIDEncoder = ConformerEncoder class FireRedLIDPositionalEmbedding(nn.Module): """Absolute sinusoidal positional embedding indexed by `positions`.""" def __init__(self, d_model: int, max_len: int = 5000): super().__init__() assert d_model % 2 == 0 pe = torch.zeros(max_len, d_model, requires_grad=False) position = torch.arange(0, max_len).unsqueeze(1).float() div_term = torch.exp( torch.arange(0, d_model, 2).float() * -(torch.log(torch.tensor(10000.0)).item() / d_model) ) pe[:, 0::2] = torch.sin(position * div_term) pe[:, 1::2] = torch.cos(position * div_term) self.register_buffer("pe", pe, persistent=False) def forward(self, position_ids: torch.Tensor) -> torch.Tensor: return self.pe[position_ids] class FireRedLIDAttention(nn.Module): """Base attention with shared QKV/FC projections for the LID decoder.""" def __init__( self, d_model: int, n_head: int, *, vllm_config: VllmConfig, prefix: str = "", ): super().__init__() tp_size = get_tensor_model_parallel_world_size() assert n_head % tp_size == 0 self.total_num_heads = n_head self.num_heads = n_head // tp_size self.num_kv_heads = max(1, n_head // tp_size) self.head_dim = d_model // n_head self.scaling = self.head_dim**-0.5 cache_config = vllm_config.cache_config quant_config = vllm_config.quant_config self.w_qs = ColumnParallelLinear( d_model, d_model, bias=True, quant_config=quant_config, prefix=f"{prefix}.w_qs", ) self.w_ks = ColumnParallelLinear( d_model, d_model, bias=False, quant_config=quant_config, prefix=f"{prefix}.w_ks", ) self.w_vs = ColumnParallelLinear( d_model, d_model, bias=True, quant_config=quant_config, prefix=f"{prefix}.w_vs", ) self.fc = RowParallelLinear( d_model, d_model, bias=True, quant_config=quant_config, prefix=f"{prefix}.fc", ) self._init_attn(cache_config, quant_config, prefix) def _init_attn(self, cache_config, quant_config, prefix: str) -> None: raise NotImplementedError class FireRedLIDSelfAttention(FireRedLIDAttention): def _init_attn(self, cache_config, quant_config, prefix: str) -> None: self.attn = Attention( self.num_heads, self.head_dim, self.scaling, num_kv_heads=self.num_kv_heads, cache_config=cache_config, quant_config=quant_config, prefix=f"{prefix}.attn", ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: q, _ = self.w_qs(hidden_states) k, _ = self.w_ks(hidden_states) v, _ = self.w_vs(hidden_states) attn_output = self.attn(q, k, v) output, _ = self.fc(attn_output) return output class FireRedLIDCrossAttention(FireRedLIDAttention): def _init_attn(self, cache_config, quant_config, prefix: str) -> None: self.attn = CrossAttention( self.num_heads, self.head_dim, self.scaling, num_kv_heads=self.num_kv_heads, cache_config=cache_config, quant_config=quant_config, prefix=f"{prefix}.attn", ) def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | None, ) -> torch.Tensor: q, _ = self.w_qs(hidden_states) if encoder_hidden_states is not None: k, _ = self.w_ks(encoder_hidden_states) v, _ = self.w_vs(encoder_hidden_states) else: k = v = None attn_output = self.attn(q, k, v) output, _ = self.fc(attn_output) return output class FireRedLIDFFN(nn.Module): def __init__(self, d_model: int, d_ff: int): super().__init__() self.w_1 = ReplicatedLinear(d_model, d_ff, bias=True) self.act = nn.GELU() self.w_2 = ReplicatedLinear(d_ff, d_model, bias=True) def forward(self, x: torch.Tensor) -> torch.Tensor: x, _ = self.w_1(x) x = self.act(x) x, _ = self.w_2(x) return x class FireRedLIDDecoderLayer(nn.Module): """vLLM-native decoder layer while preserving FireRedLID parameter names.""" def __init__( self, d_model: int, n_head: int, *, vllm_config: VllmConfig, prefix: str = "", ): super().__init__() self.self_attn_norm = nn.LayerNorm(d_model) self.self_attn = FireRedLIDSelfAttention( d_model, n_head, vllm_config=vllm_config, prefix=f"{prefix}.self_attn", ) self.cross_attn_norm = nn.LayerNorm(d_model) self.cross_attn = FireRedLIDCrossAttention( d_model, n_head, vllm_config=vllm_config, prefix=f"{prefix}.cross_attn", ) self.mlp_norm = nn.LayerNorm(d_model) self.mlp = FireRedLIDFFN(d_model, d_model * 4) def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | None, ) -> torch.Tensor: residual = hidden_states hidden_states = self.self_attn_norm(hidden_states) hidden_states = self.self_attn(hidden_states) hidden_states = residual + hidden_states residual = hidden_states hidden_states = self.cross_attn_norm(hidden_states) hidden_states = self.cross_attn(hidden_states, encoder_hidden_states) hidden_states = residual + hidden_states residual = hidden_states hidden_states = self.mlp_norm(hidden_states) hidden_states = residual + self.mlp(hidden_states) return hidden_states class FireRedLIDDecoder(nn.Module): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() config = vllm_config.model_config.hf_config self.pad_id = getattr(config, "pad_token_id", 2) self.n_layers = getattr(config, "n_layers_lid_dec", 6) self.d_model = getattr(config, "d_model", 1280) self.scale = self.d_model**0.5 self.tgt_word_emb = nn.Embedding( getattr(config, "vocab_size", 120), self.d_model, padding_idx=self.pad_id, ) self.positional_encoding = FireRedLIDPositionalEmbedding( self.d_model, max_len=getattr(config, "pe_maxlen", 5000), ) self.layer_stack = nn.ModuleList( [ FireRedLIDDecoderLayer( self.d_model, getattr(config, "n_head", 20), vllm_config=vllm_config, prefix=f"{prefix}.layer_stack.{idx}", ) for idx in range(self.n_layers) ] ) self.layer_norm_out = nn.LayerNorm(self.d_model) def forward( self, input_ids: torch.Tensor, positions: torch.Tensor, encoder_hidden_states: torch.Tensor | None, ) -> torch.Tensor: hidden_states = self.tgt_word_emb(input_ids) * self.scale hidden_states = hidden_states + self.positional_encoding(positions) for layer in self.layer_stack: hidden_states = layer(hidden_states, encoder_hidden_states) hidden_states = self.layer_norm_out(hidden_states) return hidden_states def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.tgt_word_emb(input_ids) class FireRedLIDModel(nn.Module): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() config = vllm_config.model_config.hf_config self.encoder = FireRedLIDEncoder( idim=getattr(config, "idim", 80), n_layers_enc=getattr(config, "n_layers_enc", 16), n_head=getattr(config, "n_head", 20), d_model=getattr(config, "d_model", 1280), kernel_size=getattr(config, "kernel_size", 33), pe_maxlen=getattr(config, "pe_maxlen", 5000), ) self.decoder = FireRedLIDDecoder( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "decoder"), ) def forward( self, input_ids: torch.Tensor, positions: torch.Tensor, encoder_outputs: list[torch.Tensor] | None = None, ) -> torch.Tensor: enc_states = ( torch.cat(encoder_outputs, dim=0) if encoder_outputs and len(encoder_outputs) > 0 else None ) decoder_outputs = self.decoder( input_ids=input_ids, positions=positions, encoder_hidden_states=enc_states, ) return decoder_outputs def get_encoder_outputs( self, speech: torch.Tensor | list[torch.Tensor], speech_lengths: torch.Tensor | list[torch.Tensor], ) -> tuple[torch.Tensor, torch.Tensor]: """Run the encoder and return padded outputs plus true sequence lengths.""" enc_output, enc_lengths, _ = self.encoder(speech, speech_lengths) return enc_output, enc_lengths class FireRedLIDProcessingInfo(BaseProcessingInfo): def get_hf_config(self): return self.ctx.get_hf_config() def get_supported_mm_limits(self) -> Mapping[str, int | None]: return {"audio": 1} def get_feature_extractor(self, **kwargs): hf_processor = self.get_hf_processor(**kwargs) feature_extractor = hf_processor.feature_extractor return feature_extractor def get_data_parser(self) -> MultiModalDataParser: feature_extractor = self.get_feature_extractor() return MultiModalDataParser( target_sr=feature_extractor.sampling_rate, target_channels=1, ) @property def skip_prompt_length_check(self) -> bool: return True def get_num_audio_tokens(self) -> int: # For encoder profiling – return a reasonable dummy length. # This doesn't affect actual inference since encoder processes # variable-length features. return 1 class FireRedLIDDummyInputsBuilder(BaseDummyInputsBuilder[FireRedLIDProcessingInfo]): def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: return "" def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], mm_options: Mapping[str, BaseDummyOptions], ) -> MultiModalDataDict: feature_extractor = self.info.get_feature_extractor() sampling_rate = feature_extractor.sampling_rate audio_len = feature_extractor.chunk_length * sampling_rate num_audios = mm_counts.get("audio", 0) audio_overrides = mm_options.get("audio") return { "audio": self._get_dummy_audios( length=audio_len, num_audios=num_audios, overrides=audio_overrides, ) } class FireRedLIDMultiModalProcessor( EncDecMultiModalProcessor[FireRedLIDProcessingInfo] ): def create_encoder_prompt( self, prompt: str | list[int], mm_items: MultiModalDataItems, ) -> str | list[int]: # Dummy encoder prompt for profiling (encoder only processes audio). return [0] def _call_hf_processor( self, prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], tok_kwargs: Mapping[str, object], ) -> BatchFeature: if mm_data: feature_extractor = self.info.get_feature_extractor(**mm_kwargs) mm_data = dict(audio=mm_data.pop("audios")) mm_kwargs = dict( **mm_kwargs, sampling_rate=feature_extractor.sampling_rate, ) processed_outputs = super()._call_hf_processor( prompt=prompt, mm_data=mm_data, mm_kwargs=mm_kwargs, tok_kwargs=tok_kwargs, ) if "labels" in processed_outputs: processed_outputs["input_ids"] = processed_outputs.pop("labels") return processed_outputs def _get_mm_fields_config( self, hf_inputs: BatchFeature, hf_processor_mm_kwargs: Mapping[str, object], ) -> Mapping[str, MultiModalFieldConfig]: return dict( input_features=MultiModalFieldConfig.batched("audio"), speech_lengths=MultiModalFieldConfig.batched("audio"), fake_token_lengths=MultiModalFieldConfig.batched("audio"), ) def _get_prompt_updates( self, mm_items: MultiModalDataItems, hf_processor_mm_kwargs: Mapping[str, object], out_mm_kwargs: MultiModalKwargsItems, ) -> Sequence[PromptUpdate]: out_mm_data = out_mm_kwargs.get_data() fake_token_lengths = out_mm_data.get("fake_token_lengths") if fake_token_lengths is None: # Fallback to max encoder output length if not available audio_output_lengths = [] else: assert isinstance(fake_token_lengths, torch.Tensor) audio_output_lengths = fake_token_lengths.tolist() def get_replacement(item_idx: int): if audio_output_lengths: num_tokens = int(audio_output_lengths[item_idx]) else: num_tokens = self.info.get_num_audio_tokens() return [0] * num_tokens return [ PromptReplacement( modality="audio", target=[0], replacement=get_replacement, ) ] # FireRedLID supports a wider set of languages than Whisper's shared list. # Only ISO 639-1 codes are listed; FireRedLID's dialect tokens (mandarin, # xinan, wu, …) are output tokens but not valid language *request* codes. _FIREREDLID_SUPPORTED_LANGUAGES: Mapping[str, str] = { **ISO639_1_SUPPORTED_LANGS, "am": "Amharic", "as": "Assamese", "ba": "Bashkir", "bn": "Bengali", "bo": "Tibetan", "br": "Breton", "eu": "Basque", "fo": "Faroese", "gu": "Gujarati", "ha": "Hausa", "haw": "Hawaiian", "ht": "Haitian Creole", "jw": "Javanese", "ka": "Georgian", "km": "Khmer", "la": "Latin", "lb": "Luxembourgish", "ln": "Lingala", "lo": "Lao", "mg": "Malagasy", "ml": "Malayalam", "mn": "Mongolian", "mt": "Maltese", "my": "Myanmar", "nn": "Nynorsk", "oc": "Occitan", "pa": "Panjabi", "ps": "Pashto", "sa": "Sanskrit", "sd": "Sindhi", "si": "Sinhala", "sn": "Shona", "so": "Somali", "sq": "Albanian", "su": "Sundanese", "te": "Telugu", "tg": "Tajik", "tk": "Turkmen", "tt": "Tatar", "uz": "Uzbek", "yi": "Yiddish", "yo": "Yoruba", "yue": "Cantonese", } @MULTIMODAL_REGISTRY.register_processor( FireRedLIDMultiModalProcessor, info=FireRedLIDProcessingInfo, dummy_inputs=FireRedLIDDummyInputsBuilder, ) class FireRedLIDForConditionalGeneration( nn.Module, SupportsTranscription, SupportsMultiModal ): # -- SupportsTranscription protocol attributes -- supports_transcription_only = True supported_languages = _FIREREDLID_SUPPORTED_LANGUAGES hf_to_vllm_mapper = WeightsMapper( orig_to_new_substr={ "encoder.": "model.encoder.", "lid_decoder.": "model.decoder.", # Encoder FFN: nn.Sequential indices → named children "net.0": "pre_layer_norm", "net.1": "linear_expand", "net.4": "linear_project", } ) def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() config = vllm_config.model_config.hf_config self.config = config self.dtype = vllm_config.model_config.dtype with self._mark_composite_model( vllm_config, language_targets=FireRedLIDDecoder, tower_targets={"audio": FireRedLIDEncoder}, ): self.model = FireRedLIDModel( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model"), ) self.proj_out = ParallelLMHead( getattr(config, "vocab_size", 120), getattr(config, "d_model", 1280), quant_config=vllm_config.quant_config, prefix=maybe_prefix(prefix, "proj_out"), ) self.proj_out = self.proj_out.tie_weights(self.model.decoder.tgt_word_emb) logit_scale = getattr(config, "logit_scale", 1.0) self.logits_processor = LogitsProcessor( getattr(config, "vocab_size", 120), scale=logit_scale, ) def forward( self, input_ids: torch.Tensor, positions: torch.Tensor, encoder_outputs: list[torch.Tensor] | None = None, **kwargs, ) -> torch.Tensor: if encoder_outputs is None: encoder_outputs = [] decoder_outputs = self.model( input_ids=input_ids, positions=positions, encoder_outputs=encoder_outputs, ) return decoder_outputs def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: """Run encoder on audio features and return per-item embeddings.""" audio_input = self._parse_and_validate_audio_input(**kwargs) speech = audio_input["input_features"] speech_lengths = audio_input["speech_lengths"] if speech is None or speech_lengths is None: return [] # When audio items have different time lengths, vLLM's # MultiModalBatchedField._reduce_data returns a plain # list[Tensor] instead of a stacked Tensor. The encoder # expects a padded [B, Tmax, feat_dim] Tensor, so we # normalise both speech and speech_lengths here. if isinstance(speech, (list, tuple)): # Each element: [Ti, feat_dim] (or [1, Ti, feat_dim]) tensors = [ s.squeeze(0) if s.dim() == 3 and s.size(0) == 1 else s for s in speech ] device = tensors[0].device dtype = tensors[0].dtype feat_dim = tensors[0].shape[-1] lengths = torch.tensor( [t.size(0) for t in tensors], device=device, dtype=torch.int32, ) t_max = int(lengths.max().item()) # Pre-allocate zero-padded batch tensor speech = torch.zeros( (len(tensors), t_max, feat_dim), device=device, dtype=dtype, ) for i, t in enumerate(tensors): speech[i, : t.size(0)] = t speech_lengths = lengths else: # Already a batched Tensor [B, T, feat_dim] if speech.dim() == 2: speech = speech.unsqueeze(0) speech_lengths = torch.as_tensor( speech_lengths, dtype=torch.int32, device=speech.device ) enc_output, enc_lengths = self.model.get_encoder_outputs( speech=speech, speech_lengths=speech_lengths, ) # vLLM expects one 2D tensor per multimodal item. Slice each batch entry # by the true encoder length so cross-attention never sees padded frames. return tuple( enc_output[i, : max(0, int(enc_lengths[i].item()))] for i in range(enc_output.size(0)) ) def embed_input_ids( self, input_ids: torch.Tensor, multimodal_embeddings: MultiModalEmbeddings | None = None, *, is_multimodal: torch.Tensor | None = None, ) -> torch.Tensor: return self.model.decoder.embed_input_ids(input_ids) def _parse_and_validate_audio_input( self, **kwargs: object ) -> FireRedLIDAudioInputs: input_features = kwargs.pop("input_features", None) speech_lengths = kwargs.pop("speech_lengths", None) fake_token_lengths = kwargs.pop("fake_token_lengths", None) return FireRedLIDAudioInputs( input_features=input_features, speech_lengths=speech_lengths, fake_token_lengths=fake_token_lengths, ) def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: logits = self.logits_processor(self.proj_out, hidden_states) return logits @classmethod def validate_language(cls, language: str | None) -> str | None: # FireRedLID is a language *identification* model – the caller does # not need to specify a language up-front. Accept None silently. if language is None: return None return super().validate_language(language) @classmethod def get_generation_prompt( cls, audio: np.ndarray, stt_config: SpeechToTextConfig, model_config: ModelConfig, language: str | None, task_type: Literal["transcribe", "translate"], request_prompt: str, to_language: str | None, ) -> PromptType: """Build the prompt for the FireRedLID encoder-decoder model. The decoder receives a single token; the encoder processes the raw audio waveform via the multimodal pipeline. """ prompt: PromptType = { "encoder_prompt": { "prompt": "", "multi_modal_data": { "audio": (audio, int(stt_config.sample_rate)), }, }, "decoder_prompt": { "prompt": "", }, } return prompt @classmethod def get_speech_to_text_config( cls, model_config: ModelConfig, task_type: Literal["transcribe", "translate"], ) -> SpeechToTextConfig: processor = cached_processor_from_config(model_config) return SpeechToTextConfig( max_audio_clip_s=processor.feature_extractor.chunk_length, sample_rate=processor.feature_extractor.sampling_rate, # LID output is at most 2 tokens – no chunking needed. min_energy_split_window_size=None, ) @classmethod def post_process_output(cls, text: str) -> str: # Strip any leading/trailing whitespace from the raw LID output. return text.strip() def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader( self, skip_prefixes=[ # Position encoding buffers are rebuilt at init "model.encoder.positional_encoding.pe", "model.decoder.positional_encoding.pe", # Tied output projection (shared with embedding) "model.decoder.tgt_word_prj.weight", "proj_out.", ], ) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)