From 4b044cc8c613a667042290c48368c4994f062105 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 24 Jul 2026 19:38:33 +0900 Subject: [PATCH 01/20] feat: add Qwen2.5-VL OpenXLA vision path Implement the Qwen2.5-VL windowed vision graph, checkpoint schema, and resident IREE runtime while preserving the shared Qwen-VL prefill contract. Reorder windowed patch groups and rotary positions, isolate local attention windows, restore merger output order, and reject unsupported temporal grids explicitly. Also fix the eager encoder restoration permutation used by non-self-inverse window layouts. Validation: cargo test -p mlxcel-xla qwen2_vl --lib; cargo test -p mlxcel-xla parses_and_validates_explicit_mrope_sections --lib; cargo test --lib restoration_is_the_true_inverse_for_non_self_inverse_permutation; cargo check -p mlxcel-xla --features iree --lib; cargo check --features xla-iree --lib. Refs #866 --- src/lib/mlxcel-xla/src/emitter/config.rs | 12 +- src/lib/mlxcel-xla/src/emitter/mod.rs | 3 +- src/lib/mlxcel-xla/src/emitter/qwen2_5_vl.rs | 301 ++++++++++++ src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs | 484 ++++++++++++++++--- src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs | 95 +++- src/loading/mod.rs | 6 +- src/loading/vlm.rs | 6 +- src/loading/vlm_qwen.rs | 101 ++++ src/multimodal/host_preprocessor.rs | 45 +- src/vision/encoders/qwen2_5_vl.rs | 49 +- 10 files changed, 984 insertions(+), 118 deletions(-) create mode 100644 src/lib/mlxcel-xla/src/emitter/qwen2_5_vl.rs 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..e268680b0 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; @@ -126,7 +127,7 @@ pub(crate) use qwen2_vl::emit_qwen2_vl; #[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, }; // 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. 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..f3c547561 --- /dev/null +++ b/src/lib/mlxcel-xla/src/emitter/qwen2_5_vl.rs @@ -0,0 +1,301 @@ +// 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, Ty, Val}; +use super::qwen2_vl::{ + QWEN2_VL_PATCH_BUCKETS, Qwen2VlConfig, Qwen2VlWeightSpec, QwenVlVisionVariant, +}; + +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, +) -> Val { + 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 q = apply_vision_rope(builder, &q, freqs); + let k = apply_vision_rope(builder, &k, freqs); + let q = builder.transpose(&q, &[1, 0, 2]); + let k = builder.transpose(&k, &[1, 0, 2]); + let v = builder.transpose(&v, &[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]); + linear_2d(builder, &context, &args.take(), &args.take()) +} + +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); + 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) +} + +pub(super) fn emit_qwen2_5_vl(config: &Qwen2VlConfig, patch_bucket: usize) -> 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(); + 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); + for layer in 0..config.depth { + let bias = if full_attention_blocks.contains(&layer) { + &full_bias + } else { + &window_bias + }; + hidden = encoder_layer(&mut builder, &hidden, &mut args, config, &freqs, bias); + } + 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"); + 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, + ) +} diff --git a/src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs b/src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs index 9065a41c8..3c1a7186a 100644 --- a/src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs +++ b/src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs @@ -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}" @@ -232,7 +283,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 +387,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 +417,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 +508,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 +539,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 +577,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,11 +644,20 @@ 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; @@ -482,30 +691,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,6 +872,9 @@ 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 { + if matches!(&config.variant, QwenVlVisionVariant::Qwen25 { .. }) { + return super::qwen2_5_vl::emit_qwen2_5_vl(config, patch_bucket); + } assert!( QWEN2_VL_PATCH_BUCKETS.contains(&patch_bucket), "unqualified Qwen2-VL patch bucket" @@ -728,6 +958,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(); @@ -833,4 +1087,114 @@ 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[0 * 256 + 80], 0.0); + assert_eq!(window[0 * 256 + 80], f32::NEG_INFINITY); + assert_eq!(window[144 * 256 + 144], 0.0); + assert_eq!(window[144 * 256], f32::NEG_INFINITY); + } + + #[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) + ); + } + } } diff --git a/src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs b/src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs index 2634f24b3..8e9501d58 100644 --- a/src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs +++ b/src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs @@ -35,7 +35,8 @@ 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, + prepare_qwen2_vl_host_inputs, }; 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}; @@ -135,10 +136,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 +176,18 @@ fn processor_identity(model_dir: &Path, config: &Qwen2VlConfig) -> Result "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) @@ -629,6 +649,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 +666,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 +704,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 +727,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, }, 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..e8f334de7 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 { @@ -295,7 +302,19 @@ pub struct Qwen2VlIreeHostPreprocessor { #[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:?}"), + }), + } } #[allow(clippy::too_many_arguments)] @@ -344,7 +363,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 +378,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..d048a3a6d 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, @@ -741,17 +757,10 @@ impl Qwen25VLVisionEncoder { // Merge patches h = self.merger.forward(&h); - // 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 @@ -769,3 +778,21 @@ impl super::VisionEncoder for Qwen25VLVisionEncoder { panic!("Qwen2.5-VL vision encoder requires grid_thw; use forward_with_grid() instead"); } } + +#[cfg(test)] +mod tests { + use super::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"]); + } +} From 0a3be9545be6c49e9c8941ebadcf72ff93ee8a3b Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 24 Jul 2026 19:47:17 +0900 Subject: [PATCH 02/20] fix: honor unquantized Qwen2.5 vision towers Keep the root quantization contract for text embeddings while excluding it from the resident vision contract when vision_config.skip_vision=true. This matches the pinned Qwen2.5-VL checkpoint, whose vision tensors have no affine metadata, and prevents its non-group-aligned intermediate size from being rejected without changing quantized Qwen2-VL behavior. Validation: cargo fmt --all -- --check; cargo test -p mlxcel-xla qwen2_vl --lib; cargo check -p mlxcel-xla --features iree --lib; cargo clippy -p mlxcel-xla --features iree --lib --fix --allow-dirty. Refs #866 --- src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs | 73 +++++++++++++++++++++- 1 file changed, 71 insertions(+), 2 deletions(-) diff --git a/src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs b/src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs index 3c1a7186a..decc9cfdf 100644 --- a/src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs +++ b/src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs @@ -260,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 @@ -664,7 +678,7 @@ pub(crate) fn prepare_qwen2_vl_host_inputs( 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" )); @@ -1197,4 +1211,59 @@ mod tests { ); } } + + #[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))") + ); + } } From e89057a44288781cdf554ae6d2348e35600d7b7e Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Sun, 26 Jul 2026 22:11:23 +0900 Subject: [PATCH 03/20] test(xla): add Qwen2.5-VL strict oracle Compare the eager MLX and resident IREE Qwen2.5-VL window plan, reordered patch embeddings, representative window attention, every configured full-attention layer, merger restoration, prepared M-RoPE, and final logits from one production processor payload. Add diagnostics-only negative runs that force Qwen2-style full attention, identity permutation, and zero vision positions while leaving the production graph and serving outputs unchanged. Validation: cargo fmt --all -- --check; cargo check --example xla_qwen25_vl_window_check --features xla-diagnostics; cargo test -p mlxcel-xla --lib --features diagnostics qwen2; cargo check --features xla-iree --lib; cargo test --lib restoration_is_the_true_inverse_for_non_self_inverse_permutation. Refs #866 --- Cargo.toml | 5 + examples/xla_qwen25_vl_window_check.rs | 584 +++++++++++++++++++ src/lib.rs | 6 +- src/lib/mlxcel-xla/src/batch.rs | 30 + src/lib/mlxcel-xla/src/emitter/mod.rs | 2 + src/lib/mlxcel-xla/src/emitter/qwen2_5_vl.rs | 101 +++- src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs | 21 + src/lib/mlxcel-xla/src/lib.rs | 9 +- src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs | 334 +++++++++++ src/multimodal/host_preprocessor.rs | 24 + src/vision/encoders/qwen2_5_vl.rs | 103 +++- src/vision/qwen2_5_vl.rs | 28 +- 12 files changed, 1235 insertions(+), 12 deletions(-) create mode 100644 examples/xla_qwen25_vl_window_check.rs 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..742df8ec5 --- /dev/null +++ b/examples/xla_qwen25_vl_window_check.rs @@ -0,0 +1,584 @@ +// 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_relative: f32, + 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() +} + +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_relative: 0.0, + failures: 0, + non_finite: 0, + }; + for (&observed, &reference) in actual.iter().zip(expected) { + 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.max_absolute = stats.max_absolute.max(absolute); + stats.max_relative = stats.max_relative.max(relative); + if absolute > tolerance.atol + tolerance.rtol * reference.abs() { + stats.failures += 1; + } + } + 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_relative": stats.max_relative, + "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 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 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}")); + + 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 + .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(); + compare_stage( + &mut reports, + "vision.reordered_patch_embedding", + &iree_vision.reordered_patch_embedding, + &mlx_f32(&eager_vision.reordered_patch_embedding), + tolerance, + ) + .unwrap_or_else(|error| panic!("{error}")); + assert_eq!( + iree_vision.window_layer_index, eager_vision.window_layer_index, + "representative window layer differs" + ); + compare_stage( + &mut reports, + &format!( + "vision.post_window_layer_{}", + iree_vision.window_layer_index + ), + &iree_vision.post_window_layer, + &mlx_f32(&eager_vision.post_window_layer), + tolerance, + ) + .unwrap_or_else(|error| panic!("{error}")); + 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" + ); + 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() + { + compare_stage( + &mut reports, + &format!("vision.post_full_layer_{layer}_capture_{capture_index}"), + actual, + &mlx_f32(expected), + tolerance, + ) + .unwrap_or_else(|error| panic!("{error}")); + } + compare_stage( + &mut reports, + "vision.merger_window_ordered", + &iree_vision.merger_window_ordered, + &mlx_f32(&eager_vision.merger_window_ordered), + tolerance, + ) + .unwrap_or_else(|error| panic!("{error}")); + compare_stage( + &mut reports, + "vision.restored_projection", + &iree_vision.restored_projection, + &mlx_f32(&eager_vision.restored_projection), + tolerance, + ) + .unwrap_or_else(|error| panic!("{error}")); + + let qwen2_full = iree_vision_capture( + &mut iree_projector, + &processor.patch_values, + &processor.grids, + Qwen25VlDiagnosticMutation::Qwen2FullAttention, + ); + assert!( + compare_stage( + &mut Vec::new(), + "negative.qwen2_full_attention", + &qwen2_full.post_window_layer, + &mlx_f32(&eager_vision.post_window_layer), + tolerance, + ) + .is_err(), + "negative oracle accepted Qwen2-style full attention in a window layer" + ); + 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": 1, + "model": model_path, + "images": image_paths, + "device": device, + "context_capacity": context_capacity, + "grid_thw": processor.grids, + "patch_tokens": iree_vision.patch_tokens, + "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, + "negative_qwen2_full_attention_detected": true, + "negative_identity_permutation_detected": true, + "negative_zero_vision_positions_detected": true, + "final_top1": argmax(&iree_logits), + "comparisons": reports, + "passed": true, + }); + let rendered = serde_json::to_string_pretty(&report).expect("serialize oracle report"); + if let Some(path) = report_path { + fs::write(&path, &rendered) + .unwrap_or_else(|error| panic!("write {}: {error}", path.display())); + } + println!("{rendered}"); +} + +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/mod.rs b/src/lib/mlxcel-xla/src/emitter/mod.rs index e268680b0..22293fefb 100644 --- a/src/lib/mlxcel-xla/src/emitter/mod.rs +++ b/src/lib/mlxcel-xla/src/emitter/mod.rs @@ -122,6 +122,8 @@ 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, }; +#[cfg(feature = "diagnostics")] +pub(crate) use qwen2_5_vl::{Qwen25VlVisionDiagnosticLayout, emit_qwen2_5_vl_diagnostics}; #[allow(unused_imports)] pub(crate) use qwen2_vl::emit_qwen2_vl; #[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 index f3c547561..39910375d 100644 --- a/src/lib/mlxcel-xla/src/emitter/qwen2_5_vl.rs +++ b/src/lib/mlxcel-xla/src/emitter/qwen2_5_vl.rs @@ -24,6 +24,13 @@ 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, +} + struct Args { values: Vec, declarations: Vec, @@ -237,7 +244,11 @@ fn encoder_layer( builder.add(&residual, &down) } -pub(super) fn emit_qwen2_5_vl(config: &Qwen2VlConfig, patch_bucket: usize) -> String { +fn emit_qwen2_5_vl_inner( + config: &Qwen2VlConfig, + patch_bucket: usize, + #[cfg(feature = "diagnostics")] diagnostic_layout: Option<&Qwen25VlVisionDiagnosticLayout>, +) -> String { assert!( QWEN2_VL_PATCH_BUCKETS.contains(&patch_bucket), "unqualified Qwen2.5-VL patch bucket" @@ -270,6 +281,12 @@ pub(super) fn emit_qwen2_5_vl(config: &Qwen2VlConfig, patch_bucket: usize) -> St "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(); for layer in 0..config.depth { let bias = if full_attention_blocks.contains(&layer) { &full_bias @@ -277,6 +294,15 @@ pub(super) fn emit_qwen2_5_vl(config: &Qwen2VlConfig, patch_bucket: usize) -> St &window_bias }; 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()); + } + } } 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; @@ -291,6 +317,35 @@ pub(super) fn emit_qwen2_5_vl(config: &Qwen2VlConfig, patch_bucket: usize) -> St 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); + 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(", "), @@ -299,3 +354,47 @@ pub(super) fn emit_qwen2_5_vl(config: &Qwen2VlConfig, patch_bucket: usize) -> St result = projected.name, ) } + +pub(super) fn emit_qwen2_5_vl(config: &Qwen2VlConfig, patch_bucket: usize) -> String { + emit_qwen2_5_vl_inner( + config, + patch_bucket, + #[cfg(feature = "diagnostics")] + None, + ) +} + +#[cfg(feature = "diagnostics")] +pub(crate) fn emit_qwen2_5_vl_diagnostics( + config: &Qwen2VlConfig, + patch_bucket: usize, +) -> 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 layout = Qwen25VlVisionDiagnosticLayout { + window_layer_index, + full_layer_indices: (0..config.depth) + .filter(|layer| full_attention_blocks.contains(layer)) + .collect(), + }; + Ok(( + emit_qwen2_5_vl_inner(config, patch_bucket, Some(&layout)), + layout, + )) +} diff --git a/src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs b/src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs index decc9cfdf..c6009bc5f 100644 --- a/src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs +++ b/src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs @@ -1071,6 +1071,27 @@ mod tests { assert_eq!(mlir.matches("chlo.erf").count(), 1); } + #[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).unwrap(); + assert_eq!(layout.window_layer_index, 0); + assert_eq!(layout.full_layer_indices, vec![1]); + assert_eq!( + production.matches("stablehlo.dot_general").count(), + diagnostics.matches("stablehlo.dot_general").count() + ); + assert!(diagnostics.contains( + "-> (tensor<16x1280xf32>, tensor<16x1280xf32>, tensor<16x1280xf32>, tensor<4x1536xf32>)" + )); + 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(); 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 8e9501d58..77c7a7f74 100644 --- a/src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs +++ b/src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs @@ -38,6 +38,8 @@ use crate::emitter::{ Qwen2VlConfig, Qwen2VlHostInputs, QwenVlVisionVariant, emit_qwen2_vl, prepare_qwen2_vl_host_inputs, }; +#[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}; @@ -60,11 +62,47 @@ 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, 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 merger_window_ordered: Vec, + pub restored_projection: Vec, + pub patch_tokens: usize, + pub merged_tokens: usize, + pub vision_hidden: usize, + pub text_hidden: usize, +} + struct ActiveModule { patch_bucket: usize, module: IreeAuxiliaryModule, } +#[cfg(feature = "diagnostics")] +struct ActiveDiagnosticModule { + patch_bucket: usize, + module: IreeAuxiliaryModule, + layout: Qwen25VlVisionDiagnosticLayout, +} + pub struct IreeQwen2VlProjector { model_dir: PathBuf, device: String, @@ -73,6 +111,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 { @@ -595,6 +642,45 @@ 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 (mlir, layout) = emit_qwen2_5_vl_diagnostics(config, patch_bucket)?; + 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( + ENTRY_NAME, + format!( + "{};{};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)?; @@ -738,6 +824,246 @@ 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; + 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 state_shape = [plan.patch_bucket, self.config.hidden]; + let merger_shape = [plan.patch_bucket / merge_unit, self.config.text_hidden]; + let state_bytes = state_shape.iter().product::() * std::mem::size_of::(); + let merger_bytes = merger_shape.iter().product::() * std::mem::size_of::(); + let full_layer_count = match &self.config.variant { + QwenVlVisionVariant::Qwen25 { + full_attention_blocks, + .. + } => full_attention_blocks.len(), + QwenVlVisionVariant::Qwen2 => 0, + }; + let mut output_buffers = vec![vec![0u8; state_bytes]; 2 + full_layer_count]; + output_buffers.push(vec![0u8; merger_bytes]); + 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() + .enumerate() + .map(|(index, bytes)| AuxiliaryOutput { + bytes, + dtype: AuxiliaryTensorDType::Float32, + shape: if index + 1 == 3 + full_layer_count { + &merger_shape + } else { + &state_shape + }, + }) + .collect::>(); + let active = self.ensure_bucket(plan.patch_bucket)?; + active.module.invoke(&inputs, &mut outputs)?; + let layout = active.layout.clone(); + let mut decoded = output_buffers + .into_iter() + .enumerate() + .map(|(index, bytes)| { + checked_f32_output(&format!("Qwen2.5-VL diagnostic output {index}"), bytes) + }) + .collect::, _>>()?; + let actual_state_values = plan + .actual_patches + .checked_mul(self.config.hidden) + .ok_or_else(|| "Qwen2.5-VL diagnostic state size overflowed".to_string())?; + for state in decoded.iter_mut().take(2 + full_layer_count) { + state.truncate(actual_state_values); + } + 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); + 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: decoded, + merger_window_ordered, + restored_projection, + patch_tokens: plan.actual_patches, + merged_tokens: actual_tokens, + vision_hidden: self.config.hidden, + text_hidden: self.config.text_hidden, + }) + } +} + #[cfg(test)] mod tests { use super::*; @@ -805,4 +1131,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/multimodal/host_preprocessor.rs b/src/multimodal/host_preprocessor.rs index e8f334de7..bd0461082 100644 --- a/src/multimodal/host_preprocessor.rs +++ b/src/multimodal/host_preprocessor.rs @@ -299,6 +299,13 @@ 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 { @@ -317,6 +324,23 @@ impl Qwen2VlIreeHostPreprocessor { } } + #[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)] pub(crate) fn from_parts( processor: crate::vision::processors::qwen2_vl::Qwen2VLProcessor, diff --git a/src/vision/encoders/qwen2_5_vl.rs b/src/vision/encoders/qwen2_5_vl.rs index d048a3a6d..d55326e63 100644 --- a/src/vision/encoders/qwen2_5_vl.rs +++ b/src/vision/encoders/qwen2_5_vl.rs @@ -465,6 +465,26 @@ 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 merger_window_ordered: UniquePtr, + pub restored_projection: UniquePtr, +} + +#[cfg(feature = "xla-diagnostics")] +type Qwen25VLCapture = Option; +#[cfg(not(feature = "xla-diagnostics"))] +type Qwen25VLCapture = (); + impl Qwen25VLVisionEncoder { pub fn from_weights( weights: &WeightMap, @@ -712,6 +732,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); @@ -730,6 +770,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); @@ -745,6 +788,13 @@ 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(); for (layer_num, block) in self.blocks.iter().enumerate() { let cu_seqlens_now = if self.fullatt_block_indexes.contains(&layer_num) { &cu_seqlens @@ -752,10 +802,27 @@ impl Qwen25VLVisionEncoder { &cu_window_seqlens }; 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"), + )); + } + } } // 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: destination original_position reads its corresponding // window_position. This is the inverse permutation, not window_index @@ -768,7 +835,41 @@ 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 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: (0..self.blocks.len()) + .filter(|layer| self.fullatt_block_indexes.contains(layer)) + .collect(), + post_full_layers, + 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) + } } } 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 From 572ffc28d8b6766b52e855feb9defd63800f5fd1 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Sun, 26 Jul 2026 23:57:27 +0900 Subject: [PATCH 04/20] fix(xla): align Qwen vision GPU precision Route Qwen-VL auxiliary graphs through the device precision resolver so CUDA demotes contractions to f16 while retaining f32 accumulation and sensitive math. Include the resolved precision in the artifact identity. Expand Qwen2.5-VL diagnostics across layers 24-31 and the final full-attention block substages, and emit complete RMS/error reports before failing the strict oracle. --- examples/xla_qwen25_vl_window_check.rs | 218 +++++++++++++++++-- src/lib/mlxcel-xla/src/emitter/mod.rs | 4 +- src/lib/mlxcel-xla/src/emitter/qwen2_5_vl.rs | 125 ++++++++++- src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs | 53 ++++- src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs | 80 +++++-- src/vision/encoders/qwen2_5_vl.rs | 103 ++++++++- 6 files changed, 522 insertions(+), 61 deletions(-) diff --git a/examples/xla_qwen25_vl_window_check.rs b/examples/xla_qwen25_vl_window_check.rs index 742df8ec5..9c72def36 100644 --- a/examples/xla_qwen25_vl_window_check.rs +++ b/examples/xla_qwen25_vl_window_check.rs @@ -44,6 +44,9 @@ struct Tolerance { struct ComparisonStats { max_absolute: f32, max_relative: f32, + actual_rms: f64, + expected_rms: f64, + error_rms: f64, failures: usize, non_finite: usize, } @@ -149,9 +152,13 @@ fn comparison_stats( let mut stats = ComparisonStats { max_absolute: 0.0, 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 (&observed, &reference) in actual.iter().zip(expected) { if !observed.is_finite() || !reference.is_finite() { stats.failures += 1; @@ -160,12 +167,22 @@ fn comparison_stats( } 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; stats.max_absolute = stats.max_absolute.max(absolute); 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) } @@ -185,6 +202,9 @@ fn compare_stage( "rtol": tolerance.rtol, "max_absolute": stats.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, @@ -199,6 +219,28 @@ fn compare_stage( } } +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"), @@ -351,20 +393,22 @@ fn main() { ); let mut reports = Vec::new(); - compare_stage( + 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, - ) - .unwrap_or_else(|error| panic!("{error}")); + ); assert_eq!( iree_vision.window_layer_index, eager_vision.window_layer_index, "representative window layer differs" ); - compare_stage( + record_stage( &mut reports, + &mut comparison_failures, &format!( "vision.post_window_layer_{}", iree_vision.window_layer_index @@ -372,8 +416,7 @@ fn main() { &iree_vision.post_window_layer, &mlx_f32(&eager_vision.post_window_layer), tolerance, - ) - .unwrap_or_else(|error| panic!("{error}")); + ); assert_eq!( iree_vision.full_layer_indices, eager_vision.full_layer_indices, "configured full-attention layers differ" @@ -383,6 +426,20 @@ fn main() { 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.target_full_layer_index, eager_vision.target_full_layer_index, + "target full-attention layer differs" + ); + let target_full_layer_index = iree_vision.target_full_layer_index; for (capture_index, (&layer, (actual, expected))) in iree_vision .full_layer_indices .iter() @@ -393,32 +450,154 @@ fn main() { .zip(&eager_vision.post_full_layers), ) .enumerate() + .filter(|&(_, (&layer, _))| layer != target_full_layer_index) { - compare_stage( + 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), ) - .unwrap_or_else(|error| panic!("{error}")); + .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, + ); } - compare_stage( + for (stage, actual, expected) in [ + ( + "input", + iree_vision.target_full_layer_input.as_slice(), + eager_vision + .target_full_layer_input + .as_ref() + .expect("eager target layer input"), + ), + ( + "norm1", + iree_vision.target_full_layer_norm1.as_slice(), + eager_vision + .target_full_layer_norm1 + .as_ref() + .expect("eager target layer norm1"), + ), + ( + "attention", + iree_vision.target_full_layer_attention.as_slice(), + eager_vision + .target_full_layer_attention + .as_ref() + .expect("eager target layer attention"), + ), + ( + "post_attention_residual", + iree_vision + .target_full_layer_post_attention_residual + .as_slice(), + eager_vision + .target_full_layer_post_attention_residual + .as_ref() + .expect("eager target layer post-attention residual"), + ), + ( + "norm2", + iree_vision.target_full_layer_norm2.as_slice(), + eager_vision + .target_full_layer_norm2 + .as_ref() + .expect("eager target layer norm2"), + ), + ( + "mlp", + iree_vision.target_full_layer_mlp.as_slice(), + eager_vision + .target_full_layer_mlp + .as_ref() + .expect("eager target layer MLP"), + ), + ] { + record_stage( + &mut reports, + &mut comparison_failures, + &format!("vision.layer_{target_full_layer_index}.{stage}"), + actual, + &mlx_f32(expected), + 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, - ) - .unwrap_or_else(|error| panic!("{error}")); - compare_stage( + ); + record_stage( &mut reports, + &mut comparison_failures, "vision.restored_projection", &iree_vision.restored_projection, &mlx_f32(&eager_vision.restored_projection), tolerance, - ) - .unwrap_or_else(|error| panic!("{error}")); + ); + 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_tokens": iree_vision.patch_tokens, + "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, + "target_full_layer": target_full_layer_index, + "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, @@ -545,7 +724,7 @@ fn main() { .unwrap_or_else(|error| panic!("{error}")); let report = json!({ - "schema": 1, + "schema": 2, "model": model_path, "images": image_paths, "device": device, @@ -557,6 +736,8 @@ fn main() { "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, + "target_full_layer": iree_vision.target_full_layer_index, "negative_qwen2_full_attention_detected": true, "negative_identity_permutation_detected": true, "negative_zero_vision_positions_detected": true, @@ -564,12 +745,7 @@ fn main() { "comparisons": reports, "passed": true, }); - let rendered = serde_json::to_string_pretty(&report).expect("serialize oracle report"); - if let Some(path) = report_path { - fs::write(&path, &rendered) - .unwrap_or_else(|error| panic!("write {}: {error}", path.display())); - } - println!("{rendered}"); + emit_report(report_path.as_deref(), &report); } fn iree_vision_capture( diff --git a/src/lib/mlxcel-xla/src/emitter/mod.rs b/src/lib/mlxcel-xla/src/emitter/mod.rs index 22293fefb..715600c50 100644 --- a/src/lib/mlxcel-xla/src/emitter/mod.rs +++ b/src/lib/mlxcel-xla/src/emitter/mod.rs @@ -125,12 +125,12 @@ pub(crate) use phi4_audio::{ #[cfg(feature = "diagnostics")] pub(crate) use qwen2_5_vl::{Qwen25VlVisionDiagnosticLayout, emit_qwen2_5_vl_diagnostics}; #[allow(unused_imports)] -pub(crate) use qwen2_vl::emit_qwen2_vl; -#[allow(unused_imports)] pub(crate) use qwen2_vl::{ QWEN2_VL_PATCH_BUCKETS, Qwen2VlConfig, Qwen2VlGridPlan, Qwen2VlHostInputs, Qwen2VlWeightSpec, 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 index 39910375d..ea7cc180d 100644 --- a/src/lib/mlxcel-xla/src/emitter/qwen2_5_vl.rs +++ b/src/lib/mlxcel-xla/src/emitter/qwen2_5_vl.rs @@ -19,7 +19,7 @@ //! static for one bucket while selecting the checkpoint-declared full //! attention layers during emission. -use super::builder::{Builder, Ty, Val}; +use super::builder::{Builder, Precision, Ty, Val}; use super::qwen2_vl::{ QWEN2_VL_PATCH_BUCKETS, Qwen2VlConfig, Qwen2VlWeightSpec, QwenVlVisionVariant, }; @@ -29,6 +29,19 @@ use super::qwen2_vl::{ 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) target_full_layer_index: usize, +} + +#[cfg(feature = "diagnostics")] +struct Qwen25VlBlockDiagnostics { + input: Val, + norm1: Val, + attention: Val, + post_attention_residual: Val, + norm2: Val, + mlp: Val, + output: Val, } struct Args { @@ -244,9 +257,45 @@ fn encoder_layer( 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); + let norm2 = rms_norm( + builder, + &post_attention_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 mlp = linear_2d(builder, &activated, &args.take(), &args.take()); + let output = builder.add(&post_attention_residual, &mlp); + Qwen25VlBlockDiagnostics { + input: hidden.clone(), + norm1, + attention, + post_attention_residual, + norm2, + mlp, + output, + } +} + fn emit_qwen2_5_vl_inner( config: &Qwen2VlConfig, patch_bucket: usize, + precision: Precision, #[cfg(feature = "diagnostics")] diagnostic_layout: Option<&Qwen25VlVisionDiagnosticLayout>, ) -> String { assert!( @@ -262,7 +311,7 @@ fn emit_qwen2_5_vl_inner( }; 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; @@ -287,13 +336,38 @@ fn emit_qwen2_5_vl_inner( 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 target_full_layer_state = None; for layer in 0..config.depth { let bias = if full_attention_blocks.contains(&layer) { &full_bias } else { &window_bias }; - hidden = encoder_layer(&mut builder, &hidden, &mut args, config, &freqs, bias); + #[cfg(feature = "diagnostics")] + let is_target_full_layer = + diagnostic_layout.is_some_and(|layout| layer == layout.target_full_layer_index); + #[cfg(feature = "diagnostics")] + if is_target_full_layer { + let capture = encoder_layer_with_diagnostics( + &mut builder, + &hidden, + &mut args, + config, + &freqs, + bias, + ); + hidden = capture.output.clone(); + target_full_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 { @@ -302,6 +376,9 @@ fn emit_qwen2_5_vl_inner( 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()); + } } } hidden = rms_norm(&mut builder, &hidden, &args.take(), config.layer_norm_eps); @@ -329,6 +406,21 @@ fn emit_qwen2_5_vl_inner( "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); + let target = target_full_layer_state.expect("diagnostics capture target full layer"); + outputs.extend([ + target.input, + target.norm1, + target.attention, + target.post_attention_residual, + target.norm2, + target.mlp, + ]); outputs.push(projected); let result_types = outputs .iter() @@ -355,10 +447,15 @@ fn emit_qwen2_5_vl_inner( ) } -pub(super) fn emit_qwen2_5_vl(config: &Qwen2VlConfig, patch_bucket: usize) -> String { +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, ) @@ -368,6 +465,7 @@ pub(super) fn emit_qwen2_5_vl(config: &Qwen2VlConfig, patch_bucket: usize) -> St 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, @@ -387,14 +485,25 @@ pub(crate) fn emit_qwen2_5_vl_diagnostics( .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 layout = Qwen25VlVisionDiagnosticLayout { window_layer_index, - full_layer_indices: (0..config.depth) - .filter(|layer| full_attention_blocks.contains(layer)) - .collect(), + full_layer_indices, + final_interval_layer_indices: (final_interval_start..=target_full_layer_index).collect(), + target_full_layer_index, }; Ok(( - emit_qwen2_5_vl_inner(config, patch_bucket, Some(&layout)), + emit_qwen2_5_vl_inner(config, patch_bucket, precision, Some(&layout)), layout, )) } diff --git a/src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs b/src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs index c6009bc5f..d2c232a2b 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. @@ -886,8 +886,16 @@ 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); + return super::qwen2_5_vl::emit_qwen2_5_vl(config, patch_bucket, precision); } assert!( QWEN2_VL_PATCH_BUCKETS.contains(&patch_bucket), @@ -895,7 +903,7 @@ pub(crate) fn emit_qwen2_vl(config: &Qwen2VlConfig, patch_bucket: usize) -> Stri ); 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; @@ -1071,22 +1079,53 @@ 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).unwrap(); + 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.target_full_layer_index, 1); assert_eq!( production.matches("stablehlo.dot_general").count(), diagnostics.matches("stablehlo.dot_general").count() ); - assert!(diagnostics.contains( - "-> (tensor<16x1280xf32>, tensor<16x1280xf32>, tensor<16x1280xf32>, tensor<4x1536xf32>)" - )); + let result_signature = std::iter::repeat_n("tensor<16x1280xf32>", 11) + .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\")")); diff --git a/src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs b/src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs index 77c7a7f74..dedd02d10 100644 --- a/src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs +++ b/src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs @@ -35,8 +35,8 @@ use crate::aux::{ }; use crate::aux_manifest::{AuxiliaryArtifactContract, ensure_qualified_auxiliary_artifact}; use crate::emitter::{ - Qwen2VlConfig, Qwen2VlHostInputs, QwenVlVisionVariant, 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}; @@ -83,6 +83,15 @@ pub struct Qwen25VlVisionDiagnostics { 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 target_full_layer_index: usize, + pub target_full_layer_input: Vec, + pub target_full_layer_norm1: Vec, + pub target_full_layer_attention: Vec, + pub target_full_layer_post_attention_residual: Vec, + pub target_full_layer_norm2: Vec, + pub target_full_layer_mlp: Vec, pub merger_window_ordered: Vec, pub restored_projection: Vec, pub patch_tokens: usize, @@ -610,7 +619,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())); @@ -623,7 +633,7 @@ 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()) @@ -650,7 +660,8 @@ fn compile_and_load_diagnostics( processor_identity: &str, patch_bucket: usize, ) -> Result<(IreeAuxiliaryModule, Qwen25VlVisionDiagnosticLayout), String> { - let (mlir, layout) = emit_qwen2_5_vl_diagnostics(config, patch_bucket)?; + 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())); @@ -663,7 +674,7 @@ fn compile_and_load_diagnostics( let contract = AuxiliaryArtifactContract::new( ENTRY_NAME, format!( - "{};{};diagnostics=window-full-restoration;checkpoint_schema_sha256={}", + "{};{};precision={precision:?};diagnostics=window-full-restoration;checkpoint_schema_sha256={}", config.fingerprint(patch_bucket), processor_identity, sha256_hex(checkpoint_schema.as_bytes()) @@ -955,14 +966,14 @@ impl IreeQwen25VlDiagnosticProjector { let merger_shape = [plan.patch_bucket / merge_unit, self.config.text_hidden]; let state_bytes = state_shape.iter().product::() * std::mem::size_of::(); let merger_bytes = merger_shape.iter().product::() * std::mem::size_of::(); - let full_layer_count = match &self.config.variant { - QwenVlVisionVariant::Qwen25 { - full_attention_blocks, - .. - } => full_attention_blocks.len(), - QwenVlVisionVariant::Qwen2 => 0, - }; - let mut output_buffers = vec![vec![0u8; state_bytes]; 2 + full_layer_count]; + 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 block_substage_count = 6; + let state_output_count = + 2 + full_layer_count + final_interval_layer_count + block_substage_count; + let mut output_buffers = vec![vec![0u8; state_bytes]; state_output_count]; output_buffers.push(vec![0u8; merger_bytes]); let inputs = [ AuxiliaryInput { @@ -992,16 +1003,14 @@ impl IreeQwen25VlDiagnosticProjector { .map(|(index, bytes)| AuxiliaryOutput { bytes, dtype: AuxiliaryTensorDType::Float32, - shape: if index + 1 == 3 + full_layer_count { + shape: if index == state_output_count { &merger_shape } else { &state_shape }, }) .collect::>(); - let active = self.ensure_bucket(plan.patch_bucket)?; active.module.invoke(&inputs, &mut outputs)?; - let layout = active.layout.clone(); let mut decoded = output_buffers .into_iter() .enumerate() @@ -1013,7 +1022,7 @@ impl IreeQwen25VlDiagnosticProjector { .actual_patches .checked_mul(self.config.hidden) .ok_or_else(|| "Qwen2.5-VL diagnostic state size overflowed".to_string())?; - for state in decoded.iter_mut().take(2 + full_layer_count) { + for state in decoded.iter_mut().take(state_output_count) { state.truncate(actual_state_values); } let all_projected = decoded @@ -1042,6 +1051,30 @@ impl IreeQwen25VlDiagnosticProjector { } 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 mut target_full_layer_substages = decoded.into_iter(); + let target_full_layer_input = target_full_layer_substages + .next() + .expect("diagnostic target full layer input output"); + let target_full_layer_norm1 = target_full_layer_substages + .next() + .expect("diagnostic target full layer norm1 output"); + let target_full_layer_attention = target_full_layer_substages + .next() + .expect("diagnostic target full layer attention output"); + let target_full_layer_post_attention_residual = target_full_layer_substages + .next() + .expect("diagnostic target full layer residual output"); + let target_full_layer_norm2 = target_full_layer_substages + .next() + .expect("diagnostic target full layer norm2 output"); + let target_full_layer_mlp = target_full_layer_substages + .next() + .expect("diagnostic target full layer MLP output"); + debug_assert!(target_full_layer_substages.next().is_none()); Ok(Qwen25VlVisionDiagnostics { window_index: plan.window_index, restore_indices: plan.restore_indices, @@ -1053,7 +1086,16 @@ impl IreeQwen25VlDiagnosticProjector { window_layer_index: layout.window_layer_index, post_window_layer, full_layer_indices: layout.full_layer_indices, - post_full_layers: decoded, + post_full_layers, + final_interval_layer_indices: layout.final_interval_layer_indices, + post_final_interval_layers, + target_full_layer_index: layout.target_full_layer_index, + target_full_layer_input, + target_full_layer_norm1, + target_full_layer_attention, + target_full_layer_post_attention_residual, + target_full_layer_norm2, + target_full_layer_mlp, merger_window_ordered, restored_projection, patch_tokens: plan.actual_patches, diff --git a/src/vision/encoders/qwen2_5_vl.rs b/src/vision/encoders/qwen2_5_vl.rs index d55326e63..c796215b0 100644 --- a/src/vision/encoders/qwen2_5_vl.rs +++ b/src/vision/encoders/qwen2_5_vl.rs @@ -384,6 +384,16 @@ struct VisionBlock { mlp: VisionMLP, } +#[cfg(feature = "xla-diagnostics")] +struct Qwen25VLBlockDiagnostics { + input: UniquePtr, + norm1: UniquePtr, + attention: UniquePtr, + post_attention_residual: UniquePtr, + norm2: UniquePtr, + mlp: UniquePtr, +} + impl VisionBlock { fn from_weights( weights: &WeightMap, @@ -413,6 +423,33 @@ 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) { + let input = mlxcel_core::copy(hidden_states); + let norm1 = self.norm1.forward(hidden_states); + let attention = self.attn.forward(&norm1, cu_seqlens, rotary_pos_emb); + let post_attention_residual = mlxcel_core::add(hidden_states, &attention); + let norm2 = self.norm2.forward(&post_attention_residual); + let mlp = self.mlp.forward(&norm2); + let output = mlxcel_core::add(&post_attention_residual, &mlp); + ( + output, + Qwen25VLBlockDiagnostics { + input, + norm1, + attention, + post_attention_residual, + norm2, + mlp, + }, + ) + } } // PatchMerger - RMSNorm + GELU MLP (projection to text hidden size). @@ -476,6 +513,15 @@ pub struct Qwen25VLVisionDiagnostics { 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 target_full_layer_index: usize, + pub target_full_layer_input: UniquePtr, + pub target_full_layer_norm1: UniquePtr, + pub target_full_layer_attention: UniquePtr, + pub target_full_layer_post_attention_residual: UniquePtr, + pub target_full_layer_norm2: UniquePtr, + pub target_full_layer_mlp: UniquePtr, pub merger_window_ordered: UniquePtr, pub restored_projection: UniquePtr, } @@ -795,13 +841,45 @@ impl Qwen25VLVisionEncoder { 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(); + #[cfg(feature = "xla-diagnostics")] + let mut post_final_interval_layers = Vec::new(); + #[cfg(feature = "xla-diagnostics")] + let mut target_full_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) == target_full_layer_index { + let (output, capture) = + block.forward_with_diagnostics(&h, cu_seqlens_now, &rotary_pos_emb); + h = output; + target_full_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 { @@ -814,6 +892,11 @@ impl Qwen25VLVisionEncoder { 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"), + )); + } } } @@ -837,6 +920,10 @@ impl Qwen25VLVisionEncoder { #[cfg(feature = "xla-diagnostics")] if CAPTURE { + let target_full_layer_index = + target_full_layer_index.expect("Qwen2.5-VL diagnostics require a full layer"); + let target_full_layer_state = + target_full_layer_state.expect("Qwen2.5-VL target full layer capture"); let restored_projection = mlxcel_core::copy(h.as_ref().expect("Qwen2.5-VL restored projection output")); let output = VisionEncoderOutput { hidden_states: h }; @@ -852,10 +939,18 @@ impl Qwen25VLVisionEncoder { 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: (0..self.blocks.len()) - .filter(|layer| self.fullatt_block_indexes.contains(layer)) - .collect(), + full_layer_indices, post_full_layers, + final_interval_layer_indices, + post_final_interval_layers, + target_full_layer_index, + target_full_layer_input: target_full_layer_state.input, + target_full_layer_norm1: target_full_layer_state.norm1, + target_full_layer_attention: target_full_layer_state.attention, + target_full_layer_post_attention_residual: target_full_layer_state + .post_attention_residual, + target_full_layer_norm2: target_full_layer_state.norm2, + target_full_layer_mlp: target_full_layer_state.mlp, merger_window_ordered: merger_window_ordered .expect("Qwen2.5-VL merger capture"), restored_projection, From ead807e144552cfca1e8db5ab51cad08e6d86508 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 00:28:07 +0900 Subject: [PATCH 05/20] test(xla): probe Qwen layer 29 substages Move the Qwen2.5-VL block-substage diagnostic seam from the already-divergent final full-attention block to the third-to-last final-interval layer, which is layer 29 for the pinned checkpoint and the first failing boundary in the actual report. Keep final full-layer output coverage unchanged. --- examples/xla_qwen25_vl_window_check.rs | 55 ++++++++++-------- src/lib/mlxcel-xla/src/emitter/qwen2_5_vl.rs | 59 +++++++++++++++----- src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs | 2 +- src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs | 56 +++++++++---------- src/vision/encoders/qwen2_5_vl.rs | 54 +++++++++++------- 5 files changed, 139 insertions(+), 87 deletions(-) diff --git a/examples/xla_qwen25_vl_window_check.rs b/examples/xla_qwen25_vl_window_check.rs index 9c72def36..ab078fe73 100644 --- a/examples/xla_qwen25_vl_window_check.rs +++ b/examples/xla_qwen25_vl_window_check.rs @@ -435,11 +435,20 @@ fn main() { eager_vision.post_final_interval_layers.len(), "captured final diagnostic interval count differs" ); + let target_full_layer_index = *iree_vision + .full_layer_indices + .last() + .expect("IREE diagnostics require a full-attention layer"); assert_eq!( - iree_vision.target_full_layer_index, eager_vision.target_full_layer_index, + Some(&target_full_layer_index), + eager_vision.full_layer_indices.last(), "target full-attention layer differs" ); - let target_full_layer_index = iree_vision.target_full_layer_index; + 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() @@ -484,59 +493,59 @@ fn main() { for (stage, actual, expected) in [ ( "input", - iree_vision.target_full_layer_input.as_slice(), + iree_vision.substage_probe_layer_input.as_slice(), eager_vision - .target_full_layer_input + .substage_probe_layer_input .as_ref() - .expect("eager target layer input"), + .expect("eager substage probe layer input"), ), ( "norm1", - iree_vision.target_full_layer_norm1.as_slice(), + iree_vision.substage_probe_layer_norm1.as_slice(), eager_vision - .target_full_layer_norm1 + .substage_probe_layer_norm1 .as_ref() - .expect("eager target layer norm1"), + .expect("eager substage probe layer norm1"), ), ( "attention", - iree_vision.target_full_layer_attention.as_slice(), + iree_vision.substage_probe_layer_attention.as_slice(), eager_vision - .target_full_layer_attention + .substage_probe_layer_attention .as_ref() - .expect("eager target layer attention"), + .expect("eager substage probe layer attention"), ), ( "post_attention_residual", iree_vision - .target_full_layer_post_attention_residual + .substage_probe_layer_post_attention_residual .as_slice(), eager_vision - .target_full_layer_post_attention_residual + .substage_probe_layer_post_attention_residual .as_ref() - .expect("eager target layer post-attention residual"), + .expect("eager substage probe layer post-attention residual"), ), ( "norm2", - iree_vision.target_full_layer_norm2.as_slice(), + iree_vision.substage_probe_layer_norm2.as_slice(), eager_vision - .target_full_layer_norm2 + .substage_probe_layer_norm2 .as_ref() - .expect("eager target layer norm2"), + .expect("eager substage probe layer norm2"), ), ( "mlp", - iree_vision.target_full_layer_mlp.as_slice(), + iree_vision.substage_probe_layer_mlp.as_slice(), eager_vision - .target_full_layer_mlp + .substage_probe_layer_mlp .as_ref() - .expect("eager target layer MLP"), + .expect("eager substage probe layer MLP"), ), ] { record_stage( &mut reports, &mut comparison_failures, - &format!("vision.layer_{target_full_layer_index}.{stage}"), + &format!("vision.layer_{substage_probe_layer_index}.{stage}"), actual, &mlx_f32(expected), tolerance, @@ -587,6 +596,7 @@ fn main() { "full_attention_layers": iree_vision.full_layer_indices, "final_interval_layers": iree_vision.final_interval_layer_indices, "target_full_layer": target_full_layer_index, + "substage_probe_layer": substage_probe_layer_index, "failed_phase": "vision", "failures": comparison_failures, "comparisons": reports, @@ -737,7 +747,8 @@ fn main() { "window_layer": iree_vision.window_layer_index, "full_attention_layers": iree_vision.full_layer_indices, "final_interval_layers": iree_vision.final_interval_layer_indices, - "target_full_layer": iree_vision.target_full_layer_index, + "target_full_layer": iree_vision.full_layer_indices.last(), + "substage_probe_layer": iree_vision.substage_probe_layer_index, "negative_qwen2_full_attention_detected": true, "negative_identity_permutation_detected": true, "negative_zero_vision_positions_detected": true, diff --git a/src/lib/mlxcel-xla/src/emitter/qwen2_5_vl.rs b/src/lib/mlxcel-xla/src/emitter/qwen2_5_vl.rs index ea7cc180d..07743fa26 100644 --- a/src/lib/mlxcel-xla/src/emitter/qwen2_5_vl.rs +++ b/src/lib/mlxcel-xla/src/emitter/qwen2_5_vl.rs @@ -30,7 +30,7 @@ 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) target_full_layer_index: usize, + pub(crate) substage_probe_layer_index: usize, } #[cfg(feature = "diagnostics")] @@ -44,6 +44,16 @@ struct Qwen25VlBlockDiagnostics { output: Val, } +#[cfg(feature = "diagnostics")] +fn substage_probe_layer(final_interval_layer_indices: &[usize]) -> Option { + final_interval_layer_indices + .iter() + .rev() + .nth(2) + .or_else(|| final_interval_layer_indices.last()) + .copied() +} + struct Args { values: Vec, declarations: Vec, @@ -339,7 +349,7 @@ fn emit_qwen2_5_vl_inner( #[cfg(feature = "diagnostics")] let mut final_interval_layer_states = Vec::new(); #[cfg(feature = "diagnostics")] - let mut target_full_layer_state = None; + let mut substage_probe_layer_state = None; for layer in 0..config.depth { let bias = if full_attention_blocks.contains(&layer) { &full_bias @@ -347,10 +357,10 @@ fn emit_qwen2_5_vl_inner( &window_bias }; #[cfg(feature = "diagnostics")] - let is_target_full_layer = - diagnostic_layout.is_some_and(|layout| layer == layout.target_full_layer_index); + let is_substage_probe_layer = + diagnostic_layout.is_some_and(|layout| layer == layout.substage_probe_layer_index); #[cfg(feature = "diagnostics")] - if is_target_full_layer { + if is_substage_probe_layer { let capture = encoder_layer_with_diagnostics( &mut builder, &hidden, @@ -360,7 +370,7 @@ fn emit_qwen2_5_vl_inner( bias, ); hidden = capture.output.clone(); - target_full_layer_state = Some(capture); + substage_probe_layer_state = Some(capture); } else { hidden = encoder_layer(&mut builder, &hidden, &mut args, config, &freqs, bias); } @@ -412,14 +422,14 @@ fn emit_qwen2_5_vl_inner( "diagnostics capture every layer after the penultimate full-attention boundary" ); outputs.extend(final_interval_layer_states); - let target = target_full_layer_state.expect("diagnostics capture target full layer"); + let probe = substage_probe_layer_state.expect("diagnostics capture substage probe layer"); outputs.extend([ - target.input, - target.norm1, - target.attention, - target.post_attention_residual, - target.norm2, - target.mlp, + probe.input, + probe.norm1, + probe.attention, + probe.post_attention_residual, + probe.norm2, + probe.mlp, ]); outputs.push(projected); let result_types = outputs @@ -496,14 +506,33 @@ pub(crate) fn emit_qwen2_5_vl_diagnostics( .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(&final_interval_layer_indices) + .expect("diagnostic final interval contains the target full-attention layer"); let layout = Qwen25VlVisionDiagnosticLayout { window_layer_index, full_layer_indices, - final_interval_layer_indices: (final_interval_start..=target_full_layer_index).collect(), - target_full_layer_index, + final_interval_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::substage_probe_layer; + + #[test] + fn substage_probe_selects_the_first_observed_actual_divergence_boundary() { + assert_eq!( + substage_probe_layer(&(24..=31).collect::>()), + Some(29) + ); + assert_eq!(substage_probe_layer(&[0, 1]), Some(1)); + assert_eq!(substage_probe_layer(&[]), None); + } +} diff --git a/src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs b/src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs index d2c232a2b..cea6793a2 100644 --- a/src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs +++ b/src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs @@ -1116,7 +1116,7 @@ mod tests { 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.target_full_layer_index, 1); + assert_eq!(layout.substage_probe_layer_index, 1); assert_eq!( production.matches("stablehlo.dot_general").count(), diagnostics.matches("stablehlo.dot_general").count() diff --git a/src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs b/src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs index dedd02d10..842ac43ce 100644 --- a/src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs +++ b/src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs @@ -85,13 +85,13 @@ pub struct Qwen25VlVisionDiagnostics { pub post_full_layers: Vec>, pub final_interval_layer_indices: Vec, pub post_final_interval_layers: Vec>, - pub target_full_layer_index: usize, - pub target_full_layer_input: Vec, - pub target_full_layer_norm1: Vec, - pub target_full_layer_attention: Vec, - pub target_full_layer_post_attention_residual: Vec, - pub target_full_layer_norm2: Vec, - pub target_full_layer_mlp: Vec, + pub substage_probe_layer_index: usize, + pub substage_probe_layer_input: Vec, + pub substage_probe_layer_norm1: 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: Vec, pub merger_window_ordered: Vec, pub restored_projection: Vec, pub patch_tokens: usize, @@ -1055,26 +1055,26 @@ impl IreeQwen25VlDiagnosticProjector { let post_final_interval_layers = decoded .drain(..final_interval_layer_count) .collect::>(); - let mut target_full_layer_substages = decoded.into_iter(); - let target_full_layer_input = target_full_layer_substages + let mut substage_probe_layer_substages = decoded.into_iter(); + let substage_probe_layer_input = substage_probe_layer_substages .next() - .expect("diagnostic target full layer input output"); - let target_full_layer_norm1 = target_full_layer_substages + .expect("diagnostic substage probe layer input output"); + let substage_probe_layer_norm1 = substage_probe_layer_substages .next() - .expect("diagnostic target full layer norm1 output"); - let target_full_layer_attention = target_full_layer_substages + .expect("diagnostic substage probe layer norm1 output"); + let substage_probe_layer_attention = substage_probe_layer_substages .next() - .expect("diagnostic target full layer attention output"); - let target_full_layer_post_attention_residual = target_full_layer_substages + .expect("diagnostic substage probe layer attention output"); + let substage_probe_layer_post_attention_residual = substage_probe_layer_substages .next() - .expect("diagnostic target full layer residual output"); - let target_full_layer_norm2 = target_full_layer_substages + .expect("diagnostic substage probe layer residual output"); + let substage_probe_layer_norm2 = substage_probe_layer_substages .next() - .expect("diagnostic target full layer norm2 output"); - let target_full_layer_mlp = target_full_layer_substages + .expect("diagnostic substage probe layer norm2 output"); + let substage_probe_layer_mlp = substage_probe_layer_substages .next() - .expect("diagnostic target full layer MLP output"); - debug_assert!(target_full_layer_substages.next().is_none()); + .expect("diagnostic substage probe layer MLP output"); + debug_assert!(substage_probe_layer_substages.next().is_none()); Ok(Qwen25VlVisionDiagnostics { window_index: plan.window_index, restore_indices: plan.restore_indices, @@ -1089,13 +1089,13 @@ impl IreeQwen25VlDiagnosticProjector { post_full_layers, final_interval_layer_indices: layout.final_interval_layer_indices, post_final_interval_layers, - target_full_layer_index: layout.target_full_layer_index, - target_full_layer_input, - target_full_layer_norm1, - target_full_layer_attention, - target_full_layer_post_attention_residual, - target_full_layer_norm2, - target_full_layer_mlp, + substage_probe_layer_index: layout.substage_probe_layer_index, + substage_probe_layer_input, + substage_probe_layer_norm1, + substage_probe_layer_attention, + substage_probe_layer_post_attention_residual, + substage_probe_layer_norm2, + substage_probe_layer_mlp, merger_window_ordered, restored_projection, patch_tokens: plan.actual_patches, diff --git a/src/vision/encoders/qwen2_5_vl.rs b/src/vision/encoders/qwen2_5_vl.rs index c796215b0..e5e1caea8 100644 --- a/src/vision/encoders/qwen2_5_vl.rs +++ b/src/vision/encoders/qwen2_5_vl.rs @@ -515,13 +515,13 @@ pub struct Qwen25VLVisionDiagnostics { pub post_full_layers: Vec>, pub final_interval_layer_indices: Vec, pub post_final_interval_layers: Vec>, - pub target_full_layer_index: usize, - pub target_full_layer_input: UniquePtr, - pub target_full_layer_norm1: UniquePtr, - pub target_full_layer_attention: UniquePtr, - pub target_full_layer_post_attention_residual: UniquePtr, - pub target_full_layer_norm2: UniquePtr, - pub target_full_layer_mlp: UniquePtr, + pub substage_probe_layer_index: usize, + pub substage_probe_layer_input: UniquePtr, + pub substage_probe_layer_norm1: UniquePtr, + pub substage_probe_layer_attention: UniquePtr, + pub substage_probe_layer_post_attention_residual: UniquePtr, + pub substage_probe_layer_norm2: UniquePtr, + pub substage_probe_layer_mlp: UniquePtr, pub merger_window_ordered: UniquePtr, pub restored_projection: UniquePtr, } @@ -857,10 +857,22 @@ impl Qwen25VLVisionEncoder { let final_interval_layer_indices = target_full_layer_index .map(|target| (final_interval_start..=target).collect::>()) .unwrap_or_default(); + // The pinned actual-checkpoint oracle first diverges at layer 29 while + // layer 28 still passes. Probe the third-to-last layer in the final + // interval (29 for the standard 32-layer checkpoint) so the substage + // capture localizes that boundary instead of reporting only the already + // divergent final full-attention block. + #[cfg(feature = "xla-diagnostics")] + let substage_probe_layer_index = final_interval_layer_indices + .iter() + .rev() + .nth(2) + .copied() + .or(target_full_layer_index); #[cfg(feature = "xla-diagnostics")] let mut post_final_interval_layers = Vec::new(); #[cfg(feature = "xla-diagnostics")] - let mut target_full_layer_state = None; + 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 @@ -868,11 +880,11 @@ impl Qwen25VLVisionEncoder { &cu_window_seqlens }; #[cfg(feature = "xla-diagnostics")] - if CAPTURE && Some(layer_num) == target_full_layer_index { + 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; - target_full_layer_state = Some(capture); + substage_probe_layer_state = Some(capture); } else { h = block.forward(&h, cu_seqlens_now, &rotary_pos_emb); } @@ -920,10 +932,10 @@ impl Qwen25VLVisionEncoder { #[cfg(feature = "xla-diagnostics")] if CAPTURE { - let target_full_layer_index = - target_full_layer_index.expect("Qwen2.5-VL diagnostics require a full layer"); - let target_full_layer_state = - target_full_layer_state.expect("Qwen2.5-VL target full layer 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 }; @@ -943,14 +955,14 @@ impl Qwen25VLVisionEncoder { post_full_layers, final_interval_layer_indices, post_final_interval_layers, - target_full_layer_index, - target_full_layer_input: target_full_layer_state.input, - target_full_layer_norm1: target_full_layer_state.norm1, - target_full_layer_attention: target_full_layer_state.attention, - target_full_layer_post_attention_residual: target_full_layer_state + 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_attention: substage_probe_layer_state.attention, + substage_probe_layer_post_attention_residual: substage_probe_layer_state .post_attention_residual, - target_full_layer_norm2: target_full_layer_state.norm2, - target_full_layer_mlp: target_full_layer_state.mlp, + substage_probe_layer_norm2: substage_probe_layer_state.norm2, + substage_probe_layer_mlp: substage_probe_layer_state.mlp, merger_window_ordered: merger_window_ordered .expect("Qwen2.5-VL merger capture"), restored_projection, From 5b3d16a7e1e7fce7190a8402ee201cb92152cf70 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 00:29:04 +0900 Subject: [PATCH 06/20] test(xla): materialize Qwen eager probe Snapshot the eager block substages before loading and invoking the mixed IREE/CUDA runtime. Reject vacuous all-zero RMSNorm references before comparison so the strict report cannot misclassify an invalid lazy capture as a numeric divergence. --- examples/xla_qwen25_vl_window_check.rs | 124 ++++++++++++++++--------- 1 file changed, 81 insertions(+), 43 deletions(-) diff --git a/examples/xla_qwen25_vl_window_check.rs b/examples/xla_qwen25_vl_window_check.rs index ab078fe73..0c55bce99 100644 --- a/examples/xla_qwen25_vl_window_check.rs +++ b/examples/xla_qwen25_vl_window_check.rs @@ -348,6 +348,76 @@ fn main() { let eager_vision = model .vision_encoder .forward_with_grid_diagnostics(&pixels, &processor.grids); + // MLX arrays are lazy. Materialize the eager substage oracle before + // entering the IREE/CUDA runtime so the comparison cannot observe stale or + // vacuous intermediates after the mixed-runtime launch. + let eager_substage_probe = [ + ( + "input", + mlx_f32( + eager_vision + .substage_probe_layer_input + .as_ref() + .expect("eager substage probe layer input"), + ), + ), + ( + "norm1", + mlx_f32( + eager_vision + .substage_probe_layer_norm1 + .as_ref() + .expect("eager substage probe layer norm1"), + ), + ), + ( + "attention", + mlx_f32( + eager_vision + .substage_probe_layer_attention + .as_ref() + .expect("eager substage probe layer attention"), + ), + ), + ( + "post_attention_residual", + mlx_f32( + eager_vision + .substage_probe_layer_post_attention_residual + .as_ref() + .expect("eager substage probe layer post-attention residual"), + ), + ), + ( + "norm2", + mlx_f32( + eager_vision + .substage_probe_layer_norm2 + .as_ref() + .expect("eager substage probe layer norm2"), + ), + ), + ( + "mlp", + mlx_f32( + eager_vision + .substage_probe_layer_mlp + .as_ref() + .expect("eager substage probe layer MLP"), + ), + ), + ]; + for (stage, values) in eager_substage_probe + .iter() + .filter(|(stage, _)| matches!(*stage, "norm1" | "norm2")) + { + assert!( + values + .iter() + .any(|value| value.is_finite() && *value != 0.0), + "eager substage probe {stage} is vacuously zero before IREE launch" + ); + } let mut iree_projector = IreeQwen25VlDiagnosticProjector::load(&model_path, &device) .unwrap_or_else(|error| panic!("load Qwen2.5-VL IREE diagnostics: {error}")); @@ -490,64 +560,32 @@ fn main() { tolerance, ); } - for (stage, actual, expected) in [ - ( - "input", - iree_vision.substage_probe_layer_input.as_slice(), - eager_vision - .substage_probe_layer_input - .as_ref() - .expect("eager substage probe layer input"), - ), - ( - "norm1", - iree_vision.substage_probe_layer_norm1.as_slice(), - eager_vision - .substage_probe_layer_norm1 - .as_ref() - .expect("eager substage probe layer norm1"), - ), + 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()), ( "attention", iree_vision.substage_probe_layer_attention.as_slice(), - eager_vision - .substage_probe_layer_attention - .as_ref() - .expect("eager substage probe layer attention"), ), ( "post_attention_residual", iree_vision .substage_probe_layer_post_attention_residual .as_slice(), - eager_vision - .substage_probe_layer_post_attention_residual - .as_ref() - .expect("eager substage probe layer post-attention residual"), ), - ( - "norm2", - iree_vision.substage_probe_layer_norm2.as_slice(), - eager_vision - .substage_probe_layer_norm2 - .as_ref() - .expect("eager substage probe layer norm2"), - ), - ( - "mlp", - iree_vision.substage_probe_layer_mlp.as_slice(), - eager_vision - .substage_probe_layer_mlp - .as_ref() - .expect("eager substage probe layer MLP"), - ), - ] { + ("norm2", iree_vision.substage_probe_layer_norm2.as_slice()), + ("mlp", iree_vision.substage_probe_layer_mlp.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, - &mlx_f32(expected), + expected, tolerance, ); } From 11b331335f1189836d696abf331162894827f329 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 00:54:23 +0900 Subject: [PATCH 07/20] fix(xla): snapshot Qwen diagnostic substages Materialize private diagnostic copies at the layer-29 producer boundary before each source feeds the next block substage. This prevents later MLX graph evaluation from reusing intermediate storage while preserving the production graph and strict oracle tolerances. --- examples/xla_qwen25_vl_window_check.rs | 6 +++--- src/vision/encoders/qwen2_5_vl.rs | 29 ++++++++++++++++++++------ 2 files changed, 26 insertions(+), 9 deletions(-) diff --git a/examples/xla_qwen25_vl_window_check.rs b/examples/xla_qwen25_vl_window_check.rs index 0c55bce99..1d2a0f17a 100644 --- a/examples/xla_qwen25_vl_window_check.rs +++ b/examples/xla_qwen25_vl_window_check.rs @@ -348,9 +348,9 @@ fn main() { let eager_vision = model .vision_encoder .forward_with_grid_diagnostics(&pixels, &processor.grids); - // MLX arrays are lazy. Materialize the eager substage oracle before - // entering the IREE/CUDA runtime so the comparison cannot observe stale or - // vacuous intermediates after the mixed-runtime launch. + // The producer materializes private copies at each substage boundary. + // Export those snapshots before IREE launch and reject a vacuous capture + // immediately so the strict oracle cannot report a numeric false positive. let eager_substage_probe = [ ( "input", diff --git a/src/vision/encoders/qwen2_5_vl.rs b/src/vision/encoders/qwen2_5_vl.rs index e5e1caea8..8331027c7 100644 --- a/src/vision/encoders/qwen2_5_vl.rs +++ b/src/vision/encoders/qwen2_5_vl.rs @@ -394,6 +394,13 @@ struct Qwen25VLBlockDiagnostics { mlp: UniquePtr, } +#[cfg(feature = "xla-diagnostics")] +fn materialized_diagnostic_snapshot(array: &MlxArray) -> UniquePtr { + let snapshot = mlxcel_core::copy(array); + mlxcel_core::eval(&snapshot); + snapshot +} + impl VisionBlock { fn from_weights( weights: &WeightMap, @@ -431,22 +438,32 @@ impl VisionBlock { cu_seqlens: &[i32], rotary_pos_emb: &MlxArray, ) -> (UniquePtr, Qwen25VLBlockDiagnostics) { - let input = mlxcel_core::copy(hidden_states); + // These arrays are diagnostic values, not additional lazy graph + // handles. Materialize each private copy before its source feeds the + // next substage so later graph evaluation cannot donate or reuse the + // intermediate storage observed by the oracle. + let input = materialized_diagnostic_snapshot(hidden_states); let norm1 = self.norm1.forward(hidden_states); + let norm1_snapshot = materialized_diagnostic_snapshot(&norm1); let attention = self.attn.forward(&norm1, cu_seqlens, rotary_pos_emb); + let attention_snapshot = materialized_diagnostic_snapshot(&attention); let post_attention_residual = mlxcel_core::add(hidden_states, &attention); + let post_attention_residual_snapshot = + materialized_diagnostic_snapshot(&post_attention_residual); let norm2 = self.norm2.forward(&post_attention_residual); + let norm2_snapshot = materialized_diagnostic_snapshot(&norm2); let mlp = self.mlp.forward(&norm2); + let mlp_snapshot = materialized_diagnostic_snapshot(&mlp); let output = mlxcel_core::add(&post_attention_residual, &mlp); ( output, Qwen25VLBlockDiagnostics { input, - norm1, - attention, - post_attention_residual, - norm2, - mlp, + norm1: norm1_snapshot, + attention: attention_snapshot, + post_attention_residual: post_attention_residual_snapshot, + norm2: norm2_snapshot, + mlp: mlp_snapshot, }, ) } From fb4764658c56e7b5c8d2af83e14ef227fbf4dc14 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 01:26:58 +0900 Subject: [PATCH 08/20] fix(xla): own Qwen diagnostic snapshots Read each Qwen2.5-VL block diagnostic into host-owned F32 memory at its producer boundary so later MLX graph evaluation cannot alias or donate the oracle values. Keep the production path and comparison thresholds unchanged, and cover snapshot survival with a synthetic RMSNorm regression. --- examples/xla_qwen25_vl_window_check.rs | 62 +++-------------- src/vision/encoders/qwen2_5_vl.rs | 92 ++++++++++++++++++-------- 2 files changed, 75 insertions(+), 79 deletions(-) diff --git a/examples/xla_qwen25_vl_window_check.rs b/examples/xla_qwen25_vl_window_check.rs index 1d2a0f17a..3e6db9037 100644 --- a/examples/xla_qwen25_vl_window_check.rs +++ b/examples/xla_qwen25_vl_window_check.rs @@ -348,64 +348,24 @@ fn main() { let eager_vision = model .vision_encoder .forward_with_grid_diagnostics(&pixels, &processor.grids); - // The producer materializes private copies at each substage boundary. - // Export those snapshots before IREE launch and reject a vacuous capture - // immediately so the strict oracle cannot report a numeric false positive. + // 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", - mlx_f32( - eager_vision - .substage_probe_layer_input - .as_ref() - .expect("eager substage probe layer input"), - ), - ), - ( - "norm1", - mlx_f32( - eager_vision - .substage_probe_layer_norm1 - .as_ref() - .expect("eager substage probe layer norm1"), - ), - ), + ("input", eager_vision.substage_probe_layer_input.as_slice()), + ("norm1", eager_vision.substage_probe_layer_norm1.as_slice()), ( "attention", - mlx_f32( - eager_vision - .substage_probe_layer_attention - .as_ref() - .expect("eager substage probe layer attention"), - ), + eager_vision.substage_probe_layer_attention.as_slice(), ), ( "post_attention_residual", - mlx_f32( - eager_vision - .substage_probe_layer_post_attention_residual - .as_ref() - .expect("eager substage probe layer post-attention residual"), - ), - ), - ( - "norm2", - mlx_f32( - eager_vision - .substage_probe_layer_norm2 - .as_ref() - .expect("eager substage probe layer norm2"), - ), - ), - ( - "mlp", - mlx_f32( - eager_vision - .substage_probe_layer_mlp - .as_ref() - .expect("eager substage probe layer MLP"), - ), + eager_vision + .substage_probe_layer_post_attention_residual + .as_slice(), ), + ("norm2", eager_vision.substage_probe_layer_norm2.as_slice()), + ("mlp", eager_vision.substage_probe_layer_mlp.as_slice()), ]; for (stage, values) in eager_substage_probe .iter() diff --git a/src/vision/encoders/qwen2_5_vl.rs b/src/vision/encoders/qwen2_5_vl.rs index 8331027c7..331381796 100644 --- a/src/vision/encoders/qwen2_5_vl.rs +++ b/src/vision/encoders/qwen2_5_vl.rs @@ -386,19 +386,27 @@ struct VisionBlock { #[cfg(feature = "xla-diagnostics")] struct Qwen25VLBlockDiagnostics { - input: UniquePtr, - norm1: UniquePtr, - attention: UniquePtr, - post_attention_residual: UniquePtr, - norm2: UniquePtr, - mlp: UniquePtr, + input: Vec, + norm1: Vec, + attention: Vec, + post_attention_residual: Vec, + norm2: Vec, + mlp: Vec, } -#[cfg(feature = "xla-diagnostics")] -fn materialized_diagnostic_snapshot(array: &MlxArray) -> UniquePtr { - let snapshot = mlxcel_core::copy(array); - mlxcel_core::eval(&snapshot); - snapshot +#[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 { @@ -438,22 +446,22 @@ impl VisionBlock { cu_seqlens: &[i32], rotary_pos_emb: &MlxArray, ) -> (UniquePtr, Qwen25VLBlockDiagnostics) { - // These arrays are diagnostic values, not additional lazy graph - // handles. Materialize each private copy before its source feeds the - // next substage so later graph evaluation cannot donate or reuse the - // intermediate storage observed by the oracle. - let input = materialized_diagnostic_snapshot(hidden_states); + // 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 = materialized_diagnostic_snapshot(&norm1); + let norm1_snapshot = host_f32_diagnostic_snapshot(&norm1); let attention = self.attn.forward(&norm1, cu_seqlens, rotary_pos_emb); - let attention_snapshot = materialized_diagnostic_snapshot(&attention); + let attention_snapshot = host_f32_diagnostic_snapshot(&attention); let post_attention_residual = mlxcel_core::add(hidden_states, &attention); let post_attention_residual_snapshot = - materialized_diagnostic_snapshot(&post_attention_residual); + host_f32_diagnostic_snapshot(&post_attention_residual); let norm2 = self.norm2.forward(&post_attention_residual); - let norm2_snapshot = materialized_diagnostic_snapshot(&norm2); + let norm2_snapshot = host_f32_diagnostic_snapshot(&norm2); let mlp = self.mlp.forward(&norm2); - let mlp_snapshot = materialized_diagnostic_snapshot(&mlp); + let mlp_snapshot = host_f32_diagnostic_snapshot(&mlp); let output = mlxcel_core::add(&post_attention_residual, &mlp); ( output, @@ -533,12 +541,12 @@ pub struct Qwen25VLVisionDiagnostics { pub final_interval_layer_indices: Vec, pub post_final_interval_layers: Vec>, pub substage_probe_layer_index: usize, - pub substage_probe_layer_input: UniquePtr, - pub substage_probe_layer_norm1: UniquePtr, - pub substage_probe_layer_attention: UniquePtr, - pub substage_probe_layer_post_attention_residual: UniquePtr, - pub substage_probe_layer_norm2: UniquePtr, - pub substage_probe_layer_mlp: UniquePtr, + pub substage_probe_layer_input: Vec, + pub substage_probe_layer_norm1: 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: Vec, pub merger_window_ordered: UniquePtr, pub restored_projection: UniquePtr, } @@ -1006,7 +1014,7 @@ impl super::VisionEncoder for Qwen25VLVisionEncoder { #[cfg(test)] mod tests { - use super::restoration_indices; + use super::{host_f32_diagnostic_snapshot, restoration_indices}; #[test] fn restoration_is_the_true_inverse_for_non_self_inverse_permutation() { @@ -1020,4 +1028,32 @@ mod tests { .collect::>(); assert_eq!(original, vec!["zero", "one", "two", "three"]); } + + #[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}" + ); + } + } } From 2e7a19a933ca45218ae34bfcb4f7f3bf188b100f Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 01:41:18 +0900 Subject: [PATCH 09/20] test(xla): expose Qwen F16 norm overflow Characterize the manual eager RMSNorm square-overflow signature and report input/norm statistics at the pre-IREE guard. Production graphs and comparison thresholds remain unchanged. --- examples/xla_qwen25_vl_window_check.rs | 49 +++++++++++++++++++++++--- src/vision/encoders/qwen2_5_vl.rs | 25 +++++++++++++ 2 files changed, 70 insertions(+), 4 deletions(-) diff --git a/examples/xla_qwen25_vl_window_check.rs b/examples/xla_qwen25_vl_window_check.rs index 3e6db9037..c63bd752f 100644 --- a/examples/xla_qwen25_vl_window_check.rs +++ b/examples/xla_qwen25_vl_window_check.rs @@ -111,6 +111,45 @@ fn mlx_i32(array: &mlxcel_core::MlxArray) -> Vec { .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, @@ -371,11 +410,13 @@ fn main() { .iter() .filter(|(stage, _)| matches!(*stage, "norm1" | "norm2")) { + let stats = snapshot_stats(values); assert!( - values - .iter() - .any(|value| value.is_finite() && *value != 0.0), - "eager substage probe {stage} is vacuously zero before IREE launch" + 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()) ); } diff --git a/src/vision/encoders/qwen2_5_vl.rs b/src/vision/encoders/qwen2_5_vl.rs index 331381796..1226175bc 100644 --- a/src/vision/encoders/qwen2_5_vl.rs +++ b/src/vision/encoders/qwen2_5_vl.rs @@ -1056,4 +1056,29 @@ mod tests { ); } } + + #[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 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" + ); + } } From fd320719e4a00b04559a3e2c07fafaa11063db92 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 01:47:54 +0900 Subject: [PATCH 10/20] fix(xla): stabilize Qwen vision RMSNorm --- src/vision/encoders/qwen2_5_vl.rs | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/src/vision/encoders/qwen2_5_vl.rs b/src/vision/encoders/qwen2_5_vl.rs index 1226175bc..57585e3ba 100644 --- a/src/vision/encoders/qwen2_5_vl.rs +++ b/src/vision/encoders/qwen2_5_vl.rs @@ -131,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) } } @@ -1070,6 +1070,13 @@ mod tests { "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); @@ -1080,5 +1087,11 @@ mod tests { .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}" + ); + } } } From 015e59cbcc0c142c20d2a1e1a5ad524a73217504 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 02:11:07 +0900 Subject: [PATCH 11/20] test(xla): localize Qwen full attention drift --- examples/xla_qwen25_vl_window_check.rs | 80 +++++++++++-- src/lib/mlxcel-xla/src/emitter/qwen2_5_vl.rs | 68 +++++++---- src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs | 34 +++++- src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs | 119 ++++++++++++++++++- src/vision/encoders/qwen2_5_vl.rs | 75 ++++++++++-- 5 files changed, 330 insertions(+), 46 deletions(-) diff --git a/examples/xla_qwen25_vl_window_check.rs b/examples/xla_qwen25_vl_window_check.rs index c63bd752f..c47610191 100644 --- a/examples/xla_qwen25_vl_window_check.rs +++ b/examples/xla_qwen25_vl_window_check.rs @@ -393,6 +393,15 @@ fn main() { 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(), @@ -450,6 +459,28 @@ fn main() { 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 @@ -564,6 +595,15 @@ fn main() { 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(), @@ -627,7 +667,16 @@ fn main() { "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, @@ -654,17 +703,14 @@ fn main() { &processor.grids, Qwen25VlDiagnosticMutation::Qwen2FullAttention, ); - assert!( - compare_stage( - &mut Vec::new(), - "negative.qwen2_full_attention", - &qwen2_full.post_window_layer, - &mlx_f32(&eager_vision.post_window_layer), - tolerance, - ) - .is_err(), - "negative oracle accepted Qwen2-style full attention in a window layer" - ); + 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, @@ -779,7 +825,16 @@ fn main() { "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, @@ -788,7 +843,8 @@ fn main() { "final_interval_layers": iree_vision.final_interval_layer_indices, "target_full_layer": iree_vision.full_layer_indices.last(), "substage_probe_layer": iree_vision.substage_probe_layer_index, - "negative_qwen2_full_attention_detected": true, + "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), diff --git a/src/lib/mlxcel-xla/src/emitter/qwen2_5_vl.rs b/src/lib/mlxcel-xla/src/emitter/qwen2_5_vl.rs index 07743fa26..e3c58c487 100644 --- a/src/lib/mlxcel-xla/src/emitter/qwen2_5_vl.rs +++ b/src/lib/mlxcel-xla/src/emitter/qwen2_5_vl.rs @@ -37,6 +37,10 @@ pub(crate) struct Qwen25VlVisionDiagnosticLayout { struct Qwen25VlBlockDiagnostics { input: Val, norm1: Val, + query: Val, + key: Val, + value: Val, + attention_context: Val, attention: Val, post_attention_residual: Val, norm2: Val, @@ -45,15 +49,23 @@ struct Qwen25VlBlockDiagnostics { } #[cfg(feature = "diagnostics")] -fn substage_probe_layer(final_interval_layer_indices: &[usize]) -> Option { - final_interval_layer_indices +fn substage_probe_layer(full_layer_indices: &[usize]) -> Option { + full_layer_indices .iter() .rev() - .nth(2) - .or_else(|| final_interval_layer_indices.last()) + .nth(1) + .or_else(|| full_layer_indices.last()) .copied() } +struct Qwen25VlAttentionValues { + query: Val, + key: Val, + value: Val, + context: Val, + output: Val, +} + struct Args { values: Vec, declarations: Vec, @@ -185,7 +197,7 @@ fn attention( config: &Qwen2VlConfig, freqs: &Val, attention_bias: &Val, -) -> 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()); @@ -205,11 +217,12 @@ fn attention( 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 q = apply_vision_rope(builder, &q, freqs); - let k = apply_vision_rope(builder, &k, freqs); - let q = builder.transpose(&q, &[1, 0, 2]); - let k = builder.transpose(&k, &[1, 0, 2]); - let v = builder.transpose(&v, &[1, 0, 2]); + 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, @@ -244,7 +257,14 @@ fn attention( ); let context = builder.transpose(&context, &[1, 0, 2]); let context = builder.reshape(&context, vec![tokens, config.hidden]); - linear_2d(builder, &context, &args.take(), &args.take()) + let output = linear_2d(builder, &context, &args.take(), &args.take()); + Qwen25VlAttentionValues { + query, + key, + value, + context, + output, + } } fn encoder_layer( @@ -257,7 +277,7 @@ fn encoder_layer( ) -> 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); + 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); @@ -278,7 +298,7 @@ fn encoder_layer_with_diagnostics( ) -> 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); + let post_attention_residual = builder.add(hidden, &attention.output); let norm2 = rms_norm( builder, &post_attention_residual, @@ -294,7 +314,11 @@ fn encoder_layer_with_diagnostics( Qwen25VlBlockDiagnostics { input: hidden.clone(), norm1, - attention, + query: attention.query, + key: attention.key, + value: attention.value, + attention_context: attention.context, + attention: attention.output, post_attention_residual, norm2, mlp, @@ -426,6 +450,10 @@ fn emit_qwen2_5_vl_inner( outputs.extend([ probe.input, probe.norm1, + probe.query, + probe.key, + probe.value, + probe.attention_context, probe.attention, probe.post_attention_residual, probe.norm2, @@ -508,8 +536,8 @@ pub(crate) fn emit_qwen2_5_vl_diagnostics( .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(&final_interval_layer_indices) - .expect("diagnostic final interval contains the target full-attention layer"); + let substage_probe_layer_index = substage_probe_layer(&full_layer_indices) + .expect("diagnostics contain a configured full-attention layer"); let layout = Qwen25VlVisionDiagnosticLayout { window_layer_index, full_layer_indices, @@ -528,11 +556,9 @@ mod tests { #[test] fn substage_probe_selects_the_first_observed_actual_divergence_boundary() { - assert_eq!( - substage_probe_layer(&(24..=31).collect::>()), - Some(29) - ); - assert_eq!(substage_probe_layer(&[0, 1]), Some(1)); + assert_eq!(substage_probe_layer(&[7, 15, 23, 31]), Some(23)); + assert_eq!(substage_probe_layer(&[0, 1]), Some(0)); + assert_eq!(substage_probe_layer(&[0]), Some(0)); assert_eq!(substage_probe_layer(&[]), None); } } diff --git a/src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs b/src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs index cea6793a2..c4e965198 100644 --- a/src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs +++ b/src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs @@ -1121,7 +1121,9 @@ mod tests { production.matches("stablehlo.dot_general").count(), diagnostics.matches("stablehlo.dot_general").count() ); - let result_signature = std::iter::repeat_n("tensor<16x1280xf32>", 11) + 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>", 5)) .chain(std::iter::once("tensor<4x1536xf32>")) .collect::>() .join(", "); @@ -1152,7 +1154,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!( @@ -1208,12 +1210,36 @@ mod tests { ); let full = &inputs.packed_attention_bias; let window = inputs.window_attention_bias.as_ref().unwrap(); - assert_eq!(full[0 * 256 + 80], 0.0); - assert_eq!(window[0 * 256 + 80], f32::NEG_INFINITY); + 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(); diff --git a/src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs b/src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs index 842ac43ce..9d499df16 100644 --- a/src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs +++ b/src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs @@ -71,6 +71,15 @@ pub enum Qwen25VlDiagnosticMutation { 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 { @@ -88,16 +97,22 @@ pub struct Qwen25VlVisionDiagnostics { 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: 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 { @@ -612,6 +627,54 @@ fn checked_f32_output(label: &str, bytes: Vec) -> 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, @@ -930,6 +993,15 @@ impl IreeQwen25VlDiagnosticProjector { * 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 => { @@ -970,7 +1042,7 @@ impl IreeQwen25VlDiagnosticProjector { 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 block_substage_count = 6; + let block_substage_count = 10; let state_output_count = 2 + full_layer_count + final_interval_layer_count + block_substage_count; let mut output_buffers = vec![vec![0u8; state_bytes]; state_output_count]; @@ -1062,6 +1134,18 @@ impl IreeQwen25VlDiagnosticProjector { 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"); @@ -1092,16 +1176,22 @@ impl IreeQwen25VlDiagnosticProjector { 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, 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, }) } } @@ -1110,6 +1200,33 @@ impl IreeQwen25VlDiagnosticProjector { 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); + } + fn tiny_config() -> Qwen2VlConfig { Qwen2VlConfig::from_json_str( r#"{ diff --git a/src/vision/encoders/qwen2_5_vl.rs b/src/vision/encoders/qwen2_5_vl.rs index 57585e3ba..ca2112147 100644 --- a/src/vision/encoders/qwen2_5_vl.rs +++ b/src/vision/encoders/qwen2_5_vl.rs @@ -219,6 +219,14 @@ struct VisionAttention { scale: f32, } +#[derive(Clone, Copy)] +enum VisionAttentionDiagnosticStage { + Query, + Key, + Value, + Context, +} + impl VisionAttention { fn from_weights( weights: &WeightMap, @@ -247,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]; @@ -275,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]); @@ -331,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) } @@ -388,6 +413,10 @@ struct VisionBlock { struct Qwen25VLBlockDiagnostics { input: Vec, norm1: Vec, + query: Vec, + key: Vec, + value: Vec, + attention_context: Vec, attention: Vec, post_attention_residual: Vec, norm2: Vec, @@ -453,7 +482,23 @@ impl VisionBlock { let input = host_f32_diagnostic_snapshot(hidden_states); let norm1 = self.norm1.forward(hidden_states); let norm1_snapshot = host_f32_diagnostic_snapshot(&norm1); - let attention = self.attn.forward(&norm1, cu_seqlens, rotary_pos_emb); + 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 = @@ -468,6 +513,11 @@ impl VisionBlock { 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, @@ -543,6 +593,10 @@ pub struct Qwen25VLVisionDiagnostics { 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, @@ -882,16 +936,16 @@ impl Qwen25VLVisionEncoder { let final_interval_layer_indices = target_full_layer_index .map(|target| (final_interval_start..=target).collect::>()) .unwrap_or_default(); - // The pinned actual-checkpoint oracle first diverges at layer 29 while - // layer 28 still passes. Probe the third-to-last layer in the final - // interval (29 for the standard 32-layer checkpoint) so the substage - // capture localizes that boundary instead of reporting only the already - // divergent final full-attention block. + // The multi-media actual-checkpoint oracle passes full-attention layers + // 7 and 15 and first diverges at full-attention layer 23. Probe the + // penultimate configured full-attention 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 substage_probe_layer_index = final_interval_layer_indices + let substage_probe_layer_index = full_layer_indices .iter() .rev() - .nth(2) + .nth(1) .copied() .or(target_full_layer_index); #[cfg(feature = "xla-diagnostics")] @@ -983,6 +1037,11 @@ impl Qwen25VLVisionEncoder { 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, From 0024800dc4675385c8957fde4b886bcb486d1427 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 02:31:16 +0900 Subject: [PATCH 12/20] fix(xla): preserve Qwen diagnostic output ranks The layer-23 Q/K/V diagnostics are emitted as rank-3 patch, head, and head-dimension tensors, but the auxiliary runtime declared every non-merger result as a rank-2 hidden-state tensor. IREE therefore rejected output 16 even though the flattened element count matched. Generate the complete ordered diagnostic output descriptor schema from the captured layer layout, retain rank 3 for Q/K/V at invocation, and keep the existing flat owned vectors only after transfer for comparison API compatibility. A regression pins all 25 actual-checkpoint output semantics, ordinals, ranks, and shapes. Validation: Qwen2.5-VL diagnostics tests (9 passed), actual descriptor regression, diagnostics example check, Clippy, formatting, and diff checks. Refs #866 --- src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs | 190 ++++++++++++++++++--- 1 file changed, 170 insertions(+), 20 deletions(-) diff --git a/src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs b/src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs index 9d499df16..4b9336f47 100644 --- a/src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs +++ b/src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs @@ -127,6 +127,74 @@ struct ActiveDiagnosticModule { 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 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}"))), + ); + 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")), + state(format!("substage_probe_layer.{probe}.mlp")), + ]); + 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, @@ -1034,19 +1102,20 @@ impl IreeQwen25VlDiagnosticProjector { 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 state_shape = [plan.patch_bucket, self.config.hidden]; - let merger_shape = [plan.patch_bucket / merge_unit, self.config.text_hidden]; - let state_bytes = state_shape.iter().product::() * std::mem::size_of::(); - let merger_bytes = merger_shape.iter().product::() * std::mem::size_of::(); + 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 block_substage_count = 10; - let state_output_count = - 2 + full_layer_count + final_interval_layer_count + block_substage_count; - let mut output_buffers = vec![vec![0u8; state_bytes]; state_output_count]; - output_buffers.push(vec![0u8; merger_bytes]); + 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), @@ -1071,30 +1140,32 @@ impl IreeQwen25VlDiagnosticProjector { ]; let mut outputs = output_buffers .iter_mut() - .enumerate() - .map(|(index, bytes)| AuxiliaryOutput { + .zip(&output_specs) + .map(|(bytes, spec)| AuxiliaryOutput { bytes, dtype: AuxiliaryTensorDType::Float32, - shape: if index == state_output_count { - &merger_shape - } else { - &state_shape - }, + shape: &spec.shape, }) .collect::>(); active.module.invoke(&inputs, &mut outputs)?; let mut decoded = output_buffers .into_iter() - .enumerate() - .map(|(index, bytes)| { - checked_f32_output(&format!("Qwen2.5-VL diagnostic output {index}"), bytes) + .zip(&output_specs) + .map(|(bytes, spec)| { + checked_f32_output( + &format!("Qwen2.5-VL diagnostic output {}", spec.label), + bytes, + ) }) .collect::, _>>()?; let actual_state_values = plan .actual_patches .checked_mul(self.config.hidden) .ok_or_else(|| "Qwen2.5-VL diagnostic state size overflowed".to_string())?; - for state in decoded.iter_mut().take(state_output_count) { + // 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. + for state in decoded.iter_mut().take(merger_output_index) { state.truncate(actual_state_values); } let all_projected = decoded @@ -1227,6 +1298,85 @@ mod tests { 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", + "substage_probe_layer.23.input", + "substage_probe_layer.23.norm1", + "substage_probe_layer.23.query", + "substage_probe_layer.23.key", + "substage_probe_layer.23.value", + "substage_probe_layer.23.attention_context", + "substage_probe_layer.23.attention", + "substage_probe_layer.23.post_attention_residual", + "substage_probe_layer.23.norm2", + "substage_probe_layer.23.mlp", + "merger_window_ordered", + ] + ); + assert_eq!( + specs + .iter() + .map(|spec| spec.shape.len()) + .collect::>(), + [vec![2; 16], vec![3; 3], vec![2; 6],].concat() + ); + for spec in &specs[16..=18] { + assert_eq!(spec.shape, vec![256, 16, 80]); + } + assert_eq!(specs[15].shape, vec![256, 1280]); + assert_eq!(specs[19].shape, vec![256, 1280]); + assert_eq!(specs[24].shape, vec![64, 1536]); + } + fn tiny_config() -> Qwen2VlConfig { Qwen2VlConfig::from_json_str( r#"{ From 28234ca370f5eb266c6ee06dbc70dfbce49bb5bf Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 02:43:28 +0900 Subject: [PATCH 13/20] test(xla): trace Qwen pre-probe layer outputs Context: The successful 2d13e470 two-media diagnostic run passes layer 15 but reaches layer 23 with its input already outside tolerance, leaving the first divergence somewhere in layers 16 through 22. Changes: - capture the seven ordered layer outputs between the preceding full-attention boundary and the layer 23 substage probe - expose matching eager and IREE diagnostics and compare them in the strict oracle report - pin the actual-model descriptor labels, ordinals, ranks, and shapes Validation: - cargo fmt --all -- --check - cargo test -p mlxcel-xla --features diagnostics qwen25 -- --nocapture - cargo test -p mlxcel-xla --features diagnostics pre_probe_layers -- --nocapture - cargo check --features xla-diagnostics --example xla_qwen25_vl_window_check - cargo clippy -p mlxcel-xla --features diagnostics --tests --no-deps Refs #866 --- examples/xla_qwen25_vl_window_check.rs | 26 ++++++++++++++ src/lib/mlxcel-xla/src/emitter/qwen2_5_vl.rs | 38 +++++++++++++++++++- src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs | 29 ++++++++++++--- src/vision/encoders/qwen2_5_vl.rs | 22 ++++++++++++ 4 files changed, 109 insertions(+), 6 deletions(-) diff --git a/examples/xla_qwen25_vl_window_check.rs b/examples/xla_qwen25_vl_window_check.rs index c47610191..b93a6bd1f 100644 --- a/examples/xla_qwen25_vl_window_check.rs +++ b/examples/xla_qwen25_vl_window_check.rs @@ -537,6 +537,15 @@ fn main() { eager_vision.post_final_interval_layers.len(), "captured final diagnostic interval count differs" ); + assert_eq!( + iree_vision.pre_probe_layer_indices, eager_vision.pre_probe_layer_indices, + "pre-probe diagnostic interval differs" + ); + assert_eq!( + iree_vision.post_pre_probe_layers.len(), + eager_vision.post_pre_probe_layers.len(), + "captured pre-probe diagnostic interval count differs" + ); let target_full_layer_index = *iree_vision .full_layer_indices .last() @@ -592,6 +601,21 @@ fn main() { tolerance, ); } + for (&layer, (actual, expected)) in iree_vision.pre_probe_layer_indices.iter().zip( + iree_vision + .post_pre_probe_layers + .iter() + .zip(&eager_vision.post_pre_probe_layers), + ) { + record_stage( + &mut reports, + &mut comparison_failures, + &format!("vision.post_layer_{layer}_pre_probe"), + 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()), @@ -683,6 +707,7 @@ fn main() { "window_layer": iree_vision.window_layer_index, "full_attention_layers": iree_vision.full_layer_indices, "final_interval_layers": iree_vision.final_interval_layer_indices, + "pre_probe_layers": iree_vision.pre_probe_layer_indices, "target_full_layer": target_full_layer_index, "substage_probe_layer": substage_probe_layer_index, "failed_phase": "vision", @@ -841,6 +866,7 @@ fn main() { "window_layer": iree_vision.window_layer_index, "full_attention_layers": iree_vision.full_layer_indices, "final_interval_layers": iree_vision.final_interval_layer_indices, + "pre_probe_layers": iree_vision.pre_probe_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, diff --git a/src/lib/mlxcel-xla/src/emitter/qwen2_5_vl.rs b/src/lib/mlxcel-xla/src/emitter/qwen2_5_vl.rs index e3c58c487..00699bd69 100644 --- a/src/lib/mlxcel-xla/src/emitter/qwen2_5_vl.rs +++ b/src/lib/mlxcel-xla/src/emitter/qwen2_5_vl.rs @@ -30,6 +30,7 @@ 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) pre_probe_layer_indices: Vec, pub(crate) substage_probe_layer_index: usize, } @@ -58,6 +59,17 @@ fn substage_probe_layer(full_layer_indices: &[usize]) -> Option { .copied() } +#[cfg(feature = "diagnostics")] +fn pre_probe_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, @@ -373,6 +385,8 @@ fn emit_qwen2_5_vl_inner( #[cfg(feature = "diagnostics")] let mut final_interval_layer_states = Vec::new(); #[cfg(feature = "diagnostics")] + let mut pre_probe_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) { @@ -413,6 +427,9 @@ fn emit_qwen2_5_vl_inner( if layout.final_interval_layer_indices.contains(&layer) { final_interval_layer_states.push(hidden.clone()); } + if layout.pre_probe_layer_indices.contains(&layer) { + pre_probe_layer_states.push(hidden.clone()); + } } } hidden = rms_norm(&mut builder, &hidden, &args.take(), config.layer_norm_eps); @@ -446,6 +463,13 @@ fn emit_qwen2_5_vl_inner( "diagnostics capture every layer after the penultimate full-attention boundary" ); outputs.extend(final_interval_layer_states); + assert_eq!( + pre_probe_layer_states.len(), + layout.pre_probe_layer_indices.len(), + "diagnostics capture every layer between the preceding full-attention boundary and \ + the substage probe" + ); + outputs.extend(pre_probe_layer_states); let probe = substage_probe_layer_state.expect("diagnostics capture substage probe layer"); outputs.extend([ probe.input, @@ -538,10 +562,12 @@ pub(crate) fn emit_qwen2_5_vl_diagnostics( (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 pre_probe_layer_indices = pre_probe_layers(&full_layer_indices, substage_probe_layer_index); let layout = Qwen25VlVisionDiagnosticLayout { window_layer_index, full_layer_indices, final_interval_layer_indices, + pre_probe_layer_indices, substage_probe_layer_index, }; Ok(( @@ -552,7 +578,7 @@ pub(crate) fn emit_qwen2_5_vl_diagnostics( #[cfg(all(test, feature = "diagnostics"))] mod tests { - use super::substage_probe_layer; + use super::{pre_probe_layers, substage_probe_layer}; #[test] fn substage_probe_selects_the_first_observed_actual_divergence_boundary() { @@ -561,4 +587,14 @@ mod tests { assert_eq!(substage_probe_layer(&[0]), Some(0)); assert_eq!(substage_probe_layer(&[]), None); } + + #[test] + fn pre_probe_layers_cover_only_the_gap_after_the_preceding_full_layer() { + assert_eq!( + pre_probe_layers(&[7, 15, 23, 31], 23), + (16..23).collect::>() + ); + assert_eq!(pre_probe_layers(&[0, 1], 0), Vec::::new()); + assert_eq!(pre_probe_layers(&[0], 0), Vec::::new()); + } } diff --git a/src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs b/src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs index 4b9336f47..7d555c10c 100644 --- a/src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs +++ b/src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs @@ -94,6 +94,8 @@ pub struct Qwen25VlVisionDiagnostics { pub post_full_layers: Vec>, pub final_interval_layer_indices: Vec, pub post_final_interval_layers: Vec>, + pub pre_probe_layer_indices: Vec, + pub post_pre_probe_layers: Vec>, pub substage_probe_layer_index: usize, pub substage_probe_layer_input: Vec, pub substage_probe_layer_norm1: Vec, @@ -161,6 +163,12 @@ fn qwen25_diagnostic_output_specs( .iter() .map(|layer| state(format!("post_final_interval_layer.{layer}"))), ); + specs.extend( + layout + .pre_probe_layer_indices + .iter() + .map(|layer| state(format!("post_pre_probe_layer.{layer}"))), + ); let probe = layout.substage_probe_layer_index; specs.extend([ state(format!("substage_probe_layer.{probe}.input")), @@ -1107,6 +1115,7 @@ impl IreeQwen25VlDiagnosticProjector { 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 pre_probe_layer_count = layout.pre_probe_layer_indices.len(); let output_specs = qwen25_diagnostic_output_specs(&diagnostic_config, &layout, plan.patch_bucket); let merger_output_index = output_specs.len() - 1; @@ -1198,6 +1207,7 @@ impl IreeQwen25VlDiagnosticProjector { let post_final_interval_layers = decoded .drain(..final_interval_layer_count) .collect::>(); + let post_pre_probe_layers = decoded.drain(..pre_probe_layer_count).collect::>(); let mut substage_probe_layer_substages = decoded.into_iter(); let substage_probe_layer_input = substage_probe_layer_substages .next() @@ -1244,6 +1254,8 @@ impl IreeQwen25VlDiagnosticProjector { post_full_layers, final_interval_layer_indices: layout.final_interval_layer_indices, post_final_interval_layers, + pre_probe_layer_indices: layout.pre_probe_layer_indices, + post_pre_probe_layers, substage_probe_layer_index: layout.substage_probe_layer_index, substage_probe_layer_input, substage_probe_layer_norm1, @@ -1349,6 +1361,13 @@ mod tests { "post_final_interval_layer.29", "post_final_interval_layer.30", "post_final_interval_layer.31", + "post_pre_probe_layer.16", + "post_pre_probe_layer.17", + "post_pre_probe_layer.18", + "post_pre_probe_layer.19", + "post_pre_probe_layer.20", + "post_pre_probe_layer.21", + "post_pre_probe_layer.22", "substage_probe_layer.23.input", "substage_probe_layer.23.norm1", "substage_probe_layer.23.query", @@ -1367,14 +1386,14 @@ mod tests { .iter() .map(|spec| spec.shape.len()) .collect::>(), - [vec![2; 16], vec![3; 3], vec![2; 6],].concat() + [vec![2; 23], vec![3; 3], vec![2; 6],].concat() ); - for spec in &specs[16..=18] { + for spec in &specs[23..=25] { assert_eq!(spec.shape, vec![256, 16, 80]); } - assert_eq!(specs[15].shape, vec![256, 1280]); - assert_eq!(specs[19].shape, vec![256, 1280]); - assert_eq!(specs[24].shape, vec![64, 1536]); + assert_eq!(specs[22].shape, vec![256, 1280]); + assert_eq!(specs[26].shape, vec![256, 1280]); + assert_eq!(specs[31].shape, vec![64, 1536]); } fn tiny_config() -> Qwen2VlConfig { diff --git a/src/vision/encoders/qwen2_5_vl.rs b/src/vision/encoders/qwen2_5_vl.rs index ca2112147..c2e726cc4 100644 --- a/src/vision/encoders/qwen2_5_vl.rs +++ b/src/vision/encoders/qwen2_5_vl.rs @@ -590,6 +590,8 @@ pub struct Qwen25VLVisionDiagnostics { pub post_full_layers: Vec>, pub final_interval_layer_indices: Vec, pub post_final_interval_layers: Vec>, + pub pre_probe_layer_indices: Vec, + pub post_pre_probe_layers: Vec>, pub substage_probe_layer_index: usize, pub substage_probe_layer_input: Vec, pub substage_probe_layer_norm1: Vec, @@ -949,8 +951,21 @@ impl Qwen25VLVisionEncoder { .copied() .or(target_full_layer_index); #[cfg(feature = "xla-diagnostics")] + let pre_probe_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_pre_probe_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) { @@ -988,6 +1003,11 @@ impl Qwen25VLVisionEncoder { h.as_ref().expect("Qwen2.5-VL final interval layer state"), )); } + if pre_probe_layer_indices.contains(&layer_num) { + post_pre_probe_layers.push(mlxcel_core::copy( + h.as_ref().expect("Qwen2.5-VL pre-probe layer state"), + )); + } } } @@ -1034,6 +1054,8 @@ impl Qwen25VLVisionEncoder { post_full_layers, final_interval_layer_indices, post_final_interval_layers, + pre_probe_layer_indices, + post_pre_probe_layers, substage_probe_layer_index, substage_probe_layer_input: substage_probe_layer_state.input, substage_probe_layer_norm1: substage_probe_layer_state.norm1, From 9cf66222699398393676d100551831bdf264530d Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 02:54:20 +0900 Subject: [PATCH 14/20] test(xla): probe Qwen layer 17 substages The matching IREE 3.12 CUDA report at 43f7da74 passes layer 16 but first fails catastrophically at layer 17, so the layer-23 probe is no longer the honest diagnostic boundary. Move the existing ordered substage capture to layer 17, retain only the layer 16 and 17 boundary outputs, rename their public/report schema generically, and pin the actual 27-output descriptor order and ranks. Validation: focused Qwen25 tests, diagnostic boundary test, diagnostics example check, formatting, diff checks, and XLA Clippy. Refs #866 --- examples/xla_qwen25_vl_window_check.rs | 22 ++++---- src/lib/mlxcel-xla/src/emitter/qwen2_5_vl.rs | 51 ++++++++++-------- src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs | 55 +++++++++----------- src/vision/encoders/qwen2_5_vl.rs | 38 +++++++++----- 4 files changed, 88 insertions(+), 78 deletions(-) diff --git a/examples/xla_qwen25_vl_window_check.rs b/examples/xla_qwen25_vl_window_check.rs index b93a6bd1f..752421f22 100644 --- a/examples/xla_qwen25_vl_window_check.rs +++ b/examples/xla_qwen25_vl_window_check.rs @@ -538,13 +538,13 @@ fn main() { "captured final diagnostic interval count differs" ); assert_eq!( - iree_vision.pre_probe_layer_indices, eager_vision.pre_probe_layer_indices, - "pre-probe diagnostic interval differs" + iree_vision.diagnostic_layer_indices, eager_vision.diagnostic_layer_indices, + "diagnostic layer interval differs" ); assert_eq!( - iree_vision.post_pre_probe_layers.len(), - eager_vision.post_pre_probe_layers.len(), - "captured pre-probe diagnostic interval count differs" + 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 @@ -601,16 +601,16 @@ fn main() { tolerance, ); } - for (&layer, (actual, expected)) in iree_vision.pre_probe_layer_indices.iter().zip( + for (&layer, (actual, expected)) in iree_vision.diagnostic_layer_indices.iter().zip( iree_vision - .post_pre_probe_layers + .post_diagnostic_layers .iter() - .zip(&eager_vision.post_pre_probe_layers), + .zip(&eager_vision.post_diagnostic_layers), ) { record_stage( &mut reports, &mut comparison_failures, - &format!("vision.post_layer_{layer}_pre_probe"), + &format!("vision.post_layer_{layer}_diagnostic"), actual, &mlx_f32(expected), tolerance, @@ -707,7 +707,7 @@ fn main() { "window_layer": iree_vision.window_layer_index, "full_attention_layers": iree_vision.full_layer_indices, "final_interval_layers": iree_vision.final_interval_layer_indices, - "pre_probe_layers": iree_vision.pre_probe_layer_indices, + "diagnostic_layers": iree_vision.diagnostic_layer_indices, "target_full_layer": target_full_layer_index, "substage_probe_layer": substage_probe_layer_index, "failed_phase": "vision", @@ -866,7 +866,7 @@ fn main() { "window_layer": iree_vision.window_layer_index, "full_attention_layers": iree_vision.full_layer_indices, "final_interval_layers": iree_vision.final_interval_layer_indices, - "pre_probe_layers": iree_vision.pre_probe_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, diff --git a/src/lib/mlxcel-xla/src/emitter/qwen2_5_vl.rs b/src/lib/mlxcel-xla/src/emitter/qwen2_5_vl.rs index 00699bd69..e5add77b5 100644 --- a/src/lib/mlxcel-xla/src/emitter/qwen2_5_vl.rs +++ b/src/lib/mlxcel-xla/src/emitter/qwen2_5_vl.rs @@ -30,7 +30,7 @@ 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) pre_probe_layer_indices: Vec, + pub(crate) diagnostic_layer_indices: Vec, pub(crate) substage_probe_layer_index: usize, } @@ -51,22 +51,30 @@ struct Qwen25VlBlockDiagnostics { #[cfg(feature = "diagnostics")] fn substage_probe_layer(full_layer_indices: &[usize]) -> Option { - full_layer_indices + 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 pre_probe_layers(full_layer_indices: &[usize], probe: usize) -> Vec { +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()) + .map(|previous_full| ((previous_full + 1)..=probe).collect()) .unwrap_or_default() } @@ -385,7 +393,7 @@ fn emit_qwen2_5_vl_inner( #[cfg(feature = "diagnostics")] let mut final_interval_layer_states = Vec::new(); #[cfg(feature = "diagnostics")] - let mut pre_probe_layer_states = Vec::new(); + let mut diagnostic_layer_states = Vec::new(); #[cfg(feature = "diagnostics")] let mut substage_probe_layer_state = None; for layer in 0..config.depth { @@ -427,8 +435,8 @@ fn emit_qwen2_5_vl_inner( if layout.final_interval_layer_indices.contains(&layer) { final_interval_layer_states.push(hidden.clone()); } - if layout.pre_probe_layer_indices.contains(&layer) { - pre_probe_layer_states.push(hidden.clone()); + if layout.diagnostic_layer_indices.contains(&layer) { + diagnostic_layer_states.push(hidden.clone()); } } } @@ -464,12 +472,11 @@ fn emit_qwen2_5_vl_inner( ); outputs.extend(final_interval_layer_states); assert_eq!( - pre_probe_layer_states.len(), - layout.pre_probe_layer_indices.len(), - "diagnostics capture every layer between the preceding full-attention boundary and \ - the substage probe" + diagnostic_layer_states.len(), + layout.diagnostic_layer_indices.len(), + "diagnostics capture every requested boundary layer" ); - outputs.extend(pre_probe_layer_states); + outputs.extend(diagnostic_layer_states); let probe = substage_probe_layer_state.expect("diagnostics capture substage probe layer"); outputs.extend([ probe.input, @@ -562,12 +569,13 @@ pub(crate) fn emit_qwen2_5_vl_diagnostics( (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 pre_probe_layer_indices = pre_probe_layers(&full_layer_indices, substage_probe_layer_index); + 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, - pre_probe_layer_indices, + diagnostic_layer_indices, substage_probe_layer_index, }; Ok(( @@ -578,23 +586,20 @@ pub(crate) fn emit_qwen2_5_vl_diagnostics( #[cfg(all(test, feature = "diagnostics"))] mod tests { - use super::{pre_probe_layers, substage_probe_layer}; + 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(23)); + 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 pre_probe_layers_cover_only_the_gap_after_the_preceding_full_layer() { - assert_eq!( - pre_probe_layers(&[7, 15, 23, 31], 23), - (16..23).collect::>() - ); - assert_eq!(pre_probe_layers(&[0, 1], 0), Vec::::new()); - assert_eq!(pre_probe_layers(&[0], 0), Vec::::new()); + 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/qwen2_vl_runtime.rs b/src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs index 7d555c10c..f685620c3 100644 --- a/src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs +++ b/src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs @@ -94,8 +94,8 @@ pub struct Qwen25VlVisionDiagnostics { pub post_full_layers: Vec>, pub final_interval_layer_indices: Vec, pub post_final_interval_layers: Vec>, - pub pre_probe_layer_indices: Vec, - pub post_pre_probe_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, @@ -165,9 +165,9 @@ fn qwen25_diagnostic_output_specs( ); specs.extend( layout - .pre_probe_layer_indices + .diagnostic_layer_indices .iter() - .map(|layer| state(format!("post_pre_probe_layer.{layer}"))), + .map(|layer| state(format!("post_diagnostic_layer.{layer}"))), ); let probe = layout.substage_probe_layer_index; specs.extend([ @@ -1115,7 +1115,7 @@ impl IreeQwen25VlDiagnosticProjector { 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 pre_probe_layer_count = layout.pre_probe_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; @@ -1207,7 +1207,7 @@ impl IreeQwen25VlDiagnosticProjector { let post_final_interval_layers = decoded .drain(..final_interval_layer_count) .collect::>(); - let post_pre_probe_layers = decoded.drain(..pre_probe_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() @@ -1254,8 +1254,8 @@ impl IreeQwen25VlDiagnosticProjector { post_full_layers, final_interval_layer_indices: layout.final_interval_layer_indices, post_final_interval_layers, - pre_probe_layer_indices: layout.pre_probe_layer_indices, - post_pre_probe_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, @@ -1361,23 +1361,18 @@ mod tests { "post_final_interval_layer.29", "post_final_interval_layer.30", "post_final_interval_layer.31", - "post_pre_probe_layer.16", - "post_pre_probe_layer.17", - "post_pre_probe_layer.18", - "post_pre_probe_layer.19", - "post_pre_probe_layer.20", - "post_pre_probe_layer.21", - "post_pre_probe_layer.22", - "substage_probe_layer.23.input", - "substage_probe_layer.23.norm1", - "substage_probe_layer.23.query", - "substage_probe_layer.23.key", - "substage_probe_layer.23.value", - "substage_probe_layer.23.attention_context", - "substage_probe_layer.23.attention", - "substage_probe_layer.23.post_attention_residual", - "substage_probe_layer.23.norm2", - "substage_probe_layer.23.mlp", + "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", "merger_window_ordered", ] ); @@ -1386,14 +1381,14 @@ mod tests { .iter() .map(|spec| spec.shape.len()) .collect::>(), - [vec![2; 23], vec![3; 3], vec![2; 6],].concat() + [vec![2; 18], vec![3; 3], vec![2; 6],].concat() ); - for spec in &specs[23..=25] { + for spec in &specs[18..=20] { assert_eq!(spec.shape, vec![256, 16, 80]); } - assert_eq!(specs[22].shape, vec![256, 1280]); - assert_eq!(specs[26].shape, vec![256, 1280]); - assert_eq!(specs[31].shape, vec![64, 1536]); + assert_eq!(specs[17].shape, vec![256, 1280]); + assert_eq!(specs[21].shape, vec![256, 1280]); + assert_eq!(specs[26].shape, vec![64, 1536]); } fn tiny_config() -> Qwen2VlConfig { diff --git a/src/vision/encoders/qwen2_5_vl.rs b/src/vision/encoders/qwen2_5_vl.rs index c2e726cc4..a583ee643 100644 --- a/src/vision/encoders/qwen2_5_vl.rs +++ b/src/vision/encoders/qwen2_5_vl.rs @@ -590,8 +590,8 @@ pub struct Qwen25VLVisionDiagnostics { pub post_full_layers: Vec>, pub final_interval_layer_indices: Vec, pub post_final_interval_layers: Vec>, - pub pre_probe_layer_indices: Vec, - pub post_pre_probe_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, @@ -938,33 +938,43 @@ impl Qwen25VLVisionEncoder { let final_interval_layer_indices = target_full_layer_index .map(|target| (final_interval_start..=target).collect::>()) .unwrap_or_default(); - // The multi-media actual-checkpoint oracle passes full-attention layers - // 7 and 15 and first diverges at full-attention layer 23. Probe the - // penultimate configured full-attention layer so the next bounded run + // 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 substage_probe_layer_index = full_layer_indices + let target_probe_boundary = full_layer_indices .iter() .rev() .nth(1) .copied() .or(target_full_layer_index); #[cfg(feature = "xla-diagnostics")] - let pre_probe_layer_indices = substage_probe_layer_index + 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::>()) + .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_pre_probe_layers = Vec::new(); + 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() { @@ -1003,9 +1013,9 @@ impl Qwen25VLVisionEncoder { h.as_ref().expect("Qwen2.5-VL final interval layer state"), )); } - if pre_probe_layer_indices.contains(&layer_num) { - post_pre_probe_layers.push(mlxcel_core::copy( - h.as_ref().expect("Qwen2.5-VL pre-probe 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"), )); } } @@ -1054,8 +1064,8 @@ impl Qwen25VLVisionEncoder { post_full_layers, final_interval_layer_indices, post_final_interval_layers, - pre_probe_layer_indices, - post_pre_probe_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, From d7c67d9471793c43378298ee1c8ba83a8ce5e86a Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 03:01:08 +0900 Subject: [PATCH 15/20] test(xla): trace Qwen layer 17 SwiGLU stages The 93cc1e2f CUDA report passes the layer-17 attention, residual, and second normalization boundaries but first diverges catastrophically at the MLP output, so the next bounded run needs the exact SwiGLU internal seams. Expose gate projection, SiLU gate activation, up projection, gated product, and down projection on eager and IREE diagnostics, derive active output lengths from descriptor shapes, and pin the actual output order and ranks. Validation: focused Qwen25 tests, diagnostic boundary test, diagnostics example check, formatting, diff checks, and XLA Clippy. Refs #866 --- examples/xla_qwen25_vl_window_check.rs | 62 ++++++++++++++- src/lib/mlxcel-xla/src/emitter/qwen2_5_vl.rs | 30 +++++--- src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs | 4 +- src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs | 79 ++++++++++++++++---- src/vision/encoders/qwen2_5_vl.rs | 76 +++++++++++++++++-- 5 files changed, 217 insertions(+), 34 deletions(-) diff --git a/examples/xla_qwen25_vl_window_check.rs b/examples/xla_qwen25_vl_window_check.rs index 752421f22..ad14b34ca 100644 --- a/examples/xla_qwen25_vl_window_check.rs +++ b/examples/xla_qwen25_vl_window_check.rs @@ -413,7 +413,36 @@ fn main() { .as_slice(), ), ("norm2", eager_vision.substage_probe_layer_norm2.as_slice()), - ("mlp", eager_vision.substage_probe_layer_mlp.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() @@ -639,7 +668,36 @@ fn main() { .as_slice(), ), ("norm2", iree_vision.substage_probe_layer_norm2.as_slice()), - ("mlp", iree_vision.substage_probe_layer_mlp.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) diff --git a/src/lib/mlxcel-xla/src/emitter/qwen2_5_vl.rs b/src/lib/mlxcel-xla/src/emitter/qwen2_5_vl.rs index e5add77b5..61705a232 100644 --- a/src/lib/mlxcel-xla/src/emitter/qwen2_5_vl.rs +++ b/src/lib/mlxcel-xla/src/emitter/qwen2_5_vl.rs @@ -45,7 +45,11 @@ struct Qwen25VlBlockDiagnostics { attention: Val, post_attention_residual: Val, norm2: Val, - mlp: Val, + mlp_gate_projection: Val, + mlp_gate_activation: Val, + mlp_up_projection: Val, + mlp_gated_product: Val, + mlp_down_projection: Val, output: Val, } @@ -325,12 +329,12 @@ fn encoder_layer_with_diagnostics( &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 mlp = linear_2d(builder, &activated, &args.take(), &args.take()); - let output = builder.add(&post_attention_residual, &mlp); + 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, @@ -341,7 +345,11 @@ fn encoder_layer_with_diagnostics( attention: attention.output, post_attention_residual, norm2, - mlp, + mlp_gate_projection, + mlp_gate_activation, + mlp_up_projection, + mlp_gated_product, + mlp_down_projection, output, } } @@ -488,7 +496,11 @@ fn emit_qwen2_5_vl_inner( probe.attention, probe.post_attention_residual, probe.norm2, - probe.mlp, + 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 diff --git a/src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs b/src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs index c4e965198..a967a5cc1 100644 --- a/src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs +++ b/src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs @@ -1123,7 +1123,9 @@ mod tests { ); 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>", 5)) + .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(", "); diff --git a/src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs b/src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs index f685620c3..38a0fe876 100644 --- a/src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs +++ b/src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs @@ -106,7 +106,11 @@ pub struct Qwen25VlVisionDiagnostics { 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: 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, @@ -147,6 +151,10 @@ fn qwen25_diagnostic_output_specs( 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)), @@ -191,7 +199,11 @@ fn qwen25_diagnostic_output_specs( "substage_probe_layer.{probe}.post_attention_residual" )), state(format!("substage_probe_layer.{probe}.norm2")), - state(format!("substage_probe_layer.{probe}.mlp")), + 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(), @@ -1167,15 +1179,28 @@ impl IreeQwen25VlDiagnosticProjector { ) }) .collect::, _>>()?; - let actual_state_values = plan - .actual_patches - .checked_mul(self.config.hidden) - .ok_or_else(|| "Qwen2.5-VL diagnostic state size overflowed".to_string())?; // 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. - for state in decoded.iter_mut().take(merger_output_index) { - state.truncate(actual_state_values); + // [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() @@ -1236,9 +1261,21 @@ impl IreeQwen25VlDiagnosticProjector { let substage_probe_layer_norm2 = substage_probe_layer_substages .next() .expect("diagnostic substage probe layer norm2 output"); - let substage_probe_layer_mlp = substage_probe_layer_substages + 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 output"); + .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, @@ -1266,7 +1303,11 @@ impl IreeQwen25VlDiagnosticProjector { substage_probe_layer_attention, substage_probe_layer_post_attention_residual, substage_probe_layer_norm2, - substage_probe_layer_mlp, + 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, @@ -1372,7 +1413,11 @@ mod tests { "substage_probe_layer.17.attention", "substage_probe_layer.17.post_attention_residual", "substage_probe_layer.17.norm2", - "substage_probe_layer.17.mlp", + "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", ] ); @@ -1381,14 +1426,18 @@ mod tests { .iter() .map(|spec| spec.shape.len()) .collect::>(), - [vec![2; 18], vec![3; 3], vec![2; 6],].concat() + [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]); - assert_eq!(specs[26].shape, vec![64, 1536]); + 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 { diff --git a/src/vision/encoders/qwen2_5_vl.rs b/src/vision/encoders/qwen2_5_vl.rs index a583ee643..31ec445c0 100644 --- a/src/vision/encoders/qwen2_5_vl.rs +++ b/src/vision/encoders/qwen2_5_vl.rs @@ -368,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 { @@ -399,6 +408,22 @@ impl VisionMLP { let h = mlxcel_core::multiply(&gate, &up); self.down_proj.forward(&h) } + + #[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. @@ -420,7 +445,11 @@ struct Qwen25VLBlockDiagnostics { attention: Vec, post_attention_residual: Vec, norm2: Vec, - mlp: Vec, + mlp_gate_projection: Vec, + mlp_gate_activation: Vec, + mlp_up_projection: Vec, + mlp_gated_product: Vec, + mlp_down_projection: Vec, } #[cfg(any(test, feature = "xla-diagnostics"))] @@ -505,9 +534,21 @@ impl VisionBlock { 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 = self.mlp.forward(&norm2); - let mlp_snapshot = host_f32_diagnostic_snapshot(&mlp); - let output = mlxcel_core::add(&post_attention_residual, &mlp); + 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 { @@ -521,7 +562,15 @@ impl VisionBlock { attention: attention_snapshot, post_attention_residual: post_attention_residual_snapshot, norm2: norm2_snapshot, - mlp: mlp_snapshot, + mlp_gate_projection: mlp_gate_projection + .expect("diagnostics capture MLP gate projection"), + 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_gated_product: mlp_gated_product + .expect("diagnostics capture MLP gated product"), + mlp_down_projection: mlp_down_projection_snapshot, }, ) } @@ -602,7 +651,11 @@ pub struct Qwen25VLVisionDiagnostics { 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: 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: UniquePtr, pub restored_projection: UniquePtr, } @@ -1078,7 +1131,16 @@ impl Qwen25VLVisionEncoder { 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: substage_probe_layer_state.mlp, + substage_probe_layer_mlp_gate_projection: substage_probe_layer_state + .mlp_gate_projection, + 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_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, From 41fee352e3d34dff3f5cc1b586c73f534821cee0 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 03:17:43 +0900 Subject: [PATCH 16/20] test(xla): distinguish Qwen MLP precision paths The 9173a394 CUDA report shows that layer-17 gate and up projections contain sparse outliers before the gated product amplifies them, but the existing evidence does not prove a production correction. Expose diagnostics-only eager F32 dense controls for the gate and up projections using the same norm2 values and checkpoint weights, and persist every comparison's maximum-error index and values so one bounded run can distinguish the native F16 projection from the IREE accumulation contract. Validation: Qwen25 XLA tests, diagnostic boundary test, diagnostics example check, formatting, and diff checks. Refs #866 --- examples/xla_qwen25_vl_window_check.rs | 47 +++++++++++++++- src/vision/encoders/qwen2_5_vl.rs | 75 ++++++++++++++++++++++++++ 2 files changed, 120 insertions(+), 2 deletions(-) diff --git a/examples/xla_qwen25_vl_window_check.rs b/examples/xla_qwen25_vl_window_check.rs index ad14b34ca..34701210d 100644 --- a/examples/xla_qwen25_vl_window_check.rs +++ b/examples/xla_qwen25_vl_window_check.rs @@ -43,6 +43,9 @@ struct Tolerance { #[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, @@ -190,6 +193,9 @@ fn comparison_stats( } 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, @@ -198,7 +204,7 @@ fn comparison_stats( non_finite: 0, }; let mut finite = 0usize; - for (&observed, &reference) in actual.iter().zip(expected) { + for (index, (&observed, &reference)) in actual.iter().zip(expected).enumerate() { if !observed.is_finite() || !reference.is_finite() { stats.failures += 1; stats.non_finite += 1; @@ -210,7 +216,12 @@ fn comparison_stats( stats.expected_rms += f64::from(reference).powi(2); stats.error_rms += f64::from(observed - reference).powi(2); finite += 1; - stats.max_absolute = stats.max_absolute.max(absolute); + 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; @@ -240,6 +251,9 @@ fn compare_stage( "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, @@ -387,6 +401,12 @@ fn main() { 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. @@ -712,6 +732,22 @@ fn main() { 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, + ); let target_capture_index = iree_vision .full_layer_indices .iter() @@ -768,6 +804,13 @@ fn main() { "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", + }, "failed_phase": "vision", "failures": comparison_failures, "comparisons": reports, diff --git a/src/vision/encoders/qwen2_5_vl.rs b/src/vision/encoders/qwen2_5_vl.rs index 31ec445c0..70c410964 100644 --- a/src/vision/encoders/qwen2_5_vl.rs +++ b/src/vision/encoders/qwen2_5_vl.rs @@ -409,6 +409,41 @@ impl VisionMLP { 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 @@ -446,8 +481,10 @@ struct Qwen25VLBlockDiagnostics { 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, } @@ -534,6 +571,12 @@ impl VisionBlock { 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; @@ -564,10 +607,12 @@ impl VisionBlock { 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, @@ -652,8 +697,10 @@ pub struct Qwen25VLVisionDiagnostics { 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, @@ -1133,10 +1180,14 @@ impl Qwen25VLVisionEncoder { 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 @@ -1167,6 +1218,8 @@ impl super::VisionEncoder for Qwen25VLVisionEncoder { #[cfg(test)] mod tests { + #[cfg(feature = "xla-diagnostics")] + use super::VisionMLP; use super::{host_f32_diagnostic_snapshot, restoration_indices}; #[test] @@ -1182,6 +1235,28 @@ mod tests { 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]); From 961b4366db5f972e41c2b05cd7b59c085d58199a Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 03:34:16 +0900 Subject: [PATCH 17/20] test(xla): replay Qwen MLP from IREE norm The 790c5dad CUDA report shows native F16 and eager dense-F32 gate/up controls remain close to each other while both retain the same 8/23 failures against IREE, disproving projection result precision as the root cause. Replay the captured IREE layer-17 norm2 values through the same checkpoint gate/up weights after matching the graph's F16 input demotion, then compare those eager dense-F32 controls with the existing IREE outputs without changing the IREE ABI. Validation: Qwen25 XLA tests, diagnostic boundary test, diagnostics example check, formatting, and diff checks. Refs #866 --- examples/xla_qwen25_vl_window_check.rs | 37 +++++++++++++++++++ src/vision/encoders/qwen2_5_vl.rs | 51 ++++++++++++++++++++++++++ 2 files changed, 88 insertions(+) diff --git a/examples/xla_qwen25_vl_window_check.rs b/examples/xla_qwen25_vl_window_check.rs index 34701210d..03d76a48c 100644 --- a/examples/xla_qwen25_vl_window_check.rs +++ b/examples/xla_qwen25_vl_window_check.rs @@ -487,6 +487,16 @@ fn main() { 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, @@ -748,6 +758,26 @@ fn main() { 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() @@ -811,6 +841,13 @@ fn main() { "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, diff --git a/src/vision/encoders/qwen2_5_vl.rs b/src/vision/encoders/qwen2_5_vl.rs index 70c410964..745ad83c3 100644 --- a/src/vision/encoders/qwen2_5_vl.rs +++ b/src/vision/encoders/qwen2_5_vl.rs @@ -707,6 +707,12 @@ pub struct Qwen25VLVisionDiagnostics { 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"))] @@ -758,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(); From 786f77c8fd4fb8149761f3b841708ecd2dd82e69 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 03:46:47 +0900 Subject: [PATCH 18/20] fix(xla): preserve Qwen vision F16 state boundaries The 89bed569 CUDA report proves that replaying captured IREE norm2 through the same checkpoint gate and up projections passes with zero failures, excluding projection, reduction, and weight semantics and confirming amplification of inherited norm2 drift. The actual checkpoint keeps its patch and block 0 through 17 weights in F16, and eager MLX preserves F16 block inputs, normalization outputs, and residual states. Keep IREE normalization and contraction accumulation in F32, but explicitly round the patch state, both normalization results, and both residual outputs through F16 while retaining the public F32 ABI. Validation: Qwen25 XLA tests, the diagnostic boundary test, the eager 18-block dtype regression, the diagnostics example check, focused Clippy, formatting, and diff checks. Refs #866 --- src/lib/mlxcel-xla/src/emitter/qwen2_5_vl.rs | 27 +++++++++++-- src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs | 42 ++++++++++++++++++++ src/vision/encoders/qwen2_5_vl.rs | 40 +++++++++++++++++++ 3 files changed, 106 insertions(+), 3 deletions(-) diff --git a/src/lib/mlxcel-xla/src/emitter/qwen2_5_vl.rs b/src/lib/mlxcel-xla/src/emitter/qwen2_5_vl.rs index 61705a232..7b3f43107 100644 --- a/src/lib/mlxcel-xla/src/emitter/qwen2_5_vl.rs +++ b/src/lib/mlxcel-xla/src/emitter/qwen2_5_vl.rs @@ -144,6 +144,21 @@ fn linear_2d(builder: &mut Builder, value: &Val, weight: &Val, bias: &Val) -> Va bias_2d(builder, &value, bias) } +/// Preserve the eager MLX F16 activation boundary while carrying tensors +/// through the IREE ABI as F32. +/// +/// Qwen2.5-VL checkpoint weights and eager block states are F16. The IREE +/// precision policy intentionally computes normalization in F32, but the eager +/// result is rounded back to the activation dtype before the next operation. +/// Keep that observable boundary explicit without demoting the sensitive math. +fn round_f16_activation_boundary(builder: &mut Builder, value: &Val) -> Val { + if builder.precision() != Precision::F16 { + return value.clone(); + } + let narrowed = builder.convert(value, "f16"); + builder.convert(&narrowed, "f32") +} + 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]; @@ -160,7 +175,8 @@ fn rms_norm(builder: &mut Builder, value: &Val, weight: &Val, epsilon: f32) -> V 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) + let normalized = builder.multiply(&normalized, &weight); + round_f16_activation_boundary(builder, &normalized) } fn silu(builder: &mut Builder, value: &Val) -> Val { @@ -302,13 +318,15 @@ fn encoder_layer( 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 residual = round_f16_activation_boundary(builder, &residual); 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) + let output = builder.add(&residual, &down); + round_f16_activation_boundary(builder, &output) } #[cfg(feature = "diagnostics")] @@ -323,6 +341,7 @@ fn encoder_layer_with_diagnostics( 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 post_attention_residual = round_f16_activation_boundary(builder, &post_attention_residual); let norm2 = rms_norm( builder, &post_attention_residual, @@ -335,6 +354,7 @@ fn encoder_layer_with_diagnostics( 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); + let output = round_f16_activation_boundary(builder, &output); Qwen25VlBlockDiagnostics { input: hidden.clone(), norm1, @@ -391,7 +411,8 @@ fn emit_qwen2_5_vl_inner( Ty::f32(vec![patch_bucket, patch_bucket]), "window_attention.bias", ); - let mut hidden = builder.linear_seq(&patches, &patch_weight); + let hidden = builder.linear_seq(&patches, &patch_weight); + let mut hidden = round_f16_activation_boundary(&mut builder, &hidden); #[cfg(feature = "diagnostics")] let reordered_patch_embedding = diagnostic_layout.map(|_| hidden.clone()); #[cfg(feature = "diagnostics")] diff --git a/src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs b/src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs index a967a5cc1..6c880c29e 100644 --- a/src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs +++ b/src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs @@ -1106,6 +1106,48 @@ mod tests { } } + #[test] + fn qwen25_f16_precision_rounds_block_state_boundaries_through_layer17() { + let mut config = qwen25_config(); + config.depth = 18; + + let f32 = emit_qwen2_vl_with(&config, 16, Precision::F32); + let f16 = emit_qwen2_vl_with(&config, 16, Precision::F16); + let f16_to_f32 = |module: &str| { + module + .lines() + .filter(|line| { + line.contains("stablehlo.convert") + && line.contains("xf16>) -> tensor<") + && line.ends_with("xf32>") + }) + .count() + }; + + assert_eq!(f16_to_f32(&f32), 0); + // One patch-state boundary, four boundaries per block (norm1, + // post-attention residual, norm2, and block output) across blocks + // 0..=17, plus the merger norm boundary. + assert_eq!(f16_to_f32(&f16), 1 + 4 * 18 + 1); + 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.ends_with("f32>"), + "contraction result must remain F32 before its activation boundary: {line}" + ); + } else { + assert!( + !line.contains("f16"), + "sensitive math must remain F32 before its activation boundary: {line}" + ); + } + } + } + #[cfg(feature = "diagnostics")] #[test] fn qwen25_diagnostics_share_the_production_graph_and_expose_ordered_seams() { diff --git a/src/vision/encoders/qwen2_5_vl.rs b/src/vision/encoders/qwen2_5_vl.rs index 745ad83c3..febcf2df5 100644 --- a/src/vision/encoders/qwen2_5_vl.rs +++ b/src/vision/encoders/qwen2_5_vl.rs @@ -1286,6 +1286,46 @@ mod tests { assert_eq!(original, vec!["zero", "one", "two", "three"]); } + #[test] + fn eager_f16_block_state_norm_and_residual_dtypes_hold_through_layer17() { + let input = mlxcel_core::from_slice_f32(&[1.0, -2.0, 3.0, -4.0], &[1, 4]); + let mut state = mlxcel_core::astype(&input, mlxcel_core::dtype::FLOAT16); + let weight = mlxcel_core::ones(&[4], mlxcel_core::dtype::FLOAT16); + let branch = mlxcel_core::zeros(&[1, 4], mlxcel_core::dtype::FLOAT16); + + for layer in 0..=17 { + assert_eq!( + mlxcel_core::array_dtype(&state), + mlxcel_core::dtype::FLOAT16, + "block {layer} input must remain F16" + ); + let norm1 = mlxcel_core::fast_rms_norm(&state, &weight, 1e-6); + assert_eq!( + mlxcel_core::array_dtype(&norm1), + mlxcel_core::dtype::FLOAT16, + "block {layer} norm1 must restore the activation dtype" + ); + let post_attention_residual = mlxcel_core::add(&state, &branch); + assert_eq!( + mlxcel_core::array_dtype(&post_attention_residual), + mlxcel_core::dtype::FLOAT16, + "block {layer} attention residual must remain F16" + ); + let norm2 = mlxcel_core::fast_rms_norm(&post_attention_residual, &weight, 1e-6); + assert_eq!( + mlxcel_core::array_dtype(&norm2), + mlxcel_core::dtype::FLOAT16, + "block {layer} norm2 must restore the activation dtype" + ); + state = mlxcel_core::add(&post_attention_residual, &branch); + assert_eq!( + mlxcel_core::array_dtype(&state), + mlxcel_core::dtype::FLOAT16, + "block {layer} output must remain F16" + ); + } + } + #[cfg(feature = "xla-diagnostics")] #[test] fn dense_f32_projection_control_preserves_linear_order_and_dtype() { From 3756d19c8558d42b63eac849a7fb97a14e29026d Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 03:54:16 +0900 Subject: [PATCH 19/20] fix(xla): revert unsafe Qwen F16 state rounding The 2b1c1f62 CUDA run preserved patch and layer-15 acceptance but regressed layer 17 from 392 failures at max_abs 104.7239 to 1,380 failures at max_abs 320, the down projection from 338/104.0575 to 1,504/319.2441, and the merger from 170/0.8468 to 974/1.5832. Revert only the coarse patch, normalization, and residual F16 round-trip boundaries while preserving the diagnostics through 89bed569. Matching eager F16 safely requires an operation-by-operation backend rounding contract; no further production approximation or tolerance waiver is justified by current evidence. Validation: Qwen25 XLA tests, diagnostic boundary test, diagnostics example check, formatting, and diff checks. Refs #866 --- src/lib/mlxcel-xla/src/emitter/qwen2_5_vl.rs | 27 ++----------- src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs | 42 -------------------- src/vision/encoders/qwen2_5_vl.rs | 40 ------------------- 3 files changed, 3 insertions(+), 106 deletions(-) diff --git a/src/lib/mlxcel-xla/src/emitter/qwen2_5_vl.rs b/src/lib/mlxcel-xla/src/emitter/qwen2_5_vl.rs index 7b3f43107..61705a232 100644 --- a/src/lib/mlxcel-xla/src/emitter/qwen2_5_vl.rs +++ b/src/lib/mlxcel-xla/src/emitter/qwen2_5_vl.rs @@ -144,21 +144,6 @@ fn linear_2d(builder: &mut Builder, value: &Val, weight: &Val, bias: &Val) -> Va bias_2d(builder, &value, bias) } -/// Preserve the eager MLX F16 activation boundary while carrying tensors -/// through the IREE ABI as F32. -/// -/// Qwen2.5-VL checkpoint weights and eager block states are F16. The IREE -/// precision policy intentionally computes normalization in F32, but the eager -/// result is rounded back to the activation dtype before the next operation. -/// Keep that observable boundary explicit without demoting the sensitive math. -fn round_f16_activation_boundary(builder: &mut Builder, value: &Val) -> Val { - if builder.precision() != Precision::F16 { - return value.clone(); - } - let narrowed = builder.convert(value, "f16"); - builder.convert(&narrowed, "f32") -} - 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]; @@ -175,8 +160,7 @@ fn rms_norm(builder: &mut Builder, value: &Val, weight: &Val, epsilon: f32) -> V let inverse = builder.broadcast(&inverse, &[0], vec![rows, width]); let normalized = builder.multiply(value, &inverse); let weight = builder.broadcast(weight, &[1], vec![rows, width]); - let normalized = builder.multiply(&normalized, &weight); - round_f16_activation_boundary(builder, &normalized) + builder.multiply(&normalized, &weight) } fn silu(builder: &mut Builder, value: &Val) -> Val { @@ -318,15 +302,13 @@ fn encoder_layer( 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 residual = round_f16_activation_boundary(builder, &residual); 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()); - let output = builder.add(&residual, &down); - round_f16_activation_boundary(builder, &output) + builder.add(&residual, &down) } #[cfg(feature = "diagnostics")] @@ -341,7 +323,6 @@ fn encoder_layer_with_diagnostics( 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 post_attention_residual = round_f16_activation_boundary(builder, &post_attention_residual); let norm2 = rms_norm( builder, &post_attention_residual, @@ -354,7 +335,6 @@ fn encoder_layer_with_diagnostics( 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); - let output = round_f16_activation_boundary(builder, &output); Qwen25VlBlockDiagnostics { input: hidden.clone(), norm1, @@ -411,8 +391,7 @@ fn emit_qwen2_5_vl_inner( Ty::f32(vec![patch_bucket, patch_bucket]), "window_attention.bias", ); - let hidden = builder.linear_seq(&patches, &patch_weight); - let mut hidden = round_f16_activation_boundary(&mut builder, &hidden); + let mut hidden = builder.linear_seq(&patches, &patch_weight); #[cfg(feature = "diagnostics")] let reordered_patch_embedding = diagnostic_layout.map(|_| hidden.clone()); #[cfg(feature = "diagnostics")] diff --git a/src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs b/src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs index 6c880c29e..a967a5cc1 100644 --- a/src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs +++ b/src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs @@ -1106,48 +1106,6 @@ mod tests { } } - #[test] - fn qwen25_f16_precision_rounds_block_state_boundaries_through_layer17() { - let mut config = qwen25_config(); - config.depth = 18; - - let f32 = emit_qwen2_vl_with(&config, 16, Precision::F32); - let f16 = emit_qwen2_vl_with(&config, 16, Precision::F16); - let f16_to_f32 = |module: &str| { - module - .lines() - .filter(|line| { - line.contains("stablehlo.convert") - && line.contains("xf16>) -> tensor<") - && line.ends_with("xf32>") - }) - .count() - }; - - assert_eq!(f16_to_f32(&f32), 0); - // One patch-state boundary, four boundaries per block (norm1, - // post-attention residual, norm2, and block output) across blocks - // 0..=17, plus the merger norm boundary. - assert_eq!(f16_to_f32(&f16), 1 + 4 * 18 + 1); - 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.ends_with("f32>"), - "contraction result must remain F32 before its activation boundary: {line}" - ); - } else { - assert!( - !line.contains("f16"), - "sensitive math must remain F32 before its activation boundary: {line}" - ); - } - } - } - #[cfg(feature = "diagnostics")] #[test] fn qwen25_diagnostics_share_the_production_graph_and_expose_ordered_seams() { diff --git a/src/vision/encoders/qwen2_5_vl.rs b/src/vision/encoders/qwen2_5_vl.rs index febcf2df5..745ad83c3 100644 --- a/src/vision/encoders/qwen2_5_vl.rs +++ b/src/vision/encoders/qwen2_5_vl.rs @@ -1286,46 +1286,6 @@ mod tests { assert_eq!(original, vec!["zero", "one", "two", "three"]); } - #[test] - fn eager_f16_block_state_norm_and_residual_dtypes_hold_through_layer17() { - let input = mlxcel_core::from_slice_f32(&[1.0, -2.0, 3.0, -4.0], &[1, 4]); - let mut state = mlxcel_core::astype(&input, mlxcel_core::dtype::FLOAT16); - let weight = mlxcel_core::ones(&[4], mlxcel_core::dtype::FLOAT16); - let branch = mlxcel_core::zeros(&[1, 4], mlxcel_core::dtype::FLOAT16); - - for layer in 0..=17 { - assert_eq!( - mlxcel_core::array_dtype(&state), - mlxcel_core::dtype::FLOAT16, - "block {layer} input must remain F16" - ); - let norm1 = mlxcel_core::fast_rms_norm(&state, &weight, 1e-6); - assert_eq!( - mlxcel_core::array_dtype(&norm1), - mlxcel_core::dtype::FLOAT16, - "block {layer} norm1 must restore the activation dtype" - ); - let post_attention_residual = mlxcel_core::add(&state, &branch); - assert_eq!( - mlxcel_core::array_dtype(&post_attention_residual), - mlxcel_core::dtype::FLOAT16, - "block {layer} attention residual must remain F16" - ); - let norm2 = mlxcel_core::fast_rms_norm(&post_attention_residual, &weight, 1e-6); - assert_eq!( - mlxcel_core::array_dtype(&norm2), - mlxcel_core::dtype::FLOAT16, - "block {layer} norm2 must restore the activation dtype" - ); - state = mlxcel_core::add(&post_attention_residual, &branch); - assert_eq!( - mlxcel_core::array_dtype(&state), - mlxcel_core::dtype::FLOAT16, - "block {layer} output must remain F16" - ); - } - } - #[cfg(feature = "xla-diagnostics")] #[test] fn dense_f32_projection_control_preserves_linear_order_and_dtype() { From 02d41a18178278c0fa1457bfc28d7eeebb3e7e25 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 09:09:09 +0900 Subject: [PATCH 20/20] fix(xla): mark Qwen diagnostics unqualified Keep the rebased Qwen2.5-VL diagnostics on the explicit legacy artifact path until the full numeric contract and post-CUDA-kernel oracle are qualified.\n\nRefs #866 --- src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs b/src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs index 38a0fe876..bb993dca8 100644 --- a/src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs +++ b/src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs @@ -822,7 +822,7 @@ fn compile_and_load_diagnostics( 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( + let contract = AuxiliaryArtifactContract::new_legacy_unqualified( ENTRY_NAME, format!( "{};{};precision={precision:?};diagnostics=window-full-restoration;checkpoint_schema_sha256={}",