# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Inference-only MOSS-Audio model compatible with HuggingFace weights."""

import math
from collections.abc import Iterable, Mapping, Sequence
from dataclasses import dataclass, field
from typing import Annotated, Any

import numpy as np
import regex as re
import torch
import torch.nn.functional as F
from torch import nn
from transformers import BatchFeature, PretrainedConfig, Qwen3Config
from transformers.models.whisper import WhisperFeatureExtractor

from vllm.compilation.decorators import support_torch_compile
from vllm.config import VllmConfig
from vllm.config.multimodal import BaseDummyOptions
from vllm.distributed import (
    get_pp_group,
    get_tensor_model_parallel_world_size,
)
from vllm.inputs import ModalityData, MultiModalDataDict
from vllm.model_executor.layers.activation import _ACTIVATION_REGISTRY, SiluAndMul
from vllm.model_executor.layers.linear import (
    ColumnParallelLinear,
    MergedColumnParallelLinear,
    ReplicatedLinear,
    RowParallelLinear,
)
from vllm.model_executor.layers.logits_processor import LogitsProcessor
from vllm.model_executor.layers.quantization import QuantizationConfig
from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
from vllm.multimodal import MULTIMODAL_REGISTRY
from vllm.multimodal.inputs import (
    AudioItem,
    MultiModalFieldConfig,
    MultiModalKwargsItems,
)
from vllm.multimodal.parse import (
    DictEmbeddingItems,
    ModalityDataItems,
    MultiModalDataItems,
    MultiModalDataParser,
)
from vllm.multimodal.processing import (
    BaseDummyInputsBuilder,
    BaseMultiModalProcessor,
    BaseProcessingInfo,
    PromptReplacement,
    PromptUpdate,
    PromptUpdateDetails,
)
from vllm.sequence import IntermediateTensors
from vllm.transformers_utils.repo_utils import get_hf_file_to_dict
from vllm.utils.tensor_schema import TensorSchema, TensorShape

from .interfaces import (
    MultiModalEmbeddings,
    SupportsLoRA,
    SupportsMultiModal,
    SupportsPP,
    _require_is_multimodal,
)
from .module_mapping import MultiModelKeys
from .qwen3 import Qwen3ForCausalLM, Qwen3Model
from .utils import (
    AutoWeightsLoader,
    WeightsMapper,
    _merge_multimodal_embeddings,
    maybe_prefix,
)

MOSS_AUDIO_TOKEN = "<|AUDIO|>"
MOSS_AUDIO_BOS_TOKEN = "<|audio_bos|>"
MOSS_AUDIO_EOS_TOKEN = "<|audio_eos|>"
MOSS_AUDIO_TOKEN_ID = 151654
MOSS_AUDIO_BOS_TOKEN_ID = 151669
MOSS_AUDIO_EOS_TOKEN_ID = 151670
DEFAULT_MAX_AUDIO_SECONDS = 30
DEFAULT_MOSS_AUDIO_MEL_CONFIG = {
    "mel_dim": 128,
    "mel_sr": 16000,
    "mel_hop_length": 160,
    "mel_n_fft": 400,
}
MOSS_AUDIO_PLACEHOLDER = (
    f"{MOSS_AUDIO_BOS_TOKEN}{MOSS_AUDIO_TOKEN}{MOSS_AUDIO_EOS_TOKEN}"
)
MOSS_AUDIO_SPAN_RE = re.compile(
    f"{re.escape(MOSS_AUDIO_BOS_TOKEN)}"
    f"(?:{re.escape(MOSS_AUDIO_TOKEN)})+"
    f"{re.escape(MOSS_AUDIO_EOS_TOKEN)}"
)
MOSS_AUDIO_PROCESSOR_CONFIG_KEYS = {
    "audio_token_id",
    "audio_start_id",
    "audio_end_id",
    "enable_time_marker",
    "mel_config",
}


class MossAudioAudioInputs(TensorSchema):
    """
    Dimensions:
        - b: Batch size
        - nmb: Number of mel bins
        - t: Time frames
    """

    audio_data: Annotated[torch.Tensor, TensorShape("b", "nmb", "t")]
    audio_data_seqlens: Annotated[torch.Tensor, TensorShape("b")]


def _normalize_moss_audio_mel_config(
    mel_config: Mapping[str, object] | None = None,
) -> dict[str, int]:
    config = dict(DEFAULT_MOSS_AUDIO_MEL_CONFIG)
    config.update(_extract_moss_audio_mel_config(mel_config))
    return config


def _extract_moss_audio_mel_config(
    mel_config: Mapping[str, object] | None = None,
) -> dict[str, int]:
    config: dict[str, int] = {}
    if mel_config is None:
        return config

    aliases = {
        "mel_dim": ("mel_dim", "feature_size", "n_mels", "num_mel_bins"),
        "mel_sr": ("mel_sr", "sampling_rate", "sample_rate"),
        "mel_hop_length": ("mel_hop_length", "hop_length"),
        "mel_n_fft": ("mel_n_fft", "n_fft"),
    }
    for target_key, source_keys in aliases.items():
        for source_key in source_keys:
            if source_key in mel_config:
                config[target_key] = int(mel_config[source_key])
                break

    return config


def _filter_moss_audio_processor_config(
    config: Mapping[str, object] | None,
) -> dict[str, object]:
    if not config:
        return {}

    return {
        key: value
        for key, value in config.items()
        if key in MOSS_AUDIO_PROCESSOR_CONFIG_KEYS
    }


def _merge_moss_audio_processor_configs(
    *configs: Mapping[str, object] | None,
) -> dict[str, object]:
    merged: dict[str, object] = {}
    merged_mel_config: dict[str, int] = {}
    for config in configs:
        filtered = _filter_moss_audio_processor_config(config)
        mel_config = filtered.pop("mel_config", None)
        merged.update(filtered)
        if isinstance(mel_config, Mapping):
            merged_mel_config.update(_extract_moss_audio_mel_config(mel_config))

    if merged_mel_config:
        merged["mel_config"] = merged_mel_config
    return merged


@dataclass
class MossAudioEncoderConfig:
    d_model: int = 1280
    output_dim: int = 1280
    num_mel_bins: int = 128
    encoder_layers: int = 32
    encoder_attention_heads: int = 20
    encoder_ffn_dim: int = 5120
    downsample_rate: int = 8
    downsample_hidden_size: int = 480
    encoder_attention_window_size: int = 100
    max_source_positions: int = 1500
    dropout: float = 0.1
    attention_dropout: float = 0.1
    activation_dropout: float = 0.0
    activation_function: str = "gelu"
    layer_norm_eps: float = 1e-5
    _attn_implementation: str = "eager"
    pretrained_path: str = ""
    n_window: int = 200
    conv_chunksize: int = 64
    deepstack_encoder_layer_indexes: list[int] = field(
        default_factory=lambda: [8, 16, 24]
    )

    @classmethod
    def from_config(cls, config: object) -> "MossAudioEncoderConfig":
        if isinstance(config, cls):
            return config
        if isinstance(config, Mapping):
            values = {
                key: value
                for key, value in config.items()
                if key in cls.__dataclass_fields__
            }
        else:
            values = {
                key: getattr(config, key)
                for key in cls.__dataclass_fields__
                if hasattr(config, key)
            }
        return cls(**values)


class MossAudioConfig(PretrainedConfig):
    model_type = "moss_audio"
    is_composition = True

    def __init__(
        self,
        audio_config: Mapping[str, object] | MossAudioEncoderConfig | None = None,
        language_config: Mapping[str, object] | Qwen3Config | None = None,
        adapter_hidden_size: int = 8192,
        ignore_index: int = -100,
        deepstack_num_inject_layers: int | None = None,
        **kwargs: object,
    ) -> None:
        self.audio_config = MossAudioEncoderConfig.from_config(audio_config or {})
        if isinstance(language_config, Qwen3Config):
            self.language_config = language_config
        else:
            self.language_config = Qwen3Config(**(language_config or {}))

        self.adapter_hidden_size = adapter_hidden_size
        self.ignore_index = ignore_index
        self.deepstack_num_inject_layers = deepstack_num_inject_layers

        for key in ("num_hidden_layers", "eos_token_id", "bos_token_id", "vocab_size"):
            kwargs.setdefault(key, getattr(self.language_config, key, None))
        kwargs.setdefault("tie_word_embeddings", False)
        super().__init__(**kwargs)

        for key in (
            "hidden_size",
            "num_attention_heads",
            "num_key_value_heads",
            "head_dim",
            "max_position_embeddings",
            "rms_norm_eps",
        ):
            if hasattr(self.language_config, key):
                setattr(self, key, getattr(self.language_config, key))

    def get_text_config(self, decoder: bool = False) -> Qwen3Config:
        return self.language_config


