sgl-project--sglang/python/sglang/srt/models/voxtral.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# Adapted from:
# https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/models/voxtral.py
# https://huggingface.co/mistralai/Voxtral-Mini-3B-2507
#
# Copyright 2025 Mistral AI and the HuggingFace Inc. team.
# Licensed under the Apache License, Version 2.0.
"""Inference-only Voxtral (speech-to-text) model."""

import math
from typing import Any, Iterable, List, Optional, Tuple

import torch
import torch.nn as nn
from transformers import PretrainedConfig

from sglang.srt.layers.activation import get_act_fn
from sglang.srt.layers.linear import (
    ColumnParallelLinear,
    QKVParallelLinear,
    RowParallelLinear,
)
from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.managers.mm_utils import (
    MultiModalityDataPaddingPatternMultimodalTokens,
    general_mm_embed_routine,
)
from sglang.srt.managers.schedule_batch import (
    Modality,
    MultimodalDataItem,
    MultimodalInputs,
)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.models.llama import LlamaForCausalLM


class AudioLanguageAdapter(nn.Module):
    """MLP projector: Linear -> GELU -> Linear (no bias)."""

    def __init__(self, hidden_size: int, dim: int) -> None:
        super().__init__()
        self.w_in = nn.Linear(hidden_size, dim, bias=False)
        self.gelu = nn.GELU()
        self.w_out = nn.Linear(dim, dim, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.w_out(self.gelu(self.w_in(x)))


class VoxtralWhisperAttention(nn.Module):
    """Multi-headed self-attention using plain SDPA (no KV cache).

    Note: HF Voxtral has bias on q_proj, v_proj, out_proj but NOT on k_proj.
    We use QKVParallelLinear with bias=True and create a zero bias for k_proj
    during weight loading.
    """

    def __init__(
        self,
        embed_dim: int,
        num_heads: int,
        quant_config: Optional[QuantizationConfig] = None,
    ):
        super().__init__()
        self.head_dim = embed_dim // num_heads
        self.scaling = self.head_dim**-0.5

        self.qkv_proj = QKVParallelLinear(
            embed_dim, self.head_dim, num_heads, quant_config=quant_config
        )
        # After TP split, the local head count lives on the linear layer
        self.num_heads = self.qkv_proj.num_heads
        self.out_proj = RowParallelLinear(
            embed_dim, embed_dim, bias=True, quant_config=quant_config
        )

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        batch_size, seq_len, _ = hidden_states.shape
        qkv, _ = self.qkv_proj(hidden_states)
        q, k, v = qkv.chunk(3, dim=-1)
        q = q * self.scaling

        q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).permute(
            0, 2, 1, 3
        )
        k = k.view(batch_size, seq_len, self.num_heads, self.head_dim).permute(
            0, 2, 1, 3
        )
        v = v.view(batch_size, seq_len, self.num_heads, self.head_dim).permute(
            0, 2, 1, 3
        )

        attn_output = torch.nn.functional.scaled_dot_product_attention(
            q, k, v, scale=1.0
        )
        attn_output = attn_output.permute(0, 2, 1, 3).reshape(
            batch_size, seq_len, self.num_heads * self.head_dim
        )
        attn_output, _ = self.out_proj(attn_output)
        return attn_output


