# Copyright (c) 2026 LightSeek Foundation # # Permission is hereby granted, free of charge, to any person obtaining a copy # of this software and associated documentation files (the "Software"), to deal # in the Software without restriction, including without limitation the rights # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell # copies of the Software, and to permit persons to whom the Software is # furnished to do so, subject to the following conditions: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. """Shared Qwen3-ASR/Qwen3-Omni audio encoder. The gateway extracts log-Mel features with the Whisper-compatible Qwen3 frontend. Each multimodal item contains one ``[num_mel_bins, num_frames]`` tensor (a leading singleton batch dimension is also accepted) and optionally carries ``audio_feature_lengths`` or ``feature_attention_mask`` in ``model_specific_data``. This module packs those request-local tensors and implements the common Qwen3 audio tower used by both Qwen3-ASR and the Qwen3-Omni thinker. """ from __future__ import annotations from collections.abc import Iterable, Sequence from numbers import Integral from typing import Any import torch import torch.nn.functional as F from torch import nn from transformers.activations import ACT2FN from tokenspeed.runtime.distributed.mapping import Mapping from tokenspeed.runtime.layers.attention.mm_encoder_attention import ( MultimodalEncoderAttention, ) from tokenspeed.runtime.layers.linear import ( ColumnParallelLinear, ReplicatedLinear, RowParallelLinear, ) from tokenspeed.runtime.layers.quantization.base_config import QuantizationConfig from tokenspeed.runtime.model_loader.weight_utils import default_weight_loader from tokenspeed.runtime.multimodal.inputs import MultimodalDataItem from tokenspeed.runtime.utils import add_prefix def _cnn_output_lengths(input_lengths: torch.Tensor) -> torch.Tensor: """Lengths after the tower's three stride-2, padding-1 convolutions.""" lengths = input_lengths for _ in range(3): lengths = (lengths - 1) // 2 + 1 return lengths def qwen3_audio_output_lengths( input_lengths: torch.Tensor | Sequence[int] | int, *, n_window: int = 50, ) -> torch.Tensor: """Return language-token counts produced for log-Mel frame lengths. Qwen processes audio in chunks of ``2 * n_window`` frames. Convolution padding is applied independently to every chunk, so applying a single ``ceil(length / 8)`` to a long clip would under-count tokens at chunk boundaries. Args: input_lengths: Scalar or one-dimensional Mel-frame lengths. n_window: Half of the frontend chunk size from the audio config. Returns: A one-dimensional ``torch.long`` tensor with one token count per input. """ if n_window <= 0: raise ValueError(f"n_window must be positive, got {n_window}") if isinstance(input_lengths, Integral): lengths = torch.tensor([int(input_lengths)], dtype=torch.long) elif isinstance(input_lengths, torch.Tensor): lengths = input_lengths.to(dtype=torch.long) if lengths.ndim == 0: lengths = lengths.unsqueeze(0) else: lengths = torch.as_tensor(list(input_lengths), dtype=torch.long) if lengths.ndim != 1: raise ValueError( f"audio feature lengths must be one-dimensional, got {lengths.shape}" ) if torch.any(lengths < 0): raise ValueError("audio feature lengths cannot be negative") chunk_size = 2 * n_window full_chunks = lengths // chunk_size tail = lengths % chunk_size full_chunk_output = int( _cnn_output_lengths(torch.tensor([chunk_size], dtype=torch.long)).item() ) tail_output = _cnn_output_lengths(tail) tail_output = torch.where(tail == 0, torch.zeros_like(tail), tail_output) return full_chunks * full_chunk_output + tail_output def _one_item_feature_length(item: MultimodalDataItem, max_frames: int) -> int: explicit_length = item.model_specific_data.get("audio_feature_lengths") if explicit_length is not None: if isinstance(explicit_length, torch.Tensor): if explicit_length.numel() != 1: raise ValueError( "one audio item must carry exactly one audio_feature_lengths " f"value, got shape {explicit_length.shape}" ) length = int(explicit_length.reshape(-1)[0].item()) elif isinstance(explicit_length, Integral): length = int(explicit_length) else: raise TypeError( "audio_feature_lengths must be an integer or tensor, got " f"{type(explicit_length).__name__}" ) else: attention_mask = item.model_specific_data.get("feature_attention_mask") if attention_mask is None: length = max_frames else: attention_mask = torch.as_tensor(attention_mask) if attention_mask.ndim == 2 and attention_mask.shape[0] == 1: attention_mask = attention_mask.squeeze(0) if attention_mask.ndim != 1: raise ValueError( "one audio item must carry a one-dimensional " f"feature_attention_mask, got {attention_mask.shape}" ) if attention_mask.numel() != max_frames: raise ValueError( "feature_attention_mask length does not match the audio " f"feature: {attention_mask.numel()} != {max_frames}" ) # Whisper masks are right-padded. Requiring a contiguous prefix # prevents silently packing frames from a malformed sparse mask. mask = attention_mask.to(dtype=torch.bool) length = int(mask.sum().item()) expected = torch.arange(max_frames, device=mask.device) < length if not torch.equal(mask, expected): raise ValueError("feature_attention_mask must be right-padded") if length <= 0 or length > max_frames: raise ValueError( f"audio feature length must be in [1, {max_frames}], got {length}" ) return length def pack_qwen3_audio_features( items: Sequence[MultimodalDataItem], *, num_mel_bins: int, device: torch.device, dtype: torch.dtype, ) -> tuple[torch.Tensor, torch.Tensor]: """Pack itemized log-Mel tensors for :class:`Qwen3AudioEncoder`. Returns a feature tensor shaped ``[num_mel_bins, sum(valid_frames)]`` and a length tensor shaped ``[num_items]``. Padding is removed before packing. """ if not items: raise ValueError("at least one audio item is required") features: list[torch.Tensor] = [] lengths: list[int] = [] for item in items: feature = item.feature if not isinstance(feature, torch.Tensor): raise TypeError( "audio item feature must be materialized as a torch.Tensor, got " f"{type(feature).__name__}" ) if feature.ndim == 3 and feature.shape[0] == 1: feature = feature.squeeze(0) if feature.ndim != 2: raise ValueError( "audio feature must have shape [mel, frames] or [1, mel, frames], " f"got {feature.shape}" ) if feature.shape[0] != num_mel_bins: raise ValueError( f"expected {num_mel_bins} Mel bins, got feature shape {feature.shape}" ) length = _one_item_feature_length(item, feature.shape[1]) features.append(feature[:, :length].to(device=device, dtype=dtype)) lengths.append(length) return ( torch.cat(features, dim=1).contiguous(), torch.tensor(lengths, dtype=torch.long, device=device), ) class SinusoidsPositionEmbedding(nn.Module): """Non-trainable absolute position embedding used by the audio tower.""" def __init__(self, length: int, channels: int, max_timescale: int = 10000): super().__init__() if channels % 2: raise ValueError("sinusoidal position embeddings need even channels") log_increment = torch.log(torch.tensor(float(max_timescale))) / ( channels // 2 - 1 ) inv_timescales = torch.exp(-log_increment * torch.arange(channels // 2)) positions = torch.arange(length).unsqueeze(1) * inv_timescales.unsqueeze(0) self.register_buffer( "positional_embedding", torch.cat([positions.sin(), positions.cos()], dim=1), persistent=False, ) def forward(self, sequence_length: int) -> torch.Tensor: return self.positional_embedding[:sequence_length] class Qwen3AudioEncoderLayer(nn.Module): def __init__( self, config: Any, mapping: Mapping, quant_config: QuantizationConfig | None = None, prefix: str = "", mm_attention_backend: str | None = None, ) -> None: super().__init__() self.self_attn = MultimodalEncoderAttention( embed_dim=config.d_model, num_heads=config.encoder_attention_heads, mapping=mapping, qkv_bias=True, proj_bias=True, quant_config=quant_config, prefix=add_prefix("self_attn", prefix), mm_attention_backend=mm_attention_backend, ) self.self_attn_layer_norm = nn.LayerNorm(config.d_model) self.final_layer_norm = nn.LayerNorm(config.d_model) self.activation_fn = ACT2FN[config.activation_function] vision = mapping.vision self.fc1 = ColumnParallelLinear( config.d_model, config.encoder_ffn_dim, bias=True, quant_config=quant_config, prefix=add_prefix("fc1", prefix), tp_rank=vision.tp_rank, tp_size=vision.tp_size, tp_group=vision.tp_group, ) self.fc2 = RowParallelLinear( config.encoder_ffn_dim, config.d_model, bias=True, quant_config=quant_config, prefix=add_prefix("fc2", prefix), tp_rank=vision.tp_rank, tp_size=vision.tp_size, tp_group=vision.tp_group, reduce_results=True, ) def forward( self, hidden_states: torch.Tensor, cu_seqlens: torch.Tensor, max_seqlen: int, ) -> torch.Tensor: residual = hidden_states normed = self.self_attn_layer_norm(hidden_states) attended = self.self_attn( normed.unsqueeze(0), cu_seqlens=cu_seqlens, max_seqlen=max_seqlen, ).squeeze(0) hidden_states = residual + attended residual = hidden_states hidden_states = self.final_layer_norm(hidden_states) hidden_states, _ = self.fc1(hidden_states) hidden_states = self.activation_fn(hidden_states) hidden_states, _ = self.fc2(hidden_states) hidden_states = residual + hidden_states if hidden_states.dtype == torch.float16: limit = torch.finfo(hidden_states.dtype).max - 1000 hidden_states = hidden_states.clamp(min=-limit, max=limit) return hidden_states class Qwen3AudioEncoder(nn.Module): """Parameterizable Qwen audio tower shared by ASR and Omni thinker.""" @staticmethod def _normalize_attention_backend(name: str | None) -> str | None: # The cuDNN path consumes vision-specific sequence metadata. Audio can # still use the platform default while an Omni vision tower uses cuDNN. return None if name == "flashinfer_cudnn" else name def __init__( self, config: Any, mapping: Mapping, quant_config: QuantizationConfig | None = None, prefix: str = "", mm_attention_backend: str | None = None, ) -> None: super().__init__() self.config = config self.num_mel_bins = int(config.num_mel_bins) self.n_window = int(config.n_window) self.n_window_infer = int(config.n_window_infer) self.conv_chunksize = int(config.conv_chunksize) mm_attention_backend = self._normalize_attention_backend(mm_attention_backend) self.positional_embedding = SinusoidsPositionEmbedding( int(config.max_source_positions), int(config.d_model) ) self.conv2d1 = nn.Conv2d( 1, config.downsample_hidden_size, kernel_size=3, stride=2, padding=1 ) self.conv2d2 = nn.Conv2d( config.downsample_hidden_size, config.downsample_hidden_size, kernel_size=3, stride=2, padding=1, ) self.conv2d3 = nn.Conv2d( config.downsample_hidden_size, config.downsample_hidden_size, kernel_size=3, stride=2, padding=1, ) conv_out_dim = config.downsample_hidden_size * ( (((config.num_mel_bins + 1) // 2 + 1) // 2 + 1) // 2 ) self.conv_out = ReplicatedLinear( conv_out_dim, config.d_model, bias=False, quant_config=quant_config, prefix=add_prefix("conv_out", prefix), ) self.layers = nn.ModuleList( [ Qwen3AudioEncoderLayer( config, mapping, quant_config=quant_config, prefix=add_prefix(f"layers.{index}", prefix), mm_attention_backend=mm_attention_backend, ) for index in range(config.encoder_layers) ] ) self.ln_post = nn.LayerNorm(config.d_model) self.proj1 = ReplicatedLinear( config.d_model, config.d_model, bias=True, quant_config=quant_config, prefix=add_prefix("proj1", prefix), ) self.act = ACT2FN[config.activation_function] self.proj2 = ReplicatedLinear( config.d_model, config.output_dim, bias=True, quant_config=quant_config, prefix=add_prefix("proj2", prefix), ) @property def dtype(self) -> torch.dtype: return self.conv2d1.weight.dtype @property def device(self) -> torch.device: return self.conv2d1.weight.device def encode(self, items: Sequence[MultimodalDataItem]) -> torch.Tensor: input_features, feature_lengths = pack_qwen3_audio_features( items, num_mel_bins=self.num_mel_bins, device=self.device, dtype=self.dtype, ) return self(input_features, feature_lengths) def forward( self, input_features: torch.Tensor, feature_lengths: torch.Tensor, ) -> torch.Tensor: """Encode packed features into concatenated LM-space audio tokens. Args: input_features: Packed log-Mel features shaped ``[num_mel_bins, sum(feature_lengths)]``. feature_lengths: Valid frame count for every original audio item. Returns: Audio embeddings shaped ``[sum(output_lengths), output_dim]`` in the same item order as ``feature_lengths``. """ input_features = input_features.to(device=self.device, dtype=self.dtype) feature_lengths = feature_lengths.to(device=self.device, dtype=torch.long) if input_features.ndim != 2 or input_features.shape[0] != self.num_mel_bins: raise ValueError( f"expected packed features [{self.num_mel_bins}, frames], got " f"{input_features.shape}" ) if feature_lengths.ndim != 1 or feature_lengths.numel() == 0: raise ValueError("feature_lengths must be a non-empty vector") if torch.any(feature_lengths <= 0): raise ValueError("feature_lengths must be positive") if int(feature_lengths.sum().item()) != input_features.shape[1]: raise ValueError( "packed feature width does not match feature_lengths: " f"{input_features.shape[1]} != {int(feature_lengths.sum().item())}" ) chunk_size = self.n_window * 2 chunk_lengths_list: list[int] = [] for length in feature_lengths.tolist(): full_chunks, tail = divmod(int(length), chunk_size) chunk_lengths_list.extend([chunk_size] * full_chunks) if tail: chunk_lengths_list.append(tail) chunk_lengths = torch.tensor( chunk_lengths_list, dtype=torch.long, device=self.device ) chunks = input_features.transpose(0, 1).split(chunk_lengths_list, dim=0) padded_features = nn.utils.rnn.pad_sequence(chunks, batch_first=True).transpose( 1, 2 ) chunk_output_lengths = _cnn_output_lengths(chunk_lengths) max_chunk_output = int(chunk_output_lengths.max().item()) output_mask = torch.arange(max_chunk_output, device=self.device).unsqueeze( 0 ) < chunk_output_lengths.unsqueeze(1) padded_features = padded_features.unsqueeze(1) encoded_chunks: list[torch.Tensor] = [] for conv_input in padded_features.split(self.conv_chunksize, dim=0): hidden = F.gelu(self.conv2d1(conv_input)) hidden = F.gelu(self.conv2d2(hidden)) hidden = F.gelu(self.conv2d3(hidden)) encoded_chunks.append(hidden) hidden = torch.cat(encoded_chunks, dim=0) batch, channels, frequency, time = hidden.shape hidden, _ = self.conv_out( hidden.permute(0, 3, 1, 2) .contiguous() .view(batch, time, channels * frequency) ) if hidden.shape[1] > self.positional_embedding.positional_embedding.shape[0]: raise ValueError( "audio chunk exceeds max_source_positions after convolution: " f"{hidden.shape[1]}" ) hidden = hidden + self.positional_embedding(hidden.shape[1]).to(hidden.dtype) hidden_states = hidden[output_mask] output_lengths = qwen3_audio_output_lengths( feature_lengths, n_window=self.n_window ).to(device=self.device) chunks_per_attention_window = max(1, self.n_window_infer // chunk_size) attention_window = max_chunk_output * chunks_per_attention_window attention_lengths: list[int] = [] for length in output_lengths.tolist(): full_windows, tail = divmod(int(length), attention_window) attention_lengths.extend([attention_window] * full_windows) if tail: attention_lengths.append(tail) cu_seqlens = torch.tensor( [0, *attention_lengths], dtype=torch.int32, device=self.device ).cumsum(0, dtype=torch.int32) max_seqlen = max(attention_lengths) for layer in self.layers: hidden_states = layer(hidden_states, cu_seqlens, max_seqlen) hidden_states = self.ln_post(hidden_states) hidden_states, _ = self.proj1(hidden_states) hidden_states = self.act(hidden_states) hidden_states, _ = self.proj2(hidden_states) return hidden_states @staticmethod def map_weight_name(name: str) -> tuple[str, str | None]: """Map an HF audio-tower name to ``(TokenSpeed name, shard id)``.""" marker = "audio_tower." if marker in name: name = name.split(marker, 1)[1] name = name.replace("self_attn.out_proj.", "self_attn.proj.") for checkpoint_name, shard_id in ( ("q_proj", "q"), ("k_proj", "k"), ("v_proj", "v"), ): needle = f"self_attn.{checkpoint_name}." if needle in name: return name.replace(needle, "self_attn.qkv_proj."), shard_id return name, None def load_weight(self, name: str, loaded_weight: torch.Tensor) -> str | None: """Load one audio-tower checkpoint tensor. ``name`` may be relative to the tower or retain a ``thinker.audio_tower``/``audio_tower`` prefix. Q/K/V checkpoint projections are fused into TokenSpeed's ``qkv_proj`` parameter and the HF ``out_proj`` name is mapped to ``MultimodalEncoderAttention.proj``. """ name, shard_id = self.map_weight_name(name) params = dict(self.named_parameters(remove_duplicate=False)) if name not in params: return None param = params[name] if shard_id is not None: param.weight_loader(param, loaded_weight, shard_id) else: loader = getattr(param, "weight_loader", default_weight_loader) loader(param, loaded_weight) return name def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loaded: set[str] = set() for name, tensor in weights: loaded_name = self.load_weight(name, tensor) if loaded_name is not None: loaded.add(loaded_name) return loaded __all__ = [ "Qwen3AudioEncoder", "Qwen3AudioEncoderLayer", "pack_qwen3_audio_features", "qwen3_audio_output_lengths", ]