class SinusoidsPositionEmbedding(nn.Module):
    def __init__(self, num_positions: int, embedding_dim: int) -> None:
        super().__init__()
        del num_positions  # Kept for config compatibility.
        max_timescale = 10000.0
        log_timescale_increment = math.log(max_timescale) / (embedding_dim // 2 - 1)
        inv_timescales = torch.exp(
            -log_timescale_increment * torch.arange(embedding_dim // 2).float()
        )
        self.register_buffer("inv_timescales", inv_timescales, persistent=False)

    def forward(self, seq_len: int, device: torch.device) -> torch.Tensor:
        scaled_time = (
            torch.arange(seq_len, device=device, dtype=self.inv_timescales.dtype)[
                :, None
            ]
            * self.inv_timescales[None, :]
        )
        return torch.cat([scaled_time.sin(), scaled_time.cos()], dim=1).unsqueeze(0)


class MossAudioAttention(nn.Module):
    def __init__(
        self,
        config: MossAudioEncoderConfig,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ) -> None:
        super().__init__()
        self.embed_dim = config.d_model
        self.num_heads = config.encoder_attention_heads
        self.head_dim = self.embed_dim // self.num_heads
        if self.head_dim * self.num_heads != self.embed_dim:
            raise ValueError(
                f"d_model ({self.embed_dim}) must be divisible by "
                f"encoder_attention_heads ({self.num_heads})."
            )

        tp_size = get_tensor_model_parallel_world_size()
        if self.num_heads % tp_size != 0:
            raise ValueError(
                "MOSS-Audio audio encoder attention heads must be divisible by "
                f"tensor parallel size. Got {self.num_heads=} and {tp_size=}."
            )
        self.num_local_heads = self.num_heads // tp_size
        # TODO: can use QKVParallelLinear
        self.q_proj = ColumnParallelLinear(
            input_size=self.embed_dim,
            output_size=self.embed_dim,
            bias=True,
            quant_config=quant_config,
            prefix=f"{prefix}.q_proj",
        )
        self.k_proj = ColumnParallelLinear(
            input_size=self.embed_dim,
            output_size=self.embed_dim,
            bias=False,
            quant_config=quant_config,
            prefix=f"{prefix}.k_proj",
        )
        self.v_proj = ColumnParallelLinear(
            input_size=self.embed_dim,
            output_size=self.embed_dim,
            bias=True,
            quant_config=quant_config,
            prefix=f"{prefix}.v_proj",
        )
        self.out_proj = RowParallelLinear(
            input_size=self.embed_dim,
            output_size=self.embed_dim,
            bias=True,
            quant_config=quant_config,
            prefix=f"{prefix}.out_proj",
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor,
    ) -> torch.Tensor:
        batch_size, seq_len, _ = hidden_states.shape
        q, _ = self.q_proj(hidden_states)
        k, _ = self.k_proj(hidden_states)
        v, _ = self.v_proj(hidden_states)
        q = q.view(batch_size, seq_len, -1, self.head_dim).transpose(1, 2)
        k = k.view(batch_size, seq_len, -1, self.head_dim).transpose(1, 2)
        v = v.view(batch_size, seq_len, -1, self.head_dim).transpose(1, 2)
        attn_output = F.scaled_dot_product_attention(
            q,
            k,
            v,
            attn_mask=attention_mask[:, None, None, :],
            dropout_p=0.0,
            scale=self.head_dim**-0.5,
        )
        output, _ = self.out_proj(
            attn_output.transpose(1, 2).reshape(
                batch_size,
                seq_len,
                -1,
            )
        )
        return output


class MossAudioEncoderLayer(nn.Module):
    def __init__(
        self,
        config: MossAudioEncoderConfig,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ) -> None:
        super().__init__()
        self.self_attn = MossAudioAttention(
            config,
            quant_config=quant_config,
            prefix=f"{prefix}.self_attn",
        )
        self.self_attn_layer_norm = nn.LayerNorm(
            config.d_model, eps=config.layer_norm_eps
        )
        self.activation_fn = _ACTIVATION_REGISTRY[config.activation_function]
        self.activation_dropout = config.activation_dropout
        self.dropout = config.dropout
        self.fc1 = ColumnParallelLinear(
            config.d_model,
            config.encoder_ffn_dim,
            bias=True,
            quant_config=quant_config,
            prefix=f"{prefix}.fc1",
        )
        self.fc2 = RowParallelLinear(
            config.encoder_ffn_dim,
            config.d_model,
            bias=True,
            quant_config=quant_config,
            prefix=f"{prefix}.fc2",
        )
        self.final_layer_norm = nn.LayerNorm(config.d_model, eps=config.layer_norm_eps)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor,
    ) -> torch.Tensor:
        residual = hidden_states
        hidden_states = self.self_attn_layer_norm(hidden_states)
        hidden_states = self.self_attn(hidden_states, attention_mask)
        hidden_states = residual + F.dropout(
            hidden_states, p=self.dropout, training=self.training
        )

        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 = F.dropout(
            hidden_states, p=self.activation_dropout, training=self.training
        )
        hidden_states, _ = self.fc2(hidden_states)
        hidden_states = residual + F.dropout(
            hidden_states, p=self.dropout, training=self.training
        )
        return hidden_states


class MossAudioEncoder(nn.Module):
    def __init__(
        self,
        config: MossAudioEncoderConfig,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ) -> None:
        super().__init__()
        self.config = config
        self.gelu = nn.GELU()
        self.conv1 = nn.Conv2d(
            1,
            config.downsample_hidden_size,
            kernel_size=(3, 3),
            stride=(2, 2),
            padding=(1, 1),
        )
        self.conv2 = nn.Conv2d(
            config.downsample_hidden_size,
            config.downsample_hidden_size,
            kernel_size=(3, 3),
            stride=(2, 2),
            padding=(1, 1),
        )
        self.conv3 = nn.Conv2d(
            config.downsample_hidden_size,
            config.downsample_hidden_size,
            kernel_size=(3, 3),
            stride=(2, 2),
            padding=(1, 1),
        )

        conv_freq = self._compute_downsampled_length(
            torch.tensor(config.num_mel_bins)
        ).item()
        self.stem_proj = ReplicatedLinear(
            config.downsample_hidden_size * int(conv_freq),
            config.d_model,
            bias=True,
            quant_config=quant_config,
            return_bias=False,
            prefix=f"{prefix}.stem_proj",
        )
        self.embed_positions = SinusoidsPositionEmbedding(
            config.max_source_positions, config.d_model
        )
        self.layers = nn.ModuleList(
            [
                MossAudioEncoderLayer(
                    config,
                    quant_config=quant_config,
                    prefix=f"{prefix}.layers.{layer_idx}",
                )
                for layer_idx in range(config.encoder_layers)
            ]
        )
        self.layer_norm = nn.LayerNorm(config.d_model, eps=config.layer_norm_eps)
        if config.output_dim != config.d_model:
            self.out_proj = ReplicatedLinear(
                config.d_model,
                config.output_dim,
                bias=False,
                quant_config=quant_config,
                return_bias=False,
                prefix=f"{prefix}.out_proj",
            )
        else:
            self.out_proj = nn.Identity()

        self.deepstack_encoder_layer_indexes = list(
            config.deepstack_encoder_layer_indexes or []
        )
        self._deepstack_capture_map = {
            layer_idx: capture_idx
            for capture_idx, layer_idx in enumerate(
                self.deepstack_encoder_layer_indexes
            )
        }
        self.n_window = int(config.n_window)
        self.chunk_frames = int(self.n_window * 2)
        self.conv_chunksize = int(config.conv_chunksize)

    @property
    def dtype(self) -> torch.dtype:
        return self.conv1.weight.dtype

    @staticmethod
    def _compute_downsampled_length(lengths: torch.Tensor) -> torch.Tensor:
        def conv_out_len(length: torch.Tensor) -> torch.Tensor:
            return (length - 1) // 2 + 1

        return conv_out_len(conv_out_len(conv_out_len(lengths)))

    @staticmethod
    def compute_num_audio_tokens(raw_mel_len: int) -> int:
        lengths = torch.tensor(raw_mel_len, dtype=torch.long)
        return int(MossAudioEncoder._compute_downsampled_length(lengths).item())

    def _encode_chunk_batch(
        self,
        input_features: torch.Tensor,
        seq_lengths: torch.Tensor,
        output_deepstack_hidden_states: bool = True,
    ) -> tuple[torch.Tensor, list[torch.Tensor]]:
        if input_features.dim() == 2:
            input_features = input_features.unsqueeze(0)

        downsampled_lengths = self._compute_downsampled_length(seq_lengths)
        x = input_features.unsqueeze(1)
        x = self.gelu(self.conv1(x))
        x = self.gelu(self.conv2(x))
        x = self.gelu(self.conv3(x))
        x = x.permute(0, 3, 1, 2).contiguous().flatten(2)
        x = self.stem_proj(x)

        max_len = int(downsampled_lengths.max().item())
        if x.size(1) > max_len:
            x = x[:, :max_len, :]
        x = x + self.embed_positions(x.shape[1], x.device).to(x.dtype)

        attention_mask = (
            torch.arange(x.size(1), device=x.device)[None, :]
            < downsampled_lengths[:, None]
        )

        deepstack_hidden_states: list[torch.Tensor | None] = []
        if output_deepstack_hidden_states:
            deepstack_hidden_states = [None] * len(self.deepstack_encoder_layer_indexes)
        for layer_idx, layer in enumerate(self.layers):
            x = layer(x, attention_mask)
            if output_deepstack_hidden_states:
                capture_idx = self._deepstack_capture_map.get(layer_idx)
                if capture_idx is not None:
                    deepstack_hidden_states[capture_idx] = x

        x = self.layer_norm(x)
        x = self.out_proj(x)

        if not output_deepstack_hidden_states:
            return x, []

        ordered_deepstack_hidden_states = [
            hidden_states
            for hidden_states in deepstack_hidden_states
            if hidden_states is not None
        ]
        ordered_deepstack_hidden_states = [
            self.out_proj(hidden_states)
            for hidden_states in ordered_deepstack_hidden_states
        ]
        return x, ordered_deepstack_hidden_states

    def forward(
        self,
        input_features: torch.Tensor,
        feature_lens: torch.Tensor | None = None,
        output_deepstack_hidden_states: bool = True,
    ) -> tuple[torch.Tensor, tuple[torch.Tensor, ...] | None]:
        if input_features.dim() == 3:
            if feature_lens is None:
                feature_lens = torch.full(
                    (input_features.size(0),),
                    input_features.size(-1),
                    dtype=torch.long,
                    device=input_features.device,
                )
            else:
                feature_lens = feature_lens.to(
                    device=input_features.device, dtype=torch.long
                )
            valid_chunks = [
                input_features[i, :, : int(feature_lens[i].item())]
                for i in range(int(input_features.shape[0]))
            ]
            input_features = torch.cat(valid_chunks, dim=1)
        elif input_features.dim() != 2:
            raise ValueError(
                f"Expected [n_mels, T] or [B, n_mels, T], got "
                f"{tuple(input_features.shape)}."
            )

        if feature_lens is None:
            feature_lens = torch.tensor(
                [int(input_features.shape[1])],
                device=input_features.device,
                dtype=torch.long,
            )
        else:
            feature_lens = feature_lens.to(
                device=input_features.device, dtype=torch.long
            )

        chunk_num = torch.ceil(
            feature_lens.to(torch.float32) / self.chunk_frames
        ).long()
        chunk_lengths = torch.full(
            (int(chunk_num.sum().item()),),
            self.chunk_frames,
            dtype=torch.long,
            device=feature_lens.device,
        )
        tail_chunk_index = F.pad(chunk_num, (1, 0), value=-1).cumsum(0)[1:]
        chunk_lengths[tail_chunk_index] = feature_lens % self.chunk_frames
        chunk_lengths[chunk_lengths == 0] = self.chunk_frames

        chunk_list = input_features.T.split(chunk_lengths.tolist(), dim=0)
        padded_feature = nn.utils.rnn.pad_sequence(
            chunk_list, batch_first=True
        ).transpose(1, 2)

        feature_lens_after_cnn = self._compute_downsampled_length(chunk_lengths)
        t_down_max = (
            int(feature_lens_after_cnn.max().item())
            if feature_lens_after_cnn.numel() > 0
            else 0
        )
        indices = torch.arange(t_down_max, device=padded_feature.device)
        padded_mask_after_cnn = indices[None, :] < feature_lens_after_cnn[:, None]

        num_deepstack = len(self.deepstack_encoder_layer_indexes)
        should_output_deepstack = output_deepstack_hidden_states and num_deepstack > 0
        padded_embeds: list[torch.Tensor] = []
        deepstack_padded_embeds: list[list[torch.Tensor]] = [
            [] for _ in range(num_deepstack if should_output_deepstack else 0)
        ]
        for feat_chunk, len_chunk in zip(
            padded_feature.split(self.conv_chunksize, dim=0),
            chunk_lengths.split(self.conv_chunksize, dim=0),
        ):
            out, deepstack_outs = self._encode_chunk_batch(
                feat_chunk,
                len_chunk,
                output_deepstack_hidden_states=should_output_deepstack,
            )
            if out.shape[1] < t_down_max:
                out = F.pad(out, (0, 0, 0, t_down_max - out.shape[1]))
            padded_embeds.append(out)

            if should_output_deepstack:
                if len(deepstack_outs) != num_deepstack:
                    raise RuntimeError(
                        "DeepStack output count does not match configured "
                        "layer indexes."
                    )
                for capture_idx, ds in enumerate(deepstack_outs):
                    if ds.shape[1] < t_down_max:
                        ds = F.pad(ds, (0, 0, 0, t_down_max - ds.shape[1]))
                    deepstack_padded_embeds[capture_idx].append(ds)

        if padded_embeds:
            padded_embed = torch.cat(padded_embeds, dim=0)
        else:
            padded_embed = torch.empty(
                (0, t_down_max, self.config.output_dim),
                device=padded_feature.device,
                dtype=padded_feature.dtype,
            )

        last_hidden_state = padded_embed[padded_mask_after_cnn].unsqueeze(0)

        deepstack_states: tuple[torch.Tensor, ...] | None = None
        if should_output_deepstack:
            collected: list[torch.Tensor] = []
            for chunks_list in deepstack_padded_embeds:
                if chunks_list:
                    ds = torch.cat(chunks_list, dim=0)
                    collected.append(ds[padded_mask_after_cnn].unsqueeze(0))
                else:
                    collected.append(
                        torch.empty(
                            (1, 0, self.config.output_dim),
                            device=padded_feature.device,
                            dtype=padded_embed.dtype,
                        )
                    )
            deepstack_states = tuple(collected)

        return last_hidden_state, deepstack_states


class GatedMLP(nn.Module):
    def __init__(
        self,
        input_size: int,
        hidden_size: int,
        output_size: int,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ) -> None:
        super().__init__()
        self.gate_up_proj = MergedColumnParallelLinear(
            input_size,
            [hidden_size, hidden_size],
            bias=False,
            quant_config=quant_config,
            prefix=f"{prefix}.gate_up_proj",
        )
        self.down_proj = RowParallelLinear(
            hidden_size,
            output_size,
            bias=False,
            input_is_parallel=True,
            quant_config=quant_config,
            prefix=f"{prefix}.down_proj",
        )
        self.act_fn = SiluAndMul()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        gate_up, _ = self.gate_up_proj(x)
        x = self.act_fn(gate_up)
        x, _ = self.down_proj(x)
        return x

    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
        stacked_params_mapping = [
            ("gate_up_proj", "gate_proj", 0),
            ("gate_up_proj", "up_proj", 1),
        ]
        params_dict = dict(self.named_parameters())
        loaded_params: set[str] = set()

        for name, loaded_weight in weights:
            target_name = name
            for param_name, weight_name, shard_id in stacked_params_mapping:
                components = target_name.split(".")
                if weight_name not in components:
                    continue

                target_name = ".".join(
                    param_name if component == weight_name else component
                    for component in components
                )
                param = params_dict[target_name]
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
                break
            else:
                param = params_dict[target_name]
                weight_loader = getattr(param, "weight_loader", default_weight_loader)
                weight_loader(param, loaded_weight)

            loaded_params.add(target_name)
        return loaded_params


@support_torch_compile(
    dynamic_arg_dims={
        "input_ids": 0,
        "positions": -1,
        "intermediate_tensors": 0,
        "inputs_embeds": 0,
        "deepstack_input_embeds": 0,
    }
)
class MossQwen3Model(Qwen3Model):
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
        super().__init__(vllm_config=vllm_config, prefix=prefix)
        self.deepstack_inject_layer_indices: Iterable[int] = range(0)

    def forward(
        self,
        input_ids: torch.Tensor | None,
        positions: torch.Tensor,
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
        deepstack_input_embeds: IntermediateTensors | None = None,
    ) -> torch.Tensor | IntermediateTensors:
        if get_pp_group().is_first_rank:
            if inputs_embeds is not None:
                hidden_states = inputs_embeds
            else:
                hidden_states = self.embed_input_ids(input_ids)
            residual = None
        else:
            assert intermediate_tensors is not None
            hidden_states = intermediate_tensors["hidden_states"]
            residual = intermediate_tensors["residual"]

        aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual)
        for layer_idx, layer in enumerate(
            self.layers[self.start_layer : self.end_layer],
            start=self.start_layer,
        ):
            hidden_states, residual = layer(positions, hidden_states, residual)
            deepstack_key = f"deepstack_input_embeds_{layer_idx}"
            if (
                deepstack_input_embeds is not None
                and deepstack_key in deepstack_input_embeds.tensors
            ):
                hidden_states = hidden_states + deepstack_input_embeds[deepstack_key]
            self._maybe_add_hidden_state(
                aux_hidden_states,
                layer_idx - self.start_layer + 1,
                hidden_states,
                residual,
            )

        if not get_pp_group().is_last_rank:
            tensors = {"hidden_states": hidden_states, "residual": residual}
            # Keep the DeepStack PP schema config-driven, but only carry
            # payloads needed by downstream injection points across this rank.
            # Missing downstream payloads are zero-filled below to clear
            # receive buffers instead of leaving stale tensors.
            for layer_idx in self.deepstack_inject_layer_indices:
                if layer_idx < self.end_layer:
                    continue
                deepstack_key = f"deepstack_input_embeds_{layer_idx}"
                if (
                    deepstack_input_embeds is not None
                    and deepstack_key in deepstack_input_embeds.tensors
                ):
                    tensors[deepstack_key] = deepstack_input_embeds[deepstack_key]
                else:
                    tensors[deepstack_key] = hidden_states.new_zeros(
                        hidden_states.shape
                    )
            return IntermediateTensors(tensors)

        hidden_states, _ = self.norm(hidden_states, residual)
        if len(aux_hidden_states) > 0:
            return hidden_states, aux_hidden_states
        return hidden_states


