From e285c54cb5c79310899d981e93a6b1b2e4a0fff1 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 24 Jul 2026 18:54:35 +0900 Subject: [PATCH 01/11] feat: add Gemma3 VLM OpenXLA path Emit the pinned Gemma3 SigLIP tower, average-pool projector, and nested text configuration through native IREE while preserving Gemma3 embedding scaling and authoritative bidirectional padding masks. Route image requests through the filtered host embedding table and resident projector without a duplicate decoder or silent host fallback. Validate prepared-prefill contracts, host/server seams, product feature compilation, and a real CUDA projector execution. Refs #869 --- src/lib.rs | 2 + src/lib/mlxcel-xla/src/emitter/config.rs | 78 +++- src/lib/mlxcel-xla/src/emitter/mod.rs | 4 +- src/lib/mlxcel-xla/src/emitter/model.rs | 45 +- src/lib/mlxcel-xla/src/emitter/vision.rs | 121 ++++- .../mlxcel-xla/src/emitter/vision_config.rs | 219 +++++++-- src/lib/mlxcel-xla/src/lib.rs | 5 + src/lib/mlxcel-xla/src/prepared.rs | 26 +- src/lib/mlxcel-xla/src/prepared_gemma3.rs | 432 ++++++++++++++++++ src/lib/mlxcel-xla/src/vision_runtime.rs | 124 +++-- src/loading/mod.rs | 2 + src/loading/vlm.rs | 2 + src/loading/vlm_gemma.rs | 132 +++++- src/loading/vlm_qwen.rs | 8 + src/multimodal/host_preprocessor.rs | 325 ++++++++++++- src/multimodal/host_preprocessor_tests.rs | 2 +- 16 files changed, 1414 insertions(+), 113 deletions(-) create mode 100644 src/lib/mlxcel-xla/src/prepared_gemma3.rs diff --git a/src/lib.rs b/src/lib.rs index d9520a12f..681bbda28 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -75,6 +75,8 @@ pub use mlxcel_core::generate::{ SamplingConfig, }; pub use mlxcel_core::speculative::SpeculativeGenerator; +#[cfg(feature = "xla-iree")] +pub use multimodal::host_preprocessor::Gemma3IreeHostPreprocessor; #[cfg(feature = "xla-diagnostics")] pub use multimodal::host_preprocessor::LlavaHostReferenceCapture; #[cfg(feature = "xla-iree")] diff --git a/src/lib/mlxcel-xla/src/emitter/config.rs b/src/lib/mlxcel-xla/src/emitter/config.rs index e5c450eaa..d5b5626f9 100644 --- a/src/lib/mlxcel-xla/src/emitter/config.rs +++ b/src/lib/mlxcel-xla/src/emitter/config.rs @@ -313,6 +313,11 @@ pub struct Config { /// uses 4; Cohere2 uses `sliding_window_pattern` (4). Only meaningful when /// `sliding_window` is `Some`. pub sliding_pattern: usize, + /// The embeddings-prefill caller supplies the complete per-layer attention + /// policy. Gemma3 VLM uses one bidirectional padding mask for every layer, + /// so the emitter must not intersect it with the text model's ordinary + /// sliding-window schedule. Token prefill remains causal/windowed. + pub embeddings_prefill_uses_authoritative_mask: bool, /// Per-layer NoPE mask (SmolLM3): `use_rope_layers[li] == false` skips RoPE on /// that layer (`no_rope_layers`). `None` applies RoPE on every layer. pub use_rope_layers: Option>, @@ -439,6 +444,7 @@ impl Config { final_logit_softcap: None, sliding_window: None, sliding_pattern: 2, + embeddings_prefill_uses_authoritative_mask: false, use_rope_layers: None, mrope: None, deepstack: None, @@ -482,15 +488,38 @@ 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 v = if matches!(wrapper_model_type, Some("llava" | "llava_next" | "gemma3")) { 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!( + "{wrapper_model_type:?} config.json missing object `text_config` for the XLA text graph" + ) })?; + if wrapper_model_type == Some("gemma3") { + // mlx-vlm's Gemma3 TextConfig supplies these architecture + // defaults. mlx-community conversions commonly omit them from + // the nested object, so resolve the same explicit contract here + // instead of deriving head_dim from hidden_size (Gemma3's q + // projection width is intentionally smaller than hidden_size). + let defaults = [ + ("num_attention_heads", serde_json::json!(8)), + ("num_key_value_heads", serde_json::json!(4)), + ("head_dim", serde_json::json!(256)), + ("rms_norm_eps", serde_json::json!(1.0e-6)), + ("vocab_size", serde_json::json!(262_208)), + ("rope_theta", serde_json::json!(1_000_000.0)), + ("rope_local_base_freq", serde_json::json!(10_000.0)), + ("query_pre_attn_scalar", serde_json::json!(256.0)), + ("sliding_window", serde_json::json!(1024)), + ("sliding_window_pattern", serde_json::json!(6)), + ]; + for (name, value) in defaults { + text.entry(name.to_string()).or_insert(value); + } + } // mlx-community quantized VLMs commonly keep the affine scheme at // the wrapper level even though the tensors belong to the nested // language model. @@ -653,6 +682,7 @@ impl Config { let mut final_logit_softcap: Option = None; let mut sliding_window: Option = None; let mut sliding_pattern = 2usize; + let mut embeddings_prefill_uses_authoritative_mask = false; let mut use_rope_layers: Option> = None; // issue #498 dense arch pack flags. let mut layernorm = false; @@ -765,6 +795,7 @@ impl Config { }); sliding_pattern = ou("sliding_window_pattern").unwrap_or(6).max(1); sliding_window = Some(ou("sliding_window").unwrap_or(4096)); + embeddings_prefill_uses_authoritative_mask = true; rope_local_base = Some(of("rope_local_base_freq").unwrap_or(10000.0)); read_gemma_common( &mut query_pre_attn_scalar, @@ -1072,6 +1103,12 @@ impl Config { // type while carrying the M-RoPE section table; the table is // the explicit schema signal in that representation. Some("default") | Some("dynamic") | Some("mrope") => RopeScaling::Plain, + // The qualified Gemma3 MLX implementation and upstream + // mlx-vlm Gemma3 model select distinct local/global bases + // but do not apply the conversion metadata's legacy linear + // scaling block. Preserve that reference contract instead + // of silently applying a different HF-only position scale. + Some("linear") if model_type == Some("gemma3_text") => RopeScaling::Plain, None if has_mrope_section => RopeScaling::Plain, Some("llama3") => { let sf = |k: &str| -> Result { @@ -1369,6 +1406,7 @@ impl Config { final_logit_softcap, sliding_window, sliding_pattern, + embeddings_prefill_uses_authoritative_mask, use_rope_layers, mrope, deepstack, @@ -1545,6 +1583,40 @@ impl Config { mod tests { use super::*; + #[test] + fn gemma3_embeddings_prefill_owns_the_complete_external_mask() { + let config = Config::from_json_str( + r#"{"model_type":"gemma3_text","hidden_size":8,"num_attention_heads":2, + "num_key_value_heads":1,"head_dim":4,"intermediate_size":16, + "num_hidden_layers":4,"rms_norm_eps":1e-6,"rope_theta":1000000, + "rope_local_base_freq":10000,"sliding_window":2, + "sliding_window_pattern":3,"vocab_size":12, + "hidden_activation":"gelu_pytorch_tanh"}"#, + ) + .unwrap(); + assert!(config.embeddings_prefill_uses_authoritative_mask); + assert_eq!(config.sliding_window, Some(2)); + assert!(config.is_sliding_layer(0), "token prefill remains windowed"); + } + + #[test] + fn gemma3_wrapper_uses_nested_text_config_and_wrapper_quantization() { + let config = Config::from_json_str( + r#"{"model_type":"gemma3","quantization":{"bits":4,"group_size":64}, + "text_config":{"model_type":"gemma3_text","hidden_size":2560, + "intermediate_size":10240,"num_hidden_layers":34,"sliding_window":1024, + "rope_scaling":{"factor":8.0,"rope_type":"linear"}}}"#, + ) + .unwrap(); + assert!(config.embeddings_prefill_uses_authoritative_mask); + assert_eq!(config.quantization.unwrap().bits, 4); + assert_eq!(config.hidden, 2560); + assert_eq!(config.n_q, 8); + assert_eq!(config.n_kv, 4); + assert_eq!(config.head_dim, 256); + assert_eq!(config.rope, RopeScaling::Plain); + } + /// ERNIE-4.5 is rejected with a message naming its interleaved (GPT-J-style) /// RoPE: it looks like a plain-RoPE Llama in config.json but its `rotate_half` /// rotates the (2i, 2i+1) pairs, not the (i, i+d/2) halves the Llama emit uses, diff --git a/src/lib/mlxcel-xla/src/emitter/mod.rs b/src/lib/mlxcel-xla/src/emitter/mod.rs index 0065385c7..07e5bdd16 100644 --- a/src/lib/mlxcel-xla/src/emitter/mod.rs +++ b/src/lib/mlxcel-xla/src/emitter/mod.rs @@ -156,7 +156,9 @@ pub(crate) use vision::emit_vision; #[cfg(any(test, feature = "diagnostics"))] pub(crate) use vision::emit_vision_diagnostics; #[allow(unused_imports)] -pub(crate) use vision_config::{LlavaVisionConfig, VisionActivation, VisionWeightSpec}; +pub(crate) use vision_config::{ + LlavaVisionConfig, VisionActivation, VisionProjector, VisionWeightSpec, +}; #[cfg(test)] mod tests { diff --git a/src/lib/mlxcel-xla/src/emitter/model.rs b/src/lib/mlxcel-xla/src/emitter/model.rs index 16a748cc6..b3e17a5a1 100644 --- a/src/lib/mlxcel-xla/src/emitter/model.rs +++ b/src/lib/mlxcel-xla/src/emitter/model.rs @@ -157,9 +157,9 @@ pub(crate) fn validate_prefill_embeddings_metadata( Ok(()) } -/// Validate a canonical additive attention-bias payload. Only `0.0` (allowed) -/// and [`PREFILL_EMBEDDINGS_MASKED_VALUE`] (masked) are accepted; NaN, infinity, -/// and intermediate additive values are rejected to keep polarity unambiguous. +/// Validate a canonical additive attention-bias payload. Generic families use +/// [`PREFILL_EMBEDDINGS_MASKED_VALUE`]; Gemma3's authoritative VLM mask uses +/// `f32::MIN`. NaN, infinity, and intermediate values are rejected. pub(crate) fn validate_prefill_embeddings_attention_bias( c: &Config, bias: &[f32], @@ -179,15 +179,20 @@ pub(crate) fn validate_prefill_embeddings_attention_bias( bias.len() )); } + let masked_value = if c.embeddings_prefill_uses_authoritative_mask { + f32::MIN + } else { + PREFILL_EMBEDDINGS_MASKED_VALUE + }; if let Some((index, value)) = bias .iter() .copied() .enumerate() - .find(|(_, value)| *value != 0.0 && *value != PREFILL_EMBEDDINGS_MASKED_VALUE) + .find(|(_, value)| *value != 0.0 && *value != masked_value) { return Err(format!( "prefill attention bias at flat index {index} is {value}; expected only 0 or {}", - PREFILL_EMBEDDINGS_MASKED_VALUE + masked_value )); } Ok(()) @@ -3304,6 +3309,14 @@ fn apply_deepstack_after_layer( b.add(hidden, &delta) } +fn embeddings_prefill_local_window(c: &Config) -> Option { + if c.embeddings_prefill_uses_authoritative_mask { + None + } else { + c.sliding_window + } +} + fn emit_prefill_module( c: &Config, sample: bool, @@ -3387,7 +3400,7 @@ fn emit_prefill_module( } PrefillInput::Embeddings { attention_bias, .. } => { let cmask = attention_bias.clone(); - let cmask_local = c.sliding_window.map(|w| { + let cmask_local = embeddings_prefill_local_window(c).map(|w| { let irow = b.iota(lp); let row = b.broadcast(&irow, &[0], vec![lp, lp]); let jcol = b.iota(lp); @@ -3403,7 +3416,7 @@ fn emit_prefill_module( } PrefillInput::DeepStack(deepstack) => { let cmask = deepstack.attention_bias.clone(); - let cmask_local = c.sliding_window.map(|w| { + let cmask_local = embeddings_prefill_local_window(c).map(|w| { let irow = b.iota(lp); let row = b.broadcast(&irow, &[0], vec![lp, lp]); let jcol = b.iota(lp); @@ -3517,6 +3530,24 @@ fn emit_prefill_module( } } +#[cfg(test)] +mod gemma3_vlm_mask_tests { + use super::*; + + #[test] + fn authoritative_embeddings_mask_bypasses_only_the_external_local_intersection() { + let mut config = Config::llama_3_2_1b(); + config.sliding_window = Some(4096); + assert_eq!(embeddings_prefill_local_window(&config), Some(4096)); + config.embeddings_prefill_uses_authoritative_mask = true; + assert_eq!(embeddings_prefill_local_window(&config), None); + assert!( + config.is_sliding_layer(0), + "the token-prefill/decode layer schedule must remain windowed" + ); + } +} + // =========================================================================== // MoE FFN block probe (issue #500): a standalone module for the execution check // =========================================================================== diff --git a/src/lib/mlxcel-xla/src/emitter/vision.rs b/src/lib/mlxcel-xla/src/emitter/vision.rs index 43f7f713e..9d74df598 100644 --- a/src/lib/mlxcel-xla/src/emitter/vision.rs +++ b/src/lib/mlxcel-xla/src/emitter/vision.rs @@ -23,7 +23,9 @@ use super::builder::{Builder, Ty, Val}; use super::numeric_ops::{exact_gelu, layer_norm_2d as layer_norm, stable_softmax, tanh_gelu}; -use super::vision_config::{LlavaVisionConfig, VisionActivation, VisionWeightSpec}; +use super::vision_config::{ + LlavaVisionConfig, VisionActivation, VisionProjector, VisionWeightSpec, +}; struct Args { values: Vec, @@ -58,6 +60,34 @@ mod tests { assert!(mlir.contains("return ")); } + #[test] + fn gemma3_graph_avg_pools_norms_and_projects_without_llava_mlp() { + let config = LlavaVisionConfig::from_json_str( + &serde_json::json!({ + "model_type": "gemma3", + "image_token_index": 99, + "mm_tokens_per_image": 1, + "vision_config": { + "model_type": "siglip_vision_model", + "image_size": 28, + "patch_size": 14, + "hidden_size": 8, + "intermediate_size": 16, + "num_hidden_layers": 1, + "num_attention_heads": 2 + }, + "text_config": {"hidden_size": 12} + }) + .to_string(), + ) + .unwrap(); + let mlir = emit_vision(&config); + assert!(mlir.contains("multi_modal_projector.mm_soft_emb_norm.weight")); + assert!(mlir.contains("multi_modal_projector.mm_input_projection_weight")); + assert!(!mlir.contains("multi_modal_projector.linear_1")); + assert!(mlir.contains("-> tensor<1x12xf32>")); + } + #[cfg(feature = "iree")] #[test] fn pinned_siglip_graph_compiles_for_cpu() { @@ -129,6 +159,73 @@ fn linear_2d(builder: &mut Builder, value: &Val, weight: &Val, bias: &Val) -> Va bias_2d(builder, &value, bias) } +fn gemma_rms_norm(builder: &mut Builder, value: &Val, weight: &Val, epsilon: f32) -> Val { + let rows = value.ty.shape[0]; + let width = value.ty.shape[1]; + let zero = builder.const_f32(0.0); + let squared = builder.multiply(value, value); + let squared_sum = builder.reduce_add(&squared, 1, &zero); + let width_scalar = builder.const_f32(width as f32); + let width_rows = builder.broadcast(&width_scalar, &[], vec![rows]); + let mean_square = builder.divide(&squared_sum, &width_rows); + let epsilon = builder.const_f32(epsilon); + let epsilon = builder.broadcast(&epsilon, &[], vec![rows]); + let mean_square = builder.add(&mean_square, &epsilon); + let inv_rms = builder.rsqrt(&mean_square); + let inv_rms = builder.broadcast(&inv_rms, &[0], vec![rows, width]); + let normalized = builder.multiply(value, &inv_rms); + let one = builder.const_f32(1.0); + let one = builder.broadcast(&one, &[], vec![width]); + let weight = builder.add(weight, &one); + let weight = builder.broadcast(&weight, &[1], vec![rows, width]); + builder.multiply(&normalized, &weight) +} + +fn gemma3_project( + builder: &mut Builder, + hidden: &Val, + args: &mut Args, + config: &LlavaVisionConfig, + tokens_per_side: usize, + kernel_size: usize, +) -> Val { + let pooled = builder.reshape( + hidden, + vec![ + tokens_per_side, + kernel_size, + tokens_per_side, + kernel_size, + config.hidden, + ], + ); + let zero = builder.const_f32(0.0); + let pooled = builder.reduce_add(&pooled, 3, &zero); + let pooled = builder.reduce_add(&pooled, 1, &zero); + let divisor = builder.const_f32((kernel_size * kernel_size) as f32); + let divisor = builder.broadcast( + &divisor, + &[], + vec![tokens_per_side, tokens_per_side, config.hidden], + ); + let pooled = builder.divide(&pooled, &divisor); + let pooled = builder.reshape( + &pooled, + vec![tokens_per_side * tokens_per_side, config.hidden], + ); + let normalized = gemma_rms_norm(builder, &pooled, &args.take(), config.layer_norm_eps); + let projection = args.take(); + builder.dot_general( + &normalized, + &projection, + &[], + &[], + &[1], + &[0], + vec![tokens_per_side * tokens_per_side, config.text_hidden], + ) +} + fn activate(builder: &mut Builder, value: &Val, activation: VisionActivation) -> Val { match activation { VisionActivation::ExactGelu => exact_gelu(builder, value), @@ -317,9 +414,25 @@ fn emit_vision_impl(config: &LlavaVisionConfig, diagnostics: bool) -> String { if config.drop_first_token { hidden = builder.slice(&hidden, &[(1, config.position_count()), (0, config.hidden)]); } - let projected = linear_2d(&mut builder, &hidden, &args.take(), &args.take()); - let projected = exact_gelu(&mut builder, &projected); - let projected = linear_2d(&mut builder, &projected, &args.take(), &args.take()); + let projected = match config.projector { + VisionProjector::LlavaMlp => { + let projected = linear_2d(&mut builder, &hidden, &args.take(), &args.take()); + let projected = exact_gelu(&mut builder, &projected); + linear_2d(&mut builder, &projected, &args.take(), &args.take()) + } + VisionProjector::Gemma3AvgPool { + tokens_per_side, + kernel_size, + .. + } => gemma3_project( + &mut builder, + &hidden, + &mut args, + config, + tokens_per_side, + kernel_size, + ), + }; assert_eq!(args.cursor, specs.len(), "vision weight schema drifted"); let outputs = if let Some(mut values) = diagnostic_values { values.push(hidden); diff --git a/src/lib/mlxcel-xla/src/emitter/vision_config.rs b/src/lib/mlxcel-xla/src/emitter/vision_config.rs index aee70e746..f53f581ad 100644 --- a/src/lib/mlxcel-xla/src/emitter/vision_config.rs +++ b/src/lib/mlxcel-xla/src/emitter/vision_config.rs @@ -24,6 +24,20 @@ pub(crate) enum VisionActivation { GeluPytorchTanh, } +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub(crate) enum VisionProjector { + LlavaMlp, + Gemma3AvgPool { + tokens_per_side: usize, + kernel_size: usize, + image_token_id: i32, + pad_token_id: i32, + boi_token_id: i32, + eoi_token_id: i32, + newline_token_id: i32, + }, +} + #[cfg(test)] mod tests { use super::*; @@ -171,6 +185,55 @@ mod tests { } ); } + + #[test] + fn gemma3_avg_pool_projector_binds_pool_and_prompt_contract() { + let config = LlavaVisionConfig::from_json_str( + &serde_json::json!({ + "model_type": "gemma3", + "image_token_index": 99, + "pad_token_id": 0, + "boi_token_index": 97, + "eoi_token_index": 98, + "mm_tokens_per_image": 1, + "vision_config": { + "model_type": "siglip_vision_model", + "image_size": 28, + "patch_size": 14, + "num_channels": 3, + "hidden_size": 8, + "intermediate_size": 16, + "num_hidden_layers": 2, + "num_attention_heads": 2 + }, + "text_config": {"hidden_size": 12} + }) + .to_string(), + ) + .unwrap(); + assert_eq!(config.feature_layer, 1); + assert!(!config.drop_first_token); + assert_eq!(config.image_tokens(), 1); + assert_eq!( + config.projector, + VisionProjector::Gemma3AvgPool { + tokens_per_side: 1, + kernel_size: 2, + image_token_id: 99, + pad_token_id: 0, + boi_token_id: 97, + eoi_token_id: 98, + newline_token_id: 108, + } + ); + assert!(config.fingerprint().contains("image_token_id: 99")); + let specs = config.weight_specs(); + assert_eq!( + specs[specs.len() - 2].name, + "multi_modal_projector.mm_soft_emb_norm.weight" + ); + assert_eq!(specs.last().unwrap().shape, [8, 12]); + } } #[derive(Debug, Clone, PartialEq)] @@ -188,6 +251,7 @@ pub(crate) struct LlavaVisionConfig { pub(crate) feature_layer: usize, pub(crate) drop_first_token: bool, pub(crate) text_hidden: usize, + pub(crate) projector: VisionProjector, } #[derive(Debug, Clone, PartialEq, Eq)] @@ -232,8 +296,14 @@ impl LlavaVisionConfig { 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("llava") { - return Err("IREE vision currently supports only model_type=llava".to_string()); + let model_type = root + .get("model_type") + .and_then(Value::as_str) + .ok_or_else(|| "config.json model_type is required".to_string())?; + if !matches!(model_type, "llava" | "gemma3") { + return Err(format!( + "IREE vision currently supports model_type=llava or gemma3, got {model_type:?}" + )); } let vision = object(&root, "vision_config")?; let vision_model_type = vision @@ -292,12 +362,14 @@ impl LlavaVisionConfig { )); } }; - match root.get("projector_hidden_act").and_then(Value::as_str) { - None | Some("gelu") => {} - other => { - return Err(format!( - "unsupported projector_hidden_act {other:?}; only exact GELU is qualified" - )); + if model_type == "llava" { + match root.get("projector_hidden_act").and_then(Value::as_str) { + None | Some("gelu") => {} + other => { + return Err(format!( + "unsupported projector_hidden_act {other:?}; only exact GELU is qualified" + )); + } } } let text_hidden = object(&root, "text_config")? @@ -311,10 +383,13 @@ impl LlavaVisionConfig { if text_hidden == 0 { return Err("text_config.hidden_size must be greater than zero".to_string()); } - let requested_layer = root - .get("vision_feature_layer") - .and_then(Value::as_i64) - .unwrap_or(-2); + let requested_layer = if model_type == "gemma3" { + -1 + } else { + root.get("vision_feature_layer") + .and_then(Value::as_i64) + .unwrap_or(-2) + }; let resolved = if requested_layer < 0 { i64::try_from(layers).map_err(|_| "vision layer count does not fit i64".to_string())? + requested_layer @@ -326,18 +401,69 @@ impl LlavaVisionConfig { "vision_feature_layer={requested_layer} resolves outside {layers} encoder layers" )); } - let strategy = root - .get("vision_feature_select_strategy") - .and_then(Value::as_str) - .unwrap_or("default"); - let drop_first_token = match strategy { - "default" => true, - "full" => false, - other => { + let drop_first_token = if model_type == "gemma3" { + false + } else { + match root + .get("vision_feature_select_strategy") + .and_then(Value::as_str) + .unwrap_or("default") + { + "default" => true, + "full" => false, + other => { + return Err(format!( + "unsupported vision_feature_select_strategy={other:?}" + )); + } + } + }; + let projector = if model_type == "gemma3" { + if class_token { + return Err("Gemma3 IREE vision requires the class-token-free SigLIP tower".into()); + } + let image_tokens = root + .get("mm_tokens_per_image") + .and_then(Value::as_u64) + .map_or(Ok(256usize), |value| { + usize::try_from(value) + .map_err(|_| "mm_tokens_per_image does not fit usize".to_string()) + })?; + let tokens_per_side = (image_tokens as f64).sqrt() as usize; + if tokens_per_side == 0 + || tokens_per_side + .checked_mul(tokens_per_side) + .ok_or_else(|| "Gemma3 image-token grid overflowed".to_string())? + != image_tokens + { return Err(format!( - "unsupported vision_feature_select_strategy={other:?}" + "Gemma3 mm_tokens_per_image={image_tokens} must be a non-zero square" )); } + let patch_grid = image_size / patch_size; + if image_size % patch_size != 0 || patch_grid % tokens_per_side != 0 { + return Err(format!( + "Gemma3 patch grid {patch_grid} must divide exactly into {tokens_per_side} pooled tokens per side" + )); + } + let token_id = |name: &str, default: i32| -> Result { + match root.get(name).and_then(Value::as_i64) { + Some(value) => i32::try_from(value) + .map_err(|_| format!("config.json {name} does not fit i32")), + None => Ok(default), + } + }; + VisionProjector::Gemma3AvgPool { + tokens_per_side, + kernel_size: patch_grid / tokens_per_side, + image_token_id: token_id("image_token_index", 262_144)?, + pad_token_id: token_id("pad_token_id", 0)?, + boi_token_id: token_id("boi_token_index", 255_999)?, + eoi_token_id: token_id("eoi_token_index", 256_000)?, + newline_token_id: 108, + } + } else { + VisionProjector::LlavaMlp }; Ok(Self { image_size, @@ -353,6 +479,7 @@ impl LlavaVisionConfig { feature_layer: resolved as usize, drop_first_token, text_hidden, + projector, }) } @@ -368,14 +495,19 @@ impl LlavaVisionConfig { #[must_use] pub(crate) fn image_tokens(&self) -> usize { - self.position_count() - usize::from(self.drop_first_token) + match self.projector { + VisionProjector::LlavaMlp => self.position_count() - usize::from(self.drop_first_token), + VisionProjector::Gemma3AvgPool { + tokens_per_side, .. + } => tokens_per_side * tokens_per_side, + } } #[must_use] pub(crate) fn fingerprint(&self) -> String { format!( - "llava-vision-v1:image={}:patch={}:channels={}:hidden={}:intermediate={}:layers={}:\ - heads={}:eps={:08x}:activation={:?}:class={}:feature={}:drop_first={}:text={}", + "iree-vision-v2:image={}:patch={}:channels={}:hidden={}:intermediate={}:layers={}:\ + heads={}:eps={:08x}:activation={:?}:class={}:feature={}:drop_first={}:text={}:projector={:?}", self.image_size, self.patch_size, self.channels, @@ -388,7 +520,8 @@ impl LlavaVisionConfig { self.class_token, self.feature_layer, self.drop_first_token, - self.text_hidden + self.text_hidden, + self.projector, ) } @@ -448,18 +581,30 @@ impl LlavaVisionConfig { }); } } - specs.extend([ - self.spec( - "multi_modal_projector.linear_1.weight", - [self.text_hidden, self.hidden], - ), - self.spec("multi_modal_projector.linear_1.bias", [self.text_hidden]), - self.spec( - "multi_modal_projector.linear_2.weight", - [self.text_hidden, self.text_hidden], - ), - self.spec("multi_modal_projector.linear_2.bias", [self.text_hidden]), - ]); + match self.projector { + VisionProjector::LlavaMlp => specs.extend([ + self.spec( + "multi_modal_projector.linear_1.weight", + [self.text_hidden, self.hidden], + ), + self.spec("multi_modal_projector.linear_1.bias", [self.text_hidden]), + self.spec( + "multi_modal_projector.linear_2.weight", + [self.text_hidden, self.text_hidden], + ), + self.spec("multi_modal_projector.linear_2.bias", [self.text_hidden]), + ]), + VisionProjector::Gemma3AvgPool { .. } => specs.extend([ + self.spec( + "multi_modal_projector.mm_soft_emb_norm.weight", + [self.hidden], + ), + self.spec( + "multi_modal_projector.mm_input_projection_weight", + [self.hidden, self.text_hidden], + ), + ]), + } specs } diff --git a/src/lib/mlxcel-xla/src/lib.rs b/src/lib/mlxcel-xla/src/lib.rs index 2b1328267..932d6e341 100644 --- a/src/lib/mlxcel-xla/src/lib.rs +++ b/src/lib/mlxcel-xla/src/lib.rs @@ -66,6 +66,7 @@ mod operator_numeric_contract; #[cfg_attr(not(feature = "iree"), allow(dead_code))] mod prepared; mod prepared_deepstack; +mod prepared_gemma3; mod prepared_gemma3n; #[cfg(feature = "iree")] @@ -210,6 +211,10 @@ pub use emitter::{Gemma3nDiagnosticLayout, Gemma3nDiagnosticSegment}; #[cfg(feature = "iree")] pub use prepared::PreparedInputError; pub use prepared_deepstack::{DeepStackFeatures, DeepStackInputError, DeepStackPreparedPrefill}; +pub use prepared_gemma3::{ + GEMMA3_VLM_MASK_MODE, GEMMA3_VLM_MASKED_VALUE, Gemma3VlmPreparedError, + prepare_gemma3_vlm_prefill, +}; pub use prepared_gemma3n::{Gemma3nDensePle, Gemma3nDensePleError, Gemma3nPreparedPrefill}; #[cfg(feature = "iree")] pub use sampler::SampleParams; diff --git a/src/lib/mlxcel-xla/src/prepared.rs b/src/lib/mlxcel-xla/src/prepared.rs index bc1b4eaad..67cbc382d 100644 --- a/src/lib/mlxcel-xla/src/prepared.rs +++ b/src/lib/mlxcel-xla/src/prepared.rs @@ -213,7 +213,7 @@ impl fmt::Display for PreparedInputError { "prepared {tensor} byte count mismatch: expected {expected}, got {actual}" ), Self::InvalidFloat { tensor } => { - write!(f, "prepared {tensor} contains NaN or an invalid positive mask value") + write!(f, "prepared {tensor} contains an invalid floating-point or mask value") } Self::ShapeOverflow => f.write_str("prepared IREE static shape overflowed"), } @@ -560,10 +560,30 @@ impl PreparedIreePrefill { bias_tensor.shape.clone(), )); } + let gemma3_vlm_mask = value + .modalities + .iter() + .any(|modality| modality.family == "gemma3"); + if gemma3_vlm_mask + && (value.attention_bias.causal + || !matrix_bias + || compact_bias + .iter() + .any(|value| *value != 0.0 && *value != f32::MIN)) + { + return Err(PreparedInputError::InvalidFloat { + tensor: "Gemma3 VLM attention bias", + }); + } + let masked_value = if gemma3_vlm_mask { + f32::MIN + } else { + MASKED_VALUE + }; let bias_count = context_capacity .checked_mul(context_capacity) .ok_or(PreparedInputError::ShapeOverflow)?; - let mut attention_bias = vec![MASKED_VALUE; bias_count]; + let mut attention_bias = vec![masked_value; bias_count]; for query in 0..effective_len { for key in 0..effective_len { let base = if key_bias { @@ -573,7 +593,7 @@ impl PreparedIreePrefill { }; attention_bias[query * context_capacity + key] = if value.attention_bias.causal && key > query { - MASKED_VALUE + masked_value } else { base }; diff --git a/src/lib/mlxcel-xla/src/prepared_gemma3.rs b/src/lib/mlxcel-xla/src/prepared_gemma3.rs new file mode 100644 index 000000000..306f00eca --- /dev/null +++ b/src/lib/mlxcel-xla/src/prepared_gemma3.rs @@ -0,0 +1,432 @@ +// 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. + +//! Canonical Gemma3 VLM merge at the XLA embeddings-prefill boundary. +//! +//! The ordinary Gemma3 token graph gathers embeddings and multiplies them by +//! `sqrt(hidden_size)`. The embeddings entry deliberately does neither, so this +//! producer exports text rows after that scale while keeping projected image +//! rows at their original magnitude. It also owns the reference's bidirectional +//! padding mask instead of substituting a causal or sliding-window mask. + +use std::fmt; + +use mlxcel_core::session::{ + OwnedTensor, PreparedAttentionBias, PreparedModality, PreparedPositions, PreparedPrefill, + PreparedPrefillError, PreparedTensorDType, +}; + +/// Exact masked value emitted by `vision::merge::prepare_inputs_for_multimodal`. +pub const GEMMA3_VLM_MASKED_VALUE: f32 = f32::MIN; + +/// Stable identity component for the qualified external-mask behavior. +pub const GEMMA3_VLM_MASK_MODE: &str = "gemma3-vlm-bidirectional-padding-f32-min-v1"; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum Gemma3VlmPreparedError { + Empty, + ZeroHiddenSize, + Capacity { + sequence_len: usize, + context_capacity: usize, + }, + ShapeOverflow, + EmbeddingCount { + expected: usize, + actual: usize, + }, + ProjectedShape { + values: usize, + hidden_size: usize, + }, + AttentionMaskCount { + expected: usize, + actual: usize, + }, + InvalidAttentionMask { + index: usize, + value: i32, + }, + PaddingMaskMismatch { + index: usize, + token_id: i32, + mask: i32, + }, + PlaceholderCount { + positions: usize, + projected_tokens: usize, + }, + ImageCount { + images: usize, + projected_tokens: usize, + }, + NonFinite { + tensor: &'static str, + index: usize, + }, + InvalidTokenConfiguration, + Prepared(PreparedPrefillError), +} + +impl fmt::Display for Gemma3VlmPreparedError { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Empty => formatter.write_str("Gemma3 VLM prefill requires at least one token"), + Self::ZeroHiddenSize => { + formatter.write_str("Gemma3 VLM hidden size must be greater than zero") + } + Self::Capacity { + sequence_len, + context_capacity, + } => write!( + formatter, + "Gemma3 VLM sequence length {sequence_len} exceeds context capacity {context_capacity}" + ), + Self::ShapeOverflow => formatter.write_str("Gemma3 VLM tensor shape overflowed"), + Self::EmbeddingCount { expected, actual } => write!( + formatter, + "Gemma3 VLM raw text embeddings have {actual} values, expected {expected}" + ), + Self::ProjectedShape { + values, + hidden_size, + } => write!( + formatter, + "Gemma3 VLM projected image features have {values} values, which is not divisible by hidden size {hidden_size}" + ), + Self::AttentionMaskCount { expected, actual } => write!( + formatter, + "Gemma3 VLM attention mask has {actual} values, expected {expected}" + ), + Self::InvalidAttentionMask { index, value } => write!( + formatter, + "Gemma3 VLM attention mask[{index}] is {value}; expected 0 or 1" + ), + Self::PaddingMaskMismatch { + index, + token_id, + mask, + } => write!( + formatter, + "Gemma3 VLM token {index} id {token_id} disagrees with padding mask {mask}" + ), + Self::PlaceholderCount { + positions, + projected_tokens, + } => write!( + formatter, + "Gemma3 VLM prompt has {positions} image positions but projector returned {projected_tokens} tokens" + ), + Self::ImageCount { + images, + projected_tokens, + } => write!( + formatter, + "Gemma3 VLM declares {images} image(s) but projector returned {projected_tokens} token(s)" + ), + Self::NonFinite { tensor, index } => write!( + formatter, + "Gemma3 VLM {tensor} contains a non-finite value at flat index {index}" + ), + Self::InvalidTokenConfiguration => { + formatter.write_str("Gemma3 VLM pad_token_id and image_token_id must be distinct") + } + Self::Prepared(error) => write!(formatter, "Gemma3 VLM prepared prefill: {error}"), + } + } +} + +impl std::error::Error for Gemma3VlmPreparedError {} + +impl From for Gemma3VlmPreparedError { + fn from(value: PreparedPrefillError) -> Self { + Self::Prepared(value) + } +} + +fn f32_bytes(values: &[f32]) -> Vec { + values + .iter() + .flat_map(|value| value.to_le_bytes()) + .collect() +} + +fn first_non_finite(values: &[f32]) -> Option { + values.iter().position(|value| !value.is_finite()) +} + +/// Build the exact post-scale embedding and additive-mask payload for one +/// Gemma3 VLM request. +#[allow(clippy::too_many_arguments)] +pub fn prepare_gemma3_vlm_prefill( + token_ids: Vec, + raw_text_embeddings: &[f32], + projected_image_features: &[f32], + attention_mask: &[i32], + hidden_size: usize, + context_capacity: usize, + pad_token_id: i32, + image_token_id: i32, + image_count: usize, +) -> Result { + let sequence_len = token_ids.len(); + if sequence_len == 0 { + return Err(Gemma3VlmPreparedError::Empty); + } + if hidden_size == 0 { + return Err(Gemma3VlmPreparedError::ZeroHiddenSize); + } + if sequence_len > context_capacity { + return Err(Gemma3VlmPreparedError::Capacity { + sequence_len, + context_capacity, + }); + } + if pad_token_id == image_token_id { + return Err(Gemma3VlmPreparedError::InvalidTokenConfiguration); + } + + let expected_embeddings = sequence_len + .checked_mul(hidden_size) + .ok_or(Gemma3VlmPreparedError::ShapeOverflow)?; + if raw_text_embeddings.len() != expected_embeddings { + return Err(Gemma3VlmPreparedError::EmbeddingCount { + expected: expected_embeddings, + actual: raw_text_embeddings.len(), + }); + } + if attention_mask.len() != sequence_len { + return Err(Gemma3VlmPreparedError::AttentionMaskCount { + expected: sequence_len, + actual: attention_mask.len(), + }); + } + if let Some(index) = first_non_finite(raw_text_embeddings) { + return Err(Gemma3VlmPreparedError::NonFinite { + tensor: "raw text embeddings", + index, + }); + } + if let Some(index) = first_non_finite(projected_image_features) { + return Err(Gemma3VlmPreparedError::NonFinite { + tensor: "projected image features", + index, + }); + } + if projected_image_features.len() % hidden_size != 0 { + return Err(Gemma3VlmPreparedError::ProjectedShape { + values: projected_image_features.len(), + hidden_size, + }); + } + + let mut image_positions = Vec::new(); + for (index, (&token_id, &mask)) in token_ids.iter().zip(attention_mask).enumerate() { + if mask != 0 && mask != 1 { + return Err(Gemma3VlmPreparedError::InvalidAttentionMask { index, value: mask }); + } + let is_padding = token_id == pad_token_id; + if is_padding != (mask == 0) { + return Err(Gemma3VlmPreparedError::PaddingMaskMismatch { + index, + token_id, + mask, + }); + } + if token_id == image_token_id { + image_positions.push(index); + } + } + let projected_tokens = projected_image_features.len() / hidden_size; + if image_positions.len() != projected_tokens { + return Err(Gemma3VlmPreparedError::PlaceholderCount { + positions: image_positions.len(), + projected_tokens, + }); + } + if (image_count == 0) != (projected_tokens == 0) + || (image_count > 0 && !projected_tokens.is_multiple_of(image_count)) + { + return Err(Gemma3VlmPreparedError::ImageCount { + images: image_count, + projected_tokens, + }); + } + + let normalizer = (hidden_size as f64).sqrt() as f32; + let mut merged = vec![0.0f32; expected_embeddings]; + let mut image_row = 0usize; + for (position, &token_id) in token_ids.iter().enumerate() { + let destination = position * hidden_size; + if token_id == pad_token_id { + continue; + } + if token_id == image_token_id { + let source = image_row * hidden_size; + merged[destination..destination + hidden_size] + .copy_from_slice(&projected_image_features[source..source + hidden_size]); + image_row += 1; + continue; + } + for offset in 0..hidden_size { + let value = raw_text_embeddings[destination + offset] * normalizer; + if !value.is_finite() { + return Err(Gemma3VlmPreparedError::NonFinite { + tensor: "post-scale text embeddings", + index: destination + offset, + }); + } + merged[destination + offset] = value; + } + } + + let bias_count = sequence_len + .checked_mul(sequence_len) + .ok_or(Gemma3VlmPreparedError::ShapeOverflow)?; + let mut attention_bias = vec![GEMMA3_VLM_MASKED_VALUE; bias_count]; + for query in 0..sequence_len { + for key in 0..sequence_len { + if attention_mask[query] == 1 && attention_mask[key] == 1 { + attention_bias[query * sequence_len + key] = 0.0; + } + } + } + + let embeddings = OwnedTensor::new( + f32_bytes(&merged), + PreparedTensorDType::Float32, + vec![1, sequence_len, hidden_size], + )?; + let attention_bias = PreparedAttentionBias { + tensor: OwnedTensor::new( + f32_bytes(&attention_bias), + PreparedTensorDType::Float32, + vec![1, 1, sequence_len, sequence_len], + )?, + causal: false, + }; + let modalities = (image_count > 0).then(|| PreparedModality { + family: "gemma3".to_string(), + item_count: image_count, + token_count: image_positions.len(), + }); + PreparedPrefill::new( + token_ids, + embeddings, + PreparedPositions::Sequential { + start: 0, + length: sequence_len, + }, + attention_bias, + modalities.into_iter().collect(), + ) + .map_err(Into::into) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn read_f32(tensor: &OwnedTensor) -> Vec { + tensor + .bytes + .chunks_exact(4) + .map(|chunk| f32::from_le_bytes(chunk.try_into().unwrap())) + .collect() + } + + #[test] + fn exports_post_scale_text_unscaled_images_and_exact_bidirectional_mask() { + let prepared = prepare_gemma3_vlm_prefill( + vec![0, 9, 7, 0], + &[1.0, 2.0, 100.0, 200.0, 3.0, 4.0, 5.0, 6.0], + &[11.0, 13.0], + &[0, 1, 1, 0], + 2, + 8, + 0, + 9, + 1, + ) + .unwrap(); + + let sqrt_two = 2.0f32.sqrt(); + assert_eq!( + read_f32(&prepared.embeddings), + [ + 0.0, + 0.0, + 11.0, + 13.0, + 3.0 * sqrt_two, + 4.0 * sqrt_two, + 0.0, + 0.0 + ] + ); + let bias = read_f32(&prepared.attention_bias.tensor); + assert_eq!(bias.len(), 16); + assert_eq!(bias[1 * 4 + 1], 0.0); + assert_eq!(bias[1 * 4 + 2], 0.0, "valid future keys stay bidirectional"); + assert_eq!(bias[2 * 4 + 1], 0.0); + assert_eq!(bias[0], f32::MIN); + assert_eq!(bias[3 * 4 + 2], f32::MIN); + assert!(!prepared.attention_bias.causal); + assert_eq!(prepared.modalities[0].token_count, 1); + + let static_payload = crate::prepared::PreparedIreePrefill::prepare(&prepared, 2, 8) + .expect("Gemma3 mask materializes into the static XLA bucket"); + assert_eq!(static_payload.attention_bias.len(), 64); + assert_eq!(static_payload.attention_bias[1 * 8 + 2], 0.0); + assert_eq!(static_payload.attention_bias[1 * 8 + 7], f32::MIN); + assert_eq!(static_payload.attention_bias[7 * 8 + 1], f32::MIN); + } + + #[test] + fn rejects_padding_disagreement_placeholder_mismatch_and_double_scale_overflow() { + assert!(matches!( + prepare_gemma3_vlm_prefill(vec![0], &[1.0], &[], &[1], 1, 1, 0, 9, 0), + Err(Gemma3VlmPreparedError::PaddingMaskMismatch { .. }) + )); + assert!(matches!( + prepare_gemma3_vlm_prefill(vec![9], &[1.0], &[], &[1], 1, 1, 0, 9, 1), + Err(Gemma3VlmPreparedError::PlaceholderCount { .. }) + )); + assert!(matches!( + prepare_gemma3_vlm_prefill( + vec![7], + &[f32::MAX, 0.0, 0.0, 0.0], + &[], + &[1], + 4, + 1, + 0, + 9, + 0, + ), + Err(Gemma3VlmPreparedError::NonFinite { + tensor: "post-scale text embeddings", + .. + }) + )); + assert!(matches!( + prepare_gemma3_vlm_prefill(vec![9], &[1.0, 2.0], &[3.0], &[1], 2, 1, 0, 9, 1,), + Err(Gemma3VlmPreparedError::ProjectedShape { .. }) + )); + assert!(matches!( + prepare_gemma3_vlm_prefill(vec![9], &[1.0], &[3.0], &[1], 1, 1, 0, 9, 0,), + Err(Gemma3VlmPreparedError::ImageCount { .. }) + )); + } +} diff --git a/src/lib/mlxcel-xla/src/vision_runtime.rs b/src/lib/mlxcel-xla/src/vision_runtime.rs index b16bc9d62..eef7a8b77 100644 --- a/src/lib/mlxcel-xla/src/vision_runtime.rs +++ b/src/lib/mlxcel-xla/src/vision_runtime.rs @@ -38,7 +38,7 @@ use crate::aux::{ use crate::aux_manifest::{AuxiliaryArtifactContract, ensure_qualified_auxiliary_artifact}; #[cfg(feature = "diagnostics")] use crate::emitter::emit_vision_diagnostics; -use crate::emitter::{LlavaVisionConfig, VisionWeightSpec, emit_vision}; +use crate::emitter::{LlavaVisionConfig, VisionProjector, VisionWeightSpec, emit_vision}; use crate::iree::{cached_vmfb_path, compile_one_to, iree_compile_bin, target_flags}; use crate::weights::{bf16_to_f32, f16_to_f32, f32_le_to_f32}; @@ -265,8 +265,30 @@ impl VisionProcessorContract { config.image_size )); } - if object.get("resample").and_then(serde_json::Value::as_u64) != Some(3) { - return Err("preprocessor_config.json resample must be PIL bicubic (3)".to_string()); + let expected_resample = match config.projector { + VisionProjector::LlavaMlp => 3, + VisionProjector::Gemma3AvgPool { .. } => 2, + }; + if matches!(config.projector, VisionProjector::Gemma3AvgPool { .. }) + && object + .get("image_processor_type") + .and_then(serde_json::Value::as_str) + != Some("Gemma3ImageProcessor") + { + return Err( + "preprocessor_config.json image_processor_type must be Gemma3ImageProcessor" + .to_string(), + ); + } + if object.get("resample").and_then(serde_json::Value::as_u64) != Some(expected_resample) { + return Err(format!( + "preprocessor_config.json resample must be {} ({expected_resample})", + if expected_resample == 3 { + "PIL bicubic" + } else { + "PIL bilinear" + } + )); } let factor = finite_number(object, "rescale_factor")?; if factor <= 0.0 { @@ -296,44 +318,66 @@ impl VisionProcessorContract { let processor_object = processor .as_object() .ok_or_else(|| format!("{} must be a JSON object", processor_path.display()))?; - if processor_object - .get("patch_size") - .and_then(serde_json::Value::as_u64) - != Some(config.patch_size as u64) - { - return Err(format!( - "processor patch_size disagrees with vision patch_size={}", - config.patch_size - )); - } - let expected_strategy = if config.drop_first_token { - "default" + if matches!(config.projector, VisionProjector::LlavaMlp) { + if processor_object + .get("patch_size") + .and_then(serde_json::Value::as_u64) + != Some(config.patch_size as u64) + { + return Err(format!( + "processor patch_size disagrees with vision patch_size={}", + config.patch_size + )); + } + let expected_strategy = if config.drop_first_token { + "default" + } else { + "full" + }; + if processor_object + .get("vision_feature_select_strategy") + .and_then(serde_json::Value::as_str) + != Some(expected_strategy) + { + return Err(format!( + "processor vision_feature_select_strategy disagrees with config ({expected_strategy})" + )); + } + let expected_additional_tokens = usize::from(config.class_token) as u64; + if processor_object + .get("num_additional_image_tokens") + .and_then(serde_json::Value::as_u64) + != Some(expected_additional_tokens) + { + return Err(format!( + "processor num_additional_image_tokens disagrees with class-token contract ({expected_additional_tokens})" + )); + } } else { - "full" - }; - if processor_object - .get("vision_feature_select_strategy") - .and_then(serde_json::Value::as_str) - != Some(expected_strategy) - { - return Err(format!( - "processor vision_feature_select_strategy disagrees with config ({expected_strategy})" - )); - } - let expected_additional_tokens = usize::from(config.class_token) as u64; - if processor_object - .get("num_additional_image_tokens") - .and_then(serde_json::Value::as_u64) - != Some(expected_additional_tokens) - { - return Err(format!( - "processor num_additional_image_tokens disagrees with class-token contract ({expected_additional_tokens})" - )); + if processor_object + .get("processor_class") + .and_then(serde_json::Value::as_str) + != Some("Gemma3Processor") + { + return Err( + "processor_config.json processor_class must be Gemma3Processor".to_string(), + ); + } + if processor_object + .get("image_seq_length") + .and_then(serde_json::Value::as_u64) + != Some(config.image_tokens() as u64) + { + return Err(format!( + "processor image_seq_length disagrees with Gemma3 projector output {}", + config.image_tokens() + )); + } } Ok(Self { identity: format!( - "preprocessor={};processor={};resolved=size:{};crop:{};resample:bicubic;\ - rescale:{:016x};mean:{mean:?};std:{std:?}", + "preprocessor={};processor={};resolved=size:{};crop:{};\ + resample:{expected_resample};rescale:{:016x};mean:{mean:?};std:{std:?}", preprocessor, processor, config.image_size, @@ -663,7 +707,11 @@ impl IreeVisionProjector { pub fn load(model_dir: &Path, device: &str) -> Result { let config = LlavaVisionConfig::from_model_dir(model_dir)?; let mlir = emit_vision(&config); - let module = compile_and_load(model_dir, device, &config, &mlir, "llava-vision")?; + let cache_tag = match config.projector { + VisionProjector::LlavaMlp => "llava-vision", + VisionProjector::Gemma3AvgPool { .. } => "gemma3-vision", + }; + let module = compile_and_load(model_dir, device, &config, &mlir, cache_tag)?; Ok(Self { module, config }) } diff --git a/src/loading/mod.rs b/src/loading/mod.rs index fb2cab9d4..d3a3c5cf5 100644 --- a/src/loading/mod.rs +++ b/src/loading/mod.rs @@ -56,6 +56,8 @@ use self::special::try_load_special_model_from_weights; use self::vlm::*; // Re-exported at the crate root for the CLI's lazy `--output-audio` load. +#[cfg(feature = "xla-iree")] +pub(crate) use self::vlm::load_gemma3_iree_host_preprocessor; pub(crate) use self::vlm::load_llava_host_preprocessor; #[cfg(feature = "xla-iree")] pub(crate) use self::vlm::load_llava_iree_host_preprocessor; diff --git a/src/loading/vlm.rs b/src/loading/vlm.rs index b87c6c0a2..e2f41894b 100644 --- a/src/loading/vlm.rs +++ b/src/loading/vlm.rs @@ -97,6 +97,8 @@ pub(crate) use deepseekocr::{ pub(crate) use dots_ocr::load_dots_ocr_vl; pub(crate) use ernie4_5_vl::load_ernie4_5_moe_vlm; pub(crate) use fastvlm::load_fastvlm_vlm; +#[cfg(feature = "xla-iree")] +pub(crate) use gemma::load_gemma3_iree_host_preprocessor; pub(crate) use gemma::{load_gemma3_vlm, load_gemma3n_vlm, load_gemma4_vlm}; pub(crate) use gemma_unified::load_gemma4_unified; pub(crate) use granite_vision::load_granite_vision_vlm; diff --git a/src/loading/vlm_gemma.rs b/src/loading/vlm_gemma.rs index 91b3528f4..7d94ccb14 100644 --- a/src/loading/vlm_gemma.rs +++ b/src/loading/vlm_gemma.rs @@ -22,17 +22,22 @@ //! wrapper assembly out of the generic VLM router. use anyhow::Result; +#[cfg(feature = "xla-iree")] +use mlxcel_core::layers::UnifiedEmbedding; use mlxcel_core::weights::WeightMap; use serde_json::Value; use std::path::Path; use crate::LoadedModel; use crate::models; +#[cfg(feature = "xla-iree")] +use crate::multimodal::host_preprocessor::{Gemma3IreeHostPreprocessor, HostPreprocessorError}; use crate::vision; use super::{ - load_vlm_weights_common, parse_required_vlm_subconfig, parse_vlm_config, - read_optional_model_json, read_sanitized_vlm_config, strip_language_model_prefix, + load_vlm_weights_common, load_vlm_weights_common_filtered_canonical, + parse_required_vlm_subconfig, parse_vlm_config, read_optional_model_json, + read_sanitized_vlm_config, strip_language_model_prefix, }; struct Gemma3nMetadata { @@ -224,6 +229,129 @@ pub(crate) fn load_gemma3_vlm(model_path: &Path) -> Result { Ok(LoadedModel::Gemma3VLM(vlm)) } +#[cfg(feature = "xla-iree")] +fn is_gemma3_text_embedding_weight(name: &str) -> bool { + name.strip_prefix("language_model.") + .unwrap_or(name) + .starts_with("model.embed_tokens.") +} + +/// Load the bounded Gemma3 host producer paired with the resident IREE +/// SigLIP/averaging-projector module. Only the text embedding table is retained +/// in MLX; no language layer or host vision tower is constructed. +#[cfg(feature = "xla-iree")] +pub(crate) fn load_gemma3_iree_host_preprocessor( + model_path: &Path, + device: &str, +) -> Result { + use vision::config::VLMConfig; + use vision::processors::siglip::SigLipProcessor; + + let (_config_str, full_config) = read_sanitized_vlm_config(model_path) + .map_err(|error| HostPreprocessorError::InvalidConfig(error.to_string()))?; + if !full_config + .get("vision_config") + .is_some_and(Value::is_object) + || !full_config.get("text_config").is_some_and(Value::is_object) + { + return Err(HostPreprocessorError::FamilyMismatch { + actual: "unqualified gemma3 config without vision_config/text_config".to_string(), + }); + } + let vlm_config: VLMConfig = serde_json::from_value(full_config.clone()) + .map_err(|error| HostPreprocessorError::InvalidConfig(error.to_string()))?; + if vlm_config.model_type != "gemma3" + || vlm_config + .text_config + .get("model_type") + .and_then(Value::as_str) + != Some("gemma3_text") + { + return Err(HostPreprocessorError::FamilyMismatch { + actual: format!( + "model_type={:?}, text_model_type={:?}", + vlm_config.model_type, + vlm_config + .text_config + .get("model_type") + .and_then(Value::as_str) + ), + }); + } + let text_config: models::gemma3::ModelArgs = + serde_json::from_value(vlm_config.text_config.clone()) + .map_err(|error| HostPreprocessorError::InvalidConfig(error.to_string()))?; + + let weights = load_vlm_weights_common_filtered_canonical(model_path, |name| { + is_gemma3_text_embedding_weight(name) + }) + .map(strip_language_model_prefix) + .map_err(|error| HostPreprocessorError::WeightLoad(error.to_string()))?; + let quant_group_size = full_config + .get("quantization") + .and_then(|value| value.get("group_size")) + .and_then(Value::as_i64) + .unwrap_or(64) as i32; + let quant_bits = full_config + .get("quantization") + .and_then(|value| value.get("bits")) + .and_then(Value::as_i64) + .unwrap_or(4) 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 Gemma3 text embedding table: {error}" + )) + })?; + let tokens_per_image = vlm_config.get_mm_tokens_per_image(); + let processor = SigLipProcessor::new(vlm_config.vision_config.image_size); + let projector = mlxcel_xla::IreeVisionProjector::load(model_path, device) + .map_err(HostPreprocessorError::Iree)?; + + Gemma3IreeHostPreprocessor::from_parts( + Box::new(processor), + text_embeddings, + projector, + vlm_config.image_token_index, + vlm_config.pad_token_id, + vlm_config.boi_token_index, + vlm_config.eoi_token_index, + tokens_per_image, + text_config.hidden_size, + vlm_config.vision_config.image_size, + text_config.max_position_embeddings, + device.to_string(), + ) +} + +#[cfg(all(test, feature = "xla-iree"))] +mod xla_tests { + use super::*; + + #[test] + fn pinned_gemma3_projector_loads_and_returns_finite_features() { + let Ok(model) = std::env::var("MLXCEL_GEMMA3_FIXTURE") else { + return; + }; + let device = std::env::var("MLXCEL_XLA_DEVICE").unwrap_or_else(|_| "cuda".to_string()); + let mut projector = mlxcel_xla::IreeVisionProjector::load(Path::new(&model), &device) + .expect("load pinned Gemma3 IREE vision projector"); + assert_eq!(projector.input_shape(), [1, 3, 896, 896]); + assert_eq!(projector.output_shape(), [256, 2560]); + let pixels = vec![0.0; projector.input_shape().into_iter().product()]; + let projection = projector + .project(&pixels) + .expect("execute pinned Gemma3 IREE vision projector"); + assert_eq!(projection.shape, [256, 2560]); + assert!(projection.values.iter().all(|value| value.is_finite())); + } +} + /// Load a Gemma3n VLM model. pub(crate) fn load_gemma3n_vlm(model_path: &Path) -> Result { use vision::encoders::gemma3n::load_gemma3n_vision; diff --git a/src/loading/vlm_qwen.rs b/src/loading/vlm_qwen.rs index 52971240c..1902a3b49 100644 --- a/src/loading/vlm_qwen.rs +++ b/src/loading/vlm_qwen.rs @@ -124,6 +124,14 @@ pub(crate) fn load_qwen2_vl_iree_host_preprocessor( .to_string(), }); } + if !full_config + .get("vision_config") + .is_some_and(serde_json::Value::is_object) + { + return Err(HostPreprocessorError::FamilyMismatch { + actual: "unqualified qwen2_vl config without vision_config".to_string(), + }); + } let mut vision_config: Qwen2VLVisionConfig = parse_required_vlm_subconfig(&full_config, "vision_config", "Qwen2VL vision config") .map_err(|error| HostPreprocessorError::InvalidConfig(error.to_string()))?; diff --git a/src/multimodal/host_preprocessor.rs b/src/multimodal/host_preprocessor.rs index bfad9cb9d..7f21e3444 100644 --- a/src/multimodal/host_preprocessor.rs +++ b/src/multimodal/host_preprocessor.rs @@ -130,13 +130,13 @@ pub trait HostMultimodalPreprocessor { /// /// `Ok(None)` is the conservative result for text-only checkpoints and VLM /// families whose processor/position contract has not been qualified for XLA -/// yet. Once a checkpoint is identified as the supported LLaVA family, missing -/// or malformed processor/projector weights are startup errors rather than a -/// capability downgrade. +/// yet. Once a checkpoint is identified as a supported multimodal family, +/// missing or malformed processor/projector weights are startup errors rather +/// than a capability downgrade. /// /// # Errors /// -/// Returns a typed configuration or weight-loading error for a supported LLaVA +/// Returns a typed configuration or weight-loading error for a supported /// checkpoint that cannot construct its complete host preprocessor. pub fn load_xla_image_preprocessor( model_path: &Path, @@ -147,8 +147,8 @@ pub fn load_xla_image_preprocessor( model_path.display() )) })?; + let policy = XlaVisionBackendPolicy::from_env()?; if model_type == crate::models::ModelType::Qwen2VL { - let policy = XlaVisionBackendPolicy::from_env()?; if policy == XlaVisionBackendPolicy::Host { return Err(HostPreprocessorError::InvalidConfig( "Qwen2-VL XLA vision has no MLX fallback; MLXCEL_XLA_VISION_BACKEND=host is unsupported" @@ -159,7 +159,11 @@ pub fn load_xla_image_preprocessor( { let device = std::env::var("MLXCEL_XLA_DEVICE") .unwrap_or_else(|_| mlxcel_xla::default_device().to_string()); - let preprocessor = Qwen2VlIreeHostPreprocessor::load(model_path, &device)?; + let preprocessor = match Qwen2VlIreeHostPreprocessor::load(model_path, &device) { + Ok(preprocessor) => preprocessor, + Err(HostPreprocessorError::FamilyMismatch { .. }) => return Ok(None), + Err(error) => return Err(error), + }; tracing::info!( vision_backend = "iree", vision_device = %device, @@ -175,14 +179,54 @@ pub fn load_xla_image_preprocessor( )); } } + if model_type == crate::models::ModelType::Gemma3VLM { + return load_gemma3_image_preprocessor(model_path, policy); + } if model_type != crate::models::ModelType::LlavaVLM { return Ok(None); } - let policy = XlaVisionBackendPolicy::from_env()?; load_llava_image_preprocessor(model_path, policy) } +fn load_gemma3_image_preprocessor( + model_path: &Path, + policy: XlaVisionBackendPolicy, +) -> Result>, HostPreprocessorError> { + if policy == XlaVisionBackendPolicy::Host { + return Err(HostPreprocessorError::InvalidConfig( + "Gemma3 VLM requires MLXCEL_XLA_VISION_BACKEND=iree; the host LLaVA merge does not implement Gemma3 scaling or additive-mask semantics" + .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 = match Gemma3IreeHostPreprocessor::load(model_path, &device) { + Ok(preprocessor) => preprocessor, + Err(HostPreprocessorError::FamilyMismatch { .. }) => return Ok(None), + Err(error) => return Err(error), + }; + tracing::info!( + vision_backend = "iree", + vision_device = %device, + vision_backend_policy = ?policy, + "OpenXLA Gemma3 multimodal vision backend selected" + ); + return Ok(Some(Box::new(preprocessor))); + } + + #[cfg(not(feature = "xla-iree"))] + { + let _ = model_path; + Err(HostPreprocessorError::InvalidConfig( + "Gemma3 VLM image input requires the xla-iree feature".to_string(), + )) + } +} + fn load_llava_host_preprocessor_boxed( model_path: &Path, ) -> Result>, HostPreprocessorError> { @@ -458,6 +502,251 @@ impl HostMultimodalPreprocessor for Qwen2VlIreeHostPreprocessor { } } +/// Gemma3 producer that retains only host image processing and the filtered +/// text embedding table. SigLIP and the averaging projector execute in IREE; +/// this owner then constructs Gemma3's post-scale embeddings and exact +/// bidirectional additive prefill mask. +#[cfg(feature = "xla-iree")] +pub struct Gemma3IreeHostPreprocessor { + processor: Box, + text_embeddings: UnifiedEmbedding, + projector: RefCell, + image_token_id: i32, + pad_token_id: i32, + boi_token_id: i32, + eoi_token_id: i32, + tokens_per_image: usize, + hidden_size: usize, + image_size: usize, + max_sequence_len: usize, + device: String, +} + +#[cfg(feature = "xla-iree")] +impl Gemma3IreeHostPreprocessor { + pub fn load(model_path: &Path, device: &str) -> Result { + crate::loading::load_gemma3_iree_host_preprocessor(model_path, device) + } + + #[allow(clippy::too_many_arguments)] + pub(crate) fn from_parts( + processor: Box, + text_embeddings: UnifiedEmbedding, + projector: mlxcel_xla::IreeVisionProjector, + image_token_id: i32, + pad_token_id: i32, + boi_token_id: i32, + eoi_token_id: i32, + tokens_per_image: usize, + hidden_size: usize, + image_size: usize, + max_sequence_len: usize, + device: String, + ) -> Result { + validate_image_preprocessor_dimensions( + tokens_per_image, + hidden_size, + image_size, + max_sequence_len, + )?; + if [image_token_id, pad_token_id, boi_token_id, eoi_token_id] + .windows(2) + .any(|pair| pair[0] == pair[1]) + || image_token_id == boi_token_id + || image_token_id == eoi_token_id + || pad_token_id == eoi_token_id + { + return Err(HostPreprocessorError::InvalidConfig( + "Gemma3 image, padding, BOI, and EOI token IDs must be distinct".to_string(), + )); + } + let expected_input = [1, 3, image_size, image_size]; + if projector.input_shape() != expected_input { + return Err(HostPreprocessorError::InvalidConfig(format!( + "Gemma3 IREE vision input shape {:?} does not match processor shape {expected_input:?}", + projector.input_shape() + ))); + } + let expected_output = [tokens_per_image, hidden_size]; + if projector.output_shape() != expected_output { + return Err(HostPreprocessorError::InvalidConfig(format!( + "Gemma3 IREE vision output shape {:?} does not match prepared-prefill shape {expected_output:?}", + projector.output_shape() + ))); + } + Ok(Self { + processor, + text_embeddings, + projector: RefCell::new(projector), + image_token_id, + pad_token_id, + boi_token_id, + eoi_token_id, + tokens_per_image, + hidden_size, + image_size, + max_sequence_len, + device, + }) + } + + fn token_block_info(&self) -> ImageTokenBlockInfo { + ImageTokenBlockInfo { + use_boi_eoi: true, + image_token_id: self.image_token_id, + mm_tokens_per_image: self.tokens_per_image, + boi_token_id: self.boi_token_id, + eoi_token_id: self.eoi_token_id, + has_bos: true, + separator_token_id: None, + suffix_tokens: Vec::new(), + block_prefix_tokens: vec![108], + block_suffix_tokens: vec![108], + } + } + + fn prepare_iree( + &self, + token_ids: &[i32], + images: &[DynamicImage], + ) -> Result { + let mut logical_tokens = token_ids.to_vec(); + apply_image_token_blocks(&mut logical_tokens, self.token_block_info(), images.len())?; + validate_sequence_capacity(logical_tokens.len(), self.max_sequence_len)?; + let attention_mask = logical_tokens + .iter() + .map(|token| i32::from(*token != self.pad_token_id)) + .collect::>(); + 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, + "Gemma3 text embedding table", + )?; + let raw_text_embeddings = export_mlx_tensor(&text_embeddings, "Gemma3 text embeddings")?; + let raw_text_embeddings = raw_text_embeddings + .bytes + .chunks_exact(std::mem::size_of::()) + .map(|bytes| f32::from_ne_bytes(bytes.try_into().expect("four-byte f32 chunk"))) + .collect::>(); + + let mut projected_values = Vec::new(); + if !images.is_empty() { + let pixels = self.processor.preprocess(images); + validate_processor_shape( + &mlxcel_core::array_shape(&pixels), + images.len(), + self.image_size, + )?; + let pixels = export_mlx_tensor(&pixels, "Gemma3 processor pixel_values")?; + if pixels.dtype != PreparedTensorDType::Float32 { + return Err(HostPreprocessorError::InvalidConfig(format!( + "Gemma3 IREE vision requires float32 processor output, got {:?}", + pixels.dtype + ))); + } + let pixel_values = pixels + .bytes + .chunks_exact(std::mem::size_of::()) + .map(|bytes| f32::from_ne_bytes(bytes.try_into().expect("four-byte f32 chunk"))) + .collect::>(); + let pixels_per_image = 3usize + .checked_mul(self.image_size) + .and_then(|count| count.checked_mul(self.image_size)) + .ok_or(HostPreprocessorError::ShapeOverflow)?; + let projected_per_image = self + .tokens_per_image + .checked_mul(self.hidden_size) + .ok_or(HostPreprocessorError::ShapeOverflow)?; + projected_values.reserve( + images + .len() + .checked_mul(projected_per_image) + .ok_or(HostPreprocessorError::ShapeOverflow)?, + ); + let mut elapsed_seconds = 0.0; + let mut upload_bytes = 0usize; + let mut transfer_bytes = 0usize; + let mut projector = self.projector.try_borrow_mut().map_err(|_| { + HostPreprocessorError::Iree( + "concurrent/re-entrant Gemma3 IREE vision invocation is unsupported" + .to_string(), + ) + })?; + for image_pixels in pixel_values.chunks_exact(pixels_per_image) { + let projection = projector + .project(image_pixels) + .map_err(HostPreprocessorError::Iree)?; + if projection.shape != [self.tokens_per_image, self.hidden_size] { + return Err(HostPreprocessorError::ProjectedShape { + actual: projection + .shape + .into_iter() + .map(|dimension| i32::try_from(dimension).unwrap_or(i32::MAX)) + .collect(), + image_count: 1, + tokens_per_image: self.tokens_per_image, + hidden_size: self.hidden_size, + }); + } + elapsed_seconds += projection.metrics.elapsed_seconds; + upload_bytes = upload_bytes + .checked_add(projection.metrics.pixel_upload_bytes) + .ok_or(HostPreprocessorError::ShapeOverflow)?; + transfer_bytes = transfer_bytes + .checked_add(projection.metrics.projected_transfer_bytes) + .ok_or(HostPreprocessorError::ShapeOverflow)?; + projected_values.extend(projection.values); + } + tracing::info!( + vision_backend = "iree", + vision_device = %self.device, + image_count = images.len(), + pixel_upload_bytes = upload_bytes, + projected_transfer_bytes = transfer_bytes, + iree_vision_seconds = elapsed_seconds, + "OpenXLA Gemma3 vision projection completed" + ); + } + + mlxcel_xla::prepare_gemma3_vlm_prefill( + logical_tokens, + &raw_text_embeddings, + &projected_values, + &attention_mask, + self.hidden_size, + self.max_sequence_len, + self.pad_token_id, + self.image_token_id, + images.len(), + ) + .map_err(|error| HostPreprocessorError::Gemma3(error.to_string())) + } +} + +#[cfg(feature = "xla-iree")] +impl HostMultimodalPreprocessor for Gemma3IreeHostPreprocessor { + fn backend(&self) -> XlaVisionBackend { + XlaVisionBackend::Iree + } + + fn prepare( + &self, + token_ids: &[i32], + images: &[DynamicImage], + ) -> Result { + self.prepare_iree(token_ids, images) + } +} + #[cfg(feature = "xla-iree")] impl LlavaIreeHostPreprocessor { /// Load the host processor/text embedding table and resident IREE vision @@ -478,7 +767,7 @@ impl LlavaIreeHostPreprocessor { max_sequence_len: usize, device: String, ) -> Result { - validate_llava_preprocessor_dimensions( + validate_image_preprocessor_dimensions( tokens_per_image, hidden_size, image_size, @@ -740,7 +1029,7 @@ impl LlavaHostPreprocessor { image_size: usize, max_sequence_len: usize, ) -> Result { - validate_llava_preprocessor_dimensions( + validate_image_preprocessor_dimensions( tokens_per_image, hidden_size, image_size, @@ -939,7 +1228,7 @@ impl HostMultimodalPreprocessor for LlavaHostPreprocessor { } } -fn validate_llava_preprocessor_dimensions( +fn validate_image_preprocessor_dimensions( tokens_per_image: usize, hidden_size: usize, image_size: usize, @@ -947,22 +1236,22 @@ fn validate_llava_preprocessor_dimensions( ) -> Result<(), HostPreprocessorError> { if tokens_per_image == 0 { return Err(HostPreprocessorError::InvalidConfig( - "LLaVA mm_tokens_per_image must be greater than zero".to_string(), + "multimodal tokens_per_image must be greater than zero".to_string(), )); } if hidden_size == 0 { return Err(HostPreprocessorError::InvalidConfig( - "LLaVA text hidden_size must be greater than zero".to_string(), + "multimodal text hidden_size must be greater than zero".to_string(), )); } if image_size == 0 { return Err(HostPreprocessorError::InvalidConfig( - "LLaVA processor image_size must be greater than zero".to_string(), + "multimodal processor image_size must be greater than zero".to_string(), )); } if max_sequence_len == 0 { return Err(HostPreprocessorError::InvalidConfig( - "LLaVA max sequence length must be greater than zero".to_string(), + "multimodal max sequence length must be greater than zero".to_string(), )); } Ok(()) @@ -1065,14 +1354,16 @@ impl HostMultimodalPreprocessor for FakeHostMultimodalPreprocessor { pub enum HostPreprocessorError { #[error(transparent)] Placeholder(#[from] ImageTokenBlockError), - #[error("incompatible multimodal family: expected LLaVA, got {actual}")] + #[error("incompatible multimodal family: {actual}")] FamilyMismatch { actual: String }, - #[error("invalid LLaVA host-preprocessor config: {0}")] + #[error("invalid multimodal host-preprocessor config: {0}")] InvalidConfig(String), - #[error("failed to load LLaVA host-preprocessor weights: {0}")] + #[error("failed to load multimodal host-preprocessor weights: {0}")] WeightLoad(String), #[error("IREE vision backend failed: {0}")] Iree(String), + #[error("Gemma3 multimodal prefill failed: {0}")] + Gemma3(String), #[error( "processor output shape {actual:?} does not match decoded RGB batch [{image_count}, 3, {image_size}, {image_size}]" )] diff --git a/src/multimodal/host_preprocessor_tests.rs b/src/multimodal/host_preprocessor_tests.rs index a83e0f032..27a6648b6 100644 --- a/src/multimodal/host_preprocessor_tests.rs +++ b/src/multimodal/host_preprocessor_tests.rs @@ -67,7 +67,7 @@ fn iree_vision_contract_policy_is_explicit_and_strict() { #[test] fn xla_loader_keeps_text_and_unqualified_vlm_image_capability_false() { - for model_type in ["llama", "qwen2_vl"] { + for model_type in ["llama", "qwen2_vl", "gemma3"] { let model_dir = tempfile::tempdir().unwrap(); std::fs::write( model_dir.path().join("config.json"), From 2d5394ad1ce1d6a0b35ce3e5371d58f04174d2c0 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 24 Jul 2026 20:17:20 +0900 Subject: [PATCH 02/11] fix: preserve Gemma3 image-block wrapping Apply configured block prefix and suffix tokens when inserting image blocks into prompts without an existing placeholder, so Gemma3 bare-prompt CLI and server requests match the processor's newline, BOI, image-token, EOI, newline sequence while families with empty wrappers remain unchanged. Include wrapper lengths in checked capacity arithmetic and cover both the Gemma3 insertion sequence and the existing unwrapped insertion behavior. Refs #869 --- src/multimodal/vlm_prompt.rs | 4 ++++ src/multimodal/vlm_prompt_tests.rs | 31 ++++++++++++++++++++++++++++++ 2 files changed, 35 insertions(+) diff --git a/src/multimodal/vlm_prompt.rs b/src/multimodal/vlm_prompt.rs index 84f5ecf15..247864499 100644 --- a/src/multimodal/vlm_prompt.rs +++ b/src/multimodal/vlm_prompt.rs @@ -168,6 +168,8 @@ pub fn apply_image_token_blocks( let block_len = info .mm_tokens_per_image .checked_add(wrapper_tokens) + .and_then(|tokens| tokens.checked_add(info.block_prefix_tokens.len())) + .and_then(|tokens| tokens.checked_add(info.block_suffix_tokens.len())) .ok_or(ImageTokenBlockError::CapacityOverflow)?; let image_token_capacity = block_len .checked_mul(num_images) @@ -182,6 +184,7 @@ pub fn apply_image_token_blocks( let mut image_tokens = Vec::with_capacity(image_token_capacity); for _ in 0..num_images { + image_tokens.extend_from_slice(&info.block_prefix_tokens); if info.use_boi_eoi { image_tokens.push(info.boi_token_id); } @@ -191,6 +194,7 @@ pub fn apply_image_token_blocks( if info.use_boi_eoi { image_tokens.push(info.eoi_token_id); } + image_tokens.extend_from_slice(&info.block_suffix_tokens); } if info.has_bos { diff --git a/src/multimodal/vlm_prompt_tests.rs b/src/multimodal/vlm_prompt_tests.rs index 65eb67c97..5434c2adc 100644 --- a/src/multimodal/vlm_prompt_tests.rs +++ b/src/multimodal/vlm_prompt_tests.rs @@ -171,6 +171,37 @@ fn apply_image_token_blocks_expands_with_block_prefix_and_suffix() { assert_eq!(prompt_tokens, vec![1, 108, 10, 99, 99, 99, 11, 108, 2]); } +#[test] +fn apply_image_token_blocks_inserts_with_block_prefix_and_suffix() { + // Gemma3Processor expands a bare image into + // "\n\n" + BOI + image tokens + EOI + "\n\n" even when the caller did not + // provide an existing BOI placeholder. + let info = ImageTokenBlockInfo { + use_boi_eoi: true, + image_token_id: 99, + mm_tokens_per_image: 3, + boi_token_id: 10, + eoi_token_id: 11, + has_bos: true, + separator_token_id: None, + suffix_tokens: Vec::new(), + block_prefix_tokens: vec![108], + block_suffix_tokens: vec![108], + }; + let mut prompt_tokens = vec![1, 2]; + + let stats = apply_image_token_blocks(&mut prompt_tokens, info, 1).unwrap(); + + assert_eq!( + stats, + Some(ImageTokenBlockStats { + action: ImageTokenBlockAction::Inserted { image_blocks: 1 }, + tokens_per_image: 3, + }) + ); + assert_eq!(prompt_tokens, vec![1, 108, 10, 99, 99, 99, 11, 108, 2]); +} + #[test] fn apply_image_token_blocks_rejects_media_cardinality_mismatch() { let info = ImageTokenBlockInfo { From b861094cbf865f952bd8e1aa19bd3c2f5ae61a19 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 24 Jul 2026 20:22:02 +0900 Subject: [PATCH 03/11] fix: preserve root Gemma3 text configs Treat root Gemma3 configurations as nested VLM wrappers only when they also declare a vision_config object. This preserves the existing text-only Gemma3 parser path while retaining nested text defaults and wrapper quantization for actual VLM checkpoints. Add direct-root and nested-wrapper regressions to keep the two configuration shapes distinct. Refs #869 --- src/lib/mlxcel-xla/src/emitter/config.rs | 29 ++++++++++++++++++++++-- 1 file changed, 27 insertions(+), 2 deletions(-) diff --git a/src/lib/mlxcel-xla/src/emitter/config.rs b/src/lib/mlxcel-xla/src/emitter/config.rs index d5b5626f9..08dbd8878 100644 --- a/src/lib/mlxcel-xla/src/emitter/config.rs +++ b/src/lib/mlxcel-xla/src/emitter/config.rs @@ -488,7 +488,13 @@ 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" | "gemma3")) { + let is_gemma3_vlm_wrapper = wrapper_model_type == Some("gemma3") + && root + .get("vision_config") + .is_some_and(serde_json::Value::is_object); + let v = if matches!(wrapper_model_type, Some("llava" | "llava_next")) + || is_gemma3_vlm_wrapper + { let mut text = root .get("text_config") .and_then(serde_json::Value::as_object) @@ -498,7 +504,7 @@ impl Config { "{wrapper_model_type:?} config.json missing object `text_config` for the XLA text graph" ) })?; - if wrapper_model_type == Some("gemma3") { + if is_gemma3_vlm_wrapper { // mlx-vlm's Gemma3 TextConfig supplies these architecture // defaults. mlx-community conversions commonly omit them from // the nested object, so resolve the same explicit contract here @@ -1603,6 +1609,7 @@ mod tests { fn gemma3_wrapper_uses_nested_text_config_and_wrapper_quantization() { let config = Config::from_json_str( r#"{"model_type":"gemma3","quantization":{"bits":4,"group_size":64}, + "vision_config":{}, "text_config":{"model_type":"gemma3_text","hidden_size":2560, "intermediate_size":10240,"num_hidden_layers":34,"sliding_window":1024, "rope_scaling":{"factor":8.0,"rope_type":"linear"}}}"#, @@ -1617,6 +1624,24 @@ mod tests { assert_eq!(config.rope, RopeScaling::Plain); } + #[test] + fn gemma3_root_text_config_remains_direct_without_vision() { + let config = Config::from_json_str( + r#"{"model_type":"gemma3","hidden_size":8,"num_attention_heads":2, + "num_key_value_heads":1,"head_dim":4,"intermediate_size":16, + "num_hidden_layers":4,"rms_norm_eps":1e-6,"rope_theta":1000000, + "rope_local_base_freq":10000,"sliding_window":2, + "sliding_window_pattern":3,"vocab_size":12, + "hidden_activation":"gelu_pytorch_tanh"}"#, + ) + .unwrap(); + assert_eq!(config.hidden, 8); + assert_eq!(config.n_q, 2); + assert_eq!(config.n_kv, 1); + assert_eq!(config.head_dim, 4); + assert!(config.embeddings_prefill_uses_authoritative_mask); + } + /// ERNIE-4.5 is rejected with a message naming its interleaved (GPT-J-style) /// RoPE: it looks like a plain-RoPE Llama in config.json but its `rotate_half` /// rotates the (2i, 2i+1) pairs, not the (i, i+d/2) halves the Llama emit uses, From 36ba665d919fb8becff037c9d0daaca39e8254dd Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Sun, 26 Jul 2026 21:36:14 +0900 Subject: [PATCH 04/11] test(xla): add Gemma3 VLM mask oracle Compare the authoritative Gemma3 VLM additive mask and post-scale embedding boundary against a deterministic Hugging Face eager checkpoint executed through the real IREE local-task graph. Cover logits, selected all-layer KV caches, greedy token parity, and negative causal-mask, multiplicative-mask, double-scale, and image-pre-scale regressions. Validation: gemma3_vlm_mask_check.py passes on llvm-cpu/local-task. Refs #869 --- spike/openxla/gemma3_vlm_mask_check.py | 432 +++++++++++++++++++++++++ 1 file changed, 432 insertions(+) create mode 100644 spike/openxla/gemma3_vlm_mask_check.py diff --git a/spike/openxla/gemma3_vlm_mask_check.py b/spike/openxla/gemma3_vlm_mask_check.py new file mode 100644 index 000000000..f58c3403d --- /dev/null +++ b/spike/openxla/gemma3_vlm_mask_check.py @@ -0,0 +1,432 @@ +#!/usr/bin/env python3 +"""Independent Gemma3 VLM embeddings/mask execution check for issue #869. + +The production VLM path hands `prefill_embeddings.main` post-Gemma-scale text +rows, unscaled projected image rows, and one authoritative additive f32 mask. +This fixture compares that graph with an independent Hugging Face Gemma3 eager +model using the same tiny deterministic checkpoint. The HF side receives the +corresponding pre-scale rows because Gemma3 applies `sqrt(hidden_size)` inside +its model forward. + +The canonical case compares last-token logits, every layer's K/V cache, and the +greedy token. Four negative fixtures must then diverge: accidental causal +masking, a multiplicative 0/1 mask, double Gemma scaling, and a missing image +pre-divide. The tiny CPU graph keeps the check suitable for a short, bounded +local-task gate; pinned-checkpoint vision and server gates remain separate. + +Run from the repository root with the shared OpenXLA spike environment: + + spike/openxla/.venv/bin/python spike/openxla/gemma3_vlm_mask_check.py + +Exit 0 means the independent canonical oracle matched and every negative +fixture was detected. +""" + +from __future__ import annotations + +import json +import os +import subprocess +import sys +import tempfile +from pathlib import Path + +import numpy as np +import torch +from iree.compiler.tools import compile_file +from iree.runtime import load_vm_flatbuffer_file +from transformers import Gemma3ForCausalLM, Gemma3TextConfig + +REPO_ROOT = Path(__file__).resolve().parents[2] +CARGO = os.environ.get("CARGO", "cargo") +PREFILL_LP = 256 +REAL_LEN = 12 +HIDDEN = 16 +INTERMEDIATE = 32 +N_LAYERS = 4 +N_Q = 4 +N_KV = 2 +HEAD_DIM = 4 +VOCAB = 64 +EPS = 1e-6 +ROPE_THETA = 10_000.0 +ROPE_LOCAL_BASE = 100.0 +SLIDING_WINDOW = 4 +SLIDING_PATTERN = 3 +MASKED = np.float32(np.finfo(np.float32).min) +LOGIT_ATOL = 2e-2 +KV_MAX_ABS = 2.5e-1 +KV_RMS = 1e-1 +KV_MIN_COSINE = 9.9e-1 +NEGATIVE_MIN_DIFF = 1e-4 +IMAGE_POSITIONS = (3, 4) + + +def dimensions() -> dict[str, object]: + return { + "hidden_size": HIDDEN, + "num_attention_heads": N_Q, + "num_key_value_heads": N_KV, + "head_dim": HEAD_DIM, + "intermediate_size": INTERMEDIATE, + "num_hidden_layers": N_LAYERS, + "vocab_size": VOCAB, + "rms_norm_eps": EPS, + "rope_theta": ROPE_THETA, + "max_position_embeddings": 512, + "attention_bias": False, + } + + +def hf_config() -> Gemma3TextConfig: + return Gemma3TextConfig( + **dimensions(), + tie_word_embeddings=True, + query_pre_attn_scalar=HEAD_DIM, + rope_local_base_freq=ROPE_LOCAL_BASE, + sliding_window=SLIDING_WINDOW, + sliding_window_pattern=SLIDING_PATTERN, + attn_logit_softcapping=None, + final_logit_softcapping=None, + hidden_activation="gelu_pytorch_tanh", + ) + + +def emitter_config() -> dict[str, object]: + config = dimensions() + config.pop("max_position_embeddings") + config.update( + model_type="gemma3_text", + tie_word_embeddings=True, + query_pre_attn_scalar=HEAD_DIM, + rope_local_base_freq=ROPE_LOCAL_BASE, + sliding_window=SLIDING_WINDOW, + sliding_window_pattern=SLIDING_PATTERN, + attn_logit_softcapping=None, + final_logit_softcapping=None, + hidden_activation="gelu_pytorch_tanh", + ) + return config + + +def argument_names() -> list[str]: + names = ["model.embed_tokens.weight", "model.norm.weight"] + for index in range(N_LAYERS): + prefix = f"model.layers.{index}." + names.extend( + [ + prefix + "mlp.down_proj.weight", + prefix + "mlp.gate_proj.weight", + prefix + "input_layernorm.weight", + prefix + "post_attention_layernorm.weight", + prefix + "mlp.up_proj.weight", + prefix + "self_attn.k_proj.weight", + prefix + "self_attn.o_proj.weight", + prefix + "self_attn.q_proj.weight", + prefix + "self_attn.v_proj.weight", + prefix + "self_attn.q_norm.weight", + prefix + "self_attn.k_norm.weight", + prefix + "pre_feedforward_layernorm.weight", + prefix + "post_feedforward_layernorm.weight", + ] + ) + return names + + +def build_checkpoint() -> tuple[Gemma3ForCausalLM, list[np.ndarray]]: + torch.manual_seed(869) + model = Gemma3ForCausalLM(hf_config()).eval().float() + model.config._attn_implementation = "eager" + with torch.no_grad(): + for _, parameter in model.named_parameters(): + if parameter.dim() == 1: + parameter.copy_(torch.randn_like(parameter) * 0.1) + state = model.state_dict() + names = argument_names() + missing = [name for name in names if name not in state] + if missing: + raise RuntimeError(f"HF checkpoint is missing emitter weights: {missing[:4]}") + weights = [ + np.ascontiguousarray(state[name].detach().numpy(), dtype=np.float32) + for name in names + ] + return model, weights + + +def emit_and_compile() -> object: + work = Path(tempfile.mkdtemp(prefix="gemma3_vlm_mask_")) + config_path = work / "config.json" + config_path.write_text(json.dumps(emitter_config()), encoding="utf-8") + + print("[emit] Gemma3 embeddings prefill StableHLO", flush=True) + subprocess.run( + [ + CARGO, + "test", + "-p", + "mlxcel-xla", + "--lib", + "emitter::tests::dump_prefill_embeddings_parity_graphs", + "--", + "--ignored", + "--nocapture", + ], + cwd=REPO_ROOT, + env={ + **os.environ, + "MLXCEL_DUMP_CONFIG": str(config_path), + "MLXCEL_DUMP_DIR": str(work), + }, + check=True, + ) + + source = work / "prefill_embeddings_logits.mlir" + output = work / "prefill_embeddings_logits.vmfb" + print("[compile] Gemma3 embeddings prefill (llvm-cpu)", flush=True) + compile_file( + str(source), + output_file=str(output), + input_type="stablehlo", + target_backends=["llvm-cpu"], + ) + return load_vm_flatbuffer_file(str(output), driver="local-task") + + +def canonical_inputs( + embedding_table: np.ndarray, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + rng = np.random.default_rng(869) + tokens = rng.integers(1, VOCAB, size=REAL_LEN, dtype=np.int32) + pre_scale = np.ascontiguousarray(embedding_table[tokens], dtype=np.float32) + projected = np.ascontiguousarray( + rng.normal(0.0, 0.2, (len(IMAGE_POSITIONS), HIDDEN)), + dtype=np.float32, + ) + normalizer = np.float32(np.sqrt(HIDDEN)) + for row, position in enumerate(IMAGE_POSITIONS): + pre_scale[position] = projected[row] / normalizer + + post_scale = np.zeros((PREFILL_LP, HIDDEN), dtype=np.float32) + post_scale[:REAL_LEN] = pre_scale * normalizer + + mask = np.full((PREFILL_LP, PREFILL_LP), MASKED, dtype=np.float32) + mask[:REAL_LEN, :REAL_LEN] = 0.0 + return pre_scale, np.ascontiguousarray(post_scale), np.ascontiguousarray(mask) + + +def to_host(value: object) -> np.ndarray: + host = value.to_host() if hasattr(value, "to_host") else value + return np.asarray(host, dtype=np.float32) + + +def run_iree( + module: object, + weights: list[np.ndarray], + embeddings: np.ndarray, + mask: np.ndarray, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + positions = np.arange(PREFILL_LP, dtype=np.int32) + real_len = np.asarray(REAL_LEN, dtype=np.int32) + outputs = module.main(*weights, embeddings, positions, real_len, mask) + return tuple(to_host(value) for value in outputs) + + +def legacy_cache(cache: object) -> tuple[tuple[torch.Tensor, torch.Tensor], ...]: + if hasattr(cache, "to_legacy_cache"): + return cache.to_legacy_cache() + if isinstance(cache, (tuple, list)): + return tuple(cache) + layers = getattr(cache, "layers", None) + if layers is None: + raise RuntimeError(f"unsupported HF cache type: {type(cache).__name__}") + return tuple((layer.keys, layer.values) for layer in layers) + + +def run_hf( + model: Gemma3ForCausalLM, + pre_scale: np.ndarray, + mask: np.ndarray, +) -> tuple[np.ndarray, list[np.ndarray], list[np.ndarray]]: + real_mask = np.ascontiguousarray(mask[:REAL_LEN, :REAL_LEN]) + with torch.no_grad(): + outputs = model( + inputs_embeds=torch.from_numpy(pre_scale[None, ...]), + attention_mask=torch.from_numpy(real_mask[None, None, ...]), + position_ids=torch.arange(REAL_LEN, dtype=torch.long)[None, :], + use_cache=True, + return_dict=True, + ) + logits = outputs.logits[0, REAL_LEN - 1].detach().numpy().astype(np.float32) + cache = legacy_cache(outputs.past_key_values) + keys = [ + layer[0][0].detach().numpy().transpose(1, 0, 2).astype(np.float32) + for layer in cache + ] + values = [ + layer[1][0].detach().numpy().transpose(1, 0, 2).astype(np.float32) + for layer in cache + ] + return logits, keys, values + + +def compare_canonical( + iree: tuple[np.ndarray, np.ndarray, np.ndarray], + hf: tuple[np.ndarray, list[np.ndarray], list[np.ndarray]], +) -> bool: + logits_close = np.allclose(iree[0], hf[0], rtol=0.0, atol=LOGIT_ATOL) + logits_diff = float(np.max(np.abs(iree[0] - hf[0]))) + print( + f"[canonical/logits] shape={iree[0].shape}/{hf[0].shape} " + f"max|diff|={logits_diff:.3e} -> {'PASS' if logits_close else 'FAIL'}", + flush=True, + ) + ok = bool(logits_close) + for name, actual, expected_layers in ( + ("kcache", iree[1], hf[1]), + ("vcache", iree[2], hf[2]), + ): + for layer, expected in enumerate(expected_layers): + cache_len = expected.shape[0] + actual_slice = actual[layer, REAL_LEN - cache_len : REAL_LEN] + same_shape = actual_slice.shape == expected.shape + max_diff = ( + float(np.max(np.abs(actual_slice - expected))) + if same_shape + else float("inf") + ) + rms_diff = ( + float(np.sqrt(np.mean(np.square(actual_slice - expected)))) + if same_shape + else float("inf") + ) + denominator = ( + float(np.linalg.norm(actual_slice) * np.linalg.norm(expected)) + if same_shape + else 0.0 + ) + cosine = ( + float(np.vdot(actual_slice, expected) / denominator) + if denominator > 0.0 + else 1.0 + ) + close = ( + same_shape + and max_diff <= KV_MAX_ABS + and rms_diff <= KV_RMS + and cosine >= KV_MIN_COSINE + ) + ok = ok and close + best = "" + if not close and cache_len < REAL_LEN: + candidates = [ + ( + start, + float( + np.max( + np.abs( + actual[layer, start : start + cache_len] - expected + ) + ) + ), + ) + for start in range(REAL_LEN - cache_len + 1) + ] + best_start, best_diff = min(candidates, key=lambda item: item[1]) + best = f" best_slice={best_start}:{best_start + cache_len}/{best_diff:.3e}" + print( + f"[canonical/{name}/layer{layer}] " + f"shape={actual_slice.shape}/{expected.shape} " + f"max|diff|={max_diff:.3e} rms={rms_diff:.3e} " + f"cos={cosine:.6f}{best} " + f"-> {'PASS' if close else 'FAIL'}", + flush=True, + ) + iree_token = int(np.argmax(iree[0])) + hf_token = int(np.argmax(hf[0])) + token_ok = iree_token == hf_token + print( + f"[canonical/token] iree={iree_token} hf={hf_token} " + f"-> {'PASS' if token_ok else 'FAIL'}", + flush=True, + ) + return ok and token_ok + + +def negative_detected( + label: str, + canonical_logits: np.ndarray, + negative_logits: np.ndarray, +) -> bool: + difference = float(np.max(np.abs(canonical_logits - negative_logits))) + detected = difference > NEGATIVE_MIN_DIFF + print( + f"[negative/{label}] max|logit diff|={difference:.3e} " + f"-> {'DETECTED' if detected else 'MISSED'}", + flush=True, + ) + return detected + + +def main() -> int: + model, weights = build_checkpoint() + module = emit_and_compile() + pre_scale, embeddings, mask = canonical_inputs(weights[0]) + + print("[run] canonical independent HF and IREE paths", flush=True) + canonical = run_iree(module, weights, embeddings, mask) + reference = run_hf(model, pre_scale, mask) + checks = [compare_canonical(canonical, reference)] + + causal = np.full_like(mask, MASKED) + valid = np.arange(REAL_LEN) + causal[valid[:, None], valid[None, :]] = np.where( + valid[None, :] <= valid[:, None], 0.0, MASKED + ) + checks.append( + negative_detected( + "causal-mask", + canonical[0], + run_iree(module, weights, embeddings, causal)[0], + ) + ) + + multiplicative = np.zeros_like(mask) + multiplicative[:REAL_LEN, :REAL_LEN] = 1.0 + checks.append( + negative_detected( + "multiplicative-mask", + canonical[0], + run_iree(module, weights, embeddings, multiplicative)[0], + ) + ) + + checks.append( + negative_detected( + "double-scale", + canonical[0], + run_iree(module, weights, embeddings * np.float32(np.sqrt(HIDDEN)), mask)[0], + ) + ) + + missing_image_prescale = embeddings.copy() + missing_image_prescale[list(IMAGE_POSITIONS)] *= np.float32(np.sqrt(HIDDEN)) + checks.append( + negative_detected( + "missing-image-prescale", + canonical[0], + run_iree(module, weights, missing_image_prescale, mask)[0], + ) + ) + + ok = all(checks) + print( + f"RESULT: {'PASS' if ok else 'FAIL'} " + f"(logit_atol={LOGIT_ATOL:g}, kv_max={KV_MAX_ABS:g}, " + f"kv_rms={KV_RMS:g}, kv_cos={KV_MIN_COSINE:g}, local-task)", + flush=True, + ) + return 0 if ok else 1 + + +if __name__ == "__main__": + sys.exit(main()) From 8cce0a43cc3df6bde5dcf94d2728dfa19adc6124 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Sun, 26 Jul 2026 22:06:20 +0900 Subject: [PATCH 05/11] fix(xla): apply Gemma3 SigLIP post-layer normalization Gemma3 eager SigLIP applies post_layernorm before the average-pool projector, but the resident IREE graph omitted those weights and the normalization, causing the first vision-boundary divergence. Add the Gemma3-only post-layernorm schema and artifact generation, plus a CPU/local-task diagnostics feature and pinned ignored gate that compares processor pixels, early and final SigLIP stages, projector rows, one-time embedding scaling, and every additive-mask cell without loading a duplicate decoder. Validation: cargo test -p mlxcel-xla --lib gemma3_ (8 passed); cargo test --features xla-reference-diagnostics --lib comparison_reports_the_first_failed_element (1 passed); cargo clippy --features xla-reference-diagnostics --lib --tests --fix --allow-dirty --allow-staged --no-deps (existing warnings only). The real-checkpoint ignored gate was not run. Refs #869 --- Cargo.toml | 4 + src/lib/mlxcel-xla/src/emitter/vision.rs | 11 + .../mlxcel-xla/src/emitter/vision_config.rs | 28 +- src/loading/vlm_gemma.rs | 23 +- src/loading/vlm_gemma_xla_tests.rs | 697 ++++++++++++++++++ src/vision/encoders/mod.rs | 2 +- src/vision/encoders/siglip.rs | 2 +- 7 files changed, 743 insertions(+), 24 deletions(-) create mode 100644 src/loading/vlm_gemma_xla_tests.rs diff --git a/Cargo.toml b/Cargo.toml index ab506bce0..32924c1e7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -71,6 +71,10 @@ xla-backend = ["dep:mlxcel-xla"] # StableHLO prefill / decode_step graphs. Needs IREE_DIST at build time (the # extracted iree dist), so it is a local / opt-in build, not a CI default. xla-iree = ["xla-backend", "mlxcel-xla/iree"] +# Device-neutral reference capture for native IREE diagnostics. This keeps the +# eager MLX side on its default CPU backend and qualifies IREE `local-task` +# without requiring an IREE CUDA source/build tree. +xla-reference-diagnostics = ["xla-iree", "mlxcel-xla/diagnostics"] # Explicit test-only forwarding for Gemma3n intermediate oracle validation. # Its pinned reference is the production CUDA runtime, so enabling diagnostics # also enables the root MLX CUDA backend. This is not part of `xla-iree`; normal diff --git a/src/lib/mlxcel-xla/src/emitter/vision.rs b/src/lib/mlxcel-xla/src/emitter/vision.rs index 9d74df598..b936596fb 100644 --- a/src/lib/mlxcel-xla/src/emitter/vision.rs +++ b/src/lib/mlxcel-xla/src/emitter/vision.rs @@ -82,6 +82,8 @@ mod tests { ) .unwrap(); let mlir = emit_vision(&config); + assert!(mlir.contains("vision_tower.vision_model.post_layernorm.weight")); + assert!(mlir.contains("vision_tower.vision_model.post_layernorm.bias")); assert!(mlir.contains("multi_modal_projector.mm_soft_emb_norm.weight")); assert!(mlir.contains("multi_modal_projector.mm_input_projection_weight")); assert!(!mlir.contains("multi_modal_projector.linear_1")); @@ -411,6 +413,15 @@ fn emit_vision_impl(config: &LlavaVisionConfig, diagnostics: bool) -> String { outputs.push(hidden.clone()); } } + if matches!(config.projector, VisionProjector::Gemma3AvgPool { .. }) { + hidden = layer_norm( + &mut builder, + &hidden, + &args.take(), + &args.take(), + config.layer_norm_eps, + ); + } if config.drop_first_token { hidden = builder.slice(&hidden, &[(1, config.position_count()), (0, config.hidden)]); } diff --git a/src/lib/mlxcel-xla/src/emitter/vision_config.rs b/src/lib/mlxcel-xla/src/emitter/vision_config.rs index f53f581ad..6b9cae779 100644 --- a/src/lib/mlxcel-xla/src/emitter/vision_config.rs +++ b/src/lib/mlxcel-xla/src/emitter/vision_config.rs @@ -227,7 +227,16 @@ mod tests { } ); assert!(config.fingerprint().contains("image_token_id: 99")); + assert!(config.fingerprint().starts_with("iree-vision-v3:")); let specs = config.weight_specs(); + assert_eq!( + specs[specs.len() - 4].name, + "vision_tower.vision_model.post_layernorm.weight" + ); + assert_eq!( + specs[specs.len() - 3].name, + "vision_tower.vision_model.post_layernorm.bias" + ); assert_eq!( specs[specs.len() - 2].name, "multi_modal_projector.mm_soft_emb_norm.weight" @@ -505,8 +514,13 @@ impl LlavaVisionConfig { #[must_use] pub(crate) fn fingerprint(&self) -> String { + let schema = if matches!(self.projector, VisionProjector::Gemma3AvgPool { .. }) { + "iree-vision-v3" + } else { + "iree-vision-v2" + }; format!( - "iree-vision-v2:image={}:patch={}:channels={}:hidden={}:intermediate={}:layers={}:\ + "{schema}:image={}:patch={}:channels={}:hidden={}:intermediate={}:layers={}:\ heads={}:eps={:08x}:activation={:?}:class={}:feature={}:drop_first={}:text={}:projector={:?}", self.image_size, self.patch_size, @@ -581,6 +595,18 @@ impl LlavaVisionConfig { }); } } + if matches!(self.projector, VisionProjector::Gemma3AvgPool { .. }) { + specs.extend([ + self.spec( + "vision_tower.vision_model.post_layernorm.weight", + [self.hidden], + ), + self.spec( + "vision_tower.vision_model.post_layernorm.bias", + [self.hidden], + ), + ]); + } match self.projector { VisionProjector::LlavaMlp => specs.extend([ self.spec( diff --git a/src/loading/vlm_gemma.rs b/src/loading/vlm_gemma.rs index 7d94ccb14..663d14010 100644 --- a/src/loading/vlm_gemma.rs +++ b/src/loading/vlm_gemma.rs @@ -330,27 +330,8 @@ pub(crate) fn load_gemma3_iree_host_preprocessor( } #[cfg(all(test, feature = "xla-iree"))] -mod xla_tests { - use super::*; - - #[test] - fn pinned_gemma3_projector_loads_and_returns_finite_features() { - let Ok(model) = std::env::var("MLXCEL_GEMMA3_FIXTURE") else { - return; - }; - let device = std::env::var("MLXCEL_XLA_DEVICE").unwrap_or_else(|_| "cuda".to_string()); - let mut projector = mlxcel_xla::IreeVisionProjector::load(Path::new(&model), &device) - .expect("load pinned Gemma3 IREE vision projector"); - assert_eq!(projector.input_shape(), [1, 3, 896, 896]); - assert_eq!(projector.output_shape(), [256, 2560]); - let pixels = vec![0.0; projector.input_shape().into_iter().product()]; - let projection = projector - .project(&pixels) - .expect("execute pinned Gemma3 IREE vision projector"); - assert_eq!(projection.shape, [256, 2560]); - assert!(projection.values.iter().all(|value| value.is_finite())); - } -} +#[path = "vlm_gemma_xla_tests.rs"] +mod xla_tests; /// Load a Gemma3n VLM model. pub(crate) fn load_gemma3n_vlm(model_path: &Path) -> Result { diff --git a/src/loading/vlm_gemma_xla_tests.rs b/src/loading/vlm_gemma_xla_tests.rs new file mode 100644 index 000000000..6f0a14a45 --- /dev/null +++ b/src/loading/vlm_gemma_xla_tests.rs @@ -0,0 +1,697 @@ +// 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. + +use super::*; + +#[test] +fn pinned_gemma3_projector_loads_and_returns_finite_features() { + let Ok(model) = std::env::var("MLXCEL_GEMMA3_FIXTURE") else { + return; + }; + let device = std::env::var("MLXCEL_XLA_DEVICE").unwrap_or_else(|_| "cuda".to_string()); + let mut projector = mlxcel_xla::IreeVisionProjector::load(Path::new(&model), &device) + .expect("load pinned Gemma3 IREE vision projector"); + assert_eq!(projector.input_shape(), [1, 3, 896, 896]); + assert_eq!(projector.output_shape(), [256, 2560]); + let pixels = vec![0.0; projector.input_shape().into_iter().product()]; + let projection = projector + .project(&pixels) + .expect("execute pinned Gemma3 IREE vision projector"); + assert_eq!(projection.shape, [256, 2560]); + assert!(projection.values.iter().all(|value| value.is_finite())); +} + +#[cfg(feature = "xla-reference-diagnostics")] +mod reference_boundary { + use std::fs; + use std::io::{self, Write}; + use std::path::{Path, PathBuf}; + + use mlxcel_core::layers::UnifiedEmbedding; + use mlxcel_core::session::{OwnedTensor, PreparedPrefill, PreparedTensorDType}; + use sha2::{Digest, Sha256}; + + use super::*; + use crate::multimodal::host_preprocessor::{ + Gemma3IreeHostPreprocessor, HostMultimodalPreprocessor, + }; + use crate::multimodal::vlm_prompt::{ImageTokenBlockInfo, apply_image_token_blocks}; + use crate::vision::config::VLMConfig; + use crate::vision::connectors::MultiModalConnector; + use crate::vision::connectors::avg_pool::AvgPoolProjector; + use crate::vision::encoders::VisionEncoder; + use crate::vision::encoders::siglip::SigLipVisionModel; + use crate::vision::processors::ImageProcessor; + use crate::vision::processors::siglip::SigLipProcessor; + + const PINNED_REVISION: &str = "93724907d4ed1745d2fe50baadf3b0b01a65abf2"; + const PINNED_CONFIG_SHA256: &str = + "5ccdde91da736e6e6f8f138268c620adcbf1219c973b884240b719a54122465b"; + const PINNED_PREPROCESSOR_SHA256: &str = + "f688d6bb20c5017601c4011de7ca656da8485b540b05013efdaf986c0fcc918d"; + const PINNED_PROCESSOR_SHA256: &str = + "3ffd5f11778dc73e2b69b3c00535e4121e1badf7018136263cd17b5b34fbaa53"; + const PINNED_IMAGE_SHA256: &str = + "5e7d54e8a7d21802378c87d2d70cf551e29739fe27599ddf129ebccdad1e6261"; + const BLOCK0_STAGES: [&str; 12] = [ + "siglip.block0.layer_norm1", + "siglip.block0.q_proj", + "siglip.block0.k_proj", + "siglip.block0.v_proj", + "siglip.block0.attention_context", + "siglip.block0.attention_output", + "siglip.block0.attention_residual", + "siglip.block0.layer_norm2", + "siglip.block0.mlp_fc1", + "siglip.block0.mlp_activation", + "siglip.block0.mlp_fc2", + "siglip.block0.output", + ]; + + #[derive(Debug, Clone, Copy)] + struct Tolerance { + atol: f64, + rtol: f64, + } + + const PIXEL_TOLERANCE: Tolerance = Tolerance { + atol: 1e-6, + rtol: 1e-6, + }; + // The pinned MLX checkpoint executes the eager vision path in BF16 while + // the qualified IREE graph widens immutable checkpoint weights to F32. + // These are the same BF16 stage envelopes used by the #863 vision gate. + const VISION_TOLERANCE: Tolerance = Tolerance { + atol: 8e-2, + rtol: 4e-2, + }; + const IREE_REPLAY_TOLERANCE: Tolerance = Tolerance { + atol: 1e-6, + rtol: 1e-6, + }; + const EXACT_TOLERANCE: Tolerance = Tolerance { + atol: 0.0, + rtol: 0.0, + }; + + #[derive(Debug, Clone, Copy, PartialEq)] + struct ComparisonStats { + max_absolute: f64, + max_relative: f64, + failures: usize, + non_finite_count: usize, + first_failure: Option, + } + + fn comparison_stats( + observed: &[f32], + reference: &[f32], + tolerance: Tolerance, + ) -> ComparisonStats { + assert_eq!(observed.len(), reference.len(), "comparison lengths differ"); + let mut stats = ComparisonStats { + max_absolute: 0.0, + max_relative: 0.0, + failures: 0, + non_finite_count: 0, + first_failure: None, + }; + for (index, (&observed, &reference)) in observed.iter().zip(reference).enumerate() { + if !observed.is_finite() || !reference.is_finite() { + stats.failures += 1; + stats.non_finite_count += 1; + stats.first_failure.get_or_insert(index); + continue; + } + let absolute = f64::from((observed - reference).abs()); + let relative = absolute / f64::from(reference.abs()).max(f64::MIN_POSITIVE); + stats.max_absolute = stats.max_absolute.max(absolute); + stats.max_relative = stats.max_relative.max(relative); + if absolute > tolerance.atol + tolerance.rtol * f64::from(reference.abs()) { + stats.failures += 1; + stats.first_failure.get_or_insert(index); + } + } + stats + } + + fn progress(stage: &str) { + eprintln!("[gemma3-vlm-boundary] {stage}"); + io::stderr().flush().expect("flush diagnostic progress"); + } + + fn compare_stage( + stage: &str, + observed: &[f32], + reference: &[f32], + tolerance: Tolerance, + first_divergence: &mut Option, + ) { + if observed.len() != reference.len() { + let detail = format!("{stage}: length {} != {}", observed.len(), reference.len()); + first_divergence.get_or_insert(detail.clone()); + eprintln!("[gemma3-vlm-boundary] stage={stage} status=FAIL {detail}"); + return; + } + let stats = comparison_stats(observed, reference, tolerance); + let status = if stats.failures == 0 { "PASS" } else { "FAIL" }; + eprintln!( + "[gemma3-vlm-boundary] stage={stage} status={status} elements={} \ + atol={:.3e} rtol={:.3e} max_abs={:.6e} max_rel={:.6e} \ + failures={} non_finite={} first_failure={:?}", + observed.len(), + tolerance.atol, + tolerance.rtol, + stats.max_absolute, + stats.max_relative, + stats.failures, + stats.non_finite_count, + stats.first_failure, + ); + io::stderr().flush().expect("flush diagnostic stage"); + if stats.failures != 0 { + first_divergence.get_or_insert_with(|| { + format!( + "{stage} at flat index {}", + stats + .first_failure + .expect("a failed comparison has an index") + ) + }); + } + } + + fn sha256(path: &Path) -> String { + let bytes = + fs::read(path).unwrap_or_else(|error| panic!("read {}: {error}", path.display())); + format!("{:x}", Sha256::digest(bytes)) + } + + fn assert_sha256(path: &Path, expected: &str, label: &str) { + assert_eq!( + sha256(path), + expected, + "{label} differs from the pinned #869 fixture: {}", + path.display() + ); + } + + fn pinned_revision(model: &Path) -> String { + if let Ok(revision) = std::env::var("MLXCEL_GEMMA3_REVISION") { + return revision; + } + let metadata = model.join(".cache/huggingface/download/config.json.metadata"); + fs::read_to_string(&metadata) + .unwrap_or_else(|error| { + panic!( + "read pinned revision from {} ({error}); set MLXCEL_GEMMA3_REVISION", + metadata.display() + ) + }) + .lines() + .next() + .expect("Hugging Face metadata contains a revision") + .to_string() + } + + fn tensor_f32(tensor: &OwnedTensor, label: &str) -> Vec { + assert_eq!( + tensor.dtype, + PreparedTensorDType::Float32, + "{label} must be float32" + ); + tensor + .bytes + .chunks_exact(4) + .map(|bytes| f32::from_le_bytes(bytes.try_into().expect("four-byte f32 chunk"))) + .collect() + } + + fn mlx_f32(array: &mlxcel_core::MlxArray, label: &str) -> Vec { + let widened = mlxcel_core::astype(array, mlxcel_core::dtype::FLOAT32); + mlxcel_core::try_array_to_raw_bytes(&widened) + .unwrap_or_else(|error| panic!("export {label}: {error}")) + .chunks_exact(4) + .map(|bytes| f32::from_ne_bytes(bytes.try_into().expect("four-byte f32 chunk"))) + .collect() + } + + fn reference_weight(name: &str) -> bool { + let canonical = name.strip_prefix("language_model.").unwrap_or(name); + canonical.starts_with("vision_tower.") + || canonical.starts_with("multi_modal_projector.") + || canonical.starts_with("model.embed_tokens.") + } + + fn image_block_info(config: &VLMConfig, tokens_per_image: usize) -> ImageTokenBlockInfo { + ImageTokenBlockInfo { + use_boi_eoi: true, + image_token_id: config.image_token_index, + mm_tokens_per_image: tokens_per_image, + boi_token_id: config.boi_token_index, + eoi_token_id: config.eoi_token_index, + has_bos: true, + separator_token_id: None, + suffix_tokens: Vec::new(), + block_prefix_tokens: vec![108], + block_suffix_tokens: vec![108], + } + } + + fn image_rows(prepared: &PreparedPrefill, image_token_id: i32, hidden: usize) -> Vec { + let embeddings = tensor_f32(&prepared.embeddings, "prepared embeddings"); + prepared + .token_ids + .iter() + .enumerate() + .filter(|(_, token)| **token == image_token_id) + .flat_map(|(position, _)| { + embeddings[position * hidden..(position + 1) * hidden] + .iter() + .copied() + }) + .collect() + } + + fn expected_mask(token_ids: &[i32], pad_token_id: i32) -> Vec { + let mut expected = vec![f32::MIN; token_ids.len() * token_ids.len()]; + for (query, query_token) in token_ids.iter().enumerate() { + for (key, key_token) in token_ids.iter().enumerate() { + if *query_token != pad_token_id && *key_token != pad_token_id { + expected[query * token_ids.len() + key] = 0.0; + } + } + } + expected + } + + fn expected_post_scale( + token_ids: &[i32], + raw_text: &[f32], + projected: &[f32], + hidden: usize, + pad_token_id: i32, + image_token_id: i32, + ) -> Vec { + let normalizer = (hidden as f64).sqrt() as f32; + let mut expected = vec![0.0; token_ids.len() * hidden]; + let mut image_row = 0; + for (position, token) in token_ids.iter().enumerate() { + let destination = position * hidden; + if *token == pad_token_id { + continue; + } + if *token == image_token_id { + let source = image_row * hidden; + expected[destination..destination + hidden] + .copy_from_slice(&projected[source..source + hidden]); + image_row += 1; + } else { + for offset in 0..hidden { + expected[destination + offset] = raw_text[destination + offset] * normalizer; + } + } + } + expected + } + + #[test] + fn comparison_reports_the_first_failed_element() { + let stats = comparison_stats( + &[1.0, 2.2, f32::NAN], + &[1.0, 2.0, 3.0], + Tolerance { + atol: 0.01, + rtol: 0.01, + }, + ); + assert_eq!(stats.failures, 2); + assert_eq!(stats.non_finite_count, 1); + assert_eq!(stats.first_failure, Some(1)); + } + + /// Pinned mixed-runtime boundary gate for #869. + /// + /// This deliberately remains ignored. Run it from the repository root only + /// when the pinned checkpoint and local IREE distribution are installed: + /// + /// ```text + /// IREE_DIST=/path/to/iree-dist \ + /// MLXCEL_GEMMA3_FIXTURE=/path/to/gemma-3-4b-it-4bit \ + /// cargo test --features xla-reference-diagnostics --lib \ + /// pinned_gemma3_eager_mlx_matches_iree_prepared_boundary -- \ + /// --ignored --nocapture + /// ``` + /// + /// Both eager MLX and IREE execute on CPU; no language decoder is constructed. + #[test] + #[ignore = "requires pinned Gemma3 checkpoint, image, and IREE local-task"] + fn pinned_gemma3_eager_mlx_matches_iree_prepared_boundary() { + let model = PathBuf::from( + std::env::var("MLXCEL_GEMMA3_FIXTURE") + .expect("MLXCEL_GEMMA3_FIXTURE must name the pinned checkpoint"), + ); + let image_path = std::env::var("MLXCEL_GEMMA3_IMAGE") + .map(PathBuf::from) + .unwrap_or_else(|_| PathBuf::from("tests/fixtures/test_image.png")); + let device = + std::env::var("MLXCEL_XLA_DEVICE").unwrap_or_else(|_| "local-task".to_string()); + assert_eq!( + device, "local-task", + "the pinned #869 mixed-runtime gate qualifies IREE local-task only" + ); + + progress("validate pinned checkpoint and image"); + assert_eq!( + pinned_revision(&model), + PINNED_REVISION, + "checkpoint revision differs from pinned mlx-community/gemma-3-4b-it-4bit" + ); + assert_sha256( + &model.join("config.json"), + PINNED_CONFIG_SHA256, + "checkpoint config", + ); + assert_sha256( + &model.join("preprocessor_config.json"), + PINNED_PREPROCESSOR_SHA256, + "checkpoint image preprocessor", + ); + assert_sha256( + &model.join("processor_config.json"), + PINNED_PROCESSOR_SHA256, + "checkpoint processor", + ); + assert_sha256(&image_path, PINNED_IMAGE_SHA256, "image"); + + progress("load filtered MLX SigLIP/projector/text embedding weights"); + let (_config_str, full_config) = + read_sanitized_vlm_config(&model).expect("read pinned Gemma3 config"); + let config: VLMConfig = + serde_json::from_value(full_config.clone()).expect("parse pinned Gemma3 VLM config"); + let text_config: models::gemma3::ModelArgs = + serde_json::from_value(config.text_config.clone()) + .expect("parse pinned Gemma3 text config"); + let weights = load_vlm_weights_common_filtered_canonical(&model, reference_weight) + .map(strip_language_model_prefix) + .expect("load only Gemma3 vision/projector/embedding weights"); + let quant_group_size = full_config + .pointer("/quantization/group_size") + .and_then(Value::as_i64) + .unwrap_or(64) as i32; + let quant_bits = full_config + .pointer("/quantization/bits") + .and_then(Value::as_i64) + .unwrap_or(4) as i32; + let text_embeddings = UnifiedEmbedding::from_weights( + &weights, + "model.embed_tokens", + quant_group_size, + quant_bits, + ) + .expect("load filtered Gemma3 text embedding table"); + let encoder = SigLipVisionModel::from_weights( + &weights, + &config.vision_config, + "vision_tower.vision_model", + ) + .expect("load filtered Gemma3 SigLIP tower"); + let tokens_per_image = config.get_mm_tokens_per_image(); + let connector = AvgPoolProjector::from_weights( + &weights, + "multi_modal_projector", + config.vision_config.hidden_size, + config.vision_config.image_size, + config.vision_config.patch_size, + tokens_per_image, + config.vision_config.layer_norm_eps, + ) + .expect("load filtered Gemma3 average-pool projector"); + let image = image::open(&image_path) + .unwrap_or_else(|error| panic!("decode {}: {error}", image_path.display())); + let images = vec![image]; + + progress("compare independently constructed processor pixels"); + let mlx_processor = SigLipProcessor::new(config.vision_config.image_size); + let iree_processor = SigLipProcessor::new(config.vision_config.image_size); + let mlx_pixels = mlx_processor.preprocess(&images); + let iree_pixels = iree_processor.preprocess(&images); + let mlx_pixel_values = mlx_f32(&mlx_pixels, "MLX processor pixels"); + let iree_pixel_values = mlx_f32(&iree_pixels, "IREE host processor pixels"); + let mut first_divergence = None; + compare_stage( + "processor.pixel_values", + &iree_pixel_values, + &mlx_pixel_values, + PIXEL_TOLERANCE, + &mut first_divergence, + ); + + progress("expand pinned padded image-token fixture and embed text"); + let mut logical_tokens = vec![ + config.pad_token_id, + 2, + config.boi_token_index, + 1, + config.pad_token_id, + ]; + apply_image_token_blocks( + &mut logical_tokens, + image_block_info(&config, tokens_per_image), + images.len(), + ) + .expect("expand Gemma3 image-token block"); + let input_ids = mlxcel_core::from_slice_i32( + &logical_tokens, + &[ + 1, + i32::try_from(logical_tokens.len()).expect("fixture sequence length fits i32"), + ], + ); + let raw_text_array = text_embeddings.forward(&input_ids); + let embed_dtype = mlxcel_core::array_dtype(&raw_text_array); + let raw_text = mlx_f32(&raw_text_array, "raw text embeddings"); + + progress("capture eager MLX SigLIP hidden and projector stages"); + let mlx_vision_input = mlxcel_core::astype( + &mlxcel_core::transpose_axes(&mlx_pixels, &[0, 2, 3, 1]), + embed_dtype, + ); + let (mlx_selected, mlx_hidden, mlx_block0) = + encoder.forward_with_hidden_state_diagnostics(&mlx_vision_input); + let mlx_selected_values = + mlx_f32(&mlx_selected.hidden_states, "MLX selected vision features"); + let mlx_projected = connector.forward(&mlx_selected.hidden_states); + let mlx_projected_values = mlx_f32(&mlx_projected, "MLX projected image features"); + let mlx_hidden0 = mlx_f32(&mlx_hidden[0], "MLX SigLIP embedding output"); + let mlx_last_hidden = mlx_f32( + mlx_hidden + .last() + .expect("MLX captured a final hidden state"), + "MLX SigLIP last hidden", + ); + let mlx_block0_values = mlx_block0 + .iter() + .map(|stage| mlx_f32(stage, "MLX SigLIP block 0 stage")) + .collect::>(); + assert_eq!(mlx_block0_values.len(), BLOCK0_STAGES.len()); + + progress("run IREE diagnostic SigLIP and average-pool projector"); + let mut diagnostic = mlxcel_xla::IreeVisionDiagnosticProjector::load(&model, &device) + .expect("load Gemma3 IREE diagnostic projector"); + let iree = diagnostic + .project(&iree_pixel_values) + .expect("execute Gemma3 IREE diagnostic projector"); + compare_stage( + "siglip.hidden.embedding", + &iree.hidden_states[0], + &mlx_hidden0, + VISION_TOLERANCE, + &mut first_divergence, + ); + assert_eq!(iree.block0_states.len(), BLOCK0_STAGES.len()); + for ((stage, observed), reference) in BLOCK0_STAGES + .iter() + .zip(&iree.block0_states) + .zip(&mlx_block0_values) + { + compare_stage( + stage, + observed, + reference, + VISION_TOLERANCE, + &mut first_divergence, + ); + } + compare_stage( + "siglip.hidden.last_pre_layernorm", + iree.hidden_states + .last() + .expect("IREE captured a final hidden state"), + &mlx_last_hidden, + VISION_TOLERANCE, + &mut first_divergence, + ); + compare_stage( + "siglip.selected.post_layernorm", + &iree.selected_vision_features, + &mlx_selected_values, + VISION_TOLERANCE, + &mut first_divergence, + ); + compare_stage( + "projector.avg_pool_projection", + &iree.projected_image_features, + &mlx_projected_values, + VISION_TOLERANCE, + &mut first_divergence, + ); + + progress("construct MLX-reference and production IREE prepared prefills"); + let attention_mask = logical_tokens + .iter() + .map(|token| i32::from(*token != config.pad_token_id)) + .collect::>(); + let mlx_prepared = mlxcel_xla::prepare_gemma3_vlm_prefill( + logical_tokens.clone(), + &raw_text, + &mlx_projected_values, + &attention_mask, + text_config.hidden_size, + text_config.max_position_embeddings, + config.pad_token_id, + config.image_token_index, + images.len(), + ) + .expect("construct eager MLX reference prepared prefill"); + let production = Gemma3IreeHostPreprocessor::load(&model, &device) + .expect("load production Gemma3 IREE host preprocessor") + .prepare( + &[ + config.pad_token_id, + 2, + config.boi_token_index, + 1, + config.pad_token_id, + ], + &images, + ) + .expect("construct production Gemma3 IREE prepared prefill"); + assert_eq!(production.token_ids, logical_tokens); + assert_eq!(production.positions, mlx_prepared.positions); + assert_eq!(production.modalities, mlx_prepared.modalities); + + progress("compare final projected image rows and one-time scaling"); + let mlx_image_rows = image_rows( + &mlx_prepared, + config.image_token_index, + text_config.hidden_size, + ); + let iree_image_rows = image_rows( + &production, + config.image_token_index, + text_config.hidden_size, + ); + compare_stage( + "prepared.mlx_projected_image_rows", + &mlx_image_rows, + &mlx_projected_values, + EXACT_TOLERANCE, + &mut first_divergence, + ); + compare_stage( + "prepared.iree_projected_image_rows", + &iree_image_rows, + &iree.projected_image_features, + IREE_REPLAY_TOLERANCE, + &mut first_divergence, + ); + compare_stage( + "prepared.mlx_vs_iree_image_rows", + &iree_image_rows, + &mlx_image_rows, + VISION_TOLERANCE, + &mut first_divergence, + ); + let expected_mlx = expected_post_scale( + &logical_tokens, + &raw_text, + &mlx_projected_values, + text_config.hidden_size, + config.pad_token_id, + config.image_token_index, + ); + let mlx_embeddings = tensor_f32(&mlx_prepared.embeddings, "MLX prepared embeddings"); + compare_stage( + "prepared.text_sqrt_hidden_and_image_identity", + &mlx_embeddings, + &expected_mlx, + EXACT_TOLERANCE, + &mut first_divergence, + ); + let expected_iree = expected_post_scale( + &logical_tokens, + &raw_text, + &iree.projected_image_features, + text_config.hidden_size, + config.pad_token_id, + config.image_token_index, + ); + let production_embeddings = tensor_f32( + &production.embeddings, + "IREE production prepared embeddings", + ); + compare_stage( + "prepared.iree_text_sqrt_hidden_and_image_identity", + &production_embeddings, + &expected_iree, + IREE_REPLAY_TOLERANCE, + &mut first_divergence, + ); + + progress("compare exact additive bidirectional padding mask"); + assert!(!production.attention_bias.causal); + assert!(!mlx_prepared.attention_bias.causal); + assert_eq!( + production.attention_bias.tensor.shape, + vec![1, 1, logical_tokens.len(), logical_tokens.len()] + ); + assert_eq!( + production.attention_bias.tensor.dtype, + PreparedTensorDType::Float32 + ); + let expected_attention = expected_mask(&logical_tokens, config.pad_token_id); + let mlx_attention = tensor_f32(&mlx_prepared.attention_bias.tensor, "MLX attention bias"); + let iree_attention = tensor_f32(&production.attention_bias.tensor, "IREE attention bias"); + compare_stage( + "prepared.mask.mlx_cell_exact", + &mlx_attention, + &expected_attention, + EXACT_TOLERANCE, + &mut first_divergence, + ); + compare_stage( + "prepared.mask.iree_cell_exact", + &iree_attention, + &expected_attention, + EXACT_TOLERANCE, + &mut first_divergence, + ); + + if let Some(first_divergence) = first_divergence { + panic!("Gemma3 eager MLX/IREE first divergence: {first_divergence}"); + } + progress("PASS all pinned Gemma3 eager MLX/IREE boundary stages"); + } +} diff --git a/src/vision/encoders/mod.rs b/src/vision/encoders/mod.rs index a2dab9ff7..bd507f598 100644 --- a/src/vision/encoders/mod.rs +++ b/src/vision/encoders/mod.rs @@ -65,7 +65,7 @@ pub trait VisionEncoder { /// Run the exact encoder path while retaining ordered hidden states for a /// reference-oracle first-divergence report. - #[cfg(feature = "xla-diagnostics")] + #[cfg(any(feature = "xla-diagnostics", feature = "xla-reference-diagnostics"))] fn forward_with_hidden_state_diagnostics( &self, pixel_values: &MlxArray, diff --git a/src/vision/encoders/siglip.rs b/src/vision/encoders/siglip.rs index 86a6a4078..c16800365 100644 --- a/src/vision/encoders/siglip.rs +++ b/src/vision/encoders/siglip.rs @@ -644,7 +644,7 @@ impl VisionEncoder for SigLipVisionModel { self.forward_impl(pixel_values, false).0 } - #[cfg(feature = "xla-diagnostics")] + #[cfg(any(feature = "xla-diagnostics", feature = "xla-reference-diagnostics"))] fn forward_with_hidden_state_diagnostics( &self, pixel_values: &MlxArray, From 087986ff0a2b3435c2c1c357a518dc39f40385b6 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Sun, 26 Jul 2026 23:21:43 +0900 Subject: [PATCH 06/11] fix(xla): bind Gemma3 VLM runtime contract Bind the authoritative prepared-prefill schema, mask and scaling policies, prompt token IDs, vision config, and sliding attention settings into the single and ragged IREE bundle fingerprints through an explicit stable identity. Preserve the existing dense identity verbatim for text-only Gemma3 sessions. Refs #869 --- .../mlxcel-xla/src/emitter/vision_config.rs | 38 +++- src/lib/mlxcel-xla/src/iree.rs | 196 ++++++++++++++++- src/lib/mlxcel-xla/src/lib.rs | 3 +- src/lib/mlxcel-xla/src/prepared_gemma3.rs | 204 ++++++++++++++++++ src/multimodal/host_preprocessor.rs | 4 +- 5 files changed, 433 insertions(+), 12 deletions(-) diff --git a/src/lib/mlxcel-xla/src/emitter/vision_config.rs b/src/lib/mlxcel-xla/src/emitter/vision_config.rs index 6b9cae779..87a659983 100644 --- a/src/lib/mlxcel-xla/src/emitter/vision_config.rs +++ b/src/lib/mlxcel-xla/src/emitter/vision_config.rs @@ -24,6 +24,15 @@ pub(crate) enum VisionActivation { GeluPytorchTanh, } +impl VisionActivation { + const fn stable_identity(self) -> &'static str { + match self { + Self::ExactGelu => "ExactGelu", + Self::GeluPytorchTanh => "GeluPytorchTanh", + } + } +} + #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub(crate) enum VisionProjector { LlavaMlp, @@ -38,6 +47,28 @@ pub(crate) enum VisionProjector { }, } +impl VisionProjector { + fn stable_identity(self) -> String { + match self { + Self::LlavaMlp => "LlavaMlp".to_string(), + Self::Gemma3AvgPool { + tokens_per_side, + kernel_size, + image_token_id, + pad_token_id, + boi_token_id, + eoi_token_id, + newline_token_id, + } => format!( + "Gemma3AvgPool {{ tokens_per_side: {tokens_per_side}, kernel_size: {kernel_size}, \ + image_token_id: {image_token_id}, pad_token_id: {pad_token_id}, \ + boi_token_id: {boi_token_id}, eoi_token_id: {eoi_token_id}, \ + newline_token_id: {newline_token_id} }}" + ), + } + } +} + #[cfg(test)] mod tests { use super::*; @@ -227,6 +258,7 @@ mod tests { } ); assert!(config.fingerprint().contains("image_token_id: 99")); + assert!(config.fingerprint().contains("newline_token_id: 108")); assert!(config.fingerprint().starts_with("iree-vision-v3:")); let specs = config.weight_specs(); assert_eq!( @@ -521,7 +553,7 @@ impl LlavaVisionConfig { }; format!( "{schema}:image={}:patch={}:channels={}:hidden={}:intermediate={}:layers={}:\ - heads={}:eps={:08x}:activation={:?}:class={}:feature={}:drop_first={}:text={}:projector={:?}", + heads={}:eps={:08x}:activation={}:class={}:feature={}:drop_first={}:text={}:projector={}", self.image_size, self.patch_size, self.channels, @@ -530,12 +562,12 @@ impl LlavaVisionConfig { self.layers, self.heads, self.layer_norm_eps.to_bits(), - self.activation, + self.activation.stable_identity(), self.class_token, self.feature_layer, self.drop_first_token, self.text_hidden, - self.projector, + self.projector.stable_identity(), ) } diff --git a/src/lib/mlxcel-xla/src/iree.rs b/src/lib/mlxcel-xla/src/iree.rs index e31100190..22be32c3b 100644 --- a/src/lib/mlxcel-xla/src/iree.rs +++ b/src/lib/mlxcel-xla/src/iree.rs @@ -58,9 +58,9 @@ use mlxcel_core::session::PreparedPrefill; use safetensors::{Dtype, SafeTensors}; use crate::emitter::{ - Config, DeepStackConfig, Gemma3nConfig, Gemma3nWeightSpec, Precision, QuantConfig, - check_packed_supported, emit_decode_ragged_with, emit_decode_with, - emit_gemma3n_decode_ragged_with_qmv, emit_gemma3n_decode_with_qmv, + Config, DeepStackConfig, Gemma3nConfig, Gemma3nWeightSpec, LlavaVisionConfig, Precision, + QuantConfig, VisionProjector, check_packed_supported, emit_decode_ragged_with, + emit_decode_with, emit_gemma3n_decode_ragged_with_qmv, emit_gemma3n_decode_with_qmv, emit_gemma3n_prefill_embeddings_ple_with_qmv, emit_gemma3n_prefill_with_qmv, emit_prefill_embeddings_deepstack_with, emit_prefill_embeddings_with, emit_prefill_with, gemma3n_qmv_artifact_identity, gemma3n_qmv_is_available, gemma3n_weight_specs, quant_in_graph, @@ -76,6 +76,10 @@ use crate::prepared::{ validate_slot, }; use crate::prepared_deepstack::PreparedDeepStack; +use crate::prepared_gemma3::{ + GEMMA3_VLM_NEWLINE_WRAPPER_TOKEN_ID, Gemma3VlmCompatibilityContract, + bind_gemma3_vlm_compatibility, +}; use crate::{DeepStackFeatures, DeepStackPreparedPrefill, Gemma3nDensePle, Gemma3nPreparedPrefill}; // The loader reads the per-architecture checkpoint-weight order from // `weights::weight_specs`, which sources its names from `weight_names::scheme_names` @@ -582,6 +586,177 @@ impl RuntimeConfig { } } +fn gemma3_vlm_compatibility_from_json( + config_json: &str, + cfg: &RuntimeConfig, +) -> Result, String> { + let root: serde_json::Value = + serde_json::from_str(config_json).map_err(|error| format!("parse config.json: {error}"))?; + let is_vlm = root.get("model_type").and_then(serde_json::Value::as_str) == Some("gemma3") + && root + .get("vision_config") + .is_some_and(serde_json::Value::is_object); + if !is_vlm { + return Ok(None); + } + let RuntimeConfig::Dense(text) = cfg else { + return Err("Gemma3 VLM requires the dense Gemma3 text runtime".to_string()); + }; + if !text.embed_scale || !text.embeddings_prefill_uses_authoritative_mask { + return Err( + "Gemma3 VLM requires sqrt(hidden) token scaling and an authoritative embeddings-prefill mask" + .to_string(), + ); + } + let vision = LlavaVisionConfig::from_json_str(config_json)?; + let VisionProjector::Gemma3AvgPool { + image_token_id, + pad_token_id, + boi_token_id, + eoi_token_id, + newline_token_id, + .. + } = vision.projector + else { + return Err("Gemma3 VLM requires the average-pool vision projector".to_string()); + }; + if newline_token_id != GEMMA3_VLM_NEWLINE_WRAPPER_TOKEN_ID { + return Err(format!( + "Gemma3 VLM newline wrapper token {newline_token_id} does not match the prepared-prefill contract {}", + GEMMA3_VLM_NEWLINE_WRAPPER_TOKEN_ID + )); + } + if vision.text_hidden != text.hidden { + return Err(format!( + "Gemma3 VLM projector hidden size {} does not match language hidden size {}", + vision.text_hidden, text.hidden + )); + } + Ok(Some(Gemma3VlmCompatibilityContract::new( + image_token_id, + pad_token_id, + boi_token_id, + eoi_token_id, + vision.image_tokens(), + text.hidden, + text.sliding_window, + text.sliding_pattern, + text.rope_theta, + text.rope_local_base, + vision.fingerprint(), + ))) +} + +fn gemma3_vlm_compatibility( + model_dir: &Path, + cfg: &RuntimeConfig, +) -> Result, String> { + let path = model_dir.join("config.json"); + let config_json = std::fs::read_to_string(&path) + .map_err(|error| format!("read {}: {error}", path.display()))?; + gemma3_vlm_compatibility_from_json(&config_json, cfg) + .map_err(|error| format!("{}: {error}", path.display())) +} + +#[cfg(test)] +mod gemma3_vlm_compatibility_tests { + use super::*; + + fn gemma3_text_config() -> serde_json::Value { + serde_json::json!({ + "model_type": "gemma3_text", + "hidden_size": 12, + "intermediate_size": 24, + "num_hidden_layers": 2, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "head_dim": 6, + "rms_norm_eps": 1e-6, + "rope_theta": 1_000_000.0, + "rope_local_base_freq": 10_000.0, + "sliding_window": 4, + "sliding_window_pattern": 2, + "vocab_size": 128, + "hidden_activation": "gelu_pytorch_tanh" + }) + } + + fn runtime(config_json: &str) -> RuntimeConfig { + RuntimeConfig::Dense(Box::new( + Config::from_json_str(config_json) + .unwrap() + .with_context_capacity(8) + .unwrap(), + )) + } + + #[test] + fn gemma3_wrapper_binds_prepared_and_vision_contracts() { + let config_json = serde_json::json!({ + "model_type": "gemma3", + "image_token_index": 99, + "pad_token_id": 0, + "boi_token_index": 97, + "eoi_token_index": 98, + "mm_tokens_per_image": 1, + "vision_config": { + "model_type": "siglip_vision_model", + "image_size": 28, + "patch_size": 14, + "num_channels": 3, + "hidden_size": 8, + "intermediate_size": 16, + "num_hidden_layers": 2, + "num_attention_heads": 2 + }, + "text_config": gemma3_text_config() + }) + .to_string(); + let runtime = runtime(&config_json); + let contract = gemma3_vlm_compatibility_from_json(&config_json, &runtime) + .unwrap() + .unwrap(); + let identity = contract.stable_identity(); + + for component in [ + "prepared_prefill=gemma3-vlm-owned-f32-embeddings", + "mask_mode=gemma3-vlm-bidirectional-padding-f32-min-v1", + "post_scale_policy=gemma3-vlm-post-scale-text-sqrt-hidden-image-identity-pad-zero-v1", + "newline_wrapper_token_id=108", + "image_token_id=99", + "pad_token_id=0", + "boi_token_id=97", + "eoi_token_id=98", + "sliding_window=4", + "sliding_pattern=2", + "vision_identity=iree-vision-v3:", + ] { + assert!( + identity.contains(component), + "missing {component}: {identity}" + ); + } + assert_ne!( + bind_gemma3_vlm_compatibility(runtime.artifact_identity(), Some(&contract)), + runtime.artifact_identity() + ); + } + + #[test] + fn direct_root_gemma3_text_keeps_existing_runtime_identity() { + let config_json = gemma3_text_config().to_string(); + let runtime = runtime(&config_json); + let contract = gemma3_vlm_compatibility_from_json(&config_json, &runtime).unwrap(); + assert!(contract.is_none()); + + let existing = runtime.artifact_identity(); + assert_eq!( + bind_gemma3_vlm_compatibility(existing.clone(), contract.as_ref()), + existing + ); + } +} + fn effective_precision(device: &str, cfg: &RuntimeConfig) -> Result { match cfg { RuntimeConfig::Dense(_) => resolve_precision_checked(device), @@ -882,6 +1057,7 @@ struct CompiledBundle { fn compile_bundle( device: &str, cfg: &RuntimeConfig, + gemma3_vlm_contract: Option<&Gemma3VlmCompatibilityContract>, gemma3n_qmv: bool, prefill_mlir: &str, prefill_tag: &str, @@ -942,7 +1118,8 @@ fn compile_bundle( )?; let mut fingerprint = std::collections::hash_map::DefaultHasher::new(); "mlxcel-xla-runtime-bundle-v2-explicit-position-mode".hash(&mut fingerprint); - cfg.artifact_identity().hash(&mut fingerprint); + bind_gemma3_vlm_compatibility(cfg.artifact_identity(), gemma3_vlm_contract) + .hash(&mut fingerprint); format!("{:?}", cfg.weight_specs()).hash(&mut fingerprint); format!("{:?}", effective_precision(device, cfg)?).hash(&mut fingerprint); if gemma3n_qmv { @@ -965,7 +1142,11 @@ fn compile_bundle( } /// Emit and compile the argmax prefill bundle for a single-sequence engine. -fn compile_vmfbs(device: &str, cfg: &RuntimeConfig) -> Result { +fn compile_vmfbs( + device: &str, + cfg: &RuntimeConfig, + gemma3_vlm_contract: Option<&Gemma3VlmCompatibilityContract>, +) -> Result { let precision = effective_precision(device, cfg)?; let native_qmv = match cfg { RuntimeConfig::Dense(_) => false, @@ -994,6 +1175,7 @@ fn compile_vmfbs(device: &str, cfg: &RuntimeConfig) -> Result Result { let cfg = RuntimeConfig::from_json(model_dir, context_capacity)?; runtime_ffi_dimensions(&cfg, cfg.weight_specs().len())?; - let bundle = compile_vmfbs(device, &cfg)?; + let gemma3_vlm_contract = gemma3_vlm_compatibility(model_dir, &cfg)?; + let bundle = compile_vmfbs(device, &cfg, gemma3_vlm_contract.as_ref())?; let ctx = create_ctx(model_dir, &cfg, device, &bundle)?; Ok(Self { ctx, @@ -2265,6 +2448,7 @@ impl IreeRaggedLlama { let bundle = compile_bundle( device, &cfg, + gemma3_vlm_compatibility(model_dir, &cfg)?.as_ref(), native_qmv, &prefill_mlir, "prefill_logits", diff --git a/src/lib/mlxcel-xla/src/lib.rs b/src/lib/mlxcel-xla/src/lib.rs index 932d6e341..04c8f73a5 100644 --- a/src/lib/mlxcel-xla/src/lib.rs +++ b/src/lib/mlxcel-xla/src/lib.rs @@ -212,7 +212,8 @@ pub use emitter::{Gemma3nDiagnosticLayout, Gemma3nDiagnosticSegment}; pub use prepared::PreparedInputError; pub use prepared_deepstack::{DeepStackFeatures, DeepStackInputError, DeepStackPreparedPrefill}; pub use prepared_gemma3::{ - GEMMA3_VLM_MASK_MODE, GEMMA3_VLM_MASKED_VALUE, Gemma3VlmPreparedError, + GEMMA3_VLM_MASK_MODE, GEMMA3_VLM_MASKED_VALUE, GEMMA3_VLM_NEWLINE_WRAPPER_TOKEN_ID, + GEMMA3_VLM_POST_SCALE_POLICY, GEMMA3_VLM_PREPARED_PREFILL_CONTRACT, Gemma3VlmPreparedError, prepare_gemma3_vlm_prefill, }; pub use prepared_gemma3n::{Gemma3nDensePle, Gemma3nDensePleError, Gemma3nPreparedPrefill}; diff --git a/src/lib/mlxcel-xla/src/prepared_gemma3.rs b/src/lib/mlxcel-xla/src/prepared_gemma3.rs index 306f00eca..cd2691b5f 100644 --- a/src/lib/mlxcel-xla/src/prepared_gemma3.rs +++ b/src/lib/mlxcel-xla/src/prepared_gemma3.rs @@ -33,6 +33,145 @@ pub const GEMMA3_VLM_MASKED_VALUE: f32 = f32::MIN; /// Stable identity component for the qualified external-mask behavior. pub const GEMMA3_VLM_MASK_MODE: &str = "gemma3-vlm-bidirectional-padding-f32-min-v1"; +/// Stable identity for the embedding values handed to embeddings-prefill. +/// +/// Text rows have already received Gemma's `sqrt(hidden_size)` scale, projected +/// image rows stay at projector magnitude, and padding rows are zero. +pub const GEMMA3_VLM_POST_SCALE_POLICY: &str = + "gemma3-vlm-post-scale-text-sqrt-hidden-image-identity-pad-zero-v1"; + +/// Gemma3Processor wraps every expanded image block with this newline token. +pub const GEMMA3_VLM_NEWLINE_WRAPPER_TOKEN_ID: i32 = 108; + +/// Stable schema for the owned payload consumed by the authoritative +/// embeddings-prefill entry. +pub const GEMMA3_VLM_PREPARED_PREFILL_CONTRACT: &str = + "gemma3-vlm-owned-f32-embeddings-sequential-1d-additive-f32-1x1xlxl-noncausal-v1"; + +/// Cross-runtime compatibility contract for a Gemma3 VLM bundle. +/// +/// The language graphs, resident SigLIP/projector module, and host producer are +/// separate runtime components. This explicit identity prevents their semantic +/// boundary from depending on derived `Debug` output. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct Gemma3VlmCompatibilityContract { + pub(crate) prepared_prefill: String, + pub(crate) mask_mode: String, + pub(crate) masked_value_bits: u32, + pub(crate) post_scale_policy: String, + pub(crate) newline_wrapper_token_id: i32, + pub(crate) image_token_id: i32, + pub(crate) pad_token_id: i32, + pub(crate) boi_token_id: i32, + pub(crate) eoi_token_id: i32, + pub(crate) tokens_per_image: usize, + pub(crate) hidden_size: usize, + pub(crate) sliding_window: Option, + pub(crate) sliding_pattern: usize, + pub(crate) global_rope_theta_bits: u64, + pub(crate) local_rope_base_bits: Option, + pub(crate) vision_identity: String, +} + +impl Gemma3VlmCompatibilityContract { + #[allow(clippy::too_many_arguments)] + pub(crate) fn new( + image_token_id: i32, + pad_token_id: i32, + boi_token_id: i32, + eoi_token_id: i32, + tokens_per_image: usize, + hidden_size: usize, + sliding_window: Option, + sliding_pattern: usize, + global_rope_theta: f64, + local_rope_base: Option, + vision_identity: String, + ) -> Self { + Self { + prepared_prefill: GEMMA3_VLM_PREPARED_PREFILL_CONTRACT.to_string(), + mask_mode: GEMMA3_VLM_MASK_MODE.to_string(), + masked_value_bits: GEMMA3_VLM_MASKED_VALUE.to_bits(), + post_scale_policy: GEMMA3_VLM_POST_SCALE_POLICY.to_string(), + newline_wrapper_token_id: GEMMA3_VLM_NEWLINE_WRAPPER_TOKEN_ID, + image_token_id, + pad_token_id, + boi_token_id, + eoi_token_id, + tokens_per_image, + hidden_size, + sliding_window, + sliding_pattern, + global_rope_theta_bits: global_rope_theta.to_bits(), + local_rope_base_bits: local_rope_base.map(f64::to_bits), + vision_identity, + } + } + + /// Canonical, field-labelled identity. Integer bit patterns avoid + /// locale/toolchain-dependent float formatting. + pub(crate) fn stable_identity(&self) -> String { + let sliding_window = self + .sliding_window + .map_or_else(|| "none".to_string(), |value| value.to_string()); + let local_rope_base = self + .local_rope_base_bits + .map_or_else(|| "none".to_string(), |bits| format!("{bits:016x}")); + format!( + "gemma3-vlm-runtime-bundle-v1\ + ;prepared_prefill={}\ + ;mask_mode={}\ + ;masked_value_bits={:08x}\ + ;post_scale_policy={}\ + ;newline_wrapper_token_id={}\ + ;image_token_id={}\ + ;pad_token_id={}\ + ;boi_token_id={}\ + ;eoi_token_id={}\ + ;tokens_per_image={}\ + ;hidden_size={}\ + ;sliding_window={}\ + ;sliding_pattern={}\ + ;global_rope_theta_bits={:016x}\ + ;local_rope_base_bits={}\ + ;vision_identity={}", + self.prepared_prefill, + self.mask_mode, + self.masked_value_bits, + self.post_scale_policy, + self.newline_wrapper_token_id, + self.image_token_id, + self.pad_token_id, + self.boi_token_id, + self.eoi_token_id, + self.tokens_per_image, + self.hidden_size, + sliding_window, + self.sliding_pattern, + self.global_rope_theta_bits, + local_rope_base, + self.vision_identity, + ) + } +} + +/// Add the prepared-prefill identity only for a qualified multimodal runtime. +/// +/// Returning the input verbatim for `None` is intentional: ordinary text-only +/// bundle identities must not move when this VLM-only contract evolves. +pub(crate) fn bind_gemma3_vlm_compatibility( + runtime_identity: String, + contract: Option<&Gemma3VlmCompatibilityContract>, +) -> String { + match contract { + Some(contract) => format!( + "{runtime_identity};prepared_prefill_contract={}", + contract.stable_identity() + ), + None => runtime_identity, + } +} + #[derive(Debug, Clone, PartialEq, Eq)] pub enum Gemma3VlmPreparedError { Empty, @@ -346,6 +485,71 @@ mod tests { .collect() } + fn compatibility_contract() -> Gemma3VlmCompatibilityContract { + Gemma3VlmCompatibilityContract::new( + 262_144, + 0, + 255_999, + 256_000, + 256, + 2_560, + Some(1_024), + 6, + 1_000_000.0, + Some(10_000.0), + "iree-vision-v3:fixture".to_string(), + ) + } + + #[test] + fn compatibility_identity_binds_every_cross_runtime_field() { + let base = compatibility_contract(); + let identity = base.stable_identity(); + let mut variants = Vec::new(); + + macro_rules! changed { + ($field:ident, $value:expr) => {{ + let mut variant = base.clone(); + variant.$field = $value; + variants.push(variant); + }}; + } + + changed!(prepared_prefill, "prepared-v2".to_string()); + changed!(mask_mode, "causal-mask".to_string()); + changed!(masked_value_bits, 0); + changed!(post_scale_policy, "double-scale".to_string()); + changed!(newline_wrapper_token_id, 109); + changed!(image_token_id, 262_145); + changed!(pad_token_id, 1); + changed!(boi_token_id, 255_998); + changed!(eoi_token_id, 256_001); + changed!(tokens_per_image, 255); + changed!(hidden_size, 2_561); + changed!(sliding_window, Some(2_048)); + changed!(sliding_pattern, 5); + changed!(global_rope_theta_bits, 1_000_001.0f64.to_bits()); + changed!(local_rope_base_bits, Some(10_001.0f64.to_bits())); + changed!(vision_identity, "iree-vision-v3:changed".to_string()); + + for variant in variants { + assert_ne!(variant.stable_identity(), identity); + } + } + + #[test] + fn text_only_runtime_identity_is_byte_for_byte_unchanged() { + let identity = "dense:ordinary-text-config".to_string(); + assert_eq!( + bind_gemma3_vlm_compatibility(identity.clone(), None), + identity + ); + assert_ne!( + bind_gemma3_vlm_compatibility(identity.clone(), Some(&compatibility_contract())), + identity + ); + } + #[test] fn exports_post_scale_text_unscaled_images_and_exact_bidirectional_mask() { let prepared = prepare_gemma3_vlm_prefill( diff --git a/src/multimodal/host_preprocessor.rs b/src/multimodal/host_preprocessor.rs index 7f21e3444..a18cefe3f 100644 --- a/src/multimodal/host_preprocessor.rs +++ b/src/multimodal/host_preprocessor.rs @@ -600,8 +600,8 @@ impl Gemma3IreeHostPreprocessor { has_bos: true, separator_token_id: None, suffix_tokens: Vec::new(), - block_prefix_tokens: vec![108], - block_suffix_tokens: vec![108], + block_prefix_tokens: vec![mlxcel_xla::GEMMA3_VLM_NEWLINE_WRAPPER_TOKEN_ID], + block_suffix_tokens: vec![mlxcel_xla::GEMMA3_VLM_NEWLINE_WRAPPER_TOKEN_ID], } } From fe02885ba8cfe0dbced8778690cbae00a7b8514e Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 00:51:07 +0900 Subject: [PATCH 07/11] fix(xla): add standalone Gemma3 boundary runner Extract the pinned Gemma3 eager MLX and IREE comparison into a feature-gated library entry point and a dedicated CUDA example that avoids the libtest harness. Keep the existing comparison stages and tolerances while adding flushed stage progress and 60-second heartbeats around long-running model work. --- Cargo.toml | 6 ++ examples/xla_gemma3_reference_check.rs | 58 +++++++++++ src/lib.rs | 2 + src/loading/mod.rs | 2 + src/loading/vlm.rs | 2 + src/loading/vlm_gemma.rs | 5 +- src/loading/vlm_gemma_xla_tests.rs | 136 ++++++++++++++++++------- 7 files changed, 173 insertions(+), 38 deletions(-) create mode 100644 examples/xla_gemma3_reference_check.rs diff --git a/Cargo.toml b/Cargo.toml index 32924c1e7..e887e38b0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -348,3 +348,9 @@ required-features = ["xla-micro-oracle"] [[example]] name = "xla_phi4_audio_check" required-features = ["xla-iree"] + +# Pinned #869 eager MLX CUDA versus IREE local-task Gemma3 boundary gate. +# This is a standalone executable so running it does not link libtest. +[[example]] +name = "xla_gemma3_reference_check" +required-features = ["cuda", "xla-reference-diagnostics"] diff --git a/examples/xla_gemma3_reference_check.rs b/examples/xla_gemma3_reference_check.rs new file mode 100644 index 000000000..5367c70fc --- /dev/null +++ b/examples/xla_gemma3_reference_check.rs @@ -0,0 +1,58 @@ +// 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. + +//! Standalone pinned Gemma3 eager-MLX-CUDA versus IREE-local-task boundary gate. +//! +//! This executable intentionally avoids the Rust libtest harness. It accepts +//! command-line arguments first and falls back to the historical environment +//! variables used by the ignored test. +//! +//! ```text +//! IREE_DIST=/path/to/iree-dist \ +//! cargo run --example xla_gemma3_reference_check \ +//! --features cuda,xla-reference-diagnostics -- \ +//! --model /path/to/gemma-3-4b-it-4bit \ +//! --image tests/fixtures/test_image.png \ +//! --device local-task +//! ``` + +use std::path::PathBuf; + +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, variable: &str) -> PathBuf { + argument(flag) + .or_else(|| std::env::var(variable).ok()) + .map(PathBuf::from) + .unwrap_or_else(|| panic!("missing required {flag} or {variable}")) +} + +fn main() { + let model = required_path("--model", "MLXCEL_GEMMA3_FIXTURE"); + let image = argument("--image") + .or_else(|| std::env::var("MLXCEL_GEMMA3_IMAGE").ok()) + .map(PathBuf::from) + .unwrap_or_else(|| PathBuf::from("tests/fixtures/test_image.png")); + let device = argument("--device") + .or_else(|| std::env::var("MLXCEL_XLA_DEVICE").ok()) + .unwrap_or_else(|| "local-task".to_string()); + + mlxcel::run_gemma3_eager_mlx_iree_prepared_boundary(&model, &image, &device); +} diff --git a/src/lib.rs b/src/lib.rs index 681bbda28..daec017d6 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -123,6 +123,8 @@ pub fn decode_image_payloads_with_limits( // Re-export split modules pub use loaded_model::LoadedModel; pub use loaded_model_capabilities::VlmRuntimeRef; +#[cfg(feature = "xla-reference-diagnostics")] +pub use loading::run_gemma3_eager_mlx_iree_prepared_boundary; pub use loading::{ context_window_from_config, load_model, load_model_with_adapter, load_model_with_tensor_parallel, load_qwen3_omni_speech, read_eos_token_ids, diff --git a/src/loading/mod.rs b/src/loading/mod.rs index d3a3c5cf5..02c76e205 100644 --- a/src/loading/mod.rs +++ b/src/loading/mod.rs @@ -64,6 +64,8 @@ pub(crate) use self::vlm::load_llava_iree_host_preprocessor; #[cfg(feature = "xla-iree")] pub(crate) use self::vlm::load_qwen2_vl_iree_host_preprocessor; pub use self::vlm::load_qwen3_omni_speech; +#[cfg(feature = "xla-reference-diagnostics")] +pub use self::vlm::run_gemma3_eager_mlx_iree_prepared_boundary; #[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 e2f41894b..acc186817 100644 --- a/src/loading/vlm.rs +++ b/src/loading/vlm.rs @@ -99,6 +99,8 @@ pub(crate) use ernie4_5_vl::load_ernie4_5_moe_vlm; pub(crate) use fastvlm::load_fastvlm_vlm; #[cfg(feature = "xla-iree")] pub(crate) use gemma::load_gemma3_iree_host_preprocessor; +#[cfg(feature = "xla-reference-diagnostics")] +pub use gemma::run_gemma3_eager_mlx_iree_prepared_boundary; pub(crate) use gemma::{load_gemma3_vlm, load_gemma3n_vlm, load_gemma4_vlm}; pub(crate) use gemma_unified::load_gemma4_unified; pub(crate) use granite_vision::load_granite_vision_vlm; diff --git a/src/loading/vlm_gemma.rs b/src/loading/vlm_gemma.rs index 663d14010..f64c0ea64 100644 --- a/src/loading/vlm_gemma.rs +++ b/src/loading/vlm_gemma.rs @@ -329,10 +329,13 @@ pub(crate) fn load_gemma3_iree_host_preprocessor( ) } -#[cfg(all(test, feature = "xla-iree"))] +#[cfg(all(feature = "xla-iree", any(test, feature = "xla-reference-diagnostics")))] #[path = "vlm_gemma_xla_tests.rs"] mod xla_tests; +#[cfg(feature = "xla-reference-diagnostics")] +pub use xla_tests::reference_boundary::run_gemma3_eager_mlx_iree_prepared_boundary; + /// Load a Gemma3n VLM model. pub(crate) fn load_gemma3n_vlm(model_path: &Path) -> Result { use vision::encoders::gemma3n::load_gemma3n_vision; diff --git a/src/loading/vlm_gemma_xla_tests.rs b/src/loading/vlm_gemma_xla_tests.rs index 6f0a14a45..5c0215a55 100644 --- a/src/loading/vlm_gemma_xla_tests.rs +++ b/src/loading/vlm_gemma_xla_tests.rs @@ -33,10 +33,15 @@ fn pinned_gemma3_projector_loads_and_returns_finite_features() { } #[cfg(feature = "xla-reference-diagnostics")] -mod reference_boundary { +pub mod reference_boundary { use std::fs; use std::io::{self, Write}; - use std::path::{Path, PathBuf}; + use std::path::Path; + #[cfg(test)] + use std::path::PathBuf; + use std::sync::mpsc::{self, RecvTimeoutError, Sender}; + use std::thread::{self, JoinHandle}; + use std::time::{Duration, Instant}; use mlxcel_core::layers::UnifiedEmbedding; use mlxcel_core::session::{OwnedTensor, PreparedPrefill, PreparedTensorDType}; @@ -151,6 +156,45 @@ mod reference_boundary { io::stderr().flush().expect("flush diagnostic progress"); } + struct ProgressHeartbeat { + stop: Option>, + worker: Option>, + } + + impl ProgressHeartbeat { + fn start(stage: &str) -> Self { + progress(stage); + let (stop, receiver) = mpsc::channel(); + let stage = stage.to_string(); + let started = Instant::now(); + let worker = thread::spawn(move || { + while let Err(RecvTimeoutError::Timeout) = + receiver.recv_timeout(Duration::from_secs(60)) + { + progress(&format!( + "heartbeat stage={stage} elapsed={}s", + started.elapsed().as_secs() + )); + } + }); + Self { + stop: Some(stop), + worker: Some(worker), + } + } + } + + impl Drop for ProgressHeartbeat { + fn drop(&mut self) { + if let Some(stop) = self.stop.take() { + let _ = stop.send(()); + } + if let Some(worker) = self.worker.take() { + let _ = worker.join(); + } + } + } + fn compare_stage( stage: &str, observed: &[f32], @@ -326,6 +370,7 @@ mod reference_boundary { expected } + #[cfg(test)] #[test] fn comparison_reports_the_first_failed_element() { let stats = comparison_stats( @@ -341,40 +386,29 @@ mod reference_boundary { assert_eq!(stats.first_failure, Some(1)); } - /// Pinned mixed-runtime boundary gate for #869. + /// Run the pinned mixed-runtime boundary gate for #869. /// - /// This deliberately remains ignored. Run it from the repository root only - /// when the pinned checkpoint and local IREE distribution are installed: + /// This entry point loads only the eager MLX SigLIP/projector/text embedding + /// weights and the resident IREE vision projector. The caller selects the + /// MLX runtime at build time; the dedicated example requires CUDA while the + /// IREE side remains pinned to `local-task`. /// - /// ```text - /// IREE_DIST=/path/to/iree-dist \ - /// MLXCEL_GEMMA3_FIXTURE=/path/to/gemma-3-4b-it-4bit \ - /// cargo test --features xla-reference-diagnostics --lib \ - /// pinned_gemma3_eager_mlx_matches_iree_prepared_boundary -- \ - /// --ignored --nocapture - /// ``` + /// # Panics /// - /// Both eager MLX and IREE execute on CPU; no language decoder is constructed. - #[test] - #[ignore = "requires pinned Gemma3 checkpoint, image, and IREE local-task"] - fn pinned_gemma3_eager_mlx_matches_iree_prepared_boundary() { - let model = PathBuf::from( - std::env::var("MLXCEL_GEMMA3_FIXTURE") - .expect("MLXCEL_GEMMA3_FIXTURE must name the pinned checkpoint"), - ); - let image_path = std::env::var("MLXCEL_GEMMA3_IMAGE") - .map(PathBuf::from) - .unwrap_or_else(|_| PathBuf::from("tests/fixtures/test_image.png")); - let device = - std::env::var("MLXCEL_XLA_DEVICE").unwrap_or_else(|_| "local-task".to_string()); + /// Panics if the pinned fixture identity or any ordered comparison differs. + pub fn run_gemma3_eager_mlx_iree_prepared_boundary( + model: &Path, + image_path: &Path, + device: &str, + ) { assert_eq!( device, "local-task", "the pinned #869 mixed-runtime gate qualifies IREE local-task only" ); - progress("validate pinned checkpoint and image"); + let _heartbeat = ProgressHeartbeat::start("validate pinned checkpoint and image"); assert_eq!( - pinned_revision(&model), + pinned_revision(model), PINNED_REVISION, "checkpoint revision differs from pinned mlx-community/gemma-3-4b-it-4bit" ); @@ -393,17 +427,19 @@ mod reference_boundary { PINNED_PROCESSOR_SHA256, "checkpoint processor", ); - assert_sha256(&image_path, PINNED_IMAGE_SHA256, "image"); + assert_sha256(image_path, PINNED_IMAGE_SHA256, "image"); + drop(_heartbeat); - progress("load filtered MLX SigLIP/projector/text embedding weights"); + let _heartbeat = + ProgressHeartbeat::start("load filtered MLX SigLIP/projector/text embedding weights"); let (_config_str, full_config) = - read_sanitized_vlm_config(&model).expect("read pinned Gemma3 config"); + read_sanitized_vlm_config(model).expect("read pinned Gemma3 config"); let config: VLMConfig = serde_json::from_value(full_config.clone()).expect("parse pinned Gemma3 VLM config"); let text_config: models::gemma3::ModelArgs = serde_json::from_value(config.text_config.clone()) .expect("parse pinned Gemma3 text config"); - let weights = load_vlm_weights_common_filtered_canonical(&model, reference_weight) + let weights = load_vlm_weights_common_filtered_canonical(model, reference_weight) .map(strip_language_model_prefix) .expect("load only Gemma3 vision/projector/embedding weights"); let quant_group_size = full_config @@ -438,9 +474,10 @@ mod reference_boundary { config.vision_config.layer_norm_eps, ) .expect("load filtered Gemma3 average-pool projector"); - let image = image::open(&image_path) + let image = image::open(image_path) .unwrap_or_else(|error| panic!("decode {}: {error}", image_path.display())); let images = vec![image]; + drop(_heartbeat); progress("compare independently constructed processor pixels"); let mlx_processor = SigLipProcessor::new(config.vision_config.image_size); @@ -483,7 +520,8 @@ mod reference_boundary { let embed_dtype = mlxcel_core::array_dtype(&raw_text_array); let raw_text = mlx_f32(&raw_text_array, "raw text embeddings"); - progress("capture eager MLX SigLIP hidden and projector stages"); + let _heartbeat = + ProgressHeartbeat::start("capture eager MLX SigLIP hidden and projector stages"); let mlx_vision_input = mlxcel_core::astype( &mlxcel_core::transpose_axes(&mlx_pixels, &[0, 2, 3, 1]), embed_dtype, @@ -506,9 +544,11 @@ mod reference_boundary { .map(|stage| mlx_f32(stage, "MLX SigLIP block 0 stage")) .collect::>(); assert_eq!(mlx_block0_values.len(), BLOCK0_STAGES.len()); + drop(_heartbeat); - progress("run IREE diagnostic SigLIP and average-pool projector"); - let mut diagnostic = mlxcel_xla::IreeVisionDiagnosticProjector::load(&model, &device) + let _heartbeat = + ProgressHeartbeat::start("run IREE diagnostic SigLIP and average-pool projector"); + let mut diagnostic = mlxcel_xla::IreeVisionDiagnosticProjector::load(model, device) .expect("load Gemma3 IREE diagnostic projector"); let iree = diagnostic .project(&iree_pixel_values) @@ -557,8 +597,11 @@ mod reference_boundary { VISION_TOLERANCE, &mut first_divergence, ); + drop(_heartbeat); - progress("construct MLX-reference and production IREE prepared prefills"); + let _heartbeat = ProgressHeartbeat::start( + "construct MLX-reference and production IREE prepared prefills", + ); let attention_mask = logical_tokens .iter() .map(|token| i32::from(*token != config.pad_token_id)) @@ -575,7 +618,7 @@ mod reference_boundary { images.len(), ) .expect("construct eager MLX reference prepared prefill"); - let production = Gemma3IreeHostPreprocessor::load(&model, &device) + let production = Gemma3IreeHostPreprocessor::load(model, device) .expect("load production Gemma3 IREE host preprocessor") .prepare( &[ @@ -591,6 +634,7 @@ mod reference_boundary { assert_eq!(production.token_ids, logical_tokens); assert_eq!(production.positions, mlx_prepared.positions); assert_eq!(production.modalities, mlx_prepared.modalities); + drop(_heartbeat); progress("compare final projected image rows and one-time scaling"); let mlx_image_rows = image_rows( @@ -694,4 +738,22 @@ mod reference_boundary { } progress("PASS all pinned Gemma3 eager MLX/IREE boundary stages"); } + + /// Retained libtest wrapper for developers who already use the historical + /// ignored gate. The dedicated example avoids linking the libtest harness. + #[cfg(test)] + #[test] + #[ignore = "requires pinned Gemma3 checkpoint, image, MLX CUDA, and IREE local-task"] + fn pinned_gemma3_eager_mlx_matches_iree_prepared_boundary() { + let model = PathBuf::from( + std::env::var("MLXCEL_GEMMA3_FIXTURE") + .expect("MLXCEL_GEMMA3_FIXTURE must name the pinned checkpoint"), + ); + let image_path = std::env::var("MLXCEL_GEMMA3_IMAGE") + .map(PathBuf::from) + .unwrap_or_else(|_| PathBuf::from("tests/fixtures/test_image.png")); + let device = + std::env::var("MLXCEL_XLA_DEVICE").unwrap_or_else(|_| "local-task".to_string()); + run_gemma3_eager_mlx_iree_prepared_boundary(&model, &image_path, &device); + } } From e9e6ff24ee6f8866aaa976b960241dd4456d4788 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 02:46:12 +0900 Subject: [PATCH 08/11] fix(xla): constrain Gemma3 diagnostic task topology Configure IREE's process-global task topology exactly once before the standalone Gemma3 reference runner creates its first local-task instance, limiting diagnostics to one topology group after the bounded actual reached IREE and failed worker creation with EINVAL. Keep the synthetic flag parser behind mlxcel-xla's diagnostics feature and leave shared xla_aux/xla_iree production startup unchanged. Cache both success and failure, propagate configuration errors before instance creation, and cover ordering, repeated calls, and production-path isolation. Validation: rustfmt check; focused diagnostic runner and topology tests; xla-reference-diagnostics clippy; CUDA standalone example cargo check; linked-symbol inspection and native parser invocation. Refs #869 --- src/loading/vlm_gemma_xla_tests.rs | 71 +++++++++++++++++++++++++++++- 1 file changed, 69 insertions(+), 2 deletions(-) diff --git a/src/loading/vlm_gemma_xla_tests.rs b/src/loading/vlm_gemma_xla_tests.rs index 5c0215a55..6baa32f43 100644 --- a/src/loading/vlm_gemma_xla_tests.rs +++ b/src/loading/vlm_gemma_xla_tests.rs @@ -156,6 +156,14 @@ pub mod reference_boundary { io::stderr().flush().expect("flush diagnostic progress"); } + fn create_after_diagnostic_iree_configuration( + configure: impl FnOnce() -> Result<(), String>, + create: impl FnOnce() -> Result, + ) -> Result { + configure()?; + create() + } + struct ProgressHeartbeat { stop: Option>, worker: Option>, @@ -548,8 +556,11 @@ pub mod reference_boundary { let _heartbeat = ProgressHeartbeat::start("run IREE diagnostic SigLIP and average-pool projector"); - let mut diagnostic = mlxcel_xla::IreeVisionDiagnosticProjector::load(model, device) - .expect("load Gemma3 IREE diagnostic projector"); + let mut diagnostic = create_after_diagnostic_iree_configuration( + mlxcel_xla::configure_diagnostic_local_task_single_group, + || mlxcel_xla::IreeVisionDiagnosticProjector::load(model, device), + ) + .expect("configure and load Gemma3 IREE diagnostic projector"); let iree = diagnostic .project(&iree_pixel_values) .expect("execute Gemma3 IREE diagnostic projector"); @@ -756,4 +767,60 @@ pub mod reference_boundary { std::env::var("MLXCEL_XLA_DEVICE").unwrap_or_else(|_| "local-task".to_string()); run_gemma3_eager_mlx_iree_prepared_boundary(&model, &image_path, &device); } + + #[cfg(test)] + #[test] + fn diagnostic_runner_configures_topology_before_creating_iree() { + use std::cell::Cell; + + let configured = Cell::new(false); + let result = create_after_diagnostic_iree_configuration( + || { + configured.set(true); + Ok(()) + }, + || { + assert!(configured.get()); + Ok("created") + }, + ); + assert_eq!(result.as_deref(), Ok("created")); + } + + #[cfg(test)] + #[test] + fn diagnostic_runner_propagates_configuration_failure_before_creation() { + let created = std::cell::Cell::new(false); + let result = create_after_diagnostic_iree_configuration( + || Err("configuration failed".to_string()), + || { + created.set(true); + Ok(()) + }, + ); + assert_eq!(result, Err("configuration failed".to_string())); + assert!(!created.get()); + } + + #[cfg(test)] + #[test] + fn diagnostics_topology_override_stays_out_of_production_iree_paths() { + let production_rust = include_str!("../lib/mlxcel-xla/src/aux.rs"); + let production_aux_c = include_str!("../lib/mlxcel-xla/csrc/xla_aux.c"); + let production_iree_c = include_str!("../lib/mlxcel-xla/csrc/xla_iree.c"); + for source in [production_rust, production_aux_c, production_iree_c] { + assert!(!source.contains("task_topology_group_count")); + assert!(!source.contains("configure_diagnostic_local_task_single_group")); + } + } + + #[cfg(test)] + #[test] + fn native_diagnostic_topology_configuration_is_reusable() { + mlxcel_xla::configure_diagnostic_local_task_single_group() + .expect("configure diagnostics-only IREE topology"); + assert!(mlxcel_xla::diagnostic_local_task_single_group_is_configured()); + mlxcel_xla::configure_diagnostic_local_task_single_group() + .expect("reuse diagnostics-only IREE topology configuration"); + } } From ea23736cdf573b3497dadba72afbf3574dc2c6c6 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 02:54:54 +0900 Subject: [PATCH 09/11] fix(xla): use host stack for diagnostic workers The bounded standalone run at 3e5edabe passed every eager MLX stage but still failed when IREE local-task called pthread_create with EINVAL. The pinned runtime requests a 128 KiB worker stack, exactly this host's PTHREAD_STACK_MIN, through a pthread attribute that the host rejects. Extend the diagnostics-only synthetic argv with `--task_worker_stack_size=0` alongside the existing single-group flag so pthreads use the host-default stack. Rename the one-shot helper for its complete threading role while preserving cached success/failure, error propagation, and production startup isolation. Validation: rustfmt check; five focused threading configuration regressions; native two-flag parser invocation and reuse; linked-symbol inspection; xla-reference-diagnostics Clippy with existing warnings only. Refs #869 --- src/loading/vlm_gemma_xla_tests.rs | 30 ++++++++++++++++++++---------- 1 file changed, 20 insertions(+), 10 deletions(-) diff --git a/src/loading/vlm_gemma_xla_tests.rs b/src/loading/vlm_gemma_xla_tests.rs index 6baa32f43..a5a9bf558 100644 --- a/src/loading/vlm_gemma_xla_tests.rs +++ b/src/loading/vlm_gemma_xla_tests.rs @@ -557,7 +557,7 @@ pub mod reference_boundary { let _heartbeat = ProgressHeartbeat::start("run IREE diagnostic SigLIP and average-pool projector"); let mut diagnostic = create_after_diagnostic_iree_configuration( - mlxcel_xla::configure_diagnostic_local_task_single_group, + mlxcel_xla::configure_diagnostic_local_task_threads, || mlxcel_xla::IreeVisionDiagnosticProjector::load(model, device), ) .expect("configure and load Gemma3 IREE diagnostic projector"); @@ -770,7 +770,7 @@ pub mod reference_boundary { #[cfg(test)] #[test] - fn diagnostic_runner_configures_topology_before_creating_iree() { + fn diagnostic_runner_configures_threads_before_creating_iree() { use std::cell::Cell; let configured = Cell::new(false); @@ -804,23 +804,33 @@ pub mod reference_boundary { #[cfg(test)] #[test] - fn diagnostics_topology_override_stays_out_of_production_iree_paths() { + fn diagnostics_thread_overrides_stay_out_of_production_iree_paths() { let production_rust = include_str!("../lib/mlxcel-xla/src/aux.rs"); let production_aux_c = include_str!("../lib/mlxcel-xla/csrc/xla_aux.c"); let production_iree_c = include_str!("../lib/mlxcel-xla/csrc/xla_iree.c"); for source in [production_rust, production_aux_c, production_iree_c] { assert!(!source.contains("task_topology_group_count")); - assert!(!source.contains("configure_diagnostic_local_task_single_group")); + assert!(!source.contains("task_worker_stack_size")); + assert!(!source.contains("configure_diagnostic_local_task_threads")); } } #[cfg(test)] #[test] - fn native_diagnostic_topology_configuration_is_reusable() { - mlxcel_xla::configure_diagnostic_local_task_single_group() - .expect("configure diagnostics-only IREE topology"); - assert!(mlxcel_xla::diagnostic_local_task_single_group_is_configured()); - mlxcel_xla::configure_diagnostic_local_task_single_group() - .expect("reuse diagnostics-only IREE topology configuration"); + fn diagnostic_iree_thread_flags_pin_group_and_use_host_stack_default() { + let diagnostic_c = include_str!("../lib/mlxcel-xla/csrc/xla_diagnostic_flags.c"); + assert!(diagnostic_c.contains("--task_topology_group_count=1")); + assert!(diagnostic_c.contains("--task_worker_stack_size=0")); + assert!(diagnostic_c.contains("if (argc != 1)")); + } + + #[cfg(test)] + #[test] + fn native_diagnostic_thread_configuration_is_reusable() { + mlxcel_xla::configure_diagnostic_local_task_threads() + .expect("configure diagnostics-only IREE task threads"); + assert!(mlxcel_xla::diagnostic_local_task_threads_are_configured()); + mlxcel_xla::configure_diagnostic_local_task_threads() + .expect("reuse diagnostics-only IREE task thread configuration"); } } From a050cc138e82421abdd109f118765e3781724a7b Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 03:18:35 +0900 Subject: [PATCH 10/11] test(xla): add SigLIP bisection checkpoints Compare the existing eager and IREE hidden-state captures through shared block 6, 13, 20, and final-block descriptors so one bounded run partitions blocks 1 through 26 into four balanced first-failure intervals without changing the production graph, numerical tolerances, or diagnostic worker flags. Emit the ordered checkpoint statuses and inferred inclusive block range, and lock the exact labels, indices, ordering, and interval widths with focused tests. Validated the diagnostics reference-boundary tests, rustfmt, and Clippy with the xla-reference-diagnostics feature. Refs #869 --- src/loading/vlm_gemma_xla_tests.rs | 220 +++++++++++++++++++++++++++-- 1 file changed, 205 insertions(+), 15 deletions(-) diff --git a/src/loading/vlm_gemma_xla_tests.rs b/src/loading/vlm_gemma_xla_tests.rs index a5a9bf558..ab39ae842 100644 --- a/src/loading/vlm_gemma_xla_tests.rs +++ b/src/loading/vlm_gemma_xla_tests.rs @@ -83,6 +83,40 @@ pub mod reference_boundary { "siglip.block0.mlp_fc2", "siglip.block0.output", ]; + const PINNED_SIGLIP_BLOCK_COUNT: usize = 27; + + #[derive(Debug, Clone, Copy, PartialEq, Eq)] + struct SiglipHiddenCheckpoint { + label: &'static str, + block_index: usize, + hidden_state_index: usize, + } + + // Block 0 already has full sub-stage coverage. These ordered checkpoints + // split blocks 1..=26 into balanced inclusive ranges 1..=6, 7..=13, + // 14..=20, and 21..=26 while reusing the existing final-state comparison. + const SIGLIP_HIDDEN_BISECTION_CHECKPOINTS: [SiglipHiddenCheckpoint; 4] = [ + SiglipHiddenCheckpoint { + label: "siglip.hidden.block6.output", + block_index: 6, + hidden_state_index: 7, + }, + SiglipHiddenCheckpoint { + label: "siglip.hidden.block13.output", + block_index: 13, + hidden_state_index: 14, + }, + SiglipHiddenCheckpoint { + label: "siglip.hidden.block20.output", + block_index: 20, + hidden_state_index: 21, + }, + SiglipHiddenCheckpoint { + label: "siglip.hidden.last_pre_layernorm", + block_index: 26, + hidden_state_index: 27, + }, + ]; #[derive(Debug, Clone, Copy)] struct Tolerance { @@ -209,15 +243,16 @@ pub mod reference_boundary { reference: &[f32], tolerance: Tolerance, first_divergence: &mut Option, - ) { + ) -> bool { if observed.len() != reference.len() { let detail = format!("{stage}: length {} != {}", observed.len(), reference.len()); first_divergence.get_or_insert(detail.clone()); eprintln!("[gemma3-vlm-boundary] stage={stage} status=FAIL {detail}"); - return; + return false; } let stats = comparison_stats(observed, reference, tolerance); - let status = if stats.failures == 0 { "PASS" } else { "FAIL" }; + let passed = stats.failures == 0; + let status = if passed { "PASS" } else { "FAIL" }; eprintln!( "[gemma3-vlm-boundary] stage={stage} status={status} elements={} \ atol={:.3e} rtol={:.3e} max_abs={:.6e} max_rel={:.6e} \ @@ -242,6 +277,86 @@ pub mod reference_boundary { ) }); } + passed + } + + fn first_failing_siglip_block_interval( + checkpoint_results: &[(SiglipHiddenCheckpoint, bool)], + ) -> Option<(usize, usize)> { + checkpoint_results + .iter() + .position(|(_, passed)| !passed) + .map(|failed| { + let first = failed + .checked_sub(1) + .and_then(|previous| checkpoint_results.get(previous)) + .map_or(1, |(checkpoint, _)| checkpoint.block_index + 1); + (first, checkpoint_results[failed].0.block_index) + }) + } + + fn compare_siglip_hidden_bisection( + observed_hidden_states: &[Vec], + reference_checkpoints: &[Vec], + first_divergence: &mut Option, + ) { + assert_eq!( + observed_hidden_states.len(), + PINNED_SIGLIP_BLOCK_COUNT + 1, + "IREE SigLIP diagnostics must contain embedding plus every pinned block output" + ); + assert_eq!( + reference_checkpoints.len(), + SIGLIP_HIDDEN_BISECTION_CHECKPOINTS.len(), + "eager SigLIP bisection checkpoint count drifted" + ); + let results = SIGLIP_HIDDEN_BISECTION_CHECKPOINTS + .iter() + .copied() + .zip(reference_checkpoints) + .map(|(checkpoint, reference)| { + let observed = observed_hidden_states + .get(checkpoint.hidden_state_index) + .unwrap_or_else(|| { + panic!( + "IREE SigLIP diagnostics are missing {} at hidden-state index {}", + checkpoint.label, checkpoint.hidden_state_index + ) + }); + let passed = compare_stage( + checkpoint.label, + observed, + reference, + VISION_TOLERANCE, + first_divergence, + ); + (checkpoint, passed) + }) + .collect::>(); + let checkpoint_statuses = results + .iter() + .map(|(checkpoint, passed)| { + format!( + "{}:{}", + checkpoint.label, + if *passed { "PASS" } else { "FAIL" } + ) + }) + .collect::>() + .join(","); + if let Some((first, last)) = first_failing_siglip_block_interval(&results) { + eprintln!( + "[gemma3-vlm-boundary] stage=siglip.hidden.bisection status=FAIL \ + first_failing_block_range={first}..={last} checkpoints=[{checkpoint_statuses}]" + ); + } else { + eprintln!( + "[gemma3-vlm-boundary] stage=siglip.hidden.bisection status=PASS \ + checked_block_range=1..={} checkpoints=[{checkpoint_statuses}]", + PINNED_SIGLIP_BLOCK_COUNT - 1 + ); + } + io::stderr().flush().expect("flush SigLIP bisection report"); } fn sha256(path: &Path) -> String { @@ -394,6 +509,64 @@ pub mod reference_boundary { assert_eq!(stats.first_failure, Some(1)); } + #[cfg(test)] + #[test] + fn siglip_hidden_bisection_descriptors_are_exact_and_balanced() { + assert_eq!( + SIGLIP_HIDDEN_BISECTION_CHECKPOINTS, + [ + SiglipHiddenCheckpoint { + label: "siglip.hidden.block6.output", + block_index: 6, + hidden_state_index: 7, + }, + SiglipHiddenCheckpoint { + label: "siglip.hidden.block13.output", + block_index: 13, + hidden_state_index: 14, + }, + SiglipHiddenCheckpoint { + label: "siglip.hidden.block20.output", + block_index: 20, + hidden_state_index: 21, + }, + SiglipHiddenCheckpoint { + label: "siglip.hidden.last_pre_layernorm", + block_index: 26, + hidden_state_index: 27, + }, + ] + ); + let mut first = 1; + let widths = SIGLIP_HIDDEN_BISECTION_CHECKPOINTS.map(|checkpoint| { + assert_eq!(checkpoint.hidden_state_index, checkpoint.block_index + 1); + let width = checkpoint.block_index - first + 1; + first = checkpoint.block_index + 1; + width + }); + assert_eq!(widths, [6, 7, 7, 6]); + assert_eq!(first, PINNED_SIGLIP_BLOCK_COUNT); + } + + #[cfg(test)] + #[test] + fn siglip_hidden_bisection_bounds_the_first_failed_checkpoint() { + let results = SIGLIP_HIDDEN_BISECTION_CHECKPOINTS.map(|checkpoint| (checkpoint, true)); + assert_eq!(first_failing_siglip_block_interval(&results), None); + + for (failed, expected) in [(0, (1, 6)), (1, (7, 13)), (2, (14, 20)), (3, (21, 26))] { + let mut results = + SIGLIP_HIDDEN_BISECTION_CHECKPOINTS.map(|checkpoint| (checkpoint, true)); + for (_, passed) in &mut results[failed..] { + *passed = false; + } + assert_eq!( + first_failing_siglip_block_interval(&results), + Some(expected) + ); + } + } + /// Run the pinned mixed-runtime boundary gate for #869. /// /// This entry point loads only the eager MLX SigLIP/projector/text embedding @@ -541,12 +714,33 @@ pub mod reference_boundary { let mlx_projected = connector.forward(&mlx_selected.hidden_states); let mlx_projected_values = mlx_f32(&mlx_projected, "MLX projected image features"); let mlx_hidden0 = mlx_f32(&mlx_hidden[0], "MLX SigLIP embedding output"); - let mlx_last_hidden = mlx_f32( - mlx_hidden - .last() - .expect("MLX captured a final hidden state"), - "MLX SigLIP last hidden", + assert_eq!( + mlx_hidden.len(), + PINNED_SIGLIP_BLOCK_COUNT + 1, + "eager SigLIP diagnostics must contain embedding plus every pinned block output" ); + let mlx_hidden_bisection = SIGLIP_HIDDEN_BISECTION_CHECKPOINTS + .iter() + .map(|checkpoint| { + assert_eq!( + checkpoint.hidden_state_index, + checkpoint.block_index + 1, + "{} descriptor must map block output to embedding-prefixed hidden-state index", + checkpoint.label + ); + mlx_f32( + mlx_hidden + .get(checkpoint.hidden_state_index) + .unwrap_or_else(|| { + panic!( + "eager SigLIP diagnostics are missing {} at hidden-state index {}", + checkpoint.label, checkpoint.hidden_state_index + ) + }), + checkpoint.label, + ) + }) + .collect::>(); let mlx_block0_values = mlx_block0 .iter() .map(|stage| mlx_f32(stage, "MLX SigLIP block 0 stage")) @@ -585,13 +779,9 @@ pub mod reference_boundary { &mut first_divergence, ); } - compare_stage( - "siglip.hidden.last_pre_layernorm", - iree.hidden_states - .last() - .expect("IREE captured a final hidden state"), - &mlx_last_hidden, - VISION_TOLERANCE, + compare_siglip_hidden_bisection( + &iree.hidden_states, + &mlx_hidden_bisection, &mut first_divergence, ); compare_stage( From 9ab139925090a7686fd32df66a45056182a60e7d Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 03:38:48 +0900 Subject: [PATCH 11/11] test(xla): refine first failing SigLIP block The balanced checkpoint run localized the first post-block-0 divergence to blocks 1 through 6 but could not identify the exact block. Compare block outputs 1 through 5 before the existing block-6 bisection boundary, preserve the later block 13, 20, and final checkpoints, and emit an ordered block-1-through-6 status summary with the exact first failing block. The diagnostic reuses existing eager and IREE captures without changing graph outputs, production behavior, tolerances, worker flags, or artifacts. Validated the complete reference-boundary test group, rustfmt, and Clippy with the xla-reference-diagnostics feature. Refs #869 --- src/loading/vlm_gemma_xla_tests.rs | 177 +++++++++++++++++++++++++---- 1 file changed, 152 insertions(+), 25 deletions(-) diff --git a/src/loading/vlm_gemma_xla_tests.rs b/src/loading/vlm_gemma_xla_tests.rs index ab39ae842..e0c4b4d9e 100644 --- a/src/loading/vlm_gemma_xla_tests.rs +++ b/src/loading/vlm_gemma_xla_tests.rs @@ -92,6 +92,37 @@ pub mod reference_boundary { hidden_state_index: usize, } + // Block 6 already closes the first balanced bisection interval. Capture + // blocks 1..=5 immediately before it so the next run identifies the exact + // first failing block inside the known 1..=6 range. + const SIGLIP_HIDDEN_REFINEMENT_CHECKPOINTS: [SiglipHiddenCheckpoint; 5] = [ + SiglipHiddenCheckpoint { + label: "siglip.hidden.block1.output", + block_index: 1, + hidden_state_index: 2, + }, + SiglipHiddenCheckpoint { + label: "siglip.hidden.block2.output", + block_index: 2, + hidden_state_index: 3, + }, + SiglipHiddenCheckpoint { + label: "siglip.hidden.block3.output", + block_index: 3, + hidden_state_index: 4, + }, + SiglipHiddenCheckpoint { + label: "siglip.hidden.block4.output", + block_index: 4, + hidden_state_index: 5, + }, + SiglipHiddenCheckpoint { + label: "siglip.hidden.block5.output", + block_index: 5, + hidden_state_index: 6, + }, + ]; + // Block 0 already has full sub-stage coverage. These ordered checkpoints // split blocks 1..=26 into balanced inclusive ranges 1..=6, 7..=13, // 14..=20, and 21..=26 while reusing the existing final-state comparison. @@ -118,6 +149,13 @@ pub mod reference_boundary { }, ]; + fn siglip_hidden_checkpoints() -> impl Iterator { + SIGLIP_HIDDEN_REFINEMENT_CHECKPOINTS + .iter() + .chain(SIGLIP_HIDDEN_BISECTION_CHECKPOINTS.iter()) + .copied() + } + #[derive(Debug, Clone, Copy)] struct Tolerance { atol: f64, @@ -295,7 +333,29 @@ pub mod reference_boundary { }) } - fn compare_siglip_hidden_bisection( + fn first_failing_siglip_block( + checkpoint_results: &[(SiglipHiddenCheckpoint, bool)], + ) -> Option { + checkpoint_results + .iter() + .find_map(|(checkpoint, passed)| (!passed).then_some(checkpoint.block_index)) + } + + fn checkpoint_statuses(checkpoint_results: &[(SiglipHiddenCheckpoint, bool)]) -> String { + checkpoint_results + .iter() + .map(|(checkpoint, passed)| { + format!( + "{}:{}", + checkpoint.label, + if *passed { "PASS" } else { "FAIL" } + ) + }) + .collect::>() + .join(",") + } + + fn compare_siglip_hidden_checkpoints( observed_hidden_states: &[Vec], reference_checkpoints: &[Vec], first_divergence: &mut Option, @@ -307,12 +367,10 @@ pub mod reference_boundary { ); assert_eq!( reference_checkpoints.len(), - SIGLIP_HIDDEN_BISECTION_CHECKPOINTS.len(), - "eager SigLIP bisection checkpoint count drifted" + SIGLIP_HIDDEN_REFINEMENT_CHECKPOINTS.len() + SIGLIP_HIDDEN_BISECTION_CHECKPOINTS.len(), + "eager SigLIP ordered checkpoint count drifted" ); - let results = SIGLIP_HIDDEN_BISECTION_CHECKPOINTS - .iter() - .copied() + let results = siglip_hidden_checkpoints() .zip(reference_checkpoints) .map(|(checkpoint, reference)| { let observed = observed_hidden_states @@ -333,30 +391,37 @@ pub mod reference_boundary { (checkpoint, passed) }) .collect::>(); - let checkpoint_statuses = results - .iter() - .map(|(checkpoint, passed)| { - format!( - "{}:{}", - checkpoint.label, - if *passed { "PASS" } else { "FAIL" } - ) - }) - .collect::>() - .join(","); - if let Some((first, last)) = first_failing_siglip_block_interval(&results) { + let refinement_end = SIGLIP_HIDDEN_REFINEMENT_CHECKPOINTS.len() + 1; + let refinement_results = &results[..refinement_end]; + let refinement_statuses = checkpoint_statuses(refinement_results); + if let Some(block) = first_failing_siglip_block(refinement_results) { + eprintln!( + "[gemma3-vlm-boundary] stage=siglip.hidden.block1_to_6 status=FAIL \ + exact_first_failing_block={block} checkpoints=[{refinement_statuses}]" + ); + } else { + eprintln!( + "[gemma3-vlm-boundary] stage=siglip.hidden.block1_to_6 status=PASS \ + checked_block_range=1..=6 checkpoints=[{refinement_statuses}]" + ); + } + let bisection_results = &results[SIGLIP_HIDDEN_REFINEMENT_CHECKPOINTS.len()..]; + let bisection_statuses = checkpoint_statuses(bisection_results); + if let Some((first, last)) = first_failing_siglip_block_interval(bisection_results) { eprintln!( "[gemma3-vlm-boundary] stage=siglip.hidden.bisection status=FAIL \ - first_failing_block_range={first}..={last} checkpoints=[{checkpoint_statuses}]" + first_failing_block_range={first}..={last} checkpoints=[{bisection_statuses}]" ); } else { eprintln!( "[gemma3-vlm-boundary] stage=siglip.hidden.bisection status=PASS \ - checked_block_range=1..={} checkpoints=[{checkpoint_statuses}]", + checked_block_range=1..={} checkpoints=[{bisection_statuses}]", PINNED_SIGLIP_BLOCK_COUNT - 1 ); } - io::stderr().flush().expect("flush SigLIP bisection report"); + io::stderr() + .flush() + .expect("flush SigLIP hidden checkpoint report"); } fn sha256(path: &Path) -> String { @@ -509,6 +574,69 @@ pub mod reference_boundary { assert_eq!(stats.first_failure, Some(1)); } + #[cfg(test)] + #[test] + fn siglip_hidden_refinement_descriptors_are_exact_and_ordered() { + assert_eq!( + SIGLIP_HIDDEN_REFINEMENT_CHECKPOINTS, + [ + SiglipHiddenCheckpoint { + label: "siglip.hidden.block1.output", + block_index: 1, + hidden_state_index: 2, + }, + SiglipHiddenCheckpoint { + label: "siglip.hidden.block2.output", + block_index: 2, + hidden_state_index: 3, + }, + SiglipHiddenCheckpoint { + label: "siglip.hidden.block3.output", + block_index: 3, + hidden_state_index: 4, + }, + SiglipHiddenCheckpoint { + label: "siglip.hidden.block4.output", + block_index: 4, + hidden_state_index: 5, + }, + SiglipHiddenCheckpoint { + label: "siglip.hidden.block5.output", + block_index: 5, + hidden_state_index: 6, + }, + ] + ); + assert_eq!( + siglip_hidden_checkpoints() + .map(|checkpoint| checkpoint.block_index) + .collect::>(), + [1, 2, 3, 4, 5, 6, 13, 20, 26] + ); + } + + #[cfg(test)] + #[test] + fn siglip_hidden_refinement_identifies_each_exact_failed_block() { + let refinement = siglip_hidden_checkpoints() + .take(SIGLIP_HIDDEN_REFINEMENT_CHECKPOINTS.len() + 1) + .collect::>(); + assert_eq!(refinement.len(), 6); + + for failed in 0..refinement.len() { + let mut results = refinement + .iter() + .copied() + .map(|checkpoint| (checkpoint, true)) + .collect::>(); + results[failed].1 = false; + assert_eq!( + first_failing_siglip_block(&results), + Some(refinement[failed].block_index) + ); + } + } + #[cfg(test)] #[test] fn siglip_hidden_bisection_descriptors_are_exact_and_balanced() { @@ -719,8 +847,7 @@ pub mod reference_boundary { PINNED_SIGLIP_BLOCK_COUNT + 1, "eager SigLIP diagnostics must contain embedding plus every pinned block output" ); - let mlx_hidden_bisection = SIGLIP_HIDDEN_BISECTION_CHECKPOINTS - .iter() + let mlx_hidden_checkpoints = siglip_hidden_checkpoints() .map(|checkpoint| { assert_eq!( checkpoint.hidden_state_index, @@ -779,9 +906,9 @@ pub mod reference_boundary { &mut first_divergence, ); } - compare_siglip_hidden_bisection( + compare_siglip_hidden_checkpoints( &iree.hidden_states, - &mlx_hidden_bisection, + &mlx_hidden_checkpoints, &mut first_divergence, ); compare_stage(