class VoxtralWhisperEncoderLayer(nn.Module):
    def __init__(
        self,
        config: PretrainedConfig,
        quant_config: Optional[QuantizationConfig] = None,
    ):
        super().__init__()
        embed_dim = config.d_model
        self.self_attn = VoxtralWhisperAttention(
            embed_dim=embed_dim,
            num_heads=config.encoder_attention_heads,
            quant_config=quant_config,
        )
        self.self_attn_layer_norm = nn.LayerNorm(embed_dim)
        self.activation_fn = get_act_fn(
            getattr(config, "activation_function", "gelu"),
            quant_config=quant_config,
        )
        self.fc1 = ColumnParallelLinear(embed_dim, config.encoder_ffn_dim)
        self.fc2 = RowParallelLinear(config.encoder_ffn_dim, embed_dim)
        self.final_layer_norm = nn.LayerNorm(embed_dim)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        residual = hidden_states
        hidden_states = self.self_attn_layer_norm(hidden_states)
        hidden_states = self.self_attn(hidden_states)
        hidden_states = residual + hidden_states

        residual = hidden_states
        hidden_states = self.final_layer_norm(hidden_states)
        hidden_states, _ = self.fc1(hidden_states)
        hidden_states = self.activation_fn(hidden_states)
        hidden_states, _ = self.fc2(hidden_states)
        hidden_states = residual + hidden_states

        if hidden_states.dtype == torch.float16:
            clamp_value = torch.finfo(hidden_states.dtype).max - 1000
            hidden_states = torch.clamp(
                hidden_states, min=-clamp_value, max=clamp_value
            )
        return hidden_states


class VoxtralWhisperEncoder(nn.Module):
    """Whisper encoder (Conv1d + positional embed + transformer + layer norm)."""

    def __init__(
        self,
        config: PretrainedConfig,
        quant_config: Optional[QuantizationConfig] = None,
    ):
        super().__init__()
        embed_dim = config.d_model

        self.conv1 = nn.Conv1d(config.num_mel_bins, embed_dim, kernel_size=3, padding=1)
        self.conv2 = nn.Conv1d(embed_dim, embed_dim, kernel_size=3, stride=2, padding=1)
        self.embed_positions = nn.Embedding(config.max_source_positions, embed_dim)
        self.layers = nn.ModuleList(
            [
                VoxtralWhisperEncoderLayer(config, quant_config)
                for _ in range(config.encoder_layers)
            ]
        )
        self.layer_norm = nn.LayerNorm(embed_dim)

    def forward(self, input_features: torch.Tensor) -> torch.Tensor:
        """
        Args:
            input_features: [batch, num_mel_bins, seq_len]
        Returns:
            [batch, seq_len // 2, d_model]
        """
        inputs_embeds = torch.nn.functional.gelu(self.conv1(input_features))
        inputs_embeds = torch.nn.functional.gelu(self.conv2(inputs_embeds))
        inputs_embeds = inputs_embeds.permute(0, 2, 1)

        seq_len = inputs_embeds.shape[1]
        position_ids = torch.arange(seq_len, device=inputs_embeds.device)
        hidden_states = inputs_embeds + self.embed_positions(position_ids)

        for layer in self.layers:
            hidden_states = layer(hidden_states)

        hidden_states = self.layer_norm(hidden_states)
        return hidden_states