class MossQwen3ForCausalLM(Qwen3ForCausalLM):
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
        super(Qwen3ForCausalLM, self).__init__()
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config

        self.config = config
        self.vllm_config = vllm_config
        self.quant_config = quant_config
        self.model = MossQwen3Model(
            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
        )

        if get_pp_group().is_last_rank:
            if config.tie_word_embeddings:
                self.lm_head = self.model.embed_tokens
            else:
                self.lm_head = ParallelLMHead(
                    config.vocab_size,
                    config.hidden_size,
                    quant_config=quant_config,
                    prefix=maybe_prefix(prefix, "lm_head"),
                )
        else:
            from .utils import PPMissingLayer

            self.lm_head = PPMissingLayer()

        self.logits_processor = LogitsProcessor(config.vocab_size)
        self.deepstack_inject_layer_indices: Iterable[int] = range(0)

    def make_empty_intermediate_tensors(
        self,
        batch_size: int,
        dtype: torch.dtype,
        device: torch.device,
    ) -> IntermediateTensors:
        intermediate_tensors = self.model.make_empty_intermediate_tensors(
            batch_size, dtype, device
        )
        for layer_idx in self.deepstack_inject_layer_indices:
            intermediate_tensors[f"deepstack_input_embeds_{layer_idx}"] = torch.zeros(
                (batch_size, self.config.hidden_size),
                dtype=dtype,
                device=device,
            )
        return intermediate_tensors

    def forward(
        self,
        input_ids: torch.Tensor | None,
        positions: torch.Tensor,
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
        deepstack_input_embeds: IntermediateTensors | None = None,
    ) -> torch.Tensor | IntermediateTensors:
        return self.model(
            input_ids,
            positions,
            intermediate_tensors,
            inputs_embeds,
            deepstack_input_embeds=deepstack_input_embeds,
        )


