From 1d66bc7e209d9efd0e4a3c0411a354660fb38562 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 24 Jul 2026 19:57:54 +0900 Subject: [PATCH 01/10] feat: add Qwen3-VL XLA DeepStack vision path Add a resident IREE Qwen3-VL vision emitter and runtime with learned-position interpolation, packed nonwindowed attention, the main merger, and every ordered DeepStack branch from the immutable checkpoint. Derive the dense nested text contract, preserve interleaved M-RoPE, route CLI and server image requests through owned DeepStack prefill payloads, and fail closed when an ordinary prefill would omit side features. Validate emitter/runtime fixtures, host and server routing boundaries, the XLA-IREE build and clippy paths, and one pinned CUDA/IREE image generation that returned `Orange`. Refs #868 --- src/backend/session.rs | 28 + src/commands/generate.rs | 5 +- src/lib/mlxcel-xla/src/emitter/config.rs | 67 +- src/lib/mlxcel-xla/src/emitter/mod.rs | 65 ++ src/lib/mlxcel-xla/src/emitter/qwen3_vl.rs | 1137 ++++++++++++++++++++ src/lib/mlxcel-xla/src/lib.rs | 6 + src/lib/mlxcel-xla/src/qwen3_vl_runtime.rs | 839 +++++++++++++++ src/loading/mod.rs | 2 + src/loading/vlm.rs | 4 +- src/loading/vlm_qwen.rs | 121 ++- src/multimodal/host_preprocessor.rs | 329 +++++- src/multimodal/host_preprocessor_export.rs | 101 +- src/multimodal/host_preprocessor_tests.rs | 20 + src/server/batch/xla_preprocess.rs | 35 +- src/server/batch/xla_worker.rs | 21 +- src/server/batch/xla_worker_admission.rs | 21 +- src/server/batch/xla_worker_tests.rs | 121 ++- 17 files changed, 2864 insertions(+), 58 deletions(-) create mode 100644 src/lib/mlxcel-xla/src/emitter/qwen3_vl.rs create mode 100644 src/lib/mlxcel-xla/src/qwen3_vl_runtime.rs diff --git a/src/backend/session.rs b/src/backend/session.rs index f09655e7..0bae0a99 100644 --- a/src/backend/session.rs +++ b/src/backend/session.rs @@ -133,6 +133,34 @@ impl XlaBackendSession { .prepare(token_ids, images) .map_err(|error| error.to_string()) } + + /// Prepare and execute an image request through its architecture-specific + /// owned prefill entry. Qwen3-VL is routed through the distinct DeepStack + /// entry before the ordinary fallback is considered. + pub fn generate_images_greedy( + &mut self, + token_ids: &[i32], + images: &[image::DynamicImage], + max_tokens: usize, + ) -> Result, String> { + let preprocessor = self.image_preprocessor.as_ref().ok_or_else(|| { + "the loaded OpenXLA model/runtime bundle does not support image input".to_string() + })?; + let eos = self.engine.eos_token_ids().to_vec(); + if let Some(request) = preprocessor + .prepare_deepstack(token_ids, images) + .map_err(|error| error.to_string())? + { + return self + .engine + .generate_deepstack_prepared_greedy(&request, max_tokens, &eos); + } + let prepared = preprocessor + .prepare(token_ids, images) + .map_err(|error| error.to_string())?; + self.engine + .generate_prepared_greedy(&prepared, max_tokens, &eos) + } } #[cfg(feature = "xla-backend")] diff --git a/src/commands/generate.rs b/src/commands/generate.rs index e4fe4dd1..eab9536d 100644 --- a/src/commands/generate.rs +++ b/src/commands/generate.rs @@ -1205,10 +1205,7 @@ fn generate_xla( "the loaded OpenXLA model/runtime bundle does not support image input" ); let images = decode_xla_cli_images(image_paths)?; - let prepared = s - .prepare_images(prompt_tokens, &images) - .map_err(|error| anyhow!("OpenXLA image preprocessing failed: {error}"))?; - s.generate_prepared_greedy(&prepared, max_tokens, &eos) + s.generate_images_greedy(prompt_tokens, &images, max_tokens) .map_err(|error| anyhow!("OpenXLA image generation failed: {error}"))? } } diff --git a/src/lib/mlxcel-xla/src/emitter/config.rs b/src/lib/mlxcel-xla/src/emitter/config.rs index e5c450ea..47f9dc45 100644 --- a/src/lib/mlxcel-xla/src/emitter/config.rs +++ b/src/lib/mlxcel-xla/src/emitter/config.rs @@ -482,14 +482,49 @@ impl Config { serde_json::from_str(s).map_err(|e| format!("parse config.json: {e}"))?; let wrapper_model_type = root.get("model_type").and_then(serde_json::Value::as_str); let is_phi4mm = wrapper_model_type == Some("phi4mm"); - let v = if matches!(wrapper_model_type, Some("llava" | "llava_next")) { + let inferred_qwen3_vl_deepstack = if wrapper_model_type == Some("qwen3_vl") { + let vision = root + .get("vision_config") + .and_then(serde_json::Value::as_object) + .ok_or_else(|| { + "Qwen3-VL config.json missing object `vision_config` for DeepStack".to_string() + })?; + let visual_layers = vision + .get("deepstack_visual_indexes") + .and_then(serde_json::Value::as_array) + .ok_or_else(|| { + "Qwen3-VL config.json missing array \ + `vision_config.deepstack_visual_indexes`" + .to_string() + })?; + if visual_layers.is_empty() { + return Err( + "Qwen3-VL vision_config.deepstack_visual_indexes must not be empty".to_string(), + ); + } + Some(DeepStackConfig { + // Qwen3-VL injects the ordered vision branches into the first + // language layers. The actual vision layer indices remain part + // of the vision artifact identity. + target_layer_indices: (0..visual_layers.len()).collect(), + max_visual_positions: super::qwen3_vl::QWEN3_VL_MAX_MERGED_VISUAL_POSITIONS, + }) + } else { + None + }; + let v = if matches!( + wrapper_model_type, + Some("llava" | "llava_next" | "qwen3_vl") + ) { let mut text = root .get("text_config") .and_then(serde_json::Value::as_object) .cloned() .ok_or_else(|| { - "LLaVA config.json missing object `text_config` for the XLA text graph" - .to_string() + format!( + "{} config.json missing object `text_config` for the XLA text graph", + wrapper_model_type.unwrap_or("VLM") + ) })?; // mlx-community quantized VLMs commonly keep the affine scheme at // the wrapper level even though the tensors belong to the nested @@ -580,7 +615,7 @@ impl Config { let head_dim = ou("head_dim").unwrap_or(hidden / n_q.max(1)); // ExaOne 3.x uses `num_layers` in place of `num_hidden_layers`. let n_layers = u_any(&["num_hidden_layers", "num_layers"])?; - let deepstack = match ( + let declared_deepstack = match ( v.get("deepstack_language_layer_indices"), v.get("deepstack_max_visual_positions"), ) { @@ -624,6 +659,20 @@ impl Config { ); } }; + let deepstack = match (inferred_qwen3_vl_deepstack, declared_deepstack) { + (Some(inferred), None) => { + inferred.validate_structure(n_layers)?; + Some(inferred) + } + (Some(_), Some(_)) => { + return Err( + "Qwen3-VL DeepStack mapping is architecture-defined and must not be \ + overridden in text_config" + .to_string(), + ); + } + (None, declared) => declared, + }; // Norm epsilon: `rms_norm_eps` (RMSNorm archs), else `layer_norm_eps` // (Cohere / StableLM LayerNorm), else `layer_norm_epsilon` (ExaOne), else @@ -723,7 +772,7 @@ impl Config { // is therefore the ordinary Qwen2 graph with M-RoPE positions. qkv_bias = true; } - Some("qwen3") => { + Some("qwen3" | "qwen3_vl_text") => { // Qwen3 drops the Qwen2 bias and adds a per-head q/k RMSNorm (raw // weight, over head_dim) before RoPE. qk_norm = Some(QkNorm { @@ -1197,8 +1246,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" | "qwen2_vl" | "qwen3" | "qwen3_vl_text") + ); if !supported_family { return Err(format!( "M-RoPE sections are only supported by the Qwen2/Qwen3 language backbone, got model_type={model_type:?}" @@ -1207,7 +1258,7 @@ impl Config { let interleaved = rope_scaling .and_then(|scaling| scaling.get("mrope_interleaved")) .and_then(serde_json::Value::as_bool) - .unwrap_or(model_type == Some("qwen3")); + .unwrap_or(matches!(model_type, Some("qwen3" | "qwen3_vl_text"))); Some(MropeConfig { sections, layout: if interleaved { diff --git a/src/lib/mlxcel-xla/src/emitter/mod.rs b/src/lib/mlxcel-xla/src/emitter/mod.rs index 0065385c..e9f68508 100644 --- a/src/lib/mlxcel-xla/src/emitter/mod.rs +++ b/src/lib/mlxcel-xla/src/emitter/mod.rs @@ -60,6 +60,7 @@ mod moe; pub(crate) mod numeric_ops; mod phi4_audio; mod qwen2_vl; +mod qwen3_vl; mod rope; mod vision; mod vision_config; @@ -128,6 +129,13 @@ pub(crate) use qwen2_vl::{ QWEN2_VL_PATCH_BUCKETS, Qwen2VlConfig, Qwen2VlGridPlan, Qwen2VlHostInputs, Qwen2VlWeightSpec, prepare_qwen2_vl_host_inputs, }; +#[allow(unused_imports)] +pub(crate) use qwen3_vl::emit_qwen3_vl; +#[allow(unused_imports)] +pub(crate) use qwen3_vl::{ + QWEN3_VL_PATCH_BUCKETS, Qwen3VlConfig, Qwen3VlGridPlan, Qwen3VlHostInputs, Qwen3VlWeightSpec, + prepare_qwen3_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. #[allow(unused_imports)] @@ -1196,6 +1204,63 @@ mod tests { } } + #[test] + fn qwen3_vl_wrapper_derives_dense_deepstack_text_contract() { + let wrapper = r#"{ + "model_type":"qwen3_vl", + "quantization":{"bits":4,"group_size":64}, + "text_config":{ + "model_type":"qwen3_vl_text", + "hidden_size":12, + "num_attention_heads":2, + "num_key_value_heads":1, + "head_dim":6, + "intermediate_size":16, + "num_hidden_layers":4, + "rms_norm_eps":1e-6, + "rope_theta":1e6, + "rope_scaling":{ + "rope_type":"default", + "mrope_interleaved":true, + "mrope_section":[1,1,1] + }, + "vocab_size":10, + "tie_word_embeddings":true + }, + "vision_config":{"deepstack_visual_indexes":[5,11,17]} + }"#; + let cfg = Config::from_json_str(wrapper).expect("Qwen3-VL wrapper text graph"); + assert_eq!( + cfg.deepstack, + Some(DeepStackConfig { + target_layer_indices: vec![0, 1, 2], + max_visual_positions: 64, + }) + ); + assert_eq!( + cfg.quantization, + Some(QuantConfig { + bits: 4, + group_size: 64 + }) + ); + assert_eq!( + cfg.mrope.expect("Qwen3-VL M-RoPE").layout, + MropeLayout::Interleaved + ); + assert!(cfg.qk_norm.is_some(), "Qwen3-VL uses Qwen3 q/k RMSNorm"); + + let overridden = wrapper.replace( + "\"tie_word_embeddings\":true", + "\"tie_word_embeddings\":true,\n\ + \"deepstack_language_layer_indices\":[0,1,2],\n\ + \"deepstack_max_visual_positions\":64", + ); + let error = Config::from_json_str(&overridden) + .expect_err("architecture-defined mapping must reject overrides"); + assert!(error.contains("architecture-defined")); + } + #[test] fn prefill_embeddings_metadata_and_bias_fail_closed() { let cfg = qwen_like(false); diff --git a/src/lib/mlxcel-xla/src/emitter/qwen3_vl.rs b/src/lib/mlxcel-xla/src/emitter/qwen3_vl.rs new file mode 100644 index 00000000..2ea3bfa0 --- /dev/null +++ b/src/lib/mlxcel-xla/src/emitter/qwen3_vl.rs @@ -0,0 +1,1137 @@ +// 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. + +//! Static-shape contract for the Qwen3-VL vision tower. +//! +//! The processor keeps the checkpoint's dynamic smart-resize policy, while the +//! compiled graph admits only this finite set of patch capacities. Inputs are +//! padded to the selected capacity and carry their actual patch count plus +//! packed-sequence boundaries. No graph is compiled for an arbitrary image +//! resolution. + +use std::path::Path; + +use serde_json::Value; + +use super::builder::{Builder, Ty, Val}; + +/// Qualified flattened-patch capacities for the pinned Qwen3-VL image path. +/// +/// Every capacity is divisible by `spatial_merge_size²` for the checkpoint's +/// 2x2 merger. Larger grids fail explicitly instead of triggering an unbounded +/// compile or falling back to the MLX vision encoder. +pub(crate) const QWEN3_VL_PATCH_BUCKETS: [usize; 3] = [16, 64, 256]; + +/// Compact DeepStack side-input capacity after the required 2x2 merge. +pub(crate) const QWEN3_VL_MAX_MERGED_VISUAL_POSITIONS: usize = 256 / 4; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct Qwen3VlWeightSpec { + pub(crate) name: String, + /// Logical dequantized F32 shape consumed by StableHLO. + pub(crate) shape: Vec, +} + +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct Qwen3VlConfig { + pub(crate) depth: usize, + pub(crate) hidden: usize, + pub(crate) intermediate: usize, + pub(crate) heads: usize, + pub(crate) patch_size: usize, + pub(crate) temporal_patch_size: usize, + pub(crate) spatial_merge_size: usize, + pub(crate) channels: usize, + pub(crate) text_hidden: usize, + pub(crate) num_position_embeddings: usize, + pub(crate) deepstack_visual_indexes: Vec, + pub(crate) layer_norm_eps: f32, + pub(crate) quantization: Option<(usize, usize)>, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct Qwen3VlGridPlan { + pub(crate) patch_bucket: usize, + pub(crate) actual_patches: usize, + pub(crate) packed_cu_seqlens: Vec, + pub(crate) merged_tokens_per_image: Vec, +} + +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct Qwen3VlHostInputs { + pub(crate) plan: Qwen3VlGridPlan, + pub(crate) patches: Vec, + pub(crate) vision_rope_freqs: Vec, + pub(crate) packed_attention_bias: Vec, + pub(crate) position_indices: Vec, + pub(crate) position_weights: Vec, +} + +struct Args { + values: Vec, + declarations: Vec, + cursor: usize, +} + +impl Args { + fn new(specs: &[Qwen3VlWeightSpec]) -> Self { + let mut values = Vec::with_capacity(specs.len() + 3); + let mut declarations = Vec::with_capacity(specs.len() + 3); + 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 positive_usize(object: &serde_json::Map, field: &str) -> Result { + let value = object + .get(field) + .and_then(Value::as_u64) + .ok_or_else(|| format!("config.json vision_config.{field} must be a positive integer"))?; + let value = usize::try_from(value) + .map_err(|_| format!("config.json vision_config.{field} does not fit usize"))?; + if value == 0 { + return Err(format!( + "config.json vision_config.{field} must be greater than zero" + )); + } + Ok(value) +} + +impl Qwen3VlConfig { + 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) + .map_err(|error| format!("{}: {error}", path.display()))?; + Self::from_json_str(&text) + } + + 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("qwen3_vl") { + return Err("Qwen3-VL IREE vision requires model_type=qwen3_vl".to_string()); + } + 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, "hidden_size")?; + let heads = positive_usize(vision, "num_heads")?; + if hidden % heads != 0 { + return Err(format!( + "Qwen3-VL vision hidden_size={hidden} is not divisible by num_heads={heads}" + )); + } + let intermediate = positive_usize(vision, "intermediate_size")?; + 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")?; + let channels = vision + .get("in_chans") + .or_else(|| vision.get("in_channels")) + .and_then(Value::as_u64) + .map_or(Ok(3usize), |value| { + usize::try_from(value).map_err(|_| { + "config.json vision_config.in_chans does not fit usize".to_string() + }) + })?; + if channels == 0 { + return Err("config.json vision_config.in_chans must be positive".to_string()); + } + let text = root + .get("text_config") + .and_then(Value::as_object) + .ok_or_else(|| "config.json text_config must be an object".to_string())?; + if text.get("model_type").and_then(Value::as_str) != Some("qwen3_vl_text") { + return Err("Qwen3-VL IREE supports the dense qwen3_vl_text decoder only".to_string()); + } + let text_hidden = text + .get("hidden_size") + .and_then(Value::as_u64) + .ok_or_else(|| "config.json hidden_size must be a positive integer".to_string()) + .and_then(|value| { + usize::try_from(value) + .map_err(|_| "config.json hidden_size does not fit usize".to_string()) + })?; + if text_hidden == 0 { + return Err("config.json text_config.hidden_size must be positive".to_string()); + } + let out_hidden = positive_usize(vision, "out_hidden_size")?; + if out_hidden != text_hidden { + return Err(format!( + "Qwen3-VL vision out_hidden_size={out_hidden} disagrees with text hidden_size={text_hidden}" + )); + } + let num_position_embeddings = positive_usize(vision, "num_position_embeddings")?; + let position_side = (num_position_embeddings as f64).sqrt() as usize; + if position_side.checked_mul(position_side) != Some(num_position_embeddings) { + return Err(format!( + "Qwen3-VL num_position_embeddings={num_position_embeddings} must be a square table" + )); + } + let deepstack_visual_indexes = vision + .get("deepstack_visual_indexes") + .and_then(Value::as_array) + .ok_or_else(|| { + "config.json vision_config.deepstack_visual_indexes must be an array".to_string() + })? + .iter() + .map(|value| { + value + .as_u64() + .and_then(|value| usize::try_from(value).ok()) + .ok_or_else(|| { + "Qwen3-VL deepstack_visual_indexes must contain non-negative integers" + .to_string() + }) + }) + .collect::, _>>()?; + if deepstack_visual_indexes.is_empty() + || deepstack_visual_indexes + .windows(2) + .any(|pair| pair[0] >= pair[1]) + || deepstack_visual_indexes.iter().any(|&index| index >= depth) + { + return Err(format!( + "Qwen3-VL deepstack_visual_indexes={deepstack_visual_indexes:?} must be non-empty, strictly increasing, and within depth={depth}" + )); + } + if spatial_merge_size != 2 { + return Err(format!( + "Qwen3-VL IREE vision currently qualifies spatial_merge_size=2, got {spatial_merge_size}" + )); + } + if temporal_patch_size != 2 { + return Err(format!( + "Qwen3-VL IREE image vision currently qualifies temporal_patch_size=2, got {temporal_patch_size}" + )); + } + let quantization = match root.get("quantization") { + None => None, + Some(value) => { + let object = value + .as_object() + .ok_or_else(|| "config.json quantization must be an object".to_string())?; + let bits = positive_usize(object, "bits") + .map_err(|error| error.replace("vision_config.", "quantization."))?; + let group_size = positive_usize(object, "group_size") + .map_err(|error| error.replace("vision_config.", "quantization."))?; + if !matches!(bits, 4 | 8) { + return Err(format!( + "Qwen3-VL IREE vision supports 4-bit or 8-bit affine weights, got {bits}" + )); + } + if hidden % group_size != 0 || intermediate % group_size != 0 { + return Err(format!( + "Qwen3-VL vision dimensions hidden={hidden}, intermediate={intermediate} must be divisible by quantization group_size={group_size}" + )); + } + Some((bits, group_size)) + } + }; + Ok(Self { + depth, + hidden, + intermediate, + heads, + patch_size, + temporal_patch_size, + spatial_merge_size, + channels, + text_hidden, + num_position_embeddings, + deepstack_visual_indexes, + layer_norm_eps: 1e-6, + quantization, + }) + } + + pub(crate) fn bucket_for_patches(&self, actual_patches: usize) -> Result { + if actual_patches == 0 { + return Err("Qwen3-VL image grid must contain at least one patch".to_string()); + } + let merge_area = self + .spatial_merge_size + .checked_mul(self.spatial_merge_size) + .ok_or_else(|| "Qwen3-VL merge area overflowed".to_string())?; + if !actual_patches.is_multiple_of(merge_area) { + return Err(format!( + "Qwen3-VL image grid has {actual_patches} patches, which is not divisible by spatial_merge_size²={merge_area}" + )); + } + QWEN3_VL_PATCH_BUCKETS + .iter() + .copied() + .find(|&bucket| actual_patches <= bucket) + .ok_or_else(|| { + format!( + "Qwen3-VL image grid has {actual_patches} patches, exceeding qualified capacity {}", + QWEN3_VL_PATCH_BUCKETS[QWEN3_VL_PATCH_BUCKETS.len() - 1] + ) + }) + } + + /// Validate each image grid before selecting one packed static bucket. + /// + /// Per-axis merge divisibility is stronger than checking only the total: + /// `1x16` has a merge-area-divisible product but cannot form 2x2 groups. + /// Temporal grids are rejected by the image-only issue scope. + pub(crate) fn plan_image_grids( + &self, + grids: &[(i32, i32, i32)], + ) -> Result { + if grids.is_empty() { + return Err("Qwen3-VL image execution requires at least one grid".to_string()); + } + let mut actual_patches = 0usize; + let mut packed_cu_seqlens = vec![0usize]; + let mut merged_tokens_per_image = Vec::with_capacity(grids.len()); + for (index, &(temporal, height, width)) in grids.iter().enumerate() { + if temporal != 1 { + return Err(format!( + "Qwen3-VL grid {index} has temporal size {temporal}; video/temporal grids are unsupported by the image-only XLA path" + )); + } + let height = usize::try_from(height) + .map_err(|_| format!("Qwen3-VL grid {index} has non-positive height {height}"))?; + let width = usize::try_from(width) + .map_err(|_| format!("Qwen3-VL grid {index} has non-positive width {width}"))?; + if height == 0 || width == 0 { + return Err(format!( + "Qwen3-VL grid {index} dimensions must be positive, got {height}x{width}" + )); + } + if !height.is_multiple_of(self.spatial_merge_size) + || !width.is_multiple_of(self.spatial_merge_size) + { + return Err(format!( + "Qwen3-VL grid {index} {height}x{width} must be divisible on each spatial axis by spatial_merge_size={}", + self.spatial_merge_size + )); + } + let patches = height + .checked_mul(width) + .ok_or_else(|| format!("Qwen3-VL grid {index} patch count overflowed"))?; + actual_patches = actual_patches + .checked_add(patches) + .ok_or_else(|| "Qwen3-VL packed patch count overflowed".to_string())?; + packed_cu_seqlens.push(actual_patches); + merged_tokens_per_image.push(self.merged_tokens(patches)); + } + Ok(Qwen3VlGridPlan { + patch_bucket: self.bucket_for_patches(actual_patches)?, + actual_patches, + packed_cu_seqlens, + merged_tokens_per_image, + }) + } + + #[must_use] + pub(crate) fn merged_tokens(&self, actual_patches: usize) -> usize { + actual_patches / (self.spatial_merge_size * self.spatial_merge_size) + } + + #[must_use] + pub(crate) fn fingerprint(&self, patch_bucket: usize) -> String { + format!( + "qwen3-vl-vision-v1:bucket={patch_bucket}:patch={}:temporal={}:merge={}:channels={}:hidden={}:intermediate={}:heads={}:depth={}:text={}:positions={}:deepstack={:?}:dtype=f32:source_quant={:?}", + self.patch_size, + self.temporal_patch_size, + self.spatial_merge_size, + self.channels, + self.hidden, + self.intermediate, + self.heads, + self.depth, + self.text_hidden, + self.num_position_embeddings, + self.deepstack_visual_indexes, + self.quantization, + ) + } + + pub(crate) fn weight_specs(&self) -> Vec { + let patch_width = + self.channels * self.temporal_patch_size * self.patch_size * self.patch_size; + let mut specs = vec![ + self.spec( + "vision_tower.patch_embed.proj.weight", + [self.hidden, patch_width], + ), + self.spec("vision_tower.patch_embed.proj.bias", [self.hidden]), + self.spec( + "vision_tower.pos_embed.weight", + [self.num_position_embeddings, self.hidden], + ), + ]; + 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.linear_fc1.weight", + vec![self.intermediate, self.hidden], + ), + ("mlp.linear_fc1.bias", vec![self.intermediate]), + ( + "mlp.linear_fc2.weight", + vec![self.hidden, self.intermediate], + ), + ("mlp.linear_fc2.bias", vec![self.hidden]), + ] { + specs.push(Qwen3VlWeightSpec { + name: format!("{prefix}.{suffix}"), + shape, + }); + } + } + let merge_width = self.hidden * self.spatial_merge_size * self.spatial_merge_size; + specs.extend([ + self.spec("vision_tower.merger.norm.weight", [self.hidden]), + self.spec("vision_tower.merger.norm.bias", [self.hidden]), + self.spec( + "vision_tower.merger.linear_fc1.weight", + [merge_width, merge_width], + ), + self.spec("vision_tower.merger.linear_fc1.bias", [merge_width]), + self.spec( + "vision_tower.merger.linear_fc2.weight", + [self.text_hidden, merge_width], + ), + self.spec("vision_tower.merger.linear_fc2.bias", [self.text_hidden]), + ]); + for merger in 0..self.deepstack_visual_indexes.len() { + let prefix = format!("vision_tower.deepstack_merger_list.{merger}"); + specs.extend([ + Qwen3VlWeightSpec { + name: format!("{prefix}.norm.weight"), + shape: vec![merge_width], + }, + Qwen3VlWeightSpec { + name: format!("{prefix}.norm.bias"), + shape: vec![merge_width], + }, + Qwen3VlWeightSpec { + name: format!("{prefix}.linear_fc1.weight"), + shape: vec![merge_width, merge_width], + }, + Qwen3VlWeightSpec { + name: format!("{prefix}.linear_fc1.bias"), + shape: vec![merge_width], + }, + Qwen3VlWeightSpec { + name: format!("{prefix}.linear_fc2.weight"), + shape: vec![self.text_hidden, merge_width], + }, + Qwen3VlWeightSpec { + name: format!("{prefix}.linear_fc2.bias"), + shape: vec![self.text_hidden], + }, + ]); + } + specs + } + + fn spec(&self, name: &str, shape: [usize; N]) -> Qwen3VlWeightSpec { + Qwen3VlWeightSpec { + name: name.to_string(), + shape: shape.into(), + } + } +} + +pub(crate) fn prepare_qwen3_vl_host_inputs( + config: &Qwen3VlConfig, + grids: &[(i32, i32, i32)], + temporal_patch_rows: &[f32], +) -> Result { + let plan = config.plan_image_grids(grids)?; + let row_width = config + .channels + .checked_mul(config.patch_size) + .and_then(|value| value.checked_mul(config.patch_size)) + .ok_or_else(|| "Qwen3-VL processor row width overflowed".to_string())?; + let actual_values = plan + .actual_patches + .checked_mul(config.temporal_patch_size) + .and_then(|value| value.checked_mul(row_width)) + .ok_or_else(|| "Qwen3-VL processor tensor size overflowed".to_string())?; + if temporal_patch_rows.len() != actual_values { + return Err(format!( + "Qwen3-VL processor produced {} values, expected {actual_values} for {} patches x temporal {} x row width {row_width}", + temporal_patch_rows.len(), + plan.actual_patches, + config.temporal_patch_size + )); + } + if let Some((index, value)) = temporal_patch_rows + .iter() + .enumerate() + .find(|(_, value)| !value.is_finite()) + { + return Err(format!( + "Qwen3-VL processor values contain non-finite value {value} at flat index {index}" + )); + } + let patch_width = row_width + .checked_mul(config.temporal_patch_size) + .ok_or_else(|| "Qwen3-VL temporal patch width overflowed".to_string())?; + let padded_values = plan + .patch_bucket + .checked_mul(patch_width) + .ok_or_else(|| "Qwen3-VL padded patch tensor size overflowed".to_string())?; + 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); + patches.resize(padded_values, 0.0); + + let head_dim = config.hidden / config.heads; + let rotary_dim = head_dim / 2; + let frequency_width = rotary_dim; + let axis_width = rotary_dim / 2; + if !head_dim.is_multiple_of(4) { + return Err(format!( + "Qwen3-VL vision head_dim={head_dim} must be divisible by 4 for 2D RoPE" + )); + } + let inv_freq = (0..axis_width) + .map(|index| 1.0 / 10_000.0f32.powf((2 * index) as f32 / rotary_dim as f32)) + .collect::>(); + let mut vision_rope_freqs = Vec::with_capacity(plan.patch_bucket * frequency_width); + let merge = config.spatial_merge_size; + for &(temporal, height, width) in grids { + debug_assert_eq!(temporal, 1); + let height = height as usize; + let width = width as usize; + for block_h in 0..height / merge { + for block_w in 0..width / merge { + for inner_h in 0..merge { + for inner_w in 0..merge { + let h = (block_h * merge + inner_h) as f32; + let w = (block_w * merge + inner_w) as f32; + vision_rope_freqs.extend(inv_freq.iter().map(|frequency| h * frequency)); + vision_rope_freqs.extend(inv_freq.iter().map(|frequency| w * frequency)); + } + } + } + } + } + vision_rope_freqs.resize(plan.patch_bucket * frequency_width, 0.0); + + let position_side = (config.num_position_embeddings as f64).sqrt() as usize; + let mut position_indices = + std::array::from_fn::<_, 4, _>(|_| Vec::::with_capacity(plan.patch_bucket)); + let mut position_weights = + std::array::from_fn::<_, 4, _>(|_| Vec::::with_capacity(plan.patch_bucket)); + for &(temporal, height, width) in grids { + debug_assert_eq!(temporal, 1); + let height = height as usize; + let width = width as usize; + let height_step = if height > 1 { + (position_side - 1) as f32 / (height - 1) as f32 + } else { + 0.0 + }; + let width_step = if width > 1 { + (position_side - 1) as f32 / (width - 1) as f32 + } else { + 0.0 + }; + for block_h in 0..height / merge { + for block_w in 0..width / merge { + for inner_h in 0..merge { + for inner_w in 0..merge { + let h = block_h * merge + inner_h; + let w = block_w * merge + inner_w; + let h_coordinate = h as f32 * height_step; + let w_coordinate = w as f32 * width_step; + let h_floor = h_coordinate.floor() as usize; + let w_floor = w_coordinate.floor() as usize; + let h_ceil = (h_floor + 1).min(position_side - 1); + let w_ceil = (w_floor + 1).min(position_side - 1); + let dh = h_coordinate - h_floor as f32; + let dw = w_coordinate - w_floor as f32; + for (corner, (row, column, weight)) in [ + (h_floor, w_floor, (1.0 - dh) * (1.0 - dw)), + (h_floor, w_ceil, (1.0 - dh) * dw), + (h_ceil, w_floor, dh * (1.0 - dw)), + (h_ceil, w_ceil, dh * dw), + ] + .into_iter() + .enumerate() + { + let index = row + .checked_mul(position_side) + .and_then(|value| value.checked_add(column)) + .ok_or_else(|| { + "Qwen3-VL position-table index overflowed".to_string() + })?; + position_indices[corner].push(i32::try_from(index).map_err(|_| { + "Qwen3-VL position-table index does not fit i32".to_string() + })?); + position_weights[corner].push(weight); + } + } + } + } + } + } + for corner in 0..4 { + position_indices[corner].resize(plan.patch_bucket, 0); + position_weights[corner].resize(plan.patch_bucket, 0.0); + } + let position_indices = position_indices.into_iter().flatten().collect(); + let position_weights = position_weights.into_iter().flatten().collect(); + + let bias_values = plan + .patch_bucket + .checked_mul(plan.patch_bucket) + .ok_or_else(|| "Qwen3-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; + } + } + } + // 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; + } + Ok(Qwen3VlHostInputs { + plan, + patches, + vision_rope_freqs, + packed_attention_bias, + position_indices, + position_weights, + }) +} + +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 layer_norm(builder: &mut Builder, value: &Val, weight: &Val, bias: &Val, epsilon: f32) -> Val { + let rows = value.ty.shape[0]; + let width = value.ty.shape[1]; + let zero = builder.const_f32(0.0); + let width_scalar = builder.const_f32(width as f32); + let width_rows = builder.broadcast(&width_scalar, &[], vec![rows]); + let sum = builder.reduce_add(value, 1, &zero); + let mean = builder.divide(&sum, &width_rows); + let mean = builder.broadcast(&mean, &[0], vec![rows, width]); + let centered = builder.subtract(value, &mean); + let squared = builder.multiply(¢ered, ¢ered); + let squared_sum = builder.reduce_add(&squared, 1, &zero); + let variance = builder.divide(&squared_sum, &width_rows); + let epsilon = builder.const_f32(epsilon); + let epsilon = builder.broadcast(&epsilon, &[], vec![rows]); + let variance = builder.add(&variance, &epsilon); + let inv_std = builder.rsqrt(&variance); + let inv_std = builder.broadcast(&inv_std, &[0], vec![rows, width]); + let normalized = builder.multiply(¢ered, &inv_std); + let weight = builder.broadcast(weight, &[1], vec![rows, width]); + let bias = builder.broadcast(bias, &[1], vec![rows, width]); + let normalized = builder.multiply(&normalized, &weight); + builder.add(&normalized, &bias) +} + +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 tanh_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 coefficient = builder.const_f32(0.044_715); + let coefficient = builder.broadcast(&coefficient, &[], shape.clone()); + let scale = builder.const_f32(0.797_884_6); + let scale = builder.broadcast(&scale, &[], shape); + let squared = builder.multiply(value, value); + let cubed = builder.multiply(&squared, value); + let nonlinear = builder.multiply(&coefficient, &cubed); + let inner = builder.add(value, &nonlinear); + let scaled = builder.multiply(&scale, &inner); + let tanh = builder.tanh(&scaled); + let cdf = builder.add(&one, &tanh); + 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: &Qwen3VlConfig, + 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 attention_bias = + builder.broadcast(attention_bias, &[1, 2], vec![config.heads, tokens, tokens]); + let scores = builder.add(&scores, &attention_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: &Qwen3VlConfig, + freqs: &Val, + attention_bias: &Val, +) -> Val { + let norm1 = layer_norm( + builder, + hidden, + &args.take(), + &args.take(), + config.layer_norm_eps, + ); + let attention = attention(builder, &norm1, args, config, freqs, attention_bias); + let residual = builder.add(hidden, &attention); + let norm2 = layer_norm( + builder, + &residual, + &args.take(), + &args.take(), + config.layer_norm_eps, + ); + let fc1 = linear_2d(builder, &norm2, &args.take(), &args.take()); + let activated = tanh_gelu(builder, &fc1); + let fc2 = linear_2d(builder, &activated, &args.take(), &args.take()); + builder.add(&residual, &fc2) +} + +fn interpolated_position_embeddings( + builder: &mut Builder, + table: &Val, + indices: &Val, + weights: &Val, +) -> Val { + let corners = indices.ty.shape[0]; + let tokens = indices.ty.shape[1]; + let hidden = table.ty.shape[1]; + let gathered = builder.gather_rows_nd(table, indices); + let weights = builder.reshape(weights, vec![corners, tokens, 1]); + let weights = builder.broadcast(&weights, &[0, 1, 2], vec![corners, tokens, hidden]); + let weighted = builder.multiply(&gathered, &weights); + let zero = builder.const_f32(0.0); + builder.reduce_add(&weighted, 0, &zero) +} + +fn patch_merger( + builder: &mut Builder, + hidden: &Val, + args: &mut Args, + config: &Qwen3VlConfig, + postshuffle_norm: bool, +) -> Val { + let merge_width = config.hidden * config.spatial_merge_size * config.spatial_merge_size; + let merged_tokens = + hidden.ty.shape[0] / (config.spatial_merge_size * config.spatial_merge_size); + let merged = if postshuffle_norm { + let reshaped = builder.reshape(hidden, vec![merged_tokens, merge_width]); + layer_norm( + builder, + &reshaped, + &args.take(), + &args.take(), + config.layer_norm_eps, + ) + } else { + let normalized = layer_norm( + builder, + hidden, + &args.take(), + &args.take(), + config.layer_norm_eps, + ); + builder.reshape(&normalized, vec![merged_tokens, merge_width]) + }; + let merged = linear_2d(builder, &merged, &args.take(), &args.take()); + let merged = exact_gelu(builder, &merged); + linear_2d(builder, &merged, &args.take(), &args.take()) +} + +/// Emit one Qwen3-VL bucket. `patches` are already in the processor's +/// post-smart-resize, spatial-merge-grouped order. `vision_rope` and +/// `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_qwen3_vl(config: &Qwen3VlConfig, patch_bucket: usize) -> String { + assert!( + QWEN3_VL_PATCH_BUCKETS.contains(&patch_bucket), + "unqualified Qwen3-VL patch bucket" + ); + let specs = config.weight_specs(); + let mut args = Args::new(&specs); + let mut builder = Builder::new(); + let patch_weight = args.take(); + let patch_bias = args.take(); + let position_table = 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.grouped"); + let freqs = args.push_input( + Ty::f32(vec![patch_bucket, head_dim / 2]), + "vision_rope.freqs", + ); + let attention_bias = args.push_input( + Ty::f32(vec![patch_bucket, patch_bucket]), + "packed_attention.bias", + ); + let position_indices = + args.push_input(Ty::new(vec![4, patch_bucket, 1], "i32"), "position.indices"); + let position_weights = args.push_input(Ty::f32(vec![4, patch_bucket]), "position.weights"); + let mut hidden = builder.linear_seq(&patches, &patch_weight); + hidden = bias_2d(&mut builder, &hidden, &patch_bias); + let position_embeddings = interpolated_position_embeddings( + &mut builder, + &position_table, + &position_indices, + &position_weights, + ); + hidden = builder.add(&hidden, &position_embeddings); + let mut deepstack_hidden = Vec::with_capacity(config.deepstack_visual_indexes.len()); + for layer in 0..config.depth { + hidden = encoder_layer( + &mut builder, + &hidden, + &mut args, + config, + &freqs, + &attention_bias, + ); + if config.deepstack_visual_indexes.contains(&layer) { + deepstack_hidden.push(hidden.clone()); + } + } + let projected = patch_merger(&mut builder, &hidden, &mut args, config, false); + let mut outputs = vec![projected]; + for branch in deepstack_hidden { + outputs.push(patch_merger(&mut builder, &branch, &mut args, config, true)); + } + assert_eq!(args.cursor, specs.len(), "Qwen3-VL weight schema drifted"); + let result_types = outputs + .iter() + .map(|output| output.ty.render()) + .collect::>() + .join(", "); + let result_names = outputs + .iter() + .map(|output| output.name.as_str()) + .collect::>() + .join(", "); + format!( + "module @qwen3_vl_vision {{\n func.func public @main({signature}) -> ({result_types}) {{\n{body} return {result_names} : {result_types}\n }}\n}}\n", + signature = args.declarations.join(", "), + body = builder.body(), + ) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn config() -> Qwen3VlConfig { + Qwen3VlConfig::from_json_str( + r#"{ + "model_type": "qwen3_vl", + "text_config": { + "model_type": "qwen3_vl_text", + "hidden_size": 12 + }, + "vision_config": { + "depth": 2, + "hidden_size": 8, + "intermediate_size": 16, + "out_hidden_size": 12, + "num_heads": 2, + "in_channels": 3, + "patch_size": 2, + "spatial_merge_size": 2, + "temporal_patch_size": 2, + "num_position_embeddings": 16, + "deepstack_visual_indexes": [0, 1] + } + }"#, + ) + .unwrap() + } + + #[test] + fn static_buckets_bound_qwen3_vl_grids_without_per_shape_compilation() { + let config = config(); + assert_eq!(config.bucket_for_patches(16), Ok(16)); + assert_eq!(config.bucket_for_patches(20), Ok(64)); + assert_eq!(config.bucket_for_patches(256), Ok(256)); + assert!( + config + .bucket_for_patches(18) + .unwrap_err() + .contains("divisible") + ); + assert!( + config + .bucket_for_patches(260) + .unwrap_err() + .contains("exceeding qualified capacity") + ); + assert_eq!(config.merged_tokens(64), 16); + assert!(config.fingerprint(64).contains("bucket=64")); + assert!( + config + .plan_image_grids(&[(1, 1, 16)]) + .unwrap_err() + .contains("each spatial axis") + ); + assert!( + config + .plan_image_grids(&[(2, 4, 4)]) + .unwrap_err() + .contains("video/temporal grids") + ); + let packed = config.plan_image_grids(&[(1, 4, 4), (1, 4, 8)]).unwrap(); + assert_eq!(packed.patch_bucket, 64); + assert_eq!(packed.actual_patches, 48); + assert_eq!(packed.packed_cu_seqlens, vec![0, 16, 48]); + assert_eq!(packed.merged_tokens_per_image, vec![4, 8]); + } + + #[test] + fn weight_schema_covers_patch_blocks_and_merger_in_checkpoint_order() { + let config = config(); + let specs = config.weight_specs(); + assert_eq!( + specs.first().unwrap(), + &Qwen3VlWeightSpec { + name: "vision_tower.patch_embed.proj.weight".to_string(), + shape: vec![8, 24], + } + ); + assert_eq!(specs.len(), 3 + 2 * 12 + 3 * 6); + assert_eq!( + specs.last().unwrap(), + &Qwen3VlWeightSpec { + name: "vision_tower.deepstack_merger_list.1.linear_fc2.bias".to_string(), + shape: vec![12], + } + ); + } + + #[test] + fn emitted_bucket_has_packed_bias_rope_and_merged_output_contract() { + let config = config(); + let mlir = emit_qwen3_vl(&config, 16); + assert!(mlir.contains("loc(\"patches.grouped\")")); + assert!(mlir.contains("tensor<16x24xf32>")); + assert!(mlir.contains("loc(\"vision_rope.freqs\")")); + assert!(mlir.contains("tensor<16x2xf32>")); + assert!(mlir.contains("loc(\"packed_attention.bias\")")); + assert!(mlir.contains("tensor<16x16xf32>")); + assert!(mlir.contains("loc(\"position.indices\")")); + assert!(mlir.contains("tensor<4x16x1xi32>")); + assert!(mlir.contains("-> (tensor<4x12xf32>, tensor<4x12xf32>, tensor<4x12xf32>)")); + assert_eq!(mlir.matches("chlo.erf").count(), 3); + } + + #[test] + fn host_inputs_preserve_packed_media_isolation_and_finite_padding_rows() { + let config = config(); + let grids = [(1, 4, 4), (1, 4, 8)]; + let row_width = 3 * 2 * 2; + let values = vec![0.25; (16 + 32) * 2 * row_width]; + let inputs = prepare_qwen3_vl_host_inputs(&config, &grids, &values).unwrap(); + assert_eq!(inputs.plan.patch_bucket, 64); + assert_eq!(inputs.patches.len(), 64 * 24); + assert_eq!(inputs.vision_rope_freqs.len(), 64 * 2); + assert_eq!(inputs.packed_attention_bias.len(), 64 * 64); + assert_eq!(inputs.position_indices.len(), 4 * 64); + assert_eq!(inputs.position_weights.len(), 4 * 64); + for patch in 0..48 { + let sum = (0..4) + .map(|corner| inputs.position_weights[corner * 64 + patch]) + .sum::(); + assert!((sum - 1.0).abs() < 1e-6); + } + assert!( + inputs.position_weights[48..64] + .iter() + .all(|&value| value == 0.0) + ); + assert_eq!(inputs.packed_attention_bias[15 * 64 + 15], 0.0); + assert_eq!( + inputs.packed_attention_bias[15 * 64 + 16], + f32::NEG_INFINITY + ); + assert_eq!( + inputs.packed_attention_bias[16 * 64 + 15], + 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); + let mut non_finite = values; + non_finite[7] = f32::NAN; + assert!( + prepare_qwen3_vl_host_inputs(&config, &grids, &non_finite) + .unwrap_err() + .contains("flat index 7") + ); + } +} diff --git a/src/lib/mlxcel-xla/src/lib.rs b/src/lib/mlxcel-xla/src/lib.rs index 2b132826..854bf439 100644 --- a/src/lib/mlxcel-xla/src/lib.rs +++ b/src/lib/mlxcel-xla/src/lib.rs @@ -81,6 +81,8 @@ mod phi4_audio; #[cfg(feature = "iree")] mod qwen2_vl_runtime; #[cfg(feature = "iree")] +mod qwen3_vl_runtime; +#[cfg(feature = "iree")] mod vision_runtime; #[cfg(feature = "micro-oracle")] @@ -180,6 +182,10 @@ pub use phi4_audio::{Phi4AudioCheckpoint, Phi4AudioDiagnosticRuntime, Phi4AudioD pub use qwen2_vl_runtime::{ IreeQwen2VlProjector, Qwen2VlVisionExecutionMetrics, Qwen2VlVisionProjection, }; +#[cfg(feature = "iree")] +pub use qwen3_vl_runtime::{ + IreeQwen3VlProjector, Qwen3VlVisionExecutionMetrics, Qwen3VlVisionProjection, +}; #[cfg(feature = "diagnostics")] pub use vision_runtime::{IreeVisionDiagnosticProjector, VisionDiagnosticProjection}; #[cfg(feature = "iree")] diff --git a/src/lib/mlxcel-xla/src/qwen3_vl_runtime.rs b/src/lib/mlxcel-xla/src/qwen3_vl_runtime.rs new file mode 100644 index 00000000..4352e8af --- /dev/null +++ b/src/lib/mlxcel-xla/src/qwen3_vl_runtime.rs @@ -0,0 +1,839 @@ +// 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. + +//! Resident IREE execution for the Qwen3-VL vision tower and patch merger. +//! +//! Only normalized patch extraction and grid metadata remain on the host. The +//! full vision tower, main merger, and every DeepStack merger execute in IREE; +//! this module never constructs or calls the MLX vision encoder. + +use std::collections::{BTreeMap, BTreeSet}; +use std::fs::File; +use std::io::Read; +use std::path::{Path, PathBuf}; +use std::process::Command; +use std::time::Instant; + +use memmap2::Mmap; +use safetensors::{Dtype, SafeTensors}; +use sha2::{Digest, Sha256}; + +use crate::aux::{ + AuxiliaryInput, AuxiliaryOutput, AuxiliaryTensorDType, AuxiliaryWeight, AuxiliaryWeightDType, + IreeAuxiliaryModule, +}; +use crate::aux_manifest::{AuxiliaryArtifactContract, ensure_qualified_auxiliary_artifact}; +use crate::emitter::{ + Qwen3VlConfig, Qwen3VlHostInputs, emit_qwen3_vl, prepare_qwen3_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}; + +const ENTRY_NAME: &str = "qwen3_vl_vision.main"; + +#[derive(Debug, Clone, PartialEq)] +pub struct Qwen3VlVisionExecutionMetrics { + pub patch_upload_bytes: usize, + pub metadata_upload_bytes: usize, + pub projected_transfer_bytes: usize, + pub elapsed_seconds: f64, +} + +#[derive(Debug, Clone, PartialEq)] +pub struct Qwen3VlVisionProjection { + pub values: Vec, + pub deepstack_values: Vec>, + pub shape: [usize; 2], + pub merged_tokens_per_image: Vec, + pub packed_cu_seqlens: Vec, + pub metrics: Qwen3VlVisionExecutionMetrics, +} + +struct ActiveModule { + patch_bucket: usize, + module: IreeAuxiliaryModule, +} + +pub struct IreeQwen3VlProjector { + model_dir: PathBuf, + device: String, + config: Qwen3VlConfig, + processor_identity: String, + active: Option, +} + +fn hex_bytes(bytes: &[u8]) -> String { + let mut output = String::with_capacity(bytes.len() * 2); + for byte in bytes { + use std::fmt::Write; + write!(output, "{byte:02x}").expect("writing to String cannot fail"); + } + output +} + +fn sha256_hex(bytes: &[u8]) -> String { + hex_bytes(&Sha256::digest(bytes)) +} + +fn sha256_file(path: &Path) -> Result { + let mut file = File::open(path).map_err(|error| format!("open {}: {error}", path.display()))?; + let mut digest = Sha256::new(); + let mut buffer = [0u8; 64 * 1024]; + loop { + let read = file + .read(&mut buffer) + .map_err(|error| format!("read {}: {error}", path.display()))?; + if read == 0 { + break; + } + digest.update(&buffer[..read]); + } + Ok(hex_bytes(&digest.finalize())) +} + +fn compiler_generation_identity( + compiler: &Path, + flags: &[&str], + mlir: &str, +) -> Result { + let version = Command::new(compiler) + .arg("--version") + .output() + .map_err(|error| format!("run {} --version: {error}", compiler.display()))?; + if !version.status.success() { + return Err(format!( + "{} --version failed: {}", + compiler.display(), + String::from_utf8_lossy(&version.stderr) + )); + } + Ok(format!( + "compiler={};compiler_sha256={};version={};flags={flags:?};mlir_sha256={}", + compiler.display(), + sha256_file(compiler)?, + String::from_utf8_lossy(&version.stdout).trim(), + sha256_hex(mlir.as_bytes()) + )) +} + +fn processor_identity(model_dir: &Path, config: &Qwen3VlConfig) -> Result { + let path = model_dir.join("preprocessor_config.json"); + let bytes = std::fs::read(&path).map_err(|error| format!("{}: {error}", path.display()))?; + let value: serde_json::Value = serde_json::from_slice(&bytes) + .map_err(|error| format!("parse {}: {error}", path.display()))?; + let object = value + .as_object() + .ok_or_else(|| format!("{} must contain an object", path.display()))?; + for field in ["do_resize", "do_rescale", "do_normalize", "do_convert_rgb"] { + if object.get(field).and_then(serde_json::Value::as_bool) != Some(true) { + return Err(format!( + "preprocessor_config.json {field}=true is required by Qwen3-VL IREE vision" + )); + } + } + let positive = |field: &str| -> Result { + object + .get(field) + .and_then(serde_json::Value::as_u64) + .and_then(|value| usize::try_from(value).ok()) + .filter(|&value| value > 0) + .ok_or_else(|| format!("preprocessor_config.json {field} must be a positive integer")) + }; + for (field, expected) in [ + ("patch_size", config.patch_size), + ("temporal_patch_size", config.temporal_patch_size), + ("merge_size", config.spatial_merge_size), + ] { + let actual = positive(field)?; + if actual != expected { + return Err(format!( + "preprocessor_config.json {field}={actual} disagrees with config value {expected}" + )); + } + } + if object + .get("image_processor_type") + .and_then(serde_json::Value::as_str) + != Some("Qwen2VLImageProcessorFast") + { + return Err( + "Qwen3-VL IREE vision requires image_processor_type=Qwen2VLImageProcessorFast" + .to_string(), + ); + } + for field in ["image_mean", "image_std"] { + let values = object + .get(field) + .and_then(serde_json::Value::as_array) + .ok_or_else(|| format!("preprocessor_config.json {field} must be an array"))?; + if values.len() != 3 || values.iter().any(|value| value.as_f64() != Some(0.5)) { + return Err(format!( + "Qwen3-VL IREE vision requires preprocessor_config.json {field}=[0.5,0.5,0.5]" + )); + } + } + Ok(format!( + "qwen3-vl-processor-v1:sha256={}:patch={}:temporal={}:merge={}", + sha256_hex(&bytes), + config.patch_size, + config.temporal_patch_size, + config.spatial_merge_size + )) +} + +fn model_shards(model_dir: &Path) -> Result, String> { + let mut shards = std::fs::read_dir(model_dir) + .map_err(|error| format!("read {}: {error}", model_dir.display()))? + .map(|entry| { + entry + .map(|entry| entry.path()) + .map_err(|error| format!("read {} entry: {error}", model_dir.display())) + }) + .collect::, String>>()?; + shards.retain(|path| { + path.file_name() + .and_then(|name| name.to_str()) + .is_some_and(|name| name.ends_with(".safetensors")) + }); + shards.sort(); + if shards.is_empty() { + return Err(format!( + "no safetensors checkpoint shards found in {}", + model_dir.display() + )); + } + Ok(shards) +} + +fn tensor_locations(model_dir: &Path) -> Result, String> { + let mut locations = BTreeMap::new(); + for shard in model_shards(model_dir)? { + let file = + File::open(&shard).map_err(|error| format!("open {}: {error}", shard.display()))?; + // Safety: the map is read-only and lives through the header scan. + let mmap = unsafe { Mmap::map(&file) } + .map_err(|error| format!("mmap {}: {error}", shard.display()))?; + let tensors = SafeTensors::deserialize(&mmap) + .map_err(|error| format!("parse {}: {error}", shard.display()))?; + for name in tensors + .names() + .into_iter() + .filter(|name| name.starts_with("vision_tower.")) + { + if let Some(previous) = locations.insert(name.to_string(), shard.clone()) { + return Err(format!( + "Qwen3-VL vision tensor {name:?} is duplicated in {} and {}", + previous.display(), + shard.display() + )); + } + } + } + Ok(locations) +} + +fn native_f32_bytes(values: Vec) -> Vec { + let mut bytes = Vec::with_capacity(values.len() * std::mem::size_of::()); + for value in values { + bytes.extend_from_slice(&value.to_ne_bytes()); + } + bytes +} + +fn validate_finite(label: &str, values: &[f32]) -> Result<(), String> { + if let Some((index, value)) = values + .iter() + .enumerate() + .find(|(_, value)| !value.is_finite()) + { + return Err(format!( + "{label} contains non-finite value {value} at flat index {index}" + )); + } + Ok(()) +} + +fn transpose_patch_t_h_w_c_to_t_c_h_w(values: Vec, config: &Qwen3VlConfig) -> Vec { + let out = config.hidden; + let temporal = config.temporal_patch_size; + let height = config.patch_size; + let width = config.patch_size; + let channels = config.channels; + let mut transposed = vec![0.0; values.len()]; + for output in 0..out { + for t in 0..temporal { + for channel in 0..channels { + for y in 0..height { + for x in 0..width { + let source = ((((output * temporal + t) * height + y) * width + x) + * channels) + + channel; + let target = + ((((output * temporal + t) * channels + channel) * height + y) * width) + + x; + transposed[target] = values[source]; + } + } + } + } + } + transposed +} + +fn decode_direct_f32( + label: &str, + dtype: Dtype, + data: &[u8], + expected_shape: &[usize], + source_shape: &[usize], +) -> Result, String> { + if source_shape != expected_shape { + return Err(format!( + "{label} has shape {source_shape:?}, expected {expected_shape:?}" + )); + } + let values = match dtype { + Dtype::BF16 => bf16_to_f32(data), + Dtype::F16 => f16_to_f32(data), + Dtype::F32 => f32_le_to_f32(data), + other => { + return Err(format!( + "{label} has unsupported dtype {other:?}; expected BF16, F16, F32, or an affine U32 projection" + )); + } + }; + validate_finite(label, &values)?; + Ok(values) +} + +fn quant_sibling(name: &str, suffix: &str) -> Result { + let stem = name + .strip_suffix(".weight") + .ok_or_else(|| format!("quantized Qwen3-VL tensor {name} must end in .weight"))?; + Ok(format!("{stem}.{suffix}")) +} + +fn load_weights( + model_dir: &Path, + config: &Qwen3VlConfig, +) -> Result<(Vec, String), String> { + let specs = config.weight_specs(); + let locations = tensor_locations(model_dir)?; + let required = specs + .iter() + .map(|spec| spec.name.as_str()) + .collect::>(); + let missing = required + .iter() + .filter(|name| !locations.contains_key(**name)) + .copied() + .collect::>(); + if !missing.is_empty() { + return Err(format!( + "checkpoint is missing {} Qwen3-VL vision tensor(s): {}", + missing.len(), + missing + .iter() + .take(8) + .copied() + .collect::>() + .join(", ") + )); + } + let mut by_shard = BTreeMap::<&Path, Vec>::new(); + for (index, spec) in specs.iter().enumerate() { + by_shard + .entry( + locations + .get(&spec.name) + .expect("missing specs rejected") + .as_path(), + ) + .or_default() + .push(index); + } + let mut loaded = (0..specs.len()) + .map(|_| None) + .collect::>>(); + let mut source_schema = vec![String::new(); specs.len()]; + for (shard, indices) in by_shard { + let file = + File::open(shard).map_err(|error| format!("open {}: {error}", shard.display()))?; + // Safety: the map is read-only and lives while selected tensors copy. + let mmap = unsafe { Mmap::map(&file) } + .map_err(|error| format!("mmap {}: {error}", shard.display()))?; + let tensors = SafeTensors::deserialize(&mmap) + .map_err(|error| format!("parse {}: {error}", shard.display()))?; + for index in indices { + let spec = &specs[index]; + let tensor = tensors.tensor(&spec.name).map_err(|error| { + format!( + "resolved Qwen3-VL tensor {} in {}: {error}", + spec.name, + shard.display() + ) + })?; + let source_dtype = tensor.dtype(); + let source_shape = tensor.shape().to_vec(); + let (values, transform) = if spec.name == "vision_tower.patch_embed.proj.weight" { + let expected_source = [ + config.hidden, + config.temporal_patch_size, + config.patch_size, + config.patch_size, + config.channels, + ]; + if source_shape != expected_source { + return Err(format!( + "{} has source shape {source_shape:?}, expected Conv3d [out,t,h,w,c] {expected_source:?}", + spec.name + )); + } + let values = decode_direct_f32( + &spec.name, + source_dtype, + tensor.data(), + &expected_source, + &source_shape, + )?; + ( + transpose_patch_t_h_w_c_to_t_c_h_w(values, config), + format!( + "conv3d-source={source_shape:?}:layout=O,T,H,W,C->O,T,C,H,W->O,{}", + spec.shape[1] + ), + ) + } else if source_dtype == Dtype::U32 { + let Some((bits, group_size)) = config.quantization else { + return Err(format!( + "{} is U32 affine-quantized but config.json has no quantization contract", + spec.name + )); + }; + if spec.shape.len() != 2 { + return Err(format!( + "{} is quantized but logical shape {:?} is not rank 2", + spec.name, spec.shape + )); + } + let out = spec.shape[0]; + let in_packed = spec.shape[1] + .checked_mul(bits) + .ok_or_else(|| format!("{} packed width overflowed", spec.name))? + / 32; + if source_shape != [out, in_packed] { + return Err(format!( + "{} has packed shape {source_shape:?}, expected [{out}, {in_packed}]", + spec.name + )); + } + let scales_name = quant_sibling(&spec.name, "scales")?; + let biases_name = quant_sibling(&spec.name, "biases")?; + for sibling in [&scales_name, &biases_name] { + if locations.get(sibling).map(PathBuf::as_path) != Some(shard) { + return Err(format!( + "quantized Qwen3-VL tensor {} and sibling {sibling} must share one immutable shard", + spec.name + )); + } + } + let scales = tensors + .tensor(&scales_name) + .map_err(|error| format!("load {scales_name}: {error}"))?; + let biases = tensors + .tensor(&biases_name) + .map_err(|error| format!("load {biases_name}: {error}"))?; + let scales_bf16 = match (scales.dtype(), biases.dtype()) { + (Dtype::F16, Dtype::F16) => false, + (Dtype::BF16, Dtype::BF16) => true, + (left, right) => { + return Err(format!( + "{} affine metadata dtypes {left:?}/{right:?} must be matching F16 or BF16", + spec.name + )); + } + }; + let values = dequantize_affine( + tensor.data(), + scales.data(), + biases.data(), + out, + in_packed, + bits, + group_size, + scales_bf16, + )?; + validate_finite(&format!("Qwen3-VL tensor {}", spec.name), &values)?; + ( + values, + format!( + "affine-source={source_dtype:?}:{source_shape:?};scales={:?}:{:?};biases={:?}:{:?};bits={bits};group={group_size}->f32", + scales.dtype(), + scales.shape(), + biases.dtype(), + biases.shape() + ), + ) + } else { + ( + decode_direct_f32( + &spec.name, + source_dtype, + tensor.data(), + &spec.shape, + &source_shape, + )?, + "identity->f32".to_string(), + ) + }; + let expected = spec.shape.iter().try_fold(1usize, |count, dimension| { + count + .checked_mul(*dimension) + .ok_or_else(|| format!("{} logical element count overflowed", spec.name)) + })?; + if values.len() != expected { + return Err(format!( + "{} decoded {} values, expected {expected} for logical shape {:?}", + spec.name, + values.len(), + spec.shape + )); + } + loaded[index] = Some(AuxiliaryWeight { + name: spec.name.clone(), + bytes: native_f32_bytes(values), + dtype: AuxiliaryWeightDType::Float32, + shape: spec.shape.clone(), + }); + source_schema[index] = format!( + "{}:{source_dtype:?}:{source_shape:?}:{transform}", + spec.name + ); + } + } + Ok(( + loaded + .into_iter() + .map(|weight| weight.expect("all Qwen3-VL specs loaded")) + .collect(), + source_schema.join("\n"), + )) +} + +fn f32_as_bytes(values: &[f32]) -> &[u8] { + // Safety: f32 has no invalid bit patterns and the view cannot outlive input. + unsafe { + std::slice::from_raw_parts(values.as_ptr().cast::(), std::mem::size_of_val(values)) + } +} + +fn i32_as_bytes(values: &[i32]) -> &[u8] { + // Safety: i32 has no invalid bit patterns and the view cannot outlive input. + unsafe { + std::slice::from_raw_parts(values.as_ptr().cast::(), std::mem::size_of_val(values)) + } +} + +fn checked_f32_output(label: &str, bytes: Vec) -> Result, String> { + if !bytes.len().is_multiple_of(std::mem::size_of::()) { + return Err(format!( + "{label} returned {} bytes, not a whole number of f32 values", + bytes.len() + )); + } + let values = bytes + .chunks_exact(std::mem::size_of::()) + .map(|chunk| f32::from_ne_bytes(chunk.try_into().expect("four-byte f32 chunk"))) + .collect::>(); + validate_finite(label, &values)?; + Ok(values) +} + +fn compile_and_load( + model_dir: &Path, + device: &str, + config: &Qwen3VlConfig, + processor_identity: &str, + patch_bucket: usize, +) -> Result { + let mlir = emit_qwen3_vl(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-qwen3-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!( + "{};{};checkpoint_schema_sha256={}", + config.fingerprint(patch_bucket), + processor_identity, + sha256_hex(checkpoint_schema.as_bytes()) + ), + compiler_generation_identity(&compiler, flags, &mlir)?, + )?; + let tag = format!("qwen3-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) + })?; + IreeAuxiliaryModule::load(device, &vmfb, &contract, weights) +} + +impl IreeQwen3VlProjector { + pub fn load(model_dir: &Path, device: &str) -> Result { + let config = Qwen3VlConfig::from_model_dir(model_dir)?; + 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, + }) + } + + #[must_use] + pub fn active_bucket(&self) -> Option { + self.active.as_ref().map(|active| active.patch_bucket) + } + + #[must_use] + pub fn text_hidden(&self) -> usize { + self.config.text_hidden + } + + #[must_use] + pub fn deepstack_layer_count(&self) -> usize { + self.config.deepstack_visual_indexes.len() + } + + fn ensure_bucket(&mut self, patch_bucket: usize) -> Result<&mut IreeAuxiliaryModule, String> { + if self.active.as_ref().map(|active| active.patch_bucket) != Some(patch_bucket) { + let module = compile_and_load( + &self.model_dir, + &self.device, + &self.config, + &self.processor_identity, + patch_bucket, + )?; + self.active = Some(ActiveModule { + patch_bucket, + module, + }); + } + Ok(&mut self + .active + .as_mut() + .expect("active module installed") + .module) + } + + pub fn project( + &mut self, + temporal_patch_rows: &[f32], + grids: &[(i32, i32, i32)], + ) -> Result { + let Qwen3VlHostInputs { + plan, + patches, + vision_rope_freqs, + packed_attention_bias, + position_indices, + position_weights, + } = prepare_qwen3_vl_host_inputs(&self.config, grids, temporal_patch_rows)?; + let patch_width = self.config.channels + * self.config.temporal_patch_size + * self.config.patch_size + * self.config.patch_size; + let frequency_width = self.config.hidden / self.config.heads / 2; + let 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 position_index_shape = [4, plan.patch_bucket, 1]; + let position_weight_shape = [4, plan.patch_bucket]; + let full_output_shape = [ + plan.patch_bucket / (self.config.spatial_merge_size * self.config.spatial_merge_size), + self.config.text_hidden, + ]; + let output_bytes = full_output_shape.iter().product::() * std::mem::size_of::(); + let mut output_buffers = + vec![vec![0u8; output_bytes]; 1 + self.config.deepstack_visual_indexes.len()]; + let started = Instant::now(); + { + let mut outputs = output_buffers + .iter_mut() + .map(|output| AuxiliaryOutput { + bytes: output.as_mut_slice(), + dtype: AuxiliaryTensorDType::Float32, + shape: &full_output_shape, + }) + .collect::>(); + 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, + }, + AuxiliaryInput { + bytes: i32_as_bytes(&position_indices), + dtype: AuxiliaryTensorDType::Int32, + shape: &position_index_shape, + }, + AuxiliaryInput { + bytes: f32_as_bytes(&position_weights), + dtype: AuxiliaryTensorDType::Float32, + shape: &position_weight_shape, + }, + ], + &mut outputs, + )?; + } + let elapsed_seconds = started.elapsed().as_secs_f64(); + let mut decoded_outputs = output_buffers + .into_iter() + .enumerate() + .map(|(index, output)| { + checked_f32_output(&format!("Qwen3-VL IREE output {index}"), output) + }) + .collect::, _>>()?; + let all_values = decoded_outputs.remove(0); + let actual_tokens = plan.merged_tokens_per_image.iter().sum::(); + let actual_values = actual_tokens + .checked_mul(self.config.text_hidden) + .ok_or_else(|| "Qwen3-VL projected output size overflowed".to_string())?; + let values = all_values + .get(..actual_values) + .ok_or_else(|| "Qwen3-VL IREE output is shorter than the actual grid".to_string())? + .to_vec(); + let deepstack_values = decoded_outputs + .into_iter() + .map(|output| { + output + .get(..actual_values) + .ok_or_else(|| { + "Qwen3-VL IREE DeepStack output is shorter than the actual grid".to_string() + }) + .map(<[f32]>::to_vec) + }) + .collect::, _>>()?; + Ok(Qwen3VlVisionProjection { + values, + deepstack_values, + shape: [actual_tokens, self.config.text_hidden], + merged_tokens_per_image: plan.merged_tokens_per_image, + packed_cu_seqlens: plan.packed_cu_seqlens, + metrics: Qwen3VlVisionExecutionMetrics { + 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(position_indices.as_slice()) + + std::mem::size_of_val(position_weights.as_slice()), + projected_transfer_bytes: output_bytes + * (1 + self.config.deepstack_visual_indexes.len()), + elapsed_seconds, + }, + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn tiny_config() -> Qwen3VlConfig { + Qwen3VlConfig::from_json_str( + r#"{ + "model_type":"qwen3_vl", + "text_config":{ + "model_type":"qwen3_vl_text", + "hidden_size":12 + }, + "vision_config":{ + "depth":1, + "hidden_size":8, + "intermediate_size":16, + "out_hidden_size":12, + "num_heads":2, + "in_channels":3, + "patch_size":2, + "spatial_merge_size":2, + "temporal_patch_size":2, + "num_position_embeddings":16, + "deepstack_visual_indexes":[0] + } + }"#, + ) + .unwrap() + } + + #[test] + fn conv3d_layout_transform_matches_o_t_c_h_w_flattening() { + let config = tiny_config(); + let source = (0..config.hidden + * config.temporal_patch_size + * config.patch_size + * config.patch_size + * config.channels) + .map(|value| value as f32) + .collect::>(); + let transformed = transpose_patch_t_h_w_c_to_t_c_h_w(source.clone(), &config); + let source_index = |o: usize, t: usize, y: usize, x: usize, c: usize| { + ((((o * 2 + t) * 2 + y) * 2 + x) * 3) + c + }; + let target_index = |o: usize, t: usize, c: usize, y: usize, x: usize| { + ((((o * 2 + t) * 3 + c) * 2 + y) * 2) + x + }; + for o in 0..8 { + for t in 0..2 { + for c in 0..3 { + for y in 0..2 { + for x in 0..2 { + assert_eq!( + transformed[target_index(o, t, c, y, x)], + source[source_index(o, t, y, x, c)] + ); + } + } + } + } + } + } + + #[test] + fn output_decoder_rejects_non_finite_values() { + let mut bytes = 1.0f32.to_ne_bytes().to_vec(); + bytes.extend_from_slice(&f32::NAN.to_ne_bytes()); + assert!( + checked_f32_output("qwen vision", bytes) + .unwrap_err() + .contains("flat index 1") + ); + } +} diff --git a/src/loading/mod.rs b/src/loading/mod.rs index fb2cab9d..13706a90 100644 --- a/src/loading/mod.rs +++ b/src/loading/mod.rs @@ -63,6 +63,8 @@ pub(crate) use self::vlm::load_llava_iree_host_preprocessor; 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::load_qwen3_vl_iree_host_preprocessor; +#[cfg(feature = "xla-iree")] pub(crate) use self::vlm::{ Phi4MMXlaVisionComponents, load_phi4mm_xla_media_components, load_phi4mm_xla_text_embeddings, }; diff --git a/src/loading/vlm.rs b/src/loading/vlm.rs index b87c6c0a..fd21349a 100644 --- a/src/loading/vlm.rs +++ b/src/loading/vlm.rs @@ -114,12 +114,12 @@ 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_vl_iree_host_preprocessor, load_qwen3_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 52971240..4f35453c 100644 --- a/src/loading/vlm_qwen.rs +++ b/src/loading/vlm_qwen.rs @@ -30,7 +30,9 @@ use std::path::Path; use crate::LoadedModel; use crate::models; #[cfg(feature = "xla-iree")] -use crate::multimodal::host_preprocessor::{HostPreprocessorError, Qwen2VlIreeHostPreprocessor}; +use crate::multimodal::host_preprocessor::{ + HostPreprocessorError, Qwen2VlIreeHostPreprocessor, Qwen3VlIreeHostPreprocessor, +}; use crate::vision; #[cfg(feature = "xla-iree")] @@ -200,6 +202,123 @@ pub(crate) fn load_qwen2_vl_iree_host_preprocessor( ) } +/// Load only Qwen3-VL image patch preprocessing and the dense text embedding +/// table. The full vision tower and all DeepStack mergers remain resident in +/// IREE; MoE checkpoints are deliberately rejected by the runtime config. +#[cfg(feature = "xla-iree")] +pub(crate) fn load_qwen3_vl_iree_host_preprocessor( + model_path: &Path, + device: &str, +) -> Result { + use mlxcel_core::layers::UnifiedEmbedding; + use vision::encoders::qwen3_vl::Qwen3VLVisionConfig; + + 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("qwen3_vl") + { + return Err(HostPreprocessorError::FamilyMismatch { + actual: full_config + .get("model_type") + .and_then(serde_json::Value::as_str) + .unwrap_or("") + .to_string(), + }); + } + let text_config = full_config + .get("text_config") + .and_then(serde_json::Value::as_object) + .ok_or_else(|| { + HostPreprocessorError::InvalidConfig( + "Qwen3-VL text_config must be an object".to_string(), + ) + })?; + if text_config + .get("model_type") + .and_then(serde_json::Value::as_str) + != Some("qwen3_vl_text") + { + return Err(HostPreprocessorError::FamilyMismatch { + actual: text_config + .get("model_type") + .and_then(serde_json::Value::as_str) + .unwrap_or("") + .to_string(), + }); + } + let mut vision_config: Qwen3VLVisionConfig = + parse_required_vlm_subconfig(&full_config, "vision_config", "Qwen3VL 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| { + name.starts_with("language_model.model.embed_tokens.") + || name.starts_with("model.language_model.embed_tokens.") + }) + .map(|weights| remap_qwen3_vl_weights(weights, false)) + .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 Qwen3-VL text embedding table: {error}" + )) + })?; + let positive_text_field = |field: &str| { + text_config + .get(field) + .and_then(serde_json::Value::as_u64) + .and_then(|value| usize::try_from(value).ok()) + .filter(|&value| value > 0) + .ok_or_else(|| { + HostPreprocessorError::InvalidConfig(format!( + "Qwen3-VL text_config.{field} must be a positive integer" + )) + }) + }; + let hidden_size = positive_text_field("hidden_size")?; + let max_sequence_len = positive_text_field("max_position_embeddings")?; + 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_with_norm(&vision_config, [0.5, 0.5, 0.5], [0.5, 0.5, 0.5]); + let projector = mlxcel_xla::IreeQwen3VlProjector::load(model_path, device) + .map_err(HostPreprocessorError::Iree)?; + Qwen3VlIreeHostPreprocessor::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 bfad9cb9..ddd8c23f 100644 --- a/src/multimodal/host_preprocessor.rs +++ b/src/multimodal/host_preprocessor.rs @@ -44,8 +44,8 @@ mod export; use super::qwen_vl::insert_qwen_vl_image_tokens; use export::export_mlx_tensor; use export::{ - build_prepared_prefill, export_llava_prefill, export_qwen2_vl_prefill, usize_to_i32, - validate_embedding_shape, validate_processor_shape, validate_projected_shape, + build_prepared_prefill, export_llava_prefill, export_qwen2_vl_prefill, export_qwen3_vl_prefill, + usize_to_i32, validate_embedding_shape, validate_processor_shape, validate_projected_shape, validate_sequence_capacity, }; @@ -124,6 +124,20 @@ pub trait HostMultimodalPreprocessor { token_ids: &[i32], images: &[DynamicImage], ) -> Result; + + /// Prepare the distinct DeepStack payload used by Qwen3-VL. + /// + /// `None` is the safe default for every ordinary image family. Callers must + /// check this entry before invoking [`Self::prepare`], so a DeepStack model + /// can never be silently downgraded to main-merger-only embeddings. + #[cfg(feature = "xla-backend")] + fn prepare_deepstack( + &self, + _token_ids: &[i32], + _images: &[DynamicImage], + ) -> Result, HostPreprocessorError> { + Ok(None) + } } /// Load the image preprocessor supported by the OpenXLA host-first path. @@ -175,6 +189,36 @@ pub fn load_xla_image_preprocessor( )); } } + if model_type == crate::models::ModelType::Qwen3VL { + let policy = XlaVisionBackendPolicy::from_env()?; + if policy == XlaVisionBackendPolicy::Host { + return Err(HostPreprocessorError::InvalidConfig( + "Qwen3-VL XLA vision has no MLX fallback; MLXCEL_XLA_VISION_BACKEND=host is unsupported" + .to_string(), + )); + } + #[cfg(feature = "xla-iree")] + { + let device = std::env::var("MLXCEL_XLA_DEVICE") + .unwrap_or_else(|_| mlxcel_xla::default_device().to_string()); + let preprocessor = Qwen3VlIreeHostPreprocessor::load(model_path, &device)?; + tracing::info!( + vision_backend = "iree", + vision_device = %device, + family = "qwen3_vl", + media = "image", + video = false, + "OpenXLA multimodal vision backend selected" + ); + return Ok(Some(Box::new(preprocessor))); + } + #[cfg(not(feature = "xla-iree"))] + { + return Err(HostPreprocessorError::InvalidConfig( + "Qwen3-VL XLA image execution requires the xla-iree feature".to_string(), + )); + } + } if model_type != crate::models::ModelType::LlavaVLM { return Ok(None); } @@ -292,6 +336,23 @@ pub struct Qwen2VlIreeHostPreprocessor { device: String, } +/// Dense Qwen3-VL producer. The resident IREE graph returns the main merger and +/// every ordered DeepStack branch; no MLX vision or MoE fallback is admitted. +#[cfg(feature = "xla-iree")] +pub struct Qwen3VlIreeHostPreprocessor { + processor: crate::vision::processors::qwen2_vl::Qwen2VLProcessor, + text_embeddings: UnifiedEmbedding, + projector: RefCell, + image_token_id: i32, + video_token_id: i32, + vision_start_token_id: i32, + spatial_merge_size: usize, + hidden_size: usize, + max_sequence_len: usize, + deepstack_layer_count: usize, + device: String, +} + #[cfg(feature = "xla-iree")] impl Qwen2VlIreeHostPreprocessor { pub fn load(model_path: &Path, device: &str) -> Result { @@ -458,6 +519,270 @@ impl HostMultimodalPreprocessor for Qwen2VlIreeHostPreprocessor { } } +#[cfg(feature = "xla-iree")] +impl Qwen3VlIreeHostPreprocessor { + pub fn load(model_path: &Path, device: &str) -> Result { + crate::loading::load_qwen3_vl_iree_host_preprocessor(model_path, device) + } + + #[allow(clippy::too_many_arguments)] + pub(crate) fn from_parts( + processor: crate::vision::processors::qwen2_vl::Qwen2VLProcessor, + text_embeddings: UnifiedEmbedding, + projector: mlxcel_xla::IreeQwen3VlProjector, + image_token_id: i32, + video_token_id: i32, + vision_start_token_id: i32, + spatial_merge_size: usize, + hidden_size: usize, + max_sequence_len: usize, + device: String, + ) -> Result { + if hidden_size == 0 || max_sequence_len == 0 || spatial_merge_size == 0 { + return Err(HostPreprocessorError::InvalidConfig( + "Qwen3-VL hidden size, sequence capacity, and spatial merge size must be positive" + .to_string(), + )); + } + if projector.text_hidden() != hidden_size { + return Err(HostPreprocessorError::InvalidConfig(format!( + "Qwen3-VL IREE projector hidden size {} disagrees with text hidden size {hidden_size}", + projector.text_hidden() + ))); + } + let deepstack_layer_count = projector.deepstack_layer_count(); + if deepstack_layer_count == 0 { + return Err(HostPreprocessorError::InvalidConfig( + "Qwen3-VL IREE requires at least one DeepStack branch".to_string(), + )); + } + Ok(Self { + processor, + text_embeddings, + projector: RefCell::new(projector), + image_token_id, + video_token_id, + vision_start_token_id, + spatial_merge_size, + hidden_size, + max_sequence_len, + deepstack_layer_count, + device, + }) + } + + fn prepare_iree_deepstack( + &self, + token_ids: &[i32], + images: &[DynamicImage], + ) -> Result { + if images.is_empty() { + return Err(HostPreprocessorError::InvalidConfig( + "Qwen3-VL image preprocessing requires at least one decoded image; text-only requests use ordinary XLA token prefill" + .to_string(), + )); + } + let (patch_values, grids) = self.processor.preprocess_values_with_grid(images); + let mut logical_tokens = token_ids.to_vec(); + let inserted = insert_qwen_vl_image_tokens( + &mut logical_tokens, + &grids, + self.spatial_merge_size, + self.vision_start_token_id, + self.image_token_id, + ) + .ok_or_else(|| { + HostPreprocessorError::InvalidConfig(format!( + "Qwen3-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!( + "Qwen3-VL expanded {} image blocks for {} decoded images", + inserted.image_blocks, + images.len() + ))); + } + validate_sequence_capacity(logical_tokens.len(), self.max_sequence_len)?; + let input_ids = mlxcel_core::from_slice_i32( + &logical_tokens, + &[1, usize_to_i32(logical_tokens.len(), "sequence length")?], + ); + let text_embeddings = mlxcel_core::astype( + &self.text_embeddings.forward(&input_ids), + mlxcel_core::dtype::FLOAT32, + ); + validate_embedding_shape( + &mlxcel_core::array_shape(&text_embeddings), + logical_tokens.len(), + self.hidden_size, + "Qwen3-VL text embedding table", + )?; + let mut projector = self.projector.try_borrow_mut().map_err(|_| { + HostPreprocessorError::Iree( + "concurrent/re-entrant Qwen3-VL IREE vision invocation is unsupported".to_string(), + ) + })?; + let projection = projector + .project(&patch_values, &grids) + .map_err(HostPreprocessorError::Iree)?; + let expected_projected_tokens = usize::try_from(inserted.total_image_tokens) + .map_err(|_| HostPreprocessorError::ShapeOverflow)?; + if projection.shape != [expected_projected_tokens, self.hidden_size] { + return Err(HostPreprocessorError::InvalidConfig(format!( + "Qwen3-VL IREE projection shape {:?} disagrees with expanded image-token shape [{expected_projected_tokens}, {}]", + projection.shape, self.hidden_size + ))); + } + if projection.deepstack_values.len() != self.deepstack_layer_count + || projection + .deepstack_values + .iter() + .any(|values| values.len() != expected_projected_tokens * self.hidden_size) + { + return Err(HostPreprocessorError::InvalidConfig(format!( + "Qwen3-VL IREE returned {} DeepStack branch(es) with inconsistent visual shapes; expected {} branch(es) of [{expected_projected_tokens}, {}]", + projection.deepstack_values.len(), + self.deepstack_layer_count, + self.hidden_size + ))); + } + tracing::info!( + vision_backend = "iree", + vision_device = %self.device, + image_count = images.len(), + deepstack_layers = self.deepstack_layer_count, + patch_upload_bytes = projection.metrics.patch_upload_bytes, + metadata_upload_bytes = projection.metrics.metadata_upload_bytes, + projected_transfer_bytes = projection.metrics.projected_transfer_bytes, + iree_vision_seconds = projection.metrics.elapsed_seconds, + packed_cu_seqlens = ?projection.packed_cu_seqlens, + "Qwen3-VL OpenXLA vision and DeepStack projection completed" + ); + let projected = mlxcel_core::from_slice_f32( + &projection.values, + &[ + usize_to_i32(expected_projected_tokens, "Qwen3-VL projected token count")?, + usize_to_i32(self.hidden_size, "hidden size")?, + ], + ); + let merged = merge_llava( + self.image_token_id, + &projected, + &text_embeddings, + &input_ids, + ); + let prepared = export_qwen3_vl_prefill( + logical_tokens.clone(), + InputEmbeddings { + inputs_embeds: mlxcel_core::astype( + &merged.inputs_embeds, + mlxcel_core::dtype::FLOAT32, + ), + attention_mask_4d: merged.attention_mask_4d, + }, + &grids, + self.image_token_id, + self.video_token_id, + self.spatial_merge_size, + self.hidden_size, + )?; + let visual_positions = logical_tokens + .iter() + .enumerate() + .filter_map(|(position, &token)| (token == self.image_token_id).then_some(position)) + .map(|position| { + i32::try_from(position).map_err(|_| HostPreprocessorError::ShapeOverflow) + }) + .collect::, _>>()?; + if visual_positions.len() != expected_projected_tokens { + return Err(HostPreprocessorError::InvalidConfig(format!( + "Qwen3-VL visual position count {} disagrees with projected token count {expected_projected_tokens}", + visual_positions.len() + ))); + } + let position_bytes = visual_positions + .into_iter() + .flat_map(i32::to_le_bytes) + .collect::>(); + let feature_values = projection + .deepstack_values + .into_iter() + .flatten() + .collect::>(); + let feature_bytes = feature_values + .into_iter() + .flat_map(f32::to_le_bytes) + .collect::>(); + let layer_indices = (0..self.deepstack_layer_count) + .map(|layer| { + i32::try_from(layer) + .map(i32::to_le_bytes) + .map_err(|_| HostPreprocessorError::ShapeOverflow) + }) + .collect::, _>>()? + .into_iter() + .flatten() + .collect::>(); + let features = mlxcel_xla::DeepStackFeatures::new( + OwnedTensor::new( + position_bytes, + PreparedTensorDType::Int32, + vec![expected_projected_tokens], + )?, + OwnedTensor::new( + feature_bytes, + PreparedTensorDType::Float32, + vec![ + self.deepstack_layer_count, + expected_projected_tokens, + self.hidden_size, + ], + )?, + OwnedTensor::new( + layer_indices, + PreparedTensorDType::Int32, + vec![self.deepstack_layer_count], + )?, + ) + .map_err(|error| HostPreprocessorError::InvalidConfig(error.to_string()))?; + mlxcel_xla::DeepStackPreparedPrefill::new(prepared, features) + .map_err(|error| HostPreprocessorError::InvalidConfig(error.to_string())) + } +} + +#[cfg(feature = "xla-iree")] +impl HostMultimodalPreprocessor for Qwen3VlIreeHostPreprocessor { + fn backend(&self) -> XlaVisionBackend { + XlaVisionBackend::Iree + } + + fn prepare( + &self, + _token_ids: &[i32], + _images: &[DynamicImage], + ) -> Result { + reject_qwen3_standard_prefill() + } + + fn prepare_deepstack( + &self, + token_ids: &[i32], + images: &[DynamicImage], + ) -> Result, HostPreprocessorError> { + self.prepare_iree_deepstack(token_ids, images).map(Some) + } +} + +#[cfg(feature = "xla-iree")] +pub(super) fn reject_qwen3_standard_prefill() -> Result { + Err(HostPreprocessorError::InvalidConfig( + "Qwen3-VL requires the DeepStack prepared-prefill entry; refusing to omit side features" + .to_string(), + )) +} + #[cfg(feature = "xla-iree")] impl LlavaIreeHostPreprocessor { /// Load the host processor/text embedding table and resident IREE vision diff --git a/src/multimodal/host_preprocessor_export.rs b/src/multimodal/host_preprocessor_export.rs index e50c288f..a2ebc824 100644 --- a/src/multimodal/host_preprocessor_export.rs +++ b/src/multimodal/host_preprocessor_export.rs @@ -159,32 +159,79 @@ pub(super) fn export_qwen2_vl_prefill( video_token_id: i32, spatial_merge_size: usize, hidden_size: usize, +) -> Result { + export_qwen_vl_prefill( + logical_tokens, + merged, + grids, + image_token_id, + video_token_id, + spatial_merge_size, + hidden_size, + "Qwen2-VL", + "qwen2_vl", + ) +} + +pub(super) fn export_qwen3_vl_prefill( + logical_tokens: Vec, + merged: InputEmbeddings, + grids: &[(i32, i32, i32)], + image_token_id: i32, + video_token_id: i32, + spatial_merge_size: usize, + hidden_size: usize, +) -> Result { + export_qwen_vl_prefill( + logical_tokens, + merged, + grids, + image_token_id, + video_token_id, + spatial_merge_size, + hidden_size, + "Qwen3-VL", + "qwen3_vl", + ) +} + +#[allow(clippy::too_many_arguments)] +fn export_qwen_vl_prefill( + logical_tokens: Vec, + merged: InputEmbeddings, + grids: &[(i32, i32, i32)], + image_token_id: i32, + video_token_id: i32, + spatial_merge_size: usize, + hidden_size: usize, + family_label: &'static str, + family_key: &'static str, ) -> Result { if spatial_merge_size == 0 { - return Err(HostPreprocessorError::InvalidConfig( - "Qwen2-VL spatial_merge_size must be positive".to_string(), - )); + return Err(HostPreprocessorError::InvalidConfig(format!( + "{family_label} spatial_merge_size must be positive" + ))); } if logical_tokens.contains(&video_token_id) { - return Err(HostPreprocessorError::InvalidConfig( - "Qwen2-VL XLA image preprocessing does not support video placeholders".to_string(), - )); + return Err(HostPreprocessorError::InvalidConfig(format!( + "{family_label} XLA image preprocessing does not support video placeholders" + ))); } validate_embedding_shape( &mlxcel_core::array_shape(&merged.inputs_embeds), logical_tokens.len(), hidden_size, - "Qwen2-VL merged embedding", + "Qwen-VL merged embedding", )?; if merged.attention_mask_4d.is_some() { - return Err(HostPreprocessorError::InvalidConfig( - "Qwen2-VL prepared prefill requires the standard causal mask".to_string(), - )); + return Err(HostPreprocessorError::InvalidConfig(format!( + "{family_label} prepared prefill requires the standard causal mask" + ))); } let merge = i32::try_from(spatial_merge_size).map_err(|_| { - HostPreprocessorError::InvalidConfig( - "Qwen2-VL spatial_merge_size does not fit i32".to_string(), - ) + HostPreprocessorError::InvalidConfig(format!( + "{family_label} spatial_merge_size does not fit i32" + )) })?; let expected_per_image = grids .iter() @@ -197,7 +244,7 @@ pub(super) fn export_qwen2_vl_prefill( || width % merge != 0 { return Err(HostPreprocessorError::InvalidConfig(format!( - "invalid Qwen2-VL image grid {index}: ({temporal},{height},{width})" + "invalid {family_label} image grid {index}: ({temporal},{height},{width})" ))); } let count = temporal @@ -234,9 +281,9 @@ pub(super) fn export_qwen2_vl_prefill( cursor += 1; } if image_index >= grids.len() { - return Err(HostPreprocessorError::InvalidConfig( - "Qwen2-VL prompt has more image-token runs than decoded images".to_string(), - )); + return Err(HostPreprocessorError::InvalidConfig(format!( + "{family_label} prompt has more image-token runs than decoded images" + ))); } let text_length = i32::try_from(vision_start - segment_start) .map_err(|_| HostPreprocessorError::ShapeOverflow)?; @@ -255,7 +302,7 @@ pub(super) fn export_qwen2_vl_prefill( let run_length = cursor - vision_start; if run_length != expected_per_image[image_index] { return Err(HostPreprocessorError::InvalidConfig(format!( - "Qwen2-VL image-token run {image_index} has {run_length} tokens, expected {}", + "{family_label} image-token run {image_index} has {run_length} tokens, expected {}", expected_per_image[image_index] ))); } @@ -291,7 +338,7 @@ pub(super) fn export_qwen2_vl_prefill( } if image_index != grids.len() { return Err(HostPreprocessorError::InvalidConfig(format!( - "Qwen2-VL prompt has {image_index} image-token runs, expected {}", + "{family_label} prompt has {image_index} image-token runs, expected {}", grids.len() ))); } @@ -305,9 +352,9 @@ pub(super) fn export_qwen2_vl_prefill( } } if axes.iter().any(|axis| axis.len() != logical_tokens.len()) { - return Err(HostPreprocessorError::InvalidConfig( - "Qwen2-VL M-RoPE position construction did not cover the logical prompt".to_string(), - )); + return Err(HostPreprocessorError::InvalidConfig(format!( + "{family_label} M-RoPE position construction did not cover the logical prompt" + ))); } let maximum = axes .iter() @@ -315,9 +362,9 @@ pub(super) fn export_qwen2_vl_prefill( .copied() .max() .ok_or_else(|| { - HostPreprocessorError::InvalidConfig( - "Qwen2-VL logical prompt must not be empty".to_string(), - ) + HostPreprocessorError::InvalidConfig(format!( + "{family_label} logical prompt must not be empty" + )) })?; let sequence_len_i32 = i32::try_from(logical_tokens.len()).map_err(|_| HostPreprocessorError::ShapeOverflow)?; @@ -331,7 +378,7 @@ pub(super) fn export_qwen2_vl_prefill( .flat_map(i32::to_ne_bytes) .collect::>(); let sequence_len = logical_tokens.len(); - let embeddings = export_mlx_tensor(&merged.inputs_embeds, "Qwen2-VL merged embedding")?; + let embeddings = export_mlx_tensor(&merged.inputs_embeds, "Qwen-VL merged embedding")?; let attention_bytes = vec![ 0u8; sequence_len @@ -358,7 +405,7 @@ pub(super) fn export_qwen2_vl_prefill( causal: true, }, vec![PreparedModality { - family: "qwen2_vl".to_string(), + family: family_key.to_string(), item_count: grids.len(), token_count: expected_image_tokens, }], diff --git a/src/multimodal/host_preprocessor_tests.rs b/src/multimodal/host_preprocessor_tests.rs index a83e0f03..e811940f 100644 --- a/src/multimodal/host_preprocessor_tests.rs +++ b/src/multimodal/host_preprocessor_tests.rs @@ -17,6 +17,8 @@ use std::sync::Once; use image::DynamicImage; use mlxcel_core::dtype; +#[cfg(feature = "xla-iree")] +use super::reject_qwen3_standard_prefill; use super::{ FakeHostMultimodalPreprocessor, HostMultimodalPreprocessor, HostPreprocessorError, XlaVisionBackend, XlaVisionBackendPolicy, export_llava_prefill, export_mlx_tensor, @@ -65,6 +67,24 @@ fn iree_vision_contract_policy_is_explicit_and_strict() { assert_eq!(fake().backend(), XlaVisionBackend::Host); } +#[cfg(feature = "xla-backend")] +#[test] +fn ordinary_preprocessors_do_not_claim_the_deepstack_entry() { + assert!( + fake() + .prepare_deepstack(&[1, -200, 2], &images(1)) + .unwrap() + .is_none() + ); +} + +#[cfg(feature = "xla-iree")] +#[test] +fn qwen3_standard_prefill_is_rejected_instead_of_dropping_deepstack() { + let error = reject_qwen3_standard_prefill().unwrap_err(); + assert!(error.to_string().contains("refusing to omit side features")); +} + #[test] fn xla_loader_keeps_text_and_unqualified_vlm_image_capability_false() { for model_type in ["llama", "qwen2_vl"] { diff --git a/src/server/batch/xla_preprocess.rs b/src/server/batch/xla_preprocess.rs index 1210156a..c8a45631 100644 --- a/src/server/batch/xla_preprocess.rs +++ b/src/server/batch/xla_preprocess.rs @@ -25,6 +25,7 @@ use std::sync::{Arc, mpsc}; use std::thread; use mlxcel_core::session::PreparedPrefill; +use mlxcel_xla::DeepStackPreparedPrefill; use crate::{ HostMultimodalPreprocessor, XlaVisionBackend, load_xla_image_preprocessor, @@ -43,8 +44,23 @@ pub(super) struct ImagePreprocessJob { pub cancelled: Arc, } +pub(super) enum PreparedImagePrefill { + Standard(PreparedPrefill), + DeepStack(DeepStackPreparedPrefill), +} + +impl PreparedImagePrefill { + #[must_use] + pub(super) fn sequence_len(&self) -> usize { + match self { + Self::Standard(prepared) => prepared.sequence_len, + Self::DeepStack(request) => request.prepared().sequence_len, + } + } +} + pub(super) enum ImagePreprocessOutcome { - Prepared(PreparedPrefill), + Prepared(PreparedImagePrefill), Cancelled, Failed(String), } @@ -188,9 +204,17 @@ fn process_job( if cancelled.load(Ordering::Acquire) { return Err("request cancelled during image decoding".to_string()); } - preprocessor - .prepare(&token_ids, &decoded) - .map_err(|error| error.to_string()) + if let Some(request) = preprocessor + .prepare_deepstack(&token_ids, &decoded) + .map_err(|error| error.to_string())? + { + Ok(PreparedImagePrefill::DeepStack(request)) + } else { + preprocessor + .prepare(&token_ids, &decoded) + .map(PreparedImagePrefill::Standard) + .map_err(|error| error.to_string()) + } })); let outcome = if cancelled.load(Ordering::Acquire) { @@ -332,6 +356,9 @@ mod tests { let ImagePreprocessOutcome::Prepared(prepared) = result.outcome else { panic!("expected prepared payload"); }; + let PreparedImagePrefill::Standard(prepared) = prepared else { + panic!("fake preprocessor must use the standard payload"); + }; assert_eq!(prepared.token_ids, vec![1, -200, -200, 2]); assert_eq!(prepared.sequence_len, 4); } diff --git a/src/server/batch/xla_worker.rs b/src/server/batch/xla_worker.rs index 81105bd6..94249257 100644 --- a/src/server/batch/xla_worker.rs +++ b/src/server/batch/xla_worker.rs @@ -50,7 +50,10 @@ use std::sync::mpsc; use std::time::Instant; use mlxcel_core::session::PreparedPrefill; -use mlxcel_xla::{EngineEvent, FinishReason as XlaFinishReason, SampleParams, XlaBatchEngine}; +use mlxcel_xla::{ + DeepStackPreparedPrefill, EngineEvent, FinishReason as XlaFinishReason, SampleParams, + XlaBatchEngine, +}; use super::BatchEngine; use super::observability::BatchObservability; @@ -87,6 +90,12 @@ pub(crate) trait XlaServingEngine { max_new_tokens: usize, params: SampleParams, ) -> Result; + fn submit_deepstack_prepared( + &mut self, + request: DeepStackPreparedPrefill, + max_new_tokens: usize, + params: SampleParams, + ) -> Result; fn cancel(&mut self, req_id: u64) -> bool; fn pump(&mut self) -> Result, String>; } @@ -128,6 +137,16 @@ impl XlaServingEngine for XlaBatchEngine { .map_err(|error| error.to_string()) } + fn submit_deepstack_prepared( + &mut self, + request: DeepStackPreparedPrefill, + max_new_tokens: usize, + params: SampleParams, + ) -> Result { + XlaBatchEngine::submit_deepstack_prepared(self, request, max_new_tokens, params) + .map_err(|error| error.to_string()) + } + fn cancel(&mut self, req_id: u64) -> bool { XlaBatchEngine::cancel(self, req_id) } diff --git a/src/server/batch/xla_worker_admission.rs b/src/server/batch/xla_worker_admission.rs index 567132af..b63c0978 100644 --- a/src/server/batch/xla_worker_admission.rs +++ b/src/server/batch/xla_worker_admission.rs @@ -20,7 +20,7 @@ use super::super::xla_audio_preprocess::{ AudioPreprocessJob, AudioPreprocessOutcome, AudioPreprocessResult, AudioQueueError, }; use super::super::xla_preprocess::{ - ImagePreprocessJob, ImagePreprocessOutcome, ImagePreprocessResult, + ImagePreprocessJob, ImagePreprocessOutcome, ImagePreprocessResult, PreparedImagePrefill, }; use super::*; @@ -327,7 +327,7 @@ impl XlaServeWorker { let _ = response_tx.send(GenerateEvent::Done(detok.finish(start, prompt_len, 0))); } - fn handle_preprocessed(&mut self, result: ImagePreprocessResult) { + pub(super) fn handle_preprocessed(&mut self, result: ImagePreprocessResult) { let Some(state) = self.pending_images.remove(&result.job_id) else { return; }; @@ -349,12 +349,19 @@ impl XlaServeWorker { return; } - let effective_prefill_len = prepared.sequence_len; + let effective_prefill_len = prepared.sequence_len(); let prompt_token_count = state.prompt_tokens.len(); - match self - .engine - .submit_prepared(prepared, state.max_tokens, state.params) - { + let submitted = match prepared { + PreparedImagePrefill::Standard(prepared) => { + self.engine + .submit_prepared(prepared, state.max_tokens, state.params) + } + PreparedImagePrefill::DeepStack(request) => { + self.engine + .submit_deepstack_prepared(request, state.max_tokens, state.params) + } + }; + match submitted { Ok(req_id) => { // Internal throughput/KV accounting uses the expanded prepared // length; the API result below retains the logical prompt count. diff --git a/src/server/batch/xla_worker_tests.rs b/src/server/batch/xla_worker_tests.rs index 39307d3d..bd8b3238 100644 --- a/src/server/batch/xla_worker_tests.rs +++ b/src/server/batch/xla_worker_tests.rs @@ -16,11 +16,12 @@ use std::collections::BTreeSet; use std::io::Cursor; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, mpsc}; -use std::time::Duration; +use std::time::{Duration, Instant}; use image::{DynamicImage, ImageFormat}; use mlxcel_core::session::{ - OwnedTensor, PreparedAdapterMode, PreparedAttentionBias, PreparedPositions, PreparedTensorDType, + OwnedTensor, PreparedAdapterMode, PreparedAttentionBias, PreparedModality, PreparedPositions, + PreparedPrefill, PreparedTensorDType, }; use super::*; @@ -39,6 +40,7 @@ struct FakeServingEngine { text_submissions: Vec>, prepared_submissions: Vec<(Vec, usize)>, prepared_modes: Vec, + deepstack_submissions: usize, } impl XlaServingEngine for FakeServingEngine { @@ -80,6 +82,20 @@ impl XlaServingEngine for FakeServingEngine { Ok(self.activate()) } + fn submit_deepstack_prepared( + &mut self, + request: DeepStackPreparedPrefill, + _max_new_tokens: usize, + _params: SampleParams, + ) -> Result { + self.prepared_submissions.push(( + request.prepared().token_ids.clone(), + request.prepared().sequence_len, + )); + self.deepstack_submissions += 1; + Ok(self.activate()) + } + fn cancel(&mut self, req_id: u64) -> bool { self.active.remove(&req_id) } @@ -116,6 +132,59 @@ fn png_bytes() -> Vec { bytes } +fn deepstack_prepared() -> DeepStackPreparedPrefill { + let prepared = PreparedPrefill::new( + vec![1, -200, 2], + OwnedTensor::new( + vec![0; 3 * 2 * std::mem::size_of::()], + PreparedTensorDType::Float32, + vec![1, 3, 2], + ) + .unwrap(), + PreparedPositions::Sequential { + start: 0, + length: 3, + }, + PreparedAttentionBias { + tensor: OwnedTensor::new( + vec![0; 3 * std::mem::size_of::()], + PreparedTensorDType::Float32, + vec![1, 1, 1, 3], + ) + .unwrap(), + causal: true, + }, + vec![PreparedModality { + family: "qwen3_vl".to_string(), + item_count: 1, + token_count: 1, + }], + ) + .unwrap(); + let features = mlxcel_xla::DeepStackFeatures::new( + OwnedTensor::new( + 1i32.to_le_bytes().to_vec(), + PreparedTensorDType::Int32, + vec![1], + ) + .unwrap(), + OwnedTensor::new( + vec![0; 2 * std::mem::size_of::()], + PreparedTensorDType::Float32, + vec![1, 1, 2], + ) + .unwrap(), + OwnedTensor::new( + 0i32.to_le_bytes().to_vec(), + PreparedTensorDType::Int32, + vec![1], + ) + .unwrap(), + ) + .unwrap(); + DeepStackPreparedPrefill::new(prepared, features).unwrap() +} + fn wav_bytes() -> Vec { let samples = [0i16; 320]; let data_len = samples.len() * 2; @@ -515,6 +584,54 @@ fn cancelled_image_does_not_disturb_active_text_or_prevent_slot_reuse() { assert_eq!(receive_done(&reuse_rx).prompt_tokens, 3); } +#[test] +fn deepstack_outcomes_route_distinctly_and_cancelled_payloads_leave_no_slot_state() { + use crate::server::batch::xla_preprocess::{ + ImagePreprocessOutcome, ImagePreprocessResult, PreparedImagePrefill, + }; + + let pending = |cancelled: Arc, response_tx| PendingImageState { + response_tx, + cancelled, + prompt_tokens: vec![1, -200, 2], + params: SampleParams::greedy(), + stop_sequences: Vec::new(), + max_tokens: 1, + start: Instant::now(), + }; + + let mut routed = worker(None); + let (routed_tx, _routed_rx) = mpsc::channel(); + routed + .pending_images + .insert(41, pending(Arc::new(AtomicBool::new(false)), routed_tx)); + routed.handle_preprocessed(ImagePreprocessResult { + job_id: 41, + outcome: ImagePreprocessOutcome::Prepared(PreparedImagePrefill::DeepStack( + deepstack_prepared(), + )), + }); + assert_eq!(routed.engine.deepstack_submissions, 1); + assert_eq!(routed.engine.prepared_submissions.len(), 1); + assert_eq!(routed.states.len(), 1); + + let mut cancelled = worker(None); + let (cancelled_tx, _cancelled_rx) = mpsc::channel(); + cancelled + .pending_images + .insert(42, pending(Arc::new(AtomicBool::new(true)), cancelled_tx)); + cancelled.handle_preprocessed(ImagePreprocessResult { + job_id: 42, + outcome: ImagePreprocessOutcome::Prepared(PreparedImagePrefill::DeepStack( + deepstack_prepared(), + )), + }); + assert_eq!(cancelled.engine.deepstack_submissions, 0); + assert!(cancelled.engine.prepared_submissions.is_empty()); + assert!(cancelled.states.is_empty()); + assert!(cancelled.pending_images.is_empty()); +} + #[test] fn unsupported_output_features_are_rejected_before_engine_admission() { assert!( From ddd793f56737ea205a41fccaf2f04697d422184e Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 24 Jul 2026 20:13:23 +0900 Subject: [PATCH 02/10] fix(xla): honor Qwen3-VL processor contract Read the checkpoint pixel bounds and PIL bicubic resampler instead of using the historical Qwen2 defaults, and include the effective policy in the compiled artifact identity so incompatible processor artifacts cannot be reused. Qualify the 512-patch bucket required by two canonical 256x256 images, retain per-image packed attention isolation, cap the corresponding DeepStack payload at 128 merged positions, and cover processor, overflow, ABI, and context-capacity boundaries. Refs #868 --- src/lib/mlxcel-xla/src/emitter/mod.rs | 2 +- src/lib/mlxcel-xla/src/emitter/qwen3_vl.rs | 53 +++++-- src/lib/mlxcel-xla/src/qwen3_vl_runtime.rs | 31 +++- src/loading/vlm_qwen.rs | 171 ++++++++++++++++++++- src/multimodal/host_preprocessor_tests.rs | 16 ++ src/vision/processors/qwen2_vl.rs | 32 +++- 6 files changed, 292 insertions(+), 13 deletions(-) diff --git a/src/lib/mlxcel-xla/src/emitter/mod.rs b/src/lib/mlxcel-xla/src/emitter/mod.rs index e9f68508..01042a11 100644 --- a/src/lib/mlxcel-xla/src/emitter/mod.rs +++ b/src/lib/mlxcel-xla/src/emitter/mod.rs @@ -1234,7 +1234,7 @@ mod tests { cfg.deepstack, Some(DeepStackConfig { target_layer_indices: vec![0, 1, 2], - max_visual_positions: 64, + max_visual_positions: 128, }) ); assert_eq!( diff --git a/src/lib/mlxcel-xla/src/emitter/qwen3_vl.rs b/src/lib/mlxcel-xla/src/emitter/qwen3_vl.rs index 2ea3bfa0..386f633b 100644 --- a/src/lib/mlxcel-xla/src/emitter/qwen3_vl.rs +++ b/src/lib/mlxcel-xla/src/emitter/qwen3_vl.rs @@ -31,10 +31,10 @@ use super::builder::{Builder, Ty, Val}; /// Every capacity is divisible by `spatial_merge_size²` for the checkpoint's /// 2x2 merger. Larger grids fail explicitly instead of triggering an unbounded /// compile or falling back to the MLX vision encoder. -pub(crate) const QWEN3_VL_PATCH_BUCKETS: [usize; 3] = [16, 64, 256]; +pub(crate) const QWEN3_VL_PATCH_BUCKETS: [usize; 4] = [16, 64, 256, 512]; /// Compact DeepStack side-input capacity after the required 2x2 merge. -pub(crate) const QWEN3_VL_MAX_MERGED_VISUAL_POSITIONS: usize = 256 / 4; +pub(crate) const QWEN3_VL_MAX_MERGED_VISUAL_POSITIONS: usize = 512 / 4; #[derive(Debug, Clone, PartialEq, Eq)] pub(crate) struct Qwen3VlWeightSpec { @@ -289,6 +289,12 @@ impl Qwen3VlConfig { if actual_patches == 0 { return Err("Qwen3-VL image grid must contain at least one patch".to_string()); } + let maximum = QWEN3_VL_PATCH_BUCKETS[QWEN3_VL_PATCH_BUCKETS.len() - 1]; + if actual_patches > maximum { + return Err(format!( + "Qwen3-VL image grid has {actual_patches} patches, exceeding qualified capacity {maximum}" + )); + } let merge_area = self .spatial_merge_size .checked_mul(self.spatial_merge_size) @@ -302,12 +308,7 @@ impl Qwen3VlConfig { .iter() .copied() .find(|&bucket| actual_patches <= bucket) - .ok_or_else(|| { - format!( - "Qwen3-VL image grid has {actual_patches} patches, exceeding qualified capacity {}", - QWEN3_VL_PATCH_BUCKETS[QWEN3_VL_PATCH_BUCKETS.len() - 1] - ) - }) + .ok_or_else(|| "Qwen3-VL patch bucket selection failed".to_string()) } /// Validate each image grid before selecting one packed static bucket. @@ -1021,6 +1022,7 @@ mod tests { assert_eq!(config.bucket_for_patches(16), Ok(16)); assert_eq!(config.bucket_for_patches(20), Ok(64)); assert_eq!(config.bucket_for_patches(256), Ok(256)); + assert_eq!(config.bucket_for_patches(512), Ok(512)); assert!( config .bucket_for_patches(18) @@ -1029,7 +1031,7 @@ mod tests { ); assert!( config - .bucket_for_patches(260) + .bucket_for_patches(513) .unwrap_err() .contains("exceeding qualified capacity") ); @@ -1052,6 +1054,14 @@ mod tests { assert_eq!(packed.actual_patches, 48); assert_eq!(packed.packed_cu_seqlens, vec![0, 16, 48]); assert_eq!(packed.merged_tokens_per_image, vec![4, 8]); + + let canonical_two_image = config + .plan_image_grids(&[(1, 16, 16), (1, 16, 16)]) + .unwrap(); + assert_eq!(canonical_two_image.patch_bucket, 512); + assert_eq!(canonical_two_image.actual_patches, 512); + assert_eq!(canonical_two_image.packed_cu_seqlens, vec![0, 256, 512]); + assert_eq!(canonical_two_image.merged_tokens_per_image, vec![64, 64]); } #[test] @@ -1089,6 +1099,14 @@ mod tests { assert!(mlir.contains("tensor<4x16x1xi32>")); assert!(mlir.contains("-> (tensor<4x12xf32>, tensor<4x12xf32>, tensor<4x12xf32>)")); assert_eq!(mlir.matches("chlo.erf").count(), 3); + + let two_image_mlir = emit_qwen3_vl(&config, 512); + assert!(two_image_mlir.contains("loc(\"patches.grouped\")")); + assert!(two_image_mlir.contains("tensor<512x24xf32>")); + assert!( + two_image_mlir + .contains("-> (tensor<128x12xf32>, tensor<128x12xf32>, tensor<128x12xf32>)") + ); } #[test] @@ -1133,5 +1151,22 @@ mod tests { .unwrap_err() .contains("flat index 7") ); + + let canonical_grids = [(1, 16, 16), (1, 16, 16)]; + let canonical_values = vec![0.25; 512 * 2 * row_width]; + let canonical = + prepare_qwen3_vl_host_inputs(&config, &canonical_grids, &canonical_values).unwrap(); + assert_eq!(canonical.plan.patch_bucket, 512); + assert_eq!(canonical.plan.merged_tokens_per_image, vec![64, 64]); + assert_eq!(canonical.patches.len(), 512 * 24); + assert_eq!(canonical.position_indices.len(), 4 * 512); + assert_eq!( + canonical.packed_attention_bias[255 * 512 + 256], + f32::NEG_INFINITY + ); + assert_eq!( + canonical.packed_attention_bias[256 * 512 + 255], + f32::NEG_INFINITY + ); } } diff --git a/src/lib/mlxcel-xla/src/qwen3_vl_runtime.rs b/src/lib/mlxcel-xla/src/qwen3_vl_runtime.rs index 4352e8af..5ecf245f 100644 --- a/src/lib/mlxcel-xla/src/qwen3_vl_runtime.rs +++ b/src/lib/mlxcel-xla/src/qwen3_vl_runtime.rs @@ -183,8 +183,37 @@ fn processor_identity(model_dir: &Path, config: &Qwen3VlConfig) -> Result Result { + size.and_then(|size| size.get(nested)) + .or_else(|| object.get(fallback)) + .and_then(serde_json::Value::as_u64) + .and_then(|value| usize::try_from(value).ok()) + .filter(|&value| value > 0) + .ok_or_else(|| { + format!( + "preprocessor_config.json requires positive integer size.{nested} or {fallback}" + ) + }) + }; + let min_pixels = positive_bound("shortest_edge", "min_pixels")?; + let max_pixels = positive_bound("longest_edge", "max_pixels")?; + if min_pixels > max_pixels { + return Err(format!( + "preprocessor_config.json pixel bounds are inverted: min={min_pixels}, max={max_pixels}" + )); + } + let resample = object + .get("resample") + .and_then(serde_json::Value::as_u64) + .ok_or_else(|| "preprocessor_config.json resample must be integer 3".to_string())?; + if resample != 3 { + return Err(format!( + "Qwen3-VL IREE vision requires PIL BICUBIC resample=3, got {resample}" + )); + } Ok(format!( - "qwen3-vl-processor-v1:sha256={}:patch={}:temporal={}:merge={}", + "qwen3-vl-processor-v2:sha256={}:patch={}:temporal={}:merge={}:min_pixels={min_pixels}:max_pixels={max_pixels}:resample={resample}", sha256_hex(&bytes), config.patch_size, config.temporal_patch_size, diff --git a/src/loading/vlm_qwen.rs b/src/loading/vlm_qwen.rs index 4f35453c..887f2bf5 100644 --- a/src/loading/vlm_qwen.rs +++ b/src/loading/vlm_qwen.rs @@ -25,6 +25,8 @@ //! about Qwen family details. use anyhow::Result; +#[cfg(feature = "xla-iree")] +use image::imageops::FilterType; use std::path::Path; use crate::LoadedModel; @@ -302,7 +304,10 @@ pub(crate) fn load_qwen3_vl_iree_host_preprocessor( vision_start_token_id: 151652, }, ); - let processor = qwen_vl_processor_with_norm(&vision_config, [0.5, 0.5, 0.5], [0.5, 0.5, 0.5]); + let processor_policy = load_qwen3_vl_processor_policy(model_path)?; + let processor = qwen_vl_processor_with_norm(&vision_config, [0.5, 0.5, 0.5], [0.5, 0.5, 0.5]) + .with_pixel_bounds(processor_policy.min_pixels, processor_policy.max_pixels) + .with_resize_filter(FilterType::CatmullRom); let projector = mlxcel_xla::IreeQwen3VlProjector::load(model_path, device) .map_err(HostPreprocessorError::Iree)?; Qwen3VlIreeHostPreprocessor::from_parts( @@ -319,6 +324,170 @@ pub(crate) fn load_qwen3_vl_iree_host_preprocessor( ) } +#[cfg(feature = "xla-iree")] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct Qwen3VlProcessorPolicy { + min_pixels: usize, + max_pixels: usize, +} + +#[cfg(feature = "xla-iree")] +fn load_qwen3_vl_processor_policy( + model_path: &Path, +) -> Result { + let path = model_path.join("preprocessor_config.json"); + let bytes = std::fs::read(&path).map_err(|error| { + HostPreprocessorError::InvalidConfig(format!("{}: {error}", path.display())) + })?; + let root: serde_json::Value = serde_json::from_slice(&bytes).map_err(|error| { + HostPreprocessorError::InvalidConfig(format!("parse {}: {error}", path.display())) + })?; + let object = root.as_object().ok_or_else(|| { + HostPreprocessorError::InvalidConfig(format!("{} must contain an object", path.display())) + })?; + let size = match object.get("size") { + None | Some(serde_json::Value::Null) => None, + Some(value) => Some(value.as_object().ok_or_else(|| { + HostPreprocessorError::InvalidConfig( + "Qwen3-VL preprocessor_config.json size must be an object".to_string(), + ) + })?), + }; + let positive_bound = |nested: &str, fallback: &str| -> Result { + let value = size + .and_then(|size| size.get(nested)) + .or_else(|| object.get(fallback)) + .ok_or_else(|| { + HostPreprocessorError::InvalidConfig(format!( + "Qwen3-VL preprocessor_config.json requires size.{nested} or {fallback}" + )) + })?; + value + .as_u64() + .and_then(|value| usize::try_from(value).ok()) + .filter(|&value| value > 0) + .ok_or_else(|| { + HostPreprocessorError::InvalidConfig(format!( + "Qwen3-VL preprocessor_config.json size.{nested}/{fallback} must be a positive integer" + )) + }) + }; + let min_pixels = positive_bound("shortest_edge", "min_pixels")?; + let max_pixels = positive_bound("longest_edge", "max_pixels")?; + if min_pixels > max_pixels { + return Err(HostPreprocessorError::InvalidConfig(format!( + "Qwen3-VL preprocessor pixel bounds are inverted: min={min_pixels}, max={max_pixels}" + ))); + } + match object.get("resample").and_then(serde_json::Value::as_u64) { + Some(3) => {} + Some(other) => { + return Err(HostPreprocessorError::InvalidConfig(format!( + "Qwen3-VL IREE vision requires PIL BICUBIC resample=3, got {other}" + ))); + } + None => { + return Err(HostPreprocessorError::InvalidConfig( + "Qwen3-VL preprocessor_config.json resample must be integer 3 (PIL BICUBIC)" + .to_string(), + )); + } + } + Ok(Qwen3VlProcessorPolicy { + min_pixels, + max_pixels, + }) +} + +#[cfg(all(test, feature = "xla-iree"))] +mod qwen3_vl_iree_processor_policy_tests { + use super::{Qwen3VlProcessorPolicy, load_qwen3_vl_processor_policy}; + + fn write_policy(value: serde_json::Value) -> tempfile::TempDir { + let model = tempfile::tempdir().unwrap(); + std::fs::write( + model.path().join("preprocessor_config.json"), + serde_json::to_vec(&value).unwrap(), + ) + .unwrap(); + model + } + + #[test] + fn reads_canonical_size_bounds_and_bicubic_resampler() { + let model = write_policy(serde_json::json!({ + "size": { + "shortest_edge": 65_536, + "longest_edge": 16_777_216 + }, + "resample": 3 + })); + assert_eq!( + load_qwen3_vl_processor_policy(model.path()).unwrap(), + Qwen3VlProcessorPolicy { + min_pixels: 65_536, + max_pixels: 16_777_216, + } + ); + } + + #[test] + fn accepts_explicit_legacy_pixel_bound_keys() { + let model = write_policy(serde_json::json!({ + "min_pixels": 65_536, + "max_pixels": 16_777_216, + "resample": 3 + })); + assert_eq!( + load_qwen3_vl_processor_policy(model.path()).unwrap(), + Qwen3VlProcessorPolicy { + min_pixels: 65_536, + max_pixels: 16_777_216, + } + ); + } + + #[test] + fn missing_or_malformed_policy_fails_closed() { + let missing = write_policy(serde_json::json!({"resample": 3})); + assert!( + load_qwen3_vl_processor_policy(missing.path()) + .unwrap_err() + .to_string() + .contains("shortest_edge") + ); + + let malformed = write_policy(serde_json::json!({ + "size": { + "shortest_edge": "65536", + "longest_edge": 16_777_216 + }, + "min_pixels": 65_536, + "resample": 3 + })); + assert!( + load_qwen3_vl_processor_policy(malformed.path()) + .unwrap_err() + .to_string() + .contains("positive integer") + ); + + let wrong_resampler = write_policy(serde_json::json!({ + "size": { + "shortest_edge": 65_536, + "longest_edge": 16_777_216 + }, + "resample": 1 + })); + assert!( + load_qwen3_vl_processor_policy(wrong_resampler.path()) + .unwrap_err() + .to_string() + .contains("BICUBIC") + ); + } +} + /// 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_tests.rs b/src/multimodal/host_preprocessor_tests.rs index e811940f..8d6dd6a5 100644 --- a/src/multimodal/host_preprocessor_tests.rs +++ b/src/multimodal/host_preprocessor_tests.rs @@ -17,6 +17,7 @@ use std::sync::Once; use image::DynamicImage; use mlxcel_core::dtype; +use super::export::validate_sequence_capacity; #[cfg(feature = "xla-iree")] use super::reject_qwen3_standard_prefill; use super::{ @@ -183,6 +184,21 @@ fn fake_preprocessor_rejects_expanded_sequence_over_capacity() { )); } +#[test] +fn canonical_two_image_expansion_still_obeys_model_context_capacity() { + let visual_tokens = 2 * 64; + let prompt_tokens = 2; + let expanded = visual_tokens + prompt_tokens; + let error = validate_sequence_capacity(expanded, expanded - 1).unwrap_err(); + assert!(matches!( + error, + HostPreprocessorError::SequenceCapacity { + actual, + maximum + } if actual == expanded && maximum == expanded - 1 + )); +} + #[test] fn processor_shape_validation_rejects_layout_and_size_drift() { let error = validate_processor_shape(&[1, 224, 224, 3], 1, 224).unwrap_err(); diff --git a/src/vision/processors/qwen2_vl.rs b/src/vision/processors/qwen2_vl.rs index 6924ccbb..a27adf97 100644 --- a/src/vision/processors/qwen2_vl.rs +++ b/src/vision/processors/qwen2_vl.rs @@ -34,6 +34,7 @@ pub struct Qwen2VLProcessor { pub max_pixels: usize, pub mean: [f32; 3], pub std: [f32; 3], + resize_filter: FilterType, } impl Qwen2VLProcessor { @@ -47,6 +48,7 @@ impl Qwen2VLProcessor { max_pixels: 16384 * 28 * 28, // large limit mean: [0.48145466, 0.4578275, 0.40821073], std: [0.26862954, 0.261_302_6, 0.275_777_1], + resize_filter: FilterType::Lanczos3, } } @@ -66,9 +68,26 @@ impl Qwen2VLProcessor { max_pixels: 16384 * 28 * 28, mean, std, + resize_filter: FilterType::Lanczos3, } } + /// Override the checkpoint's effective smart-resize pixel bounds. + #[must_use] + pub fn with_pixel_bounds(mut self, min_pixels: usize, max_pixels: usize) -> Self { + self.min_pixels = min_pixels; + self.max_pixels = max_pixels; + self + } + + /// Override the image resampler while preserving the historical Qwen2 + /// default for callers that do not opt in. + #[must_use] + pub fn with_resize_filter(mut self, resize_filter: FilterType) -> Self { + self.resize_filter = resize_filter; + self + } + /// Compute target size that satisfies constraints /// Returns (height, width) padded to multiples of factor pub(crate) fn smart_resize(&self, orig_h: u32, orig_w: u32) -> (u32, u32) { @@ -156,7 +175,7 @@ impl Qwen2VLProcessor { let target_w = w_patches as u32 * self.patch_size as u32; // Resize image - let resized = img.resize_exact(target_w, target_h, FilterType::Lanczos3); + let resized = img.resize_exact(target_w, target_h, self.resize_filter); let rgb = resized.to_rgb8(); // Normalize: (pixel / 255.0 - mean) / std @@ -268,4 +287,15 @@ mod tests { assert_eq!(mlx_grids, grids); assert_eq!(mlxcel_core::array_shape(&mlx), vec![32, row_width as i32]); } + + #[test] + fn qwen3_checkpoint_policy_lifts_224_square_to_sixteen_by_sixteen_grid() { + let processor = Qwen2VLProcessor::new_with_norm(16, 2, 2, [0.5; 3], [0.5; 3]) + .with_pixel_bounds(65_536, 16_777_216) + .with_resize_filter(FilterType::CatmullRom); + assert_eq!( + processor.compute_grid_thw(&[DynamicImage::new_rgb8(224, 224)]), + vec![(1, 16, 16)] + ); + } } From 00c28cacff3f6635bb88254e0f70c06fa37dd673 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Sun, 26 Jul 2026 21:46:56 +0900 Subject: [PATCH 03/10] test(xla): add Qwen3-VL DeepStack oracle Compare the eager MLX and resident IREE main vision merger and every ordered DeepStack branch from one qualified processor payload. Reuse the production DeepStack prefill loop to expose post-injection hidden states and final logits without changing serving outputs, and reject dropped or zeroed actual-checkpoint branches. Refs #868 --- Cargo.toml | 6 + examples/xla_qwen3_vl_deepstack_check.rs | 363 +++++++++++++++++++++++ src/lib.rs | 6 +- src/lib/mlxcel-xla/csrc/xla_iree.c | 60 +++- src/lib/mlxcel-xla/src/batch.rs | 42 +++ src/lib/mlxcel-xla/src/emitter/mod.rs | 6 + src/lib/mlxcel-xla/src/iree.rs | 208 ++++++++++++- src/lib/mlxcel-xla/src/lib.rs | 7 +- src/models/qwen3_vl.rs | 59 +++- src/multimodal/host_preprocessor.rs | 40 +++ 10 files changed, 775 insertions(+), 22 deletions(-) create mode 100644 examples/xla_qwen3_vl_deepstack_check.rs diff --git a/Cargo.toml b/Cargo.toml index ab506bce..9382cc96 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -329,6 +329,12 @@ required-features = ["xla-diagnostics"] name = "xla_vision_reference_check" required-features = ["xla-diagnostics"] +# Actual-checkpoint eager MLX vs IREE Qwen3-VL oracle. Compares the main +# merger, every DeepStack branch, post-injection hidden states, and logits. +[[example]] +name = "xla_qwen3_vl_deepstack_check" +required-features = ["xla-diagnostics"] + [[example]] name = "xla_aux_smoke" required-features = ["xla-iree"] diff --git a/examples/xla_qwen3_vl_deepstack_check.rs b/examples/xla_qwen3_vl_deepstack_check.rs new file mode 100644 index 00000000..36a24496 --- /dev/null +++ b/examples/xla_qwen3_vl_deepstack_check.rs @@ -0,0 +1,363 @@ +// 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 Qwen3-VL DeepStack oracle. +//! +//! The qualified production host processor creates one patch payload that is +//! consumed by both vision implementations. The gate compares the main merger, +//! every ordered DeepStack merger, post-injection language hidden states, and +//! final-position logits. It also mutates the actual captured branches to prove +//! that dropping or zeroing one branch is detected. + +use std::fs; +use std::path::{Path, PathBuf}; + +use image::DynamicImage; +use mlxcel::{ + HostMultimodalPreprocessor, LoadedModel, Qwen3VlIreeHostPreprocessor, initialize_runtime, + load_model, +}; +use mlxcel_xla::Qwen3VlDeepStackDiagnosticEngine; +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 { + let args = std::env::args().collect::>(); + args.iter() + .position(|argument| argument == flag) + .and_then(|index| args.get(index + 1)) + .cloned() +} + +fn required_path(flag: &str) -> PathBuf { + argument(flag) + .map(PathBuf::from) + .unwrap_or_else(|| panic!("missing required {flag}")) +} + +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 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_le_bytes(bytes.try_into().expect("four-byte f32 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 compare_vision_branches( + reports: &mut Vec, + actual: &[Vec], + expected: &[Vec], + tolerance: Tolerance, +) -> Result<(), String> { + if actual.len() != expected.len() { + return Err(format!( + "DeepStack branch count differs: actual={}, expected={}", + actual.len(), + expected.len() + )); + } + for (index, (actual, expected)) in actual.iter().zip(expected).enumerate() { + compare_stage( + reports, + &format!("vision.deepstack_merger_{index}"), + actual, + expected, + tolerance, + )?; + } + Ok(()) +} + +fn patch_size(model: &Path) -> usize { + let config: Value = serde_json::from_slice( + &fs::read(model.join("config.json")).expect("read Qwen3-VL config.json"), + ) + .expect("parse Qwen3-VL config.json"); + config["vision_config"]["patch_size"] + .as_u64() + .and_then(|value| usize::try_from(value).ok()) + .filter(|value| *value > 0) + .expect("vision_config.patch_size must be a positive integer") +} + +fn main() { + let model_path = required_path("--model"); + let image_path = required_path("--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 image = image::open(&image_path) + .unwrap_or_else(|error| panic!("open {}: {error}", image_path.display())) + .into_rgb8(); + let images = [DynamicImage::ImageRgb8(image)]; + + let production_preprocessor = Qwen3VlIreeHostPreprocessor::load(&model_path, &device) + .unwrap_or_else(|error| panic!("load Qwen3-VL IREE host preprocessor: {error}")); + let vision_capture = production_preprocessor + .capture_vision(&images) + .unwrap_or_else(|error| panic!("capture Qwen3-VL IREE vision: {error}")); + + let (loaded, _) = load_model(&model_path) + .unwrap_or_else(|error| panic!("load eager Qwen3-VL checkpoint: {error}")); + let LoadedModel::Qwen3VL(model) = loaded else { + panic!("checkpoint did not load as dense Qwen3-VL"); + }; + let patch_width = 3 * patch_size(&model_path).pow(2); + assert_eq!( + vision_capture.patch_values.len() % patch_width, + 0, + "processor patch payload must have complete rows" + ); + let patch_rows = vision_capture.patch_values.len() / patch_width; + let pixels = mlxcel_core::from_slice_f32( + &vision_capture.patch_values, + &[patch_rows as i32, patch_width as i32], + ); + let eager_vision = model + .vision_encoder + .forward_with_grid(&pixels, &vision_capture.grids); + let eager_main = mlx_f32(&eager_vision.hidden_states); + let eager_branches = eager_vision + .deepstack_features + .iter() + .map(|feature| mlx_f32(feature)) + .collect::>(); + + let mut reports = Vec::new(); + compare_stage( + &mut reports, + "vision.main_merger", + &vision_capture.projection.values, + &eager_main, + tolerance, + ) + .unwrap_or_else(|error| panic!("{error}")); + compare_vision_branches( + &mut reports, + &vision_capture.projection.deepstack_values, + &eager_branches, + tolerance, + ) + .unwrap_or_else(|error| panic!("{error}")); + + let mut dropped = vision_capture.projection.deepstack_values.clone(); + dropped + .pop() + .expect("actual checkpoint must expose DeepStack"); + assert!( + compare_vision_branches(&mut Vec::new(), &dropped, &eager_branches, tolerance).is_err(), + "negative oracle accepted a dropped actual-checkpoint DeepStack branch" + ); + let mut zeroed = vision_capture.projection.deepstack_values.clone(); + zeroed[0].fill(0.0); + assert!( + compare_vision_branches(&mut Vec::new(), &zeroed, &eager_branches, tolerance).is_err(), + "negative oracle accepted a zeroed actual-checkpoint DeepStack branch" + ); + + let unexpanded_tokens = vec![1, model.vision_start_token_id, model.image_token_id, 2]; + let prepared = production_preprocessor + .prepare_deepstack(&unexpanded_tokens, &images) + .unwrap_or_else(|error| panic!("prepare production DeepStack prefill: {error}")) + .expect("Qwen3-VL preprocessor must return a DeepStack payload"); + assert!( + prepared.prepared().sequence_len <= context_capacity, + "prepared sequence length {} exceeds diagnostic context capacity {context_capacity}", + prepared.prepared().sequence_len + ); + let input_ids = mlxcel_core::from_slice_i32( + &prepared.prepared().token_ids, + &[1, prepared.prepared().sequence_len as i32], + ); + let eager_embeddings = model.get_input_embeddings(&input_ids, &pixels, &vision_capture.grids); + let mut caches = model.text_model.make_caches(); + let eager_language = model.text_model.forward_deepstack_diagnostics( + &input_ids, + &eager_embeddings.inputs_embeds, + &mut caches, + eager_embeddings.attention_mask_4d.as_deref(), + ); + + let mut iree_language = + Qwen3VlDeepStackDiagnosticEngine::load(&model_path, &device, context_capacity) + .unwrap_or_else(|error| panic!("load Qwen3-VL DeepStack diagnostics: {error}")); + let iree_language = iree_language + .capture(&prepared) + .unwrap_or_else(|error| panic!("capture Qwen3-VL DeepStack diagnostics: {error}")); + assert_eq!( + eager_language.post_injection_hidden_states.len(), + iree_language.deepstack_layers, + "eager and IREE post-injection layer counts differ" + ); + let sequence_len = prepared.prepared().sequence_len; + let hidden_size = iree_language.hidden_size; + for (layer, eager) in eager_language + .post_injection_hidden_states + .iter() + .enumerate() + { + let target_layer = iree_language.target_layer_indices[layer]; + let eager = mlx_f32(eager); + let layer_start = layer * iree_language.context_capacity * hidden_size; + let actual = &iree_language.post_injection_hidden_states + [layer_start..layer_start + sequence_len * hidden_size]; + compare_stage( + &mut reports, + &format!("language.post_injection_layer_{target_layer}"), + actual, + &eager, + tolerance, + ) + .unwrap_or_else(|error| panic!("{error}")); + } + + let eager_logits = mlx_f32(&eager_language.logits); + let vocab = iree_language.logits.len(); + let eager_last = &eager_logits[(sequence_len - 1) * vocab..sequence_len * vocab]; + compare_stage( + &mut reports, + "language.final_position_logits", + &iree_language.logits, + eager_last, + tolerance, + ) + .unwrap_or_else(|error| panic!("{error}")); + + let report = json!({ + "schema": 1, + "model": model_path, + "image": image_path, + "device": device, + "context_capacity": context_capacity, + "grid_thw": vision_capture.grids, + "vision_shape": vision_capture.projection.shape, + "deepstack_branches": eager_branches.len(), + "deepstack_target_language_layers": iree_language.target_layer_indices, + "negative_dropped_branch_detected": true, + "negative_zeroed_branch_detected": true, + "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}"); +} diff --git a/src/lib.rs b/src/lib.rs index d9520a12..e273a7e4 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, Qwen3VlIreeHostPreprocessor, Qwen3VlIreeVisionCapture, +}; 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/csrc/xla_iree.c b/src/lib/mlxcel-xla/csrc/xla_iree.c index 32578ed4..4c309b27 100644 --- a/src/lib/mlxcel-xla/csrc/xla_iree.c +++ b/src/lib/mlxcel-xla/csrc/xla_iree.c @@ -567,8 +567,12 @@ static iree_status_t xla_llama_create_impl( } if (c->has_prefill_diagnostics) { iree_runtime_call_t probe; + const char* diagnostic_entry = + c->has_deepstack + ? "prefill_embeddings_deepstack_diagnostics.main" + : "prefill_diagnostics.main"; IREE_RETURN_IF_ERROR(iree_runtime_call_initialize_by_name( - c->session, iree_make_cstring_view("prefill_diagnostics.main"), &probe)); + c->session, iree_make_cstring_view(diagnostic_entry), &probe)); iree_runtime_call_deinitialize(&probe); } @@ -1367,7 +1371,8 @@ static int xla_llama_prefill_embeddings_deepstack_impl( const xla_tensor_desc* layer_features, const xla_tensor_desc* layer_indices, int32_t actual_layer_count, int32_t actual_visual_count, int32_t real_len, int32_t vocab, - int32_t* out_token, float* out_logits) { + int32_t* out_token, float* out_logits, int32_t state_count, + float* out_states) { if (!c || !c->has_deepstack || position_mode != c->position_mode || real_len <= 0 || real_len > c->context_capacity || @@ -1457,8 +1462,16 @@ static int xla_llama_prefill_embeddings_deepstack_impl( } bool batched = slot >= 0; + bool diagnostics = out_states != NULL; + size_t expected_state_count = + (size_t)c->deepstack_layers * (size_t)c->context_capacity * + (size_t)c->hidden_size; if ((!batched && !out_token) || - (batched && (slot >= c->rg_bsz || vocab <= 0 || !out_logits))) { + (batched && (slot >= c->rg_bsz || vocab <= 0 || !out_logits)) || + (diagnostics && + (!c->has_prefill_diagnostics || state_count <= 0 || + (size_t)state_count != expected_state_count)) || + (!diagnostics && state_count != 0)) { fprintf(stderr, "xla_llama_prefill_embeddings_deepstack: invalid slot/output/vocab contract\n"); return 1; @@ -1480,12 +1493,13 @@ static int xla_llama_prefill_embeddings_deepstack_impl( iree_hal_buffer_view_t* output = NULL; iree_hal_buffer_view_t* kc = NULL; iree_hal_buffer_view_t* vc = NULL; + iree_hal_buffer_view_t* states = NULL; + const char* entry = diagnostics + ? "prefill_embeddings_deepstack_diagnostics.main" + : "prefill_embeddings_deepstack.main"; XLA_CHECK_GOTO(iree_runtime_call_initialize_by_name( - c->session, - iree_make_cstring_view( - "prefill_embeddings_deepstack.main"), - &call), + c->session, iree_make_cstring_view(entry), &call), cleanup, rc); call_initialized = true; for (int32_t i = 0; i < c->n_weights; ++i) { @@ -1531,6 +1545,11 @@ static int xla_llama_prefill_embeddings_deepstack_impl( cleanup, rc); XLA_CHECK_GOTO(iree_runtime_call_outputs_pop_front_buffer_view(&call, &vc), cleanup, rc); + if (diagnostics) { + XLA_CHECK_GOTO( + iree_runtime_call_outputs_pop_front_buffer_view(&call, &states), + cleanup, rc); + } if (xla_validate_kv_pair(c, kc, vc, 4, 1) != 0) { rc = 1; goto cleanup; @@ -1539,6 +1558,11 @@ static int xla_llama_prefill_embeddings_deepstack_impl( XLA_CHECK_GOTO( xla_read_logits(c, output, out_logits, (iree_host_size_t)vocab), cleanup, rc); + if (diagnostics) { + XLA_CHECK_GOTO( + xla_read_logits(c, states, out_states, (iree_host_size_t)state_count), + cleanup, rc); + } rc = xla_store_prefill_kv_slot(c, slot, kc, vc); } else { XLA_CHECK_GOTO(xla_read_token(c, output, out_token), cleanup, rc); @@ -1557,6 +1581,7 @@ static int xla_llama_prefill_embeddings_deepstack_impl( if (output) iree_hal_buffer_view_release(output); if (kc) iree_hal_buffer_view_release(kc); if (vc) iree_hal_buffer_view_release(vc); + if (states) iree_hal_buffer_view_release(states); if (embeddings_bv) iree_hal_buffer_view_release(embeddings_bv); if (positions_bv) iree_hal_buffer_view_release(positions_bv); if (len_bv) iree_hal_buffer_view_release(len_bv); @@ -1581,7 +1606,8 @@ int xla_llama_prefill_embeddings_deepstack( return xla_llama_prefill_embeddings_deepstack_impl( c, -1, adapter_mode, position_mode, embeddings, positions, attention_bias, visual_positions, layer_features, layer_indices, - actual_layer_count, actual_visual_count, real_len, 0, out_token, NULL); + actual_layer_count, actual_visual_count, real_len, 0, out_token, NULL, 0, + NULL); } int xla_llama_prefill_embeddings_deepstack_slot_logits( @@ -1597,7 +1623,23 @@ int xla_llama_prefill_embeddings_deepstack_slot_logits( c, slot, adapter_mode, position_mode, embeddings, positions, attention_bias, visual_positions, layer_features, layer_indices, actual_layer_count, actual_visual_count, real_len, vocab, NULL, - out_logits); + out_logits, 0, NULL); +} + +int xla_llama_prefill_embeddings_deepstack_slot_diagnostics( + xla_ctx* c, int32_t slot, int32_t adapter_mode, int32_t position_mode, + const xla_tensor_desc* embeddings, + const xla_tensor_desc* positions, const xla_tensor_desc* attention_bias, + const xla_tensor_desc* visual_positions, + const xla_tensor_desc* layer_features, + const xla_tensor_desc* layer_indices, int32_t actual_layer_count, + int32_t actual_visual_count, int32_t real_len, int32_t vocab, + int32_t state_count, float* out_logits, float* out_states) { + return xla_llama_prefill_embeddings_deepstack_impl( + c, slot, adapter_mode, position_mode, embeddings, positions, + attention_bias, visual_positions, layer_features, layer_indices, + actual_layer_count, actual_visual_count, real_len, vocab, NULL, + out_logits, state_count, out_states); } // Ragged decode_step: token[B], pos[B], cache_len[B] (per row), rank-5 KV -> diff --git a/src/lib/mlxcel-xla/src/batch.rs b/src/lib/mlxcel-xla/src/batch.rs index b08cb4b2..43395776 100644 --- a/src/lib/mlxcel-xla/src/batch.rs +++ b/src/lib/mlxcel-xla/src/batch.rs @@ -901,6 +901,48 @@ pub struct LlavaReferenceDiagnosticRun { pub decode_seconds: f64, } +/// Diagnostics-only runner for the actual-checkpoint Qwen3-VL DeepStack +/// language oracle. +#[cfg(feature = "diagnostics")] +pub struct Qwen3VlDeepStackDiagnosticEngine { + engine: IreeRaggedLlama, +} + +#[cfg(feature = "diagnostics")] +impl Qwen3VlDeepStackDiagnosticEngine { + pub fn load(model_path: &Path, device: &str, context_capacity: usize) -> Result { + Ok(Self { + engine: IreeRaggedLlama::load_with_deepstack_diagnostics( + model_path, + device, + 4, + context_capacity, + )?, + }) + } + + #[must_use] + pub fn context_capacity(&self) -> usize { + self.engine.context_capacity() + } + + pub fn capture( + &mut self, + prepared: &DeepStackPreparedPrefill, + ) -> Result { + self.engine.reset_slots_for_diagnostics()?; + let prepared_iree = PreparedIreePrefill::prepare( + prepared.prepared(), + self.engine.hidden_size(), + self.engine.context_capacity(), + ) + .map_err(|error| error.to_string())?; + let deepstack = self.engine.prepare_deepstack(prepared.deepstack())?; + self.engine + .prefill_deepstack_prepared_slot_diagnostics(0, &prepared_iree, &deepstack) + } +} + #[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 01042a11..7474b3a9 100644 --- a/src/lib/mlxcel-xla/src/emitter/mod.rs +++ b/src/lib/mlxcel-xla/src/emitter/mod.rs @@ -1133,6 +1133,12 @@ mod tests { ), "ragged DeepStack ABI must carry the explicit position mode" ); + assert!( + source.contains( + "int xla_llama_prefill_embeddings_deepstack_slot_diagnostics(\n xla_ctx* c, int32_t slot, int32_t adapter_mode, int32_t position_mode," + ), + "diagnostic DeepStack ABI must preserve the production mode contract" + ); assert!( source.contains("desc->byte_length != expected_bytes"), "the shared descriptor validator must reject rank-correct truncated buffers" diff --git a/src/lib/mlxcel-xla/src/iree.rs b/src/lib/mlxcel-xla/src/iree.rs index e3110019..5854b676 100644 --- a/src/lib/mlxcel-xla/src/iree.rs +++ b/src/lib/mlxcel-xla/src/iree.rs @@ -69,7 +69,7 @@ use crate::emitter::{ #[cfg(feature = "diagnostics")] use crate::emitter::{ Gemma3nDiagnosticLayout, emit_gemma3n_all_layer_diagnostics_with_qmv, - emit_gemma3n_prefill_diagnostics_with_qmv, + emit_gemma3n_prefill_diagnostics_with_qmv, emit_prefill_embeddings_deepstack_diagnostics_with, }; use crate::prepared::{ DecodePositionState, PreparedIreePrefill, PreparedPositionMode, canonical_text_positions, @@ -374,6 +374,26 @@ unsafe extern "C" { vocab: c_int, out_logits: *mut f32, ) -> c_int; + #[cfg(feature = "diagnostics")] + fn xla_llama_prefill_embeddings_deepstack_slot_diagnostics( + c: *mut XlaCtx, + slot: c_int, + adapter_mode: c_int, + position_mode: c_int, + embeddings: *const XlaTensorDesc, + positions: *const XlaTensorDesc, + attention_bias: *const XlaTensorDesc, + visual_positions: *const XlaTensorDesc, + layer_features: *const XlaTensorDesc, + layer_indices: *const XlaTensorDesc, + actual_layer_count: c_int, + actual_visual_count: c_int, + real_len: c_int, + vocab: c_int, + state_count: c_int, + out_logits: *mut f32, + out_states: *mut f32, + ) -> c_int; fn xla_llama_decode_ragged_logits( c: *mut XlaCtx, bsz: c_int, @@ -1050,6 +1070,33 @@ fn compile_gemma3n_diagnostics( Ok((vmfb, layout)) } +#[cfg(feature = "diagnostics")] +fn compile_deepstack_diagnostics(device: &str, cfg: &RuntimeConfig) -> Result { + let RuntimeConfig::Dense(config) = cfg else { + return Err("DeepStack diagnostics require a dense runtime config".to_string()); + }; + if config.deepstack.is_none() { + return Err("DeepStack diagnostics require a declared DeepStack schema".to_string()); + } + let precision = effective_precision(device, cfg)?; + let mlir = emit_prefill_embeddings_deepstack_diagnostics_with(config, precision); + let compiler = iree_compile_bin()?; + if !compiler.exists() { + return Err(format!("iree-compile not found at {}", compiler.display())); + } + let cache = std::env::temp_dir().join("mlxcel-xla-vmfb"); + std::fs::create_dir_all(&cache) + .map_err(|error| format!("mkdir {}: {error}", cache.display()))?; + compile_one( + &compiler, + &mlir, + target_flags(device)?, + &cache, + "prefill_embeddings_deepstack_diagnostics", + cfg.context_capacity(), + ) +} + /// Map each needed weight name to the safetensors file that holds it, in the same /// order as `names`. A single-file checkpoint (`model.safetensors` present) maps /// every name to that one file; a sharded checkpoint (no `model.safetensors`, only @@ -2157,6 +2204,18 @@ pub struct PreparedPrefillDiagnostics { pub kv_width: usize, } +#[cfg(feature = "diagnostics")] +#[derive(Debug, Clone, PartialEq)] +pub struct DeepStackPrefillDiagnostics { + pub logits: Vec, + /// Flattened `[deepstack_layers, context_capacity, hidden_size]` states. + pub post_injection_hidden_states: Vec, + pub target_layer_indices: Vec, + pub deepstack_layers: usize, + pub context_capacity: usize, + pub hidden_size: usize, +} + pub struct IreeRaggedLlama { ctx: *mut XlaCtx, b_max: usize, @@ -2187,7 +2246,15 @@ impl IreeRaggedLlama { b_max: usize, context_capacity: usize, ) -> Result { - Self::load_inner(model_dir, device, b_max, context_capacity, false, false) + Self::load_inner( + model_dir, + device, + b_max, + context_capacity, + false, + false, + false, + ) } #[cfg(feature = "diagnostics")] @@ -2197,7 +2264,15 @@ impl IreeRaggedLlama { b_max: usize, context_capacity: usize, ) -> Result { - Self::load_inner(model_dir, device, b_max, context_capacity, true, false) + Self::load_inner( + model_dir, + device, + b_max, + context_capacity, + true, + false, + false, + ) } #[cfg(feature = "diagnostics")] @@ -2207,7 +2282,33 @@ impl IreeRaggedLlama { b_max: usize, context_capacity: usize, ) -> Result { - Self::load_inner(model_dir, device, b_max, context_capacity, false, true) + Self::load_inner( + model_dir, + device, + b_max, + context_capacity, + false, + true, + false, + ) + } + + #[cfg(feature = "diagnostics")] + pub fn load_with_deepstack_diagnostics( + model_dir: &Path, + device: &str, + b_max: usize, + context_capacity: usize, + ) -> Result { + Self::load_inner( + model_dir, + device, + b_max, + context_capacity, + false, + false, + true, + ) } fn load_inner( @@ -2217,8 +2318,15 @@ impl IreeRaggedLlama { context_capacity: usize, diagnostics: bool, all_layer_diagnostics: bool, + deepstack_diagnostics: bool, ) -> Result { - debug_assert!(!(diagnostics && all_layer_diagnostics)); + debug_assert!( + [diagnostics, all_layer_diagnostics, deepstack_diagnostics] + .into_iter() + .filter(|enabled| *enabled) + .count() + <= 1 + ); let cfg = RuntimeConfig::from_json(model_dir, context_capacity)?; runtime_ffi_dimensions(&cfg, cfg.weight_specs().len())?; checked_ffi_int(b_max, "b_max")?; @@ -2280,11 +2388,13 @@ impl IreeRaggedLlama { let (diagnostic_vmfb, diagnostic_layout) = if diagnostics || all_layer_diagnostics { let (vmfb, layout) = compile_gemma3n_diagnostics(device, &cfg, all_layer_diagnostics)?; (Some(vmfb), Some(layout)) + } else if deepstack_diagnostics { + (Some(compile_deepstack_diagnostics(device, &cfg)?), None) } else { (None, None) }; #[cfg(not(feature = "diagnostics"))] - debug_assert!(!diagnostics && !all_layer_diagnostics); + debug_assert!(!diagnostics && !all_layer_diagnostics && !deepstack_diagnostics); #[cfg(feature = "diagnostics")] let ctx = create_ctx_with_diagnostics( model_dir, @@ -2706,6 +2816,92 @@ impl IreeRaggedLlama { Ok(logits) } + /// Invoke the diagnostics-only DeepStack graph and read back every + /// post-injection hidden state without changing the production graph ABI. + #[cfg(feature = "diagnostics")] + pub fn prefill_deepstack_prepared_slot_diagnostics( + &mut self, + slot: usize, + prepared: &PreparedIreePrefill, + deepstack: &PreparedDeepStack, + ) -> Result { + let schema = self.deepstack_schema.as_ref().ok_or_else(|| { + "DeepStack diagnostics were requested from a runtime without that capability" + .to_string() + })?; + validate_slot(slot, self.b_max).map_err(|error| error.to_string())?; + if prepared.hidden_size != self.hidden_size + || prepared.context_capacity != self.context_capacity + || prepared.effective_len == 0 + || prepared.effective_len > self.context_capacity + || prepared.positions.mode() != self.position_mode + || deepstack.hidden_size != self.hidden_size + || deepstack.max_layer_count != schema.target_layer_indices.len() + || deepstack.max_visual_count != schema.max_visual_positions + { + return Err( + "DeepStack diagnostic payload is incompatible with this runtime bundle".to_string(), + ); + } + let (embeddings, positions, attention_bias) = prepared_descriptors(prepared)?; + let (visual_positions, layer_features, layer_indices) = deepstack_descriptors(deepstack)?; + let actual_layer_count = + checked_ffi_int(deepstack.actual_layer_count, "DeepStack actual layer count")?; + let actual_visual_count = checked_ffi_int( + deepstack.actual_visual_count, + "DeepStack actual visual count", + )?; + let real_len = checked_ffi_int(prepared.effective_len, "prepared effective length")?; + let slot = checked_ffi_int(slot, "slot")?; + let vocab = checked_ffi_int(self.vocab, "vocab_size")?; + let state_count = schema + .target_layer_indices + .len() + .checked_mul(self.context_capacity) + .and_then(|count| count.checked_mul(self.hidden_size)) + .ok_or_else(|| "DeepStack diagnostic state count overflowed".to_string())?; + let state_count_ffi = checked_ffi_int(state_count, "DeepStack diagnostic state count")?; + let mut logits = vec![0.0; self.vocab]; + let mut post_injection_hidden_states = vec![0.0; state_count]; + // Safety: all typed descriptors and output buffers remain live through + // the call. The C shim repeats the production DeepStack validation and + // verifies the static diagnostics output element count before readback. + let rc = unsafe { + xla_llama_prefill_embeddings_deepstack_slot_diagnostics( + self.ctx, + slot, + 0, + prepared.positions.mode().ffi_code(), + &embeddings, + &positions, + &attention_bias, + &visual_positions, + &layer_features, + &layer_indices, + actual_layer_count, + actual_visual_count, + real_len, + vocab, + state_count_ffi, + logits.as_mut_ptr(), + post_injection_hidden_states.as_mut_ptr(), + ) + }; + if rc != 0 { + return Err(format!( + "xla_llama_prefill_embeddings_deepstack_slot_diagnostics failed (status {rc})" + )); + } + Ok(DeepStackPrefillDiagnostics { + logits, + post_injection_hidden_states, + target_layer_indices: schema.target_layer_indices.clone(), + deepstack_layers: schema.target_layer_indices.len(), + context_capacity: self.context_capacity, + hidden_size: self.hidden_size, + }) + } + /// Seed one Gemma3n batch slot from owned prepared embeddings and dense PLE. pub fn prefill_gemma3n_prepared_slot_logits( &mut self, diff --git a/src/lib/mlxcel-xla/src/lib.rs b/src/lib/mlxcel-xla/src/lib.rs index 854bf439..434058c0 100644 --- a/src/lib/mlxcel-xla/src/lib.rs +++ b/src/lib/mlxcel-xla/src/lib.rs @@ -147,8 +147,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, Qwen3VlDeepStackDiagnosticEngine, + run_gemma3n_all_layer_diagnostics, run_gemma3n_canonical_diagnostics, + run_gemma3n_prefix_decode_diagnostic, }; pub use context::{ CONTEXT_CAPACITY_ENV, ContextCapacityError, DEFAULT_CONTEXT_CAPACITY, @@ -169,7 +170,7 @@ pub use emitter::{ run_gemma3n_sdpa_vector_context_diagnostic_probe, }; #[cfg(feature = "diagnostics")] -pub use iree::PreparedPrefillDiagnostics; +pub use iree::{DeepStackPrefillDiagnostics, PreparedPrefillDiagnostics}; #[cfg(feature = "iree")] pub use phi4_audio::{ PHI4MM_AUDIO_CHECKPOINT_REVISION, PHI4MM_AUDIO_FRAME_BUCKETS, Phi4AudioOutput, diff --git a/src/models/qwen3_vl.rs b/src/models/qwen3_vl.rs index b4573be0..b8261967 100644 --- a/src/models/qwen3_vl.rs +++ b/src/models/qwen3_vl.rs @@ -602,6 +602,16 @@ pub struct Qwen3VLModel { deepstack_visual_embeds: RefCell>>>, } +/// Diagnostics-only result from the eager Qwen3-VL language path. +/// +/// The snapshots are taken immediately after each configured DeepStack +/// residual injection. Normal generation never constructs this value. +#[cfg(feature = "xla-diagnostics")] +pub struct Qwen3VLDeepStackDiagnostics { + pub logits: UniquePtr, + pub post_injection_hidden_states: Vec>, +} + impl Qwen3VLModel { pub fn from_weights(weights: &WeightMap, config: &Qwen3VLConfig) -> Result { let gs = config.group_size(); @@ -836,6 +846,33 @@ impl Qwen3VLModel { self.forward_for_sequence(input_ids, input_embeddings, caches, mask, None) } + /// Run the eager production layer loop and expose the states immediately + /// after each DeepStack injection. + /// + /// This entry is compiled only for the actual-checkpoint XLA oracle. It + /// shares every operation with [`Self::forward_for_sequence`] and merely + /// copies the selected post-hook tensors before final norm and LM head. + #[cfg(feature = "xla-diagnostics")] + pub fn forward_deepstack_diagnostics( + &self, + input_ids: &MlxArray, + input_embeddings: &MlxArray, + caches: &mut [KVCache], + mask: Option<&MlxArray>, + ) -> Qwen3VLDeepStackDiagnostics { + let (logits, post_injection_hidden_states) = self.forward_for_sequence_inner::( + input_ids, + Some(input_embeddings), + caches, + mask, + None, + ); + Qwen3VLDeepStackDiagnostics { + logits, + post_injection_hidden_states, + } + } + /// Internal forward path that takes an optional `SequenceId` so the /// cached MRoPE state is resolved per row. pub(crate) fn forward_for_sequence( @@ -846,6 +883,18 @@ impl Qwen3VLModel { mask: Option<&MlxArray>, seq_id: Option, ) -> UniquePtr { + self.forward_for_sequence_inner::(input_ids, input_embeddings, caches, mask, seq_id) + .0 + } + + fn forward_for_sequence_inner( + &self, + input_ids: &MlxArray, + input_embeddings: Option<&MlxArray>, + caches: &mut [KVCache], + mask: Option<&MlxArray>, + seq_id: Option, + ) -> (UniquePtr, Vec>) { let cache_offset = caches[0].offset; if input_embeddings.is_none() && cache_offset == 0 { self.clear_mrope_state(); @@ -853,7 +902,7 @@ impl Qwen3VLModel { } if input_embeddings.is_none() && self.can_use_text_only_fast_path(seq_id) { - return self.forward_text_only(input_ids, caches, mask); + return (self.forward_text_only(input_ids, caches, mask), Vec::new()); } let mut h = if let Some(embeds) = input_embeddings { @@ -931,6 +980,7 @@ impl Qwen3VLModel { // Get deepstack state references let ds_masks = self.visual_pos_masks.borrow(); let ds_embeds = self.deepstack_visual_embeds.borrow(); + let mut post_injection_hidden_states = Vec::new(); for (layer_idx, layer) in self.layers.iter().enumerate() { h = layer.forward(&h, &mut caches[layer_idx], mask, &position_ids); @@ -941,11 +991,16 @@ impl Qwen3VLModel { && cache_offset == 0 { h = Self::deepstack_process(&h, masks, &embeds[layer_idx]); + if CAPTURE_DEEPSTACK { + post_injection_hidden_states.push(mlxcel_core::copy( + h.as_ref().expect("DeepStack hidden state"), + )); + } } } h = self.norm.forward(&h); - self.lm_head.forward(&h) + (self.lm_head.forward(&h), post_injection_hidden_states) } fn forward_text_only( diff --git a/src/multimodal/host_preprocessor.rs b/src/multimodal/host_preprocessor.rs index ddd8c23f..c0f32d17 100644 --- a/src/multimodal/host_preprocessor.rs +++ b/src/multimodal/host_preprocessor.rs @@ -353,6 +353,15 @@ pub struct Qwen3VlIreeHostPreprocessor { device: String, } +/// Exact shared processor payload and IREE vision result used by the +/// actual-checkpoint Qwen3-VL oracle. +#[cfg(feature = "xla-diagnostics")] +pub struct Qwen3VlIreeVisionCapture { + pub patch_values: Vec, + pub grids: Vec<(i32, i32, i32)>, + pub projection: mlxcel_xla::Qwen3VlVisionProjection, +} + #[cfg(feature = "xla-iree")] impl Qwen2VlIreeHostPreprocessor { pub fn load(model_path: &Path, device: &str) -> Result { @@ -525,6 +534,37 @@ impl Qwen3VlIreeHostPreprocessor { crate::loading::load_qwen3_vl_iree_host_preprocessor(model_path, device) } + /// Run the qualified host processor and resident IREE vision graph while + /// retaining the exact patch payload for an eager MLX comparison. + #[cfg(feature = "xla-diagnostics")] + pub fn capture_vision( + &self, + images: &[DynamicImage], + ) -> Result { + if images.is_empty() { + return Err(HostPreprocessorError::InvalidConfig( + "Qwen3-VL vision diagnostics require at least one image".to_string(), + )); + } + let (patch_values, grids) = self.processor.preprocess_values_with_grid(images); + let projection = self + .projector + .try_borrow_mut() + .map_err(|_| { + HostPreprocessorError::Iree( + "concurrent/re-entrant Qwen3-VL IREE vision invocation is unsupported" + .to_string(), + ) + })? + .project(&patch_values, &grids) + .map_err(HostPreprocessorError::Iree)?; + Ok(Qwen3VlIreeVisionCapture { + patch_values, + grids, + projection, + }) + } + #[allow(clippy::too_many_arguments)] pub(crate) fn from_parts( processor: crate::vision::processors::qwen2_vl::Qwen2VLProcessor, From e38f8e66f3a3a853bcf3c7f036bffc57a7c37813 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 00:16:42 +0900 Subject: [PATCH 04/10] test(xla): localize Qwen3-VL vision divergence --- examples/xla_qwen3_vl_deepstack_check.rs | 112 ++++++++++++++++- src/lib/mlxcel-xla/src/emitter/mod.rs | 2 + src/lib/mlxcel-xla/src/emitter/qwen3_vl.rs | 70 +++++++++-- src/lib/mlxcel-xla/src/lib.rs | 2 + src/lib/mlxcel-xla/src/qwen3_vl_runtime.rs | 135 +++++++++++++++++++-- src/vision/encoders/qwen3_vl.rs | 129 +++++++++++++++++++- 6 files changed, 422 insertions(+), 28 deletions(-) diff --git a/examples/xla_qwen3_vl_deepstack_check.rs b/examples/xla_qwen3_vl_deepstack_check.rs index 36a24496..39e20f67 100644 --- a/examples/xla_qwen3_vl_deepstack_check.rs +++ b/examples/xla_qwen3_vl_deepstack_check.rs @@ -15,10 +15,12 @@ //! Actual-checkpoint eager MLX vs IREE Qwen3-VL DeepStack oracle. //! //! The qualified production host processor creates one patch payload that is -//! consumed by both vision implementations. The gate compares the main merger, -//! every ordered DeepStack merger, post-injection language hidden states, and -//! final-position logits. It also mutates the actual captured branches to prove -//! that dropping or zeroing one branch is detected. +//! consumed by both vision implementations. The gate locates the first +//! divergence across patch projection, interpolated position embeddings, every +//! encoder block, the main merger, every ordered DeepStack merger, +//! post-injection language hidden states, and final-position logits. It also +//! mutates the actual captured branches to prove that dropping or zeroing one +//! branch is detected. use std::fs; use std::path::{Path, PathBuf}; @@ -231,15 +233,82 @@ fn main() { ); let eager_vision = model .vision_encoder - .forward_with_grid(&pixels, &vision_capture.grids); - let eager_main = mlx_f32(&eager_vision.hidden_states); + .forward_with_grid_diagnostics(&pixels, &vision_capture.grids); + let eager_main = mlx_f32(&eager_vision.output.hidden_states); let eager_branches = eager_vision + .output .deepstack_features .iter() .map(|feature| mlx_f32(feature)) .collect::>(); let mut reports = Vec::new(); + let iree_diagnostics = &vision_capture.projection.diagnostics; + assert_eq!( + iree_diagnostics.block_hidden_states.len(), + eager_vision.block_hidden_states.len(), + "eager and IREE vision block capture counts differ" + ); + compare_stage( + &mut reports, + "vision.patch_embedding", + &iree_diagnostics.patch_embeddings, + &mlx_f32(&eager_vision.patch_embeddings), + tolerance, + ) + .unwrap_or_else(|error| panic!("{error}")); + compare_stage( + &mut reports, + "vision.position_embedding", + &iree_diagnostics.position_embeddings, + &mlx_f32(&eager_vision.position_embeddings), + tolerance, + ) + .unwrap_or_else(|error| panic!("{error}")); + compare_stage( + &mut reports, + "vision.positioned_embedding", + &iree_diagnostics.positioned_embeddings, + &mlx_f32(&eager_vision.positioned_embeddings), + tolerance, + ) + .unwrap_or_else(|error| panic!("{error}")); + for (layer, (actual, expected)) in iree_diagnostics + .block_hidden_states + .iter() + .zip(&eager_vision.block_hidden_states) + .enumerate() + { + compare_stage( + &mut reports, + &format!("vision.block_{layer}"), + actual, + &mlx_f32(expected), + tolerance, + ) + .unwrap_or_else(|error| panic!("{error}")); + } + let merger_stage_names = ["normalized_shuffled", "fc1", "activation"]; + assert_eq!( + iree_diagnostics.main_merger_states.len(), + eager_vision.main_merger_states.len(), + "eager and IREE main merger capture counts differ" + ); + for ((actual, expected), name) in iree_diagnostics + .main_merger_states + .iter() + .zip(&eager_vision.main_merger_states) + .zip(merger_stage_names) + { + compare_stage( + &mut reports, + &format!("vision.main_merger.{name}"), + actual, + &mlx_f32(expected), + tolerance, + ) + .unwrap_or_else(|error| panic!("{error}")); + } compare_stage( &mut reports, "vision.main_merger", @@ -248,6 +317,37 @@ fn main() { tolerance, ) .unwrap_or_else(|error| panic!("{error}")); + assert_eq!( + iree_diagnostics.deepstack_merger_states.len(), + eager_vision.deepstack_merger_states.len(), + "eager and IREE DeepStack merger capture counts differ" + ); + for (branch, (actual_stages, expected_stages)) in iree_diagnostics + .deepstack_merger_states + .iter() + .zip(&eager_vision.deepstack_merger_states) + .enumerate() + { + assert_eq!( + actual_stages.len(), + expected_stages.len(), + "eager and IREE DeepStack merger stage counts differ for branch {branch}" + ); + for ((actual, expected), name) in actual_stages + .iter() + .zip(expected_stages) + .zip(merger_stage_names) + { + compare_stage( + &mut reports, + &format!("vision.deepstack_merger_{branch}.{name}"), + actual, + &mlx_f32(expected), + tolerance, + ) + .unwrap_or_else(|error| panic!("{error}")); + } + } compare_vision_branches( &mut reports, &vision_capture.projection.deepstack_values, diff --git a/src/lib/mlxcel-xla/src/emitter/mod.rs b/src/lib/mlxcel-xla/src/emitter/mod.rs index 7474b3a9..3b4e4c02 100644 --- a/src/lib/mlxcel-xla/src/emitter/mod.rs +++ b/src/lib/mlxcel-xla/src/emitter/mod.rs @@ -131,6 +131,8 @@ pub(crate) use qwen2_vl::{ }; #[allow(unused_imports)] pub(crate) use qwen3_vl::emit_qwen3_vl; +#[cfg(feature = "diagnostics")] +pub(crate) use qwen3_vl::emit_qwen3_vl_diagnostics; #[allow(unused_imports)] pub(crate) use qwen3_vl::{ QWEN3_VL_PATCH_BUCKETS, Qwen3VlConfig, Qwen3VlGridPlan, Qwen3VlHostInputs, Qwen3VlWeightSpec, diff --git a/src/lib/mlxcel-xla/src/emitter/qwen3_vl.rs b/src/lib/mlxcel-xla/src/emitter/qwen3_vl.rs index 386f633b..3a451700 100644 --- a/src/lib/mlxcel-xla/src/emitter/qwen3_vl.rs +++ b/src/lib/mlxcel-xla/src/emitter/qwen3_vl.rs @@ -876,13 +876,14 @@ fn interpolated_position_embeddings( builder.reduce_add(&weighted, 0, &zero) } -fn patch_merger( +fn patch_merger_impl( builder: &mut Builder, hidden: &Val, args: &mut Args, config: &Qwen3VlConfig, postshuffle_norm: bool, -) -> Val { + diagnostics: bool, +) -> (Val, Option<[Val; 3]>) { let merge_width = config.hidden * config.spatial_merge_size * config.spatial_merge_size; let merged_tokens = hidden.ty.shape[0] / (config.spatial_merge_size * config.spatial_merge_size); @@ -905,16 +906,18 @@ fn patch_merger( ); builder.reshape(&normalized, vec![merged_tokens, merge_width]) }; - let merged = linear_2d(builder, &merged, &args.take(), &args.take()); - let merged = exact_gelu(builder, &merged); - linear_2d(builder, &merged, &args.take(), &args.take()) + let fc1 = linear_2d(builder, &merged, &args.take(), &args.take()); + let activated = exact_gelu(builder, &fc1); + let projected = linear_2d(builder, &activated, &args.take(), &args.take()); + let diagnostic_values = diagnostics.then(|| [merged.clone(), fc1.clone(), activated.clone()]); + (projected, diagnostic_values) } /// Emit one Qwen3-VL bucket. `patches` are already in the processor's /// post-smart-resize, spatial-merge-grouped order. `vision_rope` and /// `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_qwen3_vl(config: &Qwen3VlConfig, patch_bucket: usize) -> String { +fn emit_qwen3_vl_impl(config: &Qwen3VlConfig, patch_bucket: usize, diagnostics: bool) -> String { assert!( QWEN3_VL_PATCH_BUCKETS.contains(&patch_bucket), "unqualified Qwen3-VL patch bucket" @@ -942,14 +945,18 @@ pub(crate) fn emit_qwen3_vl(config: &Qwen3VlConfig, patch_bucket: usize) -> Stri let position_weights = args.push_input(Ty::f32(vec![4, patch_bucket]), "position.weights"); let mut hidden = builder.linear_seq(&patches, &patch_weight); hidden = bias_2d(&mut builder, &hidden, &patch_bias); + let patch_embeddings = diagnostics.then(|| hidden.clone()); let position_embeddings = interpolated_position_embeddings( &mut builder, &position_table, &position_indices, &position_weights, ); + let diagnostic_position_embeddings = diagnostics.then(|| position_embeddings.clone()); hidden = builder.add(&hidden, &position_embeddings); + let positioned_embeddings = diagnostics.then(|| hidden.clone()); let mut deepstack_hidden = Vec::with_capacity(config.deepstack_visual_indexes.len()); + let mut block_hidden = diagnostics.then(|| Vec::with_capacity(config.depth)); for layer in 0..config.depth { hidden = encoder_layer( &mut builder, @@ -959,14 +966,41 @@ pub(crate) fn emit_qwen3_vl(config: &Qwen3VlConfig, patch_bucket: usize) -> Stri &freqs, &attention_bias, ); + if let Some(outputs) = &mut block_hidden { + outputs.push(hidden.clone()); + } if config.deepstack_visual_indexes.contains(&layer) { deepstack_hidden.push(hidden.clone()); } } - let projected = patch_merger(&mut builder, &hidden, &mut args, config, false); + let (projected, main_merger_diagnostics) = + patch_merger_impl(&mut builder, &hidden, &mut args, config, false, diagnostics); let mut outputs = vec![projected]; + let mut deepstack_merger_diagnostics = + diagnostics.then(|| Vec::with_capacity(deepstack_hidden.len())); for branch in deepstack_hidden { - outputs.push(patch_merger(&mut builder, &branch, &mut args, config, true)); + let (projected, branch_diagnostics) = + patch_merger_impl(&mut builder, &branch, &mut args, config, true, diagnostics); + outputs.push(projected); + if let Some(all_diagnostics) = &mut deepstack_merger_diagnostics { + all_diagnostics.push(branch_diagnostics.expect("diagnostics includes merger stages")); + } + } + if let Some(block_hidden) = block_hidden { + outputs.push(patch_embeddings.expect("diagnostics includes patch embeddings")); + outputs.push( + diagnostic_position_embeddings + .expect("diagnostics includes interpolated position embeddings"), + ); + outputs.push(positioned_embeddings.expect("diagnostics includes positioned embeddings")); + outputs.extend(block_hidden); + outputs.extend(main_merger_diagnostics.expect("diagnostics includes main merger stages")); + outputs.extend( + deepstack_merger_diagnostics + .expect("diagnostics includes DeepStack merger stages") + .into_iter() + .flatten(), + ); } assert_eq!(args.cursor, specs.len(), "Qwen3-VL weight schema drifted"); let result_types = outputs @@ -986,6 +1020,15 @@ pub(crate) fn emit_qwen3_vl(config: &Qwen3VlConfig, patch_bucket: usize) -> Stri ) } +pub(crate) fn emit_qwen3_vl(config: &Qwen3VlConfig, patch_bucket: usize) -> String { + emit_qwen3_vl_impl(config, patch_bucket, false) +} + +#[cfg(feature = "diagnostics")] +pub(crate) fn emit_qwen3_vl_diagnostics(config: &Qwen3VlConfig, patch_bucket: usize) -> String { + emit_qwen3_vl_impl(config, patch_bucket, true) +} + #[cfg(test)] mod tests { use super::*; @@ -1109,6 +1152,17 @@ mod tests { ); } + #[test] + fn diagnostic_bucket_appends_ordered_pre_merger_stages() { + let config = config(); + let mlir = emit_qwen3_vl_impl(&config, 16, true); + let mut output_types = vec!["tensor<4x12xf32>"; 3]; + output_types.extend(vec!["tensor<16x8xf32>"; 5]); + output_types.extend(vec!["tensor<4x32xf32>"; 9]); + assert!(mlir.contains(&format!("-> ({})", output_types.join(", ")))); + assert_eq!(mlir.matches("chlo.erf").count(), 3); + } + #[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 434058c0..72f40767 100644 --- a/src/lib/mlxcel-xla/src/lib.rs +++ b/src/lib/mlxcel-xla/src/lib.rs @@ -183,6 +183,8 @@ pub use phi4_audio::{Phi4AudioCheckpoint, Phi4AudioDiagnosticRuntime, Phi4AudioD pub use qwen2_vl_runtime::{ IreeQwen2VlProjector, Qwen2VlVisionExecutionMetrics, Qwen2VlVisionProjection, }; +#[cfg(feature = "diagnostics")] +pub use qwen3_vl_runtime::Qwen3VlVisionDiagnostics; #[cfg(feature = "iree")] pub use qwen3_vl_runtime::{ IreeQwen3VlProjector, Qwen3VlVisionExecutionMetrics, Qwen3VlVisionProjection, diff --git a/src/lib/mlxcel-xla/src/qwen3_vl_runtime.rs b/src/lib/mlxcel-xla/src/qwen3_vl_runtime.rs index 5ecf245f..8e54b71e 100644 --- a/src/lib/mlxcel-xla/src/qwen3_vl_runtime.rs +++ b/src/lib/mlxcel-xla/src/qwen3_vl_runtime.rs @@ -34,9 +34,11 @@ use crate::aux::{ IreeAuxiliaryModule, }; use crate::aux_manifest::{AuxiliaryArtifactContract, ensure_qualified_auxiliary_artifact}; -use crate::emitter::{ - Qwen3VlConfig, Qwen3VlHostInputs, emit_qwen3_vl, prepare_qwen3_vl_host_inputs, -}; +#[cfg(not(feature = "diagnostics"))] +use crate::emitter::emit_qwen3_vl; +#[cfg(feature = "diagnostics")] +use crate::emitter::emit_qwen3_vl_diagnostics; +use crate::emitter::{Qwen3VlConfig, Qwen3VlHostInputs, prepare_qwen3_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}; @@ -54,12 +56,32 @@ pub struct Qwen3VlVisionExecutionMetrics { pub struct Qwen3VlVisionProjection { pub values: Vec, pub deepstack_values: Vec>, + #[cfg(feature = "diagnostics")] + pub diagnostics: Qwen3VlVisionDiagnostics, pub shape: [usize; 2], pub merged_tokens_per_image: Vec, pub packed_cu_seqlens: Vec, pub metrics: Qwen3VlVisionExecutionMetrics, } +#[cfg(feature = "diagnostics")] +#[derive(Debug, Clone, PartialEq)] +pub struct Qwen3VlVisionDiagnostics { + pub patch_embeddings: Vec, + pub position_embeddings: Vec, + pub positioned_embeddings: Vec, + /// Encoder block outputs in layer order. The final entry is the main + /// merger input; configured DeepStack inputs are the corresponding + /// `deepstack_visual_indexes` entries. + pub block_hidden_states: Vec>, + pub shape: [usize; 2], + /// Normalized/shuffled input, first linear output, and GELU output. + pub main_merger_states: Vec>, + /// The same ordered merger stages for each DeepStack branch. + pub deepstack_merger_states: Vec>>, + pub merger_shape: [usize; 2], +} + struct ActiveModule { patch_bucket: usize, module: IreeAuxiliaryModule, @@ -596,7 +618,10 @@ fn compile_and_load( processor_identity: &str, patch_bucket: usize, ) -> Result { + #[cfg(not(feature = "diagnostics"))] let mlir = emit_qwen3_vl(config, patch_bucket); + #[cfg(feature = "diagnostics")] + let mlir = emit_qwen3_vl_diagnostics(config, patch_bucket); let compiler = iree_compile_bin()?; if !compiler.is_file() { return Err(format!("iree-compile not found at {}", compiler.display())); @@ -700,17 +725,41 @@ impl IreeQwen3VlProjector { plan.patch_bucket / (self.config.spatial_merge_size * self.config.spatial_merge_size), self.config.text_hidden, ]; - let output_bytes = full_output_shape.iter().product::() * std::mem::size_of::(); - let mut output_buffers = - vec![vec![0u8; output_bytes]; 1 + self.config.deepstack_visual_indexes.len()]; + let standard_output_count = 1 + self.config.deepstack_visual_indexes.len(); + let mut output_shapes = vec![full_output_shape.to_vec(); standard_output_count]; + #[cfg(feature = "diagnostics")] + { + output_shapes.extend( + (0..3 + self.config.depth).map(|_| vec![plan.patch_bucket, self.config.hidden]), + ); + let merge_width = self.config.hidden + * self.config.spatial_merge_size + * self.config.spatial_merge_size; + output_shapes.extend((0..3 * standard_output_count).map(|_| { + vec![ + plan.patch_bucket + / (self.config.spatial_merge_size * self.config.spatial_merge_size), + merge_width, + ] + })); + } + let projected_transfer_bytes = output_shapes + .iter() + .map(|shape| shape.iter().product::() * std::mem::size_of::()) + .sum(); + let mut output_buffers = output_shapes + .iter() + .map(|shape| vec![0u8; shape.iter().product::() * std::mem::size_of::()]) + .collect::>(); let started = Instant::now(); { let mut outputs = output_buffers .iter_mut() - .map(|output| AuxiliaryOutput { + .zip(&output_shapes) + .map(|(output, shape)| AuxiliaryOutput { bytes: output.as_mut_slice(), dtype: AuxiliaryTensorDType::Float32, - shape: &full_output_shape, + shape, }) .collect::>(); self.ensure_bucket(plan.patch_bucket)?.invoke( @@ -762,7 +811,7 @@ impl IreeQwen3VlProjector { .ok_or_else(|| "Qwen3-VL IREE output is shorter than the actual grid".to_string())? .to_vec(); let deepstack_values = decoded_outputs - .into_iter() + .drain(..self.config.deepstack_visual_indexes.len()) .map(|output| { output .get(..actual_values) @@ -772,9 +821,74 @@ impl IreeQwen3VlProjector { .map(<[f32]>::to_vec) }) .collect::, _>>()?; + #[cfg(feature = "diagnostics")] + let diagnostics = { + let actual_hidden_values = plan + .actual_patches + .checked_mul(self.config.hidden) + .ok_or_else(|| "Qwen3-VL diagnostic hidden size overflowed".to_string())?; + let mut trim = |label: &str| -> Result, String> { + decoded_outputs + .remove(0) + .get(..actual_hidden_values) + .ok_or_else(|| { + format!("Qwen3-VL IREE {label} output is shorter than the actual grid") + }) + .map(<[f32]>::to_vec) + }; + let patch_embeddings = trim("patch embedding")?; + let position_embeddings = trim("position embedding")?; + let positioned_embeddings = trim("positioned embedding")?; + let block_hidden_states = (0..self.config.depth) + .map(|layer| trim(&format!("block {layer}"))) + .collect::, _>>()?; + let actual_merged_tokens = plan.merged_tokens_per_image.iter().sum::(); + let merge_width = self.config.hidden + * self.config.spatial_merge_size + * self.config.spatial_merge_size; + let actual_merger_values = actual_merged_tokens + .checked_mul(merge_width) + .ok_or_else(|| "Qwen3-VL diagnostic merger size overflowed".to_string())?; + let mut trim_merger = |label: &str| -> Result, String> { + decoded_outputs + .remove(0) + .get(..actual_merger_values) + .ok_or_else(|| { + format!("Qwen3-VL IREE {label} output is shorter than the actual grid") + }) + .map(<[f32]>::to_vec) + }; + let main_merger_states = (0..3) + .map(|stage| trim_merger(&format!("main merger stage {stage}"))) + .collect::, _>>()?; + let deepstack_merger_states = (0..self.config.deepstack_visual_indexes.len()) + .map(|branch| { + (0..3) + .map(|stage| { + trim_merger(&format!("DeepStack merger branch {branch} stage {stage}")) + }) + .collect::, _>>() + }) + .collect::, _>>()?; + debug_assert!(decoded_outputs.is_empty()); + Qwen3VlVisionDiagnostics { + patch_embeddings, + position_embeddings, + positioned_embeddings, + block_hidden_states, + shape: [plan.actual_patches, self.config.hidden], + main_merger_states, + deepstack_merger_states, + merger_shape: [actual_merged_tokens, merge_width], + } + }; + #[cfg(not(feature = "diagnostics"))] + debug_assert!(decoded_outputs.is_empty()); Ok(Qwen3VlVisionProjection { values, deepstack_values, + #[cfg(feature = "diagnostics")] + diagnostics, shape: [actual_tokens, self.config.text_hidden], merged_tokens_per_image: plan.merged_tokens_per_image, packed_cu_seqlens: plan.packed_cu_seqlens, @@ -784,8 +898,7 @@ impl IreeQwen3VlProjector { + std::mem::size_of_val(packed_attention_bias.as_slice()) + std::mem::size_of_val(position_indices.as_slice()) + std::mem::size_of_val(position_weights.as_slice()), - projected_transfer_bytes: output_bytes - * (1 + self.config.deepstack_visual_indexes.len()), + projected_transfer_bytes, elapsed_seconds, }, }) diff --git a/src/vision/encoders/qwen3_vl.rs b/src/vision/encoders/qwen3_vl.rs index bc93c6e7..42310d83 100644 --- a/src/vision/encoders/qwen3_vl.rs +++ b/src/vision/encoders/qwen3_vl.rs @@ -630,7 +630,10 @@ impl PatchMerger { }) } - fn forward(&self, x: &MlxArray) -> UniquePtr { + fn forward_with_observer(&self, x: &MlxArray, mut observe: F) -> UniquePtr + where + F: FnMut(PatchMergerStage, &MlxArray), + { let h = if self.use_postshuffle_norm { // DeepStack: reshape first, then norm let reshaped = mlxcel_core::reshape(x, &[-1, self.hidden_size as i32]); @@ -640,12 +643,32 @@ impl PatchMerger { let normed = self.norm.forward(x); mlxcel_core::reshape(&normed, &[-1, self.hidden_size as i32]) }; + observe(PatchMergerStage::NormalizedShuffled, &h); let h = self.linear_fc1.forward(&h); + observe(PatchMergerStage::Fc1, &h); let h = mlxcel_core::gelu(&h); + observe(PatchMergerStage::Activation, &h); self.linear_fc2.forward(&h) } } +#[derive(Clone, Copy)] +enum PatchMergerStage { + NormalizedShuffled, + Fc1, + Activation, +} + +impl PatchMergerStage { + fn index(self) -> usize { + match self { + Self::NormalizedShuffled => 0, + Self::Fc1 => 1, + Self::Activation => 2, + } + } +} + // Qwen3-VL Vision Encoder Output (includes deepstack features). /// Output from Qwen3-VL vision encoder including DeepStack features pub struct Qwen3VLVisionEncoderOutput { @@ -653,6 +676,32 @@ pub struct Qwen3VLVisionEncoderOutput { pub deepstack_features: Vec>, } +#[cfg(feature = "xla-diagnostics")] +pub struct Qwen3VLVisionEncoderDiagnosticOutput { + pub output: Qwen3VLVisionEncoderOutput, + pub patch_embeddings: UniquePtr, + pub position_embeddings: UniquePtr, + pub positioned_embeddings: UniquePtr, + /// Encoder block outputs in layer order. The final entry is the main + /// merger input; configured DeepStack inputs are the corresponding + /// `deepstack_visual_indexes` entries. + pub block_hidden_states: Vec>, + /// Normalized/shuffled input, first linear output, and GELU output. + pub main_merger_states: Vec>, + /// The same ordered merger stages for each DeepStack branch. + pub deepstack_merger_states: Vec>>, +} + +#[derive(Clone, Copy)] +enum VisionStage { + PatchEmbedding, + PositionEmbedding, + PositionedEmbedding, + Block(usize), + MainMerger(PatchMergerStage), + DeepStackMerger(usize, PatchMergerStage), +} + // Qwen3-VL Vision Encoder. /// Qwen3-VL Vision Model with learned position embeddings and DeepStack /// @@ -819,14 +868,80 @@ impl Qwen3VLVisionEncoder { hidden_states: &MlxArray, grid_thw: &[(i32, i32, i32)], ) -> Qwen3VLVisionEncoderOutput { + self.forward_with_grid_observer(hidden_states, grid_thw, |_, _| {}) + } + + #[cfg(feature = "xla-diagnostics")] + pub fn forward_with_grid_diagnostics( + &self, + hidden_states: &MlxArray, + grid_thw: &[(i32, i32, i32)], + ) -> Qwen3VLVisionEncoderDiagnosticOutput { + let mut patch_embeddings = None; + let mut position_embeddings = None; + let mut positioned_embeddings = None; + let mut block_hidden_states = Vec::with_capacity(self.blocks.len()); + let mut main_merger_states = Vec::with_capacity(3); + let mut deepstack_merger_states = (0..self.deepstack_merger_list.len()) + .map(|_| Vec::with_capacity(3)) + .collect::>(); + let output = + self.forward_with_grid_observer(hidden_states, grid_thw, |stage, hidden| match stage { + VisionStage::PatchEmbedding => { + patch_embeddings = Some(mlxcel_core::copy(hidden)); + } + VisionStage::PositionEmbedding => { + position_embeddings = Some(mlxcel_core::copy(hidden)); + } + VisionStage::PositionedEmbedding => { + positioned_embeddings = Some(mlxcel_core::copy(hidden)); + } + VisionStage::Block(layer) => { + debug_assert_eq!(layer, block_hidden_states.len()); + block_hidden_states.push(mlxcel_core::copy(hidden)); + } + VisionStage::MainMerger(stage) => { + debug_assert_eq!(stage.index(), main_merger_states.len()); + main_merger_states.push(mlxcel_core::copy(hidden)); + } + VisionStage::DeepStackMerger(branch, stage) => { + debug_assert_eq!(stage.index(), deepstack_merger_states[branch].len()); + deepstack_merger_states[branch].push(mlxcel_core::copy(hidden)); + } + }); + Qwen3VLVisionEncoderDiagnosticOutput { + output, + patch_embeddings: patch_embeddings.expect("observer captures patch embeddings"), + position_embeddings: position_embeddings + .expect("observer captures interpolated position embeddings"), + positioned_embeddings: positioned_embeddings + .expect("observer captures positioned embeddings"), + block_hidden_states, + main_merger_states, + deepstack_merger_states, + } + } + + fn forward_with_grid_observer( + &self, + hidden_states: &MlxArray, + grid_thw: &[(i32, i32, i32)], + mut observe: F, + ) -> Qwen3VLVisionEncoderOutput + where + F: FnMut(VisionStage, &MlxArray), + { // 1. Patch embedding let mut h = self.patch_embed.forward(hidden_states); + observe(VisionStage::PatchEmbedding, &h); // 2. Add learned position embeddings with bilinear interpolation let pos_embeds = self .pos_embed .fast_pos_embed_interpolate(grid_thw, self.spatial_merge_size); + observe(VisionStage::PositionEmbedding, &pos_embeds); h = mlxcel_core::add(&h, &pos_embeds); + observe(VisionStage::PositionedEmbedding, &h); // 3. Compute rotary position embeddings let rotary_pos_emb = self.rot_pos_emb(grid_thw); @@ -846,19 +961,27 @@ impl Qwen3VLVisionEncoder { for (layer_num, block) in self.blocks.iter().enumerate() { h = block.forward(&h, &cu_seqlens, &rotary_pos_emb); + observe(VisionStage::Block(layer_num), &h); if let Some(ds_idx) = self .deepstack_visual_indexes .iter() .position(|&idx| idx == layer_num) { - let ds_feature = self.deepstack_merger_list[ds_idx].forward(&h); + let ds_feature = self.deepstack_merger_list[ds_idx].forward_with_observer( + &h, + |stage, hidden| { + observe(VisionStage::DeepStackMerger(ds_idx, stage), hidden); + }, + ); deepstack_features.push(ds_feature); } } // 6. Apply main merger - h = self.merger.forward(&h); + h = self.merger.forward_with_observer(&h, |stage, hidden| { + observe(VisionStage::MainMerger(stage), hidden); + }); Qwen3VLVisionEncoderOutput { hidden_states: h, From 2414576e37d9afb4d2e3625dfa531f65582cb82d Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 01:04:02 +0900 Subject: [PATCH 05/10] fix(xla): align Qwen3-VL block GELU Use the checkpoint's gelu_pytorch_tanh contract in the eager vision blocks while preserving exact GELU in the main and DeepStack patch mergers. Add eager numerical and emitted-graph regression coverage for the activation split. --- src/lib/mlxcel-xla/src/emitter/qwen3_vl.rs | 4 +++- src/vision/encoders/qwen3_vl.rs | 28 +++++++++++++++++++++- src/vision/encoders/qwen3_vl_tests.rs | 18 +++++++++++++- 3 files changed, 47 insertions(+), 3 deletions(-) diff --git a/src/lib/mlxcel-xla/src/emitter/qwen3_vl.rs b/src/lib/mlxcel-xla/src/emitter/qwen3_vl.rs index 3a451700..989e43b5 100644 --- a/src/lib/mlxcel-xla/src/emitter/qwen3_vl.rs +++ b/src/lib/mlxcel-xla/src/emitter/qwen3_vl.rs @@ -1052,7 +1052,8 @@ mod tests { "spatial_merge_size": 2, "temporal_patch_size": 2, "num_position_embeddings": 16, - "deepstack_visual_indexes": [0, 1] + "deepstack_visual_indexes": [0, 1], + "hidden_act": "gelu_pytorch_tanh" } }"#, ) @@ -1141,6 +1142,7 @@ mod tests { assert!(mlir.contains("loc(\"position.indices\")")); assert!(mlir.contains("tensor<4x16x1xi32>")); assert!(mlir.contains("-> (tensor<4x12xf32>, tensor<4x12xf32>, tensor<4x12xf32>)")); + assert_eq!(mlir.matches("stablehlo.tanh").count(), config.depth); assert_eq!(mlir.matches("chlo.erf").count(), 3); let two_image_mlir = emit_qwen3_vl(&config, 512); diff --git a/src/vision/encoders/qwen3_vl.rs b/src/vision/encoders/qwen3_vl.rs index 42310d83..6a358608 100644 --- a/src/vision/encoders/qwen3_vl.rs +++ b/src/vision/encoders/qwen3_vl.rs @@ -520,6 +520,32 @@ struct VisionMLP { linear_fc2: UnifiedLinear, } +/// Hugging Face's `gelu_pytorch_tanh`, which Qwen3-VL vision checkpoints use +/// inside transformer blocks. The patch mergers intentionally keep exact GELU. +/// +/// Evaluate in F32 to match the qualified IREE graph and cast back so the eager +/// encoder preserves its caller-visible dtype. +fn gelu_pytorch_tanh(x: &MlxArray) -> UniquePtr { + let output_dtype = mlxcel_core::array_dtype(x); + let x = mlxcel_core::astype(x, mlxcel_core::dtype::FLOAT32); + let half = mlxcel_core::full_f32(&[1], 0.5, mlxcel_core::dtype::FLOAT32); + let one = mlxcel_core::full_f32(&[1], 1.0, mlxcel_core::dtype::FLOAT32); + let sqrt_two_over_pi = mlxcel_core::full_f32(&[1], 0.797_884_6, mlxcel_core::dtype::FLOAT32); + let cubic_coefficient = mlxcel_core::full_f32(&[1], 0.044_715, mlxcel_core::dtype::FLOAT32); + + let squared = mlxcel_core::multiply(&x, &x); + let cubed = mlxcel_core::multiply(&squared, &x); + let cubic = mlxcel_core::multiply(&cubic_coefficient, &cubed); + let inner = mlxcel_core::multiply(&sqrt_two_over_pi, &mlxcel_core::add(&x, &cubic)); + let cdf = mlxcel_core::multiply(&half, &mlxcel_core::add(&one, &mlxcel_core::tanh(&inner))); + let activated = mlxcel_core::multiply(&x, &cdf); + if output_dtype == mlxcel_core::dtype::FLOAT32 { + activated + } else { + mlxcel_core::astype(&activated, output_dtype) + } +} + impl VisionMLP { fn from_weights(weights: &WeightMap, prefix: &str, gs: i32, bits: i32) -> Result { Ok(Self { @@ -540,7 +566,7 @@ impl VisionMLP { fn forward(&self, x: &MlxArray) -> UniquePtr { let h = self.linear_fc1.forward(x); - let h = mlxcel_core::gelu_approx(&h); + let h = gelu_pytorch_tanh(&h); self.linear_fc2.forward(&h) } } diff --git a/src/vision/encoders/qwen3_vl_tests.rs b/src/vision/encoders/qwen3_vl_tests.rs index 4b2769f1..ee0da0e5 100644 --- a/src/vision/encoders/qwen3_vl_tests.rs +++ b/src/vision/encoders/qwen3_vl_tests.rs @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use super::{ensure_fused_sdpa, fused_sdpa_target_dim, sdpa_pad_width}; +use super::{ensure_fused_sdpa, fused_sdpa_target_dim, gelu_pytorch_tanh, sdpa_pad_width}; use mlxcel_core::dtype; #[test] @@ -44,3 +44,19 @@ fn ensure_fused_sdpa_restores_original_output_shape() { assert_eq!(mlxcel_core::array_shape(&output), vec![1, 2, 4, 96]); assert_eq!(mlxcel_core::array_dtype(&output), dtype::FLOAT32); } + +#[test] +fn block_gelu_matches_qwen3_vl_pytorch_tanh_contract() { + let input = mlxcel_core::from_slice_f32(&[-3.0, -1.0, 0.0, 1.0, 3.0], &[5]); + let output = gelu_pytorch_tanh(&input); + mlxcel_core::eval(&output); + + let expected = [-0.003_637_433, -0.158_808, 0.0, 0.841_192, 2.996_362_7]; + for (index, expected) in expected.into_iter().enumerate() { + let value = mlxcel_core::slice(&output, &[index as i32], &[index as i32 + 1]); + assert!( + (mlxcel_core::item_f32(&value) - expected).abs() <= 2.0e-6, + "Qwen3-VL tanh GELU mismatch at {index}" + ); + } +} From 23a8fba16d1ef720b490e1303e47f46ca2e01a89 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 01:34:12 +0900 Subject: [PATCH 06/10] test(xla): expose Qwen3-VL block 2 substages --- examples/xla_qwen3_vl_deepstack_check.rs | 304 +++++++++++++++------ src/lib/mlxcel-xla/src/emitter/mod.rs | 3 +- src/lib/mlxcel-xla/src/emitter/qwen3_vl.rs | 70 ++++- src/lib/mlxcel-xla/src/qwen3_vl_runtime.rs | 22 +- src/vision/encoders/qwen3_vl.rs | 80 +++++- src/vision/encoders/qwen3_vl_tests.rs | 22 +- 6 files changed, 404 insertions(+), 97 deletions(-) diff --git a/examples/xla_qwen3_vl_deepstack_check.rs b/examples/xla_qwen3_vl_deepstack_check.rs index 39e20f67..e255af60 100644 --- a/examples/xla_qwen3_vl_deepstack_check.rs +++ b/examples/xla_qwen3_vl_deepstack_check.rs @@ -156,6 +156,22 @@ fn compare_stage( } } +fn record_comparison(failures: &mut Vec, result: Result<(), String>) { + if let Err(error) = result { + failures.push(error); + } +} + +fn render_and_persist_report(report: &Value, report_path: Option<&Path>) -> String { + 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}"); + rendered +} + fn compare_vision_branches( reports: &mut Vec, actual: &[Vec], @@ -243,50 +259,95 @@ fn main() { .collect::>(); let mut reports = Vec::new(); + let mut failures = Vec::new(); let iree_diagnostics = &vision_capture.projection.diagnostics; assert_eq!( iree_diagnostics.block_hidden_states.len(), eager_vision.block_hidden_states.len(), "eager and IREE vision block capture counts differ" ); - compare_stage( - &mut reports, - "vision.patch_embedding", - &iree_diagnostics.patch_embeddings, - &mlx_f32(&eager_vision.patch_embeddings), - tolerance, - ) - .unwrap_or_else(|error| panic!("{error}")); - compare_stage( - &mut reports, - "vision.position_embedding", - &iree_diagnostics.position_embeddings, - &mlx_f32(&eager_vision.position_embeddings), - tolerance, - ) - .unwrap_or_else(|error| panic!("{error}")); - compare_stage( - &mut reports, - "vision.positioned_embedding", - &iree_diagnostics.positioned_embeddings, - &mlx_f32(&eager_vision.positioned_embeddings), - tolerance, - ) - .unwrap_or_else(|error| panic!("{error}")); + record_comparison( + &mut failures, + compare_stage( + &mut reports, + "vision.patch_embedding", + &iree_diagnostics.patch_embeddings, + &mlx_f32(&eager_vision.patch_embeddings), + tolerance, + ), + ); + record_comparison( + &mut failures, + compare_stage( + &mut reports, + "vision.position_embedding", + &iree_diagnostics.position_embeddings, + &mlx_f32(&eager_vision.position_embeddings), + tolerance, + ), + ); + record_comparison( + &mut failures, + compare_stage( + &mut reports, + "vision.positioned_embedding", + &iree_diagnostics.positioned_embeddings, + &mlx_f32(&eager_vision.positioned_embeddings), + tolerance, + ), + ); for (layer, (actual, expected)) in iree_diagnostics .block_hidden_states .iter() .zip(&eager_vision.block_hidden_states) .enumerate() { - compare_stage( - &mut reports, - &format!("vision.block_{layer}"), - actual, - &mlx_f32(expected), - tolerance, - ) - .unwrap_or_else(|error| panic!("{error}")); + record_comparison( + &mut failures, + compare_stage( + &mut reports, + &format!("vision.block_{layer}"), + actual, + &mlx_f32(expected), + tolerance, + ), + ); + } + let block_2_stage_names = [ + "input", + "norm1", + "attention", + "post_attention_residual", + "norm2", + "mlp", + "output", + ]; + assert_eq!( + iree_diagnostics.block_2_states.len(), + eager_vision.block_2_states.len(), + "eager and IREE block 2 substage capture counts differ" + ); + assert_eq!( + iree_diagnostics.block_2_states.len(), + block_2_stage_names.len(), + "actual Qwen3-VL diagnostics must capture every block 2 substage" + ); + for ((actual, expected), name) in iree_diagnostics + .block_2_states + .iter() + .zip(&eager_vision.block_2_states) + .zip(block_2_stage_names) + { + record_comparison( + &mut failures, + compare_stage( + &mut reports, + &format!("vision.block_2.{name}"), + actual, + &mlx_f32(expected), + tolerance, + ), + ); } let merger_stage_names = ["normalized_shuffled", "fc1", "activation"]; assert_eq!( @@ -300,23 +361,27 @@ fn main() { .zip(&eager_vision.main_merger_states) .zip(merger_stage_names) { + record_comparison( + &mut failures, + compare_stage( + &mut reports, + &format!("vision.main_merger.{name}"), + actual, + &mlx_f32(expected), + tolerance, + ), + ); + } + record_comparison( + &mut failures, compare_stage( &mut reports, - &format!("vision.main_merger.{name}"), - actual, - &mlx_f32(expected), + "vision.main_merger", + &vision_capture.projection.values, + &eager_main, tolerance, - ) - .unwrap_or_else(|error| panic!("{error}")); - } - compare_stage( - &mut reports, - "vision.main_merger", - &vision_capture.projection.values, - &eager_main, - tolerance, - ) - .unwrap_or_else(|error| panic!("{error}")); + ), + ); assert_eq!( iree_diagnostics.deepstack_merger_states.len(), eager_vision.deepstack_merger_states.len(), @@ -338,23 +403,27 @@ fn main() { .zip(expected_stages) .zip(merger_stage_names) { - compare_stage( - &mut reports, - &format!("vision.deepstack_merger_{branch}.{name}"), - actual, - &mlx_f32(expected), - tolerance, - ) - .unwrap_or_else(|error| panic!("{error}")); + record_comparison( + &mut failures, + compare_stage( + &mut reports, + &format!("vision.deepstack_merger_{branch}.{name}"), + actual, + &mlx_f32(expected), + tolerance, + ), + ); } } - compare_vision_branches( - &mut reports, - &vision_capture.projection.deepstack_values, - &eager_branches, - tolerance, - ) - .unwrap_or_else(|error| panic!("{error}")); + record_comparison( + &mut failures, + compare_vision_branches( + &mut reports, + &vision_capture.projection.deepstack_values, + &eager_branches, + tolerance, + ), + ); let mut dropped = vision_capture.projection.deepstack_values.clone(); dropped @@ -370,6 +439,26 @@ fn main() { compare_vision_branches(&mut Vec::new(), &zeroed, &eager_branches, tolerance).is_err(), "negative oracle accepted a zeroed actual-checkpoint DeepStack branch" ); + if let Some(first_failure) = failures.first().cloned() { + let report = json!({ + "schema": 1, + "model": model_path, + "image": image_path, + "device": device, + "context_capacity": context_capacity, + "completed_phase": "vision", + "grid_thw": vision_capture.grids, + "vision_shape": vision_capture.projection.shape, + "deepstack_branches": eager_branches.len(), + "negative_dropped_branch_detected": true, + "negative_zeroed_branch_detected": true, + "comparisons": reports, + "comparison_failures": failures, + "passed": false, + }); + render_and_persist_report(&report, report_path.as_deref()); + panic!("Qwen3-VL DeepStack oracle failed after rendering its report: {first_failure}"); + } let unexpanded_tokens = vec![1, model.vision_start_token_id, model.image_token_id, 2]; let prepared = production_preprocessor @@ -417,34 +506,40 @@ fn main() { let layer_start = layer * iree_language.context_capacity * hidden_size; let actual = &iree_language.post_injection_hidden_states [layer_start..layer_start + sequence_len * hidden_size]; - compare_stage( - &mut reports, - &format!("language.post_injection_layer_{target_layer}"), - actual, - &eager, - tolerance, - ) - .unwrap_or_else(|error| panic!("{error}")); + record_comparison( + &mut failures, + compare_stage( + &mut reports, + &format!("language.post_injection_layer_{target_layer}"), + actual, + &eager, + tolerance, + ), + ); } let eager_logits = mlx_f32(&eager_language.logits); let vocab = iree_language.logits.len(); let eager_last = &eager_logits[(sequence_len - 1) * vocab..sequence_len * vocab]; - compare_stage( - &mut reports, - "language.final_position_logits", - &iree_language.logits, - eager_last, - tolerance, - ) - .unwrap_or_else(|error| panic!("{error}")); + record_comparison( + &mut failures, + compare_stage( + &mut reports, + "language.final_position_logits", + &iree_language.logits, + eager_last, + tolerance, + ), + ); + let first_failure = failures.first().cloned(); let report = json!({ "schema": 1, "model": model_path, "image": image_path, "device": device, "context_capacity": context_capacity, + "completed_phase": "language", "grid_thw": vision_capture.grids, "vision_shape": vision_capture.projection.shape, "deepstack_branches": eager_branches.len(), @@ -452,12 +547,55 @@ fn main() { "negative_dropped_branch_detected": true, "negative_zeroed_branch_detected": true, "comparisons": reports, - "passed": true, + "comparison_failures": failures, + "passed": first_failure.is_none(), }); - 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())); + render_and_persist_report(&report, report_path.as_deref()); + if let Some(error) = first_failure { + panic!("Qwen3-VL DeepStack oracle failed after rendering its report: {error}"); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn comparisons_continue_after_failure_and_persist_the_report() { + let mut reports = Vec::new(); + let mut failures = Vec::new(); + let tolerance = Tolerance { + atol: 0.0, + rtol: 0.0, + }; + record_comparison( + &mut failures, + compare_stage(&mut reports, "first", &[1.0], &[0.0], tolerance), + ); + record_comparison( + &mut failures, + compare_stage(&mut reports, "second", &[2.0], &[2.0], tolerance), + ); + assert_eq!(reports.len(), 2); + assert_eq!(failures.len(), 1); + assert_eq!(reports[0]["passed"], false); + assert_eq!(reports[1]["passed"], true); + + let path = std::env::temp_dir().join(format!( + "mlxcel-qwen3-vl-report-persistence-{}.json", + std::process::id() + )); + let report = json!({ + "comparisons": reports, + "comparison_failures": failures, + "passed": false, + }); + render_and_persist_report(&report, Some(&path)); + let persisted: Value = + serde_json::from_slice(&fs::read(&path).expect("read persisted report")) + .expect("parse persisted report"); + let _ = fs::remove_file(path); + assert_eq!(persisted["comparisons"].as_array().unwrap().len(), 2); + assert_eq!(persisted["passed"], false); } - println!("{rendered}"); } diff --git a/src/lib/mlxcel-xla/src/emitter/mod.rs b/src/lib/mlxcel-xla/src/emitter/mod.rs index 3b4e4c02..b4b40309 100644 --- a/src/lib/mlxcel-xla/src/emitter/mod.rs +++ b/src/lib/mlxcel-xla/src/emitter/mod.rs @@ -135,7 +135,8 @@ pub(crate) use qwen3_vl::emit_qwen3_vl; pub(crate) use qwen3_vl::emit_qwen3_vl_diagnostics; #[allow(unused_imports)] pub(crate) use qwen3_vl::{ - QWEN3_VL_PATCH_BUCKETS, Qwen3VlConfig, Qwen3VlGridPlan, Qwen3VlHostInputs, Qwen3VlWeightSpec, + QWEN3_VL_BLOCK_DIAGNOSTIC_STAGES, QWEN3_VL_DIAGNOSTIC_BLOCK, QWEN3_VL_PATCH_BUCKETS, + Qwen3VlConfig, Qwen3VlGridPlan, Qwen3VlHostInputs, Qwen3VlWeightSpec, prepare_qwen3_vl_host_inputs, }; // MoE FFN config types (issue #500), read by the weight loader (`iree.rs`) and the diff --git a/src/lib/mlxcel-xla/src/emitter/qwen3_vl.rs b/src/lib/mlxcel-xla/src/emitter/qwen3_vl.rs index 989e43b5..b1957ea0 100644 --- a/src/lib/mlxcel-xla/src/emitter/qwen3_vl.rs +++ b/src/lib/mlxcel-xla/src/emitter/qwen3_vl.rs @@ -35,6 +35,8 @@ pub(crate) const QWEN3_VL_PATCH_BUCKETS: [usize; 4] = [16, 64, 256, 512]; /// Compact DeepStack side-input capacity after the required 2x2 merge. pub(crate) const QWEN3_VL_MAX_MERGED_VISUAL_POSITIONS: usize = 512 / 4; +pub(crate) const QWEN3_VL_DIAGNOSTIC_BLOCK: usize = 2; +pub(crate) const QWEN3_VL_BLOCK_DIAGNOSTIC_STAGES: usize = 7; #[derive(Debug, Clone, PartialEq, Eq)] pub(crate) struct Qwen3VlWeightSpec { @@ -829,6 +831,30 @@ fn attention( linear_2d(builder, &context, &args.take(), &args.take()) } +struct EncoderLayerDiagnostics { + input: Val, + norm1: Val, + attention: Val, + post_attention_residual: Val, + norm2: Val, + mlp: Val, + output: Val, +} + +impl EncoderLayerDiagnostics { + fn into_values(self) -> [Val; QWEN3_VL_BLOCK_DIAGNOSTIC_STAGES] { + [ + self.input, + self.norm1, + self.attention, + self.post_attention_residual, + self.norm2, + self.mlp, + self.output, + ] + } +} + fn encoder_layer( builder: &mut Builder, hidden: &Val, @@ -836,7 +862,8 @@ fn encoder_layer( config: &Qwen3VlConfig, freqs: &Val, attention_bias: &Val, -) -> Val { + diagnostics: bool, +) -> (Val, Option) { let norm1 = layer_norm( builder, hidden, @@ -855,8 +882,18 @@ fn encoder_layer( ); let fc1 = linear_2d(builder, &norm2, &args.take(), &args.take()); let activated = tanh_gelu(builder, &fc1); - let fc2 = linear_2d(builder, &activated, &args.take(), &args.take()); - builder.add(&residual, &fc2) + let mlp = linear_2d(builder, &activated, &args.take(), &args.take()); + let output = builder.add(&residual, &mlp); + let diagnostic_values = diagnostics.then(|| EncoderLayerDiagnostics { + input: hidden.clone(), + norm1, + attention, + post_attention_residual: residual, + norm2, + mlp, + output: output.clone(), + }); + (output, diagnostic_values) } fn interpolated_position_embeddings( @@ -957,15 +994,22 @@ fn emit_qwen3_vl_impl(config: &Qwen3VlConfig, patch_bucket: usize, diagnostics: let positioned_embeddings = diagnostics.then(|| hidden.clone()); let mut deepstack_hidden = Vec::with_capacity(config.deepstack_visual_indexes.len()); let mut block_hidden = diagnostics.then(|| Vec::with_capacity(config.depth)); + let mut block_2_diagnostics = None; for layer in 0..config.depth { - hidden = encoder_layer( + let (next_hidden, layer_diagnostics) = encoder_layer( &mut builder, &hidden, &mut args, config, &freqs, &attention_bias, + diagnostics && layer == QWEN3_VL_DIAGNOSTIC_BLOCK, ); + hidden = next_hidden; + if layer_diagnostics.is_some() { + debug_assert!(block_2_diagnostics.is_none()); + block_2_diagnostics = layer_diagnostics; + } if let Some(outputs) = &mut block_hidden { outputs.push(hidden.clone()); } @@ -994,6 +1038,9 @@ fn emit_qwen3_vl_impl(config: &Qwen3VlConfig, patch_bucket: usize, diagnostics: ); outputs.push(positioned_embeddings.expect("diagnostics includes positioned embeddings")); outputs.extend(block_hidden); + if let Some(block_2_diagnostics) = block_2_diagnostics { + outputs.extend(block_2_diagnostics.into_values()); + } outputs.extend(main_merger_diagnostics.expect("diagnostics includes main merger stages")); outputs.extend( deepstack_merger_diagnostics @@ -1165,6 +1212,21 @@ mod tests { assert_eq!(mlir.matches("chlo.erf").count(), 3); } + #[test] + fn diagnostic_bucket_appends_all_block_2_substages_in_order() { + let mut config = config(); + config.depth = 3; + let diagnostic = emit_qwen3_vl_impl(&config, 16, true); + let mut output_types = vec!["tensor<4x12xf32>"; 3]; + output_types.extend(vec!["tensor<16x8xf32>"; 3 + 3 + 7]); + output_types.extend(vec!["tensor<4x32xf32>"; 9]); + assert!(diagnostic.contains(&format!("-> ({})", output_types.join(", ")))); + + let production = emit_qwen3_vl_impl(&config, 16, false); + assert!(production.contains("-> (tensor<4x12xf32>, tensor<4x12xf32>, tensor<4x12xf32>)")); + assert!(!production.contains(&format!("-> ({})", output_types.join(", ")))); + } + #[test] fn host_inputs_preserve_packed_media_isolation_and_finite_padding_rows() { let config = config(); diff --git a/src/lib/mlxcel-xla/src/qwen3_vl_runtime.rs b/src/lib/mlxcel-xla/src/qwen3_vl_runtime.rs index 8e54b71e..dbeb92b0 100644 --- a/src/lib/mlxcel-xla/src/qwen3_vl_runtime.rs +++ b/src/lib/mlxcel-xla/src/qwen3_vl_runtime.rs @@ -38,7 +38,10 @@ use crate::aux_manifest::{AuxiliaryArtifactContract, ensure_qualified_auxiliary_ use crate::emitter::emit_qwen3_vl; #[cfg(feature = "diagnostics")] use crate::emitter::emit_qwen3_vl_diagnostics; -use crate::emitter::{Qwen3VlConfig, Qwen3VlHostInputs, prepare_qwen3_vl_host_inputs}; +use crate::emitter::{ + QWEN3_VL_BLOCK_DIAGNOSTIC_STAGES, QWEN3_VL_DIAGNOSTIC_BLOCK, Qwen3VlConfig, Qwen3VlHostInputs, + prepare_qwen3_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}; @@ -74,6 +77,9 @@ pub struct Qwen3VlVisionDiagnostics { /// merger input; configured DeepStack inputs are the corresponding /// `deepstack_visual_indexes` entries. pub block_hidden_states: Vec>, + /// Input, norm1, attention, post-attention residual, norm2, MLP, and + /// output for the first divergent encoder block. + pub block_2_states: Vec>, pub shape: [usize; 2], /// Normalized/shuffled input, first linear output, and GELU output. pub main_merger_states: Vec>, @@ -732,6 +738,12 @@ impl IreeQwen3VlProjector { output_shapes.extend( (0..3 + self.config.depth).map(|_| vec![plan.patch_bucket, self.config.hidden]), ); + if self.config.depth > QWEN3_VL_DIAGNOSTIC_BLOCK { + output_shapes.extend( + (0..QWEN3_VL_BLOCK_DIAGNOSTIC_STAGES) + .map(|_| vec![plan.patch_bucket, self.config.hidden]), + ); + } let merge_width = self.config.hidden * self.config.spatial_merge_size * self.config.spatial_merge_size; @@ -842,6 +854,13 @@ impl IreeQwen3VlProjector { let block_hidden_states = (0..self.config.depth) .map(|layer| trim(&format!("block {layer}"))) .collect::, _>>()?; + let block_2_states = if self.config.depth > QWEN3_VL_DIAGNOSTIC_BLOCK { + (0..QWEN3_VL_BLOCK_DIAGNOSTIC_STAGES) + .map(|stage| trim(&format!("block 2 stage {stage}"))) + .collect::, _>>()? + } else { + Vec::new() + }; let actual_merged_tokens = plan.merged_tokens_per_image.iter().sum::(); let merge_width = self.config.hidden * self.config.spatial_merge_size @@ -876,6 +895,7 @@ impl IreeQwen3VlProjector { position_embeddings, positioned_embeddings, block_hidden_states, + block_2_states, shape: [plan.actual_patches, self.config.hidden], main_merger_states, deepstack_merger_states, diff --git a/src/vision/encoders/qwen3_vl.rs b/src/vision/encoders/qwen3_vl.rs index 6a358608..9f513c89 100644 --- a/src/vision/encoders/qwen3_vl.rs +++ b/src/vision/encoders/qwen3_vl.rs @@ -601,12 +601,60 @@ impl VisionBlock { cu_seqlens: &[i32], rotary_pos_emb: &MlxArray, ) -> UniquePtr { - let normed = self.norm1.forward(hidden_states); - let attn_out = self.attn.forward(&normed, cu_seqlens, rotary_pos_emb); - let h = mlxcel_core::add(hidden_states, &attn_out); - let normed = self.norm2.forward(&h); - let mlp_out = self.mlp.forward(&normed); - mlxcel_core::add(&h, &mlp_out) + self.forward_with_observer(hidden_states, cu_seqlens, rotary_pos_emb, |_, _| {}) + } + + fn forward_with_observer( + &self, + hidden_states: &MlxArray, + cu_seqlens: &[i32], + rotary_pos_emb: &MlxArray, + mut observe: F, + ) -> UniquePtr + where + F: FnMut(VisionBlockDiagnosticStage, &MlxArray), + { + observe(VisionBlockDiagnosticStage::Input, hidden_states); + let norm1 = self.norm1.forward(hidden_states); + observe(VisionBlockDiagnosticStage::Norm1, &norm1); + let attention = self.attn.forward(&norm1, cu_seqlens, rotary_pos_emb); + observe(VisionBlockDiagnosticStage::Attention, &attention); + let residual = mlxcel_core::add(hidden_states, &attention); + observe(VisionBlockDiagnosticStage::PostAttentionResidual, &residual); + let norm2 = self.norm2.forward(&residual); + observe(VisionBlockDiagnosticStage::Norm2, &norm2); + let mlp = self.mlp.forward(&norm2); + observe(VisionBlockDiagnosticStage::Mlp, &mlp); + let output = mlxcel_core::add(&residual, &mlp); + observe(VisionBlockDiagnosticStage::Output, &output); + output + } +} + +#[derive(Clone, Copy)] +enum VisionBlockDiagnosticStage { + Input, + Norm1, + Attention, + PostAttentionResidual, + Norm2, + Mlp, + Output, +} + +impl VisionBlockDiagnosticStage { + const COUNT: usize = 7; + + fn index(self) -> usize { + match self { + Self::Input => 0, + Self::Norm1 => 1, + Self::Attention => 2, + Self::PostAttentionResidual => 3, + Self::Norm2 => 4, + Self::Mlp => 5, + Self::Output => 6, + } } } @@ -712,6 +760,9 @@ pub struct Qwen3VLVisionEncoderDiagnosticOutput { /// merger input; configured DeepStack inputs are the corresponding /// `deepstack_visual_indexes` entries. pub block_hidden_states: Vec>, + /// Input, norm1, attention, post-attention residual, norm2, MLP, and + /// output for encoder block 2. + pub block_2_states: Vec>, /// Normalized/shuffled input, first linear output, and GELU output. pub main_merger_states: Vec>, /// The same ordered merger stages for each DeepStack branch. @@ -724,6 +775,7 @@ enum VisionStage { PositionEmbedding, PositionedEmbedding, Block(usize), + Block2(VisionBlockDiagnosticStage), MainMerger(PatchMergerStage), DeepStackMerger(usize, PatchMergerStage), } @@ -907,6 +959,7 @@ impl Qwen3VLVisionEncoder { let mut position_embeddings = None; let mut positioned_embeddings = None; let mut block_hidden_states = Vec::with_capacity(self.blocks.len()); + let mut block_2_states = Vec::with_capacity(VisionBlockDiagnosticStage::COUNT); let mut main_merger_states = Vec::with_capacity(3); let mut deepstack_merger_states = (0..self.deepstack_merger_list.len()) .map(|_| Vec::with_capacity(3)) @@ -926,6 +979,12 @@ impl Qwen3VLVisionEncoder { debug_assert_eq!(layer, block_hidden_states.len()); block_hidden_states.push(mlxcel_core::copy(hidden)); } + VisionStage::Block2(stage) => { + debug_assert_eq!(stage.index(), block_2_states.len()); + let snapshot = mlxcel_core::copy(hidden); + mlxcel_core::eval(&snapshot); + block_2_states.push(snapshot); + } VisionStage::MainMerger(stage) => { debug_assert_eq!(stage.index(), main_merger_states.len()); main_merger_states.push(mlxcel_core::copy(hidden)); @@ -943,6 +1002,7 @@ impl Qwen3VLVisionEncoder { positioned_embeddings: positioned_embeddings .expect("observer captures positioned embeddings"), block_hidden_states, + block_2_states, main_merger_states, deepstack_merger_states, } @@ -986,7 +1046,13 @@ impl Qwen3VLVisionEncoder { let mut deepstack_features: Vec> = Vec::new(); for (layer_num, block) in self.blocks.iter().enumerate() { - h = block.forward(&h, &cu_seqlens, &rotary_pos_emb); + h = if layer_num == 2 { + block.forward_with_observer(&h, &cu_seqlens, &rotary_pos_emb, |stage, hidden| { + observe(VisionStage::Block2(stage), hidden) + }) + } else { + block.forward(&h, &cu_seqlens, &rotary_pos_emb) + }; observe(VisionStage::Block(layer_num), &h); if let Some(ds_idx) = self diff --git a/src/vision/encoders/qwen3_vl_tests.rs b/src/vision/encoders/qwen3_vl_tests.rs index ee0da0e5..bc793665 100644 --- a/src/vision/encoders/qwen3_vl_tests.rs +++ b/src/vision/encoders/qwen3_vl_tests.rs @@ -12,7 +12,10 @@ // See the License for the specific language governing permissions and // limitations under the License. -use super::{ensure_fused_sdpa, fused_sdpa_target_dim, gelu_pytorch_tanh, sdpa_pad_width}; +use super::{ + VisionBlockDiagnosticStage, ensure_fused_sdpa, fused_sdpa_target_dim, gelu_pytorch_tanh, + sdpa_pad_width, +}; use mlxcel_core::dtype; #[test] @@ -60,3 +63,20 @@ fn block_gelu_matches_qwen3_vl_pytorch_tanh_contract() { ); } } + +#[test] +fn block_2_diagnostic_stage_order_is_stable() { + let stages = [ + VisionBlockDiagnosticStage::Input, + VisionBlockDiagnosticStage::Norm1, + VisionBlockDiagnosticStage::Attention, + VisionBlockDiagnosticStage::PostAttentionResidual, + VisionBlockDiagnosticStage::Norm2, + VisionBlockDiagnosticStage::Mlp, + VisionBlockDiagnosticStage::Output, + ]; + assert_eq!(stages.len(), VisionBlockDiagnosticStage::COUNT); + for (index, stage) in stages.into_iter().enumerate() { + assert_eq!(stage.index(), index); + } +} From b7e6bb8c010f66fe23da9a435e784f6d534872d4 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 01:49:10 +0900 Subject: [PATCH 07/10] test(xla): isolate Qwen3-VL LayerNorm drift Block-level comparisons can identify LayerNorm as the first failing stage even when its input already contains tolerated quantized-matmul drift. Re-run the eager norm1 and norm2 operators on the exact IREE block input and post-attention residual, then report dedicated same-input comparisons. Add deterministic eager and StableHLO contract tests for the centered-variance F32 LayerNorm formula. Validated with rustfmt, focused eager and emitter tests, and the xla-diagnostics example check. Refs #868 --- examples/xla_qwen3_vl_deepstack_check.rs | 24 ++++++++++ src/lib/mlxcel-xla/src/emitter/qwen3_vl.rs | 24 ++++++++++ src/vision/encoders/qwen3_vl.rs | 53 ++++++++++++++++++++++ src/vision/encoders/qwen3_vl_tests.rs | 42 +++++++++++++++++ 4 files changed, 143 insertions(+) diff --git a/examples/xla_qwen3_vl_deepstack_check.rs b/examples/xla_qwen3_vl_deepstack_check.rs index e255af60..ff8292aa 100644 --- a/examples/xla_qwen3_vl_deepstack_check.rs +++ b/examples/xla_qwen3_vl_deepstack_check.rs @@ -349,6 +349,30 @@ fn main() { ), ); } + let same_input_norms = model + .vision_encoder + .block_layer_norms_from_f32( + 2, + &iree_diagnostics.block_2_states[0], + &iree_diagnostics.block_2_states[3], + iree_diagnostics.shape[0], + ) + .unwrap_or_else(|error| panic!("run Qwen3-VL same-input LayerNorm probe: {error}")); + for (name, actual_index, expected) in [ + ("norm1", 1, &same_input_norms[0]), + ("norm2", 4, &same_input_norms[1]), + ] { + record_comparison( + &mut failures, + compare_stage( + &mut reports, + &format!("vision.block_2.{name}.same_input"), + &iree_diagnostics.block_2_states[actual_index], + &mlx_f32(expected), + tolerance, + ), + ); + } let merger_stage_names = ["normalized_shuffled", "fc1", "activation"]; assert_eq!( iree_diagnostics.main_merger_states.len(), diff --git a/src/lib/mlxcel-xla/src/emitter/qwen3_vl.rs b/src/lib/mlxcel-xla/src/emitter/qwen3_vl.rs index b1957ea0..343ef374 100644 --- a/src/lib/mlxcel-xla/src/emitter/qwen3_vl.rs +++ b/src/lib/mlxcel-xla/src/emitter/qwen3_vl.rs @@ -1227,6 +1227,30 @@ mod tests { assert!(!production.contains(&format!("-> ({})", output_types.join(", ")))); } + #[test] + fn layer_norm_emits_centered_variance_f32_contract() { + let mut builder = Builder::new(); + let value = Builder::arg(0, Ty::f32(vec![2, 8])); + let weight = Builder::arg(1, Ty::f32(vec![8])); + let bias = Builder::arg(2, Ty::f32(vec![8])); + let output = layer_norm(&mut builder, &value, &weight, &bias, 1.0e-6); + let body = builder.body(); + + assert_eq!(output.ty.render(), "tensor<2x8xf32>"); + assert_eq!(body.matches("stablehlo.reduce").count(), 2); + assert_eq!(body.matches("stablehlo.subtract").count(), 1); + assert_eq!(body.matches("stablehlo.rsqrt").count(), 1); + let reductions = body + .match_indices("stablehlo.reduce") + .map(|(index, _)| index) + .collect::>(); + let centered = body.find("stablehlo.subtract").unwrap(); + assert!( + reductions[0] < centered && centered < reductions[1], + "variance reduction must consume centered values" + ); + } + #[test] fn host_inputs_preserve_packed_media_isolation_and_finite_padding_rows() { let config = config(); diff --git a/src/vision/encoders/qwen3_vl.rs b/src/vision/encoders/qwen3_vl.rs index 9f513c89..0ddac3ac 100644 --- a/src/vision/encoders/qwen3_vl.rs +++ b/src/vision/encoders/qwen3_vl.rs @@ -1008,6 +1008,59 @@ impl Qwen3VLVisionEncoder { } } + /// Re-run only the two LayerNorm operators of one vision block on supplied + /// F32 activations. The Qwen3-VL oracle uses this to compare eager MLX and + /// IREE with identical inputs, separating normalization arithmetic from + /// upstream quantized-matmul drift. + #[cfg(feature = "xla-diagnostics")] + pub fn block_layer_norms_from_f32( + &self, + layer: usize, + norm1_input: &[f32], + norm2_input: &[f32], + rows: usize, + ) -> Result<[UniquePtr; 2], String> { + let block = self + .blocks + .get(layer) + .ok_or_else(|| format!("Qwen3-VL diagnostic block {layer} is out of range"))?; + let weight_shape = mlxcel_core::array_shape(&block.norm1.weight); + let [width] = weight_shape.as_slice() else { + return Err(format!( + "Qwen3-VL block {layer} norm1 weight must be rank 1, got {weight_shape:?}" + )); + }; + let norm2_weight_shape = mlxcel_core::array_shape(&block.norm2.weight); + if norm2_weight_shape != weight_shape { + return Err(format!( + "Qwen3-VL block {layer} norm widths disagree: norm1={weight_shape:?}, norm2={norm2_weight_shape:?}" + )); + } + let width = usize::try_from(*width) + .map_err(|_| format!("Qwen3-VL block {layer} norm width must be non-negative"))?; + let expected = rows + .checked_mul(width) + .ok_or_else(|| "Qwen3-VL diagnostic LayerNorm shape overflowed".to_string())?; + for (name, input) in [("norm1", norm1_input), ("norm2", norm2_input)] { + if input.len() != expected { + return Err(format!( + "Qwen3-VL block {layer} {name} diagnostic input has {} values, expected {expected}", + input.len() + )); + } + } + let rows = i32::try_from(rows) + .map_err(|_| "Qwen3-VL diagnostic LayerNorm row count does not fit i32".to_string())?; + let width = i32::try_from(width) + .map_err(|_| "Qwen3-VL diagnostic LayerNorm width does not fit i32".to_string())?; + let norm1_input = mlxcel_core::from_slice_f32(norm1_input, &[rows, width]); + let norm2_input = mlxcel_core::from_slice_f32(norm2_input, &[rows, width]); + Ok([ + block.norm1.forward(&norm1_input), + block.norm2.forward(&norm2_input), + ]) + } + fn forward_with_grid_observer( &self, hidden_states: &MlxArray, diff --git a/src/vision/encoders/qwen3_vl_tests.rs b/src/vision/encoders/qwen3_vl_tests.rs index bc793665..83055bc0 100644 --- a/src/vision/encoders/qwen3_vl_tests.rs +++ b/src/vision/encoders/qwen3_vl_tests.rs @@ -17,6 +17,7 @@ use super::{ sdpa_pad_width, }; use mlxcel_core::dtype; +use mlxcel_core::layers::LayerNorm; #[test] fn fused_sdpa_target_dim_uses_next_supported_kernel_width() { @@ -64,6 +65,47 @@ fn block_gelu_matches_qwen3_vl_pytorch_tanh_contract() { } } +#[test] +fn block_layer_norm_matches_centered_variance_f32_contract() { + let values = [3.5_f32, -2.0, 0.25, 7.0, -4.0, 1.5, 8.0, -0.75]; + let weights = [1.25_f32, 0.5, -0.75, 2.0]; + let biases = [-0.25_f32, 0.125, 0.5, -1.0]; + let input = mlxcel_core::from_slice_f32(&values, &[2, 4]); + let layer = LayerNorm::new( + mlxcel_core::from_slice_f32(&weights, &[4]), + Some(mlxcel_core::from_slice_f32(&biases, &[4])), + 1.0e-6, + ); + let output = layer.forward(&input); + mlxcel_core::eval(&output); + + for row in 0..2 { + let row_values = &values[row * 4..row * 4 + 4]; + let mean = row_values.iter().sum::() / 4.0; + let variance = row_values + .iter() + .map(|value| { + let centered = value - mean; + centered * centered + }) + .sum::() + / 4.0; + let inv_std = (variance + 1.0e-6).sqrt().recip(); + for column in 0..4 { + let expected = (row_values[column] - mean) * inv_std * weights[column] + biases[column]; + let actual = mlxcel_core::slice( + &output, + &[row as i32, column as i32], + &[row as i32 + 1, column as i32 + 1], + ); + assert!( + (mlxcel_core::item_f32(&actual) - expected).abs() <= 2.0e-5, + "Qwen3-VL LayerNorm mismatch at row {row}, column {column}" + ); + } + } +} + #[test] fn block_2_diagnostic_stage_order_is_stable() { let stages = [ From b21753e8e3fbc737dbd8eb7c1a6ee25a79c08e7b Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 02:07:32 +0900 Subject: [PATCH 08/10] test(xla): expose Qwen3-VL block 0 substages The same-input LayerNorm probe proves block 2 normalization is within tolerance when both runtimes consume identical activations, so the remaining divergence is accumulated upstream rather than a LayerNorm precision defect. Capture block 0 input, norm1, attention, residual, norm2, MLP, and output on both eager MLX and IREE diagnostic paths. The oracle now reports these stages before block 2 without changing production precision, tolerances, or checkpoint execution. Validated with the focused Qwen3-VL emitter tests, diagnostics library checks, the diagnostics-enabled oracle example check, and Clippy for mlxcel-xla. Refs #868 --- examples/xla_qwen3_vl_deepstack_check.rs | 33 ++++++++++++++++++++-- src/lib/mlxcel-xla/src/emitter/mod.rs | 4 +-- src/lib/mlxcel-xla/src/emitter/qwen3_vl.rs | 24 +++++++++++----- src/lib/mlxcel-xla/src/qwen3_vl_runtime.rs | 22 +++++++++++++-- src/vision/encoders/qwen3_vl.rs | 33 +++++++++++++++++----- 5 files changed, 94 insertions(+), 22 deletions(-) diff --git a/examples/xla_qwen3_vl_deepstack_check.rs b/examples/xla_qwen3_vl_deepstack_check.rs index ff8292aa..1f4176f4 100644 --- a/examples/xla_qwen3_vl_deepstack_check.rs +++ b/examples/xla_qwen3_vl_deepstack_check.rs @@ -313,7 +313,7 @@ fn main() { ), ); } - let block_2_stage_names = [ + let block_stage_names = [ "input", "norm1", "attention", @@ -322,6 +322,33 @@ fn main() { "mlp", "output", ]; + assert_eq!( + iree_diagnostics.block_0_states.len(), + eager_vision.block_0_states.len(), + "eager and IREE block 0 substage capture counts differ" + ); + assert_eq!( + iree_diagnostics.block_0_states.len(), + block_stage_names.len(), + "actual Qwen3-VL diagnostics must capture every block 0 substage" + ); + for ((actual, expected), name) in iree_diagnostics + .block_0_states + .iter() + .zip(&eager_vision.block_0_states) + .zip(block_stage_names) + { + record_comparison( + &mut failures, + compare_stage( + &mut reports, + &format!("vision.block_0.{name}"), + actual, + &mlx_f32(expected), + tolerance, + ), + ); + } assert_eq!( iree_diagnostics.block_2_states.len(), eager_vision.block_2_states.len(), @@ -329,14 +356,14 @@ fn main() { ); assert_eq!( iree_diagnostics.block_2_states.len(), - block_2_stage_names.len(), + block_stage_names.len(), "actual Qwen3-VL diagnostics must capture every block 2 substage" ); for ((actual, expected), name) in iree_diagnostics .block_2_states .iter() .zip(&eager_vision.block_2_states) - .zip(block_2_stage_names) + .zip(block_stage_names) { record_comparison( &mut failures, diff --git a/src/lib/mlxcel-xla/src/emitter/mod.rs b/src/lib/mlxcel-xla/src/emitter/mod.rs index b4b40309..5382b692 100644 --- a/src/lib/mlxcel-xla/src/emitter/mod.rs +++ b/src/lib/mlxcel-xla/src/emitter/mod.rs @@ -135,8 +135,8 @@ pub(crate) use qwen3_vl::emit_qwen3_vl; pub(crate) use qwen3_vl::emit_qwen3_vl_diagnostics; #[allow(unused_imports)] pub(crate) use qwen3_vl::{ - QWEN3_VL_BLOCK_DIAGNOSTIC_STAGES, QWEN3_VL_DIAGNOSTIC_BLOCK, QWEN3_VL_PATCH_BUCKETS, - Qwen3VlConfig, Qwen3VlGridPlan, Qwen3VlHostInputs, Qwen3VlWeightSpec, + QWEN3_VL_BLOCK_DIAGNOSTIC_STAGES, QWEN3_VL_DIAGNOSTIC_BLOCK, QWEN3_VL_FIRST_DIAGNOSTIC_BLOCK, + QWEN3_VL_PATCH_BUCKETS, Qwen3VlConfig, Qwen3VlGridPlan, Qwen3VlHostInputs, Qwen3VlWeightSpec, prepare_qwen3_vl_host_inputs, }; // MoE FFN config types (issue #500), read by the weight loader (`iree.rs`) and the diff --git a/src/lib/mlxcel-xla/src/emitter/qwen3_vl.rs b/src/lib/mlxcel-xla/src/emitter/qwen3_vl.rs index 343ef374..70a0f5e5 100644 --- a/src/lib/mlxcel-xla/src/emitter/qwen3_vl.rs +++ b/src/lib/mlxcel-xla/src/emitter/qwen3_vl.rs @@ -35,6 +35,7 @@ pub(crate) const QWEN3_VL_PATCH_BUCKETS: [usize; 4] = [16, 64, 256, 512]; /// Compact DeepStack side-input capacity after the required 2x2 merge. pub(crate) const QWEN3_VL_MAX_MERGED_VISUAL_POSITIONS: usize = 512 / 4; +pub(crate) const QWEN3_VL_FIRST_DIAGNOSTIC_BLOCK: usize = 0; pub(crate) const QWEN3_VL_DIAGNOSTIC_BLOCK: usize = 2; pub(crate) const QWEN3_VL_BLOCK_DIAGNOSTIC_STAGES: usize = 7; @@ -994,6 +995,7 @@ fn emit_qwen3_vl_impl(config: &Qwen3VlConfig, patch_bucket: usize, diagnostics: let positioned_embeddings = diagnostics.then(|| hidden.clone()); let mut deepstack_hidden = Vec::with_capacity(config.deepstack_visual_indexes.len()); let mut block_hidden = diagnostics.then(|| Vec::with_capacity(config.depth)); + let mut block_0_diagnostics = None; let mut block_2_diagnostics = None; for layer in 0..config.depth { let (next_hidden, layer_diagnostics) = encoder_layer( @@ -1003,12 +1005,17 @@ fn emit_qwen3_vl_impl(config: &Qwen3VlConfig, patch_bucket: usize, diagnostics: config, &freqs, &attention_bias, - diagnostics && layer == QWEN3_VL_DIAGNOSTIC_BLOCK, + diagnostics + && matches!( + layer, + QWEN3_VL_FIRST_DIAGNOSTIC_BLOCK | QWEN3_VL_DIAGNOSTIC_BLOCK + ), ); hidden = next_hidden; - if layer_diagnostics.is_some() { - debug_assert!(block_2_diagnostics.is_none()); - block_2_diagnostics = layer_diagnostics; + match layer { + QWEN3_VL_FIRST_DIAGNOSTIC_BLOCK => block_0_diagnostics = layer_diagnostics, + QWEN3_VL_DIAGNOSTIC_BLOCK => block_2_diagnostics = layer_diagnostics, + _ => debug_assert!(layer_diagnostics.is_none()), } if let Some(outputs) = &mut block_hidden { outputs.push(hidden.clone()); @@ -1038,6 +1045,9 @@ fn emit_qwen3_vl_impl(config: &Qwen3VlConfig, patch_bucket: usize, diagnostics: ); outputs.push(positioned_embeddings.expect("diagnostics includes positioned embeddings")); outputs.extend(block_hidden); + if let Some(block_0_diagnostics) = block_0_diagnostics { + outputs.extend(block_0_diagnostics.into_values()); + } if let Some(block_2_diagnostics) = block_2_diagnostics { outputs.extend(block_2_diagnostics.into_values()); } @@ -1206,19 +1216,19 @@ mod tests { let config = config(); let mlir = emit_qwen3_vl_impl(&config, 16, true); let mut output_types = vec!["tensor<4x12xf32>"; 3]; - output_types.extend(vec!["tensor<16x8xf32>"; 5]); + output_types.extend(vec!["tensor<16x8xf32>"; 5 + 7]); output_types.extend(vec!["tensor<4x32xf32>"; 9]); assert!(mlir.contains(&format!("-> ({})", output_types.join(", ")))); assert_eq!(mlir.matches("chlo.erf").count(), 3); } #[test] - fn diagnostic_bucket_appends_all_block_2_substages_in_order() { + fn diagnostic_bucket_appends_all_block_0_and_block_2_substages_in_order() { let mut config = config(); config.depth = 3; let diagnostic = emit_qwen3_vl_impl(&config, 16, true); let mut output_types = vec!["tensor<4x12xf32>"; 3]; - output_types.extend(vec!["tensor<16x8xf32>"; 3 + 3 + 7]); + output_types.extend(vec!["tensor<16x8xf32>"; 3 + 3 + 2 * 7]); output_types.extend(vec!["tensor<4x32xf32>"; 9]); assert!(diagnostic.contains(&format!("-> ({})", output_types.join(", ")))); diff --git a/src/lib/mlxcel-xla/src/qwen3_vl_runtime.rs b/src/lib/mlxcel-xla/src/qwen3_vl_runtime.rs index dbeb92b0..a6f0941a 100644 --- a/src/lib/mlxcel-xla/src/qwen3_vl_runtime.rs +++ b/src/lib/mlxcel-xla/src/qwen3_vl_runtime.rs @@ -39,8 +39,8 @@ use crate::emitter::emit_qwen3_vl; #[cfg(feature = "diagnostics")] use crate::emitter::emit_qwen3_vl_diagnostics; use crate::emitter::{ - QWEN3_VL_BLOCK_DIAGNOSTIC_STAGES, QWEN3_VL_DIAGNOSTIC_BLOCK, Qwen3VlConfig, Qwen3VlHostInputs, - prepare_qwen3_vl_host_inputs, + QWEN3_VL_BLOCK_DIAGNOSTIC_STAGES, QWEN3_VL_DIAGNOSTIC_BLOCK, QWEN3_VL_FIRST_DIAGNOSTIC_BLOCK, + Qwen3VlConfig, Qwen3VlHostInputs, prepare_qwen3_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}; @@ -78,7 +78,9 @@ pub struct Qwen3VlVisionDiagnostics { /// `deepstack_visual_indexes` entries. pub block_hidden_states: Vec>, /// Input, norm1, attention, post-attention residual, norm2, MLP, and - /// output for the first divergent encoder block. + /// output for encoder block 0. + pub block_0_states: Vec>, + /// The same ordered stages for encoder block 2. pub block_2_states: Vec>, pub shape: [usize; 2], /// Normalized/shuffled input, first linear output, and GELU output. @@ -738,6 +740,12 @@ impl IreeQwen3VlProjector { output_shapes.extend( (0..3 + self.config.depth).map(|_| vec![plan.patch_bucket, self.config.hidden]), ); + if self.config.depth > QWEN3_VL_FIRST_DIAGNOSTIC_BLOCK { + output_shapes.extend( + (0..QWEN3_VL_BLOCK_DIAGNOSTIC_STAGES) + .map(|_| vec![plan.patch_bucket, self.config.hidden]), + ); + } if self.config.depth > QWEN3_VL_DIAGNOSTIC_BLOCK { output_shapes.extend( (0..QWEN3_VL_BLOCK_DIAGNOSTIC_STAGES) @@ -854,6 +862,13 @@ impl IreeQwen3VlProjector { let block_hidden_states = (0..self.config.depth) .map(|layer| trim(&format!("block {layer}"))) .collect::, _>>()?; + let block_0_states = if self.config.depth > QWEN3_VL_FIRST_DIAGNOSTIC_BLOCK { + (0..QWEN3_VL_BLOCK_DIAGNOSTIC_STAGES) + .map(|stage| trim(&format!("block 0 stage {stage}"))) + .collect::, _>>()? + } else { + Vec::new() + }; let block_2_states = if self.config.depth > QWEN3_VL_DIAGNOSTIC_BLOCK { (0..QWEN3_VL_BLOCK_DIAGNOSTIC_STAGES) .map(|stage| trim(&format!("block 2 stage {stage}"))) @@ -895,6 +910,7 @@ impl IreeQwen3VlProjector { position_embeddings, positioned_embeddings, block_hidden_states, + block_0_states, block_2_states, shape: [plan.actual_patches, self.config.hidden], main_merger_states, diff --git a/src/vision/encoders/qwen3_vl.rs b/src/vision/encoders/qwen3_vl.rs index 0ddac3ac..6e038280 100644 --- a/src/vision/encoders/qwen3_vl.rs +++ b/src/vision/encoders/qwen3_vl.rs @@ -761,7 +761,9 @@ pub struct Qwen3VLVisionEncoderDiagnosticOutput { /// `deepstack_visual_indexes` entries. pub block_hidden_states: Vec>, /// Input, norm1, attention, post-attention residual, norm2, MLP, and - /// output for encoder block 2. + /// output for encoder block 0. + pub block_0_states: Vec>, + /// The same ordered stages for encoder block 2. pub block_2_states: Vec>, /// Normalized/shuffled input, first linear output, and GELU output. pub main_merger_states: Vec>, @@ -775,6 +777,7 @@ enum VisionStage { PositionEmbedding, PositionedEmbedding, Block(usize), + Block0(VisionBlockDiagnosticStage), Block2(VisionBlockDiagnosticStage), MainMerger(PatchMergerStage), DeepStackMerger(usize, PatchMergerStage), @@ -959,6 +962,7 @@ impl Qwen3VLVisionEncoder { let mut position_embeddings = None; let mut positioned_embeddings = None; let mut block_hidden_states = Vec::with_capacity(self.blocks.len()); + let mut block_0_states = Vec::with_capacity(VisionBlockDiagnosticStage::COUNT); let mut block_2_states = Vec::with_capacity(VisionBlockDiagnosticStage::COUNT); let mut main_merger_states = Vec::with_capacity(3); let mut deepstack_merger_states = (0..self.deepstack_merger_list.len()) @@ -979,6 +983,12 @@ impl Qwen3VLVisionEncoder { debug_assert_eq!(layer, block_hidden_states.len()); block_hidden_states.push(mlxcel_core::copy(hidden)); } + VisionStage::Block0(stage) => { + debug_assert_eq!(stage.index(), block_0_states.len()); + let snapshot = mlxcel_core::copy(hidden); + mlxcel_core::eval(&snapshot); + block_0_states.push(snapshot); + } VisionStage::Block2(stage) => { debug_assert_eq!(stage.index(), block_2_states.len()); let snapshot = mlxcel_core::copy(hidden); @@ -1002,6 +1012,7 @@ impl Qwen3VLVisionEncoder { positioned_embeddings: positioned_embeddings .expect("observer captures positioned embeddings"), block_hidden_states, + block_0_states, block_2_states, main_merger_states, deepstack_merger_states, @@ -1099,12 +1110,20 @@ impl Qwen3VLVisionEncoder { let mut deepstack_features: Vec> = Vec::new(); for (layer_num, block) in self.blocks.iter().enumerate() { - h = if layer_num == 2 { - block.forward_with_observer(&h, &cu_seqlens, &rotary_pos_emb, |stage, hidden| { - observe(VisionStage::Block2(stage), hidden) - }) - } else { - block.forward(&h, &cu_seqlens, &rotary_pos_emb) + h = match layer_num { + 0 => block.forward_with_observer( + &h, + &cu_seqlens, + &rotary_pos_emb, + |stage, hidden| observe(VisionStage::Block0(stage), hidden), + ), + 2 => block.forward_with_observer( + &h, + &cu_seqlens, + &rotary_pos_emb, + |stage, hidden| observe(VisionStage::Block2(stage), hidden), + ), + _ => block.forward(&h, &cu_seqlens, &rotary_pos_emb), }; observe(VisionStage::Block(layer_num), &h); From 9814b93ab9efc61e0c7fbd665973ff0812c536f5 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 03:37:54 +0900 Subject: [PATCH 09/10] test(xla): isolate Qwen3-VL projection drift Qwen3-VL vision parity first crosses tolerance at block 2 after block 0 and block 1 remain within tolerance, leaving fused quantized projections and IREE's host-dequantized dense F32 weights as the primary unresolved arithmetic boundary. Add diagnostics-only same-input block 0 and block 1 controls that run the production fused-QMM path beside a dense F32 path using the exact affine host widening shared with the IREE checkpoint loader. Record both outputs in the actual-checkpoint oracle without changing production execution or tolerance policy. Validated with cargo fmt, the focused dense-F32 control unit test, cargo check for the xla-diagnostics example on the GB10/IREE CUDA setup, and cargo clippy for the same example. Refs #868 --- examples/xla_qwen3_vl_deepstack_check.rs | 27 +++ src/lib/mlxcel-xla/src/lib.rs | 25 +++ src/vision/encoders/qwen3_vl.rs | 271 +++++++++++++++++++++++ src/vision/encoders/qwen3_vl_tests.rs | 30 ++- 4 files changed, 351 insertions(+), 2 deletions(-) diff --git a/examples/xla_qwen3_vl_deepstack_check.rs b/examples/xla_qwen3_vl_deepstack_check.rs index 1f4176f4..4c12276a 100644 --- a/examples/xla_qwen3_vl_deepstack_check.rs +++ b/examples/xla_qwen3_vl_deepstack_check.rs @@ -376,6 +376,33 @@ fn main() { ), ); } + for (layer, input) in [ + (0, &iree_diagnostics.block_0_states[0]), + (1, &iree_diagnostics.block_hidden_states[0]), + ] { + let controls = model + .vision_encoder + .block_projection_controls_from_f32(layer, input, &vision_capture.grids) + .unwrap_or_else(|error| { + panic!("run Qwen3-VL block {layer} same-input projection controls: {error}") + }); + let actual = &iree_diagnostics.block_hidden_states[layer]; + for (name, expected) in [ + ("fused_qmm", &controls.fused_qmm), + ("host_dequant_dense_f32", &controls.host_dequant_dense_f32), + ] { + record_comparison( + &mut failures, + compare_stage( + &mut reports, + &format!("vision.block_{layer}.output.same_input.{name}"), + actual, + &mlx_f32(expected), + tolerance, + ), + ); + } + } let same_input_norms = model .vision_encoder .block_layer_norms_from_f32( diff --git a/src/lib/mlxcel-xla/src/lib.rs b/src/lib/mlxcel-xla/src/lib.rs index 72f40767..651a0a68 100644 --- a/src/lib/mlxcel-xla/src/lib.rs +++ b/src/lib/mlxcel-xla/src/lib.rs @@ -214,6 +214,31 @@ pub fn dequantize_gemma3n_affine_diagnostic( ) -> Result, String> { weights::dequantize_affine_bf16_fused(packed, scales, biases, out, in_packed, bits, group_size) } + +/// Expose the exact host-side affine widening used by the Qwen3-VL IREE +/// checkpoint loader to diagnostics without making it part of the runtime API. +#[cfg(feature = "diagnostics")] +pub fn dequantize_affine_f32_diagnostic( + packed: &[u8], + scales: &[u8], + biases: &[u8], + out: usize, + in_packed: usize, + bits: usize, + group_size: usize, + scales_bf16: bool, +) -> Result, String> { + weights::dequantize_affine( + packed, + scales, + biases, + out, + in_packed, + bits, + group_size, + scales_bf16, + ) +} #[cfg(feature = "diagnostics")] pub use emitter::{Gemma3nDiagnosticLayout, Gemma3nDiagnosticSegment}; #[cfg(feature = "iree")] diff --git a/src/vision/encoders/qwen3_vl.rs b/src/vision/encoders/qwen3_vl.rs index 6e038280..faf13fdd 100644 --- a/src/vision/encoders/qwen3_vl.rs +++ b/src/vision/encoders/qwen3_vl.rs @@ -402,6 +402,111 @@ struct VisionAttention { scale: f32, } +#[cfg(any(test, feature = "xla-diagnostics"))] +fn diagnostic_dense_linear_from_weight_f32( + input: &MlxArray, + weight: &MlxArray, + bias: Option<&MlxArray>, +) -> UniquePtr { + let input = mlxcel_core::astype(input, mlxcel_core::dtype::FLOAT32); + let weight = mlxcel_core::astype(weight, mlxcel_core::dtype::FLOAT32); + let weight = mlxcel_core::transpose(&weight); + let output = mlxcel_core::matmul(&input, &weight); + match bias { + Some(bias) => { + let bias = mlxcel_core::astype(bias, mlxcel_core::dtype::FLOAT32); + mlxcel_core::add(&output, &bias) + } + None => output, + } +} + +/// Run one checkpoint projection with the exact affine-to-F32 host widening +/// used by the IREE loader, followed by an MLX dense F32 matmul. This is +/// diagnostics-only: production eager inference keeps its fused QMM path. +#[cfg(feature = "xla-diagnostics")] +fn diagnostic_dense_linear_f32( + linear: &UnifiedLinear, + input: &MlxArray, +) -> Result, String> { + let (weight, bias) = match linear { + UnifiedLinear::Quantized { weight, bias } => { + if weight.mode != "affine" { + return Err(format!( + "Qwen3-VL host-dequant control requires affine weights, got {}", + weight.mode + )); + } + if weight.global_scale.is_some() { + return Err( + "Qwen3-VL IREE host-dequant control does not accept a global-scale sidecar" + .to_string(), + ); + } + let weight_shape = mlxcel_core::array_shape(&weight.weight); + let [out, in_packed] = weight_shape.as_slice() else { + return Err(format!( + "Qwen3-VL packed projection must be rank 2, got {weight_shape:?}" + )); + }; + let out = usize::try_from(*out) + .map_err(|_| "Qwen3-VL packed projection rows must be non-negative".to_string())?; + let in_packed = usize::try_from(*in_packed) + .map_err(|_| "Qwen3-VL packed projection width must be non-negative".to_string())?; + let bits = usize::try_from(weight.bits) + .map_err(|_| "Qwen3-VL quantization bits must be non-negative".to_string())?; + let group_size = usize::try_from(weight.group_size) + .map_err(|_| "Qwen3-VL quantization group size must be non-negative".to_string())?; + let biases = weight.biases.as_deref().ok_or_else(|| { + "Qwen3-VL affine projection is missing quantization biases".to_string() + })?; + let scales_dtype = mlxcel_core::array_dtype(&weight.scales); + let biases_dtype = mlxcel_core::array_dtype(biases); + let scales_bf16 = match (scales_dtype, biases_dtype) { + (mlxcel_core::dtype::BFLOAT16, mlxcel_core::dtype::BFLOAT16) => true, + (mlxcel_core::dtype::FLOAT16, mlxcel_core::dtype::FLOAT16) => false, + _ => { + return Err(format!( + "Qwen3-VL affine metadata must use matching F16 or BF16, got {scales_dtype}/{biases_dtype}" + )); + } + }; + let values = mlxcel_xla::dequantize_affine_f32_diagnostic( + &mlxcel_core::array_to_raw_bytes(&weight.weight), + &mlxcel_core::array_to_raw_bytes(&weight.scales), + &mlxcel_core::array_to_raw_bytes(biases), + out, + in_packed, + bits, + group_size, + scales_bf16, + )?; + let values_per_word = 32usize + .checked_div(bits) + .filter(|value| *value > 0) + .ok_or_else(|| format!("invalid Qwen3-VL affine bit width {bits}"))?; + let width = in_packed + .checked_mul(values_per_word) + .ok_or_else(|| "Qwen3-VL dequantized projection width overflowed".to_string())?; + let out_i32 = i32::try_from(out) + .map_err(|_| "Qwen3-VL projection rows do not fit i32".to_string())?; + let width_i32 = i32::try_from(width) + .map_err(|_| "Qwen3-VL projection width does not fit i32".to_string())?; + ( + mlxcel_core::from_slice_f32(&values, &[out_i32, width_i32]), + bias.as_deref(), + ) + } + UnifiedLinear::Regular(linear) => ( + mlxcel_core::astype(&linear.weight, mlxcel_core::dtype::FLOAT32), + linear.bias.as_deref(), + ), + }; + Ok(diagnostic_dense_linear_from_weight_f32( + input, &weight, bias, + )) +} + impl VisionAttention { fn from_weights( weights: &WeightMap, @@ -507,6 +612,84 @@ impl VisionAttention { self.proj.forward(&output) } + + #[cfg(feature = "xla-diagnostics")] + fn forward_dequantized_f32( + &self, + x: &MlxArray, + cu_seqlens: &[i32], + rotary_pos_emb: &MlxArray, + ) -> Result, String> { + let shape = mlxcel_core::array_shape(x); + let seq_length = shape[0]; + + let qkv = diagnostic_dense_linear_f32(&self.qkv, x)?; + let qkv = mlxcel_core::reshape(&qkv, &[seq_length, 3, self.num_heads, self.head_dim]); + let qkv = mlxcel_core::transpose_axes(&qkv, &[1, 0, 2, 3]); + + let q = mlxcel_core::slice( + &qkv, + &[0, 0, 0, 0], + &[1, seq_length, self.num_heads, self.head_dim], + ); + let k = mlxcel_core::slice( + &qkv, + &[1, 0, 0, 0], + &[2, seq_length, self.num_heads, self.head_dim], + ); + let v = mlxcel_core::slice( + &qkv, + &[2, 0, 0, 0], + &[3, seq_length, self.num_heads, self.head_dim], + ); + let q = mlxcel_core::squeeze_axis(&q, 0); + let k = mlxcel_core::squeeze_axis(&k, 0); + let v = mlxcel_core::squeeze_axis(&v, 0); + + let q = apply_rotary_pos_emb_vision(&q, rotary_pos_emb); + let k = apply_rotary_pos_emb_vision(&k, rotary_pos_emb); + + let q = mlxcel_core::expand_dims(&mlxcel_core::transpose_axes(&q, &[1, 0, 2]), 0); + let k = mlxcel_core::expand_dims(&mlxcel_core::transpose_axes(&k, &[1, 0, 2]), 0); + let v = mlxcel_core::expand_dims(&mlxcel_core::transpose_axes(&v, &[1, 0, 2]), 0); + let mut attn_outputs = Vec::with_capacity(cu_seqlens.len().saturating_sub(1)); + for segment in cu_seqlens.windows(2) { + let [start, end] = segment else { + unreachable!("windows(2) always yields two boundaries"); + }; + let q_segment = mlxcel_core::slice( + &q, + &[0, 0, *start, 0], + &[1, self.num_heads, *end, self.head_dim], + ); + let k_segment = mlxcel_core::slice( + &k, + &[0, 0, *start, 0], + &[1, self.num_heads, *end, self.head_dim], + ); + let v_segment = mlxcel_core::slice( + &v, + &[0, 0, *start, 0], + &[1, self.num_heads, *end, self.head_dim], + ); + attn_outputs.push(ensure_fused_sdpa( + &q_segment, &k_segment, &v_segment, self.scale, None, + )); + } + + let output = if attn_outputs.len() == 1 { + attn_outputs + .into_iter() + .next() + .expect("one attention segment") + } else { + concat_many(&attn_outputs, 2) + }; + 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]); + diagnostic_dense_linear_f32(&self.proj, &output) + } } #[cfg(test)] @@ -569,6 +752,13 @@ impl VisionMLP { let h = gelu_pytorch_tanh(&h); self.linear_fc2.forward(&h) } + + #[cfg(feature = "xla-diagnostics")] + fn forward_dequantized_f32(&self, x: &MlxArray) -> Result, String> { + let hidden = diagnostic_dense_linear_f32(&self.linear_fc1, x)?; + let hidden = gelu_pytorch_tanh(&hidden); + diagnostic_dense_linear_f32(&self.linear_fc2, &hidden) + } } // VisionBlock - LayerNorm + GELU MLP. @@ -629,6 +819,23 @@ impl VisionBlock { observe(VisionBlockDiagnosticStage::Output, &output); output } + + #[cfg(feature = "xla-diagnostics")] + fn forward_dequantized_f32( + &self, + hidden_states: &MlxArray, + cu_seqlens: &[i32], + rotary_pos_emb: &MlxArray, + ) -> Result, String> { + let norm1 = self.norm1.forward(hidden_states); + let attention = self + .attn + .forward_dequantized_f32(&norm1, cu_seqlens, rotary_pos_emb)?; + let residual = mlxcel_core::add(hidden_states, &attention); + let norm2 = self.norm2.forward(&residual); + let mlp = self.mlp.forward_dequantized_f32(&norm2)?; + Ok(mlxcel_core::add(&residual, &mlp)) + } } #[derive(Clone, Copy)] @@ -771,6 +978,14 @@ pub struct Qwen3VLVisionEncoderDiagnosticOutput { pub deepstack_merger_states: Vec>>, } +/// Same-input eager controls for distinguishing fused QMM accumulation from +/// explicitly dequantized dense F32 projection arithmetic. +#[cfg(feature = "xla-diagnostics")] +pub struct Qwen3VLVisionBlockProjectionControls { + pub fused_qmm: UniquePtr, + pub host_dequant_dense_f32: UniquePtr, +} + #[derive(Clone, Copy)] enum VisionStage { PatchEmbedding, @@ -1072,6 +1287,62 @@ impl Qwen3VLVisionEncoder { ]) } + /// Re-run one vision block on an exact supplied F32 input through both the + /// production eager fused-QMM path and a diagnostics-only host-dequantized + /// dense-F32 path. All non-linear operators are shared, so the two outputs + /// isolate the block's four checkpoint projections. + #[cfg(feature = "xla-diagnostics")] + pub fn block_projection_controls_from_f32( + &self, + layer: usize, + input: &[f32], + grid_thw: &[(i32, i32, i32)], + ) -> Result { + let block = self + .blocks + .get(layer) + .ok_or_else(|| format!("Qwen3-VL diagnostic block {layer} is out of range"))?; + let weight_shape = mlxcel_core::array_shape(&block.norm1.weight); + let [width] = weight_shape.as_slice() else { + return Err(format!( + "Qwen3-VL block {layer} norm1 weight must be rank 1, got {weight_shape:?}" + )); + }; + let width = usize::try_from(*width) + .map_err(|_| format!("Qwen3-VL block {layer} width must be non-negative"))?; + if width == 0 || !input.len().is_multiple_of(width) { + return Err(format!( + "Qwen3-VL block {layer} diagnostic input has {} values, not complete rows of width {width}", + input.len() + )); + } + let rows = input.len() / width; + let rows_i32 = i32::try_from(rows) + .map_err(|_| "Qwen3-VL diagnostic block row count does not fit i32".to_string())?; + let width_i32 = i32::try_from(width) + .map_err(|_| "Qwen3-VL diagnostic block width does not fit i32".to_string())?; + let cu_seqlens = Self::compute_cu_seqlens(grid_thw); + let grid_rows = cu_seqlens.last().copied().unwrap_or_default(); + if grid_rows != rows_i32 { + return Err(format!( + "Qwen3-VL block {layer} diagnostic grid covers {grid_rows} rows, input has {rows}" + )); + } + let rotary_pos_emb = self.rot_pos_emb(grid_thw); + let rotary_shape = mlxcel_core::array_shape(&rotary_pos_emb); + let rotary_pos_emb = mlxcel_core::reshape(&rotary_pos_emb, &[rotary_shape[0], -1]); + let input = mlxcel_core::from_slice_f32(input, &[rows_i32, width_i32]); + + Ok(Qwen3VLVisionBlockProjectionControls { + fused_qmm: block.forward(&input, &cu_seqlens, &rotary_pos_emb), + host_dequant_dense_f32: block.forward_dequantized_f32( + &input, + &cu_seqlens, + &rotary_pos_emb, + )?, + }) + } + fn forward_with_grid_observer( &self, hidden_states: &MlxArray, diff --git a/src/vision/encoders/qwen3_vl_tests.rs b/src/vision/encoders/qwen3_vl_tests.rs index 83055bc0..be4d92cd 100644 --- a/src/vision/encoders/qwen3_vl_tests.rs +++ b/src/vision/encoders/qwen3_vl_tests.rs @@ -13,11 +13,17 @@ // limitations under the License. use super::{ - VisionBlockDiagnosticStage, ensure_fused_sdpa, fused_sdpa_target_dim, gelu_pytorch_tanh, - sdpa_pad_width, + VisionBlockDiagnosticStage, diagnostic_dense_linear_from_weight_f32, ensure_fused_sdpa, + fused_sdpa_target_dim, gelu_pytorch_tanh, sdpa_pad_width, }; use mlxcel_core::dtype; use mlxcel_core::layers::LayerNorm; +use std::sync::Once; + +fn ensure_cpu_device() { + static INIT: Once = Once::new(); + INIT.call_once(|| mlxcel_core::set_default_device(false)); +} #[test] fn fused_sdpa_target_dim_uses_next_supported_kernel_width() { @@ -106,6 +112,26 @@ fn block_layer_norm_matches_centered_variance_f32_contract() { } } +#[test] +fn diagnostic_dense_linear_runs_transposed_weight_and_bias_in_f32() { + ensure_cpu_device(); + let input = mlxcel_core::from_slice_f32(&[1.0, 2.0, 3.0, 4.0], &[2, 2]); + let weight = mlxcel_core::from_slice_f32(&[2.0, 0.0, 0.0, 3.0], &[2, 2]); + let bias = mlxcel_core::from_slice_f32(&[0.5, -1.0], &[2]); + + let output = diagnostic_dense_linear_from_weight_f32(&input, &weight, Some(&bias)); + mlxcel_core::eval(&output); + + assert_eq!(mlxcel_core::array_dtype(&output), dtype::FLOAT32); + assert_eq!(mlxcel_core::array_shape(&output), vec![2, 2]); + for (index, expected) in [2.5_f32, 5.0, 6.5, 11.0].into_iter().enumerate() { + let row = i32::try_from(index / 2).expect("small test row"); + let column = i32::try_from(index % 2).expect("small test column"); + let value = mlxcel_core::slice(&output, &[row, column], &[row + 1, column + 1]); + assert_eq!(mlxcel_core::item_f32(&value), expected); + } +} + #[test] fn block_2_diagnostic_stage_order_is_stable() { let stages = [ From a4b641e6a538e559d0b9c9f18a0b9f4beb3e11b2 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 09:19:29 +0900 Subject: [PATCH 10/10] fix(xla): mark Qwen3-VL artifact unqualified Keep the rebased Qwen3-VL runtime on the explicit legacy artifact path until its Q4 numeric oracle is rerun against the current CUDA kernel.\n\nRefs #868 --- src/lib/mlxcel-xla/src/qwen3_vl_runtime.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/lib/mlxcel-xla/src/qwen3_vl_runtime.rs b/src/lib/mlxcel-xla/src/qwen3_vl_runtime.rs index a6f0941a..26ccbd6a 100644 --- a/src/lib/mlxcel-xla/src/qwen3_vl_runtime.rs +++ b/src/lib/mlxcel-xla/src/qwen3_vl_runtime.rs @@ -639,7 +639,7 @@ fn compile_and_load( 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!( "{};{};checkpoint_schema_sha256={}",