diff --git a/Cargo.toml b/Cargo.toml index ab506bce0..e1080c5d2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -329,6 +329,11 @@ required-features = ["xla-diagnostics"] name = "xla_vision_reference_check" required-features = ["xla-diagnostics"] +# Actual-checkpoint eager MLX vs IREE Qwen2.5-VL window-attention oracle. +[[example]] +name = "xla_qwen25_vl_window_check" +required-features = ["xla-diagnostics"] + [[example]] name = "xla_aux_smoke" required-features = ["xla-iree"] diff --git a/examples/xla_qwen25_vl_window_check.rs b/examples/xla_qwen25_vl_window_check.rs new file mode 100644 index 000000000..03d76a48c --- /dev/null +++ b/examples/xla_qwen25_vl_window_check.rs @@ -0,0 +1,1030 @@ +// 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. + +//! Actual-checkpoint eager MLX vs IREE Qwen2.5-VL strict oracle. +//! +//! One production processor payload drives both vision implementations. The +//! gate compares window planning, reordered patch embeddings, representative +//! window attention, every configured full-attention layer, restoration, +//! prepared M-RoPE, and final-position logits. Negative runs force Qwen2-style +//! full attention, identity permutation, and zero vision positions. + +use std::fs; +use std::path::{Path, PathBuf}; + +use image::DynamicImage; +use mlxcel::{ + HostMultimodalPreprocessor, LoadedModel, Qwen2VlIreeHostPreprocessor, initialize_runtime, + load_model, +}; +use mlxcel_core::session::{OwnedTensor, PreparedPositions, PreparedTensorDType}; +use mlxcel_xla::{ + IreeQwen25VlDiagnosticProjector, Qwen25VlDiagnosticMutation, Qwen25VlLanguageDiagnosticEngine, +}; +use serde_json::{Value, json}; + +#[derive(Debug, Clone, Copy)] +struct Tolerance { + atol: f32, + rtol: f32, +} + +#[derive(Debug, Clone, Copy)] +struct ComparisonStats { + max_absolute: f32, + max_absolute_index: Option, + actual_at_max_absolute: Option, + expected_at_max_absolute: Option, + max_relative: f32, + actual_rms: f64, + expected_rms: f64, + error_rms: f64, + failures: usize, + non_finite: usize, +} + +fn argument(flag: &str) -> Option { + arguments(flag).into_iter().next() +} + +fn arguments(flag: &str) -> Vec { + let args = std::env::args().collect::>(); + args.iter() + .enumerate() + .filter(|(_, argument)| argument.as_str() == flag) + .filter_map(|(index, _)| args.get(index + 1).cloned()) + .collect() +} + +fn required_paths(flag: &str) -> Vec { + let paths = arguments(flag) + .into_iter() + .map(PathBuf::from) + .collect::>(); + assert!(!paths.is_empty(), "missing required {flag}"); + paths +} + +fn usize_argument(flag: &str, default: usize) -> usize { + argument(flag) + .map(|value| { + value + .parse::() + .unwrap_or_else(|_| panic!("{flag} must be an unsigned integer")) + }) + .unwrap_or(default) +} + +fn f32_argument(flag: &str, default: f32) -> f32 { + argument(flag) + .map(|value| { + value + .parse::() + .unwrap_or_else(|_| panic!("{flag} must be a finite non-negative number")) + }) + .filter(|value| value.is_finite() && *value >= 0.0) + .unwrap_or(default) +} + +fn mlx_f32(array: &mlxcel_core::MlxArray) -> Vec { + let array = mlxcel_core::astype(array, mlxcel_core::dtype::FLOAT32); + mlxcel_core::eval(&array); + mlxcel_core::array_to_raw_bytes(&array) + .chunks_exact(4) + .map(|bytes| f32::from_ne_bytes(bytes.try_into().expect("four-byte f32 chunk"))) + .collect() +} + +fn mlx_i32(array: &mlxcel_core::MlxArray) -> Vec { + mlxcel_core::eval(array); + mlxcel_core::array_to_raw_bytes(array) + .chunks_exact(4) + .map(|bytes| i32::from_ne_bytes(bytes.try_into().expect("four-byte i32 chunk"))) + .collect() +} + +#[derive(Debug, Clone, Copy)] +struct SnapshotStats { + max_abs: f32, + finite_nonzero: usize, + zero: usize, + non_finite: usize, + f16_square_overflow_candidates: usize, +} + +fn snapshot_stats(values: &[f32]) -> SnapshotStats { + // Largest finite F16 value is 65,504, so squaring a value above its + // square root in F16 produces infinity before the RMS mean is evaluated. + const F16_SQRT_MAX: f32 = 255.9375; + let mut stats = SnapshotStats { + max_abs: 0.0, + finite_nonzero: 0, + zero: 0, + non_finite: 0, + f16_square_overflow_candidates: 0, + }; + for &value in values { + if !value.is_finite() { + stats.non_finite += 1; + continue; + } + let absolute = value.abs(); + stats.max_abs = stats.max_abs.max(absolute); + if value == 0.0 { + stats.zero += 1; + } else { + stats.finite_nonzero += 1; + } + if absolute > F16_SQRT_MAX { + stats.f16_square_overflow_candidates += 1; + } + } + stats +} + +fn owned_f32(tensor: &OwnedTensor, label: &str) -> Vec { + assert_eq!( + tensor.dtype, + PreparedTensorDType::Float32, + "{label} must be float32" + ); + tensor + .bytes + .chunks_exact(4) + .map(|bytes| f32::from_ne_bytes(bytes.try_into().expect("four-byte f32 chunk"))) + .collect() +} + +fn owned_i32(tensor: &OwnedTensor, label: &str) -> Vec { + assert_eq!( + tensor.dtype, + PreparedTensorDType::Int32, + "{label} must be int32" + ); + tensor + .bytes + .chunks_exact(4) + .map(|bytes| i32::from_ne_bytes(bytes.try_into().expect("four-byte i32 chunk"))) + .collect() +} + +fn comparison_stats( + actual: &[f32], + expected: &[f32], + tolerance: Tolerance, +) -> Result { + if actual.len() != expected.len() { + return Err(format!( + "element count differs: actual={}, expected={}", + actual.len(), + expected.len() + )); + } + let mut stats = ComparisonStats { + max_absolute: 0.0, + max_absolute_index: None, + actual_at_max_absolute: None, + expected_at_max_absolute: None, + max_relative: 0.0, + actual_rms: 0.0, + expected_rms: 0.0, + error_rms: 0.0, + failures: 0, + non_finite: 0, + }; + let mut finite = 0usize; + for (index, (&observed, &reference)) in actual.iter().zip(expected).enumerate() { + if !observed.is_finite() || !reference.is_finite() { + stats.failures += 1; + stats.non_finite += 1; + continue; + } + let absolute = (observed - reference).abs(); + let relative = absolute / reference.abs().max(f32::MIN_POSITIVE); + stats.actual_rms += f64::from(observed).powi(2); + stats.expected_rms += f64::from(reference).powi(2); + stats.error_rms += f64::from(observed - reference).powi(2); + finite += 1; + if stats.max_absolute_index.is_none() || absolute > stats.max_absolute { + stats.max_absolute = absolute; + stats.max_absolute_index = Some(index); + stats.actual_at_max_absolute = Some(observed); + stats.expected_at_max_absolute = Some(reference); + } + stats.max_relative = stats.max_relative.max(relative); + if absolute > tolerance.atol + tolerance.rtol * reference.abs() { + stats.failures += 1; + } + } + if finite > 0 { + let denominator = finite as f64; + stats.actual_rms = (stats.actual_rms / denominator).sqrt(); + stats.expected_rms = (stats.expected_rms / denominator).sqrt(); + stats.error_rms = (stats.error_rms / denominator).sqrt(); + } + Ok(stats) +} + +fn compare_stage( + reports: &mut Vec, + stage: &str, + actual: &[f32], + expected: &[f32], + tolerance: Tolerance, +) -> Result<(), String> { + let stats = comparison_stats(actual, expected, tolerance) + .map_err(|error| format!("{stage}: {error}"))?; + reports.push(json!({ + "stage": stage, + "elements": actual.len(), + "atol": tolerance.atol, + "rtol": tolerance.rtol, + "max_absolute": stats.max_absolute, + "max_absolute_index": stats.max_absolute_index, + "actual_at_max_absolute": stats.actual_at_max_absolute, + "expected_at_max_absolute": stats.expected_at_max_absolute, + "max_relative": stats.max_relative, + "actual_rms": stats.actual_rms, + "expected_rms": stats.expected_rms, + "error_rms": stats.error_rms, + "failures": stats.failures, + "non_finite": stats.non_finite, + "passed": stats.failures == 0, + })); + if stats.failures == 0 { + Ok(()) + } else { + Err(format!( + "{stage}: {} values exceeded tolerance (max_abs={}, max_rel={})", + stats.failures, stats.max_absolute, stats.max_relative + )) + } +} + +fn record_stage( + reports: &mut Vec, + failures: &mut Vec, + stage: &str, + actual: &[f32], + expected: &[f32], + tolerance: Tolerance, +) { + if let Err(error) = compare_stage(reports, stage, actual, expected, tolerance) { + failures.push(error); + } +} + +fn emit_report(path: Option<&Path>, report: &Value) { + let rendered = serde_json::to_string_pretty(report).expect("serialize oracle report"); + if let Some(path) = path { + fs::write(path, &rendered) + .unwrap_or_else(|error| panic!("write {}: {error}", path.display())); + } + println!("{rendered}"); +} + +fn vision_patch_shape(model: &Path, values: &[f32], grids: &[(i32, i32, i32)]) -> [i32; 2] { + let config: Value = serde_json::from_slice( + &fs::read(model.join("config.json")).expect("read Qwen2.5-VL config.json"), + ) + .expect("parse Qwen2.5-VL config.json"); + let vision = &config["vision_config"]; + let temporal = vision["temporal_patch_size"] + .as_u64() + .and_then(|value| usize::try_from(value).ok()) + .expect("vision temporal_patch_size"); + let rows = grids + .iter() + .map(|&(t, h, w)| { + usize::try_from(t) + .expect("positive grid temporal") + .checked_mul(usize::try_from(h).expect("positive grid height")) + .and_then(|count| { + count.checked_mul(usize::try_from(w).expect("positive grid width")) + }) + .and_then(|count| count.checked_mul(temporal)) + .expect("vision patch row count") + }) + .sum::(); + assert_eq!(values.len() % rows, 0, "processor patch rows must be dense"); + [ + i32::try_from(rows).expect("patch rows fit i32"), + i32::try_from(values.len() / rows).expect("patch width fits i32"), + ] +} + +fn usize_plan(values: &[i32], label: &str) -> Vec { + values + .iter() + .map(|&value| usize::try_from(value).unwrap_or_else(|_| panic!("{label} contains {value}"))) + .collect() +} + +fn argmax(values: &[f32]) -> usize { + values + .iter() + .enumerate() + .max_by(|left, right| left.1.total_cmp(right.1)) + .map_or(0, |(index, _)| index) +} + +fn main() { + let model_path = required_paths("--model") + .into_iter() + .next() + .expect("--model was checked"); + let image_paths = required_paths("--image"); + let device = argument("--device").unwrap_or_else(|| "local-task".to_string()); + let context_capacity = usize_argument("--context-capacity", 256); + let tolerance = Tolerance { + atol: f32_argument("--atol", 0.08), + rtol: f32_argument("--rtol", 0.08), + }; + let report_path = argument("--report").map(PathBuf::from); + let _runtime = initialize_runtime(); + let images = image_paths + .iter() + .map(|path| { + image::open(path) + .unwrap_or_else(|error| panic!("open {}: {error}", path.display())) + .into_rgb8() + }) + .map(DynamicImage::ImageRgb8) + .collect::>(); + + let production_preprocessor = Qwen2VlIreeHostPreprocessor::load(&model_path, &device) + .unwrap_or_else(|error| panic!("load Qwen2.5-VL production preprocessor: {error}")); + let processor = production_preprocessor + .capture_processor_inputs(&images) + .unwrap_or_else(|error| panic!("capture production processor inputs: {error}")); + + let (loaded, _) = load_model(&model_path) + .unwrap_or_else(|error| panic!("load eager Qwen2.5-VL checkpoint: {error}")); + let LoadedModel::Qwen25VL(model) = loaded else { + panic!("checkpoint did not load as Qwen2.5-VL"); + }; + let mut unexpanded_tokens = vec![1]; + for _ in &images { + unexpanded_tokens.extend([model.vision_start_token_id, model.image_token_id]); + } + unexpanded_tokens.push(2); + let prepared = production_preprocessor + .prepare(&unexpanded_tokens, &images) + .unwrap_or_else(|error| panic!("prepare production Qwen2.5-VL prefill: {error}")); + assert!( + prepared.sequence_len <= context_capacity, + "prepared sequence length {} exceeds diagnostic capacity {context_capacity}", + prepared.sequence_len + ); + let input_ids = mlxcel_core::from_slice_i32( + &prepared.token_ids, + &[ + 1, + i32::try_from(prepared.sequence_len).expect("sequence length fits i32"), + ], + ); + let text_embeddings = model.text_model.get_embed_tokens(&input_ids); + let patch_shape = vision_patch_shape(&model_path, &processor.patch_values, &processor.grids); + let pixels = mlxcel_core::from_slice_f32(&processor.patch_values, &patch_shape); + let pixels = mlxcel_core::astype(&pixels, mlxcel_core::array_dtype(&text_embeddings)); + let eager_vision = model + .vision_encoder + .forward_with_grid_diagnostics(&pixels, &processor.grids); + let eager_mlp_gate_projection_dense_f32_control = eager_vision + .substage_probe_layer_mlp_gate_projection_dense_f32_control + .as_slice(); + let eager_mlp_up_projection_dense_f32_control = eager_vision + .substage_probe_layer_mlp_up_projection_dense_f32_control + .as_slice(); + // The producer exports host-owned F32 values at each substage boundary. + // Reject a vacuous capture immediately so the strict oracle cannot report + // a numeric false positive. + let eager_substage_probe = [ + ("input", eager_vision.substage_probe_layer_input.as_slice()), + ("norm1", eager_vision.substage_probe_layer_norm1.as_slice()), + ("query", eager_vision.substage_probe_layer_query.as_slice()), + ("key", eager_vision.substage_probe_layer_key.as_slice()), + ("value", eager_vision.substage_probe_layer_value.as_slice()), + ( + "attention_context", + eager_vision + .substage_probe_layer_attention_context + .as_slice(), + ), + ( + "attention", + eager_vision.substage_probe_layer_attention.as_slice(), + ), + ( + "post_attention_residual", + eager_vision + .substage_probe_layer_post_attention_residual + .as_slice(), + ), + ("norm2", eager_vision.substage_probe_layer_norm2.as_slice()), + ( + "mlp_gate_projection", + eager_vision + .substage_probe_layer_mlp_gate_projection + .as_slice(), + ), + ( + "mlp_gate_activation", + eager_vision + .substage_probe_layer_mlp_gate_activation + .as_slice(), + ), + ( + "mlp_up_projection", + eager_vision + .substage_probe_layer_mlp_up_projection + .as_slice(), + ), + ( + "mlp_gated_product", + eager_vision + .substage_probe_layer_mlp_gated_product + .as_slice(), + ), + ( + "mlp_down_projection", + eager_vision + .substage_probe_layer_mlp_down_projection + .as_slice(), + ), + ]; + for (stage, values) in eager_substage_probe + .iter() + .filter(|(stage, _)| matches!(*stage, "norm1" | "norm2")) + { + let stats = snapshot_stats(values); + assert!( + stats.finite_nonzero > 0, + "eager substage probe {stage} has no finite non-zero values before IREE launch: \ + {stats:?}; layer input stats: {:?}. F16 RMSNorm squares values before its mean, \ + so any input magnitude above sqrt(65504) can zero an entire normalized row", + snapshot_stats(eager_vision.substage_probe_layer_input.as_slice()) + ); + } + + let mut iree_projector = IreeQwen25VlDiagnosticProjector::load(&model_path, &device) + .unwrap_or_else(|error| panic!("load Qwen2.5-VL IREE diagnostics: {error}")); + let iree_vision = iree_projector + .capture( + &processor.patch_values, + &processor.grids, + Qwen25VlDiagnosticMutation::None, + ) + .unwrap_or_else(|error| panic!("capture Qwen2.5-VL IREE diagnostics: {error}")); + let iree_norm2_mlp_projection_controls = model + .vision_encoder + .f16_input_dense_f32_mlp_projection_controls( + iree_vision.substage_probe_layer_index, + &iree_vision.substage_probe_layer_norm2, + iree_vision.patch_tokens, + ) + .unwrap_or_else(|error| { + panic!("project captured IREE norm2 through eager Qwen2.5-VL MLP weights: {error}") + }); + + assert_eq!( + iree_vision.window_index, + usize_plan(&eager_vision.window_index, "eager window_index"), + "window_index differs" + ); + assert_eq!( + iree_vision.restore_indices, + usize_plan(&eager_vision.restore_indices, "eager restore_indices"), + "restoration permutation differs" + ); + assert_eq!( + iree_vision.packed_cu_seqlens, + usize_plan(&eager_vision.packed_cu_seqlens, "eager packed cu lengths"), + "full-attention cumulative lengths differ" + ); + assert_eq!( + iree_vision.window_cu_seqlens, + usize_plan(&eager_vision.window_cu_seqlens, "eager window cu lengths"), + "window-attention cumulative lengths differ" + ); + assert!( + iree_vision.attention_mask_audit.full_window_differences > 0, + "strict oracle fixture does not structurally distinguish window and full attention" + ); + assert_eq!( + iree_vision.attention_mask_audit.cross_media_full_bias_leaks, 0, + "full-attention bias leaks across media" + ); + assert_eq!( + iree_vision + .attention_mask_audit + .active_padding_full_bias_leaks, + 0, + "full-attention bias connects active and padded bucket rows" + ); + assert_eq!( + iree_vision + .attention_mask_audit + .invalid_padded_full_bias_entries, + 0, + "padded bucket queries must expose only one finite self key" + ); + assert!( + iree_vision + .window_index + .iter() + .enumerate() + .any(|(index, &position)| index != position), + "strict oracle fixture must exercise a non-identity window permutation" + ); + assert!( + iree_vision.window_cu_seqlens.len() > iree_vision.packed_cu_seqlens.len(), + "strict oracle fixture must span multiple local attention windows" + ); + + let mut reports = Vec::new(); + let mut comparison_failures = Vec::new(); + record_stage( + &mut reports, + &mut comparison_failures, + "vision.reordered_patch_embedding", + &iree_vision.reordered_patch_embedding, + &mlx_f32(&eager_vision.reordered_patch_embedding), + tolerance, + ); + assert_eq!( + iree_vision.window_layer_index, eager_vision.window_layer_index, + "representative window layer differs" + ); + record_stage( + &mut reports, + &mut comparison_failures, + &format!( + "vision.post_window_layer_{}", + iree_vision.window_layer_index + ), + &iree_vision.post_window_layer, + &mlx_f32(&eager_vision.post_window_layer), + tolerance, + ); + assert_eq!( + iree_vision.full_layer_indices, eager_vision.full_layer_indices, + "configured full-attention layers differ" + ); + assert_eq!( + iree_vision.post_full_layers.len(), + eager_vision.post_full_layers.len(), + "captured full-attention layer count differs" + ); + assert_eq!( + iree_vision.final_interval_layer_indices, eager_vision.final_interval_layer_indices, + "final diagnostic interval differs" + ); + assert_eq!( + iree_vision.post_final_interval_layers.len(), + eager_vision.post_final_interval_layers.len(), + "captured final diagnostic interval count differs" + ); + assert_eq!( + iree_vision.diagnostic_layer_indices, eager_vision.diagnostic_layer_indices, + "diagnostic layer interval differs" + ); + assert_eq!( + iree_vision.post_diagnostic_layers.len(), + eager_vision.post_diagnostic_layers.len(), + "captured diagnostic layer count differs" + ); + let target_full_layer_index = *iree_vision + .full_layer_indices + .last() + .expect("IREE diagnostics require a full-attention layer"); + assert_eq!( + Some(&target_full_layer_index), + eager_vision.full_layer_indices.last(), + "target full-attention layer differs" + ); + assert_eq!( + iree_vision.substage_probe_layer_index, eager_vision.substage_probe_layer_index, + "substage probe layer differs" + ); + let substage_probe_layer_index = iree_vision.substage_probe_layer_index; + for (capture_index, (&layer, (actual, expected))) in iree_vision + .full_layer_indices + .iter() + .zip( + iree_vision + .post_full_layers + .iter() + .zip(&eager_vision.post_full_layers), + ) + .enumerate() + .filter(|&(_, (&layer, _))| layer != target_full_layer_index) + { + record_stage( + &mut reports, + &mut comparison_failures, + &format!("vision.post_full_layer_{layer}_capture_{capture_index}"), + actual, + &mlx_f32(expected), + tolerance, + ); + } + for (&layer, (actual, expected)) in iree_vision + .final_interval_layer_indices + .iter() + .zip( + iree_vision + .post_final_interval_layers + .iter() + .zip(&eager_vision.post_final_interval_layers), + ) + .filter(|&(&layer, _)| layer != target_full_layer_index) + { + record_stage( + &mut reports, + &mut comparison_failures, + &format!("vision.post_layer_{layer}_final_interval"), + actual, + &mlx_f32(expected), + tolerance, + ); + } + for (&layer, (actual, expected)) in iree_vision.diagnostic_layer_indices.iter().zip( + iree_vision + .post_diagnostic_layers + .iter() + .zip(&eager_vision.post_diagnostic_layers), + ) { + record_stage( + &mut reports, + &mut comparison_failures, + &format!("vision.post_layer_{layer}_diagnostic"), + actual, + &mlx_f32(expected), + tolerance, + ); + } + for ((stage, actual), (eager_stage, expected)) in [ + ("input", iree_vision.substage_probe_layer_input.as_slice()), + ("norm1", iree_vision.substage_probe_layer_norm1.as_slice()), + ("query", iree_vision.substage_probe_layer_query.as_slice()), + ("key", iree_vision.substage_probe_layer_key.as_slice()), + ("value", iree_vision.substage_probe_layer_value.as_slice()), + ( + "attention_context", + iree_vision + .substage_probe_layer_attention_context + .as_slice(), + ), + ( + "attention", + iree_vision.substage_probe_layer_attention.as_slice(), + ), + ( + "post_attention_residual", + iree_vision + .substage_probe_layer_post_attention_residual + .as_slice(), + ), + ("norm2", iree_vision.substage_probe_layer_norm2.as_slice()), + ( + "mlp_gate_projection", + iree_vision + .substage_probe_layer_mlp_gate_projection + .as_slice(), + ), + ( + "mlp_gate_activation", + iree_vision + .substage_probe_layer_mlp_gate_activation + .as_slice(), + ), + ( + "mlp_up_projection", + iree_vision + .substage_probe_layer_mlp_up_projection + .as_slice(), + ), + ( + "mlp_gated_product", + iree_vision + .substage_probe_layer_mlp_gated_product + .as_slice(), + ), + ( + "mlp_down_projection", + iree_vision + .substage_probe_layer_mlp_down_projection + .as_slice(), + ), + ] + .into_iter() + .zip(&eager_substage_probe) + { + assert_eq!(stage, *eager_stage, "substage probe order differs"); + record_stage( + &mut reports, + &mut comparison_failures, + &format!("vision.layer_{substage_probe_layer_index}.{stage}"), + actual, + expected, + tolerance, + ); + } + record_stage( + &mut reports, + &mut comparison_failures, + &format!("vision.layer_{substage_probe_layer_index}.mlp_gate_projection_dense_f32_control"), + &iree_vision.substage_probe_layer_mlp_gate_projection, + eager_mlp_gate_projection_dense_f32_control, + tolerance, + ); + record_stage( + &mut reports, + &mut comparison_failures, + &format!("vision.layer_{substage_probe_layer_index}.mlp_up_projection_dense_f32_control"), + &iree_vision.substage_probe_layer_mlp_up_projection, + eager_mlp_up_projection_dense_f32_control, + tolerance, + ); + record_stage( + &mut reports, + &mut comparison_failures, + &format!( + "vision.layer_{substage_probe_layer_index}.mlp_gate_projection_iree_norm2_dense_f32_control" + ), + &iree_vision.substage_probe_layer_mlp_gate_projection, + &iree_norm2_mlp_projection_controls.gate_projection, + tolerance, + ); + record_stage( + &mut reports, + &mut comparison_failures, + &format!( + "vision.layer_{substage_probe_layer_index}.mlp_up_projection_iree_norm2_dense_f32_control" + ), + &iree_vision.substage_probe_layer_mlp_up_projection, + &iree_norm2_mlp_projection_controls.up_projection, + tolerance, + ); + let target_capture_index = iree_vision + .full_layer_indices + .iter() + .position(|&layer| layer == target_full_layer_index) + .expect("target full layer is present in configured captures"); + record_stage( + &mut reports, + &mut comparison_failures, + &format!("vision.post_full_layer_{target_full_layer_index}_capture_{target_capture_index}"), + &iree_vision.post_full_layers[target_capture_index], + &mlx_f32(&eager_vision.post_full_layers[target_capture_index]), + tolerance, + ); + record_stage( + &mut reports, + &mut comparison_failures, + "vision.merger_window_ordered", + &iree_vision.merger_window_ordered, + &mlx_f32(&eager_vision.merger_window_ordered), + tolerance, + ); + record_stage( + &mut reports, + &mut comparison_failures, + "vision.restored_projection", + &iree_vision.restored_projection, + &mlx_f32(&eager_vision.restored_projection), + tolerance, + ); + if !comparison_failures.is_empty() { + let report = json!({ + "schema": 2, + "model": model_path, + "images": image_paths, + "device": device, + "context_capacity": context_capacity, + "grid_thw": processor.grids, + "patch_bucket": iree_vision.patch_bucket, + "patch_tokens": iree_vision.patch_tokens, + "padded_bucket_rows": iree_vision.patch_bucket - iree_vision.patch_tokens, + "attention_mask_audit": { + "full_window_differences": iree_vision.attention_mask_audit.full_window_differences, + "cross_media_full_bias_leaks": iree_vision.attention_mask_audit.cross_media_full_bias_leaks, + "active_padding_full_bias_leaks": iree_vision.attention_mask_audit.active_padding_full_bias_leaks, + "invalid_padded_full_bias_entries": iree_vision.attention_mask_audit.invalid_padded_full_bias_entries, + "all_diagnostic_outputs_finite_including_padding": true, + }, + "merged_tokens": iree_vision.merged_tokens, + "vision_hidden": iree_vision.vision_hidden, + "text_hidden": iree_vision.text_hidden, + "window_layer": iree_vision.window_layer_index, + "full_attention_layers": iree_vision.full_layer_indices, + "final_interval_layers": iree_vision.final_interval_layer_indices, + "diagnostic_layers": iree_vision.diagnostic_layer_indices, + "target_full_layer": target_full_layer_index, + "substage_probe_layer": substage_probe_layer_index, + "mlp_projection_control": { + "runtime": "eager_mlx", + "input": "same_eager_norm2", + "weight": "same_checkpoint_weight_cast_to_f32", + "operation": "dense_f32_matmul_then_f32_bias", + "purpose": "distinguish native_f16_projection_from_f32_accumulation_contract", + }, + "iree_norm2_mlp_projection_control": { + "runtime": "eager_mlx", + "input": "captured_iree_norm2_rounded_to_f16", + "weight": "same_checkpoint_weight_cast_to_f32", + "operation": "dense_f32_matmul_then_f32_bias", + "purpose": "distinguish inherited_norm2_drift_from_projection_reduction_behavior", + }, + "failed_phase": "vision", + "failures": comparison_failures, + "comparisons": reports, + "passed": false, + }); + emit_report(report_path.as_deref(), &report); + panic!( + "strict oracle vision comparison failed: {}", + report["failures"] + ); + } + + let qwen2_full = iree_vision_capture( + &mut iree_projector, + &processor.patch_values, + &processor.grids, + Qwen25VlDiagnosticMutation::Qwen2FullAttention, + ); + let qwen2_full_attention_numerically_detected = compare_stage( + &mut Vec::new(), + "negative.qwen2_full_attention", + &qwen2_full.post_window_layer, + &mlx_f32(&eager_vision.post_window_layer), + tolerance, + ) + .is_err(); + let identity_mutation = iree_vision_capture( + &mut iree_projector, + &processor.patch_values, + &processor.grids, + Qwen25VlDiagnosticMutation::IdentityPermutation, + ); + assert!( + compare_stage( + &mut Vec::new(), + "negative.identity_permutation", + &identity_mutation.reordered_patch_embedding, + &mlx_f32(&eager_vision.reordered_patch_embedding), + tolerance, + ) + .is_err(), + "negative oracle accepted an identity patch permutation" + ); + let zero_positions = iree_vision_capture( + &mut iree_projector, + &processor.patch_values, + &processor.grids, + Qwen25VlDiagnosticMutation::ZeroVisionPositions, + ); + assert!( + compare_stage( + &mut Vec::new(), + "negative.zero_vision_positions", + &zero_positions.post_window_layer, + &mlx_f32(&eager_vision.post_window_layer), + tolerance, + ) + .is_err(), + "negative oracle accepted zeroed vision positions" + ); + + let eager_embeddings = model.get_input_embeddings(&input_ids, &pixels, &processor.grids); + let prepared_embedding_values = owned_f32(&prepared.embeddings, "prepared embeddings"); + let prepared_hidden = prepared.embeddings.shape[2]; + let mut production_image_rows = Vec::new(); + for (position, &token) in prepared.token_ids.iter().enumerate() { + if token == model.image_token_id { + let start = position * prepared_hidden; + production_image_rows + .extend_from_slice(&prepared_embedding_values[start..start + prepared_hidden]); + } + } + compare_stage( + &mut reports, + "production.restored_projection", + &iree_vision.restored_projection, + &production_image_rows, + tolerance, + ) + .unwrap_or_else(|error| panic!("{error}")); + compare_stage( + &mut reports, + "prepared.merged_embeddings", + &prepared_embedding_values, + &mlx_f32(&eager_embeddings.inputs_embeds), + tolerance, + ) + .unwrap_or_else(|error| panic!("{error}")); + let (eager_positions, eager_rope_delta) = model.mrope_diagnostics(&input_ids, &processor.grids); + let PreparedPositions::Mrope3D { tensor, rope_delta } = &prepared.positions else { + panic!("production Qwen2.5-VL prefill did not export M-RoPE positions"); + }; + assert_eq!( + owned_i32(tensor, "prepared M-RoPE"), + mlx_i32(&eager_positions), + "prepared M-RoPE coordinates differ" + ); + assert_eq!( + *rope_delta, eager_rope_delta, + "prepared M-RoPE delta differs" + ); + + let mut eager_caches = model.text_model.make_caches(); + let eager_logits = model.text_model.forward_impl( + &input_ids, + Some(&eager_embeddings.inputs_embeds), + &mut eager_caches, + eager_embeddings.attention_mask_4d.as_deref(), + ); + let eager_logits = mlx_f32(&eager_logits); + let mut iree_language = + Qwen25VlLanguageDiagnosticEngine::load(&model_path, &device, context_capacity) + .unwrap_or_else(|error| panic!("load Qwen2.5-VL language diagnostics: {error}")); + let iree_logits = iree_language + .capture(&prepared) + .unwrap_or_else(|error| panic!("capture Qwen2.5-VL final logits: {error}")); + let vocab = iree_logits.len(); + let eager_last = + &eager_logits[(prepared.sequence_len - 1) * vocab..prepared.sequence_len * vocab]; + assert_eq!( + argmax(&iree_logits), + argmax(eager_last), + "final-position greedy token differs" + ); + compare_stage( + &mut reports, + "language.final_position_logits", + &iree_logits, + eager_last, + tolerance, + ) + .unwrap_or_else(|error| panic!("{error}")); + + let report = json!({ + "schema": 2, + "model": model_path, + "images": image_paths, + "device": device, + "context_capacity": context_capacity, + "grid_thw": processor.grids, + "patch_bucket": iree_vision.patch_bucket, + "patch_tokens": iree_vision.patch_tokens, + "padded_bucket_rows": iree_vision.patch_bucket - iree_vision.patch_tokens, + "attention_mask_audit": { + "full_window_differences": iree_vision.attention_mask_audit.full_window_differences, + "cross_media_full_bias_leaks": iree_vision.attention_mask_audit.cross_media_full_bias_leaks, + "active_padding_full_bias_leaks": iree_vision.attention_mask_audit.active_padding_full_bias_leaks, + "invalid_padded_full_bias_entries": iree_vision.attention_mask_audit.invalid_padded_full_bias_entries, + "all_diagnostic_outputs_finite_including_padding": true, + }, + "merged_tokens": iree_vision.merged_tokens, + "vision_hidden": iree_vision.vision_hidden, + "text_hidden": iree_vision.text_hidden, + "window_layer": iree_vision.window_layer_index, + "full_attention_layers": iree_vision.full_layer_indices, + "final_interval_layers": iree_vision.final_interval_layer_indices, + "diagnostic_layers": iree_vision.diagnostic_layer_indices, + "target_full_layer": iree_vision.full_layer_indices.last(), + "substage_probe_layer": iree_vision.substage_probe_layer_index, + "negative_qwen2_full_attention_structurally_detected": true, + "negative_qwen2_full_attention_numerically_detected": qwen2_full_attention_numerically_detected, + "negative_identity_permutation_detected": true, + "negative_zero_vision_positions_detected": true, + "final_top1": argmax(&iree_logits), + "comparisons": reports, + "passed": true, + }); + emit_report(report_path.as_deref(), &report); +} + +fn iree_vision_capture( + projector: &mut IreeQwen25VlDiagnosticProjector, + patch_values: &[f32], + grids: &[(i32, i32, i32)], + mutation: Qwen25VlDiagnosticMutation, +) -> mlxcel_xla::Qwen25VlVisionDiagnostics { + projector + .capture(patch_values, grids, mutation) + .unwrap_or_else(|error| panic!("capture Qwen2.5-VL negative diagnostic: {error}")) +} diff --git a/src/lib.rs b/src/lib.rs index d9520a12f..686879a0f 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -75,14 +75,16 @@ pub use mlxcel_core::generate::{ SamplingConfig, }; pub use mlxcel_core::speculative::SpeculativeGenerator; -#[cfg(feature = "xla-diagnostics")] -pub use multimodal::host_preprocessor::LlavaHostReferenceCapture; #[cfg(feature = "xla-iree")] pub use multimodal::host_preprocessor::LlavaIreeHostPreprocessor; pub use multimodal::host_preprocessor::{ FakeHostMultimodalPreprocessor, HostMultimodalPreprocessor, HostPreprocessorError, LlavaHostPreprocessor, XlaVisionBackend, load_xla_image_preprocessor, }; +#[cfg(feature = "xla-diagnostics")] +pub use multimodal::host_preprocessor::{ + LlavaHostReferenceCapture, Qwen2VlIreeHostPreprocessor, QwenVlProcessorCapture, +}; pub use multimodal::{ internvl_prompt, kimi_vl_prompt, minicpmo_prompt, moondream2_prompt, moondream3_prompt, phi3v_prompt, phi4_siglip_prompt, phi4mm_prompt, pixtral_prompt, qwen_vl, smolvlm_prompt, diff --git a/src/lib/mlxcel-xla/src/batch.rs b/src/lib/mlxcel-xla/src/batch.rs index b08cb4b2f..0edb85ec8 100644 --- a/src/lib/mlxcel-xla/src/batch.rs +++ b/src/lib/mlxcel-xla/src/batch.rs @@ -901,6 +901,36 @@ pub struct LlavaReferenceDiagnosticRun { pub decode_seconds: f64, } +/// Diagnostics-only Qwen2.5-VL language runner that invokes the ordinary +/// embeddings-seeded M-RoPE prefill and reads its final-position logits. +#[cfg(feature = "diagnostics")] +pub struct Qwen25VlLanguageDiagnosticEngine { + engine: IreeRaggedLlama, +} + +#[cfg(feature = "diagnostics")] +impl Qwen25VlLanguageDiagnosticEngine { + pub fn load(model_path: &Path, device: &str, context_capacity: usize) -> Result { + Ok(Self { + engine: IreeRaggedLlama::load(model_path, device, 4, context_capacity)?, + }) + } + + pub fn capture(&mut self, prepared: &PreparedPrefill) -> Result, String> { + self.engine.reset_slots_for_diagnostics()?; + let prepared_iree = PreparedIreePrefill::prepare( + prepared, + self.engine.hidden_size(), + self.engine.context_capacity(), + ) + .map_err(|error| error.to_string())?; + if prepared_iree.positions.mode() != PreparedPositionMode::Mrope3D { + return Err("Qwen2.5-VL diagnostics require prepared M-RoPE positions".to_string()); + } + self.engine.prefill_prepared_slot_logits(0, &prepared_iree) + } +} + #[cfg(feature = "diagnostics")] pub fn run_gemma3n_all_layer_diagnostics( model_dir: &Path, diff --git a/src/lib/mlxcel-xla/src/emitter/config.rs b/src/lib/mlxcel-xla/src/emitter/config.rs index e5c450eaa..e67327c75 100644 --- a/src/lib/mlxcel-xla/src/emitter/config.rs +++ b/src/lib/mlxcel-xla/src/emitter/config.rs @@ -717,7 +717,7 @@ impl Config { } } } - Some("qwen2") | Some("qwen2_vl") => { + Some("qwen2") | Some("qwen2_vl") | Some("qwen2_5_vl") => { // Qwen2-VL keeps the Qwen2 language configuration at the wrapper // root (unlike LLaVA's nested `text_config`). Its language graph // is therefore the ordinary Qwen2 graph with M-RoPE positions. @@ -1197,8 +1197,10 @@ impl Config { rotary_width / 2 )); } - let supported_family = - matches!(model_type, Some("qwen2") | Some("qwen2_vl") | Some("qwen3")); + let supported_family = matches!( + model_type, + Some("qwen2") | Some("qwen2_vl") | Some("qwen2_5_vl") | Some("qwen3") + ); if !supported_family { return Err(format!( "M-RoPE sections are only supported by the Qwen2/Qwen3 language backbone, got model_type={model_type:?}" @@ -1712,6 +1714,10 @@ mod tests { .expect("Qwen2-VL root language config parses"); assert_eq!(qwen2_vl.mrope, qwen2.mrope); assert!(qwen2_vl.qkv_bias); + let qwen25_vl = Config::from_json_str(&base("qwen2_5_vl", "[1,1,1]", "")) + .expect("Qwen2.5-VL root language config parses"); + assert_eq!(qwen25_vl.mrope, qwen2.mrope); + assert!(qwen25_vl.qkv_bias); let qwen3 = Config::from_json_str(&base("qwen3", "[1,1,1]", "")).expect("Qwen3 parses"); assert_eq!( qwen3.mrope.unwrap().layout, diff --git a/src/lib/mlxcel-xla/src/emitter/mod.rs b/src/lib/mlxcel-xla/src/emitter/mod.rs index 0065385c7..715600c50 100644 --- a/src/lib/mlxcel-xla/src/emitter/mod.rs +++ b/src/lib/mlxcel-xla/src/emitter/mod.rs @@ -59,6 +59,7 @@ mod model; mod moe; pub(crate) mod numeric_ops; mod phi4_audio; +mod qwen2_5_vl; mod qwen2_vl; mod rope; mod vision; @@ -121,13 +122,15 @@ pub(crate) use phi4_audio::{ emit_phi4_audio_diagnostic_with, emit_phi4_audio_with, phi4_audio_diagnostic_specs, phi4_audio_weight_specs, validate_phi4_audio_weight_shapes, }; -#[allow(unused_imports)] -pub(crate) use qwen2_vl::emit_qwen2_vl; +#[cfg(feature = "diagnostics")] +pub(crate) use qwen2_5_vl::{Qwen25VlVisionDiagnosticLayout, emit_qwen2_5_vl_diagnostics}; #[allow(unused_imports)] pub(crate) use qwen2_vl::{ QWEN2_VL_PATCH_BUCKETS, Qwen2VlConfig, Qwen2VlGridPlan, Qwen2VlHostInputs, Qwen2VlWeightSpec, - prepare_qwen2_vl_host_inputs, + QwenVlVisionVariant, prepare_qwen2_vl_host_inputs, }; +#[allow(unused_imports)] +pub(crate) use qwen2_vl::{emit_qwen2_vl, emit_qwen2_vl_with}; // MoE FFN config types (issue #500), read by the weight loader (`iree.rs`) and the // validation harness. `SharedExpertConfig` is only named in some build cfgs. #[allow(unused_imports)] diff --git a/src/lib/mlxcel-xla/src/emitter/qwen2_5_vl.rs b/src/lib/mlxcel-xla/src/emitter/qwen2_5_vl.rs new file mode 100644 index 000000000..61705a232 --- /dev/null +++ b/src/lib/mlxcel-xla/src/emitter/qwen2_5_vl.rs @@ -0,0 +1,617 @@ +// 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. + +//! Qwen2.5-VL StableHLO vision tower. +//! +//! Host preprocessing supplies the exact window permutation plus separate +//! full-media and window-isolated masks. The graph keeps the permutation +//! static for one bucket while selecting the checkpoint-declared full +//! attention layers during emission. + +use super::builder::{Builder, Precision, Ty, Val}; +use super::qwen2_vl::{ + QWEN2_VL_PATCH_BUCKETS, Qwen2VlConfig, Qwen2VlWeightSpec, QwenVlVisionVariant, +}; + +#[cfg(feature = "diagnostics")] +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct Qwen25VlVisionDiagnosticLayout { + pub(crate) window_layer_index: usize, + pub(crate) full_layer_indices: Vec, + pub(crate) final_interval_layer_indices: Vec, + pub(crate) diagnostic_layer_indices: Vec, + pub(crate) substage_probe_layer_index: usize, +} + +#[cfg(feature = "diagnostics")] +struct Qwen25VlBlockDiagnostics { + input: Val, + norm1: Val, + query: Val, + key: Val, + value: Val, + attention_context: Val, + attention: Val, + post_attention_residual: Val, + norm2: Val, + mlp_gate_projection: Val, + mlp_gate_activation: Val, + mlp_up_projection: Val, + mlp_gated_product: Val, + mlp_down_projection: Val, + output: Val, +} + +#[cfg(feature = "diagnostics")] +fn substage_probe_layer(full_layer_indices: &[usize]) -> Option { + let target = full_layer_indices + .iter() + .rev() + .nth(1) + .or_else(|| full_layer_indices.last()) + .copied()?; + full_layer_indices + .iter() + .copied() + .take_while(|&layer| layer < target) + .last() + .and_then(|previous_full| previous_full.checked_add(2)) + .filter(|&layer| layer <= target) + .or(Some(target)) +} + +#[cfg(feature = "diagnostics")] +fn diagnostic_layers(full_layer_indices: &[usize], probe: usize) -> Vec { + full_layer_indices + .iter() + .copied() + .take_while(|&layer| layer < probe) + .last() + .map(|previous_full| ((previous_full + 1)..=probe).collect()) + .unwrap_or_default() +} + +struct Qwen25VlAttentionValues { + query: Val, + key: Val, + value: Val, + context: Val, + output: Val, +} + +struct Args { + values: Vec, + declarations: Vec, + cursor: usize, +} + +impl Args { + fn new(specs: &[Qwen2VlWeightSpec]) -> Self { + let mut values = Vec::with_capacity(specs.len() + 4); + let mut declarations = Vec::with_capacity(specs.len() + 4); + for (index, spec) in specs.iter().enumerate() { + let ty = Ty::f32(spec.shape.clone()); + declarations.push(format!( + "%arg{index}: {} loc(\"{}\")", + ty.render(), + spec.name + )); + values.push(Builder::arg(index, ty)); + } + Self { + values, + declarations, + cursor: 0, + } + } + + fn take(&mut self) -> Val { + let value = self.values[self.cursor].clone(); + self.cursor += 1; + value + } + + fn push_input(&mut self, ty: Ty, name: &str) -> Val { + let index = self.values.len(); + self.declarations + .push(format!("%arg{index}: {} loc(\"{name}\")", ty.render())); + let value = Builder::arg(index, ty); + self.values.push(value.clone()); + value + } +} + +fn bias_2d(builder: &mut Builder, value: &Val, bias: &Val) -> Val { + let rows = value.ty.shape[0]; + let width = value.ty.shape[1]; + let bias = builder.broadcast(bias, &[1], vec![rows, width]); + builder.add(value, &bias) +} + +fn linear_2d(builder: &mut Builder, value: &Val, weight: &Val, bias: &Val) -> Val { + let value = builder.linear_seq(value, weight); + bias_2d(builder, &value, bias) +} + +fn rms_norm(builder: &mut Builder, value: &Val, weight: &Val, epsilon: f32) -> Val { + let rows = value.ty.shape[0]; + let width = value.ty.shape[1]; + let zero = builder.const_f32(0.0); + let squared = builder.multiply(value, value); + let sum = builder.reduce_add(&squared, 1, &zero); + let width_scalar = builder.const_f32(width as f32); + let width_rows = builder.broadcast(&width_scalar, &[], vec![rows]); + let mean = builder.divide(&sum, &width_rows); + let epsilon = builder.const_f32(epsilon); + let epsilon = builder.broadcast(&epsilon, &[], vec![rows]); + let mean = builder.add(&mean, &epsilon); + let inverse = builder.rsqrt(&mean); + let inverse = builder.broadcast(&inverse, &[0], vec![rows, width]); + let normalized = builder.multiply(value, &inverse); + let weight = builder.broadcast(weight, &[1], vec![rows, width]); + builder.multiply(&normalized, &weight) +} + +fn silu(builder: &mut Builder, value: &Val) -> Val { + let one = builder.const_f32(1.0); + let one = builder.broadcast(&one, &[], value.ty.shape.clone()); + let negative = builder.negate(value); + let exponential = builder.exponential(&negative); + let denominator = builder.add(&one, &exponential); + builder.divide(value, &denominator) +} + +fn exact_gelu(builder: &mut Builder, value: &Val) -> Val { + let shape = value.ty.shape.clone(); + let half = builder.const_f32(0.5); + let half = builder.broadcast(&half, &[], shape.clone()); + let one = builder.const_f32(1.0); + let one = builder.broadcast(&one, &[], shape.clone()); + let inv_sqrt_two = builder.const_f32(std::f32::consts::FRAC_1_SQRT_2); + let inv_sqrt_two = builder.broadcast(&inv_sqrt_two, &[], shape); + let scaled = builder.multiply(value, &inv_sqrt_two); + let erf = builder.erf(&scaled); + let cdf = builder.add(&one, &erf); + let half_value = builder.multiply(value, &half); + builder.multiply(&half_value, &cdf) +} + +fn rotate_half(builder: &mut Builder, value: &Val) -> Val { + let tokens = value.ty.shape[0]; + let heads = value.ty.shape[1]; + let width = value.ty.shape[2]; + let half = width / 2; + let first = builder.slice(value, &[(0, tokens), (0, heads), (0, half)]); + let second = builder.slice(value, &[(0, tokens), (0, heads), (half, width)]); + let second = builder.negate(&second); + builder.concatenate(&second, &first, 2) +} + +fn apply_vision_rope(builder: &mut Builder, value: &Val, freqs: &Val) -> Val { + let tokens = value.ty.shape[0]; + let heads = value.ty.shape[1]; + let half = freqs.ty.shape[1]; + let cos = builder.cosine(freqs); + let sin = builder.sine(freqs); + let cos = builder.concatenate(&cos, &cos, 1); + let sin = builder.concatenate(&sin, &sin, 1); + let cos = builder.broadcast(&cos, &[0, 2], vec![tokens, heads, half * 2]); + let sin = builder.broadcast(&sin, &[0, 2], vec![tokens, heads, half * 2]); + let rotated = rotate_half(builder, value); + let direct = builder.multiply(value, &cos); + let rotated = builder.multiply(&rotated, &sin); + builder.add(&direct, &rotated) +} + +fn attention( + builder: &mut Builder, + hidden: &Val, + args: &mut Args, + config: &Qwen2VlConfig, + freqs: &Val, + attention_bias: &Val, +) -> Qwen25VlAttentionValues { + let tokens = hidden.ty.shape[0]; + let head_dim = config.hidden / config.heads; + let qkv = linear_2d(builder, hidden, &args.take(), &args.take()); + let qkv = builder.reshape(&qkv, vec![tokens, 3, config.heads, head_dim]); + let q = builder.slice( + &qkv, + &[(0, tokens), (0, 1), (0, config.heads), (0, head_dim)], + ); + let k = builder.slice( + &qkv, + &[(0, tokens), (1, 2), (0, config.heads), (0, head_dim)], + ); + let v = builder.slice( + &qkv, + &[(0, tokens), (2, 3), (0, config.heads), (0, head_dim)], + ); + let q = builder.reshape(&q, vec![tokens, config.heads, head_dim]); + let k = builder.reshape(&k, vec![tokens, config.heads, head_dim]); + let v = builder.reshape(&v, vec![tokens, config.heads, head_dim]); + let query = apply_vision_rope(builder, &q, freqs); + let key = apply_vision_rope(builder, &k, freqs); + let value = v; + let q = builder.transpose(&query, &[1, 0, 2]); + let k = builder.transpose(&key, &[1, 0, 2]); + let v = builder.transpose(&value, &[1, 0, 2]); + let scores = builder.dot_general( + &q, + &k, + &[0], + &[0], + &[2], + &[2], + vec![config.heads, tokens, tokens], + ); + let scale = builder.const_f32((head_dim as f32).powf(-0.5)); + let scale = builder.broadcast(&scale, &[], vec![config.heads, tokens, tokens]); + let scores = builder.multiply(&scores, &scale); + let bias = builder.broadcast(attention_bias, &[1, 2], vec![config.heads, tokens, tokens]); + let scores = builder.add(&scores, &bias); + let negative_infinity = builder.const_f32(f32::NEG_INFINITY); + let maximum = builder.reduce_max(&scores, 2, &negative_infinity); + let maximum = builder.broadcast(&maximum, &[0, 1], vec![config.heads, tokens, tokens]); + let shifted = builder.subtract(&scores, &maximum); + let exponentials = builder.exponential(&shifted); + let zero = builder.const_f32(0.0); + let denominator = builder.reduce_add(&exponentials, 2, &zero); + let denominator = builder.broadcast(&denominator, &[0, 1], vec![config.heads, tokens, tokens]); + let probabilities = builder.divide(&exponentials, &denominator); + let context = builder.dot_general( + &probabilities, + &v, + &[0], + &[0], + &[2], + &[1], + vec![config.heads, tokens, head_dim], + ); + let context = builder.transpose(&context, &[1, 0, 2]); + let context = builder.reshape(&context, vec![tokens, config.hidden]); + let output = linear_2d(builder, &context, &args.take(), &args.take()); + Qwen25VlAttentionValues { + query, + key, + value, + context, + output, + } +} + +fn encoder_layer( + builder: &mut Builder, + hidden: &Val, + args: &mut Args, + config: &Qwen2VlConfig, + freqs: &Val, + attention_bias: &Val, +) -> Val { + let norm1 = rms_norm(builder, hidden, &args.take(), config.layer_norm_eps); + let attention = attention(builder, &norm1, args, config, freqs, attention_bias); + let residual = builder.add(hidden, &attention.output); + let norm2 = rms_norm(builder, &residual, &args.take(), config.layer_norm_eps); + let gate = linear_2d(builder, &norm2, &args.take(), &args.take()); + let gate = silu(builder, &gate); + let up = linear_2d(builder, &norm2, &args.take(), &args.take()); + let activated = builder.multiply(&gate, &up); + let down = linear_2d(builder, &activated, &args.take(), &args.take()); + builder.add(&residual, &down) +} + +#[cfg(feature = "diagnostics")] +fn encoder_layer_with_diagnostics( + builder: &mut Builder, + hidden: &Val, + args: &mut Args, + config: &Qwen2VlConfig, + freqs: &Val, + attention_bias: &Val, +) -> Qwen25VlBlockDiagnostics { + let norm1 = rms_norm(builder, hidden, &args.take(), config.layer_norm_eps); + let attention = attention(builder, &norm1, args, config, freqs, attention_bias); + let post_attention_residual = builder.add(hidden, &attention.output); + let norm2 = rms_norm( + builder, + &post_attention_residual, + &args.take(), + config.layer_norm_eps, + ); + let mlp_gate_projection = linear_2d(builder, &norm2, &args.take(), &args.take()); + let mlp_gate_activation = silu(builder, &mlp_gate_projection); + let mlp_up_projection = linear_2d(builder, &norm2, &args.take(), &args.take()); + let mlp_gated_product = builder.multiply(&mlp_gate_activation, &mlp_up_projection); + let mlp_down_projection = linear_2d(builder, &mlp_gated_product, &args.take(), &args.take()); + let output = builder.add(&post_attention_residual, &mlp_down_projection); + Qwen25VlBlockDiagnostics { + input: hidden.clone(), + norm1, + query: attention.query, + key: attention.key, + value: attention.value, + attention_context: attention.context, + attention: attention.output, + post_attention_residual, + norm2, + mlp_gate_projection, + mlp_gate_activation, + mlp_up_projection, + mlp_gated_product, + mlp_down_projection, + output, + } +} + +fn emit_qwen2_5_vl_inner( + config: &Qwen2VlConfig, + patch_bucket: usize, + precision: Precision, + #[cfg(feature = "diagnostics")] diagnostic_layout: Option<&Qwen25VlVisionDiagnosticLayout>, +) -> String { + assert!( + QWEN2_VL_PATCH_BUCKETS.contains(&patch_bucket), + "unqualified Qwen2.5-VL patch bucket" + ); + let QwenVlVisionVariant::Qwen25 { + full_attention_blocks, + .. + } = &config.variant + else { + panic!("Qwen2.5-VL emitter requires the Qwen2.5 variant"); + }; + let specs = config.weight_specs(); + let mut args = Args::new(&specs); + let mut builder = Builder::new().with_precision(precision); + let patch_weight = args.take(); + let patch_width = + config.channels * config.temporal_patch_size * config.patch_size * config.patch_size; + let head_dim = config.hidden / config.heads; + let patches = args.push_input(Ty::f32(vec![patch_bucket, patch_width]), "patches.windowed"); + let freqs = args.push_input( + Ty::f32(vec![patch_bucket, head_dim / 2]), + "vision_rope.windowed_freqs", + ); + let full_bias = args.push_input( + Ty::f32(vec![patch_bucket, patch_bucket]), + "full_attention.bias", + ); + let window_bias = args.push_input( + Ty::f32(vec![patch_bucket, patch_bucket]), + "window_attention.bias", + ); + let mut hidden = builder.linear_seq(&patches, &patch_weight); + #[cfg(feature = "diagnostics")] + let reordered_patch_embedding = diagnostic_layout.map(|_| hidden.clone()); + #[cfg(feature = "diagnostics")] + let mut window_layer_state = None; + #[cfg(feature = "diagnostics")] + let mut full_layer_states = Vec::new(); + #[cfg(feature = "diagnostics")] + let mut final_interval_layer_states = Vec::new(); + #[cfg(feature = "diagnostics")] + let mut diagnostic_layer_states = Vec::new(); + #[cfg(feature = "diagnostics")] + let mut substage_probe_layer_state = None; + for layer in 0..config.depth { + let bias = if full_attention_blocks.contains(&layer) { + &full_bias + } else { + &window_bias + }; + #[cfg(feature = "diagnostics")] + let is_substage_probe_layer = + diagnostic_layout.is_some_and(|layout| layer == layout.substage_probe_layer_index); + #[cfg(feature = "diagnostics")] + if is_substage_probe_layer { + let capture = encoder_layer_with_diagnostics( + &mut builder, + &hidden, + &mut args, + config, + &freqs, + bias, + ); + hidden = capture.output.clone(); + substage_probe_layer_state = Some(capture); + } else { + hidden = encoder_layer(&mut builder, &hidden, &mut args, config, &freqs, bias); + } + #[cfg(not(feature = "diagnostics"))] + { + hidden = encoder_layer(&mut builder, &hidden, &mut args, config, &freqs, bias); + } + #[cfg(feature = "diagnostics")] + if let Some(layout) = diagnostic_layout { + if layer == layout.window_layer_index { + window_layer_state = Some(hidden.clone()); + } + if layout.full_layer_indices.contains(&layer) { + full_layer_states.push(hidden.clone()); + } + if layout.final_interval_layer_indices.contains(&layer) { + final_interval_layer_states.push(hidden.clone()); + } + if layout.diagnostic_layer_indices.contains(&layer) { + diagnostic_layer_states.push(hidden.clone()); + } + } + } + hidden = rms_norm(&mut builder, &hidden, &args.take(), config.layer_norm_eps); + let merge_width = config.hidden * config.spatial_merge_size * config.spatial_merge_size; + let merged = builder.reshape( + &hidden, + vec![ + patch_bucket / (config.spatial_merge_size * config.spatial_merge_size), + merge_width, + ], + ); + let merged = linear_2d(&mut builder, &merged, &args.take(), &args.take()); + let merged = exact_gelu(&mut builder, &merged); + let projected = linear_2d(&mut builder, &merged, &args.take(), &args.take()); + assert_eq!(args.cursor, specs.len(), "Qwen2.5-VL weight schema drifted"); + #[cfg(feature = "diagnostics")] + if let Some(layout) = diagnostic_layout { + let mut outputs = vec![ + reordered_patch_embedding.expect("diagnostics capture patch embedding"), + window_layer_state.expect("diagnostics capture window layer"), + ]; + assert_eq!( + full_layer_states.len(), + layout.full_layer_indices.len(), + "diagnostics capture every configured full-attention layer" + ); + outputs.extend(full_layer_states); + assert_eq!( + final_interval_layer_states.len(), + layout.final_interval_layer_indices.len(), + "diagnostics capture every layer after the penultimate full-attention boundary" + ); + outputs.extend(final_interval_layer_states); + assert_eq!( + diagnostic_layer_states.len(), + layout.diagnostic_layer_indices.len(), + "diagnostics capture every requested boundary layer" + ); + outputs.extend(diagnostic_layer_states); + let probe = substage_probe_layer_state.expect("diagnostics capture substage probe layer"); + outputs.extend([ + probe.input, + probe.norm1, + probe.query, + probe.key, + probe.value, + probe.attention_context, + probe.attention, + probe.post_attention_residual, + probe.norm2, + probe.mlp_gate_projection, + probe.mlp_gate_activation, + probe.mlp_up_projection, + probe.mlp_gated_product, + probe.mlp_down_projection, + ]); + outputs.push(projected); + let result_types = outputs + .iter() + .map(|value| value.ty.render()) + .collect::>(); + return format!( + "module @qwen2_vl_vision {{\n func.func public @main({signature}) -> ({result_signature}) {{\n{body} return {results} : {result_signature}\n }}\n}}\n", + signature = args.declarations.join(", "), + result_signature = result_types.join(", "), + body = builder.body(), + results = outputs + .iter() + .map(|value| value.name.as_str()) + .collect::>() + .join(", "), + ); + } + format!( + "module @qwen2_vl_vision {{\n func.func public @main({signature}) -> {result_type} {{\n{body} return {result} : {result_type}\n }}\n}}\n", + signature = args.declarations.join(", "), + result_type = projected.ty.render(), + body = builder.body(), + result = projected.name, + ) +} + +pub(super) fn emit_qwen2_5_vl( + config: &Qwen2VlConfig, + patch_bucket: usize, + precision: Precision, +) -> String { + emit_qwen2_5_vl_inner( + config, + patch_bucket, + precision, + #[cfg(feature = "diagnostics")] + None, + ) +} + +#[cfg(feature = "diagnostics")] +pub(crate) fn emit_qwen2_5_vl_diagnostics( + config: &Qwen2VlConfig, + patch_bucket: usize, + precision: Precision, +) -> Result<(String, Qwen25VlVisionDiagnosticLayout), String> { + let QwenVlVisionVariant::Qwen25 { + full_attention_blocks, + .. + } = &config.variant + else { + return Err("Qwen2.5-VL diagnostics require the Qwen2.5 vision variant".to_string()); + }; + let window_layer_index = (0..config.depth) + .find(|layer| !full_attention_blocks.contains(layer)) + .ok_or_else(|| { + "Qwen2.5-VL diagnostics require at least one window-attention layer".to_string() + })?; + if full_attention_blocks.is_empty() { + return Err( + "Qwen2.5-VL diagnostics require at least one configured full-attention layer" + .to_string(), + ); + } + let full_layer_indices = (0..config.depth) + .filter(|layer| full_attention_blocks.contains(layer)) + .collect::>(); + let target_full_layer_index = *full_layer_indices + .last() + .expect("non-empty full-attention layers checked above"); + let final_interval_start = full_layer_indices + .iter() + .rev() + .nth(1) + .map_or(0, |layer| layer + 1); + let final_interval_layer_indices = + (final_interval_start..=target_full_layer_index).collect::>(); + let substage_probe_layer_index = substage_probe_layer(&full_layer_indices) + .expect("diagnostics contain a configured full-attention layer"); + let diagnostic_layer_indices = + diagnostic_layers(&full_layer_indices, substage_probe_layer_index); + let layout = Qwen25VlVisionDiagnosticLayout { + window_layer_index, + full_layer_indices, + final_interval_layer_indices, + diagnostic_layer_indices, + substage_probe_layer_index, + }; + Ok(( + emit_qwen2_5_vl_inner(config, patch_bucket, precision, Some(&layout)), + layout, + )) +} + +#[cfg(all(test, feature = "diagnostics"))] +mod tests { + use super::{diagnostic_layers, substage_probe_layer}; + + #[test] + fn substage_probe_selects_the_first_observed_actual_divergence_boundary() { + assert_eq!(substage_probe_layer(&[7, 15, 23, 31]), Some(17)); + assert_eq!(substage_probe_layer(&[0, 1]), Some(0)); + assert_eq!(substage_probe_layer(&[0]), Some(0)); + assert_eq!(substage_probe_layer(&[]), None); + } + + #[test] + fn diagnostic_layers_cover_the_passing_and_first_failing_boundary() { + assert_eq!(diagnostic_layers(&[7, 15, 23, 31], 17), vec![16, 17]); + assert_eq!(diagnostic_layers(&[0, 1], 0), Vec::::new()); + assert_eq!(diagnostic_layers(&[0], 0), Vec::::new()); + } +} diff --git a/src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs b/src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs index 9065a41c8..a967a5cc1 100644 --- a/src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs +++ b/src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs @@ -24,7 +24,7 @@ use std::path::Path; use serde_json::Value; -use super::builder::{Builder, Ty, Val}; +use super::builder::{Builder, Precision, Ty, Val}; use super::numeric_ops::{exact_gelu, layer_norm_2d as layer_norm, stable_softmax, tanh_gelu}; /// Qualified flattened-patch capacities for the pinned Qwen2-VL image path. @@ -34,6 +34,15 @@ use super::numeric_ops::{exact_gelu, layer_norm_2d as layer_norm, stable_softmax /// compile or falling back to the MLX vision encoder. pub(crate) const QWEN2_VL_PATCH_BUCKETS: [usize; 3] = [16, 64, 256]; +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum QwenVlVisionVariant { + Qwen2, + Qwen25 { + window_size: usize, + full_attention_blocks: Vec, + }, +} + #[derive(Debug, Clone, PartialEq, Eq)] pub(crate) struct Qwen2VlWeightSpec { pub(crate) name: String, @@ -43,6 +52,7 @@ pub(crate) struct Qwen2VlWeightSpec { #[derive(Debug, Clone, PartialEq)] pub(crate) struct Qwen2VlConfig { + pub(crate) variant: QwenVlVisionVariant, pub(crate) depth: usize, pub(crate) hidden: usize, pub(crate) intermediate: usize, @@ -61,6 +71,9 @@ pub(crate) struct Qwen2VlGridPlan { pub(crate) patch_bucket: usize, pub(crate) actual_patches: usize, pub(crate) packed_cu_seqlens: Vec, + pub(crate) window_cu_seqlens: Option>, + pub(crate) window_index: Vec, + pub(crate) restore_indices: Vec, pub(crate) merged_tokens_per_image: Vec, } @@ -70,6 +83,7 @@ pub(crate) struct Qwen2VlHostInputs { pub(crate) patches: Vec, pub(crate) vision_rope_freqs: Vec, pub(crate) packed_attention_bias: Vec, + pub(crate) window_attention_bias: Option>, } struct Args { @@ -130,6 +144,13 @@ fn positive_usize(object: &serde_json::Map, field: &str) -> Resul } impl Qwen2VlConfig { + fn family_name(&self) -> &'static str { + match &self.variant { + QwenVlVisionVariant::Qwen2 => "Qwen2-VL", + QwenVlVisionVariant::Qwen25 { .. } => "Qwen2.5-VL", + } + } + pub(crate) fn from_model_dir(model_dir: &Path) -> Result { let path = model_dir.join("config.json"); let text = std::fs::read_to_string(&path) @@ -140,39 +161,56 @@ impl Qwen2VlConfig { pub(crate) fn from_json_str(text: &str) -> Result { let root: Value = serde_json::from_str(text).map_err(|error| format!("parse config.json: {error}"))?; - if root.get("model_type").and_then(Value::as_str) != Some("qwen2_vl") { - return Err("Qwen2-VL IREE vision requires model_type=qwen2_vl".to_string()); + let model_type = root + .get("model_type") + .and_then(Value::as_str) + .ok_or_else(|| "Qwen-VL IREE vision requires config.json model_type".to_string())?; + if !matches!(model_type, "qwen2_vl" | "qwen2_5_vl") { + return Err(format!( + "Qwen-VL IREE vision requires model_type=qwen2_vl or qwen2_5_vl, got {model_type}" + )); } let vision = root .get("vision_config") .and_then(Value::as_object) .ok_or_else(|| "config.json vision_config must be an object".to_string())?; let depth = positive_usize(vision, "depth")?; - let hidden = positive_usize(vision, "embed_dim")?; + let hidden = positive_usize( + vision, + if model_type == "qwen2_5_vl" { + "hidden_size" + } else { + "embed_dim" + }, + )?; let heads = positive_usize(vision, "num_heads")?; if hidden % heads != 0 { return Err(format!( "Qwen2-VL vision embed_dim={hidden} is not divisible by num_heads={heads}" )); } - let mlp_ratio = vision - .get("mlp_ratio") - .and_then(Value::as_f64) - .ok_or_else(|| { - "config.json vision_config.mlp_ratio must be a positive number".to_string() - })?; - if !mlp_ratio.is_finite() || mlp_ratio <= 0.0 { - return Err( - "config.json vision_config.mlp_ratio must be finite and positive".to_string(), - ); - } - let intermediate_f64 = hidden as f64 * mlp_ratio; - if intermediate_f64.fract() != 0.0 || intermediate_f64 > usize::MAX as f64 { - return Err(format!( - "Qwen2-VL vision intermediate size {hidden} * {mlp_ratio} is not an integer" - )); - } - let intermediate = intermediate_f64 as usize; + let intermediate = if model_type == "qwen2_5_vl" { + positive_usize(vision, "intermediate_size")? + } else { + let mlp_ratio = vision + .get("mlp_ratio") + .and_then(Value::as_f64) + .ok_or_else(|| { + "config.json vision_config.mlp_ratio must be a positive number".to_string() + })?; + if !mlp_ratio.is_finite() || mlp_ratio <= 0.0 { + return Err( + "config.json vision_config.mlp_ratio must be finite and positive".to_string(), + ); + } + let intermediate_f64 = hidden as f64 * mlp_ratio; + if intermediate_f64.fract() != 0.0 || intermediate_f64 > usize::MAX as f64 { + return Err(format!( + "Qwen2-VL vision intermediate size {hidden} * {mlp_ratio} is not an integer" + )); + } + intermediate_f64 as usize + }; let patch_size = positive_usize(vision, "patch_size")?; let temporal_patch_size = positive_usize(vision, "temporal_patch_size")?; let spatial_merge_size = positive_usize(vision, "spatial_merge_size")?; @@ -199,6 +237,19 @@ impl Qwen2VlConfig { if text_hidden == 0 { return Err("config.json hidden_size must be positive".to_string()); } + if model_type == "qwen2_5_vl" { + let out_hidden = positive_usize(vision, "out_hidden_size")?; + if out_hidden != text_hidden { + return Err(format!( + "Qwen2.5-VL vision out_hidden_size={out_hidden} disagrees with text hidden_size={text_hidden}" + )); + } + if vision.get("hidden_act").and_then(Value::as_str) != Some("silu") { + return Err( + "Qwen2.5-VL IREE vision requires vision_config.hidden_act=silu".to_string(), + ); + } + } if spatial_merge_size != 2 { return Err(format!( "Qwen2-VL IREE vision currently qualifies spatial_merge_size=2, got {spatial_merge_size}" @@ -209,7 +260,21 @@ impl Qwen2VlConfig { "Qwen2-VL IREE image vision currently qualifies temporal_patch_size=2, got {temporal_patch_size}" )); } - let quantization = match root.get("quantization") { + // mlx-community conversion can quantize the language tower while + // deliberately leaving the vision tower in its source precision. The + // Qwen2.5-VL converter records that split as + // `vision_config.skip_vision=true`; treating the root text contract as + // a vision contract rejects valid dimensions and would later expect + // affine metadata that the vision checkpoint does not contain. + let vision_quantization = + if vision.get("skip_vision").and_then(Value::as_bool) == Some(true) { + None + } else { + vision + .get("quantization") + .or_else(|| root.get("quantization")) + }; + let quantization = match vision_quantization { None => None, Some(value) => { let object = value @@ -232,7 +297,53 @@ impl Qwen2VlConfig { Some((bits, group_size)) } }; + let variant = if model_type == "qwen2_5_vl" { + let window_size = positive_usize(vision, "window_size")?; + let window_divisor = spatial_merge_size + .checked_mul(patch_size) + .ok_or_else(|| "Qwen2.5-VL window divisor overflowed".to_string())?; + if !window_size.is_multiple_of(window_divisor) { + return Err(format!( + "Qwen2.5-VL window_size={window_size} must be divisible by spatial_merge_size*patch_size={window_divisor}" + )); + } + let full_attention_blocks = vision + .get("fullatt_block_indexes") + .and_then(Value::as_array) + .ok_or_else(|| "Qwen2.5-VL IREE vision requires fullatt_block_indexes".to_string())? + .iter() + .map(|value| { + value + .as_u64() + .and_then(|value| usize::try_from(value).ok()) + .ok_or_else(|| { + "Qwen2.5-VL fullatt_block_indexes must contain non-negative integers" + .to_string() + }) + }) + .collect::, _>>()?; + let mut unique = std::collections::BTreeSet::new(); + for &index in &full_attention_blocks { + if index >= depth { + return Err(format!( + "Qwen2.5-VL full-attention block index {index} is outside depth {depth}" + )); + } + if !unique.insert(index) { + return Err(format!( + "Qwen2.5-VL full-attention block index {index} is duplicated" + )); + } + } + QwenVlVisionVariant::Qwen25 { + window_size, + full_attention_blocks, + } + } else { + QwenVlVisionVariant::Qwen2 + }; Ok(Self { + variant, depth, hidden, intermediate, @@ -290,7 +401,8 @@ impl Qwen2VlConfig { for (index, &(temporal, height, width)) in grids.iter().enumerate() { if temporal != 1 { return Err(format!( - "Qwen2-VL grid {index} has temporal size {temporal}; video/temporal grids are unsupported by the image-only XLA path" + "{} grid {index} has temporal size {temporal}; video/temporal grids are unsupported by the image-only XLA path", + self.family_name() )); } let height = usize::try_from(height) @@ -319,12 +431,88 @@ impl Qwen2VlConfig { packed_cu_seqlens.push(actual_patches); merged_tokens_per_image.push(self.merged_tokens(patches)); } - Ok(Qwen2VlGridPlan { + let mut plan = Qwen2VlGridPlan { patch_bucket: self.bucket_for_patches(actual_patches)?, actual_patches, packed_cu_seqlens, + window_cu_seqlens: None, + window_index: (0..self.merged_tokens(actual_patches)).collect(), + restore_indices: (0..self.merged_tokens(actual_patches)).collect(), merged_tokens_per_image, - }) + }; + if let QwenVlVisionVariant::Qwen25 { window_size, .. } = &self.variant { + let window_groups = *window_size / self.spatial_merge_size / self.patch_size; + let merge_unit = self.spatial_merge_size * self.spatial_merge_size; + let mut window_index = Vec::with_capacity(self.merged_tokens(actual_patches)); + let mut window_cu_seqlens = vec![0usize]; + let mut merged_offset = 0usize; + for &(temporal, height, width) in grids { + let temporal = temporal as usize; + let merged_height = height as usize / self.spatial_merge_size; + let merged_width = width as usize / self.spatial_merge_size; + let windows_height = merged_height.div_ceil(window_groups); + let windows_width = merged_width.div_ceil(window_groups); + for time in 0..temporal { + for window_h in 0..windows_height { + for window_w in 0..windows_width { + let mut valid_groups = 0usize; + for inner_h in 0..window_groups { + for inner_w in 0..window_groups { + let row = window_h * window_groups + inner_h; + let column = window_w * window_groups + inner_w; + if row < merged_height && column < merged_width { + let local = time * merged_height * merged_width + + row * merged_width + + column; + window_index.push(merged_offset + local); + valid_groups += 1; + } + } + } + if valid_groups > 0 { + let next = window_cu_seqlens + .last() + .copied() + .ok_or_else(|| { + "Qwen2.5-VL window boundary state is empty".to_string() + })? + .checked_add(valid_groups * merge_unit) + .ok_or_else(|| { + "Qwen2.5-VL window cumulative length overflowed".to_string() + })?; + window_cu_seqlens.push(next); + } + } + } + } + merged_offset += temporal * merged_height * merged_width; + } + if window_index.len() != self.merged_tokens(actual_patches) + || window_cu_seqlens.last().copied() != Some(actual_patches) + { + return Err(format!( + "Qwen2.5-VL window reordering produced {} merged indices and final boundary {:?}, expected {} indices and {actual_patches} patches", + window_index.len(), + window_cu_seqlens.last(), + self.merged_tokens(actual_patches) + )); + } + let mut seen = vec![false; window_index.len()]; + let mut restore_indices = vec![0usize; window_index.len()]; + for (window_position, &original_position) in window_index.iter().enumerate() { + if original_position >= seen.len() || seen[original_position] { + return Err(format!( + "Qwen2.5-VL window index {original_position} is out of range or duplicated" + )); + } + seen[original_position] = true; + restore_indices[original_position] = window_position; + } + plan.window_cu_seqlens = Some(window_cu_seqlens); + plan.window_index = window_index; + plan.restore_indices = restore_indices; + } + Ok(plan) } #[must_use] @@ -334,8 +522,15 @@ impl Qwen2VlConfig { #[must_use] pub(crate) fn fingerprint(&self, patch_bucket: usize) -> String { + let variant = match &self.variant { + QwenVlVisionVariant::Qwen2 => "family=qwen2".to_string(), + QwenVlVisionVariant::Qwen25 { + window_size, + full_attention_blocks, + } => format!("family=qwen2_5:window={window_size}:full={full_attention_blocks:?}"), + }; format!( - "qwen2-vl-vision-v1:bucket={patch_bucket}:patch={}:temporal={}:merge={}:channels={}:hidden={}:intermediate={}:heads={}:depth={}:text={}:dtype=f32:source_quant={:?}", + "qwen-vl-vision-v2:{variant}:bucket={patch_bucket}:patch={}:temporal={}:merge={}:channels={}:hidden={}:intermediate={}:heads={}:depth={}:text={}:dtype=f32:source_quant={:?}", self.patch_size, self.temporal_patch_size, self.spatial_merge_size, @@ -358,20 +553,37 @@ impl Qwen2VlConfig { )]; for layer in 0..self.depth { let prefix = format!("vision_tower.blocks.{layer}"); - for (suffix, shape) in [ - ("norm1.weight", vec![self.hidden]), - ("norm1.bias", vec![self.hidden]), - ("attn.qkv.weight", vec![self.hidden * 3, self.hidden]), - ("attn.qkv.bias", vec![self.hidden * 3]), - ("attn.proj.weight", vec![self.hidden, self.hidden]), - ("attn.proj.bias", vec![self.hidden]), - ("norm2.weight", vec![self.hidden]), - ("norm2.bias", vec![self.hidden]), - ("mlp.fc1.weight", vec![self.intermediate, self.hidden]), - ("mlp.fc1.bias", vec![self.intermediate]), - ("mlp.fc2.weight", vec![self.hidden, self.intermediate]), - ("mlp.fc2.bias", vec![self.hidden]), - ] { + let layer_specs = match &self.variant { + QwenVlVisionVariant::Qwen2 => vec![ + ("norm1.weight", vec![self.hidden]), + ("norm1.bias", vec![self.hidden]), + ("attn.qkv.weight", vec![self.hidden * 3, self.hidden]), + ("attn.qkv.bias", vec![self.hidden * 3]), + ("attn.proj.weight", vec![self.hidden, self.hidden]), + ("attn.proj.bias", vec![self.hidden]), + ("norm2.weight", vec![self.hidden]), + ("norm2.bias", vec![self.hidden]), + ("mlp.fc1.weight", vec![self.intermediate, self.hidden]), + ("mlp.fc1.bias", vec![self.intermediate]), + ("mlp.fc2.weight", vec![self.hidden, self.intermediate]), + ("mlp.fc2.bias", vec![self.hidden]), + ], + QwenVlVisionVariant::Qwen25 { .. } => vec![ + ("norm1.weight", vec![self.hidden]), + ("attn.qkv.weight", vec![self.hidden * 3, self.hidden]), + ("attn.qkv.bias", vec![self.hidden * 3]), + ("attn.proj.weight", vec![self.hidden, self.hidden]), + ("attn.proj.bias", vec![self.hidden]), + ("norm2.weight", vec![self.hidden]), + ("mlp.gate_proj.weight", vec![self.intermediate, self.hidden]), + ("mlp.gate_proj.bias", vec![self.intermediate]), + ("mlp.up_proj.weight", vec![self.intermediate, self.hidden]), + ("mlp.up_proj.bias", vec![self.intermediate]), + ("mlp.down_proj.weight", vec![self.hidden, self.intermediate]), + ("mlp.down_proj.bias", vec![self.hidden]), + ], + }; + for (suffix, shape) in layer_specs { specs.push(Qwen2VlWeightSpec { name: format!("{prefix}.{suffix}"), shape, @@ -379,9 +591,11 @@ impl Qwen2VlConfig { } } let merge_width = self.hidden * self.spatial_merge_size * self.spatial_merge_size; + specs.push(self.spec("vision_tower.merger.ln_q.weight", [self.hidden])); + if matches!(&self.variant, QwenVlVisionVariant::Qwen2) { + specs.push(self.spec("vision_tower.merger.ln_q.bias", [self.hidden])); + } specs.extend([ - self.spec("vision_tower.merger.ln_q.weight", [self.hidden]), - self.spec("vision_tower.merger.ln_q.bias", [self.hidden]), self.spec( "vision_tower.merger.mlp.0.weight", [merge_width, merge_width], @@ -444,18 +658,27 @@ pub(crate) fn prepare_qwen2_vl_host_inputs( .patch_bucket .checked_mul(patch_width) .ok_or_else(|| "Qwen2-VL padded patch tensor size overflowed".to_string())?; + let merge_unit = config.spatial_merge_size * config.spatial_merge_size; + let patch_order = plan + .window_index + .iter() + .flat_map(|&group| (0..merge_unit).map(move |inner| group * merge_unit + inner)) + .collect::>(); let mut patches = Vec::with_capacity(padded_values); - // The shared processor emits the temporal rows contiguously for each - // spatial patch. Flattening those adjacent rows is exactly the - // `[patch, temporal*channels*height*width]` patch-projection contract. - patches.extend_from_slice(temporal_patch_rows); + // The shared processor emits temporal rows contiguously for each patch. + // Qwen2 uses identity order; Qwen2.5 expands its merged-group window + // permutation back into patch rows before native execution. + for &patch in &patch_order { + let start = patch * patch_width; + patches.extend_from_slice(&temporal_patch_rows[start..start + patch_width]); + } patches.resize(padded_values, 0.0); let head_dim = config.hidden / config.heads; let rotary_dim = head_dim / 2; let frequency_width = rotary_dim; let axis_width = rotary_dim / 2; - if head_dim % 4 != 0 { + if !head_dim.is_multiple_of(4) { return Err(format!( "Qwen2-VL vision head_dim={head_dim} must be divisible by 4 for 2D RoPE" )); @@ -482,30 +705,48 @@ pub(crate) fn prepare_qwen2_vl_host_inputs( } } } + if matches!(&config.variant, QwenVlVisionVariant::Qwen25 { .. }) { + let original = vision_rope_freqs; + vision_rope_freqs = Vec::with_capacity(plan.patch_bucket * frequency_width); + for &patch in &patch_order { + let start = patch * frequency_width; + vision_rope_freqs.extend_from_slice(&original[start..start + frequency_width]); + } + } vision_rope_freqs.resize(plan.patch_bucket * frequency_width, 0.0); - let bias_values = plan - .patch_bucket - .checked_mul(plan.patch_bucket) - .ok_or_else(|| "Qwen2-VL packed attention bias size overflowed".to_string())?; - let mut packed_attention_bias = vec![f32::NEG_INFINITY; bias_values]; - for segment in plan.packed_cu_seqlens.windows(2) { - for row in segment[0]..segment[1] { - for column in segment[0]..segment[1] { - packed_attention_bias[row * plan.patch_bucket + column] = 0.0; + let attention_bias = |boundaries: &[usize]| -> Result, String> { + let bias_values = plan + .patch_bucket + .checked_mul(plan.patch_bucket) + .ok_or_else(|| "Qwen-VL packed attention bias size overflowed".to_string())?; + let mut bias = vec![f32::NEG_INFINITY; bias_values]; + for segment in boundaries.windows(2) { + for row in segment[0]..segment[1] { + for column in segment[0]..segment[1] { + bias[row * plan.patch_bucket + column] = 0.0; + } } } - } - // Give every padded query one finite self key so its softmax cannot create - // NaNs. Padded rows remain unreachable from every real media segment. - for index in plan.actual_patches..plan.patch_bucket { - packed_attention_bias[index * plan.patch_bucket + index] = 0.0; - } + // Give every padded query one finite self key so its softmax cannot + // create NaNs. Padding stays unreachable from real media segments. + for index in plan.actual_patches..plan.patch_bucket { + bias[index * plan.patch_bucket + index] = 0.0; + } + Ok(bias) + }; + let packed_attention_bias = attention_bias(&plan.packed_cu_seqlens)?; + let window_attention_bias = plan + .window_cu_seqlens + .as_deref() + .map(attention_bias) + .transpose()?; Ok(Qwen2VlHostInputs { plan, patches, vision_rope_freqs, packed_attention_bias, + window_attention_bias, }) } @@ -645,13 +886,24 @@ fn encoder_layer( /// `attention_bias` are host-built from `grid_thw`/packed boundaries, keeping /// cross-image isolation explicit while preserving one finite static graph. pub(crate) fn emit_qwen2_vl(config: &Qwen2VlConfig, patch_bucket: usize) -> String { + emit_qwen2_vl_with(config, patch_bucket, Precision::F32) +} + +pub(crate) fn emit_qwen2_vl_with( + config: &Qwen2VlConfig, + patch_bucket: usize, + precision: Precision, +) -> String { + if matches!(&config.variant, QwenVlVisionVariant::Qwen25 { .. }) { + return super::qwen2_5_vl::emit_qwen2_5_vl(config, patch_bucket, precision); + } assert!( QWEN2_VL_PATCH_BUCKETS.contains(&patch_bucket), "unqualified Qwen2-VL patch bucket" ); let specs = config.weight_specs(); let mut args = Args::new(&specs); - let mut builder = Builder::new(); + let mut builder = Builder::new().with_precision(precision); let patch_weight = args.take(); let patch_width = config.channels * config.temporal_patch_size * config.patch_size * config.patch_size; @@ -728,6 +980,30 @@ mod tests { .unwrap() } + fn qwen25_config() -> Qwen2VlConfig { + Qwen2VlConfig::from_json_str( + r#"{ + "model_type": "qwen2_5_vl", + "hidden_size": 1536, + "vision_config": { + "depth": 2, + "hidden_act": "silu", + "hidden_size": 1280, + "intermediate_size": 3420, + "out_hidden_size": 1536, + "num_heads": 16, + "in_chans": 3, + "patch_size": 14, + "spatial_merge_size": 2, + "temporal_patch_size": 2, + "window_size": 112, + "fullatt_block_indexes": [1] + } + }"#, + ) + .unwrap() + } + #[test] fn static_buckets_bound_qwen2_vl_grids_without_per_shape_compilation() { let config = config(); @@ -803,6 +1079,62 @@ mod tests { assert_eq!(mlir.matches("chlo.erf").count(), 1); } + #[test] + fn qwen25_gpu_precision_demotes_contractions_but_keeps_sensitive_math_f32() { + let config = qwen25_config(); + let f32 = emit_qwen2_vl_with(&config, 16, Precision::F32); + let f16 = emit_qwen2_vl_with(&config, 16, Precision::F16); + assert!(!f32.contains("xf16")); + assert!(f16.contains("xf16"), "f16 graph emitted no demotion"); + for line in f16.lines().filter(|line| { + line.contains("stablehlo.dot_general") + || line.contains("stablehlo.exponential") + || line.contains("stablehlo.rsqrt") + }) { + if line.contains("stablehlo.dot_general") { + assert!(line.contains("f16>"), "matmul was not demoted: {line}"); + assert!( + line.ends_with("f32>"), + "matmul must accumulate to f32: {line}" + ); + } else { + assert!( + !line.contains("f16"), + "norm/softmax sensitive math was demoted: {line}" + ); + } + } + } + + #[cfg(feature = "diagnostics")] + #[test] + fn qwen25_diagnostics_share_the_production_graph_and_expose_ordered_seams() { + let config = qwen25_config(); + let production = emit_qwen2_vl(&config, 16); + let (diagnostics, layout) = + crate::emitter::emit_qwen2_5_vl_diagnostics(&config, 16, Precision::F32).unwrap(); + assert_eq!(layout.window_layer_index, 0); + assert_eq!(layout.full_layer_indices, vec![1]); + assert_eq!(layout.final_interval_layer_indices, vec![0, 1]); + assert_eq!(layout.substage_probe_layer_index, 1); + assert_eq!( + production.matches("stablehlo.dot_general").count(), + diagnostics.matches("stablehlo.dot_general").count() + ); + let result_signature = std::iter::repeat_n("tensor<16x1280xf32>", 7) + .chain(std::iter::repeat_n("tensor<16x16x80xf32>", 3)) + .chain(std::iter::repeat_n("tensor<16x1280xf32>", 4)) + .chain(std::iter::repeat_n("tensor<16x3420xf32>", 4)) + .chain(std::iter::once("tensor<16x1280xf32>")) + .chain(std::iter::once("tensor<4x1536xf32>")) + .collect::>() + .join(", "); + assert!(diagnostics.contains(&format!("-> ({result_signature})"))); + assert!(diagnostics.contains("return %")); + assert!(diagnostics.contains("loc(\"window_attention.bias\")")); + assert!(diagnostics.contains("loc(\"full_attention.bias\")")); + } + #[test] fn host_inputs_preserve_packed_media_isolation_and_finite_padding_rows() { let config = config(); @@ -824,7 +1156,7 @@ mod tests { f32::NEG_INFINITY ); assert_eq!(inputs.packed_attention_bias[48 * 64 + 48], 0.0); - assert_eq!(inputs.packed_attention_bias[48 * 64 + 0], f32::NEG_INFINITY); + assert_eq!(inputs.packed_attention_bias[48 * 64], f32::NEG_INFINITY); let mut non_finite = values; non_finite[7] = f32::NAN; assert!( @@ -833,4 +1165,193 @@ mod tests { .contains("flat index 7") ); } + + #[test] + fn qwen25_window_plan_matches_padded_reference_partition_and_inverse() { + let config = qwen25_config(); + assert!( + config + .plan_image_grids(&[(2, 4, 4)]) + .unwrap_err() + .contains("Qwen2.5-VL") + ); + let plan = config.plan_image_grids(&[(1, 12, 12)]).unwrap(); + assert_eq!(plan.patch_bucket, 256); + assert_eq!(plan.window_cu_seqlens, Some(vec![0, 64, 96, 128, 144])); + assert_eq!( + &plan.window_index[..16], + &[0, 1, 2, 3, 6, 7, 8, 9, 12, 13, 14, 15, 18, 19, 20, 21] + ); + for (original, &window_position) in plan.restore_indices.iter().enumerate() { + assert_eq!(plan.window_index[window_position], original); + } + let multi = config.plan_image_grids(&[(1, 8, 8), (1, 4, 4)]).unwrap(); + assert_eq!(multi.packed_cu_seqlens, vec![0, 64, 80]); + assert_eq!(multi.window_cu_seqlens, Some(vec![0, 64, 80])); + assert_eq!(multi.merged_tokens_per_image, vec![16, 4]); + } + + #[test] + fn qwen25_host_inputs_reorder_groups_and_isolate_windows() { + let config = qwen25_config(); + let grids = [(1, 12, 12)]; + let patch_width = 3 * 2 * 14 * 14; + let mut values = vec![0.0; 144 * patch_width]; + for patch in 0..144 { + values[patch * patch_width] = patch as f32; + } + let inputs = prepare_qwen2_vl_host_inputs(&config, &grids, &values).unwrap(); + let observed = (0..20) + .map(|patch| inputs.patches[patch * patch_width] as usize) + .collect::>(); + assert_eq!( + observed, + vec![ + 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 24, 25, 26, 27 + ] + ); + let full = &inputs.packed_attention_bias; + let window = inputs.window_attention_bias.as_ref().unwrap(); + assert_eq!(full[80], 0.0); + assert_eq!(window[80], f32::NEG_INFINITY); + assert_eq!(window[144 * 256 + 144], 0.0); + assert_eq!(window[144 * 256], f32::NEG_INFINITY); + } + + #[test] + fn qwen25_two_media_bucket_bias_isolates_media_and_padding() { + let config = qwen25_config(); + let grids = [(1, 10, 10), (1, 10, 10)]; + let patch_width = 3 * 2 * 14 * 14; + let inputs = + prepare_qwen2_vl_host_inputs(&config, &grids, &vec![0.0; 200 * patch_width]).unwrap(); + + assert_eq!(inputs.plan.actual_patches, 200); + assert_eq!(inputs.plan.patch_bucket, 256); + assert_eq!(inputs.plan.packed_cu_seqlens, vec![0, 100, 200]); + let full = &inputs.packed_attention_bias; + let window = inputs.window_attention_bias.as_ref().unwrap(); + + assert_eq!(full[99 * 256 + 100], f32::NEG_INFINITY); + assert_eq!(full[100 * 256 + 99], f32::NEG_INFINITY); + assert_eq!(full[199], f32::NEG_INFINITY); + assert_eq!(full[199 * 256 + 200], f32::NEG_INFINITY); + assert_eq!(full[200 * 256 + 199], f32::NEG_INFINITY); + assert_eq!(full[200 * 256 + 200], 0.0); + assert_eq!(full[200 * 256 + 201], f32::NEG_INFINITY); + assert_ne!(full, window); + } + + #[test] + fn qwen25_schema_and_emitter_are_family_specific_and_fingerprinted() { + let config = qwen25_config(); + let specs = config.weight_specs(); + assert_eq!(specs.len(), 1 + 2 * 12 + 5); + assert!( + specs + .iter() + .any(|spec| spec.name.ends_with("mlp.gate_proj.weight")) + ); + assert!(!specs.iter().any(|spec| spec.name.ends_with("norm1.bias"))); + let fingerprint = config.fingerprint(256); + assert!(fingerprint.contains("family=qwen2_5")); + assert!(fingerprint.contains("window=112")); + assert!(fingerprint.contains("full=[1]")); + let mlir = emit_qwen2_vl(&config, 16); + assert!(mlir.contains("loc(\"patches.windowed\")")); + assert!(mlir.contains("loc(\"full_attention.bias\")")); + assert!(mlir.contains("loc(\"window_attention.bias\")")); + assert!(mlir.contains("mlp.gate_proj.weight")); + assert!(!mlir.contains("norm1.bias")); + } + + #[test] + fn qwen25_rejects_unknown_or_ambiguous_window_configuration() { + for (fragment, expected) in [ + ( + r#""window_size":112,"fullatt_block_indexes":[2]"#, + "outside depth", + ), + ( + r#""window_size":112,"fullatt_block_indexes":[1,1]"#, + "block index 1 is duplicated", + ), + ( + r#""window_size":101,"fullatt_block_indexes":[1]"#, + "must be divisible", + ), + ] { + let json = format!( + r#"{{ + "model_type":"qwen2_5_vl","hidden_size":12, + "vision_config":{{ + "depth":2,"hidden_act":"silu","hidden_size":8, + "intermediate_size":16,"out_hidden_size":12,"num_heads":2, + "patch_size":2,"spatial_merge_size":2,"temporal_patch_size":2, + {fragment} + }} + }}"# + ); + assert!( + Qwen2VlConfig::from_json_str(&json) + .unwrap_err() + .contains(expected) + ); + } + } + + #[test] + fn qwen25_skip_vision_keeps_text_quantization_out_of_vision_contract() { + let config = Qwen2VlConfig::from_json_str( + r#"{ + "model_type":"qwen2_5_vl", + "hidden_size":2048, + "quantization":{"group_size":64,"bits":4}, + "vision_config":{ + "depth":32, + "hidden_act":"silu", + "hidden_size":1280, + "intermediate_size":3420, + "out_hidden_size":2048, + "num_heads":16, + "patch_size":14, + "spatial_merge_size":2, + "temporal_patch_size":2, + "window_size":112, + "fullatt_block_indexes":[7,15,23,31], + "skip_vision":true + } + }"#, + ) + .unwrap(); + assert_eq!(config.quantization, None); + assert!(config.fingerprint(256).contains("source_quant=None")); + } + + #[test] + fn qwen2_still_inherits_the_root_quantization_contract() { + let config = Qwen2VlConfig::from_json_str( + r#"{ + "model_type":"qwen2_vl", + "hidden_size":1536, + "quantization":{"group_size":64,"bits":4}, + "vision_config":{ + "depth":2, + "embed_dim":1280, + "mlp_ratio":4, + "num_heads":16, + "patch_size":14, + "spatial_merge_size":2, + "temporal_patch_size":2 + } + }"#, + ) + .unwrap(); + assert_eq!(config.quantization, Some((4, 64))); + assert!( + config + .fingerprint(256) + .contains("source_quant=Some((4, 64))") + ); + } } diff --git a/src/lib/mlxcel-xla/src/lib.rs b/src/lib/mlxcel-xla/src/lib.rs index 2b1328267..79f76f9a4 100644 --- a/src/lib/mlxcel-xla/src/lib.rs +++ b/src/lib/mlxcel-xla/src/lib.rs @@ -145,8 +145,9 @@ pub use batch::{EngineEvent, FinishReason, XlaAdmissionError, XlaBatchEngine, Xl #[cfg(feature = "diagnostics")] pub use batch::{ Gemma3nAllLayerDiagnosticRun, Gemma3nCanonicalDiagnosticRun, Gemma3nPrefixDecodeDiagnosticRun, - LlavaReferenceDiagnosticEngine, LlavaReferenceDiagnosticRun, run_gemma3n_all_layer_diagnostics, - run_gemma3n_canonical_diagnostics, run_gemma3n_prefix_decode_diagnostic, + LlavaReferenceDiagnosticEngine, LlavaReferenceDiagnosticRun, Qwen25VlLanguageDiagnosticEngine, + run_gemma3n_all_layer_diagnostics, run_gemma3n_canonical_diagnostics, + run_gemma3n_prefix_decode_diagnostic, }; pub use context::{ CONTEXT_CAPACITY_ENV, ContextCapacityError, DEFAULT_CONTEXT_CAPACITY, @@ -181,6 +182,10 @@ pub use qwen2_vl_runtime::{ IreeQwen2VlProjector, Qwen2VlVisionExecutionMetrics, Qwen2VlVisionProjection, }; #[cfg(feature = "diagnostics")] +pub use qwen2_vl_runtime::{ + IreeQwen25VlDiagnosticProjector, Qwen25VlDiagnosticMutation, Qwen25VlVisionDiagnostics, +}; +#[cfg(feature = "diagnostics")] pub use vision_runtime::{IreeVisionDiagnosticProjector, VisionDiagnosticProjection}; #[cfg(feature = "iree")] pub use vision_runtime::{IreeVisionProjector, VisionExecutionMetrics, VisionProjection}; diff --git a/src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs b/src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs index 2634f24b3..bb993dca8 100644 --- a/src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs +++ b/src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs @@ -35,8 +35,11 @@ use crate::aux::{ }; use crate::aux_manifest::{AuxiliaryArtifactContract, ensure_qualified_auxiliary_artifact}; use crate::emitter::{ - Qwen2VlConfig, Qwen2VlHostInputs, emit_qwen2_vl, prepare_qwen2_vl_host_inputs, + Qwen2VlConfig, Qwen2VlHostInputs, QwenVlVisionVariant, emit_qwen2_vl_with, + prepare_qwen2_vl_host_inputs, resolve_precision_checked, }; +#[cfg(feature = "diagnostics")] +use crate::emitter::{Qwen25VlVisionDiagnosticLayout, emit_qwen2_5_vl_diagnostics}; use crate::iree::{cached_vmfb_path, compile_one_to, iree_compile_bin, target_flags}; use crate::weights::{bf16_to_f32, dequantize_affine, f16_to_f32, f32_le_to_f32}; @@ -59,11 +62,159 @@ pub struct Qwen2VlVisionProjection { pub metrics: Qwen2VlVisionExecutionMetrics, } +#[cfg(feature = "diagnostics")] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Qwen25VlDiagnosticMutation { + None, + Qwen2FullAttention, + IdentityPermutation, + ZeroVisionPositions, +} + +#[cfg(feature = "diagnostics")] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct Qwen25VlAttentionMaskAudit { + pub full_window_differences: usize, + pub cross_media_full_bias_leaks: usize, + pub active_padding_full_bias_leaks: usize, + pub invalid_padded_full_bias_entries: usize, +} + +#[cfg(feature = "diagnostics")] +#[derive(Debug, Clone, PartialEq)] +pub struct Qwen25VlVisionDiagnostics { + pub window_index: Vec, + pub restore_indices: Vec, + pub packed_cu_seqlens: Vec, + pub window_cu_seqlens: Vec, + pub reordered_patch_embedding: Vec, + pub window_layer_index: usize, + pub post_window_layer: Vec, + pub full_layer_indices: Vec, + pub post_full_layers: Vec>, + pub final_interval_layer_indices: Vec, + pub post_final_interval_layers: Vec>, + pub diagnostic_layer_indices: Vec, + pub post_diagnostic_layers: Vec>, + pub substage_probe_layer_index: usize, + pub substage_probe_layer_input: Vec, + pub substage_probe_layer_norm1: Vec, + pub substage_probe_layer_query: Vec, + pub substage_probe_layer_key: Vec, + pub substage_probe_layer_value: Vec, + pub substage_probe_layer_attention_context: Vec, + pub substage_probe_layer_attention: Vec, + pub substage_probe_layer_post_attention_residual: Vec, + pub substage_probe_layer_norm2: Vec, + pub substage_probe_layer_mlp_gate_projection: Vec, + pub substage_probe_layer_mlp_gate_activation: Vec, + pub substage_probe_layer_mlp_up_projection: Vec, + pub substage_probe_layer_mlp_gated_product: Vec, + pub substage_probe_layer_mlp_down_projection: Vec, + pub merger_window_ordered: Vec, + pub restored_projection: Vec, + pub patch_bucket: usize, + pub patch_tokens: usize, + pub merged_tokens: usize, + pub vision_hidden: usize, + pub text_hidden: usize, + pub attention_mask_audit: Qwen25VlAttentionMaskAudit, +} + struct ActiveModule { patch_bucket: usize, module: IreeAuxiliaryModule, } +#[cfg(feature = "diagnostics")] +struct ActiveDiagnosticModule { + patch_bucket: usize, + module: IreeAuxiliaryModule, + layout: Qwen25VlVisionDiagnosticLayout, +} + +#[cfg(feature = "diagnostics")] +#[derive(Debug, Clone, PartialEq, Eq)] +struct Qwen25VlDiagnosticOutputSpec { + label: String, + shape: Vec, +} + +#[cfg(feature = "diagnostics")] +fn qwen25_diagnostic_output_specs( + config: &Qwen2VlConfig, + layout: &Qwen25VlVisionDiagnosticLayout, + patch_bucket: usize, +) -> Vec { + let state_shape = || vec![patch_bucket, config.hidden]; + let state = |label: String| Qwen25VlDiagnosticOutputSpec { + label, + shape: state_shape(), + }; + let intermediate = |label: String| Qwen25VlDiagnosticOutputSpec { + label, + shape: vec![patch_bucket, config.intermediate], + }; + let mut specs = vec![ + state("reordered_patch_embedding".to_string()), + state(format!("post_window_layer.{}", layout.window_layer_index)), + ]; + specs.extend( + layout + .full_layer_indices + .iter() + .map(|layer| state(format!("post_full_layer.{layer}"))), + ); + specs.extend( + layout + .final_interval_layer_indices + .iter() + .map(|layer| state(format!("post_final_interval_layer.{layer}"))), + ); + specs.extend( + layout + .diagnostic_layer_indices + .iter() + .map(|layer| state(format!("post_diagnostic_layer.{layer}"))), + ); + let probe = layout.substage_probe_layer_index; + specs.extend([ + state(format!("substage_probe_layer.{probe}.input")), + state(format!("substage_probe_layer.{probe}.norm1")), + Qwen25VlDiagnosticOutputSpec { + label: format!("substage_probe_layer.{probe}.query"), + shape: vec![patch_bucket, config.heads, config.hidden / config.heads], + }, + Qwen25VlDiagnosticOutputSpec { + label: format!("substage_probe_layer.{probe}.key"), + shape: vec![patch_bucket, config.heads, config.hidden / config.heads], + }, + Qwen25VlDiagnosticOutputSpec { + label: format!("substage_probe_layer.{probe}.value"), + shape: vec![patch_bucket, config.heads, config.hidden / config.heads], + }, + state(format!("substage_probe_layer.{probe}.attention_context")), + state(format!("substage_probe_layer.{probe}.attention")), + state(format!( + "substage_probe_layer.{probe}.post_attention_residual" + )), + state(format!("substage_probe_layer.{probe}.norm2")), + intermediate(format!("substage_probe_layer.{probe}.mlp_gate_projection")), + intermediate(format!("substage_probe_layer.{probe}.mlp_gate_activation")), + intermediate(format!("substage_probe_layer.{probe}.mlp_up_projection")), + intermediate(format!("substage_probe_layer.{probe}.mlp_gated_product")), + state(format!("substage_probe_layer.{probe}.mlp_down_projection")), + ]); + specs.push(Qwen25VlDiagnosticOutputSpec { + label: "merger_window_ordered".to_string(), + shape: vec![ + patch_bucket / (config.spatial_merge_size * config.spatial_merge_size), + config.text_hidden, + ], + }); + specs +} + pub struct IreeQwen2VlProjector { model_dir: PathBuf, device: String, @@ -72,6 +223,15 @@ pub struct IreeQwen2VlProjector { active: Option, } +#[cfg(feature = "diagnostics")] +pub struct IreeQwen25VlDiagnosticProjector { + model_dir: PathBuf, + device: String, + config: Qwen2VlConfig, + processor_identity: String, + active: Option, +} + fn hex_bytes(bytes: &[u8]) -> String { let mut output = String::with_capacity(bytes.len() * 2); for byte in bytes { @@ -135,10 +295,15 @@ fn processor_identity(model_dir: &Path, config: &Qwen2VlConfig) -> Result {} + _ => { + return Err(format!( + "preprocessor_config.json {field}=true is required by Qwen-VL IREE vision" + )); + } } } let positive = |field: &str| -> Result { @@ -170,8 +335,18 @@ fn processor_identity(model_dir: &Path, config: &Qwen2VlConfig) -> Result) -> Result, String> { Ok(values) } +#[cfg(feature = "diagnostics")] +fn audit_attention_masks( + patch_bucket: usize, + actual_patches: usize, + packed_cu_seqlens: &[usize], + full_bias: &[f32], + window_bias: &[f32], +) -> Qwen25VlAttentionMaskAudit { + let media_for = |index: usize| { + packed_cu_seqlens + .windows(2) + .position(|range| range[0] <= index && index < range[1]) + }; + let mut audit = Qwen25VlAttentionMaskAudit { + full_window_differences: 0, + cross_media_full_bias_leaks: 0, + active_padding_full_bias_leaks: 0, + invalid_padded_full_bias_entries: 0, + }; + for row in 0..patch_bucket { + for column in 0..patch_bucket { + let index = row * patch_bucket + column; + if full_bias[index] != window_bias[index] { + audit.full_window_differences += 1; + } + let finite = full_bias[index].is_finite(); + match (row < actual_patches, column < actual_patches) { + (true, true) => { + if media_for(row) != media_for(column) && finite { + audit.cross_media_full_bias_leaks += 1; + } + } + (true, false) | (false, true) => { + if finite { + audit.active_padding_full_bias_leaks += 1; + } + } + (false, false) => { + if finite != (row == column) { + audit.invalid_padded_full_bias_entries += 1; + } + } + } + } + } + audit +} + fn compile_and_load( model_dir: &Path, device: &str, @@ -547,7 +770,8 @@ fn compile_and_load( processor_identity: &str, patch_bucket: usize, ) -> Result { - let mlir = emit_qwen2_vl(config, patch_bucket); + let precision = resolve_precision_checked(device)?; + let mlir = emit_qwen2_vl_with(config, patch_bucket, precision); let compiler = iree_compile_bin()?; if !compiler.is_file() { return Err(format!("iree-compile not found at {}", compiler.display())); @@ -560,14 +784,18 @@ fn compile_and_load( let contract = AuxiliaryArtifactContract::new_legacy_unqualified( ENTRY_NAME, format!( - "{};{};checkpoint_schema_sha256={}", + "{};{};precision={precision:?};checkpoint_schema_sha256={}", config.fingerprint(patch_bucket), processor_identity, sha256_hex(checkpoint_schema.as_bytes()) ), compiler_generation_identity(&compiler, flags, &mlir)?, )?; - let tag = format!("qwen2-vl-vision-{patch_bucket}"); + let family = match &config.variant { + QwenVlVisionVariant::Qwen2 => "qwen2", + QwenVlVisionVariant::Qwen25 { .. } => "qwen2-5", + }; + let tag = format!("{family}-vl-vision-{patch_bucket}"); let vmfb = cached_vmfb_path(&compiler, &mlir, flags, &cache, &tag, 0); ensure_qualified_auxiliary_artifact(&vmfb, &contract, &weights, |temporary| { compile_one_to(&compiler, &mlir, flags, &cache, &tag, 0, temporary) @@ -575,6 +803,46 @@ fn compile_and_load( IreeAuxiliaryModule::load(device, &vmfb, &contract, weights) } +#[cfg(feature = "diagnostics")] +fn compile_and_load_diagnostics( + model_dir: &Path, + device: &str, + config: &Qwen2VlConfig, + processor_identity: &str, + patch_bucket: usize, +) -> Result<(IreeAuxiliaryModule, Qwen25VlVisionDiagnosticLayout), String> { + let precision = resolve_precision_checked(device)?; + let (mlir, layout) = emit_qwen2_5_vl_diagnostics(config, patch_bucket, precision)?; + let compiler = iree_compile_bin()?; + if !compiler.is_file() { + return Err(format!("iree-compile not found at {}", compiler.display())); + } + let flags = target_flags(device)?; + let cache = std::env::temp_dir().join("mlxcel-xla-qwen2-vl-vmfb"); + std::fs::create_dir_all(&cache) + .map_err(|error| format!("mkdir {}: {error}", cache.display()))?; + let (weights, checkpoint_schema) = load_weights(model_dir, config)?; + let contract = AuxiliaryArtifactContract::new_legacy_unqualified( + ENTRY_NAME, + format!( + "{};{};precision={precision:?};diagnostics=window-full-restoration;checkpoint_schema_sha256={}", + config.fingerprint(patch_bucket), + processor_identity, + sha256_hex(checkpoint_schema.as_bytes()) + ), + compiler_generation_identity(&compiler, flags, &mlir)?, + )?; + let tag = format!("qwen2-5-vl-vision-diagnostics-{patch_bucket}"); + let vmfb = cached_vmfb_path(&compiler, &mlir, flags, &cache, &tag, 0); + ensure_qualified_auxiliary_artifact(&vmfb, &contract, &weights, |temporary| { + compile_one_to(&compiler, &mlir, flags, &cache, &tag, 0, temporary) + })?; + Ok(( + IreeAuxiliaryModule::load(device, &vmfb, &contract, weights)?, + layout, + )) +} + impl IreeQwen2VlProjector { pub fn load(model_dir: &Path, device: &str) -> Result { let config = Qwen2VlConfig::from_model_dir(model_dir)?; @@ -629,6 +897,7 @@ impl IreeQwen2VlProjector { patches, vision_rope_freqs, packed_attention_bias, + window_attention_bias, } = prepare_qwen2_vl_host_inputs(&self.config, grids, temporal_patch_rows)?; let patch_width = self.config.channels * self.config.temporal_patch_size @@ -645,24 +914,32 @@ impl IreeQwen2VlProjector { let mut output = vec![0u8; full_output_shape.iter().product::() * std::mem::size_of::()]; let started = Instant::now(); + let mut inputs = vec![ + AuxiliaryInput { + bytes: f32_as_bytes(&patches), + dtype: AuxiliaryTensorDType::Float32, + shape: &patch_shape, + }, + AuxiliaryInput { + bytes: f32_as_bytes(&vision_rope_freqs), + dtype: AuxiliaryTensorDType::Float32, + shape: &frequency_shape, + }, + AuxiliaryInput { + bytes: f32_as_bytes(&packed_attention_bias), + dtype: AuxiliaryTensorDType::Float32, + shape: &bias_shape, + }, + ]; + if let Some(window_attention_bias) = &window_attention_bias { + inputs.push(AuxiliaryInput { + bytes: f32_as_bytes(window_attention_bias), + dtype: AuxiliaryTensorDType::Float32, + shape: &bias_shape, + }); + } self.ensure_bucket(plan.patch_bucket)?.invoke( - &[ - AuxiliaryInput { - bytes: f32_as_bytes(&patches), - dtype: AuxiliaryTensorDType::Float32, - shape: &patch_shape, - }, - AuxiliaryInput { - bytes: f32_as_bytes(&vision_rope_freqs), - dtype: AuxiliaryTensorDType::Float32, - shape: &frequency_shape, - }, - AuxiliaryInput { - bytes: f32_as_bytes(&packed_attention_bias), - dtype: AuxiliaryTensorDType::Float32, - shape: &bias_shape, - }, - ], + &inputs, &mut [AuxiliaryOutput { bytes: &mut output, dtype: AuxiliaryTensorDType::Float32, @@ -675,10 +952,21 @@ impl IreeQwen2VlProjector { let actual_values = actual_tokens .checked_mul(self.config.text_hidden) .ok_or_else(|| "Qwen2-VL projected output size overflowed".to_string())?; - let values = all_values + let window_ordered = all_values .get(..actual_values) .ok_or_else(|| "Qwen2-VL IREE output is shorter than the actual grid".to_string())? .to_vec(); + let mut values = Vec::with_capacity(actual_values); + for &window_position in &plan.restore_indices { + let start = window_position + .checked_mul(self.config.text_hidden) + .ok_or_else(|| "Qwen-VL restoration offset overflowed".to_string())?; + values.extend_from_slice( + window_ordered + .get(start..start + self.config.text_hidden) + .ok_or_else(|| "Qwen-VL restoration index is out of range".to_string())?, + ); + } Ok(Qwen2VlVisionProjection { values, shape: [actual_tokens, self.config.text_hidden], @@ -687,7 +975,10 @@ impl IreeQwen2VlProjector { metrics: Qwen2VlVisionExecutionMetrics { patch_upload_bytes: std::mem::size_of_val(patches.as_slice()), metadata_upload_bytes: std::mem::size_of_val(vision_rope_freqs.as_slice()) - + std::mem::size_of_val(packed_attention_bias.as_slice()), + + std::mem::size_of_val(packed_attention_bias.as_slice()) + + window_attention_bias + .as_ref() + .map_or(0, |bias| std::mem::size_of_val(bias.as_slice())), projected_transfer_bytes: std::mem::size_of_val(all_values.as_slice()), elapsed_seconds, }, @@ -695,10 +986,460 @@ impl IreeQwen2VlProjector { } } +#[cfg(feature = "diagnostics")] +fn restore_grouped_rows( + values: &[f32], + row_width: usize, + merge_unit: usize, + patch_bucket: usize, + actual_patches: usize, + restore_indices: &[usize], +) -> Result, String> { + let expected = patch_bucket + .checked_mul(row_width) + .ok_or_else(|| "Qwen2.5-VL diagnostic row count overflowed".to_string())?; + if values.len() != expected || actual_patches != restore_indices.len() * merge_unit { + return Err("Qwen2.5-VL diagnostic permutation dimensions are inconsistent".to_string()); + } + let mut restored = vec![0.0; expected]; + for (original_group, &window_group) in restore_indices.iter().enumerate() { + let source = window_group + .checked_mul(merge_unit) + .and_then(|row| row.checked_mul(row_width)) + .ok_or_else(|| "Qwen2.5-VL diagnostic source offset overflowed".to_string())?; + let destination = original_group + .checked_mul(merge_unit) + .and_then(|row| row.checked_mul(row_width)) + .ok_or_else(|| "Qwen2.5-VL diagnostic destination offset overflowed".to_string())?; + let count = merge_unit + .checked_mul(row_width) + .ok_or_else(|| "Qwen2.5-VL diagnostic group width overflowed".to_string())?; + let source_rows = values + .get(source..source + count) + .ok_or_else(|| "Qwen2.5-VL diagnostic source group is out of range".to_string())?; + restored[destination..destination + count].copy_from_slice(source_rows); + } + Ok(restored) +} + +#[cfg(feature = "diagnostics")] +impl IreeQwen25VlDiagnosticProjector { + pub fn load(model_dir: &Path, device: &str) -> Result { + let config = Qwen2VlConfig::from_model_dir(model_dir)?; + if !matches!(&config.variant, QwenVlVisionVariant::Qwen25 { .. }) { + return Err("Qwen2.5-VL vision diagnostics require model_type=qwen2_5_vl".to_string()); + } + let processor_identity = processor_identity(model_dir, &config)?; + Ok(Self { + model_dir: model_dir.to_path_buf(), + device: device.to_string(), + config, + processor_identity, + active: None, + }) + } + + fn ensure_bucket( + &mut self, + patch_bucket: usize, + ) -> Result<&mut ActiveDiagnosticModule, String> { + if self.active.as_ref().map(|active| active.patch_bucket) != Some(patch_bucket) { + let (module, layout) = compile_and_load_diagnostics( + &self.model_dir, + &self.device, + &self.config, + &self.processor_identity, + patch_bucket, + )?; + self.active = Some(ActiveDiagnosticModule { + patch_bucket, + module, + layout, + }); + } + self.active + .as_mut() + .ok_or_else(|| "Qwen2.5-VL diagnostic module was not installed".to_string()) + } + + pub fn capture( + &mut self, + temporal_patch_rows: &[f32], + grids: &[(i32, i32, i32)], + mutation: Qwen25VlDiagnosticMutation, + ) -> Result { + let Qwen2VlHostInputs { + plan, + mut patches, + mut vision_rope_freqs, + packed_attention_bias, + mut window_attention_bias, + } = prepare_qwen2_vl_host_inputs(&self.config, grids, temporal_patch_rows)?; + let merge_unit = self.config.spatial_merge_size * self.config.spatial_merge_size; + let patch_width = self.config.channels + * self.config.temporal_patch_size + * self.config.patch_size + * self.config.patch_size; + let frequency_width = self.config.hidden / self.config.heads / 2; + let attention_mask_audit = audit_attention_masks( + plan.patch_bucket, + plan.actual_patches, + &plan.packed_cu_seqlens, + &packed_attention_bias, + window_attention_bias.as_deref().ok_or_else(|| { + "Qwen2.5-VL diagnostics require a window-attention bias".to_string() + })?, + ); + match mutation { + Qwen25VlDiagnosticMutation::None => {} + Qwen25VlDiagnosticMutation::Qwen2FullAttention => { + window_attention_bias = Some(packed_attention_bias.clone()); + } + Qwen25VlDiagnosticMutation::IdentityPermutation => { + patches = restore_grouped_rows( + &patches, + patch_width, + merge_unit, + plan.patch_bucket, + plan.actual_patches, + &plan.restore_indices, + )?; + vision_rope_freqs = restore_grouped_rows( + &vision_rope_freqs, + frequency_width, + merge_unit, + plan.patch_bucket, + plan.actual_patches, + &plan.restore_indices, + )?; + } + Qwen25VlDiagnosticMutation::ZeroVisionPositions => { + vision_rope_freqs.fill(0.0); + } + } + let window_attention_bias = window_attention_bias + .ok_or_else(|| "Qwen2.5-VL diagnostics require a window-attention bias".to_string())?; + let patch_shape = [plan.patch_bucket, patch_width]; + let frequency_shape = [plan.patch_bucket, frequency_width]; + let bias_shape = [plan.patch_bucket, plan.patch_bucket]; + let diagnostic_config = self.config.clone(); + let active = self.ensure_bucket(plan.patch_bucket)?; + let layout = active.layout.clone(); + let full_layer_count = layout.full_layer_indices.len(); + let final_interval_layer_count = layout.final_interval_layer_indices.len(); + let diagnostic_layer_count = layout.diagnostic_layer_indices.len(); + let output_specs = + qwen25_diagnostic_output_specs(&diagnostic_config, &layout, plan.patch_bucket); + let merger_output_index = output_specs.len() - 1; + let mut output_buffers = output_specs + .iter() + .map(|spec| { + vec![0u8; spec.shape.iter().product::() * std::mem::size_of::()] + }) + .collect::>(); + let inputs = [ + AuxiliaryInput { + bytes: f32_as_bytes(&patches), + dtype: AuxiliaryTensorDType::Float32, + shape: &patch_shape, + }, + AuxiliaryInput { + bytes: f32_as_bytes(&vision_rope_freqs), + dtype: AuxiliaryTensorDType::Float32, + shape: &frequency_shape, + }, + AuxiliaryInput { + bytes: f32_as_bytes(&packed_attention_bias), + dtype: AuxiliaryTensorDType::Float32, + shape: &bias_shape, + }, + AuxiliaryInput { + bytes: f32_as_bytes(&window_attention_bias), + dtype: AuxiliaryTensorDType::Float32, + shape: &bias_shape, + }, + ]; + let mut outputs = output_buffers + .iter_mut() + .zip(&output_specs) + .map(|(bytes, spec)| AuxiliaryOutput { + bytes, + dtype: AuxiliaryTensorDType::Float32, + shape: &spec.shape, + }) + .collect::>(); + active.module.invoke(&inputs, &mut outputs)?; + let mut decoded = output_buffers + .into_iter() + .zip(&output_specs) + .map(|(bytes, spec)| { + checked_f32_output( + &format!("Qwen2.5-VL diagnostic output {}", spec.label), + bytes, + ) + }) + .collect::, _>>()?; + // The public diagnostics API intentionally owns flat comparison + // vectors, but the IREE descriptors above retain the emitted rank-3 + // [patch, head, head_dim] axes for Q/K/V and the wider MLP + // intermediate axis until after validated device transfer. + for (output, spec) in decoded + .iter_mut() + .zip(&output_specs) + .take(merger_output_index) + { + let actual_values = spec + .shape + .iter() + .skip(1) + .try_fold(plan.actual_patches, |values, dimension| { + values.checked_mul(*dimension) + }); + output.truncate(actual_values.ok_or_else(|| { + format!( + "Qwen2.5-VL diagnostic output {} size overflowed", + spec.label + ) + })?); + } + let all_projected = decoded + .pop() + .ok_or_else(|| "Qwen2.5-VL diagnostic merger output is missing".to_string())?; + let actual_tokens = plan.merged_tokens_per_image.iter().sum::(); + let actual_projected_values = actual_tokens + .checked_mul(self.config.text_hidden) + .ok_or_else(|| "Qwen2.5-VL diagnostic merger size overflowed".to_string())?; + let merger_window_ordered = all_projected + .get(..actual_projected_values) + .ok_or_else(|| "Qwen2.5-VL diagnostic merger output is truncated".to_string())? + .to_vec(); + let mut restored_projection = Vec::with_capacity(actual_projected_values); + for &window_position in &plan.restore_indices { + let start = window_position + .checked_mul(self.config.text_hidden) + .ok_or_else(|| "Qwen2.5-VL diagnostic restore offset overflowed".to_string())?; + restored_projection.extend_from_slice( + merger_window_ordered + .get(start..start + self.config.text_hidden) + .ok_or_else(|| { + "Qwen2.5-VL diagnostic restore index is out of range".to_string() + })?, + ); + } + let reordered_patch_embedding = decoded.remove(0); + let post_window_layer = decoded.remove(0); + let post_full_layers = decoded.drain(..full_layer_count).collect::>(); + let post_final_interval_layers = decoded + .drain(..final_interval_layer_count) + .collect::>(); + let post_diagnostic_layers = decoded.drain(..diagnostic_layer_count).collect::>(); + let mut substage_probe_layer_substages = decoded.into_iter(); + let substage_probe_layer_input = substage_probe_layer_substages + .next() + .expect("diagnostic substage probe layer input output"); + let substage_probe_layer_norm1 = substage_probe_layer_substages + .next() + .expect("diagnostic substage probe layer norm1 output"); + let substage_probe_layer_query = substage_probe_layer_substages + .next() + .expect("diagnostic substage probe layer query output"); + let substage_probe_layer_key = substage_probe_layer_substages + .next() + .expect("diagnostic substage probe layer key output"); + let substage_probe_layer_value = substage_probe_layer_substages + .next() + .expect("diagnostic substage probe layer value output"); + let substage_probe_layer_attention_context = substage_probe_layer_substages + .next() + .expect("diagnostic substage probe layer attention context output"); + let substage_probe_layer_attention = substage_probe_layer_substages + .next() + .expect("diagnostic substage probe layer attention output"); + let substage_probe_layer_post_attention_residual = substage_probe_layer_substages + .next() + .expect("diagnostic substage probe layer residual output"); + let substage_probe_layer_norm2 = substage_probe_layer_substages + .next() + .expect("diagnostic substage probe layer norm2 output"); + let substage_probe_layer_mlp_gate_projection = substage_probe_layer_substages + .next() + .expect("diagnostic substage probe layer MLP gate projection output"); + let substage_probe_layer_mlp_gate_activation = substage_probe_layer_substages + .next() + .expect("diagnostic substage probe layer MLP gate activation output"); + let substage_probe_layer_mlp_up_projection = substage_probe_layer_substages + .next() + .expect("diagnostic substage probe layer MLP up projection output"); + let substage_probe_layer_mlp_gated_product = substage_probe_layer_substages + .next() + .expect("diagnostic substage probe layer MLP gated product output"); + let substage_probe_layer_mlp_down_projection = substage_probe_layer_substages + .next() + .expect("diagnostic substage probe layer MLP down projection output"); + debug_assert!(substage_probe_layer_substages.next().is_none()); + Ok(Qwen25VlVisionDiagnostics { + window_index: plan.window_index, + restore_indices: plan.restore_indices, + packed_cu_seqlens: plan.packed_cu_seqlens, + window_cu_seqlens: plan + .window_cu_seqlens + .ok_or_else(|| "Qwen2.5-VL diagnostic window lengths are missing".to_string())?, + reordered_patch_embedding, + window_layer_index: layout.window_layer_index, + post_window_layer, + full_layer_indices: layout.full_layer_indices, + post_full_layers, + final_interval_layer_indices: layout.final_interval_layer_indices, + post_final_interval_layers, + diagnostic_layer_indices: layout.diagnostic_layer_indices, + post_diagnostic_layers, + substage_probe_layer_index: layout.substage_probe_layer_index, + substage_probe_layer_input, + substage_probe_layer_norm1, + substage_probe_layer_query, + substage_probe_layer_key, + substage_probe_layer_value, + substage_probe_layer_attention_context, + substage_probe_layer_attention, + substage_probe_layer_post_attention_residual, + substage_probe_layer_norm2, + substage_probe_layer_mlp_gate_projection, + substage_probe_layer_mlp_gate_activation, + substage_probe_layer_mlp_up_projection, + substage_probe_layer_mlp_gated_product, + substage_probe_layer_mlp_down_projection, + merger_window_ordered, + restored_projection, + patch_bucket: plan.patch_bucket, + patch_tokens: plan.actual_patches, + merged_tokens: actual_tokens, + vision_hidden: self.config.hidden, + text_hidden: self.config.text_hidden, + attention_mask_audit, + }) + } +} + #[cfg(test)] mod tests { use super::*; + #[cfg(feature = "diagnostics")] + #[test] + fn attention_mask_audit_pins_media_and_padding_isolation() { + let bucket = 8; + let actual = 4; + let boundaries = [0, 2, 4]; + let mut full = vec![f32::NEG_INFINITY; bucket * bucket]; + for range in boundaries.windows(2) { + for row in range[0]..range[1] { + for column in range[0]..range[1] { + full[row * bucket + column] = 0.0; + } + } + } + for row in actual..bucket { + full[row * bucket + row] = 0.0; + } + let mut window = full.clone(); + window[1] = f32::NEG_INFINITY; + window[bucket] = f32::NEG_INFINITY; + let audit = audit_attention_masks(bucket, actual, &boundaries, &full, &window); + assert_eq!(audit.full_window_differences, 2); + assert_eq!(audit.cross_media_full_bias_leaks, 0); + assert_eq!(audit.active_padding_full_bias_leaks, 0); + assert_eq!(audit.invalid_padded_full_bias_entries, 0); + } + + #[cfg(feature = "diagnostics")] + #[test] + fn qwen25_actual_output_descriptors_pin_semantics_order_and_rank() { + let config = Qwen2VlConfig::from_json_str( + r#"{ + "model_type":"qwen2_5_vl", + "hidden_size":1536, + "vision_config":{ + "depth":32, + "hidden_act":"silu", + "hidden_size":1280, + "intermediate_size":3420, + "out_hidden_size":1536, + "num_heads":16, + "in_chans":3, + "patch_size":14, + "spatial_merge_size":2, + "temporal_patch_size":2, + "window_size":112, + "fullatt_block_indexes":[7,15,23,31] + } + }"#, + ) + .unwrap(); + let (_, layout) = emit_qwen2_5_vl_diagnostics( + &config, + 256, + resolve_precision_checked("local-sync").unwrap(), + ) + .unwrap(); + let specs = qwen25_diagnostic_output_specs(&config, &layout, 256); + assert_eq!( + specs + .iter() + .map(|spec| spec.label.as_str()) + .collect::>(), + vec![ + "reordered_patch_embedding", + "post_window_layer.0", + "post_full_layer.7", + "post_full_layer.15", + "post_full_layer.23", + "post_full_layer.31", + "post_final_interval_layer.24", + "post_final_interval_layer.25", + "post_final_interval_layer.26", + "post_final_interval_layer.27", + "post_final_interval_layer.28", + "post_final_interval_layer.29", + "post_final_interval_layer.30", + "post_final_interval_layer.31", + "post_diagnostic_layer.16", + "post_diagnostic_layer.17", + "substage_probe_layer.17.input", + "substage_probe_layer.17.norm1", + "substage_probe_layer.17.query", + "substage_probe_layer.17.key", + "substage_probe_layer.17.value", + "substage_probe_layer.17.attention_context", + "substage_probe_layer.17.attention", + "substage_probe_layer.17.post_attention_residual", + "substage_probe_layer.17.norm2", + "substage_probe_layer.17.mlp_gate_projection", + "substage_probe_layer.17.mlp_gate_activation", + "substage_probe_layer.17.mlp_up_projection", + "substage_probe_layer.17.mlp_gated_product", + "substage_probe_layer.17.mlp_down_projection", + "merger_window_ordered", + ] + ); + assert_eq!( + specs + .iter() + .map(|spec| spec.shape.len()) + .collect::>(), + [vec![2; 18], vec![3; 3], vec![2; 10],].concat() + ); + for spec in &specs[18..=20] { + assert_eq!(spec.shape, vec![256, 16, 80]); + } + assert_eq!(specs[17].shape, vec![256, 1280]); + assert_eq!(specs[21].shape, vec![256, 1280]); + for spec in &specs[25..=28] { + assert_eq!(spec.shape, vec![256, 3420]); + } + assert_eq!(specs[29].shape, vec![256, 1280]); + assert_eq!(specs[30].shape, vec![64, 1536]); + } + fn tiny_config() -> Qwen2VlConfig { Qwen2VlConfig::from_json_str( r#"{ @@ -762,4 +1503,12 @@ mod tests { .contains("flat index 1") ); } + + #[cfg(feature = "diagnostics")] + #[test] + fn diagnostic_identity_mutation_restores_window_groups_to_original_order() { + let window_ordered = [2.0, 0.0, 3.0, 1.0]; + let restored = restore_grouped_rows(&window_ordered, 1, 1, 4, 4, &[1, 3, 0, 2]).unwrap(); + assert_eq!(restored, vec![0.0, 1.0, 2.0, 3.0]); + } } diff --git a/src/loading/mod.rs b/src/loading/mod.rs index fb2cab9d4..46b4fbb43 100644 --- a/src/loading/mod.rs +++ b/src/loading/mod.rs @@ -59,13 +59,15 @@ use self::vlm::*; pub(crate) use self::vlm::load_llava_host_preprocessor; #[cfg(feature = "xla-iree")] pub(crate) use self::vlm::load_llava_iree_host_preprocessor; -#[cfg(feature = "xla-iree")] -pub(crate) use self::vlm::load_qwen2_vl_iree_host_preprocessor; pub use self::vlm::load_qwen3_omni_speech; #[cfg(feature = "xla-iree")] pub(crate) use self::vlm::{ Phi4MMXlaVisionComponents, load_phi4mm_xla_media_components, load_phi4mm_xla_text_embeddings, }; +#[cfg(feature = "xla-iree")] +pub(crate) use self::vlm::{ + load_qwen2_5_vl_iree_host_preprocessor, load_qwen2_vl_iree_host_preprocessor, +}; /// Resolve model path: if a file is given, use its parent directory. /// diff --git a/src/loading/vlm.rs b/src/loading/vlm.rs index b87c6c0a2..8040fb2f5 100644 --- a/src/loading/vlm.rs +++ b/src/loading/vlm.rs @@ -114,12 +114,14 @@ pub(crate) use mllama::load_mllama_vlm; pub(crate) use nemotron_h_nano_omni::load_nemotron_h_nano_omni_vlm; pub(crate) use paddleocr::load_paddleocr_vl; pub(crate) use pixtral::{load_mistral3_vlm, load_pixtral_vlm}; -#[cfg(feature = "xla-iree")] -pub(crate) use qwen::load_qwen2_vl_iree_host_preprocessor; pub(crate) use qwen::{ load_glm_ocr, load_glm4v, load_glm4v_moe, load_qwen2_5_vl, load_qwen2_vl, load_qwen3_5_moe_vlm, load_qwen3_5_vlm, load_qwen3_omni_moe, load_qwen3_vl, load_qwen3_vl_moe, }; +#[cfg(feature = "xla-iree")] +pub(crate) use qwen::{ + load_qwen2_5_vl_iree_host_preprocessor, load_qwen2_vl_iree_host_preprocessor, +}; // Public (re-exported at the crate root): the CLI loads the speech stack // lazily, outside the LoadedModel dispatch. pub use qwen::load_qwen3_omni_speech; diff --git a/src/loading/vlm_qwen.rs b/src/loading/vlm_qwen.rs index 52971240c..803d22c0a 100644 --- a/src/loading/vlm_qwen.rs +++ b/src/loading/vlm_qwen.rs @@ -200,6 +200,107 @@ pub(crate) fn load_qwen2_vl_iree_host_preprocessor( ) } +/// Load the Qwen2.5-VL processor and text embedding table while keeping its +/// windowed vision tower and merger resident in IREE. +#[cfg(feature = "xla-iree")] +pub(crate) fn load_qwen2_5_vl_iree_host_preprocessor( + model_path: &Path, + device: &str, +) -> Result { + use mlxcel_core::layers::UnifiedEmbedding; + use vision::encoders::qwen2_5_vl::Qwen25VLVisionConfig; + + let (_config_str, full_config) = read_sanitized_vlm_config(model_path) + .map_err(|error| HostPreprocessorError::InvalidConfig(error.to_string()))?; + if full_config + .get("model_type") + .and_then(serde_json::Value::as_str) + != Some("qwen2_5_vl") + { + return Err(HostPreprocessorError::FamilyMismatch { + actual: full_config + .get("model_type") + .and_then(serde_json::Value::as_str) + .unwrap_or("") + .to_string(), + }); + } + let mut vision_config: Qwen25VLVisionConfig = + parse_required_vlm_subconfig(&full_config, "vision_config", "Qwen2.5VL vision config") + .map_err(|error| HostPreprocessorError::InvalidConfig(error.to_string()))?; + inherit_qwen_vision_quantization(&mut vision_config, &full_config); + let weights = load_vlm_weights_common_filtered_canonical(model_path, |name| { + let canonical = name.strip_prefix("language_model.").unwrap_or(name); + canonical.starts_with("model.embed_tokens.") + }) + .map(strip_language_model_prefix) + .map_err(|error| HostPreprocessorError::WeightLoad(error.to_string()))?; + let quant_group_size = full_config + .get("quantization") + .and_then(|value| value.get("group_size")) + .and_then(serde_json::Value::as_i64) + .unwrap_or(0) as i32; + let quant_bits = full_config + .get("quantization") + .and_then(|value| value.get("bits")) + .and_then(serde_json::Value::as_i64) + .unwrap_or(0) as i32; + let text_embeddings = UnifiedEmbedding::from_weights( + &weights, + "model.embed_tokens", + quant_group_size, + quant_bits, + ) + .map_err(|error| { + HostPreprocessorError::WeightLoad(format!( + "missing or invalid Qwen2.5-VL text embedding table: {error}" + )) + })?; + let hidden_size = full_config + .get("hidden_size") + .and_then(serde_json::Value::as_u64) + .and_then(|value| usize::try_from(value).ok()) + .filter(|&value| value > 0) + .ok_or_else(|| { + HostPreprocessorError::InvalidConfig( + "Qwen2.5-VL hidden_size must be a positive integer".to_string(), + ) + })?; + let max_sequence_len = full_config + .get("max_position_embeddings") + .and_then(serde_json::Value::as_u64) + .and_then(|value| usize::try_from(value).ok()) + .filter(|&value| value > 0) + .ok_or_else(|| { + HostPreprocessorError::InvalidConfig( + "Qwen2.5-VL max_position_embeddings must be a positive integer".to_string(), + ) + })?; + let token_ids = qwen_vl_token_ids( + &full_config, + QwenVisionTokenIds { + image_token_id: 151655, + video_token_id: 151656, + vision_start_token_id: 151652, + }, + ); + let processor = qwen_vl_processor(&vision_config); + let projector = mlxcel_xla::IreeQwen2VlProjector::load(model_path, device) + .map_err(HostPreprocessorError::Iree)?; + Qwen2VlIreeHostPreprocessor::from_parts( + processor, + text_embeddings, + projector, + token_ids.image_token_id, + token_ids.video_token_id, + token_ids.vision_start_token_id, + vision_config.spatial_merge_size, + hidden_size, + max_sequence_len, + device.to_string(), + ) +} + /// Load a Qwen2.5-VL model (windowed ViT + Qwen2 language model with MRoPE) pub(crate) fn load_qwen2_5_vl(model_path: &Path) -> Result { use vision::encoders::qwen2_5_vl::{Qwen25VLVisionConfig, Qwen25VLVisionEncoder}; diff --git a/src/multimodal/host_preprocessor.rs b/src/multimodal/host_preprocessor.rs index bfad9cb9d..bd0461082 100644 --- a/src/multimodal/host_preprocessor.rs +++ b/src/multimodal/host_preprocessor.rs @@ -147,13 +147,20 @@ pub fn load_xla_image_preprocessor( model_path.display() )) })?; - if model_type == crate::models::ModelType::Qwen2VL { + if matches!( + model_type, + crate::models::ModelType::Qwen2VL | crate::models::ModelType::Qwen25VL + ) { + let family = if model_type == crate::models::ModelType::Qwen25VL { + "Qwen2.5-VL" + } else { + "Qwen2-VL" + }; let policy = XlaVisionBackendPolicy::from_env()?; if policy == XlaVisionBackendPolicy::Host { - return Err(HostPreprocessorError::InvalidConfig( - "Qwen2-VL XLA vision has no MLX fallback; MLXCEL_XLA_VISION_BACKEND=host is unsupported" - .to_string(), - )); + return Err(HostPreprocessorError::InvalidConfig(format!( + "{family} XLA vision has no MLX fallback; MLXCEL_XLA_VISION_BACKEND=host is unsupported" + ))); } #[cfg(feature = "xla-iree")] { @@ -163,16 +170,16 @@ pub fn load_xla_image_preprocessor( tracing::info!( vision_backend = "iree", vision_device = %device, - family = "qwen2_vl", + family, "OpenXLA multimodal vision backend selected" ); return Ok(Some(Box::new(preprocessor))); } #[cfg(not(feature = "xla-iree"))] { - return Err(HostPreprocessorError::InvalidConfig( - "Qwen2-VL XLA image execution requires the xla-iree feature".to_string(), - )); + return Err(HostPreprocessorError::InvalidConfig(format!( + "{family} XLA image execution requires the xla-iree feature" + ))); } } if model_type != crate::models::ModelType::LlavaVLM { @@ -292,10 +299,46 @@ pub struct Qwen2VlIreeHostPreprocessor { device: String, } +#[cfg(feature = "xla-diagnostics")] +#[derive(Debug, Clone, PartialEq)] +pub struct QwenVlProcessorCapture { + pub patch_values: Vec, + pub grids: Vec<(i32, i32, i32)>, +} + #[cfg(feature = "xla-iree")] impl Qwen2VlIreeHostPreprocessor { pub fn load(model_path: &Path, device: &str) -> Result { - crate::loading::load_qwen2_vl_iree_host_preprocessor(model_path, device) + match crate::models::get_model_type(model_path) + .map_err(|error| HostPreprocessorError::InvalidConfig(error.to_string()))? + { + crate::models::ModelType::Qwen2VL => { + crate::loading::load_qwen2_vl_iree_host_preprocessor(model_path, device) + } + crate::models::ModelType::Qwen25VL => { + crate::loading::load_qwen2_5_vl_iree_host_preprocessor(model_path, device) + } + actual => Err(HostPreprocessorError::FamilyMismatch { + actual: format!("{actual:?}"), + }), + } + } + + #[cfg(feature = "xla-diagnostics")] + pub fn capture_processor_inputs( + &self, + images: &[DynamicImage], + ) -> Result { + if images.is_empty() { + return Err(HostPreprocessorError::InvalidConfig( + "Qwen-VL vision diagnostics require at least one image".to_string(), + )); + } + let (patch_values, grids) = self.processor.preprocess_values_with_grid(images); + Ok(QwenVlProcessorCapture { + patch_values, + grids, + }) } #[allow(clippy::too_many_arguments)] @@ -344,7 +387,7 @@ impl Qwen2VlIreeHostPreprocessor { ) -> Result { if images.is_empty() { return Err(HostPreprocessorError::InvalidConfig( - "Qwen2-VL image preprocessing requires at least one decoded image; text-only requests use ordinary XLA token prefill" + "Qwen-VL image preprocessing requires at least one decoded image; text-only requests use ordinary XLA token prefill" .to_string(), )); } @@ -359,13 +402,13 @@ impl Qwen2VlIreeHostPreprocessor { ) .ok_or_else(|| { HostPreprocessorError::InvalidConfig(format!( - "Qwen2-VL image placeholder count does not match {} decoded image(s), or the prompt is already expanded", + "Qwen-VL image placeholder count does not match {} decoded image(s), or the prompt is already expanded", images.len() )) })?; if inserted.image_blocks != images.len() { return Err(HostPreprocessorError::InvalidConfig(format!( - "Qwen2-VL expanded {} image blocks for {} decoded images", + "Qwen-VL expanded {} image blocks for {} decoded images", inserted.image_blocks, images.len() ))); diff --git a/src/vision/encoders/qwen2_5_vl.rs b/src/vision/encoders/qwen2_5_vl.rs index 16af1f280..745ad83c3 100644 --- a/src/vision/encoders/qwen2_5_vl.rs +++ b/src/vision/encoders/qwen2_5_vl.rs @@ -98,6 +98,22 @@ fn default_fullatt_block_indexes() -> Vec { vec![7, 15, 23, 31] } +fn restoration_indices(window_index: &[i32]) -> Vec { + let mut restore = vec![0i32; window_index.len()]; + let mut seen = vec![false; window_index.len()]; + for (window_position, &original_position) in window_index.iter().enumerate() { + let original_position = usize::try_from(original_position) + .expect("Qwen2.5-VL window indices are validated as non-negative"); + assert!( + original_position < restore.len() && !seen[original_position], + "Qwen2.5-VL window indices must be an in-range permutation" + ); + seen[original_position] = true; + restore[original_position] = window_position as i32; + } + restore +} + // RMSNorm for vision encoder. struct VisionRMSNorm { weight: UniquePtr, @@ -115,7 +131,7 @@ impl VisionRMSNorm { } fn forward(&self, x: &MlxArray) -> UniquePtr { - mlxcel_core::rms_norm(x, &self.weight, self.eps) + mlxcel_core::fast_rms_norm(x, &self.weight, self.eps) } } @@ -203,6 +219,14 @@ struct VisionAttention { scale: f32, } +#[derive(Clone, Copy)] +enum VisionAttentionDiagnosticStage { + Query, + Key, + Value, + Context, +} + impl VisionAttention { fn from_weights( weights: &WeightMap, @@ -231,6 +255,19 @@ impl VisionAttention { cu_seqlens: &[i32], rotary_pos_emb: &MlxArray, ) -> UniquePtr { + self.forward_with_observer(x, cu_seqlens, rotary_pos_emb, |_, _| {}) + } + + fn forward_with_observer( + &self, + x: &MlxArray, + cu_seqlens: &[i32], + rotary_pos_emb: &MlxArray, + mut observe: F, + ) -> UniquePtr + where + F: FnMut(VisionAttentionDiagnosticStage, &MlxArray), + { let shape = mlxcel_core::array_shape(x); let seq_length = shape[0]; @@ -259,6 +296,9 @@ impl VisionAttention { let q = apply_rotary_pos_emb_vision(&q, rotary_pos_emb); let k = apply_rotary_pos_emb_vision(&k, rotary_pos_emb); + observe(VisionAttentionDiagnosticStage::Query, &q); + observe(VisionAttentionDiagnosticStage::Key, &k); + observe(VisionAttentionDiagnosticStage::Value, &v); // [seq, heads, head_dim] -> [1, heads, seq, head_dim] let q = mlxcel_core::transpose_axes(&q, &[1, 0, 2]); @@ -315,6 +355,7 @@ impl VisionAttention { let output = mlxcel_core::squeeze_axis(&output, 0); let output = mlxcel_core::transpose_axes(&output, &[1, 0, 2]); let output = mlxcel_core::reshape(&output, &[seq_length, -1]); + observe(VisionAttentionDiagnosticStage::Context, &output); self.proj.forward(&output) } @@ -327,6 +368,15 @@ struct VisionMLP { down_proj: UnifiedLinear, } +#[cfg(feature = "xla-diagnostics")] +#[derive(Clone, Copy)] +enum VisionMLPDiagnosticStage { + GateProjection, + GateActivation, + UpProjection, + GatedProduct, +} + impl VisionMLP { fn from_weights(weights: &WeightMap, prefix: &str, gs: i32, bits: i32) -> Result { Ok(Self { @@ -358,6 +408,57 @@ impl VisionMLP { let h = mlxcel_core::multiply(&gate, &up); self.down_proj.forward(&h) } + + #[cfg(feature = "xla-diagnostics")] + fn dense_f32_projection(linear: &UnifiedLinear, x: &MlxArray) -> UniquePtr { + let x = mlxcel_core::astype(x, mlxcel_core::dtype::FLOAT32); + let weight = linear.dequantized_weight(); + let weight = mlxcel_core::astype(&weight, mlxcel_core::dtype::FLOAT32); + let weight = mlxcel_core::transpose(&weight); + let mut output = mlxcel_core::matmul(&x, &weight); + let (global_scale, bias) = match linear { + UnifiedLinear::Quantized { weight, bias } => { + (weight.global_scale.as_ref(), bias.as_ref()) + } + UnifiedLinear::Regular(linear) => (None, linear.bias.as_ref()), + }; + if let Some(global_scale) = global_scale { + let global_scale = mlxcel_core::astype(global_scale, mlxcel_core::dtype::FLOAT32); + output = mlxcel_core::multiply(&output, &global_scale); + } + if let Some(bias) = bias { + let bias = mlxcel_core::astype(bias, mlxcel_core::dtype::FLOAT32); + output = mlxcel_core::add(&output, &bias); + } + output + } + + #[cfg(feature = "xla-diagnostics")] + fn dense_f32_gate_up_controls( + &self, + x: &MlxArray, + ) -> (UniquePtr, UniquePtr) { + ( + Self::dense_f32_projection(&self.gate_proj, x), + Self::dense_f32_projection(&self.up_proj, x), + ) + } + + #[cfg(feature = "xla-diagnostics")] + fn forward_with_observer(&self, x: &MlxArray, mut observe: F) -> UniquePtr + where + F: FnMut(VisionMLPDiagnosticStage, &MlxArray), + { + let gate = self.gate_proj.forward(x); + observe(VisionMLPDiagnosticStage::GateProjection, &gate); + let gate = mlxcel_core::silu(&gate); + observe(VisionMLPDiagnosticStage::GateActivation, &gate); + let up = self.up_proj.forward(x); + observe(VisionMLPDiagnosticStage::UpProjection, &up); + let gated_product = mlxcel_core::multiply(&gate, &up); + observe(VisionMLPDiagnosticStage::GatedProduct, &gated_product); + self.down_proj.forward(&gated_product) + } } // VisionBlock - RMSNorm + SwiGLU MLP. @@ -368,6 +469,41 @@ struct VisionBlock { mlp: VisionMLP, } +#[cfg(feature = "xla-diagnostics")] +struct Qwen25VLBlockDiagnostics { + input: Vec, + norm1: Vec, + query: Vec, + key: Vec, + value: Vec, + attention_context: Vec, + attention: Vec, + post_attention_residual: Vec, + norm2: Vec, + mlp_gate_projection: Vec, + mlp_gate_projection_dense_f32_control: Vec, + mlp_gate_activation: Vec, + mlp_up_projection: Vec, + mlp_up_projection_dense_f32_control: Vec, + mlp_gated_product: Vec, + mlp_down_projection: Vec, +} + +#[cfg(any(test, feature = "xla-diagnostics"))] +fn host_f32_diagnostic_snapshot(array: &MlxArray) -> Vec { + let array = mlxcel_core::astype(array, mlxcel_core::dtype::FLOAT32); + mlxcel_core::eval(&array); + let bytes = mlxcel_core::array_to_raw_bytes(&array); + assert!( + bytes.len().is_multiple_of(std::mem::size_of::()), + "Qwen2.5-VL diagnostic F32 snapshot has a partial element" + ); + bytes + .chunks_exact(std::mem::size_of::()) + .map(|bytes| f32::from_ne_bytes([bytes[0], bytes[1], bytes[2], bytes[3]])) + .collect() +} + impl VisionBlock { fn from_weights( weights: &WeightMap, @@ -397,6 +533,92 @@ impl VisionBlock { let mlp_out = self.mlp.forward(&normed); mlxcel_core::add(&h, &mlp_out) } + + #[cfg(feature = "xla-diagnostics")] + fn forward_with_diagnostics( + &self, + hidden_states: &MlxArray, + cu_seqlens: &[i32], + rotary_pos_emb: &MlxArray, + ) -> (UniquePtr, Qwen25VLBlockDiagnostics) { + // Copy each diagnostic value to host-owned F32 memory before its source + // feeds the next substage. Retaining even an evaluated MlxArray handle + // is insufficient because later graph evaluation may donate or reuse + // its backing storage. + let input = host_f32_diagnostic_snapshot(hidden_states); + let norm1 = self.norm1.forward(hidden_states); + let norm1_snapshot = host_f32_diagnostic_snapshot(&norm1); + let mut query = None; + let mut key = None; + let mut value = None; + let mut attention_context = None; + let attention = + self.attn + .forward_with_observer(&norm1, cu_seqlens, rotary_pos_emb, |stage, array| { + let snapshot = host_f32_diagnostic_snapshot(array); + match stage { + VisionAttentionDiagnosticStage::Query => query = Some(snapshot), + VisionAttentionDiagnosticStage::Key => key = Some(snapshot), + VisionAttentionDiagnosticStage::Value => value = Some(snapshot), + VisionAttentionDiagnosticStage::Context => { + attention_context = Some(snapshot) + } + } + }); + let attention_snapshot = host_f32_diagnostic_snapshot(&attention); + let post_attention_residual = mlxcel_core::add(hidden_states, &attention); + let post_attention_residual_snapshot = + host_f32_diagnostic_snapshot(&post_attention_residual); + let norm2 = self.norm2.forward(&post_attention_residual); + let norm2_snapshot = host_f32_diagnostic_snapshot(&norm2); + let (mlp_gate_projection_dense_f32_control, mlp_up_projection_dense_f32_control) = + self.mlp.dense_f32_gate_up_controls(&norm2); + let mlp_gate_projection_dense_f32_control = + host_f32_diagnostic_snapshot(&mlp_gate_projection_dense_f32_control); + let mlp_up_projection_dense_f32_control = + host_f32_diagnostic_snapshot(&mlp_up_projection_dense_f32_control); + let mut mlp_gate_projection = None; + let mut mlp_gate_activation = None; + let mut mlp_up_projection = None; + let mut mlp_gated_product = None; + let mlp_down_projection = self.mlp.forward_with_observer(&norm2, |stage, array| { + let snapshot = host_f32_diagnostic_snapshot(array); + match stage { + VisionMLPDiagnosticStage::GateProjection => mlp_gate_projection = Some(snapshot), + VisionMLPDiagnosticStage::GateActivation => mlp_gate_activation = Some(snapshot), + VisionMLPDiagnosticStage::UpProjection => mlp_up_projection = Some(snapshot), + VisionMLPDiagnosticStage::GatedProduct => mlp_gated_product = Some(snapshot), + } + }); + let mlp_down_projection_snapshot = host_f32_diagnostic_snapshot(&mlp_down_projection); + let output = mlxcel_core::add(&post_attention_residual, &mlp_down_projection); + ( + output, + Qwen25VLBlockDiagnostics { + input, + norm1: norm1_snapshot, + query: query.expect("diagnostics capture attention query"), + key: key.expect("diagnostics capture attention key"), + value: value.expect("diagnostics capture attention value"), + attention_context: attention_context + .expect("diagnostics capture pre-projection attention context"), + attention: attention_snapshot, + post_attention_residual: post_attention_residual_snapshot, + norm2: norm2_snapshot, + mlp_gate_projection: mlp_gate_projection + .expect("diagnostics capture MLP gate projection"), + mlp_gate_projection_dense_f32_control, + mlp_gate_activation: mlp_gate_activation + .expect("diagnostics capture MLP gate activation"), + mlp_up_projection: mlp_up_projection + .expect("diagnostics capture MLP up projection"), + mlp_up_projection_dense_f32_control, + mlp_gated_product: mlp_gated_product + .expect("diagnostics capture MLP gated product"), + mlp_down_projection: mlp_down_projection_snapshot, + }, + ) + } } // PatchMerger - RMSNorm + GELU MLP (projection to text hidden size). @@ -449,6 +671,53 @@ pub struct Qwen25VLVisionEncoder { fullatt_block_indexes: Vec, } +#[cfg(feature = "xla-diagnostics")] +pub struct Qwen25VLVisionDiagnostics { + pub window_index: Vec, + pub restore_indices: Vec, + pub packed_cu_seqlens: Vec, + pub window_cu_seqlens: Vec, + pub reordered_patch_embedding: UniquePtr, + pub window_layer_index: usize, + pub post_window_layer: UniquePtr, + pub full_layer_indices: Vec, + pub post_full_layers: Vec>, + pub final_interval_layer_indices: Vec, + pub post_final_interval_layers: Vec>, + pub diagnostic_layer_indices: Vec, + pub post_diagnostic_layers: Vec>, + pub substage_probe_layer_index: usize, + pub substage_probe_layer_input: Vec, + pub substage_probe_layer_norm1: Vec, + pub substage_probe_layer_query: Vec, + pub substage_probe_layer_key: Vec, + pub substage_probe_layer_value: Vec, + pub substage_probe_layer_attention_context: Vec, + pub substage_probe_layer_attention: Vec, + pub substage_probe_layer_post_attention_residual: Vec, + pub substage_probe_layer_norm2: Vec, + pub substage_probe_layer_mlp_gate_projection: Vec, + pub substage_probe_layer_mlp_gate_projection_dense_f32_control: Vec, + pub substage_probe_layer_mlp_gate_activation: Vec, + pub substage_probe_layer_mlp_up_projection: Vec, + pub substage_probe_layer_mlp_up_projection_dense_f32_control: Vec, + pub substage_probe_layer_mlp_gated_product: Vec, + pub substage_probe_layer_mlp_down_projection: Vec, + pub merger_window_ordered: UniquePtr, + pub restored_projection: UniquePtr, +} + +#[cfg(feature = "xla-diagnostics")] +pub struct Qwen25VLMlpProjectionControls { + pub gate_projection: Vec, + pub up_projection: Vec, +} + +#[cfg(feature = "xla-diagnostics")] +type Qwen25VLCapture = Option; +#[cfg(not(feature = "xla-diagnostics"))] +type Qwen25VLCapture = (); + impl Qwen25VLVisionEncoder { pub fn from_weights( weights: &WeightMap, @@ -495,6 +764,51 @@ impl Qwen25VLVisionEncoder { }) } + #[cfg(feature = "xla-diagnostics")] + pub fn f16_input_dense_f32_mlp_projection_controls( + &self, + layer_index: usize, + input: &[f32], + rows: usize, + ) -> Result { + let block = self + .blocks + .get(layer_index) + .ok_or_else(|| format!("Qwen2.5-VL diagnostic layer {layer_index} is out of range"))?; + if rows == 0 || input.len() % rows != 0 { + return Err(format!( + "Qwen2.5-VL diagnostic projection input has {} values for {rows} rows", + input.len() + )); + } + let width = input.len() / rows; + let width_i32 = i32::try_from(width) + .map_err(|_| "Qwen2.5-VL diagnostic projection width exceeds i32".to_string())?; + for (name, linear) in [("gate", &block.mlp.gate_proj), ("up", &block.mlp.up_proj)] { + let weight = linear.dequantized_weight(); + let shape = mlxcel_core::array_shape(&weight); + if shape.len() != 2 || shape[1] != width_i32 { + return Err(format!( + "Qwen2.5-VL layer {layer_index} {name} projection weight shape {shape:?} \ + does not accept diagnostic input width {width}" + )); + } + } + let rows = i32::try_from(rows) + .map_err(|_| "Qwen2.5-VL diagnostic projection row count exceeds i32".to_string())?; + let input = mlxcel_core::from_slice_f32(input, &[rows, width_i32]); + // The CUDA IREE graph demotes its captured F32 norm2 value to F16 + // immediately before dot_general. Round through the same input type, + // then use the existing dense-F32 control to isolate accumulation and + // reduction behavior from inherited norm2 drift. + let input = mlxcel_core::astype(&input, mlxcel_core::dtype::FLOAT16); + let (gate_projection, up_projection) = block.mlp.dense_f32_gate_up_controls(&input); + Ok(Qwen25VLMlpProjectionControls { + gate_projection: host_f32_diagnostic_snapshot(&gate_projection), + up_projection: host_f32_diagnostic_snapshot(&up_projection), + }) + } + /// Compute 2D rotary position embeddings from grid_thw (same as Qwen2-VL) fn rot_pos_emb(&self, grid_thw: &[(i32, i32, i32)]) -> UniquePtr { let mut all_pos_ids: Vec> = Vec::new(); @@ -696,6 +1010,26 @@ impl Qwen25VLVisionEncoder { hidden_states: &MlxArray, grid_thw: &[(i32, i32, i32)], ) -> VisionEncoderOutput { + self.forward_with_grid_inner::(hidden_states, grid_thw) + .0 + } + + #[cfg(feature = "xla-diagnostics")] + pub fn forward_with_grid_diagnostics( + &self, + hidden_states: &MlxArray, + grid_thw: &[(i32, i32, i32)], + ) -> Qwen25VLVisionDiagnostics { + self.forward_with_grid_inner::(hidden_states, grid_thw) + .1 + .expect("Qwen2.5-VL diagnostics requested a capture") + } + + fn forward_with_grid_inner( + &self, + hidden_states: &MlxArray, + grid_thw: &[(i32, i32, i32)], + ) -> (VisionEncoderOutput, Qwen25VLCapture) { let mut h = self.patch_embed.forward(hidden_states); let rotary_pos_emb = self.rot_pos_emb(grid_thw); @@ -714,6 +1048,9 @@ impl Qwen25VLVisionEncoder { mlxcel_core::from_slice_i32(&window_index, &[window_index.len() as i32]); let h_reordered = mlxcel_core::take(&h_grouped, &window_idx_arr, 0); h = mlxcel_core::reshape(&h_reordered, &[-1, dim]); + #[cfg(feature = "xla-diagnostics")] + let reordered_patch_embedding = CAPTURE + .then(|| mlxcel_core::copy(h.as_ref().expect("Qwen2.5-VL reordered patch embedding"))); // Reorder rotary_pos_emb similarly let rope_shape = mlxcel_core::array_shape(&rotary_pos_emb); @@ -729,29 +1066,123 @@ impl Qwen25VLVisionEncoder { let cu_seqlens = Self::compute_cu_seqlens(grid_thw, self.spatial_merge_size as i32); // Run blocks with windowed or full attention + #[cfg(feature = "xla-diagnostics")] + let window_layer_index = + (0..self.blocks.len()).find(|layer| !self.fullatt_block_indexes.contains(layer)); + #[cfg(feature = "xla-diagnostics")] + let mut post_window_layer = None; + #[cfg(feature = "xla-diagnostics")] + let mut post_full_layers = Vec::new(); + #[cfg(feature = "xla-diagnostics")] + let full_layer_indices = (0..self.blocks.len()) + .filter(|layer| self.fullatt_block_indexes.contains(layer)) + .collect::>(); + #[cfg(feature = "xla-diagnostics")] + let target_full_layer_index = full_layer_indices.last().copied(); + #[cfg(feature = "xla-diagnostics")] + let final_interval_start = full_layer_indices + .iter() + .rev() + .nth(1) + .map_or(0, |layer| layer + 1); + #[cfg(feature = "xla-diagnostics")] + let final_interval_layer_indices = target_full_layer_index + .map(|target| (final_interval_start..=target).collect::>()) + .unwrap_or_default(); + // The two-media CUDA oracle passes layer 16 and first diverges at layer + // 17. Probe that first failing layer so the next bounded run + // distinguishes Q/K/V, masked attention context, projection, residual, + // norm, and MLP without changing production execution. + #[cfg(feature = "xla-diagnostics")] + let target_probe_boundary = full_layer_indices + .iter() + .rev() + .nth(1) + .copied() + .or(target_full_layer_index); + #[cfg(feature = "xla-diagnostics")] + let substage_probe_layer_index = target_probe_boundary.map(|target| { + full_layer_indices + .iter() + .copied() + .take_while(|&layer| layer < target) + .last() + .and_then(|previous_full| previous_full.checked_add(2)) + .filter(|&layer| layer <= target) + .unwrap_or(target) + }); + #[cfg(feature = "xla-diagnostics")] + let diagnostic_layer_indices = substage_probe_layer_index + .and_then(|probe| { + full_layer_indices + .iter() + .copied() + .take_while(|&layer| layer < probe) + .last() + .map(|previous_full| ((previous_full + 1)..=probe).collect::>()) + }) + .unwrap_or_default(); + #[cfg(feature = "xla-diagnostics")] + let mut post_final_interval_layers = Vec::new(); + #[cfg(feature = "xla-diagnostics")] + let mut post_diagnostic_layers = Vec::new(); + #[cfg(feature = "xla-diagnostics")] + let mut substage_probe_layer_state = None; for (layer_num, block) in self.blocks.iter().enumerate() { let cu_seqlens_now = if self.fullatt_block_indexes.contains(&layer_num) { &cu_seqlens } else { &cu_window_seqlens }; - h = block.forward(&h, cu_seqlens_now, &rotary_pos_emb); + #[cfg(feature = "xla-diagnostics")] + if CAPTURE && Some(layer_num) == substage_probe_layer_index { + let (output, capture) = + block.forward_with_diagnostics(&h, cu_seqlens_now, &rotary_pos_emb); + h = output; + substage_probe_layer_state = Some(capture); + } else { + h = block.forward(&h, cu_seqlens_now, &rotary_pos_emb); + } + #[cfg(not(feature = "xla-diagnostics"))] + { + h = block.forward(&h, cu_seqlens_now, &rotary_pos_emb); + } + #[cfg(feature = "xla-diagnostics")] + if CAPTURE { + if Some(layer_num) == window_layer_index { + post_window_layer = Some(mlxcel_core::copy( + h.as_ref().expect("Qwen2.5-VL window layer state"), + )); + } + if self.fullatt_block_indexes.contains(&layer_num) { + post_full_layers.push(mlxcel_core::copy( + h.as_ref().expect("Qwen2.5-VL full layer state"), + )); + } + if final_interval_layer_indices.contains(&layer_num) { + post_final_interval_layers.push(mlxcel_core::copy( + h.as_ref().expect("Qwen2.5-VL final interval layer state"), + )); + } + if diagnostic_layer_indices.contains(&layer_num) { + post_diagnostic_layers.push(mlxcel_core::copy( + h.as_ref().expect("Qwen2.5-VL diagnostic layer state"), + )); + } + } } // Merge patches h = self.merger.forward(&h); + #[cfg(feature = "xla-diagnostics")] + let merger_window_ordered = CAPTURE.then(|| { + mlxcel_core::copy(h.as_ref().expect("Qwen2.5-VL window-ordered merger output")) + }); - // Un-reorder: apply reverse indices from argsort(window_index) - let mut reverse_indices = vec![0i32; window_index.len()]; - let mut indexed: Vec<(i32, usize)> = window_index - .iter() - .enumerate() - .map(|(i, &v)| (v, i)) - .collect(); - indexed.sort_by_key(|&(v, _)| v); - for (rank, &(_, orig_idx)) in indexed.iter().enumerate() { - reverse_indices[orig_idx] = rank as i32; - } + // Un-reorder: destination original_position reads its corresponding + // window_position. This is the inverse permutation, not window_index + // itself (the two differ for non-self-inverse window layouts). + let reverse_indices = restoration_indices(&window_index); // After merger: h has shape [total_merged, out_hidden_size] // reverse_indices maps from window-ordered to original order @@ -759,7 +1190,73 @@ impl Qwen25VLVisionEncoder { mlxcel_core::from_slice_i32(&reverse_indices, &[reverse_indices.len() as i32]); h = mlxcel_core::take(&h, &reverse_arr, 0); - VisionEncoderOutput { hidden_states: h } + #[cfg(feature = "xla-diagnostics")] + if CAPTURE { + let substage_probe_layer_index = substage_probe_layer_index + .expect("Qwen2.5-VL diagnostics require a substage probe layer"); + let substage_probe_layer_state = + substage_probe_layer_state.expect("Qwen2.5-VL substage probe layer capture"); + let restored_projection = + mlxcel_core::copy(h.as_ref().expect("Qwen2.5-VL restored projection output")); + let output = VisionEncoderOutput { hidden_states: h }; + return ( + output, + Some(Qwen25VLVisionDiagnostics { + window_index, + restore_indices: reverse_indices, + packed_cu_seqlens: cu_seqlens, + window_cu_seqlens: cu_window_seqlens, + reordered_patch_embedding: reordered_patch_embedding + .expect("Qwen2.5-VL patch capture"), + window_layer_index: window_layer_index + .expect("Qwen2.5-VL diagnostics require a window layer"), + post_window_layer: post_window_layer.expect("Qwen2.5-VL window layer capture"), + full_layer_indices, + post_full_layers, + final_interval_layer_indices, + post_final_interval_layers, + diagnostic_layer_indices, + post_diagnostic_layers, + substage_probe_layer_index, + substage_probe_layer_input: substage_probe_layer_state.input, + substage_probe_layer_norm1: substage_probe_layer_state.norm1, + substage_probe_layer_query: substage_probe_layer_state.query, + substage_probe_layer_key: substage_probe_layer_state.key, + substage_probe_layer_value: substage_probe_layer_state.value, + substage_probe_layer_attention_context: substage_probe_layer_state + .attention_context, + substage_probe_layer_attention: substage_probe_layer_state.attention, + substage_probe_layer_post_attention_residual: substage_probe_layer_state + .post_attention_residual, + substage_probe_layer_norm2: substage_probe_layer_state.norm2, + substage_probe_layer_mlp_gate_projection: substage_probe_layer_state + .mlp_gate_projection, + substage_probe_layer_mlp_gate_projection_dense_f32_control: + substage_probe_layer_state.mlp_gate_projection_dense_f32_control, + substage_probe_layer_mlp_gate_activation: substage_probe_layer_state + .mlp_gate_activation, + substage_probe_layer_mlp_up_projection: substage_probe_layer_state + .mlp_up_projection, + substage_probe_layer_mlp_up_projection_dense_f32_control: + substage_probe_layer_state.mlp_up_projection_dense_f32_control, + substage_probe_layer_mlp_gated_product: substage_probe_layer_state + .mlp_gated_product, + substage_probe_layer_mlp_down_projection: substage_probe_layer_state + .mlp_down_projection, + merger_window_ordered: merger_window_ordered + .expect("Qwen2.5-VL merger capture"), + restored_projection, + }), + ); + } + #[cfg(not(feature = "xla-diagnostics"))] + { + (VisionEncoderOutput { hidden_states: h }, ()) + } + #[cfg(feature = "xla-diagnostics")] + { + (VisionEncoderOutput { hidden_states: h }, None) + } } } @@ -769,3 +1266,111 @@ impl super::VisionEncoder for Qwen25VLVisionEncoder { panic!("Qwen2.5-VL vision encoder requires grid_thw; use forward_with_grid() instead"); } } + +#[cfg(test)] +mod tests { + #[cfg(feature = "xla-diagnostics")] + use super::VisionMLP; + use super::{host_f32_diagnostic_snapshot, restoration_indices}; + + #[test] + fn restoration_is_the_true_inverse_for_non_self_inverse_permutation() { + let window_index = [2, 0, 3, 1]; + let restore = restoration_indices(&window_index); + assert_eq!(restore, vec![1, 3, 0, 2]); + let window_ordered = ["two", "zero", "three", "one"]; + let original = restore + .into_iter() + .map(|index| window_ordered[index as usize]) + .collect::>(); + assert_eq!(original, vec!["zero", "one", "two", "three"]); + } + + #[cfg(feature = "xla-diagnostics")] + #[test] + fn dense_f32_projection_control_preserves_linear_order_and_dtype() { + use mlxcel_core::layers::{Linear, UnifiedLinear}; + + let input = mlxcel_core::from_slice_f32(&[1.0, 2.0, -1.0, 0.5], &[2, 2]); + let input = mlxcel_core::astype(&input, mlxcel_core::dtype::FLOAT16); + let weight = mlxcel_core::from_slice_f32(&[3.0, 4.0, -2.0, 5.0], &[2, 2]); + let weight = mlxcel_core::astype(&weight, mlxcel_core::dtype::FLOAT16); + let bias = mlxcel_core::from_slice_f32(&[0.5, -1.0], &[2]); + let bias = mlxcel_core::astype(&bias, mlxcel_core::dtype::FLOAT16); + let linear = UnifiedLinear::Regular(Linear::new(weight, Some(bias))); + + let output = VisionMLP::dense_f32_projection(&linear, &input); + assert_eq!( + mlxcel_core::array_dtype(&output), + mlxcel_core::dtype::FLOAT32 + ); + let output = host_f32_diagnostic_snapshot(&output); + assert_eq!(output, vec![11.5, 7.0, -0.5, 3.5]); + } + + #[test] + fn host_snapshot_preserves_nonzero_rms_norm_after_later_graph_evaluation() { + let input = mlxcel_core::from_slice_f32(&[1.0, -2.0, 3.0, -4.0], &[1, 4]); + let weight = mlxcel_core::ones(&[4], mlxcel_core::dtype::FLOAT32); + let normalized = mlxcel_core::rms_norm(&input, &weight, 1e-6); + let snapshot = host_f32_diagnostic_snapshot(&normalized); + + assert!( + snapshot + .iter() + .any(|value| value.is_finite() && *value != 0.0), + "RMSNorm snapshot must be non-vacuous" + ); + + let residual = mlxcel_core::add(&input, &normalized); + let donated_candidate = mlxcel_core::multiply(&residual, &residual); + mlxcel_core::eval(&donated_candidate); + + let expected = [0.365_148_37, -0.730_296_73, 1.095_445_2, -1.460_593_5]; + assert_eq!(snapshot.len(), expected.len()); + for (index, (actual, expected)) in snapshot.iter().zip(expected).enumerate() { + assert!( + (*actual - expected).abs() <= 2.0e-6, + "host snapshot changed at index {index}: actual={actual}, expected={expected}" + ); + } + } + + #[test] + fn f16_manual_rms_norm_characterizes_square_overflow() { + let input = mlxcel_core::from_slice_f32(&[300.0, 1.0, -2.0, 3.0], &[1, 4]); + let input = mlxcel_core::astype(&input, mlxcel_core::dtype::FLOAT16); + let weight = mlxcel_core::ones(&[4], mlxcel_core::dtype::FLOAT16); + + let manual = mlxcel_core::rms_norm(&input, &weight, 1e-6); + let manual = host_f32_diagnostic_snapshot(&manual); + assert!( + manual.iter().all(|value| *value == 0.0), + "the current F16 square/mean path should expose its overflow signature" + ); + + let fast = mlxcel_core::fast_rms_norm(&input, &weight, 1e-6); + let fast = host_f32_diagnostic_snapshot(&fast); + assert!( + fast.iter().any(|value| value.is_finite() && *value != 0.0), + "the production fast RMSNorm must accumulate F16 squares in F32" + ); + + let promoted = mlxcel_core::astype(&input, mlxcel_core::dtype::FLOAT32); + let promoted_weight = mlxcel_core::astype(&weight, mlxcel_core::dtype::FLOAT32); + let promoted = mlxcel_core::rms_norm(&promoted, &promoted_weight, 1e-6); + let promoted = host_f32_diagnostic_snapshot(&promoted); + assert!( + promoted + .iter() + .any(|value| value.is_finite() && *value != 0.0), + "F32 square/mean must preserve a non-zero normalized row" + ); + for (index, (fast, promoted)) in fast.iter().zip(promoted).enumerate() { + assert!( + (*fast - promoted).abs() <= 2.0e-3, + "production fast RMSNorm diverged from F32 accumulation at {index}: fast={fast}, promoted={promoted}" + ); + } + } +} diff --git a/src/vision/qwen2_5_vl.rs b/src/vision/qwen2_5_vl.rs index 97317b045..29b07e584 100644 --- a/src/vision/qwen2_5_vl.rs +++ b/src/vision/qwen2_5_vl.rs @@ -115,19 +115,35 @@ impl Qwen25VLModel { ); // Compute MRoPE position IDs (same logic as Qwen2-VL) + let (position_ids, rope_deltas) = self.compute_rope_state(input_ids, grid_thw); + + self.text_model.set_mrope_state(position_ids, rope_deltas); + + merged + } + + #[cfg(feature = "xla-diagnostics")] + pub fn mrope_diagnostics( + &self, + input_ids: &MlxArray, + grid_thw: &[(i32, i32, i32)], + ) -> (UniquePtr, i32) { + self.compute_rope_state(input_ids, grid_thw) + } + + fn compute_rope_state( + &self, + input_ids: &MlxArray, + grid_thw: &[(i32, i32, i32)], + ) -> (UniquePtr, i32) { let position_ids = self.compute_rope_index(input_ids, grid_thw); let ids_shape = mlxcel_core::array_shape(input_ids); let seq_len = ids_shape[1]; - mlxcel_core::eval(&position_ids); let max_pos = mlxcel_core::max_all(&position_ids); mlxcel_core::eval(&max_pos); let max_pos_val = mlxcel_core::item_i32(&max_pos); - let rope_deltas = max_pos_val + 1 - seq_len; - - self.text_model.set_mrope_state(position_ids, rope_deltas); - - merged + (position_ids, max_pos_val + 1 - seq_len) } /// Compute 3D position IDs [T, H, W] for mixed text+image sequences