def _moss_audio_field_config(
    hf_inputs: Mapping[str, torch.Tensor],
) -> Mapping[str, MultiModalFieldConfig]:
    return {
        "audio_data": MultiModalFieldConfig.batched("audio"),
        "audio_data_seqlens": MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
    }


class MossAudioMultiModalDataParser(MultiModalDataParser):
    def _parse_audio_data(
        self,
        data: dict[str, torch.Tensor] | ModalityData[AudioItem],
    ) -> ModalityDataItems[Any, Any] | None:
        if isinstance(data, dict):
            return DictEmbeddingItems(
                data,
                modality="audio",
                required_fields={"audio_data", "audio_data_seqlens"},
                fields_factory=_moss_audio_field_config,
            )

        return super()._parse_audio_data(data)


class MossAudioProcessor:
    model_input_names = [
        "input_ids",
        "attention_mask",
        "audio_data",
        "audio_data_seqlens",
    ]

    def __init__(
        self,
        tokenizer: object,
        *,
        audio_token_id: int = MOSS_AUDIO_TOKEN_ID,
        audio_start_id: int = MOSS_AUDIO_BOS_TOKEN_ID,
        audio_end_id: int = MOSS_AUDIO_EOS_TOKEN_ID,
        enable_time_marker: bool = False,
        mel_config: Mapping[str, object] | None = None,
    ) -> None:
        self.tokenizer = tokenizer
        self.audio_token_id = int(audio_token_id)
        self.audio_start_id = int(audio_start_id)
        self.audio_end_id = int(audio_end_id)
        self.enable_time_marker = bool(enable_time_marker)
        self.mel_config = _normalize_moss_audio_mel_config(mel_config)
        self.feature_extractor = WhisperFeatureExtractor(
            feature_size=self.mel_config["mel_dim"],
            sampling_rate=self.mel_config["mel_sr"],
            hop_length=self.mel_config["mel_hop_length"],
            n_fft=self.mel_config["mel_n_fft"],
        )
        self.audio_tokens_per_second = self.mel_config["mel_sr"] / (
            self.mel_config["mel_hop_length"] * 8
        )
        self.time_marker_every_seconds = 2
        self.time_marker_every_audio_tokens = int(
            self.audio_tokens_per_second * self.time_marker_every_seconds
        )
        self._digit_token_ids = {
            "0": 15,
            "1": 16,
            "2": 17,
            "3": 18,
            "4": 19,
            "5": 20,
            "6": 21,
            "7": 22,
            "8": 23,
            "9": 24,
        }

    @staticmethod
    def conv3_downsample_len(raw_mel_len: int) -> int:
        return MossAudioEncoder.compute_num_audio_tokens(raw_mel_len)

    def _extract_mel(self, audio: np.ndarray | torch.Tensor) -> torch.Tensor:
        if isinstance(audio, torch.Tensor):
            wav = audio.detach().to("cpu", dtype=torch.float32).numpy()
        else:
            wav = np.asarray(audio, dtype=np.float32)
        if wav.size == 0:
            raise ValueError("The audio is too short to be represented.")
        if wav.ndim == 2:
            wav = wav[0]
        feats = self.feature_extractor._np_extract_fbank_features(
            wav[None, ...], device="cpu"
        )
        return torch.from_numpy(feats[0])

    def _get_default_audio_prompt(self) -> str:
        return MOSS_AUDIO_PLACEHOLDER

    def _ensure_audio_placeholders(
        self,
        prompt_text: str,
        num_audios: int,
    ) -> str:
        if num_audios == 0 or MOSS_AUDIO_SPAN_RE.search(prompt_text):
            return prompt_text

        audio_prompt = self._get_default_audio_prompt() * num_audios
        if prompt_text:
            return f"{audio_prompt}\n{prompt_text}"
        return audio_prompt

    def _build_audio_tokens_with_time_markers(self, audio_seq_len: int) -> list[int]:
        total_duration_seconds = audio_seq_len / self.audio_tokens_per_second
        num_full_seconds = int(total_duration_seconds)
        token_ids: list[int] = []
        audio_tokens_consumed = 0
        for second in range(
            self.time_marker_every_seconds,
            num_full_seconds + 1,
            self.time_marker_every_seconds,
        ):
            marker_pos = (
                second // self.time_marker_every_seconds
            ) * self.time_marker_every_audio_tokens
            audio_segment_len = marker_pos - audio_tokens_consumed
            if audio_segment_len > 0:
                token_ids.extend([self.audio_token_id] * audio_segment_len)
                audio_tokens_consumed += audio_segment_len
            token_ids.extend(self._digit_token_ids[digit] for digit in str(second))

        remaining = audio_seq_len - audio_tokens_consumed
        if remaining > 0:
            token_ids.extend([self.audio_token_id] * remaining)
        return token_ids

    def build_audio_placeholder_ids(self, num_audio_tokens: int) -> list[int]:
        if self.enable_time_marker:
            return self._build_audio_tokens_with_time_markers(num_audio_tokens)
        return [self.audio_token_id] * num_audio_tokens

    def __call__(
        self,
        text: str | Sequence[str] | None = None,
        audios: Sequence[np.ndarray | torch.Tensor] | None = None,
        audio: Sequence[np.ndarray | torch.Tensor] | None = None,
        return_tensors: str = "pt",
        **kwargs: object,
    ) -> BatchFeature:
        """Build text tokens and audio tensors for one MossAudio prompt.

        Example:
            text="Describe this.", audio=[waveform]
            -> input_ids contains audio_start, N audio tokens, audio_end
            -> audio_data has shape [1, mel_dim, max_time]
            -> mel_dim is the number of mel filter-bank bins, 128 by default
            -> audio_data_seqlens stores the unpadded mel length
        """
        del kwargs

        # Step 1. Normalize text input; this processor handles one prompt.
        if isinstance(text, (list, tuple)):
            if len(text) != 1:
                raise ValueError(f"Expected text batch size 1, got {len(text)}")
            prompt_text = text[0]
        elif text is None:
            prompt_text = ""
        else:
            prompt_text = text

        # Step 2. Accept either `audios` or `audio` and normalize to a list.
        audio_list = audios if audios is not None else audio
        audio_list = [] if audio_list is None else list(audio_list)

        # Step 3. Convert waveforms to [mel_dim, time] mel features and token
        # counts. mel_dim is the number of mel filter-bank bins.
        mels: list[torch.Tensor] = []
        raw_lengths: list[int] = []
        token_lens: list[int] = []
        for one_audio in audio_list:
            mel = self._extract_mel(one_audio)
            raw_len = int(mel.shape[-1])
            num_tokens = self.conv3_downsample_len(raw_len)
            if raw_len <= 0 or num_tokens <= 0:
                raise ValueError("The audio is too short to be represented.")
            mels.append(mel)
            raw_lengths.append(raw_len)
            token_lens.append(num_tokens)

        # Step 4. Pad variable-length mel features into a batch tensor.
        if mels:
            max_length = max(raw_lengths)
            audio_batch = torch.zeros(
                (len(mels), self.mel_config["mel_dim"], max_length),
                dtype=torch.float32,
            )
            for index, mel in enumerate(mels):
                audio_batch[index, :, : mel.shape[-1]] = mel
            audio_data_seqlens = torch.tensor(raw_lengths, dtype=torch.long)
        else:
            audio_batch = None
            audio_data_seqlens = None

        # Step 5. Ensure each audio item has a placeholder span in the prompt.
        prompt_text = self._ensure_audio_placeholders(prompt_text, len(audio_list))
        input_ids = []
        cursor = 0

        # Step 6. Text-only path: tokenize and preserve placeholder spans.
        if not audio_list:
            for match in MOSS_AUDIO_SPAN_RE.finditer(prompt_text):
                prefix = prompt_text[cursor : match.start()]
                input_ids.extend(
                    self.tokenizer.encode(prefix, add_special_tokens=False)
                )
                input_ids.extend(
                    [self.audio_start_id, self.audio_token_id, self.audio_end_id]
                )
                cursor = match.end()
            suffix = prompt_text[cursor:]
            input_ids.extend(self.tokenizer.encode(suffix, add_special_tokens=False))
            data: dict[str, torch.Tensor] = {
                "input_ids": torch.tensor([input_ids], dtype=torch.long),
                "attention_mask": torch.ones((1, len(input_ids)), dtype=torch.long),
            }
            return BatchFeature(data=data, tensor_type=return_tensors)

        # Step 7. Audio path: expand each placeholder to its audio-token count.
        span_iter = iter(MOSS_AUDIO_SPAN_RE.finditer(prompt_text))
        for item_idx, _ in enumerate(audio_list):
            match = next(span_iter, None)
            if match is None:
                raise ValueError(
                    "Audio placeholder count mismatch: expected one "
                    f"{MOSS_AUDIO_PLACEHOLDER!r} span per audio item."
                )
            prefix = prompt_text[cursor : match.start()]
            input_ids.extend(self.tokenizer.encode(prefix, add_special_tokens=False))
            input_ids.append(self.audio_start_id)
            input_ids.extend(self.build_audio_placeholder_ids(token_lens[item_idx]))
            input_ids.append(self.audio_end_id)
            cursor = match.end()

        # Step 8. Reject extra placeholder spans after all audio items are used.
        suffix = prompt_text[cursor:]
        if MOSS_AUDIO_SPAN_RE.search(suffix):
            raise ValueError(
                "Audio placeholder count mismatch: found more placeholder spans "
                "than audio items."
            )
        input_ids.extend(self.tokenizer.encode(suffix, add_special_tokens=False))

        # Step 9. Return tokenizer output plus audio tensors for embed_multimodal.
        data = {
            "input_ids": torch.tensor([input_ids], dtype=torch.long),
            "attention_mask": torch.ones((1, len(input_ids)), dtype=torch.long),
        }
        if audio_batch is not None and audio_data_seqlens is not None:
            data["audio_data"] = audio_batch
            data["audio_data_seqlens"] = audio_data_seqlens
        return BatchFeature(data=data, tensor_type=return_tensors)

    def decode(self, *args: object, **kwargs: object) -> str:
        return self.tokenizer.decode(*args, **kwargs)

    def batch_decode(self, *args: object, **kwargs: object) -> list[str]:
        return self.tokenizer.batch_decode(*args, **kwargs)


