# 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.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 @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: # Non-first PP ranks only receive DeepStack payloads for layers # at or after their local start layer. if layer_idx < self.model.start_layer: continue 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.", }, orig_to_new_stacked={ ".gate_proj": (".gate_up_proj", 0), ".up_proj": (".gate_up_proj", 1), }, ) 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)