diff --git a/Cargo.toml b/Cargo.toml index ab506bce0..c30116456 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -343,4 +343,15 @@ 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"] + +# 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_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/examples/xla_gemma3n_audio_reference_check.rs b/examples/xla_gemma3n_audio_reference_check.rs new file mode 100644 index 000000000..59092ed89 --- /dev/null +++ b/examples/xla_gemma3n_audio_reference_check.rs @@ -0,0 +1,482 @@ +// 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; +const MAX_BF16_OUTLIER_PROBES: usize = 32; + +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, + 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 { + 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 { + 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; + let mut max_bf16_ulp = 0u32; + let mut non_bf16_values = 0usize; + let mut over_one_bf16_ulp = 0usize; + let mut bf16_outliers = Vec::new(); + 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; + } + match bf16_ulp_distance(*reference, *actual) { + Some(distance) => { + max_bf16_ulp = max_bf16_ulp.max(distance); + 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, + } + } + 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, + ]; + let max_reference_bits = display_bf16_bits(reference[max_index]); + let max_actual_bits = display_bf16_bits(actual[max_index]); + let max_pair_ulp = bf16_ulp_distance(reference[max_index], actual[max_index]); + eprintln!( + "gemma3n-audio-reference: stage={name} max_abs={max_abs:.9e} rms={rms:.9e} \ + max_index={max_index} max_pair_mlx={:.9e} max_pair_iree={:.9e} \ + max_pair_mlx_bf16={max_reference_bits} max_pair_iree_bf16={max_actual_bits} \ + max_pair_bf16_ulp={max_pair_ulp:?} max_bf16_ulp={max_bf16_ulp} \ + non_bf16_values={non_bf16_values} over_one_bf16_ulp={over_one_bf16_ulp}", + reference[max_index], actual[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() + ); + } + for outlier in &bf16_outliers { + eprintln!( + "gemma3n-audio-reference: bf16-outlier stage={name} index={} mlx={:.9e} \ + iree={:.9e} abs={:.9e} mlx_bf16={} iree_bf16={} bf16_ulp={}", + outlier.index, + outlier.reference, + outlier.actual, + (outlier.actual - outlier.reference).abs(), + display_bf16_bits(outlier.reference), + display_bf16_bits(outlier.actual), + outlier.distance, + ); + } + if bf16_outliers.len() < over_one_bf16_ulp { + eprintln!( + "gemma3n-audio-reference: bf16-outlier stage={name} captured={} total={} truncated=true", + bf16_outliers.len(), + over_one_bf16_ulp, + ); + } + Ok(Diff { + max_abs, + rms, + max_index, + max_bf16_ulp, + non_bf16_values, + over_one_bf16_ulp, + bf16_outliers, + }) +} + +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 = [ + ( + "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_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()), + ( + "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()), + ( + "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()), + ("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={} \ + max_bf16_ulp={} non_bf16_values={} over_one_bf16_ulp={} \ + captured_bf16_outliers={}", + names[index], + diff.max_abs, + diff.rms, + diff.max_index, + 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"); + Ok(()) + } + } +} + +#[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); + 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/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::()); + } +} diff --git a/src/audio/gemma3n/diagnostics.rs b/src/audio/gemma3n/diagnostics.rs new file mode 100644 index 000000000..17d5f8eac --- /dev/null +++ b/src/audio/gemma3n/diagnostics.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. + +//! 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 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, + 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, + 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, encoder_stages) = + tower.forward_with_audio_diagnostics(&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 { + 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), + 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), + 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/encoder.rs b/src/audio/gemma3n/encoder.rs index 6fc8f5e46..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, + } } } @@ -123,6 +167,21 @@ struct SscpConvBlock { stride_frequency: i32, } +#[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, +} + impl SscpConvBlock { fn from_weights( weights: &WeightMap, @@ -157,7 +216,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 +226,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 +237,32 @@ 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_stages(&convolution); + let activated = mlxcel_core::relu(&norm.output); + Ok(SscpConvDiagnosticStages { + convolution, + 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, + }) } } @@ -189,6 +272,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 @@ -246,6 +336,25 @@ 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_diagnostics( + &self, + audio_mel: &MlxArray, + ) -> Result { + let x = mlxcel_core::expand_dims(audio_mel, -1); + 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, + }) + } } fn clip_gradient(x: &MlxArray, limit: f32) -> UniquePtr { @@ -529,14 +638,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), @@ -554,6 +662,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 { @@ -562,6 +696,26 @@ pub struct Gemma3nAudioEncoder { conformer: Vec, } +#[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, + 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, @@ -624,11 +778,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,18 +794,26 @@ impl Gemma3nAudioEncoder { self.config.input_feat_size )); } + Ok(()) + } - let mut encodings = self.subsample.forward(audio_mel)?; - let mut invalid_mask = Self::stride_mask( + fn initial_invalid_mask( + &self, + encodings: &MlxArray, + invalid_mel_mask: &MlxArray, + ) -> 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 @@ -674,6 +836,82 @@ 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, + 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_audio_diagnostics( + &self, + audio_mel: &MlxArray, + invalid_mel_mask: &MlxArray, + ) -> Result< + ( + UniquePtr, + UniquePtr, + Gemma3nAudioEncoderDiagnosticStages, + ), + String, + > { + self.validate_input_shapes(audio_mel, invalid_mel_mask)?; + 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_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, + sscp_conv_1_norm: subsample.conv1.norm, + sscp_conv_1: subsample.conv1.activated, + input_projection: subsample.input_projection, + conformer_0_feed_forward_start, + }, + )) + } } #[cfg(test)] diff --git a/src/audio/gemma3n/mod.rs b/src/audio/gemma3n/mod.rs index 4f657d71d..331613067 100644 --- a/src/audio/gemma3n/mod.rs +++ b/src/audio/gemma3n/mod.rs @@ -20,8 +20,12 @@ mod attention; mod config; +#[cfg(feature = "xla-diagnostics")] +mod diagnostics; mod encoder; mod feature_extractor; +#[cfg(feature = "xla-backend")] +mod xla; use mlxcel_core::layers::UnifiedLinear; use mlxcel_core::weights::WeightMap; @@ -84,8 +88,12 @@ 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, 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..6b75f36e0 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), + 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/builder.rs b/src/lib/mlxcel-xla/src/emitter/builder.rs index ae182c154..89e6eda45 100644 --- a/src/lib/mlxcel-xla/src/emitter/builder.rs +++ b/src/lib/mlxcel-xla/src/emitter/builder.rs @@ -715,6 +715,18 @@ 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"), + } + } + /// 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 +999,29 @@ 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 } + } + 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]; @@ -1137,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!( @@ -1167,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() { @@ -1186,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() @@ -1269,6 +1333,12 @@ impl Builder { pub fn divide(&mut self, a: &Val, b: &Val) -> Val { self.binary("divide", 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 +1374,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 +1456,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 +1683,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 +1727,18 @@ 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) + } } #[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..70a8c8328 --- /dev/null +++ b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit.rs @@ -0,0 +1,524 @@ +//! 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, + #[cfg(feature = "diagnostics")] + diagnostics: Vec<(String, 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_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_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); + #[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); + 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, + trace: &mut Trace, +) -> (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); + 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, + 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)], + ); + + 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); + trace.push("soft_projection", &projected_audio); + trace.push("hard_projection", &hard_audio); + #[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_inverse_stddev_sqrt_reciprocal", + "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", + "conformer.0.feed_forward_start", + "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, + } +} + +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(); + 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}) -> ({output_types}) \ + {{\n{body} return {output_values} : {output_types}\n }}\n}}\n", + signature = signature(decls), + body = b.body(), + ) +} + +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 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(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..8894ef2c1 --- /dev/null +++ b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_emit_tests.rs @@ -0,0 +1,458 @@ +use super::*; +use crate::emitter::builder::Ty; + +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() +} + +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(); + 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.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 + 11 + 8 + usize::from(cfg!(any(feature = "diagnostics", test))) + ); + 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_inverse_stddev_sqrt_reciprocal", + "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", + "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", + ] { + 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\")")); + 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"] + ); +} + +#[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 { + 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); +} + +#[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>" + ), + "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] +#[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); +} + +#[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] +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() { + 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")); + 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); +} + +#[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_audio_math.rs b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_math.rs new file mode 100644 index 000000000..4295537f3 --- /dev/null +++ b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_math.rs @@ -0,0 +1,228 @@ +//! 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); + 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") +} + +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) +} + +/// 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_eq!( + x.ty.shape[2..], + [1, 1], + "Gemma3n cumulative sum expects singleton feature axes" + ); + 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 { + 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 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 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..4688e5a6f --- /dev/null +++ b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_ops.rs @@ -0,0 +1,391 @@ +//! SSCP, feed-forward, and light-convolution audio graph operations. + +use super::builder::{Builder, Val}; +use super::gemma3n_audio_math::{ + 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; + +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_inverse_stddev_sqrt_reciprocal: Val, + 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, + pub valid: Val, +} + +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, + norm_squared_at_time: Val, + norm_cumulative_squared: Val, + norm_variance: Val, + norm_stabilized_variance: Val, + norm_inverse_stddev: Val, + norm_inverse_stddev_sqrt_reciprocal: 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, + inverse_stddev_sqrt_reciprocal: Val, + output: 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, +) -> CumulativeGroupNormStages { + 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 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 = mlx_cuda_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_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); + 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 = 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 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()); + let weighted = b.multiply(&normalized, &weight); + 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, + inverse_stddev_sqrt_reciprocal, + output, + } +} + +#[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, + _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, + // 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.output); + 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, + 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_inverse_stddev_sqrt_reciprocal: normalized.inverse_stddev_sqrt_reciprocal, + norm: normalized.output, + activated, + } +} + +fn subsample_convs( + b: &mut Builder, + args: &EncoderArgs, + config: &Gemma3nXlaAudioConfig, +) -> (SscpConvStages, SscpConvStages) { + 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]); + // 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( + 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, + false, + ); + let conv1 = sscp_conv( + b, + &conv0.activated, + 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, + true, + ); + (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.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")), + ); + 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.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")), + ); + let hidden = round_bf16(b, &projected); + let valid = stride_time( + b, + &args.valid_mask, + config.time_stride_product().expect("validated stride"), + ); + 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_inverse_stddev_sqrt_reciprocal: conv0.norm_inverse_stddev_sqrt_reciprocal, + 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, + 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..ce13fb2a1 --- /dev/null +++ b/src/lib/mlxcel-xla/src/emitter/gemma3n_audio_schema.rs @@ -0,0 +1,286 @@ +//! Stable argument order for the split Gemma3n audio artifacts. + +use std::collections::BTreeMap; + +use super::builder::{Builder, Ty, Val}; +use super::gemma3n::Gemma3nConfig; +#[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, + pub loc: String, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub(crate) 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, + } +} + +pub(crate) 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) +} + +pub(crate) 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..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,167 @@ 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 +/// 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 = 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, + &[], + 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/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 0065385c7..dbd040805 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; @@ -77,6 +82,14 @@ pub(crate) use gemma3n::{ Gemma3nConfig, Gemma3nLayerType, Gemma3nPleDType, Gemma3nPleMetadata, validate_gemma3n_ple, }; #[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_config.rs b/src/lib/mlxcel-xla/src/gemma3n_audio_config.rs new file mode 100644 index 000000000..ee7198be0 --- /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 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-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; +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/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/gemma3n_audio_runtime.rs b/src/lib/mlxcel-xla/src/gemma3n_audio_runtime.rs new file mode 100644 index 000000000..fe03d0100 --- /dev/null +++ b/src/lib/mlxcel-xla/src/gemma3n_audio_runtime.rs @@ -0,0 +1,1073 @@ +//! 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. + +#[cfg(feature = "diagnostics")] +use std::collections::BTreeMap; +use std::fs::File; +use std::io::Read; +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_sequential, 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, + #[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, + 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 + } + + #[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] + } + + #[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, + #[cfg(feature = "diagnostics")] + diagnostic_shapes: Vec<(String, Vec)>, +} + +fn sha256_hex(bytes: &[u8]) -> String { + 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"); + } + 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) + )); + } + generation_identity_from_version( + compiler, + flags, + mlir, + String::from_utf8_lossy(&version.stdout).trim(), + ) +} + +fn generation_identity_from_version( + compiler: &Path, + flags: &[&str], + mlir: &str, + version: &str, +) -> Result { + Ok(format!( + "compiler={};compiler_sha256={};version={version};flags={flags:?};mlir_sha256={}", + compiler.display(), + 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}")) +} + +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_sequential( + 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 + )); + } + 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], + 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 { + #[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, + 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, 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); + 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_legacy_unqualified( + "audio_encode.main", + format!("encode:{identity}:language={language_fingerprint:016x}"), + generation_identity(&compiler, &flags, &encode_mlir)?, + )?; + let merge_contract = AuxiliaryArtifactContract::new_legacy_unqualified( + "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); + #[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_inverse_stddev_sqrt_reciprocal", + "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", + "conformer.0.feed_forward_start", + "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, + audio, + frame_bucket, + clips, + context_capacity, + hidden_size: text.hidden, + layers: text.n_layers, + hidden_per_layer: text.hidden_per_layer_input, + language_fingerprint, + #[cfg(feature = "diagnostics")] + diagnostic_shapes, + }) + } + + #[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(); + 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( + &encode_inputs, + &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], + }, + ], + )?; + #[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() + .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, + #[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, + 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::*; + 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() { + 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], + #[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, + 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")); + } + + #[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_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; + 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_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") + .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_sequential( + &[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/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/iree.rs b/src/lib/mlxcel-xla/src/iree.rs index e31100190..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, }; @@ -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; } @@ -1547,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(), @@ -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 2b1328267..99c288f2a 100644 --- a/src/lib/mlxcel-xla/src/lib.rs +++ b/src/lib/mlxcel-xla/src/lib.rs @@ -52,6 +52,11 @@ use mlxcel_core::session::{InferenceSession, PreparedPrefill, SessionCapabilitie mod context; #[cfg(any(feature = "diagnostics", test))] 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)] mod numeric_dtype_contract; @@ -67,10 +72,13 @@ mod operator_numeric_contract; mod prepared; 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; @@ -203,14 +211,33 @@ 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}; +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, +}; +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, + validate_gemma3n_audio_checkpoint, +}; #[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; @@ -265,6 +292,8 @@ pub struct XlaInferenceSession { #[cfg(feature = "iree")] engine: iree::IreeLlama, #[cfg(feature = "iree")] + device: String, + #[cfg(feature = "iree")] cache_len: i32, } @@ -327,6 +356,7 @@ impl XlaInferenceSession { context_capacity, eos_token_ids, engine, + device, cache_len: 0, }) } @@ -366,6 +396,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/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], }]; 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..3e1872821 --- /dev/null +++ b/src/lib/mlxcel-xla/src/prepared_gemma3n_audio.rs @@ -0,0 +1,514 @@ +//! Owned Gemma3n audio inputs and prepared-prefill outputs. +//! +//! 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. + +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, +}; +use crate::gemma3n_audio_rows::build_gemma3n_audio_row_indices; + +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, + NonZeroBucketPadding, +} + +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::NonZeroBucketPadding => { + f.write_str("Gemma3n audio static bucket tail must contain only zeros") + } + } + } +} + +impl std::error::Error for Gemma3nAudioInputError {} + +/// Canonical, zero-padded input to one static `audio.encode` 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]; + 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, + bucket: frame_bucket, + }); + } + } + 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, + 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 + } +} + +/// 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, +} + +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() + )); + } + 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, + }) + } + + #[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 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 + } + + 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_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( + vec![0.0; bucket * GEMMA3N_AUDIO_MEL_BINS], + valid, + vec![2], + bucket, + ) + .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], + 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::NonZeroBucketPadding) + )); + } + + #[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 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); + 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") + ); + } +} 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(), }); 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], }]; 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); 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)] 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,