class MossAudioProcessingInfo(BaseProcessingInfo):
    def get_hf_config(self) -> MossAudioConfig:
        config = self.ctx.get_hf_config()
        if isinstance(config, MossAudioConfig):
            return config
        return MossAudioConfig(
            audio_config=getattr(config, "audio_config", None),
            language_config=getattr(config, "language_config", None),
            adapter_hidden_size=getattr(config, "adapter_hidden_size", 8192),
            ignore_index=getattr(config, "ignore_index", -100),
            deepstack_num_inject_layers=getattr(
                config, "deepstack_num_inject_layers", None
            ),
        )

    def _get_processor_config_defaults(self) -> dict[str, object]:
        cached_defaults = getattr(self, "_processor_config_defaults", None)
        if cached_defaults is not None:
            return cached_defaults

        model_config = self.ctx.model_config
        for file_name in ("processor_config.json", "preprocessor_config.json"):
            config = get_hf_file_to_dict(
                file_name,
                model_config.model,
                model_config.revision,
            )
            defaults = _filter_moss_audio_processor_config(config)
            if defaults:
                self._processor_config_defaults = defaults
                return defaults

        defaults = {}
        self._processor_config_defaults = defaults
        return defaults

    @staticmethod
    def _get_processor_cache_key(kwargs: Mapping[str, object]) -> tuple[object, ...]:
        mel_config = _normalize_moss_audio_mel_config(
            kwargs.get("mel_config")
            if isinstance(kwargs.get("mel_config"), Mapping)
            else None
        )
        return (
            int(kwargs.get("audio_token_id", MOSS_AUDIO_TOKEN_ID)),
            int(kwargs.get("audio_start_id", MOSS_AUDIO_BOS_TOKEN_ID)),
            int(kwargs.get("audio_end_id", MOSS_AUDIO_EOS_TOKEN_ID)),
            bool(kwargs.get("enable_time_marker", False)),
            tuple(sorted(mel_config.items())),
        )

    def get_hf_processor(self, **kwargs: object) -> MossAudioProcessor:
        merged_kwargs = _merge_moss_audio_processor_configs(
            self._get_processor_config_defaults(),
            self.ctx.get_merged_mm_kwargs({}),
            kwargs,
        )
        mel_config = _normalize_moss_audio_mel_config(
            merged_kwargs.get("mel_config")
            if isinstance(merged_kwargs.get("mel_config"), Mapping)
            else None
        )
        processor_kwargs = {
            "audio_token_id": int(
                merged_kwargs.get("audio_token_id", MOSS_AUDIO_TOKEN_ID)
            ),
            "audio_start_id": int(
                merged_kwargs.get("audio_start_id", MOSS_AUDIO_BOS_TOKEN_ID)
            ),
            "audio_end_id": int(
                merged_kwargs.get("audio_end_id", MOSS_AUDIO_EOS_TOKEN_ID)
            ),
            "enable_time_marker": bool(merged_kwargs.get("enable_time_marker", False)),
            "mel_config": mel_config,
        }

        cache = getattr(self, "_hf_processor_cache", None)
        if cache is None:
            cache = {}
            self._hf_processor_cache = cache

        cache_key = self._get_processor_cache_key(processor_kwargs)
        processor = cache.get(cache_key)
        if processor is not None:
            return processor

        processor = MossAudioProcessor(
            self.get_tokenizer(),
            **processor_kwargs,
        )
        cache[cache_key] = processor
        return processor

    def get_feature_extractor(self, **kwargs: object) -> WhisperFeatureExtractor:
        return self.get_hf_processor(**kwargs).feature_extractor

    def get_supported_mm_limits(self) -> Mapping[str, int | None]:
        return {"audio": None}

    def get_data_parser(self) -> MultiModalDataParser:
        processor = self.get_hf_processor()
        return MossAudioMultiModalDataParser(
            target_sr=processor.mel_config["mel_sr"],
            target_channels=1,
            expected_hidden_size=self._get_expected_hidden_size(),
        )

    def get_mm_max_tokens_per_item(
        self,
        seq_len: int,
        mm_counts: Mapping[str, int],
    ) -> Mapping[str, int] | None:
        if mm_counts.get("audio", 0) <= 0:
            return {}
        processor = self.get_hf_processor()
        raw_mel_len = math.ceil(
            (processor.mel_config["mel_sr"] * DEFAULT_MAX_AUDIO_SECONDS)
            / processor.mel_config["mel_hop_length"]
        )
        return {"audio": MossAudioEncoder.compute_num_audio_tokens(raw_mel_len)}


