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..4c12276a --- /dev/null +++ b/examples/xla_qwen3_vl_deepstack_check.rs @@ -0,0 +1,679 @@ +// 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 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}; + +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 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], + 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_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 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" + ); + 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() + { + record_comparison( + &mut failures, + compare_stage( + &mut reports, + &format!("vision.block_{layer}"), + actual, + &mlx_f32(expected), + tolerance, + ), + ); + } + let block_stage_names = [ + "input", + "norm1", + "attention", + "post_attention_residual", + "norm2", + "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(), + "eager and IREE block 2 substage capture counts differ" + ); + assert_eq!( + iree_diagnostics.block_2_states.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_stage_names) + { + record_comparison( + &mut failures, + compare_stage( + &mut reports, + &format!("vision.block_2.{name}"), + actual, + &mlx_f32(expected), + tolerance, + ), + ); + } + 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( + 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(), + 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) + { + 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, + "vision.main_merger", + &vision_capture.projection.values, + &eager_main, + tolerance, + ), + ); + 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) + { + record_comparison( + &mut failures, + compare_stage( + &mut reports, + &format!("vision.deepstack_merger_{branch}.{name}"), + actual, + &mlx_f32(expected), + tolerance, + ), + ); + } + } + 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 + .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" + ); + 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 + .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]; + 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]; + 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(), + "deepstack_target_language_layers": iree_language.target_layer_indices, + "negative_dropped_branch_detected": true, + "negative_zeroed_branch_detected": true, + "comparisons": reports, + "comparison_failures": failures, + "passed": first_failure.is_none(), + }); + 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); + } +} 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.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/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..5382b692 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,16 @@ 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; +#[cfg(feature = "diagnostics")] +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_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 // validation harness. `SharedExpertConfig` is only named in some build cfgs. #[allow(unused_imports)] @@ -1125,6 +1136,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" @@ -1196,6 +1213,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: 128, + }) + ); + 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..70a0f5e5 --- /dev/null +++ b/src/lib/mlxcel-xla/src/emitter/qwen3_vl.rs @@ -0,0 +1,1324 @@ +// 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; 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; + +#[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 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) + .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(|| "Qwen3-VL patch bucket selection failed".to_string()) + } + + /// 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()) +} + +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, + args: &mut Args, + config: &Qwen3VlConfig, + freqs: &Val, + attention_bias: &Val, + diagnostics: bool, +) -> (Val, Option) { + 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 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( + 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_impl( + builder: &mut Builder, + hidden: &Val, + args: &mut Args, + config: &Qwen3VlConfig, + postshuffle_norm: bool, + 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); + 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 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. +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" + ); + 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 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)); + 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( + &mut builder, + &hidden, + &mut args, + config, + &freqs, + &attention_bias, + diagnostics + && matches!( + layer, + QWEN3_VL_FIRST_DIAGNOSTIC_BLOCK | QWEN3_VL_DIAGNOSTIC_BLOCK + ), + ); + hidden = next_hidden; + 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()); + } + if config.deepstack_visual_indexes.contains(&layer) { + deepstack_hidden.push(hidden.clone()); + } + } + 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 { + 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); + 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()); + } + 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 + .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(), + ) +} + +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::*; + + 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], + "hidden_act": "gelu_pytorch_tanh" + } + }"#, + ) + .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_eq!(config.bucket_for_patches(512), Ok(512)); + assert!( + config + .bucket_for_patches(18) + .unwrap_err() + .contains("divisible") + ); + assert!( + config + .bucket_for_patches(513) + .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]); + + 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] + 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("stablehlo.tanh").count(), config.depth); + 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] + 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 + 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_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 + 2 * 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 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(); + 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") + ); + + 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/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 2b132826..651a0a68 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")] @@ -145,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, @@ -167,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, @@ -181,6 +184,12 @@ 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, +}; +#[cfg(feature = "diagnostics")] pub use vision_runtime::{IreeVisionDiagnosticProjector, VisionDiagnosticProjection}; #[cfg(feature = "iree")] pub use vision_runtime::{IreeVisionProjector, VisionExecutionMetrics, VisionProjection}; @@ -205,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/lib/mlxcel-xla/src/qwen3_vl_runtime.rs b/src/lib/mlxcel-xla/src/qwen3_vl_runtime.rs new file mode 100644 index 00000000..26ccbd6a --- /dev/null +++ b/src/lib/mlxcel-xla/src/qwen3_vl_runtime.rs @@ -0,0 +1,1017 @@ +// 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}; +#[cfg(not(feature = "diagnostics"))] +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, 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}; + +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>, + #[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>, + /// Input, norm1, attention, post-attention residual, norm2, MLP, and + /// 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. + 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, +} + +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]" + )); + } + } + let size = object.get("size").and_then(serde_json::Value::as_object); + let positive_bound = |nested: &str, fallback: &str| -> 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-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, + 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 { + #[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())); + } + 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_legacy_unqualified( + 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 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]), + ); + 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) + .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() + .zip(&output_shapes) + .map(|(output, shape)| AuxiliaryOutput { + bytes: output.as_mut_slice(), + dtype: AuxiliaryTensorDType::Float32, + 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 + .drain(..self.config.deepstack_visual_indexes.len()) + .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::, _>>()?; + #[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 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}"))) + .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 + * 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, + block_0_states, + block_2_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, + 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, + 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..887f2bf5 100644 --- a/src/loading/vlm_qwen.rs +++ b/src/loading/vlm_qwen.rs @@ -25,12 +25,16 @@ //! about Qwen family details. use anyhow::Result; +#[cfg(feature = "xla-iree")] +use image::imageops::FilterType; 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 +204,290 @@ 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_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( + 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(), + ) +} + +#[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/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 bfad9cb9..c0f32d17 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,32 @@ 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, +} + +/// 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 { @@ -458,6 +528,301 @@ 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) + } + + /// 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, + 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..8d6dd6a5 100644 --- a/src/multimodal/host_preprocessor_tests.rs +++ b/src/multimodal/host_preprocessor_tests.rs @@ -17,6 +17,9 @@ 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::{ FakeHostMultimodalPreprocessor, HostMultimodalPreprocessor, HostPreprocessorError, XlaVisionBackend, XlaVisionBackendPolicy, export_llava_prefill, export_mlx_tensor, @@ -65,6 +68,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"] { @@ -163,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/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!( diff --git a/src/vision/encoders/qwen3_vl.rs b/src/vision/encoders/qwen3_vl.rs index bc93c6e7..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)] @@ -520,6 +703,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,9 +749,16 @@ 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) } + + #[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. @@ -575,12 +791,77 @@ 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 + } + + #[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)] +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, + } } } @@ -630,7 +911,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 +924,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 +957,47 @@ 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>, + /// Input, norm1, attention, post-attention residual, norm2, MLP, and + /// 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>, + /// The same ordered merger stages for each DeepStack branch. + 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, + PositionEmbedding, + PositionedEmbedding, + Block(usize), + Block0(VisionBlockDiagnosticStage), + Block2(VisionBlockDiagnosticStage), + MainMerger(PatchMergerStage), + DeepStackMerger(usize, PatchMergerStage), +} + // Qwen3-VL Vision Encoder. /// Qwen3-VL Vision Model with learned position embeddings and DeepStack /// @@ -819,14 +1164,205 @@ 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 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()) + .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::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); + 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)); + } + 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, + block_0_states, + block_2_states, + main_merger_states, + deepstack_merger_states, + } + } + + /// 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), + ]) + } + + /// 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, + 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); @@ -845,20 +1381,42 @@ 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 = 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); 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, diff --git a/src/vision/encoders/qwen3_vl_tests.rs b/src/vision/encoders/qwen3_vl_tests.rs index 4b2769f1..be4d92cd 100644 --- a/src/vision/encoders/qwen3_vl_tests.rs +++ b/src/vision/encoders/qwen3_vl_tests.rs @@ -12,8 +12,18 @@ // 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::{ + 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() { @@ -44,3 +54,97 @@ 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}" + ); + } +} + +#[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 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 = [ + 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); + } +} 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)] + ); + } }