From 5f475332038d8c5bfac034d2addc9d95fb60b85c Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 24 Jul 2026 14:01:36 +0900 Subject: [PATCH 01/19] feat(xla): define Gemma3n audio graph contract Add fail-closed audio configuration and artifact identities for static frame buckets. Validate canonical mel masks, fixed placeholder spans, and owned prepared-prefill metadata before requests cross the IREE boundary. --- .../mlxcel-xla/src/gemma3n_audio_config.rs | 416 ++++++++++++++++ src/lib/mlxcel-xla/src/lib.rs | 11 + .../mlxcel-xla/src/prepared_gemma3n_audio.rs | 467 ++++++++++++++++++ 3 files changed, 894 insertions(+) create mode 100644 src/lib/mlxcel-xla/src/gemma3n_audio_config.rs create mode 100644 src/lib/mlxcel-xla/src/prepared_gemma3n_audio.rs diff --git a/src/lib/mlxcel-xla/src/gemma3n_audio_config.rs b/src/lib/mlxcel-xla/src/gemma3n_audio_config.rs new file mode 100644 index 000000000..23d104cb8 --- /dev/null +++ b/src/lib/mlxcel-xla/src/gemma3n_audio_config.rs @@ -0,0 +1,416 @@ +//! Gemma3n audio graph configuration and artifact identity. +//! +//! Waveform decoding, resampling, and log-mel extraction stay on the bounded +//! host stage. This module defines the exact static shape and numeric policy +//! compiled into `audio.main`. + +use std::path::Path; + +use serde::Deserialize; + +pub const GEMMA3N_AUDIO_GRAPH_ABI: &str = "gemma3n-audio-main-v1"; +pub const GEMMA3N_AUDIO_MODALITY_FAMILY: &str = "gemma3n_audio"; +pub const GEMMA3N_AUDIO_MEL_BINS: usize = 128; +pub const GEMMA3N_AUDIO_MAX_FRAMES: usize = 2_997; +pub const GEMMA3N_AUDIO_SOFT_TOKENS: usize = 188; +pub const GEMMA3N_AUDIO_MAX_CLIPS: usize = 4; +pub const GEMMA3N_AUDIO_FRAME_BUCKETS: &[usize] = &[8, 32, 128, 512, 1_024, 2_048, 2_997]; + +fn default_hidden_size() -> usize { + 1_536 +} + +fn default_input_feat_size() -> usize { + GEMMA3N_AUDIO_MEL_BINS +} + +fn default_vocab_size() -> usize { + 128 +} + +fn default_vocab_offset() -> i32 { + 262_272 +} + +fn default_rms_norm_eps() -> f32 { + 1e-6 +} + +fn default_gradient_clipping() -> f32 { + 1e10 +} + +fn default_chunk_size() -> usize { + 12 +} + +fn default_context_left() -> usize { + 13 +} + +fn default_logit_cap() -> f32 { + 50.0 +} + +fn default_attention_heads() -> usize { + 8 +} + +fn default_hidden_layers() -> usize { + 12 +} + +fn default_conv_kernel_size() -> usize { + 5 +} + +fn default_reduction_factor() -> usize { + 4 +} + +fn default_residual_weight() -> f32 { + 0.5 +} + +fn default_conv_channels() -> Vec { + vec![128, 32] +} + +fn default_group_norm_eps() -> f32 { + 1e-3 +} + +fn default_conv_kernels() -> Vec<[usize; 2]> { + vec![[3, 3], [3, 3]] +} + +fn default_conv_strides() -> Vec<[usize; 2]> { + vec![[2, 2], [2, 2]] +} + +/// Audio sub-configuration consumed by the StableHLO emitter. +#[derive(Clone, Debug, Deserialize, PartialEq)] +#[serde(default)] +pub struct Gemma3nXlaAudioConfig { + pub input_feat_size: usize, + pub hidden_size: usize, + pub vocab_size: usize, + pub vocab_offset: i32, + 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 Gemma3nXlaAudioConfig { + fn default() -> Self { + Self { + input_feat_size: default_input_feat_size(), + hidden_size: default_hidden_size(), + vocab_size: default_vocab_size(), + vocab_offset: default_vocab_offset(), + rms_norm_eps: default_rms_norm_eps(), + gradient_clipping: default_gradient_clipping(), + conf_attention_chunk_size: default_chunk_size(), + conf_attention_context_left: default_context_left(), + conf_attention_context_right: 0, + conf_attention_logit_cap: default_logit_cap(), + conf_num_attention_heads: default_attention_heads(), + conf_num_hidden_layers: default_hidden_layers(), + conf_conv_kernel_size: default_conv_kernel_size(), + conf_reduction_factor: default_reduction_factor(), + conf_residual_weight: default_residual_weight(), + sscp_conv_channel_size: default_conv_channels(), + sscp_conv_group_norm_eps: default_group_norm_eps(), + sscp_conv_kernel_size: default_conv_kernels(), + sscp_conv_stride_size: default_conv_strides(), + } + } +} + +impl Gemma3nXlaAudioConfig { + pub fn from_model_dir(model_dir: &Path) -> Result, String> { + let path = model_dir.join("config.json"); + let text = std::fs::read_to_string(&path) + .map_err(|error| format!("read {}: {error}", path.display()))?; + Self::from_json_str(&text).map_err(|error| format!("{}: {error}", path.display())) + } + + pub fn from_json_str(text: &str) -> Result, String> { + let root: serde_json::Value = + serde_json::from_str(text).map_err(|error| format!("parse config JSON: {error}"))?; + let Some(value) = root.get("audio_config") else { + return Ok(None); + }; + if value.is_null() { + return Ok(None); + } + if value.get("model_type").and_then(serde_json::Value::as_str) != Some("gemma3n_audio") { + return Err("Gemma3n audio_config.model_type must be `gemma3n_audio`".into()); + } + let config: Self = serde_json::from_value(value.clone()) + .map_err(|error| format!("parse Gemma3n audio_config: {error}"))?; + config.validate()?; + Ok(Some(config)) + } + + pub fn validate(&self) -> Result<(), String> { + if self.input_feat_size != GEMMA3N_AUDIO_MEL_BINS { + return Err(format!( + "Gemma3n audio input_feat_size={} is unsupported; expected {GEMMA3N_AUDIO_MEL_BINS}", + self.input_feat_size + )); + } + 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 + || self.conf_num_hidden_layers == 0 + { + return Err( + "Gemma3n audio chunk, reduction, convolution, and layer counts 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.sscp_conv_channel_size.contains(&0) + || self + .sscp_conv_kernel_size + .iter() + .flatten() + .any(|&value| value == 0) + || self + .sscp_conv_stride_size + .iter() + .flatten() + .any(|&value| value == 0) + { + return Err("Gemma3n audio SSCP dimensions and strides must be positive".into()); + } + if self.vocab_size == 0 + || self.vocab_offset < 0 + || !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 scalar configuration is invalid".into()); + } + let output_frequency = self.output_frequency()?; + self.sscp_conv_channel_size[1] + .checked_mul(output_frequency) + .ok_or_else(|| "Gemma3n audio SSCP projection width overflows".to_string())?; + Ok(()) + } + + #[must_use] + pub fn head_dim(&self) -> usize { + self.hidden_size / self.conf_num_attention_heads + } + + #[must_use] + pub fn context_size(&self) -> usize { + self.conf_attention_context_left.saturating_sub(1) + + self.conf_attention_chunk_size + + self.conf_attention_context_right + } + + pub fn time_stride_product(&self) -> Result { + self.sscp_conv_stride_size + .iter() + .try_fold(1usize, |product, stride| { + product + .checked_mul(stride[0]) + .ok_or_else(|| "Gemma3n audio time stride product overflows".to_string()) + }) + } + + pub fn output_frequency(&self) -> Result { + let mut frequency = self.input_feat_size; + for (kernel, stride) in self + .sscp_conv_kernel_size + .iter() + .zip(&self.sscp_conv_stride_size) + { + frequency = frequency + .checked_add(2) + .filter(|&padded| padded >= kernel[1]) + .ok_or_else(|| { + "Gemma3n audio SSCP frequency kernel exceeds padded input".to_string() + })?; + frequency = (frequency - kernel[1]) / stride[1] + 1; + } + Ok(frequency) + } + + pub fn encoded_frames(&self, frame_bucket: usize) -> Result { + validate_gemma3n_audio_frame_bucket(frame_bucket)?; + let mut time = frame_bucket; + for (kernel, stride) in self + .sscp_conv_kernel_size + .iter() + .zip(&self.sscp_conv_stride_size) + { + time = time + .checked_add(2) + .filter(|&padded| padded >= kernel[0]) + .ok_or_else(|| "Gemma3n audio SSCP time kernel exceeds padded input".to_string())?; + time = (time - kernel[0]) / stride[0] + 1; + } + Ok(time) + } + + pub fn projected_frames(&self, frame_bucket: usize) -> Result { + Ok(self + .encoded_frames(frame_bucket)? + .div_ceil(self.conf_reduction_factor)) + } + + pub fn artifact_identity( + &self, + frame_bucket: usize, + clips: usize, + text_hidden: usize, + text_layers: usize, + text_ple_hidden: usize, + ) -> Result { + validate_gemma3n_audio_batch(frame_bucket, clips)?; + if text_hidden == 0 || text_layers == 0 || text_ple_hidden == 0 { + return Err("Gemma3n audio language dimensions must be positive".into()); + } + Ok(format!( + "{GEMMA3N_AUDIO_GRAPH_ABI}:dtype=bf16-accum-f32:mel={}:bucket={frame_bucket}:clips={clips}:\ + sscp-kernels={:?}:strides={:?}:channels={:?}:group-eps={:08x}:hidden={}:heads={}:layers={}:\ + chunk={}:left={}:right={}:softcap={:08x}:conv={}:reduction={}:residual={:08x}:\ + clip={:08x}:rms={:08x}:audio-vocab={}:offset={}:soft-tokens={}:text={}:{}/{}", + self.input_feat_size, + self.sscp_conv_kernel_size, + self.sscp_conv_stride_size, + self.sscp_conv_channel_size, + self.sscp_conv_group_norm_eps.to_bits(), + self.hidden_size, + self.conf_num_attention_heads, + self.conf_num_hidden_layers, + self.conf_attention_chunk_size, + self.conf_attention_context_left, + self.conf_attention_context_right, + self.conf_attention_logit_cap.to_bits(), + self.conf_conv_kernel_size, + self.conf_reduction_factor, + self.conf_residual_weight.to_bits(), + self.gradient_clipping.to_bits(), + self.rms_norm_eps.to_bits(), + self.vocab_size, + self.vocab_offset, + GEMMA3N_AUDIO_SOFT_TOKENS, + text_hidden, + text_layers, + text_ple_hidden, + )) + } +} + +pub(crate) fn validate_gemma3n_audio_frame_bucket(frame_bucket: usize) -> Result<(), String> { + if GEMMA3N_AUDIO_FRAME_BUCKETS.contains(&frame_bucket) { + Ok(()) + } else { + Err(format!( + "Gemma3n audio frame bucket {frame_bucket} is not one of {GEMMA3N_AUDIO_FRAME_BUCKETS:?}" + )) + } +} + +pub(crate) fn validate_gemma3n_audio_batch( + frame_bucket: usize, + clips: usize, +) -> Result<(), String> { + validate_gemma3n_audio_frame_bucket(frame_bucket)?; + if !(1..=GEMMA3N_AUDIO_MAX_CLIPS).contains(&clips) { + return Err(format!( + "Gemma3n audio clip count {clips} is outside 1..={GEMMA3N_AUDIO_MAX_CLIPS}" + )); + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn config_json(audio: serde_json::Value) -> String { + serde_json::json!({ + "model_type": "gemma3n", + "audio_config": audio, + }) + .to_string() + } + + #[test] + fn config_identity_covers_every_graph_shape_and_numeric_policy() { + let mut config = Gemma3nXlaAudioConfig::default(); + config.validate().unwrap(); + assert_eq!(config.encoded_frames(2_997).unwrap(), 750); + assert_eq!(config.projected_frames(2_997).unwrap(), 188); + assert_eq!(config.output_frequency().unwrap(), 32); + let first = config.artifact_identity(2_997, 1, 2_048, 30, 256).unwrap(); + config.conf_attention_logit_cap = 49.0; + let second = config.artifact_identity(2_997, 1, 2_048, 30, 256).unwrap(); + assert_ne!(first, second); + assert!(first.contains(GEMMA3N_AUDIO_GRAPH_ABI)); + assert!(first.contains("bucket=2997")); + assert!(first.contains("dtype=bf16-accum-f32")); + } + + #[test] + fn capability_requires_gemma3n_audio_config() { + assert!( + Gemma3nXlaAudioConfig::from_json_str( + &serde_json::json!({"model_type": "gemma3n"}).to_string() + ) + .unwrap() + .is_none() + ); + assert!( + Gemma3nXlaAudioConfig::from_json_str(&config_json(serde_json::json!({ + "model_type": "gemma4_audio" + }))) + .unwrap_err() + .contains("gemma3n_audio") + ); + assert!( + Gemma3nXlaAudioConfig::from_json_str(&config_json(serde_json::json!({ + "model_type": "gemma3n_audio" + }))) + .unwrap() + .is_some() + ); + } +} diff --git a/src/lib/mlxcel-xla/src/lib.rs b/src/lib/mlxcel-xla/src/lib.rs index 2b1328267..629b6ffa1 100644 --- a/src/lib/mlxcel-xla/src/lib.rs +++ b/src/lib/mlxcel-xla/src/lib.rs @@ -52,6 +52,7 @@ use mlxcel_core::session::{InferenceSession, PreparedPrefill, SessionCapabilitie mod context; #[cfg(any(feature = "diagnostics", test))] mod diagnostic_flags; +mod gemma3n_audio_config; #[cfg(any(feature = "iree", test))] #[allow(dead_code)] mod numeric_dtype_contract; @@ -67,6 +68,7 @@ mod operator_numeric_contract; mod prepared; mod prepared_deepstack; mod prepared_gemma3n; +mod prepared_gemma3n_audio; #[cfg(feature = "iree")] mod aux; @@ -207,10 +209,19 @@ pub fn dequantize_gemma3n_affine_diagnostic( } #[cfg(feature = "diagnostics")] pub use emitter::{Gemma3nDiagnosticLayout, Gemma3nDiagnosticSegment}; +pub use gemma3n_audio_config::{ + GEMMA3N_AUDIO_FRAME_BUCKETS, GEMMA3N_AUDIO_GRAPH_ABI, GEMMA3N_AUDIO_MAX_CLIPS, + GEMMA3N_AUDIO_MAX_FRAMES, GEMMA3N_AUDIO_MEL_BINS, GEMMA3N_AUDIO_MODALITY_FAMILY, + GEMMA3N_AUDIO_SOFT_TOKENS, Gemma3nXlaAudioConfig, +}; #[cfg(feature = "iree")] pub use prepared::PreparedInputError; pub use prepared_deepstack::{DeepStackFeatures, DeepStackInputError, DeepStackPreparedPrefill}; pub use prepared_gemma3n::{Gemma3nDensePle, Gemma3nDensePleError, Gemma3nPreparedPrefill}; +pub use prepared_gemma3n_audio::{ + Gemma3nAudioInput, Gemma3nAudioInputError, Gemma3nAudioPreparedPrefill, + select_gemma3n_audio_frame_bucket, +}; #[cfg(feature = "iree")] pub use sampler::SampleParams; diff --git a/src/lib/mlxcel-xla/src/prepared_gemma3n_audio.rs b/src/lib/mlxcel-xla/src/prepared_gemma3n_audio.rs new file mode 100644 index 000000000..ca7fe7583 --- /dev/null +++ b/src/lib/mlxcel-xla/src/prepared_gemma3n_audio.rs @@ -0,0 +1,467 @@ +//! Owned Gemma3n audio inputs and prepared-prefill outputs. +//! +//! Canonical mel features and valid-frame masks enter `audio.main`; post-scale +//! merged embeddings and dense PLE leave it. Owned buffers make normal +//! completion, cancellation, and admission failure use the same drop path. + +use std::fmt; + +use crate::Gemma3nPreparedPrefill; +use crate::gemma3n_audio_config::{ + GEMMA3N_AUDIO_FRAME_BUCKETS, GEMMA3N_AUDIO_MAX_CLIPS, GEMMA3N_AUDIO_MAX_FRAMES, + GEMMA3N_AUDIO_MEL_BINS, GEMMA3N_AUDIO_MODALITY_FAMILY, GEMMA3N_AUDIO_SOFT_TOKENS, + validate_gemma3n_audio_frame_bucket, +}; + +pub fn select_gemma3n_audio_frame_bucket(frames: usize) -> Result { + if frames == 0 { + return Err(Gemma3nAudioInputError::ZeroFrames); + } + GEMMA3N_AUDIO_FRAME_BUCKETS + .iter() + .copied() + .find(|&bucket| frames <= bucket) + .ok_or(Gemma3nAudioInputError::FrameLimit { + frames, + maximum: GEMMA3N_AUDIO_MAX_FRAMES, + }) +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum Gemma3nAudioInputError { + ZeroFrames, + FrameLimit { + frames: usize, + maximum: usize, + }, + ClipCount { + clips: usize, + maximum: usize, + }, + LengthMismatch { + kind: &'static str, + actual: usize, + expected: usize, + }, + InvalidFrameLength { + clip: usize, + frames: usize, + bucket: usize, + }, + AllPadded { + clip: usize, + }, + NonFinite, + NonZeroMaskedFeature, +} + +impl fmt::Display for Gemma3nAudioInputError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::ZeroFrames => { + f.write_str("Gemma3n audio input must contain at least one mel frame") + } + Self::FrameLimit { frames, maximum } => { + write!( + f, + "Gemma3n audio input has {frames} frames; maximum is {maximum}" + ) + } + Self::ClipCount { clips, maximum } => { + write!( + f, + "Gemma3n audio input has {clips} clips; maximum is {maximum}" + ) + } + Self::LengthMismatch { + kind, + actual, + expected, + } => { + write!( + f, + "Gemma3n audio {kind} has {actual} elements; expected {expected}" + ) + } + Self::InvalidFrameLength { + clip, + frames, + bucket, + } => { + write!( + f, + "Gemma3n audio clip {clip} has invalid frame length {frames} for bucket {bucket}" + ) + } + Self::AllPadded { clip } => { + write!(f, "Gemma3n audio clip {clip} is entirely padded") + } + Self::NonFinite => f.write_str("Gemma3n audio mel input contains a non-finite value"), + Self::NonZeroMaskedFeature => f.write_str( + "Gemma3n audio masked mel rows must be explicit zeros at the IREE boundary", + ), + } + } +} + +impl std::error::Error for Gemma3nAudioInputError {} + +/// Canonical, zero-padded input to one static `audio.main` frame bucket. +#[derive(Clone, Debug, PartialEq)] +pub struct Gemma3nAudioInput { + mel: Vec, + valid_mask: Vec, + frame_lengths: Vec, + frame_bucket: usize, +} + +impl Gemma3nAudioInput { + pub fn new( + mel: Vec, + valid_mask: Vec, + frame_lengths: Vec, + frame_bucket: usize, + ) -> Result { + let clips = frame_lengths.len(); + if !(1..=GEMMA3N_AUDIO_MAX_CLIPS).contains(&clips) { + return Err(Gemma3nAudioInputError::ClipCount { + clips, + maximum: GEMMA3N_AUDIO_MAX_CLIPS, + }); + } + if !GEMMA3N_AUDIO_FRAME_BUCKETS.contains(&frame_bucket) { + return Err(Gemma3nAudioInputError::FrameLimit { + frames: frame_bucket, + maximum: GEMMA3N_AUDIO_MAX_FRAMES, + }); + } + let mask_len = + clips + .checked_mul(frame_bucket) + .ok_or(Gemma3nAudioInputError::LengthMismatch { + kind: "mask", + actual: valid_mask.len(), + expected: usize::MAX, + })?; + let mel_len = mask_len.checked_mul(GEMMA3N_AUDIO_MEL_BINS).ok_or( + Gemma3nAudioInputError::LengthMismatch { + kind: "mel tensor", + actual: mel.len(), + expected: usize::MAX, + }, + )?; + if valid_mask.len() != mask_len { + return Err(Gemma3nAudioInputError::LengthMismatch { + kind: "mask", + actual: valid_mask.len(), + expected: mask_len, + }); + } + if mel.len() != mel_len { + return Err(Gemma3nAudioInputError::LengthMismatch { + kind: "mel tensor", + actual: mel.len(), + expected: mel_len, + }); + } + if mel.iter().any(|value| !value.is_finite()) { + return Err(Gemma3nAudioInputError::NonFinite); + } + for (clip, &frames) in frame_lengths.iter().enumerate() { + if frames == 0 { + return Err(Gemma3nAudioInputError::AllPadded { clip }); + } + if frames > frame_bucket { + return Err(Gemma3nAudioInputError::InvalidFrameLength { + clip, + frames, + bucket: frame_bucket, + }); + } + let mask = &valid_mask[clip * frame_bucket..(clip + 1) * frame_bucket]; + if mask[..frames].iter().any(|&value| value != 1) + || mask[frames..].iter().any(|&value| value != 0) + { + return Err(Gemma3nAudioInputError::InvalidFrameLength { + clip, + frames, + bucket: frame_bucket, + }); + } + } + if valid_mask.iter().enumerate().any(|(index, &valid)| { + valid == 0 + && mel[index * GEMMA3N_AUDIO_MEL_BINS..(index + 1) * GEMMA3N_AUDIO_MEL_BINS] + .iter() + .any(|&value| value != 0.0) + }) { + return Err(Gemma3nAudioInputError::NonZeroMaskedFeature); + } + Ok(Self { + mel, + valid_mask, + frame_lengths, + frame_bucket, + }) + } + + #[must_use] + pub fn mel(&self) -> &[f32] { + &self.mel + } + + #[must_use] + pub fn valid_mask(&self) -> &[u8] { + &self.valid_mask + } + + #[must_use] + pub fn frame_lengths(&self) -> &[usize] { + &self.frame_lengths + } + + #[must_use] + pub fn clips(&self) -> usize { + self.frame_lengths.len() + } + + #[must_use] + pub fn frame_bucket(&self) -> usize { + self.frame_bucket + } +} + +/// `audio.main` output ready for the existing Gemma3n dense-PLE prefill entry. +#[derive(Clone, Debug)] +pub struct Gemma3nAudioPreparedPrefill { + request: Gemma3nPreparedPrefill, + projected_lengths: Vec, + placeholder_starts: Vec, + frame_bucket: usize, +} + +impl Gemma3nAudioPreparedPrefill { + pub fn new( + request: Gemma3nPreparedPrefill, + audio_token_id: i32, + projected_lengths: Vec, + frame_bucket: usize, + ) -> Result { + validate_gemma3n_audio_frame_bucket(frame_bucket)?; + if projected_lengths.is_empty() + || projected_lengths.len() > GEMMA3N_AUDIO_MAX_CLIPS + || projected_lengths + .iter() + .any(|&length| length == 0 || length > GEMMA3N_AUDIO_SOFT_TOKENS) + { + return Err("Gemma3n audio projected lengths violate the fixed clip contract".into()); + } + let clips = projected_lengths.len(); + let expected_tokens = clips + .checked_mul(GEMMA3N_AUDIO_SOFT_TOKENS) + .ok_or_else(|| "Gemma3n audio placeholder count overflows".to_string())?; + let prepared = request.prepared(); + let audio_modalities: Vec<_> = prepared + .modalities + .iter() + .filter(|modality| modality.family == GEMMA3N_AUDIO_MODALITY_FAMILY) + .collect(); + if audio_modalities.len() != 1 + || audio_modalities[0].item_count != clips + || audio_modalities[0].token_count != expected_tokens + { + return Err(format!( + "Gemma3n audio modality metadata must declare {clips} clips and \ + {expected_tokens} placeholder tokens" + )); + } + let mut placeholder_starts = Vec::with_capacity(clips); + let mut index = 0; + while index < prepared.token_ids.len() { + if prepared.token_ids[index] != audio_token_id { + index += 1; + continue; + } + let start = index; + while index < prepared.token_ids.len() && prepared.token_ids[index] == audio_token_id { + index += 1; + } + let length = index - start; + if length != GEMMA3N_AUDIO_SOFT_TOKENS { + return Err(format!( + "Gemma3n audio placeholder run at token {start} has {length} rows; \ + expected {GEMMA3N_AUDIO_SOFT_TOKENS}" + )); + } + placeholder_starts.push(start); + } + if placeholder_starts.len() != clips { + return Err(format!( + "Gemma3n audio placeholder run count {} does not match {clips} clips", + placeholder_starts.len() + )); + } + Ok(Self { + request, + projected_lengths, + placeholder_starts, + frame_bucket, + }) + } + + #[must_use] + pub fn request(&self) -> &Gemma3nPreparedPrefill { + &self.request + } + + #[must_use] + pub fn projected_lengths(&self) -> &[usize] { + &self.projected_lengths + } + + #[must_use] + pub fn placeholder_starts(&self) -> &[usize] { + &self.placeholder_starts + } + + #[must_use] + pub fn frame_bucket(&self) -> usize { + self.frame_bucket + } + + pub fn into_request(self) -> Gemma3nPreparedPrefill { + self.request + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::Gemma3nDensePle; + use mlxcel_core::session::{ + OwnedTensor, PreparedAttentionBias, PreparedModality, PreparedPositions, PreparedPrefill, + PreparedTensorDType, + }; + + fn tensor(shape: Vec, values: Vec) -> OwnedTensor { + OwnedTensor::new( + values.into_iter().flat_map(f32::to_le_bytes).collect(), + PreparedTensorDType::Float32, + shape, + ) + .unwrap() + } + + fn prepared_audio_request( + token_ids: Vec, + clips: usize, + audio_tokens: usize, + ) -> Gemma3nPreparedPrefill { + let sequence_len = token_ids.len(); + let prepared = PreparedPrefill::new( + token_ids, + tensor(vec![1, sequence_len, 2], vec![0.0; sequence_len * 2]), + PreparedPositions::Sequential { + start: 0, + length: sequence_len, + }, + PreparedAttentionBias { + tensor: tensor(vec![1, 1, 1, sequence_len], vec![0.0; sequence_len]), + causal: true, + }, + vec![PreparedModality { + family: GEMMA3N_AUDIO_MODALITY_FAMILY.into(), + item_count: clips, + token_count: audio_tokens, + }], + ) + .unwrap(); + let ple = Gemma3nDensePle::new(vec![0.0; sequence_len * 2], sequence_len, 1, 2).unwrap(); + Gemma3nPreparedPrefill::new(prepared, ple).unwrap() + } + + #[test] + fn frame_buckets_are_exact_and_fail_closed() { + assert_eq!(select_gemma3n_audio_frame_bucket(1).unwrap(), 8); + assert_eq!(select_gemma3n_audio_frame_bucket(8).unwrap(), 8); + assert_eq!(select_gemma3n_audio_frame_bucket(9).unwrap(), 32); + assert_eq!( + select_gemma3n_audio_frame_bucket(GEMMA3N_AUDIO_MAX_FRAMES).unwrap(), + GEMMA3N_AUDIO_MAX_FRAMES + ); + assert!(matches!( + select_gemma3n_audio_frame_bucket(GEMMA3N_AUDIO_MAX_FRAMES + 1), + Err(Gemma3nAudioInputError::FrameLimit { .. }) + )); + } + + #[test] + fn input_rejects_all_padded_nonfinite_and_dirty_padding() { + let bucket = 8; + let valid = vec![1, 1, 0, 0, 0, 0, 0, 0]; + let input = Gemma3nAudioInput::new( + vec![0.0; bucket * GEMMA3N_AUDIO_MEL_BINS], + valid, + vec![2], + bucket, + ) + .unwrap(); + assert_eq!(input.clips(), 1); + assert!(matches!( + Gemma3nAudioInput::new( + vec![0.0; bucket * GEMMA3N_AUDIO_MEL_BINS], + vec![1, 0, 1, 0, 0, 0, 0, 0], + vec![2], + bucket, + ), + Err(Gemma3nAudioInputError::InvalidFrameLength { .. }) + )); + assert!(matches!( + Gemma3nAudioInput::new( + vec![0.0; bucket * GEMMA3N_AUDIO_MEL_BINS], + vec![0; bucket], + vec![0], + bucket, + ), + Err(Gemma3nAudioInputError::AllPadded { .. }) + )); + let mut nonfinite = vec![0.0; bucket * GEMMA3N_AUDIO_MEL_BINS]; + nonfinite[0] = f32::NAN; + assert!(matches!( + Gemma3nAudioInput::new(nonfinite, vec![1; bucket], vec![bucket], bucket), + Err(Gemma3nAudioInputError::NonFinite) + )); + let mut dirty = vec![0.0; bucket * GEMMA3N_AUDIO_MEL_BINS]; + dirty[2 * GEMMA3N_AUDIO_MEL_BINS] = 1.0; + assert!(matches!( + Gemma3nAudioInput::new(dirty, vec![1, 1, 0, 0, 0, 0, 0, 0], vec![2], bucket,), + Err(Gemma3nAudioInputError::NonZeroMaskedFeature) + )); + } + + #[test] + fn prepared_audio_requires_exact_clip_spans_and_modality_metadata() { + const AUDIO: i32 = 262_273; + let mut tokens = vec![1]; + tokens.extend(std::iter::repeat_n(AUDIO, GEMMA3N_AUDIO_SOFT_TOKENS)); + tokens.push(2); + tokens.extend(std::iter::repeat_n(AUDIO, GEMMA3N_AUDIO_SOFT_TOKENS)); + tokens.push(3); + let request = prepared_audio_request(tokens.clone(), 2, 2 * GEMMA3N_AUDIO_SOFT_TOKENS); + let prepared = + Gemma3nAudioPreparedPrefill::new(request, AUDIO, vec![188, 17], 2_997).unwrap(); + assert_eq!( + prepared.placeholder_starts(), + &[1, GEMMA3N_AUDIO_SOFT_TOKENS + 2] + ); + + let mut malformed = tokens; + malformed.remove(1); + let request = prepared_audio_request(malformed, 2, 2 * GEMMA3N_AUDIO_SOFT_TOKENS); + assert!( + Gemma3nAudioPreparedPrefill::new(request, AUDIO, vec![188, 17], 2_997) + .unwrap_err() + .contains("placeholder run") + ); + } +} From 1a33c339764cb2e8dc104d0edb09ddcc77cb332c Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 24 Jul 2026 14:06:30 +0900 Subject: [PATCH 02/19] feat(xla): validate Gemma3n audio weights Pin the converted Q4 checkpoint's 277 audio tensor names, dtypes, shapes, and deterministic order. Reject missing, unexpected, or incompatible tensors before enabling the audio capability. --- .../mlxcel-xla/src/gemma3n_audio_weights.rs | 493 ++++++++++++++++++ src/lib/mlxcel-xla/src/lib.rs | 6 + 2 files changed, 499 insertions(+) create mode 100644 src/lib/mlxcel-xla/src/gemma3n_audio_weights.rs diff --git a/src/lib/mlxcel-xla/src/gemma3n_audio_weights.rs b/src/lib/mlxcel-xla/src/gemma3n_audio_weights.rs new file mode 100644 index 000000000..1fe787c47 --- /dev/null +++ b/src/lib/mlxcel-xla/src/gemma3n_audio_weights.rs @@ -0,0 +1,493 @@ +//! Exact converted-checkpoint schema for the Gemma3n audio graph. + +use std::collections::{BTreeMap, BTreeSet}; +use std::fmt; +use std::io::{Read, Seek}; +use std::path::{Path, PathBuf}; + +use serde::Deserialize; + +use crate::Gemma3nXlaAudioConfig; + +pub const GEMMA3N_AUDIO_CHECKPOINT_TENSOR_COUNT: usize = 277; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Gemma3nAudioCheckpointDType { + Bf16, + U32, +} + +impl Gemma3nAudioCheckpointDType { + fn name(self) -> &'static str { + match self { + Self::Bf16 => "BF16", + Self::U32 => "U32", + } + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct Gemma3nAudioCheckpointTensorSpec { + pub name: String, + pub dtype: Gemma3nAudioCheckpointDType, + pub shape: Vec, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum Gemma3nAudioCheckpointError { + Io(String), + InvalidHeader(String), + Missing(Vec), + Unexpected(Vec), + DType { + name: String, + actual: String, + expected: &'static str, + }, + Shape { + name: String, + actual: Vec, + expected: Vec, + }, +} + +impl fmt::Display for Gemma3nAudioCheckpointError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Io(error) | Self::InvalidHeader(error) => f.write_str(error), + Self::Missing(names) => { + write!(f, "Gemma3n audio checkpoint is missing tensors: {names:?}") + } + Self::Unexpected(names) => { + write!( + f, + "Gemma3n audio checkpoint has unexpected audio tensors: {names:?}" + ) + } + Self::DType { + name, + actual, + expected, + } => write!( + f, + "Gemma3n audio tensor {name} has dtype {actual}; expected {expected}" + ), + Self::Shape { + name, + actual, + expected, + } => write!( + f, + "Gemma3n audio tensor {name} has shape {actual:?}; expected {expected:?}" + ), + } + } +} + +impl std::error::Error for Gemma3nAudioCheckpointError {} + +fn tensor( + out: &mut Vec, + name: impl Into, + dtype: Gemma3nAudioCheckpointDType, + shape: impl Into>, +) { + out.push(Gemma3nAudioCheckpointTensorSpec { + name: name.into(), + dtype, + shape: shape.into(), + }); +} + +fn dense( + out: &mut Vec, + name: impl Into, + shape: impl Into>, +) { + tensor(out, name, Gemma3nAudioCheckpointDType::Bf16, shape); +} + +fn feed_forward(out: &mut Vec, prefix: &str, hidden: usize) { + dense(out, format!("{prefix}.pre_layer_norm.weight"), [hidden]); + dense( + out, + format!("{prefix}.ffw_layer_1.weight"), + [hidden * 4, hidden], + ); + dense( + out, + format!("{prefix}.ffw_layer_2.weight"), + [hidden, hidden * 4], + ); + dense(out, format!("{prefix}.post_layer_norm.weight"), [hidden]); +} + +fn affine_q4( + out: &mut Vec, + prefix: &str, + rows: usize, + columns: usize, +) { + debug_assert!(columns.is_multiple_of(64)); + tensor( + out, + format!("{prefix}.weight"), + Gemma3nAudioCheckpointDType::U32, + [rows, columns / 8], + ); + for suffix in ["scales", "biases"] { + tensor( + out, + format!("{prefix}.{suffix}"), + Gemma3nAudioCheckpointDType::Bf16, + [rows, columns / 64], + ); + } +} + +/// The exact 277 raw tensors consumed from the pinned converted Q4 checkpoint. +pub fn gemma3n_audio_checkpoint_specs( + config: &Gemma3nXlaAudioConfig, + text_hidden: usize, +) -> Result, String> { + config.validate()?; + if text_hidden == 0 || !config.hidden_size.is_multiple_of(64) { + return Err("Gemma3n audio Q4 dimensions must be positive multiples of 64".into()); + } + let hidden = config.hidden_size; + let channels = &config.sscp_conv_channel_size; + let kernels = &config.sscp_conv_kernel_size; + let mut out = Vec::with_capacity(GEMMA3N_AUDIO_CHECKPOINT_TENSOR_COUNT); + dense( + &mut out, + "audio_tower.subsample_conv_projection.conv_0.conv.weight", + [channels[0], kernels[0][0], kernels[0][1], 1], + ); + dense( + &mut out, + "audio_tower.subsample_conv_projection.conv_0.norm.weight", + [channels[0]], + ); + dense( + &mut out, + "audio_tower.subsample_conv_projection.conv_1.conv.weight", + [channels[1], kernels[1][0], kernels[1][1], channels[0]], + ); + dense( + &mut out, + "audio_tower.subsample_conv_projection.conv_1.norm.weight", + [channels[1]], + ); + dense( + &mut out, + "audio_tower.subsample_conv_projection.input_proj_linear.weight", + [hidden, config.output_frequency()? * channels[1]], + ); + for layer in 0..config.conf_num_hidden_layers { + let prefix = format!("audio_tower.conformer.{layer}"); + feed_forward(&mut out, &format!("{prefix}.ffw_layer_start"), hidden); + let attention = format!("{prefix}.attention"); + dense( + &mut out, + format!("{attention}.pre_attn_norm.weight"), + [hidden], + ); + for projection in ["q_proj", "k_proj", "v_proj"] { + dense( + &mut out, + format!("{attention}.attn.{projection}.weight"), + [hidden, hidden], + ); + } + dense( + &mut out, + format!("{attention}.attn.per_dim_scale"), + [config.head_dim()], + ); + dense( + &mut out, + format!("{attention}.attn.relative_position_embedding.pos_proj.weight"), + [hidden, hidden], + ); + dense( + &mut out, + format!("{attention}.post.weight"), + [hidden, hidden], + ); + dense(&mut out, format!("{attention}.post_norm.weight"), [hidden]); + let light = format!("{prefix}.lconv1d"); + dense(&mut out, format!("{light}.pre_layer_norm.weight"), [hidden]); + dense( + &mut out, + format!("{light}.linear_start.weight"), + [hidden * 2, hidden], + ); + dense( + &mut out, + format!("{light}.depthwise_conv1d.weight"), + [hidden, config.conf_conv_kernel_size, 1], + ); + dense(&mut out, format!("{light}.conv_norm.weight"), [hidden]); + dense( + &mut out, + format!("{light}.linear_end.weight"), + [hidden, hidden], + ); + feed_forward(&mut out, &format!("{prefix}.ffw_layer_end"), hidden); + dense(&mut out, format!("{prefix}.norm.weight"), [hidden]); + } + affine_q4(&mut out, "embed_audio.embedding", config.vocab_size, hidden); + affine_q4( + &mut out, + "embed_audio.embedding_projection", + text_hidden, + hidden, + ); + dense(&mut out, "embed_audio.hard_embedding_norm.weight", [hidden]); + dense(&mut out, "embed_audio.soft_embedding_norm.weight", [hidden]); + if out.len() != GEMMA3N_AUDIO_CHECKPOINT_TENSOR_COUNT { + return Err(format!( + "Gemma3n audio schema generated {} tensors; expected {GEMMA3N_AUDIO_CHECKPOINT_TENSOR_COUNT}", + out.len() + )); + } + Ok(out) +} + +#[derive(Deserialize)] +struct HeaderTensor { + dtype: String, + shape: Vec, +} + +fn safetensor_files(model_dir: &Path) -> Result, Gemma3nAudioCheckpointError> { + let mut paths = std::fs::read_dir(model_dir) + .map_err(|error| { + Gemma3nAudioCheckpointError::Io(format!("read {}: {error}", model_dir.display())) + })? + .filter_map(Result::ok) + .map(|entry| entry.path()) + .filter(|path| { + path.extension() + .is_some_and(|extension| extension == "safetensors") + }) + .collect::>(); + paths.sort(); + if paths.is_empty() { + return Err(Gemma3nAudioCheckpointError::Io(format!( + "{} contains no safetensors files", + model_dir.display() + ))); + } + Ok(paths) +} + +fn read_header(path: &Path) -> Result, Gemma3nAudioCheckpointError> { + let mut file = std::fs::File::open(path).map_err(|error| { + Gemma3nAudioCheckpointError::Io(format!("open {}: {error}", path.display())) + })?; + let mut size = [0_u8; 8]; + file.read_exact(&mut size).map_err(|error| { + Gemma3nAudioCheckpointError::InvalidHeader(format!( + "read {} header size: {error}", + path.display() + )) + })?; + let size = usize::try_from(u64::from_le_bytes(size)).map_err(|_| { + Gemma3nAudioCheckpointError::InvalidHeader(format!( + "{} header size exceeds this platform", + path.display() + )) + })?; + if size == 0 || size > 128 * 1024 * 1024 { + return Err(Gemma3nAudioCheckpointError::InvalidHeader(format!( + "{} has invalid safetensors header size {size}", + path.display() + ))); + } + let mut bytes = vec![0_u8; size]; + file.seek(std::io::SeekFrom::Start(8)) + .and_then(|_| file.read_exact(&mut bytes)) + .map_err(|error| { + Gemma3nAudioCheckpointError::InvalidHeader(format!( + "read {} header: {error}", + path.display() + )) + })?; + let value: serde_json::Value = serde_json::from_slice(&bytes).map_err(|error| { + Gemma3nAudioCheckpointError::InvalidHeader(format!( + "parse {} header: {error}", + path.display() + )) + })?; + let object = value.as_object().ok_or_else(|| { + Gemma3nAudioCheckpointError::InvalidHeader(format!( + "{} safetensors header must be an object", + path.display() + )) + })?; + object + .iter() + .filter(|(name, _)| name.as_str() != "__metadata__") + .map(|(name, value)| { + serde_json::from_value(value.clone()) + .map(|metadata| (name.clone(), metadata)) + .map_err(|error| { + Gemma3nAudioCheckpointError::InvalidHeader(format!( + "{} tensor {name}: {error}", + path.display() + )) + }) + }) + .collect() +} + +pub fn validate_gemma3n_audio_checkpoint( + model_dir: &Path, + config: &Gemma3nXlaAudioConfig, + text_hidden: usize, +) -> Result<(), Gemma3nAudioCheckpointError> { + let mut actual = BTreeMap::new(); + for path in safetensor_files(model_dir)? { + actual.extend(read_header(&path)?); + } + validate_gemma3n_audio_tensor_map(&actual, config, text_hidden) +} + +fn validate_gemma3n_audio_tensor_map( + actual: &BTreeMap, + config: &Gemma3nXlaAudioConfig, + text_hidden: usize, +) -> Result<(), Gemma3nAudioCheckpointError> { + let specs = gemma3n_audio_checkpoint_specs(config, text_hidden) + .map_err(Gemma3nAudioCheckpointError::InvalidHeader)?; + let expected_names = specs + .iter() + .map(|spec| spec.name.as_str()) + .collect::>(); + let missing = expected_names + .iter() + .filter(|name| !actual.contains_key(**name)) + .map(|name| (*name).to_string()) + .collect::>(); + if !missing.is_empty() { + return Err(Gemma3nAudioCheckpointError::Missing(missing)); + } + let unexpected = actual + .keys() + .filter(|name| { + (name.starts_with("audio_tower.") || name.starts_with("embed_audio.")) + && !expected_names.contains(name.as_str()) + }) + .cloned() + .collect::>(); + if !unexpected.is_empty() { + return Err(Gemma3nAudioCheckpointError::Unexpected(unexpected)); + } + for spec in specs { + let metadata = &actual[&spec.name]; + if metadata.dtype != spec.dtype.name() { + return Err(Gemma3nAudioCheckpointError::DType { + name: spec.name, + actual: metadata.dtype.clone(), + expected: spec.dtype.name(), + }); + } + if metadata.shape != spec.shape { + return Err(Gemma3nAudioCheckpointError::Shape { + name: spec.name, + actual: metadata.shape.clone(), + expected: spec.shape, + }); + } + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn map_from_specs( + specs: &[Gemma3nAudioCheckpointTensorSpec], + ) -> BTreeMap { + specs + .iter() + .map(|spec| { + ( + spec.name.clone(), + HeaderTensor { + dtype: spec.dtype.name().into(), + shape: spec.shape.clone(), + }, + ) + }) + .collect() + } + + #[test] + fn pinned_q4_schema_has_exact_order_count_and_boundaries() { + let specs = + gemma3n_audio_checkpoint_specs(&Gemma3nXlaAudioConfig::default(), 2_048).unwrap(); + assert_eq!(specs.len(), GEMMA3N_AUDIO_CHECKPOINT_TENSOR_COUNT); + assert_eq!( + specs[0].name, + "audio_tower.subsample_conv_projection.conv_0.conv.weight" + ); + assert_eq!(specs[0].shape, [128, 3, 3, 1]); + assert_eq!(specs[268].name, "audio_tower.conformer.11.norm.weight"); + assert_eq!(specs[269].name, "embed_audio.embedding.weight"); + assert_eq!(specs[269].shape, [128, 192]); + assert_eq!( + specs.last().unwrap().name, + "embed_audio.soft_embedding_norm.weight" + ); + } + + #[test] + fn validation_rejects_missing_extra_dtype_and_shape() { + let config = Gemma3nXlaAudioConfig::default(); + let specs = gemma3n_audio_checkpoint_specs(&config, 2_048).unwrap(); + let mut actual = map_from_specs(&specs); + actual.remove(&specs[7].name); + assert!(matches!( + validate_gemma3n_audio_tensor_map(&actual, &config, 2_048), + Err(Gemma3nAudioCheckpointError::Missing(_)) + )); + let mut actual = map_from_specs(&specs); + actual.insert( + "audio_tower.unqualified.weight".into(), + HeaderTensor { + dtype: "BF16".into(), + shape: vec![1], + }, + ); + assert!(matches!( + validate_gemma3n_audio_tensor_map(&actual, &config, 2_048), + Err(Gemma3nAudioCheckpointError::Unexpected(_)) + )); + let mut actual = map_from_specs(&specs); + actual.get_mut(&specs[0].name).unwrap().dtype = "F32".into(); + assert!(matches!( + validate_gemma3n_audio_tensor_map(&actual, &config, 2_048), + Err(Gemma3nAudioCheckpointError::DType { .. }) + )); + let mut actual = map_from_specs(&specs); + actual.get_mut(&specs[0].name).unwrap().shape[0] = 127; + assert!(matches!( + validate_gemma3n_audio_tensor_map(&actual, &config, 2_048), + Err(Gemma3nAudioCheckpointError::Shape { .. }) + )); + } + + #[test] + #[ignore = "requires GEMMA3N_MODEL_DIR pointing at a converted Q4 checkpoint"] + fn real_checkpoint_has_exact_audio_schema() { + let model_dir = + PathBuf::from(std::env::var_os("GEMMA3N_MODEL_DIR").expect("model directory")); + let config = Gemma3nXlaAudioConfig::from_model_dir(&model_dir) + .unwrap() + .unwrap(); + validate_gemma3n_audio_checkpoint(&model_dir, &config, 2_048).unwrap(); + } +} diff --git a/src/lib/mlxcel-xla/src/lib.rs b/src/lib/mlxcel-xla/src/lib.rs index 629b6ffa1..257763b8f 100644 --- a/src/lib/mlxcel-xla/src/lib.rs +++ b/src/lib/mlxcel-xla/src/lib.rs @@ -53,6 +53,7 @@ mod context; #[cfg(any(feature = "diagnostics", test))] mod diagnostic_flags; mod gemma3n_audio_config; +mod gemma3n_audio_weights; #[cfg(any(feature = "iree", test))] #[allow(dead_code)] mod numeric_dtype_contract; @@ -214,6 +215,11 @@ pub use gemma3n_audio_config::{ GEMMA3N_AUDIO_MAX_FRAMES, GEMMA3N_AUDIO_MEL_BINS, GEMMA3N_AUDIO_MODALITY_FAMILY, GEMMA3N_AUDIO_SOFT_TOKENS, Gemma3nXlaAudioConfig, }; +pub use gemma3n_audio_weights::{ + GEMMA3N_AUDIO_CHECKPOINT_TENSOR_COUNT, Gemma3nAudioCheckpointDType, + Gemma3nAudioCheckpointError, Gemma3nAudioCheckpointTensorSpec, gemma3n_audio_checkpoint_specs, + validate_gemma3n_audio_checkpoint, +}; #[cfg(feature = "iree")] pub use prepared::PreparedInputError; pub use prepared_deepstack::{DeepStackFeatures, DeepStackInputError, DeepStackPreparedPrefill}; From 2322d42217bbceb616a10e088a0738c6bf56db7c Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 24 Jul 2026 14:51:12 +0900 Subject: [PATCH 03/19] feat(xla): emit split Gemma3n audio artifacts Implement the complete Gemma3n Conformer audio path with exact cumulative normalization, local relative attention, masking, projection, row mapping, and shared dense-PLE rounding semantics. Split the production graph at the natural typed boundary: audio.encode owns audio tower/embedder weights and returns projected soft rows, hard audio embeddings, and lengths; audio.merge_ple owns only the shared language input-head weights. This avoids pinned IREE monolithic-module backend failures while keeping all model math in IREE. Validate the full mlxcel-xla library and compile the maximum 2997-frame/context-256 artifacts with pinned IREE for both local CPU and CUDA targets. Refs #878 --- src/lib/mlxcel-xla/src/emitter/builder.rs | 268 ++++++++++ .../src/emitter/gemma3n_audio_attention.rs | 314 ++++++++++++ .../src/emitter/gemma3n_audio_emit.rs | 466 ++++++++++++++++++ .../src/emitter/gemma3n_audio_emit_tests.rs | 130 +++++ .../src/emitter/gemma3n_audio_math.rs | 127 +++++ .../src/emitter/gemma3n_audio_ops.rs | 259 ++++++++++ .../src/emitter/gemma3n_audio_schema.rs | 287 +++++++++++ .../mlxcel-xla/src/emitter/gemma3n_emit.rs | 50 +- .../src/emitter/gemma3n_emit_ops.rs | 56 +++ src/lib/mlxcel-xla/src/emitter/mod.rs | 11 + .../mlxcel-xla/src/gemma3n_audio_config.rs | 4 +- src/lib/mlxcel-xla/src/gemma3n_audio_rows.rs | 230 +++++++++ src/lib/mlxcel-xla/src/lib.rs | 2 + .../mlxcel-xla/src/prepared_gemma3n_audio.rs | 34 +- 14 files changed, 2194 insertions(+), 44 deletions(-) create mode 100644 src/lib/mlxcel-xla/src/emitter/gemma3n_audio_attention.rs create mode 100644 src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit.rs create mode 100644 src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit_tests.rs create mode 100644 src/lib/mlxcel-xla/src/emitter/gemma3n_audio_math.rs create mode 100644 src/lib/mlxcel-xla/src/emitter/gemma3n_audio_ops.rs create mode 100644 src/lib/mlxcel-xla/src/emitter/gemma3n_audio_schema.rs create mode 100644 src/lib/mlxcel-xla/src/gemma3n_audio_rows.rs diff --git a/src/lib/mlxcel-xla/src/emitter/builder.rs b/src/lib/mlxcel-xla/src/emitter/builder.rs index ae182c154..a5a285b8e 100644 --- a/src/lib/mlxcel-xla/src/emitter/builder.rs +++ b/src/lib/mlxcel-xla/src/emitter/builder.rs @@ -715,6 +715,35 @@ impl Builder { } } + pub fn const_i1(&mut self, value: bool) -> Val { + let r = self.fresh(); + self.line(format!( + "{} = stablehlo.constant dense<{}> : tensor", + r, value + )); + Val { + name: r, + ty: Ty::scalar("i1"), + } + } + + pub fn const_tensor_i32(&mut self, data: &[i32], shape: Vec) -> Val { + assert_eq!(data.len(), shape.iter().product::()); + assert_eq!(shape.len(), 1); + let ty = Ty::new(shape, "i32"); + let values = data + .iter() + .map(i32::to_string) + .collect::>() + .join(", "); + let r = self.fresh(); + self.line(format!( + "{r} = stablehlo.constant dense<[{values}]> : {}", + ty.render() + )); + Val { name: r, ty } + } + /// Dense f32 constant tensor from raw data (row-major), emitted as hex blob. pub fn const_tensor_f32(&mut self, data: &[f32], shape: Vec) -> Val { let ty = Ty::f32(shape); @@ -987,6 +1016,58 @@ impl Builder { Val { name: r, ty } } + /// Static slice with an explicit positive stride in each dimension. + pub fn slice_strided(&mut self, x: &Val, ranges: &[(usize, usize, usize)]) -> Val { + let shape: Vec = ranges + .iter() + .map(|(start, limit, stride)| (limit - start).div_ceil(*stride)) + .collect(); + let ty = Ty::new(shape, x.ty.elt); + let parts: Vec = ranges + .iter() + .map(|(start, limit, stride)| format!("{start}:{limit}:{stride}")) + .collect(); + let r = self.fresh(); + self.line(format!( + "{} = stablehlo.slice {} [{}] : ({}) -> {}", + r, + x.name, + parts.join(", "), + x.ty.render(), + ty.render() + )); + Val { name: r, ty } + } + + /// Zero-interior static padding. + pub fn pad(&mut self, x: &Val, value: &Val, low: &[usize], high: &[usize]) -> Val { + assert_eq!(low.len(), x.ty.shape.len()); + assert_eq!(high.len(), x.ty.shape.len()); + let shape = + x.ty.shape + .iter() + .zip(low) + .zip(high) + .map(|((&dim, &before), &after)| before + dim + after) + .collect(); + let ty = Ty::new(shape, x.ty.elt); + let interior = vec![0; x.ty.shape.len()]; + let r = self.fresh(); + self.line(format!( + "{} = stablehlo.pad {}, {}, low = {}, high = {}, interior = {} : ({}, {}) -> {}", + r, + x.name, + value.name, + dims_list(low), + dims_list(high), + dims_list(&interior), + x.ty.render(), + value.ty.render(), + ty.render() + )); + Val { name: r, ty } + } + pub fn concatenate(&mut self, a: &Val, b: &Val, dim: usize) -> Val { let mut shape = a.ty.shape.clone(); shape[dim] = a.ty.shape[dim] + b.ty.shape[dim]; @@ -1269,6 +1350,15 @@ impl Builder { pub fn divide(&mut self, a: &Val, b: &Val) -> Val { self.binary("divide", a, b) } + pub fn maximum(&mut self, a: &Val, b: &Val) -> Val { + self.binary("maximum", a, b) + } + pub fn minimum(&mut self, a: &Val, b: &Val) -> Val { + self.binary("minimum", a, b) + } + pub fn logical_and(&mut self, a: &Val, b: &Val) -> Val { + self.binary("and", a, b) + } fn unary(&mut self, op: &str, a: &Val) -> Val { let ty = a.ty.clone(); @@ -1304,6 +1394,12 @@ impl Builder { pub fn sine(&mut self, a: &Val) -> Val { self.unary("sine", a) } + pub fn logarithm(&mut self, a: &Val) -> Val { + self.unary("log", a) + } + pub fn logical_not(&mut self, a: &Val) -> Val { + self.unary("not", a) + } /// Error function from the CHLO extension accepted by StableHLO import. /// Gemma3n's sparse GELU uses the exact erf form (the ordinary GeGLU path @@ -1380,6 +1476,57 @@ impl Builder { self.reduce("maximum", x, dim, init) } + /// Inclusive prefix sum along one static axis using a left-padded window. + /// + /// This is the StableHLO cumsum lowering used by Gemma3n cumulative group + /// normalization. It avoids a quadratic triangular constant and keeps the + /// reduction in f32 on both CPU and CUDA. + pub fn reduce_window_prefix_add(&mut self, x: &Val, dim: usize) -> Val { + assert_eq!(x.ty.elt, "f32"); + assert!(dim < x.ty.shape.len()); + let rank = x.ty.shape.len(); + let width = x.ty.shape[dim]; + assert!(width > 0); + let zero = self.const_f32(0.0); + let window = (0..rank) + .map(|axis| if axis == dim { width } else { 1 }) + .map(|value| value.to_string()) + .collect::>() + .join(", "); + let strides = std::iter::repeat_n("1", rank) + .collect::>() + .join(", "); + let padding = (0..rank) + .map(|axis| { + if axis == dim { + format!("[{}, 0]", width - 1) + } else { + "[0, 0]".to_string() + } + }) + .collect::>() + .join(", "); + let r = self.fresh(); + let suffix = r.trim_start_matches('%'); + self.line(format!( + "{r} = \"stablehlo.reduce_window\"({x}, {zero}) ({{\n\ + ^bb0(%rw_l_{suffix}: tensor, %rw_r_{suffix}: tensor):\n \ + %rw_sum_{suffix} = stablehlo.add %rw_l_{suffix}, %rw_r_{suffix} : tensor\n \ + stablehlo.return %rw_sum_{suffix} : tensor\n\ + }}) {{window_dimensions = array, \ + window_strides = array, \ + padding = dense<[{padding}]> : tensor<{rank}x2xi64>}} : \ + ({ty}, tensor) -> {ty}", + x = x.name, + zero = zero.name, + ty = x.ty.render(), + )); + Val { + name: r, + ty: x.ty.clone(), + } + } + /// On-device argmax over a `[V]` vector -> scalar `i32` index (the first /// index of the max, numpy/jax semantics). Mirrors the JAX/IREE-emitted /// argmax reducer in `spike/openxla/artifacts/fp32_decode_argmax.stablehlo.mlir`: @@ -1556,6 +1703,35 @@ impl Builder { Val { name: r, ty } } + /// A contraction that always retains f32 inputs. Audio cumulative + /// normalization uses this for prefix sums, which must not inherit the + /// model's BF16 projection policy. + pub fn dot_general_f32( + &mut self, + lhs: &Val, + rhs: &Val, + lhs_contract: &[usize], + rhs_contract: &[usize], + out_shape: Vec, + ) -> Val { + assert_eq!(lhs.ty.elt, "f32"); + assert_eq!(rhs.ty.elt, "f32"); + let ty = Ty::f32(out_shape); + let r = self.fresh(); + self.line(format!( + "{} = stablehlo.dot_general {}, {}, contracting_dims = {} x {} : ({}, {}) -> {}", + r, + lhs.name, + rhs.name, + dims_list(lhs_contract), + dims_list(rhs_contract), + lhs.ty.render(), + rhs.ty.render(), + ty.render() + )); + Val { name: r, ty } + } + /// Convenience: y = x @ W^T for x:[K], W:[N,K] (weights stored [out,in]). /// Contracts x dim 0 with W dim 1, yielding [N]. No transpose needed. pub fn linear(&mut self, x: &Val, w: &Val) -> Val { @@ -1571,6 +1747,98 @@ impl Builder { let n = w.ty.shape[0]; self.dot_general(x, w, &[], &[], &[1], &[1], vec![l, n]) } + + /// Linear projection over the final axis of an arbitrary-rank tensor. + pub fn linear_last(&mut self, x: &Val, w: &Val) -> Val { + let input = *x.ty.shape.last().expect("linear input rank"); + assert_eq!(input, w.ty.shape[1]); + let rows = x.ty.shape[..x.ty.shape.len() - 1].iter().product(); + let flat = self.reshape(x, vec![rows, input]); + let projected = self.linear_seq(&flat, w); + let mut shape = x.ty.shape[..x.ty.shape.len() - 1].to_vec(); + shape.push(w.ty.shape[0]); + self.reshape(&projected, shape) + } + + /// Channels-last NWC/KIO or NHWC/HWIO convolution. + pub fn convolution( + &mut self, + input: &Val, + kernel: &Val, + strides: &[usize], + padding: &[(usize, usize)], + feature_groups: usize, + ) -> Val { + let rank = input.ty.shape.len(); + assert!(matches!(rank, 3 | 4)); + assert_eq!(kernel.ty.shape.len(), rank); + assert_eq!(input.ty.elt, kernel.ty.elt); + assert_eq!(input.ty.elt, "f32"); + assert_eq!(strides.len(), rank - 2); + assert_eq!(padding.len(), rank - 2); + assert!(strides.iter().all(|stride| *stride > 0)); + assert!(feature_groups > 0); + assert_eq!( + input.ty.shape[rank - 1], + kernel.ty.shape[rank - 2] * feature_groups + ); + assert!(kernel.ty.shape[rank - 1].is_multiple_of(feature_groups)); + let batch = input.ty.shape[0]; + let mut shape = vec![batch]; + for axis in 0..rank - 2 { + let padded = input.ty.shape[axis + 1] + padding[axis].0 + padding[axis].1; + assert!(padded >= kernel.ty.shape[axis]); + shape.push((padded - kernel.ty.shape[axis]) / strides[axis] + 1); + } + shape.push(kernel.ty.shape[rank - 1]); + let output_shape = shape; + let lhs_demoted; + let rhs_demoted; + let (input, kernel) = match self.precision.dot_elt() { + Some(elt) => { + lhs_demoted = self.convert(input, elt); + rhs_demoted = self.convert(kernel, elt); + (&lhs_demoted, &rhs_demoted) + } + None => (input, kernel), + }; + let op_ty = Ty::new(output_shape, input.ty.elt); + let r = self.fresh(); + let strides = dims_list(strides); + let padding = padding + .iter() + .map(|(low, high)| format!("[{low}, {high}]")) + .collect::>() + .join(", "); + let dimension_numbers = if rank == 3 { + "[b, 0, f]x[0, i, o]->[b, 0, f]" + } else { + "[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]" + }; + self.line(format!( + "{} = stablehlo.convolution({}, {}) dim_numbers = {}, \ + window = {{stride = {}, pad = [{}]}} \ + {{batch_group_count = 1 : i64, feature_group_count = {} : i64, \ + precision_config = [#stablehlo, #stablehlo]}} : \ + ({}, {}) -> {}", + r, + input.name, + kernel.name, + dimension_numbers, + strides, + padding, + feature_groups, + input.ty.render(), + kernel.ty.render(), + op_ty.render() + )); + let result = Val { name: r, ty: op_ty }; + if result.ty.elt == "f32" { + result + } else { + self.convert(&result, "f32") + } + } } #[cfg(test)] diff --git a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_attention.rs b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_attention.rs new file mode 100644 index 000000000..754ffdd60 --- /dev/null +++ b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_attention.rs @@ -0,0 +1,314 @@ +//! Chunked local relative attention for Gemma3n audio. + +use super::builder::{Builder, Val}; +use super::gemma3n_audio_math::{clip, rms_norm, round_bf16, scalar_like, softmax_last, softplus}; +use super::gemma3n_audio_schema::EncoderArgs; +use crate::Gemma3nXlaAudioConfig; + +fn pad_time(b: &mut Builder, x: &Val, left: usize, right: usize, value: &Val) -> Val { + let mut low = vec![0; x.ty.shape.len()]; + let mut high = vec![0; x.ty.shape.len()]; + low[1] = left; + high[1] = right; + b.pad(x, value, &low, &high) +} + +fn query_blocks(b: &mut Builder, x: &Val, chunk: usize) -> Val { + let batch = x.ty.shape[0]; + let time = x.ty.shape[1]; + let heads = x.ty.shape[2]; + let head_dim = x.ty.shape[3]; + let blocks = time.div_ceil(chunk); + let zero = b.const_f32(0.0); + let padded = pad_time(b, x, 0, blocks * chunk - time, &zero); + b.reshape(&padded, vec![batch, blocks, chunk, heads, head_dim]) +} + +fn extract_context(b: &mut Builder, x: &Val, config: &Gemma3nXlaAudioConfig) -> Val { + let time = x.ty.shape[1]; + let chunk = config.conf_attention_chunk_size; + let left = config.conf_attention_context_left - 1; + let right = config.conf_attention_context_right + chunk - 1; + let context = config.context_size(); + let blocks = time.div_ceil(chunk); + let padded = if x.ty.elt == "i1" { + let zero = b.const_i1(false); + pad_time(b, x, left, right, &zero) + } else { + let zero = b.const_f32(0.0); + pad_time(b, x, left, right, &zero) + }; + let mut time_first = vec![1, 0]; + time_first.extend(2..padded.ty.shape.len()); + let padded = b.transpose(&padded, &time_first); + let trailing = padded.ty.shape[1..].iter().product::(); + let table = b.reshape(&padded, vec![padded.ty.shape[0], trailing]); + let indices = (0..blocks) + .flat_map(|block| { + (0..context).map(move |offset| { + i32::try_from(block * chunk + offset).expect("audio context index") + }) + }) + .collect::>(); + let indices = b.const_tensor_i32(&indices, vec![blocks * context]); + let indices = b.reshape(&indices, vec![blocks * context, 1]); + let gathered = b.gather(&table, &indices); + let mut gathered_shape = vec![blocks, context]; + gathered_shape.extend_from_slice(&padded.ty.shape[1..]); + let gathered = b.reshape(&gathered, gathered_shape); + let mut batch_first = vec![2, 0, 1]; + batch_first.extend(3..gathered.ty.shape.len()); + b.transpose(&gathered, &batch_first) +} + +fn causal_valid_mask(b: &mut Builder, config: &Gemma3nXlaAudioConfig) -> Val { + let chunk = config.conf_attention_chunk_size; + let context = config.context_size(); + let diagonal = config.conf_attention_context_left - 1 + config.conf_attention_context_right; + let values = (0..chunk) + .flat_map(|query| { + (0..context).map(move |key| { + if query <= key && key <= query + diagonal { + 1.0 + } else { + 0.0 + } + }) + }) + .collect::>(); + let values = b.const_tensor_f32(&values, vec![chunk, context]); + let zero = scalar_like(b, 0.0, &values); + b.compare("GT", &values, &zero, "FLOAT") +} + +fn relative_signal( + b: &mut Builder, + args: &EncoderArgs, + config: &Gemma3nXlaAudioConfig, + prefix: &str, +) -> Val { + let past = config.conf_attention_context_left - 1; + let future = config.conf_attention_context_right; + let positions = (0..=past + future) + .rev() + .map(|value| (value as i32 - future as i32) as f32) + .collect::>(); + let timescales = config.hidden_size / 2; + let increment = 10_000.0f32.ln() / (timescales as f32 - 1.0).max(1.0); + let inverse = (0..timescales) + .map(|index| (-increment * index as f32).exp()) + .collect::>(); + let span = positions.len(); + let positions = b.const_tensor_f32(&positions, vec![1, span, 1]); + let inverse = b.const_tensor_f32(&inverse, vec![1, 1, timescales]); + let positions = b.broadcast(&positions, &[0, 1, 2], vec![1, span, timescales]); + let inverse = b.broadcast(&inverse, &[0, 1, 2], vec![1, span, timescales]); + let scaled = b.multiply(&positions, &inverse); + let sine = b.sine(&scaled); + let cosine = b.cosine(&scaled); + let signal = b.concatenate(&sine, &cosine, 2); + let signal = b.linear_last( + &signal, + args.weight(&format!( + "{prefix}.attn.relative_position_embedding.pos_proj.weight" + )), + ); + let signal = round_bf16(b, &signal); + b.reshape( + &signal, + vec![span, config.conf_num_attention_heads, config.head_dim()], + ) +} + +fn relative_logits( + b: &mut Builder, + queries: &Val, + keys: &Val, + args: &EncoderArgs, + config: &Gemma3nXlaAudioConfig, + prefix: &str, +) -> Val { + let batch = queries.ty.shape[0]; + let blocks = queries.ty.shape[1]; + let chunk = queries.ty.shape[2]; + let heads = queries.ty.shape[3]; + let head_dim = queries.ty.shape[4]; + let context = keys.ty.shape[2]; + let span = config.conf_attention_context_left + config.conf_attention_context_right; + let query = b.transpose(queries, &[0, 3, 1, 2, 4]); + let key = b.transpose(keys, &[0, 3, 1, 4, 2]); + let content = b.dot_general( + &query, + &key, + &[0, 1, 2], + &[0, 1, 2], + &[4], + &[3], + vec![batch, heads, blocks, chunk, context], + ); + let content = round_bf16(b, &content); + + let signal = relative_signal(b, args, config, prefix); + let signal = b.transpose(&signal, &[1, 2, 0]); + let signal = b.reshape(&signal, vec![1, heads, 1, head_dim, span]); + let signal = b.broadcast( + &signal, + &[0, 1, 2, 3, 4], + vec![batch, heads, blocks, head_dim, span], + ); + let position = b.dot_general( + &query, + &signal, + &[0, 1, 2], + &[0, 1, 2], + &[4], + &[3], + vec![batch, heads, blocks, chunk, span], + ); + let position = round_bf16(b, &position); + let zero = b.const_f32(0.0); + let position = b.pad( + &position, + &zero, + &[0, 0, 0, 0, 0], + &[0, 0, 0, 0, context + 1 - span], + ); + let position = b.reshape(&position, vec![batch, heads, blocks, chunk * (context + 1)]); + let position = b.slice( + &position, + &[(0, batch), (0, heads), (0, blocks), (0, chunk * context)], + ); + let position = b.reshape(&position, vec![batch, heads, blocks, chunk, context]); + let logits = b.add(&content, &position); + round_bf16(b, &logits) +} + +pub(super) fn attention( + b: &mut Builder, + args: &EncoderArgs, + config: &Gemma3nXlaAudioConfig, + prefix: &str, + input: &Val, + valid: &Val, +) -> Val { + let batch = input.ty.shape[0]; + let time = input.ty.shape[1]; + let heads = config.conf_num_attention_heads; + let head_dim = config.head_dim(); + let clipped = clip(b, input, config.gradient_clipping); + let normalized = rms_norm( + b, + &clipped, + Some(args.weight(&format!("{prefix}.pre_attn_norm.weight"))), + config.rms_norm_eps, + ); + let normalized = round_bf16(b, &normalized); + let project = |b: &mut Builder, name: &str| { + let projected = b.linear_last( + &normalized, + args.weight(&format!("{prefix}.attn.{name}.weight")), + ); + let projected = round_bf16(b, &projected); + b.reshape(&projected, vec![batch, time, heads, head_dim]) + }; + let query = project(b, "q_proj"); + let key = project(b, "k_proj"); + let value = project(b, "v_proj"); + let scale = softplus(b, args.weight(&format!("{prefix}.attn.per_dim_scale"))); + let scale = round_bf16(b, &scale); + let q_scale = (head_dim as f32).powf(-0.5) / 2.0f32.ln(); + let q_scale = scalar_like(b, q_scale, &scale); + let scale = b.multiply(&scale, &q_scale); + let scale = round_bf16(b, &scale); + let scale = b.broadcast(&scale, &[3], query.ty.shape.clone()); + let query = b.multiply(&query, &scale); + let query = round_bf16(b, &query); + + let query = query_blocks(b, &query, config.conf_attention_chunk_size); + let keys = extract_context(b, &key, config); + let values = extract_context(b, &value, config); + let blocks = query.ty.shape[1]; + let logits = relative_logits(b, &query, &keys, args, config, prefix); + let cap = scalar_like(b, config.conf_attention_logit_cap, &logits); + let scaled = b.divide(&logits, &cap); + let capped = b.tanh(&scaled); + let capped = b.multiply(&capped, &cap); + let capped = round_bf16(b, &capped); + + let context_valid = extract_context(b, valid, config); + let context_valid = b.reshape( + &context_valid, + vec![batch, 1, blocks, 1, config.context_size()], + ); + let context_valid = b.broadcast( + &context_valid, + &[0, 1, 2, 3, 4], + vec![ + batch, + heads, + blocks, + config.conf_attention_chunk_size, + config.context_size(), + ], + ); + let causal = causal_valid_mask(b, config); + let causal = b.reshape( + &causal, + vec![ + 1, + 1, + 1, + config.conf_attention_chunk_size, + config.context_size(), + ], + ); + let causal = b.broadcast(&causal, &[0, 1, 2, 3, 4], capped.ty.shape.clone()); + let condition = b.logical_and(&context_valid, &causal); + let minimum = scalar_like(b, -3.38e38, &capped); + let masked = b.select(&condition, &capped, &minimum); + let probabilities = softmax_last(b, &masked); + let probabilities = round_bf16(b, &probabilities); + let values = b.transpose(&values, &[0, 3, 1, 2, 4]); + let context = b.dot_general( + &probabilities, + &values, + &[0, 1, 2], + &[0, 1, 2], + &[4], + &[3], + vec![ + batch, + heads, + blocks, + config.conf_attention_chunk_size, + head_dim, + ], + ); + let context = round_bf16(b, &context); + let context = b.transpose(&context, &[0, 2, 3, 1, 4]); + let context = b.reshape( + &context, + vec![ + batch, + blocks * config.conf_attention_chunk_size, + config.hidden_size, + ], + ); + let context = if context.ty.shape[1] > time { + b.slice(&context, &[(0, batch), (0, time), (0, config.hidden_size)]) + } else { + context + }; + let post = b.linear_last(&context, args.weight(&format!("{prefix}.post.weight"))); + let post = round_bf16(b, &post); + let post = clip(b, &post, config.gradient_clipping); + let post = rms_norm( + b, + &post, + Some(args.weight(&format!("{prefix}.post_norm.weight"))), + config.rms_norm_eps, + ); + let post = round_bf16(b, &post); + let output = b.add(input, &post); + round_bf16(b, &output) +} diff --git a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit.rs b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit.rs new file mode 100644 index 000000000..814dac1f3 --- /dev/null +++ b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit.rs @@ -0,0 +1,466 @@ +//! Split Gemma3n `audio.encode` and `audio.merge_ple` StableHLO emitters. + +use super::builder::{Builder, Precision, Val}; +use super::gemma3n::Gemma3nConfig; +use super::gemma3n_audio_attention::attention; +use super::gemma3n_audio_math::{ + rms_norm, round_bf16, scalar_like, stride_time, zero_invalid, zeros_like, +}; +use super::gemma3n_audio_ops::{feed_forward, light_conv, subsample_with_stages}; +use super::gemma3n_audio_schema::{ + Decl, EncoderArgs, MergeArgs, build_encoder_schema, build_merge_schema, +}; +use super::gemma3n_emit_ops::{bf16_scalar, dense_ple_input_head}; +use crate::{GEMMA3N_AUDIO_FRAME_BUCKETS, GEMMA3N_AUDIO_SOFT_TOKENS, Gemma3nXlaAudioConfig}; + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct Gemma3nAudioDiagnosticStage { + pub name: String, + pub shape: Vec, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct Gemma3nAudioDiagnosticLayout { + pub stages: Vec, +} + +struct Trace { + stages: Vec<(String, Val)>, +} + +impl Trace { + fn new() -> Self { + Self { stages: Vec::new() } + } + + fn push(&mut self, name: impl Into, value: &Val) { + self.stages.push((name.into(), value.clone())); + } + + fn layout(&self) -> Gemma3nAudioDiagnosticLayout { + Gemma3nAudioDiagnosticLayout { + stages: self + .stages + .iter() + .map(|(name, value)| Gemma3nAudioDiagnosticStage { + name: name.clone(), + shape: value.ty.shape.clone(), + }) + .collect(), + } + } +} + +struct EncodeOutputs { + projected_audio: Val, + hard_audio: Val, + projected_lengths: Val, +} + +struct MergeOutputs { + embeddings: Val, + dense_ple: Val, +} + +fn final_conformer_norm( + b: &mut Builder, + args: &EncoderArgs, + audio: &Gemma3nXlaAudioConfig, + prefix: &str, + input: &Val, +) -> Val { + let clipped = super::gemma3n_audio_math::clip(b, input, audio.gradient_clipping); + let normalized = rms_norm( + b, + &clipped, + Some(args.weight(&format!("{prefix}.norm.weight"))), + audio.rms_norm_eps, + ); + round_bf16(b, &normalized) +} + +fn audio_encoder( + b: &mut Builder, + args: &EncoderArgs, + audio: &Gemma3nXlaAudioConfig, + trace: &mut Trace, +) -> (Val, Val, Val) { + let subsampled = subsample_with_stages(b, args, audio); + trace.push("sscp_conv_0", &subsampled.conv0); + trace.push("sscp_conv_1", &subsampled.conv1); + trace.push("input_projection", &subsampled.hidden); + let mut hidden = subsampled.hidden; + let mut valid = subsampled.valid; + for layer in 0..audio.conf_num_hidden_layers { + let prefix = format!("audio_tower.conformer.{layer}"); + hidden = feed_forward( + b, + args, + audio, + &format!("{prefix}.ffw_layer_start"), + &hidden, + ); + trace.push(format!("conformer.{layer}.feed_forward_start"), &hidden); + hidden = attention( + b, + args, + audio, + &format!("{prefix}.attention"), + &hidden, + &valid, + ); + trace.push(format!("conformer.{layer}.attention"), &hidden); + let masked = zero_invalid(b, &hidden, &valid); + hidden = light_conv(b, args, audio, &format!("{prefix}.lconv1d"), &masked); + trace.push(format!("conformer.{layer}.light_conv"), &hidden); + hidden = feed_forward(b, args, audio, &format!("{prefix}.ffw_layer_end"), &hidden); + trace.push(format!("conformer.{layer}.feed_forward_end"), &hidden); + hidden = final_conformer_norm(b, args, audio, &prefix, &hidden); + trace.push(format!("conformer.{layer}.final_norm"), &hidden); + } + if audio.conf_reduction_factor > 1 { + hidden = stride_time(b, &hidden, audio.conf_reduction_factor); + valid = stride_time(b, &valid, audio.conf_reduction_factor); + } + hidden = zero_invalid(b, &hidden, &valid); + trace.push("encoded_reduced", &hidden); + let valid_f32 = b.convert(&valid, "f32"); + let zero = b.const_f32(0.0); + let lengths = b.reduce_add(&valid_f32, 1, &zero); + let lengths = b.convert(&lengths, "i32"); + (hidden, valid, lengths) +} + +fn project_audio( + b: &mut Builder, + args: &EncoderArgs, + audio: &Gemma3nXlaAudioConfig, + text: &Gemma3nConfig, + encoded: &Val, + valid: &Val, +) -> (Val, Val) { + let soft = rms_norm( + b, + encoded, + Some(args.weight("embed_audio.soft_embedding_norm.weight")), + audio.rms_norm_eps, + ); + let soft = round_bf16(b, &soft); + let soft = b.linear_last( + &soft, + args.weight("embed_audio.embedding_projection.weight"), + ); + let soft = round_bf16(b, &soft); + let soft = rms_norm(b, &soft, None, audio.rms_norm_eps); + let soft = round_bf16(b, &soft); + + let hard = rms_norm( + b, + args.weight("embed_audio.embedding.weight"), + Some(args.weight("embed_audio.hard_embedding_norm.weight")), + audio.rms_norm_eps, + ); + let hard = round_bf16(b, &hard); + let hard = b.linear_last( + &hard, + args.weight("embed_audio.embedding_projection.weight"), + ); + let hard = round_bf16(b, &hard); + let hard = rms_norm(b, &hard, None, audio.rms_norm_eps); + let hard = round_bf16(b, &hard); + let padding = b.slice( + &hard, + &[(audio.vocab_size - 1, audio.vocab_size), (0, text.hidden)], + ); + + let batch = soft.ty.shape[0]; + let frames = soft.ty.shape[1]; + let valid = b.reshape(valid, vec![batch, frames, 1]); + let valid = b.broadcast(&valid, &[0, 1, 2], soft.ty.shape.clone()); + let padding = b.reshape(&padding, vec![1, 1, text.hidden]); + let padding = b.broadcast(&padding, &[0, 1, 2], soft.ty.shape.clone()); + let projected = b.select(&valid, &soft, &padding); + let projected = if frames < GEMMA3N_AUDIO_SOFT_TOKENS { + let padding = b.slice( + &hard, + &[(audio.vocab_size - 1, audio.vocab_size), (0, text.hidden)], + ); + let padding = b.reshape(&padding, vec![1, 1, text.hidden]); + let padding = b.broadcast( + &padding, + &[0, 1, 2], + vec![batch, GEMMA3N_AUDIO_SOFT_TOKENS - frames, text.hidden], + ); + b.concatenate(&projected, &padding, 1) + } else { + projected + }; + ( + b.reshape( + &projected, + vec![batch * GEMMA3N_AUDIO_SOFT_TOKENS, text.hidden], + ), + hard, + ) +} + +fn gather_rows(b: &mut Builder, table: &Val, indices: &Val) -> Val { + let indices = b.reshape(indices, vec![indices.ty.shape[0], 1]); + b.gather(table, &indices) +} + +fn token_embeddings( + b: &mut Builder, + args: &MergeArgs, + audio: &Gemma3nXlaAudioConfig, + text: &Gemma3nConfig, +) -> Val { + let capacity = text.context_capacity; + let text_table = args.weight("model.language_model.embed_tokens.weight"); + let embedded = gather_rows(b, text_table, &args.tokens); + let scale = scalar_like(b, (text.hidden as f32).sqrt(), &embedded); + let embedded = b.multiply(&embedded, &scale); + let mut merged = round_bf16(b, &embedded); + + let offset = b.const_i32(audio.vocab_offset); + let offset = b.broadcast(&offset, &[], vec![capacity]); + let hard_mask = b.compare("GE", &args.tokens, &offset, "SIGNED"); + let local = b.subtract(&args.tokens, &offset); + let dummy = b.const_i32((audio.vocab_size - 1) as i32); + let dummy = b.broadcast(&dummy, &[], vec![capacity]); + let local = b.select(&hard_mask, &local, &dummy); + let hard = gather_rows(b, &args.hard_audio, &local); + let hard_mask = b.reshape(&hard_mask, vec![capacity, 1]); + let hard_mask = b.broadcast(&hard_mask, &[0, 1], hard.ty.shape.clone()); + merged = b.select(&hard_mask, &hard, &merged); + + let zero = b.const_i32(0); + let zero = b.broadcast(&zero, &[], vec![capacity]); + let mapped = b.compare("GE", &args.audio_rows, &zero, "SIGNED"); + let audio_token = b.const_i32(audio.vocab_offset + 1); + let audio_token = b.broadcast(&audio_token, &[], vec![capacity]); + let token_is_audio = b.compare("EQ", &args.tokens, &audio_token, "SIGNED"); + let soft_mask = b.logical_and(&mapped, &token_is_audio); + let rank = b.maximum(&args.audio_rows, &zero); + let soft = gather_rows(b, &args.projected_audio, &rank); + let soft_mask = b.reshape(&soft_mask, vec![capacity, 1]); + let soft_mask = b.broadcast(&soft_mask, &[0, 1], soft.ty.shape.clone()); + b.select(&soft_mask, &soft, &merged) +} + +fn dense_ple(b: &mut Builder, args: &MergeArgs, text: &Gemma3nConfig, embeddings: &Val) -> Val { + let capacity = text.context_capacity; + let zero = b.const_f32(0.0); + let eps = b.const_f32(text.eps); + let inv_sqrt2 = bf16_scalar(b, std::f32::consts::FRAC_1_SQRT_2); + dense_ple_input_head( + b, + text, + &args.tokens, + embeddings, + args.weight("model.language_model.embed_tokens_per_layer.weight"), + args.weight("model.language_model.per_layer_model_projection.weight"), + args.weight("model.language_model.per_layer_projection_norm.weight"), + capacity, + &eps, + &zero, + &inv_sqrt2, + ) +} + +fn mask_output_rows( + b: &mut Builder, + tokens: &Val, + real_len: &Val, + embeddings: &Val, + ple: &Val, +) -> (Val, Val) { + let rows = tokens.ty.shape[0]; + let indices = b.iota(rows); + let real_len = b.broadcast(real_len, &[], vec![rows]); + let valid = b.compare("LT", &indices, &real_len, "SIGNED"); + let embeddings_valid = b.reshape(&valid, vec![rows, 1]); + let embeddings_valid = b.broadcast(&embeddings_valid, &[0, 1], embeddings.ty.shape.clone()); + let embeddings_zero = zeros_like(b, embeddings); + let embeddings = b.select(&embeddings_valid, embeddings, &embeddings_zero); + let ple_valid = b.reshape(&valid, vec![rows, 1, 1]); + let ple_valid = b.broadcast(&ple_valid, &[0, 1, 2], ple.ty.shape.clone()); + let ple_zero = zeros_like(b, ple); + let ple = b.select(&ple_valid, ple, &ple_zero); + (embeddings, ple) +} + +fn build_encode( + b: &mut Builder, + args: &EncoderArgs, + audio: &Gemma3nXlaAudioConfig, + text: &Gemma3nConfig, + trace: &mut Trace, +) -> EncodeOutputs { + let (encoded, valid, projected_lengths) = audio_encoder(b, args, audio, trace); + let (projected_audio, hard_audio) = project_audio(b, args, audio, text, &encoded, &valid); + trace.push("soft_projection", &projected_audio); + trace.push("hard_projection", &hard_audio); + EncodeOutputs { + projected_audio, + hard_audio, + projected_lengths, + } +} + +fn build_merge( + b: &mut Builder, + args: &MergeArgs, + audio: &Gemma3nXlaAudioConfig, + text: &Gemma3nConfig, + trace: &mut Trace, +) -> MergeOutputs { + let embeddings = token_embeddings(b, args, audio, text); + trace.push("merged_embeddings", &embeddings); + let ple = dense_ple(b, args, text, &embeddings); + trace.push("dense_ple", &ple); + let (embeddings, dense_ple) = + mask_output_rows(b, &args.tokens, &args.real_len, &embeddings, &ple); + MergeOutputs { + embeddings, + dense_ple, + } +} + +fn signature(decls: &[Decl]) -> String { + decls + .iter() + .enumerate() + .map(|(index, decl)| format!("%arg{index}: {} loc(\"{}\")", decl.ty.render(), decl.loc)) + .collect::>() + .join(", ") +} + +fn render_encode(decls: &[Decl], b: &Builder, outputs: &EncodeOutputs) -> String { + let projected_ty = outputs.projected_audio.ty.render(); + let hard_ty = outputs.hard_audio.ty.render(); + let lengths_ty = outputs.projected_lengths.ty.render(); + format!( + "module @audio_encode {{\n func.func public @main({signature}) -> ({projected_ty}, \ + {hard_ty}, {lengths_ty}) {{\n{body} return {projected}, {hard}, {lengths} : \ + {projected_ty}, {hard_ty}, {lengths_ty}\n }}\n}}\n", + signature = signature(decls), + body = b.body(), + projected = outputs.projected_audio.name, + hard = outputs.hard_audio.name, + lengths = outputs.projected_lengths.name, + ) +} + +fn render_merge(decls: &[Decl], b: &Builder, outputs: &MergeOutputs) -> String { + let embeddings_ty = outputs.embeddings.ty.render(); + let ple_ty = outputs.dense_ple.ty.render(); + format!( + "module @audio_merge_ple {{\n func.func public @main({signature}) -> \ + ({embeddings_ty}, {ple_ty}) {{\n{body} return {embeddings}, {ple} : \ + {embeddings_ty}, {ple_ty}\n }}\n}}\n", + signature = signature(decls), + body = b.body(), + embeddings = outputs.embeddings.name, + ple = outputs.dense_ple.name, + ) +} + +fn render_diagnostic(decls: &[Decl], b: &Builder, stage: &Val) -> String { + let stage_ty = stage.ty.render(); + format!( + "module @audio_diagnostic {{\n func.func public @main({signature}) -> {stage_ty} \ + {{\n{body} return {stage} : {stage_ty}\n }}\n}}\n", + signature = signature(decls), + body = b.body(), + stage = stage.name, + ) +} + +fn validate_artifact( + text: &Gemma3nConfig, + audio: &Gemma3nXlaAudioConfig, + frame_bucket: usize, + clips: usize, +) -> Result<(), String> { + audio.artifact_identity( + frame_bucket, + clips, + text.hidden, + text.n_layers, + text.hidden_per_layer_input, + )?; + if audio.projected_frames(frame_bucket)? > GEMMA3N_AUDIO_SOFT_TOKENS { + return Err("Gemma3n audio bucket exceeds the fixed 188-token projection".into()); + } + Ok(()) +} + +pub(crate) fn emit_gemma3n_audio_encode( + text: &Gemma3nConfig, + audio: &Gemma3nXlaAudioConfig, + frame_bucket: usize, + clips: usize, + precision: Precision, +) -> Result<(String, Gemma3nAudioDiagnosticLayout), String> { + validate_artifact(text, audio, frame_bucket, clips)?; + let (decls, args) = build_encoder_schema(audio, text, frame_bucket, clips)?; + let mut builder = Builder::new().with_precision(precision); + let mut trace = Trace::new(); + let outputs = build_encode(&mut builder, &args, audio, text, &mut trace); + Ok((render_encode(&decls, &builder, &outputs), trace.layout())) +} + +pub(crate) fn emit_gemma3n_audio_merge_ple( + text: &Gemma3nConfig, + audio: &Gemma3nXlaAudioConfig, + clips: usize, + precision: Precision, +) -> Result<(String, Gemma3nAudioDiagnosticLayout), String> { + // The merge artifact is context-bucketed and independent of mel length. + // Validate against the smallest legal frame bucket so the shared + // architecture identity is still checked without coupling this graph to a + // particular encoder bucket. + validate_artifact(text, audio, GEMMA3N_AUDIO_FRAME_BUCKETS[0], clips)?; + let (decls, args) = build_merge_schema(audio, text, clips); + let mut builder = Builder::new().with_precision(precision); + let mut trace = Trace::new(); + let outputs = build_merge(&mut builder, &args, audio, text, &mut trace); + Ok((render_merge(&decls, &builder, &outputs), trace.layout())) +} + +#[cfg(feature = "diagnostics")] +pub(crate) fn emit_gemma3n_audio_diagnostic( + text: &Gemma3nConfig, + audio: &Gemma3nXlaAudioConfig, + frame_bucket: usize, + clips: usize, + precision: Precision, + stage_name: &str, +) -> Result<(String, Gemma3nAudioDiagnosticStage), String> { + audio.artifact_identity( + frame_bucket, + clips, + text.hidden, + text.n_layers, + text.hidden_per_layer_input, + )?; + let (decls, args) = build_encoder_schema(audio, text, frame_bucket, clips)?; + let mut builder = Builder::new().with_precision(precision); + let mut trace = Trace::new(); + let _ = build_encode(&mut builder, &args, audio, text, &mut trace); + let (_, stage) = trace + .stages + .iter() + .find(|(name, _)| name == stage_name) + .ok_or_else(|| format!("unknown Gemma3n audio diagnostic stage `{stage_name}`"))?; + let metadata = Gemma3nAudioDiagnosticStage { + name: stage_name.to_string(), + shape: stage.ty.shape.clone(), + }; + Ok((render_diagnostic(&decls, &builder, stage), metadata)) +} + +#[cfg(test)] +#[path = "gemma3n_audio_emit_tests.rs"] +mod tests; diff --git a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit_tests.rs b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit_tests.rs new file mode 100644 index 000000000..0d7bd32a7 --- /dev/null +++ b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit_tests.rs @@ -0,0 +1,130 @@ +use super::*; + +fn text() -> Gemma3nConfig { + Gemma3nConfig::from_json_str( + &serde_json::json!({ + "model_type": "gemma3n_text", + "hidden_size": 8, "intermediate_size": [12, 12], + "max_position_embeddings": 4096, + "num_hidden_layers": 2, "num_attention_heads": 2, + "num_key_value_heads": 1, "head_dim": 4, "rms_norm_eps": 1e-6, + "vocab_size": 12, "vocab_size_per_layer_input": 10, + "hidden_size_per_layer_input": 2, + "layer_types": ["sliding_attention", "full_attention"], + "activation_sparsity_pattern": [0.0, 0.0], + "sliding_window": 2, "rope_theta": 1000000.0, + "rope_local_base_freq": 10000.0, "final_logit_softcapping": 30.0, + "num_kv_shared_layers": 0, "altup_num_inputs": 2, + "altup_active_idx": 0, "altup_coef_clip": 120.0, + "altup_correct_scale": true, "laurel_rank": 2 + }) + .to_string(), + ) + .unwrap() + .with_context_capacity(8) + .unwrap() +} + +#[test] +fn split_audio_graphs_preserve_stage_and_weight_boundaries() { + let text = text(); + let audio = Gemma3nXlaAudioConfig::default(); + let (encode, encode_layout) = + emit_gemma3n_audio_encode(&text, &audio, 8, 1, Precision::F32).unwrap(); + let (merge, merge_layout) = + emit_gemma3n_audio_merge_ple(&text, &audio, 1, Precision::F32).unwrap(); + + assert!(encode.starts_with("module @audio_encode {")); + assert!(encode.contains("stablehlo.convolution")); + assert!(encode.contains("\"stablehlo.reduce_window\"")); + assert!(!encode.contains("loc(\"audio.hard_embeddings\")")); + assert!(!encode.contains("model.language_model.")); + assert_eq!(encode_layout.stages.len(), 3 + 12 * 5 + 3); + for name in [ + "sscp_conv_0", + "sscp_conv_1", + "input_projection", + "conformer.0.feed_forward_start", + "conformer.11.final_norm", + "encoded_reduced", + "soft_projection", + "hard_projection", + ] { + assert!(encode_layout.stages.iter().any(|stage| stage.name == name)); + } + + assert!(merge.starts_with("module @audio_merge_ple {")); + assert!(merge.contains("loc(\"audio.projected\")")); + assert!(merge.contains("loc(\"audio.hard_embeddings\")")); + assert!(merge.contains("model.language_model.")); + assert!(!merge.contains("audio_tower.")); + assert!(!merge.contains("embed_audio.")); + assert_eq!( + merge_layout + .stages + .iter() + .map(|stage| stage.name.as_str()) + .collect::>(), + ["merged_embeddings", "dense_ple"] + ); +} + +fn compile(compiler: &std::ffi::OsStr, device: &str, stem: &str, artifact: &str, mlir: String) { + let input = std::env::temp_dir().join(format!("{stem}-{artifact}.mlir")); + let output = std::env::temp_dir().join(format!("{stem}-{artifact}.vmfb")); + std::fs::write(&input, mlir).unwrap(); + let mut command = std::process::Command::new(compiler); + command.arg("--iree-input-type=stablehlo"); + match device { + "local" => { + command + .arg("--iree-hal-target-device=local") + .arg("--iree-hal-local-target-device-backends=llvm-cpu"); + } + "cuda" => { + command + .arg("--iree-hal-target-device=cuda") + .arg("--iree-cuda-target=sm_80"); + } + _ => panic!("unsupported IREE test device {device}"), + } + let result = command.arg(&input).arg("-o").arg(&output).output().unwrap(); + let _ = std::fs::remove_file(&input); + let _ = std::fs::remove_file(&output); + assert!( + result.status.success(), + "audio.{artifact} failed IREE {device} compile: {}", + String::from_utf8_lossy(&result.stderr) + ); +} + +#[test] +#[ignore = "requires GEMMA3N_MODEL_DIR and the pinned IREE compiler"] +fn iree_compiles_split_real_audio_artifacts() { + let model_dir = + std::path::PathBuf::from(std::env::var_os("GEMMA3N_MODEL_DIR").expect("model directory")); + let compiler = std::env::var_os("MLXCEL_XLA_IREE_COMPILE").expect("IREE compiler"); + let frame_bucket = std::env::var("GEMMA3N_AUDIO_TEST_BUCKET") + .map(|value| value.parse::().expect("frame bucket")) + .unwrap_or(8); + let context_capacity = std::env::var("GEMMA3N_AUDIO_TEST_CONTEXT") + .map(|value| value.parse::().expect("context capacity")) + .unwrap_or(256); + let device = std::env::var("GEMMA3N_AUDIO_TEST_DEVICE").unwrap_or_else(|_| "local".to_string()); + let text = Gemma3nConfig::from_json(&model_dir) + .unwrap() + .with_context_capacity(context_capacity) + .unwrap(); + let audio = Gemma3nXlaAudioConfig::from_model_dir(&model_dir) + .unwrap() + .unwrap(); + let (encode, _) = + emit_gemma3n_audio_encode(&text, &audio, frame_bucket, 1, Precision::Bf16).unwrap(); + let (merge, _) = emit_gemma3n_audio_merge_ple(&text, &audio, 1, Precision::Bf16).unwrap(); + let stem = format!( + "mlxcel-gemma3n-audio-{frame_bucket}-{context_capacity}-{}", + std::process::id() + ); + compile(&compiler, &device, &stem, "encode", encode); + compile(&compiler, &device, &stem, "merge-ple", merge); +} diff --git a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_math.rs b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_math.rs new file mode 100644 index 000000000..7f9581cef --- /dev/null +++ b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_math.rs @@ -0,0 +1,127 @@ +//! Shared exact math for the Gemma3n audio StableHLO graph. + +use super::builder::{Builder, Val}; + +pub(super) fn scalar_like(b: &mut Builder, value: f32, x: &Val) -> Val { + let scalar = b.const_f32(value); + b.broadcast(&scalar, &[], x.ty.shape.clone()) +} + +pub(super) fn zeros_like(b: &mut Builder, x: &Val) -> Val { + scalar_like(b, 0.0, x) +} + +pub(super) fn round_bf16(b: &mut Builder, x: &Val) -> Val { + let narrow = b.convert(x, "bf16"); + b.convert(&narrow, "f32") +} + +pub(super) fn clip(b: &mut Builder, x: &Val, limit: f32) -> Val { + let low = scalar_like(b, -limit, x); + let high = scalar_like(b, limit, x); + let x = b.maximum(x, &low); + b.minimum(&x, &high) +} + +pub(super) fn relu(b: &mut Builder, x: &Val) -> Val { + let zero = zeros_like(b, x); + b.maximum(x, &zero) +} + +pub(super) fn sigmoid(b: &mut Builder, x: &Val) -> Val { + let negative = b.negate(x); + let exponential = b.exponential(&negative); + let one = scalar_like(b, 1.0, x); + let denominator = b.add(&one, &exponential); + b.divide(&one, &denominator) +} + +pub(super) fn silu(b: &mut Builder, x: &Val) -> Val { + let gate = sigmoid(b, x); + b.multiply(x, &gate) +} + +pub(super) fn softplus(b: &mut Builder, x: &Val) -> Val { + let exponential = b.exponential(x); + let one = scalar_like(b, 1.0, x); + let sum = b.add(&one, &exponential); + b.logarithm(&sum) +} + +pub(super) fn rms_norm(b: &mut Builder, x: &Val, weight: Option<&Val>, eps: f32) -> Val { + let last = x.ty.shape.len() - 1; + let width = x.ty.shape[last]; + let square = b.multiply(x, x); + let zero = b.const_f32(0.0); + let sum = b.reduce_add(&square, last, &zero); + let divisor = scalar_like(b, width as f32, &sum); + let mean = b.divide(&sum, &divisor); + let epsilon = scalar_like(b, eps, &mean); + let stabilized = b.add(&mean, &epsilon); + let inverse = b.rsqrt(&stabilized); + let dims: Vec = (0..last).collect(); + let inverse = b.broadcast(&inverse, &dims, x.ty.shape.clone()); + let normalized = b.multiply(x, &inverse); + match weight { + Some(weight) => { + let weight = b.broadcast(weight, &[last], x.ty.shape.clone()); + b.multiply(&normalized, &weight) + } + None => normalized, + } +} + +pub(super) fn cumsum_time_f32(b: &mut Builder, x: &Val) -> Val { + assert_eq!(x.ty.shape.len(), 4); + b.reduce_window_prefix_add(x, 1) +} + +pub(super) fn softmax_last(b: &mut Builder, x: &Val) -> Val { + let last = x.ty.shape.len() - 1; + let negative_infinity = b.const_f32(f32::NEG_INFINITY); + let maximum = b.reduce_max(x, last, &negative_infinity); + let dims: Vec = (0..last).collect(); + let maximum = b.broadcast(&maximum, &dims, x.ty.shape.clone()); + let centered = b.subtract(x, &maximum); + let exponential = b.exponential(¢ered); + let zero = b.const_f32(0.0); + let denominator = b.reduce_add(&exponential, last, &zero); + let denominator = b.broadcast(&denominator, &dims, x.ty.shape.clone()); + b.divide(&exponential, &denominator) +} + +pub(super) fn stride_time(b: &mut Builder, x: &Val, stride: usize) -> Val { + let ranges = + x.ty.shape + .iter() + .enumerate() + .map(|(axis, &limit)| { + if axis == 1 { + (0, limit, stride) + } else { + (0, limit, 1) + } + }) + .collect::>(); + b.slice_strided(x, &ranges) +} + +pub(super) fn concat_axis(b: &mut Builder, values: &[Val], axis: usize) -> Val { + let mut result = values[0].clone(); + for value in &values[1..] { + result = b.concatenate(&result, value, axis); + } + result +} + +pub(super) fn zero_invalid(b: &mut Builder, x: &Val, valid: &Val) -> Val { + let mut shape = valid.ty.shape.clone(); + while shape.len() < x.ty.shape.len() { + shape.push(1); + } + let valid = b.reshape(valid, shape); + let dims = (0..valid.ty.shape.len()).collect::>(); + let valid = b.broadcast(&valid, &dims, x.ty.shape.clone()); + let zero = zeros_like(b, x); + b.select(&valid, x, &zero) +} diff --git a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_ops.rs b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_ops.rs new file mode 100644 index 000000000..2dc924093 --- /dev/null +++ b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_ops.rs @@ -0,0 +1,259 @@ +//! SSCP, feed-forward, and light-convolution audio graph operations. + +use super::builder::{Builder, Val}; +use super::gemma3n_audio_math::{ + clip, cumsum_time_f32, relu, rms_norm, round_bf16, sigmoid, silu, stride_time, +}; +use super::gemma3n_audio_schema::EncoderArgs; +use crate::Gemma3nXlaAudioConfig; + +pub(super) struct SscpOutput { + pub conv0: Val, + pub conv1: Val, + pub hidden: Val, + pub valid: Val, +} + +pub(super) struct ConformerStages { + pub feed_forward_start: Val, + pub attention: Val, + pub light_conv: Val, + pub feed_forward_end: Val, + pub final_norm: Val, +} + +fn cumulative_group_norm(b: &mut Builder, x: &Val, weight: &Val, eps: f32) -> Val { + let batch = x.ty.shape[0]; + let time = x.ty.shape[1]; + let frequency = x.ty.shape[2]; + let channels = x.ty.shape[3]; + let zero = b.const_f32(0.0); + let sum_channels = b.reduce_add(x, 3, &zero); + let sum_frequency = b.reduce_add(&sum_channels, 2, &zero); + let sum_at_time = b.reshape(&sum_frequency, vec![batch, time, 1, 1]); + let cumulative_sum = cumsum_time_f32(b, &sum_at_time); + + let index = b.iota(time); + let index = b.convert(&index, "f32"); + let one = super::gemma3n_audio_math::scalar_like(b, 1.0, &index); + let count = b.add(&index, &one); + let width = super::gemma3n_audio_math::scalar_like(b, (frequency * channels) as f32, &count); + let count = b.multiply(&count, &width); + let count = b.reshape(&count, vec![1, time, 1, 1]); + let count = b.broadcast(&count, &[0, 1, 2, 3], vec![batch, time, 1, 1]); + let mean = b.divide(&cumulative_sum, &count); + let mean = b.broadcast(&mean, &[0, 1, 2, 3], x.ty.shape.clone()); + let centered = b.subtract(x, &mean); + + let squared = b.multiply(¢ered, ¢ered); + let squared_channels = b.reduce_add(&squared, 3, &zero); + let squared_frequency = b.reduce_add(&squared_channels, 2, &zero); + let squared_at_time = b.reshape(&squared_frequency, vec![batch, time, 1, 1]); + let cumulative_squared = cumsum_time_f32(b, &squared_at_time); + let variance = b.divide(&cumulative_squared, &count); + let epsilon = super::gemma3n_audio_math::scalar_like(b, eps, &variance); + let variance = b.add(&variance, &epsilon); + let inverse = b.rsqrt(&variance); + let inverse = b.broadcast(&inverse, &[0, 1, 2, 3], x.ty.shape.clone()); + let normalized = b.multiply(¢ered, &inverse); + let weight = b.broadcast(weight, &[3], x.ty.shape.clone()); + let weighted = b.multiply(&normalized, &weight); + round_bf16(b, &weighted) +} + +#[allow(clippy::too_many_arguments)] +fn sscp_conv( + b: &mut Builder, + x: &Val, + weight: &Val, + norm: &Val, + kernel: [usize; 2], + stride: [usize; 2], + groups: usize, + eps: f32, +) -> Val { + let zero = b.const_f32(0.0); + let padded = b.pad(x, &zero, &[0, 0, 1, 0], &[0, kernel[0] - 1, 1, 0]); + let weight = b.transpose(weight, &[1, 2, 3, 0]); + let convolved = b.convolution(&padded, &weight, &stride, &[(0, 0), (0, 0)], groups); + let convolved = round_bf16(b, &convolved); + let normalized = cumulative_group_norm(b, &convolved, norm, eps); + let activated = relu(b, &normalized); + round_bf16(b, &activated) +} + +fn subsample_convs( + b: &mut Builder, + args: &EncoderArgs, + config: &Gemma3nXlaAudioConfig, +) -> (Val, Val) { + let batch = args.mel.ty.shape[0]; + let time = args.mel.ty.shape[1]; + let frequency = args.mel.ty.shape[2]; + let x = b.reshape(&args.mel, vec![batch, time, frequency, 1]); + let root = "audio_tower.subsample_conv_projection"; + + let conv0 = sscp_conv( + b, + &x, + args.weight(&format!("{root}.conv_0.conv.weight")), + args.weight(&format!("{root}.conv_0.norm.weight")), + config.sscp_conv_kernel_size[0], + config.sscp_conv_stride_size[0], + 1, + config.sscp_conv_group_norm_eps, + ); + let conv1 = sscp_conv( + b, + &conv0, + args.weight(&format!("{root}.conv_1.conv.weight")), + args.weight(&format!("{root}.conv_1.norm.weight")), + config.sscp_conv_kernel_size[1], + config.sscp_conv_stride_size[1], + 1, + config.sscp_conv_group_norm_eps, + ); + (conv0, conv1) +} + +pub(super) fn subsample( + b: &mut Builder, + args: &EncoderArgs, + config: &Gemma3nXlaAudioConfig, +) -> Val { + let (_, conv1) = subsample_convs(b, args, config); + let root = "audio_tower.subsample_conv_projection"; + let shape = &conv1.ty.shape; + let flattened = b.reshape(&conv1, vec![shape[0], shape[1], shape[2] * shape[3]]); + let projected = b.linear_last( + &flattened, + args.weight(&format!("{root}.input_proj_linear.weight")), + ); + round_bf16(b, &projected) +} + +pub(super) fn subsample_with_stages( + b: &mut Builder, + args: &EncoderArgs, + config: &Gemma3nXlaAudioConfig, +) -> SscpOutput { + let (conv0, conv1) = subsample_convs(b, args, config); + let root = "audio_tower.subsample_conv_projection"; + let shape = &conv1.ty.shape; + let flattened = b.reshape(&conv1, vec![shape[0], shape[1], shape[2] * shape[3]]); + let projected = b.linear_last( + &flattened, + args.weight(&format!("{root}.input_proj_linear.weight")), + ); + let hidden = round_bf16(b, &projected); + let valid = stride_time( + b, + &args.valid_mask, + config.time_stride_product().expect("validated stride"), + ); + SscpOutput { + conv0, + conv1, + hidden, + valid, + } +} + +pub(super) fn feed_forward( + b: &mut Builder, + args: &EncoderArgs, + config: &Gemma3nXlaAudioConfig, + prefix: &str, + input: &Val, +) -> Val { + let clipped = clip(b, input, config.gradient_clipping); + let normalized = rms_norm( + b, + &clipped, + Some(args.weight(&format!("{prefix}.pre_layer_norm.weight"))), + config.rms_norm_eps, + ); + let normalized = round_bf16(b, &normalized); + let first = b.linear_last( + &normalized, + args.weight(&format!("{prefix}.ffw_layer_1.weight")), + ); + let first = round_bf16(b, &first); + let first = silu(b, &first); + let first = round_bf16(b, &first); + let second = b.linear_last(&first, args.weight(&format!("{prefix}.ffw_layer_2.weight"))); + let second = round_bf16(b, &second); + let second = clip(b, &second, config.gradient_clipping); + let second = rms_norm( + b, + &second, + Some(args.weight(&format!("{prefix}.post_layer_norm.weight"))), + config.rms_norm_eps, + ); + let second = round_bf16(b, &second); + let residual_weight = + super::gemma3n_audio_math::scalar_like(b, config.conf_residual_weight, &second); + let update = b.multiply(&second, &residual_weight); + let result = b.add(input, &update); + round_bf16(b, &result) +} + +pub(super) fn light_conv( + b: &mut Builder, + args: &EncoderArgs, + config: &Gemma3nXlaAudioConfig, + prefix: &str, + input: &Val, +) -> Val { + let hidden = config.hidden_size; + let normalized = rms_norm( + b, + input, + Some(args.weight(&format!("{prefix}.pre_layer_norm.weight"))), + config.rms_norm_eps, + ); + let normalized = round_bf16(b, &normalized); + let projected = b.linear_last( + &normalized, + args.weight(&format!("{prefix}.linear_start.weight")), + ); + let projected = round_bf16(b, &projected); + let batch = projected.ty.shape[0]; + let time = projected.ty.shape[1]; + let left = b.slice(&projected, &[(0, batch), (0, time), (0, hidden)]); + let right = b.slice(&projected, &[(0, batch), (0, time), (hidden, hidden * 2)]); + let gate = sigmoid(b, &right); + let gated = b.multiply(&left, &gate); + let gated = round_bf16(b, &gated); + let gated = b.reshape(&gated, vec![batch, time, 1, hidden]); + let zero = b.const_f32(0.0); + let padded = b.pad( + &gated, + &zero, + &[0, config.conf_conv_kernel_size - 1, 0, 0], + &[0, 0, 0, 0], + ); + let weight = args.weight(&format!("{prefix}.depthwise_conv1d.weight")); + let weight = b.transpose(weight, &[1, 2, 0]); + let weight = b.reshape(&weight, vec![config.conf_conv_kernel_size, 1, 1, hidden]); + let convolved = b.convolution(&padded, &weight, &[1, 1], &[(0, 0), (0, 0)], hidden); + let convolved = round_bf16(b, &convolved); + let convolved = b.reshape(&convolved, vec![batch, time, hidden]); + let convolved = clip(b, &convolved, config.gradient_clipping); + let convolved = rms_norm( + b, + &convolved, + Some(args.weight(&format!("{prefix}.conv_norm.weight"))), + config.rms_norm_eps, + ); + let convolved = round_bf16(b, &convolved); + let activated = silu(b, &convolved); + let activated = round_bf16(b, &activated); + let projected = b.linear_last( + &activated, + args.weight(&format!("{prefix}.linear_end.weight")), + ); + let projected = round_bf16(b, &projected); + let output = b.add(input, &projected); + round_bf16(b, &output) +} diff --git a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_schema.rs b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_schema.rs new file mode 100644 index 000000000..b21994933 --- /dev/null +++ b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_schema.rs @@ -0,0 +1,287 @@ +//! Stable argument order for the split Gemma3n audio artifacts. + +use std::collections::BTreeMap; + +use super::builder::{Builder, Ty, Val}; +use super::gemma3n::Gemma3nConfig; +use crate::{ + GEMMA3N_AUDIO_CHECKPOINT_TENSOR_COUNT, GEMMA3N_AUDIO_SOFT_TOKENS, Gemma3nXlaAudioConfig, + gemma3n_audio_checkpoint_specs, +}; + +pub(super) struct Decl { + pub ty: Ty, + pub loc: String, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub(super) struct WeightSpec { + pub name: String, + pub shape: Vec, + pub projection: bool, +} + +pub(super) struct EncoderArgs { + weights: BTreeMap, + pub mel: Val, + pub valid_mask: Val, +} + +impl EncoderArgs { + pub fn weight(&self, name: &str) -> &Val { + self.weights + .get(name) + .unwrap_or_else(|| panic!("Gemma3n audio encoder weight missing from schema: {name}")) + } +} + +pub(super) struct MergeArgs { + weights: BTreeMap, + pub projected_audio: Val, + pub hard_audio: Val, + pub tokens: Val, + pub audio_rows: Val, + pub real_len: Val, +} + +impl MergeArgs { + pub fn weight(&self, name: &str) -> &Val { + self.weights + .get(name) + .unwrap_or_else(|| panic!("Gemma3n audio merge weight missing from schema: {name}")) + } +} + +fn logical_weight(name: impl Into, shape: Vec) -> WeightSpec { + let name = name.into(); + WeightSpec { + projection: shape.len() == 2 && name.ends_with(".weight") && !name.contains("norm.weight"), + name, + shape, + } +} + +fn encoder_weight_specs( + audio: &Gemma3nXlaAudioConfig, + text: &Gemma3nConfig, +) -> Result, String> { + let mut specs = gemma3n_audio_checkpoint_specs(audio, text.hidden)? + .into_iter() + .filter_map(|spec| { + (!spec.name.starts_with("embed_audio.")).then(|| logical_weight(spec.name, spec.shape)) + }) + .collect::>(); + specs.extend([ + logical_weight( + "embed_audio.embedding.weight", + vec![audio.vocab_size, audio.hidden_size], + ), + logical_weight( + "embed_audio.embedding_projection.weight", + vec![text.hidden, audio.hidden_size], + ), + logical_weight( + "embed_audio.hard_embedding_norm.weight", + vec![audio.hidden_size], + ), + logical_weight( + "embed_audio.soft_embedding_norm.weight", + vec![audio.hidden_size], + ), + ]); + Ok(specs) +} + +fn merge_weight_specs(text: &Gemma3nConfig) -> Vec { + vec![ + logical_weight( + "model.language_model.embed_tokens.weight", + vec![text.vocab, text.hidden], + ), + logical_weight( + "model.language_model.embed_tokens_per_layer.weight", + vec![ + text.per_layer_vocab, + text.n_layers * text.hidden_per_layer_input, + ], + ), + logical_weight( + "model.language_model.per_layer_model_projection.weight", + vec![text.n_layers * text.hidden_per_layer_input, text.hidden], + ), + logical_weight( + "model.language_model.per_layer_projection_norm.weight", + vec![text.hidden_per_layer_input], + ), + ] +} + +fn take(decls: &mut Vec, index: &mut usize, ty: Ty, loc: impl Into) -> Val { + let value = Builder::arg(*index, ty.clone()); + decls.push(Decl { + ty, + loc: loc.into(), + }); + *index += 1; + value +} + +pub(super) fn build_encoder_schema( + audio: &Gemma3nXlaAudioConfig, + text: &Gemma3nConfig, + frame_bucket: usize, + clips: usize, +) -> Result<(Vec, EncoderArgs), String> { + let specs = encoder_weight_specs(audio, text)?; + let mut decls = Vec::with_capacity(specs.len() + 2); + let mut index = 0; + let mut weights = BTreeMap::new(); + for spec in specs { + let value = take( + &mut decls, + &mut index, + Ty::f32(spec.shape), + spec.name.clone(), + ); + weights.insert(spec.name, value); + } + let mel = take( + &mut decls, + &mut index, + Ty::f32(vec![clips, frame_bucket, audio.input_feat_size]), + "audio.mel", + ); + let valid_mask = take( + &mut decls, + &mut index, + Ty::new(vec![clips, frame_bucket], "i1"), + "audio.valid_mask", + ); + Ok(( + decls, + EncoderArgs { + weights, + mel, + valid_mask, + }, + )) +} + +pub(super) fn build_merge_schema( + audio: &Gemma3nXlaAudioConfig, + text: &Gemma3nConfig, + clips: usize, +) -> (Vec, MergeArgs) { + let specs = merge_weight_specs(text); + let mut decls = Vec::with_capacity(specs.len() + 5); + let mut index = 0; + let mut weights = BTreeMap::new(); + for spec in specs { + let value = take( + &mut decls, + &mut index, + Ty::f32(spec.shape), + spec.name.clone(), + ); + weights.insert(spec.name, value); + } + let projected_audio = take( + &mut decls, + &mut index, + Ty::f32(vec![clips * GEMMA3N_AUDIO_SOFT_TOKENS, text.hidden]), + "audio.projected", + ); + let hard_audio = take( + &mut decls, + &mut index, + Ty::f32(vec![audio.vocab_size, text.hidden]), + "audio.hard_embeddings", + ); + let tokens = take( + &mut decls, + &mut index, + Ty::new(vec![text.context_capacity], "i32"), + "tokens", + ); + let audio_rows = take( + &mut decls, + &mut index, + Ty::new(vec![text.context_capacity], "i32"), + "audio.row_indices", + ); + let real_len = take(&mut decls, &mut index, Ty::scalar("i32"), "real_len"); + ( + decls, + MergeArgs { + weights, + projected_audio, + hard_audio, + tokens, + audio_rows, + real_len, + }, + ) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn text() -> Gemma3nConfig { + Gemma3nConfig::from_json_str( + &serde_json::json!({ + "model_type": "gemma3n_text", + "hidden_size": 8, "intermediate_size": [12, 12], + "max_position_embeddings": 4096, + "num_hidden_layers": 2, "num_attention_heads": 2, + "num_key_value_heads": 1, "head_dim": 4, "rms_norm_eps": 1e-6, + "vocab_size": 12, "vocab_size_per_layer_input": 10, + "hidden_size_per_layer_input": 2, + "layer_types": ["sliding_attention", "full_attention"], + "activation_sparsity_pattern": [0.0, 0.0], + "sliding_window": 2, "rope_theta": 1000000.0, + "rope_local_base_freq": 10000.0, "final_logit_softcapping": 30.0, + "num_kv_shared_layers": 0, "altup_num_inputs": 2, + "altup_active_idx": 0, "altup_coef_clip": 120.0, + "altup_correct_scale": true, "laurel_rank": 2 + }) + .to_string(), + ) + .unwrap() + .with_context_capacity(8) + .unwrap() + } + + #[test] + fn split_schemas_partition_audio_and_language_weights() { + let audio = Gemma3nXlaAudioConfig::default(); + let text = text(); + let encoder = encoder_weight_specs(&audio, &text).unwrap(); + let merge = merge_weight_specs(&text); + assert_eq!( + encoder.len() + merge.len(), + GEMMA3N_AUDIO_CHECKPOINT_TENSOR_COUNT + ); + assert!( + encoder + .iter() + .all(|spec| !spec.name.starts_with("model.language_model.")) + ); + assert!( + merge + .iter() + .all(|spec| spec.name.starts_with("model.language_model.")) + ); + assert!( + encoder + .iter() + .any(|spec| spec.name == "embed_audio.embedding.weight" + && spec.shape == [128, 1_536] + && spec.projection) + ); + assert_eq!( + merge.last().unwrap().name, + "model.language_model.per_layer_projection_norm.weight" + ); + } +} diff --git a/src/lib/mlxcel-xla/src/emitter/gemma3n_emit.rs b/src/lib/mlxcel-xla/src/emitter/gemma3n_emit.rs index 86fa31d18..ab1cedb5d 100644 --- a/src/lib/mlxcel-xla/src/emitter/gemma3n_emit.rs +++ b/src/lib/mlxcel-xla/src/emitter/gemma3n_emit.rs @@ -8,8 +8,8 @@ use super::builder::{Builder, Precision, Val, precision_from_env}; use super::gemma3n::{Gemma3nConfig, Gemma3nLayerType}; use super::gemma3n_emit_ops::{ Constants, altup_correct, altup_predict, apply_sliding_window, attention, bf16_scalar, - causal_mask, constants, geglu, linear_bf16, linear_seq_bf16, mean_planes_bf16, normalize_to, - rms_last_bf16, round_bf16, sparse_gelu, + causal_mask, constants, dense_ple_input_head, geglu, linear_bf16, linear_seq_bf16, + mean_planes_bf16, normalize_to, rms_last_bf16, round_bf16, sparse_gelu, }; use super::gemma3n_qmv; use super::gemma3n_schema::{Args, Decl, Input, Weights, build_schema}; @@ -620,7 +620,6 @@ pub(super) fn token_input_head( tokens: &Val, lp: usize, ) -> (Val, Val) { - let ple_width = c.n_layers * c.hidden_per_layer_input; let tokens = if tokens.ty.shape.is_empty() { b.reshape(tokens, vec![1]) } else { @@ -632,47 +631,20 @@ pub(super) fn token_input_head( let scale = b.broadcast(&scale, &[], vec![lp, c.hidden]); let base = b.multiply(&base, &scale); let base = round_bf16(b, &base); - let limit = b.const_i32(c.per_layer_vocab as i32); - let limit = b.broadcast(&limit, &[], vec![lp]); - let zero_i = b.const_i32(0); - let zero_ids = b.broadcast(&zero_i, &[], vec![lp]); - let nonnegative = b.compare("GE", &tokens, &zero_ids, "SIGNED"); - let below = b.compare("LT", &tokens, &limit, "SIGNED"); - let valid = b.select(&nonnegative, &below, &nonnegative); - let safe = b.select(&valid, &tokens, &zero_ids); - let safe = b.reshape(&safe, vec![lp, 1]); - let token_ple = b.gather(&weights.token_ple, &safe); - let zeros = b.broadcast(&k.zero, &[], vec![lp, ple_width]); - let valid = b.broadcast(&valid, &[0], vec![lp, ple_width]); - let token_ple = b.select(&valid, &token_ple, &zeros); - let token_scale = bf16_scalar(b, (c.hidden_per_layer_input as f32).sqrt()); - let token_scale = b.broadcast(&token_scale, &[], vec![lp, ple_width]); - let token_ple = b.multiply(&token_ple, &token_scale); - let token_ple = round_bf16(b, &token_ple); - let projected = linear_seq_bf16(b, &base, &weights.ple_projection); - let model_scale = bf16_scalar(b, (c.hidden as f32).sqrt().recip()); - let model_scale = b.broadcast(&model_scale, &[], vec![lp, ple_width]); - let projected = b.multiply(&projected, &model_scale); - let projected = round_bf16(b, &projected); - let projected = b.reshape(&projected, vec![lp, c.n_layers, c.hidden_per_layer_input]); - let projected = rms_last_bf16( + let dense_ple = dense_ple_input_head( b, - &projected, - Some(&weights.ple_projection_norm), + c, + &tokens, + &base, + &weights.token_ple, + &weights.ple_projection, + &weights.ple_projection_norm, + lp, &k.eps, &k.zero, - ); - let token_ple = b.reshape(&token_ple, vec![lp, c.n_layers, c.hidden_per_layer_input]); - let inv = b.broadcast( &k.inv_sqrt2, - &[], - vec![lp, c.n_layers, c.hidden_per_layer_input], ); - let combined = b.add(&projected, &token_ple); - let combined = round_bf16(b, &combined); - let combined = b.multiply(&combined, &inv); - let combined = round_bf16(b, &combined); - (base, combined) + (base, dense_ple) } fn magnitude(b: &mut Builder, value: &Val, c: &Gemma3nConfig, k: &Constants) -> Val { diff --git a/src/lib/mlxcel-xla/src/emitter/gemma3n_emit_ops.rs b/src/lib/mlxcel-xla/src/emitter/gemma3n_emit_ops.rs index 719279a8b..f3a075083 100644 --- a/src/lib/mlxcel-xla/src/emitter/gemma3n_emit_ops.rs +++ b/src/lib/mlxcel-xla/src/emitter/gemma3n_emit_ops.rs @@ -206,6 +206,62 @@ pub(super) fn rms_last_bf16( } } +/// Build the dense per-layer embedding stream from post-scale model +/// embeddings. Token prefill and audio prefill must use this exact shared +/// sequence: the BF16 rounding boundaries are part of Gemma3n's numerical +/// contract. +#[allow(clippy::too_many_arguments)] +pub(super) fn dense_ple_input_head( + b: &mut Builder, + c: &Gemma3nConfig, + tokens: &Val, + embeddings: &Val, + token_ple_weight: &Val, + projection_weight: &Val, + projection_norm_weight: &Val, + rows: usize, + eps: &Val, + zero: &Val, + inv_sqrt2: &Val, +) -> Val { + let ple_width = c.n_layers * c.hidden_per_layer_input; + let limit = b.const_i32(c.per_layer_vocab as i32); + let limit = b.broadcast(&limit, &[], vec![rows]); + let zero_i = b.const_i32(0); + let zero_ids = b.broadcast(&zero_i, &[], vec![rows]); + let nonnegative = b.compare("GE", tokens, &zero_ids, "SIGNED"); + let below = b.compare("LT", tokens, &limit, "SIGNED"); + let valid = b.select(&nonnegative, &below, &nonnegative); + let safe = b.select(&valid, tokens, &zero_ids); + let safe = b.reshape(&safe, vec![rows, 1]); + let token_ple = b.gather(token_ple_weight, &safe); + let zeros = b.broadcast(zero, &[], vec![rows, ple_width]); + let valid = b.broadcast(&valid, &[0], vec![rows, ple_width]); + let token_ple = b.select(&valid, &token_ple, &zeros); + let token_scale = bf16_scalar(b, (c.hidden_per_layer_input as f32).sqrt()); + let token_scale = b.broadcast(&token_scale, &[], vec![rows, ple_width]); + let token_ple = b.multiply(&token_ple, &token_scale); + let token_ple = round_bf16(b, &token_ple); + + let projected = linear_seq_bf16(b, embeddings, projection_weight); + let model_scale = bf16_scalar(b, (c.hidden as f32).sqrt().recip()); + let model_scale = b.broadcast(&model_scale, &[], vec![rows, ple_width]); + let projected = b.multiply(&projected, &model_scale); + let projected = round_bf16(b, &projected); + let projected = b.reshape(&projected, vec![rows, c.n_layers, c.hidden_per_layer_input]); + let projected = rms_last_bf16(b, &projected, Some(projection_norm_weight), eps, zero); + let token_ple = b.reshape(&token_ple, vec![rows, c.n_layers, c.hidden_per_layer_input]); + let inv = b.broadcast( + inv_sqrt2, + &[], + vec![rows, c.n_layers, c.hidden_per_layer_input], + ); + let combined = b.add(&projected, &token_ple); + let combined = round_bf16(b, &combined); + let combined = b.multiply(&combined, &inv); + round_bf16(b, &combined) +} + pub(super) fn normalize_to(b: &mut Builder, plane: &Val, target: &Val, k: &Constants) -> Val { let axis = plane.ty.shape.len() - 1; let width = plane.ty.shape[axis]; diff --git a/src/lib/mlxcel-xla/src/emitter/mod.rs b/src/lib/mlxcel-xla/src/emitter/mod.rs index 0065385c7..685aa6920 100644 --- a/src/lib/mlxcel-xla/src/emitter/mod.rs +++ b/src/lib/mlxcel-xla/src/emitter/mod.rs @@ -48,6 +48,11 @@ mod builder; mod config; mod gemma3n; +mod gemma3n_audio_attention; +mod gemma3n_audio_emit; +mod gemma3n_audio_math; +mod gemma3n_audio_ops; +mod gemma3n_audio_schema; mod gemma3n_decode; mod gemma3n_emit; mod gemma3n_emit_ops; @@ -76,6 +81,12 @@ pub(crate) use config::{ pub(crate) use gemma3n::{ Gemma3nConfig, Gemma3nLayerType, Gemma3nPleDType, Gemma3nPleMetadata, validate_gemma3n_ple, }; +#[cfg(feature = "diagnostics")] +pub(crate) use gemma3n_audio_emit::emit_gemma3n_audio_diagnostic; +#[cfg(feature = "diagnostics")] +pub use gemma3n_audio_emit::{Gemma3nAudioDiagnosticLayout, Gemma3nAudioDiagnosticStage}; +#[allow(unused_imports)] +pub(crate) use gemma3n_audio_emit::{emit_gemma3n_audio_encode, emit_gemma3n_audio_merge_ple}; #[allow(unused_imports)] pub(crate) use gemma3n_decode::{ emit_gemma3n_decode, emit_gemma3n_decode_ragged, emit_gemma3n_decode_ragged_with, diff --git a/src/lib/mlxcel-xla/src/gemma3n_audio_config.rs b/src/lib/mlxcel-xla/src/gemma3n_audio_config.rs index 23d104cb8..ee7198be0 100644 --- a/src/lib/mlxcel-xla/src/gemma3n_audio_config.rs +++ b/src/lib/mlxcel-xla/src/gemma3n_audio_config.rs @@ -2,13 +2,13 @@ //! //! Waveform decoding, resampling, and log-mel extraction stay on the bounded //! host stage. This module defines the exact static shape and numeric policy -//! compiled into `audio.main`. +//! compiled into the split `audio.encode` and `audio.merge_ple` artifacts. use std::path::Path; use serde::Deserialize; -pub const GEMMA3N_AUDIO_GRAPH_ABI: &str = "gemma3n-audio-main-v1"; +pub const GEMMA3N_AUDIO_GRAPH_ABI: &str = "gemma3n-audio-split-v2"; pub const GEMMA3N_AUDIO_MODALITY_FAMILY: &str = "gemma3n_audio"; pub const GEMMA3N_AUDIO_MEL_BINS: usize = 128; pub const GEMMA3N_AUDIO_MAX_FRAMES: usize = 2_997; diff --git a/src/lib/mlxcel-xla/src/gemma3n_audio_rows.rs b/src/lib/mlxcel-xla/src/gemma3n_audio_rows.rs new file mode 100644 index 000000000..0060f56e3 --- /dev/null +++ b/src/lib/mlxcel-xla/src/gemma3n_audio_rows.rs @@ -0,0 +1,230 @@ +//! Fail-closed `audio.row_indices` validation at the native IREE boundary. + +use std::fmt; + +use crate::GEMMA3N_AUDIO_SOFT_TOKENS; + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum Gemma3nAudioRowMapError { + LengthMismatch { + tokens: usize, + rows: usize, + }, + RowCountOverflow, + MissingAudioRow { + token: usize, + expected: i32, + }, + UnexpectedAudioRow { + token: usize, + row: i32, + }, + OutOfRange { + token: usize, + row: i32, + rows: usize, + }, + NonCanonicalOrder { + token: usize, + row: i32, + expected: i32, + }, + PlaceholderCount { + actual: usize, + expected: usize, + }, +} + +impl fmt::Display for Gemma3nAudioRowMapError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::LengthMismatch { tokens, rows } => write!( + f, + "Gemma3n audio row map has {rows} entries for {tokens} token rows" + ), + Self::RowCountOverflow => { + f.write_str("Gemma3n audio row-map count exceeds the i32 graph ABI") + } + Self::MissingAudioRow { token, expected } => write!( + f, + "Gemma3n audio placeholder at token {token} is missing canonical row {expected}" + ), + Self::UnexpectedAudioRow { token, row } => write!( + f, + "Gemma3n non-audio token {token} carries unexpected audio row {row}" + ), + Self::OutOfRange { token, row, rows } => write!( + f, + "Gemma3n audio row {row} at token {token} is outside 0..{rows}" + ), + Self::NonCanonicalOrder { + token, + row, + expected, + } => write!( + f, + "Gemma3n audio row {row} at token {token} is duplicate or reordered; expected {expected}" + ), + Self::PlaceholderCount { actual, expected } => write!( + f, + "Gemma3n prompt has {actual} audio placeholders; expected {expected}" + ), + } + } +} + +impl std::error::Error for Gemma3nAudioRowMapError {} + +/// Validate the complete static-capacity token and row-map buffers. +/// +/// Every audio placeholder must map once, in prompt order, to the contiguous +/// flattened soft-token range. Every non-audio or padded token must carry `-1`. +pub fn validate_gemma3n_audio_row_indices( + token_ids: &[i32], + audio_token_id: i32, + clips: usize, + rows: &[i32], +) -> Result<(), Gemma3nAudioRowMapError> { + if token_ids.len() != rows.len() { + return Err(Gemma3nAudioRowMapError::LengthMismatch { + tokens: token_ids.len(), + rows: rows.len(), + }); + } + let expected_rows = clips + .checked_mul(GEMMA3N_AUDIO_SOFT_TOKENS) + .ok_or(Gemma3nAudioRowMapError::RowCountOverflow)?; + let expected_rows_i32 = + i32::try_from(expected_rows).map_err(|_| Gemma3nAudioRowMapError::RowCountOverflow)?; + let mut next = 0_i32; + for (token, (&token_id, &row)) in token_ids.iter().zip(rows).enumerate() { + if token_id != audio_token_id { + if row != -1 { + return Err(Gemma3nAudioRowMapError::UnexpectedAudioRow { token, row }); + } + continue; + } + if row < 0 { + return Err(Gemma3nAudioRowMapError::MissingAudioRow { + token, + expected: next, + }); + } + if row >= expected_rows_i32 { + return Err(Gemma3nAudioRowMapError::OutOfRange { + token, + row, + rows: expected_rows, + }); + } + if row != next { + return Err(Gemma3nAudioRowMapError::NonCanonicalOrder { + token, + row, + expected: next, + }); + } + next += 1; + } + if next != expected_rows_i32 { + return Err(Gemma3nAudioRowMapError::PlaceholderCount { + actual: usize::try_from(next).unwrap_or_default(), + expected: expected_rows, + }); + } + Ok(()) +} + +pub(crate) fn build_gemma3n_audio_row_indices( + token_ids: &[i32], + audio_token_id: i32, + clips: usize, +) -> Result, Gemma3nAudioRowMapError> { + let mut rows = Vec::with_capacity(token_ids.len()); + let mut next = 0_i32; + for &token in token_ids { + if token == audio_token_id { + rows.push(next); + next = next + .checked_add(1) + .ok_or(Gemma3nAudioRowMapError::RowCountOverflow)?; + } else { + rows.push(-1); + } + } + validate_gemma3n_audio_row_indices(token_ids, audio_token_id, clips, &rows)?; + Ok(rows) +} + +#[cfg(test)] +mod tests { + use super::*; + + const AUDIO: i32 = 99; + + fn valid() -> (Vec, Vec) { + let mut tokens = vec![1]; + tokens.extend(std::iter::repeat_n(AUDIO, GEMMA3N_AUDIO_SOFT_TOKENS)); + tokens.push(2); + let rows = build_gemma3n_audio_row_indices(&tokens, AUDIO, 1).unwrap(); + (tokens, rows) + } + + #[test] + fn accepts_exact_placeholder_order_and_padding() { + let (mut tokens, mut rows) = valid(); + tokens.resize(tokens.len() + 4, 0); + rows.resize(rows.len() + 4, -1); + validate_gemma3n_audio_row_indices(&tokens, AUDIO, 1, &rows).unwrap(); + } + + #[test] + fn rejects_length_missing_and_non_audio_mappings() { + let (tokens, rows) = valid(); + assert!(matches!( + validate_gemma3n_audio_row_indices(&tokens, AUDIO, 1, &rows[..rows.len() - 1]), + Err(Gemma3nAudioRowMapError::LengthMismatch { .. }) + )); + let mut missing = rows.clone(); + missing[1] = -1; + assert!(matches!( + validate_gemma3n_audio_row_indices(&tokens, AUDIO, 1, &missing), + Err(Gemma3nAudioRowMapError::MissingAudioRow { .. }) + )); + let mut unexpected = rows; + unexpected[0] = 0; + assert!(matches!( + validate_gemma3n_audio_row_indices(&tokens, AUDIO, 1, &unexpected), + Err(Gemma3nAudioRowMapError::UnexpectedAudioRow { .. }) + )); + } + + #[test] + fn rejects_duplicate_reordered_out_of_range_and_missing_count() { + let (tokens, rows) = valid(); + let mut duplicate = rows.clone(); + duplicate[2] = 0; + assert!(matches!( + validate_gemma3n_audio_row_indices(&tokens, AUDIO, 1, &duplicate), + Err(Gemma3nAudioRowMapError::NonCanonicalOrder { .. }) + )); + let mut reordered = rows.clone(); + reordered.swap(1, 2); + assert!(matches!( + validate_gemma3n_audio_row_indices(&tokens, AUDIO, 1, &reordered), + Err(Gemma3nAudioRowMapError::NonCanonicalOrder { .. }) + )); + let mut out_of_range = rows; + out_of_range[1] = GEMMA3N_AUDIO_SOFT_TOKENS as i32; + assert!(matches!( + validate_gemma3n_audio_row_indices(&tokens, AUDIO, 1, &out_of_range), + Err(Gemma3nAudioRowMapError::OutOfRange { .. }) + )); + let fewer_tokens = vec![AUDIO; GEMMA3N_AUDIO_SOFT_TOKENS - 1]; + let fewer_rows = (0..GEMMA3N_AUDIO_SOFT_TOKENS as i32 - 1).collect::>(); + assert!(matches!( + validate_gemma3n_audio_row_indices(&fewer_tokens, AUDIO, 1, &fewer_rows), + Err(Gemma3nAudioRowMapError::PlaceholderCount { .. }) + )); + } +} diff --git a/src/lib/mlxcel-xla/src/lib.rs b/src/lib/mlxcel-xla/src/lib.rs index 257763b8f..65034e1c3 100644 --- a/src/lib/mlxcel-xla/src/lib.rs +++ b/src/lib/mlxcel-xla/src/lib.rs @@ -53,6 +53,7 @@ mod context; #[cfg(any(feature = "diagnostics", test))] mod diagnostic_flags; mod gemma3n_audio_config; +mod gemma3n_audio_rows; mod gemma3n_audio_weights; #[cfg(any(feature = "iree", test))] #[allow(dead_code)] @@ -215,6 +216,7 @@ pub use gemma3n_audio_config::{ GEMMA3N_AUDIO_MAX_FRAMES, GEMMA3N_AUDIO_MEL_BINS, GEMMA3N_AUDIO_MODALITY_FAMILY, GEMMA3N_AUDIO_SOFT_TOKENS, Gemma3nXlaAudioConfig, }; +pub use gemma3n_audio_rows::{Gemma3nAudioRowMapError, validate_gemma3n_audio_row_indices}; pub use gemma3n_audio_weights::{ GEMMA3N_AUDIO_CHECKPOINT_TENSOR_COUNT, Gemma3nAudioCheckpointDType, Gemma3nAudioCheckpointError, Gemma3nAudioCheckpointTensorSpec, gemma3n_audio_checkpoint_specs, diff --git a/src/lib/mlxcel-xla/src/prepared_gemma3n_audio.rs b/src/lib/mlxcel-xla/src/prepared_gemma3n_audio.rs index ca7fe7583..cc9419822 100644 --- a/src/lib/mlxcel-xla/src/prepared_gemma3n_audio.rs +++ b/src/lib/mlxcel-xla/src/prepared_gemma3n_audio.rs @@ -1,6 +1,6 @@ //! Owned Gemma3n audio inputs and prepared-prefill outputs. //! -//! Canonical mel features and valid-frame masks enter `audio.main`; post-scale +//! Canonical mel features and valid-frame masks enter `audio.encode`; post-scale //! merged embeddings and dense PLE leave it. Owned buffers make normal //! completion, cancellation, and admission failure use the same drop path. @@ -12,6 +12,7 @@ use crate::gemma3n_audio_config::{ GEMMA3N_AUDIO_MEL_BINS, GEMMA3N_AUDIO_MODALITY_FAMILY, GEMMA3N_AUDIO_SOFT_TOKENS, validate_gemma3n_audio_frame_bucket, }; +use crate::gemma3n_audio_rows::build_gemma3n_audio_row_indices; pub fn select_gemma3n_audio_frame_bucket(frames: usize) -> Result { if frames == 0 { @@ -106,7 +107,7 @@ impl fmt::Display for Gemma3nAudioInputError { impl std::error::Error for Gemma3nAudioInputError {} -/// Canonical, zero-padded input to one static `audio.main` frame bucket. +/// Canonical, zero-padded input to one static `audio.encode` frame bucket. #[derive(Clone, Debug, PartialEq)] pub struct Gemma3nAudioInput { mel: Vec, @@ -231,12 +232,13 @@ impl Gemma3nAudioInput { } } -/// `audio.main` output ready for the existing Gemma3n dense-PLE prefill entry. +/// Split audio artifact output ready for the existing Gemma3n dense-PLE prefill entry. #[derive(Clone, Debug)] pub struct Gemma3nAudioPreparedPrefill { request: Gemma3nPreparedPrefill, projected_lengths: Vec, placeholder_starts: Vec, + audio_row_indices: Vec, frame_bucket: usize, } @@ -301,10 +303,14 @@ impl Gemma3nAudioPreparedPrefill { placeholder_starts.len() )); } + let audio_row_indices = + build_gemma3n_audio_row_indices(&prepared.token_ids, audio_token_id, clips) + .map_err(|error| error.to_string())?; Ok(Self { request, projected_lengths, placeholder_starts, + audio_row_indices, frame_bucket, }) } @@ -324,6 +330,23 @@ impl Gemma3nAudioPreparedPrefill { &self.placeholder_starts } + #[must_use] + pub fn audio_row_indices(&self) -> &[i32] { + &self.audio_row_indices + } + + pub fn padded_audio_row_indices(&self, capacity: usize) -> Result, String> { + if capacity < self.audio_row_indices.len() { + return Err(format!( + "Gemma3n audio row-map capacity {capacity} is smaller than sequence length {}", + self.audio_row_indices.len() + )); + } + let mut rows = self.audio_row_indices.clone(); + rows.resize(capacity, -1); + Ok(rows) + } + #[must_use] pub fn frame_bucket(&self) -> usize { self.frame_bucket @@ -454,6 +477,11 @@ mod tests { prepared.placeholder_starts(), &[1, GEMMA3N_AUDIO_SOFT_TOKENS + 2] ); + let rows = prepared.audio_row_indices(); + assert_eq!(rows[0], -1); + assert_eq!(rows[1], 0); + assert_eq!(rows[GEMMA3N_AUDIO_SOFT_TOKENS], 187); + assert_eq!(rows[GEMMA3N_AUDIO_SOFT_TOKENS + 2], 188); let mut malformed = tokens; malformed.remove(1); From 0d90dc9331dc0ab72e780da633d26adca67fd1c4 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 24 Jul 2026 16:34:26 +0900 Subject: [PATCH 04/19] feat(xla): serve Gemma3n audio requests Execute the split audio encode and merge graphs through IREE and wire them into CLI and continuous batching. Preserve pinned processor left-padding mel features while zeroing only the static bucket tail. Refs #878 --- Cargo.toml | 5 + examples/xla_gemma3n_audio_check.rs | 185 +++++ src/audio/gemma3n/mod.rs | 4 + src/audio/gemma3n/xla.rs | 263 ++++++ src/commands/generate.rs | 68 +- src/lib/mlxcel-xla/Cargo.toml | 5 + src/lib/mlxcel-xla/src/aux.rs | 31 +- src/lib/mlxcel-xla/src/aux_manifest.rs | 4 +- src/lib/mlxcel-xla/src/aux_smoke.rs | 4 +- src/lib/mlxcel-xla/src/batch.rs | 41 +- .../src/emitter/gemma3n_audio_schema.rs | 13 +- .../mlxcel-xla/src/emitter/gemma3n_schema.rs | 81 ++ src/lib/mlxcel-xla/src/emitter/mod.rs | 6 + .../mlxcel-xla/src/gemma3n_audio_runtime.rs | 786 ++++++++++++++++++ src/lib/mlxcel-xla/src/iree.rs | 124 +++ src/lib/mlxcel-xla/src/lib.rs | 31 +- .../mlxcel-xla/src/prepared_gemma3n_audio.rs | 51 +- src/server/batch/mod.rs | 6 +- src/server/batch/xla_audio_preprocess.rs | 80 +- .../batch/xla_audio_preprocess_tests.rs | 24 +- src/server/batch/xla_gemma3n_audio.rs | 145 ++++ src/server/batch/xla_worker.rs | 51 +- src/server/batch/xla_worker_admission.rs | 155 +++- src/server/batch/xla_worker_tests.rs | 200 ++++- src/server/media.rs | 16 +- 25 files changed, 2295 insertions(+), 84 deletions(-) create mode 100644 examples/xla_gemma3n_audio_check.rs create mode 100644 src/audio/gemma3n/xla.rs create mode 100644 src/lib/mlxcel-xla/src/gemma3n_audio_runtime.rs create mode 100644 src/server/batch/xla_gemma3n_audio.rs diff --git a/Cargo.toml b/Cargo.toml index ab506bce0..6068b7dba 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -343,4 +343,9 @@ required-features = ["xla-micro-oracle"] # Pinned #874 MLX-vs-IREE Phi4MM Conformer and speech/vision projection parity. [[example]] name = "xla_phi4_audio_check" + +# Deterministic Gemma3n audio split-runtime + dense-PLE language gate (#878). +# Uses the pinned #875 WAV/checkpoint fixture selected by environment variables. +[[example]] +name = "xla_gemma3n_audio_check" required-features = ["xla-iree"] diff --git a/examples/xla_gemma3n_audio_check.rs b/examples/xla_gemma3n_audio_check.rs new file mode 100644 index 000000000..48cf161e0 --- /dev/null +++ b/examples/xla_gemma3n_audio_check.rs @@ -0,0 +1,185 @@ +// 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. + +//! Real-IREE #875 fixture gate for the split Gemma3n audio runtime. + +use std::path::{Path, PathBuf}; +use std::sync::atomic::AtomicBool; + +use mlxcel::audio::gemma3n::Gemma3nAudioFeatureExtractor; +use mlxcel::audio::{AudioFamilyPolicy, preprocess_wav_file}; +use mlxcel::tokenizer::load_tokenizer; +use mlxcel_xla::{ + GEMMA3N_AUDIO_MEL_BINS, GEMMA3N_AUDIO_SOFT_TOKENS, Gemma3nAudioInput, XlaInferenceSession, + select_gemma3n_audio_frame_bucket, +}; + +const PROMPT: &str = "Transcribe the following speech segment in English:"; +const AUDIO_TOKEN_ID: i32 = 262_273; +const BOA_TOKEN_ID: i32 = 256_000; +const EOA_TOKEN_ID: i32 = 262_272; +const CONTEXT_CAPACITY: usize = 256; +const EXPECTED_TRANSCRIPT: &str = "Chapter 1. Mrs. Rachel Lind is surprised.\n\nMrs. Rachel Lind"; + +fn path_env(name: &str, default: &str) -> PathBuf { + std::env::var_os(name) + .map(PathBuf::from) + .unwrap_or_else(|| PathBuf::from(default)) +} + +fn language_layer_count(model_dir: &Path) -> Result { + let path = model_dir.join("config.json"); + let bytes = + std::fs::read(&path).map_err(|error| format!("read {}: {error}", path.display()))?; + let config: serde_json::Value = serde_json::from_slice(&bytes) + .map_err(|error| format!("parse {}: {error}", path.display()))?; + let layers = config + .get("text_config") + .and_then(|text| text.get("num_hidden_layers")) + .and_then(serde_json::Value::as_u64) + .ok_or_else(|| format!("{} has no text_config.num_hidden_layers", path.display()))?; + usize::try_from(layers) + .map_err(|_| format!("Gemma3n language layer count {layers} does not fit usize")) +} + +fn expanded_prompt( + model_dir: &Path, +) -> Result<(mlxcel::tokenizer::MlxcelTokenizer, Vec), String> { + let tokenizer = load_tokenizer(model_dir).map_err(|error| error.to_string())?; + let rendered = format!( + "user\n{PROMPT}\n\ + model\n" + ); + let mut tokens = tokenizer + .encode(&rendered, false) + .map_err(|error| error.to_string())? + .into_iter() + .map(|token| token as i32) + .collect::>(); + let wrapper = tokenizer + .encode("\n\n", false) + .map_err(|error| error.to_string())? + .into_iter() + .map(|token| token as i32) + .collect::>(); + mlxcel::vlm_runtime::expand_gemma3n_audio_tokens( + &mut tokens, + AUDIO_TOKEN_ID, + BOA_TOKEN_ID, + EOA_TOKEN_ID, + 1, + GEMMA3N_AUDIO_SOFT_TOKENS, + &wrapper, + None, + ) + .map_err(|error| error.to_string())?; + if tokens.len() != 210 { + return Err(format!( + "#875 expanded prompt has {} tokens, expected 210", + tokens.len() + )); + } + Ok((tokenizer, tokens)) +} + +fn audio_input(audio_path: &Path) -> Result { + let waveforms = preprocess_wav_file( + audio_path, + AudioFamilyPolicy::gemma3n(), + &AtomicBool::new(false), + ) + .map_err(|error| error.to_string())?; + let clips = waveforms + .clips + .into_iter() + .map(|clip| clip.samples) + .collect::>(); + let features = Gemma3nAudioFeatureExtractor::new().extract_batch(&clips)?; + if features.frames != 1_406 { + return Err(format!( + "#875 mel frontend emitted {} frames, expected 1406", + features.frames + )); + } + let bucket = + select_gemma3n_audio_frame_bucket(features.frames).map_err(|error| error.to_string())?; + if bucket != 2_048 { + return Err(format!( + "#875 selected frame bucket {bucket}, expected 2048" + )); + } + let mut mel = vec![0.0; bucket * GEMMA3N_AUDIO_MEL_BINS]; + for frame in 0..features.frames { + let start = frame * GEMMA3N_AUDIO_MEL_BINS; + mel[start..start + GEMMA3N_AUDIO_MEL_BINS] + .copy_from_slice(&features.features[start..start + GEMMA3N_AUDIO_MEL_BINS]); + } + let mut valid_mask = vec![0; bucket]; + for (output, valid) in valid_mask + .iter_mut() + .zip(features.valid_mask.iter().copied()) + { + *output = u8::from(valid); + } + let valid_frames = features.valid_mask.iter().filter(|&&valid| valid).count(); + Gemma3nAudioInput::new(mel, valid_mask, vec![valid_frames], bucket) + .map_err(|error| error.to_string()) +} + +fn main() -> Result<(), String> { + let model_dir = path_env( + "MLXCEL_GEMMA3N_AUDIO_MODEL", + "/home/inureyes/models/gemma3n-e4b-4bit", + ); + let audio_path = path_env("MLXCEL_GEMMA3N_AUDIO_WAV", "/tmp/gemma3n-103-1240-0000.wav"); + let (tokenizer, tokens) = expanded_prompt(&model_dir)?; + let input = audio_input(&audio_path)?; + eprintln!( + "gemma3n-audio-check: frontend frames=1406 bucket={} expanded={}", + input.frame_bucket(), + tokens.len() + ); + + let layers = language_layer_count(&model_dir)?; + let mut session = + XlaInferenceSession::load_with_context_capacity(&model_dir, layers, CONTEXT_CAPACITY)?; + eprintln!("gemma3n-audio-check: #876 language bundle loaded"); + let mut audio = session.load_gemma3n_audio(input.frame_bucket(), input.clips())?; + if !audio.has_capability() { + return Err("Gemma3n audio capability stayed false after verified bundle load".to_string()); + } + eprintln!("gemma3n-audio-check: split audio bundle loaded"); + let prepared = audio.invoke_prepared(&input, tokens, AUDIO_TOKEN_ID)?; + eprintln!( + "gemma3n-audio-check: encode->merge complete, projected_lengths={:?}", + prepared.projected_lengths() + ); + let eos_token_ids = session.eos_token_ids().to_vec(); + let output = + session.generate_gemma3n_prepared_greedy(prepared.request(), 16, &eos_token_ids)?; + let text = tokenizer + .decode( + &output.iter().map(|token| *token as u32).collect::>(), + true, + ) + .map_err(|error| error.to_string())?; + eprintln!("gemma3n-audio-check: generated tokens={output:?}"); + if text != EXPECTED_TRANSCRIPT { + return Err(format!( + "#875 token-exact transcript mismatch:\nexpected: {EXPECTED_TRANSCRIPT:?}\nactual: {text:?}" + )); + } + println!("{text}"); + Ok(()) +} diff --git a/src/audio/gemma3n/mod.rs b/src/audio/gemma3n/mod.rs index 4f657d71d..12ded4fd9 100644 --- a/src/audio/gemma3n/mod.rs +++ b/src/audio/gemma3n/mod.rs @@ -22,6 +22,8 @@ mod attention; mod config; mod encoder; mod feature_extractor; +#[cfg(feature = "xla-backend")] +mod xla; use mlxcel_core::layers::UnifiedLinear; use mlxcel_core::weights::WeightMap; @@ -89,3 +91,5 @@ pub use feature_extractor::{ GEMMA3N_AUDIO_SOFT_TOKENS, GEMMA3N_MAX_SAMPLES, GEMMA3N_SAMPLE_RATE, Gemma3nAudioFeatureBatch, Gemma3nAudioFeatureExtractor, }; +#[cfg(feature = "xla-backend")] +pub use xla::{Gemma3nXlaAudioPreparer, Gemma3nXlaPreparedAudioInput}; diff --git a/src/audio/gemma3n/xla.rs b/src/audio/gemma3n/xla.rs new file mode 100644 index 000000000..a93b10817 --- /dev/null +++ b/src/audio/gemma3n/xla.rs @@ -0,0 +1,263 @@ +// 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. + +//! Shared Gemma3n host-front-end contract for OpenXLA CLI and serving. + +use std::path::Path; +use std::sync::atomic::{AtomicBool, Ordering}; + +use mlxcel_xla::{ + GEMMA3N_AUDIO_MAX_CLIPS, GEMMA3N_AUDIO_MEL_BINS, GEMMA3N_AUDIO_SOFT_TOKENS, Gemma3nAudioInput, + Gemma3nXlaAudioConfig, select_gemma3n_audio_frame_bucket, validate_gemma3n_audio_checkpoint, +}; + +use super::{Gemma3nAudioFeatureBatch, Gemma3nAudioFeatureExtractor}; +use crate::audio::{AudioFamilyPolicy, AudioWaveformBatch}; +use crate::tokenizer::MlxcelTokenizer; + +/// Canonical request input produced before invoking the split IREE audio graph. +pub struct Gemma3nXlaPreparedAudioInput { + pub input: Gemma3nAudioInput, + pub token_ids: Vec, + pub clips: usize, +} + +/// Model-derived token and preprocessing policy shared by CLI and server paths. +pub struct Gemma3nXlaAudioPreparer { + context_capacity: usize, + policy: AudioFamilyPolicy, + wrapper_tokens: Vec, + audio_token_id: i32, + boa_token_id: i32, + eoa_token_id: i32, + end_of_turn_token_id: Option, +} + +fn bucket_audio_features(features: Gemma3nAudioFeatureBatch) -> Result { + let frame_bucket = + select_gemma3n_audio_frame_bucket(features.frames).map_err(|error| error.to_string())?; + let mut mel = vec![0.0f32; features.batch_size * frame_bucket * GEMMA3N_AUDIO_MEL_BINS]; + let mut valid_mask = vec![0u8; features.batch_size * frame_bucket]; + let mut frame_lengths = Vec::with_capacity(features.batch_size); + for clip in 0..features.batch_size { + let mut valid_frames = 0usize; + for frame in 0..features.frames { + let source_row = clip * features.frames + frame; + let target_row = clip * frame_bucket + frame; + valid_mask[target_row] = u8::from(features.valid_mask[source_row]); + valid_frames += usize::from(features.valid_mask[source_row]); + let source = source_row * GEMMA3N_AUDIO_MEL_BINS; + let target = target_row * GEMMA3N_AUDIO_MEL_BINS; + mel[target..target + GEMMA3N_AUDIO_MEL_BINS] + .copy_from_slice(&features.features[source..source + GEMMA3N_AUDIO_MEL_BINS]); + } + frame_lengths.push(valid_frames); + } + Gemma3nAudioInput::new(mel, valid_mask, frame_lengths, frame_bucket) + .map_err(|error| error.to_string()) +} + +impl Gemma3nXlaAudioPreparer { + /// Build an exact Gemma3n audio contract when the checkpoint declares one. + pub fn from_model( + model_path: &Path, + context_capacity: usize, + tokenizer: &MlxcelTokenizer, + ) -> Result, String> { + let Some(audio_config) = Gemma3nXlaAudioConfig::from_model_dir(model_path)? else { + return Ok(None); + }; + let config_path = model_path.join("config.json"); + let config_text = std::fs::read_to_string(&config_path) + .map_err(|error| format!("read {}: {error}", config_path.display()))?; + let config: serde_json::Value = serde_json::from_str(&config_text) + .map_err(|error| format!("parse {}: {error}", config_path.display()))?; + let text_hidden = config + .get("text_config") + .and_then(|text| text.get("hidden_size")) + .and_then(serde_json::Value::as_u64) + .and_then(|value| usize::try_from(value).ok()) + .ok_or_else(|| { + format!( + "{} has no valid text_config.hidden_size", + config_path.display() + ) + })?; + validate_gemma3n_audio_checkpoint(model_path, &audio_config, text_hidden) + .map_err(|error| error.to_string())?; + + let processor = ["preprocessor_config.json", "processor_config.json"] + .into_iter() + .find_map(|name| { + let path = model_path.join(name); + std::fs::read_to_string(path) + .ok() + .and_then(|text| serde_json::from_str::(&text).ok()) + }); + let mut policy = AudioFamilyPolicy::from_gemma3n_configs( + &config, + processor.as_ref(), + GEMMA3N_AUDIO_SOFT_TOKENS, + ) + .map_err(|error| error.to_string())?; + policy.max_clips = policy.max_clips.min(GEMMA3N_AUDIO_MAX_CLIPS); + + let wrapper_tokens = tokenizer + .encode("\n\n", false) + .map_err(|error| format!("tokenize Gemma3n audio wrapper: {error}"))? + .into_iter() + .map(|token| token as i32) + .collect::>(); + if wrapper_tokens.is_empty() { + return Err("Gemma3n audio wrapper tokenized to an empty sequence".to_string()); + } + let token = |name: &str, fallback: i32| -> Result { + config + .get(name) + .and_then(serde_json::Value::as_i64) + .map_or(Ok(fallback), |value| { + i32::try_from(value) + .map_err(|_| format!("Gemma3n {name}={value} does not fit i32")) + }) + }; + let audio_token_id = token("audio_token_id", 262_273)?; + let boa_token_id = token("boa_token_id", 256_000)?; + let eoa_token_id = token("eoa_token_id", 262_272)?; + if audio_token_id != audio_config.vocab_offset + 1 + || eoa_token_id != audio_config.vocab_offset + { + return Err(format!( + "Gemma3n audio token ids {audio_token_id}/{eoa_token_id} do not match \ + audio vocab offset {}", + audio_config.vocab_offset + )); + } + let configured_soft_tokens = config + .get("audio_soft_tokens_per_image") + .and_then(serde_json::Value::as_u64) + .unwrap_or(GEMMA3N_AUDIO_SOFT_TOKENS as u64); + if configured_soft_tokens != GEMMA3N_AUDIO_SOFT_TOKENS as u64 { + return Err(format!( + "Gemma3n audio soft-token count {configured_soft_tokens} is unsupported; \ + expected {GEMMA3N_AUDIO_SOFT_TOKENS}" + )); + } + let end_of_turn_token_id = ["", ""].iter().find_map(|marker| { + tokenizer + .encode(marker, false) + .ok() + .filter(|tokens| tokens.len() == 1) + .map(|tokens| tokens[0] as i32) + }); + + Ok(Some(Self { + context_capacity, + policy, + wrapper_tokens, + audio_token_id, + boa_token_id, + eoa_token_id, + end_of_turn_token_id, + })) + } + + #[must_use] + pub fn policy(&self) -> AudioFamilyPolicy { + self.policy + } + + #[must_use] + pub fn audio_token_id(&self) -> i32 { + self.audio_token_id + } + + /// Convert bounded normalized waveforms into the exact static IREE input. + pub fn prepare( + &self, + extractor: &Gemma3nAudioFeatureExtractor, + waveforms: AudioWaveformBatch, + mut token_ids: Vec, + cancelled: &AtomicBool, + ) -> Result { + if waveforms.family != self.policy.family { + return Err(format!( + "Gemma3n audio producer received {} waveforms", + waveforms.family + )); + } + let clips = waveforms + .clips + .into_iter() + .map(|clip| clip.samples) + .collect::>(); + let features = extractor.extract_batch_cancellable(&clips, cancelled)?; + if cancelled.load(Ordering::Acquire) { + return Err("Gemma3n audio feature extraction was cancelled".to_string()); + } + let input = bucket_audio_features(features)?; + crate::vlm_runtime::expand_gemma3n_audio_tokens( + &mut token_ids, + self.audio_token_id, + self.boa_token_id, + self.eoa_token_id, + clips.len(), + GEMMA3N_AUDIO_SOFT_TOKENS, + &self.wrapper_tokens, + self.end_of_turn_token_id, + ) + .map_err(|error| error.to_string())?; + if token_ids.len() > self.context_capacity { + return Err(format!( + "Gemma3n audio expanded prompt has {} tokens; context capacity is {}", + token_ids.len(), + self.context_capacity + )); + } + Ok(Gemma3nXlaPreparedAudioInput { + input, + token_ids, + clips: clips.len(), + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn bucket_preserves_processor_left_padding_features_and_zeros_only_static_tail() { + let mut features = vec![-6.25; 2 * GEMMA3N_AUDIO_MEL_BINS]; + features[GEMMA3N_AUDIO_MEL_BINS..].fill(1.5); + let input = bucket_audio_features(Gemma3nAudioFeatureBatch { + features, + valid_mask: vec![false, true], + batch_size: 1, + frames: 2, + }) + .unwrap(); + + assert_eq!(input.frame_bucket(), 8); + assert_eq!( + &input.mel()[..GEMMA3N_AUDIO_MEL_BINS], + vec![-6.25; GEMMA3N_AUDIO_MEL_BINS] + ); + assert!( + input.mel()[2 * GEMMA3N_AUDIO_MEL_BINS..] + .iter() + .all(|&value| value == 0.0) + ); + assert_eq!(&input.valid_mask()[..2], &[0, 1]); + } +} diff --git a/src/commands/generate.rs b/src/commands/generate.rs index e4fe4dd1e..78423b72f 100644 --- a/src/commands/generate.rs +++ b/src/commands/generate.rs @@ -1029,6 +1029,11 @@ fn xla_num_layers(model_dir: &Path) -> usize { .and_then(|v| { v.get("num_hidden_layers") .and_then(serde_json::Value::as_u64) + .or_else(|| { + v.get("text_config") + .and_then(|text| text.get("num_hidden_layers")) + .and_then(serde_json::Value::as_u64) + }) }) .map_or(0, |n| n as usize) } @@ -1096,7 +1101,9 @@ fn decode_xla_cli_images(paths: &[std::path::PathBuf]) -> Result Result<()> { args.generation.max_tokens, kv_cache_mode, token_bias, + &tokenizer, )?; let generated_text = decode_generated_text(&tokenizer, &prompt_tokens, &generated_tokens); let visible = filter_reasoning_for_display( diff --git a/src/lib/mlxcel-xla/Cargo.toml b/src/lib/mlxcel-xla/Cargo.toml index 3581799ad..6437801ec 100644 --- a/src/lib/mlxcel-xla/Cargo.toml +++ b/src/lib/mlxcel-xla/Cargo.toml @@ -53,3 +53,8 @@ serde_json = { version = "1" } [build-dependencies] cc = "1" + +[dev-dependencies] +# Auxiliary artifact qualification is pure Rust and must remain testable +# without enabling/linking the native IREE runtime. +sha2 = "0.10" diff --git a/src/lib/mlxcel-xla/src/aux.rs b/src/lib/mlxcel-xla/src/aux.rs index 184de8767..f196a6581 100644 --- a/src/lib/mlxcel-xla/src/aux.rs +++ b/src/lib/mlxcel-xla/src/aux.rs @@ -95,6 +95,7 @@ impl AuxiliaryTensorDType { } #[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[allow(dead_code)] pub(crate) enum AuxiliaryWeightDType { Float32, Float16, @@ -118,9 +119,31 @@ impl AuxiliaryWeightDType { } } +#[allow(dead_code)] +pub(crate) enum AuxiliaryWeightStorage { + Bytes(Vec), + Float32(Vec), +} + +impl AuxiliaryWeightStorage { + pub(crate) fn byte_len(&self) -> usize { + match self { + Self::Bytes(values) => values.len(), + Self::Float32(values) => values.len() * std::mem::size_of::(), + } + } + + fn as_ptr(&self) -> *const c_void { + match self { + Self::Bytes(values) => values.as_ptr().cast(), + Self::Float32(values) => values.as_ptr().cast(), + } + } +} + pub(crate) struct AuxiliaryWeight { pub(crate) name: String, - pub(crate) bytes: Vec, + pub(crate) storage: AuxiliaryWeightStorage, pub(crate) dtype: AuxiliaryWeightDType, pub(crate) shape: Vec, } @@ -202,10 +225,10 @@ impl IreeAuxiliaryModule { let label = format!("auxiliary weight {index}"); let dims = checked_shape(&weight.shape, weight.dtype.size_bytes(), &label)?; let expected = expected_bytes(&weight.shape, weight.dtype.size_bytes(), &label)?; - if weight.bytes.len() != expected { + if weight.storage.byte_len() != expected { return Err(format!( "{label} has {} bytes, expected {expected} for {:?} {:?}", - weight.bytes.len(), + weight.storage.byte_len(), weight.dtype, weight.shape )); @@ -216,7 +239,7 @@ impl IreeAuxiliaryModule { } let pointers: Vec<*const c_void> = weights .iter() - .map(|weight| weight.bytes.as_ptr().cast()) + .map(|weight| weight.storage.as_ptr()) .collect(); let device = CString::new(device).map_err(|_| "device URI has an interior nul byte".to_string())?; diff --git a/src/lib/mlxcel-xla/src/aux_manifest.rs b/src/lib/mlxcel-xla/src/aux_manifest.rs index 1c0dc638f..afb5974ff 100644 --- a/src/lib/mlxcel-xla/src/aux_manifest.rs +++ b/src/lib/mlxcel-xla/src/aux_manifest.rs @@ -658,7 +658,7 @@ mod tests { use std::sync::{Arc, Barrier}; use super::*; - use crate::aux::{AuxiliaryWeight, AuxiliaryWeightDType}; + use crate::aux::{AuxiliaryWeight, AuxiliaryWeightDType, AuxiliaryWeightStorage}; use crate::numeric_dtype_contract::{NumericDType, NumericDTypeContract, WeightExecution}; use crate::operator_numeric_contract::{ AffineDequantizationContract, AffineEvaluationOrder, AssociationPolicy, @@ -680,7 +680,7 @@ mod tests { fn weights() -> Vec { vec![AuxiliaryWeight { name: "weight".to_string(), - bytes: 1.0f32.to_ne_bytes().to_vec(), + storage: AuxiliaryWeightStorage::Float32(vec![1.0]), dtype: AuxiliaryWeightDType::Float32, shape: vec![1], }] diff --git a/src/lib/mlxcel-xla/src/aux_smoke.rs b/src/lib/mlxcel-xla/src/aux_smoke.rs index 31ab78dbb..21dc16210 100644 --- a/src/lib/mlxcel-xla/src/aux_smoke.rs +++ b/src/lib/mlxcel-xla/src/aux_smoke.rs @@ -19,7 +19,7 @@ use std::process::Command; use crate::aux::{ AuxiliaryInput, AuxiliaryOutput, AuxiliaryTensorDType, AuxiliaryWeight, AuxiliaryWeightDType, - IreeAuxiliaryModule, + AuxiliaryWeightStorage, IreeAuxiliaryModule, }; use crate::aux_manifest::{AuxiliaryArtifactContract, write_auxiliary_manifest}; use crate::iree::{compile_one, iree_compile_bin, target_flags}; @@ -52,7 +52,7 @@ pub struct AuxiliaryAbiSmokeReport { fn weight() -> AuxiliaryWeight { AuxiliaryWeight { name: "smoke.weight".to_string(), - bytes: f32_bytes(&[1.5, -2.0]), + storage: AuxiliaryWeightStorage::Float32(vec![1.5, -2.0]), dtype: AuxiliaryWeightDType::Float32, shape: vec![2], } diff --git a/src/lib/mlxcel-xla/src/batch.rs b/src/lib/mlxcel-xla/src/batch.rs index b08cb4b2f..3fee54410 100644 --- a/src/lib/mlxcel-xla/src/batch.rs +++ b/src/lib/mlxcel-xla/src/batch.rs @@ -48,7 +48,7 @@ use mlxcel_core::session::{ }; #[cfg(feature = "iree")] -use std::path::Path; +use std::path::{Path, PathBuf}; use crate::Gemma3nDensePle; #[cfg(feature = "iree")] @@ -444,6 +444,8 @@ where pub struct XlaBatchEngine { engine: IreeRaggedLlama, sched: Scheduler, + model_path: PathBuf, + device: String, } #[cfg(feature = "iree")] @@ -475,6 +477,8 @@ impl XlaBatchEngine { Ok(Self { engine, sched: Scheduler::new(b_max, eos), + model_path: model_path.to_path_buf(), + device: device.to_string(), }) } @@ -490,12 +494,47 @@ impl XlaBatchEngine { self.engine.context_capacity() } + /// Checkpoint directory backing this resident language bundle. + #[must_use] + pub fn model_path(&self) -> &Path { + &self.model_path + } + + /// IREE HAL device used by this resident language bundle. + #[must_use] + pub fn device(&self) -> &str { + &self.device + } + + /// Verified identity shared by every conditional Gemma3n auxiliary bundle. + #[must_use] + pub fn compatibility_fingerprint(&self) -> u64 { + self.engine.compatibility_fingerprint() + } + /// The model's EOS token ids (from `generation_config.json`). #[must_use] pub fn eos_token_ids(&self) -> &[i32] { &self.sched.eos } + /// Conditionally load a Gemma3n audio bucket against this exact verified + /// #876 batched language bundle. + pub fn load_gemma3n_audio( + &self, + frame_bucket: usize, + clips: usize, + ) -> Result { + crate::Gemma3nAudioIreeRuntime::load( + &self.model_path, + &self.device, + self.context_capacity(), + frame_bucket, + clips, + self.engine.compatibility_fingerprint(), + ) + } + /// No active slots and nothing queued: a driver loop can stop pumping. #[must_use] pub fn is_idle(&self) -> bool { diff --git a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_schema.rs b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_schema.rs index b21994933..ce13fb2a1 100644 --- a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_schema.rs +++ b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_schema.rs @@ -4,10 +4,9 @@ use std::collections::BTreeMap; use super::builder::{Builder, Ty, Val}; use super::gemma3n::Gemma3nConfig; -use crate::{ - GEMMA3N_AUDIO_CHECKPOINT_TENSOR_COUNT, GEMMA3N_AUDIO_SOFT_TOKENS, Gemma3nXlaAudioConfig, - gemma3n_audio_checkpoint_specs, -}; +#[cfg(test)] +use crate::GEMMA3N_AUDIO_CHECKPOINT_TENSOR_COUNT; +use crate::{GEMMA3N_AUDIO_SOFT_TOKENS, Gemma3nXlaAudioConfig, gemma3n_audio_checkpoint_specs}; pub(super) struct Decl { pub ty: Ty, @@ -15,7 +14,7 @@ pub(super) struct Decl { } #[derive(Clone, Debug, PartialEq, Eq)] -pub(super) struct WeightSpec { +pub(crate) struct WeightSpec { pub name: String, pub shape: Vec, pub projection: bool, @@ -61,7 +60,7 @@ fn logical_weight(name: impl Into, shape: Vec) -> WeightSpec { } } -fn encoder_weight_specs( +pub(crate) fn encoder_weight_specs( audio: &Gemma3nXlaAudioConfig, text: &Gemma3nConfig, ) -> Result, String> { @@ -92,7 +91,7 @@ fn encoder_weight_specs( Ok(specs) } -fn merge_weight_specs(text: &Gemma3nConfig) -> Vec { +pub(crate) fn merge_weight_specs(text: &Gemma3nConfig) -> Vec { vec![ logical_weight( "model.language_model.embed_tokens.weight", diff --git a/src/lib/mlxcel-xla/src/emitter/gemma3n_schema.rs b/src/lib/mlxcel-xla/src/emitter/gemma3n_schema.rs index d12c30a54..bc7b2db3a 100644 --- a/src/lib/mlxcel-xla/src/emitter/gemma3n_schema.rs +++ b/src/lib/mlxcel-xla/src/emitter/gemma3n_schema.rs @@ -365,3 +365,84 @@ pub(super) fn build_schema( fn norm(decls: &mut Vec, index: &mut usize, width: usize, name: String) -> Val { take(decls, index, Ty::f32(vec![width]), name) } + +#[cfg(test)] +mod tests { + use super::*; + use crate::emitter::{Gemma3nWeightSpec, gemma3n_weight_specs}; + #[cfg(xla_iree_cuda)] + use crate::emitter::{Precision, emit_gemma3n_prefill_with_qmv}; + + fn config() -> Gemma3nConfig { + Gemma3nConfig::from_json_str( + &serde_json::json!({ + "model_type": "gemma3n_text", + "hidden_size": 128, + "intermediate_size": [192, 192, 192, 192], + "max_position_embeddings": 4096, + "num_hidden_layers": 4, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "head_dim": 64, + "rms_norm_eps": 1e-6, + "vocab_size": 128, + "vocab_size_per_layer_input": 128, + "hidden_size_per_layer_input": 64, + "layer_types": [ + "sliding_attention", + "full_attention", + "sliding_attention", + "full_attention" + ], + "activation_sparsity_pattern": [0.5, 0.0, 0.0, 0.0], + "sliding_window": 2, + "rope_theta": 1000000.0, + "rope_local_base_freq": 10000.0, + "final_logit_softcapping": 30.0, + "num_kv_shared_layers": 2, + "altup_num_inputs": 4, + "altup_active_idx": 0, + "altup_coef_clip": 120.0, + "altup_correct_scale": true, + "laurel_rank": 64, + "tie_word_embeddings": true, + "quantization": {"bits": 4, "group_size": 64} + }) + .to_string(), + ) + .unwrap() + .with_context_capacity(4) + .unwrap() + } + + #[test] + fn native_qmv_keeps_one_f32_carrier_argument_per_loader_weight() { + let config = config(); + let loader_specs = gemma3n_weight_specs(&config); + let (declarations, _) = build_schema(&config, false, 4, false, false); + let graph_weights = &declarations[..loader_specs.len()]; + + assert_eq!(graph_weights.len(), loader_specs.len()); + for (declaration, spec) in graph_weights.iter().zip(&loader_specs) { + assert_eq!(declaration.loc, spec.name()); + assert_eq!(declaration.ty.elt, "f32"); + if matches!(spec, Gemma3nWeightSpec::Projection(_)) { + assert_eq!(declaration.ty.shape.len(), 2); + } + } + + #[cfg(xla_iree_cuda)] + { + let graph = emit_gemma3n_prefill_with_qmv(&config, false, Precision::Bf16, true); + assert_eq!( + graph.matches("loc(\"model.language_model").count(), + loader_specs.len() + ); + assert!(!graph.contains(".scales\")")); + assert!(!graph.contains(".biases\")")); + assert!(graph.contains( + "(i32, i32, i32, tensor<4x128xf32>, tensor<192x128xf32>) -> tensor<4x192xf32>" + )); + } + } +} diff --git a/src/lib/mlxcel-xla/src/emitter/mod.rs b/src/lib/mlxcel-xla/src/emitter/mod.rs index 685aa6920..8f56cb9e7 100644 --- a/src/lib/mlxcel-xla/src/emitter/mod.rs +++ b/src/lib/mlxcel-xla/src/emitter/mod.rs @@ -88,6 +88,12 @@ pub use gemma3n_audio_emit::{Gemma3nAudioDiagnosticLayout, Gemma3nAudioDiagnosti #[allow(unused_imports)] pub(crate) use gemma3n_audio_emit::{emit_gemma3n_audio_encode, emit_gemma3n_audio_merge_ple}; #[allow(unused_imports)] +pub(crate) use gemma3n_audio_schema::{ + WeightSpec as Gemma3nAudioGraphWeightSpec, + encoder_weight_specs as gemma3n_audio_encoder_weights, + merge_weight_specs as gemma3n_audio_merge_weights, +}; +#[allow(unused_imports)] pub(crate) use gemma3n_decode::{ emit_gemma3n_decode, emit_gemma3n_decode_ragged, emit_gemma3n_decode_ragged_with, emit_gemma3n_decode_ragged_with_qmv, emit_gemma3n_decode_with, emit_gemma3n_decode_with_qmv, diff --git a/src/lib/mlxcel-xla/src/gemma3n_audio_runtime.rs b/src/lib/mlxcel-xla/src/gemma3n_audio_runtime.rs new file mode 100644 index 000000000..ef22cd4e2 --- /dev/null +++ b/src/lib/mlxcel-xla/src/gemma3n_audio_runtime.rs @@ -0,0 +1,786 @@ +//! Conditional two-artifact IREE runtime for Gemma3n audio. +//! +//! The encoder and language merge are separate resident-weight modules. This +//! keeps the compiler at natural tensor boundaries while ensuring every model +//! operation after host mel extraction still executes through IREE. + +use std::fs::File; +use std::path::{Path, PathBuf}; +use std::process::Command; + +use memmap2::Mmap; +use safetensors::{Dtype, SafeTensors}; +use sha2::{Digest, Sha256}; + +use crate::aux::{ + AuxiliaryInput, AuxiliaryOutput, AuxiliaryTensorDType, AuxiliaryWeight, AuxiliaryWeightDType, + AuxiliaryWeightStorage, IreeAuxiliaryModule, +}; +use crate::aux_manifest::{AuxiliaryArtifactContract, ensure_qualified_auxiliary_artifact}; +use crate::emitter::{ + Gemma3nAudioGraphWeightSpec, Gemma3nConfig, Precision, emit_gemma3n_audio_encode, + emit_gemma3n_audio_merge_ple, gemma3n_audio_encoder_weights, gemma3n_audio_merge_weights, +}; +use crate::iree::{cached_vmfb_path, compile_one_to, iree_compile_bin, target_flags}; +use crate::weights::{bf16_to_f32, dequantize_affine_bf16_fused, f16_to_f32, f32_le_to_f32}; +use crate::{ + GEMMA3N_AUDIO_MODALITY_FAMILY, GEMMA3N_AUDIO_SOFT_TOKENS, Gemma3nAudioInput, + Gemma3nAudioPreparedPrefill, Gemma3nDensePle, Gemma3nPreparedPrefill, Gemma3nXlaAudioConfig, + validate_gemma3n_audio_checkpoint, validate_gemma3n_audio_row_indices, +}; +use mlxcel_core::session::{ + OwnedTensor, PreparedAttentionBias, PreparedModality, PreparedPositions, PreparedPrefill, + PreparedTensorDType, +}; + +#[derive(Debug, Clone, PartialEq)] +pub struct Gemma3nAudioGraphOutput { + embeddings: Vec, + dense_ple: Vec, + projected_lengths: Vec, + context_capacity: usize, + hidden_size: usize, + layers: usize, + hidden_per_layer: usize, +} + +impl Gemma3nAudioGraphOutput { + #[must_use] + pub fn embeddings(&self) -> &[f32] { + &self.embeddings + } + + #[must_use] + pub fn dense_ple(&self) -> &[f32] { + &self.dense_ple + } + + #[must_use] + pub fn projected_lengths(&self) -> &[usize] { + &self.projected_lengths + } + + #[must_use] + pub fn embeddings_shape(&self) -> [usize; 2] { + [self.context_capacity, self.hidden_size] + } + + #[must_use] + pub fn dense_ple_shape(&self) -> [usize; 3] { + [self.context_capacity, self.layers, self.hidden_per_layer] + } + + pub fn into_parts(self) -> (Vec, Vec, Vec) { + (self.embeddings, self.dense_ple, self.projected_lengths) + } + + /// Move a verified split-graph result into the existing #876 prepared + /// embeddings-plus-dense-PLE contract. + pub fn into_prepared( + self, + token_ids: Vec, + audio_token_id: i32, + frame_bucket: usize, + ) -> Result { + let sequence_len = token_ids.len(); + if sequence_len == 0 || sequence_len > self.context_capacity { + return Err(format!( + "Gemma3n audio prepared length {sequence_len} is outside 1..={}", + self.context_capacity + )); + } + let embedding_elements = sequence_len + .checked_mul(self.hidden_size) + .ok_or_else(|| "Gemma3n audio embedding length overflows".to_string())?; + let embeddings = OwnedTensor::new( + self.embeddings[..embedding_elements] + .iter() + .flat_map(|value| value.to_le_bytes()) + .collect(), + PreparedTensorDType::Float32, + vec![1, sequence_len, self.hidden_size], + ) + .map_err(|error| error.to_string())?; + let attention_bias = OwnedTensor::new( + vec![0; sequence_len * std::mem::size_of::()], + PreparedTensorDType::Float32, + vec![1, 1, 1, sequence_len], + ) + .map_err(|error| error.to_string())?; + let clips = self.projected_lengths.len(); + let prepared = PreparedPrefill::new( + token_ids, + embeddings, + PreparedPositions::Sequential { + start: 0, + length: sequence_len, + }, + PreparedAttentionBias { + tensor: attention_bias, + causal: true, + }, + vec![PreparedModality { + family: GEMMA3N_AUDIO_MODALITY_FAMILY.to_string(), + item_count: clips, + token_count: clips * GEMMA3N_AUDIO_SOFT_TOKENS, + }], + ) + .map_err(|error| error.to_string())?; + let dense_ple = Gemma3nDensePle::new( + self.dense_ple, + self.context_capacity, + self.layers, + self.hidden_per_layer, + ) + .map_err(|error| error.to_string())?; + let request = + Gemma3nPreparedPrefill::new(prepared, dense_ple).map_err(|error| error.to_string())?; + Gemma3nAudioPreparedPrefill::new( + request, + audio_token_id, + self.projected_lengths, + frame_bucket, + ) + } +} + +pub struct Gemma3nAudioIreeRuntime { + encode: IreeAuxiliaryModule, + merge_ple: IreeAuxiliaryModule, + audio: Gemma3nXlaAudioConfig, + frame_bucket: usize, + clips: usize, + context_capacity: usize, + hidden_size: usize, + layers: usize, + hidden_per_layer: usize, + language_fingerprint: u64, +} + +fn sha256_hex(bytes: &[u8]) -> String { + let digest = Sha256::digest(bytes); + let mut result = String::with_capacity(digest.len() * 2); + for byte in digest { + use std::fmt::Write; + write!(result, "{byte:02x}").expect("writing to String cannot fail"); + } + result +} + +fn generation_identity(compiler: &Path, flags: &[&str], mlir: &str) -> Result { + let version = Command::new(compiler) + .arg("--version") + .output() + .map_err(|error| format!("run {} --version: {error}", compiler.display()))?; + if !version.status.success() { + return Err(format!( + "{} --version failed: {}", + compiler.display(), + String::from_utf8_lossy(&version.stderr) + )); + } + Ok(format!( + "compiler={};version={};flags={flags:?};mlir-sha256={}", + compiler.display(), + String::from_utf8_lossy(&version.stdout).trim(), + sha256_hex(mlir.as_bytes()) + )) +} + +fn checkpoint_alias(name: &str) -> Option { + name.strip_prefix("model.language_model.") + .map(|suffix| format!("language_model.model.{suffix}")) +} + +fn safetensor_files(model_dir: &Path) -> Result, String> { + let mut files = std::fs::read_dir(model_dir) + .map_err(|error| format!("read {}: {error}", model_dir.display()))? + .filter_map(Result::ok) + .map(|entry| entry.path()) + .filter(|path| { + path.is_file() + && path + .file_name() + .and_then(|name| name.to_str()) + .is_some_and(|name| name.ends_with(".safetensors")) + }) + .collect::>(); + files.sort(); + if files.is_empty() { + return Err(format!( + "{} contains no safetensors files", + model_dir.display() + )); + } + Ok(files) +} + +fn decode_weight( + tensors: &SafeTensors<'_>, + checkpoint_name: &str, + spec: &Gemma3nAudioGraphWeightSpec, + quant_bits: usize, + quant_group: usize, +) -> Result, String> { + let tensor = tensors + .tensor(checkpoint_name) + .map_err(|error| format!("read {checkpoint_name}: {error}"))?; + let (values, shape) = match tensor.dtype() { + Dtype::BF16 => (bf16_to_f32(tensor.data()), tensor.shape().to_vec()), + Dtype::F16 => (f16_to_f32(tensor.data()), tensor.shape().to_vec()), + Dtype::F32 => (f32_le_to_f32(tensor.data()), tensor.shape().to_vec()), + Dtype::U32 => { + let prefix = checkpoint_name + .strip_suffix(".weight") + .ok_or_else(|| format!("quantized tensor {checkpoint_name} is not a weight"))?; + let scales_name = format!("{prefix}.scales"); + let biases_name = format!("{prefix}.biases"); + let scales = tensors + .tensor(&scales_name) + .map_err(|error| format!("read {scales_name}: {error}"))?; + let biases = tensors + .tensor(&biases_name) + .map_err(|error| format!("read {biases_name}: {error}"))?; + if scales.dtype() != Dtype::BF16 || biases.dtype() != Dtype::BF16 { + return Err(format!( + "{prefix} affine metadata must be BF16, got {:?}/{:?}", + scales.dtype(), + biases.dtype() + )); + } + let packed = tensor.shape(); + if packed.len() != 2 { + return Err(format!( + "quantized tensor {checkpoint_name} rank {} is not 2", + packed.len() + )); + } + let logical_columns = packed[1] + .checked_mul(32 / quant_bits) + .ok_or_else(|| format!("{checkpoint_name} logical width overflows"))?; + let values = dequantize_affine_bf16_fused( + tensor.data(), + scales.data(), + biases.data(), + packed[0], + packed[1], + quant_bits, + quant_group, + ) + .map_err(|error| format!("dequantize {checkpoint_name}: {error}"))?; + (values, vec![packed[0], logical_columns]) + } + dtype => { + return Err(format!( + "weight {checkpoint_name} dtype {dtype:?} is not BF16/F16/F32/U32" + )); + } + }; + if shape != spec.shape { + return Err(format!( + "weight {} has logical shape {shape:?}, expected {:?}", + spec.name, spec.shape + )); + } + Ok(values) +} + +fn load_weights( + model_dir: &Path, + specs: &[Gemma3nAudioGraphWeightSpec], + text: &Gemma3nConfig, +) -> Result, String> { + let quant = text + .quantization + .ok_or_else(|| "Gemma3n audio runtime requires the pinned Q4 config".to_string())?; + let mut loaded = (0..specs.len()).map(|_| None).collect::>(); + for path in safetensor_files(model_dir)? { + let file = + File::open(&path).map_err(|error| format!("open {}: {error}", path.display()))?; + // Safety: the file remains open and the mapping is read-only for this scope. + let mapping = unsafe { Mmap::map(&file) } + .map_err(|error| format!("mmap {}: {error}", path.display()))?; + let tensors = SafeTensors::deserialize(&mapping) + .map_err(|error| format!("parse {}: {error}", path.display()))?; + for (index, spec) in specs.iter().enumerate() { + let canonical = tensors.tensor(&spec.name).is_ok(); + let alias = checkpoint_alias(&spec.name); + let repackaged = alias + .as_deref() + .is_some_and(|name| tensors.tensor(name).is_ok()); + let checkpoint_name = match (canonical, repackaged) { + (false, false) => continue, + (true, false) => spec.name.clone(), + (false, true) => alias.expect("repackaged presence requires an alias"), + (true, true) => { + return Err(format!( + "{} contains both canonical and repackaged {}", + path.display(), + spec.name + )); + } + }; + if loaded[index].is_some() { + return Err(format!( + "Gemma3n auxiliary weight {} occurs in multiple shards", + spec.name + )); + } + let values = decode_weight( + &tensors, + &checkpoint_name, + spec, + quant.bits, + quant.group_size, + )?; + loaded[index] = Some(AuxiliaryWeight { + name: spec.name.clone(), + storage: AuxiliaryWeightStorage::Float32(values), + dtype: AuxiliaryWeightDType::Float32, + shape: spec.shape.clone(), + }); + } + } + loaded + .into_iter() + .zip(specs) + .map(|(weight, spec)| { + weight.ok_or_else(|| format!("Gemma3n auxiliary weight {} is missing", spec.name)) + }) + .collect() +} + +fn f32_bytes(values: &[f32]) -> &[u8] { + // Safety: f32 is plain data and the returned view cannot outlive `values`. + unsafe { std::slice::from_raw_parts(values.as_ptr().cast(), std::mem::size_of_val(values)) } +} + +fn f32_bytes_mut(values: &mut [f32]) -> &mut [u8] { + // Safety: f32 is plain data, the byte view covers exactly the same + // allocation, and the exclusive borrow prevents typed access until return. + unsafe { + std::slice::from_raw_parts_mut(values.as_mut_ptr().cast(), std::mem::size_of_val(values)) + } +} + +fn i32_bytes(values: &[i32]) -> &[u8] { + // Safety: i32 is plain data and the returned view cannot outlive `values`. + unsafe { std::slice::from_raw_parts(values.as_ptr().cast(), std::mem::size_of_val(values)) } +} + +fn i32_bytes_mut(values: &mut [i32]) -> &mut [u8] { + // Safety: see `f32_bytes_mut`; the same invariant holds for i32. + unsafe { + std::slice::from_raw_parts_mut(values.as_mut_ptr().cast(), std::mem::size_of_val(values)) + } +} + +fn load_timing_enabled() -> bool { + std::env::var_os("MLXCEL_XLA_LOAD_TIMING").is_some() +} + +fn current_rss_kib() -> Option { + std::fs::read_to_string("/proc/self/status") + .ok()? + .lines() + .find_map(|line| line.strip_prefix("VmRSS:")) + .and_then(|value| value.split_whitespace().next()) + .and_then(|value| value.parse().ok()) +} + +fn report_load_timing(label: &str, started: std::time::Instant) { + if load_timing_enabled() { + eprintln!( + "mlxcel-xla-load: phase={label} elapsed_ms={} rss_kib={}", + started.elapsed().as_millis(), + current_rss_kib() + .map(|value| value.to_string()) + .unwrap_or_else(|| "unavailable".to_string()) + ); + } +} + +fn report_auxiliary_weights(label: &str, started: std::time::Instant, weights: &[AuxiliaryWeight]) { + if load_timing_enabled() { + let f32_bytes = weights + .iter() + .filter(|weight| matches!(&weight.storage, AuxiliaryWeightStorage::Float32(_))) + .map(|weight| weight.storage.byte_len()) + .sum::(); + let raw_bytes = weights + .iter() + .filter(|weight| matches!(&weight.storage, AuxiliaryWeightStorage::Bytes(_))) + .map(|weight| weight.storage.byte_len()) + .sum::(); + eprintln!( + "mlxcel-xla-load: phase={label} elapsed_ms={} rss_kib={} \ + buffers={} f32_bytes={f32_bytes} raw_bytes={raw_bytes}", + started.elapsed().as_millis(), + current_rss_kib() + .map(|value| value.to_string()) + .unwrap_or_else(|| "unavailable".to_string()), + weights.len() + ); + } +} + +impl Gemma3nAudioIreeRuntime { + #[allow(clippy::too_many_arguments)] + pub fn load( + model_dir: &Path, + device: &str, + context_capacity: usize, + frame_bucket: usize, + clips: usize, + language_fingerprint: u64, + ) -> Result { + let load_started = std::time::Instant::now(); + if language_fingerprint == 0 { + return Err("Gemma3n audio requires a verified #876 language bundle".to_string()); + } + let text = Gemma3nConfig::from_json(model_dir)?.with_context_capacity(context_capacity)?; + let audio = Gemma3nXlaAudioConfig::from_model_dir(model_dir)? + .ok_or_else(|| "checkpoint has no compatible Gemma3n audio config".to_string())?; + validate_gemma3n_audio_checkpoint(model_dir, &audio, text.hidden) + .map_err(|error| error.to_string())?; + report_load_timing("audio-config-and-schema-validation", load_started); + let emit_started = std::time::Instant::now(); + let (encode_mlir, _) = + emit_gemma3n_audio_encode(&text, &audio, frame_bucket, clips, Precision::Bf16)?; + let (merge_mlir, _) = emit_gemma3n_audio_merge_ple(&text, &audio, clips, Precision::Bf16)?; + report_load_timing("audio-stablehlo-emission", emit_started); + let compiler = iree_compile_bin()?; + let mut flags = target_flags(device)?.to_vec(); + if device == "cuda" { + flags.push("--iree-cuda-target=sm_80"); + } + let cache = std::env::temp_dir().join("mlxcel-xla-gemma3n-audio"); + std::fs::create_dir_all(&cache) + .map_err(|error| format!("mkdir {}: {error}", cache.display()))?; + let encode_tag = format!("audio-encode-f{frame_bucket}-b{clips}"); + let merge_tag = format!("audio-merge-ple-c{context_capacity}-b{clips}"); + let encode_vmfb = cached_vmfb_path( + &compiler, + &encode_mlir, + &flags, + &cache, + &encode_tag, + context_capacity, + ); + let merge_vmfb = cached_vmfb_path( + &compiler, + &merge_mlir, + &flags, + &cache, + &merge_tag, + context_capacity, + ); + let identity = audio.artifact_identity( + frame_bucket, + clips, + text.hidden, + text.n_layers, + text.hidden_per_layer_input, + )?; + let encode_contract = AuxiliaryArtifactContract::new( + "audio_encode.main", + format!("encode:{identity}:language={language_fingerprint:016x}"), + generation_identity(&compiler, &flags, &encode_mlir)?, + )?; + let merge_contract = AuxiliaryArtifactContract::new( + "audio_merge_ple.main", + format!( + "merge:{identity}:context={context_capacity}:text={}:language={language_fingerprint:016x}", + text.compatibility_fingerprint() + ), + generation_identity(&compiler, &flags, &merge_mlir)?, + )?; + let encode_weights_started = std::time::Instant::now(); + let encode_weights = load_weights( + model_dir, + &gemma3n_audio_encoder_weights(&audio, &text)?, + &text, + )?; + report_auxiliary_weights( + "audio-encoder-weights", + encode_weights_started, + &encode_weights, + ); + let encode_compile_started = std::time::Instant::now(); + ensure_qualified_auxiliary_artifact( + &encode_vmfb, + &encode_contract, + &encode_weights, + |temporary| { + compile_one_to( + &compiler, + &encode_mlir, + &flags, + &cache, + &encode_tag, + context_capacity, + temporary, + ) + }, + )?; + report_load_timing("audio-encoder-vmfb-qualified", encode_compile_started); + let encode_module_started = std::time::Instant::now(); + let encode = + IreeAuxiliaryModule::load(device, &encode_vmfb, &encode_contract, encode_weights)?; + report_load_timing("audio-encoder-vmfb-loaded", encode_module_started); + let merge_weights_started = std::time::Instant::now(); + let merge_weights = load_weights(model_dir, &gemma3n_audio_merge_weights(&text), &text)?; + report_auxiliary_weights("audio-merge-weights", merge_weights_started, &merge_weights); + let merge_compile_started = std::time::Instant::now(); + ensure_qualified_auxiliary_artifact( + &merge_vmfb, + &merge_contract, + &merge_weights, + |temporary| { + compile_one_to( + &compiler, + &merge_mlir, + &flags, + &cache, + &merge_tag, + context_capacity, + temporary, + ) + }, + )?; + report_load_timing("audio-merge-vmfb-qualified", merge_compile_started); + let merge_module_started = std::time::Instant::now(); + let merge_ple = + IreeAuxiliaryModule::load(device, &merge_vmfb, &merge_contract, merge_weights)?; + report_load_timing("audio-merge-vmfb-loaded", merge_module_started); + report_load_timing("audio-total", load_started); + Ok(Self { + encode, + merge_ple, + audio, + frame_bucket, + clips, + context_capacity, + hidden_size: text.hidden, + layers: text.n_layers, + hidden_per_layer: text.hidden_per_layer_input, + language_fingerprint, + }) + } + + #[must_use] + pub fn has_capability(&self) -> bool { + self.language_fingerprint != 0 + && self.encode.fingerprint() != 0 + && self.merge_ple.fingerprint() != 0 + } + + pub fn invoke( + &mut self, + input: &Gemma3nAudioInput, + token_ids: &[i32], + audio_row_indices: &[i32], + ) -> Result { + let invoke_started = std::time::Instant::now(); + if input.frame_bucket() != self.frame_bucket || input.clips() != self.clips { + return Err(format!( + "Gemma3n audio runtime expects bucket/clips {}/{}, got {}/{}", + self.frame_bucket, + self.clips, + input.frame_bucket(), + input.clips() + )); + } + if token_ids.is_empty() || token_ids.len() > self.context_capacity { + return Err(format!( + "Gemma3n audio expanded length {} is outside 1..={}", + token_ids.len(), + self.context_capacity + )); + } + validate_gemma3n_audio_row_indices( + token_ids, + self.audio.vocab_offset + 1, + self.clips, + audio_row_indices, + ) + .map_err(|error| error.to_string())?; + + let projected_rows = self.clips * GEMMA3N_AUDIO_SOFT_TOKENS; + let mut projected = vec![0.0f32; projected_rows * self.hidden_size]; + let mut hard = vec![0.0f32; self.audio.vocab_size * self.hidden_size]; + let mut lengths = vec![0i32; self.clips]; + let encode_started = std::time::Instant::now(); + self.encode.invoke( + &[ + AuxiliaryInput { + bytes: f32_bytes(input.mel()), + dtype: AuxiliaryTensorDType::Float32, + shape: &[self.clips, self.frame_bucket, self.audio.input_feat_size], + }, + AuxiliaryInput { + bytes: input.valid_mask(), + dtype: AuxiliaryTensorDType::Bool, + shape: &[self.clips, self.frame_bucket], + }, + ], + &mut [ + AuxiliaryOutput { + bytes: f32_bytes_mut(&mut projected), + dtype: AuxiliaryTensorDType::Float32, + shape: &[projected_rows, self.hidden_size], + }, + AuxiliaryOutput { + bytes: f32_bytes_mut(&mut hard), + dtype: AuxiliaryTensorDType::Float32, + shape: &[self.audio.vocab_size, self.hidden_size], + }, + AuxiliaryOutput { + bytes: i32_bytes_mut(&mut lengths), + dtype: AuxiliaryTensorDType::Int32, + shape: &[self.clips], + }, + ], + )?; + report_load_timing("audio-encoder-invoke", encode_started); + let projected_lengths = lengths + .into_iter() + .enumerate() + .map(|(clip, length)| { + usize::try_from(length) + .ok() + .filter(|length| (1..=GEMMA3N_AUDIO_SOFT_TOKENS).contains(length)) + .ok_or_else(|| { + format!("audio.encode returned invalid length {length} for clip {clip}") + }) + }) + .collect::, _>>()?; + + let real_len = i32::try_from(token_ids.len()) + .map_err(|_| "Gemma3n audio expanded length does not fit i32".to_string())?; + let mut padded_tokens = token_ids.to_vec(); + padded_tokens.resize(self.context_capacity, 0); + let mut padded_rows = audio_row_indices.to_vec(); + padded_rows.resize(self.context_capacity, -1); + let mut embeddings = vec![0.0f32; self.context_capacity * self.hidden_size]; + let mut dense_ple = + vec![0.0f32; self.context_capacity * self.layers * self.hidden_per_layer]; + let merge_started = std::time::Instant::now(); + self.merge_ple.invoke( + &[ + AuxiliaryInput { + bytes: f32_bytes(&projected), + dtype: AuxiliaryTensorDType::Float32, + shape: &[projected_rows, self.hidden_size], + }, + AuxiliaryInput { + bytes: f32_bytes(&hard), + dtype: AuxiliaryTensorDType::Float32, + shape: &[self.audio.vocab_size, self.hidden_size], + }, + AuxiliaryInput { + bytes: i32_bytes(&padded_tokens), + dtype: AuxiliaryTensorDType::Int32, + shape: &[self.context_capacity], + }, + AuxiliaryInput { + bytes: i32_bytes(&padded_rows), + dtype: AuxiliaryTensorDType::Int32, + shape: &[self.context_capacity], + }, + AuxiliaryInput { + bytes: i32_bytes(std::slice::from_ref(&real_len)), + dtype: AuxiliaryTensorDType::Int32, + shape: &[], + }, + ], + &mut [ + AuxiliaryOutput { + bytes: f32_bytes_mut(&mut embeddings), + dtype: AuxiliaryTensorDType::Float32, + shape: &[self.context_capacity, self.hidden_size], + }, + AuxiliaryOutput { + bytes: f32_bytes_mut(&mut dense_ple), + dtype: AuxiliaryTensorDType::Float32, + shape: &[self.context_capacity, self.layers, self.hidden_per_layer], + }, + ], + )?; + report_load_timing("audio-merge-invoke", merge_started); + report_load_timing("audio-invoke-total", invoke_started); + Ok(Gemma3nAudioGraphOutput { + embeddings, + dense_ple, + projected_lengths, + context_capacity: self.context_capacity, + hidden_size: self.hidden_size, + layers: self.layers, + hidden_per_layer: self.hidden_per_layer, + }) + } + + /// Run both split audio modules and move their output directly into the + /// request-scoped #876 prefill owner. + pub fn invoke_prepared( + &mut self, + input: &Gemma3nAudioInput, + token_ids: Vec, + audio_token_id: i32, + ) -> Result { + let audio_row_indices = crate::gemma3n_audio_rows::build_gemma3n_audio_row_indices( + &token_ids, + audio_token_id, + input.clips(), + ) + .map_err(|error| error.to_string())?; + let output = self.invoke(input, &token_ids, &audio_row_indices)?; + output.into_prepared(token_ids, audio_token_id, input.frame_bucket()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn output_moves_into_request_scoped_dense_ple_owner() { + let context_capacity = GEMMA3N_AUDIO_SOFT_TOKENS + 2; + let mut token_ids = vec![7]; + token_ids.extend(std::iter::repeat_n(262_273, GEMMA3N_AUDIO_SOFT_TOKENS)); + token_ids.push(9); + let output = Gemma3nAudioGraphOutput { + embeddings: vec![0.25; context_capacity * 2], + dense_ple: vec![0.5; context_capacity * 2], + projected_lengths: vec![GEMMA3N_AUDIO_SOFT_TOKENS], + context_capacity, + hidden_size: 2, + layers: 1, + hidden_per_layer: 2, + }; + let prepared = output + .into_prepared(token_ids, 262_273, 2_048) + .expect("valid audio prepared prefill"); + assert_eq!(prepared.projected_lengths(), [GEMMA3N_AUDIO_SOFT_TOKENS]); + assert_eq!(prepared.placeholder_starts(), [1]); + assert_eq!( + prepared.request().dense_ple().shape(), + [context_capacity, 1, 2] + ); + } + + #[test] + fn zero_language_bundle_identity_fails_before_checkpoint_access() { + let error = Gemma3nAudioIreeRuntime::load( + Path::new("/path/that/must/not/be/read"), + "local-task", + 256, + 2_048, + 1, + 0, + ) + .err() + .expect("zero identity must fail"); + assert!(error.contains("verified #876 language bundle")); + } +} diff --git a/src/lib/mlxcel-xla/src/iree.rs b/src/lib/mlxcel-xla/src/iree.rs index e31100190..cb20f8703 100644 --- a/src/lib/mlxcel-xla/src/iree.rs +++ b/src/lib/mlxcel-xla/src/iree.rs @@ -194,6 +194,73 @@ impl WeightBuf { WeightBuf::Raw(v) => v.as_ptr(), } } + + fn byte_len(&self) -> usize { + match self { + Self::F32(values) => std::mem::size_of_val(values.as_slice()), + Self::F16(values) => std::mem::size_of_val(values.as_slice()), + Self::Raw(values) => values.len(), + } + } +} + +fn load_timing_enabled() -> bool { + std::env::var_os("MLXCEL_XLA_LOAD_TIMING").is_some() +} + +fn current_rss_kib() -> Option { + std::fs::read_to_string("/proc/self/status") + .ok()? + .lines() + .find_map(|line| line.strip_prefix("VmRSS:")) + .and_then(|value| value.split_whitespace().next()) + .and_then(|value| value.parse().ok()) +} + +fn report_load_timing(label: &str, started: std::time::Instant) { + if load_timing_enabled() { + eprintln!( + "mlxcel-xla-load: phase={label} elapsed_ms={} rss_kib={}", + started.elapsed().as_millis(), + current_rss_kib() + .map(|value| value.to_string()) + .unwrap_or_else(|| "unavailable".to_string()) + ); + } +} + +fn report_language_weight_progress( + buffers: &[WeightBuf], + loaded: usize, + total: usize, + started: std::time::Instant, +) { + if !load_timing_enabled() || (!loaded.is_multiple_of(32) && loaded != total) { + return; + } + let f32_bytes = buffers + .iter() + .filter(|buffer| matches!(buffer, WeightBuf::F32(_))) + .map(WeightBuf::byte_len) + .sum::(); + let f16_bytes = buffers + .iter() + .filter(|buffer| matches!(buffer, WeightBuf::F16(_))) + .map(WeightBuf::byte_len) + .sum::(); + let raw_bytes = buffers + .iter() + .filter(|buffer| matches!(buffer, WeightBuf::Raw(_))) + .map(WeightBuf::byte_len) + .sum::(); + eprintln!( + "mlxcel-xla-load: phase=language-weight-progress elapsed_ms={} rss_kib={} \ + loaded={loaded}/{total} f32_bytes={f32_bytes} f16_bytes={f16_bytes} raw_bytes={raw_bytes}", + started.elapsed().as_millis(), + current_rss_kib() + .map(|value| value.to_string()) + .unwrap_or_else(|| "unavailable".to_string()) + ); } /// Opaque handle to the C-side execution context. @@ -1390,6 +1457,7 @@ fn load_weights( cfg: &RuntimeConfig, resident_f16: bool, ) -> Result<(Vec, Vec, Vec, Vec), String> { + let load_started = std::time::Instant::now(); let specs = cfg.weight_specs(); let names: Vec = specs.iter().map(|s| s.tensor_name().to_string()).collect(); let (shard_paths, mut dense_name_scheme, gemma3n_name_scheme) = match cfg { @@ -1418,6 +1486,7 @@ fn load_weights( let mut dtypes: Vec = vec![WDT_F32; n]; let mut ranks: Vec = vec![0; n]; let mut dims: Vec = vec![0; n * 4]; + let mut loaded_count = 0usize; for (shard, idxs) in by_shard { let file = File::open(shard).map_err(|e| format!("open {}: {e}", shard.display()))?; // Safety: the file is read-only for the lifetime of the mmap below. @@ -1499,6 +1568,8 @@ fn load_weights( dims[i * 4] = checked_ffi_i64(shape[0], &format!("weight {name} dim 0"))?; dims[i * 4 + 1] = checked_ffi_i64(shape[1], &format!("weight {name} dim 1"))?; bufs[i] = WeightBuf::Raw(t.data().to_vec()); + loaded_count += 1; + report_language_weight_progress(&bufs, loaded_count, n, load_started); continue; } @@ -1662,6 +1733,8 @@ fn load_weights( unreachable!("QuantRaw is handled by the early-continue above") } } + loaded_count += 1; + report_language_weight_progress(&bufs, loaded_count, n, load_started); } } Ok((bufs, dtypes, ranks, dims)) @@ -1734,11 +1807,38 @@ fn create_ctx_with_diagnostics( bundle: &CompiledBundle, prefill_diagnostics_vmfb: Option<&Path>, ) -> Result<*mut XlaCtx, String> { + let load_started = std::time::Instant::now(); // issue #572: the f16 GPU path uploads the projection weights f16-resident to // match the emitter's f16 args. Uses the same resolve_precision(device) the graph // emit uses, so the uploaded buffer dtype always lines up with the emitted arg. let resident_f16 = resolve_precision(device) == Precision::F16 && cfg.supports_f16_resident(); let (bufs, dtypes, ranks, dims) = load_weights(model_dir, cfg, resident_f16)?; + if load_timing_enabled() { + let f32_bytes = bufs + .iter() + .filter(|buffer| matches!(buffer, WeightBuf::F32(_))) + .map(WeightBuf::byte_len) + .sum::(); + let f16_bytes = bufs + .iter() + .filter(|buffer| matches!(buffer, WeightBuf::F16(_))) + .map(WeightBuf::byte_len) + .sum::(); + let raw_bytes = bufs + .iter() + .filter(|buffer| matches!(buffer, WeightBuf::Raw(_))) + .map(WeightBuf::byte_len) + .sum::(); + eprintln!( + "mlxcel-xla-load: phase=language-weights-loaded elapsed_ms={} rss_kib={} \ + buffers={} f32_bytes={f32_bytes} f16_bytes={f16_bytes} raw_bytes={raw_bytes}", + load_started.elapsed().as_millis(), + current_rss_kib() + .map(|value| value.to_string()) + .unwrap_or_else(|| "unavailable".to_string()), + bufs.len() + ); + } let ffi = runtime_ffi_dimensions(cfg, bufs.len())?; let ptrs: Vec<*const c_void> = bufs .iter() @@ -1802,6 +1902,7 @@ fn create_ctx_with_diagnostics( i32::from(cfg.has_adapter_modes()), ) }; + report_load_timing("language-vmfb-and-device-upload", load_started); // Weights are resident on the device now; free the host copy. drop(ptrs); drop(bufs); @@ -1818,6 +1919,7 @@ fn create_ctx_with_diagnostics( /// (the raw context is single-threaded), matching the single-sequence session. pub struct IreeLlama { ctx: *mut XlaCtx, + compatibility_fingerprint: u64, context_capacity: usize, hidden_size: usize, has_adapter_modes: bool, @@ -1833,12 +1935,20 @@ impl IreeLlama { /// (`"local-task"` for CPU). Compiles the bundled graphs, uploads the /// weights resident, and readies the prefill / decode calls. pub fn load(model_dir: &Path, device: &str, context_capacity: usize) -> Result { + let load_started = std::time::Instant::now(); let cfg = RuntimeConfig::from_json(model_dir, context_capacity)?; + report_load_timing("language-config", load_started); runtime_ffi_dimensions(&cfg, cfg.weight_specs().len())?; + let compile_started = std::time::Instant::now(); let bundle = compile_vmfbs(device, &cfg)?; + report_load_timing("language-emit-and-compile", compile_started); + let context_started = std::time::Instant::now(); let ctx = create_ctx(model_dir, &cfg, device, &bundle)?; + report_load_timing("language-context", context_started); + report_load_timing("language-total", load_started); Ok(Self { ctx, + compatibility_fingerprint: bundle.compatibility_fingerprint, context_capacity: cfg.context_capacity(), hidden_size: cfg.hidden(), has_adapter_modes: cfg.has_adapter_modes(), @@ -1856,6 +1966,12 @@ impl IreeLlama { self.context_capacity } + /// Verified identity of the complete #876 language execution bundle. + #[must_use] + pub fn compatibility_fingerprint(&self) -> u64 { + self.compatibility_fingerprint + } + /// Seed the KV cache with `token_ids` (length <= the configured capacity) via the /// bucketed prefill graph. The returned token (the graph's argmax at /// `real_len-1`) is unused by the seed-then-decode drive loop. @@ -2159,6 +2275,7 @@ pub struct PreparedPrefillDiagnostics { pub struct IreeRaggedLlama { ctx: *mut XlaCtx, + compatibility_fingerprint: u64, b_max: usize, /// Vocabulary size (logits per row), from the model config; the readback /// buffers and the per-row logits slice are sized by it. @@ -2304,6 +2421,7 @@ impl IreeRaggedLlama { } Ok(Self { ctx, + compatibility_fingerprint: bundle.compatibility_fingerprint, b_max, vocab: cfg.vocab(), context_capacity: cfg.context_capacity(), @@ -2355,6 +2473,12 @@ impl IreeRaggedLlama { self.context_capacity } + /// Verified identity of the complete #876 language execution bundle. + #[must_use] + pub fn compatibility_fingerprint(&self) -> u64 { + self.compatibility_fingerprint + } + /// Model hidden width compiled into the embeddings entry. #[must_use] pub fn hidden_size(&self) -> usize { diff --git a/src/lib/mlxcel-xla/src/lib.rs b/src/lib/mlxcel-xla/src/lib.rs index 65034e1c3..1534a92c1 100644 --- a/src/lib/mlxcel-xla/src/lib.rs +++ b/src/lib/mlxcel-xla/src/lib.rs @@ -54,6 +54,8 @@ mod context; mod diagnostic_flags; mod gemma3n_audio_config; mod gemma3n_audio_rows; +#[cfg(feature = "iree")] +mod gemma3n_audio_runtime; mod gemma3n_audio_weights; #[cfg(any(feature = "iree", test))] #[allow(dead_code)] @@ -72,9 +74,11 @@ mod prepared_deepstack; mod prepared_gemma3n; mod prepared_gemma3n_audio; -#[cfg(feature = "iree")] +#[cfg(any(feature = "iree", test))] +#[cfg_attr(not(feature = "iree"), allow(dead_code))] mod aux; -#[cfg(feature = "iree")] +#[cfg(any(feature = "iree", test))] +#[cfg_attr(not(feature = "iree"), allow(dead_code))] mod aux_manifest; #[cfg(feature = "iree")] mod aux_smoke; @@ -217,6 +221,8 @@ pub use gemma3n_audio_config::{ GEMMA3N_AUDIO_SOFT_TOKENS, Gemma3nXlaAudioConfig, }; pub use gemma3n_audio_rows::{Gemma3nAudioRowMapError, validate_gemma3n_audio_row_indices}; +#[cfg(feature = "iree")] +pub use gemma3n_audio_runtime::{Gemma3nAudioGraphOutput, Gemma3nAudioIreeRuntime}; pub use gemma3n_audio_weights::{ GEMMA3N_AUDIO_CHECKPOINT_TENSOR_COUNT, Gemma3nAudioCheckpointDType, Gemma3nAudioCheckpointError, Gemma3nAudioCheckpointTensorSpec, gemma3n_audio_checkpoint_specs, @@ -284,6 +290,8 @@ pub struct XlaInferenceSession { #[cfg(feature = "iree")] engine: iree::IreeLlama, #[cfg(feature = "iree")] + device: String, + #[cfg(feature = "iree")] cache_len: i32, } @@ -346,6 +354,7 @@ impl XlaInferenceSession { context_capacity, eos_token_ids, engine, + device, cache_len: 0, }) } @@ -385,6 +394,24 @@ impl XlaInferenceSession { &self.eos_token_ids } + /// Conditionally load the Gemma3n audio split artifacts against this exact + /// verified #876 language bundle. + #[cfg(feature = "iree")] + pub fn load_gemma3n_audio( + &self, + frame_bucket: usize, + clips: usize, + ) -> Result { + Gemma3nAudioIreeRuntime::load( + &self.model_path, + &self.device, + self.context_capacity, + frame_bucket, + clips, + self.engine.compatibility_fingerprint(), + ) + } + /// Self-contained greedy generation over the object-safe contract. /// /// Seeds KV with the prompt prefix, then advances one token per step until an diff --git a/src/lib/mlxcel-xla/src/prepared_gemma3n_audio.rs b/src/lib/mlxcel-xla/src/prepared_gemma3n_audio.rs index cc9419822..3e1872821 100644 --- a/src/lib/mlxcel-xla/src/prepared_gemma3n_audio.rs +++ b/src/lib/mlxcel-xla/src/prepared_gemma3n_audio.rs @@ -53,7 +53,7 @@ pub enum Gemma3nAudioInputError { clip: usize, }, NonFinite, - NonZeroMaskedFeature, + NonZeroBucketPadding, } impl fmt::Display for Gemma3nAudioInputError { @@ -98,9 +98,9 @@ impl fmt::Display for Gemma3nAudioInputError { write!(f, "Gemma3n audio clip {clip} is entirely padded") } Self::NonFinite => f.write_str("Gemma3n audio mel input contains a non-finite value"), - Self::NonZeroMaskedFeature => f.write_str( - "Gemma3n audio masked mel rows must be explicit zeros at the IREE boundary", - ), + Self::NonZeroBucketPadding => { + f.write_str("Gemma3n audio static bucket tail must contain only zeros") + } } } } @@ -180,9 +180,13 @@ impl Gemma3nAudioInput { }); } let mask = &valid_mask[clip * frame_bucket..(clip + 1) * frame_bucket]; - if mask[..frames].iter().any(|&value| value != 1) - || mask[frames..].iter().any(|&value| value != 0) - { + let first_valid = mask.iter().position(|&value| value == 1); + let last_valid = mask.iter().rposition(|&value| value == 1); + let valid_count = mask.iter().filter(|&&value| value == 1).count(); + let contiguous = first_valid + .zip(last_valid) + .is_some_and(|(first, last)| mask[first..=last].iter().all(|&value| value == 1)); + if mask.iter().any(|&value| value > 1) || valid_count != frames || !contiguous { return Err(Gemma3nAudioInputError::InvalidFrameLength { clip, frames, @@ -190,13 +194,17 @@ impl Gemma3nAudioInput { }); } } - if valid_mask.iter().enumerate().any(|(index, &valid)| { - valid == 0 - && mel[index * GEMMA3N_AUDIO_MEL_BINS..(index + 1) * GEMMA3N_AUDIO_MEL_BINS] - .iter() - .any(|&value| value != 0.0) - }) { - return Err(Gemma3nAudioInputError::NonZeroMaskedFeature); + for clip in 0..clips { + let mask = &valid_mask[clip * frame_bucket..(clip + 1) * frame_bucket]; + let last_valid = mask + .iter() + .rposition(|&value| value == 1) + .expect("all-padded clips rejected above"); + let tail_start = (clip * frame_bucket + last_valid + 1) * GEMMA3N_AUDIO_MEL_BINS; + let tail_end = (clip + 1) * frame_bucket * GEMMA3N_AUDIO_MEL_BINS; + if mel[tail_start..tail_end].iter().any(|&value| value != 0.0) { + return Err(Gemma3nAudioInputError::NonZeroBucketPadding); + } } Ok(Self { mel, @@ -419,7 +427,7 @@ mod tests { } #[test] - fn input_rejects_all_padded_nonfinite_and_dirty_padding() { + fn input_preserves_left_padding_but_rejects_dirty_bucket_tail() { let bucket = 8; let valid = vec![1, 1, 0, 0, 0, 0, 0, 0]; let input = Gemma3nAudioInput::new( @@ -430,6 +438,17 @@ mod tests { ) .unwrap(); assert_eq!(input.clips(), 1); + Gemma3nAudioInput::new( + { + let mut mel = vec![0.0; bucket * GEMMA3N_AUDIO_MEL_BINS]; + mel[0] = -6.25; + mel + }, + vec![0, 0, 1, 1, 0, 0, 0, 0], + vec![2], + bucket, + ) + .expect("the pinned processor uses left padding"); assert!(matches!( Gemma3nAudioInput::new( vec![0.0; bucket * GEMMA3N_AUDIO_MEL_BINS], @@ -458,7 +477,7 @@ mod tests { dirty[2 * GEMMA3N_AUDIO_MEL_BINS] = 1.0; assert!(matches!( Gemma3nAudioInput::new(dirty, vec![1, 1, 0, 0, 0, 0, 0, 0], vec![2], bucket,), - Err(Gemma3nAudioInputError::NonZeroMaskedFeature) + Err(Gemma3nAudioInputError::NonZeroBucketPadding) )); } diff --git a/src/server/batch/mod.rs b/src/server/batch/mod.rs index d20c46e91..491782adb 100644 --- a/src/server/batch/mod.rs +++ b/src/server/batch/mod.rs @@ -61,10 +61,12 @@ mod speculative_slice_tests; mod stop_matcher; /// Backend-neutral bounded audio preprocessing foundation. This is compiled /// without `xla-iree` so queue/cancellation contracts remain testable on every -/// platform, but no XLA worker advertises audio until a qualified family -/// feature producer is wired. +/// platform. XLA workers advertise audio only when a qualified family feature +/// producer is constructed from the loaded checkpoint. #[cfg_attr(not(feature = "xla-iree"), allow(dead_code))] pub(crate) mod xla_audio_preprocess; +#[cfg(feature = "xla-iree")] +mod xla_gemma3n_audio; /// OpenXLA / IREE serve worker (issue #449 M3 Stage 2c): adapts the /// `mlxcel-xla` continuous-batching engine to the [`BatchEngine`] contract. /// Behind `xla-iree` (real IREE execution); the MLX serving path is unaffected. diff --git a/src/server/batch/xla_audio_preprocess.rs b/src/server/batch/xla_audio_preprocess.rs index 12e2401eb..1419c3ff3 100644 --- a/src/server/batch/xla_audio_preprocess.rs +++ b/src/server/batch/xla_audio_preprocess.rs @@ -16,8 +16,8 @@ //! //! This stage owns encoded image/audio request buffers, image and waveform //! decoding, resampling, host feature extraction, and prepared-prefill export. -//! A single thread-confined producer handles image-only, audio-only, and mixed -//! Phi4MM requests so request mode selection cannot split across workers. +//! A single thread-confined family producer handles the request's supported +//! media combination so request mode selection cannot split across workers. use std::mem::size_of; use std::panic::{AssertUnwindSafe, catch_unwind}; @@ -27,6 +27,8 @@ use std::thread; use std::time::Instant; use mlxcel_core::session::{OwnedTensor, PreparedPositions, PreparedPrefill}; +#[cfg(feature = "xla-iree")] +use mlxcel_xla::Gemma3nAudioPreparedPrefill; use thiserror::Error; use super::observability::BatchObservability; @@ -47,11 +49,33 @@ pub(crate) struct AudioPreprocessJob { } pub(crate) enum AudioPreprocessOutcome { - Prepared(PreparedPrefill), + Prepared(AudioPreparedRequest), Cancelled(AudioPreprocessCheckpoint), Failed(AudioStageError), } +pub(crate) enum AudioPreparedRequest { + Generic(PreparedPrefill), + #[cfg(feature = "xla-iree")] + Gemma3n(Gemma3nAudioPreparedPrefill), +} + +impl AudioPreparedRequest { + pub(crate) fn sequence_len(&self) -> usize { + match self { + Self::Generic(prepared) => prepared.sequence_len, + #[cfg(feature = "xla-iree")] + Self::Gemma3n(prepared) => prepared.request().prepared().sequence_len, + } + } +} + +impl From for AudioPreparedRequest { + fn from(value: PreparedPrefill) -> Self { + Self::Generic(value) + } +} + pub(crate) struct AudioPreprocessResult { pub job_id: u64, pub outcome: AudioPreprocessOutcome, @@ -129,7 +153,7 @@ pub(crate) trait AudioFeatureProducer: 'static { token_ids: Vec, images: Vec, cancelled: &AtomicBool, - ) -> Result; + ) -> Result; } #[cfg(feature = "xla-iree")] @@ -140,8 +164,9 @@ impl AudioFeatureProducer for crate::multimodal::phi4mm_xla_audio::Phi4MmXlaAudi token_ids: Vec, images: Vec, cancelled: &AtomicBool, - ) -> Result { + ) -> Result { self.prepare_media(waveforms, token_ids, &images, cancelled) + .map(Into::into) } } @@ -220,6 +245,7 @@ impl AudioPreprocessMetrics { } } + #[cfg(test)] pub fn snapshot(&self) -> AudioPreprocessMetricsSnapshot { AudioPreprocessMetricsSnapshot { accepted: self.accepted.load(Ordering::Relaxed), @@ -339,6 +365,7 @@ enum AudioRejectionReason { } #[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[cfg(test)] pub(crate) struct AudioPreprocessMetricsSnapshot { pub accepted: u64, pub completed: u64, @@ -409,6 +436,7 @@ pub(crate) struct AudioPreprocessStage { sender: Option>, result_rx: Option>, metrics: Arc, + #[cfg_attr(not(test), allow(dead_code))] healthy: Arc, max_queued_encoded_bytes: usize, max_in_flight_host_bytes: usize, @@ -634,14 +662,17 @@ impl AudioPreprocessStage { .try_recv() } + #[cfg_attr(not(test), allow(dead_code))] pub fn recv(&self) -> Result { self.result_rx.as_ref().ok_or(mpsc::RecvError)?.recv() } + #[cfg_attr(not(test), allow(dead_code))] pub fn metrics(&self) -> &Arc { &self.metrics } + #[cfg_attr(not(test), allow(dead_code))] pub fn is_healthy(&self) -> bool { self.healthy.load(Ordering::Acquire) } @@ -727,10 +758,10 @@ fn process_job( AudioPreprocessCheckpoint::Feature, ) } - Ok(Ok(prepared)) if prepared.sequence_len > max_prefill_tokens => { + Ok(Ok(prepared)) if prepared.sequence_len() > max_prefill_tokens => { AudioPreprocessOutcome::Failed(AudioStageError::ContextLimit { family, - actual: prepared.sequence_len, + actual: prepared.sequence_len(), maximum: max_prefill_tokens, }) } @@ -746,7 +777,7 @@ fn process_job( }, ) } else { - metrics.record_effective_prefill(prepared.sequence_len); + metrics.record_effective_prefill(prepared.sequence_len()); AudioPreprocessOutcome::Prepared(prepared) } } @@ -804,7 +835,38 @@ fn process_job( } } -fn prepared_prefill_host_bytes(prepared: &PreparedPrefill) -> Option { +fn prepared_prefill_host_bytes(prepared: &AudioPreparedRequest) -> Option { + match prepared { + AudioPreparedRequest::Generic(prepared) => generic_prepared_prefill_host_bytes(prepared), + #[cfg(feature = "xla-iree")] + AudioPreparedRequest::Gemma3n(audio) => { + let request = audio.request(); + generic_prepared_prefill_host_bytes(request.prepared())? + .checked_add(size_of::())? + .checked_add(request.dense_ple().byte_len())? + .checked_add( + audio + .projected_lengths() + .len() + .checked_mul(size_of::())?, + )? + .checked_add( + audio + .placeholder_starts() + .len() + .checked_mul(size_of::())?, + )? + .checked_add( + audio + .audio_row_indices() + .len() + .checked_mul(size_of::())?, + ) + } + } +} + +fn generic_prepared_prefill_host_bytes(prepared: &PreparedPrefill) -> Option { fn tensor_bytes(tensor: &OwnedTensor) -> Option { tensor .bytes diff --git a/src/server/batch/xla_audio_preprocess_tests.rs b/src/server/batch/xla_audio_preprocess_tests.rs index 262147747..349b47557 100644 --- a/src/server/batch/xla_audio_preprocess_tests.rs +++ b/src/server/batch/xla_audio_preprocess_tests.rs @@ -94,9 +94,9 @@ impl AudioFeatureProducer for RecordingProducer { token_ids: Vec, _images: Vec, _cancelled: &AtomicBool, - ) -> Result { + ) -> Result { self.order.lock().unwrap().push(token_ids[1] as u64); - Ok(prepared(token_ids)) + Ok(prepared(token_ids).into()) } } @@ -121,7 +121,9 @@ fn single_worker_preserves_fifo_and_records_separate_audio_metrics() { for expected in 1..=32 { let result = stage.recv().unwrap(); assert_eq!(result.job_id, expected); - let AudioPreprocessOutcome::Prepared(prefill) = result.outcome else { + let AudioPreprocessOutcome::Prepared(AudioPreparedRequest::Generic(prefill)) = + result.outcome + else { panic!("job must prepare"); }; // Logical/public prompt ids remain the producer input. Effective audio @@ -152,10 +154,10 @@ impl AudioFeatureProducer for BlockingProducer { token_ids: Vec, _images: Vec, _cancelled: &AtomicBool, - ) -> Result { + ) -> Result { let _ = self.started.send(()); let _ = self.release.recv(); - Ok(prepared(token_ids)) + Ok(prepared(token_ids).into()) } } @@ -230,9 +232,9 @@ impl AudioFeatureProducer for CancelAfterFeature { token_ids: Vec, _images: Vec, cancelled: &AtomicBool, - ) -> Result { + ) -> Result { cancelled.store(true, Ordering::Release); - Ok(prepared(token_ids)) + Ok(prepared(token_ids).into()) } } @@ -277,12 +279,12 @@ impl AudioFeatureProducer for PanicOnce { token_ids: Vec, _images: Vec, _cancelled: &AtomicBool, - ) -> Result { + ) -> Result { if !self.0 { self.0 = true; panic!("synthetic feature panic"); } - Ok(prepared(token_ids)) + Ok(prepared(token_ids).into()) } } @@ -378,9 +380,9 @@ impl AudioFeatureProducer for StartedProducer { token_ids: Vec, _images: Vec, _cancelled: &AtomicBool, - ) -> Result { + ) -> Result { self.started.send(token_ids[1] as u64).unwrap(); - Ok(prepared(token_ids)) + Ok(prepared(token_ids).into()) } } diff --git a/src/server/batch/xla_gemma3n_audio.rs b/src/server/batch/xla_gemma3n_audio.rs new file mode 100644 index 000000000..d75341886 --- /dev/null +++ b/src/server/batch/xla_gemma3n_audio.rs @@ -0,0 +1,145 @@ +// 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. + +//! Gemma3n host-mel to IREE audio producer for bounded server admission. + +use std::path::{Path, PathBuf}; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; + +use mlxcel_xla::Gemma3nAudioIreeRuntime; + +use super::observability::BatchObservability; +use super::xla_audio_preprocess::{ + AudioFeatureProducer, AudioPreparedRequest, AudioPreprocessLimits, AudioPreprocessStage, +}; +use crate::audio::gemma3n::{Gemma3nAudioFeatureExtractor, Gemma3nXlaAudioPreparer}; +use crate::audio::{AudioFamilyPolicy, AudioWaveformBatch}; +use crate::tokenizer::MlxcelTokenizer; + +pub(super) struct Gemma3nAudioStageConfig { + model_path: PathBuf, + device: String, + context_capacity: usize, + language_fingerprint: u64, + preparer: Gemma3nXlaAudioPreparer, +} + +impl Gemma3nAudioStageConfig { + pub(super) fn from_model( + model_path: &Path, + device: &str, + context_capacity: usize, + language_fingerprint: u64, + tokenizer: &MlxcelTokenizer, + ) -> Result, String> { + let Some(preparer) = + Gemma3nXlaAudioPreparer::from_model(model_path, context_capacity, tokenizer)? + else { + return Ok(None); + }; + if language_fingerprint == 0 { + return Err("Gemma3n audio requires a verified language bundle".to_string()); + } + Ok(Some(Self { + model_path: model_path.to_path_buf(), + device: device.to_string(), + context_capacity, + language_fingerprint, + preparer, + })) + } + + pub(super) fn policy(&self) -> AudioFamilyPolicy { + self.preparer.policy() + } + + pub(super) fn spawn( + self, + limits: AudioPreprocessLimits, + observability: Arc, + ) -> Result { + AudioPreprocessStage::spawn_with_loader(limits, observability, move || { + Ok(Gemma3nAudioProducer::new(self)) + }) + } +} + +struct Gemma3nAudioProducer { + config: Gemma3nAudioStageConfig, + extractor: Gemma3nAudioFeatureExtractor, + runtime: Option<((usize, usize), Gemma3nAudioIreeRuntime)>, +} + +impl Gemma3nAudioProducer { + fn new(config: Gemma3nAudioStageConfig) -> Self { + Self { + config, + extractor: Gemma3nAudioFeatureExtractor::new(), + runtime: None, + } + } + + fn runtime( + &mut self, + frame_bucket: usize, + clips: usize, + ) -> Result<&mut Gemma3nAudioIreeRuntime, String> { + let key = (frame_bucket, clips); + if self + .runtime + .as_ref() + .is_none_or(|(loaded, _)| *loaded != key) + { + let runtime = Gemma3nAudioIreeRuntime::load( + &self.config.model_path, + &self.config.device, + self.config.context_capacity, + frame_bucket, + clips, + self.config.language_fingerprint, + )?; + // Keep a single shape-specialized runtime resident. This bounds + // conditional audio weight memory even when request buckets vary. + self.runtime = Some((key, runtime)); + } + Ok(&mut self.runtime.as_mut().expect("runtime inserted").1) + } +} + +impl AudioFeatureProducer for Gemma3nAudioProducer { + fn prepare( + &mut self, + waveforms: AudioWaveformBatch, + token_ids: Vec, + images: Vec, + cancelled: &AtomicBool, + ) -> Result { + if !images.is_empty() { + return Err("Gemma3n audio producer does not accept image inputs".to_string()); + } + let prepared = + self.config + .preparer + .prepare(&self.extractor, waveforms, token_ids, cancelled)?; + if cancelled.load(Ordering::Acquire) { + return Err("Gemma3n audio request was cancelled before IREE invocation".to_string()); + } + let audio_token_id = self.config.preparer.audio_token_id(); + let prepared = self + .runtime(prepared.input.frame_bucket(), prepared.clips)? + .invoke_prepared(&prepared.input, prepared.token_ids, audio_token_id)?; + Ok(AudioPreparedRequest::Gemma3n(prepared)) + } +} diff --git a/src/server/batch/xla_worker.rs b/src/server/batch/xla_worker.rs index 81105bd6e..5c7d82028 100644 --- a/src/server/batch/xla_worker.rs +++ b/src/server/batch/xla_worker.rs @@ -38,7 +38,8 @@ //! applied incrementally so it is safe across token boundaries). Requests that //! need features the engine cannot provide are rejected with a clear error rather //! than served wrong: logprobs (no logit readback), structured / JSON-schema -//! output (no constraint masking), and unsupported audio/video inputs. +//! output (no constraint masking), video, and modalities not qualified by the +//! loaded bundle. Gemma3n audio uses a bounded worker-owned split IREE runtime. //! //! Compiled only under `xla-iree` (real IREE execution). The MLX serving path is //! untouched. @@ -50,13 +51,18 @@ use std::sync::mpsc; use std::time::Instant; use mlxcel_core::session::PreparedPrefill; -use mlxcel_xla::{EngineEvent, FinishReason as XlaFinishReason, SampleParams, XlaBatchEngine}; +use mlxcel_xla::{ + EngineEvent, FinishReason as XlaFinishReason, Gemma3nPreparedPrefill, SampleParams, + XlaBatchEngine, +}; use super::BatchEngine; use super::observability::BatchObservability; use super::stop_matcher::StopMatcher; use super::xla_audio_preprocess::{AudioPreprocessLimits, AudioPreprocessStage}; +use super::xla_gemma3n_audio::Gemma3nAudioStageConfig; use super::xla_preprocess::ImagePreprocessStage; +use crate::audio::AudioFamilyPolicy; use crate::server::ServerGenerateOptions; use crate::server::media::MediaRequestMetadata; use crate::server::model_provider::model_worker::StreamingDecodeState; @@ -72,6 +78,7 @@ mod tests; pub(crate) trait XlaServingEngine { fn b_max(&self) -> usize; + fn context_capacity(&self) -> usize; fn is_idle(&self) -> bool; fn pending_len(&self) -> usize; fn active_len(&self) -> usize; @@ -87,6 +94,12 @@ pub(crate) trait XlaServingEngine { max_new_tokens: usize, params: SampleParams, ) -> Result; + fn submit_gemma3n_prepared( + &mut self, + prepared: Gemma3nPreparedPrefill, + max_new_tokens: usize, + params: SampleParams, + ) -> Result; fn cancel(&mut self, req_id: u64) -> bool; fn pump(&mut self) -> Result, String>; } @@ -96,6 +109,10 @@ impl XlaServingEngine for XlaBatchEngine { XlaBatchEngine::b_max(self) } + fn context_capacity(&self) -> usize { + XlaBatchEngine::context_capacity(self) + } + fn is_idle(&self) -> bool { XlaBatchEngine::is_idle(self) } @@ -128,6 +145,16 @@ impl XlaServingEngine for XlaBatchEngine { .map_err(|error| error.to_string()) } + fn submit_gemma3n_prepared( + &mut self, + prepared: Gemma3nPreparedPrefill, + max_new_tokens: usize, + params: SampleParams, + ) -> Result { + XlaBatchEngine::submit_gemma3n_prepared(self, prepared, max_new_tokens, params) + .map_err(|error| error.to_string()) + } + fn cancel(&mut self, req_id: u64) -> bool { XlaBatchEngine::cancel(self, req_id) } @@ -202,7 +229,7 @@ pub(crate) struct XlaServeWorker { /// The audio stage is the unified Phi4MM image/audio producer rather than /// an audio-only family producer. phi4mm_media_preprocessor: bool, - audio_policy: Option, + audio_policy: Option, pending_audio: HashMap, next_audio_job_id: u64, context_capacity: usize, @@ -242,7 +269,23 @@ impl XlaServeWorker { )?; (Some(stage), Some(policy)) } else { - (None, None) + let audio_config = Gemma3nAudioStageConfig::from_model( + &model_path, + &device, + context_capacity, + engine.compatibility_fingerprint(), + &tokenizer, + )?; + let audio_policy = audio_config.as_ref().map(Gemma3nAudioStageConfig::policy); + let audio_preprocessor = audio_config + .map(|config| { + config.spawn( + AudioPreprocessLimits::default(), + batch_observability.clone(), + ) + }) + .transpose()?; + (audio_preprocessor, audio_policy) }; Ok(Self { engine, diff --git a/src/server/batch/xla_worker_admission.rs b/src/server/batch/xla_worker_admission.rs index 567132af8..960d040af 100644 --- a/src/server/batch/xla_worker_admission.rs +++ b/src/server/batch/xla_worker_admission.rs @@ -17,12 +17,14 @@ use std::time::Duration; use super::super::xla_audio_preprocess::{ - AudioPreprocessJob, AudioPreprocessOutcome, AudioPreprocessResult, AudioQueueError, + AudioPreparedRequest, AudioPreprocessJob, AudioPreprocessOutcome, AudioPreprocessResult, + AudioQueueError, }; use super::super::xla_preprocess::{ ImagePreprocessJob, ImagePreprocessOutcome, ImagePreprocessResult, }; use super::*; +use crate::audio::{AudioEncodedClip, AudioSourceKind}; pub(super) fn validate_xla_output_features(logprobs: bool, structured: bool) -> Result<(), String> { if logprobs { @@ -212,7 +214,14 @@ impl XlaServeWorker { return; } - if images.is_empty() { + if !images.is_empty() && !audio.is_empty() { + let _ = response_tx.send(GenerateEvent::Error( + "mixed image/audio input is not qualified for the loaded OpenXLA bundle" + .to_string(), + )); + return; + } + if images.is_empty() && audio.is_empty() { self.submit_text( prompt_tokens, options.max_tokens, @@ -224,6 +233,19 @@ impl XlaServeWorker { ); return; } + if !audio.is_empty() { + self.submit_audio( + prompt_tokens, + options.max_tokens, + params, + stop_sequences, + start, + audio, + response_tx, + cancelled, + ); + return; + } let Some(stage) = self.image_preprocessor.as_ref() else { let _ = response_tx.send(GenerateEvent::Error( @@ -274,6 +296,81 @@ impl XlaServeWorker { } } + #[allow(clippy::too_many_arguments)] + fn submit_audio( + &mut self, + prompt_tokens: Vec, + max_tokens: usize, + params: SampleParams, + stop_sequences: Vec, + start: Instant, + audio: Vec>, + response_tx: mpsc::Sender, + cancelled: Arc, + ) { + let (Some(stage), Some(policy)) = (self.audio_preprocessor.as_ref(), self.audio_policy) + else { + let _ = response_tx.send(GenerateEvent::Error( + "the loaded OpenXLA model/runtime bundle does not support audio input".to_string(), + )); + return; + }; + let Some(next_job_id) = self.next_audio_job_id.checked_add(1) else { + let _ = response_tx.send(GenerateEvent::Error( + "OpenXLA audio preprocessing request id overflowed".to_string(), + )); + return; + }; + let job_id = self.next_audio_job_id; + self.next_audio_job_id = next_job_id; + let clips = audio + .into_iter() + .enumerate() + .map(|(index, bytes)| AudioEncodedClip { + bytes, + source: AudioSourceKind::ServerInline, + placeholder_ordinal: index + 1, + }) + .collect(); + let job = AudioPreprocessJob { + job_id, + token_ids: prompt_tokens.clone(), + max_prefill_tokens: self.engine.context_capacity(), + expected_image_count: 0, + images: Vec::new(), + clips, + policy, + cancelled: cancelled.clone(), + }; + match stage.try_submit(job) { + Ok(()) => { + self.pending_audio.insert( + job_id, + PendingAudioState { + response_tx, + cancelled, + prompt_tokens, + params, + stop_sequences, + max_tokens, + start, + }, + ); + } + Err(AudioQueueError::Cancelled { .. }) => {} + Err(AudioQueueError::Full) => { + let _ = response_tx.send(GenerateEvent::Error( + "OpenXLA audio preprocessing queue is full; retry later".to_string(), + )); + } + Err(error) => { + let _ = response_tx.send(GenerateEvent::Error(format!( + "OpenXLA audio preprocessing admission failed: {error}" + ))); + } + } + } + #[allow(clippy::too_many_arguments)] fn submit_text( &mut self, @@ -392,7 +489,13 @@ impl XlaServeWorker { } let prepared = match result.outcome { AudioPreprocessOutcome::Prepared(prepared) => prepared, - AudioPreprocessOutcome::Cancelled(_) => return, + AudioPreprocessOutcome::Cancelled(checkpoint) => { + tracing::debug!( + ?checkpoint, + "OpenXLA audio preprocessing cancelled before admission" + ); + return; + } AudioPreprocessOutcome::Failed(error) => { let _ = state.response_tx.send(GenerateEvent::Error(format!( "OpenXLA audio preprocessing failed: {error}" @@ -404,12 +507,20 @@ impl XlaServeWorker { self.finish_zero_budget(state.prompt_tokens, state.start, state.response_tx); return; } - let effective_prefill_len = prepared.sequence_len; + let effective_prefill_len = prepared.sequence_len(); let prompt_token_count = state.prompt_tokens.len(); - match self - .engine - .submit_prepared(prepared, state.max_tokens, state.params) - { + let submitted = match prepared { + AudioPreparedRequest::Generic(prepared) => { + self.engine + .submit_prepared(prepared, state.max_tokens, state.params) + } + AudioPreparedRequest::Gemma3n(prepared) => self.engine.submit_gemma3n_prepared( + prepared.into_request(), + state.max_tokens, + state.params, + ), + }; + match submitted { Ok(req_id) => { self.batch_observability .record_prefill_start(effective_prefill_len); @@ -437,6 +548,11 @@ impl XlaServeWorker { } pub(super) fn drain_preprocessed(&mut self) { + self.drain_image_preprocessed(); + self.drain_audio_preprocessed(); + } + + fn drain_image_preprocessed(&mut self) { loop { let received = match self.image_preprocessor.as_ref() { Some(stage) => stage.try_recv(), @@ -478,6 +594,29 @@ impl XlaServeWorker { } } + fn drain_audio_preprocessed(&mut self) { + loop { + let received = match self.audio_preprocessor.as_ref() { + Some(stage) => stage.try_recv(), + None => return, + }; + match received { + Ok(result) => self.handle_audio_preprocessed(result), + Err(mpsc::TryRecvError::Empty) => return, + Err(mpsc::TryRecvError::Disconnected) => { + self.audio_preprocessor = None; + self.audio_policy = None; + for (_, state) in self.pending_audio.drain() { + let _ = state.response_tx.send(GenerateEvent::Error( + "OpenXLA audio preprocessor stopped unexpectedly".to_string(), + )); + } + return; + } + } + } + } + fn handle(&mut self, request: ModelRequest) { match request { ModelRequest::Generate { diff --git a/src/server/batch/xla_worker_tests.rs b/src/server/batch/xla_worker_tests.rs index 39307d3da..f5f5a5546 100644 --- a/src/server/batch/xla_worker_tests.rs +++ b/src/server/batch/xla_worker_tests.rs @@ -20,14 +20,19 @@ use std::time::Duration; use image::{DynamicImage, ImageFormat}; use mlxcel_core::session::{ - OwnedTensor, PreparedAdapterMode, PreparedAttentionBias, PreparedPositions, PreparedTensorDType, + OwnedTensor, PreparedAdapterMode, PreparedAttentionBias, PreparedModality, PreparedPositions, + PreparedTensorDType, +}; +use mlxcel_xla::{ + GEMMA3N_AUDIO_MODALITY_FAMILY, GEMMA3N_AUDIO_SOFT_TOKENS, Gemma3nAudioPreparedPrefill, + Gemma3nDensePle, }; use super::*; use crate::FakeHostMultimodalPreprocessor; use crate::audio::{AudioFamilyPolicy, AudioWaveformBatch}; use crate::server::batch::xla_audio_preprocess::{ - AudioFeatureProducer, AudioPreprocessLimits, AudioPreprocessStage, + AudioFeatureProducer, AudioPreparedRequest, AudioPreprocessLimits, AudioPreprocessStage, }; use crate::server::request_options::{RequestOptionOverrides, build_server_generate_options}; use crate::server::{GenerationResult, ServerConfig}; @@ -39,6 +44,7 @@ struct FakeServingEngine { text_submissions: Vec>, prepared_submissions: Vec<(Vec, usize)>, prepared_modes: Vec, + gemma3n_submissions: Vec<(Vec, usize)>, } impl XlaServingEngine for FakeServingEngine { @@ -46,6 +52,10 @@ impl XlaServingEngine for FakeServingEngine { 4 } + fn context_capacity(&self) -> usize { + 4_096 + } + fn is_idle(&self) -> bool { self.active.is_empty() } @@ -80,6 +90,19 @@ impl XlaServingEngine for FakeServingEngine { Ok(self.activate()) } + fn submit_gemma3n_prepared( + &mut self, + prepared: Gemma3nPreparedPrefill, + _max_new_tokens: usize, + _params: SampleParams, + ) -> Result { + self.gemma3n_submissions.push(( + prepared.prepared().token_ids.clone(), + prepared.prepared().sequence_len, + )); + Ok(self.activate()) + } + fn cancel(&mut self, req_id: u64) -> bool { self.active.remove(&req_id) } @@ -118,6 +141,10 @@ fn png_bytes() -> Vec { fn wav_bytes() -> Vec { let samples = [0i16; 320]; + wav_pcm16(&samples) +} + +fn wav_pcm16(samples: &[i16]) -> Vec { let data_len = samples.len() * 2; let mut wav = Vec::with_capacity(44 + data_len); wav.extend_from_slice(b"RIFF"); @@ -138,6 +165,72 @@ fn wav_bytes() -> Vec { wav } +fn owned_tensor(shape: Vec, values: Vec) -> OwnedTensor { + OwnedTensor::new( + values.into_iter().flat_map(f32::to_le_bytes).collect(), + PreparedTensorDType::Float32, + shape, + ) + .unwrap() +} + +struct FakeGemma3nAudioProducer; + +impl AudioFeatureProducer for FakeGemma3nAudioProducer { + fn prepare( + &mut self, + _waveforms: AudioWaveformBatch, + token_ids: Vec, + _images: Vec, + _cancelled: &AtomicBool, + ) -> Result { + const AUDIO: i32 = 262_273; + let placeholder = token_ids + .iter() + .position(|&token| token == AUDIO) + .ok_or_else(|| "test prompt has no audio placeholder".to_string())?; + let mut expanded = token_ids; + expanded.splice( + placeholder..=placeholder, + std::iter::repeat_n(AUDIO, GEMMA3N_AUDIO_SOFT_TOKENS), + ); + let sequence_len = expanded.len(); + let prepared = PreparedPrefill::new( + expanded, + owned_tensor(vec![1, sequence_len, 2], vec![0.0; sequence_len * 2]), + PreparedPositions::Sequential { + start: 0, + length: sequence_len, + }, + PreparedAttentionBias { + tensor: owned_tensor(vec![1, 1, 1, sequence_len], vec![0.0; sequence_len]), + causal: true, + }, + vec![PreparedModality { + family: GEMMA3N_AUDIO_MODALITY_FAMILY.to_string(), + item_count: 1, + token_count: GEMMA3N_AUDIO_SOFT_TOKENS, + }], + ) + .map_err(|error| error.to_string())?; + let dense_ple = Gemma3nDensePle::new(vec![0.0; sequence_len * 2], sequence_len, 1, 2) + .map_err(|error| error.to_string())?; + let prepared = + Gemma3nPreparedPrefill::new(prepared, dense_ple).map_err(|error| error.to_string())?; + Gemma3nAudioPreparedPrefill::new(prepared, AUDIO, vec![GEMMA3N_AUDIO_SOFT_TOKENS], 8) + .map(AudioPreparedRequest::Gemma3n) + } +} + +fn audio_stage() -> AudioPreprocessStage { + AudioPreprocessStage::spawn( + FakeGemma3nAudioProducer, + AudioPreprocessLimits::default(), + Arc::new(BatchObservability::new()), + ) + .unwrap() +} + fn options(max_tokens: usize) -> ServerGenerateOptions { build_server_generate_options( &ServerConfig::default(), @@ -196,7 +289,7 @@ impl AudioFeatureProducer for FakePhi4MediaProducer { token_ids: Vec, images: Vec, _cancelled: &AtomicBool, - ) -> Result { + ) -> Result { let sequence = token_ids.len(); let embeddings = OwnedTensor::new( vec![0; sequence * 2 * std::mem::size_of::()], @@ -231,6 +324,7 @@ impl AudioFeatureProducer for FakePhi4MediaProducer { Vec::new(), ) .map(|prepared| prepared.with_adapter_mode(mode)) + .map(Into::into) .map_err(|error| error.to_string()) } } @@ -330,6 +424,17 @@ fn phi4mm_media_admission_selects_speech_or_vision_per_request() { } } +fn wait_for_audio_preprocessing(worker: &mut XlaServeWorker) { + for _ in 0..500 { + worker.drain_preprocessed(); + if worker.pending_audio.is_empty() { + return; + } + std::thread::sleep(Duration::from_millis(1)); + } + panic!("timed out waiting for audio preprocessing"); +} + fn receive_done(rx: &mpsc::Receiver) -> GenerationResult { loop { match rx.recv_timeout(Duration::from_secs(1)).unwrap() { @@ -402,6 +507,95 @@ fn mixed_text_and_image_admission_keeps_public_and_effective_lengths_distinct() ); } +#[test] +fn text_and_gemma3n_audio_batch_cancel_and_reuse_keep_request_ownership() { + let mut worker = worker(None); + worker.audio_preprocessor = Some(audio_stage()); + worker.audio_policy = Some(AudioFamilyPolicy::gemma3n()); + let (text_tx, text_rx) = mpsc::channel(); + let (audio_tx, audio_rx) = mpsc::channel(); + + worker.admit( + String::new(), + Some(vec![10, 11]), + options(1), + Vec::new(), + Vec::new(), + Vec::new(), + media(0, 0, 0), + text_tx, + Arc::new(AtomicBool::new(false)), + ); + worker.admit( + String::new(), + Some(vec![20, 262_273, 21]), + options(1), + Vec::new(), + vec![wav_pcm16(&vec![0; 800])], + Vec::new(), + media(0, 1, 0), + audio_tx, + Arc::new(AtomicBool::new(false)), + ); + + assert_eq!(worker.engine.text_submissions, vec![vec![10, 11]]); + wait_for_audio_preprocessing(&mut worker); + assert_eq!(worker.engine.gemma3n_submissions.len(), 1); + assert_eq!( + worker.engine.gemma3n_submissions[0].1, + 2 + GEMMA3N_AUDIO_SOFT_TOKENS + ); + assert_eq!(worker.states.len(), 2); + + let events = worker.engine.pump().unwrap(); + worker.dispatch(events); + assert_eq!(receive_done(&text_rx).prompt_tokens, 2); + assert_eq!(receive_done(&audio_rx).prompt_tokens, 3); + + let cancelled = Arc::new(AtomicBool::new(true)); + let (cancelled_tx, cancelled_rx) = mpsc::channel(); + worker.admit( + String::new(), + Some(vec![30, 262_273, 31]), + options(1), + Vec::new(), + vec![wav_pcm16(&vec![0; 800])], + Vec::new(), + media(0, 1, 0), + cancelled_tx, + cancelled, + ); + worker.evict_cancelled(); + wait_for_audio_preprocessing(&mut worker); + assert!(cancelled_rx.try_recv().is_err()); + assert!(worker.states.is_empty()); + + let (reuse_tx, reuse_rx) = mpsc::channel(); + worker.admit( + String::new(), + Some(vec![40, 262_273, 41]), + options(1), + Vec::new(), + vec![wav_pcm16(&vec![0; 800])], + Vec::new(), + media(0, 1, 0), + reuse_tx, + Arc::new(AtomicBool::new(false)), + ); + wait_for_audio_preprocessing(&mut worker); + assert_eq!(worker.engine.gemma3n_submissions.len(), 2); + let events = worker.engine.pump().unwrap(); + worker.dispatch(events); + assert_eq!(receive_done(&reuse_rx).prompt_tokens, 3); + assert_eq!( + worker + .batch_metrics + .total_sequences_processed + .load(Ordering::Relaxed), + 3 + ); +} + #[test] fn malformed_image_failure_does_not_disturb_active_text_request() { let mut worker = worker(Some(image_stage())); diff --git a/src/server/media.rs b/src/server/media.rs index 088528adf..b4b79a2b2 100644 --- a/src/server/media.rs +++ b/src/server/media.rs @@ -43,11 +43,9 @@ pub(crate) trait AudioAcquisitionCancellation: Sync { /// admission. Its cancellation boundary is the request task itself: dropping /// the Axum handler future drops this acquisition future and any in-flight /// reqwest byte stream. The explicit probe is reserved for callers that already -/// own a cooperative token (including the future XLA audio admission path). -/// -/// Do not interpret this fallback as a live client-disconnect probe. XLA audio -/// admission remains disabled until a family feature producer and token wiring -/// are both present. +/// own a cooperative token. Do not interpret this fallback as a live +/// client-disconnect probe; after acquisition, XLA audio admission uses its +/// request-owned cancellation token at every bounded preprocessing checkpoint. struct NeverCancelAudioAcquisition; impl AudioAcquisitionCancellation for NeverCancelAudioAcquisition { @@ -169,10 +167,10 @@ impl MediaRequestMetadata { /// Validate the XLA request seam without changing tolerant MLX resolution. /// - /// Audio/video declarations are rejected from the declared counts so a - /// failed resolver cannot erase an unsupported modality. Images require a - /// one-to-one declaration → raw payload handoff; decoded cardinality is - /// checked later by the bounded preprocessing worker. + /// Video declarations are rejected from the declared count so a failed + /// resolver cannot erase an unsupported modality. Images and audio require + /// a one-to-one declaration → raw payload handoff; decoded cardinality is + /// checked later by their bounded preprocessing workers. #[cfg_attr(not(feature = "xla-iree"), allow(dead_code))] pub(crate) fn validate_xla_raw_counts( self, From 600e8a2c20223efede9baccfdc45892a941f2310 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 24 Jul 2026 16:39:32 +0900 Subject: [PATCH 05/19] fix(xla): align vision auxiliary weight storage Keep the rebased LLaVA vision runtime on the shared aligned auxiliary-weight storage API used by Gemma3n audio. Refs #878 --- src/lib/mlxcel-xla/src/vision_runtime.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/lib/mlxcel-xla/src/vision_runtime.rs b/src/lib/mlxcel-xla/src/vision_runtime.rs index b16bc9d62..c09cee3bc 100644 --- a/src/lib/mlxcel-xla/src/vision_runtime.rs +++ b/src/lib/mlxcel-xla/src/vision_runtime.rs @@ -33,7 +33,7 @@ use sha2::{Digest, Sha256}; use crate::aux::{ AuxiliaryInput, AuxiliaryOutput, AuxiliaryTensorDType, AuxiliaryWeight, AuxiliaryWeightDType, - IreeAuxiliaryModule, + AuxiliaryWeightStorage, IreeAuxiliaryModule, }; use crate::aux_manifest::{AuxiliaryArtifactContract, ensure_qualified_auxiliary_artifact}; #[cfg(feature = "diagnostics")] @@ -539,7 +539,7 @@ fn load_weights( validate_finite_values(&format!("vision tensor {}", spec.name), &values)?; loaded[index] = Some(AuxiliaryWeight { name: spec.name.clone(), - bytes: native_f32_bytes(values), + storage: AuxiliaryWeightStorage::Bytes(native_f32_bytes(values)), dtype: AuxiliaryWeightDType::Float32, shape: spec.shape.clone(), }); @@ -910,7 +910,7 @@ mod tests { let vmfb = temp_path("compiler-digest").with_extension("vmfb"); let weights = vec![AuxiliaryWeight { name: "weight".to_string(), - bytes: 1.0f32.to_ne_bytes().to_vec(), + storage: AuxiliaryWeightStorage::Bytes(1.0f32.to_ne_bytes().to_vec()), dtype: AuxiliaryWeightDType::Float32, shape: vec![1], }]; From ca87b5cd6bffd11e87aebed132957dbaec887ac1 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 24 Jul 2026 17:24:54 +0900 Subject: [PATCH 06/19] fix(xla): harden Gemma3n audio diagnostics --- Cargo.toml | 6 + examples/xla_gemma3n_audio_reference_check.rs | 286 +++++++++++++++ src/audio/gemma3n/diagnostics.rs | 334 ++++++++++++++++++ src/audio/gemma3n/mod.rs | 4 + .../src/emitter/gemma3n_audio_emit.rs | 103 +++--- .../src/emitter/gemma3n_audio_emit_tests.rs | 9 +- src/lib/mlxcel-xla/src/emitter/mod.rs | 4 - .../mlxcel-xla/src/gemma3n_audio_runtime.rs | 290 ++++++++++++++- src/models/gemma3n.rs | 33 ++ 9 files changed, 995 insertions(+), 74 deletions(-) create mode 100644 examples/xla_gemma3n_audio_reference_check.rs create mode 100644 src/audio/gemma3n/diagnostics.rs diff --git a/Cargo.toml b/Cargo.toml index 6068b7dba..c30116456 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -349,3 +349,9 @@ name = "xla_phi4_audio_check" [[example]] name = "xla_gemma3n_audio_check" required-features = ["xla-iree"] + +# Audio-only MLX↔IREE first-divergence gate for #878. This intentionally does +# not construct the 30 GB language bundle. +[[example]] +name = "xla_gemma3n_audio_reference_check" +required-features = ["xla-diagnostics"] diff --git a/examples/xla_gemma3n_audio_reference_check.rs b/examples/xla_gemma3n_audio_reference_check.rs new file mode 100644 index 000000000..6b20795b2 --- /dev/null +++ b/examples/xla_gemma3n_audio_reference_check.rs @@ -0,0 +1,286 @@ +// 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. + +//! Audio-only MLX↔IREE intermediate gate for the pinned Gemma3n #875 fixture. + +use std::path::{Path, PathBuf}; +use std::sync::atomic::AtomicBool; + +use mlxcel::audio::gemma3n::{Gemma3nAudioFeatureExtractor, run_gemma3n_audio_mlx_diagnostics}; +use mlxcel::audio::{AudioFamilyPolicy, preprocess_wav_file}; +use mlxcel::tokenizer::load_tokenizer; +use mlxcel_xla::{ + GEMMA3N_AUDIO_MEL_BINS, GEMMA3N_AUDIO_SOFT_TOKENS, Gemma3nAudioInput, Gemma3nAudioIreeRuntime, + select_gemma3n_audio_frame_bucket, +}; + +const PROMPT: &str = "Transcribe the following speech segment in English:"; +const AUDIO_TOKEN_ID: i32 = 262_273; +const BOA_TOKEN_ID: i32 = 256_000; +const EOA_TOKEN_ID: i32 = 262_272; +const CONTEXT_CAPACITY: usize = 256; + +fn path_env(name: &str, default: &str) -> PathBuf { + std::env::var_os(name) + .map(PathBuf::from) + .unwrap_or_else(|| PathBuf::from(default)) +} + +fn expanded_prompt(model_dir: &Path) -> Result, String> { + let tokenizer = load_tokenizer(model_dir).map_err(|error| error.to_string())?; + let rendered = format!( + "user\n{PROMPT}\n\ + model\n" + ); + let mut tokens = tokenizer + .encode(&rendered, false) + .map_err(|error| error.to_string())? + .into_iter() + .map(|token| token as i32) + .collect::>(); + let wrapper = tokenizer + .encode("\n\n", false) + .map_err(|error| error.to_string())? + .into_iter() + .map(|token| token as i32) + .collect::>(); + mlxcel::vlm_runtime::expand_gemma3n_audio_tokens( + &mut tokens, + AUDIO_TOKEN_ID, + BOA_TOKEN_ID, + EOA_TOKEN_ID, + 1, + GEMMA3N_AUDIO_SOFT_TOKENS, + &wrapper, + None, + ) + .map_err(|error| error.to_string())?; + Ok(tokens) +} + +fn audio_input(audio_path: &Path) -> Result { + let waveforms = preprocess_wav_file( + audio_path, + AudioFamilyPolicy::gemma3n(), + &AtomicBool::new(false), + ) + .map_err(|error| error.to_string())?; + let clips = waveforms + .clips + .into_iter() + .map(|clip| clip.samples) + .collect::>(); + let features = Gemma3nAudioFeatureExtractor::new().extract_batch(&clips)?; + let bucket = + select_gemma3n_audio_frame_bucket(features.frames).map_err(|error| error.to_string())?; + let mut mel = vec![0.0; features.batch_size * bucket * GEMMA3N_AUDIO_MEL_BINS]; + let mut valid_mask = vec![0; features.batch_size * bucket]; + let mut frame_lengths = Vec::with_capacity(features.batch_size); + for clip in 0..features.batch_size { + let mut valid_frames = 0; + for frame in 0..features.frames { + let source_row = clip * features.frames + frame; + let target_row = clip * bucket + frame; + valid_mask[target_row] = u8::from(features.valid_mask[source_row]); + valid_frames += usize::from(features.valid_mask[source_row]); + let source = source_row * GEMMA3N_AUDIO_MEL_BINS; + let target = target_row * GEMMA3N_AUDIO_MEL_BINS; + mel[target..target + GEMMA3N_AUDIO_MEL_BINS] + .copy_from_slice(&features.features[source..source + GEMMA3N_AUDIO_MEL_BINS]); + } + frame_lengths.push(valid_frames); + } + Gemma3nAudioInput::new(mel, valid_mask, frame_lengths, bucket) + .map_err(|error| error.to_string()) +} + +fn audio_rows(tokens: &[i32]) -> Result, String> { + let mut rank = 0i32; + let rows = tokens + .iter() + .map(|token| { + if *token == AUDIO_TOKEN_ID { + let row = rank; + rank += 1; + row + } else { + -1 + } + }) + .collect::>(); + if rank as usize != GEMMA3N_AUDIO_SOFT_TOKENS { + return Err(format!( + "expanded prompt has {rank} audio rows, expected {GEMMA3N_AUDIO_SOFT_TOKENS}" + )); + } + Ok(rows) +} + +struct Diff { + max_abs: f64, + rms: f64, + max_index: usize, +} + +fn compare(name: &str, reference: &[f32], actual: &[f32]) -> Result { + if reference.len() != actual.len() { + return Err(format!( + "{name} length mismatch: MLX={} IREE={}", + reference.len(), + actual.len() + )); + } + let mut squared = 0.0f64; + let mut max_abs = 0.0f64; + let mut max_index = 0usize; + for (index, (reference, actual)) in reference.iter().zip(actual).enumerate() { + let difference = f64::from(*actual) - f64::from(*reference); + squared += difference * difference; + if difference.abs() > max_abs { + max_abs = difference.abs(); + max_index = index; + } + } + let rms = (squared / reference.len().max(1) as f64).sqrt(); + let probes = [ + 0usize, + 1.min(reference.len().saturating_sub(1)), + reference.len() / 3, + reference.len() / 2, + reference.len().saturating_sub(1), + max_index, + ]; + eprintln!( + "gemma3n-audio-reference: stage={name} max_abs={max_abs:.9e} rms={rms:.9e} max_index={max_index}" + ); + for index in probes { + eprintln!( + "gemma3n-audio-reference: probe stage={name} index={index} mlx={:.9e} iree={:.9e} abs={:.9e}", + reference[index], + actual[index], + (actual[index] - reference[index]).abs() + ); + } + Ok(Diff { + max_abs, + rms, + max_index, + }) +} + +fn main() -> Result<(), String> { + let model_dir = path_env( + "MLXCEL_GEMMA3N_AUDIO_MODEL", + "/home/inureyes/models/gemma3n-e4b-4bit", + ); + let audio_path = path_env("MLXCEL_GEMMA3N_AUDIO_WAV", "/tmp/gemma3n-103-1240-0000.wav"); + let device = std::env::var("MLXCEL_XLA_DEVICE") + .map_err(|_| "MLXCEL_XLA_DEVICE must be set explicitly for this gate".to_string())?; + let tokens = expanded_prompt(&model_dir)?; + let rows = audio_rows(&tokens)?; + let input = audio_input(&audio_path)?; + eprintln!( + "gemma3n-audio-reference: frontend bucket={} clips={} valid={:?} tokens={}", + input.frame_bucket(), + input.clips(), + input.frame_lengths(), + tokens.len() + ); + + let mlx = run_gemma3n_audio_mlx_diagnostics( + &model_dir, + input.mel(), + input.valid_mask(), + input.frame_bucket(), + input.clips(), + &tokens, + AUDIO_TOKEN_ID, + CONTEXT_CAPACITY, + )?; + mlxcel_core::memory::clear_cache(); + eprintln!( + "gemma3n-audio-reference: MLX audio-only oracle complete projected_lengths={:?}", + mlx.projected_lengths + ); + + let mut iree = Gemma3nAudioIreeRuntime::load_audio_only_diagnostic( + &model_dir, + &device, + CONTEXT_CAPACITY, + input.frame_bucket(), + input.clips(), + )?; + let iree = iree.invoke(&input, &tokens, &rows)?; + eprintln!( + "gemma3n-audio-reference: IREE audio-only path complete projected_lengths={:?}", + iree.projected_lengths() + ); + if mlx.projected_lengths != iree.projected_lengths() { + return Err(format!( + "projected length mismatch: MLX={:?} IREE={:?}", + mlx.projected_lengths, + iree.projected_lengths() + )); + } + + let stage_values = [ + ("encoded_reduced", mlx.encoded_reduced.as_slice()), + ("soft_norm", mlx.soft_norm.as_slice()), + ("soft_linear", mlx.soft_linear.as_slice()), + ("soft_post_norm", mlx.soft_post_norm.as_slice()), + ("hard_embedding", mlx.hard_embedding.as_slice()), + ("hard_norm", mlx.hard_norm.as_slice()), + ("hard_linear", mlx.hard_linear.as_slice()), + ("hard_post_norm", mlx.hard_post_norm.as_slice()), + ]; + let mut names = Vec::with_capacity(stage_values.len() + 4); + let mut stages = Vec::with_capacity(stage_values.len() + 4); + for (name, reference) in stage_values { + let actual = iree + .diagnostic_stage(name) + .ok_or_else(|| format!("IREE diagnostic output is missing stage {name}"))?; + names.push(name); + stages.push(compare(name, reference, actual)?); + } + names.extend([ + "projected_audio", + "hard_audio", + "merged_embeddings", + "dense_ple", + ]); + stages.extend([ + compare( + "projected_audio", + &mlx.projected_audio, + iree.projected_audio(), + )?, + compare("hard_audio", &mlx.hard_audio, iree.hard_audio())?, + compare("merged_embeddings", &mlx.embeddings, iree.embeddings())?, + compare("dense_ple", &mlx.dense_ple, iree.dense_ple())?, + ]); + let first = stages + .iter() + .enumerate() + .find(|(_, diff)| diff.max_abs > 5e-3 || diff.rms > 5e-4); + match first { + Some((index, diff)) => Err(format!( + "first divergent stage={} max_abs={:.9e} rms={:.9e} max_index={}", + names[index], diff.max_abs, diff.rms, diff.max_index + )), + None => { + println!("Gemma3n audio-only MLX↔IREE intermediate gate PASS"); + Ok(()) + } + } +} diff --git a/src/audio/gemma3n/diagnostics.rs b/src/audio/gemma3n/diagnostics.rs new file mode 100644 index 000000000..a5885ff83 --- /dev/null +++ b/src/audio/gemma3n/diagnostics.rs @@ -0,0 +1,334 @@ +// 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. + +//! Audio-only MLX oracle used by the opt-in Gemma3n XLA first-divergence gate. + +use std::path::Path; + +use mlxcel_core::layers::{RMSNorm, UnifiedEmbedding, UnifiedLinear}; +use mlxcel_core::weights::WeightMap; +use mlxcel_core::{MlxArray, UniquePtr}; + +use super::{GEMMA3N_AUDIO_SOFT_TOKENS, Gemma3nAudioConfig, Gemma3nAudioEncoder}; +use crate::models::gemma3n::{Gemma3nAudioEmbedder, ModelArgs, TextConfig}; + +#[derive(Debug, Clone)] +pub struct Gemma3nAudioMlxDiagnosticOutput { + pub encoded_reduced: Vec, + pub soft_norm: Vec, + pub soft_linear: Vec, + pub soft_post_norm: Vec, + pub hard_embedding: Vec, + pub hard_norm: Vec, + pub hard_linear: Vec, + pub hard_post_norm: Vec, + pub projected_audio: Vec, + pub hard_audio: Vec, + pub embeddings: Vec, + pub dense_ple: Vec, + pub projected_lengths: Vec, + pub hidden_size: usize, + pub layers: usize, + pub hidden_per_layer: usize, +} + +fn canonical_name(name: &str) -> &str { + name.strip_prefix("model.").unwrap_or(name) +} + +fn keep_diagnostic_weight(name: &str) -> bool { + let name = canonical_name(name); + name.starts_with("audio_tower.") + || name.starts_with("embed_audio.") + || [ + "language_model.embed_tokens.", + "language_model.embed_tokens_per_layer.", + "language_model.per_layer_model_projection.", + "language_model.per_layer_projection_norm.", + "language_model.model.embed_tokens.", + "language_model.model.embed_tokens_per_layer.", + "language_model.model.per_layer_model_projection.", + "language_model.model.per_layer_projection_norm.", + ] + .iter() + .any(|prefix| name.starts_with(prefix)) +} + +fn load_diagnostic_weights(model_dir: &Path) -> Result { + let raw = + mlxcel_core::weights::load_weights_from_dir_filtered(model_dir, keep_diagnostic_weight)?; + let mut weights = WeightMap::new(); + for (name, value) in raw { + weights.insert(canonical_name(&name).to_string(), mlxcel_core::copy(&value)); + } + Ok(weights) +} + +fn read_config(model_dir: &Path) -> Result<(TextConfig, Gemma3nAudioConfig), String> { + let path = model_dir.join("config.json"); + let text = std::fs::read_to_string(&path) + .map_err(|error| format!("read {}: {error}", path.display()))?; + let root: ModelArgs = serde_json::from_str(&text) + .map_err(|error| format!("parse Gemma3n config {}: {error}", path.display()))?; + let value: serde_json::Value = serde_json::from_str(&text) + .map_err(|error| format!("parse Gemma3n config {}: {error}", path.display()))?; + let audio = value + .get("audio_config") + .ok_or_else(|| format!("{} has no audio_config", path.display())) + .and_then(|value| { + serde_json::from_value(value.clone()) + .map_err(|error| format!("parse Gemma3n audio_config: {error}")) + })?; + Ok((root.text_args(), audio)) +} + +fn array_f32(array: &MlxArray) -> Vec { + let values = mlxcel_core::astype(array, mlxcel_core::dtype::FLOAT32); + mlxcel_core::eval(&values); + mlxcel_core::array_to_raw_bytes(&values) + .chunks_exact(std::mem::size_of::()) + .map(|bytes| f32::from_ne_bytes(bytes.try_into().expect("f32 byte width"))) + .collect() +} + +fn copy_weight(weights: &WeightMap, key: &str) -> Result, String> { + weights + .get(key) + .map(|weight| mlxcel_core::copy(weight)) + .ok_or_else(|| format!("Gemma3n diagnostic weight not found: {key}")) +} + +fn language_prefix(weights: &WeightMap) -> &'static str { + if weights.contains_key("language_model.model.embed_tokens.weight") { + "language_model.model" + } else { + "language_model" + } +} + +#[allow(clippy::too_many_arguments)] +pub fn run_gemma3n_audio_mlx_diagnostics( + model_dir: &Path, + mel: &[f32], + valid_mask: &[u8], + frame_bucket: usize, + clips: usize, + token_ids: &[i32], + audio_token_id: i32, + context_capacity: usize, +) -> Result { + if token_ids.is_empty() || token_ids.len() > context_capacity { + return Err(format!( + "Gemma3n diagnostic token length {} is outside 1..={context_capacity}", + token_ids.len() + )); + } + let (text, audio) = read_config(model_dir)?; + let expected_mel = clips + .checked_mul(frame_bucket) + .and_then(|value| value.checked_mul(audio.input_feat_size)) + .ok_or_else(|| "Gemma3n diagnostic mel size overflow".to_string())?; + if mel.len() != expected_mel || valid_mask.len() != clips * frame_bucket { + return Err(format!( + "Gemma3n diagnostic input lengths {}/{} do not match clips/bucket/features {clips}/{frame_bucket}/{}", + mel.len(), + valid_mask.len(), + audio.input_feat_size + )); + } + + let weights = load_diagnostic_weights(model_dir)?; + let group_size = text + .quantization + .as_ref() + .map(|quant| quant.group_size as i32) + .unwrap_or(64); + let bits = text + .quantization + .as_ref() + .map(|quant| quant.bits as i32) + .unwrap_or(4); + let tower = + Gemma3nAudioEncoder::from_weights(&weights, "audio_tower", &audio, group_size, bits)?; + let embed_audio = Gemma3nAudioEmbedder::from_weights( + &weights, + "embed_audio", + &audio, + text.hidden_size, + group_size, + bits, + )?; + + let mel = mlxcel_core::from_slice_f32( + mel, + &[ + clips as i32, + frame_bucket as i32, + audio.input_feat_size as i32, + ], + ); + let invalid = valid_mask + .iter() + .map(|valid| i32::from(*valid == 0)) + .collect::>(); + let invalid = mlxcel_core::from_slice_i32(&invalid, &[clips as i32, frame_bucket as i32]); + let invalid = mlxcel_core::astype(&invalid, mlxcel_core::dtype::BOOL); + let (encoded, encoded_invalid) = tower.forward(&mel, &invalid)?; + let encoded_shape = mlxcel_core::array_shape(&encoded); + let projected_lengths = (0..clips) + .map(|clip| { + let row = mlxcel_core::slice( + &encoded_invalid, + &[clip as i32, 0], + &[(clip + 1) as i32, encoded_shape[1]], + ); + let valid = mlxcel_core::logical_not(&row); + let valid = mlxcel_core::astype(&valid, mlxcel_core::dtype::INT32); + mlxcel_core::item_i32(&mlxcel_core::sum_all(&valid)) as usize + }) + .collect::>(); + + let padding = embed_audio.padding_embedding(); + let (soft_norm, soft_linear, soft_post_norm) = + embed_audio.diagnostic_soft_projection_stages(&encoded); + let mut projected = mlxcel_core::copy(&soft_post_norm); + let projected_shape = mlxcel_core::array_shape(&projected); + let invalid = mlxcel_core::reshape( + &encoded_invalid, + &[projected_shape[0], projected_shape[1], 1], + ); + projected = mlxcel_core::where_cond(&invalid, &padding, &projected); + let missing = GEMMA3N_AUDIO_SOFT_TOKENS as i32 - projected_shape[1]; + if missing < 0 { + return Err(format!( + "Gemma3n MLX diagnostic projected {} rows, exceeding {}", + projected_shape[1], GEMMA3N_AUDIO_SOFT_TOKENS + )); + } + if missing > 0 { + let tail = + mlxcel_core::broadcast_to(&padding, &[clips as i32, missing, text.hidden_size as i32]); + projected = mlxcel_core::concatenate(&projected, &tail, 1); + } + let (hard_embedding, hard_norm, hard_linear, hard_post_norm) = + embed_audio.diagnostic_hard_projection_stages(); + + let language = language_prefix(&weights); + let text_embedding = UnifiedEmbedding::from_weights( + &weights, + &format!("{language}.embed_tokens"), + group_size, + bits, + )?; + let token_ple = UnifiedEmbedding::from_weights( + &weights, + &format!("{language}.embed_tokens_per_layer"), + group_size, + bits, + )?; + let model_projection = UnifiedLinear::from_weights( + &weights, + &format!("{language}.per_layer_model_projection"), + group_size, + bits, + )?; + let projection_norm = RMSNorm::new( + copy_weight( + &weights, + &format!("{language}.per_layer_projection_norm.weight"), + )?, + text.rms_norm_eps, + ); + + let mut padded_tokens = token_ids.to_vec(); + padded_tokens.resize(context_capacity, 0); + let token_array = mlxcel_core::from_slice_i32(&padded_tokens, &[1, context_capacity as i32]); + let scaled_text = mlxcel_core::multiply_scalar( + &text_embedding.forward(&token_array), + (text.hidden_size as f32).sqrt(), + ); + let hard_merged = embed_audio.merge_hard_tokens(&token_array, &scaled_text); + let projected = mlxcel_core::reshape( + &projected, + &[ + clips as i32 * GEMMA3N_AUDIO_SOFT_TOKENS as i32, + text.hidden_size as i32, + ], + ); + let merged = + crate::vision::merge::merge_llava(audio_token_id, &projected, &hard_merged, &token_array) + .inputs_embeds; + + let per_layer_limit = + mlxcel_core::from_slice_i32(&[text.vocab_size_per_layer_input as i32], &[1]); + let per_layer_mask = mlxcel_core::less(&token_array, &per_layer_limit); + let per_layer_tokens = mlxcel_core::where_cond( + &per_layer_mask, + &token_array, + &mlxcel_core::zeros(&[1, context_capacity as i32], mlxcel_core::dtype::INT32), + ); + let token_ple = mlxcel_core::multiply_scalar( + &token_ple.forward(&per_layer_tokens), + (text.hidden_size_per_layer_input as f32).sqrt(), + ); + let token_ple = mlxcel_core::reshape( + &token_ple, + &[ + context_capacity as i32, + text.num_hidden_layers as i32, + text.hidden_size_per_layer_input as i32, + ], + ); + let projected_ple = model_projection.forward(&merged); + let projected_ple = + mlxcel_core::multiply_scalar(&projected_ple, (text.hidden_size as f32).powf(-0.5)); + let projected_ple = mlxcel_core::reshape( + &projected_ple, + &[ + context_capacity as i32, + text.num_hidden_layers as i32, + text.hidden_size_per_layer_input as i32, + ], + ); + let projected_ple = projection_norm.forward(&projected_ple); + let dense_ple = mlxcel_core::multiply_scalar( + &mlxcel_core::add(&projected_ple, &token_ple), + std::f32::consts::FRAC_1_SQRT_2, + ); + + let mut embeddings = array_f32(&merged); + embeddings[token_ids.len() * text.hidden_size..].fill(0.0); + let mut dense_ple = array_f32(&dense_ple); + dense_ple[token_ids.len() * text.num_hidden_layers * text.hidden_size_per_layer_input..] + .fill(0.0); + + Ok(Gemma3nAudioMlxDiagnosticOutput { + encoded_reduced: array_f32(&encoded), + soft_norm: array_f32(&soft_norm), + soft_linear: array_f32(&soft_linear), + soft_post_norm: array_f32(&soft_post_norm), + hard_embedding: array_f32(&hard_embedding), + hard_norm: array_f32(&hard_norm), + hard_linear: array_f32(&hard_linear), + hard_post_norm: array_f32(&hard_post_norm), + projected_audio: array_f32(&projected), + hard_audio: array_f32(&hard_post_norm), + embeddings, + dense_ple, + projected_lengths, + hidden_size: text.hidden_size, + layers: text.num_hidden_layers, + hidden_per_layer: text.hidden_size_per_layer_input, + }) +} diff --git a/src/audio/gemma3n/mod.rs b/src/audio/gemma3n/mod.rs index 12ded4fd9..331613067 100644 --- a/src/audio/gemma3n/mod.rs +++ b/src/audio/gemma3n/mod.rs @@ -20,6 +20,8 @@ mod attention; mod config; +#[cfg(feature = "xla-diagnostics")] +mod diagnostics; mod encoder; mod feature_extractor; #[cfg(feature = "xla-backend")] @@ -86,6 +88,8 @@ pub(crate) fn checked_unified_linear( } pub use config::Gemma3nAudioConfig; +#[cfg(feature = "xla-diagnostics")] +pub use diagnostics::{Gemma3nAudioMlxDiagnosticOutput, run_gemma3n_audio_mlx_diagnostics}; pub use encoder::Gemma3nAudioEncoder; pub use feature_extractor::{ GEMMA3N_AUDIO_SOFT_TOKENS, GEMMA3N_MAX_SAMPLES, GEMMA3N_SAMPLE_RATE, Gemma3nAudioFeatureBatch, diff --git a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit.rs b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit.rs index 814dac1f3..7e4412fea 100644 --- a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit.rs +++ b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit.rs @@ -55,6 +55,8 @@ struct EncodeOutputs { projected_audio: Val, hard_audio: Val, projected_lengths: Val, + #[cfg(feature = "diagnostics")] + diagnostics: Vec<(String, Val)>, } struct MergeOutputs { @@ -138,6 +140,7 @@ fn project_audio( text: &Gemma3nConfig, encoded: &Val, valid: &Val, + trace: &mut Trace, ) -> (Val, Val) { let soft = rms_norm( b, @@ -146,28 +149,36 @@ fn project_audio( audio.rms_norm_eps, ); let soft = round_bf16(b, &soft); + trace.push("soft_norm", &soft); let soft = b.linear_last( &soft, args.weight("embed_audio.embedding_projection.weight"), ); let soft = round_bf16(b, &soft); + trace.push("soft_linear", &soft); let soft = rms_norm(b, &soft, None, audio.rms_norm_eps); let soft = round_bf16(b, &soft); + trace.push("soft_post_norm", &soft); + let hard_embedding = args.weight("embed_audio.embedding.weight"); + trace.push("hard_embedding", hard_embedding); let hard = rms_norm( b, - args.weight("embed_audio.embedding.weight"), + hard_embedding, Some(args.weight("embed_audio.hard_embedding_norm.weight")), audio.rms_norm_eps, ); let hard = round_bf16(b, &hard); + trace.push("hard_norm", &hard); let hard = b.linear_last( &hard, args.weight("embed_audio.embedding_projection.weight"), ); let hard = round_bf16(b, &hard); + trace.push("hard_linear", &hard); let hard = rms_norm(b, &hard, None, audio.rms_norm_eps); let hard = round_bf16(b, &hard); + trace.push("hard_post_norm", &hard); let padding = b.slice( &hard, &[(audio.vocab_size - 1, audio.vocab_size), (0, text.hidden)], @@ -298,13 +309,37 @@ fn build_encode( trace: &mut Trace, ) -> EncodeOutputs { let (encoded, valid, projected_lengths) = audio_encoder(b, args, audio, trace); - let (projected_audio, hard_audio) = project_audio(b, args, audio, text, &encoded, &valid); + let (projected_audio, hard_audio) = + project_audio(b, args, audio, text, &encoded, &valid, trace); trace.push("soft_projection", &projected_audio); trace.push("hard_projection", &hard_audio); + #[cfg(feature = "diagnostics")] + let diagnostics = [ + "encoded_reduced", + "soft_norm", + "soft_linear", + "soft_post_norm", + "hard_embedding", + "hard_norm", + "hard_linear", + "hard_post_norm", + ] + .into_iter() + .map(|name| { + trace + .stages + .iter() + .find(|(stage, _)| stage == name) + .cloned() + .expect("selected Gemma3n audio diagnostic stage must exist") + }) + .collect(); EncodeOutputs { projected_audio, hard_audio, projected_lengths, + #[cfg(feature = "diagnostics")] + diagnostics, } } @@ -340,15 +375,24 @@ fn render_encode(decls: &[Decl], b: &Builder, outputs: &EncodeOutputs) -> String let projected_ty = outputs.projected_audio.ty.render(); let hard_ty = outputs.hard_audio.ty.render(); let lengths_ty = outputs.projected_lengths.ty.render(); + let mut output_types = vec![projected_ty, hard_ty, lengths_ty]; + let mut output_values = vec![ + outputs.projected_audio.name.clone(), + outputs.hard_audio.name.clone(), + outputs.projected_lengths.name.clone(), + ]; + #[cfg(feature = "diagnostics")] + for (_, stage) in &outputs.diagnostics { + output_types.push(stage.ty.render()); + output_values.push(stage.name.clone()); + } + let output_types = output_types.join(", "); + let output_values = output_values.join(", "); format!( - "module @audio_encode {{\n func.func public @main({signature}) -> ({projected_ty}, \ - {hard_ty}, {lengths_ty}) {{\n{body} return {projected}, {hard}, {lengths} : \ - {projected_ty}, {hard_ty}, {lengths_ty}\n }}\n}}\n", + "module @audio_encode {{\n func.func public @main({signature}) -> ({output_types}) \ + {{\n{body} return {output_values} : {output_types}\n }}\n}}\n", signature = signature(decls), body = b.body(), - projected = outputs.projected_audio.name, - hard = outputs.hard_audio.name, - lengths = outputs.projected_lengths.name, ) } @@ -366,17 +410,6 @@ fn render_merge(decls: &[Decl], b: &Builder, outputs: &MergeOutputs) -> String { ) } -fn render_diagnostic(decls: &[Decl], b: &Builder, stage: &Val) -> String { - let stage_ty = stage.ty.render(); - format!( - "module @audio_diagnostic {{\n func.func public @main({signature}) -> {stage_ty} \ - {{\n{body} return {stage} : {stage_ty}\n }}\n}}\n", - signature = signature(decls), - body = b.body(), - stage = stage.name, - ) -} - fn validate_artifact( text: &Gemma3nConfig, audio: &Gemma3nXlaAudioConfig, @@ -429,38 +462,6 @@ pub(crate) fn emit_gemma3n_audio_merge_ple( Ok((render_merge(&decls, &builder, &outputs), trace.layout())) } -#[cfg(feature = "diagnostics")] -pub(crate) fn emit_gemma3n_audio_diagnostic( - text: &Gemma3nConfig, - audio: &Gemma3nXlaAudioConfig, - frame_bucket: usize, - clips: usize, - precision: Precision, - stage_name: &str, -) -> Result<(String, Gemma3nAudioDiagnosticStage), String> { - audio.artifact_identity( - frame_bucket, - clips, - text.hidden, - text.n_layers, - text.hidden_per_layer_input, - )?; - let (decls, args) = build_encoder_schema(audio, text, frame_bucket, clips)?; - let mut builder = Builder::new().with_precision(precision); - let mut trace = Trace::new(); - let _ = build_encode(&mut builder, &args, audio, text, &mut trace); - let (_, stage) = trace - .stages - .iter() - .find(|(name, _)| name == stage_name) - .ok_or_else(|| format!("unknown Gemma3n audio diagnostic stage `{stage_name}`"))?; - let metadata = Gemma3nAudioDiagnosticStage { - name: stage_name.to_string(), - shape: stage.ty.shape.clone(), - }; - Ok((render_diagnostic(&decls, &builder, stage), metadata)) -} - #[cfg(test)] #[path = "gemma3n_audio_emit_tests.rs"] mod tests; diff --git a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit_tests.rs b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit_tests.rs index 0d7bd32a7..0f8c7d4e0 100644 --- a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit_tests.rs +++ b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit_tests.rs @@ -39,7 +39,7 @@ fn split_audio_graphs_preserve_stage_and_weight_boundaries() { assert!(encode.contains("\"stablehlo.reduce_window\"")); assert!(!encode.contains("loc(\"audio.hard_embeddings\")")); assert!(!encode.contains("model.language_model.")); - assert_eq!(encode_layout.stages.len(), 3 + 12 * 5 + 3); + assert_eq!(encode_layout.stages.len(), 3 + 12 * 5 + 10); for name in [ "sscp_conv_0", "sscp_conv_1", @@ -47,6 +47,13 @@ fn split_audio_graphs_preserve_stage_and_weight_boundaries() { "conformer.0.feed_forward_start", "conformer.11.final_norm", "encoded_reduced", + "soft_norm", + "soft_linear", + "soft_post_norm", + "hard_embedding", + "hard_norm", + "hard_linear", + "hard_post_norm", "soft_projection", "hard_projection", ] { diff --git a/src/lib/mlxcel-xla/src/emitter/mod.rs b/src/lib/mlxcel-xla/src/emitter/mod.rs index 8f56cb9e7..dbd040805 100644 --- a/src/lib/mlxcel-xla/src/emitter/mod.rs +++ b/src/lib/mlxcel-xla/src/emitter/mod.rs @@ -81,10 +81,6 @@ pub(crate) use config::{ pub(crate) use gemma3n::{ Gemma3nConfig, Gemma3nLayerType, Gemma3nPleDType, Gemma3nPleMetadata, validate_gemma3n_ple, }; -#[cfg(feature = "diagnostics")] -pub(crate) use gemma3n_audio_emit::emit_gemma3n_audio_diagnostic; -#[cfg(feature = "diagnostics")] -pub use gemma3n_audio_emit::{Gemma3nAudioDiagnosticLayout, Gemma3nAudioDiagnosticStage}; #[allow(unused_imports)] pub(crate) use gemma3n_audio_emit::{emit_gemma3n_audio_encode, emit_gemma3n_audio_merge_ple}; #[allow(unused_imports)] diff --git a/src/lib/mlxcel-xla/src/gemma3n_audio_runtime.rs b/src/lib/mlxcel-xla/src/gemma3n_audio_runtime.rs index ef22cd4e2..df420b6ea 100644 --- a/src/lib/mlxcel-xla/src/gemma3n_audio_runtime.rs +++ b/src/lib/mlxcel-xla/src/gemma3n_audio_runtime.rs @@ -4,7 +4,10 @@ //! keeps the compiler at natural tensor boundaries while ensuring every model //! operation after host mel extraction still executes through IREE. +#[cfg(feature = "diagnostics")] +use std::collections::BTreeMap; use std::fs::File; +use std::io::Read; use std::path::{Path, PathBuf}; use std::process::Command; @@ -38,6 +41,12 @@ pub struct Gemma3nAudioGraphOutput { embeddings: Vec, dense_ple: Vec, projected_lengths: Vec, + #[cfg(feature = "diagnostics")] + projected_audio: Vec, + #[cfg(feature = "diagnostics")] + hard_audio: Vec, + #[cfg(feature = "diagnostics")] + diagnostics: BTreeMap>, context_capacity: usize, hidden_size: usize, layers: usize, @@ -60,6 +69,24 @@ impl Gemma3nAudioGraphOutput { &self.projected_lengths } + #[cfg(feature = "diagnostics")] + #[must_use] + pub fn projected_audio(&self) -> &[f32] { + &self.projected_audio + } + + #[cfg(feature = "diagnostics")] + #[must_use] + pub fn hard_audio(&self) -> &[f32] { + &self.hard_audio + } + + #[cfg(feature = "diagnostics")] + #[must_use] + pub fn diagnostic_stage(&self, name: &str) -> Option<&[f32]> { + self.diagnostics.get(name).map(Vec::as_slice) + } + #[must_use] pub fn embeddings_shape(&self) -> [usize; 2] { [self.context_capacity, self.hidden_size] @@ -155,12 +182,17 @@ pub struct Gemma3nAudioIreeRuntime { layers: usize, hidden_per_layer: usize, language_fingerprint: u64, + #[cfg(feature = "diagnostics")] + diagnostic_shapes: Vec<(String, Vec)>, } fn sha256_hex(bytes: &[u8]) -> String { - let digest = Sha256::digest(bytes); - let mut result = String::with_capacity(digest.len() * 2); - for byte in digest { + hex_bytes(&Sha256::digest(bytes)) +} + +fn hex_bytes(bytes: &[u8]) -> String { + let mut result = String::with_capacity(bytes.len() * 2); + for byte in bytes { use std::fmt::Write; write!(result, "{byte:02x}").expect("writing to String cannot fail"); } @@ -179,14 +211,44 @@ fn generation_identity(compiler: &Path, flags: &[&str], mlir: &str) -> Result Result { Ok(format!( - "compiler={};version={};flags={flags:?};mlir-sha256={}", + "compiler={};compiler_sha256={};version={version};flags={flags:?};mlir_sha256={}", compiler.display(), - String::from_utf8_lossy(&version.stdout).trim(), + sha256_file(compiler)?, sha256_hex(mlir.as_bytes()) )) } +fn sha256_file(path: &Path) -> Result { + let mut file = File::open(path).map_err(|error| format!("open {}: {error}", path.display()))?; + let mut digest = Sha256::new(); + let mut buffer = [0u8; 64 * 1024]; + loop { + let read = file + .read(&mut buffer) + .map_err(|error| format!("read {}: {error}", path.display()))?; + if read == 0 { + break; + } + digest.update(&buffer[..read]); + } + Ok(hex_bytes(&digest.finalize())) +} + fn checkpoint_alias(name: &str) -> Option { name.strip_prefix("model.language_model.") .map(|suffix| format!("language_model.model.{suffix}")) @@ -282,9 +344,23 @@ fn decode_weight( spec.name, spec.shape )); } + validate_finite_values(&format!("Gemma3n auxiliary weight {}", spec.name), &values)?; Ok(values) } +fn validate_finite_values(label: &str, values: &[f32]) -> Result<(), String> { + if let Some((index, value)) = values + .iter() + .enumerate() + .find(|(_, value)| !value.is_finite()) + { + return Err(format!( + "{label} contains non-finite value {value} at flat index {index}" + )); + } + Ok(()) +} + fn load_weights( model_dir: &Path, specs: &[Gemma3nAudioGraphWeightSpec], @@ -425,6 +501,25 @@ fn report_auxiliary_weights(label: &str, started: std::time::Instant, weights: & } impl Gemma3nAudioIreeRuntime { + #[cfg(feature = "diagnostics")] + pub fn load_audio_only_diagnostic( + model_dir: &Path, + device: &str, + context_capacity: usize, + frame_bucket: usize, + clips: usize, + ) -> Result { + const AUDIO_ONLY_DIAGNOSTIC_IDENTITY: u64 = 0x6175_6469_6f2d_6469; + Self::load( + model_dir, + device, + context_capacity, + frame_bucket, + clips, + AUDIO_ONLY_DIAGNOSTIC_IDENTITY, + ) + } + #[allow(clippy::too_many_arguments)] pub fn load( model_dir: &Path, @@ -445,7 +540,7 @@ impl Gemma3nAudioIreeRuntime { .map_err(|error| error.to_string())?; report_load_timing("audio-config-and-schema-validation", load_started); let emit_started = std::time::Instant::now(); - let (encode_mlir, _) = + let (encode_mlir, encode_layout) = emit_gemma3n_audio_encode(&text, &audio, frame_bucket, clips, Precision::Bf16)?; let (merge_mlir, _) = emit_gemma3n_audio_merge_ple(&text, &audio, clips, Precision::Bf16)?; report_load_timing("audio-stablehlo-emission", emit_started); @@ -554,6 +649,29 @@ impl Gemma3nAudioIreeRuntime { IreeAuxiliaryModule::load(device, &merge_vmfb, &merge_contract, merge_weights)?; report_load_timing("audio-merge-vmfb-loaded", merge_module_started); report_load_timing("audio-total", load_started); + #[cfg(feature = "diagnostics")] + let diagnostic_shapes = [ + "encoded_reduced", + "soft_norm", + "soft_linear", + "soft_post_norm", + "hard_embedding", + "hard_norm", + "hard_linear", + "hard_post_norm", + ] + .into_iter() + .map(|name| { + encode_layout + .stages + .iter() + .find(|stage| stage.name == name) + .map(|stage| (name.to_string(), stage.shape.clone())) + .ok_or_else(|| format!("Gemma3n audio diagnostic stage {name} is missing")) + }) + .collect::, _>>()?; + #[cfg(not(feature = "diagnostics"))] + let _ = encode_layout; Ok(Self { encode, merge_ple, @@ -565,6 +683,8 @@ impl Gemma3nAudioIreeRuntime { layers: text.n_layers, hidden_per_layer: text.hidden_per_layer_input, language_fingerprint, + #[cfg(feature = "diagnostics")] + diagnostic_shapes, }) } @@ -611,19 +731,21 @@ impl Gemma3nAudioIreeRuntime { let mut hard = vec![0.0f32; self.audio.vocab_size * self.hidden_size]; let mut lengths = vec![0i32; self.clips]; let encode_started = std::time::Instant::now(); + let encode_inputs = [ + AuxiliaryInput { + bytes: f32_bytes(input.mel()), + dtype: AuxiliaryTensorDType::Float32, + shape: &[self.clips, self.frame_bucket, self.audio.input_feat_size], + }, + AuxiliaryInput { + bytes: input.valid_mask(), + dtype: AuxiliaryTensorDType::Bool, + shape: &[self.clips, self.frame_bucket], + }, + ]; + #[cfg(not(feature = "diagnostics"))] self.encode.invoke( - &[ - AuxiliaryInput { - bytes: f32_bytes(input.mel()), - dtype: AuxiliaryTensorDType::Float32, - shape: &[self.clips, self.frame_bucket, self.audio.input_feat_size], - }, - AuxiliaryInput { - bytes: input.valid_mask(), - dtype: AuxiliaryTensorDType::Bool, - shape: &[self.clips, self.frame_bucket], - }, - ], + &encode_inputs, &mut [ AuxiliaryOutput { bytes: f32_bytes_mut(&mut projected), @@ -642,6 +764,48 @@ impl Gemma3nAudioIreeRuntime { }, ], )?; + #[cfg(feature = "diagnostics")] + let diagnostics = { + let diagnostic_shapes = self.diagnostic_shapes.clone(); + let mut diagnostic_values = diagnostic_shapes + .iter() + .map(|(_, shape)| vec![0.0f32; shape.iter().product()]) + .collect::>(); + let projected_shape = [projected_rows, self.hidden_size]; + let hard_shape = [self.audio.vocab_size, self.hidden_size]; + let lengths_shape = [self.clips]; + let mut outputs = vec![ + AuxiliaryOutput { + bytes: f32_bytes_mut(&mut projected), + dtype: AuxiliaryTensorDType::Float32, + shape: &projected_shape, + }, + AuxiliaryOutput { + bytes: f32_bytes_mut(&mut hard), + dtype: AuxiliaryTensorDType::Float32, + shape: &hard_shape, + }, + AuxiliaryOutput { + bytes: i32_bytes_mut(&mut lengths), + dtype: AuxiliaryTensorDType::Int32, + shape: &lengths_shape, + }, + ]; + for (values, (_, shape)) in diagnostic_values.iter_mut().zip(diagnostic_shapes.iter()) { + outputs.push(AuxiliaryOutput { + bytes: f32_bytes_mut(values), + dtype: AuxiliaryTensorDType::Float32, + shape, + }); + } + self.encode.invoke(&encode_inputs, &mut outputs)?; + drop(outputs); + diagnostic_shapes + .into_iter() + .map(|(name, _)| name) + .zip(diagnostic_values) + .collect::>() + }; report_load_timing("audio-encoder-invoke", encode_started); let projected_lengths = lengths .into_iter() @@ -713,6 +877,12 @@ impl Gemma3nAudioIreeRuntime { embeddings, dense_ple, projected_lengths, + #[cfg(feature = "diagnostics")] + projected_audio: projected, + #[cfg(feature = "diagnostics")] + hard_audio: hard, + #[cfg(feature = "diagnostics")] + diagnostics, context_capacity: self.context_capacity, hidden_size: self.hidden_size, layers: self.layers, @@ -742,6 +912,18 @@ impl Gemma3nAudioIreeRuntime { #[cfg(test)] mod tests { use super::*; + use crate::aux_manifest::auxiliary_manifest_path; + + fn temp_path(tag: &str) -> PathBuf { + let nonce = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .expect("system clock must be after the Unix epoch") + .as_nanos(); + std::env::temp_dir().join(format!( + "mlxcel-xla-audio-{tag}-{}-{nonce}", + std::process::id() + )) + } #[test] fn output_moves_into_request_scoped_dense_ple_owner() { @@ -753,6 +935,12 @@ mod tests { embeddings: vec![0.25; context_capacity * 2], dense_ple: vec![0.5; context_capacity * 2], projected_lengths: vec![GEMMA3N_AUDIO_SOFT_TOKENS], + #[cfg(feature = "diagnostics")] + projected_audio: Vec::new(), + #[cfg(feature = "diagnostics")] + hard_audio: Vec::new(), + #[cfg(feature = "diagnostics")] + diagnostics: BTreeMap::new(), context_capacity, hidden_size: 2, layers: 1, @@ -783,4 +971,70 @@ mod tests { .expect("zero identity must fail"); assert!(error.contains("verified #876 language bundle")); } + + #[test] + fn compiler_binary_replacement_at_same_path_and_version_rebuilds_cache() { + let compiler = temp_path("compiler"); + let vmfb = temp_path("compiler-digest").with_extension("vmfb"); + let weights = vec![AuxiliaryWeight { + name: "weight".to_string(), + storage: AuxiliaryWeightStorage::Bytes(1.0f32.to_ne_bytes().to_vec()), + dtype: AuxiliaryWeightDType::Float32, + shape: vec![1], + }]; + std::fs::write(&compiler, b"compiler-build-a").unwrap(); + let first_generation = + generation_identity_from_version(&compiler, &["--target=cuda"], "mlir", "v1").unwrap(); + assert!(first_generation.contains(&sha256_hex(b"compiler-build-a"))); + let first_contract = + AuxiliaryArtifactContract::new("audio.main", "config=v1", &first_generation).unwrap(); + let mut compile_count = 0usize; + ensure_qualified_auxiliary_artifact(&vmfb, &first_contract, &weights, |temporary| { + compile_count += 1; + std::fs::write(temporary, b"vmfb-a") + .map_err(|error| format!("write test VMFB: {error}")) + }) + .unwrap(); + + std::fs::write(&compiler, b"compiler-build-b").unwrap(); + let second_generation = + generation_identity_from_version(&compiler, &["--target=cuda"], "mlir", "v1").unwrap(); + assert_ne!(first_generation, second_generation); + let second_contract = + AuxiliaryArtifactContract::new("audio.main", "config=v1", second_generation).unwrap(); + ensure_qualified_auxiliary_artifact(&vmfb, &second_contract, &weights, |temporary| { + compile_count += 1; + std::fs::write(temporary, b"vmfb-b") + .map_err(|error| format!("write test VMFB: {error}")) + }) + .unwrap(); + assert_eq!(compile_count, 2); + assert_eq!(std::fs::read(&vmfb).unwrap(), b"vmfb-b"); + + std::fs::remove_file(auxiliary_manifest_path(&vmfb)).ok(); + std::fs::remove_file(vmfb).ok(); + std::fs::remove_file(compiler).ok(); + } + + #[test] + fn decoded_audio_weights_reject_non_finite_values_before_native_upload() { + let bf16_nan = bf16_to_f32(&0x7fc0u16.to_le_bytes()); + let f16_nan = f16_to_f32(&0x7e00u16.to_le_bytes()); + let f32_nan = f32_le_to_f32(&f32::NAN.to_le_bytes()); + let q4_infinite = + dequantize_affine_bf16_fused(&[0; 4], &0x7f80u16.to_le_bytes(), &[0; 2], 1, 1, 4, 8) + .unwrap(); + + for (dtype, values) in [ + ("BF16", bf16_nan), + ("F16", f16_nan), + ("F32", f32_nan), + ("Q4", q4_infinite), + ] { + let error = validate_finite_values(dtype, &values).unwrap_err(); + assert!(error.contains(dtype)); + assert!(error.contains("non-finite")); + assert!(error.contains("flat index 0")); + } + } } diff --git a/src/models/gemma3n.rs b/src/models/gemma3n.rs index d481bcc9f..f45828a50 100644 --- a/src/models/gemma3n.rs +++ b/src/models/gemma3n.rs @@ -5544,6 +5544,39 @@ impl Gemma3nAudioEmbedder { let embedded = self.embedding.forward(&token); self.project(&self.hard_embedding_norm.forward(&embedded)) } + + #[cfg(feature = "xla-diagnostics")] + pub fn diagnostic_soft_projection_stages( + &self, + inputs: &MlxArray, + ) -> ( + UniquePtr, + UniquePtr, + UniquePtr, + ) { + let normalized = self.soft_embedding_norm.forward(inputs); + let linear = self.embedding_projection.forward(&normalized); + let projected = self.post_projection_norm.forward(&linear); + (normalized, linear, projected) + } + + #[cfg(feature = "xla-diagnostics")] + pub fn diagnostic_hard_projection_stages( + &self, + ) -> ( + UniquePtr, + UniquePtr, + UniquePtr, + UniquePtr, + ) { + let tokens = (0..self.vocab_size as i32).collect::>(); + let tokens = mlxcel_core::from_slice_i32(&tokens, &[self.vocab_size as i32]); + let embedded = self.embedding.forward(&tokens); + let normalized = self.hard_embedding_norm.forward(&embedded); + let linear = self.embedding_projection.forward(&normalized); + let projected = self.post_projection_norm.forward(&linear); + (embedded, normalized, linear, projected) + } } #[cfg(test)] From 367ce167572b2874f8cc0bf9f0416357c0640cd6 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 24 Jul 2026 17:46:10 +0900 Subject: [PATCH 07/19] fix(xla): reconcile audio rebase contracts --- src/commands/generate.rs | 1 - src/lib/mlxcel-xla/src/emitter/builder.rs | 129 ---------------------- src/lib/mlxcel-xla/src/phi4_audio.rs | 8 +- 3 files changed, 4 insertions(+), 134 deletions(-) diff --git a/src/commands/generate.rs b/src/commands/generate.rs index 78423b72f..6b75f36e0 100644 --- a/src/commands/generate.rs +++ b/src/commands/generate.rs @@ -2046,7 +2046,6 @@ fn run_generate_once(mut args: GenerateArgs) -> Result<()> { args.generation.max_tokens, kv_cache_mode, token_bias, - &tokenizer, )?; let generated_text = decode_generated_text(&tokenizer, &prompt_tokens, &generated_tokens); let visible = filter_reasoning_for_display( diff --git a/src/lib/mlxcel-xla/src/emitter/builder.rs b/src/lib/mlxcel-xla/src/emitter/builder.rs index a5a285b8e..b8898cfab 100644 --- a/src/lib/mlxcel-xla/src/emitter/builder.rs +++ b/src/lib/mlxcel-xla/src/emitter/builder.rs @@ -727,23 +727,6 @@ impl Builder { } } - pub fn const_tensor_i32(&mut self, data: &[i32], shape: Vec) -> Val { - assert_eq!(data.len(), shape.iter().product::()); - assert_eq!(shape.len(), 1); - let ty = Ty::new(shape, "i32"); - let values = data - .iter() - .map(i32::to_string) - .collect::>() - .join(", "); - let r = self.fresh(); - self.line(format!( - "{r} = stablehlo.constant dense<[{values}]> : {}", - ty.render() - )); - Val { name: r, ty } - } - /// Dense f32 constant tensor from raw data (row-major), emitted as hex blob. pub fn const_tensor_f32(&mut self, data: &[f32], shape: Vec) -> Val { let ty = Ty::f32(shape); @@ -1039,35 +1022,6 @@ impl Builder { Val { name: r, ty } } - /// Zero-interior static padding. - pub fn pad(&mut self, x: &Val, value: &Val, low: &[usize], high: &[usize]) -> Val { - assert_eq!(low.len(), x.ty.shape.len()); - assert_eq!(high.len(), x.ty.shape.len()); - let shape = - x.ty.shape - .iter() - .zip(low) - .zip(high) - .map(|((&dim, &before), &after)| before + dim + after) - .collect(); - let ty = Ty::new(shape, x.ty.elt); - let interior = vec![0; x.ty.shape.len()]; - let r = self.fresh(); - self.line(format!( - "{} = stablehlo.pad {}, {}, low = {}, high = {}, interior = {} : ({}, {}) -> {}", - r, - x.name, - value.name, - dims_list(low), - dims_list(high), - dims_list(&interior), - x.ty.render(), - value.ty.render(), - ty.render() - )); - Val { name: r, ty } - } - pub fn concatenate(&mut self, a: &Val, b: &Val, dim: usize) -> Val { let mut shape = a.ty.shape.clone(); shape[dim] = a.ty.shape[dim] + b.ty.shape[dim]; @@ -1350,9 +1304,6 @@ impl Builder { pub fn divide(&mut self, a: &Val, b: &Val) -> Val { self.binary("divide", a, b) } - pub fn maximum(&mut self, a: &Val, b: &Val) -> Val { - self.binary("maximum", a, b) - } pub fn minimum(&mut self, a: &Val, b: &Val) -> Val { self.binary("minimum", a, b) } @@ -1759,86 +1710,6 @@ impl Builder { shape.push(w.ty.shape[0]); self.reshape(&projected, shape) } - - /// Channels-last NWC/KIO or NHWC/HWIO convolution. - pub fn convolution( - &mut self, - input: &Val, - kernel: &Val, - strides: &[usize], - padding: &[(usize, usize)], - feature_groups: usize, - ) -> Val { - let rank = input.ty.shape.len(); - assert!(matches!(rank, 3 | 4)); - assert_eq!(kernel.ty.shape.len(), rank); - assert_eq!(input.ty.elt, kernel.ty.elt); - assert_eq!(input.ty.elt, "f32"); - assert_eq!(strides.len(), rank - 2); - assert_eq!(padding.len(), rank - 2); - assert!(strides.iter().all(|stride| *stride > 0)); - assert!(feature_groups > 0); - assert_eq!( - input.ty.shape[rank - 1], - kernel.ty.shape[rank - 2] * feature_groups - ); - assert!(kernel.ty.shape[rank - 1].is_multiple_of(feature_groups)); - let batch = input.ty.shape[0]; - let mut shape = vec![batch]; - for axis in 0..rank - 2 { - let padded = input.ty.shape[axis + 1] + padding[axis].0 + padding[axis].1; - assert!(padded >= kernel.ty.shape[axis]); - shape.push((padded - kernel.ty.shape[axis]) / strides[axis] + 1); - } - shape.push(kernel.ty.shape[rank - 1]); - let output_shape = shape; - let lhs_demoted; - let rhs_demoted; - let (input, kernel) = match self.precision.dot_elt() { - Some(elt) => { - lhs_demoted = self.convert(input, elt); - rhs_demoted = self.convert(kernel, elt); - (&lhs_demoted, &rhs_demoted) - } - None => (input, kernel), - }; - let op_ty = Ty::new(output_shape, input.ty.elt); - let r = self.fresh(); - let strides = dims_list(strides); - let padding = padding - .iter() - .map(|(low, high)| format!("[{low}, {high}]")) - .collect::>() - .join(", "); - let dimension_numbers = if rank == 3 { - "[b, 0, f]x[0, i, o]->[b, 0, f]" - } else { - "[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]" - }; - self.line(format!( - "{} = stablehlo.convolution({}, {}) dim_numbers = {}, \ - window = {{stride = {}, pad = [{}]}} \ - {{batch_group_count = 1 : i64, feature_group_count = {} : i64, \ - precision_config = [#stablehlo, #stablehlo]}} : \ - ({}, {}) -> {}", - r, - input.name, - kernel.name, - dimension_numbers, - strides, - padding, - feature_groups, - input.ty.render(), - kernel.ty.render(), - op_ty.render() - )); - let result = Val { name: r, ty: op_ty }; - if result.ty.elt == "f32" { - result - } else { - self.convert(&result, "f32") - } - } } #[cfg(test)] diff --git a/src/lib/mlxcel-xla/src/phi4_audio.rs b/src/lib/mlxcel-xla/src/phi4_audio.rs index 3c1ebae3d..6199a5efd 100644 --- a/src/lib/mlxcel-xla/src/phi4_audio.rs +++ b/src/lib/mlxcel-xla/src/phi4_audio.rs @@ -31,7 +31,7 @@ use sha2::{Digest, Sha256}; use crate::aux::{ AuxiliaryInput, AuxiliaryOutput, AuxiliaryTensorDType, AuxiliaryWeight, AuxiliaryWeightDType, - IreeAuxiliaryModule, + AuxiliaryWeightStorage, IreeAuxiliaryModule, }; use crate::aux_manifest::{AuxiliaryArtifactContract, ensure_qualified_auxiliary_artifact}; use crate::emitter::{ @@ -306,7 +306,7 @@ fn load_audio_weights( validate_finite_values(&format!("Phi4MM audio weight `{}`", spec.name), &values)?; loaded[index] = Some(AuxiliaryWeight { name: spec.name.clone(), - bytes: f32_bytes(&values), + storage: AuxiliaryWeightStorage::Bytes(f32_bytes(&values)), dtype: AuxiliaryWeightDType::Float32, shape: spec.shape.clone(), }); @@ -752,7 +752,7 @@ mod tests { let vmfb = temporary_audio_path("vmfb"); let weights = vec![AuxiliaryWeight { name: "model.embed_tokens_extend.audio_embed.weight".to_string(), - bytes: 1.0f32.to_ne_bytes().to_vec(), + storage: AuxiliaryWeightStorage::Bytes(1.0f32.to_ne_bytes().to_vec()), dtype: AuxiliaryWeightDType::Float32, shape: vec![1], }]; @@ -805,7 +805,7 @@ mod tests { let vmfb = temporary_audio_path("compiler-digest.vmfb"); let weights = vec![AuxiliaryWeight { name: "weight".to_string(), - bytes: 1.0f32.to_ne_bytes().to_vec(), + storage: AuxiliaryWeightStorage::Bytes(1.0f32.to_ne_bytes().to_vec()), dtype: AuxiliaryWeightDType::Float32, shape: vec![1], }]; From b2015f6d7c7eb1f690cfc86b4a7b4a659313f11c Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 24 Jul 2026 20:25:43 +0900 Subject: [PATCH 08/19] fix(xla): match MLX BF16 affine dequantization MLX CUDA materializes the quantized product as BF16 before adding the BF16 affine bias, while the XLA loader previously fused both operations before a single rounding step. This changed roughly half of the pinned Gemma3n audio embedding and projection weights and explains the tower-independent hard-audio mismatch. Use sequential BF16 rounding in both Gemma3n language and audio weight loaders, rename the helper to reflect its semantics, and cover a real projection double-rounding counterexample. Validation: cargo test -p mlxcel-xla weights::tests --lib; cargo check -p mlxcel-xla --features diagnostics; cargo clippy -p mlxcel-xla --features diagnostics --lib --fix --allow-dirty --allow-staged; cargo fmt --all -- --check. Refs #878 --- .../mlxcel-xla/src/gemma3n_audio_runtime.rs | 17 ++++-- src/lib/mlxcel-xla/src/iree.rs | 4 +- src/lib/mlxcel-xla/src/lib.rs | 4 +- src/lib/mlxcel-xla/src/weights.rs | 57 ++++++++++++------- 4 files changed, 52 insertions(+), 30 deletions(-) diff --git a/src/lib/mlxcel-xla/src/gemma3n_audio_runtime.rs b/src/lib/mlxcel-xla/src/gemma3n_audio_runtime.rs index df420b6ea..ba465852e 100644 --- a/src/lib/mlxcel-xla/src/gemma3n_audio_runtime.rs +++ b/src/lib/mlxcel-xla/src/gemma3n_audio_runtime.rs @@ -25,7 +25,7 @@ use crate::emitter::{ emit_gemma3n_audio_merge_ple, gemma3n_audio_encoder_weights, gemma3n_audio_merge_weights, }; use crate::iree::{cached_vmfb_path, compile_one_to, iree_compile_bin, target_flags}; -use crate::weights::{bf16_to_f32, dequantize_affine_bf16_fused, f16_to_f32, f32_le_to_f32}; +use crate::weights::{bf16_to_f32, dequantize_affine_bf16_sequential, f16_to_f32, f32_le_to_f32}; use crate::{ GEMMA3N_AUDIO_MODALITY_FAMILY, GEMMA3N_AUDIO_SOFT_TOKENS, Gemma3nAudioInput, Gemma3nAudioPreparedPrefill, Gemma3nDensePle, Gemma3nPreparedPrefill, Gemma3nXlaAudioConfig, @@ -320,7 +320,7 @@ fn decode_weight( let logical_columns = packed[1] .checked_mul(32 / quant_bits) .ok_or_else(|| format!("{checkpoint_name} logical width overflows"))?; - let values = dequantize_affine_bf16_fused( + let values = dequantize_affine_bf16_sequential( tensor.data(), scales.data(), biases.data(), @@ -1021,9 +1021,16 @@ mod tests { let bf16_nan = bf16_to_f32(&0x7fc0u16.to_le_bytes()); let f16_nan = f16_to_f32(&0x7e00u16.to_le_bytes()); let f32_nan = f32_le_to_f32(&f32::NAN.to_le_bytes()); - let q4_infinite = - dequantize_affine_bf16_fused(&[0; 4], &0x7f80u16.to_le_bytes(), &[0; 2], 1, 1, 4, 8) - .unwrap(); + let q4_infinite = dequantize_affine_bf16_sequential( + &[0; 4], + &0x7f80u16.to_le_bytes(), + &[0; 2], + 1, + 1, + 4, + 8, + ) + .unwrap(); for (dtype, values) in [ ("BF16", bf16_nan), diff --git a/src/lib/mlxcel-xla/src/iree.rs b/src/lib/mlxcel-xla/src/iree.rs index cb20f8703..8202e0a8f 100644 --- a/src/lib/mlxcel-xla/src/iree.rs +++ b/src/lib/mlxcel-xla/src/iree.rs @@ -84,7 +84,7 @@ use crate::{DeepStackFeatures, DeepStackPreparedPrefill, Gemma3nDensePle, Gemma3 // and the #500 MoE expert bank (the stacked `switch_mlp` weights, dequantized with // `dequantize_affine_stacked`). Both are pure-Rust and unit-tested without `iree`. use crate::weights::{ - QuantPart, WeightSpec, bf16_to_f32, dequantize_affine, dequantize_affine_bf16_fused, + QuantPart, WeightSpec, bf16_to_f32, dequantize_affine, dequantize_affine_bf16_sequential, dequantize_affine_stacked, f16_to_f32, f32_le_to_f32, pack_f16, slice_rows, weight_specs, }; @@ -1618,7 +1618,7 @@ fn load_weights( "{prefix} uses F16 affine scales/biases, but Gemma3n requires BF16 affine metadata" )); } - dequantize_affine_bf16_fused( + dequantize_affine_bf16_sequential( t.data(), scales.data(), biases.data(), diff --git a/src/lib/mlxcel-xla/src/lib.rs b/src/lib/mlxcel-xla/src/lib.rs index 1534a92c1..99c288f2a 100644 --- a/src/lib/mlxcel-xla/src/lib.rs +++ b/src/lib/mlxcel-xla/src/lib.rs @@ -211,7 +211,9 @@ pub fn dequantize_gemma3n_affine_diagnostic( bits: usize, group_size: usize, ) -> Result, String> { - weights::dequantize_affine_bf16_fused(packed, scales, biases, out, in_packed, bits, group_size) + weights::dequantize_affine_bf16_sequential( + packed, scales, biases, out, in_packed, bits, group_size, + ) } #[cfg(feature = "diagnostics")] pub use emitter::{Gemma3nDiagnosticLayout, Gemma3nDiagnosticSegment}; diff --git a/src/lib/mlxcel-xla/src/weights.rs b/src/lib/mlxcel-xla/src/weights.rs index 22338ab7e..d28d30018 100644 --- a/src/lib/mlxcel-xla/src/weights.rs +++ b/src/lib/mlxcel-xla/src/weights.rs @@ -575,11 +575,14 @@ pub(crate) fn dequantize_affine( Ok(w) } -/// Gemma3n's MLX CUDA affine-dequant kernel evaluates `scale * q + bias` as a -/// fused expression and rounds the result once to the BF16 scales dtype. -/// Preserve those BF16 values in f32 carriers for the public IREE graph ABI. +/// Match MLX CUDA's affine dequantization for BF16 scales and biases. +/// +/// The CUDA kernel materializes `q * scale` as BF16 before adding the BF16 +/// zero point, so both operations round independently: +/// `bf16(bf16(q * scale) + bias)`. Preserve those values in f32 carriers for +/// the public IREE graph ABI. #[allow(clippy::too_many_arguments)] -pub(crate) fn dequantize_affine_bf16_fused( +pub(crate) fn dequantize_affine_bf16_sequential( packed: &[u8], scales: &[u8], biases: &[u8], @@ -623,19 +626,21 @@ pub(crate) fn dequantize_affine_bf16_fused( if bits == 4 { // Reuse one LUT bank across all output rows. The embedding table has // hundreds of thousands of rows, so allocating it inside this loop - // would erase most of the fused-arithmetic speedup. + // would erase most of the dequantization speedup. let mut luts = vec![[0.0f32; 16]; n_groups]; for output in 0..out { let row = &packed[output * in_packed * 4..(output + 1) * in_packed * 4]; let group_row = output * n_groups; let weight_row = output * in_; // A Gemma3n E2B group has 64 lanes but only 16 possible nibble - // values. Compute the fused BF16 affine result once per value and - // group, then decode the packed row through the tiny hot LUT. + // values. Compute the sequentially-rounded BF16 affine result once + // per value and group, then decode the packed row through the tiny + // hot LUT. for (group_index, lut) in luts.iter_mut().enumerate() { let group = group_row + group_index; for (quantized, value) in lut.iter_mut().enumerate() { - *value = round_bf16_f32(quantized as f32 * scales[group] + biases[group]); + let scaled = round_bf16_f32(quantized as f32 * scales[group]); + *value = round_bf16_f32(scaled + biases[group]); } } for packed_index in 0..in_packed { @@ -650,7 +655,8 @@ pub(crate) fn dequantize_affine_bf16_fused( } } else { // A 256-entry table can cost more than the usual 8-bit group saves, so - // retain the direct fused-round calculation for uncommon 8-bit checkpoints. + // retain the direct sequential-round calculation for uncommon 8-bit + // checkpoints. for output in 0..out { let row = &packed[output * in_packed * 4..(output + 1) * in_packed * 4]; let group_row = output * n_groups; @@ -662,8 +668,8 @@ pub(crate) fn dequantize_affine_bf16_fused( let input = packed_index * per_u32 + lane; let quantized = ((word >> (bits * lane)) & mask) as f32; let group = group_row + input / group_size; - weights[weight_row + input] = - round_bf16_f32(quantized * scales[group] + biases[group]); + let scaled = round_bf16_f32(quantized * scales[group]); + weights[weight_row + input] = round_bf16_f32(scaled + biases[group]); } } } @@ -1145,15 +1151,21 @@ mod tests { } #[test] - fn gemma3n_bf16_affine_fuses_multiply_and_bias_before_rounding() { - // Low nibble q=9 with the actual E2B token-2 embedding group values. - let packed = 9u32.to_le_bytes(); - let scale = 0xbbfb_u16.to_le_bytes(); - let bias = 0x3d7b_u16.to_le_bytes(); - let fused = dequantize_affine_bf16_fused(&packed, &scale, &bias, 1, 1, 4, 8).unwrap(); - assert_eq!(fused[0], -0.007_659_912); - let ordinary = dequantize_affine(&packed, &scale, &bias, 1, 1, 4, 8, true).unwrap(); - assert_eq!(fused, ordinary); + fn gemma3n_bf16_affine_rounds_multiply_before_adding_bias() { + // This projection group is a double-rounding counterexample: rounding + // only after q*scale+bias gives BF16 0xbc52, while MLX CUDA stores the + // BF16 product before the BF16 addition and therefore yields 0xbc50. + let packed = 10u32.to_le_bytes(); + let scale = 0xbbd2_u16.to_le_bytes(); + let bias = 0x3d52_u16.to_le_bytes(); + let sequential = + dequantize_affine_bf16_sequential(&packed, &scale, &bias, 1, 1, 4, 8).unwrap(); + assert_eq!(sequential[0].to_bits(), 0xbc50_0000); + + let fused = + round_bf16_f32(10.0 * f32::from_bits(0xbbd2_0000) + f32::from_bits(0x3d52_0000)); + assert_eq!(fused.to_bits(), 0xbc52_0000); + assert_ne!(sequential[0], fused); } #[test] @@ -1169,12 +1181,13 @@ mod tests { }; let scales = bf16_bytes(&scale_values); let biases = bf16_bytes(&bias_values); - let got = dequantize_affine_bf16_fused(&packed, &scales, &biases, 1, 1, 4, 4).unwrap(); + let got = dequantize_affine_bf16_sequential(&packed, &scales, &biases, 1, 1, 4, 4).unwrap(); let expected = (1..=8) .enumerate() .map(|(input, quantized)| { let group = input / 4; - round_bf16_f32(quantized as f32 * scale_values[group] + bias_values[group]) + let scaled = round_bf16_f32(quantized as f32 * scale_values[group]); + round_bf16_f32(scaled + bias_values[group]) }) .collect::>(); assert_eq!(got, expected); From 7c57beb0bd718d41f656db8547b87309629e94db Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 24 Jul 2026 20:51:20 +0900 Subject: [PATCH 09/19] fix(xla): cast Gemma3n mel before SSCP convolution Make the MLX checkpoint-dtype boundary explicit by rounding processor-produced F32 mel features to BF16 before the first SSCP convolution, independent of the graph-wide contraction precision. Add a synthetic regression whose pre-cast and post-convolution-only results differ, and expose the first SSCP output through the diagnostics-only MLX and IREE bundles for one bounded parity check. Validation: cargo fmt --all -- --check; git diff --check; focused Gemma3n MLX CPU test; focused Gemma3n audio emitter tests; mlxcel-xla diagnostics cargo check. Refs #878 --- examples/xla_gemma3n_audio_reference_check.rs | 1 + src/audio/gemma3n/diagnostics.rs | 5 +- src/audio/gemma3n/encoder.rs | 54 +++++++++++++++- .../src/emitter/gemma3n_audio_emit.rs | 1 + .../src/emitter/gemma3n_audio_emit_tests.rs | 62 +++++++++++++++++++ .../src/emitter/gemma3n_audio_ops.rs | 6 ++ .../mlxcel-xla/src/gemma3n_audio_runtime.rs | 1 + 7 files changed, 126 insertions(+), 4 deletions(-) diff --git a/examples/xla_gemma3n_audio_reference_check.rs b/examples/xla_gemma3n_audio_reference_check.rs index 6b20795b2..0da9b0bf8 100644 --- a/examples/xla_gemma3n_audio_reference_check.rs +++ b/examples/xla_gemma3n_audio_reference_check.rs @@ -235,6 +235,7 @@ fn main() -> Result<(), String> { } let stage_values = [ + ("sscp_conv_0", mlx.sscp_conv_0.as_slice()), ("encoded_reduced", mlx.encoded_reduced.as_slice()), ("soft_norm", mlx.soft_norm.as_slice()), ("soft_linear", mlx.soft_linear.as_slice()), diff --git a/src/audio/gemma3n/diagnostics.rs b/src/audio/gemma3n/diagnostics.rs index a5885ff83..de3465451 100644 --- a/src/audio/gemma3n/diagnostics.rs +++ b/src/audio/gemma3n/diagnostics.rs @@ -25,6 +25,7 @@ use crate::models::gemma3n::{Gemma3nAudioEmbedder, ModelArgs, TextConfig}; #[derive(Debug, Clone)] pub struct Gemma3nAudioMlxDiagnosticOutput { + pub sscp_conv_0: Vec, pub encoded_reduced: Vec, pub soft_norm: Vec, pub soft_linear: Vec, @@ -184,7 +185,8 @@ pub fn run_gemma3n_audio_mlx_diagnostics( .collect::>(); let invalid = mlxcel_core::from_slice_i32(&invalid, &[clips as i32, frame_bucket as i32]); let invalid = mlxcel_core::astype(&invalid, mlxcel_core::dtype::BOOL); - let (encoded, encoded_invalid) = tower.forward(&mel, &invalid)?; + let (encoded, encoded_invalid, sscp_conv_0) = + tower.forward_with_sscp_conv0_diagnostic(&mel, &invalid)?; let encoded_shape = mlxcel_core::array_shape(&encoded); let projected_lengths = (0..clips) .map(|clip| { @@ -314,6 +316,7 @@ pub fn run_gemma3n_audio_mlx_diagnostics( .fill(0.0); Ok(Gemma3nAudioMlxDiagnosticOutput { + sscp_conv_0: array_f32(&sscp_conv_0), encoded_reduced: array_f32(&encoded), soft_norm: array_f32(&soft_norm), soft_linear: array_f32(&soft_linear), diff --git a/src/audio/gemma3n/encoder.rs b/src/audio/gemma3n/encoder.rs index 6fc8f5e46..c75ca5e29 100644 --- a/src/audio/gemma3n/encoder.rs +++ b/src/audio/gemma3n/encoder.rs @@ -246,6 +246,19 @@ impl SubSampleConvProjection { let x = mlxcel_core::reshape(&x, &[shape[0], shape[1], shape[2] * shape[3]]); Ok(self.input_projection.forward(&x)) } + + #[cfg(feature = "xla-diagnostics")] + fn forward_with_conv0( + &self, + audio_mel: &MlxArray, + ) -> Result<(UniquePtr, UniquePtr), String> { + let x = mlxcel_core::expand_dims(audio_mel, -1); + let conv0 = self.conv0.forward(&x)?; + let x = self.conv1.forward(&conv0)?; + 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), conv0)) + } } fn clip_gradient(x: &MlxArray, limit: f32) -> UniquePtr { @@ -624,11 +637,11 @@ impl Gemma3nAudioEncoder { mlxcel_core::take(mask, &mlxcel_core::from_slice_i32(&indices, &[length]), 1) } - pub fn forward( + fn validate_input_shapes( &self, audio_mel: &MlxArray, invalid_mel_mask: &MlxArray, - ) -> Result<(UniquePtr, UniquePtr), String> { + ) -> Result<(), 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 @@ -640,8 +653,14 @@ impl Gemma3nAudioEncoder { self.config.input_feat_size )); } + Ok(()) + } - let mut encodings = self.subsample.forward(audio_mel)?; + fn finish_forward( + &self, + mut encodings: UniquePtr, + invalid_mel_mask: &MlxArray, + ) -> Result<(UniquePtr, UniquePtr), String> { let mut invalid_mask = Self::stride_mask( invalid_mel_mask, self.config.time_stride_product(), @@ -674,6 +693,35 @@ impl Gemma3nAudioEncoder { ); Ok((encodings, invalid_mask)) } + + pub fn forward( + &self, + audio_mel: &MlxArray, + invalid_mel_mask: &MlxArray, + ) -> Result<(UniquePtr, UniquePtr), String> { + self.validate_input_shapes(audio_mel, invalid_mel_mask)?; + let encodings = self.subsample.forward(audio_mel)?; + self.finish_forward(encodings, invalid_mel_mask) + } + + #[cfg(feature = "xla-diagnostics")] + pub(super) fn forward_with_sscp_conv0_diagnostic( + &self, + audio_mel: &MlxArray, + invalid_mel_mask: &MlxArray, + ) -> Result< + ( + UniquePtr, + UniquePtr, + UniquePtr, + ), + String, + > { + self.validate_input_shapes(audio_mel, invalid_mel_mask)?; + let (encodings, sscp_conv0) = self.subsample.forward_with_conv0(audio_mel)?; + let (encoded, invalid) = self.finish_forward(encodings, invalid_mel_mask)?; + Ok((encoded, invalid, sscp_conv0)) + } } #[cfg(test)] diff --git a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit.rs b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit.rs index 7e4412fea..4aab7405d 100644 --- a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit.rs +++ b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit.rs @@ -315,6 +315,7 @@ fn build_encode( trace.push("hard_projection", &hard_audio); #[cfg(feature = "diagnostics")] let diagnostics = [ + "sscp_conv_0", "encoded_reduced", "soft_norm", "soft_linear", diff --git a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit_tests.rs b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit_tests.rs index 0f8c7d4e0..27457bee8 100644 --- a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit_tests.rs +++ b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit_tests.rs @@ -76,6 +76,68 @@ fn split_audio_graphs_preserve_stage_and_weight_boundaries() { ); } +#[test] +fn first_sscp_convolution_rounds_mel_inputs_before_accumulation() { + fn synthetic_convolution(input: &[f32], weights: &[f32], pre_cast: bool) -> f32 { + let sum = input + .iter() + .zip(weights) + .map(|(&value, &weight)| { + let value = if pre_cast { + crate::weights::round_bf16_f32(value) + } else { + value + }; + value * crate::weights::round_bf16_f32(weight) + }) + .sum::(); + crate::weights::round_bf16_f32(sum) + } + + // Both values map to 1.0 in BF16, so the maintained MLX pre-cast cancels + // them exactly. Rounding only after the convolution preserves the small + // F32 difference and therefore produces a distinct non-zero output. + let input = [1.003, 1.0]; + let weights = [1.0, -1.0]; + let pre_cast = synthetic_convolution(&input, &weights, true); + let post_conv_only = synthetic_convolution(&input, &weights, false); + assert_eq!(pre_cast, 0.0); + assert_ne!(post_conv_only, 0.0); + + // Pin that same semantic boundary in the emitted graph even when global + // contraction precision is F32: the reshaped mel is narrowed to BF16 and + // widened back before padding feeds the first convolution. + let (encode, _) = emit_gemma3n_audio_encode( + &text(), + &Gemma3nXlaAudioConfig::default(), + 8, + 1, + Precision::F32, + ) + .unwrap(); + let first_convolution = encode + .find("\"stablehlo.convolution\"") + .expect("Gemma3n audio graph must contain SSCP convolution"); + let prefix = &encode[..first_convolution]; + let narrow = prefix + .find( + "stablehlo.convert %0 : (tensor<1x8x128x1xf32>) -> \ + tensor<1x8x128x1xbf16>", + ) + .expect("mel input must narrow to BF16 before the first SSCP convolution"); + let widen = prefix[narrow..] + .find( + "stablehlo.convert %1 : (tensor<1x8x128x1xbf16>) -> \ + tensor<1x8x128x1xf32>", + ) + .expect("BF16 mel input must widen to the graph's F32 carrier"); + let pad = prefix[narrow + widen..] + .find("\"stablehlo.pad\"(%2") + .expect("the first SSCP padding must consume BF16-rounded mel"); + assert!(widen > 0); + assert!(pad > 0); +} + fn compile(compiler: &std::ffi::OsStr, device: &str, stem: &str, artifact: &str, mlir: String) { let input = std::env::temp_dir().join(format!("{stem}-{artifact}.mlir")); let output = std::env::temp_dir().join(format!("{stem}-{artifact}.vmfb")); diff --git a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_ops.rs b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_ops.rs index 2dc924093..e87a4eb54 100644 --- a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_ops.rs +++ b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_ops.rs @@ -91,6 +91,12 @@ fn subsample_convs( let time = args.mel.ty.shape[1]; let frequency = args.mel.ty.shape[2]; let x = b.reshape(&args.mel, vec![batch, time, frequency, 1]); + // The maintained MLX path explicitly casts processor-produced F32 mel + // features to `conv_0.weight.dtype` (BF16 in released checkpoints) before + // the first SSCP convolution. Keep this boundary explicit instead of + // relying on the graph-wide contraction precision: post-convolution + // rounding cannot recover the products formed from BF16-rounded inputs. + let x = round_bf16(b, &x); let root = "audio_tower.subsample_conv_projection"; let conv0 = sscp_conv( diff --git a/src/lib/mlxcel-xla/src/gemma3n_audio_runtime.rs b/src/lib/mlxcel-xla/src/gemma3n_audio_runtime.rs index ba465852e..45a4942e0 100644 --- a/src/lib/mlxcel-xla/src/gemma3n_audio_runtime.rs +++ b/src/lib/mlxcel-xla/src/gemma3n_audio_runtime.rs @@ -651,6 +651,7 @@ impl Gemma3nAudioIreeRuntime { report_load_timing("audio-total", load_started); #[cfg(feature = "diagnostics")] let diagnostic_shapes = [ + "sscp_conv_0", "encoded_reduced", "soft_norm", "soft_linear", From 313768b9e162c75828f6dda41e8e78b8e63564de Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Sun, 26 Jul 2026 21:40:46 +0900 Subject: [PATCH 10/19] test(xla): isolate Gemma3n SSCP convolution drift Add a deterministic first-convolution MLX/IREE CPU probe that distinguishes backend reduction schedules within the production BF16 input and output boundary. Extend the real-checkpoint audio gate with max-pair values, BF16 bit patterns, ULP counts, and non-BF16 counts without relaxing its absolute or RMS failure thresholds. Validation: oracle unit tests 3/3, CPU local-task probe pass, rustfmt and diff check. Refs #878 --- examples/xla_gemma3n_audio_reference_check.rs | 86 +++- examples/xla_gemma3n_sscp_conv_oracle.rs | 373 ++++++++++++++++++ 2 files changed, 456 insertions(+), 3 deletions(-) create mode 100644 examples/xla_gemma3n_sscp_conv_oracle.rs diff --git a/examples/xla_gemma3n_audio_reference_check.rs b/examples/xla_gemma3n_audio_reference_check.rs index 0da9b0bf8..e196a80fb 100644 --- a/examples/xla_gemma3n_audio_reference_check.rs +++ b/examples/xla_gemma3n_audio_reference_check.rs @@ -131,6 +131,34 @@ struct Diff { max_abs: f64, rms: f64, max_index: usize, + max_bf16_ulp: u32, + non_bf16_values: usize, + over_one_bf16_ulp: usize, +} + +fn bf16_bits(value: f32) -> Option { + let bits = value.to_bits(); + (bits & 0xffff == 0).then_some((bits >> 16) as u16) +} + +fn ordered_bf16(bits: u16) -> i32 { + if bits & 0x8000 == 0 { + 0x8000 + i32::from(bits) + } else { + 0x8000 - i32::from(bits & 0x7fff) + } +} + +fn bf16_ulp_distance(left: f32, right: f32) -> Option { + let left = ordered_bf16(bf16_bits(left)?); + let right = ordered_bf16(bf16_bits(right)?); + Some(left.abs_diff(right)) +} + +fn display_bf16_bits(value: f32) -> String { + bf16_bits(value) + .map(|bits| format!("0x{bits:04x}")) + .unwrap_or_else(|| "not-bf16".to_string()) } fn compare(name: &str, reference: &[f32], actual: &[f32]) -> Result { @@ -144,6 +172,9 @@ fn compare(name: &str, reference: &[f32], actual: &[f32]) -> Result Result { + max_bf16_ulp = max_bf16_ulp.max(distance); + over_one_bf16_ulp += usize::from(distance > 1); + } + None => non_bf16_values += 1, + } } let rms = (squared / reference.len().max(1) as f64).sqrt(); let probes = [ @@ -161,8 +199,16 @@ fn compare(name: &str, reference: &[f32], actual: &[f32]) -> Result Result Result<(), String> { .find(|(_, diff)| diff.max_abs > 5e-3 || diff.rms > 5e-4); match first { Some((index, diff)) => Err(format!( - "first divergent stage={} max_abs={:.9e} rms={:.9e} max_index={}", - names[index], diff.max_abs, diff.rms, diff.max_index + "first divergent stage={} max_abs={:.9e} rms={:.9e} max_index={} \ + max_bf16_ulp={} non_bf16_values={} over_one_bf16_ulp={}", + names[index], + diff.max_abs, + diff.rms, + diff.max_index, + diff.max_bf16_ulp, + diff.non_bf16_values, + diff.over_one_bf16_ulp, )), None => { println!("Gemma3n audio-only MLX↔IREE intermediate gate PASS"); @@ -285,3 +341,27 @@ fn main() -> Result<(), String> { } } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn bf16_ulp_distance_handles_sign_zero_and_one_unit_steps() { + assert_eq!(bf16_ulp_distance(0.0, -0.0), Some(0)); + assert_eq!(bf16_ulp_distance(128.0, 129.0), Some(1)); + assert_eq!(bf16_ulp_distance(-128.0, -129.0), Some(1)); + assert_eq!(bf16_ulp_distance(128.0, 130.0), Some(2)); + assert_eq!(bf16_ulp_distance(1.0, 1.003), None); + } + + #[test] + fn compare_keeps_absolute_failure_data_and_bf16_drift_data_separate() { + let diff = compare("fixture", &[128.0, -128.0], &[129.0, -130.0]).unwrap(); + assert_eq!(diff.max_abs, 2.0); + assert_eq!(diff.max_index, 1); + assert_eq!(diff.max_bf16_ulp, 2); + assert_eq!(diff.non_bf16_values, 0); + assert_eq!(diff.over_one_bf16_ulp, 1); + } +} diff --git a/examples/xla_gemma3n_sscp_conv_oracle.rs b/examples/xla_gemma3n_sscp_conv_oracle.rs new file mode 100644 index 000000000..dbdf371e5 --- /dev/null +++ b/examples/xla_gemma3n_sscp_conv_oracle.rs @@ -0,0 +1,373 @@ +// 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. + +//! Small, deterministic MLX↔IREE oracle for the Gemma3n first SSCP convolution. +//! +//! This deliberately excludes the cumulative group norm and every later audio +//! operation. Gemma3n's first SSCP convolution has no bias; its exact dtype +//! boundary is F32 processor features -> BF16 input, BF16 checkpoint weights, +//! a 3x3 contraction, and a materialized BF16 output. + +use std::path::{Path, PathBuf}; +use std::process::Command; + +const INPUT_SHAPE: [usize; 4] = [1, 3, 3, 1]; +const OUTPUT_CHANNELS: usize = 4; +const OUTPUT_SHAPE: [usize; 4] = [1, 1, 1, OUTPUT_CHANNELS]; +const INPUT: [f32; 9] = [16384.0, 1.0, -16384.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]; + +// MLX stores convolution weights as [O, H, W, I]. Each row makes the three +// non-zero products a cancellation-sensitive permutation of [2^24, 1, -2^24]. +const MLX_WEIGHTS: [[f32; 9]; OUTPUT_CHANNELS] = [ + [1024.0, 1.0, 1024.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [ + 1024.0, + -16777216.0, + -0.00006103515625, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + ], + [ + 0.00006103515625, + 16777216.0, + 1024.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + ], + [-1024.0, 1.0, -1024.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], +]; + +const MLIR: &str = r#"module @gemma3n_sscp_conv_oracle { + func.func public @main( + %input: tensor<1x3x3x1xf32>, + %kernel: tensor<3x3x1x4xf32> + ) -> tensor<1x1x1x4xf32> { + %input_bf16 = stablehlo.convert %input : (tensor<1x3x3x1xf32>) -> tensor<1x3x3x1xbf16> + %kernel_bf16 = stablehlo.convert %kernel : (tensor<3x3x1x4xf32>) -> tensor<3x3x1x4xbf16> + %convolved = "stablehlo.convolution"(%input_bf16, %kernel_bf16) { + window_strides = array, + padding = dense<[[0, 0], [0, 0]]> : tensor<2x2xi64>, + lhs_dilation = array, + rhs_dilation = array, + window_reversal = array, + dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]>, + batch_group_count = 1 : i64, + feature_group_count = 1 : i64, + precision_config = [#stablehlo, #stablehlo] + } : (tensor<1x3x3x1xbf16>, tensor<3x3x1x4xbf16>) -> tensor<1x1x1x4xbf16> + %output = stablehlo.convert %convolved : (tensor<1x1x1x4xbf16>) -> tensor<1x1x1x4xf32> + return %output : tensor<1x1x1x4xf32> + } +} +"#; + +fn round_bf16(value: f32) -> f32 { + if !value.is_finite() { + return value; + } + let bits = value.to_bits(); + let bias = 0x7fff + ((bits >> 16) & 1); + f32::from_bits(bits.wrapping_add(bias) & 0xffff_0000) +} + +fn bf16_bits(value: f32) -> Option { + let bits = value.to_bits(); + (bits & 0xffff == 0).then_some((bits >> 16) as u16) +} + +fn ordered_bf16(bits: u16) -> i32 { + if bits & 0x8000 == 0 { + 0x8000 + i32::from(bits) + } else { + 0x8000 - i32::from(bits & 0x7fff) + } +} + +fn bf16_ulp_distance(left: f32, right: f32) -> Option { + Some(ordered_bf16(bf16_bits(left)?).abs_diff(ordered_bf16(bf16_bits(right)?))) +} + +fn mlx_output() -> Result, String> { + let input_shape = INPUT_SHAPE.map(|dimension| dimension as i32); + let input = mlxcel_core::astype( + &mlxcel_core::from_slice_f32(&INPUT, &input_shape), + mlxcel_core::dtype::BFLOAT16, + ); + let weights = MLX_WEIGHTS + .iter() + .flat_map(|row| row.iter().copied()) + .collect::>(); + let weights = mlxcel_core::astype( + &mlxcel_core::from_slice_f32(&weights, &[OUTPUT_CHANNELS as i32, 3, 3, 1]), + mlxcel_core::dtype::BFLOAT16, + ); + let output = mlxcel_core::try_conv2d(&input, &weights, 1, 1, 0, 0, 1, 1, 1) + .map_err(|error| format!("MLX first-conv probe failed: {error}"))?; + if mlxcel_core::array_dtype(&output) != mlxcel_core::dtype::BFLOAT16 { + return Err("MLX first-conv probe did not materialize BF16 output".to_string()); + } + let output = mlxcel_core::astype(&output, mlxcel_core::dtype::FLOAT32); + mlxcel_core::eval(&output); + Ok(mlxcel_core::array_to_raw_bytes(&output) + .chunks_exact(4) + .map(|bytes| f32::from_ne_bytes(bytes.try_into().expect("f32 byte width"))) + .collect()) +} + +fn iree_kernel() -> Vec { + let mut kernel = vec![0.0; 3 * 3 * OUTPUT_CHANNELS]; + for output in 0..OUTPUT_CHANNELS { + for spatial in 0..9 { + kernel[spatial * OUTPUT_CHANNELS + output] = MLX_WEIGHTS[output][spatial]; + } + } + kernel +} + +fn write_f32(path: &Path, values: &[f32]) -> Result<(), String> { + let bytes = values + .iter() + .flat_map(|value| value.to_le_bytes()) + .collect::>(); + std::fs::write(path, bytes).map_err(|error| format!("write {}: {error}", path.display())) +} + +fn required_path( + name: &str, + fallback: impl FnOnce() -> Option, +) -> Result { + let path = std::env::var_os(name) + .map(PathBuf::from) + .or_else(fallback) + .ok_or_else(|| format!("set {name} for the first-conv oracle"))?; + if !path.is_file() { + return Err(format!("{name} is not a file: {}", path.display())); + } + Ok(path) +} + +fn iree_output() -> Result, String> { + let compiler = required_path("MLXCEL_XLA_IREE_COMPILE", || { + std::env::var_os("IREE_CUDA_HOME") + .map(PathBuf::from) + .map(|home| home.join("venv/bin/iree-compile")) + })?; + let runner = required_path("MLXCEL_XLA_IREE_RUN_MODULE", || { + std::env::var_os("IREE_CUDA_HOME") + .map(PathBuf::from) + .map(|home| home.join("build/tools/iree-run-module")) + })?; + let stem = format!("mlxcel-gemma3n-sscp-conv-oracle-{}", std::process::id()); + let temporary = std::env::temp_dir(); + let mlir_path = temporary.join(format!("{stem}.mlir")); + let vmfb_path = temporary.join(format!("{stem}.vmfb")); + let input_path = temporary.join(format!("{stem}-input.bin")); + let kernel_path = temporary.join(format!("{stem}-kernel.bin")); + let output_path = temporary.join(format!("{stem}-output.bin")); + let paths = [ + mlir_path.as_path(), + vmfb_path.as_path(), + input_path.as_path(), + kernel_path.as_path(), + output_path.as_path(), + ]; + let result = (|| { + std::fs::write(&mlir_path, MLIR) + .map_err(|error| format!("write {}: {error}", mlir_path.display()))?; + write_f32(&input_path, &INPUT)?; + write_f32(&kernel_path, &iree_kernel())?; + let compiled = Command::new(&compiler) + .arg("--iree-input-type=stablehlo") + .arg("--iree-hal-target-device=local") + .arg("--iree-hal-local-target-device-backends=llvm-cpu") + .arg(&mlir_path) + .arg("-o") + .arg(&vmfb_path) + .output() + .map_err(|error| format!("run {}: {error}", compiler.display()))?; + if !compiled.status.success() { + return Err(format!( + "iree-compile failed:\n{}", + String::from_utf8_lossy(&compiled.stderr) + )); + } + let ran = Command::new(&runner) + .arg("--device=local-task") + .arg(format!("--module={}", vmfb_path.display())) + .arg("--function=main") + .arg(format!("--input=1x3x3x1xf32=@{}", input_path.display())) + .arg(format!( + "--input=3x3x1x{OUTPUT_CHANNELS}xf32=@{}", + kernel_path.display() + )) + .arg(format!("--output=@{}", output_path.display())) + .output() + .map_err(|error| format!("run {}: {error}", runner.display()))?; + if !ran.status.success() { + return Err(format!( + "iree-run-module failed:\nstdout:\n{}\nstderr:\n{}", + String::from_utf8_lossy(&ran.stdout), + String::from_utf8_lossy(&ran.stderr) + )); + } + let bytes = std::fs::read(&output_path) + .map_err(|error| format!("read {}: {error}", output_path.display()))?; + if bytes.len() != OUTPUT_CHANNELS * 4 { + return Err(format!( + "IREE first-conv probe returned {} bytes, expected {}", + bytes.len(), + OUTPUT_CHANNELS * 4 + )); + } + Ok(bytes + .chunks_exact(4) + .map(|bytes| f32::from_le_bytes(bytes.try_into().expect("f32 byte width"))) + .collect()) + })(); + for path in paths { + let _ = std::fs::remove_file(path); + } + result +} + +fn schedule_outcomes(input: &[f32], weights: &[f32]) -> Vec { + fn permutations(values: &mut [f32], start: usize, outputs: &mut Vec) { + if start == values.len() { + let left = round_bf16(values.iter().copied().fold(0.0, |sum, value| sum + value)); + let right = round_bf16( + values + .iter() + .rev() + .copied() + .fold(0.0, |sum, value| sum + value), + ); + outputs.extend([left, right]); + return; + } + for index in start..values.len() { + values.swap(start, index); + permutations(values, start + 1, outputs); + values.swap(start, index); + } + } + + let mut products = input + .iter() + .zip(weights) + .map(|(&input, &weight)| round_bf16(input) * round_bf16(weight)) + .filter(|value| *value != 0.0) + .collect::>(); + let mut outcomes = Vec::new(); + permutations(&mut products, 0, &mut outcomes); + outcomes.sort_by_key(|value| value.to_bits()); + outcomes.dedup_by_key(|value| value.to_bits()); + outcomes +} + +fn main() -> Result<(), String> { + let mlx = mlx_output()?; + let iree = iree_output()?; + if mlx.len() != OUTPUT_SHAPE.iter().product::() + || iree.len() != OUTPUT_SHAPE.iter().product::() + { + return Err(format!( + "first-conv output shape mismatch: MLX={} IREE={} expected={OUTPUT_SHAPE:?}", + mlx.len(), + iree.len() + )); + } + + let mut exact = true; + for output in 0..OUTPUT_CHANNELS { + let mlx_value = mlx[output]; + let iree_value = iree[output]; + let outcomes = schedule_outcomes(&INPUT, &MLX_WEIGHTS[output]); + if !outcomes + .iter() + .any(|value| value.to_bits() == mlx_value.to_bits()) + || !outcomes + .iter() + .any(|value| value.to_bits() == iree_value.to_bits()) + { + return Err(format!( + "channel {output} escaped the BF16-input/F32-accumulator schedule envelope: \ + MLX={mlx_value:?} IREE={iree_value:?} outcomes={outcomes:?}" + )); + } + let distance = bf16_ulp_distance(mlx_value, iree_value) + .ok_or_else(|| format!("channel {output} did not materialize exact BF16 values"))?; + eprintln!( + "gemma3n-sscp-conv-oracle: channel={output} mlx={mlx_value:.9e} \ + iree={iree_value:.9e} mlx_bf16=0x{:04x} iree_bf16=0x{:04x} \ + bf16_ulp={distance} schedule_outcomes={outcomes:?}", + bf16_bits(mlx_value).expect("checked BF16"), + bf16_bits(iree_value).expect("checked BF16"), + ); + exact &= mlx_value.to_bits() == iree_value.to_bits(); + } + + println!( + "Gemma3n first SSCP convolution MLX↔IREE CPU oracle PASS ({})", + if exact { + "exact schedule match" + } else { + "backend reduction schedules differ within the enumerated F32 accumulation envelope" + } + ); + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn probe_pins_the_production_first_convolution_dtype_boundary() { + assert!(MLIR.contains("tensor<1x3x3x1xbf16>")); + assert!(MLIR.contains("tensor<3x3x1x4xbf16>")); + assert!(MLIR.contains("-> tensor<1x1x1x4xbf16>")); + assert!(MLIR.contains("precision DEFAULT")); + assert!(!MLIR.contains("stablehlo.add")); + } + + #[test] + fn cancellation_fixture_distinguishes_valid_f32_reduction_schedules() { + let outcomes = schedule_outcomes(&INPUT, &MLX_WEIGHTS[0]); + assert_eq!(outcomes, [0.0, 1.0]); + assert_eq!(round_bf16(128.0), 128.0); + assert_eq!(bf16_ulp_distance(128.0, 129.0), Some(1)); + } + + #[test] + fn iree_kernel_transposes_mlx_output_major_weights() { + let kernel = iree_kernel(); + for output in 0..OUTPUT_CHANNELS { + for spatial in 0..9 { + assert_eq!( + kernel[spatial * OUTPUT_CHANNELS + output], + MLX_WEIGHTS[output][spatial] + ); + } + } + assert_eq!(INPUT.len(), INPUT_SHAPE.iter().product::()); + } +} From 5cd16cd66ff5572f770d240aa196426d0f324c87 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Sun, 26 Jul 2026 22:51:49 +0900 Subject: [PATCH 11/19] fix(xla): keep Gemma3n SSCP accumulation in f32 The bounded audio-only oracle localized broad BF16 drift to the first SSCP output: max abs 1.0, RMS 0.02212289754, and 1,429,092 values beyond one BF16 ULP. The graph declared bf16-accum-f32 but the convolution builder gave BF16-demoted operands a BF16 result, allowing cancellation-sensitive low-precision accumulation before normalization. Add an SSCP-only convolution contract that keeps BF16 operands while producing an F32 StableHLO result, then applies the existing explicit BF16 output rounding. Padding, NHWC/HWIO layout, weights, and failure thresholds remain unchanged. Validation: focused Gemma3n audio emitter tests, builder tests, a minimal IREE local LLVM CPU compile, cargo check with diagnostics, cargo fmt, and cargo clippy. The pinned actual-checkpoint gate was not rerun. Refs #878 --- src/lib/mlxcel-xla/src/emitter/builder.rs | 33 ++++++++++++- .../src/emitter/gemma3n_audio_emit_tests.rs | 46 +++++++++++++++++++ .../src/emitter/gemma3n_audio_ops.rs | 7 ++- 3 files changed, 83 insertions(+), 3 deletions(-) diff --git a/src/lib/mlxcel-xla/src/emitter/builder.rs b/src/lib/mlxcel-xla/src/emitter/builder.rs index b8898cfab..89e6eda45 100644 --- a/src/lib/mlxcel-xla/src/emitter/builder.rs +++ b/src/lib/mlxcel-xla/src/emitter/builder.rs @@ -1172,6 +1172,32 @@ impl Builder { strides: &[usize], padding: &[(usize, usize)], feature_groups: usize, + ) -> Val { + self.convolution_with_output(input, kernel, strides, padding, feature_groups, false) + } + + /// Channels-last convolution whose low-precision inputs accumulate into an + /// explicit F32 result before the caller applies its output rounding. + pub fn convolution_f32_accumulate( + &mut self, + input: &Val, + kernel: &Val, + strides: &[usize], + padding: &[(usize, usize)], + feature_groups: usize, + ) -> Val { + self.convolution_with_output(input, kernel, strides, padding, feature_groups, true) + } + + #[allow(clippy::too_many_arguments)] + fn convolution_with_output( + &mut self, + input: &Val, + kernel: &Val, + strides: &[usize], + padding: &[(usize, usize)], + feature_groups: usize, + f32_output: bool, ) -> Val { assert!(feature_groups > 0); assert_eq!( @@ -1202,7 +1228,7 @@ impl Builder { shape.push((padded - kernel_size) / strides[axis] + 1); } shape.push(kernel.ty.shape[spatial_rank + 1]); - let output_ty = Ty::new(shape, input.ty.elt); + let output_ty = Ty::new(shape, if f32_output { "f32" } else { input.ty.elt }); let input_demoted; let kernel_demoted; let (input, kernel) = match self.precision.dot_elt() { @@ -1221,7 +1247,10 @@ impl Builder { } None => (input, kernel), }; - let convolution_ty = Ty::new(output_ty.shape.clone(), input.ty.elt); + let convolution_ty = Ty::new( + output_ty.shape.clone(), + if f32_output { "f32" } else { input.ty.elt }, + ); let array = |values: &[usize]| { values .iter() diff --git a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit_tests.rs b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit_tests.rs index 27457bee8..56207f287 100644 --- a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit_tests.rs +++ b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit_tests.rs @@ -1,4 +1,5 @@ use super::*; +use crate::emitter::builder::Ty; fn text() -> Gemma3nConfig { Gemma3nConfig::from_json_str( @@ -138,6 +139,51 @@ fn first_sscp_convolution_rounds_mel_inputs_before_accumulation() { assert!(pad > 0); } +#[test] +fn sscp_convolutions_accumulate_bf16_operands_into_f32_results() { + let (encode, _) = emit_gemma3n_audio_encode( + &text(), + &Gemma3nXlaAudioConfig::default(), + 8, + 1, + Precision::Bf16, + ) + .unwrap(); + assert!(encode.contains( + "(tensor<1x10x130x1xbf16>, tensor<3x3x1x128xbf16>) -> \ + tensor<1x4x64x128xf32>" + )); + assert!(encode.contains( + "(tensor<1x6x66x128xbf16>, tensor<3x3x128x32xbf16>) -> \ + tensor<1x2x32x32xf32>" + )); + assert!(!encode.contains( + "(tensor<1x10x130x1xbf16>, tensor<3x3x1x128xbf16>) -> \ + tensor<1x4x64x128xbf16>" + )); +} + +#[test] +#[ignore = "requires IREE_DIST"] +fn sscp_f32_accumulator_contract_compiles_for_cpu() { + let mut builder = Builder::new().with_precision(Precision::Bf16); + let input = Builder::arg(0, Ty::f32(vec![1, 3, 3, 1])); + let kernel = Builder::arg(1, Ty::f32(vec![3, 3, 1, 4])); + let output = builder.convolution_f32_accumulate(&input, &kernel, &[1, 1], &[(0, 0); 2], 1); + let output_ty = output.ty.render(); + let mlir = format!( + "module @sscp_f32_accumulator {{\n func.func public @main(%arg0: \ + tensor<1x3x3x1xf32>, %arg1: tensor<3x3x1x4xf32>) -> {output_ty} {{\n{} \ + return {} : {output_ty}\n }}\n}}\n", + builder.body(), + output.name, + ); + let compiler = std::path::PathBuf::from(std::env::var_os("IREE_DIST").expect("IREE_DIST")) + .join("bin/iree-compile"); + let stem = format!("mlxcel-gemma3n-sscp-f32-accumulator-{}", std::process::id()); + compile(compiler.as_os_str(), "local", &stem, "contract", mlir); +} + fn compile(compiler: &std::ffi::OsStr, device: &str, stem: &str, artifact: &str, mlir: String) { let input = std::env::temp_dir().join(format!("{stem}-{artifact}.mlir")); let output = std::env::temp_dir().join(format!("{stem}-{artifact}.vmfb")); diff --git a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_ops.rs b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_ops.rs index e87a4eb54..d435f39a5 100644 --- a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_ops.rs +++ b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_ops.rs @@ -75,7 +75,12 @@ fn sscp_conv( let zero = b.const_f32(0.0); let padded = b.pad(x, &zero, &[0, 0, 1, 0], &[0, kernel[0] - 1, 1, 0]); let weight = b.transpose(weight, &[1, 2, 3, 0]); - let convolved = b.convolution(&padded, &weight, &stride, &[(0, 0), (0, 0)], groups); + // MLX stores both operands as BF16 but accumulates the convolution into an + // F32 carrier before materializing the BF16 output tensor. The graph-wide + // BF16 contraction mode otherwise gives StableHLO a BF16 result type, + // allowing the local CPU backend to lose cancellation-sensitive lanes. + let convolved = + b.convolution_f32_accumulate(&padded, &weight, &stride, &[(0, 0), (0, 0)], groups); let convolved = round_bf16(b, &convolved); let normalized = cumulative_group_norm(b, &convolved, norm, eps); let activated = relu(b, &normalized); From 9160313da85265487412502bc2934e061a976c5c Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Sun, 26 Jul 2026 23:28:00 +0900 Subject: [PATCH 12/19] test(xla): pinpoint Gemma3n audio drift boundaries The F32 SSCP accumulator fix reduced the bounded audio-only conv0 drift from RMS 0.02212289754 and 1,429,092 values beyond one BF16 ULP to RMS 0.00005109118683 and four values beyond one ULP, but the unchanged absolute threshold still fails and later encoder stages remain broadly divergent. Expose the BF16 convolution materialization, cumulative norm, activation, conv1, input projection, and first Conformer feed-forward boundaries in both MLX and IREE diagnostics. Report each over-one-ULP index with values, absolute difference, BF16 bits, and distance while capping pathological logs at 32 entries. Validation: Gemma3n audio emitter tests, xla-diagnostics example check, three focused oracle comparison tests, cargo fmt, and cargo clippy. The bounded actual-checkpoint oracle was not rerun after this diagnostics-only change. Refs #878 --- examples/xla_gemma3n_audio_reference_check.rs | 79 +++++++- src/audio/gemma3n/diagnostics.rs | 20 +- src/audio/gemma3n/encoder.rs | 172 +++++++++++++++--- .../src/emitter/gemma3n_audio_emit.rs | 11 ++ .../src/emitter/gemma3n_audio_emit_tests.rs | 6 +- .../src/emitter/gemma3n_audio_ops.rs | 45 ++++- .../mlxcel-xla/src/gemma3n_audio_runtime.rs | 7 + 7 files changed, 295 insertions(+), 45 deletions(-) diff --git a/examples/xla_gemma3n_audio_reference_check.rs b/examples/xla_gemma3n_audio_reference_check.rs index e196a80fb..4d48534dc 100644 --- a/examples/xla_gemma3n_audio_reference_check.rs +++ b/examples/xla_gemma3n_audio_reference_check.rs @@ -30,6 +30,7 @@ const AUDIO_TOKEN_ID: i32 = 262_273; const BOA_TOKEN_ID: i32 = 256_000; const EOA_TOKEN_ID: i32 = 262_272; const CONTEXT_CAPACITY: usize = 256; +const MAX_BF16_OUTLIER_PROBES: usize = 32; fn path_env(name: &str, default: &str) -> PathBuf { std::env::var_os(name) @@ -134,6 +135,14 @@ struct Diff { max_bf16_ulp: u32, non_bf16_values: usize, over_one_bf16_ulp: usize, + bf16_outliers: Vec, +} + +struct Bf16Outlier { + index: usize, + reference: f32, + actual: f32, + distance: u32, } fn bf16_bits(value: f32) -> Option { @@ -175,6 +184,7 @@ fn compare(name: &str, reference: &[f32], actual: &[f32]) -> Result Result { max_bf16_ulp = max_bf16_ulp.max(distance); - over_one_bf16_ulp += usize::from(distance > 1); + if distance > 1 { + over_one_bf16_ulp += 1; + if bf16_outliers.len() < MAX_BF16_OUTLIER_PROBES { + bf16_outliers.push(Bf16Outlier { + index, + reference: *reference, + actual: *actual, + distance, + }); + } + } } None => non_bf16_values += 1, } @@ -218,6 +238,26 @@ fn compare(name: &str, reference: &[f32], actual: &[f32]) -> Result Result Result<(), String> { } let stage_values = [ + ( + "sscp_conv_0_convolution", + mlx.sscp_conv_0_convolution.as_slice(), + ), + ("sscp_conv_0_norm", mlx.sscp_conv_0_norm.as_slice()), ("sscp_conv_0", mlx.sscp_conv_0.as_slice()), + ( + "sscp_conv_1_convolution", + mlx.sscp_conv_1_convolution.as_slice(), + ), + ("sscp_conv_1_norm", mlx.sscp_conv_1_norm.as_slice()), + ("sscp_conv_1", mlx.sscp_conv_1.as_slice()), + ("input_projection", mlx.input_projection.as_slice()), + ( + "conformer.0.feed_forward_start", + mlx.conformer_0_feed_forward_start.as_slice(), + ), ("encoded_reduced", mlx.encoded_reduced.as_slice()), ("soft_norm", mlx.soft_norm.as_slice()), ("soft_linear", mlx.soft_linear.as_slice()), @@ -326,7 +383,8 @@ fn main() -> Result<(), String> { match first { Some((index, diff)) => Err(format!( "first divergent stage={} max_abs={:.9e} rms={:.9e} max_index={} \ - max_bf16_ulp={} non_bf16_values={} over_one_bf16_ulp={}", + max_bf16_ulp={} non_bf16_values={} over_one_bf16_ulp={} \ + captured_bf16_outliers={}", names[index], diff.max_abs, diff.rms, @@ -334,6 +392,7 @@ fn main() -> Result<(), String> { diff.max_bf16_ulp, diff.non_bf16_values, diff.over_one_bf16_ulp, + diff.bf16_outliers.len(), )), None => { println!("Gemma3n audio-only MLX↔IREE intermediate gate PASS"); @@ -363,5 +422,21 @@ mod tests { assert_eq!(diff.max_bf16_ulp, 2); assert_eq!(diff.non_bf16_values, 0); assert_eq!(diff.over_one_bf16_ulp, 1); + assert_eq!(diff.bf16_outliers.len(), 1); + assert_eq!(diff.bf16_outliers[0].index, 1); + assert_eq!(diff.bf16_outliers[0].distance, 2); + } + + #[test] + fn compare_bounds_detailed_bf16_outlier_reporting() { + let reference = vec![128.0; MAX_BF16_OUTLIER_PROBES + 8]; + let actual = vec![130.0; reference.len()]; + let diff = compare("bounded-outliers", &reference, &actual).unwrap(); + assert_eq!(diff.over_one_bf16_ulp, reference.len()); + assert_eq!(diff.bf16_outliers.len(), MAX_BF16_OUTLIER_PROBES); + assert_eq!( + diff.bf16_outliers.last().map(|outlier| outlier.index), + Some(MAX_BF16_OUTLIER_PROBES - 1), + ); } } diff --git a/src/audio/gemma3n/diagnostics.rs b/src/audio/gemma3n/diagnostics.rs index de3465451..a5ca90f50 100644 --- a/src/audio/gemma3n/diagnostics.rs +++ b/src/audio/gemma3n/diagnostics.rs @@ -25,7 +25,14 @@ use crate::models::gemma3n::{Gemma3nAudioEmbedder, ModelArgs, TextConfig}; #[derive(Debug, Clone)] pub struct Gemma3nAudioMlxDiagnosticOutput { + pub sscp_conv_0_convolution: Vec, + pub sscp_conv_0_norm: Vec, pub sscp_conv_0: Vec, + pub sscp_conv_1_convolution: Vec, + pub sscp_conv_1_norm: Vec, + pub sscp_conv_1: Vec, + pub input_projection: Vec, + pub conformer_0_feed_forward_start: Vec, pub encoded_reduced: Vec, pub soft_norm: Vec, pub soft_linear: Vec, @@ -185,8 +192,8 @@ pub fn run_gemma3n_audio_mlx_diagnostics( .collect::>(); let invalid = mlxcel_core::from_slice_i32(&invalid, &[clips as i32, frame_bucket as i32]); let invalid = mlxcel_core::astype(&invalid, mlxcel_core::dtype::BOOL); - let (encoded, encoded_invalid, sscp_conv_0) = - tower.forward_with_sscp_conv0_diagnostic(&mel, &invalid)?; + let (encoded, encoded_invalid, encoder_stages) = + tower.forward_with_audio_diagnostics(&mel, &invalid)?; let encoded_shape = mlxcel_core::array_shape(&encoded); let projected_lengths = (0..clips) .map(|clip| { @@ -316,7 +323,14 @@ pub fn run_gemma3n_audio_mlx_diagnostics( .fill(0.0); Ok(Gemma3nAudioMlxDiagnosticOutput { - sscp_conv_0: array_f32(&sscp_conv_0), + sscp_conv_0_convolution: array_f32(&encoder_stages.sscp_conv_0_convolution), + sscp_conv_0_norm: array_f32(&encoder_stages.sscp_conv_0_norm), + sscp_conv_0: array_f32(&encoder_stages.sscp_conv_0), + sscp_conv_1_convolution: array_f32(&encoder_stages.sscp_conv_1_convolution), + sscp_conv_1_norm: array_f32(&encoder_stages.sscp_conv_1_norm), + sscp_conv_1: array_f32(&encoder_stages.sscp_conv_1), + input_projection: array_f32(&encoder_stages.input_projection), + conformer_0_feed_forward_start: array_f32(&encoder_stages.conformer_0_feed_forward_start), encoded_reduced: array_f32(&encoded), soft_norm: array_f32(&soft_norm), soft_linear: array_f32(&soft_linear), diff --git a/src/audio/gemma3n/encoder.rs b/src/audio/gemma3n/encoder.rs index c75ca5e29..c0bbd7f53 100644 --- a/src/audio/gemma3n/encoder.rs +++ b/src/audio/gemma3n/encoder.rs @@ -123,6 +123,13 @@ struct SscpConvBlock { stride_frequency: i32, } +#[cfg(feature = "xla-diagnostics")] +struct SscpConvDiagnosticStages { + convolution: UniquePtr, + norm: UniquePtr, + activated: UniquePtr, +} + impl SscpConvBlock { fn from_weights( weights: &WeightMap, @@ -157,7 +164,7 @@ impl SscpConvBlock { }) } - fn forward(&self, x: &MlxArray) -> Result, String> { + fn convolution(&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 @@ -167,7 +174,7 @@ impl SscpConvBlock { // 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( + mlxcel_core::try_conv2d( &x, &self.conv_weight, self.stride_time, @@ -178,8 +185,24 @@ impl SscpConvBlock { 1, 1, ) - .map_err(|error| format!("Gemma3n SSCP conv2d failed: {error}"))?; - Ok(mlxcel_core::relu(&self.norm.forward(&x))) + .map_err(|error| format!("Gemma3n SSCP conv2d failed: {error}")) + } + + fn forward(&self, x: &MlxArray) -> Result, String> { + let convolution = self.convolution(x)?; + Ok(mlxcel_core::relu(&self.norm.forward(&convolution))) + } + + #[cfg(feature = "xla-diagnostics")] + fn forward_with_diagnostics(&self, x: &MlxArray) -> Result { + let convolution = self.convolution(x)?; + let norm = self.norm.forward(&convolution); + let activated = mlxcel_core::relu(&norm); + Ok(SscpConvDiagnosticStages { + convolution, + norm, + activated, + }) } } @@ -189,6 +212,13 @@ struct SubSampleConvProjection { input_projection: UnifiedLinear, } +#[cfg(feature = "xla-diagnostics")] +struct SubsampleDiagnosticStages { + conv0: SscpConvDiagnosticStages, + conv1: SscpConvDiagnosticStages, + input_projection: UniquePtr, +} + fn final_frequency_size(config: &Gemma3nAudioConfig) -> usize { config .sscp_conv_kernel_size @@ -248,16 +278,22 @@ impl SubSampleConvProjection { } #[cfg(feature = "xla-diagnostics")] - fn forward_with_conv0( + fn forward_with_diagnostics( &self, audio_mel: &MlxArray, - ) -> Result<(UniquePtr, UniquePtr), String> { + ) -> Result { let x = mlxcel_core::expand_dims(audio_mel, -1); - let conv0 = self.conv0.forward(&x)?; - let x = self.conv1.forward(&conv0)?; - 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), conv0)) + let conv0 = self.conv0.forward_with_diagnostics(&x)?; + let conv1 = self.conv1.forward_with_diagnostics(&conv0.activated)?; + let shape = mlxcel_core::array_shape(&conv1.activated); + let flattened = + mlxcel_core::reshape(&conv1.activated, &[shape[0], shape[1], shape[2] * shape[3]]); + let input_projection = self.input_projection.forward(&flattened); + Ok(SubsampleDiagnosticStages { + conv0, + conv1, + input_projection, + }) } } @@ -542,14 +578,13 @@ impl ConformerBlock { }) } - fn forward( + fn forward_after_feed_forward_start( &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 x = self.attention.forward(x, invalid_mask, causal_valid_mask); let valid = mlxcel_core::astype( &mlxcel_core::reshape( &mlxcel_core::logical_not(invalid_mask), @@ -567,6 +602,32 @@ impl ConformerBlock { let x = self.feed_forward_end.forward(&x); Ok(self.norm.forward(&clip_gradient(&x, self.clipping))) } + + fn forward( + &self, + x: &MlxArray, + invalid_mask: &MlxArray, + causal_valid_mask: &MlxArray, + ) -> Result, String> { + let feed_forward_start = self.feed_forward_start.forward(x); + self.forward_after_feed_forward_start(&feed_forward_start, invalid_mask, causal_valid_mask) + } + + #[cfg(feature = "xla-diagnostics")] + fn forward_with_feed_forward_start( + &self, + x: &MlxArray, + invalid_mask: &MlxArray, + causal_valid_mask: &MlxArray, + ) -> Result<(UniquePtr, UniquePtr), String> { + let feed_forward_start = self.feed_forward_start.forward(x); + let output = self.forward_after_feed_forward_start( + &feed_forward_start, + invalid_mask, + causal_valid_mask, + )?; + Ok((output, feed_forward_start)) + } } pub struct Gemma3nAudioEncoder { @@ -575,6 +636,18 @@ pub struct Gemma3nAudioEncoder { conformer: Vec, } +#[cfg(feature = "xla-diagnostics")] +pub(super) struct Gemma3nAudioEncoderDiagnosticStages { + pub sscp_conv_0_convolution: UniquePtr, + pub sscp_conv_0_norm: UniquePtr, + pub sscp_conv_0: UniquePtr, + pub sscp_conv_1_convolution: UniquePtr, + pub sscp_conv_1_norm: UniquePtr, + pub sscp_conv_1: UniquePtr, + pub input_projection: UniquePtr, + pub conformer_0_feed_forward_start: UniquePtr, +} + impl Gemma3nAudioEncoder { pub fn from_weights( weights: &WeightMap, @@ -656,21 +729,23 @@ impl Gemma3nAudioEncoder { Ok(()) } - fn finish_forward( + fn initial_invalid_mask( &self, - mut encodings: UniquePtr, + encodings: &MlxArray, invalid_mel_mask: &MlxArray, - ) -> Result<(UniquePtr, UniquePtr), String> { - let mut invalid_mask = Self::stride_mask( + ) -> UniquePtr { + 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)?; - } + mlxcel_core::array_shape(encodings)[1], + ) + } + fn reduce_and_mask( + &self, + mut encodings: UniquePtr, + mut invalid_mask: UniquePtr, + ) -> Result<(UniquePtr, UniquePtr), String> { if self.config.conf_reduction_factor > 1 { let reduced_len = (mlxcel_core::array_shape(&encodings)[1] + self.config.conf_reduction_factor as i32 @@ -694,6 +769,19 @@ impl Gemma3nAudioEncoder { Ok((encodings, invalid_mask)) } + fn finish_forward( + &self, + mut encodings: UniquePtr, + invalid_mel_mask: &MlxArray, + ) -> Result<(UniquePtr, UniquePtr), String> { + let invalid_mask = self.initial_invalid_mask(&encodings, invalid_mel_mask); + let causal_valid_mask = self.causal_valid_mask(); + for block in &self.conformer { + encodings = block.forward(&encodings, &invalid_mask, &causal_valid_mask)?; + } + self.reduce_and_mask(encodings, invalid_mask) + } + pub fn forward( &self, audio_mel: &MlxArray, @@ -705,7 +793,7 @@ impl Gemma3nAudioEncoder { } #[cfg(feature = "xla-diagnostics")] - pub(super) fn forward_with_sscp_conv0_diagnostic( + pub(super) fn forward_with_audio_diagnostics( &self, audio_mel: &MlxArray, invalid_mel_mask: &MlxArray, @@ -713,14 +801,40 @@ impl Gemma3nAudioEncoder { ( UniquePtr, UniquePtr, - UniquePtr, + Gemma3nAudioEncoderDiagnosticStages, ), String, > { self.validate_input_shapes(audio_mel, invalid_mel_mask)?; - let (encodings, sscp_conv0) = self.subsample.forward_with_conv0(audio_mel)?; - let (encoded, invalid) = self.finish_forward(encodings, invalid_mel_mask)?; - Ok((encoded, invalid, sscp_conv0)) + let subsample = self.subsample.forward_with_diagnostics(audio_mel)?; + let mut encodings = mlxcel_core::copy(&subsample.input_projection); + let invalid_mask = self.initial_invalid_mask(&encodings, invalid_mel_mask); + let causal_valid_mask = self.causal_valid_mask(); + let first = self + .conformer + .first() + .ok_or_else(|| "Gemma3n audio diagnostic requires a Conformer block".to_string())?; + let (first_output, conformer_0_feed_forward_start) = + first.forward_with_feed_forward_start(&encodings, &invalid_mask, &causal_valid_mask)?; + encodings = first_output; + for block in &self.conformer[1..] { + encodings = block.forward(&encodings, &invalid_mask, &causal_valid_mask)?; + } + let (encoded, invalid) = self.reduce_and_mask(encodings, invalid_mask)?; + Ok(( + encoded, + invalid, + Gemma3nAudioEncoderDiagnosticStages { + sscp_conv_0_convolution: subsample.conv0.convolution, + sscp_conv_0_norm: subsample.conv0.norm, + sscp_conv_0: subsample.conv0.activated, + sscp_conv_1_convolution: subsample.conv1.convolution, + sscp_conv_1_norm: subsample.conv1.norm, + sscp_conv_1: subsample.conv1.activated, + input_projection: subsample.input_projection, + conformer_0_feed_forward_start, + }, + )) } } diff --git a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit.rs b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit.rs index 4aab7405d..e915f90e9 100644 --- a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit.rs +++ b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit.rs @@ -88,7 +88,11 @@ fn audio_encoder( trace: &mut Trace, ) -> (Val, Val, Val) { let subsampled = subsample_with_stages(b, args, audio); + trace.push("sscp_conv_0_convolution", &subsampled.conv0_convolution); + trace.push("sscp_conv_0_norm", &subsampled.conv0_norm); trace.push("sscp_conv_0", &subsampled.conv0); + trace.push("sscp_conv_1_convolution", &subsampled.conv1_convolution); + trace.push("sscp_conv_1_norm", &subsampled.conv1_norm); trace.push("sscp_conv_1", &subsampled.conv1); trace.push("input_projection", &subsampled.hidden); let mut hidden = subsampled.hidden; @@ -315,7 +319,14 @@ fn build_encode( trace.push("hard_projection", &hard_audio); #[cfg(feature = "diagnostics")] let diagnostics = [ + "sscp_conv_0_convolution", + "sscp_conv_0_norm", "sscp_conv_0", + "sscp_conv_1_convolution", + "sscp_conv_1_norm", + "sscp_conv_1", + "input_projection", + "conformer.0.feed_forward_start", "encoded_reduced", "soft_norm", "soft_linear", diff --git a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit_tests.rs b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit_tests.rs index 56207f287..93119092e 100644 --- a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit_tests.rs +++ b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit_tests.rs @@ -40,9 +40,13 @@ fn split_audio_graphs_preserve_stage_and_weight_boundaries() { assert!(encode.contains("\"stablehlo.reduce_window\"")); assert!(!encode.contains("loc(\"audio.hard_embeddings\")")); assert!(!encode.contains("model.language_model.")); - assert_eq!(encode_layout.stages.len(), 3 + 12 * 5 + 10); + assert_eq!(encode_layout.stages.len(), 3 + 4 + 12 * 5 + 10); for name in [ + "sscp_conv_0_convolution", + "sscp_conv_0_norm", "sscp_conv_0", + "sscp_conv_1_convolution", + "sscp_conv_1_norm", "sscp_conv_1", "input_projection", "conformer.0.feed_forward_start", diff --git a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_ops.rs b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_ops.rs index d435f39a5..d7d20491b 100644 --- a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_ops.rs +++ b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_ops.rs @@ -8,12 +8,22 @@ use super::gemma3n_audio_schema::EncoderArgs; use crate::Gemma3nXlaAudioConfig; pub(super) struct SscpOutput { + pub conv0_convolution: Val, + pub conv0_norm: Val, pub conv0: Val, + pub conv1_convolution: Val, + pub conv1_norm: Val, pub conv1: Val, pub hidden: Val, pub valid: Val, } +struct SscpConvStages { + convolution: Val, + norm: Val, + activated: Val, +} + pub(super) struct ConformerStages { pub feed_forward_start: Val, pub attention: Val, @@ -71,7 +81,7 @@ fn sscp_conv( stride: [usize; 2], groups: usize, eps: f32, -) -> Val { +) -> SscpConvStages { let zero = b.const_f32(0.0); let padded = b.pad(x, &zero, &[0, 0, 1, 0], &[0, kernel[0] - 1, 1, 0]); let weight = b.transpose(weight, &[1, 2, 3, 0]); @@ -84,14 +94,19 @@ fn sscp_conv( let convolved = round_bf16(b, &convolved); let normalized = cumulative_group_norm(b, &convolved, norm, eps); let activated = relu(b, &normalized); - round_bf16(b, &activated) + let activated = round_bf16(b, &activated); + SscpConvStages { + convolution: convolved, + norm: normalized, + activated, + } } fn subsample_convs( b: &mut Builder, args: &EncoderArgs, config: &Gemma3nXlaAudioConfig, -) -> (Val, Val) { +) -> (SscpConvStages, SscpConvStages) { let batch = args.mel.ty.shape[0]; let time = args.mel.ty.shape[1]; let frequency = args.mel.ty.shape[2]; @@ -116,7 +131,7 @@ fn subsample_convs( ); let conv1 = sscp_conv( b, - &conv0, + &conv0.activated, args.weight(&format!("{root}.conv_1.conv.weight")), args.weight(&format!("{root}.conv_1.norm.weight")), config.sscp_conv_kernel_size[1], @@ -134,8 +149,11 @@ pub(super) fn subsample( ) -> Val { let (_, conv1) = subsample_convs(b, args, config); let root = "audio_tower.subsample_conv_projection"; - let shape = &conv1.ty.shape; - let flattened = b.reshape(&conv1, vec![shape[0], shape[1], shape[2] * shape[3]]); + let shape = &conv1.activated.ty.shape; + let flattened = b.reshape( + &conv1.activated, + vec![shape[0], shape[1], shape[2] * shape[3]], + ); let projected = b.linear_last( &flattened, args.weight(&format!("{root}.input_proj_linear.weight")), @@ -150,8 +168,11 @@ pub(super) fn subsample_with_stages( ) -> SscpOutput { let (conv0, conv1) = subsample_convs(b, args, config); let root = "audio_tower.subsample_conv_projection"; - let shape = &conv1.ty.shape; - let flattened = b.reshape(&conv1, vec![shape[0], shape[1], shape[2] * shape[3]]); + let shape = &conv1.activated.ty.shape; + let flattened = b.reshape( + &conv1.activated, + vec![shape[0], shape[1], shape[2] * shape[3]], + ); let projected = b.linear_last( &flattened, args.weight(&format!("{root}.input_proj_linear.weight")), @@ -163,8 +184,12 @@ pub(super) fn subsample_with_stages( config.time_stride_product().expect("validated stride"), ); SscpOutput { - conv0, - conv1, + conv0_convolution: conv0.convolution, + conv0_norm: conv0.norm, + conv0: conv0.activated, + conv1_convolution: conv1.convolution, + conv1_norm: conv1.norm, + conv1: conv1.activated, hidden, valid, } diff --git a/src/lib/mlxcel-xla/src/gemma3n_audio_runtime.rs b/src/lib/mlxcel-xla/src/gemma3n_audio_runtime.rs index 45a4942e0..da73f27e1 100644 --- a/src/lib/mlxcel-xla/src/gemma3n_audio_runtime.rs +++ b/src/lib/mlxcel-xla/src/gemma3n_audio_runtime.rs @@ -651,7 +651,14 @@ impl Gemma3nAudioIreeRuntime { report_load_timing("audio-total", load_started); #[cfg(feature = "diagnostics")] let diagnostic_shapes = [ + "sscp_conv_0_convolution", + "sscp_conv_0_norm", "sscp_conv_0", + "sscp_conv_1_convolution", + "sscp_conv_1_norm", + "sscp_conv_1", + "input_projection", + "conformer.0.feed_forward_start", "encoded_reduced", "soft_norm", "soft_linear", From bef382691c8471f172699c87895138ae962c3c0c Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Sun, 26 Jul 2026 23:49:01 +0900 Subject: [PATCH 13/19] fix(xla): match MLX CUDA cumulative norm schedule The modified actual diagnostic is bit-exact through the first SSCP convolution and first diverges in cumulative normalization. StableHLO's generic reductions and reduce-window prefix sum select backend-dependent association trees, while the pinned MLX CUDA oracle uses four-value lane accumulators, XOR warp reductions, and hierarchical warp scans. Lower only the Gemma3n SSCP cumulative norm sums and scans with the MLX CUDA association order. Keep convolution, F32 arithmetic, epsilon, BF16 materialization, and gate thresholds unchanged, and add a focused local IREE compile regression for the schedule. Validation: focused Gemma3n audio emitter tests; mlxcel-xla check; xla-diagnostics example check and three tests; focused IREE local LLVM CPU compile; rustfmt and diff check. Strict clippy remains blocked by pre-existing unrelated warnings; regular package clippy reports no new warning. The actual-checkpoint gate was not rerun. Refs #878 --- .../src/emitter/gemma3n_audio_emit_tests.rs | 31 ++- .../src/emitter/gemma3n_audio_math.rs | 192 +++++++++++++++++- .../src/emitter/gemma3n_audio_ops.rs | 16 +- 3 files changed, 225 insertions(+), 14 deletions(-) diff --git a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit_tests.rs b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit_tests.rs index 93119092e..a23974b6b 100644 --- a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit_tests.rs +++ b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit_tests.rs @@ -37,7 +37,11 @@ fn split_audio_graphs_preserve_stage_and_weight_boundaries() { assert!(encode.starts_with("module @audio_encode {")); assert!(encode.contains("stablehlo.convolution")); - assert!(encode.contains("\"stablehlo.reduce_window\"")); + assert!(!encode.contains("\"stablehlo.reduce_window\"")); + assert!( + encode.matches("\"stablehlo.gather\"").count() >= 40, + "both SSCP norms must pin four MLX CUDA XOR reduction trees" + ); assert!(!encode.contains("loc(\"audio.hard_embeddings\")")); assert!(!encode.contains("model.language_model.")); assert_eq!(encode_layout.stages.len(), 3 + 4 + 12 * 5 + 10); @@ -188,6 +192,31 @@ fn sscp_f32_accumulator_contract_compiles_for_cpu() { compile(compiler.as_os_str(), "local", &stem, "contract", mlir); } +#[test] +#[ignore = "requires IREE_DIST"] +fn sscp_cuda_norm_schedule_compiles_for_cpu() { + use crate::emitter::gemma3n_audio_math::{mlx_cuda_cumsum_time_f32, mlx_cuda_row_sum_f32}; + + let mut builder = Builder::new().with_precision(Precision::Bf16); + let input = Builder::arg(0, Ty::f32(vec![1, 8, 64, 128])); + let channels = mlx_cuda_row_sum_f32(&mut builder, &input, 3); + let frequency = mlx_cuda_row_sum_f32(&mut builder, &channels, 2); + let per_time = builder.reshape(&frequency, vec![1, 8, 1, 1]); + let output = mlx_cuda_cumsum_time_f32(&mut builder, &per_time); + let output_ty = output.ty.render(); + let mlir = format!( + "module @sscp_cuda_norm_schedule {{\n func.func public @main(%arg0: \ + tensor<1x8x64x128xf32>) -> {output_ty} {{\n{} return {} : {output_ty}\n \ + }}\n}}\n", + builder.body(), + output.name, + ); + let compiler = std::path::PathBuf::from(std::env::var_os("IREE_DIST").expect("IREE_DIST")) + .join("bin/iree-compile"); + let stem = format!("mlxcel-gemma3n-sscp-cuda-norm-{}", std::process::id()); + compile(compiler.as_os_str(), "local", &stem, "contract", mlir); +} + fn compile(compiler: &std::ffi::OsStr, device: &str, stem: &str, artifact: &str, mlir: String) { let input = std::env::temp_dir().join(format!("{stem}-{artifact}.mlir")); let output = std::env::temp_dir().join(format!("{stem}-{artifact}.vmfb")); diff --git a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_math.rs b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_math.rs index 7f9581cef..cad1dcb10 100644 --- a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_math.rs +++ b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_math.rs @@ -16,6 +16,193 @@ pub(super) fn round_bf16(b: &mut Builder, x: &Val) -> Val { b.convert(&narrow, "f32") } +fn slice_axis(b: &mut Builder, x: &Val, axis: usize, start: usize, limit: usize) -> Val { + let ranges = + x.ty.shape + .iter() + .enumerate() + .map(|(current, &size)| { + if current == axis { + (start, limit) + } else { + (0, size) + } + }) + .collect::>(); + b.slice(x, &ranges) +} + +fn gather_axis(b: &mut Builder, x: &Val, axis: usize, indices: &[usize]) -> Val { + assert_eq!(indices.len(), x.ty.shape[axis]); + let mut permutation = vec![axis]; + permutation.extend((0..x.ty.shape.len()).filter(|¤t| current != axis)); + let transposed = b.transpose(x, &permutation); + let width = transposed.ty.shape[0]; + let remainder = transposed.ty.shape[1..].iter().product::(); + let flattened = b.reshape(&transposed, vec![width, remainder]); + let indices = indices + .iter() + .map(|&index| i32::try_from(index).expect("static axis index must fit i32")) + .collect::>(); + let indices = b.const_tensor_i32(&indices, vec![width, 1]); + let gathered = b.gather(&flattened, &indices); + let gathered = b.reshape(&gathered, transposed.ty.shape); + let mut inverse = vec![0; permutation.len()]; + for (transposed_axis, original_axis) in permutation.into_iter().enumerate() { + inverse[original_axis] = transposed_axis; + } + b.transpose(&gathered, &inverse) +} + +/// Match MLX CUDA's f32 `row_reduce_simple` schedule for the SSCP normalization +/// widths. Each CUDA lane accumulates four adjacent values serially, then the +/// 32-lane tile combines its partials through XOR shuffles (16, 8, 4, 2, 1). +/// +/// StableHLO reductions deliberately leave the tree to the backend. That is a +/// poor fit for the CUDA-produced MLX oracle: cancellation in cumulative group +/// normalization can move a BF16-rounded result when the CPU backend selects a +/// different tree. Keep the CUDA schedule local to this normalization instead +/// of changing graph-wide reductions. +pub(super) fn mlx_cuda_row_sum_f32(b: &mut Builder, x: &Val, axis: usize) -> Val { + assert_eq!(x.ty.elt, "f32"); + assert!(axis < x.ty.shape.len()); + let width = x.ty.shape[axis]; + assert!( + (1..=128).contains(&width), + "MLX CUDA row-sum emulation supports one 32-lane block" + ); + + let padded = if width == 128 { + x.clone() + } else { + let zero = b.const_f32(0.0); + let low = vec![0; x.ty.shape.len()]; + let mut high = low.clone(); + high[axis] = 128 - width; + b.pad(x, &zero, &low, &high) + }; + let mut grouped_shape = padded.ty.shape.clone(); + grouped_shape[axis] = 32; + grouped_shape.insert(axis + 1, 4); + let grouped = b.reshape(&padded, grouped_shape); + let lane0 = slice_axis(b, &grouped, axis + 1, 0, 1); + let lane1 = slice_axis(b, &grouped, axis + 1, 1, 2); + let lane2 = slice_axis(b, &grouped, axis + 1, 2, 3); + let lane3 = slice_axis(b, &grouped, axis + 1, 3, 4); + let zero = zeros_like(b, &lane0); + let local0 = b.add(&zero, &lane0); + let local1 = b.add(&local0, &lane1); + let local2 = b.add(&local1, &lane2); + let local3 = b.add(&local2, &lane3); + + let mut partial_shape = padded.ty.shape.clone(); + partial_shape[axis] = 32; + let mut partial = b.reshape(&local3, partial_shape); + for mask in [16, 8, 4, 2, 1] { + let shuffled_indices = (0..32).map(|lane| lane ^ mask).collect::>(); + let shuffled = gather_axis(b, &partial, axis, &shuffled_indices); + partial = b.add(&partial, &shuffled); + } + + let reduced = slice_axis(b, &partial, axis, 0, 1); + let mut output_shape = x.ty.shape.clone(); + output_shape.remove(axis); + b.reshape(&reduced, output_shape) +} + +fn inclusive_scan_32_f32(b: &mut Builder, x: &Val, axis: usize) -> Val { + assert_eq!(x.ty.elt, "f32"); + assert_eq!(x.ty.shape[axis], 32); + let zero = b.const_f32(0.0); + let mut result = x.clone(); + for offset in [1, 2, 4, 8, 16] { + let prior = slice_axis(b, &result, axis, 0, 32 - offset); + let mut low = vec![0; result.ty.shape.len()]; + let high = low.clone(); + low[axis] = offset; + let prior = b.pad(&prior, &zero, &low, &high); + // CUDA cooperative_groups uses `out = op(out, shfl_up(out))`. + result = b.add(&result, &prior); + } + result +} + +fn exclusive_from_inclusive_32_f32(b: &mut Builder, inclusive: &Val, axis: usize) -> Val { + let prior = slice_axis(b, inclusive, axis, 0, 31); + let zero = b.const_f32(0.0); + let mut low = vec![0; inclusive.ty.shape.len()]; + let high = low.clone(); + low[axis] = 1; + b.pad(&prior, &zero, &low, &high) +} + +/// Match the one-block MLX CUDA contiguous inclusive-scan schedule used by the +/// maintained Gemma3n oracle. Released frame buckets produce at most 1,499 +/// values after the first SSCP convolution, below CUDA's 4,096-value block +/// limit, so no inter-block prefix is needed. +pub(super) fn mlx_cuda_cumsum_time_f32(b: &mut Builder, x: &Val) -> Val { + assert_eq!(x.ty.elt, "f32"); + assert_eq!(x.ty.shape.len(), 4); + let batch = x.ty.shape[0]; + let time = x.ty.shape[1]; + assert!( + (1..=4096).contains(&time), + "MLX CUDA cumsum emulation supports one scan block" + ); + let padded_time = time.div_ceil(128) * 128; + let padded = if padded_time == time { + x.clone() + } else { + let zero = b.const_f32(0.0); + b.pad(x, &zero, &[0, 0, 0, 0], &[0, padded_time - time, 0, 0]) + }; + let threads = padded_time / 4; + let warps = threads / 32; + let grouped = b.reshape(&padded, vec![batch, threads, 4, 1, 1]); + let value0 = slice_axis(b, &grouped, 2, 0, 1); + let value1 = slice_axis(b, &grouped, 2, 1, 2); + let value2 = slice_axis(b, &grouped, 2, 2, 3); + let value3 = slice_axis(b, &grouped, 2, 3, 4); + let local0 = value0; + let local1 = b.add(&value1, &local0); + let local2 = b.add(&value2, &local1); + let local3 = b.add(&value3, &local2); + let local = concat_axis(b, &[local0, local1, local2, local3.clone()], 2); + + let thread_sums = b.reshape(&local3, vec![batch, warps, 32, 1, 1]); + let thread_inclusive = inclusive_scan_32_f32(b, &thread_sums, 2); + let thread_exclusive = exclusive_from_inclusive_32_f32(b, &thread_inclusive, 2); + + // The MLX kernel forms each warp total as + // `exclusive_scan(last_lane) + last_lane_value`, rather than taking the + // inclusive result of lane 31. Preserve that distinct association. + let prior_last = slice_axis(b, &thread_exclusive, 2, 31, 32); + let current_last = slice_axis(b, &thread_sums, 2, 31, 32); + let warp_totals = b.add(&prior_last, ¤t_last); + let warp_totals = b.reshape(&warp_totals, vec![batch, warps, 1, 1]); + let zero = b.const_f32(0.0); + let warp_totals_32 = b.pad(&warp_totals, &zero, &[0, 0, 0, 0], &[0, 32 - warps, 0, 0]); + let warp_inclusive = inclusive_scan_32_f32(b, &warp_totals_32, 1); + let warp_exclusive = exclusive_from_inclusive_32_f32(b, &warp_inclusive, 1); + let warp_exclusive = slice_axis(b, &warp_exclusive, 1, 0, warps); + + let local = b.reshape(&local, vec![batch, warps, 32, 4, 1, 1]); + let prefix = zeros_like(b, &local); + let local = b.add(&local, &prefix); + let warp_exclusive = b.reshape(&warp_exclusive, vec![batch, warps, 1, 1, 1, 1]); + let warp_exclusive = b.broadcast(&warp_exclusive, &[0, 1, 2, 3, 4, 5], local.ty.shape.clone()); + let with_warp_prefix = b.add(&local, &warp_exclusive); + let thread_exclusive = b.reshape(&thread_exclusive, vec![batch, warps, 32, 1, 1, 1]); + let thread_exclusive = b.broadcast( + &thread_exclusive, + &[0, 1, 2, 3, 4, 5], + with_warp_prefix.ty.shape.clone(), + ); + let scanned = b.add(&with_warp_prefix, &thread_exclusive); + let scanned = b.reshape(&scanned, vec![batch, padded_time, 1, 1]); + slice_axis(b, &scanned, 1, 0, time) +} + pub(super) fn clip(b: &mut Builder, x: &Val, limit: f32) -> Val { let low = scalar_like(b, -limit, x); let high = scalar_like(b, limit, x); @@ -71,11 +258,6 @@ pub(super) fn rms_norm(b: &mut Builder, x: &Val, weight: Option<&Val>, eps: f32) } } -pub(super) fn cumsum_time_f32(b: &mut Builder, x: &Val) -> Val { - assert_eq!(x.ty.shape.len(), 4); - b.reduce_window_prefix_add(x, 1) -} - pub(super) fn softmax_last(b: &mut Builder, x: &Val) -> Val { let last = x.ty.shape.len() - 1; let negative_infinity = b.const_f32(f32::NEG_INFINITY); diff --git a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_ops.rs b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_ops.rs index d7d20491b..fe727ba0a 100644 --- a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_ops.rs +++ b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_ops.rs @@ -2,7 +2,8 @@ use super::builder::{Builder, Val}; use super::gemma3n_audio_math::{ - clip, cumsum_time_f32, relu, rms_norm, round_bf16, sigmoid, silu, stride_time, + clip, mlx_cuda_cumsum_time_f32, mlx_cuda_row_sum_f32, relu, rms_norm, round_bf16, sigmoid, + silu, stride_time, }; use super::gemma3n_audio_schema::EncoderArgs; use crate::Gemma3nXlaAudioConfig; @@ -37,11 +38,10 @@ fn cumulative_group_norm(b: &mut Builder, x: &Val, weight: &Val, eps: f32) -> Va let time = x.ty.shape[1]; let frequency = x.ty.shape[2]; let channels = x.ty.shape[3]; - let zero = b.const_f32(0.0); - let sum_channels = b.reduce_add(x, 3, &zero); - let sum_frequency = b.reduce_add(&sum_channels, 2, &zero); + let sum_channels = mlx_cuda_row_sum_f32(b, x, 3); + let sum_frequency = mlx_cuda_row_sum_f32(b, &sum_channels, 2); let sum_at_time = b.reshape(&sum_frequency, vec![batch, time, 1, 1]); - let cumulative_sum = cumsum_time_f32(b, &sum_at_time); + let cumulative_sum = mlx_cuda_cumsum_time_f32(b, &sum_at_time); let index = b.iota(time); let index = b.convert(&index, "f32"); @@ -56,10 +56,10 @@ fn cumulative_group_norm(b: &mut Builder, x: &Val, weight: &Val, eps: f32) -> Va let centered = b.subtract(x, &mean); let squared = b.multiply(¢ered, ¢ered); - let squared_channels = b.reduce_add(&squared, 3, &zero); - let squared_frequency = b.reduce_add(&squared_channels, 2, &zero); + let squared_channels = mlx_cuda_row_sum_f32(b, &squared, 3); + let squared_frequency = mlx_cuda_row_sum_f32(b, &squared_channels, 2); let squared_at_time = b.reshape(&squared_frequency, vec![batch, time, 1, 1]); - let cumulative_squared = cumsum_time_f32(b, &squared_at_time); + let cumulative_squared = mlx_cuda_cumsum_time_f32(b, &squared_at_time); let variance = b.divide(&cumulative_squared, &count); let epsilon = super::gemma3n_audio_math::scalar_like(b, eps, &variance); let variance = b.add(&variance, &epsilon); From e7f92ecb635fd767929101269b11dc6b19367cf8 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 00:06:50 +0900 Subject: [PATCH 14/19] test(xla): expose cumulative norm arithmetic stages The MLX CUDA reduction/scan schedule fix reduced the first SSCP norm drift from RMS 7.032964031e-5, max BF16 ULP 16, and six over-one-ULP values to RMS 2.938648021e-5, max BF16 ULP one, and no over-one-ULP values. The unchanged absolute gate still fails at a one-ULP pair, while the preceding convolution remains bit-exact. Expose the first norm's per-time sum, cumulative sum, mean, squared sum, cumulative squared sum, variance, epsilon-stabilized variance, and inverse standard deviation on both MLX and IREE diagnostic paths. These bounded tensors distinguish remaining reduction/scan drift from epsilon addition or rsqrt without adding another full-size activation capture. Validation: focused Gemma3n audio emitter tests; xla-diagnostics example check and three tests; focused IREE local LLVM CPU compile; package and diagnostic-example clippy; rustfmt and diff check. No actual-checkpoint run was performed. Refs #878 --- examples/xla_gemma3n_audio_reference_check.rs | 32 +++++++ src/audio/gemma3n/diagnostics.rs | 22 +++++ src/audio/gemma3n/encoder.rs | 86 +++++++++++++++++-- .../src/emitter/gemma3n_audio_emit.rs | 34 ++++++++ .../src/emitter/gemma3n_audio_emit_tests.rs | 10 ++- .../src/emitter/gemma3n_audio_ops.rs | 78 +++++++++++++++-- .../mlxcel-xla/src/gemma3n_audio_runtime.rs | 8 ++ 7 files changed, 255 insertions(+), 15 deletions(-) diff --git a/examples/xla_gemma3n_audio_reference_check.rs b/examples/xla_gemma3n_audio_reference_check.rs index 4d48534dc..e93c779d8 100644 --- a/examples/xla_gemma3n_audio_reference_check.rs +++ b/examples/xla_gemma3n_audio_reference_check.rs @@ -329,6 +329,38 @@ fn main() -> Result<(), String> { "sscp_conv_0_convolution", mlx.sscp_conv_0_convolution.as_slice(), ), + ( + "sscp_conv_0_norm_sum_at_time", + mlx.sscp_conv_0_norm_sum_at_time.as_slice(), + ), + ( + "sscp_conv_0_norm_cumulative_sum", + mlx.sscp_conv_0_norm_cumulative_sum.as_slice(), + ), + ( + "sscp_conv_0_norm_mean", + mlx.sscp_conv_0_norm_mean.as_slice(), + ), + ( + "sscp_conv_0_norm_squared_at_time", + mlx.sscp_conv_0_norm_squared_at_time.as_slice(), + ), + ( + "sscp_conv_0_norm_cumulative_squared", + mlx.sscp_conv_0_norm_cumulative_squared.as_slice(), + ), + ( + "sscp_conv_0_norm_variance", + mlx.sscp_conv_0_norm_variance.as_slice(), + ), + ( + "sscp_conv_0_norm_stabilized_variance", + mlx.sscp_conv_0_norm_stabilized_variance.as_slice(), + ), + ( + "sscp_conv_0_norm_inverse_stddev", + mlx.sscp_conv_0_norm_inverse_stddev.as_slice(), + ), ("sscp_conv_0_norm", mlx.sscp_conv_0_norm.as_slice()), ("sscp_conv_0", mlx.sscp_conv_0.as_slice()), ( diff --git a/src/audio/gemma3n/diagnostics.rs b/src/audio/gemma3n/diagnostics.rs index a5ca90f50..17d5f8eac 100644 --- a/src/audio/gemma3n/diagnostics.rs +++ b/src/audio/gemma3n/diagnostics.rs @@ -26,6 +26,14 @@ use crate::models::gemma3n::{Gemma3nAudioEmbedder, ModelArgs, TextConfig}; #[derive(Debug, Clone)] pub struct Gemma3nAudioMlxDiagnosticOutput { pub sscp_conv_0_convolution: Vec, + pub sscp_conv_0_norm_sum_at_time: Vec, + pub sscp_conv_0_norm_cumulative_sum: Vec, + pub sscp_conv_0_norm_mean: Vec, + pub sscp_conv_0_norm_squared_at_time: Vec, + pub sscp_conv_0_norm_cumulative_squared: Vec, + pub sscp_conv_0_norm_variance: Vec, + pub sscp_conv_0_norm_stabilized_variance: Vec, + pub sscp_conv_0_norm_inverse_stddev: Vec, pub sscp_conv_0_norm: Vec, pub sscp_conv_0: Vec, pub sscp_conv_1_convolution: Vec, @@ -324,6 +332,20 @@ pub fn run_gemma3n_audio_mlx_diagnostics( Ok(Gemma3nAudioMlxDiagnosticOutput { sscp_conv_0_convolution: array_f32(&encoder_stages.sscp_conv_0_convolution), + sscp_conv_0_norm_sum_at_time: array_f32(&encoder_stages.sscp_conv_0_norm_sum_at_time), + sscp_conv_0_norm_cumulative_sum: array_f32(&encoder_stages.sscp_conv_0_norm_cumulative_sum), + sscp_conv_0_norm_mean: array_f32(&encoder_stages.sscp_conv_0_norm_mean), + sscp_conv_0_norm_squared_at_time: array_f32( + &encoder_stages.sscp_conv_0_norm_squared_at_time, + ), + sscp_conv_0_norm_cumulative_squared: array_f32( + &encoder_stages.sscp_conv_0_norm_cumulative_squared, + ), + sscp_conv_0_norm_variance: array_f32(&encoder_stages.sscp_conv_0_norm_variance), + sscp_conv_0_norm_stabilized_variance: array_f32( + &encoder_stages.sscp_conv_0_norm_stabilized_variance, + ), + sscp_conv_0_norm_inverse_stddev: array_f32(&encoder_stages.sscp_conv_0_norm_inverse_stddev), sscp_conv_0_norm: array_f32(&encoder_stages.sscp_conv_0_norm), sscp_conv_0: array_f32(&encoder_stages.sscp_conv_0), sscp_conv_1_convolution: array_f32(&encoder_stages.sscp_conv_1_convolution), diff --git a/src/audio/gemma3n/encoder.rs b/src/audio/gemma3n/encoder.rs index c0bbd7f53..68c507020 100644 --- a/src/audio/gemma3n/encoder.rs +++ b/src/audio/gemma3n/encoder.rs @@ -69,6 +69,26 @@ struct CumulativeGroupNorm { eps: f32, } +struct CumulativeGroupNormStages { + output: UniquePtr, + #[cfg(feature = "xla-diagnostics")] + sum_at_time: UniquePtr, + #[cfg(feature = "xla-diagnostics")] + cumulative_sum: UniquePtr, + #[cfg(feature = "xla-diagnostics")] + mean: UniquePtr, + #[cfg(feature = "xla-diagnostics")] + squared_at_time: UniquePtr, + #[cfg(feature = "xla-diagnostics")] + cumulative_squared: UniquePtr, + #[cfg(feature = "xla-diagnostics")] + variance: UniquePtr, + #[cfg(feature = "xla-diagnostics")] + stabilized_variance: UniquePtr, + #[cfg(feature = "xla-diagnostics")] + inverse_stddev: UniquePtr, +} + impl CumulativeGroupNorm { fn from_weights( weights: &WeightMap, @@ -85,6 +105,10 @@ impl CumulativeGroupNorm { /// 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 { + self.forward_stages(hidden_states).output + } + + fn forward_stages(&self, hidden_states: &MlxArray) -> CumulativeGroupNormStages { 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); @@ -108,10 +132,30 @@ impl CumulativeGroupNorm { 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 stabilized_variance = mlxcel_core::add(&variance, &eps); + let inverse_stddev = mlxcel_core::rsqrt(&stabilized_variance); let normalized = mlxcel_core::multiply(¢ered, &inverse_stddev); let normalized = mlxcel_core::multiply(&normalized, &self.weight); - mlxcel_core::astype(&normalized, input_dtype) + let output = mlxcel_core::astype(&normalized, input_dtype); + CumulativeGroupNormStages { + output, + #[cfg(feature = "xla-diagnostics")] + sum_at_time, + #[cfg(feature = "xla-diagnostics")] + cumulative_sum, + #[cfg(feature = "xla-diagnostics")] + mean: cumulative_mean, + #[cfg(feature = "xla-diagnostics")] + squared_at_time, + #[cfg(feature = "xla-diagnostics")] + cumulative_squared, + #[cfg(feature = "xla-diagnostics")] + variance, + #[cfg(feature = "xla-diagnostics")] + stabilized_variance, + #[cfg(feature = "xla-diagnostics")] + inverse_stddev, + } } } @@ -126,6 +170,14 @@ struct SscpConvBlock { #[cfg(feature = "xla-diagnostics")] struct SscpConvDiagnosticStages { convolution: UniquePtr, + norm_sum_at_time: UniquePtr, + norm_cumulative_sum: UniquePtr, + norm_mean: UniquePtr, + norm_squared_at_time: UniquePtr, + norm_cumulative_squared: UniquePtr, + norm_variance: UniquePtr, + norm_stabilized_variance: UniquePtr, + norm_inverse_stddev: UniquePtr, norm: UniquePtr, activated: UniquePtr, } @@ -196,11 +248,19 @@ impl SscpConvBlock { #[cfg(feature = "xla-diagnostics")] fn forward_with_diagnostics(&self, x: &MlxArray) -> Result { let convolution = self.convolution(x)?; - let norm = self.norm.forward(&convolution); - let activated = mlxcel_core::relu(&norm); + let norm = self.norm.forward_stages(&convolution); + let activated = mlxcel_core::relu(&norm.output); Ok(SscpConvDiagnosticStages { convolution, - norm, + norm_sum_at_time: norm.sum_at_time, + norm_cumulative_sum: norm.cumulative_sum, + norm_mean: norm.mean, + norm_squared_at_time: norm.squared_at_time, + norm_cumulative_squared: norm.cumulative_squared, + norm_variance: norm.variance, + norm_stabilized_variance: norm.stabilized_variance, + norm_inverse_stddev: norm.inverse_stddev, + norm: norm.output, activated, }) } @@ -639,6 +699,14 @@ pub struct Gemma3nAudioEncoder { #[cfg(feature = "xla-diagnostics")] pub(super) struct Gemma3nAudioEncoderDiagnosticStages { pub sscp_conv_0_convolution: UniquePtr, + pub sscp_conv_0_norm_sum_at_time: UniquePtr, + pub sscp_conv_0_norm_cumulative_sum: UniquePtr, + pub sscp_conv_0_norm_mean: UniquePtr, + pub sscp_conv_0_norm_squared_at_time: UniquePtr, + pub sscp_conv_0_norm_cumulative_squared: UniquePtr, + pub sscp_conv_0_norm_variance: UniquePtr, + pub sscp_conv_0_norm_stabilized_variance: UniquePtr, + pub sscp_conv_0_norm_inverse_stddev: UniquePtr, pub sscp_conv_0_norm: UniquePtr, pub sscp_conv_0: UniquePtr, pub sscp_conv_1_convolution: UniquePtr, @@ -826,6 +894,14 @@ impl Gemma3nAudioEncoder { invalid, Gemma3nAudioEncoderDiagnosticStages { sscp_conv_0_convolution: subsample.conv0.convolution, + sscp_conv_0_norm_sum_at_time: subsample.conv0.norm_sum_at_time, + sscp_conv_0_norm_cumulative_sum: subsample.conv0.norm_cumulative_sum, + sscp_conv_0_norm_mean: subsample.conv0.norm_mean, + sscp_conv_0_norm_squared_at_time: subsample.conv0.norm_squared_at_time, + sscp_conv_0_norm_cumulative_squared: subsample.conv0.norm_cumulative_squared, + sscp_conv_0_norm_variance: subsample.conv0.norm_variance, + sscp_conv_0_norm_stabilized_variance: subsample.conv0.norm_stabilized_variance, + sscp_conv_0_norm_inverse_stddev: subsample.conv0.norm_inverse_stddev, sscp_conv_0_norm: subsample.conv0.norm, sscp_conv_0: subsample.conv0.activated, sscp_conv_1_convolution: subsample.conv1.convolution, diff --git a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit.rs b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit.rs index e915f90e9..bcc809217 100644 --- a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit.rs +++ b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit.rs @@ -89,6 +89,32 @@ fn audio_encoder( ) -> (Val, Val, Val) { let subsampled = subsample_with_stages(b, args, audio); trace.push("sscp_conv_0_convolution", &subsampled.conv0_convolution); + trace.push( + "sscp_conv_0_norm_sum_at_time", + &subsampled.conv0_norm_sum_at_time, + ); + trace.push( + "sscp_conv_0_norm_cumulative_sum", + &subsampled.conv0_norm_cumulative_sum, + ); + trace.push("sscp_conv_0_norm_mean", &subsampled.conv0_norm_mean); + trace.push( + "sscp_conv_0_norm_squared_at_time", + &subsampled.conv0_norm_squared_at_time, + ); + trace.push( + "sscp_conv_0_norm_cumulative_squared", + &subsampled.conv0_norm_cumulative_squared, + ); + trace.push("sscp_conv_0_norm_variance", &subsampled.conv0_norm_variance); + trace.push( + "sscp_conv_0_norm_stabilized_variance", + &subsampled.conv0_norm_stabilized_variance, + ); + trace.push( + "sscp_conv_0_norm_inverse_stddev", + &subsampled.conv0_norm_inverse_stddev, + ); trace.push("sscp_conv_0_norm", &subsampled.conv0_norm); trace.push("sscp_conv_0", &subsampled.conv0); trace.push("sscp_conv_1_convolution", &subsampled.conv1_convolution); @@ -320,6 +346,14 @@ fn build_encode( #[cfg(feature = "diagnostics")] let diagnostics = [ "sscp_conv_0_convolution", + "sscp_conv_0_norm_sum_at_time", + "sscp_conv_0_norm_cumulative_sum", + "sscp_conv_0_norm_mean", + "sscp_conv_0_norm_squared_at_time", + "sscp_conv_0_norm_cumulative_squared", + "sscp_conv_0_norm_variance", + "sscp_conv_0_norm_stabilized_variance", + "sscp_conv_0_norm_inverse_stddev", "sscp_conv_0_norm", "sscp_conv_0", "sscp_conv_1_convolution", diff --git a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit_tests.rs b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit_tests.rs index a23974b6b..5ae7e5ead 100644 --- a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit_tests.rs +++ b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit_tests.rs @@ -44,9 +44,17 @@ fn split_audio_graphs_preserve_stage_and_weight_boundaries() { ); assert!(!encode.contains("loc(\"audio.hard_embeddings\")")); assert!(!encode.contains("model.language_model.")); - assert_eq!(encode_layout.stages.len(), 3 + 4 + 12 * 5 + 10); + assert_eq!(encode_layout.stages.len(), 3 + 4 + 12 * 5 + 10 + 8); for name in [ "sscp_conv_0_convolution", + "sscp_conv_0_norm_sum_at_time", + "sscp_conv_0_norm_cumulative_sum", + "sscp_conv_0_norm_mean", + "sscp_conv_0_norm_squared_at_time", + "sscp_conv_0_norm_cumulative_squared", + "sscp_conv_0_norm_variance", + "sscp_conv_0_norm_stabilized_variance", + "sscp_conv_0_norm_inverse_stddev", "sscp_conv_0_norm", "sscp_conv_0", "sscp_conv_1_convolution", diff --git a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_ops.rs b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_ops.rs index fe727ba0a..471c50d3c 100644 --- a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_ops.rs +++ b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_ops.rs @@ -10,6 +10,14 @@ use crate::Gemma3nXlaAudioConfig; pub(super) struct SscpOutput { pub conv0_convolution: Val, + pub conv0_norm_sum_at_time: Val, + pub conv0_norm_cumulative_sum: Val, + pub conv0_norm_mean: Val, + pub conv0_norm_squared_at_time: Val, + pub conv0_norm_cumulative_squared: Val, + pub conv0_norm_variance: Val, + pub conv0_norm_stabilized_variance: Val, + pub conv0_norm_inverse_stddev: Val, pub conv0_norm: Val, pub conv0: Val, pub conv1_convolution: Val, @@ -21,10 +29,30 @@ pub(super) struct SscpOutput { struct SscpConvStages { convolution: Val, + norm_sum_at_time: Val, + norm_cumulative_sum: Val, + norm_mean: Val, + norm_squared_at_time: Val, + norm_cumulative_squared: Val, + norm_variance: Val, + norm_stabilized_variance: Val, + norm_inverse_stddev: Val, norm: Val, activated: Val, } +struct CumulativeGroupNormStages { + sum_at_time: Val, + cumulative_sum: Val, + mean: Val, + squared_at_time: Val, + cumulative_squared: Val, + variance: Val, + stabilized_variance: Val, + inverse_stddev: Val, + output: Val, +} + pub(super) struct ConformerStages { pub feed_forward_start: Val, pub attention: Val, @@ -33,7 +61,12 @@ pub(super) struct ConformerStages { pub final_norm: Val, } -fn cumulative_group_norm(b: &mut Builder, x: &Val, weight: &Val, eps: f32) -> Val { +fn cumulative_group_norm( + b: &mut Builder, + x: &Val, + weight: &Val, + eps: f32, +) -> CumulativeGroupNormStages { let batch = x.ty.shape[0]; let time = x.ty.shape[1]; let frequency = x.ty.shape[2]; @@ -51,8 +84,8 @@ fn cumulative_group_norm(b: &mut Builder, x: &Val, weight: &Val, eps: f32) -> Va let count = b.multiply(&count, &width); let count = b.reshape(&count, vec![1, time, 1, 1]); let count = b.broadcast(&count, &[0, 1, 2, 3], vec![batch, time, 1, 1]); - let mean = b.divide(&cumulative_sum, &count); - let mean = b.broadcast(&mean, &[0, 1, 2, 3], x.ty.shape.clone()); + let mean_at_time = b.divide(&cumulative_sum, &count); + let mean = b.broadcast(&mean_at_time, &[0, 1, 2, 3], x.ty.shape.clone()); let centered = b.subtract(x, &mean); let squared = b.multiply(¢ered, ¢ered); @@ -62,13 +95,24 @@ fn cumulative_group_norm(b: &mut Builder, x: &Val, weight: &Val, eps: f32) -> Va let cumulative_squared = mlx_cuda_cumsum_time_f32(b, &squared_at_time); let variance = b.divide(&cumulative_squared, &count); let epsilon = super::gemma3n_audio_math::scalar_like(b, eps, &variance); - let variance = b.add(&variance, &epsilon); - let inverse = b.rsqrt(&variance); - let inverse = b.broadcast(&inverse, &[0, 1, 2, 3], x.ty.shape.clone()); + let stabilized_variance = b.add(&variance, &epsilon); + let inverse_stddev = b.rsqrt(&stabilized_variance); + let inverse = b.broadcast(&inverse_stddev, &[0, 1, 2, 3], x.ty.shape.clone()); let normalized = b.multiply(¢ered, &inverse); let weight = b.broadcast(weight, &[3], x.ty.shape.clone()); let weighted = b.multiply(&normalized, &weight); - round_bf16(b, &weighted) + let output = round_bf16(b, &weighted); + CumulativeGroupNormStages { + sum_at_time, + cumulative_sum, + mean: mean_at_time, + squared_at_time, + cumulative_squared, + variance, + stabilized_variance, + inverse_stddev, + output, + } } #[allow(clippy::too_many_arguments)] @@ -93,11 +137,19 @@ fn sscp_conv( b.convolution_f32_accumulate(&padded, &weight, &stride, &[(0, 0), (0, 0)], groups); let convolved = round_bf16(b, &convolved); let normalized = cumulative_group_norm(b, &convolved, norm, eps); - let activated = relu(b, &normalized); + let activated = relu(b, &normalized.output); let activated = round_bf16(b, &activated); SscpConvStages { convolution: convolved, - norm: normalized, + norm_sum_at_time: normalized.sum_at_time, + norm_cumulative_sum: normalized.cumulative_sum, + norm_mean: normalized.mean, + norm_squared_at_time: normalized.squared_at_time, + norm_cumulative_squared: normalized.cumulative_squared, + norm_variance: normalized.variance, + norm_stabilized_variance: normalized.stabilized_variance, + norm_inverse_stddev: normalized.inverse_stddev, + norm: normalized.output, activated, } } @@ -185,6 +237,14 @@ pub(super) fn subsample_with_stages( ); SscpOutput { conv0_convolution: conv0.convolution, + conv0_norm_sum_at_time: conv0.norm_sum_at_time, + conv0_norm_cumulative_sum: conv0.norm_cumulative_sum, + conv0_norm_mean: conv0.norm_mean, + conv0_norm_squared_at_time: conv0.norm_squared_at_time, + conv0_norm_cumulative_squared: conv0.norm_cumulative_squared, + conv0_norm_variance: conv0.norm_variance, + conv0_norm_stabilized_variance: conv0.norm_stabilized_variance, + conv0_norm_inverse_stddev: conv0.norm_inverse_stddev, conv0_norm: conv0.norm, conv0: conv0.activated, conv1_convolution: conv1.convolution, diff --git a/src/lib/mlxcel-xla/src/gemma3n_audio_runtime.rs b/src/lib/mlxcel-xla/src/gemma3n_audio_runtime.rs index da73f27e1..7bcc6188e 100644 --- a/src/lib/mlxcel-xla/src/gemma3n_audio_runtime.rs +++ b/src/lib/mlxcel-xla/src/gemma3n_audio_runtime.rs @@ -652,6 +652,14 @@ impl Gemma3nAudioIreeRuntime { #[cfg(feature = "diagnostics")] let diagnostic_shapes = [ "sscp_conv_0_convolution", + "sscp_conv_0_norm_sum_at_time", + "sscp_conv_0_norm_cumulative_sum", + "sscp_conv_0_norm_mean", + "sscp_conv_0_norm_squared_at_time", + "sscp_conv_0_norm_cumulative_squared", + "sscp_conv_0_norm_variance", + "sscp_conv_0_norm_stabilized_variance", + "sscp_conv_0_norm_inverse_stddev", "sscp_conv_0_norm", "sscp_conv_0", "sscp_conv_1_convolution", From ca1b9639759ff68a6348b2a2712b17deff92266c Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 00:25:56 +0900 Subject: [PATCH 15/19] test(xla): probe Gemma3n rsqrt lowering --- examples/xla_gemma3n_audio_reference_check.rs | 4 ++++ .../mlxcel-xla/src/emitter/gemma3n_audio_emit.rs | 5 +++++ .../src/emitter/gemma3n_audio_emit_tests.rs | 3 ++- .../mlxcel-xla/src/emitter/gemma3n_audio_ops.rs | 14 ++++++++++++++ src/lib/mlxcel-xla/src/gemma3n_audio_runtime.rs | 1 + 5 files changed, 26 insertions(+), 1 deletion(-) diff --git a/examples/xla_gemma3n_audio_reference_check.rs b/examples/xla_gemma3n_audio_reference_check.rs index e93c779d8..dd5e38b23 100644 --- a/examples/xla_gemma3n_audio_reference_check.rs +++ b/examples/xla_gemma3n_audio_reference_check.rs @@ -361,6 +361,10 @@ fn main() -> Result<(), String> { "sscp_conv_0_norm_inverse_stddev", mlx.sscp_conv_0_norm_inverse_stddev.as_slice(), ), + ( + "sscp_conv_0_norm_inverse_stddev_sqrt_reciprocal", + mlx.sscp_conv_0_norm_inverse_stddev.as_slice(), + ), ("sscp_conv_0_norm", mlx.sscp_conv_0_norm.as_slice()), ("sscp_conv_0", mlx.sscp_conv_0.as_slice()), ( diff --git a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit.rs b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit.rs index bcc809217..248e4168f 100644 --- a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit.rs +++ b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit.rs @@ -115,6 +115,10 @@ fn audio_encoder( "sscp_conv_0_norm_inverse_stddev", &subsampled.conv0_norm_inverse_stddev, ); + trace.push( + "sscp_conv_0_norm_inverse_stddev_sqrt_reciprocal", + &subsampled.conv0_norm_inverse_stddev_sqrt_reciprocal, + ); trace.push("sscp_conv_0_norm", &subsampled.conv0_norm); trace.push("sscp_conv_0", &subsampled.conv0); trace.push("sscp_conv_1_convolution", &subsampled.conv1_convolution); @@ -354,6 +358,7 @@ fn build_encode( "sscp_conv_0_norm_variance", "sscp_conv_0_norm_stabilized_variance", "sscp_conv_0_norm_inverse_stddev", + "sscp_conv_0_norm_inverse_stddev_sqrt_reciprocal", "sscp_conv_0_norm", "sscp_conv_0", "sscp_conv_1_convolution", diff --git a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit_tests.rs b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit_tests.rs index 5ae7e5ead..5060f20e4 100644 --- a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit_tests.rs +++ b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit_tests.rs @@ -44,7 +44,7 @@ fn split_audio_graphs_preserve_stage_and_weight_boundaries() { ); assert!(!encode.contains("loc(\"audio.hard_embeddings\")")); assert!(!encode.contains("model.language_model.")); - assert_eq!(encode_layout.stages.len(), 3 + 4 + 12 * 5 + 10 + 8); + assert_eq!(encode_layout.stages.len(), 3 + 4 + 12 * 5 + 11 + 8); for name in [ "sscp_conv_0_convolution", "sscp_conv_0_norm_sum_at_time", @@ -55,6 +55,7 @@ fn split_audio_graphs_preserve_stage_and_weight_boundaries() { "sscp_conv_0_norm_variance", "sscp_conv_0_norm_stabilized_variance", "sscp_conv_0_norm_inverse_stddev", + "sscp_conv_0_norm_inverse_stddev_sqrt_reciprocal", "sscp_conv_0_norm", "sscp_conv_0", "sscp_conv_1_convolution", diff --git a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_ops.rs b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_ops.rs index 471c50d3c..801e4e05d 100644 --- a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_ops.rs +++ b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_ops.rs @@ -18,6 +18,7 @@ pub(super) struct SscpOutput { pub conv0_norm_variance: Val, pub conv0_norm_stabilized_variance: Val, pub conv0_norm_inverse_stddev: Val, + pub conv0_norm_inverse_stddev_sqrt_reciprocal: Val, pub conv0_norm: Val, pub conv0: Val, pub conv1_convolution: Val, @@ -37,6 +38,7 @@ struct SscpConvStages { norm_variance: Val, norm_stabilized_variance: Val, norm_inverse_stddev: Val, + norm_inverse_stddev_sqrt_reciprocal: Val, norm: Val, activated: Val, } @@ -50,6 +52,7 @@ struct CumulativeGroupNormStages { variance: Val, stabilized_variance: Val, inverse_stddev: Val, + inverse_stddev_sqrt_reciprocal: Val, output: Val, } @@ -97,6 +100,14 @@ fn cumulative_group_norm( let epsilon = super::gemma3n_audio_math::scalar_like(b, eps, &variance); let stabilized_variance = b.add(&variance, &epsilon); let inverse_stddev = b.rsqrt(&stabilized_variance); + // CUDA MLX lowers `rsqrtf(f32)` to the target-specific + // `rsqrt.approx.f32` instruction. StableHLO has no portable spelling for + // that approximation. Keep the production path on `stablehlo.rsqrt`, and + // expose the precise sqrt-then-reciprocal ordering as a diagnostic + // candidate so one actual run can distinguish the two lowering semantics. + let sqrt = b.sqrt(&stabilized_variance); + let one = super::gemma3n_audio_math::scalar_like(b, 1.0, &sqrt); + let inverse_stddev_sqrt_reciprocal = b.divide(&one, &sqrt); let inverse = b.broadcast(&inverse_stddev, &[0, 1, 2, 3], x.ty.shape.clone()); let normalized = b.multiply(¢ered, &inverse); let weight = b.broadcast(weight, &[3], x.ty.shape.clone()); @@ -111,6 +122,7 @@ fn cumulative_group_norm( variance, stabilized_variance, inverse_stddev, + inverse_stddev_sqrt_reciprocal, output, } } @@ -149,6 +161,7 @@ fn sscp_conv( norm_variance: normalized.variance, norm_stabilized_variance: normalized.stabilized_variance, norm_inverse_stddev: normalized.inverse_stddev, + norm_inverse_stddev_sqrt_reciprocal: normalized.inverse_stddev_sqrt_reciprocal, norm: normalized.output, activated, } @@ -245,6 +258,7 @@ pub(super) fn subsample_with_stages( conv0_norm_variance: conv0.norm_variance, conv0_norm_stabilized_variance: conv0.norm_stabilized_variance, conv0_norm_inverse_stddev: conv0.norm_inverse_stddev, + conv0_norm_inverse_stddev_sqrt_reciprocal: conv0.norm_inverse_stddev_sqrt_reciprocal, conv0_norm: conv0.norm, conv0: conv0.activated, conv1_convolution: conv1.convolution, diff --git a/src/lib/mlxcel-xla/src/gemma3n_audio_runtime.rs b/src/lib/mlxcel-xla/src/gemma3n_audio_runtime.rs index 7bcc6188e..2df969bee 100644 --- a/src/lib/mlxcel-xla/src/gemma3n_audio_runtime.rs +++ b/src/lib/mlxcel-xla/src/gemma3n_audio_runtime.rs @@ -660,6 +660,7 @@ impl Gemma3nAudioIreeRuntime { "sscp_conv_0_norm_variance", "sscp_conv_0_norm_stabilized_variance", "sscp_conv_0_norm_inverse_stddev", + "sscp_conv_0_norm_inverse_stddev_sqrt_reciprocal", "sscp_conv_0_norm", "sscp_conv_0", "sscp_conv_1_convolution", From 190cb80e9b036de3059d952126679ec0a9ecdfa5 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 00:41:03 +0900 Subject: [PATCH 16/19] fix(xla): unblock Gemma3n CUDA PLE compile --- .../src/emitter/gemma3n_audio_emit_tests.rs | 47 ++++++++ .../src/emitter/gemma3n_emit_ops.rs | 107 +++++++++++++++++- 2 files changed, 153 insertions(+), 1 deletion(-) diff --git a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit_tests.rs b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit_tests.rs index 5060f20e4..98971eb40 100644 --- a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit_tests.rs +++ b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit_tests.rs @@ -26,6 +26,14 @@ fn text() -> Gemma3nConfig { .unwrap() } +fn released_merge_shape() -> Gemma3nConfig { + let mut text = text(); + text.context_capacity = 256; + text.n_layers = 35; + text.hidden_per_layer_input = 256; + text +} + #[test] fn split_audio_graphs_preserve_stage_and_weight_boundaries() { let text = text(); @@ -94,6 +102,30 @@ fn split_audio_graphs_preserve_stage_and_weight_boundaries() { ); } +#[test] +fn merge_ple_pins_released_cuda_rms_reduction_tree() { + let (merge, _) = emit_gemma3n_audio_merge_ple( + &released_merge_shape(), + &Gemma3nXlaAudioConfig::default(), + 1, + Precision::Bf16, + ) + .unwrap(); + + assert!( + !merge.contains("stablehlo.reduce("), + "audio merge-PLE must not emit the IREE v3.11 CUDA-blocking reduction" + ); + assert!( + merge.contains("tensor<256x35x32x8xf32>"), + "the released dense-PLE RMSNorm must group eight BF16 values per CUDA lane" + ); + assert!( + merge.matches("\"stablehlo.gather\"").count() >= 5, + "the released dense-PLE RMSNorm must pin the five 32-lane XOR stages" + ); +} + #[test] fn first_sscp_convolution_rounds_mel_inputs_before_accumulation() { fn synthetic_convolution(input: &[f32], weights: &[f32], pre_cast: bool) -> f32 { @@ -285,3 +317,18 @@ fn iree_compiles_split_real_audio_artifacts() { compile(&compiler, &device, &stem, "encode", encode); compile(&compiler, &device, &stem, "merge-ple", merge); } + +#[test] +#[ignore = "requires the version-matched IREE v3.11 CUDA compiler"] +fn iree_v311_cuda_compiles_released_audio_merge_ple() { + let compiler = std::env::var_os("MLXCEL_XLA_IREE_COMPILE").expect("IREE compiler"); + let (merge, _) = emit_gemma3n_audio_merge_ple( + &released_merge_shape(), + &Gemma3nXlaAudioConfig::default(), + 1, + Precision::Bf16, + ) + .unwrap(); + let stem = format!("mlxcel-gemma3n-audio-merge-cuda-{}", std::process::id()); + compile(&compiler, "cuda", &stem, "merge-ple", merge); +} diff --git a/src/lib/mlxcel-xla/src/emitter/gemma3n_emit_ops.rs b/src/lib/mlxcel-xla/src/emitter/gemma3n_emit_ops.rs index f3a075083..c4440aa91 100644 --- a/src/lib/mlxcel-xla/src/emitter/gemma3n_emit_ops.rs +++ b/src/lib/mlxcel-xla/src/emitter/gemma3n_emit_ops.rs @@ -206,6 +206,107 @@ pub(super) fn rms_last_bf16( } } +fn slice_axis(b: &mut Builder, x: &Val, axis: usize, start: usize, limit: usize) -> Val { + let ranges = + x.ty.shape + .iter() + .enumerate() + .map(|(current, &size)| { + if current == axis { + (start, limit) + } else { + (0, size) + } + }) + .collect::>(); + b.slice(x, &ranges) +} + +fn gather_axis(b: &mut Builder, x: &Val, axis: usize, indices: &[usize]) -> Val { + assert_eq!(indices.len(), x.ty.shape[axis]); + let mut permutation = vec![axis]; + permutation.extend((0..x.ty.shape.len()).filter(|¤t| current != axis)); + let transposed = b.transpose(x, &permutation); + let width = transposed.ty.shape[0]; + let remainder = transposed.ty.shape[1..].iter().product::(); + let flattened = b.reshape(&transposed, vec![width, remainder]); + let indices = indices + .iter() + .map(|&index| i32::try_from(index).expect("static axis index must fit i32")) + .collect::>(); + let indices = b.const_tensor_i32(&indices, vec![width, 1]); + let gathered = b.gather(&flattened, &indices); + let gathered = b.reshape(&gathered, transposed.ty.shape); + let mut inverse = vec![0; permutation.len()]; + for (transposed_axis, original_axis) in permutation.into_iter().enumerate() { + inverse[original_axis] = transposed_axis; + } + b.transpose(&gathered, &inverse) +} + +/// Match the maintained MLX CUDA BF16 RMSNorm reduction for the released +/// Gemma3n per-layer input width. CUDA assigns eight adjacent BF16 values to +/// each of 32 lanes, accumulates their F32 squares serially, then broadcasts +/// the warp sum through XOR shuffles (16, 8, 4, 2, 1). +/// +/// Besides pinning the reference association order, spelling out the 32-lane +/// tree avoids IREE v3.11 CUDA's `LLVMGPUVectorDistribute` failure on the +/// otherwise generated 256-wide `stablehlo.reduce`. +fn mlx_cuda_rms_sum_256(b: &mut Builder, x: &Val) -> Val { + const LANES: usize = 32; + const READS: usize = 8; + + assert_eq!(x.ty.elt, "f32"); + let axis = x.ty.shape.len() - 1; + assert_eq!(x.ty.shape[axis], LANES * READS); + + let mut grouped_shape = x.ty.shape.clone(); + grouped_shape[axis] = LANES; + grouped_shape.insert(axis + 1, READS); + let grouped = b.reshape(x, grouped_shape); + let first = slice_axis(b, &grouped, axis + 1, 0, 1); + let zero = b.const_f32(0.0); + let mut partial = b.broadcast(&zero, &[], first.ty.shape.clone()); + for read in 0..READS { + let value = slice_axis(b, &grouped, axis + 1, read, read + 1); + let square = b.multiply(&value, &value); + partial = b.add(&partial, &square); + } + + let mut lane_shape = x.ty.shape.clone(); + lane_shape[axis] = LANES; + let mut partial = b.reshape(&partial, lane_shape); + for mask in [16, 8, 4, 2, 1] { + let shuffled_indices = (0..LANES).map(|lane| lane ^ mask).collect::>(); + let shuffled = gather_axis(b, &partial, axis, &shuffled_indices); + partial = b.add(&partial, &shuffled); + } + + let reduced = slice_axis(b, &partial, axis, 0, 1); + let mut reduced_shape = x.ty.shape.clone(); + reduced_shape.remove(axis); + b.reshape(&reduced, reduced_shape) +} + +fn rms_last_bf16_mlx_cuda_256(b: &mut Builder, x: &Val, weight: &Val, eps: &Val) -> Val { + let axis = x.ty.shape.len() - 1; + let reduced_shape = x.ty.shape[..axis].to_vec(); + let width = b.const_f32(256.0); + let width = b.broadcast(&width, &[], reduced_shape.clone()); + let sum = mlx_cuda_rms_sum_256(b, x); + let mean = b.divide(&sum, &width); + let eps = b.broadcast(eps, &[], reduced_shape); + let stabilized = b.add(&mean, &eps); + let inverse = b.rsqrt(&stabilized); + let keep = (0..axis).collect::>(); + let inverse = b.broadcast(&inverse, &keep, x.ty.shape.clone()); + let normalized = b.multiply(x, &inverse); + let normalized = round_bf16(b, &normalized); + let weight = b.broadcast(weight, &[axis], x.ty.shape.clone()); + let weighted = b.multiply(&normalized, &weight); + round_bf16(b, &weighted) +} + /// Build the dense per-layer embedding stream from post-scale model /// embeddings. Token prefill and audio prefill must use this exact shared /// sequence: the BF16 rounding boundaries are part of Gemma3n's numerical @@ -249,7 +350,11 @@ pub(super) fn dense_ple_input_head( let projected = b.multiply(&projected, &model_scale); let projected = round_bf16(b, &projected); let projected = b.reshape(&projected, vec![rows, c.n_layers, c.hidden_per_layer_input]); - let projected = rms_last_bf16(b, &projected, Some(projection_norm_weight), eps, zero); + let projected = if c.hidden_per_layer_input == 256 { + rms_last_bf16_mlx_cuda_256(b, &projected, projection_norm_weight, eps) + } else { + rms_last_bf16(b, &projected, Some(projection_norm_weight), eps, zero) + }; let token_ple = b.reshape(&token_ple, vec![rows, c.n_layers, c.hidden_per_layer_input]); let inv = b.broadcast( inv_sqrt2, From dde69494148e53e622a21d445eb3c3f38608976d Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 01:12:03 +0900 Subject: [PATCH 17/19] test(xla): isolate Gemma3n conv1 result typing --- examples/xla_gemma3n_audio_reference_check.rs | 4 + .../src/emitter/gemma3n_audio_emit.rs | 6 + .../src/emitter/gemma3n_audio_emit_tests.rs | 118 +++++++++++++++++- .../src/emitter/gemma3n_audio_ops.rs | 22 ++++ .../mlxcel-xla/src/gemma3n_audio_runtime.rs | 1 + 5 files changed, 146 insertions(+), 5 deletions(-) diff --git a/examples/xla_gemma3n_audio_reference_check.rs b/examples/xla_gemma3n_audio_reference_check.rs index dd5e38b23..59092ed89 100644 --- a/examples/xla_gemma3n_audio_reference_check.rs +++ b/examples/xla_gemma3n_audio_reference_check.rs @@ -371,6 +371,10 @@ fn main() -> Result<(), String> { "sscp_conv_1_convolution", mlx.sscp_conv_1_convolution.as_slice(), ), + ( + "sscp_conv_1_convolution_bf16_result", + mlx.sscp_conv_1_convolution.as_slice(), + ), ("sscp_conv_1_norm", mlx.sscp_conv_1_norm.as_slice()), ("sscp_conv_1", mlx.sscp_conv_1.as_slice()), ("input_projection", mlx.input_projection.as_slice()), diff --git a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit.rs b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit.rs index 248e4168f..70a8c8328 100644 --- a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit.rs +++ b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit.rs @@ -122,6 +122,11 @@ fn audio_encoder( trace.push("sscp_conv_0_norm", &subsampled.conv0_norm); trace.push("sscp_conv_0", &subsampled.conv0); trace.push("sscp_conv_1_convolution", &subsampled.conv1_convolution); + #[cfg(any(feature = "diagnostics", test))] + trace.push( + "sscp_conv_1_convolution_bf16_result", + &subsampled.conv1_convolution_bf16_result, + ); trace.push("sscp_conv_1_norm", &subsampled.conv1_norm); trace.push("sscp_conv_1", &subsampled.conv1); trace.push("input_projection", &subsampled.hidden); @@ -362,6 +367,7 @@ fn build_encode( "sscp_conv_0_norm", "sscp_conv_0", "sscp_conv_1_convolution", + "sscp_conv_1_convolution_bf16_result", "sscp_conv_1_norm", "sscp_conv_1", "input_projection", diff --git a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit_tests.rs b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit_tests.rs index 98971eb40..dccf8b940 100644 --- a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit_tests.rs +++ b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit_tests.rs @@ -52,7 +52,10 @@ fn split_audio_graphs_preserve_stage_and_weight_boundaries() { ); assert!(!encode.contains("loc(\"audio.hard_embeddings\")")); assert!(!encode.contains("model.language_model.")); - assert_eq!(encode_layout.stages.len(), 3 + 4 + 12 * 5 + 11 + 8); + assert_eq!( + encode_layout.stages.len(), + 3 + 4 + 12 * 5 + 11 + 8 + usize::from(cfg!(any(feature = "diagnostics", test))) + ); for name in [ "sscp_conv_0_convolution", "sscp_conv_0_norm_sum_at_time", @@ -85,6 +88,12 @@ fn split_audio_graphs_preserve_stage_and_weight_boundaries() { ] { assert!(encode_layout.stages.iter().any(|stage| stage.name == name)); } + assert!( + encode_layout + .stages + .iter() + .any(|stage| stage.name == "sscp_conv_1_convolution_bf16_result") + ); assert!(merge.starts_with("module @audio_merge_ple {")); assert!(merge.contains("loc(\"audio.projected\")")); @@ -206,10 +215,84 @@ fn sscp_convolutions_accumulate_bf16_operands_into_f32_results() { "(tensor<1x6x66x128xbf16>, tensor<3x3x128x32xbf16>) -> \ tensor<1x2x32x32xf32>" )); - assert!(!encode.contains( - "(tensor<1x10x130x1xbf16>, tensor<3x3x1x128xbf16>) -> \ - tensor<1x4x64x128xbf16>" - )); + assert!( + !encode.contains( + "(tensor<1x10x130x1xbf16>, tensor<3x3x1x128xbf16>) -> \ + tensor<1x4x64x128xbf16>" + ), + "the conv0 production path must remain on the explicit F32 result" + ); +} + +#[test] +fn second_sscp_convolution_pins_result_type_as_the_only_candidate_variable() { + let (encode, layout) = emit_gemma3n_audio_encode( + &text(), + &Gemma3nXlaAudioConfig::default(), + 8, + 1, + Precision::Bf16, + ) + .unwrap(); + let conv1_lines = encode + .lines() + .filter(|line| { + line.contains("\"stablehlo.convolution\"") + && line.contains("(tensor<1x6x66x128xbf16>, tensor<3x3x128x32xbf16>)") + }) + .collect::>(); + assert_eq!( + conv1_lines.len(), + 2, + "diagnostics must emit exactly the production and BF16-result conv1 forms" + ); + + let demoted_operands = |line: &str| { + line.split_once("\"stablehlo.convolution\"(") + .and_then(|(_, suffix)| suffix.split_once(')')) + .map(|(operands, _)| operands.to_string()) + .expect("convolution line must expose its operands") + }; + let demotion_sources = |line: &str| { + demoted_operands(line) + .split(", ") + .map(|operand| { + encode + .lines() + .find(|candidate| candidate.trim_start().starts_with(&format!("{operand} ="))) + .and_then(|conversion| conversion.split_once("stablehlo.convert ")) + .and_then(|(_, suffix)| suffix.split_once(" :")) + .map(|(source, _)| source.to_string()) + .expect("BF16 convolution operand must come from an explicit conversion") + }) + .collect::>() + }; + assert_eq!( + demotion_sources(conv1_lines[0]), + demotion_sources(conv1_lines[1]), + "both forms must consume the same padded activation and HWIO weight" + ); + for line in &conv1_lines { + assert!(line.contains("window_strides = array")); + assert!(line.contains( + "dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]>" + )); + } + assert!(conv1_lines.iter().any(|line| { + line.contains( + "(tensor<1x6x66x128xbf16>, tensor<3x3x128x32xbf16>) -> \ + tensor<1x2x32x32xf32>", + ) + })); + assert!(conv1_lines.iter().any(|line| { + line.contains( + "(tensor<1x6x66x128xbf16>, tensor<3x3x128x32xbf16>) -> \ + tensor<1x2x32x32xbf16>", + ) + })); + assert!(layout.stages.iter().any(|stage| { + stage.name == "sscp_conv_1_convolution_bf16_result" && stage.shape == [1, 2, 32, 32] + })); } #[test] @@ -233,6 +316,31 @@ fn sscp_f32_accumulator_contract_compiles_for_cpu() { compile(compiler.as_os_str(), "local", &stem, "contract", mlir); } +#[test] +#[ignore = "requires the version-matched pinned IREE CUDA compiler"] +fn sscp_conv1_result_type_candidates_compile_for_cuda() { + let mut builder = Builder::new().with_precision(Precision::Bf16); + let input = Builder::arg(0, Ty::f32(vec![1, 4, 64, 128])); + let kernel = Builder::arg(1, Ty::f32(vec![32, 3, 3, 128])); + let zero = builder.const_f32(0.0); + let padded = builder.pad(&input, &zero, &[0, 0, 1, 0], &[0, 2, 1, 0]); + let kernel = builder.transpose(&kernel, &[1, 2, 3, 0]); + let f32_result = builder.convolution_f32_accumulate(&padded, &kernel, &[2, 2], &[(0, 0); 2], 1); + let bf16_result = builder.convolution(&padded, &kernel, &[2, 2], &[(0, 0); 2], 1); + let output_ty = f32_result.ty.render(); + let mlir = format!( + "module @sscp_conv1_result_candidates {{\n func.func public @main(%arg0: \ + tensor<1x4x64x128xf32>, %arg1: tensor<32x3x3x128xf32>) -> ({output_ty}, \ + {output_ty}) {{\n{} return {}, {} : {output_ty}, {output_ty}\n }}\n}}\n", + builder.body(), + f32_result.name, + bf16_result.name, + ); + let compiler = std::env::var_os("MLXCEL_XLA_IREE_COMPILE").expect("IREE compiler"); + let stem = format!("mlxcel-gemma3n-sscp-conv1-cuda-{}", std::process::id()); + compile(&compiler, "cuda", &stem, "result-candidates", mlir); +} + #[test] #[ignore = "requires IREE_DIST"] fn sscp_cuda_norm_schedule_compiles_for_cpu() { diff --git a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_ops.rs b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_ops.rs index 801e4e05d..4688e5a6f 100644 --- a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_ops.rs +++ b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_ops.rs @@ -22,6 +22,8 @@ pub(super) struct SscpOutput { pub conv0_norm: Val, pub conv0: Val, pub conv1_convolution: Val, + #[cfg(any(feature = "diagnostics", test))] + pub conv1_convolution_bf16_result: Val, pub conv1_norm: Val, pub conv1: Val, pub hidden: Val, @@ -30,6 +32,8 @@ pub(super) struct SscpOutput { struct SscpConvStages { convolution: Val, + #[cfg(any(feature = "diagnostics", test))] + convolution_bf16_result: Option, norm_sum_at_time: Val, norm_cumulative_sum: Val, norm_mean: Val, @@ -137,10 +141,20 @@ fn sscp_conv( stride: [usize; 2], groups: usize, eps: f32, + _diagnostic_bf16_result: bool, ) -> SscpConvStages { let zero = b.const_f32(0.0); let padded = b.pad(x, &zero, &[0, 0, 1, 0], &[0, kernel[0] - 1, 1, 0]); let weight = b.transpose(weight, &[1, 2, 3, 0]); + // IREE CUDA may select a different contraction lowering when StableHLO + // carries a BF16 result instead of an explicit F32 result followed by a + // convert. Keep this diagnostics-only candidate on the exact same padded + // input and HWIO weight so an actual CUDA run can distinguish result + // typing from input, layout, and stride differences without changing the + // production convolution. + #[cfg(any(feature = "diagnostics", test))] + let convolution_bf16_result = _diagnostic_bf16_result + .then(|| b.convolution(&padded, &weight, &stride, &[(0, 0), (0, 0)], groups)); // MLX stores both operands as BF16 but accumulates the convolution into an // F32 carrier before materializing the BF16 output tensor. The graph-wide // BF16 contraction mode otherwise gives StableHLO a BF16 result type, @@ -153,6 +167,8 @@ fn sscp_conv( let activated = round_bf16(b, &activated); SscpConvStages { convolution: convolved, + #[cfg(any(feature = "diagnostics", test))] + convolution_bf16_result, norm_sum_at_time: normalized.sum_at_time, norm_cumulative_sum: normalized.cumulative_sum, norm_mean: normalized.mean, @@ -193,6 +209,7 @@ fn subsample_convs( config.sscp_conv_stride_size[0], 1, config.sscp_conv_group_norm_eps, + false, ); let conv1 = sscp_conv( b, @@ -203,6 +220,7 @@ fn subsample_convs( config.sscp_conv_stride_size[1], 1, config.sscp_conv_group_norm_eps, + true, ); (conv0, conv1) } @@ -262,6 +280,10 @@ pub(super) fn subsample_with_stages( conv0_norm: conv0.norm, conv0: conv0.activated, conv1_convolution: conv1.convolution, + #[cfg(any(feature = "diagnostics", test))] + conv1_convolution_bf16_result: conv1 + .convolution_bf16_result + .expect("the second SSCP convolution must expose its BF16-result candidate"), conv1_norm: conv1.norm, conv1: conv1.activated, hidden, diff --git a/src/lib/mlxcel-xla/src/gemma3n_audio_runtime.rs b/src/lib/mlxcel-xla/src/gemma3n_audio_runtime.rs index 2df969bee..149854c7a 100644 --- a/src/lib/mlxcel-xla/src/gemma3n_audio_runtime.rs +++ b/src/lib/mlxcel-xla/src/gemma3n_audio_runtime.rs @@ -664,6 +664,7 @@ impl Gemma3nAudioIreeRuntime { "sscp_conv_0_norm", "sscp_conv_0", "sscp_conv_1_convolution", + "sscp_conv_1_convolution_bf16_result", "sscp_conv_1_norm", "sscp_conv_1", "input_projection", From 385718132f338db0447086f8636d52695a743a81 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 09:28:13 +0900 Subject: [PATCH 18/19] fix(xla): rebase Gemma3n audio contracts Keep Gemma3n audio artifacts explicitly legacy-unqualified and adapt the newer Qwen2-VL loader to the typed auxiliary weight storage introduced by this branch.\n\nRefs #878 --- .../mlxcel-xla/src/gemma3n_audio_runtime.rs | 20 +++++++++++++------ src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs | 4 ++-- 2 files changed, 16 insertions(+), 8 deletions(-) diff --git a/src/lib/mlxcel-xla/src/gemma3n_audio_runtime.rs b/src/lib/mlxcel-xla/src/gemma3n_audio_runtime.rs index 149854c7a..fe03d0100 100644 --- a/src/lib/mlxcel-xla/src/gemma3n_audio_runtime.rs +++ b/src/lib/mlxcel-xla/src/gemma3n_audio_runtime.rs @@ -577,12 +577,12 @@ impl Gemma3nAudioIreeRuntime { text.n_layers, text.hidden_per_layer_input, )?; - let encode_contract = AuxiliaryArtifactContract::new( + let encode_contract = AuxiliaryArtifactContract::new_legacy_unqualified( "audio_encode.main", format!("encode:{identity}:language={language_fingerprint:016x}"), generation_identity(&compiler, &flags, &encode_mlir)?, )?; - let merge_contract = AuxiliaryArtifactContract::new( + let merge_contract = AuxiliaryArtifactContract::new_legacy_unqualified( "audio_merge_ple.main", format!( "merge:{identity}:context={context_capacity}:text={}:language={language_fingerprint:016x}", @@ -1004,8 +1004,12 @@ mod tests { let first_generation = generation_identity_from_version(&compiler, &["--target=cuda"], "mlir", "v1").unwrap(); assert!(first_generation.contains(&sha256_hex(b"compiler-build-a"))); - let first_contract = - AuxiliaryArtifactContract::new("audio.main", "config=v1", &first_generation).unwrap(); + let first_contract = AuxiliaryArtifactContract::new_legacy_unqualified( + "audio.main", + "config=v1", + &first_generation, + ) + .unwrap(); let mut compile_count = 0usize; ensure_qualified_auxiliary_artifact(&vmfb, &first_contract, &weights, |temporary| { compile_count += 1; @@ -1018,8 +1022,12 @@ mod tests { let second_generation = generation_identity_from_version(&compiler, &["--target=cuda"], "mlir", "v1").unwrap(); assert_ne!(first_generation, second_generation); - let second_contract = - AuxiliaryArtifactContract::new("audio.main", "config=v1", second_generation).unwrap(); + let second_contract = AuxiliaryArtifactContract::new_legacy_unqualified( + "audio.main", + "config=v1", + second_generation, + ) + .unwrap(); ensure_qualified_auxiliary_artifact(&vmfb, &second_contract, &weights, |temporary| { compile_count += 1; std::fs::write(temporary, b"vmfb-b") diff --git a/src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs b/src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs index 2634f24b3..e26ccd78c 100644 --- a/src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs +++ b/src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs @@ -31,7 +31,7 @@ use sha2::{Digest, Sha256}; use crate::aux::{ AuxiliaryInput, AuxiliaryOutput, AuxiliaryTensorDType, AuxiliaryWeight, AuxiliaryWeightDType, - IreeAuxiliaryModule, + AuxiliaryWeightStorage, IreeAuxiliaryModule, }; use crate::aux_manifest::{AuxiliaryArtifactContract, ensure_qualified_auxiliary_artifact}; use crate::emitter::{ @@ -499,7 +499,7 @@ fn load_weights( } loaded[index] = Some(AuxiliaryWeight { name: spec.name.clone(), - bytes: native_f32_bytes(values), + storage: AuxiliaryWeightStorage::Bytes(native_f32_bytes(values)), dtype: AuxiliaryWeightDType::Float32, shape: spec.shape.clone(), }); From 0a5aa90d1496b3a3b3fe9e2b2f5014738281dd45 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 10:41:32 +0900 Subject: [PATCH 19/19] refactor(xla): share Gemma3n CUDA prefix scan --- .../src/emitter/gemma3n_audio_emit_tests.rs | 16 +++ .../src/emitter/gemma3n_audio_math.rs | 97 ++----------------- 2 files changed, 24 insertions(+), 89 deletions(-) diff --git a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit_tests.rs b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit_tests.rs index dccf8b940..8894ef2c1 100644 --- a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit_tests.rs +++ b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit_tests.rs @@ -341,6 +341,22 @@ fn sscp_conv1_result_type_candidates_compile_for_cuda() { compile(&compiler, "cuda", &stem, "result-candidates", mlir); } +#[test] +fn sscp_cumsum_uses_shared_cuda_prefix_schedule() { + use crate::emitter::gemma3n_audio_math::mlx_cuda_cumsum_time_f32; + + let mut builder = Builder::new().with_precision(Precision::Bf16); + let input = Builder::arg(0, Ty::f32(vec![2, 130, 1, 1])); + let output = mlx_cuda_cumsum_time_f32(&mut builder, &input); + let body = builder.body(); + + assert_eq!(output.ty.shape, [2, 130, 1, 1]); + assert_eq!(body.matches("stablehlo.add").count(), 17); + assert_eq!(body.matches("stablehlo.concatenate").count(), 3); + assert!(!body.contains("stablehlo.reduce")); + assert!(!body.contains("stablehlo.reduce_window")); +} + #[test] #[ignore = "requires IREE_DIST"] fn sscp_cuda_norm_schedule_compiles_for_cpu() { diff --git a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_math.rs b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_math.rs index cad1dcb10..4295537f3 100644 --- a/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_math.rs +++ b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_math.rs @@ -1,6 +1,7 @@ //! Shared exact math for the Gemma3n audio StableHLO graph. use super::builder::{Builder, Val}; +use super::numeric_ops::mlx_cuda_prefix_sum_2d; pub(super) fn scalar_like(b: &mut Builder, value: f32, x: &Val) -> Val { let scalar = b.const_f32(value); @@ -110,32 +111,6 @@ pub(super) fn mlx_cuda_row_sum_f32(b: &mut Builder, x: &Val, axis: usize) -> Val b.reshape(&reduced, output_shape) } -fn inclusive_scan_32_f32(b: &mut Builder, x: &Val, axis: usize) -> Val { - assert_eq!(x.ty.elt, "f32"); - assert_eq!(x.ty.shape[axis], 32); - let zero = b.const_f32(0.0); - let mut result = x.clone(); - for offset in [1, 2, 4, 8, 16] { - let prior = slice_axis(b, &result, axis, 0, 32 - offset); - let mut low = vec![0; result.ty.shape.len()]; - let high = low.clone(); - low[axis] = offset; - let prior = b.pad(&prior, &zero, &low, &high); - // CUDA cooperative_groups uses `out = op(out, shfl_up(out))`. - result = b.add(&result, &prior); - } - result -} - -fn exclusive_from_inclusive_32_f32(b: &mut Builder, inclusive: &Val, axis: usize) -> Val { - let prior = slice_axis(b, inclusive, axis, 0, 31); - let zero = b.const_f32(0.0); - let mut low = vec![0; inclusive.ty.shape.len()]; - let high = low.clone(); - low[axis] = 1; - b.pad(&prior, &zero, &low, &high) -} - /// Match the one-block MLX CUDA contiguous inclusive-scan schedule used by the /// maintained Gemma3n oracle. Released frame buckets produce at most 1,499 /// values after the first SSCP convolution, below CUDA's 4,096-value block @@ -145,62 +120,14 @@ pub(super) fn mlx_cuda_cumsum_time_f32(b: &mut Builder, x: &Val) -> Val { assert_eq!(x.ty.shape.len(), 4); let batch = x.ty.shape[0]; let time = x.ty.shape[1]; - assert!( - (1..=4096).contains(&time), - "MLX CUDA cumsum emulation supports one scan block" + assert_eq!( + x.ty.shape[2..], + [1, 1], + "Gemma3n cumulative sum expects singleton feature axes" ); - let padded_time = time.div_ceil(128) * 128; - let padded = if padded_time == time { - x.clone() - } else { - let zero = b.const_f32(0.0); - b.pad(x, &zero, &[0, 0, 0, 0], &[0, padded_time - time, 0, 0]) - }; - let threads = padded_time / 4; - let warps = threads / 32; - let grouped = b.reshape(&padded, vec![batch, threads, 4, 1, 1]); - let value0 = slice_axis(b, &grouped, 2, 0, 1); - let value1 = slice_axis(b, &grouped, 2, 1, 2); - let value2 = slice_axis(b, &grouped, 2, 2, 3); - let value3 = slice_axis(b, &grouped, 2, 3, 4); - let local0 = value0; - let local1 = b.add(&value1, &local0); - let local2 = b.add(&value2, &local1); - let local3 = b.add(&value3, &local2); - let local = concat_axis(b, &[local0, local1, local2, local3.clone()], 2); - - let thread_sums = b.reshape(&local3, vec![batch, warps, 32, 1, 1]); - let thread_inclusive = inclusive_scan_32_f32(b, &thread_sums, 2); - let thread_exclusive = exclusive_from_inclusive_32_f32(b, &thread_inclusive, 2); - - // The MLX kernel forms each warp total as - // `exclusive_scan(last_lane) + last_lane_value`, rather than taking the - // inclusive result of lane 31. Preserve that distinct association. - let prior_last = slice_axis(b, &thread_exclusive, 2, 31, 32); - let current_last = slice_axis(b, &thread_sums, 2, 31, 32); - let warp_totals = b.add(&prior_last, ¤t_last); - let warp_totals = b.reshape(&warp_totals, vec![batch, warps, 1, 1]); - let zero = b.const_f32(0.0); - let warp_totals_32 = b.pad(&warp_totals, &zero, &[0, 0, 0, 0], &[0, 32 - warps, 0, 0]); - let warp_inclusive = inclusive_scan_32_f32(b, &warp_totals_32, 1); - let warp_exclusive = exclusive_from_inclusive_32_f32(b, &warp_inclusive, 1); - let warp_exclusive = slice_axis(b, &warp_exclusive, 1, 0, warps); - - let local = b.reshape(&local, vec![batch, warps, 32, 4, 1, 1]); - let prefix = zeros_like(b, &local); - let local = b.add(&local, &prefix); - let warp_exclusive = b.reshape(&warp_exclusive, vec![batch, warps, 1, 1, 1, 1]); - let warp_exclusive = b.broadcast(&warp_exclusive, &[0, 1, 2, 3, 4, 5], local.ty.shape.clone()); - let with_warp_prefix = b.add(&local, &warp_exclusive); - let thread_exclusive = b.reshape(&thread_exclusive, vec![batch, warps, 32, 1, 1, 1]); - let thread_exclusive = b.broadcast( - &thread_exclusive, - &[0, 1, 2, 3, 4, 5], - with_warp_prefix.ty.shape.clone(), - ); - let scanned = b.add(&with_warp_prefix, &thread_exclusive); - let scanned = b.reshape(&scanned, vec![batch, padded_time, 1, 1]); - slice_axis(b, &scanned, 1, 0, time) + let flattened = b.reshape(x, vec![batch, time]); + let scanned = mlx_cuda_prefix_sum_2d(b, &flattened); + b.reshape(&scanned, x.ty.shape.clone()) } pub(super) fn clip(b: &mut Builder, x: &Val, limit: f32) -> Val { @@ -288,14 +215,6 @@ pub(super) fn stride_time(b: &mut Builder, x: &Val, stride: usize) -> Val { b.slice_strided(x, &ranges) } -pub(super) fn concat_axis(b: &mut Builder, values: &[Val], axis: usize) -> Val { - let mut result = values[0].clone(); - for value in &values[1..] { - result = b.concatenate(&result, value, axis); - } - result -} - pub(super) fn zero_invalid(b: &mut Builder, x: &Val, valid: &Val) -> Val { let mut shape = valid.ty.shape.clone(); while shape.len() < x.ty.shape.len() {