class MossAudioDummyInputsBuilder(BaseDummyInputsBuilder[MossAudioProcessingInfo]):
    def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
        num_audios = mm_counts.get("audio", 0)
        return MOSS_AUDIO_PLACEHOLDER * num_audios

    def get_dummy_mm_data(
        self,
        seq_len: int,
        mm_counts: Mapping[str, int],
        mm_options: Mapping[str, BaseDummyOptions],
    ) -> MultiModalDataDict:
        num_audios = mm_counts.get("audio", 0)
        audio_overrides = mm_options.get("audio")
        return {
            "audio": self._get_dummy_audios(
                length=16000,
                num_audios=num_audios,
                overrides=audio_overrides,
            )
        }


class MossAudioMultiModalProcessor(BaseMultiModalProcessor[MossAudioProcessingInfo]):
    def _call_hf_processor(
        self,
        prompt: str,
        mm_data: Mapping[str, object],
        mm_kwargs: Mapping[str, object],
        tok_kwargs: Mapping[str, object],
    ) -> BatchFeature:
        mm_data = dict(mm_data)
        audios = mm_data.pop("audios", [])
        if audios:
            mm_data["audio"] = audios
        mm_kwargs = dict(mm_kwargs)
        processor_kwargs = _filter_moss_audio_processor_config(mm_kwargs)
        tok_kwargs = {
            key: value
            for key, value in tok_kwargs.items()
            if key not in MOSS_AUDIO_PROCESSOR_CONFIG_KEYS
        }
        return self.info.ctx.call_hf_processor(
            self.info.get_hf_processor(**processor_kwargs),
            dict(text=prompt, **mm_data),
            dict(**tok_kwargs),
        )

    def _get_mm_fields_config(
        self,
        hf_inputs: BatchFeature,
        hf_processor_mm_kwargs: Mapping[str, object],
    ) -> Mapping[str, MultiModalFieldConfig]:
        return _moss_audio_field_config(hf_inputs)

    def _get_prompt_updates(
        self,
        mm_items: MultiModalDataItems,
        hf_processor_mm_kwargs: Mapping[str, object],
        out_mm_kwargs: MultiModalKwargsItems,
    ) -> Sequence[PromptUpdate]:
        processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
        out_mm_data = out_mm_kwargs.get_data()
        audio_data_seqlens = out_mm_data.get("audio_data_seqlens")
        if audio_data_seqlens is None:
            audio_token_lens: list[int] = []
        else:
            if isinstance(audio_data_seqlens, torch.Tensor):
                lens = audio_data_seqlens.reshape(-1).tolist()
            else:
                lens = list(audio_data_seqlens)
            audio_token_lens = [
                MossAudioEncoder.compute_num_audio_tokens(int(length))
                for length in lens
            ]

        def get_replacement(
            item_idx: int,
            suffix_token_ids: list[int] | None = None,
        ) -> PromptUpdateDetails[list[int]]:
            num_tokens = audio_token_lens[item_idx]
            if num_tokens == 0:
                raise ValueError("The audio is too short to be represented.")
            audio_token_ids = processor.build_audio_placeholder_ids(num_tokens)
            suffix_token_ids = suffix_token_ids or []
            is_embed = torch.tensor(
                [token_id == processor.audio_token_id for token_id in audio_token_ids],
                dtype=torch.bool,
            )
            return PromptUpdateDetails(
                full=[
                    processor.audio_start_id,
                    *audio_token_ids,
                    processor.audio_end_id,
                    *suffix_token_ids,
                ],
                is_embed=lambda _tokenizer, _seq: torch.cat(
                    [
                        torch.tensor([False]),
                        is_embed,
                        torch.tensor([False]),
                        torch.zeros(len(suffix_token_ids), dtype=torch.bool),
                    ]
                ),
            )

        prompt_update_specs = [
            (
                [
                    processor.audio_start_id,
                    processor.audio_token_id,
                    processor.audio_end_id,
                ],
                [],
            )
        ]
        for suffix in ("", "\n"):
            tokenizer_target = processor.tokenizer.encode(
                MOSS_AUDIO_PLACEHOLDER + suffix,
                add_special_tokens=False,
            )
            suffix_token_ids = processor.tokenizer.encode(
                suffix,
                add_special_tokens=False,
            )
            if any(target == tokenizer_target for target, _ in prompt_update_specs):
                continue
            prompt_update_specs.append((tokenizer_target, suffix_token_ids))

        return [
            PromptReplacement(
                modality="audio",
                target=target,
                replacement=(
                    lambda item_idx, suffix_token_ids=suffix_token_ids: get_replacement(
                        item_idx,
                        suffix_token_ids,
                    )
                ),
            )
            for target, suffix_token_ids in prompt_update_specs
        ]