class VoxtralForConditionalGeneration(nn.Module):
    """Voxtral: Whisper encoder + MLP projector + Llama decoder.

    HF weight prefixes:
        audio_tower.*           -> self.audio_tower (VoxtralWhisperEncoder)
        multi_modal_projector.* -> self.multi_modal_projector (AudioLanguageAdapter)
        language_model.*        -> self.language_model (LlamaForCausalLM)
    """

    def __init__(
        self,
        config: PretrainedConfig,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ) -> None:
        super().__init__()
        self.config = config

        audio_config = config.audio_config
        text_config = config.text_config

        # Ensure text_config has rope_parameters (transformers v5 compatibility)
        if not hasattr(text_config, "rope_parameters"):
            text_config.rope_parameters = {
                "rope_type": getattr(text_config, "rope_type", "default"),
                "rope_theta": getattr(text_config, "rope_theta", 10000.0),
            }
            if getattr(text_config, "rope_scaling", None):
                text_config.rope_parameters.update(text_config.rope_scaling)

        # Infer downsample_factor: intermediate_size / hidden_size for HF format
        self.downsample_factor = getattr(
            audio_config,
            "downsample_factor",
            audio_config.intermediate_size // audio_config.hidden_size,
        )

        # Encoder (named audio_tower to match HF weight prefix directly)
        self.audio_tower = VoxtralWhisperEncoder(audio_config, quant_config)

        # Projector: input = d_model * downsample_factor, output = text_hidden_size
        adapter_input_dim = audio_config.d_model * self.downsample_factor
        self.multi_modal_projector = AudioLanguageAdapter(
            hidden_size=adapter_input_dim,
            dim=text_config.hidden_size,
        )

        # Language model
        self.language_model = LlamaForCausalLM(text_config, quant_config=quant_config)

        # Mel filter bank for raw waveform -> mel spectrogram
        self._init_mel_filters(audio_config)

        self.pattern = MultiModalityDataPaddingPatternMultimodalTokens()

    def _init_mel_filters(self, audio_config: PretrainedConfig):
        """Initialize mel filter bank for mel spectrogram computation."""
        self._window_size = getattr(audio_config, "window_size", 400)
        self._hop_length = getattr(audio_config, "hop_length", 160)
        self._sampling_rate = getattr(audio_config, "sampling_rate", 16000)

        try:
            from mistral_common.audio import mel_filter_bank
        except ImportError:
            raise ImportError(
                "mistral_common is required for Voxtral. "
                "Install it with: pip install mistral_common"
            )

        mel_filters = mel_filter_bank(
            num_frequency_bins=1 + self._window_size // 2,
            num_mel_bins=audio_config.num_mel_bins,
            min_frequency=0.0,
            max_frequency=8000.0,
            sampling_rate=self._sampling_rate,
        )
        self.register_buffer(
            "mel_filters", torch.tensor(mel_filters, dtype=torch.float32)
        )

    @property
    def _conv_downsample_factor(self) -> int:
        return self.audio_tower.conv1.stride[0] * self.audio_tower.conv2.stride[0]

    @property
    def _chunk_size(self) -> int:
        return (
            self.config.audio_config.max_source_positions * self._conv_downsample_factor
        )

    def _compute_mel_spectrogram(self, audio_waveform: torch.Tensor) -> torch.Tensor:
        """Compute log-mel spectrogram from raw waveform using STFT."""
        window = torch.hann_window(self._window_size, device=audio_waveform.device)
        stft = torch.stft(
            audio_waveform,
            self._window_size,
            self._hop_length,
            window=window,
            return_complex=True,
        )
        magnitudes = stft[..., :-1].abs() ** 2
        mel_spec = self.mel_filters.T @ magnitudes
        log_spec = torch.clamp(mel_spec, min=1e-10).log10()
        log_spec_max = log_spec.max()
        log_spec = torch.maximum(log_spec, log_spec_max - 8.0)
        log_spec = (log_spec + 4.0) / 4.0
        return log_spec

    def _encode_audio(self, audio_waveforms: List[torch.Tensor]) -> List[torch.Tensor]:
        """Encode raw audio waveforms through mel spectrogram + whisper encoder."""
        dtype = self.audio_tower.conv1.weight.dtype
        device = self.audio_tower.conv1.weight.device

        chunked_features: List[torch.Tensor] = []
        chunks_per_example: List[int] = []
        chunk_size = self._chunk_size
        # Pad raw audio to a multiple of chunk_samples so that silence is
        # properly converted to mel features (matching HF VoxtralProcessor).
        chunk_samples = chunk_size * self._hop_length

        for waveform in audio_waveforms:
            waveform = waveform.to(device=device, dtype=torch.float32)
            n_samples = waveform.shape[-1]
            target_samples = chunk_samples * math.ceil(n_samples / chunk_samples)
            if target_samples > n_samples:
                waveform = torch.nn.functional.pad(
                    waveform, (0, target_samples - n_samples)
                )
            mel = self._compute_mel_spectrogram(waveform)
            chunks = mel.split(chunk_size, dim=-1)
            chunked_features.extend(chunks)
            chunks_per_example.append(len(chunks))

        if not chunked_features:
            return []

        input_embeds = torch.stack(chunked_features).to(dtype)
        encoder_out = self.audio_tower(input_embeds)

        results = []
        chunk_idx = 0
        for n_chunks in chunks_per_example:
            result = encoder_out[chunk_idx : chunk_idx + n_chunks].flatten(0, 1)
            results.append(result)
            chunk_idx += n_chunks

        return results

    def pad_input_ids(self, input_ids: List[int], mm_inputs: MultimodalInputs):
        return self.pattern.pad_input_tokens(input_ids, mm_inputs)

    def get_audio_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor:
        """Encode audio waveforms -> downsample -> project."""
        audio_waveforms = [item.feature for item in items]
        audio_embeddings = self._encode_audio(audio_waveforms)

        # Downsample: reshape to merge adjacent frames
        for i, emb in enumerate(audio_embeddings):
            seq_len, dim = emb.shape
            audio_embeddings[i] = emb.reshape(
                seq_len // self.downsample_factor,
                dim * self.downsample_factor,
            )

        # Project through adapter
        packed = torch.cat(audio_embeddings, dim=0)
        packed = self.multi_modal_projector(packed)

        return packed

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        forward_batch: ForwardBatch,
        **kwargs: Any,
    ) -> torch.Tensor:
        hidden_states = general_mm_embed_routine(
            input_ids=input_ids,
            forward_batch=forward_batch,
            language_model=self.language_model,
            data_embedding_funcs={
                Modality.AUDIO: self.get_audio_feature,
            },
            positions=positions,
        )
        return hidden_states

    def get_language_model(self) -> nn.Module:
        return self.language_model

    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
        encoder_stacked = [
            ("qkv_proj", "q_proj", "q"),
            ("qkv_proj", "k_proj", "k"),
            ("qkv_proj", "v_proj", "v"),
        ]

        encoder_dict = dict(self.audio_tower.named_parameters())
        projector_dict = dict(self.multi_modal_projector.named_parameters())

        # Collect all weights; synthesise missing k_proj bias as zeros.
        weights_list = list(weights)
        extra_weights = []
        for name, w in weights_list:
            if name.startswith("audio_tower.") and ".self_attn.k_proj.weight" in name:
                bias_name = name.replace(".weight", ".bias")
                if not any(n == bias_name for n, _ in weights_list):
                    extra_weights.append(
                        (bias_name, torch.zeros(w.shape[0], dtype=w.dtype))
                    )
        weights_list.extend(extra_weights)

        def llm_weights_generator():
            for name, w in weights_list:
                # Encoder weights
                if name.startswith("audio_tower."):
                    trimmed = name[len("audio_tower.") :]
                    loaded = False
                    for param_name, weight_name, shard_id in encoder_stacked:
                        if f".{weight_name}." in trimmed:
                            stacked_name = trimmed.replace(weight_name, param_name)
                            if stacked_name in encoder_dict:
                                param = encoder_dict[stacked_name]
                                param.weight_loader(param, w, shard_id)
                                loaded = True
                                break
                    if not loaded and trimmed in encoder_dict:
                        param = encoder_dict[trimmed]
                        weight_loader = getattr(
                            param, "weight_loader", default_weight_loader
                        )
                        weight_loader(param, w)
                    continue

                # Projector weights
                if name.startswith("multi_modal_projector."):
                    trimmed = name[len("multi_modal_projector.") :]
                    trimmed = trimmed.replace("linear_1.", "w_in.").replace(
                        "linear_2.", "w_out."
                    )
                    if trimmed in projector_dict:
                        param = projector_dict[trimmed]
                        default_weight_loader(param, w)
                    continue

                # LLM weights
                if name.startswith("language_model."):
                    name = name[len("language_model.") :]
                yield (name, w)

        self.language_model.load_weights(llm_weights_generator())


EntryClass = [VoxtralForConditionalGeneration]