diff --git a/src/audio/gemma3n/attention.rs b/src/audio/gemma3n/attention.rs new file mode 100644 index 00000000..f7529633 --- /dev/null +++ b/src/audio/gemma3n/attention.rs @@ -0,0 +1,332 @@ +// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use super::checked_unified_linear; +use super::config::Gemma3nAudioConfig; +use mlxcel_core::layers::UnifiedLinear; +use mlxcel_core::weights::WeightMap; +use mlxcel_core::{MlxArray, UniquePtr}; + +fn copy_weight(weights: &WeightMap, key: &str) -> Result, String> { + weights + .get(key) + .map(|weight| mlxcel_core::copy(weight)) + .ok_or_else(|| format!("Gemma3n audio weight not found: {key}")) +} + +pub(crate) struct Gemma3nAudioAttention { + num_heads: usize, + head_dim: usize, + chunk_size: usize, + max_past_horizon: usize, + max_future_horizon: usize, + context_size: usize, + softcap: f32, + q_scale: f32, + q_proj: UnifiedLinear, + k_proj: UnifiedLinear, + v_proj: UnifiedLinear, + per_dim_scale: UniquePtr, + pos_proj: UnifiedLinear, + inv_timescales: UniquePtr, +} + +impl Gemma3nAudioAttention { + pub(crate) fn from_weights( + weights: &WeightMap, + prefix: &str, + config: &Gemma3nAudioConfig, + group_size: i32, + bits: i32, + ) -> Result { + let head_dim = config.head_dim(); + let num_timescales = config.hidden_size / 2; + let increment = 10_000.0f32.ln() / (num_timescales as f32 - 1.0).max(1.0); + let inv_timescales: Vec = (0..num_timescales) + .map(|index| (-increment * index as f32).exp()) + .collect(); + let per_dim_scale = copy_weight(weights, &format!("{prefix}.per_dim_scale"))?; + let scale_shape = mlxcel_core::array_shape(&per_dim_scale); + if scale_shape != [head_dim as i32] { + return Err(format!( + "{prefix}.per_dim_scale has shape {scale_shape:?}; expected [{head_dim}]" + )); + } + + Ok(Self { + num_heads: config.conf_num_attention_heads, + head_dim, + chunk_size: config.conf_attention_chunk_size, + max_past_horizon: config.max_past_horizon(), + max_future_horizon: config.conf_attention_context_right, + context_size: config.context_size(), + softcap: config.conf_attention_logit_cap, + // head_dim^-0.5 / softplus(0), exactly as the official model. + q_scale: (head_dim as f32).powf(-0.5) / 2.0f32.ln(), + q_proj: checked_unified_linear( + weights, + &format!("{prefix}.q_proj"), + config.hidden_size, + config.hidden_size, + group_size, + bits, + )?, + k_proj: checked_unified_linear( + weights, + &format!("{prefix}.k_proj"), + config.hidden_size, + config.hidden_size, + group_size, + bits, + )?, + v_proj: checked_unified_linear( + weights, + &format!("{prefix}.v_proj"), + config.hidden_size, + config.hidden_size, + group_size, + bits, + )?, + per_dim_scale, + pos_proj: checked_unified_linear( + weights, + &format!("{prefix}.relative_position_embedding.pos_proj"), + config.hidden_size, + config.hidden_size, + group_size, + bits, + )?, + inv_timescales: mlxcel_core::from_slice_f32( + &inv_timescales, + &[1, 1, num_timescales as i32], + ), + }) + } + + fn pad_dim1(&self, x: &MlxArray, left: i32, right: i32) -> UniquePtr { + let mut width = vec![0; mlxcel_core::array_ndim(x) * 2]; + width[2] = left; + width[3] = right; + mlxcel_core::pad(x, &width, 0.0) + } + + fn convert_to_blocks(&self, x: &MlxArray) -> UniquePtr { + let shape = mlxcel_core::array_shape(x); + let batch = shape[0]; + let time = shape[1]; + let blocks = (time + self.chunk_size as i32 - 1) / self.chunk_size as i32; + let padding = blocks * self.chunk_size as i32 - time; + let x = if padding > 0 { + self.pad_dim1(x, 0, padding) + } else { + mlxcel_core::copy(x) + }; + let mut result_shape = vec![batch, blocks, self.chunk_size as i32]; + result_shape.extend_from_slice(&shape[2..]); + mlxcel_core::reshape(&x, &result_shape) + } + + fn extract_block_context(&self, x: &MlxArray) -> UniquePtr { + let padded = self.pad_dim1( + x, + self.max_past_horizon as i32, + (self.max_future_horizon + self.chunk_size - 1) as i32, + ); + let shape = mlxcel_core::array_shape(&padded); + let context = self.context_size as i32; + let chunk = self.chunk_size as i32; + let blocks = (shape[1] - context) / chunk + 1; + let mut indices = Vec::with_capacity((blocks * context) as usize); + for block in 0..blocks { + for offset in 0..context { + indices.push(block * chunk + offset); + } + } + let indices = mlxcel_core::from_slice_i32(&indices, &[blocks * context]); + let gathered = mlxcel_core::take(&padded, &indices, 1); + let mut result_shape = vec![shape[0], blocks, context]; + result_shape.extend_from_slice(&shape[2..]); + mlxcel_core::reshape(&gathered, &result_shape) + } + + fn timing_signal(&self, positions: &[i32], dtype: i32) -> UniquePtr { + let positions: Vec = positions.iter().map(|position| *position as f32).collect(); + let positions = mlxcel_core::from_slice_f32(&positions, &[1, positions.len() as i32, 1]); + let scaled = mlxcel_core::multiply(&positions, &self.inv_timescales); + let signal = + mlxcel_core::concatenate(&mlxcel_core::sin(&scaled), &mlxcel_core::cos(&scaled), -1); + mlxcel_core::astype(&signal, dtype) + } + + fn relative_logits(&self, queries: &MlxArray, keys: &MlxArray) -> UniquePtr { + let q_shape = mlxcel_core::array_shape(queries); + let k_shape = mlxcel_core::array_shape(keys); + let (batch, blocks, block_size, heads, head_dim) = + (q_shape[0], q_shape[1], q_shape[2], q_shape[3], q_shape[4]); + let context = k_shape[2]; + let positions: Vec = (0..=self.max_past_horizon + self.max_future_horizon) + .rev() + .map(|value| value as i32 - self.max_future_horizon as i32) + .collect(); + let span = positions.len() as i32; + let signal = self.timing_signal(&positions, mlxcel_core::array_dtype(queries)); + let signal = self.pos_proj.forward(&signal); + let signal = mlxcel_core::reshape(&signal, &[span, heads, head_dim]); + + let q = mlxcel_core::transpose_axes(queries, &[0, 3, 1, 2, 4]); + let k = mlxcel_core::transpose_axes(keys, &[0, 3, 1, 4, 2]); + let content = mlxcel_core::matmul(&q, &k); + let signal = mlxcel_core::transpose_axes(&signal, &[1, 2, 0]); + let q_flat = mlxcel_core::reshape(&q, &[batch, heads, blocks * block_size, head_dim]); + let position = mlxcel_core::matmul(&q_flat, &signal); + let position = mlxcel_core::reshape(&position, &[batch, heads, blocks, block_size, span]); + let position = relative_shift(&position, batch, heads, blocks, block_size, context, span); + mlxcel_core::add(&content, &position) + } + + pub(crate) fn forward( + &self, + hidden_states: &MlxArray, + invalid_mask: &MlxArray, + causal_valid_mask: &MlxArray, + ) -> UniquePtr { + let shape = mlxcel_core::array_shape(hidden_states); + let (batch, time) = (shape[0], shape[1]); + let heads = self.num_heads as i32; + let head_dim = self.head_dim as i32; + let qkv_shape = [batch, time, heads, head_dim]; + + let query = mlxcel_core::reshape(&self.q_proj.forward(hidden_states), &qkv_shape); + let key = mlxcel_core::reshape(&self.k_proj.forward(hidden_states), &qkv_shape); + let value = mlxcel_core::reshape(&self.v_proj.forward(hidden_states), &qkv_shape); + let scale = + mlxcel_core::multiply_scalar(&mlxcel_core::softplus(&self.per_dim_scale), self.q_scale); + let query = mlxcel_core::multiply(&query, &scale); + + let query_blocks = self.convert_to_blocks(&query); + let key_blocks = self.extract_block_context(&key); + let value_blocks = self.extract_block_context(&value); + let blocks = mlxcel_core::array_shape(&query_blocks)[1]; + + let valid = mlxcel_core::logical_not(invalid_mask); + let valid = self.extract_block_context(&valid); + let valid = mlxcel_core::reshape(&valid, &[batch, 1, blocks, 1, self.context_size as i32]); + let causal = mlxcel_core::reshape( + causal_valid_mask, + &[1, 1, 1, self.chunk_size as i32, self.context_size as i32], + ); + let condition = mlxcel_core::logical_and(&valid, &causal); + + let logits = self.relative_logits(&query_blocks, &key_blocks); + let logits = softcap_logits(&logits, self.softcap); + let logits_dtype = mlxcel_core::array_dtype(&logits); + let minimum = mlxcel_core::full_f32(&[1], dtype_min(logits_dtype), logits_dtype); + let logits = mlxcel_core::where_cond(&condition, &logits, &minimum); + let probabilities = mlxcel_core::softmax( + &mlxcel_core::astype(&logits, mlxcel_core::dtype::FLOAT32), + -1, + ); + let probabilities = + mlxcel_core::astype(&probabilities, mlxcel_core::array_dtype(&value_blocks)); + let operands: [*const MlxArray; 2] = [ + probabilities.as_ref().unwrap() as *const MlxArray, + value_blocks.as_ref().unwrap() as *const MlxArray, + ]; + let context = unsafe { mlxcel_core::einsum("bnuwc,bucnh->buwnh", &operands) }; + let context = mlxcel_core::reshape( + &context, + &[batch, blocks * self.chunk_size as i32, heads, head_dim], + ); + if blocks * self.chunk_size as i32 > time { + mlxcel_core::slice(&context, &[0, 0, 0, 0], &[batch, time, heads, head_dim]) + } else { + context + } + } +} + +fn dtype_min(dtype: i32) -> f32 { + match dtype { + d if d == mlxcel_core::dtype::FLOAT16 => -65504.0, + d if d == mlxcel_core::dtype::BFLOAT16 => -3.38e38, + _ => f32::MIN, + } +} + +#[allow(clippy::too_many_arguments)] +fn relative_shift( + value: &MlxArray, + batch: i32, + heads: i32, + blocks: i32, + block_size: i32, + context: i32, + span: i32, +) -> UniquePtr { + let mut width = vec![0; mlxcel_core::array_ndim(value) * 2]; + let last = width.len() - 1; + width[last] = context + 1 - span; + let value = mlxcel_core::pad(value, &width, 0.0); + let value = mlxcel_core::reshape(&value, &[batch, heads, blocks, block_size * (context + 1)]); + let value = mlxcel_core::slice( + &value, + &[0, 0, 0, 0], + &[batch, heads, blocks, block_size * context], + ); + mlxcel_core::reshape(&value, &[batch, heads, blocks, block_size, context]) +} + +fn softcap_logits(logits: &MlxArray, cap: f32) -> UniquePtr { + let scaled = mlxcel_core::multiply_scalar(logits, 1.0 / cap); + mlxcel_core::multiply_scalar(&mlxcel_core::tanh(&scaled), cap) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn values(array: &MlxArray) -> Vec { + mlxcel_core::eval(array); + mlxcel_core::array_to_raw_bytes(array) + .chunks_exact(4) + .map(|bytes| f32::from_ne_bytes(bytes.try_into().unwrap())) + .collect() + } + + #[test] + fn relative_shift_matches_pinned_jax_layout() { + let input = mlxcel_core::from_slice_f32(&[1.0, 2.0, 3.0, 4.0], &[1, 1, 1, 2, 2]); + let shifted = relative_shift(&input, 1, 1, 1, 2, 3, 2); + assert_eq!(mlxcel_core::array_shape(&shifted), vec![1, 1, 1, 2, 3]); + assert_eq!(values(&shifted), vec![1.0, 2.0, 0.0, 0.0, 3.0, 4.0]); + } + + #[test] + fn softcap_is_not_an_unbounded_attention_substitute() { + let logits = mlxcel_core::from_slice_f32(&[-100.0, 0.0, 100.0], &[3]); + let capped = values(&softcap_logits(&logits, 5.0)); + assert!(capped[0] >= -5.0 && capped[0] < -4.99); + assert_eq!(capped[1], 0.0); + assert!(capped[2] <= 5.0 && capped[2] > 4.99); + } + + #[test] + fn attention_mask_minimum_remains_finite_in_checkpoint_dtypes() { + for dtype in [mlxcel_core::dtype::FLOAT16, mlxcel_core::dtype::BFLOAT16] { + let minimum = mlxcel_core::full_f32(&[1], dtype_min(dtype), dtype); + let minimum = mlxcel_core::astype(&minimum, mlxcel_core::dtype::FLOAT32); + assert!(mlxcel_core::item_f32(&minimum).is_finite()); + } + } +} diff --git a/src/audio/gemma3n/config.rs b/src/audio/gemma3n/config.rs new file mode 100644 index 00000000..2228ae32 --- /dev/null +++ b/src/audio/gemma3n/config.rs @@ -0,0 +1,209 @@ +// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use serde::Deserialize; + +fn hidden_size() -> usize { + 1536 +} +fn input_feat_size() -> usize { + 128 +} +fn vocab_size() -> usize { + 128 +} +fn vocab_offset() -> i32 { + 262_272 +} +fn rms_norm_eps() -> f32 { + 1e-6 +} +fn gradient_clipping() -> f32 { + 1e10 +} +fn chunk_size() -> usize { + 12 +} +fn context_left() -> usize { + 13 +} +fn attention_logit_cap() -> f32 { + 50.0 +} +fn attention_heads() -> usize { + 8 +} +fn hidden_layers() -> usize { + 12 +} +fn conv_kernel_size() -> usize { + 5 +} +fn reduction_factor() -> usize { + 4 +} +fn residual_weight() -> f32 { + 0.5 +} +fn conv_channels() -> Vec { + vec![128, 32] +} +fn group_norm_eps() -> f32 { + 1e-3 +} +fn conv_kernel_sizes() -> Vec<[usize; 2]> { + vec![[3, 3], [3, 3]] +} +fn conv_stride_sizes() -> Vec<[usize; 2]> { + vec![[2, 2], [2, 2]] +} + +/// Official `Gemma3nAudioConfig` fields from Transformers commit +/// `181beb3ba4c47098ed8cbc97ee250d1d45ae0107`. +#[derive(Debug, Clone, Deserialize)] +#[serde(default)] +pub struct Gemma3nAudioConfig { + pub vocab_size: usize, + pub vocab_offset: i32, + pub input_feat_size: usize, + pub hidden_size: usize, + pub rms_norm_eps: f32, + pub gradient_clipping: f32, + pub conf_attention_chunk_size: usize, + pub conf_attention_context_left: usize, + pub conf_attention_context_right: usize, + pub conf_attention_logit_cap: f32, + pub conf_num_attention_heads: usize, + pub conf_num_hidden_layers: usize, + pub conf_conv_kernel_size: usize, + pub conf_reduction_factor: usize, + pub conf_residual_weight: f32, + pub sscp_conv_channel_size: Vec, + pub sscp_conv_group_norm_eps: f32, + pub sscp_conv_kernel_size: Vec<[usize; 2]>, + pub sscp_conv_stride_size: Vec<[usize; 2]>, +} + +impl Default for Gemma3nAudioConfig { + fn default() -> Self { + Self { + vocab_size: vocab_size(), + vocab_offset: vocab_offset(), + input_feat_size: input_feat_size(), + hidden_size: hidden_size(), + rms_norm_eps: rms_norm_eps(), + gradient_clipping: gradient_clipping(), + conf_attention_chunk_size: chunk_size(), + conf_attention_context_left: context_left(), + conf_attention_context_right: 0, + conf_attention_logit_cap: attention_logit_cap(), + conf_num_attention_heads: attention_heads(), + conf_num_hidden_layers: hidden_layers(), + conf_conv_kernel_size: conv_kernel_size(), + conf_reduction_factor: reduction_factor(), + conf_residual_weight: residual_weight(), + sscp_conv_channel_size: conv_channels(), + sscp_conv_group_norm_eps: group_norm_eps(), + sscp_conv_kernel_size: conv_kernel_sizes(), + sscp_conv_stride_size: conv_stride_sizes(), + } + } +} + +impl Gemma3nAudioConfig { + pub fn validate(&self) -> Result<(), String> { + if self.hidden_size == 0 + || self.conf_num_attention_heads == 0 + || !self + .hidden_size + .is_multiple_of(self.conf_num_attention_heads) + { + return Err("Gemma3n audio hidden size must be divisible by the head count".into()); + } + if self.conf_attention_chunk_size == 0 + || self.conf_reduction_factor == 0 + || self.conf_conv_kernel_size == 0 + { + return Err( + "Gemma3n audio chunk, reduction, and convolution sizes must be positive".into(), + ); + } + if self.sscp_conv_channel_size.len() != 2 + || self.sscp_conv_kernel_size.len() != 2 + || self.sscp_conv_stride_size.len() != 2 + { + return Err("Gemma3n audio SSCP requires exactly two convolution stages".into()); + } + if self.input_feat_size == 0 || self.vocab_size == 0 { + return Err("Gemma3n audio feature and vocabulary sizes must be positive".into()); + } + if self.sscp_conv_channel_size.contains(&0) + || self + .sscp_conv_kernel_size + .iter() + .flatten() + .any(|size| *size == 0) + || self + .sscp_conv_stride_size + .iter() + .flatten() + .any(|size| *size == 0) + { + return Err("Gemma3n audio SSCP dimensions and strides must be positive".into()); + } + let mut frequency = self.input_feat_size; + for (kernel, stride) in self + .sscp_conv_kernel_size + .iter() + .zip(&self.sscp_conv_stride_size) + { + if frequency + 2 < kernel[1] { + return Err("Gemma3n audio SSCP frequency kernel exceeds padded input".into()); + } + frequency = (frequency + 2 - kernel[1]) / stride[1] + 1; + } + if !self.rms_norm_eps.is_finite() + || self.rms_norm_eps <= 0.0 + || !self.sscp_conv_group_norm_eps.is_finite() + || self.sscp_conv_group_norm_eps <= 0.0 + || !self.gradient_clipping.is_finite() + || self.gradient_clipping <= 0.0 + || !self.conf_attention_logit_cap.is_finite() + || self.conf_attention_logit_cap <= 0.0 + || !self.conf_residual_weight.is_finite() + { + return Err("Gemma3n audio normalization and clipping values must be finite".into()); + } + Ok(()) + } + + pub fn head_dim(&self) -> usize { + self.hidden_size / self.conf_num_attention_heads + } + + pub fn max_past_horizon(&self) -> usize { + self.conf_attention_context_left.saturating_sub(1) + } + + pub fn context_size(&self) -> usize { + self.conf_attention_chunk_size + self.max_past_horizon() + self.conf_attention_context_right + } + + pub fn time_stride_product(&self) -> usize { + self.sscp_conv_stride_size + .iter() + .map(|pair| pair[0]) + .product() + } +} diff --git a/src/audio/gemma3n/encoder.rs b/src/audio/gemma3n/encoder.rs new file mode 100644 index 00000000..6fc8f5e4 --- /dev/null +++ b/src/audio/gemma3n/encoder.rs @@ -0,0 +1,681 @@ +// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use super::attention::Gemma3nAudioAttention; +use super::checked_unified_linear; +use super::config::Gemma3nAudioConfig; +use mlxcel_core::layers::UnifiedLinear; +use mlxcel_core::weights::WeightMap; +use mlxcel_core::{MlxArray, UniquePtr}; + +fn copy_weight(weights: &WeightMap, key: &str) -> Result, String> { + weights + .get(key) + .map(|weight| mlxcel_core::copy(weight)) + .ok_or_else(|| format!("Gemma3n audio weight not found: {key}")) +} + +fn checked_weight( + weights: &WeightMap, + key: &str, + expected: &[i32], +) -> Result, String> { + let weight = copy_weight(weights, key)?; + let actual = mlxcel_core::array_shape(&weight); + if actual != expected { + return Err(format!( + "Gemma3n audio weight {key} has shape {actual:?}; expected {expected:?}" + )); + } + Ok(weight) +} + +struct AudioRmsNorm { + weight: UniquePtr, + eps: f32, +} + +impl AudioRmsNorm { + fn from_weights( + weights: &WeightMap, + prefix: &str, + size: usize, + eps: f32, + ) -> Result { + Ok(Self { + weight: checked_weight(weights, &format!("{prefix}.weight"), &[size as i32])?, + eps, + }) + } + + fn forward(&self, x: &MlxArray) -> UniquePtr { + mlxcel_core::rms_norm(x, &self.weight, self.eps) + } +} + +struct CumulativeGroupNorm { + weight: UniquePtr, + eps: f32, +} + +impl CumulativeGroupNorm { + fn from_weights( + weights: &WeightMap, + prefix: &str, + channels: usize, + eps: f32, + ) -> Result { + Ok(Self { + weight: checked_weight(weights, &format!("{prefix}.weight"), &[channels as i32])?, + eps, + }) + } + + /// Input and output are `[B, T, F, C]`. Statistics are float32 and are + /// cumulative over time after reducing both frequency and channels. + fn forward(&self, hidden_states: &MlxArray) -> UniquePtr { + let input_dtype = mlxcel_core::array_dtype(hidden_states); + let x = mlxcel_core::astype(hidden_states, mlxcel_core::dtype::FLOAT32); + let sum_at_time = mlxcel_core::sum_axis(&mlxcel_core::sum_axis(&x, 3, true), 2, true); + let cumulative_sum = mlxcel_core::cumsum(&sum_at_time, 1, false, true); + let count_at_time = mlxcel_core::sum_axis( + &mlxcel_core::sum_axis(&mlxcel_core::ones_like(&x), 3, true), + 2, + true, + ); + let cumulative_count = mlxcel_core::cumsum(&count_at_time, 1, false, true); + let one = mlxcel_core::full_f32(&[1], 1.0, mlxcel_core::dtype::FLOAT32); + let safe_count = mlxcel_core::maximum(&cumulative_count, &one); + let cumulative_mean = mlxcel_core::divide(&cumulative_sum, &safe_count); + + // Preserve the pinned reference's cumulative squared-difference + // algorithm; replacing it with E[x^2]-E[x]^2 is not equivalent. + let centered = mlxcel_core::subtract(&x, &cumulative_mean); + let squared = mlxcel_core::multiply(¢ered, ¢ered); + let squared_at_time = + mlxcel_core::sum_axis(&mlxcel_core::sum_axis(&squared, 3, true), 2, true); + let cumulative_squared = mlxcel_core::cumsum(&squared_at_time, 1, false, true); + let variance = mlxcel_core::divide(&cumulative_squared, &safe_count); + let eps = mlxcel_core::full_f32(&[1], self.eps, mlxcel_core::dtype::FLOAT32); + let inverse_stddev = mlxcel_core::rsqrt(&mlxcel_core::add(&variance, &eps)); + let normalized = mlxcel_core::multiply(¢ered, &inverse_stddev); + let normalized = mlxcel_core::multiply(&normalized, &self.weight); + mlxcel_core::astype(&normalized, input_dtype) + } +} + +struct SscpConvBlock { + conv_weight: UniquePtr, + norm: CumulativeGroupNorm, + time_padding_after: i32, + stride_time: i32, + stride_frequency: i32, +} + +impl SscpConvBlock { + fn from_weights( + weights: &WeightMap, + prefix: &str, + input_channels: usize, + output_channels: usize, + kernel: [usize; 2], + stride: [usize; 2], + eps: f32, + ) -> Result { + let conv_weight = checked_weight( + weights, + &format!("{prefix}.conv.weight"), + &[ + output_channels as i32, + kernel[0] as i32, + kernel[1] as i32, + input_channels as i32, + ], + )?; + Ok(Self { + conv_weight, + norm: CumulativeGroupNorm::from_weights( + weights, + &format!("{prefix}.norm"), + output_channels, + eps, + )?, + time_padding_after: kernel[0] as i32 - 1, + stride_time: stride[0] as i32, + stride_frequency: stride[1] as i32, + }) + } + + fn forward(&self, x: &MlxArray) -> Result, String> { + // Match the maintained reference's explicit + // `audio_encodings_padded.to(self.conv.weight.dtype)` boundary. The + // processor emits float32 mel features, while released checkpoints use + // BF16 convolution weights; leaving the input in float32 promotes the + // entire SSCP path and changes later greedy logits. + let x = mlxcel_core::astype(x, mlxcel_core::array_dtype(&self.conv_weight)); + // Reverse-causal time padding `(0, kernel-1)` and SAME-like frequency + // padding `(1,1)`, in MLX's NHWC layout. + let x = mlxcel_core::pad(&x, &[0, 0, 0, self.time_padding_after, 1, 1, 0, 0], 0.0); + let x = mlxcel_core::try_conv2d( + &x, + &self.conv_weight, + self.stride_time, + self.stride_frequency, + 0, + 0, + 1, + 1, + 1, + ) + .map_err(|error| format!("Gemma3n SSCP conv2d failed: {error}"))?; + Ok(mlxcel_core::relu(&self.norm.forward(&x))) + } +} + +struct SubSampleConvProjection { + conv0: SscpConvBlock, + conv1: SscpConvBlock, + input_projection: UnifiedLinear, +} + +fn final_frequency_size(config: &Gemma3nAudioConfig) -> usize { + config + .sscp_conv_kernel_size + .iter() + .zip(&config.sscp_conv_stride_size) + .fold(config.input_feat_size, |frequency, (kernel, stride)| { + (frequency + 2 - kernel[1]) / stride[1] + 1 + }) +} + +impl SubSampleConvProjection { + fn from_weights( + weights: &WeightMap, + prefix: &str, + config: &Gemma3nAudioConfig, + group_size: i32, + bits: i32, + ) -> Result { + let channels = &config.sscp_conv_channel_size; + Ok(Self { + conv0: SscpConvBlock::from_weights( + weights, + &format!("{prefix}.conv_0"), + 1, + channels[0], + config.sscp_conv_kernel_size[0], + config.sscp_conv_stride_size[0], + config.sscp_conv_group_norm_eps, + )?, + conv1: SscpConvBlock::from_weights( + weights, + &format!("{prefix}.conv_1"), + channels[0], + channels[1], + config.sscp_conv_kernel_size[1], + config.sscp_conv_stride_size[1], + config.sscp_conv_group_norm_eps, + )?, + input_projection: checked_unified_linear( + weights, + &format!("{prefix}.input_proj_linear"), + final_frequency_size(config) * channels[1], + config.hidden_size, + group_size, + bits, + )?, + }) + } + + fn forward(&self, audio_mel: &MlxArray) -> Result, String> { + let x = mlxcel_core::expand_dims(audio_mel, -1); + let x = self.conv0.forward(&x)?; + let x = self.conv1.forward(&x)?; + let shape = mlxcel_core::array_shape(&x); + let x = mlxcel_core::reshape(&x, &[shape[0], shape[1], shape[2] * shape[3]]); + Ok(self.input_projection.forward(&x)) + } +} + +fn clip_gradient(x: &MlxArray, limit: f32) -> UniquePtr { + let minimum = mlxcel_core::full_f32(&[1], -limit, mlxcel_core::array_dtype(x)); + let maximum = mlxcel_core::full_f32(&[1], limit, mlxcel_core::array_dtype(x)); + mlxcel_core::clip(x, &minimum, &maximum) +} + +struct FeedForward { + clipping: f32, + residual_weight: f32, + pre_norm: AudioRmsNorm, + layer1: UnifiedLinear, + layer2: UnifiedLinear, + post_norm: AudioRmsNorm, +} + +impl FeedForward { + fn from_weights( + weights: &WeightMap, + prefix: &str, + config: &Gemma3nAudioConfig, + group_size: i32, + bits: i32, + ) -> Result { + Ok(Self { + clipping: config.gradient_clipping, + residual_weight: config.conf_residual_weight, + pre_norm: AudioRmsNorm::from_weights( + weights, + &format!("{prefix}.pre_layer_norm"), + config.hidden_size, + config.rms_norm_eps, + )?, + layer1: checked_unified_linear( + weights, + &format!("{prefix}.ffw_layer_1"), + config.hidden_size, + config.hidden_size * 4, + group_size, + bits, + )?, + layer2: checked_unified_linear( + weights, + &format!("{prefix}.ffw_layer_2"), + config.hidden_size * 4, + config.hidden_size, + group_size, + bits, + )?, + post_norm: AudioRmsNorm::from_weights( + weights, + &format!("{prefix}.post_layer_norm"), + config.hidden_size, + config.rms_norm_eps, + )?, + }) + } + + fn forward(&self, x: &MlxArray) -> UniquePtr { + let residual = mlxcel_core::copy(x); + let x = self.pre_norm.forward(&clip_gradient(x, self.clipping)); + let x = mlxcel_core::silu(&self.layer1.forward(&x)); + let x = self.layer2.forward(&x); + let x = self.post_norm.forward(&clip_gradient(&x, self.clipping)); + mlxcel_core::add( + &residual, + &mlxcel_core::multiply_scalar(&x, self.residual_weight), + ) + } +} + +struct ConformerAttention { + clipping: f32, + pre_norm: AudioRmsNorm, + attention: Gemma3nAudioAttention, + post: UnifiedLinear, + post_norm: AudioRmsNorm, +} + +impl ConformerAttention { + fn from_weights( + weights: &WeightMap, + prefix: &str, + config: &Gemma3nAudioConfig, + group_size: i32, + bits: i32, + ) -> Result { + Ok(Self { + clipping: config.gradient_clipping, + pre_norm: AudioRmsNorm::from_weights( + weights, + &format!("{prefix}.pre_attn_norm"), + config.hidden_size, + config.rms_norm_eps, + )?, + attention: Gemma3nAudioAttention::from_weights( + weights, + &format!("{prefix}.attn"), + config, + group_size, + bits, + )?, + post: checked_unified_linear( + weights, + &format!("{prefix}.post"), + config.hidden_size, + config.hidden_size, + group_size, + bits, + )?, + post_norm: AudioRmsNorm::from_weights( + weights, + &format!("{prefix}.post_norm"), + config.hidden_size, + config.rms_norm_eps, + )?, + }) + } + + fn forward( + &self, + x: &MlxArray, + invalid_mask: &MlxArray, + causal_valid_mask: &MlxArray, + ) -> UniquePtr { + let residual = mlxcel_core::copy(x); + let normalized = self.pre_norm.forward(&clip_gradient(x, self.clipping)); + let attention = self + .attention + .forward(&normalized, invalid_mask, causal_valid_mask); + let shape = mlxcel_core::array_shape(&attention); + let attention = + mlxcel_core::reshape(&attention, &[shape[0], shape[1], shape[2] * shape[3]]); + let attention = self.post.forward(&attention); + let attention = self + .post_norm + .forward(&clip_gradient(&attention, self.clipping)); + mlxcel_core::add(&residual, &attention) + } +} + +struct LightConv1d { + clipping: f32, + causal_padding: i32, + pre_norm: AudioRmsNorm, + linear_start: UnifiedLinear, + depthwise_weight: UniquePtr, + conv_norm: AudioRmsNorm, + linear_end: UnifiedLinear, +} + +impl LightConv1d { + fn from_weights( + weights: &WeightMap, + prefix: &str, + config: &Gemma3nAudioConfig, + group_size: i32, + bits: i32, + ) -> Result { + Ok(Self { + clipping: config.gradient_clipping, + causal_padding: config.conf_conv_kernel_size as i32 - 1, + pre_norm: AudioRmsNorm::from_weights( + weights, + &format!("{prefix}.pre_layer_norm"), + config.hidden_size, + config.rms_norm_eps, + )?, + linear_start: checked_unified_linear( + weights, + &format!("{prefix}.linear_start"), + config.hidden_size, + config.hidden_size * 2, + group_size, + bits, + )?, + depthwise_weight: checked_weight( + weights, + &format!("{prefix}.depthwise_conv1d.weight"), + &[ + config.hidden_size as i32, + config.conf_conv_kernel_size as i32, + 1, + ], + )?, + conv_norm: AudioRmsNorm::from_weights( + weights, + &format!("{prefix}.conv_norm"), + config.hidden_size, + config.rms_norm_eps, + )?, + linear_end: checked_unified_linear( + weights, + &format!("{prefix}.linear_end"), + config.hidden_size, + config.hidden_size, + group_size, + bits, + )?, + }) + } + + fn forward(&self, x: &MlxArray) -> Result, String> { + let residual = mlxcel_core::copy(x); + let x = self.linear_start.forward(&self.pre_norm.forward(x)); + let shape = mlxcel_core::array_shape(&x); + let half = shape[2] / 2; + let left = mlxcel_core::slice(&x, &[0, 0, 0], &[shape[0], shape[1], half]); + let right = mlxcel_core::slice(&x, &[0, 0, half], &[shape[0], shape[1], shape[2]]); + let x = mlxcel_core::multiply(&left, &mlxcel_core::sigmoid(&right)); + let x = mlxcel_core::pad(&x, &[0, 0, self.causal_padding, 0, 0, 0], 0.0); + let x = mlxcel_core::try_conv1d( + &x, + &self.depthwise_weight, + 1, + 0, + 1, + mlxcel_core::array_shape(&x)[2], + ) + .map_err(|error| format!("Gemma3n audio depthwise conv1d failed: {error}"))?; + let x = self.conv_norm.forward(&clip_gradient(&x, self.clipping)); + let x = self.linear_end.forward(&mlxcel_core::silu(&x)); + Ok(mlxcel_core::add(&residual, &x)) + } +} + +struct ConformerBlock { + clipping: f32, + feed_forward_start: FeedForward, + attention: ConformerAttention, + light_conv: LightConv1d, + feed_forward_end: FeedForward, + norm: AudioRmsNorm, +} + +impl ConformerBlock { + fn from_weights( + weights: &WeightMap, + prefix: &str, + config: &Gemma3nAudioConfig, + group_size: i32, + bits: i32, + ) -> Result { + Ok(Self { + clipping: config.gradient_clipping, + feed_forward_start: FeedForward::from_weights( + weights, + &format!("{prefix}.ffw_layer_start"), + config, + group_size, + bits, + )?, + attention: ConformerAttention::from_weights( + weights, + &format!("{prefix}.attention"), + config, + group_size, + bits, + )?, + light_conv: LightConv1d::from_weights( + weights, + &format!("{prefix}.lconv1d"), + config, + group_size, + bits, + )?, + feed_forward_end: FeedForward::from_weights( + weights, + &format!("{prefix}.ffw_layer_end"), + config, + group_size, + bits, + )?, + norm: AudioRmsNorm::from_weights( + weights, + &format!("{prefix}.norm"), + config.hidden_size, + config.rms_norm_eps, + )?, + }) + } + + fn forward( + &self, + x: &MlxArray, + invalid_mask: &MlxArray, + causal_valid_mask: &MlxArray, + ) -> Result, String> { + let x = self.feed_forward_start.forward(x); + let x = self.attention.forward(&x, invalid_mask, causal_valid_mask); + let valid = mlxcel_core::astype( + &mlxcel_core::reshape( + &mlxcel_core::logical_not(invalid_mask), + &[ + mlxcel_core::array_shape(&x)[0], + mlxcel_core::array_shape(&x)[1], + 1, + ], + ), + mlxcel_core::array_dtype(&x), + ); + let x = self + .light_conv + .forward(&mlxcel_core::multiply(&x, &valid))?; + let x = self.feed_forward_end.forward(&x); + Ok(self.norm.forward(&clip_gradient(&x, self.clipping))) + } +} + +pub struct Gemma3nAudioEncoder { + config: Gemma3nAudioConfig, + subsample: SubSampleConvProjection, + conformer: Vec, +} + +impl Gemma3nAudioEncoder { + pub fn from_weights( + weights: &WeightMap, + prefix: &str, + config: &Gemma3nAudioConfig, + group_size: i32, + bits: i32, + ) -> Result { + config.validate()?; + let mut conformer = Vec::with_capacity(config.conf_num_hidden_layers); + for index in 0..config.conf_num_hidden_layers { + conformer.push(ConformerBlock::from_weights( + weights, + &format!("{prefix}.conformer.{index}"), + config, + group_size, + bits, + )?); + } + Ok(Self { + config: config.clone(), + subsample: SubSampleConvProjection::from_weights( + weights, + &format!("{prefix}.subsample_conv_projection"), + config, + group_size, + bits, + )?, + conformer, + }) + } + + fn causal_valid_mask(&self) -> UniquePtr { + let chunk = self.config.conf_attention_chunk_size as i32; + let context = self.config.context_size() as i32; + let diagonal = + (self.config.max_past_horizon() + self.config.conf_attention_context_right) as i32; + let lower = mlxcel_core::transpose_axes( + &mlxcel_core::tril( + &mlxcel_core::ones(&[context, chunk], mlxcel_core::dtype::FLOAT32), + 0, + ), + &[1, 0], + ); + let upper = mlxcel_core::tril( + &mlxcel_core::ones(&[chunk, context], mlxcel_core::dtype::FLOAT32), + diagonal, + ); + mlxcel_core::astype( + &mlxcel_core::multiply(&lower, &upper), + mlxcel_core::dtype::BOOL, + ) + } + + fn stride_mask(mask: &MlxArray, stride: usize, length: i32) -> UniquePtr { + let source_len = mlxcel_core::array_shape(mask)[1]; + let indices: Vec = (0..length) + .map(|index| (index * stride as i32).min(source_len - 1)) + .collect(); + mlxcel_core::take(mask, &mlxcel_core::from_slice_i32(&indices, &[length]), 1) + } + + pub fn forward( + &self, + audio_mel: &MlxArray, + invalid_mel_mask: &MlxArray, + ) -> Result<(UniquePtr, UniquePtr), String> { + let mel_shape = mlxcel_core::array_shape(audio_mel); + let mask_shape = mlxcel_core::array_shape(invalid_mel_mask); + if mel_shape.len() != 3 + || mel_shape[2] != self.config.input_feat_size as i32 + || mask_shape != mel_shape[..2] + { + return Err(format!( + "Gemma3n audio input shapes must be [B,T,{}] and [B,T], got {mel_shape:?} and {mask_shape:?}", + self.config.input_feat_size + )); + } + + let mut encodings = self.subsample.forward(audio_mel)?; + let mut invalid_mask = Self::stride_mask( + invalid_mel_mask, + self.config.time_stride_product(), + mlxcel_core::array_shape(&encodings)[1], + ); + let causal_valid_mask = self.causal_valid_mask(); + for block in &self.conformer { + encodings = block.forward(&encodings, &invalid_mask, &causal_valid_mask)?; + } + + if self.config.conf_reduction_factor > 1 { + let reduced_len = (mlxcel_core::array_shape(&encodings)[1] + + self.config.conf_reduction_factor as i32 + - 1) + / self.config.conf_reduction_factor as i32; + encodings = + Self::stride_mask(&encodings, self.config.conf_reduction_factor, reduced_len); + invalid_mask = Self::stride_mask( + &invalid_mask, + self.config.conf_reduction_factor, + reduced_len, + ); + } + let shape = mlxcel_core::array_shape(&encodings); + let expanded_mask = mlxcel_core::reshape(&invalid_mask, &[shape[0], shape[1], 1]); + encodings = mlxcel_core::where_cond( + &expanded_mask, + &mlxcel_core::zeros_like(&encodings), + &encodings, + ); + Ok((encodings, invalid_mask)) + } +} + +#[cfg(test)] +#[path = "encoder_tests.rs"] +mod tests; diff --git a/src/audio/gemma3n/encoder_tests.rs b/src/audio/gemma3n/encoder_tests.rs new file mode 100644 index 00000000..abdb6491 --- /dev/null +++ b/src/audio/gemma3n/encoder_tests.rs @@ -0,0 +1,292 @@ +// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use super::*; + +fn tiny_config() -> Gemma3nAudioConfig { + Gemma3nAudioConfig { + vocab_size: 8, + vocab_offset: 100, + input_feat_size: 4, + hidden_size: 4, + rms_norm_eps: 1e-6, + gradient_clipping: 1e4, + conf_attention_chunk_size: 2, + conf_attention_context_left: 2, + conf_attention_context_right: 0, + conf_attention_logit_cap: 5.0, + conf_num_attention_heads: 2, + conf_num_hidden_layers: 1, + conf_conv_kernel_size: 3, + conf_reduction_factor: 2, + conf_residual_weight: 0.5, + sscp_conv_channel_size: vec![2, 2], + sscp_conv_group_norm_eps: 1e-3, + sscp_conv_kernel_size: vec![[3, 3], [3, 3]], + sscp_conv_stride_size: vec![[2, 2], [2, 2]], + } +} + +fn insert(weights: &mut WeightMap, key: &str, shape: &[i32], value: f32) { + weights.insert( + key.to_string(), + mlxcel_core::full_f32(shape, value, mlxcel_core::dtype::FLOAT32), + ); +} + +fn values(array: &MlxArray) -> Vec { + mlxcel_core::eval(array); + mlxcel_core::array_to_raw_bytes(array) + .chunks_exact(4) + .map(|bytes| f32::from_ne_bytes(bytes.try_into().unwrap())) + .collect() +} + +fn tiny_weights() -> WeightMap { + let mut weights = WeightMap::new(); + insert( + &mut weights, + "audio.subsample_conv_projection.conv_0.conv.weight", + &[2, 3, 3, 1], + 0.01, + ); + insert( + &mut weights, + "audio.subsample_conv_projection.conv_0.norm.weight", + &[2], + 1.0, + ); + insert( + &mut weights, + "audio.subsample_conv_projection.conv_1.conv.weight", + &[2, 3, 3, 2], + 0.01, + ); + insert( + &mut weights, + "audio.subsample_conv_projection.conv_1.norm.weight", + &[2], + 1.0, + ); + insert( + &mut weights, + "audio.subsample_conv_projection.input_proj_linear.weight", + &[4, 2], + 0.01, + ); + + for name in ["ffw_layer_start", "ffw_layer_end"] { + let prefix = format!("audio.conformer.0.{name}"); + insert( + &mut weights, + &format!("{prefix}.pre_layer_norm.weight"), + &[4], + 1.0, + ); + insert( + &mut weights, + &format!("{prefix}.ffw_layer_1.weight"), + &[16, 4], + 0.01, + ); + insert( + &mut weights, + &format!("{prefix}.ffw_layer_2.weight"), + &[4, 16], + 0.01, + ); + insert( + &mut weights, + &format!("{prefix}.post_layer_norm.weight"), + &[4], + 1.0, + ); + } + + for name in ["pre_attn_norm", "post_norm"] { + insert( + &mut weights, + &format!("audio.conformer.0.attention.{name}.weight"), + &[4], + 1.0, + ); + } + for name in ["q_proj", "k_proj", "v_proj"] { + insert( + &mut weights, + &format!("audio.conformer.0.attention.attn.{name}.weight"), + &[4, 4], + 0.01, + ); + } + insert( + &mut weights, + "audio.conformer.0.attention.attn.per_dim_scale", + &[2], + 0.0, + ); + insert( + &mut weights, + "audio.conformer.0.attention.attn.relative_position_embedding.pos_proj.weight", + &[4, 4], + 0.01, + ); + insert( + &mut weights, + "audio.conformer.0.attention.post.weight", + &[4, 4], + 0.01, + ); + + for name in ["pre_layer_norm", "conv_norm"] { + insert( + &mut weights, + &format!("audio.conformer.0.lconv1d.{name}.weight"), + &[4], + 1.0, + ); + } + insert( + &mut weights, + "audio.conformer.0.lconv1d.linear_start.weight", + &[8, 4], + 0.01, + ); + insert( + &mut weights, + "audio.conformer.0.lconv1d.depthwise_conv1d.weight", + &[4, 3, 1], + 0.01, + ); + insert( + &mut weights, + "audio.conformer.0.lconv1d.linear_end.weight", + &[4, 4], + 0.01, + ); + insert(&mut weights, "audio.conformer.0.norm.weight", &[4], 1.0); + weights +} + +#[test] +fn default_local_attention_contract_rejects_global_substitution() { + let config = Gemma3nAudioConfig::default(); + assert_eq!(config.conf_attention_chunk_size, 12); + assert_eq!(config.max_past_horizon(), 12); + assert_eq!(config.context_size(), 24); + assert_eq!(config.head_dim(), 192); + assert_eq!(config.conf_attention_logit_cap, 50.0); + + let encoder = + Gemma3nAudioEncoder::from_weights(&tiny_weights(), "audio", &tiny_config(), 64, 4).unwrap(); + let causal_mask = encoder.causal_valid_mask(); + mlxcel_core::eval(&causal_mask); + assert_eq!( + mlxcel_core::array_to_raw_bytes(&causal_mask), + vec![1, 1, 0, 0, 1, 1] + ); +} + +#[test] +fn invalid_sscp_config_is_rejected_before_loading_weights() { + let mut config = Gemma3nAudioConfig::default(); + config.sscp_conv_channel_size.pop(); + assert!(config.validate().is_err()); + let config = Gemma3nAudioConfig { + conf_attention_logit_cap: 0.0, + ..Gemma3nAudioConfig::default() + }; + assert!(config.validate().is_err()); +} + +#[test] +fn cumulative_group_norm_does_not_collapse_to_ordinary_group_norm() { + let norm = CumulativeGroupNorm { + weight: mlxcel_core::from_slice_f32(&[1.0, 2.0], &[2]), + eps: 0.0, + }; + let input = mlxcel_core::from_slice_f32(&[1.0, 3.0, 5.0, 7.0], &[1, 2, 1, 2]); + let actual = values(&norm.forward(&input)); + let expected = [-1.0, 2.0, 1.0 / 3.0f32.sqrt(), 6.0 / 3.0f32.sqrt()]; + for (actual, expected) in actual.iter().zip(expected) { + assert!((actual - expected).abs() < 1e-5, "{actual} != {expected}"); + } +} + +#[test] +fn sscp_casts_processor_features_to_checkpoint_convolution_dtype() { + let block = SscpConvBlock { + conv_weight: mlxcel_core::astype( + &mlxcel_core::full_f32(&[2, 3, 3, 1], 0.01, mlxcel_core::dtype::FLOAT32), + mlxcel_core::dtype::BFLOAT16, + ), + norm: CumulativeGroupNorm { + weight: mlxcel_core::ones(&[2], mlxcel_core::dtype::BFLOAT16), + eps: 1e-3, + }, + time_padding_after: 2, + stride_time: 2, + stride_frequency: 2, + }; + let samples: Vec = (0..16).map(|index| index as f32 / 37.0).collect(); + let input = mlxcel_core::from_slice_f32(&samples, &[1, 4, 4, 1]); + let input_bf16 = mlxcel_core::astype(&input, mlxcel_core::dtype::BFLOAT16); + let from_processor = values(&block.forward(&input).unwrap()); + let from_checkpoint_dtype = values(&block.forward(&input_bf16).unwrap()); + assert_eq!(from_processor, from_checkpoint_dtype); +} + +#[test] +fn tiny_encoder_exercises_sscp_conformer_reduction_and_mask() { + let encoder = + Gemma3nAudioEncoder::from_weights(&tiny_weights(), "audio", &tiny_config(), 64, 4).unwrap(); + let mel = mlxcel_core::from_slice_f32( + &(0..32).map(|index| index as f32 / 32.0).collect::>(), + &[1, 8, 4], + ); + let invalid = mlxcel_core::astype( + &mlxcel_core::from_slice_i32(&[0, 0, 0, 0, 0, 0, 1, 1], &[1, 8]), + mlxcel_core::dtype::BOOL, + ); + let (encoded, mask) = encoder.forward(&mel, &invalid).unwrap(); + assert_eq!(mlxcel_core::array_shape(&encoded), vec![1, 1, 4]); + assert_eq!(mlxcel_core::array_shape(&mask), vec![1, 1]); + mlxcel_core::eval(&encoded); +} + +#[test] +fn malformed_checkpoint_conv_shape_is_rejected() { + let mut weights = tiny_weights(); + weights.insert( + "audio.subsample_conv_projection.conv_0.conv.weight".into(), + mlxcel_core::zeros(&[2, 1, 3, 3], mlxcel_core::dtype::FLOAT32), + ); + let error = Gemma3nAudioEncoder::from_weights(&weights, "audio", &tiny_config(), 64, 4) + .err() + .expect("shape mismatch must fail"); + assert!(error.contains("expected")); +} + +#[test] +fn malformed_checkpoint_projection_shape_is_rejected() { + let mut weights = tiny_weights(); + weights.insert( + "audio.conformer.0.ffw_layer_start.ffw_layer_1.weight".into(), + mlxcel_core::zeros(&[4, 4], mlxcel_core::dtype::FLOAT32), + ); + let error = Gemma3nAudioEncoder::from_weights(&weights, "audio", &tiny_config(), 64, 4) + .err() + .expect("linear shape mismatch must fail"); + assert!(error.contains("expected logical")); +} diff --git a/src/audio/gemma3n/feature_extractor.rs b/src/audio/gemma3n/feature_extractor.rs new file mode 100644 index 00000000..338a2444 --- /dev/null +++ b/src/audio/gemma3n/feature_extractor.rs @@ -0,0 +1,355 @@ +// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Official Gemma 3n USM waveform processor. + +use std::f64::consts::PI; + +pub const GEMMA3N_SAMPLE_RATE: u32 = 16_000; +pub const GEMMA3N_MAX_SAMPLES: usize = 480_000; +pub const GEMMA3N_AUDIO_SOFT_TOKENS: usize = 188; +const FEATURE_SIZE: usize = 128; +const FRAME_LENGTH: usize = 512; +const HOP_LENGTH: usize = 160; +const FFT_LENGTH: usize = 1024; +const PAD_MULTIPLE: usize = 128; + +#[derive(Debug, Clone)] +pub struct Gemma3nAudioFeatureBatch { + /// Flattened `[batch, frames, 128]` log-mel values. + pub features: Vec, + /// Flattened `[batch, frames]`; `true` means a real input frame. + pub valid_mask: Vec, + pub batch_size: usize, + pub frames: usize, +} + +pub struct Gemma3nAudioFeatureExtractor { + window: Vec, + mel_filters: Vec, +} + +impl Default for Gemma3nAudioFeatureExtractor { + fn default() -> Self { + Self::new() + } +} + +impl Gemma3nAudioFeatureExtractor { + pub fn new() -> Self { + let window = (0..FRAME_LENGTH) + .map(|i| (0.5 - 0.5 * (2.0 * PI * i as f64 / FRAME_LENGTH as f64).cos()) as f32) + .collect(); + Self { + window, + mel_filters: mel_filter_bank(), + } + } + + /// Process already-normalized mono 16 kHz float32 clips as one batch. + /// + /// The pinned integration processor truncates each clip at 30 seconds, + /// left-pads the batch to its longest clip and then to a multiple of 128 + /// samples (`AutoProcessor(..., padding_side="left")`). + pub fn extract_batch(&self, clips: &[Vec]) -> Result { + if clips.is_empty() { + return Err("Gemma3n audio batch must contain at least one clip".into()); + } + for (index, clip) in clips.iter().enumerate() { + if clip.is_empty() { + return Err(format!("Gemma3n audio clip {index} is empty")); + } + if clip.iter().any(|sample| !sample.is_finite()) { + return Err(format!( + "Gemma3n audio clip {index} contains a non-finite sample" + )); + } + if clip.iter().any(|sample| !(-1.0..=1.0).contains(sample)) { + return Err(format!( + "Gemma3n audio clip {index} contains a sample outside [-1, 1]" + )); + } + } + + let longest = clips + .iter() + .map(|clip| clip.len().min(GEMMA3N_MAX_SAMPLES)) + .max() + .unwrap_or(0); + let padded_len = longest.div_ceil(PAD_MULTIPLE) * PAD_MULTIPLE; + let frames = if padded_len < FRAME_LENGTH + 1 { + 0 + } else { + (padded_len - (FRAME_LENGTH + 1)) / HOP_LENGTH + 1 + }; + let mut features = Vec::with_capacity(clips.len() * frames * FEATURE_SIZE); + let mut valid_mask = Vec::with_capacity(clips.len() * frames); + + for clip in clips { + let effective_len = clip.len().min(GEMMA3N_MAX_SAMPLES); + let mut waveform = vec![0.0f32; padded_len]; + let left_padding = padded_len - effective_len; + waveform[left_padding..left_padding + effective_len] + .copy_from_slice(&clip[..effective_len]); + + for frame_index in 0..frames { + let start = frame_index * HOP_LENGTH; + let unfolded = &waveform[start..start + FRAME_LENGTH + 1]; + let mut fft_input = vec![0.0f64; FFT_LENGTH]; + fft_input[0] = (unfolded[0] * 0.03 * self.window[0]) as f64; + for i in 1..FRAME_LENGTH { + let emphasized = unfolded[i] - 0.97 * unfolded[i - 1]; + fft_input[i] = (emphasized * self.window[i]) as f64; + } + if fft_input[..FRAME_LENGTH].iter().all(|value| *value == 0.0) { + features.extend(std::iter::repeat_n(1e-5f32.ln(), FEATURE_SIZE)); + valid_mask.push(start >= left_padding && start < left_padding + effective_len); + continue; + } + let magnitude = real_fft_magnitude_1024(&fft_input); + for mel in 0..FEATURE_SIZE { + let mut value = 0.0f64; + for (bin, magnitude) in magnitude.iter().enumerate() { + value += magnitude * self.mel_filters[bin * FEATURE_SIZE + mel] as f64; + } + features.push(value.max(1e-5).ln() as f32); + } + // The pinned processor samples the left-padded waveform + // attention mask at the frame start, not at the end of the + // analysis window. + valid_mask.push(start >= left_padding && start < left_padding + effective_len); + } + } + + Ok(Gemma3nAudioFeatureBatch { + features, + valid_mask, + batch_size: clips.len(), + frames, + }) + } +} + +/// Compute the positive-frequency magnitude spectrum for the fixed 1024-point +/// Gemma 3n transform. The in-tree generic audio helper intentionally uses a +/// direct DFT, which is useful for small reference transforms but makes a +/// 30-second Gemma 3n clip prohibitively expensive. This radix-2 FFT preserves +/// the reference math while keeping frontend preprocessing practical. +fn real_fft_magnitude_1024(input: &[f64]) -> Vec { + debug_assert_eq!(input.len(), FFT_LENGTH); + let mut real = input.to_vec(); + let mut imaginary = vec![0.0f64; FFT_LENGTH]; + + let mut reversed = 0usize; + for index in 1..FFT_LENGTH { + let mut bit = FFT_LENGTH >> 1; + while reversed & bit != 0 { + reversed ^= bit; + bit >>= 1; + } + reversed ^= bit; + if index < reversed { + real.swap(index, reversed); + imaginary.swap(index, reversed); + } + } + + let mut width = 2usize; + while width <= FFT_LENGTH { + let half = width / 2; + let angle = -2.0 * PI / width as f64; + let step_real = angle.cos(); + let step_imaginary = angle.sin(); + for start in (0..FFT_LENGTH).step_by(width) { + let mut twiddle_real = 1.0f64; + let mut twiddle_imaginary = 0.0f64; + for offset in 0..half { + let right = start + offset + half; + let product_real = + twiddle_real * real[right] - twiddle_imaginary * imaginary[right]; + let product_imaginary = + twiddle_real * imaginary[right] + twiddle_imaginary * real[right]; + let left = start + offset; + let left_real = real[left]; + let left_imaginary = imaginary[left]; + real[left] = left_real + product_real; + imaginary[left] = left_imaginary + product_imaginary; + real[right] = left_real - product_real; + imaginary[right] = left_imaginary - product_imaginary; + let next_real = twiddle_real * step_real - twiddle_imaginary * step_imaginary; + twiddle_imaginary = twiddle_real * step_imaginary + twiddle_imaginary * step_real; + twiddle_real = next_real; + } + } + width *= 2; + } + + (0..=FFT_LENGTH / 2) + .map(|index| real[index].hypot(imaginary[index])) + .collect() +} + +fn mel_filter_bank() -> Vec { + let n_freqs = FFT_LENGTH / 2 + 1; + let hz_to_mel = |hz: f64| 2595.0 * (1.0 + hz / 700.0).log10(); + let mel_to_hz = |mel: f64| 700.0 * (10.0f64.powf(mel / 2595.0) - 1.0); + let min_mel = hz_to_mel(125.0); + let max_mel = hz_to_mel(7600.0); + let points: Vec = (0..FEATURE_SIZE + 2) + .map(|i| mel_to_hz(min_mel + (max_mel - min_mel) * i as f64 / (FEATURE_SIZE + 1) as f64)) + .collect(); + let mut filters = vec![0.0; n_freqs * FEATURE_SIZE]; + for bin in 0..n_freqs { + let frequency = bin as f64 * GEMMA3N_SAMPLE_RATE as f64 / FFT_LENGTH as f64; + for mel in 0..FEATURE_SIZE { + let down = (frequency - points[mel]) / (points[mel + 1] - points[mel]); + let up = (points[mel + 2] - frequency) / (points[mel + 2] - points[mel + 1]); + filters[bin * FEATURE_SIZE + mel] = down.min(up).max(0.0) as f32; + } + } + filters +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn official_boundary_emits_2997_mel_frames_and_188_tokens() { + let batch = Gemma3nAudioFeatureExtractor::new() + .extract_batch(&[vec![0.0; GEMMA3N_MAX_SAMPLES]]) + .unwrap(); + assert_eq!(batch.frames, 2997); + let sscp_frames = batch.frames.div_ceil(2).div_ceil(2); + assert_eq!(sscp_frames.div_ceil(4), GEMMA3N_AUDIO_SOFT_TOKENS); + } + + #[test] + fn batch_padding_keeps_clip_boundaries_and_mask() { + let batch = Gemma3nAudioFeatureExtractor::new() + .extract_batch(&[vec![0.0; 16_000], vec![0.0; 8_000]]) + .unwrap(); + assert_eq!(batch.batch_size, 2); + assert_eq!(batch.frames, 97); + assert!(batch.valid_mask[..97].iter().all(|valid| *valid)); + assert!(batch.valid_mask[97..97 + 50].iter().all(|valid| !*valid)); + assert!(batch.valid_mask[97 + 50..].iter().all(|valid| *valid)); + } + + #[test] + fn rejects_empty_nonfinite_and_out_of_range_clips() { + let extractor = Gemma3nAudioFeatureExtractor::new(); + assert!(extractor.extract_batch(&[vec![]]).is_err()); + assert!(extractor.extract_batch(&[vec![f32::NAN]]).is_err()); + assert!(extractor.extract_batch(&[vec![1.01]]).is_err()); + } + + #[test] + fn sub_frame_clip_matches_reference_zero_frame_result() { + let batch = Gemma3nAudioFeatureExtractor::new() + .extract_batch(&[vec![0.25; 128]]) + .unwrap(); + assert_eq!(batch.frames, 0); + assert!(batch.features.is_empty()); + assert!(batch.valid_mask.is_empty()); + } + + #[test] + fn hard_mel_floor_is_log_one_e_minus_five() { + let batch = Gemma3nAudioFeatureExtractor::new() + .extract_batch(&[vec![0.0; 640]]) + .unwrap(); + let expected = 1e-5f32.ln(); + assert!( + batch + .features + .iter() + .all(|value| (*value - expected).abs() < 1e-6) + ); + } + + #[test] + fn clips_longer_than_thirty_seconds_are_truncated_at_the_reference_boundary() { + let extractor = Gemma3nAudioFeatureExtractor::new(); + let at_limit = extractor + .extract_batch(&[vec![0.0; GEMMA3N_MAX_SAMPLES]]) + .unwrap(); + let over_limit = extractor + .extract_batch(&[vec![ + 0.0; + GEMMA3N_MAX_SAMPLES + GEMMA3N_SAMPLE_RATE as usize + ]]) + .unwrap(); + assert_eq!(over_limit.frames, at_limit.frames); + assert_eq!(over_limit.features, at_limit.features); + assert_eq!(over_limit.valid_mask, at_limit.valid_mask); + } + + #[test] + fn deterministic_waveform_matches_pinned_transformers_mel_fixture() { + // Fixture source: Transformers commit + // 181beb3ba4c47098ed8cbc97ee250d1d45ae0107's NumPy algorithm. + let waveform: Vec = (0..1600) + .map(|index| { + let time = index as f64 / GEMMA3N_SAMPLE_RATE as f64; + (0.25 * (2.0 * PI * 440.0 * time).sin() + 0.1 * (2.0 * PI * 1000.0 * time).cos()) + as f32 + }) + .collect(); + let batch = Gemma3nAudioFeatureExtractor::new() + .extract_batch(&[waveform]) + .unwrap(); + assert_eq!(batch.frames, 8); + let fixtures = [ + ((0, 0), -4.375_432), + ((0, 10), -2.835_216_8), + ((0, 64), -2.875_424), + ((1, 20), -0.997_653_1), + ((3, 100), -11.512_925), + ((6, 127), -11.512_925), + ]; + for ((frame, bin), expected) in fixtures { + let actual = batch.features[frame * FEATURE_SIZE + bin]; + assert!( + (actual - expected).abs() < 1e-3, + "mel[{frame},{bin}] = {actual}, expected {expected}" + ); + } + + // Full-tensor checksums from the same pinned NumPy implementation. + // The Rust radix-2 FFT and NumPy's pocketfft use different reduction + // orders, so accumulated tolerances are wider than the per-bin 1e-3. + let sum: f64 = batch.features.iter().map(|value| *value as f64).sum(); + let squared_sum: f64 = batch + .features + .iter() + .map(|value| (*value as f64).powi(2)) + .sum(); + let weighted_sum: f64 = batch + .features + .iter() + .enumerate() + .map(|(index, value)| (index + 1) as f64 * *value as f64) + .sum(); + assert!((sum - -7_277.967_449_545_86).abs() < 0.1, "sum={sum}"); + assert!( + (squared_sum - 69_566.958_335_671_2).abs() < 2.0, + "squared_sum={squared_sum}" + ); + assert!( + (weighted_sum - -4_151_055.318_334_281_4).abs() < 100.0, + "weighted_sum={weighted_sum}" + ); + } +} diff --git a/src/audio/gemma3n/mod.rs b/src/audio/gemma3n/mod.rs new file mode 100644 index 00000000..4f657d71 --- /dev/null +++ b/src/audio/gemma3n/mod.rs @@ -0,0 +1,91 @@ +// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Token-exact Gemma 3n audio reference path. +//! +//! This is intentionally separate from the Gemma 4 audio implementation: +//! Gemma 3n has a different waveform front-end, cumulative SSCP group norm, +//! attention scaling, weight layout, and fixed 188-token merge contract. + +mod attention; +mod config; +mod encoder; +mod feature_extractor; + +use mlxcel_core::layers::UnifiedLinear; +use mlxcel_core::weights::WeightMap; + +/// Maintained implementation used for architecture and frontend fixtures. +pub const GEMMA3N_TRANSFORMERS_REFERENCE_REVISION: &str = + "181beb3ba4c47098ed8cbc97ee250d1d45ae0107"; +/// Official `google/gemma-3n-E4B` checkpoint revision used for weight/layout qualification. +pub const GEMMA3N_E4B_REFERENCE_REVISION: &str = "af70430f84ea4d7ac191aaa2bd8e14d2a5e8f6ee"; + +/// Load a linear tensor while checking its declared logical dimensions. +/// Quantized layouts are additionally validated by `UnifiedLinear`; here we +/// pin the output width and all dense/sidecar ranks so malformed audio exports +/// fail during model loading rather than in request-time matmul. +pub(crate) fn checked_unified_linear( + weights: &WeightMap, + prefix: &str, + input_size: usize, + output_size: usize, + group_size: i32, + bits: i32, +) -> Result { + let weight_key = format!("{prefix}.weight"); + let weight = weights + .get(&weight_key) + .ok_or_else(|| format!("Gemma3n audio weight not found: {weight_key}"))?; + let weight_shape = mlxcel_core::array_shape(weight); + let quantized = weights.contains_key(&format!("{prefix}.scales")); + let valid_weight = weight_shape.len() == 2 + && weight_shape[0] == output_size as i32 + && (quantized || weight_shape[1] == input_size as i32); + if !valid_weight { + return Err(format!( + "Gemma3n audio linear {weight_key} has shape {weight_shape:?}; expected logical [{output_size}, {input_size}]" + )); + } + if let Some(scales) = weights.get(&format!("{prefix}.scales")) { + let scales_shape = mlxcel_core::array_shape(scales); + if scales_shape.len() != 2 || scales_shape[0] != output_size as i32 { + return Err(format!( + "Gemma3n audio linear {prefix}.scales has invalid shape {scales_shape:?}" + )); + } + if let Some(biases) = weights.get(&format!("{prefix}.biases")) + && mlxcel_core::array_shape(biases) != scales_shape + { + return Err(format!( + "Gemma3n audio linear {prefix}.biases must match scales shape {scales_shape:?}" + )); + } + } + if let Some(bias) = weights.get(&format!("{prefix}.bias")) + && mlxcel_core::array_shape(bias) != [output_size as i32] + { + return Err(format!( + "Gemma3n audio linear {prefix}.bias must have {output_size} elements" + )); + } + UnifiedLinear::from_weights(weights, prefix, group_size, bits) +} + +pub use config::Gemma3nAudioConfig; +pub use encoder::Gemma3nAudioEncoder; +pub use feature_extractor::{ + GEMMA3N_AUDIO_SOFT_TOKENS, GEMMA3N_MAX_SAMPLES, GEMMA3N_SAMPLE_RATE, Gemma3nAudioFeatureBatch, + Gemma3nAudioFeatureExtractor, +}; diff --git a/src/audio/mod.rs b/src/audio/mod.rs index da674129..16ce0701 100644 --- a/src/audio/mod.rs +++ b/src/audio/mod.rs @@ -30,6 +30,7 @@ mod attention; pub mod config; pub mod encoder; pub mod feature_extractor; +pub mod gemma3n; pub mod nemotron_h_nano_omni; pub mod phi4mm; pub mod qwen3_omni_moe; diff --git a/src/commands/generate_vlm.rs b/src/commands/generate_vlm.rs index 5dcdeb09..a6f8078d 100644 --- a/src/commands/generate_vlm.rs +++ b/src/commands/generate_vlm.rs @@ -18,6 +18,7 @@ use std::path::{Path, PathBuf}; use mlxcel::LoadedModel; use mlxcel::video; use mlxcel::vision::merge::InputEmbeddings; +use mlxcel::vision::processors::ImageProcessor; use mlxcel::vlm_prompt::ImageTokenBlockAction; use mlxcel::vlm_runtime::{VlmPreparationSummary, prepare_and_compute_vlm_embeddings_with_budget}; @@ -81,6 +82,15 @@ fn print_preparation_summary(summary: VlmPreparationSummary) { audio_tokens, total_tokens ); } + VlmPreparationSummary::Gemma3nAudio { + audio_clips, + audio_tokens, + total_tokens, + } => { + println!( + "Gemma3n: expanded {audio_clips} audio clip(s) into {audio_tokens} soft tokens ({total_tokens} total tokens)" + ); + } VlmPreparationSummary::Gemma4Video { video_count, frame_slots, @@ -392,6 +402,19 @@ pub(crate) fn compute_vlm_embeddings( ); } + // Handle audio-only and image+audio modes for Gemma 3n. + if let Some(audio) = audio_path + && let LoadedModel::Gemma3nVLM(gemma3n) = model + { + return compute_gemma3n_audio_embeddings( + gemma3n, + prompt_tokens, + image_paths, + audio, + tokenizer, + ); + } + // Handle audio-only mode for Gemma4 if image_paths.is_empty() && let Some(audio) = audio_path @@ -456,7 +479,7 @@ pub(crate) fn compute_vlm_embeddings( if audio_path.is_some() { return Err(anyhow::anyhow!( "--audio input is not supported for this model family. Currently audio is wired \ - through Phi4MM, Gemma 4, Nemotron H Nano Omni, and Qwen3-Omni MoE VLMs only." + through Phi4MM, Gemma 3n, Gemma 4, Nemotron H Nano Omni, and Qwen3-Omni MoE VLMs only." )); } @@ -595,6 +618,115 @@ fn compute_phi4mm_audio_embeddings( Ok(Some(embeddings)) } +fn compute_gemma3n_audio_embeddings( + gemma3n: &mlxcel::vision::Gemma3nVLModel, + prompt_tokens: &mut Vec, + image_paths: &[PathBuf], + audio_path: &Path, + tokenizer: &MlxcelTokenizer, +) -> Result> { + if !gemma3n.supports_audio() { + return Err(anyhow::anyhow!( + "This Gemma3n checkpoint was loaded without audio parameters" + )); + } + if !image_paths.is_empty() { + let _ = mlxcel::vlm_prompt::apply_image_token_blocks( + prompt_tokens, + mlxcel::vlm_prompt::ImageTokenBlockInfo { + use_boi_eoi: true, + image_token_id: gemma3n.image_token_id, + mm_tokens_per_image: 256, + boi_token_id: gemma3n.boi_token_id, + eoi_token_id: gemma3n.eoi_token_id, + has_bos: true, + separator_token_id: None, + suffix_tokens: Vec::new(), + block_prefix_tokens: Vec::new(), + block_suffix_tokens: Vec::new(), + }, + image_paths.len(), + )?; + } + let (samples, sample_rate) = + mlxcel::audio::load_wav_file(audio_path).map_err(|error| anyhow::anyhow!(error))?; + let samples = if sample_rate == mlxcel::audio::gemma3n::GEMMA3N_SAMPLE_RATE { + samples + } else { + mlxcel::audio::whisper_mel::resample_to_16k(&samples, sample_rate) + }; + let batch = mlxcel::audio::gemma3n::Gemma3nAudioFeatureExtractor::new() + .extract_batch(&[samples]) + .map_err(|error| anyhow::anyhow!(error))?; + let wrapper_tokens: Vec = tokenizer + .encode("\n\n", false) + .map_err(|error| anyhow::anyhow!("Failed to tokenize Gemma3n audio wrapper: {error}"))? + .into_iter() + .map(|token| token as i32) + .collect(); + if wrapper_tokens.is_empty() { + return Err(anyhow::anyhow!( + "Gemma3n audio wrapper tokenized to an empty sequence" + )); + } + let end_of_turn = ["", ""].iter().find_map(|marker| { + tokenizer + .encode(marker, false) + .ok() + .filter(|tokens| tokens.len() == 1) + .map(|tokens| tokens[0] as i32) + }); + mlxcel::vlm_runtime::expand_gemma3n_audio_tokens( + prompt_tokens, + gemma3n.audio_token_id, + gemma3n.boa_token_id, + gemma3n.eoa_token_id, + 1, + gemma3n.audio_soft_tokens_per_clip, + &wrapper_tokens, + end_of_turn, + )?; + + let features = mlxcel_core::from_slice_f32(&batch.features, &[1, batch.frames as i32, 128]); + let valid: Vec = batch + .valid_mask + .iter() + .map(|valid| i32::from(*valid)) + .collect(); + let valid = mlxcel_core::astype( + &mlxcel_core::from_slice_i32(&valid, &[1, batch.frames as i32]), + mlxcel_core::dtype::BOOL, + ); + let invalid = mlxcel_core::logical_not(&valid); + let images = image_paths + .iter() + .map(|path| { + image::open(path) + .map_err(|error| anyhow::anyhow!("Failed to load image {:?}: {error}", path)) + }) + .collect::>>()?; + let pixel_values = if images.is_empty() { + None + } else { + Some(gemma3n.processor.preprocess(&images)) + }; + let input_ids = mlxcel_core::from_slice_i32(prompt_tokens, &[1, prompt_tokens.len() as i32]); + let embeddings = gemma3n + .get_input_embeddings_with_media( + &input_ids, + pixel_values.as_deref(), + Some(&features), + Some(&invalid), + ) + .map_err(|error| anyhow::anyhow!(error))?; + print_preparation_summary(VlmPreparationSummary::Gemma3nAudio { + audio_clips: 1, + audio_tokens: gemma3n.audio_soft_tokens_per_clip, + total_tokens: prompt_tokens.len(), + }); + Ok(Some(embeddings)) +} + /// Compute audio-only embeddings for Gemma4 VLM. fn compute_gemma4_audio_embeddings( gemma4_vl: &mlxcel::vision::Gemma4VLModel, diff --git a/src/loading/vlm_gemma.rs b/src/loading/vlm_gemma.rs index 2a6f510a..0afa00a3 100644 --- a/src/loading/vlm_gemma.rs +++ b/src/loading/vlm_gemma.rs @@ -42,6 +42,10 @@ struct Gemma3nMetadata { boi_token_id: i32, eoi_token_id: i32, vision_rms_eps: f32, + audio_token_id: i32, + boa_token_id: i32, + eoa_token_id: i32, + audio_soft_tokens_per_clip: usize, } fn gemma3n_metadata(config: &Value) -> Gemma3nMetadata { @@ -74,6 +78,23 @@ fn gemma3n_metadata(config: &Value) -> Gemma3nMetadata { .and_then(|vc| vc.get("rms_norm_eps")) .and_then(|v| v.as_f64()) .unwrap_or(1e-6) as f32, + audio_token_id: config + .get("audio_token_id") + .and_then(|value| value.as_i64()) + .unwrap_or(262_273) as i32, + boa_token_id: config + .get("boa_token_id") + .and_then(|value| value.as_i64()) + .unwrap_or(256_000) as i32, + eoa_token_id: config + .get("eoa_token_id") + .and_then(|value| value.as_i64()) + .unwrap_or(262_272) as i32, + audio_soft_tokens_per_clip: config + .get("audio_soft_tokens_per_image") + .and_then(|value| value.as_u64()) + .unwrap_or(crate::audio::gemma3n::GEMMA3N_AUDIO_SOFT_TOKENS as u64) + as usize, } } @@ -111,6 +132,8 @@ fn sanitize_gemma3n_weights(raw_weights: WeightMap) -> WeightMap { let shape = mlxcel_core::array_shape(&value); if shape.len() == 4 { mlxcel_core::transpose_axes(&value, &[0, 2, 3, 1]) + } else if shape.len() == 3 && new_key.ends_with(".lconv1d.depthwise_conv1d.weight") { + mlxcel_core::transpose_axes(&value, &[0, 2, 1]) } else { mlxcel_core::copy(&value) } @@ -253,7 +276,7 @@ pub(crate) fn load_gemma3n_vlm(model_path: &Path) -> Result { let processor = SigLipProcessor::new_rescale_only(metadata.image_size); - let vlm = vision::Gemma3nVLModel::new( + let mut vlm = vision::Gemma3nVLModel::new( text_model, vision_tower, embed_vision, @@ -264,6 +287,43 @@ pub(crate) fn load_gemma3n_vlm(model_path: &Path) -> Result { metadata.vision_hidden_size, ); + // Audio is optional at load time. Text/vision-only converted checkpoints + // can omit the entire audio namespace without failing model startup. + if let Some(audio_config_value) = full_config.get("audio_config") + && !audio_config_value.is_null() + && weights.contains_key("audio_tower.subsample_conv_projection.conv_0.conv.weight") + { + let audio_config: crate::audio::gemma3n::Gemma3nAudioConfig = + serde_json::from_value(audio_config_value.clone()).map_err(|error| { + anyhow::anyhow!("Failed to parse Gemma3n audio_config: {error}") + })?; + let audio_tower = crate::audio::gemma3n::Gemma3nAudioEncoder::from_weights( + &weights, + "audio_tower", + &audio_config, + group_size, + bits, + ) + .map_err(|error| anyhow::anyhow!("Failed to load Gemma3n audio tower: {error}"))?; + let embed_audio = models::gemma3n::Gemma3nAudioEmbedder::from_weights( + &weights, + "embed_audio", + &audio_config, + text_config.hidden_size, + group_size, + bits, + ) + .map_err(|error| anyhow::anyhow!("Failed to load Gemma3n audio embedder: {error}"))?; + vlm.set_audio( + audio_tower, + embed_audio, + metadata.audio_token_id, + metadata.boa_token_id, + metadata.eoa_token_id, + metadata.audio_soft_tokens_per_clip, + ); + } + Ok(LoadedModel::Gemma3nVLM(vlm)) } diff --git a/src/loading/vlm_gemma_tests.rs b/src/loading/vlm_gemma_tests.rs index 8479e250..a2a56b11 100644 --- a/src/loading/vlm_gemma_tests.rs +++ b/src/loading/vlm_gemma_tests.rs @@ -29,6 +29,10 @@ fn gemma3n_metadata_applies_defaults_and_overrides() { assert_eq!(defaults.boi_token_id, 255_999); assert_eq!(defaults.eoi_token_id, 262_144); assert!((defaults.vision_rms_eps - 1e-6).abs() < f32::EPSILON); + assert_eq!(defaults.audio_token_id, 262_273); + assert_eq!(defaults.boa_token_id, 256_000); + assert_eq!(defaults.eoa_token_id, 262_272); + assert_eq!(defaults.audio_soft_tokens_per_clip, 188); let overrides = gemma3n_metadata(&json!({ "vision_config": { @@ -39,12 +43,20 @@ fn gemma3n_metadata_applies_defaults_and_overrides() { "image_token_index": 9, "boi_token_id": 10, "eoi_token_id": 11 + ,"audio_token_id": 12, + "boa_token_id": 13, + "eoa_token_id": 14, + "audio_soft_tokens_per_image": 15 })); assert_eq!(overrides.vision_hidden_size, 3072); assert_eq!(overrides.image_size, 384); assert_eq!(overrides.image_token_id, 9); assert_eq!(overrides.boi_token_id, 10); assert_eq!(overrides.eoi_token_id, 11); + assert_eq!(overrides.audio_token_id, 12); + assert_eq!(overrides.boa_token_id, 13); + assert_eq!(overrides.eoa_token_id, 14); + assert_eq!(overrides.audio_soft_tokens_per_clip, 15); assert!((overrides.vision_rms_eps - 1e-5).abs() < f32::EPSILON); } @@ -95,6 +107,10 @@ fn sanitize_gemma3n_weights_strips_model_prefix_and_transposes_conv_weights() { "model.language_model.embed_tokens.weight".to_string(), mlxcel_core::ones(&[4, 4], dtype::FLOAT32), ); + raw_weights.insert( + "model.audio_tower.conformer.0.lconv1d.depthwise_conv1d.weight".to_string(), + mlxcel_core::ones(&[1536, 1, 5], dtype::FLOAT32), + ); let sanitized = sanitize_gemma3n_weights(raw_weights); @@ -108,6 +124,14 @@ fn sanitize_gemma3n_weights_strips_model_prefix_and_transposes_conv_weights() { ), vec![8, 3, 3, 16] ); + assert_eq!( + mlxcel_core::array_shape( + sanitized + .get("audio_tower.conformer.0.lconv1d.depthwise_conv1d.weight") + .unwrap() + ), + vec![1536, 5, 1] + ); } fn audio_conv_shape(weights: &WeightMap, key: &str) -> Vec { diff --git a/src/models/gemma3n.rs b/src/models/gemma3n.rs index 54ebe6df..9983f542 100644 --- a/src/models/gemma3n.rs +++ b/src/models/gemma3n.rs @@ -1747,3 +1747,175 @@ impl Gemma3nMultimodalEmbedder { }) } } + +/// Gemma 3n audio embedder with both the learned hard-token table and the +/// soft encoder projection. Audio padding uses the final hard-vocabulary row, +/// while real encoder frames use the separate soft normalization path. +pub struct Gemma3nAudioEmbedder { + embedding: UnifiedEmbedding, + hard_embedding_norm: RMSNorm, + soft_embedding_norm: RMSNorm, + embedding_projection: UnifiedLinear, + post_projection_norm: RMSNoScale, + vocab_size: usize, + vocab_offset: i32, +} + +impl Gemma3nAudioEmbedder { + pub fn from_weights( + weights: &WeightMap, + prefix: &str, + config: &crate::audio::gemma3n::Gemma3nAudioConfig, + text_hidden_size: usize, + group_size: i32, + bits: i32, + ) -> Result { + let embedding_prefix = format!("{prefix}.embedding"); + let weight_key = format!("{embedding_prefix}.weight"); + let weight = weights + .get(&weight_key) + .ok_or_else(|| format!("Gemma3n audio weight not found: {weight_key}"))?; + let shape = mlxcel_core::array_shape(weight); + let quantized = weights.contains_key(&format!("{embedding_prefix}.scales")); + if shape.len() != 2 + || shape[0] != config.vocab_size as i32 + || (!quantized && shape[1] != config.hidden_size as i32) + { + return Err(format!( + "{weight_key} has shape {shape:?}; expected logical [{}, {}]", + config.vocab_size, config.hidden_size + )); + } + if let Some(scales) = weights.get(&format!("{embedding_prefix}.scales")) { + let scales_shape = mlxcel_core::array_shape(scales); + if scales_shape.len() != 2 || scales_shape[0] != config.vocab_size as i32 { + return Err(format!( + "{embedding_prefix}.scales has invalid shape {scales_shape:?}" + )); + } + if let Some(biases) = weights.get(&format!("{embedding_prefix}.biases")) + && mlxcel_core::array_shape(biases) != scales_shape + { + return Err(format!( + "{embedding_prefix}.biases must match scales shape {scales_shape:?}" + )); + } + } + let embedding = + UnifiedEmbedding::from_weights(weights, &embedding_prefix, group_size, bits)?; + let hard_weight = + get_weight_copy(weights, &format!("{prefix}.hard_embedding_norm.weight"))?; + let soft_weight = + get_weight_copy(weights, &format!("{prefix}.soft_embedding_norm.weight"))?; + if mlxcel_core::array_shape(&hard_weight) != [config.hidden_size as i32] + || mlxcel_core::array_shape(&soft_weight) != [config.hidden_size as i32] + { + return Err(format!( + "{prefix} audio RMS norm weights must have {} elements", + config.hidden_size + )); + } + Ok(Self { + embedding, + hard_embedding_norm: RMSNorm::new(hard_weight, config.rms_norm_eps), + soft_embedding_norm: RMSNorm::new(soft_weight, config.rms_norm_eps), + embedding_projection: crate::audio::gemma3n::checked_unified_linear( + weights, + &format!("{prefix}.embedding_projection"), + config.hidden_size, + text_hidden_size, + group_size, + bits, + )?, + post_projection_norm: RMSNoScale::new(text_hidden_size as i32, config.rms_norm_eps), + vocab_size: config.vocab_size, + vocab_offset: config.vocab_offset, + }) + } + + fn project(&self, normalized: &MlxArray) -> UniquePtr { + self.post_projection_norm + .forward(&self.embedding_projection.forward(normalized)) + } + + pub fn forward_soft(&self, inputs: &MlxArray) -> UniquePtr { + self.project(&self.soft_embedding_norm.forward(inputs)) + } + + /// Replace the text-table rows for hard audio vocabulary tokens. The + /// official model performs this before replacing `` + /// rows with encoder output; notably `` remains a hard + /// audio embedding and must not use the text embedding table. + pub fn merge_hard_tokens( + &self, + input_ids: &MlxArray, + inputs_embeds: &MlxArray, + ) -> UniquePtr { + let offset = mlxcel_core::from_slice_i32(&[self.vocab_offset], &[1]); + let audio_mask = mlxcel_core::greater_equal(input_ids, &offset); + let dummy = + mlxcel_core::from_slice_i32(&[self.vocab_offset + self.vocab_size as i32 - 1], &[1]); + let global_ids = mlxcel_core::where_cond(&audio_mask, input_ids, &dummy); + let local_ids = mlxcel_core::subtract(&global_ids, &offset); + let hard = self.embedding.forward(&local_ids); + let hard = self.project(&self.hard_embedding_norm.forward(&hard)); + mlxcel_core::where_cond( + &mlxcel_core::expand_dims(&audio_mask, -1), + &hard, + inputs_embeds, + ) + } + + pub fn padding_embedding(&self) -> UniquePtr { + let token = mlxcel_core::from_slice_i32(&[(self.vocab_size - 1) as i32], &[1, 1]); + let embedded = self.embedding.forward(&token); + self.project(&self.hard_embedding_norm.forward(&embedded)) + } +} + +#[cfg(test)] +mod audio_embedder_tests { + use super::*; + + fn values(array: &MlxArray) -> Vec { + mlxcel_core::eval(array); + mlxcel_core::array_to_raw_bytes(array) + .chunks_exact(4) + .map(|bytes| f32::from_ne_bytes(bytes.try_into().unwrap())) + .collect() + } + + #[test] + fn hard_audio_vocabulary_replaces_only_tokens_at_or_above_offset() { + let config = crate::audio::gemma3n::Gemma3nAudioConfig { + vocab_size: 2, + vocab_offset: 100, + hidden_size: 2, + ..crate::audio::gemma3n::Gemma3nAudioConfig::default() + }; + let mut weights = WeightMap::new(); + weights.insert( + "embed.embedding.weight".into(), + mlxcel_core::from_slice_f32(&[1.0, 0.0, 0.0, 1.0], &[2, 2]), + ); + for name in ["hard_embedding_norm", "soft_embedding_norm"] { + weights.insert( + format!("embed.{name}.weight"), + mlxcel_core::ones(&[2], mlxcel_core::dtype::FLOAT32), + ); + } + weights.insert( + "embed.embedding_projection.weight".into(), + mlxcel_core::from_slice_f32(&[1.0, 0.0, 0.0, 1.0], &[2, 2]), + ); + let embedder = + Gemma3nAudioEmbedder::from_weights(&weights, "embed", &config, 2, 64, 4).unwrap(); + let ids = mlxcel_core::from_slice_i32(&[99, 100, 101], &[1, 3]); + let text = mlxcel_core::zeros(&[1, 3, 2], mlxcel_core::dtype::FLOAT32); + let merged = embedder.merge_hard_tokens(&ids, &text); + let output = values(&merged); + assert_eq!(&output[..2], &[0.0, 0.0]); + assert!(output[2] > 1.4 && output[3].abs() < 1e-6); + assert!(output[4].abs() < 1e-6 && output[5] > 1.4); + } +} diff --git a/src/multimodal/vlm_runtime.rs b/src/multimodal/vlm_runtime.rs index fbc4415f..69d3f4c8 100644 --- a/src/multimodal/vlm_runtime.rs +++ b/src/multimodal/vlm_runtime.rs @@ -77,6 +77,11 @@ pub enum VlmPreparationSummary { audio_tokens: usize, total_tokens: usize, }, + Gemma3nAudio { + audio_clips: usize, + audio_tokens: usize, + total_tokens: usize, + }, /// Gemma 4 expanded `<|video|>` placeholders for one or /// more decoded videos. `frame_slots` is the total number of frames /// across all videos (each frame consumes @@ -2029,6 +2034,67 @@ pub fn expand_gemma4_audio_tokens_for_server( } } +/// Expand Gemma 3n audio placeholders using the pinned processor's fixed +/// `boa + audio*188 + eoa` sequence for every clip. When the rendered template +/// dropped audio content parts, all clip blocks are inserted before the last +/// user-turn terminator. Existing placeholders must match the clip count +/// exactly; a mismatch is rejected instead of silently reordering media. +pub fn expand_gemma3n_audio_tokens( + prompt_tokens: &mut Vec, + audio_token_id: i32, + boa_token_id: i32, + eoa_token_id: i32, + audio_clips: usize, + soft_tokens_per_clip: usize, + wrapper_tokens: &[i32], + end_of_turn_token_id: Option, +) -> Result<()> { + if audio_clips == 0 { + return Ok(()); + } + let block = || { + let mut tokens = Vec::with_capacity(wrapper_tokens.len() * 2 + soft_tokens_per_clip + 2); + tokens.extend_from_slice(wrapper_tokens); + tokens.push(boa_token_id); + tokens.extend(std::iter::repeat_n(audio_token_id, soft_tokens_per_clip)); + tokens.push(eoa_token_id); + tokens.extend_from_slice(wrapper_tokens); + tokens + }; + let placeholders = prompt_tokens + .iter() + .filter(|token| **token == audio_token_id) + .count(); + if placeholders > 0 { + if placeholders != audio_clips { + anyhow::bail!( + "Gemma3n audio placeholder mismatch: prompt has {placeholders}, request has {audio_clips} clips" + ); + } + let mut expanded = + Vec::with_capacity(prompt_tokens.len() + audio_clips * (soft_tokens_per_clip + 2)); + for token in prompt_tokens.iter().copied() { + if token == audio_token_id { + expanded.extend(block()); + } else { + expanded.push(token); + } + } + *prompt_tokens = expanded; + return Ok(()); + } + + let insertion = end_of_turn_token_id + .and_then(|eot| prompt_tokens.iter().rposition(|token| *token == eot)) + .unwrap_or(prompt_tokens.len()); + let mut blocks = Vec::with_capacity(audio_clips * (soft_tokens_per_clip + 2)); + for _ in 0..audio_clips { + blocks.extend(block()); + } + prompt_tokens.splice(insertion..insertion, blocks); + Ok(()) +} + /// Place the Nemotron H Nano Omni sound-context block in the server prompt. /// /// The server renders chat messages text-only (see diff --git a/src/multimodal/vlm_runtime_tests.rs b/src/multimodal/vlm_runtime_tests.rs index 7688d50c..434367ad 100644 --- a/src/multimodal/vlm_runtime_tests.rs +++ b/src/multimodal/vlm_runtime_tests.rs @@ -13,11 +13,100 @@ // limitations under the License. use super::{ - VlmPreparationSummary, expand_gemma4_audio_tokens_for_server, + VlmPreparationSummary, expand_gemma3n_audio_tokens, expand_gemma4_audio_tokens_for_server, expand_gemma4_unified_video_tokens, expand_gemma4_video_tokens, expand_nemotron_h_nano_omni_audio_tokens_for_server, format_molmo_v1_prompt_for_processor, prepared_embedding_refs, shift_molmo_v1_image_input_idx_for_bos, should_prepare_vlm_embeddings, }; +use crate::vlm_prompt::{ImageTokenBlockInfo, apply_image_token_blocks}; + +#[test] +fn gemma3n_audio_expands_multiple_placeholders_in_order() { + const AUDIO: i32 = 262_273; + let mut prompt = vec![10, AUDIO, 20, AUDIO, 106]; + expand_gemma3n_audio_tokens( + &mut prompt, + AUDIO, + 256_000, + 262_272, + 2, + 3, + &[108], + Some(106), + ) + .unwrap(); + assert_eq!( + prompt, + vec![ + 10, 108, 256_000, AUDIO, AUDIO, AUDIO, 262_272, 108, 20, 108, 256_000, AUDIO, AUDIO, + AUDIO, 262_272, 108, 106, + ] + ); +} + +#[test] +fn gemma3n_audio_inserts_all_clips_inside_user_turn() { + const AUDIO: i32 = 262_273; + let mut prompt = vec![1, 2, 106, 3]; + expand_gemma3n_audio_tokens(&mut prompt, AUDIO, 256_000, 262_272, 2, 2, &[], Some(106)) + .unwrap(); + assert_eq!( + prompt, + vec![ + 1, 2, 256_000, AUDIO, AUDIO, 262_272, 256_000, AUDIO, AUDIO, 262_272, 106, 3, + ] + ); +} + +#[test] +fn gemma3n_audio_rejects_placeholder_clip_mismatch() { + let mut prompt = vec![1, 262_273, 2]; + let error = + expand_gemma3n_audio_tokens(&mut prompt, 262_273, 256_000, 262_272, 2, 188, &[], None) + .unwrap_err(); + assert!(error.to_string().contains("placeholder mismatch")); +} + +#[test] +fn gemma3n_mixed_media_preserves_image_then_audio_order() { + const IMAGE: i32 = 262_145; + const AUDIO: i32 = 262_273; + let mut prompt = vec![2, IMAGE, 9, AUDIO, 106]; + apply_image_token_blocks( + &mut prompt, + ImageTokenBlockInfo { + use_boi_eoi: true, + image_token_id: IMAGE, + mm_tokens_per_image: 2, + boi_token_id: 255_999, + eoi_token_id: 262_144, + has_bos: true, + separator_token_id: None, + suffix_tokens: Vec::new(), + block_prefix_tokens: Vec::new(), + block_suffix_tokens: Vec::new(), + }, + 1, + ) + .unwrap(); + expand_gemma3n_audio_tokens( + &mut prompt, + AUDIO, + 256_000, + 262_272, + 1, + 2, + &[108], + Some(106), + ) + .unwrap(); + assert_eq!( + prompt, + vec![ + 2, 255_999, IMAGE, IMAGE, 262_144, 9, 108, 256_000, AUDIO, AUDIO, 262_272, 108, 106, + ] + ); +} use crate::vision::merge::InputEmbeddings; use crate::vlm_prompt::{ImageTokenBlockAction, ImageTokenBlockStats}; use mlxcel_core::{self, UniquePtr, dtype}; diff --git a/src/server/model_worker.rs b/src/server/model_worker.rs index 66d2bd52..18bf2108 100644 --- a/src/server/model_worker.rs +++ b/src/server/model_worker.rs @@ -39,6 +39,7 @@ use crate::server::state::BatchMetrics; use crate::tokenizer::MlxcelTokenizer; use crate::vision::feature_cache::ModelVisionCaches; use crate::vision::merge::InputEmbeddings; +use crate::vision::processors::ImageProcessor; use crate::vlm_runtime::{ prepare_and_compute_vlm_embeddings_with_budget, prepare_and_compute_vlm_embeddings_with_cache, }; @@ -1179,6 +1180,16 @@ pub(crate) fn prepare_request_vlm_embeddings( // the last user turn instead of the model turn (which forces an // immediate EOS / 0-token output). let end_of_turn_token_id = resolve_end_of_turn_token_id(tokenizer); + if let Some(embeddings) = prepare_gemma3n_audio_embeddings( + model, + tokenizer, + prompt_tokens, + images, + audio, + end_of_turn_token_id, + )? { + return Ok(Some(embeddings)); + } if let Some(embeddings) = prepare_gemma4_unified_audio_embeddings( model, prompt_tokens, @@ -1402,6 +1413,103 @@ fn prepare_phi4mm_audio_embeddings( Ok(Some(embeddings)) } +/// Process all Gemma 3n clips as one padded batch, preserving clip order and +/// merging exactly 188 projected rows for each prompt placeholder block. +fn prepare_gemma3n_audio_embeddings( + model: &LoadedModel, + tokenizer: &MlxcelTokenizer, + prompt_tokens: &mut Vec, + images: &[Vec], + audio_data: &[Vec], + end_of_turn_token_id: Option, +) -> Result> { + let gemma3n = match model { + LoadedModel::Gemma3nVLM(model) => model, + _ => return Ok(None), + }; + if !gemma3n.supports_audio() { + return Err(anyhow!( + "This Gemma3n checkpoint was loaded without audio parameters" + )); + } + if !images.is_empty() + && let Some(info) = model.image_token_block_info() + { + let _ = crate::vlm_prompt::apply_image_token_blocks(prompt_tokens, info, images.len())?; + } + + let mut clips = Vec::with_capacity(audio_data.len()); + for (index, bytes) in audio_data.iter().enumerate() { + let (samples, sample_rate) = crate::audio::load_wav_from_bytes(bytes) + .map_err(|error| anyhow!("Failed to decode Gemma3n audio clip {index}: {error}"))?; + let samples = if sample_rate == crate::audio::gemma3n::GEMMA3N_SAMPLE_RATE { + samples + } else { + crate::audio::whisper_mel::resample_to_16k(&samples, sample_rate) + }; + clips.push(samples); + } + + let batch = crate::audio::gemma3n::Gemma3nAudioFeatureExtractor::new() + .extract_batch(&clips) + .map_err(|error| anyhow!(error))?; + let wrapper_tokens: Vec = tokenizer + .encode("\n\n", false) + .map_err(|error| anyhow!("Failed to tokenize Gemma3n audio wrapper: {error}"))? + .into_iter() + .map(|token| token as i32) + .collect(); + if wrapper_tokens.is_empty() { + return Err(anyhow!( + "Gemma3n audio wrapper tokenized to an empty sequence" + )); + } + crate::vlm_runtime::expand_gemma3n_audio_tokens( + prompt_tokens, + gemma3n.audio_token_id, + gemma3n.boa_token_id, + gemma3n.eoa_token_id, + clips.len(), + gemma3n.audio_soft_tokens_per_clip, + &wrapper_tokens, + end_of_turn_token_id, + )?; + + let features = mlxcel_core::from_slice_f32( + &batch.features, + &[batch.batch_size as i32, batch.frames as i32, 128], + ); + let valid: Vec = batch + .valid_mask + .iter() + .map(|valid| i32::from(*valid)) + .collect(); + let valid = mlxcel_core::astype( + &mlxcel_core::from_slice_i32(&valid, &[batch.batch_size as i32, batch.frames as i32]), + mlxcel_core::dtype::BOOL, + ); + let invalid = mlxcel_core::logical_not(&valid); + + let decoded_images = if images.is_empty() { + None + } else { + Some(decode_request_images(images)?) + }; + let pixel_values = decoded_images + .as_ref() + .map(|images| gemma3n.processor.preprocess(images)); + let input_ids = mlxcel_core::from_slice_i32(prompt_tokens, &[1, prompt_tokens.len() as i32]); + let embeddings = gemma3n + .get_input_embeddings_with_media( + &input_ids, + pixel_values.as_deref(), + Some(&features), + Some(&invalid), + ) + .map_err(|error| anyhow!(error))?; + Ok(Some(embeddings)) +} + /// Resolve the per-family end-of-turn token id from the tokenizer. /// /// The server flattens chat messages to text-only, so an `input_audio` request diff --git a/src/vision/gemma3n_per_layer_inputs_state.rs b/src/vision/gemma3n_per_layer_inputs_state.rs index de86a12f..c49cf690 100644 --- a/src/vision/gemma3n_per_layer_inputs_state.rs +++ b/src/vision/gemma3n_per_layer_inputs_state.rs @@ -43,11 +43,10 @@ //! into this map under the scheduler-allocated id, draining the //! fallback in the same step so the next request cannot inherit the //! previous row's tensor. -//! - **Fallback slot** preserves legacy single-instance callers (CLI -//! `mlxcel generate`, `mlxcel-bench-decode`, single-row tests). It also -//! acts as a last-resort when `take_for_sequence(id)` finds no entry -//! under `id` (e.g., a request that started before per-sequence wiring -//! landed). +//! - **Fallback slot** preserves only legacy single-instance callers (CLI +//! `mlxcel generate`, `mlxcel-bench-decode`, single-row tests). A request +//! carrying a `SequenceId` never reads this slot, because it could belong to +//! a different concurrently-prepared request. use std::cell::RefCell; use std::collections::HashMap; @@ -103,9 +102,8 @@ impl Gemma3nPerLayerInputsState { } /// Take a row's per-layer-inputs out of the map for prefill - /// consumption. Returns `None` when no entry exists — the prefill - /// path then falls back to the legacy single-row slot via - /// [`Self::take_fallback`]. + /// consumption. Returns `None` when no entry exists; sequence-scoped + /// consumers treat that as an invariant violation and never fall back. pub(crate) fn take_for_sequence(&self, seq_id: SequenceId) -> Option> { self.sequences.borrow_mut().remove(&seq_id) } diff --git a/src/vision/gemma3n_per_layer_inputs_state_tests.rs b/src/vision/gemma3n_per_layer_inputs_state_tests.rs index f677b922..0e967363 100644 --- a/src/vision/gemma3n_per_layer_inputs_state_tests.rs +++ b/src/vision/gemma3n_per_layer_inputs_state_tests.rs @@ -218,3 +218,15 @@ fn take_for_sequence_returns_none_for_unknown_id() { let unknown = SequenceId::from_raw(401); assert!(state.take_for_sequence(unknown).is_none()); } + +#[test] +#[ignore = "requires serial MLX execution"] +fn unknown_sequence_cannot_consume_another_requests_fallback() { + let state = Gemma3nPerLayerInputsState::new(); + state.set_fallback(Some(make_pli(77.0))); + assert!(state.take_for_sequence(SequenceId::from_raw(999)).is_none()); + let fallback = state + .take_fallback() + .expect("unknown sequence must not drain fallback"); + assert!((read_pli_scalar(fallback.as_ref().unwrap()) - 77.0).abs() < 1e-5); +} diff --git a/src/vision/gemma3n_vl.rs b/src/vision/gemma3n_vl.rs index c6228275..b9fff0e4 100644 --- a/src/vision/gemma3n_vl.rs +++ b/src/vision/gemma3n_vl.rs @@ -51,6 +51,12 @@ pub struct Gemma3nVLModel { pub boi_token_id: i32, // 255_999 () pub eoi_token_id: i32, // 262_144 () pub vision_hidden_size: usize, // 2048 + pub audio_tower: Option, + pub embed_audio: Option, + pub audio_token_id: i32, + pub boa_token_id: i32, + pub eoa_token_id: i32, + pub audio_soft_tokens_per_clip: usize, /// Per-sequence storage for the projected `per_layer_inputs` tensor /// that is produced during embedding prep and consumed during the /// prefill `forward_with_embeddings_and_sequence_id` call. See module @@ -78,10 +84,37 @@ impl Gemma3nVLModel { boi_token_id, eoi_token_id, vision_hidden_size, + audio_tower: None, + embed_audio: None, + audio_token_id: 262_273, + boa_token_id: 256_000, + eoa_token_id: 262_272, + audio_soft_tokens_per_clip: crate::audio::gemma3n::GEMMA3N_AUDIO_SOFT_TOKENS, per_layer_inputs_state: Gemma3nPerLayerInputsState::new(), } } + pub fn set_audio( + &mut self, + audio_tower: crate::audio::gemma3n::Gemma3nAudioEncoder, + embed_audio: crate::models::gemma3n::Gemma3nAudioEmbedder, + audio_token_id: i32, + boa_token_id: i32, + eoa_token_id: i32, + audio_soft_tokens_per_clip: usize, + ) { + self.audio_tower = Some(audio_tower); + self.embed_audio = Some(embed_audio); + self.audio_token_id = audio_token_id; + self.boa_token_id = boa_token_id; + self.eoa_token_id = eoa_token_id; + self.audio_soft_tokens_per_clip = audio_soft_tokens_per_clip; + } + + pub fn supports_audio(&self) -> bool { + self.audio_tower.is_some() && self.embed_audio.is_some() + } + /// Get input embeddings with vision features merged in. /// /// Side effect: writes the freshly projected `per_layer_inputs` tensor @@ -96,51 +129,142 @@ impl Gemma3nVLModel { input_ids: &MlxArray, pixel_values: &MlxArray, ) -> merge::InputEmbeddings { + self.get_input_embeddings_with_media(input_ids, Some(pixel_values), None, None) + .expect("Gemma3n image embedding preparation failed") + } + + /// Prepare text embeddings and merge any supplied image and/or audio + /// features. Audio masks use `true = invalid`, matching the encoder. + pub fn get_input_embeddings_with_media( + &self, + input_ids: &MlxArray, + pixel_values: Option<&MlxArray>, + audio_features: Option<&MlxArray>, + invalid_audio_mask: Option<&MlxArray>, + ) -> Result { // 1. Text embeddings - let inputs_embeds = self.text_model.language_model.get_embed_tokens(input_ids); + let mut inputs_embeds = self.text_model.language_model.get_embed_tokens(input_ids); + if let Some(embed_audio) = &self.embed_audio { + inputs_embeds = embed_audio.merge_hard_tokens(input_ids, &inputs_embeds); + } // 2. Per-layer inputs (image_token_id >= vocab_size_per_layer, auto-zeroed) let per_layer_inputs = self .text_model .language_model .get_per_layer_inputs(input_ids); + + let mut merged = merge::InputEmbeddings { + inputs_embeds, + attention_mask_4d: None, + }; + + // 3. Optional vision path. + if let Some(pixel_values) = pixel_values { + let embed_dtype = mlxcel_core::array_dtype(&merged.inputs_embeds); + let pv = mlxcel_core::astype(pixel_values, embed_dtype); + let vision_out = self.vision_tower.forward(&pv); + let vo = mlxcel_core::transpose_axes(&vision_out, &[0, 3, 1, 2]); + let vo_shape = mlxcel_core::array_shape(&vo); + let (batch, channels) = (vo_shape[0], vo_shape[1]); + let patches = vo_shape[2] * vo_shape[3]; + let vo = mlxcel_core::reshape(&vo, &[batch, channels, patches]); + let vo = mlxcel_core::transpose_axes(&vo, &[0, 2, 1]); + let vo = mlxcel_core::multiply_scalar(&vo, (self.vision_hidden_size as f32).sqrt()); + let image_features = self.embed_vision.forward_soft(&vo); + merged = merge::merge_llava( + self.image_token_id, + &image_features, + &merged.inputs_embeds, + input_ids, + ); + } + + // 4. Optional audio tower, projection, fixed-length padding and merge. + match (audio_features, invalid_audio_mask) { + (Some(audio_features), Some(invalid_audio_mask)) => { + let audio_tower = self + .audio_tower + .as_ref() + .ok_or_else(|| "Gemma3n checkpoint has no audio tower".to_string())?; + let embed_audio = self + .embed_audio + .as_ref() + .ok_or_else(|| "Gemma3n checkpoint has no audio embedder".to_string())?; + let padding_embedding = embed_audio.padding_embedding(); + let input_shape = mlxcel_core::array_shape(audio_features); + let mask_shape = mlxcel_core::array_shape(invalid_audio_mask); + if input_shape.len() != 3 || input_shape[2] != 128 || mask_shape != input_shape[..2] + { + return Err(format!( + "Gemma3n audio input shapes must be [B,T,128] and [B,T], got {input_shape:?} and {mask_shape:?}" + )); + } + let batch_size = input_shape[0]; + let projected = if input_shape[1] == 0 { + // The pinned processor returns zero mel frames for clips + // shorter than its 513-sample unfold. The reference model + // consequently emits only the fixed hard-padding rows. + let hidden = mlxcel_core::array_shape(&padding_embedding)[2]; + mlxcel_core::broadcast_to( + &padding_embedding, + &[batch_size, self.audio_soft_tokens_per_clip as i32, hidden], + ) + } else { + let (encoded, encoded_invalid_mask) = + audio_tower.forward(audio_features, invalid_audio_mask)?; + let mut projected = embed_audio.forward_soft(&encoded); + let shape = mlxcel_core::array_shape(&projected); + if shape[1] > self.audio_soft_tokens_per_clip as i32 { + return Err(format!( + "Gemma3n audio encoder emitted {} tokens, exceeding the fixed {}-token contract", + shape[1], self.audio_soft_tokens_per_clip + )); + } + let invalid = + mlxcel_core::reshape(&encoded_invalid_mask, &[shape[0], shape[1], 1]); + projected = mlxcel_core::where_cond(&invalid, &padding_embedding, &projected); + let missing = self.audio_soft_tokens_per_clip as i32 - shape[1]; + if missing > 0 { + let padding = mlxcel_core::broadcast_to( + &padding_embedding, + &[shape[0], missing, shape[2]], + ); + projected = mlxcel_core::concatenate(&projected, &padding, 1); + } + projected + }; + + let expected_tokens = batch_size as usize * self.audio_soft_tokens_per_clip; + let actual_tokens = count_i32_token(input_ids, self.audio_token_id)?; + if actual_tokens != expected_tokens { + return Err(format!( + "Gemma3n audio features and placeholders do not match: {expected_tokens} features, {actual_tokens} tokens" + )); + } + merged = merge::merge_llava( + self.audio_token_id, + &projected, + &merged.inputs_embeds, + input_ids, + ); + } + (None, None) => {} + _ => { + return Err( + "Gemma3n audio features and invalid mask must be provided together".into(), + ); + } + } + + // The official multimodal forward projects PLE after hard/soft media + // embeddings have replaced their placeholder rows. Projecting from + // the initial text-table rows gives image/audio positions the wrong + // router input. let per_layer_inputs = self .text_model .language_model - .project_per_layer_inputs(&inputs_embeds, &per_layer_inputs); - - // 3. Vision: pixel_values → VisionTower → [B, H, W, C] (NHWC) - let embed_dtype = mlxcel_core::array_dtype(&inputs_embeds); - let pv = mlxcel_core::astype(pixel_values, embed_dtype); - let vision_out = self.vision_tower.forward(&pv); - - // Reshape NHWC → [B, num_patches, hidden_size] - let vo = mlxcel_core::transpose_axes(&vision_out, &[0, 3, 1, 2]); - let vo_shape = mlxcel_core::array_shape(&vo); - let b = vo_shape[0]; - let c = vo_shape[1]; // hidden_size (2048) - let num_patches = vo_shape[2] * vo_shape[3]; // H*W - let vo = mlxcel_core::reshape(&vo, &[b, c, num_patches]); - let vo = mlxcel_core::transpose_axes(&vo, &[0, 2, 1]); // [B, num_patches, hidden_size] - - // Scale by sqrt(vision_hidden_size) - let scale = mlxcel_core::full_f32( - &[1], - (self.vision_hidden_size as f32).sqrt(), - mlxcel_core::dtype::FLOAT32, - ); - let vo = mlxcel_core::multiply(&vo, &scale); - - // 4. Multimodal embedder: → [B, num_patches, text_hidden] - let image_features = self.embed_vision.forward_soft(&vo); - - // 5. masked_scatter merge - let merged = merge::merge_llava( - self.image_token_id, - &image_features, - &inputs_embeds, - input_ids, - ); + .project_per_layer_inputs(&merged.inputs_embeds, &per_layer_inputs); // 6. Park the projected per_layer_inputs in the state container's // fallback slot. The scheduler's @@ -151,7 +275,7 @@ impl Gemma3nVLModel { self.per_layer_inputs_state .set_fallback(Some(per_layer_inputs)); - merged + Ok(merged) } // -- Per-sequence per_layer_inputs (issue #85) -------------------- @@ -201,15 +325,12 @@ impl Gemma3nVLModel { } } - /// Internal helper: resolve the active per_layer_inputs tensor for a - /// VLM prefill. Prefers the per-`SequenceId` slot (server flow) and - /// falls back to the fallback slot (CLI / bench / single-row). + /// Resolve the active per-layer tensor. A server sequence must never + /// consume the unbound fallback slot: doing so can reuse another + /// concurrent request's PLE state. fn take_per_layer_inputs(&self, seq_id: Option) -> Option> { match seq_id { - Some(id) => self - .per_layer_inputs_state - .take_for_sequence(id) - .or_else(|| self.per_layer_inputs_state.take_fallback()), + Some(id) => self.per_layer_inputs_state.take_for_sequence(id), None => self.per_layer_inputs_state.take_fallback(), } } @@ -243,6 +364,18 @@ impl Gemma3nVLModel { } } +fn count_i32_token(input_ids: &MlxArray, token_id: i32) -> Result { + if mlxcel_core::array_dtype(input_ids) != mlxcel_core::dtype::INT32 { + return Err("Gemma3n input ids must use int32 dtype".into()); + } + mlxcel_core::eval(input_ids); + let bytes = mlxcel_core::array_to_raw_bytes(input_ids); + Ok(bytes + .chunks_exact(std::mem::size_of::()) + .filter(|bytes| i32::from_ne_bytes((*bytes).try_into().unwrap()) == token_id) + .count()) +} + impl LanguageModel for Gemma3nVLModel { fn forward( &self,