@MULTIMODAL_REGISTRY.register_processor(
    MossAudioMultiModalProcessor,
    info=MossAudioProcessingInfo,
    dummy_inputs=MossAudioDummyInputsBuilder,
)
class MossAudioModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):
    packed_modules_mapping = {
        "qkv_proj": [
            "q_proj",
            "k_proj",
            "v_proj",
        ],
        "gate_up_proj": [
            "gate_proj",
            "up_proj",
        ],
    }

    embedding_modules = {
        "embed_tokens": "input_embeddings",
        "lm_head": "output_embeddings",
    }

    hf_to_vllm_mapper = WeightsMapper(
        orig_to_new_prefix={
            "lm_head.": "language_model.lm_head.",
            "language_model.embed_tokens.": "language_model.model.embed_tokens.",
            "language_model.layers.": "language_model.model.layers.",
            "language_model.norm.": "language_model.model.norm.",
        }
    )

    def get_mm_mapping(self) -> MultiModelKeys:
        return MultiModelKeys.from_string_field(
            language_model="language_model.",
        )

    @classmethod
    def get_placeholder_str(cls, modality: str, i: int) -> str | None:
        if modality.startswith("audio"):
            return MOSS_AUDIO_PLACEHOLDER
        raise ValueError("Only audio modality is supported")

    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
        super().__init__()
        self.vllm_config = vllm_config
        config = vllm_config.model_config.hf_config
        if not isinstance(config, MossAudioConfig):
            config = MossAudioConfig(
                audio_config=getattr(config, "audio_config", None),
                language_config=getattr(config, "language_config", None),
                adapter_hidden_size=getattr(config, "adapter_hidden_size", 8192),
                ignore_index=getattr(config, "ignore_index", -100),
                deepstack_num_inject_layers=getattr(
                    config, "deepstack_num_inject_layers", None
                ),
            )
        self.config = config
        self.quant_config = vllm_config.quant_config
        self.multimodal_config = vllm_config.model_config.multimodal_config

        parallel_config = vllm_config.parallel_config
        tp_size = parallel_config.tensor_parallel_size
        if self.config.adapter_hidden_size % tp_size != 0:
            raise ValueError(
                "MOSS-Audio adapter_hidden_size must be divisible by tensor "
                f"parallel size. Got adapter_hidden_size="
                f"{self.config.adapter_hidden_size} and tensor_parallel_size="
                f"{tp_size}."
            )

        audio_config = MossAudioEncoderConfig.from_config(self.config.audio_config)
        if audio_config.encoder_attention_heads % tp_size != 0:
            raise ValueError(
                "MOSS-Audio encoder_attention_heads must be divisible by "
                "tensor parallel size. Got encoder_attention_heads="
                f"{audio_config.encoder_attention_heads} and "
                f"tensor_parallel_size={tp_size}."
            )
        language_config = self.config.language_config
        self.audio_token_id = MOSS_AUDIO_TOKEN_ID
        self.deepstack_input_embeds: IntermediateTensors | None = None

        with self._mark_tower_model(vllm_config, "audio"):
            self.audio_encoder = MossAudioEncoder(
                audio_config,
                quant_config=self.quant_config,
                prefix=maybe_prefix(prefix, "audio_encoder"),
            )
            self.audio_adapter = GatedMLP(
                input_size=audio_config.output_dim,
                hidden_size=self.config.adapter_hidden_size,
                output_size=language_config.hidden_size,
                quant_config=self.quant_config,
                prefix=maybe_prefix(prefix, "audio_adapter"),
            )

            deepstack_k = len(audio_config.deepstack_encoder_layer_indexes or [])
            if self.config.deepstack_num_inject_layers is not None:
                deepstack_k = min(
                    deepstack_k,
                    int(self.config.deepstack_num_inject_layers),
                )
            self.deepstack_audio_merger_list = nn.ModuleList(
                [
                    GatedMLP(
                        input_size=audio_config.output_dim,
                        hidden_size=self.config.adapter_hidden_size,
                        output_size=language_config.hidden_size,
                        quant_config=self.quant_config,
                        prefix=maybe_prefix(
                            prefix,
                            f"deepstack_audio_merger_list.{layer_idx}",
                        ),
                    )
                    for layer_idx in range(deepstack_k)
                ]
            )

        with self._mark_language_model(vllm_config):
            self.language_model = MossQwen3ForCausalLM(
                vllm_config=vllm_config.with_hf_config(
                    language_config, architectures=["Qwen3ForCausalLM"]
                ),
                prefix=maybe_prefix(prefix, "language_model"),
            )
            self.language_model.deepstack_inject_layer_indices = range(deepstack_k)
            self.language_model.model.deepstack_inject_layer_indices = range(
                deepstack_k
            )

        self.make_empty_intermediate_tensors = (
            self.language_model.make_empty_intermediate_tensors
        )

    @staticmethod
    def _validate_audio_batch_size(
        audio_batch_size: int, audio_data_seqlens: torch.Tensor
    ) -> None:
        if audio_batch_size != audio_data_seqlens.numel():
            raise ValueError(
                "audio_data batch size does not match audio_data_seqlens: "
                f"{audio_batch_size} != {audio_data_seqlens.numel()}."
            )

    @staticmethod
    def _pad_audio_data_list(
        audio_data: list[torch.Tensor],
        audio_data_seqlens: torch.Tensor,
    ) -> torch.Tensor:
        if len(audio_data) == 0:
            raise ValueError("audio_data list must not be empty.")
        MossAudioModel._validate_audio_batch_size(len(audio_data), audio_data_seqlens)

        # pad_sequence needs every item to share the same trailing feature
        # layout, so validate the mel-major audio tensors before transposing.
        first = audio_data[0]
        if not isinstance(first, torch.Tensor):
            raise TypeError("audio_data list items must be torch.Tensor.")
        if first.ndim != 2:
            raise ValueError("audio_data list items must have shape [mel_dim, time].")

        mel_dim = first.shape[0]
        dtype = first.dtype
        device = first.device
        for item in audio_data[1:]:
            if not isinstance(item, torch.Tensor):
                raise TypeError("audio_data list items must be torch.Tensor.")
            if item.ndim != 2:
                raise ValueError(
                    "audio_data list items must have shape [mel_dim, time]."
                )
            if item.shape[0] != mel_dim:
                raise ValueError("audio_data list items must have the same mel_dim.")
            if item.dtype != dtype:
                raise TypeError("audio_data list items must have the same dtype.")
            if item.device != device:
                raise ValueError("audio_data list items must be on the same device.")

        # Each item arrives as [mel_dim, time]. pad_sequence pads along dim 1
        # after converting to [time, mel_dim], then we restore [batch, mel, time].
        time_major = [item.transpose(0, 1) for item in audio_data]
        padded = torch.nn.utils.rnn.pad_sequence(time_major, batch_first=True)
        return padded.transpose(1, 2).contiguous()

    def _parse_and_validate_audio_input(
        self, **kwargs: object
    ) -> MossAudioAudioInputs | None:
        """Normalize and validate model-side audio kwargs.

        If audio_data is provided, this checks that audio_data_seqlens is also
        present, flattens sequence lengths to a long tensor, pads list inputs
        to [batch, mel_dim, time], validates batch-size/sequence-length
        agreement, and rejects empty, non-positive, or downsampled-zero audio
        lengths.
        """
        audio_data = kwargs.pop("audio_data", None)
        audio_data_seqlens = kwargs.pop("audio_data_seqlens", None)
        if audio_data is None:
            return None
        if audio_data_seqlens is None:
            raise ValueError(
                "audio_data_seqlens is required when audio_data is provided."
            )
        if not isinstance(audio_data_seqlens, torch.Tensor):
            audio_data_seqlens = torch.tensor(audio_data_seqlens, dtype=torch.long)
        audio_data_seqlens = audio_data_seqlens.to(dtype=torch.long).reshape(-1)

        if isinstance(audio_data, list):
            audio_data = self._pad_audio_data_list(audio_data, audio_data_seqlens)
        elif isinstance(audio_data, torch.Tensor):
            if audio_data.ndim == 3:
                self._validate_audio_batch_size(audio_data.shape[0], audio_data_seqlens)
        else:
            raise TypeError("audio_data must be a torch.Tensor or list[torch.Tensor].")

        audio_token_lens = MossAudioEncoder._compute_downsampled_length(
            audio_data_seqlens
        )
        if (
            audio_data_seqlens.numel() == 0
            or torch.any(audio_data_seqlens <= 0).item()
            or torch.any(audio_token_lens <= 0).item()
        ):
            raise ValueError("The audio is too short to be represented.")
        return MossAudioAudioInputs(
            audio_data=audio_data,
            audio_data_seqlens=audio_data_seqlens,
        )

    def _process_audio_input(
        self,
        audio_input: MossAudioAudioInputs,
    ) -> tuple[torch.Tensor, ...]:
        """Run the audio encoder and return one embedding tensor per audio.

        Example:
            audio_data=[2, 128, 1200], audio_data_seqlens=[800, 1200]
            -> returns (audio0_embeds, audio1_embeds), split by token length
            -> DeepStack packs each item as [main, layer0, ...] on dim -1
        """
        audio_data = audio_input["audio_data"]
        audio_data_seqlens = audio_input["audio_data_seqlens"]
        last_hidden_state, deepstack = self.audio_encoder(
            audio_data.to(self.audio_encoder.dtype),
            feature_lens=audio_data_seqlens,
            output_deepstack_hidden_states=len(self.deepstack_audio_merger_list) > 0,
        )
        audio_embeds = self.audio_adapter(last_hidden_state)
        audio_lengths = MossAudioEncoder._compute_downsampled_length(
            audio_data_seqlens.to(device=audio_embeds.device, dtype=torch.long)
        ).tolist()
        main_embeddings = tuple(audio_embeds.squeeze(0).split(audio_lengths, dim=0))

        deepstack_embeddings: list[tuple[torch.Tensor, ...]] = []
        if deepstack is not None:
            if len(deepstack) < len(self.deepstack_audio_merger_list):
                raise RuntimeError(
                    "DeepStack output count does not match configured audio "
                    "merger count."
                )
            for idx, hidden_states in enumerate(
                deepstack[: len(self.deepstack_audio_merger_list)]
            ):
                ds_embeds = self.deepstack_audio_merger_list[idx](hidden_states)
                deepstack_embeddings.append(
                    tuple(ds_embeds.squeeze(0).split(audio_lengths, dim=0))
                )

        if not deepstack_embeddings:
            return main_embeddings

        return tuple(
            torch.cat(
                [
                    main_embedding,
                    *(
                        layer_embeddings[item_idx]
                        for layer_embeddings in deepstack_embeddings
                    ),
                ],
                dim=-1,
            )
            for item_idx, main_embedding in enumerate(main_embeddings)
        )

    def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
        audio_input = self._parse_and_validate_audio_input(**kwargs)
        if audio_input is None:
            return ()
        return self._process_audio_input(audio_input)

    def _split_multimodal_embeddings(
        self,
        multimodal_embeddings: MultiModalEmbeddings,
        hidden_size: int,
    ) -> tuple[tuple[torch.Tensor, ...], tuple[tuple[torch.Tensor, ...], ...]]:
        """Unpack audio embeddings before merging them into token embeddings.

        embed_input_ids calls this on the output of embed_multimodal. Plain
        audio embeddings already have width hidden_size and are returned as the
        main embeddings for _merge_multimodal_embeddings. When DeepStack is
        enabled, _process_audio_input packs each audio item as
        [main, layer0, layer1, ...] along the last dimension so the standard
        multimodal path can carry a single embedding object. This method splits
        that packed layout back into main embeddings plus per-layer DeepStack
        embeddings, which _cache_deepstack_input_embeds scatters and forward
        passes into MossQwen3Model for layer injection.
        """
        if isinstance(multimodal_embeddings, torch.Tensor):
            embeddings = tuple(multimodal_embeddings.unbind(0))
        else:
            embeddings = tuple(multimodal_embeddings)

        if len(embeddings) == 0:
            return (), ()

        deepstack_count = len(self.deepstack_audio_merger_list)
        if all(embedding.shape[-1] == hidden_size for embedding in embeddings):
            return embeddings, ()

        packed_hidden_size = hidden_size * (deepstack_count + 1)
        if deepstack_count == 0 or any(
            embedding.shape[-1] != packed_hidden_size for embedding in embeddings
        ):
            got = [int(embedding.shape[-1]) for embedding in embeddings]
            raise ValueError(
                "MOSS-Audio multimodal embedding width mismatch: expected "
                f"{hidden_size} or {packed_hidden_size}, got {got}."
            )

        split_by_item = [
            torch.split(embedding, hidden_size, dim=-1) for embedding in embeddings
        ]
        main_embeddings = tuple(parts[0] for parts in split_by_item)
        deepstack_embeddings = tuple(
            tuple(parts[layer_idx + 1] for parts in split_by_item)
            for layer_idx in range(deepstack_count)
        )
        return main_embeddings, deepstack_embeddings

    def _cache_deepstack_input_embeds(
        self,
        inputs_embeds: torch.Tensor,
        deepstack_embeddings: tuple[tuple[torch.Tensor, ...], ...],
        is_multimodal: torch.Tensor,
    ) -> None:
        if len(deepstack_embeddings) == 0:
            self.deepstack_input_embeds = None
            return
        flat_by_layer = [
            torch.cat(layer_embeds, dim=0).to(
                device=inputs_embeds.device, dtype=inputs_embeds.dtype
            )
            for layer_embeds in deepstack_embeddings
        ]
        num_mm_tokens = int(is_multimodal.sum().item())
        if any(layer.shape[0] != num_mm_tokens for layer in flat_by_layer):
            got = [int(layer.shape[0]) for layer in flat_by_layer]
            raise ValueError(
                "DeepStack audio token count mismatch: "
                f"expected {num_mm_tokens}, got {got}."
            )
        data = {}
        for layer_idx, layer_embeds in enumerate(flat_by_layer):
            scattered = inputs_embeds.new_zeros(inputs_embeds.shape)
            scattered[is_multimodal] = layer_embeds
            data[f"deepstack_input_embeds_{layer_idx}"] = scattered
        self.deepstack_input_embeds = IntermediateTensors(data)

    def embed_input_ids(
        self,
        input_ids: torch.Tensor,
        multimodal_embeddings: MultiModalEmbeddings | None = None,
        *,
        is_multimodal: torch.Tensor | None = None,
    ) -> torch.Tensor:
        inputs_embeds = self._embed_text_input_ids(
            input_ids,
            self.language_model.embed_input_ids,
            is_multimodal=is_multimodal,
        )

        self.deepstack_input_embeds = None
        if multimodal_embeddings is None or len(multimodal_embeddings) == 0:
            return inputs_embeds
        is_multimodal = _require_is_multimodal(is_multimodal)
        multimodal_embeddings, deepstack_embeddings = self._split_multimodal_embeddings(
            multimodal_embeddings,
            hidden_size=int(inputs_embeds.shape[-1]),
        )

        inputs_embeds = _merge_multimodal_embeddings(
            inputs_embeds=inputs_embeds,
            multimodal_embeddings=multimodal_embeddings,
            is_multimodal=is_multimodal,
        )
        self._cache_deepstack_input_embeds(
            inputs_embeds,
            deepstack_embeddings,
            is_multimodal,
        )
        return inputs_embeds

    def forward(
        self,
        input_ids: torch.Tensor | None,
        positions: torch.Tensor,
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
        **kwargs: object,
    ) -> torch.Tensor | IntermediateTensors:
        if intermediate_tensors is None:
            deepstack_input_embeds = self.deepstack_input_embeds
        else:
            # Non-first PP ranks consume hidden states from intermediate_tensors.
            # The executor may still pass dummy inputs_embeds during profiling.
            inputs_embeds = None
            deepstack_input_embeds = intermediate_tensors
        hidden_states = self.language_model(
            input_ids,
            positions,
            intermediate_tensors=intermediate_tensors,
            inputs_embeds=inputs_embeds,
            deepstack_input_embeds=deepstack_input_embeds,
        )
        self.deepstack_input_embeds = None
        return hidden_states

    def compute_logits(
        self,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor | None:
        return self.language_model.compute_logits(hidden_states)

    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
        loader = AutoWeightsLoader(
            self,
            skip_prefixes=["audio_encoder.embed_positions"],
        )
        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
