From 5378d9eba55c79155c2267e308ac7c46a3e04f07 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 24 Jul 2026 18:41:21 +0900 Subject: [PATCH 01/20] feat(xla): add Molmo2 vision runtime Emit and load the pinned Molmo2 flattened-patch ViT, indexed attention pooler, and additive projector through IREE. Extend the shared XLA text loader for Molmo2 fused weights and wire the owned prepared-prefill producer into the existing CLI and server lifecycle. --- src/lib/mlxcel-xla/src/emitter/config.rs | 34 +- src/lib/mlxcel-xla/src/emitter/mod.rs | 6 + .../mlxcel-xla/src/emitter/molmo2_config.rs | 462 +++++++++++++++++ .../src/emitter/molmo2_config_tests.rs | 81 +++ .../mlxcel-xla/src/emitter/molmo2_vision.rs | 464 +++++++++++++++++ src/lib/mlxcel-xla/src/iree.rs | 114 +++- src/lib/mlxcel-xla/src/lib.rs | 12 + src/lib/mlxcel-xla/src/molmo2.rs | 402 +++++++++++++++ .../mlxcel-xla/src/molmo2_vision_runtime.rs | 488 ++++++++++++++++++ .../src/molmo2_vision_runtime_tests.rs | 33 ++ src/lib/mlxcel-xla/src/weight_names.rs | 22 + src/lib/mlxcel-xla/src/weights.rs | 174 ++++++- src/loading/mod.rs | 3 +- src/loading/vlm.rs | 3 +- src/loading/vlm_special.rs | 41 ++ src/multimodal/host_preprocessor.rs | 110 ++-- src/multimodal/molmo2_xla_preprocessor.rs | 347 +++++++++++++ 17 files changed, 2723 insertions(+), 73 deletions(-) create mode 100644 src/lib/mlxcel-xla/src/emitter/molmo2_config.rs create mode 100644 src/lib/mlxcel-xla/src/emitter/molmo2_config_tests.rs create mode 100644 src/lib/mlxcel-xla/src/emitter/molmo2_vision.rs create mode 100644 src/lib/mlxcel-xla/src/molmo2.rs create mode 100644 src/lib/mlxcel-xla/src/molmo2_vision_runtime.rs create mode 100644 src/lib/mlxcel-xla/src/molmo2_vision_runtime_tests.rs create mode 100644 src/multimodal/molmo2_xla_preprocessor.rs diff --git a/src/lib/mlxcel-xla/src/emitter/config.rs b/src/lib/mlxcel-xla/src/emitter/config.rs index e5c450eaa..7b7aaa68b 100644 --- a/src/lib/mlxcel-xla/src/emitter/config.rs +++ b/src/lib/mlxcel-xla/src/emitter/config.rs @@ -132,6 +132,8 @@ pub enum WeightScheme { /// ExaOne 3.x GPT-2-style names (`transformer.h.{i}...`, gated MLP `c_fc_0` / /// `c_fc_1` / `c_proj`, `out_proj` attention output). Exaone, + /// Molmo2's OLMo-style names under a VLM `language_model.*` namespace. + Molmo2, } /// MLX affine weight quantization (`config.json` `quantization`). The linear / @@ -482,14 +484,17 @@ 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 is_molmo2 = wrapper_model_type == Some("molmo2"); + let v = if matches!(wrapper_model_type, Some("llava" | "llava_next" | "molmo2")) { let mut text = root .get("text_config") .and_then(serde_json::Value::as_object) .cloned() .ok_or_else(|| { - "LLaVA config.json missing object `text_config` for the XLA text graph" - .to_string() + format!( + "{} config.json missing object `text_config` for the XLA text graph", + wrapper_model_type.unwrap_or("VLM") + ) })?; // mlx-community quantized VLMs commonly keep the affine scheme at // the wrapper level even though the tensors belong to the nested @@ -509,10 +514,10 @@ impl Config { // ExaOne 3.x keeps GPT-2-style tensor names; every other supported family // uses the standard HF Llama layout. Loader-only (see [`WeightScheme`]), so // it never changes the emitted graph. - let weight_scheme = if model_type == Some("exaone") { - WeightScheme::Exaone - } else { - WeightScheme::Llama + let weight_scheme = match model_type { + Some("exaone") => WeightScheme::Exaone, + Some("molmo2_text") if is_molmo2 => WeightScheme::Molmo2, + _ => WeightScheme::Llama, }; // Interleaved (GPT-J-style) RoPE reaches the supported families only through @@ -870,6 +875,19 @@ impl Config { tie_default = false; rotary_dim = partial_rotary(1.0); } + Some("molmo2_text") if is_molmo2 => { + // Molmo2 reuses the shared dense text graph: plain pre-norm, + // fused QKV, per-head raw q/k RMSNorm, and fused SwiGLU. Its + // checkpoint names and gate/up half ordering are loader-only + // differences captured by WeightScheme::Molmo2. + fused_qkv = true; + fused_gate_up = true; + tie_default = false; + qk_norm = Some(QkNorm { + per_head: true, + one_plus: false, + }); + } Some("stablelm") => { // LayerNorm with bias, partial RoPE, optional q/k/v bias, untied. layernorm = true; @@ -1003,7 +1021,7 @@ impl Config { return Err(format!( "the OpenXLA emitter supports the dense architectures Llama, Qwen2, \ Qwen3, Gemma1/2/3, SmolLM3, OLMo2/3, Seed-OSS, MiMo, InternLM3, ExaOne, \ - Cohere, Cohere2, Phi3, Phi4MM, StableLM, StarCoder2, Granite, and MiniCPM, plus \ + Cohere, Cohere2, Phi3, Phi4MM, Molmo2, StableLM, StarCoder2, Granite, and MiniCPM, plus \ the Mixtral, Qwen2-MoE, Qwen3-MoE, and OLMoE mixture-of-experts \ architectures; config.json model_type = {other:?} (other MoE / MLA / \ novel-activation variants are follow-ups)" diff --git a/src/lib/mlxcel-xla/src/emitter/mod.rs b/src/lib/mlxcel-xla/src/emitter/mod.rs index 0592c4f76..b725295de 100644 --- a/src/lib/mlxcel-xla/src/emitter/mod.rs +++ b/src/lib/mlxcel-xla/src/emitter/mod.rs @@ -58,6 +58,8 @@ mod gemma3n_weights; mod model; mod moe; pub(crate) mod numeric_ops; +mod molmo2_config; +mod molmo2_vision; mod phi4_audio; mod qwen2_vl; mod rope; @@ -152,6 +154,10 @@ pub(crate) use model::{ validate_prefill_embeddings_metadata, }; #[allow(unused_imports)] +pub(crate) use molmo2_config::{Molmo2VisionConfig, Molmo2VisionWeightSpec}; +#[allow(unused_imports)] +pub(crate) use molmo2_vision::emit_molmo2_vision; +#[allow(unused_imports)] pub(crate) use vision::emit_vision; // The only consumer is `IreeVisionDiagnosticProjector` in `vision_runtime`, which // is gated on `diagnostics` alone. The `test` arm is kept deliberately: it is the diff --git a/src/lib/mlxcel-xla/src/emitter/molmo2_config.rs b/src/lib/mlxcel-xla/src/emitter/molmo2_config.rs new file mode 100644 index 000000000..521f324e2 --- /dev/null +++ b/src/lib/mlxcel-xla/src/emitter/molmo2_config.rs @@ -0,0 +1,462 @@ +// 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. + +//! Strict Molmo2 vision/processor contract for the pinned XLA graph. + +use std::path::Path; + +use serde_json::Value; + +const MAX_SUPPORTED_CROPS: usize = 8; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct Molmo2VisionWeightSpec { + pub(crate) name: String, + pub(crate) shape: Vec, +} + +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct Molmo2VisionConfig { + pub(crate) crop_size: usize, + pub(crate) patch_size: usize, + pub(crate) patches_per_crop: usize, + pub(crate) patch_dim: usize, + pub(crate) max_crops: usize, + pub(crate) static_crops: usize, + pub(crate) overlap: [usize; 2], + pub(crate) pool_h: usize, + pub(crate) pool_w: usize, + pub(crate) pool_size: usize, + pub(crate) static_pool_groups: usize, + pub(crate) hidden: usize, + pub(crate) intermediate: usize, + pub(crate) heads: usize, + pub(crate) head_dim: usize, + pub(crate) layers: usize, + pub(crate) emitted_layers: usize, + pub(crate) selected_layers: Vec, + pub(crate) position_count: usize, + pub(crate) layer_norm_eps: f32, + pub(crate) pool_hidden: usize, + pub(crate) pool_heads: usize, + pub(crate) pool_head_dim: usize, + pub(crate) projector_intermediate: usize, + pub(crate) text_hidden: usize, + pub(crate) pooling_attention_mask: bool, + pub(crate) image_patch_id: i32, +} + +fn object<'a>(value: &'a Value, name: &str) -> Result<&'a serde_json::Map, String> { + value + .get(name) + .and_then(Value::as_object) + .ok_or_else(|| format!("Molmo2 config missing object `{name}`")) +} + +fn usize_field(object: &serde_json::Map, name: &str) -> Result { + object + .get(name) + .and_then(Value::as_u64) + .and_then(|value| usize::try_from(value).ok()) + .filter(|value| *value > 0) + .ok_or_else(|| format!("Molmo2 `{name}` must be a positive integer")) +} + +fn usize_pair(object: &serde_json::Map, name: &str) -> Result<[usize; 2], String> { + let values = object + .get(name) + .and_then(Value::as_array) + .ok_or_else(|| format!("Molmo2 `{name}` must contain two integers"))?; + if values.len() != 2 { + return Err(format!("Molmo2 `{name}` must contain two integers")); + } + let mut output = [0usize; 2]; + for (index, value) in values.iter().enumerate() { + output[index] = value + .as_u64() + .and_then(|value| usize::try_from(value).ok()) + .filter(|value| *value > 0) + .ok_or_else(|| format!("Molmo2 `{name}[{index}]` must be positive"))?; + } + Ok(output) +} + +fn usize_pair_nonnegative( + object: &serde_json::Map, + name: &str, +) -> Result<[usize; 2], String> { + let values = object + .get(name) + .and_then(Value::as_array) + .ok_or_else(|| format!("Molmo2 `{name}` must contain two nonnegative integers"))?; + if values.len() != 2 { + return Err(format!( + "Molmo2 `{name}` must contain two nonnegative integers" + )); + } + let mut output = [0usize; 2]; + for (index, value) in values.iter().enumerate() { + output[index] = value + .as_u64() + .and_then(|value| usize::try_from(value).ok()) + .ok_or_else(|| format!("Molmo2 `{name}[{index}]` must be nonnegative"))?; + } + Ok(output) +} + +fn maximum_pool_groups(max_crops: usize, crop_patches: usize, overlap: [usize; 2]) -> usize { + let window = crop_patches - overlap[0] - overlap[1]; + let low = crop_patches.div_ceil(2).pow(2); + let high = (1..=max_crops) + .flat_map(|rows| (1..=max_crops).map(move |columns| (rows, columns))) + .filter(|(rows, columns)| rows * columns <= max_crops) + .map(|(rows, columns)| { + let height = rows * window + overlap[0] + overlap[1]; + let width = columns * window + overlap[0] + overlap[1]; + height.div_ceil(2) * width.div_ceil(2) + }) + .max() + .unwrap_or(0); + low + high +} + +impl Molmo2VisionConfig { + pub(crate) fn from_model_dir(model_dir: &Path) -> Result { + let config_path = model_dir.join("config.json"); + let processor_path = model_dir.join("preprocessor_config.json"); + let config_text = std::fs::read_to_string(&config_path) + .map_err(|error| format!("read {}: {error}", config_path.display()))?; + let processor_text = std::fs::read_to_string(&processor_path) + .map_err(|error| format!("read {}: {error}", processor_path.display()))?; + Self::from_json_strs(&config_text, &processor_text) + } + + pub(crate) fn from_json_strs(config: &str, processor: &str) -> Result { + let root: Value = + serde_json::from_str(config).map_err(|error| format!("parse config.json: {error}"))?; + if root.get("model_type").and_then(Value::as_str) != Some("molmo2") { + return Err("Molmo2 XLA vision requires config.json model_type `molmo2`".to_string()); + } + let vit = object(&root, "vit_config")?; + let adapter = object(&root, "adapter_config")?; + let processor: Value = serde_json::from_str(processor) + .map_err(|error| format!("parse preprocessor_config.json: {error}"))?; + let processor = processor + .as_object() + .ok_or_else(|| "Molmo2 preprocessor config must be an object".to_string())?; + + let input_size = usize_pair(vit, "image_default_input_size")?; + if input_size[0] != input_size[1] { + return Err("Molmo2 XLA requires square default image crops".to_string()); + } + let crop_size = input_size[0]; + let patch_size = usize_field(vit, "image_patch_size")?; + if !crop_size.is_multiple_of(patch_size) { + return Err(format!( + "Molmo2 crop size {crop_size} is not divisible by patch size {patch_size}" + )); + } + let crop_patches = crop_size / patch_size; + let patches_per_crop = crop_patches * crop_patches; + let position_count = usize_field(vit, "image_num_pos")?; + if position_count != patches_per_crop { + return Err(format!( + "Molmo2 XLA supports the pinned exact position table only: image_num_pos={position_count}, default grid has {patches_per_crop} patches" + )); + } + let preprocessor_patch = usize_field(processor, "patch_size")?; + let size = processor + .get("size") + .and_then(Value::as_object) + .ok_or_else(|| "Molmo2 preprocessor missing object `size`".to_string())?; + if preprocessor_patch != patch_size + || usize_field(size, "height")? != crop_size + || usize_field(size, "width")? != crop_size + { + return Err("Molmo2 processor and ViT crop/patch geometry disagree".to_string()); + } + let max_crops = usize_field(processor, "max_crops")?; + if max_crops > MAX_SUPPORTED_CROPS { + return Err(format!( + "Molmo2 XLA supports at most {MAX_SUPPORTED_CROPS} high-resolution crops, got {max_crops}" + )); + } + let overlap = usize_pair_nonnegative(processor, "overlap_margins")?; + if overlap[0] + overlap[1] >= crop_patches { + return Err("Molmo2 overlap margins consume the complete crop".to_string()); + } + let pooling = usize_pair(processor, "pooling_size")?; + if pooling != [2, 2] { + return Err(format!( + "Molmo2 pinned XLA graph requires 2x2 pooling, got {pooling:?}" + )); + } + + let layers = usize_field(vit, "num_hidden_layers")?.min(25); + let selected_raw = adapter + .get("vit_layers") + .and_then(Value::as_array) + .ok_or_else(|| "Molmo2 adapter `vit_layers` must be an array".to_string())?; + if selected_raw.is_empty() { + return Err("Molmo2 adapter must select at least one ViT layer".to_string()); + } + let selected_layers = selected_raw + .iter() + .enumerate() + .map(|(index, value)| { + let raw = value + .as_i64() + .ok_or_else(|| format!("Molmo2 vit_layers[{index}] must be an integer"))?; + let resolved = if raw < 0 { layers as i64 + raw } else { raw }; + usize::try_from(resolved) + .ok() + .filter(|layer| *layer < layers) + .ok_or_else(|| { + format!("Molmo2 vit_layers[{index}]={raw} resolves outside [0,{layers})") + }) + }) + .collect::, String>>()?; + let emitted_layers = selected_layers + .iter() + .copied() + .max() + .ok_or_else(|| "Molmo2 adapter vit_layers must not be empty".to_string())? + + 1; + let hidden = usize_field(vit, "hidden_size")?; + let heads = usize_field(vit, "num_attention_heads")?; + let kv_heads = vit + .get("num_key_value_heads") + .and_then(Value::as_u64) + .map(|value| value as usize) + .unwrap_or(heads); + let head_dim = usize_field(vit, "head_dim")?; + if heads * head_dim != hidden || kv_heads != heads { + return Err( + "Molmo2 XLA requires ViT MHA with heads * head_dim = hidden_size".to_string(), + ); + } + let pool_hidden = usize_field(adapter, "hidden_size")?; + let pool_heads = usize_field(adapter, "num_attention_heads")?; + let pool_kv_heads = adapter + .get("num_key_value_heads") + .and_then(Value::as_u64) + .map(|value| value as usize) + .unwrap_or(pool_heads); + let pool_head_dim = usize_field(adapter, "head_dim")?; + if pool_heads * pool_head_dim != pool_hidden || pool_kv_heads != pool_heads { + return Err( + "Molmo2 XLA requires pooling MHA with heads * head_dim = hidden_size".to_string(), + ); + } + let selected_width = hidden + .checked_mul(selected_layers.len()) + .ok_or_else(|| "Molmo2 selected feature width overflowed".to_string())?; + let pool_q_width = adapter + .get("image_feature_dim") + .and_then(Value::as_u64) + .map(|value| value as usize) + .unwrap_or(selected_width); + if pool_q_width != selected_width { + return Err(format!( + "Molmo2 pooling input width {pool_q_width} disagrees with selected layer width {selected_width}" + )); + } + + let layer_norm_eps = vit + .get("layer_norm_eps") + .and_then(Value::as_f64) + .filter(|value| value.is_finite() && *value > 0.0) + .ok_or_else(|| "Molmo2 layer_norm_eps must be positive and finite".to_string())? + as f32; + let image_patch_id = root + .get("image_patch_id") + .and_then(Value::as_i64) + .and_then(|value| i32::try_from(value).ok()) + .ok_or_else(|| "Molmo2 image_patch_id must fit i32".to_string())?; + + Ok(Self { + crop_size, + patch_size, + patches_per_crop, + patch_dim: patch_size * patch_size * 3, + max_crops, + static_crops: max_crops + 1, + overlap, + pool_h: pooling[0], + pool_w: pooling[1], + pool_size: pooling[0] * pooling[1], + static_pool_groups: maximum_pool_groups(max_crops, crop_patches, overlap), + hidden, + intermediate: usize_field(vit, "intermediate_size")?, + heads, + head_dim, + layers, + emitted_layers, + selected_layers, + position_count, + layer_norm_eps, + pool_hidden, + pool_heads, + pool_head_dim, + projector_intermediate: usize_field(adapter, "intermediate_size")?, + text_hidden: usize_field(adapter, "text_hidden_size")?, + pooling_attention_mask: adapter + .get("pooling_attention_mask") + .and_then(Value::as_bool) + .ok_or_else(|| { + "Molmo2 pooling_attention_mask must be an explicit boolean".to_string() + })?, + image_patch_id, + }) + } + + #[must_use] + pub(crate) fn selected_width(&self) -> usize { + self.hidden * self.selected_layers.len() + } + + pub(crate) fn valid_runtime_geometry(&self, crops: usize, grid: [usize; 4]) -> bool { + if !(2..=self.static_crops).contains(&crops) { + return false; + } + let crop_patches = self.crop_size / self.patch_size; + let low = crop_patches.div_ceil(self.pool_h); + if grid[0] != low || grid[1] != low { + return false; + } + let high_crops = crops - 1; + let window = crop_patches - self.overlap[0] - self.overlap[1]; + (1..=high_crops).any(|rows| { + high_crops.is_multiple_of(rows) && { + let columns = high_crops / rows; + let height = rows * window + self.overlap[0] + self.overlap[1]; + let width = columns * window + self.overlap[0] + self.overlap[1]; + grid[2] == height.div_ceil(self.pool_h) && grid[3] == width.div_ceil(self.pool_w) + } + }) + } + + pub(crate) fn fingerprint(&self) -> String { + format!( + "molmo2-vision-v1;position=exact-default;selected={:?};pool-mask={};patch-id={};crops={};overlap={:?};patches={};pool-groups={};pool={}x{};hidden={};inter={};layers={};emitted={};heads={};head-dim={};pool-hidden={};pool-heads={};pool-head-dim={};projector-inter={};text-hidden={}", + self.selected_layers, + self.pooling_attention_mask, + self.image_patch_id, + self.static_crops, + self.overlap, + self.patches_per_crop, + self.static_pool_groups, + self.pool_h, + self.pool_w, + self.hidden, + self.intermediate, + self.layers, + self.emitted_layers, + self.heads, + self.head_dim, + self.pool_hidden, + self.pool_heads, + self.pool_head_dim, + self.projector_intermediate, + self.text_hidden + ) + } + + fn spec( + &self, + name: impl Into, + shape: impl Into>, + ) -> Molmo2VisionWeightSpec { + Molmo2VisionWeightSpec { + name: name.into(), + shape: shape.into(), + } + } + + pub(crate) fn weight_specs(&self) -> Vec { + let mut specs = vec![ + self.spec( + "vision_tower.image_vit.patch_embedding.weight", + [self.hidden, self.patch_dim], + ), + self.spec("vision_tower.image_vit.patch_embedding.bias", [self.hidden]), + self.spec( + "vision_tower.image_vit.positional_embedding", + [self.position_count, self.hidden], + ), + ]; + for layer in 0..self.emitted_layers { + let prefix = format!("vision_tower.image_vit.transformer.{layer}"); + for projection in ["wq", "wk", "wv", "wo"] { + specs.push(self.spec( + format!("{prefix}.attention.{projection}.weight"), + [self.hidden, self.hidden], + )); + specs.push(self.spec( + format!("{prefix}.attention.{projection}.bias"), + [self.hidden], + )); + } + for norm in ["attention_norm", "ffn_norm"] { + specs.push(self.spec(format!("{prefix}.{norm}.weight"), [self.hidden])); + specs.push(self.spec(format!("{prefix}.{norm}.bias"), [self.hidden])); + } + specs.push(self.spec( + format!("{prefix}.feed_forward.w1.weight"), + [self.intermediate, self.hidden], + )); + specs.push(self.spec( + format!("{prefix}.feed_forward.w1.bias"), + [self.intermediate], + )); + specs.push(self.spec( + format!("{prefix}.feed_forward.w2.weight"), + [self.hidden, self.intermediate], + )); + specs.push(self.spec(format!("{prefix}.feed_forward.w2.bias"), [self.hidden])); + } + for projection in ["wq", "wk", "wv"] { + specs.push(self.spec( + format!("vision_tower.image_pooling_2d.{projection}.weight"), + [self.pool_hidden, self.selected_width()], + )); + specs.push(self.spec( + format!("vision_tower.image_pooling_2d.{projection}.bias"), + [self.pool_hidden], + )); + } + specs.push(self.spec( + "vision_tower.image_pooling_2d.wo.weight", + [self.pool_hidden, self.pool_hidden], + )); + specs.push(self.spec("vision_tower.image_pooling_2d.wo.bias", [self.pool_hidden])); + specs.push(self.spec( + "vision_tower.image_projector.w1.weight", + [self.projector_intermediate, self.pool_hidden], + )); + specs.push(self.spec( + "vision_tower.image_projector.w2.weight", + [self.text_hidden, self.projector_intermediate], + )); + specs.push(self.spec( + "vision_tower.image_projector.w3.weight", + [self.projector_intermediate, self.pool_hidden], + )); + specs + } +} + +#[cfg(test)] +#[path = "molmo2_config_tests.rs"] +mod tests; diff --git a/src/lib/mlxcel-xla/src/emitter/molmo2_config_tests.rs b/src/lib/mlxcel-xla/src/emitter/molmo2_config_tests.rs new file mode 100644 index 000000000..263387781 --- /dev/null +++ b/src/lib/mlxcel-xla/src/emitter/molmo2_config_tests.rs @@ -0,0 +1,81 @@ +// 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::*; + +fn pinned_config() -> String { + serde_json::json!({ + "model_type": "molmo2", + "image_patch_id": 151938, + "vit_config": { + "hidden_size": 8, "intermediate_size": 16, + "num_attention_heads": 2, "head_dim": 4, + "num_hidden_layers": 27, "image_default_input_size": [28, 28], + "image_patch_size": 14, "image_num_pos": 4, "layer_norm_eps": 1e-6 + }, + "adapter_config": { + "hidden_size": 8, "intermediate_size": 12, "text_hidden_size": 10, + "num_attention_heads": 2, "head_dim": 4, "vit_layers": [-3, -9], + "pooling_attention_mask": true + } + }) + .to_string() +} + +fn pinned_processor() -> String { + serde_json::json!({ + "patch_size": 14, "max_crops": 8, "overlap_margins": [0, 0], + "pooling_size": [2, 2], "size": {"height": 28, "width": 28} + }) + .to_string() +} + +#[test] +fn resolves_pinned_layers_and_static_bucket_identity() { + let config = Molmo2VisionConfig::from_json_strs(&pinned_config(), &pinned_processor()).unwrap(); + assert_eq!(config.layers, 25); + assert_eq!(config.selected_layers, vec![22, 16]); + assert_eq!(config.emitted_layers, 23); + assert_eq!(config.static_crops, 9); + assert_eq!(config.static_pool_groups, 9); + assert!(config.fingerprint().contains("position=exact-default")); + assert!(config.fingerprint().contains("selected=[22, 16]")); + assert!(config.fingerprint().contains("pool-mask=true")); +} + +#[test] +fn validates_runtime_crop_grid_relationship() { + let config = Molmo2VisionConfig::from_json_strs(&pinned_config(), &pinned_processor()).unwrap(); + assert!(config.valid_runtime_geometry(3, [1, 1, 1, 2])); + assert!(config.valid_runtime_geometry(3, [1, 1, 2, 1])); + assert!(!config.valid_runtime_geometry(4, [1, 1, 1, 2])); +} + +#[test] +fn rejects_position_grid_and_selected_layer_drift() { + let mut config: Value = serde_json::from_str(&pinned_config()).unwrap(); + config["vit_config"]["image_num_pos"] = Value::from(5); + assert!( + Molmo2VisionConfig::from_json_strs(&config.to_string(), &pinned_processor()) + .unwrap_err() + .contains("exact position") + ); + config["vit_config"]["image_num_pos"] = Value::from(4); + config["adapter_config"]["vit_layers"] = serde_json::json!([-26]); + assert!( + Molmo2VisionConfig::from_json_strs(&config.to_string(), &pinned_processor()) + .unwrap_err() + .contains("outside") + ); +} diff --git a/src/lib/mlxcel-xla/src/emitter/molmo2_vision.rs b/src/lib/mlxcel-xla/src/emitter/molmo2_vision.rs new file mode 100644 index 000000000..587f53cb6 --- /dev/null +++ b/src/lib/mlxcel-xla/src/emitter/molmo2_vision.rs @@ -0,0 +1,464 @@ +// 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. + +//! StableHLO emitter for Molmo2's flattened-patch ViT and indexed 2D pooler. + +use super::builder::{Builder, Ty, Val}; +use super::molmo2_config::{Molmo2VisionConfig, Molmo2VisionWeightSpec}; + +struct Args { + values: Vec, + declarations: Vec, + cursor: usize, +} + +impl Args { + fn new(specs: &[Molmo2VisionWeightSpec]) -> Self { + let mut values = Vec::with_capacity(specs.len()); + let mut declarations = Vec::with_capacity(specs.len() + 2); + for (index, spec) in specs.iter().enumerate() { + let ty = Ty::f32(spec.shape.clone()); + declarations.push(format!( + "%arg{index}: {} loc(\"{}\")", + ty.render(), + spec.name + )); + values.push(Builder::arg(index, ty)); + } + Self { + values, + declarations, + cursor: 0, + } + } + + fn take(&mut self) -> Val { + let value = self.values[self.cursor].clone(); + self.cursor += 1; + value + } + + fn input(&mut self, ty: Ty, name: &str) -> Val { + let index = self.values.len(); + self.declarations + .push(format!("%arg{index}: {} loc(\"{name}\")", ty.render())); + let value = Builder::arg(index, ty); + self.values.push(value.clone()); + value + } +} + +fn broadcast_bias(builder: &mut Builder, value: &Val, bias: &Val) -> Val { + let rank = value.ty.shape.len(); + let bias = builder.broadcast(bias, &[rank - 1], value.ty.shape.clone()); + builder.add(value, &bias) +} + +fn linear(builder: &mut Builder, value: &Val, weight: &Val, bias: Option<&Val>) -> Val { + let rank = value.ty.shape.len(); + let mut output_shape = value.ty.shape[..rank - 1].to_vec(); + output_shape.push(weight.ty.shape[0]); + let output = builder.dot_general(value, weight, &[], &[], &[rank - 1], &[1], output_shape); + bias.map_or(output.clone(), |bias| { + broadcast_bias(builder, &output, bias) + }) +} + +fn layer_norm(builder: &mut Builder, value: &Val, weight: &Val, bias: &Val, epsilon: f32) -> Val { + let rank = value.ty.shape.len(); + let axis = rank - 1; + let width = value.ty.shape[axis]; + let leading = value.ty.shape[..axis].to_vec(); + let zero = builder.const_f32(0.0); + let width_value = builder.const_f32(width as f32); + let width_value = builder.broadcast(&width_value, &[], leading.clone()); + let sum = builder.reduce_add(value, axis, &zero); + let mean = builder.divide(&sum, &width_value); + let mean = builder.broadcast( + &mean, + &(0..axis).collect::>(), + value.ty.shape.clone(), + ); + let centered = builder.subtract(value, &mean); + let squared = builder.multiply(¢ered, ¢ered); + let variance = builder.reduce_add(&squared, axis, &zero); + let variance = builder.divide(&variance, &width_value); + let epsilon = builder.const_f32(epsilon); + let epsilon = builder.broadcast(&epsilon, &[], leading.clone()); + let variance = builder.add(&variance, &epsilon); + let inv_std = builder.rsqrt(&variance); + let inv_std = builder.broadcast( + &inv_std, + &(0..axis).collect::>(), + value.ty.shape.clone(), + ); + let normalized = builder.multiply(¢ered, &inv_std); + let weight = builder.broadcast(weight, &[axis], value.ty.shape.clone()); + let bias = builder.broadcast(bias, &[axis], value.ty.shape.clone()); + let normalized = builder.multiply(&normalized, &weight); + builder.add(&normalized, &bias) +} + +fn tanh_gelu(builder: &mut Builder, value: &Val) -> Val { + let shape = value.ty.shape.clone(); + let half = builder.const_f32(0.5); + let half = builder.broadcast(&half, &[], shape.clone()); + let one = builder.const_f32(1.0); + let one = builder.broadcast(&one, &[], shape.clone()); + let coefficient = builder.const_f32(0.044_715); + let coefficient = builder.broadcast(&coefficient, &[], shape.clone()); + let scale = builder.const_f32(0.797_884_6); + let scale = builder.broadcast(&scale, &[], shape); + let squared = builder.multiply(value, value); + let cubed = builder.multiply(&squared, value); + let nonlinear = builder.multiply(&coefficient, &cubed); + let inner = builder.add(value, &nonlinear); + let scaled = builder.multiply(&scale, &inner); + let tanh = builder.tanh(&scaled); + let cdf = builder.add(&one, &tanh); + let half_value = builder.multiply(value, &half); + builder.multiply(&half_value, &cdf) +} + +fn silu(builder: &mut Builder, value: &Val) -> Val { + let shape = value.ty.shape.clone(); + let one = builder.const_f32(1.0); + let one = builder.broadcast(&one, &[], shape); + let negative = builder.negate(value); + let exponential = builder.exponential(&negative); + let denominator = builder.add(&one, &exponential); + builder.divide(value, &denominator) +} + +fn softmax_last(builder: &mut Builder, scores: &Val) -> Val { + let axis = scores.ty.shape.len() - 1; + let leading = scores.ty.shape[..axis].to_vec(); + let negative_infinity = builder.const_f32(f32::NEG_INFINITY); + let maximum = builder.reduce_max(scores, axis, &negative_infinity); + let maximum = builder.broadcast( + &maximum, + &(0..axis).collect::>(), + scores.ty.shape.clone(), + ); + let shifted = builder.subtract(scores, &maximum); + let exponentials = builder.exponential(&shifted); + let zero = builder.const_f32(0.0); + let denominator = builder.reduce_add(&exponentials, axis, &zero); + let denominator = builder.broadcast( + &denominator, + &(0..axis).collect::>(), + scores.ty.shape.clone(), + ); + debug_assert_eq!(denominator.ty.shape[..axis], leading); + builder.divide(&exponentials, &denominator) +} + +fn selected_slot(selected_layers: &[usize], layer: usize) -> Option { + selected_layers + .iter() + .position(|&selected| selected == layer) +} + +fn self_attention( + builder: &mut Builder, + hidden: &Val, + args: &mut Args, + config: &Molmo2VisionConfig, +) -> Val { + let q_weight = args.take(); + let q_bias = args.take(); + let k_weight = args.take(); + let k_bias = args.take(); + let v_weight = args.take(); + let v_bias = args.take(); + let o_weight = args.take(); + let o_bias = args.take(); + let q = linear(builder, hidden, &q_weight, Some(&q_bias)); + let k = linear(builder, hidden, &k_weight, Some(&k_bias)); + let v = linear(builder, hidden, &v_weight, Some(&v_bias)); + let crops = config.static_crops; + let tokens = config.patches_per_crop; + let q = builder.reshape(&q, vec![crops, tokens, config.heads, config.head_dim]); + let q = builder.transpose(&q, &[0, 2, 1, 3]); + let k = builder.reshape(&k, vec![crops, tokens, config.heads, config.head_dim]); + let k = builder.transpose(&k, &[0, 2, 1, 3]); + let v = builder.reshape(&v, vec![crops, tokens, config.heads, config.head_dim]); + let v = builder.transpose(&v, &[0, 2, 1, 3]); + let scores = builder.dot_general( + &q, + &k, + &[0, 1], + &[0, 1], + &[3], + &[3], + vec![crops, config.heads, tokens, tokens], + ); + let scale = builder.const_f32((config.head_dim as f32).powf(-0.5)); + let scale = builder.broadcast(&scale, &[], scores.ty.shape.clone()); + let scores = builder.multiply(&scores, &scale); + let probabilities = softmax_last(builder, &scores); + let context = builder.dot_general( + &probabilities, + &v, + &[0, 1], + &[0, 1], + &[3], + &[2], + vec![crops, config.heads, tokens, config.head_dim], + ); + let context = builder.transpose(&context, &[0, 2, 1, 3]); + let context = builder.reshape(&context, vec![crops, tokens, config.hidden]); + linear(builder, &context, &o_weight, Some(&o_bias)) +} + +fn indexed_pool( + builder: &mut Builder, + features: &Val, + signed_indices: &Val, + args: &mut Args, + config: &Molmo2VisionConfig, +) -> (Val, Val) { + let groups = config.static_pool_groups; + let group_size = config.pool_size; + let zero_i32 = builder.const_i32(0); + let zero_indices = builder.broadcast(&zero_i32, &[], vec![groups, group_size]); + let valid = builder.compare("GE", signed_indices, &zero_indices, "SIGNED"); + let safe = builder.maximum(signed_indices, &zero_indices); + let safe = builder.reshape(&safe, vec![groups, group_size, 1]); + let gathered = builder.gather_rows_nd(features, &safe); + let valid_f32 = builder.convert(&valid, "f32"); + let valid_features = builder.broadcast(&valid_f32, &[0, 1], gathered.ty.shape.clone()); + let gathered = builder.multiply(&gathered, &valid_features); + let zero_f32 = builder.const_f32(0.0); + let sums = builder.reduce_add(&gathered, 1, &zero_f32); + let counts = builder.reduce_add(&valid_f32, 1, &zero_f32); + let one = builder.const_f32(1.0); + let ones = builder.broadcast(&one, &[], vec![groups]); + let denominator = builder.maximum(&counts, &ones); + let denominator = builder.broadcast(&denominator, &[0], vec![groups, config.selected_width()]); + let query = builder.divide(&sums, &denominator); + + let q_weight = args.take(); + let q_bias = args.take(); + let k_weight = args.take(); + let k_bias = args.take(); + let v_weight = args.take(); + let v_bias = args.take(); + let o_weight = args.take(); + let o_bias = args.take(); + let q = linear(builder, &query, &q_weight, Some(&q_bias)); + let k = linear(builder, &gathered, &k_weight, Some(&k_bias)); + let v = linear(builder, &gathered, &v_weight, Some(&v_bias)); + let q = builder.reshape(&q, vec![groups, 1, config.pool_heads, config.pool_head_dim]); + let q = builder.transpose(&q, &[0, 2, 1, 3]); + let k = builder.reshape( + &k, + vec![groups, group_size, config.pool_heads, config.pool_head_dim], + ); + let k = builder.transpose(&k, &[0, 2, 1, 3]); + let v = builder.reshape( + &v, + vec![groups, group_size, config.pool_heads, config.pool_head_dim], + ); + let v = builder.transpose(&v, &[0, 2, 1, 3]); + let scores = builder.dot_general( + &q, + &k, + &[0, 1], + &[0, 1], + &[3], + &[3], + vec![groups, config.pool_heads, 1, group_size], + ); + let scale = builder.const_f32((config.pool_head_dim as f32).powf(-0.5)); + let scale = builder.broadcast(&scale, &[], scores.ty.shape.clone()); + let mut scores = builder.multiply(&scores, &scale); + if config.pooling_attention_mask { + let mask = builder.broadcast(&valid, &[0, 3], scores.ty.shape.clone()); + let masked = builder.const_f32(-1.0e30); + let masked = builder.broadcast(&masked, &[], scores.ty.shape.clone()); + scores = builder.select(&mask, &scores, &masked); + } + let probabilities = softmax_last(builder, &scores); + let context = builder.dot_general( + &probabilities, + &v, + &[0, 1], + &[0, 1], + &[3], + &[2], + vec![groups, config.pool_heads, 1, config.pool_head_dim], + ); + let context = builder.transpose(&context, &[0, 2, 1, 3]); + let context = builder.reshape(&context, vec![groups, config.pool_hidden]); + let pooled = linear(builder, &context, &o_weight, Some(&o_bias)); + (pooled, counts) +} + +pub(crate) fn emit_molmo2_vision(config: &Molmo2VisionConfig) -> String { + let specs = config.weight_specs(); + let mut args = Args::new(&specs); + let patch_weight = args.take(); + let patch_bias = args.take(); + let position_embedding = args.take(); + let patches = args.input( + Ty::f32(vec![ + config.static_crops, + config.patches_per_crop, + config.patch_dim, + ]), + "molmo2.patches", + ); + let signed_indices = args.input( + Ty::new(vec![config.static_pool_groups, config.pool_size], "i32"), + "molmo2.image_token_pooling.signed", + ); + let mut builder = Builder::new(); + let mut hidden = linear(&mut builder, &patches, &patch_weight, Some(&patch_bias)); + let position = builder.broadcast( + &position_embedding, + &[1, 2], + vec![config.static_crops, config.patches_per_crop, config.hidden], + ); + hidden = builder.add(&hidden, &position); + let mut selected = vec![None::; config.selected_layers.len()]; + for layer in 0..config.emitted_layers { + // Norm weights follow attention projection weights in the persisted + // schema. Pull them before emitting the attention and pass normalized + // hidden to the projection sequence. + let block_start = args.cursor; + let norm_weight = args.values[block_start + 8].clone(); + let norm_bias = args.values[block_start + 9].clone(); + let normalized = layer_norm( + &mut builder, + &hidden, + &norm_weight, + &norm_bias, + config.layer_norm_eps, + ); + let attention = self_attention(&mut builder, &normalized, &mut args, config); + // encoder_block consumes the already-taken attention schema, so finish + // this block explicitly. + let _attention_norm_weight = args.take(); + let _attention_norm_bias = args.take(); + let ffn_norm_weight = args.take(); + let ffn_norm_bias = args.take(); + let residual = builder.add(&hidden, &attention); + let normalized = layer_norm( + &mut builder, + &residual, + &ffn_norm_weight, + &ffn_norm_bias, + config.layer_norm_eps, + ); + let w1 = args.take(); + let b1 = args.take(); + let w2 = args.take(); + let b2 = args.take(); + let mlp = linear(&mut builder, &normalized, &w1, Some(&b1)); + let mlp = tanh_gelu(&mut builder, &mlp); + let mlp = linear(&mut builder, &mlp, &w2, Some(&b2)); + hidden = builder.add(&residual, &mlp); + if let Some(slot) = selected_slot(&config.selected_layers, layer) { + selected[slot] = Some(hidden.clone()); + } + } + let mut selected = selected.into_iter(); + let mut selected_features = match selected.next() { + Some(Some(feature)) => feature, + _ => unreachable!("validated Molmo2 selected layers must be emitted"), + }; + for feature in selected { + let feature = match feature { + Some(feature) => feature, + None => unreachable!("validated Molmo2 selected layer must be emitted"), + }; + selected_features = builder.concatenate(&selected_features, &feature, 2); + } + let selected_features = builder.reshape( + &selected_features, + vec![ + config.static_crops * config.patches_per_crop, + config.selected_width(), + ], + ); + let (pooled, _counts) = indexed_pool( + &mut builder, + &selected_features, + &signed_indices, + &mut args, + config, + ); + let w1 = args.take(); + let w2 = args.take(); + let w3 = args.take(); + let gate = linear(&mut builder, &pooled, &w1, None); + let gate = silu(&mut builder, &gate); + let up = linear(&mut builder, &pooled, &w3, None); + let projected = builder.multiply(&gate, &up); + let projected = linear(&mut builder, &projected, &w2, None); + assert_eq!( + args.cursor, + specs.len(), + "Molmo2 vision weight schema drifted" + ); + format!( + "module @molmo2_vision {{\n func.func public @main({signature}) -> {output} {{\n{body} return {value} : {output}\n }}\n}}\n", + signature = args.declarations.join(", "), + output = projected.ty.render(), + body = builder.body(), + value = projected.name, + ) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn emitted_pooling_clamps_then_masks_and_keeps_additive_merge_outside_graph() { + let config = Molmo2VisionConfig::from_json_strs( + &serde_json::json!({ + "model_type":"molmo2","image_patch_id":151938, + "vit_config":{"hidden_size":8,"intermediate_size":16,"num_attention_heads":2, + "head_dim":4,"num_hidden_layers":2,"image_default_input_size":[28,28], + "image_patch_size":14,"image_num_pos":4,"layer_norm_eps":1e-6}, + "adapter_config":{"hidden_size":8,"intermediate_size":12,"text_hidden_size":10, + "num_attention_heads":2,"head_dim":4,"vit_layers":[0,1], + "pooling_attention_mask":true} + }) + .to_string(), + &serde_json::json!({"patch_size":14,"max_crops":1,"overlap_margins":[0,0], + "pooling_size":[2,2],"size":{"height":28,"width":28}}) + .to_string(), + ) + .unwrap(); + let mlir = emit_molmo2_vision(&config); + assert!(mlir.contains("molmo2.image_token_pooling.signed")); + assert!(mlir.contains("stablehlo.compare GE")); + assert!(mlir.contains("stablehlo.maximum")); + assert!(mlir.contains("\"stablehlo.gather\"")); + assert!(mlir.contains("stablehlo.select")); + assert!(mlir.contains("vision_tower.image_vit.patch_embedding.weight")); + assert!(mlir.contains("vision_tower.image_projector.w3.weight")); + assert!(!mlir.contains("image_input_idx")); + } + + #[test] + fn selected_layers_keep_adapter_order_instead_of_encoder_order() { + assert_eq!(selected_slot(&[22, 16], 22), Some(0)); + assert_eq!(selected_slot(&[22, 16], 16), Some(1)); + assert_eq!(selected_slot(&[22, 16], 21), None); + } +} diff --git a/src/lib/mlxcel-xla/src/iree.rs b/src/lib/mlxcel-xla/src/iree.rs index e31100190..47bc620f0 100644 --- a/src/lib/mlxcel-xla/src/iree.rs +++ b/src/lib/mlxcel-xla/src/iree.rs @@ -59,7 +59,7 @@ use safetensors::{Dtype, SafeTensors}; use crate::emitter::{ Config, DeepStackConfig, Gemma3nConfig, Gemma3nWeightSpec, Precision, QuantConfig, - check_packed_supported, emit_decode_ragged_with, emit_decode_with, + WeightScheme, 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, @@ -85,7 +85,8 @@ use crate::{DeepStackFeatures, DeepStackPreparedPrefill, Gemma3nDensePle, Gemma3 // `dequantize_affine_stacked`). Both are pure-Rust and unit-tested without `iree`. use crate::weights::{ QuantPart, WeightSpec, bf16_to_f32, dequantize_affine, dequantize_affine_bf16_fused, - dequantize_affine_stacked, f16_to_f32, f32_le_to_f32, pack_f16, slice_rows, weight_specs, + dequantize_affine_f32, dequantize_affine_stacked, f16_to_f32, f32_le_to_f32, pack_f16, + slice_rows, weight_specs, }; /// Weight-buffer element dtype passed to the C shim (issue #516 per-weight ABI): @@ -1127,6 +1128,68 @@ fn resolve_dense_weight_sources( }) } +fn resolve_molmo2_weight_sources( + model_dir: &Path, + names: &[String], +) -> Result, String> { + let mut shards = std::fs::read_dir(model_dir) + .map_err(|error| format!("read {}: {error}", model_dir.display()))? + .filter_map(Result::ok) + .map(|entry| entry.path()) + .filter(|path| { + path.file_name() + .and_then(|name| name.to_str()) + .is_some_and(|name| name.ends_with(".safetensors")) + }) + .collect::>(); + shards.sort(); + let mut resolved = vec![None::; names.len()]; + for shard in shards { + let file = + File::open(&shard).map_err(|error| format!("open {}: {error}", shard.display()))?; + // Safety: the file is read-only for the lifetime of this header scan. + let mmap = unsafe { Mmap::map(&file) } + .map_err(|error| format!("mmap {}: {error}", shard.display()))?; + let tensors = SafeTensors::deserialize(&mmap) + .map_err(|error| format!("parse {}: {error}", shard.display()))?; + for (index, name) in names.iter().enumerate() { + let wrapped = language_model_tensor_name(name); + if tensors.tensor(&wrapped).is_err() { + continue; + } + if let Some(previous) = &resolved[index] { + return Err(format!( + "Molmo2 tensor {wrapped} occurs in {} and {}", + previous.display(), + shard.display() + )); + } + resolved[index] = Some(shard.clone()); + } + } + let missing = resolved + .iter() + .zip(names) + .filter_map(|(path, name)| path.is_none().then_some(name.as_str())) + .collect::>(); + if !missing.is_empty() { + return Err(format!( + "Molmo2 checkpoint is missing {} language graph weight(s): {}", + missing.len(), + missing + .iter() + .take(8) + .copied() + .collect::>() + .join(", ") + )); + } + resolved + .into_iter() + .map(|path| path.ok_or_else(|| "Molmo2 weight resolution drifted".to_string())) + .collect() +} + /// mlx-community Gemma3n quantized repacks retain the upstream index but rewrite /// the language backbone's safetensors prefix. Keep the graph's canonical HF /// argument order and only try this one architecture-specific alternate after @@ -1397,9 +1460,17 @@ fn load_weights( let (paths, scheme) = resolve_gemma3n_weight_sources(model_dir, &names)?; (paths, None, Some(scheme)) } - RuntimeConfig::Dense(_) => { - let (paths, scheme) = resolve_dense_weight_sources(model_dir, &names)?; - (paths, scheme, None) + RuntimeConfig::Dense(config) => { + if config.weight_scheme == WeightScheme::Molmo2 { + ( + resolve_molmo2_weight_sources(model_dir, &names)?, + Some(DenseCheckpointNames::LanguageModelPrefixed), + None, + ) + } else { + let (paths, scheme) = resolve_dense_weight_sources(model_dir, &names)?; + (paths, scheme, None) + } } }; @@ -1525,12 +1596,21 @@ fn load_weights( // mlx-lm emits the affine scales/biases as either f16 or bf16 // (Qwen3 / Gemma3-27B / Qwen3-MoE checkpoints use bf16); accept a // matching 16-bit pair and widen it to f32 in the dequantizer. - let sb_bf16 = match (scales.dtype(), biases.dtype()) { - (Dtype::F16, Dtype::F16) => false, - (Dtype::BF16, Dtype::BF16) => true, + let sb_dtype = match (scales.dtype(), biases.dtype()) { + (Dtype::F16, Dtype::F16) => Dtype::F16, + (Dtype::BF16, Dtype::BF16) => Dtype::BF16, + (Dtype::F32, Dtype::F32) + if matches!( + cfg, + RuntimeConfig::Dense(config) + if config.weight_scheme == WeightScheme::Molmo2 + ) => + { + Dtype::F32 + } (s, b) => { return Err(format!( - "{prefix} scales/biases dtype {s:?}/{b:?}, expected a matching F16 or BF16 pair" + "{prefix} scales/biases dtype {s:?}/{b:?}, expected a matching F16/BF16 pair or Molmo2 F32 pair" )); } }; @@ -1542,7 +1622,7 @@ fn load_weights( let (out, in_packed) = (shape[0], shape[1]); let in_ = in_packed * (32 / qc.bits); let d = if matches!(cfg, RuntimeConfig::Gemma3n(_)) { - if !sb_bf16 { + if sb_dtype != Dtype::BF16 { return Err(format!( "{prefix} uses F16 affine scales/biases, but Gemma3n requires BF16 affine metadata" )); @@ -1556,6 +1636,16 @@ fn load_weights( qc.bits, qc.group_size, ) + } else if sb_dtype == Dtype::F32 { + dequantize_affine_f32( + t.data(), + scales.data(), + biases.data(), + out, + in_packed, + qc.bits, + qc.group_size, + ) } else { dequantize_affine( t.data(), @@ -1565,7 +1655,7 @@ fn load_weights( in_packed, qc.bits, qc.group_size, - sb_bf16, + sb_dtype == Dtype::BF16, ) } .map_err(|e| format!("dequantize {name}: {e}"))?; @@ -1586,7 +1676,7 @@ fn load_weights( in_packed, qc.bits, qc.group_size, - sb_bf16, + sb_dtype == Dtype::BF16, ) .map_err(|e| format!("dequantize stacked {name}: {e}"))?; (d, vec![experts, out, in_]) diff --git a/src/lib/mlxcel-xla/src/lib.rs b/src/lib/mlxcel-xla/src/lib.rs index 2be9febdc..d76594964 100644 --- a/src/lib/mlxcel-xla/src/lib.rs +++ b/src/lib/mlxcel-xla/src/lib.rs @@ -76,6 +76,10 @@ mod aux_manifest; mod aux_smoke; #[cfg(feature = "iree")] mod iree; +#[cfg(any(feature = "iree", test))] +mod molmo2; +#[cfg(feature = "iree")] +mod molmo2_vision_runtime; #[cfg(feature = "iree")] mod phi4_audio; #[cfg(feature = "iree")] @@ -168,6 +172,14 @@ pub use emitter::{ }; #[cfg(feature = "diagnostics")] pub use iree::PreparedPrefillDiagnostics; +#[cfg(any(feature = "iree", test))] +pub use molmo2::{ + Molmo2InputError, Molmo2SafePooling, add_projected_features as add_molmo2_projected_features, +}; +#[cfg(feature = "iree")] +pub use molmo2_vision_runtime::{ + IreeMolmo2VisionProjector, Molmo2VisionInput, Molmo2VisionProjection, +}; #[cfg(feature = "iree")] pub use phi4_audio::{ PHI4MM_AUDIO_CHECKPOINT_REVISION, PHI4MM_AUDIO_FRAME_BUCKETS, Phi4AudioOutput, diff --git a/src/lib/mlxcel-xla/src/molmo2.rs b/src/lib/mlxcel-xla/src/molmo2.rs new file mode 100644 index 000000000..558fe304d --- /dev/null +++ b/src/lib/mlxcel-xla/src/molmo2.rs @@ -0,0 +1,402 @@ +// 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. + +//! Molmo2-specific indexed-pooling and additive-merge contracts. +//! +//! Molmo2 processor output contains negative pooling sentinels. Native graph +//! inputs retain those signed indices verbatim; this module separately derives +//! safe gather indices and a valid mask so clamped patch zero can never +//! contribute to a pooled token. Molmo v1's sparse `image_input_idx` mapping is +//! deliberately not represented here. + +use std::fmt; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum Molmo2InputError { + ZeroHiddenSize, + EmptyPoolingGroup, + PoolingShape { + values: usize, + groups: usize, + group_size: usize, + }, + PoolingIndex { + group: usize, + offset: usize, + value: i32, + patch_count: usize, + }, + EmbeddingShape { + values: usize, + tokens: usize, + hidden_size: usize, + }, + ProjectedShape { + values: usize, + tokens: usize, + hidden_size: usize, + }, + ProjectedTokenCount { + positions: usize, + projected_tokens: usize, + }, + NonFinite { + tensor: &'static str, + index: usize, + }, + ShapeOverflow, +} + +impl fmt::Display for Molmo2InputError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::ZeroHiddenSize => f.write_str("Molmo2 hidden size must be positive"), + Self::EmptyPoolingGroup => { + f.write_str("Molmo2 pooling groups must contain at least one patch index") + } + Self::PoolingShape { + values, + groups, + group_size, + } => write!( + f, + "Molmo2 pooling has {values} indices, expected {groups} * {group_size}" + ), + Self::PoolingIndex { + group, + offset, + value, + patch_count, + } => write!( + f, + "Molmo2 pooling group {group} offset {offset} nonnegative index {value} is outside [0,{patch_count})" + ), + Self::EmbeddingShape { + values, + tokens, + hidden_size, + } => write!( + f, + "Molmo2 text embeddings have {values} values, expected {tokens} * {hidden_size}" + ), + Self::ProjectedShape { + values, + tokens, + hidden_size, + } => write!( + f, + "Molmo2 projected features have {values} values, expected {tokens} * {hidden_size}" + ), + Self::ProjectedTokenCount { + positions, + projected_tokens, + } => write!( + f, + "Molmo2 prompt has {positions} image_patch_id positions but projector returned {projected_tokens} tokens" + ), + Self::NonFinite { tensor, index } => { + write!( + f, + "Molmo2 {tensor} contains a non-finite value at flat index {index}" + ) + } + Self::ShapeOverflow => f.write_str("Molmo2 tensor shape overflowed"), + } + } +} + +impl std::error::Error for Molmo2InputError {} + +/// Signed processor indices plus the separate values consumed by a safe gather. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Molmo2SafePooling { + /// Original processor output, including every negative invalid sentinel. + pub signed_indices: Vec, + /// Gather-only indices. Invalid entries are clamped to patch zero. + pub gather_indices: Vec, + /// One byte per index (`1` valid, `0` invalid), suitable for an IREE bool input. + pub valid_mask: Vec, + /// Number of valid patches in each group. Zero is retained for all-invalid groups. + pub valid_counts: Vec, + pub groups: usize, + pub group_size: usize, +} + +impl Molmo2SafePooling { + /// Validate processor indices without erasing their negative sentinel values. + /// + /// The graph must multiply gathered K/V values by `valid_mask` before every + /// reduction. Query-mean denominators use `max(valid_counts, 1)`, so an + /// all-invalid group is deterministic and cannot divide by zero. + pub fn prepare( + signed_indices: &[i32], + groups: usize, + group_size: usize, + patch_count: usize, + ) -> Result { + if group_size == 0 { + return Err(Molmo2InputError::EmptyPoolingGroup); + } + let expected = groups + .checked_mul(group_size) + .ok_or(Molmo2InputError::ShapeOverflow)?; + if signed_indices.len() != expected { + return Err(Molmo2InputError::PoolingShape { + values: signed_indices.len(), + groups, + group_size, + }); + } + + let mut gather_indices = Vec::with_capacity(expected); + let mut valid_mask = Vec::with_capacity(expected); + let mut valid_counts = vec![0i32; groups]; + for (flat_index, &value) in signed_indices.iter().enumerate() { + let group = flat_index / group_size; + let offset = flat_index % group_size; + if value < 0 { + gather_indices.push(0); + valid_mask.push(0); + continue; + } + let index = usize::try_from(value).map_err(|_| Molmo2InputError::PoolingIndex { + group, + offset, + value, + patch_count, + })?; + if index >= patch_count { + return Err(Molmo2InputError::PoolingIndex { + group, + offset, + value, + patch_count, + }); + } + gather_indices.push(value); + valid_mask.push(1); + valid_counts[group] += 1; + } + + Ok(Self { + signed_indices: signed_indices.to_vec(), + gather_indices, + valid_mask, + valid_counts, + groups, + group_size, + }) + } + + /// Safe denominators for query means. All-invalid groups use one and stay + /// zero because their gathered values are masked before reduction. + #[must_use] + pub fn mean_denominators(&self) -> Vec { + self.valid_counts + .iter() + .map(|&count| count.max(1)) + .collect() + } +} + +/// Find logical image positions and add projected Molmo2 features in order. +/// +/// This intentionally requires the caller's logical expanded token sequence; +/// it never accepts Molmo v1 `image_input_idx` coordinates. +pub fn add_projected_features( + token_ids: &[i32], + image_patch_id: i32, + text_embeddings: &mut [f32], + hidden_size: usize, + projected_features: &[f32], +) -> Result, Molmo2InputError> { + if hidden_size == 0 { + return Err(Molmo2InputError::ZeroHiddenSize); + } + let expected_embeddings = token_ids + .len() + .checked_mul(hidden_size) + .ok_or(Molmo2InputError::ShapeOverflow)?; + if text_embeddings.len() != expected_embeddings { + return Err(Molmo2InputError::EmbeddingShape { + values: text_embeddings.len(), + tokens: token_ids.len(), + hidden_size, + }); + } + if let Some((index, _)) = text_embeddings + .iter() + .enumerate() + .find(|(_, value)| !value.is_finite()) + { + return Err(Molmo2InputError::NonFinite { + tensor: "text embeddings", + index, + }); + } + if projected_features.len() % hidden_size != 0 { + return Err(Molmo2InputError::ProjectedShape { + values: projected_features.len(), + tokens: projected_features.len() / hidden_size, + hidden_size, + }); + } + if let Some((index, _)) = projected_features + .iter() + .enumerate() + .find(|(_, value)| !value.is_finite()) + { + return Err(Molmo2InputError::NonFinite { + tensor: "projected features", + index, + }); + } + + let positions = token_ids + .iter() + .enumerate() + .filter_map(|(index, &token)| (token == image_patch_id).then_some(index)) + .collect::>(); + let projected_tokens = projected_features.len() / hidden_size; + if positions.len() != projected_tokens { + return Err(Molmo2InputError::ProjectedTokenCount { + positions: positions.len(), + projected_tokens, + }); + } + + // Preflight every sum so an overflow cannot leave a partially-mutated + // prepared payload behind. + for (feature_index, &position) in positions.iter().enumerate() { + let destination = position * hidden_size; + let source = feature_index * hidden_size; + for offset in 0..hidden_size { + if !(text_embeddings[destination + offset] + projected_features[source + offset]) + .is_finite() + { + return Err(Molmo2InputError::NonFinite { + tensor: "merged embeddings", + index: destination + offset, + }); + } + } + } + for (feature_index, &position) in positions.iter().enumerate() { + let destination = position * hidden_size; + let source = feature_index * hidden_size; + for offset in 0..hidden_size { + text_embeddings[destination + offset] += projected_features[source + offset]; + } + } + Ok(positions) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn negative_pooling_sentinels_are_preserved_but_cannot_leak_patch_zero() { + let pooling = Molmo2SafePooling::prepare(&[4, -1, 2, -7, -1, -1, -1, -1], 2, 4, 5) + .expect("valid signed pooling"); + assert_eq!(pooling.signed_indices, vec![4, -1, 2, -7, -1, -1, -1, -1]); + assert_eq!(pooling.gather_indices, vec![4, 0, 2, 0, 0, 0, 0, 0]); + assert_eq!(pooling.valid_mask, vec![1, 0, 1, 0, 0, 0, 0, 0]); + assert_eq!(pooling.valid_counts, vec![2, 0]); + assert_eq!(pooling.mean_denominators(), vec![2, 1]); + + let patches = [100.0, 1.0, 2.0, 3.0, 4.0]; + let masked_sums = (0..pooling.groups) + .map(|group| { + (0..pooling.group_size) + .map(|offset| { + let flat = group * pooling.group_size + offset; + patches[pooling.gather_indices[flat] as usize] + * f32::from(pooling.valid_mask[flat]) + }) + .sum::() + }) + .collect::>(); + assert_eq!(masked_sums, vec![6.0, 0.0]); + } + + #[test] + fn pooling_rejects_shape_and_positive_out_of_range_indices() { + assert!(matches!( + Molmo2SafePooling::prepare(&[0, 1, 2], 1, 4, 3), + Err(Molmo2InputError::PoolingShape { .. }) + )); + assert!(matches!( + Molmo2SafePooling::prepare(&[0, 1, 3, -1], 1, 4, 3), + Err(Molmo2InputError::PoolingIndex { + group: 0, + offset: 2, + value: 3, + .. + }) + )); + } + + #[test] + fn image_patch_features_are_added_in_scanned_position_order() { + let tokens = [9, 151_938, 8, 151_938]; + let mut embeddings = vec![ + 1.0, 1.0, // text + 10.0, 20.0, // image slot zero + 2.0, 2.0, // text + 30.0, 40.0, // image slot one + ]; + let positions = + add_projected_features(&tokens, 151_938, &mut embeddings, 2, &[0.5, 1.5, 2.5, 3.5]) + .expect("matching image positions"); + assert_eq!(positions, vec![1, 3]); + assert_eq!(embeddings, vec![1.0, 1.0, 10.5, 21.5, 2.0, 2.0, 32.5, 43.5]); + assert_ne!(embeddings[2..4], [0.5, 1.5]); + } + + #[test] + fn additive_merge_rejects_count_mismatch_and_non_finite_values() { + let tokens = [151_938, 7]; + let mut embeddings = vec![0.0; 4]; + assert!(matches!( + add_projected_features(&tokens, 151_938, &mut embeddings, 0, &[]), + Err(Molmo2InputError::ZeroHiddenSize) + )); + assert!(matches!( + add_projected_features(&tokens, 151_938, &mut embeddings, 2, &[]), + Err(Molmo2InputError::ProjectedTokenCount { + positions: 1, + projected_tokens: 0 + }) + )); + assert!(matches!( + add_projected_features(&tokens, 151_938, &mut embeddings, 2, &[f32::NAN, 0.0]), + Err(Molmo2InputError::NonFinite { + tensor: "projected features", + index: 0 + }) + )); + + let mut overflowing = vec![f32::MAX, 7.0, 0.0, 0.0]; + let snapshot = overflowing.clone(); + assert!(matches!( + add_projected_features(&tokens, 151_938, &mut overflowing, 2, &[f32::MAX, 1.0]), + Err(Molmo2InputError::NonFinite { + tensor: "merged embeddings", + index: 0 + }) + )); + assert_eq!(overflowing, snapshot, "overflow rejection must be atomic"); + } +} diff --git a/src/lib/mlxcel-xla/src/molmo2_vision_runtime.rs b/src/lib/mlxcel-xla/src/molmo2_vision_runtime.rs new file mode 100644 index 000000000..338cfba6e --- /dev/null +++ b/src/lib/mlxcel-xla/src/molmo2_vision_runtime.rs @@ -0,0 +1,488 @@ +// 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 std::collections::{BTreeMap, BTreeSet}; +use std::fs::File; +use std::io::Read as _; +use std::path::{Path, PathBuf}; +use std::process::Command; +use std::time::Instant; + +use memmap2::Mmap; +use safetensors::{Dtype, SafeTensors}; +use sha2::{Digest, Sha256}; + +use crate::aux::{ + AuxiliaryInput, AuxiliaryOutput, AuxiliaryTensorDType, AuxiliaryWeight, AuxiliaryWeightDType, + IreeAuxiliaryModule, +}; +use crate::aux_manifest::{AuxiliaryArtifactContract, ensure_qualified_auxiliary_artifact}; +use crate::emitter::{Molmo2VisionConfig, Molmo2VisionWeightSpec, emit_molmo2_vision}; +use crate::iree::{cached_vmfb_path, compile_one_to, iree_compile_bin, target_flags}; +use crate::molmo2::Molmo2SafePooling; +use crate::weights::{bf16_to_f32, f16_to_f32, f32_le_to_f32}; + +const ENTRY_NAME: &str = "molmo2_vision.main"; + +#[derive(Debug, Clone, PartialEq)] +pub struct Molmo2VisionProjection { + pub values: Vec, + pub shape: [usize; 2], + pub signed_pooling_indices: Vec, + pub valid_pooling_counts: Vec, + pub elapsed_seconds: f64, + pub upload_bytes: usize, + pub transfer_bytes: usize, +} + +#[derive(Debug, Clone, Copy)] +pub struct Molmo2VisionInput<'a> { + pub patches: &'a [f32], + pub patches_shape: [usize; 3], + pub image_token_pooling: &'a [i32], + pub pooling_shape: [usize; 2], + pub image_grid: [i32; 4], + pub image_num_crops: usize, +} + +fn hex(bytes: &[u8]) -> String { + let mut output = String::with_capacity(bytes.len() * 2); + for byte in bytes { + use std::fmt::Write as _; + let _ = write!(output, "{byte:02x}"); + } + output +} + +fn sha256(bytes: &[u8]) -> String { + hex(&Sha256::digest(bytes)) +} + +fn sha256_file(path: &Path) -> Result { + let mut file = File::open(path).map_err(|error| format!("open {}: {error}", path.display()))?; + let mut digest = Sha256::new(); + let mut buffer = [0u8; 64 * 1024]; + loop { + let read = file + .read(&mut buffer) + .map_err(|error| format!("read {}: {error}", path.display()))?; + if read == 0 { + break; + } + digest.update(&buffer[..read]); + } + Ok(hex(&digest.finalize())) +} + +fn generation_identity(compiler: &Path, flags: &[&str], mlir: &str) -> Result { + let output = Command::new(compiler) + .arg("--version") + .output() + .map_err(|error| format!("run {} --version: {error}", compiler.display()))?; + if !output.status.success() { + return Err(format!( + "{} --version failed: {}", + compiler.display(), + String::from_utf8_lossy(&output.stderr) + )); + } + Ok(format!( + "compiler={};compiler_sha256={};version={};flags={flags:?};mlir_sha256={}", + compiler.display(), + sha256_file(compiler)?, + String::from_utf8_lossy(&output.stdout).trim(), + sha256(mlir.as_bytes()) + )) +} + +fn model_shards(model_dir: &Path) -> Result, String> { + let mut shards = std::fs::read_dir(model_dir) + .map_err(|error| format!("read {}: {error}", model_dir.display()))? + .map(|entry| { + entry + .map(|entry| entry.path()) + .map_err(|error| error.to_string()) + }) + .collect::, _>>()?; + shards.retain(|path| { + path.file_name() + .and_then(|name| name.to_str()) + .is_some_and(|name| name.ends_with(".safetensors")) + }); + shards.sort(); + if shards.is_empty() { + return Err(format!( + "no safetensors checkpoint shards found in {}", + model_dir.display() + )); + } + Ok(shards) +} + +fn resolve_shards( + model_dir: &Path, + specs: &[Molmo2VisionWeightSpec], +) -> Result, String> { + let required = specs + .iter() + .map(|spec| spec.name.clone()) + .collect::>(); + let mut locations = BTreeMap::::new(); + for shard in model_shards(model_dir)? { + let file = + File::open(&shard).map_err(|error| format!("open {}: {error}", shard.display()))?; + // Safety: read-only map remains live for the header scan. + let mmap = unsafe { Mmap::map(&file) } + .map_err(|error| format!("mmap {}: {error}", shard.display()))?; + let tensors = SafeTensors::deserialize(&mmap) + .map_err(|error| format!("parse {}: {error}", shard.display()))?; + for name in tensors.names() { + if required.contains(name) + && let Some(previous) = locations.insert(name.to_string(), shard.clone()) + { + return Err(format!( + "Molmo2 vision tensor {name} occurs in {} and {}", + previous.display(), + shard.display() + )); + } + } + } + let missing = specs + .iter() + .filter(|spec| !locations.contains_key(&spec.name)) + .map(|spec| spec.name.as_str()) + .collect::>(); + if !missing.is_empty() { + return Err(format!( + "checkpoint is missing {} Molmo2 vision tensor(s): {}", + missing.len(), + missing + .iter() + .take(8) + .copied() + .collect::>() + .join(", ") + )); + } + Ok(specs + .iter() + .map(|spec| locations[&spec.name].clone()) + .collect()) +} + +fn native_f32_bytes(values: &[f32]) -> Vec { + values + .iter() + .flat_map(|value| value.to_ne_bytes()) + .collect() +} + +fn finite(label: &str, values: &[f32]) -> Result<(), String> { + if let Some((index, value)) = values + .iter() + .enumerate() + .find(|(_, value)| !value.is_finite()) + { + return Err(format!( + "{label} contains non-finite value {value} at flat index {index}" + )); + } + Ok(()) +} + +fn load_weights( + model_dir: &Path, + specs: &[Molmo2VisionWeightSpec], +) -> Result<(Vec, String), String> { + let shards = resolve_shards(model_dir, specs)?; + let mut by_shard = BTreeMap::<&Path, Vec>::new(); + for (index, shard) in shards.iter().enumerate() { + by_shard.entry(shard).or_default().push(index); + } + let mut loaded = (0..specs.len()) + .map(|_| None) + .collect::>>(); + let mut schema = vec![String::new(); specs.len()]; + for (shard, indices) in by_shard { + let file = + File::open(shard).map_err(|error| format!("open {}: {error}", shard.display()))?; + // Safety: read-only map remains live while tensors are copied. + let mmap = unsafe { Mmap::map(&file) } + .map_err(|error| format!("mmap {}: {error}", shard.display()))?; + let tensors = SafeTensors::deserialize(&mmap) + .map_err(|error| format!("parse {}: {error}", shard.display()))?; + for index in indices { + let spec = &specs[index]; + let tensor = tensors + .tensor(&spec.name) + .map_err(|error| format!("load Molmo2 vision tensor {}: {error}", spec.name))?; + if tensor.shape() != spec.shape { + return Err(format!( + "Molmo2 vision tensor {} has shape {:?}, expected {:?}", + spec.name, + tensor.shape(), + spec.shape + )); + } + let values = match tensor.dtype() { + Dtype::BF16 => bf16_to_f32(tensor.data()), + Dtype::F16 => f16_to_f32(tensor.data()), + Dtype::F32 => f32_le_to_f32(tensor.data()), + dtype => { + return Err(format!( + "Molmo2 vision tensor {} has unsupported dtype {dtype:?}", + spec.name + )); + } + }; + finite(&format!("Molmo2 vision tensor {}", spec.name), &values)?; + schema[index] = format!("{}:{:?}:{:?}", spec.name, tensor.dtype(), tensor.shape()); + loaded[index] = Some(AuxiliaryWeight { + name: spec.name.clone(), + bytes: native_f32_bytes(&values), + dtype: AuxiliaryWeightDType::Float32, + shape: spec.shape.clone(), + }); + } + } + let loaded = loaded + .into_iter() + .map(|weight| { + weight.ok_or_else(|| "resolved Molmo2 vision tensor was not loaded".to_string()) + }) + .collect::, _>>()?; + Ok((loaded, schema.join("\n"))) +} + +fn f32_bytes(values: &[f32]) -> &[u8] { + // Safety: f32 has no invalid bit patterns and the result is borrowed. + unsafe { + std::slice::from_raw_parts(values.as_ptr().cast::(), std::mem::size_of_val(values)) + } +} + +fn i32_bytes(values: &[i32]) -> &[u8] { + // Safety: i32 has no invalid bit patterns and the result is borrowed. + unsafe { + std::slice::from_raw_parts(values.as_ptr().cast::(), std::mem::size_of_val(values)) + } +} + +fn decode_output(bytes: &[u8]) -> Result, String> { + if !bytes.len().is_multiple_of(4) { + return Err("Molmo2 IREE output byte count is not f32-aligned".to_string()); + } + let values = bytes + .chunks_exact(4) + .map(|bytes| f32::from_ne_bytes([bytes[0], bytes[1], bytes[2], bytes[3]])) + .collect::>(); + finite("Molmo2 IREE projected output", &values)?; + Ok(values) +} + +pub struct IreeMolmo2VisionProjector { + module: IreeAuxiliaryModule, + config: Molmo2VisionConfig, +} + +impl IreeMolmo2VisionProjector { + pub fn load(model_dir: &Path, device: &str) -> Result { + let config = Molmo2VisionConfig::from_model_dir(model_dir)?; + let mlir = emit_molmo2_vision(&config); + let compiler = iree_compile_bin()?; + if !compiler.is_file() { + return Err(format!("iree-compile not found at {}", compiler.display())); + } + let flags = target_flags(device)?; + let cache = std::env::temp_dir().join("mlxcel-xla-molmo2-vision-vmfb"); + std::fs::create_dir_all(&cache) + .map_err(|error| format!("mkdir {}: {error}", cache.display()))?; + let (weights, checkpoint_schema) = load_weights(model_dir, &config.weight_specs())?; + let contract = AuxiliaryArtifactContract::new( + ENTRY_NAME, + format!( + "{};checkpoint_schema_sha256={}", + config.fingerprint(), + sha256(checkpoint_schema.as_bytes()) + ), + generation_identity(&compiler, flags, &mlir)?, + )?; + let vmfb = cached_vmfb_path(&compiler, &mlir, flags, &cache, "molmo2-vision", 0); + ensure_qualified_auxiliary_artifact(&vmfb, &contract, &weights, |temporary| { + compile_one_to( + &compiler, + &mlir, + flags, + &cache, + "molmo2-vision", + 0, + temporary, + ) + })?; + let module = IreeAuxiliaryModule::load(device, &vmfb, &contract, weights)?; + Ok(Self { module, config }) + } + + pub fn image_patch_id(&self) -> i32 { + self.config.image_patch_id + } + + pub fn text_hidden_size(&self) -> usize { + self.config.text_hidden + } + + pub fn artifact_fingerprint(&self) -> u64 { + self.module.fingerprint() + } + + pub fn project( + &mut self, + input: Molmo2VisionInput<'_>, + ) -> Result { + let [crops, patches, patch_dim] = input.patches_shape; + if crops != input.image_num_crops + || crops == 0 + || crops > self.config.static_crops + || patches != self.config.patches_per_crop + || patch_dim != self.config.patch_dim + { + return Err(format!( + "Molmo2 patch shape {:?}, image_num_crops={} disagrees with static [{},{},{}]", + input.patches_shape, + input.image_num_crops, + self.config.static_crops, + self.config.patches_per_crop, + self.config.patch_dim + )); + } + let patch_values = crops + .checked_mul(patches) + .and_then(|value| value.checked_mul(patch_dim)) + .ok_or_else(|| "Molmo2 patch shape overflowed".to_string())?; + if input.patches.len() != patch_values { + return Err(format!( + "Molmo2 patch payload has {} values, expected {patch_values}", + input.patches.len() + )); + } + finite("Molmo2 patches", input.patches)?; + let [groups, group_size] = input.pooling_shape; + let grid = input + .image_grid + .iter() + .try_fold((), |(), value| { + (*value >= 0) + .then_some(()) + .ok_or_else(|| "Molmo2 image grid contains a negative dimension".to_string()) + }) + .and_then(|()| Ok(input.image_grid.map(|value| value as usize)))?; + if !self.config.valid_runtime_geometry(crops, grid) { + return Err(format!( + "Molmo2 crop count {crops} and image grid {:?} disagree with processor geometry", + input.image_grid + )); + } + let [lo_h, lo_w, hi_h, hi_w] = grid; + let grid_groups = lo_h + .checked_mul(lo_w) + .and_then(|low| { + hi_h.checked_mul(hi_w) + .and_then(|high| low.checked_add(high)) + }) + .ok_or_else(|| "Molmo2 image grid overflowed".to_string())?; + if groups != grid_groups + || groups > self.config.static_pool_groups + || group_size != self.config.pool_size + { + return Err(format!( + "Molmo2 pooling shape {:?} disagrees with grid {:?} and static [{},{}]", + input.pooling_shape, + input.image_grid, + self.config.static_pool_groups, + self.config.pool_size + )); + } + let safe = Molmo2SafePooling::prepare( + input.image_token_pooling, + groups, + group_size, + crops * patches, + ) + .map_err(|error| error.to_string())?; + let static_patch_values = + self.config.static_crops * self.config.patches_per_crop * self.config.patch_dim; + let mut padded_patches = vec![0.0f32; static_patch_values]; + padded_patches[..input.patches.len()].copy_from_slice(input.patches); + let mut signed = vec![-1i32; self.config.static_pool_groups * self.config.pool_size]; + signed[..safe.signed_indices.len()].copy_from_slice(&safe.signed_indices); + let output_shape = [self.config.static_pool_groups, self.config.text_hidden]; + let mut output = vec![0u8; output_shape.iter().product::() * 4]; + let patch_shape = [ + self.config.static_crops, + self.config.patches_per_crop, + self.config.patch_dim, + ]; + let pooling_shape = [self.config.static_pool_groups, self.config.pool_size]; + let started = Instant::now(); + self.module.invoke( + &[ + AuxiliaryInput { + bytes: f32_bytes(&padded_patches), + dtype: AuxiliaryTensorDType::Float32, + shape: &patch_shape, + }, + AuxiliaryInput { + bytes: i32_bytes(&signed), + dtype: AuxiliaryTensorDType::Int32, + shape: &pooling_shape, + }, + ], + &mut [AuxiliaryOutput { + bytes: &mut output, + dtype: AuxiliaryTensorDType::Float32, + shape: &output_shape, + }], + )?; + let all_values = decode_output(&output)?; + let valid_groups = safe + .valid_counts + .iter() + .enumerate() + .filter_map(|(index, &count)| (count > 0).then_some(index)) + .collect::>(); + let mut values = Vec::with_capacity(valid_groups.len() * self.config.text_hidden); + for group in valid_groups { + let start = group * self.config.text_hidden; + values.extend_from_slice(&all_values[start..start + self.config.text_hidden]); + } + Ok(Molmo2VisionProjection { + shape: [ + values.len() / self.config.text_hidden, + self.config.text_hidden, + ], + values, + signed_pooling_indices: safe.signed_indices, + valid_pooling_counts: safe.valid_counts, + elapsed_seconds: started.elapsed().as_secs_f64(), + upload_bytes: std::mem::size_of_val(padded_patches.as_slice()) + + std::mem::size_of_val(signed.as_slice()), + transfer_bytes: output.len(), + }) + } +} + +#[cfg(test)] +#[path = "molmo2_vision_runtime_tests.rs"] +mod tests; diff --git a/src/lib/mlxcel-xla/src/molmo2_vision_runtime_tests.rs b/src/lib/mlxcel-xla/src/molmo2_vision_runtime_tests.rs new file mode 100644 index 000000000..999ddf7df --- /dev/null +++ b/src/lib/mlxcel-xla/src/molmo2_vision_runtime_tests.rs @@ -0,0 +1,33 @@ +// 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 output_decode_rejects_non_finite_values() { + let mut bytes = 1.0f32.to_ne_bytes().to_vec(); + bytes.extend_from_slice(&f32::INFINITY.to_ne_bytes()); + assert!(decode_output(&bytes).unwrap_err().contains("flat index 1")); +} + +#[test] +fn compiler_identity_changes_when_same_path_and_version_bytes_change() { + let path = std::env::temp_dir().join(format!("mlxcel-molmo2-compiler-{}", std::process::id())); + std::fs::write(&path, b"first").unwrap(); + let first = sha256_file(&path).unwrap(); + std::fs::write(&path, b"second").unwrap(); + let second = sha256_file(&path).unwrap(); + std::fs::remove_file(path).ok(); + assert_ne!(first, second); +} diff --git a/src/lib/mlxcel-xla/src/weight_names.rs b/src/lib/mlxcel-xla/src/weight_names.rs index ab7b20dc6..94b1fc79f 100644 --- a/src/lib/mlxcel-xla/src/weight_names.rs +++ b/src/lib/mlxcel-xla/src/weight_names.rs @@ -146,6 +146,28 @@ pub(crate) fn scheme_names(scheme: WeightScheme) -> SchemeNames { pre_ff_norm: "pre_feedforward_layernorm.weight", post_ff_norm: "post_feedforward_layernorm.weight", }, + WeightScheme::Molmo2 => SchemeNames { + embed: "model.wte.embedding", + final_norm: "model.ln_f.weight", + lm_head: "lm_head.weight", + layer_stem: "model.blocks.", + down: "mlp.ff_out.weight", + gate: "mlp.ff_proj.weight", + input_layernorm: "attn_norm.weight", + post_attention_layernorm: "ff_norm.weight", + up: "mlp.ff_proj.weight", + k_proj: "self_attn.att_proj.weight", + o_proj: "self_attn.attn_out.weight", + q_proj: "self_attn.att_proj.weight", + v_proj: "self_attn.att_proj.weight", + k_bias: "self_attn.att_proj.bias", + q_bias: "self_attn.att_proj.bias", + v_bias: "self_attn.att_proj.bias", + q_norm: "self_attn.q_norm.weight", + k_norm: "self_attn.k_norm.weight", + pre_ff_norm: "ff_norm.weight", + post_ff_norm: "ff_norm.weight", + }, } } diff --git a/src/lib/mlxcel-xla/src/weights.rs b/src/lib/mlxcel-xla/src/weights.rs index 22338ab7e..87700af82 100644 --- a/src/lib/mlxcel-xla/src/weights.rs +++ b/src/lib/mlxcel-xla/src/weights.rs @@ -135,6 +135,30 @@ fn phi4_base_projection(cfg: &Config, name: String) -> String { } } +fn fused_gate_up(cfg: &Config, prefix: &str) -> (String, (usize, usize), (usize, usize)) { + if cfg.weight_scheme == crate::emitter::WeightScheme::Molmo2 { + ( + format!("{prefix}mlp.ff_proj.weight"), + (cfg.inter, 2 * cfg.inter), + (0, cfg.inter), + ) + } else { + ( + format!("{prefix}mlp.gate_up_proj.weight"), + (0, cfg.inter), + (cfg.inter, 2 * cfg.inter), + ) + } +} + +fn fused_qkv_name(cfg: &Config, prefix: &str) -> String { + if cfg.weight_scheme == crate::emitter::WeightScheme::Molmo2 { + format!("{prefix}self_attn.att_proj.weight") + } else { + format!("{prefix}self_attn.qkv_proj.weight") + } +} + fn push_phi4_lora_pair(out: &mut Vec, stem: &str, adapter: &str) { out.push(WeightSpec::Proj(format!("{stem}.lora_A.{adapter}.weight"))); out.push(WeightSpec::Proj(format!("{stem}.lora_B.{adapter}.weight"))); @@ -149,7 +173,6 @@ fn weight_specs_q(cfg: &Config, quant: bool) -> Vec { let hd = cfg.head_dim; let nq = cfg.n_q * hd; let nkv = cfg.n_kv * hd; - let inter = cfg.inter; let gated = !cfg.dense_mlp; let has_post = !cfg.parallel_block; @@ -180,10 +203,11 @@ fn weight_specs_q(cfg: &Config, quant: bool) -> Vec { // gate (gated MLP only; the first half of gate_up_proj for a fused Phi3). if gated && !moe_layer { if cfg.fused_gate_up { + let (name, (start, end), _) = fused_gate_up(cfg, &p); out.push(WeightSpec::Rows { - name: phi4_base_projection(cfg, format!("{p}mlp.gate_up_proj.weight")), - start: 0, - end: inter, + name: phi4_base_projection(cfg, name), + start, + end, }); } else { push_proj( @@ -208,10 +232,11 @@ fn weight_specs_q(cfg: &Config, quant: bool) -> Vec { // Skipped on a MoE layer (issue #500), which has no dense up projection. if !moe_layer { if cfg.fused_gate_up { + let (name, _, (start, end)) = fused_gate_up(cfg, &p); out.push(WeightSpec::Rows { - name: phi4_base_projection(cfg, format!("{p}mlp.gate_up_proj.weight")), - start: inter, - end: 2 * inter, + name: phi4_base_projection(cfg, name), + start, + end, }); } else if cfg.dense_mlp { out.push(WeightSpec::Whole(format!("{p}mlp.c_fc.weight"))); @@ -225,7 +250,7 @@ fn weight_specs_q(cfg: &Config, quant: bool) -> Vec { } // wk, wo, wq, wv (JAX-alphabetical; a fused Phi3 qkv_proj is [Q|K|V] rows). if cfg.fused_qkv { - let qkv = phi4_base_projection(cfg, format!("{p}self_attn.qkv_proj.weight")); + let qkv = phi4_base_projection(cfg, fused_qkv_name(cfg, &p)); out.push(WeightSpec::Rows { name: qkv.clone(), start: nq, @@ -519,6 +544,63 @@ pub(crate) fn dequantize_affine( bits: usize, group_size: usize, scales_bf16: bool, +) -> Result, String> { + dequantize_affine_with_metadata( + packed, + scales, + biases, + out, + in_packed, + bits, + group_size, + if scales_bf16 { + AffineMetadataDType::BFloat16 + } else { + AffineMetadataDType::Float16 + }, + ) +} + +/// Molmo2's converted 8-bit checkpoint keeps affine scales and biases as F32. +#[allow(clippy::too_many_arguments)] +pub(crate) fn dequantize_affine_f32( + packed: &[u8], + scales: &[u8], + biases: &[u8], + out: usize, + in_packed: usize, + bits: usize, + group_size: usize, +) -> Result, String> { + dequantize_affine_with_metadata( + packed, + scales, + biases, + out, + in_packed, + bits, + group_size, + AffineMetadataDType::Float32, + ) +} + +#[derive(Clone, Copy)] +enum AffineMetadataDType { + Float16, + BFloat16, + Float32, +} + +#[allow(clippy::too_many_arguments)] +fn dequantize_affine_with_metadata( + packed: &[u8], + scales: &[u8], + biases: &[u8], + out: usize, + in_packed: usize, + bits: usize, + group_size: usize, + metadata_dtype: AffineMetadataDType, ) -> Result, String> { if !(bits == 4 || bits == 8) { return Err(format!( @@ -540,12 +622,10 @@ pub(crate) fn dequantize_affine( out * in_packed * 4 )); } - // mlx-lm stores the affine scale/bias in either f16 or bf16; widen the - // matching 16-bit format to f32 (both are exact in f32). - let (scales, biases) = if scales_bf16 { - (bf16_to_f32(scales), bf16_to_f32(biases)) - } else { - (f16_to_f32(scales), f16_to_f32(biases)) + let (scales, biases) = match metadata_dtype { + AffineMetadataDType::Float16 => (f16_to_f32(scales), f16_to_f32(biases)), + AffineMetadataDType::BFloat16 => (bf16_to_f32(scales), bf16_to_f32(biases)), + AffineMetadataDType::Float32 => (f32_le_to_f32(scales), f32_le_to_f32(biases)), }; if scales.len() != out * n_groups || biases.len() != out * n_groups { return Err(format!( @@ -810,6 +890,57 @@ mod tests { ); } + #[test] + fn molmo2_uses_olmo_names_and_reversed_fused_swiglu_halves() { + let config = Config::from_json_str( + r#"{ + "model_type":"molmo2", + "quantization":{"bits":8,"group_size":64}, + "text_config":{ + "model_type":"molmo2_text","hidden_size":8,"intermediate_size":16, + "num_hidden_layers":1,"num_attention_heads":2,"num_key_value_heads":1, + "head_dim":4,"layer_norm_eps":1e-6,"rope_theta":5000000, + "max_position_embeddings":32,"vocab_size":24,"hidden_act":"silu", + "use_qk_norm":true + } + }"#, + ) + .expect("Molmo2 nested text config"); + assert_eq!(config.weight_scheme, crate::emitter::WeightScheme::Molmo2); + assert!(config.fused_qkv && config.fused_gate_up); + assert_eq!( + config.qk_norm, + Some(crate::emitter::QkNorm { + per_head: true, + one_plus: false + }) + ); + let specs = weight_specs_q(&config, false); + assert_eq!(specs[0].tensor_name(), "model.wte.embedding"); + assert_eq!(specs[1].tensor_name(), "model.ln_f.weight"); + assert_eq!(specs[2].tensor_name(), "lm_head.weight"); + assert!(matches!( + &specs[4], + WeightSpec::Rows { name, start: 16, end: 32 } + if name == "model.blocks.0.mlp.ff_proj.weight" + )); + assert!(matches!( + &specs[7], + WeightSpec::Rows { name, start: 0, end: 16 } + if name == "model.blocks.0.mlp.ff_proj.weight" + )); + assert!( + specs + .iter() + .any(|spec| spec.tensor_name() == "model.blocks.0.self_attn.att_proj.weight") + ); + assert!( + specs + .iter() + .any(|spec| spec.tensor_name() == "model.blocks.0.self_attn.q_norm.weight") + ); + } + #[test] #[ignore = "requires PHI4MM_MODEL_DIR pointing at the pinned official checkpoint"] fn real_phi4mm_index_covers_decoder_and_complete_lora_schema() { @@ -1193,6 +1324,21 @@ mod tests { assert_eq!(w, vec![30.0, 50.0, 14.0, 19.0]); } + #[test] + fn dequantize_affine_accepts_molmo2_f32_metadata() { + let packed = [0x0Au8, 0x14, 0x1E, 0x28]; + let scales = [2.0f32, 0.5] + .into_iter() + .flat_map(f32::to_le_bytes) + .collect::>(); + let biases = [10.0f32, -1.0] + .into_iter() + .flat_map(f32::to_le_bytes) + .collect::>(); + let w = dequantize_affine_f32(&packed, &scales, &biases, 1, 1, 8, 2).unwrap(); + assert_eq!(w, vec![30.0, 50.0, 14.0, 19.0]); + } + /// A packed buffer whose size disagrees with `[out, in_packed]` is rejected. #[test] fn dequantize_affine_rejects_size_mismatch() { diff --git a/src/loading/mod.rs b/src/loading/mod.rs index b69e916fc..7a7223cba 100644 --- a/src/loading/mod.rs +++ b/src/loading/mod.rs @@ -64,7 +64,8 @@ pub(crate) use self::vlm::load_qwen2_vl_iree_host_preprocessor; pub use self::vlm::load_qwen3_omni_speech; #[cfg(feature = "xla-iree")] pub(crate) use self::vlm::{ - Phi4MMXlaVisionComponents, load_phi4mm_xla_media_components, load_phi4mm_xla_text_embeddings, + Phi4MMXlaVisionComponents, load_molmo2_xla_text_embeddings, load_phi4mm_xla_media_components, + load_phi4mm_xla_text_embeddings, }; pub(crate) use self::vlm::{ensure_supported_muse_weight_map, normalize_muse_glimmer_weights}; diff --git a/src/loading/vlm.rs b/src/loading/vlm.rs index 6a1d0b2af..c3a0357c0 100644 --- a/src/loading/vlm.rs +++ b/src/loading/vlm.rs @@ -144,7 +144,8 @@ pub(crate) use siglip::{load_aya_vision_vlm, load_paligemma_vlm}; pub(crate) use smolvlm::load_smolvlm_vlm; #[cfg(feature = "xla-iree")] pub(crate) use special::{ - Phi4MMXlaVisionComponents, load_phi4mm_xla_media_components, load_phi4mm_xla_text_embeddings, + Phi4MMXlaVisionComponents, load_molmo2_xla_text_embeddings, load_phi4mm_xla_media_components, + load_phi4mm_xla_text_embeddings, }; pub(crate) use special::{ load_llama4_vlm, load_minicpmo_vlm, load_minicpmv4_6_vlm, load_molmo_point_vlm, load_molmo_vlm, diff --git a/src/loading/vlm_special.rs b/src/loading/vlm_special.rs index ed26cb0a4..1124a4972 100644 --- a/src/loading/vlm_special.rs +++ b/src/loading/vlm_special.rs @@ -1369,6 +1369,47 @@ pub(crate) fn load_phi4mm_xla_text_embeddings( )) } +/// Load only Molmo2's dual text embedding table for prepared XLA prefills. +#[cfg(feature = "xla-iree")] +pub(crate) fn load_molmo2_xla_text_embeddings( + model_path: &Path, +) -> Result<(models::molmo2::Molmo2Embedding, usize, usize)> { + let (_config_str, full_config) = read_sanitized_vlm_config(model_path)?; + if full_config.get("model_type").and_then(Value::as_str) != Some("molmo2") { + return Err(anyhow::anyhow!( + "{} is not a Molmo2 checkpoint", + model_path.display() + )); + } + let text = full_config + .get("text_config") + .ok_or_else(|| anyhow::anyhow!("Molmo2 config is missing text_config"))?; + let hidden_size = text + .get("hidden_size") + .and_then(Value::as_u64) + .ok_or_else(|| anyhow::anyhow!("Molmo2 text_config.hidden_size is required"))? + as usize; + let max_sequence_len = text + .get("max_position_embeddings") + .and_then(Value::as_u64) + .ok_or_else(|| anyhow::anyhow!("Molmo2 text_config.max_position_embeddings is required"))? + as usize; + let raw_weights = super::load_vlm_weights_common_filtered_canonical(model_path, |name| { + matches!( + name, + "language_model.model.wte.embedding" + | "language_model.model.wte.new_embedding" + | "model.transformer.wte.embedding" + | "model.transformer.wte.new_embedding" + ) + })?; + let weights = remap_molmo2_weights(raw_weights); + let embeddings = + models::molmo2::Molmo2Embedding::from_weights(&weights, "language_model.model.wte") + .map_err(|error| anyhow::anyhow!("invalid Molmo2 text embedding tables: {error}"))?; + Ok((embeddings, hidden_size, max_sequence_len)) +} + /// Filtered Phi4MM host components used by the XLA prepared-prefill producer. /// /// The struct deliberately contains no text decoder, LM head, audio encoder, or diff --git a/src/multimodal/host_preprocessor.rs b/src/multimodal/host_preprocessor.rs index f791d261d..1aa9f3caa 100644 --- a/src/multimodal/host_preprocessor.rs +++ b/src/multimodal/host_preprocessor.rs @@ -42,6 +42,9 @@ use super::vlm_prompt::{ImageTokenBlockError, ImageTokenBlockInfo, apply_image_t mod export; #[cfg(feature = "xla-iree")] use super::qwen_vl::insert_qwen_vl_image_tokens; +#[cfg(feature = "xla-iree")] +#[path = "molmo2_xla_preprocessor.rs"] +mod molmo2_xla; use export::export_mlx_tensor; use export::{ build_prepared_prefill, export_llava_prefill, usize_to_i32, validate_embedding_shape, @@ -52,6 +55,8 @@ use export::{ // turns that into a hard error. The gate has to match the single call site. #[cfg(any(feature = "xla-iree", test))] use export::export_qwen2_vl_prefill; +#[cfg(feature = "xla-iree")] +pub use molmo2_xla::Molmo2IreeHostPreprocessor; /// Vision implementation selected for OpenXLA multimodal preprocessing. #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -179,47 +184,78 @@ pub fn load_xla_image_preprocessor( .to_string(), )); } - 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" - .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 = Qwen2VlIreeHostPreprocessor::load(model_path, &device)?; - tracing::info!( - vision_backend = "iree", - vision_device = %device, - family = "qwen2_vl", - "OpenXLA multimodal vision backend selected" - ); - return Ok(Some(Box::new(preprocessor))); - } - // No MLX fallback exists for this family, so an image-capable session - // cannot be built here. Report it as a build-configuration error with - // the remedy attached: both callers surface this string verbatim, so - // the rebuild instruction is the only actionable part the operator - // gets. - #[cfg(not(feature = "xla-iree"))] - { - return Err(HostPreprocessorError::InvalidConfig( - "Qwen2-VL XLA image execution requires the xla-iree feature; rebuild mlxcel with \ - `--features xla-iree` (this family has no MLX vision fallback)" - .to_string(), - )); - } - } - if model_type != crate::models::ModelType::LlavaVLM { + if model_type != crate::models::ModelType::LlavaVLM + && model_type != crate::models::ModelType::Molmo2VLM + && model_type != crate::models::ModelType::Qwen2VL + { return Ok(None); } let policy = XlaVisionBackendPolicy::from_env()?; - load_llava_image_preprocessor(model_path, policy) + match model_type { + crate::models::ModelType::Qwen2VL => { + if policy == XlaVisionBackendPolicy::Host { + return Err(HostPreprocessorError::InvalidConfig( + "Qwen2-VL XLA vision has no MLX fallback; MLXCEL_XLA_VISION_BACKEND=host is unsupported" + .to_string(), + )); + } + #[cfg(feature = "xla-iree")] + { + let device = std::env::var("MLXCEL_XLA_DEVICE") + .unwrap_or_else(|_| mlxcel_xla::default_device().to_string()); + let preprocessor = Qwen2VlIreeHostPreprocessor::load(model_path, &device)?; + tracing::info!( + vision_backend = "iree", + vision_device = %device, + family = "qwen2_vl", + "OpenXLA multimodal vision backend selected" + ); + Ok(Some(Box::new(preprocessor))) + } + // No MLX fallback exists for this family, so an image-capable session + // cannot be built here. Report it as a build-configuration error with + // the remedy attached: both callers surface this string verbatim, so + // the rebuild instruction is the only actionable part the operator + // gets. + #[cfg(not(feature = "xla-iree"))] + { + Err(HostPreprocessorError::InvalidConfig( + "Qwen2-VL XLA image execution requires the xla-iree feature; rebuild mlxcel \ + with `--features xla-iree` (this family has no MLX vision fallback)" + .to_string(), + )) + } + } + crate::models::ModelType::LlavaVLM => load_llava_image_preprocessor(model_path, policy), + crate::models::ModelType::Molmo2VLM => { + #[cfg(feature = "xla-iree")] + { + if policy == XlaVisionBackendPolicy::Host { + return Ok(None); + } + let device = std::env::var("MLXCEL_XLA_DEVICE") + .unwrap_or_else(|_| mlxcel_xla::default_device().to_string()); + let preprocessor = Molmo2IreeHostPreprocessor::load(model_path, &device)?; + tracing::info!( + vision_backend = "iree", + vision_device = %device, + "OpenXLA Molmo2 image preprocessing stage ready" + ); + Ok(Some(Box::new(preprocessor))) + } + #[cfg(not(feature = "xla-iree"))] + { + if policy == XlaVisionBackendPolicy::Iree { + return Err(HostPreprocessorError::InvalidConfig( + "MLXCEL_XLA_VISION_BACKEND=iree requires the xla-iree feature".to_string(), + )); + } + Ok(None) + } + } + _ => Ok(None), + } } fn load_llava_host_preprocessor_boxed( diff --git a/src/multimodal/molmo2_xla_preprocessor.rs b/src/multimodal/molmo2_xla_preprocessor.rs new file mode 100644 index 000000000..7c8852560 --- /dev/null +++ b/src/multimodal/molmo2_xla_preprocessor.rs @@ -0,0 +1,347 @@ +// 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 std::cell::RefCell; +use std::path::Path; + +use image::DynamicImage; +use mlxcel_core::session::{OwnedTensor, PreparedPrefill, PreparedTensorDType}; + +use super::export::{ + build_prepared_prefill, export_mlx_tensor, usize_to_i32, validate_embedding_shape, + validate_sequence_capacity, +}; +use super::{HostMultimodalPreprocessor, HostPreprocessorError, XlaVisionBackend}; +use crate::models::molmo2::Molmo2Embedding; +use crate::vision::processors::molmo2::Molmo2Processor; + +const IMAGE_START_ID: i32 = 151936; +const IMAGE_END_ID: i32 = 151937; +const IMAGE_PATCH_ID: i32 = 151938; +const IMAGE_COL_ID: i32 = 151939; +const LOW_RES_IMAGE_START_ID: i32 = 151940; +const IMAGE_PLACEHOLDER_ID: i32 = 151941; + +pub struct Molmo2IreeHostPreprocessor { + processor: Molmo2Processor, + text_embeddings: Molmo2Embedding, + projector: RefCell, + hidden_size: usize, + max_sequence_len: usize, + device: String, +} + +impl Molmo2IreeHostPreprocessor { + pub fn load(model_path: &Path, device: &str) -> Result { + let (text_embeddings, hidden_size, max_sequence_len) = + crate::loading::load_molmo2_xla_text_embeddings(model_path) + .map_err(|error| HostPreprocessorError::WeightLoad(error.to_string()))?; + let preprocessor_path = model_path.join("preprocessor_config.json"); + let preprocessor_text = std::fs::read_to_string(&preprocessor_path).map_err(|error| { + HostPreprocessorError::InvalidConfig(format!( + "read {}: {error}", + preprocessor_path.display() + )) + })?; + let config: serde_json::Value = + serde_json::from_str(&preprocessor_text).map_err(|error| { + HostPreprocessorError::InvalidConfig(format!( + "parse {}: {error}", + preprocessor_path.display() + )) + })?; + let pair = |key: &str| -> Result<(usize, usize), HostPreprocessorError> { + let values = config + .get(key) + .and_then(serde_json::Value::as_array) + .ok_or_else(|| { + HostPreprocessorError::InvalidConfig(format!( + "Molmo2 preprocessor {key} must contain two integers" + )) + })?; + let first = values.first().and_then(serde_json::Value::as_u64); + let second = values.get(1).and_then(serde_json::Value::as_u64); + match (first, second) { + (Some(first), Some(second)) => Ok((first as usize, second as usize)), + _ => Err(HostPreprocessorError::InvalidConfig(format!( + "Molmo2 preprocessor {key} must contain two nonnegative integers" + ))), + } + }; + let size = config + .get("size") + .and_then(serde_json::Value::as_object) + .and_then(|size| { + Some(( + size.get("height")?.as_u64()? as usize, + size.get("width")?.as_u64()? as usize, + )) + }) + .ok_or_else(|| { + HostPreprocessorError::InvalidConfig( + "Molmo2 preprocessor size.height/width are required".to_string(), + ) + })?; + let max_crops = config + .get("max_crops") + .and_then(serde_json::Value::as_u64) + .ok_or_else(|| { + HostPreprocessorError::InvalidConfig( + "Molmo2 preprocessor max_crops is required".to_string(), + ) + })? as usize; + let patch_size = config + .get("patch_size") + .and_then(serde_json::Value::as_u64) + .ok_or_else(|| { + HostPreprocessorError::InvalidConfig( + "Molmo2 preprocessor patch_size is required".to_string(), + ) + })? as usize; + let processor = Molmo2Processor::new( + max_crops, + Some(pair("overlap_margins")?), + Some(patch_size), + Some(pair("pooling_size")?), + Some(size), + ); + let projector = mlxcel_xla::IreeMolmo2VisionProjector::load(model_path, device) + .map_err(HostPreprocessorError::Iree)?; + if projector.image_patch_id() != IMAGE_PATCH_ID + || projector.text_hidden_size() != hidden_size + { + return Err(HostPreprocessorError::InvalidConfig( + "Molmo2 vision/text token or hidden-size contract mismatch".to_string(), + )); + } + Ok(Self { + processor, + text_embeddings, + projector: RefCell::new(projector), + hidden_size, + max_sequence_len, + device: device.to_string(), + }) + } + + fn image_tokens(grid: [i32; 4]) -> Result, HostPreprocessorError> { + let [lo_h, lo_w, hi_h, hi_w] = grid; + let dimensions = [lo_h, lo_w, hi_h, hi_w] + .map(|value| usize::try_from(value).map_err(|_| HostPreprocessorError::ShapeOverflow)) + .into_iter() + .collect::, _>>()?; + let mut tokens = Vec::new(); + tokens.push(LOW_RES_IMAGE_START_ID); + for _ in 0..dimensions[0] { + tokens.extend(std::iter::repeat_n(IMAGE_PATCH_ID, dimensions[1])); + tokens.push(IMAGE_COL_ID); + } + tokens.push(IMAGE_END_ID); + tokens.push(IMAGE_START_ID); + for _ in 0..dimensions[2] { + tokens.extend(std::iter::repeat_n(IMAGE_PATCH_ID, dimensions[3])); + tokens.push(IMAGE_COL_ID); + } + tokens.push(IMAGE_END_ID); + Ok(tokens) + } + + fn expand_prompt( + token_ids: &[i32], + image_tokens: &[i32], + ) -> Result, HostPreprocessorError> { + let positions = token_ids + .iter() + .enumerate() + .filter_map(|(index, &token)| (token == IMAGE_PLACEHOLDER_ID).then_some(index)) + .collect::>(); + if positions.len() > 1 { + return Err(HostPreprocessorError::InvalidConfig(format!( + "Molmo2 prompt contains {} image placeholders but one image is supported", + positions.len() + ))); + } + let Some(position) = positions.first().copied() else { + let mut expanded = image_tokens.to_vec(); + expanded.extend_from_slice(token_ids); + return Ok(expanded); + }; + let mut expanded = Vec::with_capacity(token_ids.len() - 1 + image_tokens.len()); + expanded.extend_from_slice(&token_ids[..position]); + expanded.extend_from_slice(image_tokens); + expanded.extend_from_slice(&token_ids[position + 1..]); + Ok(expanded) + } + + fn prepare_one( + &self, + token_ids: &[i32], + image: &DynamicImage, + ) -> Result { + let processed = self.processor.preprocess_image(image); + let image_tokens = Self::image_tokens(processed.image_grid)?; + let logical_tokens = Self::expand_prompt(token_ids, &image_tokens)?; + validate_sequence_capacity(logical_tokens.len(), self.max_sequence_len)?; + let input_ids = mlxcel_core::from_slice_i32( + &logical_tokens, + &[1, usize_to_i32(logical_tokens.len(), "sequence length")?], + ); + let text = mlxcel_core::astype( + &self.text_embeddings.forward(&input_ids), + mlxcel_core::dtype::FLOAT32, + ); + validate_embedding_shape( + &mlxcel_core::array_shape(&text), + logical_tokens.len(), + self.hidden_size, + "Molmo2 dual embedding table", + )?; + let text = export_mlx_tensor(&text, "Molmo2 text embeddings")?; + if text.dtype != PreparedTensorDType::Float32 { + return Err(HostPreprocessorError::InvalidConfig( + "Molmo2 prepared embeddings must be Float32".to_string(), + )); + } + let mut text_values = text + .bytes + .chunks_exact(4) + .map(|bytes| f32::from_ne_bytes([bytes[0], bytes[1], bytes[2], bytes[3]])) + .collect::>(); + let patches_shape = processed + .pixel_values_shape + .map(|value| usize::try_from(value).map_err(|_| HostPreprocessorError::ShapeOverflow)) + .into_iter() + .collect::, _>>()?; + let pooling_shape = processed + .image_token_pooling_shape + .map(|value| usize::try_from(value).map_err(|_| HostPreprocessorError::ShapeOverflow)) + .into_iter() + .collect::, _>>()?; + let mut projector = self.projector.try_borrow_mut().map_err(|_| { + HostPreprocessorError::Iree( + "concurrent/re-entrant Molmo2 vision invocation is unsupported".to_string(), + ) + })?; + let projection = projector + .project(mlxcel_xla::Molmo2VisionInput { + patches: &processed.pixel_values, + patches_shape: [patches_shape[0], patches_shape[1], patches_shape[2]], + image_token_pooling: &processed.image_token_pooling, + pooling_shape: [pooling_shape[0], pooling_shape[1]], + image_grid: processed.image_grid, + image_num_crops: usize::try_from(processed.image_num_crops) + .map_err(|_| HostPreprocessorError::ShapeOverflow)?, + }) + .map_err(HostPreprocessorError::Iree)?; + let positions = mlxcel_xla::add_molmo2_projected_features( + &logical_tokens, + IMAGE_PATCH_ID, + &mut text_values, + self.hidden_size, + &projection.values, + ) + .map_err(|error| HostPreprocessorError::InvalidConfig(error.to_string()))?; + let bytes = text_values + .into_iter() + .flat_map(f32::to_ne_bytes) + .collect::>(); + let embeddings = OwnedTensor::new( + bytes, + PreparedTensorDType::Float32, + vec![1, logical_tokens.len(), self.hidden_size], + )?; + tracing::info!( + vision_backend = "iree", + vision_device = %self.device, + image_crops = processed.image_num_crops, + image_tokens = positions.len(), + iree_vision_seconds = projection.elapsed_seconds, + "OpenXLA Molmo2 projection completed" + ); + build_prepared_prefill(logical_tokens, embeddings, 1, positions.len(), "molmo2") + } +} + +impl HostMultimodalPreprocessor for Molmo2IreeHostPreprocessor { + fn backend(&self) -> XlaVisionBackend { + XlaVisionBackend::Iree + } + + fn prepare( + &self, + token_ids: &[i32], + images: &[DynamicImage], + ) -> Result { + match images { + [image] => self.prepare_one(token_ids, image), + _ => Err(HostPreprocessorError::InvalidConfig(format!( + "Molmo2 XLA requires exactly one image, got {}", + images.len() + ))), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn prompt_expansion_preserves_framing_and_patch_order() { + let image = Molmo2IreeHostPreprocessor::image_tokens([1, 2, 2, 1]).unwrap(); + assert_eq!( + image, + vec![ + LOW_RES_IMAGE_START_ID, + IMAGE_PATCH_ID, + IMAGE_PATCH_ID, + IMAGE_COL_ID, + IMAGE_END_ID, + IMAGE_START_ID, + IMAGE_PATCH_ID, + IMAGE_COL_ID, + IMAGE_PATCH_ID, + IMAGE_COL_ID, + IMAGE_END_ID, + ] + ); + let expanded = + Molmo2IreeHostPreprocessor::expand_prompt(&[7, IMAGE_PLACEHOLDER_ID, 8], &image) + .unwrap(); + assert_eq!(expanded.first(), Some(&7)); + assert_eq!(expanded.last(), Some(&8)); + assert_eq!( + expanded + .iter() + .filter(|&&token| token == IMAGE_PATCH_ID) + .count(), + 4 + ); + } + + #[test] + fn prompt_without_placeholder_prefixes_image_and_rejects_duplicates() { + assert_eq!( + Molmo2IreeHostPreprocessor::expand_prompt(&[7, 8], &[1, 2]).unwrap(), + vec![1, 2, 7, 8] + ); + assert!( + Molmo2IreeHostPreprocessor::expand_prompt( + &[IMAGE_PLACEHOLDER_ID, IMAGE_PLACEHOLDER_ID], + &[1] + ) + .is_err() + ); + } +} From 482d56aa812e42bfe0bc9cf473cae25d872e2299 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 24 Jul 2026 18:52:06 +0900 Subject: [PATCH 02/20] fix(xla): satisfy Molmo2 runtime lints --- src/lib/mlxcel-xla/src/molmo2.rs | 2 +- src/lib/mlxcel-xla/src/molmo2_vision_runtime.rs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/lib/mlxcel-xla/src/molmo2.rs b/src/lib/mlxcel-xla/src/molmo2.rs index 558fe304d..d77017cbd 100644 --- a/src/lib/mlxcel-xla/src/molmo2.rs +++ b/src/lib/mlxcel-xla/src/molmo2.rs @@ -245,7 +245,7 @@ pub fn add_projected_features( index, }); } - if projected_features.len() % hidden_size != 0 { + if !projected_features.len().is_multiple_of(hidden_size) { return Err(Molmo2InputError::ProjectedShape { values: projected_features.len(), tokens: projected_features.len() / hidden_size, diff --git a/src/lib/mlxcel-xla/src/molmo2_vision_runtime.rs b/src/lib/mlxcel-xla/src/molmo2_vision_runtime.rs index 338cfba6e..2640a7fb3 100644 --- a/src/lib/mlxcel-xla/src/molmo2_vision_runtime.rs +++ b/src/lib/mlxcel-xla/src/molmo2_vision_runtime.rs @@ -387,7 +387,7 @@ impl IreeMolmo2VisionProjector { .then_some(()) .ok_or_else(|| "Molmo2 image grid contains a negative dimension".to_string()) }) - .and_then(|()| Ok(input.image_grid.map(|value| value as usize)))?; + .map(|()| input.image_grid.map(|value| value as usize))?; if !self.config.valid_runtime_geometry(crops, grid) { return Err(format!( "Molmo2 crop count {crops} and image grid {:?} disagree with processor geometry", From b9c4d09fa0821b207c3144fb98c6e41f183e1295 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 24 Jul 2026 20:11:45 +0900 Subject: [PATCH 03/20] fix(xla): harden Molmo2 pooling preflight --- .../mlxcel-xla/src/emitter/molmo2_vision.rs | 34 +++++-- src/lib/mlxcel-xla/src/molmo2.rs | 94 +++++++++++++++++-- .../mlxcel-xla/src/molmo2_vision_runtime.rs | 14 ++- src/multimodal/molmo2_xla_preprocessor.rs | 5 + 4 files changed, 124 insertions(+), 23 deletions(-) diff --git a/src/lib/mlxcel-xla/src/emitter/molmo2_vision.rs b/src/lib/mlxcel-xla/src/emitter/molmo2_vision.rs index 587f53cb6..fc9fce343 100644 --- a/src/lib/mlxcel-xla/src/emitter/molmo2_vision.rs +++ b/src/lib/mlxcel-xla/src/emitter/molmo2_vision.rs @@ -243,9 +243,16 @@ fn indexed_pool( let zero_f32 = builder.const_f32(0.0); let sums = builder.reduce_add(&gathered, 1, &zero_f32); let counts = builder.reduce_add(&valid_f32, 1, &zero_f32); - let one = builder.const_f32(1.0); - let ones = builder.broadcast(&one, &[], vec![groups]); - let denominator = builder.maximum(&counts, &ones); + let denominator = if config.pooling_attention_mask { + let one = builder.const_f32(1.0); + let ones = builder.broadcast(&one, &[], vec![groups]); + builder.maximum(&counts, &ones) + } else { + // Match the MLX reference: invalid entries are zeroed above, but an + // unmasked pooling query is the mean over the full fixed-size window. + let group_size = builder.const_f32(group_size as f32); + builder.broadcast(&group_size, &[], vec![groups]) + }; let denominator = builder.broadcast(&denominator, &[0], vec![groups, config.selected_width()]); let query = builder.divide(&sums, &denominator); @@ -426,9 +433,8 @@ pub(crate) fn emit_molmo2_vision(config: &Molmo2VisionConfig) -> String { mod tests { use super::*; - #[test] - fn emitted_pooling_clamps_then_masks_and_keeps_additive_merge_outside_graph() { - let config = Molmo2VisionConfig::from_json_strs( + fn test_config(pooling_attention_mask: bool) -> Molmo2VisionConfig { + Molmo2VisionConfig::from_json_strs( &serde_json::json!({ "model_type":"molmo2","image_patch_id":151938, "vit_config":{"hidden_size":8,"intermediate_size":16,"num_attention_heads":2, @@ -436,14 +442,19 @@ mod tests { "image_patch_size":14,"image_num_pos":4,"layer_norm_eps":1e-6}, "adapter_config":{"hidden_size":8,"intermediate_size":12,"text_hidden_size":10, "num_attention_heads":2,"head_dim":4,"vit_layers":[0,1], - "pooling_attention_mask":true} + "pooling_attention_mask":pooling_attention_mask} }) .to_string(), &serde_json::json!({"patch_size":14,"max_crops":1,"overlap_margins":[0,0], "pooling_size":[2,2],"size":{"height":28,"width":28}}) .to_string(), ) - .unwrap(); + .unwrap() + } + + #[test] + fn emitted_pooling_clamps_then_masks_and_keeps_additive_merge_outside_graph() { + let config = test_config(true); let mlir = emit_molmo2_vision(&config); assert!(mlir.contains("molmo2.image_token_pooling.signed")); assert!(mlir.contains("stablehlo.compare GE")); @@ -455,6 +466,13 @@ mod tests { assert!(!mlir.contains("image_input_idx")); } + #[test] + fn unmasked_pooling_uses_full_window_denominator() { + let mlir = emit_molmo2_vision(&test_config(false)); + assert!(mlir.contains("stablehlo.constant dense<0x40800000> : tensor")); + assert!(!mlir.contains("stablehlo.select")); + } + #[test] fn selected_layers_keep_adapter_order_instead_of_encoder_order() { assert_eq!(selected_slot(&[22, 16], 22), Some(0)); diff --git a/src/lib/mlxcel-xla/src/molmo2.rs b/src/lib/mlxcel-xla/src/molmo2.rs index d77017cbd..24b7d0bb1 100644 --- a/src/lib/mlxcel-xla/src/molmo2.rs +++ b/src/lib/mlxcel-xla/src/molmo2.rs @@ -51,6 +51,12 @@ pub enum Molmo2InputError { positions: usize, projected_tokens: usize, }, + ActiveTokenCount { + active_groups: usize, + grid_groups: usize, + prompt_positions: usize, + all_invalid_groups: Vec, + }, NonFinite { tensor: &'static str, index: usize, @@ -105,6 +111,15 @@ impl fmt::Display for Molmo2InputError { f, "Molmo2 prompt has {positions} image_patch_id positions but projector returned {projected_tokens} tokens" ), + Self::ActiveTokenCount { + active_groups, + grid_groups, + prompt_positions, + all_invalid_groups, + } => write!( + f, + "Molmo2 projected active rows ({active_groups}) and grid rows ({grid_groups}) must match prompt image_patch_id positions ({prompt_positions}) before native invocation; all-invalid pooling groups: {all_invalid_groups:?}" + ), Self::NonFinite { tensor, index } => { write!( f, @@ -199,14 +214,49 @@ impl Molmo2SafePooling { }) } - /// Safe denominators for query means. All-invalid groups use one and stay - /// zero because their gathered values are masked before reduction. + /// Safe denominators for query means under the configured attention-mask policy. + /// + /// Masked pooling averages only valid patches and clamps all-invalid groups + /// to one. Unmasked pooling matches the MLX reference by averaging the + /// zero-masked values over the full fixed-size pooling window. #[must_use] - pub fn mean_denominators(&self) -> Vec { - self.valid_counts + pub fn mean_denominators(&self, pooling_attention_mask: bool) -> Vec { + if pooling_attention_mask { + self.valid_counts + .iter() + .map(|&count| count.max(1)) + .collect() + } else { + vec![i32::try_from(self.group_size).unwrap_or(i32::MAX); self.groups] + } + } + + pub(crate) fn active_groups_for_prompt( + &self, + grid_groups: usize, + prompt_positions: usize, + ) -> Result, Molmo2InputError> { + let active = self + .valid_counts .iter() - .map(|&count| count.max(1)) - .collect() + .enumerate() + .filter_map(|(index, &count)| (count > 0).then_some(index)) + .collect::>(); + if grid_groups != prompt_positions || active.len() != prompt_positions { + let all_invalid_groups = self + .valid_counts + .iter() + .enumerate() + .filter_map(|(index, &count)| (count == 0).then_some(index)) + .collect::>(); + return Err(Molmo2InputError::ActiveTokenCount { + active_groups: active.len(), + grid_groups, + prompt_positions, + all_invalid_groups, + }); + } + Ok(active) } } @@ -314,7 +364,8 @@ mod tests { assert_eq!(pooling.gather_indices, vec![4, 0, 2, 0, 0, 0, 0, 0]); assert_eq!(pooling.valid_mask, vec![1, 0, 1, 0, 0, 0, 0, 0]); assert_eq!(pooling.valid_counts, vec![2, 0]); - assert_eq!(pooling.mean_denominators(), vec![2, 1]); + assert_eq!(pooling.mean_denominators(true), vec![2, 1]); + assert_eq!(pooling.mean_denominators(false), vec![4, 4]); let patches = [100.0, 1.0, 2.0, 3.0, 4.0]; let masked_sums = (0..pooling.groups) @@ -348,6 +399,35 @@ mod tests { )); } + #[test] + fn active_rows_are_rejected_before_native_invoke_when_cardinality_drifts() { + let partial = Molmo2SafePooling::prepare(&[0, -1, -1, -1, 1, -1, -1, -1], 2, 4, 2) + .expect("partially filled groups are valid"); + assert_eq!(partial.active_groups_for_prompt(2, 2).unwrap(), vec![0, 1]); + + let all_invalid = Molmo2SafePooling::prepare(&[0, -1, -1, -1, -1, -1, -1, -1], 2, 4, 2) + .expect("negative sentinels remain valid inputs"); + assert!(matches!( + all_invalid.active_groups_for_prompt(2, 2), + Err(Molmo2InputError::ActiveTokenCount { + active_groups: 1, + grid_groups: 2, + prompt_positions: 2, + all_invalid_groups, + }) if all_invalid_groups == vec![1] + )); + + assert!(matches!( + partial.active_groups_for_prompt(2, 1), + Err(Molmo2InputError::ActiveTokenCount { + active_groups: 2, + grid_groups: 2, + prompt_positions: 1, + all_invalid_groups, + }) if all_invalid_groups.is_empty() + )); + } + #[test] fn image_patch_features_are_added_in_scanned_position_order() { let tokens = [9, 151_938, 8, 151_938]; diff --git a/src/lib/mlxcel-xla/src/molmo2_vision_runtime.rs b/src/lib/mlxcel-xla/src/molmo2_vision_runtime.rs index 2640a7fb3..8468658b9 100644 --- a/src/lib/mlxcel-xla/src/molmo2_vision_runtime.rs +++ b/src/lib/mlxcel-xla/src/molmo2_vision_runtime.rs @@ -54,6 +54,7 @@ pub struct Molmo2VisionInput<'a> { pub pooling_shape: [usize; 2], pub image_grid: [i32; 4], pub image_num_crops: usize, + pub prompt_image_patch_count: usize, } fn hex(bytes: &[u8]) -> String { @@ -421,6 +422,9 @@ impl IreeMolmo2VisionProjector { crops * patches, ) .map_err(|error| error.to_string())?; + let active_groups = safe + .active_groups_for_prompt(grid_groups, input.prompt_image_patch_count) + .map_err(|error| error.to_string())?; let static_patch_values = self.config.static_crops * self.config.patches_per_crop * self.config.patch_dim; let mut padded_patches = vec![0.0f32; static_patch_values]; @@ -456,14 +460,8 @@ impl IreeMolmo2VisionProjector { }], )?; let all_values = decode_output(&output)?; - let valid_groups = safe - .valid_counts - .iter() - .enumerate() - .filter_map(|(index, &count)| (count > 0).then_some(index)) - .collect::>(); - let mut values = Vec::with_capacity(valid_groups.len() * self.config.text_hidden); - for group in valid_groups { + let mut values = Vec::with_capacity(active_groups.len() * self.config.text_hidden); + for group in active_groups { let start = group * self.config.text_hidden; values.extend_from_slice(&all_values[start..start + self.config.text_hidden]); } diff --git a/src/multimodal/molmo2_xla_preprocessor.rs b/src/multimodal/molmo2_xla_preprocessor.rs index 7c8852560..e4b63f7e5 100644 --- a/src/multimodal/molmo2_xla_preprocessor.rs +++ b/src/multimodal/molmo2_xla_preprocessor.rs @@ -192,6 +192,10 @@ impl Molmo2IreeHostPreprocessor { let processed = self.processor.preprocess_image(image); let image_tokens = Self::image_tokens(processed.image_grid)?; let logical_tokens = Self::expand_prompt(token_ids, &image_tokens)?; + let prompt_image_patch_count = logical_tokens + .iter() + .filter(|&&token| token == IMAGE_PATCH_ID) + .count(); validate_sequence_capacity(logical_tokens.len(), self.max_sequence_len)?; let input_ids = mlxcel_core::from_slice_i32( &logical_tokens, @@ -242,6 +246,7 @@ impl Molmo2IreeHostPreprocessor { image_grid: processed.image_grid, image_num_crops: usize::try_from(processed.image_num_crops) .map_err(|_| HostPreprocessorError::ShapeOverflow)?, + prompt_image_patch_count, }) .map_err(HostPreprocessorError::Iree)?; let positions = mlxcel_xla::add_molmo2_projected_features( From 0e4c2a8594724b8dc5adb9ebf977332d7b7ffbf7 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 24 Jul 2026 20:55:19 +0900 Subject: [PATCH 04/20] fix(xla): resolve Molmo2 layers against declared depth Molmo2 declares 27 logical ViT layers but persists only the 25-layer prefix needed by adapter selections `[-3, -9]`. Resolving the negative indices after truncation selected `[22, 16]` instead of the canonical `[24, 18]`. The eager loader and XLA emitter now resolve configured indices against the declared depth first, derive execution depth from the largest selected layer, and preserve configured order across negative and positive indices. Validation: `cargo test -p mlxcel-xla molmo2 --lib`; `cargo test -p mlxcel molmo2 --lib`; `cargo check -p mlxcel-xla`; `cargo check -p mlxcel --lib`; `cargo fmt --all -- --check`; `git diff --check`. Refs #871 --- .../mlxcel-xla/src/emitter/molmo2_config.rs | 10 ++- .../src/emitter/molmo2_config_tests.rs | 44 ++++++++++-- .../mlxcel-xla/src/emitter/molmo2_vision.rs | 7 +- src/loading/vlm_special.rs | 68 +++++++++++++++++-- src/loading/vlm_special_tests.rs | 33 +++++++-- src/vision/encoders/molmo2.rs | 24 +++---- 6 files changed, 148 insertions(+), 38 deletions(-) diff --git a/src/lib/mlxcel-xla/src/emitter/molmo2_config.rs b/src/lib/mlxcel-xla/src/emitter/molmo2_config.rs index 521f324e2..c2dfe3e4a 100644 --- a/src/lib/mlxcel-xla/src/emitter/molmo2_config.rs +++ b/src/lib/mlxcel-xla/src/emitter/molmo2_config.rs @@ -203,7 +203,7 @@ impl Molmo2VisionConfig { )); } - let layers = usize_field(vit, "num_hidden_layers")?.min(25); + let layers = usize_field(vit, "num_hidden_layers")?; let selected_raw = adapter .get("vit_layers") .and_then(Value::as_array) @@ -218,7 +218,13 @@ impl Molmo2VisionConfig { let raw = value .as_i64() .ok_or_else(|| format!("Molmo2 vit_layers[{index}] must be an integer"))?; - let resolved = if raw < 0 { layers as i64 + raw } else { raw }; + let resolved = if raw < 0 { + i64::try_from(layers).map_err(|_| { + "Molmo2 declared ViT layer count does not fit i64".to_string() + })? + raw + } else { + raw + }; usize::try_from(resolved) .ok() .filter(|layer| *layer < layers) diff --git a/src/lib/mlxcel-xla/src/emitter/molmo2_config_tests.rs b/src/lib/mlxcel-xla/src/emitter/molmo2_config_tests.rs index 263387781..0503b1ee7 100644 --- a/src/lib/mlxcel-xla/src/emitter/molmo2_config_tests.rs +++ b/src/lib/mlxcel-xla/src/emitter/molmo2_config_tests.rs @@ -44,14 +44,42 @@ fn pinned_processor() -> String { #[test] fn resolves_pinned_layers_and_static_bucket_identity() { let config = Molmo2VisionConfig::from_json_strs(&pinned_config(), &pinned_processor()).unwrap(); - assert_eq!(config.layers, 25); - assert_eq!(config.selected_layers, vec![22, 16]); - assert_eq!(config.emitted_layers, 23); + assert_eq!(config.layers, 27); + assert_eq!(config.selected_layers, vec![24, 18]); + assert_eq!(config.emitted_layers, 25); + let specs = config.weight_specs(); + assert!( + specs + .iter() + .any(|spec| spec.name.contains("image_vit.transformer.24.")) + ); + assert!( + !specs + .iter() + .any(|spec| spec.name.contains("image_vit.transformer.25.")) + ); assert_eq!(config.static_crops, 9); assert_eq!(config.static_pool_groups, 9); assert!(config.fingerprint().contains("position=exact-default")); - assert!(config.fingerprint().contains("selected=[22, 16]")); + assert!(config.fingerprint().contains("selected=[24, 18]")); assert!(config.fingerprint().contains("pool-mask=true")); + assert!(config.fingerprint().contains("layers=27;emitted=25")); +} + +#[test] +fn configured_layer_order_and_mixed_index_signs_are_mutation_sensitive() { + let mut config: Value = serde_json::from_str(&pinned_config()).unwrap(); + config["adapter_config"]["vit_layers"] = serde_json::json!([3, -3, 1, -9]); + let resolved = + Molmo2VisionConfig::from_json_strs(&config.to_string(), &pinned_processor()).unwrap(); + assert_eq!(resolved.selected_layers, vec![3, 24, 1, 18]); + assert_eq!(resolved.emitted_layers, 25); + + config["adapter_config"]["vit_layers"] = serde_json::json!([-9, -3, 3, 1]); + let reordered = + Molmo2VisionConfig::from_json_strs(&config.to_string(), &pinned_processor()).unwrap(); + assert_eq!(reordered.selected_layers, vec![18, 24, 3, 1]); + assert_ne!(resolved.selected_layers, reordered.selected_layers); } #[test] @@ -72,7 +100,13 @@ fn rejects_position_grid_and_selected_layer_drift() { .contains("exact position") ); config["vit_config"]["image_num_pos"] = Value::from(4); - config["adapter_config"]["vit_layers"] = serde_json::json!([-26]); + config["adapter_config"]["vit_layers"] = serde_json::json!([-28]); + assert!( + Molmo2VisionConfig::from_json_strs(&config.to_string(), &pinned_processor()) + .unwrap_err() + .contains("outside") + ); + config["adapter_config"]["vit_layers"] = serde_json::json!([27]); assert!( Molmo2VisionConfig::from_json_strs(&config.to_string(), &pinned_processor()) .unwrap_err() diff --git a/src/lib/mlxcel-xla/src/emitter/molmo2_vision.rs b/src/lib/mlxcel-xla/src/emitter/molmo2_vision.rs index fc9fce343..1ca57496b 100644 --- a/src/lib/mlxcel-xla/src/emitter/molmo2_vision.rs +++ b/src/lib/mlxcel-xla/src/emitter/molmo2_vision.rs @@ -475,8 +475,9 @@ mod tests { #[test] fn selected_layers_keep_adapter_order_instead_of_encoder_order() { - assert_eq!(selected_slot(&[22, 16], 22), Some(0)); - assert_eq!(selected_slot(&[22, 16], 16), Some(1)); - assert_eq!(selected_slot(&[22, 16], 21), None); + assert_eq!(selected_slot(&[24, 4, 18], 24), Some(0)); + assert_eq!(selected_slot(&[24, 4, 18], 4), Some(1)); + assert_eq!(selected_slot(&[24, 4, 18], 18), Some(2)); + assert_eq!(selected_slot(&[24, 4, 18], 22), None); } } diff --git a/src/loading/vlm_special.rs b/src/loading/vlm_special.rs index 1124a4972..66dde4341 100644 --- a/src/loading/vlm_special.rs +++ b/src/loading/vlm_special.rs @@ -1772,6 +1772,54 @@ pub(super) fn parse_molmo2_vit_layers(adapter_config: &Value) -> Vec { .unwrap_or_else(|| vec![-3, -9]) } +/// Resolve adapter layer indices against the checkpoint's declared ViT depth. +/// +/// Molmo2 declares 27 logical layers but the pinned adapter selects `[-3, -9]` +/// and therefore persists only the first 25 blocks. Resolving after truncating +/// to the persisted execution depth would silently select `[22, 16]` instead +/// of the canonical `[24, 18]`. +pub(super) fn resolve_molmo2_vit_layers( + declared_layers: usize, + configured_layers: &[i32], +) -> Result, String> { + if declared_layers == 0 { + return Err("Molmo2 ViT must declare at least one layer".to_string()); + } + if configured_layers.is_empty() { + return Err("Molmo2 adapter must select at least one ViT layer".to_string()); + } + configured_layers + .iter() + .enumerate() + .map(|(index, &layer)| { + let resolved = if layer < 0 { + i64::try_from(declared_layers) + .map_err(|_| "Molmo2 declared ViT layer count does not fit i64".to_string())? + + i64::from(layer) + } else { + i64::from(layer) + }; + usize::try_from(resolved) + .ok() + .filter(|&resolved| resolved < declared_layers) + .ok_or_else(|| { + format!( + "Molmo2 vit_layers[{index}]={layer} resolves outside [0,{declared_layers})" + ) + }) + }) + .collect() +} + +pub(super) fn molmo2_vit_execution_depth(selected_layers: &[usize]) -> Result { + selected_layers + .iter() + .copied() + .max() + .and_then(|layer| layer.checked_add(1)) + .ok_or_else(|| "Molmo2 adapter must select at least one ViT layer".to_string()) +} + pub(super) fn rewrite_molmo2_weight_key(key: &str) -> String { let mut new_key = key.to_string(); if new_key.starts_with("model.transformer.") { @@ -2009,12 +2057,18 @@ pub(crate) fn load_molmo2_vlm(model_path: &Path) -> Result { let vit_config = vision_config.get("vit_config").unwrap_or(vision_config); let adapter_config = vision_config.get("adapter_config").unwrap_or(vision_config); - let vit_num_layers = cap_molmo2_vit_num_layers( - vit_config - .get("num_hidden_layers") - .and_then(|v| v.as_u64()) - .unwrap_or(25) as usize, - ); + // Resolve negative adapter indices against the declared depth before + // deriving the smaller prefix of blocks that the checkpoint must execute. + let vit_declared_num_layers = vit_config + .get("num_hidden_layers") + .and_then(|v| v.as_u64()) + .unwrap_or(25) as usize; + let vit_layers = resolve_molmo2_vit_layers( + vit_declared_num_layers, + &parse_molmo2_vit_layers(adapter_config), + ) + .map_err(anyhow::Error::msg)?; + let vit_num_layers = molmo2_vit_execution_depth(&vit_layers).map_err(anyhow::Error::msg)?; let vit_hidden_size = vit_config .get("hidden_size") .and_then(|v| v.as_i64()) @@ -2116,7 +2170,7 @@ pub(crate) fn load_molmo2_vlm(model_path: &Path) -> Result { adapter_num_kv_heads, adapter_head_dim, adapter_float32_attention, - &parse_molmo2_vit_layers(adapter_config), + &vit_layers, pooling_attention_mask, ) .map_err(|e| anyhow::anyhow!("Failed to load vision model: {}", e))?; diff --git a/src/loading/vlm_special_tests.rs b/src/loading/vlm_special_tests.rs index 7dbc7e185..9c525da67 100644 --- a/src/loading/vlm_special_tests.rs +++ b/src/loading/vlm_special_tests.rs @@ -16,13 +16,14 @@ use super::{ cap_molmo2_vit_num_layers, dequantize_moondream3_weight, flatten_phi4mm_patch_embedding, inherit_quantization_if_missing, llama4_mm_tokens_per_image, llama4_quantization_params, llama4_token_ids, llama4_vision_prefix, load_minicpmv4_6_vlm, minicpmv4_6_text_weights, - molmo2_max_crops, moondream2_text_config_value, moondream3_text_config_value, - moondream3_vision_config_value, parse_molmo2_vit_layers, phi3_num_crops, - phi4_siglip_text_config_value, phi4mm_text_config_value, phi4mm_vision_config_value, - remap_minicpmo_text_weights, remap_minicpmv4_6_weights, remap_phi4mm_weights, - resolve_moondream2_eos_token_id, rewrite_molmo2_weight_key, rewrite_moondream2_weight_key, - rewrite_moondream3_weight_key, rewrite_phi3_weight_key, rewrite_phi4_siglip_weight_key, - rewrite_phi4mm_weight_key, should_transpose_phi3_patch_embedding, + molmo2_max_crops, molmo2_vit_execution_depth, moondream2_text_config_value, + moondream3_text_config_value, moondream3_vision_config_value, parse_molmo2_vit_layers, + phi3_num_crops, phi4_siglip_text_config_value, phi4mm_text_config_value, + phi4mm_vision_config_value, remap_minicpmo_text_weights, remap_minicpmv4_6_weights, + remap_phi4mm_weights, resolve_molmo2_vit_layers, resolve_moondream2_eos_token_id, + rewrite_molmo2_weight_key, rewrite_moondream2_weight_key, rewrite_moondream3_weight_key, + rewrite_phi3_weight_key, rewrite_phi4_siglip_weight_key, rewrite_phi4mm_weight_key, + should_transpose_phi3_patch_embedding, }; use crate::moondream2_prompt::Moondream2PromptStyle; use mlxcel_core::dtype; @@ -766,6 +767,24 @@ fn molmo2_helpers_clamp_layer_count_and_parse_defaults() { ); } +#[test] +fn molmo2_layer_resolution_uses_declared_depth_and_preserves_config_order() { + let configured = [-3, 4, -9, 0]; + let resolved = resolve_molmo2_vit_layers(27, &configured).unwrap(); + assert_eq!(resolved, vec![24, 4, 18, 0]); + assert_eq!(molmo2_vit_execution_depth(&resolved).unwrap(), 25); + + let mut reordered = configured; + reordered.swap(0, 2); + let reordered = resolve_molmo2_vit_layers(27, &reordered).unwrap(); + assert_eq!(reordered, vec![18, 4, 24, 0]); + assert_ne!(resolved, reordered); + + assert!(resolve_molmo2_vit_layers(27, &[-28]).is_err()); + assert!(resolve_molmo2_vit_layers(27, &[27]).is_err()); + assert!(resolve_molmo2_vit_layers(27, &[]).is_err()); +} + #[test] fn rewrite_molmo2_weight_key_maps_text_vision_and_lm_head_prefixes() { assert_eq!( diff --git a/src/vision/encoders/molmo2.rs b/src/vision/encoders/molmo2.rs index b3daae083..55dc31aaa 100644 --- a/src/vision/encoders/molmo2.rs +++ b/src/vision/encoders/molmo2.rs @@ -574,9 +574,17 @@ impl Molmo2VisionModel { adapter_num_kv_heads: i32, adapter_head_dim: i32, adapter_float32_attention: bool, - vit_layers: &[i32], + vit_layers: &[usize], pooling_attention_mask: bool, ) -> Result { + if vit_layers.is_empty() { + return Err("Molmo2 adapter must select at least one ViT layer".to_string()); + } + if let Some(&layer) = vit_layers.iter().find(|&&layer| layer >= vit_num_layers) { + return Err(format!( + "Molmo2 selected ViT layer {layer} is outside execution depth {vit_num_layers}" + )); + } let image_vit = Molmo2VisionTransformer::from_weights( weights, &format!("{}.image_vit", prefix), @@ -607,23 +615,11 @@ impl Molmo2VisionModel { let image_projector = ImageProjectorMLP::from_weights(weights, &format!("{}.image_projector", prefix))?; - // Convert negative layer indices to positive - let resolved_layers: Vec = vit_layers - .iter() - .map(|&layer| { - if layer < 0 { - (layer + vit_num_layers as i32) as usize - } else { - layer as usize - } - }) - .collect(); - Ok(Self { image_vit, image_pooling_2d, image_projector, - vit_layers: resolved_layers, + vit_layers: vit_layers.to_vec(), pooling_attention_mask, }) } From 0457293f4fffafb1b20819a0376ef348e1e1fb87 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 24 Jul 2026 21:11:41 +0900 Subject: [PATCH 05/20] test(xla): add Molmo2 vision parity gate --- Cargo.toml | 4 + src/lib.rs | 4 + src/loading/mod.rs | 4 + src/loading/vlm.rs | 4 + src/loading/vlm_special.rs | 211 +++++++++++++++++---- tests/molmo2_xla_vision_parity.rs | 301 ++++++++++++++++++++++++++++++ 6 files changed, 492 insertions(+), 36 deletions(-) create mode 100644 tests/molmo2_xla_vision_parity.rs diff --git a/Cargo.toml b/Cargo.toml index a7ecfe4b3..68876075a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -324,6 +324,10 @@ required-features = ["test-utils"] name = "xla_prepared_prefill" required-features = ["xla-iree"] +[[test]] +name = "molmo2_xla_vision_parity" +required-features = ["xla-iree"] + # Reference-equivalence + throughput harness for the OpenXLA/IREE continuous # batching engine (#449 M3 Stage 2b). Needs real IREE execution, so it builds # only under `xla-iree` (which carries the runtime link recipe); default/CI diff --git a/src/lib.rs b/src/lib.rs index 938261357..8683d3a7e 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -121,6 +121,10 @@ 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-diagnostics")] +pub use loading::{ + Molmo2XlaVisionReference, Molmo2XlaVisionReferenceProjection, load_molmo2_xla_vision_reference, +}; 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 7a7223cba..fc0e98d65 100644 --- a/src/loading/mod.rs +++ b/src/loading/mod.rs @@ -62,6 +62,10 @@ 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-diagnostics")] +pub use self::vlm::{ + Molmo2XlaVisionReference, Molmo2XlaVisionReferenceProjection, load_molmo2_xla_vision_reference, +}; #[cfg(feature = "xla-iree")] pub(crate) use self::vlm::{ Phi4MMXlaVisionComponents, load_molmo2_xla_text_embeddings, load_phi4mm_xla_media_components, diff --git a/src/loading/vlm.rs b/src/loading/vlm.rs index c3a0357c0..d6d816277 100644 --- a/src/loading/vlm.rs +++ b/src/loading/vlm.rs @@ -142,6 +142,10 @@ pub(crate) use qwen::{ pub use qwen::load_qwen3_omni_speech; pub(crate) use siglip::{load_aya_vision_vlm, load_paligemma_vlm}; pub(crate) use smolvlm::load_smolvlm_vlm; +#[cfg(feature = "xla-diagnostics")] +pub use special::{ + Molmo2XlaVisionReference, Molmo2XlaVisionReferenceProjection, load_molmo2_xla_vision_reference, +}; #[cfg(feature = "xla-iree")] pub(crate) use special::{ Phi4MMXlaVisionComponents, load_molmo2_xla_text_embeddings, load_phi4mm_xla_media_components, diff --git a/src/loading/vlm_special.rs b/src/loading/vlm_special.rs index 66dde4341..b4c7560e8 100644 --- a/src/loading/vlm_special.rs +++ b/src/loading/vlm_special.rs @@ -2038,22 +2038,14 @@ fn read_clip_triple(config: Option<&Value>, key: &str) -> Option<[f32; 3]> { }) } -pub(crate) fn load_molmo2_vlm(model_path: &Path) -> Result { +fn build_molmo2_vision_model( + weights: &WeightMap, + full_config: &Value, + text_hidden_size: usize, +) -> Result { use vision::encoders::molmo2::Molmo2VisionModel; - use vision::processors::molmo2::Molmo2Processor; - - let (_config_str, full_config) = read_sanitized_vlm_config(model_path)?; - - let mut text_config_value = full_config - .get("text_config") - .cloned() - .unwrap_or_else(|| full_config.clone()); - inherit_quantization_if_missing(&mut text_config_value, &full_config)?; - let text_config: models::molmo2::Molmo2TextConfig = - serde_json::from_value(text_config_value) - .map_err(|e| anyhow::anyhow!("Failed to parse text config: {}", e))?; - let vision_config = full_config.get("vision_config").unwrap_or(&full_config); + let vision_config = full_config.get("vision_config").unwrap_or(full_config); let vit_config = vision_config.get("vit_config").unwrap_or(vision_config); let adapter_config = vision_config.get("adapter_config").unwrap_or(vision_config); @@ -2113,7 +2105,7 @@ pub(crate) fn load_molmo2_vlm(model_path: &Path) -> Result { let adapter_text_hidden_size = adapter_config .get("text_hidden_size") .and_then(|v| v.as_i64()) - .unwrap_or(text_config.hidden_size as i64) as i32; + .unwrap_or(text_hidden_size as i64) as i32; let adapter_num_heads = adapter_config .get("num_attention_heads") .and_then(|v| v.as_i64()) @@ -2135,24 +2127,8 @@ pub(crate) fn load_molmo2_vlm(model_path: &Path) -> Result { .and_then(|v| v.as_bool()) .unwrap_or(true); - let image_patch_id = full_config - .get("image_patch_id") - .and_then(|v| v.as_i64()) - .unwrap_or(151938) as i32; - let image_end_token_id = full_config - .get("image_end_token_id") - .and_then(|v| v.as_i64()) - .unwrap_or(151937) as i32; - - let mut weights = remap_molmo2_weights(load_vlm_weights_common(model_path, None)?); - models::sanitize_tied_embeddings(&mut weights, &full_config); - - let text_model = - models::Molmo2Model::from_weights(&weights, &text_config, "language_model.model") - .map_err(|e| anyhow::anyhow!("Failed to load text model: {}", e))?; - - let vision_tower = Molmo2VisionModel::from_weights( - &weights, + Molmo2VisionModel::from_weights( + weights, "vision_tower", vit_num_layers, vit_hidden_size, @@ -2173,16 +2149,179 @@ pub(crate) fn load_molmo2_vlm(model_path: &Path) -> Result { &vit_layers, pooling_attention_mask, ) - .map_err(|e| anyhow::anyhow!("Failed to load vision model: {}", e))?; + .map_err(|e| anyhow::anyhow!("Failed to load vision model: {}", e)) +} +fn build_molmo2_processor(model_path: &Path) -> vision::processors::molmo2::Molmo2Processor { let preprocessor_config = read_optional_model_json(model_path, "preprocessor_config.json"); - let processor = Molmo2Processor::new( + vision::processors::molmo2::Molmo2Processor::new( molmo2_max_crops(preprocessor_config.as_ref()), None, None, None, None, - ); + ) +} + +/// Vision-only MLX reference used by the ignored Molmo2 XLA parity gate. +/// +/// This diagnostics surface filters the checkpoint before loading and never +/// constructs the text decoder, LM head, or text embedding tables. +#[cfg(any(test, feature = "xla-diagnostics"))] +#[cfg_attr(test, allow(dead_code))] +pub struct Molmo2XlaVisionReference { + vision_tower: vision::encoders::molmo2::Molmo2VisionModel, + processor: vision::processors::molmo2::Molmo2Processor, + image_patch_id: i32, + text_hidden_size: usize, +} + +/// Eager MLX projection and the exact processor payload that produced it. +#[cfg(any(test, feature = "xla-diagnostics"))] +#[cfg_attr(test, allow(dead_code))] +pub struct Molmo2XlaVisionReferenceProjection { + pub processed: vision::processors::molmo2::Molmo2ProcessorOutput, + pub values: Vec, + pub shape: [usize; 2], + pub active_groups: Vec, +} + +#[cfg(any(test, feature = "xla-diagnostics"))] +#[cfg_attr(test, allow(dead_code))] +impl Molmo2XlaVisionReference { + pub fn image_patch_id(&self) -> i32 { + self.image_patch_id + } + + pub fn text_hidden_size(&self) -> usize { + self.text_hidden_size + } + + pub fn project( + &self, + image: &image::DynamicImage, + ) -> Result { + let processed = self.processor.preprocess_image(image); + let [crops, patches, patch_dim] = processed.pixel_values_shape; + let [groups, group_size] = processed.image_token_pooling_shape; + let active_groups = processed + .image_token_pooling + .chunks_exact(group_size as usize) + .enumerate() + .filter_map(|(group, values)| values.iter().any(|&value| value >= 0).then_some(group)) + .collect::>(); + let images = + mlxcel_core::from_slice_f32(&processed.pixel_values, &[1, crops, patches, patch_dim]); + let pooling = + mlxcel_core::from_slice_i32(&processed.image_token_pooling, &[1, groups, group_size]); + let projected = self.vision_tower.forward(&images, &pooling); + let projected = mlxcel_core::astype(&projected, mlxcel_core::dtype::FLOAT32); + let raw = mlxcel_core::try_array_to_raw_bytes(&projected) + .map_err(|error| anyhow::anyhow!("Failed to export Molmo2 MLX projection: {error}"))?; + if !raw.len().is_multiple_of(4) { + return Err(anyhow::anyhow!( + "Molmo2 MLX projection byte count is not f32-aligned" + )); + } + let dimensions = mlxcel_core::array_shape(&projected); + if dimensions.len() != 2 { + return Err(anyhow::anyhow!( + "Molmo2 MLX projection shape must be rank 2, got {dimensions:?}" + )); + } + let shape = [ + usize::try_from(dimensions[0]) + .map_err(|_| anyhow::anyhow!("Molmo2 MLX projection has a negative row count"))?, + usize::try_from(dimensions[1]).map_err(|_| { + anyhow::anyhow!("Molmo2 MLX projection has a negative hidden dimension") + })?, + ]; + let values = raw + .chunks_exact(4) + .map(|chunk| f32::from_ne_bytes([chunk[0], chunk[1], chunk[2], chunk[3]])) + .collect::>(); + if shape != [active_groups.len(), self.text_hidden_size] + || values.len() != shape[0] * shape[1] + { + return Err(anyhow::anyhow!( + "Molmo2 MLX projection shape {shape:?} disagrees with {} active groups and hidden size {}", + active_groups.len(), + self.text_hidden_size + )); + } + Ok(Molmo2XlaVisionReferenceProjection { + processed, + values, + shape, + active_groups, + }) + } +} + +/// Load only Molmo2's eager vision encoder/projector for diagnostics. +#[cfg(any(test, feature = "xla-diagnostics"))] +#[cfg_attr(test, allow(dead_code))] +pub fn load_molmo2_xla_vision_reference(model_path: &Path) -> Result { + let (_config_str, full_config) = read_sanitized_vlm_config(model_path)?; + if full_config.get("model_type").and_then(Value::as_str) != Some("molmo2") { + return Err(anyhow::anyhow!( + "{} is not a Molmo2 checkpoint", + model_path.display() + )); + } + let text_hidden_size = full_config + .get("text_config") + .and_then(|text| text.get("hidden_size")) + .and_then(Value::as_u64) + .ok_or_else(|| anyhow::anyhow!("Molmo2 text_config.hidden_size is required"))? + as usize; + let raw_weights = super::load_vlm_weights_common_filtered_canonical(model_path, |name| { + name.starts_with("vision_tower.") || name.starts_with("model.vision_backbone.") + })?; + let weights = remap_molmo2_weights(raw_weights); + let vision_tower = build_molmo2_vision_model(&weights, &full_config, text_hidden_size)?; + let image_patch_id = full_config + .get("image_patch_id") + .and_then(Value::as_i64) + .unwrap_or(151938) as i32; + Ok(Molmo2XlaVisionReference { + vision_tower, + processor: build_molmo2_processor(model_path), + image_patch_id, + text_hidden_size, + }) +} + +pub(crate) fn load_molmo2_vlm(model_path: &Path) -> Result { + let (_config_str, full_config) = read_sanitized_vlm_config(model_path)?; + + let mut text_config_value = full_config + .get("text_config") + .cloned() + .unwrap_or_else(|| full_config.clone()); + inherit_quantization_if_missing(&mut text_config_value, &full_config)?; + let text_config: models::molmo2::Molmo2TextConfig = + serde_json::from_value(text_config_value) + .map_err(|e| anyhow::anyhow!("Failed to parse text config: {}", e))?; + + let image_patch_id = full_config + .get("image_patch_id") + .and_then(|v| v.as_i64()) + .unwrap_or(151938) as i32; + let image_end_token_id = full_config + .get("image_end_token_id") + .and_then(|v| v.as_i64()) + .unwrap_or(151937) as i32; + + let mut weights = remap_molmo2_weights(load_vlm_weights_common(model_path, None)?); + models::sanitize_tied_embeddings(&mut weights, &full_config); + + let text_model = + models::Molmo2Model::from_weights(&weights, &text_config, "language_model.model") + .map_err(|e| anyhow::anyhow!("Failed to load text model: {}", e))?; + + let vision_tower = build_molmo2_vision_model(&weights, &full_config, text_config.hidden_size)?; + let processor = build_molmo2_processor(model_path); let vlm = vision::Molmo2VLModel { text_model, diff --git a/tests/molmo2_xla_vision_parity.rs b/tests/molmo2_xla_vision_parity.rs new file mode 100644 index 000000000..182a58437 --- /dev/null +++ b/tests/molmo2_xla_vision_parity.rs @@ -0,0 +1,301 @@ +// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Ignored real-checkpoint Molmo2 vision parity gate. +//! +//! This intentionally compares only the filtered eager MLX vision path with +//! the IREE vision projector. It never loads either text decoder. + +#[cfg(feature = "xla-diagnostics")] +use std::path::PathBuf; + +use anyhow::{Result, anyhow}; +#[cfg(feature = "xla-diagnostics")] +use mlxcel::{initialize_runtime, load_molmo2_xla_vision_reference}; +use mlxcel_xla::add_molmo2_projected_features; +#[cfg(feature = "xla-diagnostics")] +use mlxcel_xla::{IreeMolmo2VisionProjector, Molmo2VisionInput}; + +#[derive(Debug, Clone, Copy)] +struct Comparison { + max_abs: f32, + max_index: usize, + rms: f32, +} + +fn compare(actual: &[f32], expected: &[f32]) -> Result { + if actual.len() != expected.len() { + return Err(anyhow!( + "comparison length mismatch: actual={}, expected={}", + actual.len(), + expected.len() + )); + } + if actual.is_empty() { + return Err(anyhow!("cannot compare empty tensors")); + } + let mut max_abs = 0.0f32; + let mut max_index = 0usize; + let mut squared = 0.0f64; + for (index, (&actual, &expected)) in actual.iter().zip(expected).enumerate() { + if !actual.is_finite() || !expected.is_finite() { + return Err(anyhow!( + "non-finite comparison value at {index}: actual={actual}, expected={expected}" + )); + } + let difference = (actual - expected).abs(); + if difference > max_abs { + max_abs = difference; + max_index = index; + } + squared += f64::from(difference) * f64::from(difference); + } + Ok(Comparison { + max_abs, + max_index, + rms: (squared / actual.len() as f64).sqrt() as f32, + }) +} + +#[cfg(feature = "xla-diagnostics")] +fn assert_within( + label: &str, + actual: &[f32], + expected: &[f32], + max_abs_limit: f32, + rms_limit: f32, +) -> Result<()> { + let comparison = compare(actual, expected)?; + if comparison.max_abs > max_abs_limit || comparison.rms > rms_limit { + return Err(anyhow!( + "{label} parity failed: max_abs={} at {}, rms={}, limits=({}, {})", + comparison.max_abs, + comparison.max_index, + comparison.rms, + max_abs_limit, + rms_limit + )); + } + eprintln!( + "{label}: max_abs={} at {}, rms={}", + comparison.max_abs, comparison.max_index, comparison.rms + ); + Ok(()) +} + +fn active_groups(pooling: &[i32], groups: usize, group_size: usize) -> Result> { + if group_size == 0 || pooling.len() != groups * group_size { + return Err(anyhow!( + "invalid pooling shape: values={}, groups={groups}, group_size={group_size}", + pooling.len() + )); + } + Ok(pooling + .chunks_exact(group_size) + .enumerate() + .filter_map(|(group, values)| values.iter().any(|&value| value >= 0).then_some(group)) + .collect()) +} + +fn independent_scatter_add( + token_ids: &[i32], + image_patch_id: i32, + base: &[f32], + hidden_size: usize, + projected: &[f32], +) -> Result> { + if hidden_size == 0 || base.len() != token_ids.len() * hidden_size { + return Err(anyhow!("invalid base embedding shape")); + } + let positions = token_ids + .iter() + .enumerate() + .filter_map(|(index, &token)| (token == image_patch_id).then_some(index)) + .collect::>(); + if projected.len() != positions.len() * hidden_size { + return Err(anyhow!( + "projected feature count does not match image positions" + )); + } + let mut merged = base.to_vec(); + for (row, position) in positions.into_iter().enumerate() { + for hidden in 0..hidden_size { + merged[position * hidden_size + hidden] += projected[row * hidden_size + hidden]; + } + } + Ok(merged) +} + +#[cfg(feature = "xla-diagnostics")] +fn tolerance(name: &str, default: f32) -> Result { + match std::env::var(name) { + Ok(value) => value + .parse::() + .map_err(|error| anyhow!("invalid {name}={value}: {error}")), + Err(std::env::VarError::NotPresent) => Ok(default), + Err(error) => Err(anyhow!("cannot read {name}: {error}")), + } +} + +#[test] +fn synthetic_active_rows_skip_all_negative_groups() { + let pooling = [0, -1, -1, -1, -1, -1, -1, -1, 4, 5, -1, -1]; + assert_eq!(active_groups(&pooling, 3, 4).unwrap(), vec![0, 2]); +} + +#[test] +fn synthetic_independent_scatter_matches_production_helper() { + let tokens = [7, 151_938, 8, 151_938, 9]; + let base = (0..15).map(|index| index as f32 * 0.25).collect::>(); + let projected = [0.5, -1.0, 1.5, 2.0, 3.0, -0.25]; + let expected = independent_scatter_add(&tokens, 151_938, &base, 3, &projected).unwrap(); + let mut actual = base; + let positions = + add_molmo2_projected_features(&tokens, 151_938, &mut actual, 3, &projected).unwrap(); + assert_eq!(positions, vec![1, 3]); + assert_eq!(actual, expected); +} + +#[test] +fn synthetic_comparison_reports_max_and_rms() { + let comparison = compare(&[1.0, 2.5, 3.0], &[1.0, 2.0, 3.0]).unwrap(); + assert_eq!(comparison.max_abs, 0.5); + assert_eq!(comparison.max_index, 1); + assert!((comparison.rms - (0.25f32 / 3.0).sqrt()).abs() < 1e-7); +} + +#[test] +#[cfg(feature = "xla-diagnostics")] +#[ignore = "requires a Molmo2 checkpoint plus configured MLX and IREE runtimes"] +fn real_checkpoint_mlx_iree_vision_and_scatter_parity() -> Result<()> { + let model = PathBuf::from( + std::env::var("MLXCEL_MOLMO2_MODEL") + .map_err(|_| anyhow!("MLXCEL_MOLMO2_MODEL is required"))?, + ); + let image_path = std::env::var("MLXCEL_MOLMO2_IMAGE") + .map(PathBuf::from) + .unwrap_or_else(|_| { + PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures/test_image.png") + }); + let device = std::env::var("MLXCEL_XLA_DEVICE").unwrap_or_else(|_| "local-task".to_string()); + let max_abs_limit = tolerance("MLXCEL_MOLMO2_MAX_ABS", 0.05)?; + let rms_limit = tolerance("MLXCEL_MOLMO2_RMS", 0.01)?; + + let _runtime = initialize_runtime(); + let reference = load_molmo2_xla_vision_reference(&model)?; + let image = image::open(&image_path) + .map_err(|error| anyhow!("open {}: {error}", image_path.display()))?; + let eager = reference.project(&image)?; + let processed = &eager.processed; + let pooling_shape = processed + .image_token_pooling_shape + .map(|value| value as usize); + let independently_active = active_groups( + &processed.image_token_pooling, + pooling_shape[0], + pooling_shape[1], + )?; + if eager.active_groups != independently_active { + return Err(anyhow!( + "eager active rows {:?} disagree with processor rows {:?}", + eager.active_groups, + independently_active + )); + } + + let mut projector = + IreeMolmo2VisionProjector::load(&model, &device).map_err(anyhow::Error::msg)?; + if projector.image_patch_id() != reference.image_patch_id() + || projector.text_hidden_size() != reference.text_hidden_size() + { + return Err(anyhow!("MLX and IREE Molmo2 metadata disagree")); + } + let iree = projector + .project(Molmo2VisionInput { + patches: &processed.pixel_values, + patches_shape: processed.pixel_values_shape.map(|value| value as usize), + image_token_pooling: &processed.image_token_pooling, + pooling_shape, + image_grid: processed.image_grid, + image_num_crops: processed.image_num_crops as usize, + prompt_image_patch_count: independently_active.len(), + }) + .map_err(anyhow::Error::msg)?; + let iree_active = iree + .valid_pooling_counts + .iter() + .enumerate() + .filter_map(|(group, &count)| (count > 0).then_some(group)) + .collect::>(); + if iree_active != independently_active || iree.shape != eager.shape { + return Err(anyhow!( + "active-row mismatch: processor={independently_active:?}, IREE={iree_active:?}, MLX shape={:?}, IREE shape={:?}", + eager.shape, + iree.shape + )); + } + assert_within( + "vision projection", + &iree.values, + &eager.values, + max_abs_limit, + rms_limit, + )?; + + let hidden = reference.text_hidden_size(); + let mut tokens = Vec::with_capacity(independently_active.len() * 2 + 1); + tokens.push(7); + for row in 0..independently_active.len() { + tokens.push(reference.image_patch_id()); + tokens.push(8 + (row % 17) as i32); + } + let base = (0..tokens.len() * hidden) + .map(|index| ((index % 31) as f32 - 15.0) * 0.001) + .collect::>(); + let eager_merged = independent_scatter_add( + &tokens, + reference.image_patch_id(), + &base, + hidden, + &eager.values, + )?; + let independently_iree_merged = independent_scatter_add( + &tokens, + reference.image_patch_id(), + &base, + hidden, + &iree.values, + )?; + let mut production_iree_merged = base; + add_molmo2_projected_features( + &tokens, + reference.image_patch_id(), + &mut production_iree_merged, + hidden, + &iree.values, + ) + .map_err(|error| anyhow!(error.to_string()))?; + if production_iree_merged != independently_iree_merged { + return Err(anyhow!( + "production scatter-add disagrees with the independent implementation" + )); + } + assert_within( + "scatter-added embeddings", + &production_iree_merged, + &eager_merged, + max_abs_limit, + rms_limit, + ) +} From 4a8724d4b0179cc4183a3f6f075a292924174e46 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Sun, 26 Jul 2026 22:08:02 +0900 Subject: [PATCH 06/20] test(xla): expose Molmo2 CPU parity diagnostics --- Cargo.toml | 4 +++ src/lib.rs | 2 +- src/lib/mlxcel-core/src/weights.rs | 48 ++++++++++++++++++++++--- src/loading/mod.rs | 2 +- src/loading/vlm.rs | 2 +- src/loading/vlm_special.rs | 18 ++++++---- src/models/mod.rs | 3 +- src/models/sanitize.rs | 9 ++++- src/models/sanitize_tests.rs | 56 +++++++++++++++++++++++++++++- tests/molmo2_xla_vision_parity.rs | 12 +++---- 10 files changed, 133 insertions(+), 23 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 68876075a..a3b338200 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -104,6 +104,10 @@ xla-diagnostics = ["cuda", "xla-iree", "mlxcel-xla/diagnostics"] # CPU-capable bounded MLX/IREE operator-oracle harness. CUDA-specific production # probes may add `cuda`, but the shared report/comparison layer does not require it. xla-micro-oracle = ["xla-iree", "mlxcel-xla/micro-oracle"] +# CPU-only companion for bounded host-reference checks on machines where the +# CUDA driver is unavailable. This deliberately does not replace the qualified +# `xla-diagnostics` production-target gate. +xla-diagnostics-cpu = ["xla-iree"] # Axis A "weight-load surgery" framework. # # On by default so production `mlxcel` and `mlxcel-server` binaries expose diff --git a/src/lib.rs b/src/lib.rs index 8683d3a7e..0646f491a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -121,7 +121,7 @@ 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-diagnostics")] +#[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] pub use loading::{ Molmo2XlaVisionReference, Molmo2XlaVisionReferenceProjection, load_molmo2_xla_vision_reference, }; diff --git a/src/lib/mlxcel-core/src/weights.rs b/src/lib/mlxcel-core/src/weights.rs index c0143c1e0..6a4eba8da 100644 --- a/src/lib/mlxcel-core/src/weights.rs +++ b/src/lib/mlxcel-core/src/weights.rs @@ -182,7 +182,7 @@ fn safetensors_dtype_itemsize(dtype: &str) -> Option { /// /// Returns `None` when the file does not exist or cannot be parsed; callers /// should fall back to the analytical estimate in that case. -fn read_safetensors_header_bytes(path: &Path) -> Option { +fn read_safetensors_header(path: &Path) -> Option> { use std::io::Read; let mut f = std::fs::File::open(path).ok()?; @@ -201,10 +201,13 @@ fn read_safetensors_header_bytes(path: &Path) -> Option { let mut header_bytes = vec![0u8; header_len as usize]; f.read_exact(&mut header_bytes).ok()?; let header_json = serde_json::from_slice::(&header_bytes).ok()?; + header_json.as_object().cloned() +} - let obj = header_json.as_object()?; +fn read_safetensors_header_bytes(path: &Path) -> Option { + let obj = read_safetensors_header(path)?; let mut total: u64 = 0; - for (key, meta) in obj { + for (key, meta) in &obj { // The special "__metadata__" key is not a tensor entry. if key == "__metadata__" { continue; @@ -551,6 +554,25 @@ where F: FnMut(&str) -> bool, { let path = path.as_ref(); + // A stale index can force the directory loader to inspect every local shard. + // Read the lightweight safetensors header first so a filtered sub-stack does + // not enter MLX's native loader for a shard that cannot contain any retained + // tensor. Besides avoiding unnecessary mmap state, this is important for + // repackaged quantized shards whose unrelated tensor dtypes/layouts may not be + // loadable by the current MLX build. + let retained_names = read_safetensors_header(path).map(|header| { + header + .into_iter() + .map(|(name, _)| name) + .filter(|name| name != "__metadata__" && keep(name)) + .collect::>() + }); + if retained_names + .as_ref() + .is_some_and(|names| names.is_empty()) + { + return Ok(HashMap::new()); + } let path_str = path .to_str() .ok_or_else(|| format!("Non-UTF8 path: {}", path.display()))?; @@ -560,7 +582,10 @@ where let mut weights = HashMap::with_capacity(len); for i in 0..len { let name = ffi::loaded_weights_name(&loaded, i); - if !keep(&name) { + let retain = retained_names + .as_ref() + .map_or_else(|| keep(&name), |names| names.contains(&name)); + if !retain { continue; } let array = ffi::loaded_weights_take(loaded.pin_mut(), i); @@ -901,6 +926,21 @@ mod tests { f.write_all(header_bytes).unwrap(); } + #[test] + fn filtered_loader_skips_nonmatching_shard_before_native_load() { + let dir = tempfile::tempdir().unwrap(); + let file = dir.path().join("model-00001-of-00002.safetensors"); + write_safetensors_stub(&file, "U32", &[4, 4]); + let mut inspected = Vec::new(); + let weights = load_safetensors_filtered(&file, |name| { + inspected.push(name.to_string()); + name.starts_with("vision_tower.") + }) + .expect("a shard with no retained header names is skipped before native loading"); + assert!(weights.is_empty()); + assert_eq!(inspected, ["test_tensor"]); + } + #[test] fn test_parse_shard_index_with_total_size_present() { let dir = tempfile::tempdir().unwrap(); diff --git a/src/loading/mod.rs b/src/loading/mod.rs index fc0e98d65..22abd7b5e 100644 --- a/src/loading/mod.rs +++ b/src/loading/mod.rs @@ -62,7 +62,7 @@ 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-diagnostics")] +#[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] pub use self::vlm::{ Molmo2XlaVisionReference, Molmo2XlaVisionReferenceProjection, load_molmo2_xla_vision_reference, }; diff --git a/src/loading/vlm.rs b/src/loading/vlm.rs index d6d816277..5ca951f83 100644 --- a/src/loading/vlm.rs +++ b/src/loading/vlm.rs @@ -142,7 +142,7 @@ pub(crate) use qwen::{ pub use qwen::load_qwen3_omni_speech; pub(crate) use siglip::{load_aya_vision_vlm, load_paligemma_vlm}; pub(crate) use smolvlm::load_smolvlm_vlm; -#[cfg(feature = "xla-diagnostics")] +#[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] pub use special::{ Molmo2XlaVisionReference, Molmo2XlaVisionReferenceProjection, load_molmo2_xla_vision_reference, }; diff --git a/src/loading/vlm_special.rs b/src/loading/vlm_special.rs index b4c7560e8..1da5d5103 100644 --- a/src/loading/vlm_special.rs +++ b/src/loading/vlm_special.rs @@ -2167,7 +2167,7 @@ fn build_molmo2_processor(model_path: &Path) -> vision::processors::molmo2::Molm /// /// This diagnostics surface filters the checkpoint before loading and never /// constructs the text decoder, LM head, or text embedding tables. -#[cfg(any(test, feature = "xla-diagnostics"))] +#[cfg(any(test, feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] #[cfg_attr(test, allow(dead_code))] pub struct Molmo2XlaVisionReference { vision_tower: vision::encoders::molmo2::Molmo2VisionModel, @@ -2177,7 +2177,7 @@ pub struct Molmo2XlaVisionReference { } /// Eager MLX projection and the exact processor payload that produced it. -#[cfg(any(test, feature = "xla-diagnostics"))] +#[cfg(any(test, feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] #[cfg_attr(test, allow(dead_code))] pub struct Molmo2XlaVisionReferenceProjection { pub processed: vision::processors::molmo2::Molmo2ProcessorOutput, @@ -2186,7 +2186,7 @@ pub struct Molmo2XlaVisionReferenceProjection { pub active_groups: Vec, } -#[cfg(any(test, feature = "xla-diagnostics"))] +#[cfg(any(test, feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] #[cfg_attr(test, allow(dead_code))] impl Molmo2XlaVisionReference { pub fn image_patch_id(&self) -> i32 { @@ -2259,7 +2259,7 @@ impl Molmo2XlaVisionReference { } /// Load only Molmo2's eager vision encoder/projector for diagnostics. -#[cfg(any(test, feature = "xla-diagnostics"))] +#[cfg(any(test, feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] #[cfg_attr(test, allow(dead_code))] pub fn load_molmo2_xla_vision_reference(model_path: &Path) -> Result { let (_config_str, full_config) = read_sanitized_vlm_config(model_path)?; @@ -2275,9 +2275,13 @@ pub fn load_molmo2_xla_vision_reference(model_path: &Path) -> Result( +/// Load a selected subset of checkpoint tensors through the safetensors parser. +/// +/// Unlike MLX's native whole-shard loader, this path parses the shard table +/// first and materializes only names accepted by `keep`. It is therefore the +/// canonical boundary for host-only sub-stacks paired with an independent +/// backend: unrelated quantized tensors in the same shard cannot make the host +/// reference fail before filtering is applied. +pub(crate) fn load_weights_from_dir_with_filter( model_dir: P, keep: F, prefer_native_full_shard_load: bool, diff --git a/src/models/sanitize_tests.rs b/src/models/sanitize_tests.rs index 01d557c08..12cc6d68e 100644 --- a/src/models/sanitize_tests.rs +++ b/src/models/sanitize_tests.rs @@ -12,7 +12,10 @@ // See the License for the specific language governing permissions and // limitations under the License. -use super::sanitize::{load_text_weights, sanitize_config_json, sanitize_tied_embeddings}; +use super::sanitize::{ + load_text_weights, load_weights_from_dir_with_filter, sanitize_config_json, + sanitize_tied_embeddings, +}; use crate::test_support::env_lock::env_lock; use mlxcel_core::weights::{WeightMap, WeightTransform}; use mlxcel_core::{self, dtype}; @@ -89,6 +92,57 @@ fn write_safetensors(path: &Path, tensors: &[(&str, OwnedTensor)]) { safetensors::serialize_to_file(&views, None, path).unwrap(); } +#[test] +fn selective_loader_materializes_only_retained_names_across_mixed_shards() { + let dir = temp_model_dir("selective_mixed_shards"); + std::fs::create_dir_all(&dir).unwrap(); + write_safetensors( + &dir.join("model-00001-of-00002.safetensors"), + &[( + "language_model.quantized.weight", + OwnedTensor { + dtype: SafeTensorDtype::U32, + shape: vec![1], + data: 7_u32.to_le_bytes().to_vec(), + }, + )], + ); + write_safetensors( + &dir.join("model-00002-of-00002.safetensors"), + &[ + ( + "vision_tower.projector.weight", + OwnedTensor { + dtype: SafeTensorDtype::F32, + shape: vec![2], + data: [1.0_f32, 2.0_f32] + .into_iter() + .flat_map(f32::to_le_bytes) + .collect(), + }, + ), + ( + "language_model.quantized.scales", + OwnedTensor { + dtype: SafeTensorDtype::U32, + shape: vec![1], + data: 11_u32.to_le_bytes().to_vec(), + }, + ), + ], + ); + + let weights = + load_weights_from_dir_with_filter(&dir, |name| name.starts_with("vision_tower."), false) + .unwrap(); + assert_eq!( + weights.keys().map(String::as_str).collect::>(), + ["vision_tower.projector.weight"] + ); + drop(weights); + std::fs::remove_dir_all(dir).unwrap(); +} + const NVFP4_REPACK_ENV_KEYS: &[&str] = &["MLXCEL_NVFP4_DENSE_REPACK", "MLXCEL_NVFP4_NATIVE_REPACK"]; struct EnvRestore { diff --git a/tests/molmo2_xla_vision_parity.rs b/tests/molmo2_xla_vision_parity.rs index 182a58437..1dde85f83 100644 --- a/tests/molmo2_xla_vision_parity.rs +++ b/tests/molmo2_xla_vision_parity.rs @@ -17,14 +17,14 @@ //! This intentionally compares only the filtered eager MLX vision path with //! the IREE vision projector. It never loads either text decoder. -#[cfg(feature = "xla-diagnostics")] +#[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] use std::path::PathBuf; use anyhow::{Result, anyhow}; -#[cfg(feature = "xla-diagnostics")] +#[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] use mlxcel::{initialize_runtime, load_molmo2_xla_vision_reference}; use mlxcel_xla::add_molmo2_projected_features; -#[cfg(feature = "xla-diagnostics")] +#[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] use mlxcel_xla::{IreeMolmo2VisionProjector, Molmo2VisionInput}; #[derive(Debug, Clone, Copy)] @@ -68,7 +68,7 @@ fn compare(actual: &[f32], expected: &[f32]) -> Result { }) } -#[cfg(feature = "xla-diagnostics")] +#[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] fn assert_within( label: &str, actual: &[f32], @@ -137,7 +137,7 @@ fn independent_scatter_add( Ok(merged) } -#[cfg(feature = "xla-diagnostics")] +#[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] fn tolerance(name: &str, default: f32) -> Result { match std::env::var(name) { Ok(value) => value @@ -176,7 +176,7 @@ fn synthetic_comparison_reports_max_and_rms() { } #[test] -#[cfg(feature = "xla-diagnostics")] +#[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] #[ignore = "requires a Molmo2 checkpoint plus configured MLX and IREE runtimes"] fn real_checkpoint_mlx_iree_vision_and_scatter_parity() -> Result<()> { let model = PathBuf::from( From 813a41f77e9e8cde1ab2417e737a8171ce46ca93 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Sun, 26 Jul 2026 22:41:36 +0900 Subject: [PATCH 07/20] fix(xla): trace Molmo2 vision first divergence Align the eager Molmo2 ViT with the checkpoint's declared gelu_pytorch_tanh activation instead of the global exact-erf helper, and reject incompatible activation metadata in the StableHLO contract. Add diagnostics-only MLX and IREE stage outputs for patch and position embeddings, the first and configured ViT blocks, selected-layer concatenation, masked gather and valid counts, pooling query and output, and the complete SwiGLU projection. The ignored checkpoint gate now compares those stages in order and reports the first divergence while preserving negative fixtures for layer order, denominator, and clamped-index leakage. Validation: cargo fmt --all -- --check; cargo check --features xla-iree --lib; cargo check --no-default-features --features xla-diagnostics-cpu --test molmo2_xla_vision_parity; cargo test -p mlxcel-xla --features diagnostics molmo2 --lib; cargo test --no-default-features --features xla-diagnostics-cpu --test molmo2_xla_vision_parity synthetic; cargo test -p mlxcel molmo2_uses_the_checkpoint_pytorch_tanh_gelu --lib; focused cargo clippy with the repository's pre-existing warning classes allowed. Refs #871 --- Cargo.toml | 2 +- src/lib.rs | 3 +- src/lib/mlxcel-xla/src/emitter/mod.rs | 2 + .../mlxcel-xla/src/emitter/molmo2_config.rs | 7 +- .../src/emitter/molmo2_config_tests.rs | 19 +- .../mlxcel-xla/src/emitter/molmo2_vision.rs | 124 ++++- src/lib/mlxcel-xla/src/lib.rs | 4 + .../mlxcel-xla/src/molmo2_vision_runtime.rs | 504 ++++++++++++++---- src/loading/mod.rs | 3 +- src/loading/vlm.rs | 3 +- src/loading/vlm_special.rs | 67 ++- src/vision/encoders/molmo2.rs | 343 +++++++++++- tests/molmo2_xla_vision_parity.rs | 66 ++- 13 files changed, 964 insertions(+), 183 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index a3b338200..834cdc86f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -107,7 +107,7 @@ xla-micro-oracle = ["xla-iree", "mlxcel-xla/micro-oracle"] # CPU-only companion for bounded host-reference checks on machines where the # CUDA driver is unavailable. This deliberately does not replace the qualified # `xla-diagnostics` production-target gate. -xla-diagnostics-cpu = ["xla-iree"] +xla-diagnostics-cpu = ["xla-iree", "mlxcel-xla/diagnostics"] # Axis A "weight-load surgery" framework. # # On by default so production `mlxcel` and `mlxcel-server` binaries expose diff --git a/src/lib.rs b/src/lib.rs index 0646f491a..f4f3a35f1 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -123,7 +123,8 @@ pub use loaded_model::LoadedModel; pub use loaded_model_capabilities::VlmRuntimeRef; #[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] pub use loading::{ - Molmo2XlaVisionReference, Molmo2XlaVisionReferenceProjection, load_molmo2_xla_vision_reference, + Molmo2XlaVisionReference, Molmo2XlaVisionReferenceProjection, Molmo2XlaVisionReferenceStage, + load_molmo2_xla_vision_reference, }; pub use loading::{ context_window_from_config, load_model, load_model_with_adapter, diff --git a/src/lib/mlxcel-xla/src/emitter/mod.rs b/src/lib/mlxcel-xla/src/emitter/mod.rs index b725295de..65d203ba9 100644 --- a/src/lib/mlxcel-xla/src/emitter/mod.rs +++ b/src/lib/mlxcel-xla/src/emitter/mod.rs @@ -157,6 +157,8 @@ pub(crate) use model::{ pub(crate) use molmo2_config::{Molmo2VisionConfig, Molmo2VisionWeightSpec}; #[allow(unused_imports)] pub(crate) use molmo2_vision::emit_molmo2_vision; +#[cfg(feature = "diagnostics")] +pub(crate) use molmo2_vision::emit_molmo2_vision_diagnostics; #[allow(unused_imports)] pub(crate) use vision::emit_vision; // The only consumer is `IreeVisionDiagnosticProjector` in `vision_runtime`, which diff --git a/src/lib/mlxcel-xla/src/emitter/molmo2_config.rs b/src/lib/mlxcel-xla/src/emitter/molmo2_config.rs index c2dfe3e4a..b1ca2ea01 100644 --- a/src/lib/mlxcel-xla/src/emitter/molmo2_config.rs +++ b/src/lib/mlxcel-xla/src/emitter/molmo2_config.rs @@ -150,6 +150,11 @@ impl Molmo2VisionConfig { } let vit = object(&root, "vit_config")?; let adapter = object(&root, "adapter_config")?; + if vit.get("hidden_act").and_then(Value::as_str) != Some("gelu_pytorch_tanh") { + return Err( + "Molmo2 XLA requires vit_config.hidden_act `gelu_pytorch_tanh`".to_string(), + ); + } let processor: Value = serde_json::from_str(processor) .map_err(|error| format!("parse preprocessor_config.json: {error}"))?; let processor = processor @@ -355,7 +360,7 @@ impl Molmo2VisionConfig { pub(crate) fn fingerprint(&self) -> String { format!( - "molmo2-vision-v1;position=exact-default;selected={:?};pool-mask={};patch-id={};crops={};overlap={:?};patches={};pool-groups={};pool={}x{};hidden={};inter={};layers={};emitted={};heads={};head-dim={};pool-hidden={};pool-heads={};pool-head-dim={};projector-inter={};text-hidden={}", + "molmo2-vision-v1;position=exact-default;activation=gelu-pytorch-tanh;selected={:?};pool-mask={};patch-id={};crops={};overlap={:?};patches={};pool-groups={};pool={}x{};hidden={};inter={};layers={};emitted={};heads={};head-dim={};pool-hidden={};pool-heads={};pool-head-dim={};projector-inter={};text-hidden={}", self.selected_layers, self.pooling_attention_mask, self.image_patch_id, diff --git a/src/lib/mlxcel-xla/src/emitter/molmo2_config_tests.rs b/src/lib/mlxcel-xla/src/emitter/molmo2_config_tests.rs index 0503b1ee7..41ca55469 100644 --- a/src/lib/mlxcel-xla/src/emitter/molmo2_config_tests.rs +++ b/src/lib/mlxcel-xla/src/emitter/molmo2_config_tests.rs @@ -22,7 +22,8 @@ fn pinned_config() -> String { "hidden_size": 8, "intermediate_size": 16, "num_attention_heads": 2, "head_dim": 4, "num_hidden_layers": 27, "image_default_input_size": [28, 28], - "image_patch_size": 14, "image_num_pos": 4, "layer_norm_eps": 1e-6 + "image_patch_size": 14, "image_num_pos": 4, "layer_norm_eps": 1e-6, + "hidden_act": "gelu_pytorch_tanh" }, "adapter_config": { "hidden_size": 8, "intermediate_size": 12, "text_hidden_size": 10, @@ -61,6 +62,11 @@ fn resolves_pinned_layers_and_static_bucket_identity() { assert_eq!(config.static_crops, 9); assert_eq!(config.static_pool_groups, 9); assert!(config.fingerprint().contains("position=exact-default")); + assert!( + config + .fingerprint() + .contains("activation=gelu-pytorch-tanh") + ); assert!(config.fingerprint().contains("selected=[24, 18]")); assert!(config.fingerprint().contains("pool-mask=true")); assert!(config.fingerprint().contains("layers=27;emitted=25")); @@ -113,3 +119,14 @@ fn rejects_position_grid_and_selected_layer_drift() { .contains("outside") ); } + +#[test] +fn rejects_noncanonical_vit_activation() { + let mut config: Value = serde_json::from_str(&pinned_config()).unwrap(); + config["vit_config"]["hidden_act"] = Value::from("gelu"); + assert!( + Molmo2VisionConfig::from_json_strs(&config.to_string(), &pinned_processor()) + .unwrap_err() + .contains("gelu_pytorch_tanh") + ); +} diff --git a/src/lib/mlxcel-xla/src/emitter/molmo2_vision.rs b/src/lib/mlxcel-xla/src/emitter/molmo2_vision.rs index 1ca57496b..9230e34f2 100644 --- a/src/lib/mlxcel-xla/src/emitter/molmo2_vision.rs +++ b/src/lib/mlxcel-xla/src/emitter/molmo2_vision.rs @@ -222,13 +222,20 @@ fn self_attention( linear(builder, &context, &o_weight, Some(&o_bias)) } +struct IndexedPoolValues { + gathered_masked: Val, + query: Val, + pooled: Val, + counts: Val, +} + fn indexed_pool( builder: &mut Builder, features: &Val, signed_indices: &Val, args: &mut Args, config: &Molmo2VisionConfig, -) -> (Val, Val) { +) -> IndexedPoolValues { let groups = config.static_pool_groups; let group_size = config.pool_size; let zero_i32 = builder.const_i32(0); @@ -310,10 +317,15 @@ fn indexed_pool( let context = builder.transpose(&context, &[0, 2, 1, 3]); let context = builder.reshape(&context, vec![groups, config.pool_hidden]); let pooled = linear(builder, &context, &o_weight, Some(&o_bias)); - (pooled, counts) + IndexedPoolValues { + gathered_masked: gathered, + query, + pooled, + counts, + } } -pub(crate) fn emit_molmo2_vision(config: &Molmo2VisionConfig) -> String { +fn emit_molmo2_vision_inner(config: &Molmo2VisionConfig, diagnostics: bool) -> String { let specs = config.weight_specs(); let mut args = Args::new(&specs); let patch_weight = args.take(); @@ -333,12 +345,15 @@ pub(crate) fn emit_molmo2_vision(config: &Molmo2VisionConfig) -> String { ); let mut builder = Builder::new(); let mut hidden = linear(&mut builder, &patches, &patch_weight, Some(&patch_bias)); + let patch_embedding = diagnostics.then(|| hidden.clone()); let position = builder.broadcast( &position_embedding, &[1, 2], vec![config.static_crops, config.patches_per_crop, config.hidden], ); hidden = builder.add(&hidden, &position); + let positioned_embedding = diagnostics.then(|| hidden.clone()); + let mut early_block = None; let mut selected = vec![None::; config.selected_layers.len()]; for layer in 0..config.emitted_layers { // Norm weights follow attention projection weights in the persisted @@ -377,21 +392,26 @@ pub(crate) fn emit_molmo2_vision(config: &Molmo2VisionConfig) -> String { let mlp = tanh_gelu(&mut builder, &mlp); let mlp = linear(&mut builder, &mlp, &w2, Some(&b2)); hidden = builder.add(&residual, &mlp); + if diagnostics && layer == 0 { + early_block = Some(hidden.clone()); + } if let Some(slot) = selected_slot(&config.selected_layers, layer) { selected[slot] = Some(hidden.clone()); } } - let mut selected = selected.into_iter(); - let mut selected_features = match selected.next() { - Some(Some(feature)) => feature, - _ => unreachable!("validated Molmo2 selected layers must be emitted"), - }; - for feature in selected { - let feature = match feature { - Some(feature) => feature, - None => unreachable!("validated Molmo2 selected layer must be emitted"), - }; - selected_features = builder.concatenate(&selected_features, &feature, 2); + let selected = selected + .into_iter() + .map(|feature| { + feature + .unwrap_or_else(|| unreachable!("validated Molmo2 selected layer must be emitted")) + }) + .collect::>(); + let mut selected_features = selected + .first() + .cloned() + .unwrap_or_else(|| unreachable!("validated Molmo2 selection must not be empty")); + for feature in &selected[1..] { + selected_features = builder.concatenate(&selected_features, feature, 2); } let selected_features = builder.reshape( &selected_features, @@ -400,7 +420,7 @@ pub(crate) fn emit_molmo2_vision(config: &Molmo2VisionConfig) -> String { config.selected_width(), ], ); - let (pooled, _counts) = indexed_pool( + let pool = indexed_pool( &mut builder, &selected_features, &signed_indices, @@ -410,9 +430,9 @@ pub(crate) fn emit_molmo2_vision(config: &Molmo2VisionConfig) -> String { let w1 = args.take(); let w2 = args.take(); let w3 = args.take(); - let gate = linear(&mut builder, &pooled, &w1, None); + let gate = linear(&mut builder, &pool.pooled, &w1, None); let gate = silu(&mut builder, &gate); - let up = linear(&mut builder, &pooled, &w3, None); + let up = linear(&mut builder, &pool.pooled, &w3, None); let projected = builder.multiply(&gate, &up); let projected = linear(&mut builder, &projected, &w2, None); assert_eq!( @@ -420,15 +440,57 @@ pub(crate) fn emit_molmo2_vision(config: &Molmo2VisionConfig) -> String { specs.len(), "Molmo2 vision weight schema drifted" ); + let outputs = if diagnostics { + let mut outputs = vec![ + patch_embedding.expect("Molmo2 diagnostic patch embedding"), + position_embedding, + positioned_embedding.expect("Molmo2 diagnostic positioned embedding"), + early_block.expect("Molmo2 diagnostics require an early block"), + ]; + outputs.extend(selected); + outputs.extend([ + selected_features, + pool.gathered_masked, + pool.counts, + pool.query, + pool.pooled, + projected, + ]); + outputs + } else { + vec![projected] + }; + let output_types = outputs + .iter() + .map(|output| output.ty.render()) + .collect::>(); + let result_type = if output_types.len() == 1 { + output_types[0].clone() + } else { + format!("({})", output_types.join(", ")) + }; + let result_values = outputs + .iter() + .map(|output| output.name.as_str()) + .collect::>() + .join(", "); format!( - "module @molmo2_vision {{\n func.func public @main({signature}) -> {output} {{\n{body} return {value} : {output}\n }}\n}}\n", + "module @molmo2_vision {{\n func.func public @main({signature}) -> {result_type} {{\n{body} return {result_values} : {return_types}\n }}\n}}\n", signature = args.declarations.join(", "), - output = projected.ty.render(), body = builder.body(), - value = projected.name, + return_types = output_types.join(", "), ) } +pub(crate) fn emit_molmo2_vision(config: &Molmo2VisionConfig) -> String { + emit_molmo2_vision_inner(config, false) +} + +#[cfg(any(test, feature = "diagnostics"))] +pub(crate) fn emit_molmo2_vision_diagnostics(config: &Molmo2VisionConfig) -> String { + emit_molmo2_vision_inner(config, true) +} + #[cfg(test)] mod tests { use super::*; @@ -439,7 +501,8 @@ mod tests { "model_type":"molmo2","image_patch_id":151938, "vit_config":{"hidden_size":8,"intermediate_size":16,"num_attention_heads":2, "head_dim":4,"num_hidden_layers":2,"image_default_input_size":[28,28], - "image_patch_size":14,"image_num_pos":4,"layer_norm_eps":1e-6}, + "image_patch_size":14,"image_num_pos":4,"layer_norm_eps":1e-6, + "hidden_act":"gelu_pytorch_tanh"}, "adapter_config":{"hidden_size":8,"intermediate_size":12,"text_hidden_size":10, "num_attention_heads":2,"head_dim":4,"vit_layers":[0,1], "pooling_attention_mask":pooling_attention_mask} @@ -480,4 +543,23 @@ mod tests { assert_eq!(selected_slot(&[24, 4, 18], 18), Some(2)); assert_eq!(selected_slot(&[24, 4, 18], 22), None); } + + #[test] + fn diagnostic_graph_shares_production_math_and_orders_first_divergence_stages() { + let config = test_config(true); + let production = emit_molmo2_vision(&config); + let diagnostics = emit_molmo2_vision_diagnostics(&config); + assert_eq!( + production.matches("stablehlo.dot_general").count(), + diagnostics.matches("stablehlo.dot_general").count() + ); + assert!(production.contains("stablehlo.tanh")); + assert!( + diagnostics.contains( + "tensor<2x4x8xf32>, tensor<4x8xf32>, tensor<2x4x8xf32>, tensor<2x4x8xf32>" + ) + ); + assert!(diagnostics.contains("tensor<2xf32>")); + assert!(diagnostics.contains("tensor<2x10xf32>")); + } } diff --git a/src/lib/mlxcel-xla/src/lib.rs b/src/lib/mlxcel-xla/src/lib.rs index d76594964..6b56dffc4 100644 --- a/src/lib/mlxcel-xla/src/lib.rs +++ b/src/lib/mlxcel-xla/src/lib.rs @@ -176,6 +176,10 @@ pub use iree::PreparedPrefillDiagnostics; pub use molmo2::{ Molmo2InputError, Molmo2SafePooling, add_projected_features as add_molmo2_projected_features, }; +#[cfg(feature = "diagnostics")] +pub use molmo2_vision_runtime::{ + IreeMolmo2VisionDiagnosticProjector, Molmo2VisionDiagnosticStage, Molmo2VisionDiagnostics, +}; #[cfg(feature = "iree")] pub use molmo2_vision_runtime::{ IreeMolmo2VisionProjector, Molmo2VisionInput, Molmo2VisionProjection, diff --git a/src/lib/mlxcel-xla/src/molmo2_vision_runtime.rs b/src/lib/mlxcel-xla/src/molmo2_vision_runtime.rs index 8468658b9..c4f892f07 100644 --- a/src/lib/mlxcel-xla/src/molmo2_vision_runtime.rs +++ b/src/lib/mlxcel-xla/src/molmo2_vision_runtime.rs @@ -28,6 +28,8 @@ use crate::aux::{ IreeAuxiliaryModule, }; use crate::aux_manifest::{AuxiliaryArtifactContract, ensure_qualified_auxiliary_artifact}; +#[cfg(feature = "diagnostics")] +use crate::emitter::emit_molmo2_vision_diagnostics; use crate::emitter::{Molmo2VisionConfig, Molmo2VisionWeightSpec, emit_molmo2_vision}; use crate::iree::{cached_vmfb_path, compile_one_to, iree_compile_bin, target_flags}; use crate::molmo2::Molmo2SafePooling; @@ -57,6 +59,40 @@ pub struct Molmo2VisionInput<'a> { pub prompt_image_patch_count: usize, } +#[cfg(feature = "diagnostics")] +#[derive(Debug, Clone, PartialEq)] +pub struct Molmo2VisionDiagnosticStage { + pub name: String, + pub values: Vec, + pub shape: Vec, +} + +#[cfg(feature = "diagnostics")] +#[derive(Debug, Clone, PartialEq)] +pub struct Molmo2VisionDiagnostics { + pub stages: Vec, + pub projected_values: Vec, + pub projected_shape: [usize; 2], + pub signed_pooling_indices: Vec, + pub valid_pooling_counts: Vec, + pub active_groups: Vec, + pub elapsed_seconds: f64, + pub upload_bytes: usize, + pub transfer_bytes: usize, +} + +struct PreparedMolmo2VisionInput { + padded_patches: Vec, + padded_signed_indices: Vec, + signed_pooling_indices: Vec, + valid_pooling_counts: Vec, + active_groups: Vec, + #[cfg(feature = "diagnostics")] + crops: usize, + #[cfg(feature = "diagnostics")] + groups: usize, +} + fn hex(bytes: &[u8]) -> String { let mut output = String::with_capacity(bytes.len() * 2); for byte in bytes { @@ -293,6 +329,131 @@ fn decode_output(bytes: &[u8]) -> Result, String> { Ok(values) } +fn prepare_vision_input( + config: &Molmo2VisionConfig, + input: Molmo2VisionInput<'_>, +) -> Result { + let [crops, patches, patch_dim] = input.patches_shape; + if crops != input.image_num_crops + || crops == 0 + || crops > config.static_crops + || patches != config.patches_per_crop + || patch_dim != config.patch_dim + { + return Err(format!( + "Molmo2 patch shape {:?}, image_num_crops={} disagrees with static [{},{},{}]", + input.patches_shape, + input.image_num_crops, + config.static_crops, + config.patches_per_crop, + config.patch_dim + )); + } + let patch_values = crops + .checked_mul(patches) + .and_then(|value| value.checked_mul(patch_dim)) + .ok_or_else(|| "Molmo2 patch shape overflowed".to_string())?; + if input.patches.len() != patch_values { + return Err(format!( + "Molmo2 patch payload has {} values, expected {patch_values}", + input.patches.len() + )); + } + finite("Molmo2 patches", input.patches)?; + let [groups, group_size] = input.pooling_shape; + let grid = input + .image_grid + .iter() + .try_fold((), |(), value| { + (*value >= 0) + .then_some(()) + .ok_or_else(|| "Molmo2 image grid contains a negative dimension".to_string()) + }) + .map(|()| input.image_grid.map(|value| value as usize))?; + if !config.valid_runtime_geometry(crops, grid) { + return Err(format!( + "Molmo2 crop count {crops} and image grid {:?} disagree with processor geometry", + input.image_grid + )); + } + let [lo_h, lo_w, hi_h, hi_w] = grid; + let grid_groups = lo_h + .checked_mul(lo_w) + .and_then(|low| { + hi_h.checked_mul(hi_w) + .and_then(|high| low.checked_add(high)) + }) + .ok_or_else(|| "Molmo2 image grid overflowed".to_string())?; + if groups != grid_groups || groups > config.static_pool_groups || group_size != config.pool_size + { + return Err(format!( + "Molmo2 pooling shape {:?} disagrees with grid {:?} and static [{},{}]", + input.pooling_shape, input.image_grid, config.static_pool_groups, config.pool_size + )); + } + let safe = Molmo2SafePooling::prepare( + input.image_token_pooling, + groups, + group_size, + crops * patches, + ) + .map_err(|error| error.to_string())?; + let active_groups = safe + .active_groups_for_prompt(grid_groups, input.prompt_image_patch_count) + .map_err(|error| error.to_string())?; + let static_patch_values = config.static_crops * config.patches_per_crop * config.patch_dim; + let mut padded_patches = vec![0.0f32; static_patch_values]; + padded_patches[..input.patches.len()].copy_from_slice(input.patches); + let mut padded_signed_indices = vec![-1i32; config.static_pool_groups * config.pool_size]; + padded_signed_indices[..safe.signed_indices.len()].copy_from_slice(&safe.signed_indices); + Ok(PreparedMolmo2VisionInput { + padded_patches, + padded_signed_indices, + signed_pooling_indices: safe.signed_indices, + valid_pooling_counts: safe.valid_counts, + active_groups, + #[cfg(feature = "diagnostics")] + crops, + #[cfg(feature = "diagnostics")] + groups, + }) +} + +fn compile_vision_module( + model_dir: &Path, + device: &str, + config: &Molmo2VisionConfig, + mlir: &str, + tag: &str, + diagnostic_identity: Option<&str>, +) -> Result { + let compiler = iree_compile_bin()?; + if !compiler.is_file() { + return Err(format!("iree-compile not found at {}", compiler.display())); + } + let flags = target_flags(device)?; + let cache = std::env::temp_dir().join("mlxcel-xla-molmo2-vision-vmfb"); + std::fs::create_dir_all(&cache) + .map_err(|error| format!("mkdir {}: {error}", cache.display()))?; + let (weights, checkpoint_schema) = load_weights(model_dir, &config.weight_specs())?; + let graph_identity = diagnostic_identity + .map(|identity| format!("{};diagnostics={identity}", config.fingerprint())) + .unwrap_or_else(|| config.fingerprint()); + let contract = AuxiliaryArtifactContract::new( + ENTRY_NAME, + format!( + "{graph_identity};checkpoint_schema_sha256={}", + sha256(checkpoint_schema.as_bytes()) + ), + generation_identity(&compiler, flags, mlir)?, + )?; + let vmfb = cached_vmfb_path(&compiler, mlir, flags, &cache, tag, 0); + ensure_qualified_auxiliary_artifact(&vmfb, &contract, &weights, |temporary| { + compile_one_to(&compiler, mlir, flags, &cache, tag, 0, temporary) + })?; + IreeAuxiliaryModule::load(device, &vmfb, &contract, weights) +} + pub struct IreeMolmo2VisionProjector { module: IreeAuxiliaryModule, config: Molmo2VisionConfig, @@ -302,37 +463,8 @@ impl IreeMolmo2VisionProjector { pub fn load(model_dir: &Path, device: &str) -> Result { let config = Molmo2VisionConfig::from_model_dir(model_dir)?; let mlir = emit_molmo2_vision(&config); - let compiler = iree_compile_bin()?; - if !compiler.is_file() { - return Err(format!("iree-compile not found at {}", compiler.display())); - } - let flags = target_flags(device)?; - let cache = std::env::temp_dir().join("mlxcel-xla-molmo2-vision-vmfb"); - std::fs::create_dir_all(&cache) - .map_err(|error| format!("mkdir {}: {error}", cache.display()))?; - let (weights, checkpoint_schema) = load_weights(model_dir, &config.weight_specs())?; - let contract = AuxiliaryArtifactContract::new( - ENTRY_NAME, - format!( - "{};checkpoint_schema_sha256={}", - config.fingerprint(), - sha256(checkpoint_schema.as_bytes()) - ), - generation_identity(&compiler, flags, &mlir)?, - )?; - let vmfb = cached_vmfb_path(&compiler, &mlir, flags, &cache, "molmo2-vision", 0); - ensure_qualified_auxiliary_artifact(&vmfb, &contract, &weights, |temporary| { - compile_one_to( - &compiler, - &mlir, - flags, - &cache, - "molmo2-vision", - 0, - temporary, - ) - })?; - let module = IreeAuxiliaryModule::load(device, &vmfb, &contract, weights)?; + let module = + compile_vision_module(model_dir, device, &config, &mlir, "molmo2-vision", None)?; Ok(Self { module, config }) } @@ -352,85 +484,7 @@ impl IreeMolmo2VisionProjector { &mut self, input: Molmo2VisionInput<'_>, ) -> Result { - let [crops, patches, patch_dim] = input.patches_shape; - if crops != input.image_num_crops - || crops == 0 - || crops > self.config.static_crops - || patches != self.config.patches_per_crop - || patch_dim != self.config.patch_dim - { - return Err(format!( - "Molmo2 patch shape {:?}, image_num_crops={} disagrees with static [{},{},{}]", - input.patches_shape, - input.image_num_crops, - self.config.static_crops, - self.config.patches_per_crop, - self.config.patch_dim - )); - } - let patch_values = crops - .checked_mul(patches) - .and_then(|value| value.checked_mul(patch_dim)) - .ok_or_else(|| "Molmo2 patch shape overflowed".to_string())?; - if input.patches.len() != patch_values { - return Err(format!( - "Molmo2 patch payload has {} values, expected {patch_values}", - input.patches.len() - )); - } - finite("Molmo2 patches", input.patches)?; - let [groups, group_size] = input.pooling_shape; - let grid = input - .image_grid - .iter() - .try_fold((), |(), value| { - (*value >= 0) - .then_some(()) - .ok_or_else(|| "Molmo2 image grid contains a negative dimension".to_string()) - }) - .map(|()| input.image_grid.map(|value| value as usize))?; - if !self.config.valid_runtime_geometry(crops, grid) { - return Err(format!( - "Molmo2 crop count {crops} and image grid {:?} disagree with processor geometry", - input.image_grid - )); - } - let [lo_h, lo_w, hi_h, hi_w] = grid; - let grid_groups = lo_h - .checked_mul(lo_w) - .and_then(|low| { - hi_h.checked_mul(hi_w) - .and_then(|high| low.checked_add(high)) - }) - .ok_or_else(|| "Molmo2 image grid overflowed".to_string())?; - if groups != grid_groups - || groups > self.config.static_pool_groups - || group_size != self.config.pool_size - { - return Err(format!( - "Molmo2 pooling shape {:?} disagrees with grid {:?} and static [{},{}]", - input.pooling_shape, - input.image_grid, - self.config.static_pool_groups, - self.config.pool_size - )); - } - let safe = Molmo2SafePooling::prepare( - input.image_token_pooling, - groups, - group_size, - crops * patches, - ) - .map_err(|error| error.to_string())?; - let active_groups = safe - .active_groups_for_prompt(grid_groups, input.prompt_image_patch_count) - .map_err(|error| error.to_string())?; - let static_patch_values = - self.config.static_crops * self.config.patches_per_crop * self.config.patch_dim; - let mut padded_patches = vec![0.0f32; static_patch_values]; - padded_patches[..input.patches.len()].copy_from_slice(input.patches); - let mut signed = vec![-1i32; self.config.static_pool_groups * self.config.pool_size]; - signed[..safe.signed_indices.len()].copy_from_slice(&safe.signed_indices); + let prepared = prepare_vision_input(&self.config, input)?; let output_shape = [self.config.static_pool_groups, self.config.text_hidden]; let mut output = vec![0u8; output_shape.iter().product::() * 4]; let patch_shape = [ @@ -443,12 +497,12 @@ impl IreeMolmo2VisionProjector { self.module.invoke( &[ AuxiliaryInput { - bytes: f32_bytes(&padded_patches), + bytes: f32_bytes(&prepared.padded_patches), dtype: AuxiliaryTensorDType::Float32, shape: &patch_shape, }, AuxiliaryInput { - bytes: i32_bytes(&signed), + bytes: i32_bytes(&prepared.padded_signed_indices), dtype: AuxiliaryTensorDType::Int32, shape: &pooling_shape, }, @@ -460,8 +514,8 @@ impl IreeMolmo2VisionProjector { }], )?; let all_values = decode_output(&output)?; - let mut values = Vec::with_capacity(active_groups.len() * self.config.text_hidden); - for group in active_groups { + let mut values = Vec::with_capacity(prepared.active_groups.len() * self.config.text_hidden); + for &group in &prepared.active_groups { let start = group * self.config.text_hidden; values.extend_from_slice(&all_values[start..start + self.config.text_hidden]); } @@ -471,16 +525,230 @@ impl IreeMolmo2VisionProjector { self.config.text_hidden, ], values, - signed_pooling_indices: safe.signed_indices, - valid_pooling_counts: safe.valid_counts, + signed_pooling_indices: prepared.signed_pooling_indices, + valid_pooling_counts: prepared.valid_pooling_counts, elapsed_seconds: started.elapsed().as_secs_f64(), - upload_bytes: std::mem::size_of_val(padded_patches.as_slice()) - + std::mem::size_of_val(signed.as_slice()), + upload_bytes: std::mem::size_of_val(prepared.padded_patches.as_slice()) + + std::mem::size_of_val(prepared.padded_signed_indices.as_slice()), transfer_bytes: output.len(), }) } } +#[cfg(feature = "diagnostics")] +struct Molmo2DiagnosticStageSpec { + name: String, + static_shape: Vec, + active_shape: Vec, +} + +#[cfg(feature = "diagnostics")] +fn diagnostic_stage_specs( + config: &Molmo2VisionConfig, + crops: usize, + groups: usize, +) -> Vec { + let hidden_static = vec![config.static_crops, config.patches_per_crop, config.hidden]; + let hidden_active = vec![crops, config.patches_per_crop, config.hidden]; + let mut specs = vec![ + Molmo2DiagnosticStageSpec { + name: "vit.patch_embedding".to_string(), + static_shape: hidden_static.clone(), + active_shape: hidden_active.clone(), + }, + Molmo2DiagnosticStageSpec { + name: "vit.position_embedding".to_string(), + static_shape: vec![config.position_count, config.hidden], + active_shape: vec![config.position_count, config.hidden], + }, + Molmo2DiagnosticStageSpec { + name: "vit.positioned_embedding".to_string(), + static_shape: hidden_static.clone(), + active_shape: hidden_active.clone(), + }, + Molmo2DiagnosticStageSpec { + name: "vit.block.0".to_string(), + static_shape: hidden_static.clone(), + active_shape: hidden_active.clone(), + }, + ]; + specs.extend( + config + .selected_layers + .iter() + .map(|layer| Molmo2DiagnosticStageSpec { + name: format!("vit.selected.{layer}"), + static_shape: hidden_static.clone(), + active_shape: hidden_active.clone(), + }), + ); + specs.extend([ + Molmo2DiagnosticStageSpec { + name: "vit.concatenated".to_string(), + static_shape: vec![ + config.static_crops * config.patches_per_crop, + config.selected_width(), + ], + active_shape: vec![crops * config.patches_per_crop, config.selected_width()], + }, + Molmo2DiagnosticStageSpec { + name: "pool.gathered_masked".to_string(), + static_shape: vec![ + config.static_pool_groups, + config.pool_size, + config.selected_width(), + ], + active_shape: vec![groups, config.pool_size, config.selected_width()], + }, + Molmo2DiagnosticStageSpec { + name: "pool.valid_counts".to_string(), + static_shape: vec![config.static_pool_groups], + active_shape: vec![groups], + }, + Molmo2DiagnosticStageSpec { + name: "pool.query".to_string(), + static_shape: vec![config.static_pool_groups, config.selected_width()], + active_shape: vec![groups, config.selected_width()], + }, + Molmo2DiagnosticStageSpec { + name: "pool.output".to_string(), + static_shape: vec![config.static_pool_groups, config.pool_hidden], + active_shape: vec![groups, config.pool_hidden], + }, + Molmo2DiagnosticStageSpec { + name: "projector.output_all".to_string(), + static_shape: vec![config.static_pool_groups, config.text_hidden], + active_shape: vec![groups, config.text_hidden], + }, + ]); + specs +} + +#[cfg(feature = "diagnostics")] +pub struct IreeMolmo2VisionDiagnosticProjector { + module: IreeAuxiliaryModule, + config: Molmo2VisionConfig, +} + +#[cfg(feature = "diagnostics")] +impl IreeMolmo2VisionDiagnosticProjector { + pub fn load(model_dir: &Path, device: &str) -> Result { + let config = Molmo2VisionConfig::from_model_dir(model_dir)?; + let mlir = emit_molmo2_vision_diagnostics(&config); + let module = compile_vision_module( + model_dir, + device, + &config, + &mlir, + "molmo2-vision-diagnostics", + Some("first-divergence-v1"), + )?; + Ok(Self { module, config }) + } + + #[must_use] + pub fn image_patch_id(&self) -> i32 { + self.config.image_patch_id + } + + #[must_use] + pub fn text_hidden_size(&self) -> usize { + self.config.text_hidden + } + + pub fn project( + &mut self, + input: Molmo2VisionInput<'_>, + ) -> Result { + let prepared = prepare_vision_input(&self.config, input)?; + let specs = diagnostic_stage_specs(&self.config, prepared.crops, prepared.groups); + let mut buffers = specs + .iter() + .map(|spec| { + vec![0u8; spec.static_shape.iter().product::() * std::mem::size_of::()] + }) + .collect::>(); + let mut outputs = buffers + .iter_mut() + .zip(&specs) + .map(|(bytes, spec)| AuxiliaryOutput { + bytes, + dtype: AuxiliaryTensorDType::Float32, + shape: &spec.static_shape, + }) + .collect::>(); + let patch_shape = [ + self.config.static_crops, + self.config.patches_per_crop, + self.config.patch_dim, + ]; + let pooling_shape = [self.config.static_pool_groups, self.config.pool_size]; + let started = Instant::now(); + self.module.invoke( + &[ + AuxiliaryInput { + bytes: f32_bytes(&prepared.padded_patches), + dtype: AuxiliaryTensorDType::Float32, + shape: &patch_shape, + }, + AuxiliaryInput { + bytes: i32_bytes(&prepared.padded_signed_indices), + dtype: AuxiliaryTensorDType::Int32, + shape: &pooling_shape, + }, + ], + &mut outputs, + )?; + let elapsed_seconds = started.elapsed().as_secs_f64(); + drop(outputs); + let transfer_bytes = buffers.iter().map(Vec::len).sum(); + let stages = buffers + .into_iter() + .zip(specs) + .map(|(bytes, spec)| { + let mut values = decode_output(&bytes)?; + values.truncate(spec.active_shape.iter().product()); + Ok(Molmo2VisionDiagnosticStage { + name: spec.name, + values, + shape: spec.active_shape, + }) + }) + .collect::, String>>()?; + let projected_all = stages + .last() + .ok_or_else(|| "Molmo2 diagnostic projector output is missing".to_string())?; + let mut projected_values = + Vec::with_capacity(prepared.active_groups.len() * self.config.text_hidden); + for &group in &prepared.active_groups { + let start = group * self.config.text_hidden; + projected_values.extend_from_slice( + projected_all + .values + .get(start..start + self.config.text_hidden) + .ok_or_else(|| { + "Molmo2 diagnostic active projection row is truncated".to_string() + })?, + ); + } + Ok(Molmo2VisionDiagnostics { + projected_shape: [ + projected_values.len() / self.config.text_hidden, + self.config.text_hidden, + ], + projected_values, + stages, + signed_pooling_indices: prepared.signed_pooling_indices, + valid_pooling_counts: prepared.valid_pooling_counts, + active_groups: prepared.active_groups, + elapsed_seconds, + upload_bytes: std::mem::size_of_val(prepared.padded_patches.as_slice()) + + std::mem::size_of_val(prepared.padded_signed_indices.as_slice()), + transfer_bytes, + }) + } +} + #[cfg(test)] #[path = "molmo2_vision_runtime_tests.rs"] mod tests; diff --git a/src/loading/mod.rs b/src/loading/mod.rs index 22abd7b5e..a7b595c78 100644 --- a/src/loading/mod.rs +++ b/src/loading/mod.rs @@ -64,7 +64,8 @@ pub(crate) use self::vlm::load_qwen2_vl_iree_host_preprocessor; pub use self::vlm::load_qwen3_omni_speech; #[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] pub use self::vlm::{ - Molmo2XlaVisionReference, Molmo2XlaVisionReferenceProjection, load_molmo2_xla_vision_reference, + Molmo2XlaVisionReference, Molmo2XlaVisionReferenceProjection, Molmo2XlaVisionReferenceStage, + load_molmo2_xla_vision_reference, }; #[cfg(feature = "xla-iree")] pub(crate) use self::vlm::{ diff --git a/src/loading/vlm.rs b/src/loading/vlm.rs index 5ca951f83..9dfe8f234 100644 --- a/src/loading/vlm.rs +++ b/src/loading/vlm.rs @@ -144,7 +144,8 @@ pub(crate) use siglip::{load_aya_vision_vlm, load_paligemma_vlm}; pub(crate) use smolvlm::load_smolvlm_vlm; #[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] pub use special::{ - Molmo2XlaVisionReference, Molmo2XlaVisionReferenceProjection, load_molmo2_xla_vision_reference, + Molmo2XlaVisionReference, Molmo2XlaVisionReferenceProjection, Molmo2XlaVisionReferenceStage, + load_molmo2_xla_vision_reference, }; #[cfg(feature = "xla-iree")] pub(crate) use special::{ diff --git a/src/loading/vlm_special.rs b/src/loading/vlm_special.rs index 1da5d5103..e6f58e503 100644 --- a/src/loading/vlm_special.rs +++ b/src/loading/vlm_special.rs @@ -2048,6 +2048,15 @@ fn build_molmo2_vision_model( let vision_config = full_config.get("vision_config").unwrap_or(full_config); let vit_config = vision_config.get("vit_config").unwrap_or(vision_config); let adapter_config = vision_config.get("adapter_config").unwrap_or(vision_config); + let vit_hidden_act = vit_config + .get("hidden_act") + .and_then(Value::as_str) + .unwrap_or("gelu_pytorch_tanh"); + if vit_hidden_act != "gelu_pytorch_tanh" { + return Err(anyhow::anyhow!( + "Molmo2 vision hidden_act `{vit_hidden_act}` is unsupported; expected `gelu_pytorch_tanh`" + )); + } // Resolve negative adapter indices against the declared depth before // deriving the smaller prefix of blocks that the checkpoint must execute. @@ -2184,6 +2193,15 @@ pub struct Molmo2XlaVisionReferenceProjection { pub values: Vec, pub shape: [usize; 2], pub active_groups: Vec, + pub stages: Vec, +} + +#[cfg(any(test, feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] +#[cfg_attr(test, allow(dead_code))] +pub struct Molmo2XlaVisionReferenceStage { + pub name: String, + pub values: Vec, + pub shape: Vec, } #[cfg(any(test, feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] @@ -2214,7 +2232,53 @@ impl Molmo2XlaVisionReference { mlxcel_core::from_slice_f32(&processed.pixel_values, &[1, crops, patches, patch_dim]); let pooling = mlxcel_core::from_slice_i32(&processed.image_token_pooling, &[1, groups, group_size]); - let projected = self.vision_tower.forward(&images, &pooling); + #[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] + let (projected, stages) = { + let (projected, diagnostics) = self.vision_tower.forward_diagnostics(&images, &pooling); + let stages = diagnostics + .stages + .into_iter() + .map(|stage| { + let tensor = mlxcel_core::astype(&stage.tensor, mlxcel_core::dtype::FLOAT32); + let shape = mlxcel_core::array_shape(&tensor) + .into_iter() + .map(|dimension| { + usize::try_from(dimension).map_err(|_| { + anyhow::anyhow!( + "Molmo2 MLX stage {} has a negative dimension", + stage.name + ) + }) + }) + .collect::>>()?; + let raw = mlxcel_core::try_array_to_raw_bytes(&tensor).map_err(|error| { + anyhow::anyhow!("Failed to export Molmo2 MLX stage {}: {error}", stage.name) + })?; + if raw.len() != shape.iter().product::() * std::mem::size_of::() { + return Err(anyhow::anyhow!( + "Molmo2 MLX stage {} byte count disagrees with shape {shape:?}", + stage.name + )); + } + Ok(Molmo2XlaVisionReferenceStage { + name: stage.name, + values: raw + .chunks_exact(4) + .map(|chunk| { + f32::from_ne_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]) + }) + .collect(), + shape, + }) + }) + .collect::>>()?; + (projected, stages) + }; + #[cfg(not(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu")))] + let (projected, stages) = ( + self.vision_tower.forward(&images, &pooling), + Vec::::new(), + ); let projected = mlxcel_core::astype(&projected, mlxcel_core::dtype::FLOAT32); let raw = mlxcel_core::try_array_to_raw_bytes(&projected) .map_err(|error| anyhow::anyhow!("Failed to export Molmo2 MLX projection: {error}"))?; @@ -2254,6 +2318,7 @@ impl Molmo2XlaVisionReference { values, shape, active_groups, + stages, }) } } diff --git a/src/vision/encoders/molmo2.rs b/src/vision/encoders/molmo2.rs index 55dc31aaa..520ef4eb9 100644 --- a/src/vision/encoders/molmo2.rs +++ b/src/vision/encoders/molmo2.rs @@ -18,7 +18,8 @@ //! - ViT: 25 transformer blocks, Linear patch embedding (not Conv2d), //! positional embedding with bicubic interpolation //! - Adapter: Attention pooling 2D + SwiGLU image projector -//! - Layer selection: [-3, -9] = [22, 16] → concatenate → pool_dim = 2*1152 +//! - Layer selection: [-3, -9] over the declared 27 layers = [24, 18], +//! then concatenate → pool_dim = 2*1152 //! //! Reference: https://github.com/Blaizzy/mlx-vlm/blob/main/mlx_vlm/models/molmo2/vision.py @@ -26,6 +27,28 @@ use mlxcel_core::layers::{LayerNorm, Linear}; use mlxcel_core::weights::WeightMap; use mlxcel_core::{MlxArray, UniquePtr}; +/// Hugging Face's `gelu_pytorch_tanh`, evaluated in F32 so the eager Molmo2 +/// reference matches both the checkpoint's declared activation and StableHLO. +fn gelu_pytorch_tanh(x: &MlxArray) -> UniquePtr { + let output_dtype = mlxcel_core::array_dtype(x); + let x = mlxcel_core::astype(x, mlxcel_core::dtype::FLOAT32); + let half = mlxcel_core::full_f32(&[1], 0.5, mlxcel_core::dtype::FLOAT32); + let one = mlxcel_core::full_f32(&[1], 1.0, mlxcel_core::dtype::FLOAT32); + let sqrt_two_over_pi = mlxcel_core::full_f32(&[1], 0.797_884_6, mlxcel_core::dtype::FLOAT32); + let cubic_coefficient = mlxcel_core::full_f32(&[1], 0.044_715, mlxcel_core::dtype::FLOAT32); + let squared = mlxcel_core::multiply(&x, &x); + let cubed = mlxcel_core::multiply(&squared, &x); + let cubic = mlxcel_core::multiply(&cubic_coefficient, &cubed); + let inner = mlxcel_core::multiply(&sqrt_two_over_pi, &mlxcel_core::add(&x, &cubic)); + let cdf = mlxcel_core::multiply(&half, &mlxcel_core::add(&one, &mlxcel_core::tanh(&inner))); + let activated = mlxcel_core::multiply(&x, &cdf); + if output_dtype == mlxcel_core::dtype::FLOAT32 { + activated + } else { + mlxcel_core::astype(&activated, output_dtype) + } +} + // ViT MLP. pub(crate) struct ViTMLP { w1: Linear, @@ -35,7 +58,7 @@ pub(crate) struct ViTMLP { impl ViTMLP { fn forward(&self, x: &MlxArray) -> UniquePtr { let h = self.w1.forward(x); - let h = mlxcel_core::gelu_approx(&h); + let h = gelu_pytorch_tanh(&h); self.w2.forward(&h) } @@ -231,13 +254,20 @@ pub(crate) struct Molmo2VisionTransformer { image_num_pos: usize, } +#[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] +struct Molmo2VitDiagnostics { + patch_embedding: UniquePtr, + position_embedding: UniquePtr, + positioned_embedding: UniquePtr, +} + +#[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] +type Molmo2VitCapture = Option; +#[cfg(not(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu")))] +type Molmo2VitCapture = (); + impl Molmo2VisionTransformer { - pub(crate) fn add_pos_emb( - &self, - x: &MlxArray, - patch_h: i32, - patch_w: i32, - ) -> UniquePtr { + fn position_embedding(&self, x: &MlxArray, patch_h: i32, patch_w: i32) -> UniquePtr { let num_pos = self.image_num_pos as i32; let hidden_size = mlxcel_core::array_shape(&self.positional_embedding)[1]; @@ -259,8 +289,7 @@ impl Molmo2VisionTransformer { // x + pos_emb[None, :, :] let pos_emb = mlxcel_core::reshape(&pos_emb, &[1, num_patches.min(num_pos), hidden_size]); - let pos_emb = mlxcel_core::astype(&pos_emb, mlxcel_core::array_dtype(x)); - mlxcel_core::add(x, &pos_emb) + mlxcel_core::astype(&pos_emb, mlxcel_core::array_dtype(x)) } pub(crate) fn forward( @@ -268,18 +297,73 @@ impl Molmo2VisionTransformer { x: &MlxArray, patch_num: Option<(i32, i32)>, ) -> Vec> { + self.forward_inner::(x, patch_num).0 + } + + #[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] + fn forward_diagnostics( + &self, + x: &MlxArray, + patch_num: Option<(i32, i32)>, + ) -> (Vec>, Molmo2VitDiagnostics) { + let (hidden_states, capture) = self.forward_inner::(x, patch_num); + ( + hidden_states, + capture.expect("Molmo2 ViT diagnostics requested a capture"), + ) + } + + fn forward_inner( + &self, + x: &MlxArray, + patch_num: Option<(i32, i32)>, + ) -> (Vec>, Molmo2VitCapture) { let default_patch_size = (self.image_num_pos as f64).sqrt() as i32; let (patch_h, patch_w) = patch_num.unwrap_or((default_patch_size, default_patch_size)); - let x = self.patch_embedding.forward(x); - let mut x = self.add_pos_emb(&x, patch_h, patch_w); + let patch_embedding = self.patch_embedding.forward(x); + let position_embedding = self.position_embedding(&patch_embedding, patch_h, patch_w); + let mut x = mlxcel_core::add(&patch_embedding, &position_embedding); + #[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] + let capture = CAPTURE.then(|| Molmo2VitDiagnostics { + patch_embedding: mlxcel_core::copy( + patch_embedding + .as_ref() + .expect("Molmo2 patch embedding must be materialized"), + ), + position_embedding: mlxcel_core::copy( + mlxcel_core::reshape( + &position_embedding, + &[ + -1, + mlxcel_core::array_shape(&position_embedding) + .last() + .copied() + .expect("Molmo2 position embedding has a hidden dimension"), + ], + ) + .as_ref() + .expect("Molmo2 position embedding must be materialized"), + ), + positioned_embedding: mlxcel_core::copy( + x.as_ref() + .expect("Molmo2 positioned embedding must be materialized"), + ), + }); let mut hidden_states = Vec::with_capacity(self.blocks.len()); for block in &self.blocks { x = block.forward(&x); hidden_states.push(mlxcel_core::copy(&x)); } - hidden_states + #[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] + { + (hidden_states, capture) + } + #[cfg(not(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu")))] + { + (hidden_states, ()) + } } pub(crate) fn from_weights( @@ -360,9 +444,40 @@ pub struct Molmo2VisionModel { pooling_attention_mask: bool, } +#[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] +pub struct Molmo2VisionDiagnosticTensor { + pub name: String, + pub tensor: UniquePtr, +} + +#[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] +pub struct Molmo2VisionDiagnostics { + pub stages: Vec, +} + +#[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] +struct Molmo2EncodeDiagnostics { + vit: Molmo2VitDiagnostics, + early_block: UniquePtr, + selected_layers: Vec<(usize, UniquePtr)>, + concatenated_features: UniquePtr, +} + +#[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] +type Molmo2EncodeCapture = Option; +#[cfg(not(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu")))] +type Molmo2EncodeCapture = (); + +#[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] +type Molmo2ForwardCapture = Option; +#[cfg(not(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu")))] +type Molmo2ForwardCapture = (); + impl Molmo2VisionModel { - /// Encode images through the ViT, extracting features from selected layers - fn encode_image(&self, images: &MlxArray) -> UniquePtr { + fn encode_image_inner( + &self, + images: &MlxArray, + ) -> (UniquePtr, Molmo2EncodeCapture) { let shape = mlxcel_core::array_shape(images); let batch_size = shape[0]; let num_crops = shape[1]; @@ -371,6 +486,14 @@ impl Molmo2VisionModel { // Reshape to [B*num_crops, num_patch, patch_dim] let flat = mlxcel_core::reshape(images, &[batch_size * num_crops, num_patch, patch_dim]); + #[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] + let (hidden_states, vit_capture) = if CAPTURE { + let (hidden_states, capture) = self.image_vit.forward_diagnostics(&flat, None); + (hidden_states, Some(capture)) + } else { + (self.image_vit.forward(&flat, None), None) + }; + #[cfg(not(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu")))] let hidden_states = self.image_vit.forward(&flat, None); // Select and concatenate features from specified layers @@ -394,18 +517,82 @@ impl Molmo2VisionModel { // Reshape back to [B, num_crops, num_patch, features_dim] let feat_dim = mlxcel_core::array_shape(&image_features); let last_dim = feat_dim[feat_dim.len() - 1]; - mlxcel_core::reshape( + let output = mlxcel_core::reshape( &image_features, &[batch_size, num_crops, num_patch, last_dim], - ) + ); + #[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] + if CAPTURE { + let early_block = hidden_states + .first() + .and_then(|state| state.as_ref()) + .map(mlxcel_core::copy) + .expect("Molmo2 diagnostics require at least one ViT block"); + let selected_layers = self + .vit_layers + .iter() + .map(|&layer| { + ( + layer, + mlxcel_core::copy( + hidden_states[layer] + .as_ref() + .expect("Molmo2 selected layer must be materialized"), + ), + ) + }) + .collect(); + let concatenated_features = mlxcel_core::reshape( + &image_features, + &[batch_size * num_crops * num_patch, last_dim], + ); + return ( + output, + Some(Molmo2EncodeDiagnostics { + vit: vit_capture.expect("Molmo2 ViT capture must exist"), + early_block, + selected_layers, + concatenated_features, + }), + ); + } + #[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] + { + (output, None) + } + #[cfg(not(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu")))] + { + (output, ()) + } } /// Full forward: encode → pool → project pub fn forward(&self, images: &MlxArray, pooled_patches_idx: &MlxArray) -> UniquePtr { + self.forward_inner::(images, pooled_patches_idx).0 + } + + #[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] + pub fn forward_diagnostics( + &self, + images: &MlxArray, + pooled_patches_idx: &MlxArray, + ) -> (UniquePtr, Molmo2VisionDiagnostics) { + let (projected, capture) = self.forward_inner::(images, pooled_patches_idx); + ( + projected, + capture.expect("Molmo2 vision diagnostics requested a capture"), + ) + } + + fn forward_inner( + &self, + images: &MlxArray, + pooled_patches_idx: &MlxArray, + ) -> (UniquePtr, Molmo2ForwardCapture) { let shape = mlxcel_core::array_shape(images); let batch_size = shape[0]; - let image_features = self.encode_image(images); + let (image_features, _encode_capture) = self.encode_image_inner::(images); let feat_shape = mlxcel_core::array_shape(&image_features); let dim = feat_shape[feat_shape.len() - 1]; @@ -421,6 +608,11 @@ impl Molmo2VisionModel { let valid = mlxcel_core::greater_equal(pooled_patches_idx, &zeros); // valid_token = any(valid, axis=-1) let valid_i32 = mlxcel_core::astype(&valid, mlxcel_core::dtype::INT32); + #[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] + let valid_counts = CAPTURE.then(|| { + let valid_f32 = mlxcel_core::astype(&valid, mlxcel_core::dtype::FLOAT32); + mlxcel_core::reshape(&mlxcel_core::sum_axis(&valid_f32, -1, false), &[-1]) + }); // Clip indices to >= 0 let idx = mlxcel_core::maximum(pooled_patches_idx, &zeros); @@ -440,6 +632,14 @@ impl Molmo2VisionModel { // Reshape for attention: [B * num_pooled, pool_size, dim] let to_pool = mlxcel_core::reshape(&to_pool, &[-1, pool_shape[2], dim]); + #[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] + let gathered_masked = CAPTURE.then(|| { + mlxcel_core::copy( + to_pool + .as_ref() + .expect("Molmo2 gathered features must be materialized"), + ) + }); // Build query: mean of valid patches per pooled position let (query, attn_mask) = if self.pooling_attention_mask { @@ -461,6 +661,8 @@ impl Molmo2VisionModel { let query = mlxcel_core::mean_axis(&to_pool, -2, true); (query, None) }; + #[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] + let pooling_query = CAPTURE.then(|| mlxcel_core::reshape(&query, &[-1, dim])); // Cross-attention pooling let pooled = self.image_pooling_2d.forward( @@ -473,6 +675,8 @@ impl Molmo2VisionModel { let pooled_shape = mlxcel_core::array_shape(&pooled); let pooled_dim = pooled_shape[pooled_shape.len() - 1]; let pooled = mlxcel_core::reshape(&pooled, &[batch_size, -1, pooled_dim]); + #[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] + let pooling_output = CAPTURE.then(|| mlxcel_core::reshape(&pooled, &[-1, pooled_dim])); // Project through SwiGLU MLP let projected = self.image_projector.forward(&pooled); @@ -481,6 +685,14 @@ impl Molmo2VisionModel { let proj_shape = mlxcel_core::array_shape(&projected); let out_dim = proj_shape[proj_shape.len() - 1]; let projected = mlxcel_core::reshape(&projected, &[-1, out_dim]); + #[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] + let projector_output = CAPTURE.then(|| { + mlxcel_core::copy( + projected + .as_ref() + .expect("Molmo2 projected features must be materialized"), + ) + }); // Filter valid tokens: valid_token = any(valid, axis=-1) // sum valid along pool_size axis, then check > 0 @@ -502,12 +714,77 @@ impl Molmo2VisionModel { } } - if valid_indices.is_empty() { - return mlxcel_core::zeros(&[0, out_dim], mlxcel_core::array_dtype(&projected)); - } + let active_projected = if valid_indices.is_empty() { + mlxcel_core::zeros(&[0, out_dim], mlxcel_core::array_dtype(&projected)) + } else { + let indices = + mlxcel_core::from_slice_i32(&valid_indices, &[valid_indices.len() as i32]); + mlxcel_core::take(&projected, &indices, 0) + }; - let indices = mlxcel_core::from_slice_i32(&valid_indices, &[valid_indices.len() as i32]); - mlxcel_core::take(&projected, &indices, 0) + #[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] + if CAPTURE { + let encode = _encode_capture.expect("Molmo2 encode capture must exist"); + let mut stages = vec![ + Molmo2VisionDiagnosticTensor { + name: "vit.patch_embedding".to_string(), + tensor: encode.vit.patch_embedding, + }, + Molmo2VisionDiagnosticTensor { + name: "vit.position_embedding".to_string(), + tensor: encode.vit.position_embedding, + }, + Molmo2VisionDiagnosticTensor { + name: "vit.positioned_embedding".to_string(), + tensor: encode.vit.positioned_embedding, + }, + Molmo2VisionDiagnosticTensor { + name: "vit.block.0".to_string(), + tensor: encode.early_block, + }, + ]; + stages.extend(encode.selected_layers.into_iter().map(|(layer, tensor)| { + Molmo2VisionDiagnosticTensor { + name: format!("vit.selected.{layer}"), + tensor, + } + })); + stages.extend([ + Molmo2VisionDiagnosticTensor { + name: "vit.concatenated".to_string(), + tensor: encode.concatenated_features, + }, + Molmo2VisionDiagnosticTensor { + name: "pool.gathered_masked".to_string(), + tensor: gathered_masked.expect("Molmo2 gather capture must exist"), + }, + Molmo2VisionDiagnosticTensor { + name: "pool.valid_counts".to_string(), + tensor: valid_counts.expect("Molmo2 valid-count capture must exist"), + }, + Molmo2VisionDiagnosticTensor { + name: "pool.query".to_string(), + tensor: pooling_query.expect("Molmo2 query capture must exist"), + }, + Molmo2VisionDiagnosticTensor { + name: "pool.output".to_string(), + tensor: pooling_output.expect("Molmo2 pool output capture must exist"), + }, + Molmo2VisionDiagnosticTensor { + name: "projector.output_all".to_string(), + tensor: projector_output.expect("Molmo2 projector capture must exist"), + }, + ]); + return (active_projected, Some(Molmo2VisionDiagnostics { stages })); + } + #[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] + { + (active_projected, None) + } + #[cfg(not(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu")))] + { + (active_projected, ()) + } } /// Batched gather: for each batch, gather from features using indices @@ -632,3 +909,23 @@ fn get_weight_copy(weights: &WeightMap, name: &str) -> Result 0.0); + + let masked_values = [2.0, 6.0, 0.0, 0.0]; + let valid_mean = masked_values.iter().sum::() / 2.0; + let wrong_fixed_window_mean = masked_values.iter().sum::() / 4.0; + assert_ne!(valid_mean, wrong_fixed_window_mean); + + let patch_zero = 100.0; + let valid_patch = 2.0; + let masked_gather = valid_patch; + let unmasked_clamped_gather = valid_patch + patch_zero; + assert_ne!(masked_gather, unmasked_clamped_gather); +} + #[test] #[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] #[ignore = "requires a Molmo2 checkpoint plus configured MLX and IREE runtimes"] @@ -215,7 +233,7 @@ fn real_checkpoint_mlx_iree_vision_and_scatter_parity() -> Result<()> { } let mut projector = - IreeMolmo2VisionProjector::load(&model, &device).map_err(anyhow::Error::msg)?; + IreeMolmo2VisionDiagnosticProjector::load(&model, &device).map_err(anyhow::Error::msg)?; if projector.image_patch_id() != reference.image_patch_id() || projector.text_hidden_size() != reference.text_hidden_size() { @@ -232,22 +250,42 @@ fn real_checkpoint_mlx_iree_vision_and_scatter_parity() -> Result<()> { prompt_image_patch_count: independently_active.len(), }) .map_err(anyhow::Error::msg)?; - let iree_active = iree - .valid_pooling_counts - .iter() - .enumerate() - .filter_map(|(group, &count)| (count > 0).then_some(group)) - .collect::>(); - if iree_active != independently_active || iree.shape != eager.shape { + if iree.active_groups != independently_active || iree.projected_shape != eager.shape { return Err(anyhow!( - "active-row mismatch: processor={independently_active:?}, IREE={iree_active:?}, MLX shape={:?}, IREE shape={:?}", + "active-row mismatch: processor={independently_active:?}, IREE={:?}, MLX shape={:?}, IREE shape={:?}", + iree.active_groups, eager.shape, - iree.shape + iree.projected_shape )); } + if iree.stages.len() != eager.stages.len() { + return Err(anyhow!( + "diagnostic stage count mismatch: MLX={}, IREE={}", + eager.stages.len(), + iree.stages.len() + )); + } + for (eager_stage, iree_stage) in eager.stages.iter().zip(&iree.stages) { + if eager_stage.name != iree_stage.name || eager_stage.shape != iree_stage.shape { + return Err(anyhow!( + "diagnostic stage layout mismatch: MLX {} {:?}, IREE {} {:?}", + eager_stage.name, + eager_stage.shape, + iree_stage.name, + iree_stage.shape + )); + } + assert_within( + &format!("first-divergence stage {}", eager_stage.name), + &iree_stage.values, + &eager_stage.values, + max_abs_limit, + rms_limit, + )?; + } assert_within( "vision projection", - &iree.values, + &iree.projected_values, &eager.values, max_abs_limit, rms_limit, @@ -275,7 +313,7 @@ fn real_checkpoint_mlx_iree_vision_and_scatter_parity() -> Result<()> { reference.image_patch_id(), &base, hidden, - &iree.values, + &iree.projected_values, )?; let mut production_iree_merged = base; add_molmo2_projected_features( @@ -283,7 +321,7 @@ fn real_checkpoint_mlx_iree_vision_and_scatter_parity() -> Result<()> { reference.image_patch_id(), &mut production_iree_merged, hidden, - &iree.values, + &iree.projected_values, ) .map_err(|error| anyhow!(error.to_string()))?; if production_iree_merged != independently_iree_merged { From 4f73c50ee6264d904d57a1fa3b2040b5b10ea6d8 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Sun, 26 Jul 2026 22:59:09 +0900 Subject: [PATCH 08/20] fix(molmo2): prefer declared vision sections Prefer authoritative top-level vit_config and adapter_config sections over the empty vision_config compatibility object shipped by the pinned Molmo2 checkpoint. This keeps the eager MLX oracle aligned with IREE by resolving [-3, -9] against the declared 27-layer depth as [24, 18] before deriving the 25-block execution prefix. Add a regression fixture for the exact empty-wrapper checkpoint layout and preserve post-layer selection and configured concatenation order. Validation: cargo fmt --all -- --check; cargo test -p mlxcel molmo2_ --lib; cargo check --no-default-features --features xla-diagnostics-cpu --test molmo2_xla_vision_parity; cargo test -p mlxcel-xla --features diagnostics molmo2 --lib; git diff --check. Refs #871 --- src/loading/vlm_special.rs | 22 +++++++++++++++++++--- src/loading/vlm_special_tests.rs | 32 ++++++++++++++++++++++++-------- 2 files changed, 43 insertions(+), 11 deletions(-) diff --git a/src/loading/vlm_special.rs b/src/loading/vlm_special.rs index e6f58e503..2ffaa67fe 100644 --- a/src/loading/vlm_special.rs +++ b/src/loading/vlm_special.rs @@ -2038,6 +2038,24 @@ fn read_clip_triple(config: Option<&Value>, key: &str) -> Option<[f32; 3]> { }) } +fn molmo2_vision_config_sections(full_config: &Value) -> (&Value, &Value) { + // Official Molmo2 exports carry an empty `vision_config` compatibility + // object alongside the authoritative top-level sections. Prefer those + // explicit sections so the empty wrapper cannot trigger fallback depths. + let nested = full_config.get("vision_config"); + let vit_config = full_config + .get("vit_config") + .or_else(|| nested.and_then(|config| config.get("vit_config"))) + .or_else(|| nested.filter(|config| config.get("num_hidden_layers").is_some())) + .unwrap_or(full_config); + let adapter_config = full_config + .get("adapter_config") + .or_else(|| nested.and_then(|config| config.get("adapter_config"))) + .or_else(|| nested.filter(|config| config.get("vit_layers").is_some())) + .unwrap_or(full_config); + (vit_config, adapter_config) +} + fn build_molmo2_vision_model( weights: &WeightMap, full_config: &Value, @@ -2045,9 +2063,7 @@ fn build_molmo2_vision_model( ) -> Result { use vision::encoders::molmo2::Molmo2VisionModel; - let vision_config = full_config.get("vision_config").unwrap_or(full_config); - let vit_config = vision_config.get("vit_config").unwrap_or(vision_config); - let adapter_config = vision_config.get("adapter_config").unwrap_or(vision_config); + let (vit_config, adapter_config) = molmo2_vision_config_sections(full_config); let vit_hidden_act = vit_config .get("hidden_act") .and_then(Value::as_str) diff --git a/src/loading/vlm_special_tests.rs b/src/loading/vlm_special_tests.rs index 9c525da67..d195ce8c8 100644 --- a/src/loading/vlm_special_tests.rs +++ b/src/loading/vlm_special_tests.rs @@ -16,14 +16,14 @@ use super::{ cap_molmo2_vit_num_layers, dequantize_moondream3_weight, flatten_phi4mm_patch_embedding, inherit_quantization_if_missing, llama4_mm_tokens_per_image, llama4_quantization_params, llama4_token_ids, llama4_vision_prefix, load_minicpmv4_6_vlm, minicpmv4_6_text_weights, - molmo2_max_crops, molmo2_vit_execution_depth, moondream2_text_config_value, - moondream3_text_config_value, moondream3_vision_config_value, parse_molmo2_vit_layers, - phi3_num_crops, phi4_siglip_text_config_value, phi4mm_text_config_value, - phi4mm_vision_config_value, remap_minicpmo_text_weights, remap_minicpmv4_6_weights, - remap_phi4mm_weights, resolve_molmo2_vit_layers, resolve_moondream2_eos_token_id, - rewrite_molmo2_weight_key, rewrite_moondream2_weight_key, rewrite_moondream3_weight_key, - rewrite_phi3_weight_key, rewrite_phi4_siglip_weight_key, rewrite_phi4mm_weight_key, - should_transpose_phi3_patch_embedding, + molmo2_max_crops, molmo2_vision_config_sections, molmo2_vit_execution_depth, + moondream2_text_config_value, moondream3_text_config_value, moondream3_vision_config_value, + parse_molmo2_vit_layers, phi3_num_crops, phi4_siglip_text_config_value, + phi4mm_text_config_value, phi4mm_vision_config_value, remap_minicpmo_text_weights, + remap_minicpmv4_6_weights, remap_phi4mm_weights, resolve_molmo2_vit_layers, + resolve_moondream2_eos_token_id, rewrite_molmo2_weight_key, rewrite_moondream2_weight_key, + rewrite_moondream3_weight_key, rewrite_phi3_weight_key, rewrite_phi4_siglip_weight_key, + rewrite_phi4mm_weight_key, should_transpose_phi3_patch_embedding, }; use crate::moondream2_prompt::Moondream2PromptStyle; use mlxcel_core::dtype; @@ -785,6 +785,22 @@ fn molmo2_layer_resolution_uses_declared_depth_and_preserves_config_order() { assert!(resolve_molmo2_vit_layers(27, &[]).is_err()); } +#[test] +fn molmo2_empty_vision_wrapper_does_not_shadow_top_level_sections() { + let config = json!({ + "vision_config": {}, + "vit_config": {"num_hidden_layers": 27}, + "adapter_config": {"vit_layers": [-3, -9]} + }); + let (vit_config, adapter_config) = molmo2_vision_config_sections(&config); + let declared_layers = vit_config["num_hidden_layers"].as_u64().unwrap() as usize; + let configured_layers = parse_molmo2_vit_layers(adapter_config); + let selected_layers = resolve_molmo2_vit_layers(declared_layers, &configured_layers).unwrap(); + + assert_eq!(selected_layers, vec![24, 18]); + assert_eq!(molmo2_vit_execution_depth(&selected_layers).unwrap(), 25); +} + #[test] fn rewrite_molmo2_weight_key_maps_text_vision_and_lm_head_prefixes() { assert_eq!( From 64a19c394af8ad5f95d396d84e24d441123e4dfe Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Sun, 26 Jul 2026 23:28:29 +0900 Subject: [PATCH 09/20] fix(xla): preserve Molmo2 projector SiLU order The post-selected-layer actual gate passes pooling output but first diverges at projector.output_all with max_abs 0.1640625. The checkpoint stores all three projector weights as F32, and eager MLX and StableHLO agree on shapes and w1/w3/product/w2 ordering; the remaining semantic drift is that StableHLO reassociated MLX's x * sigmoid(x) into x / (1 + exp(-x)), changing F32 rounding before the wide w2 projection. Compute the reciprocal sigmoid first and multiply by x, matching the eager MLX operation graph. Add a mutation-sensitive emitter regression that rejects the quotient form and pins the sigmoid-then-multiply data dependency. Validation: - cargo test -p mlxcel-xla --features diagnostics molmo2 --lib -- --nocapture (19 passed) - cargo check --no-default-features --features xla-diagnostics-cpu --test molmo2_xla_vision_parity - focused mlxcel-xla Clippy with unrelated baseline lints allowed - cargo fmt --all -- --check - git diff --check The heavyweight actual-checkpoint gate was not repeated after this fix. Refs #871 --- .../mlxcel-xla/src/emitter/molmo2_vision.rs | 49 ++++++++++++++++++- 1 file changed, 48 insertions(+), 1 deletion(-) diff --git a/src/lib/mlxcel-xla/src/emitter/molmo2_vision.rs b/src/lib/mlxcel-xla/src/emitter/molmo2_vision.rs index 9230e34f2..05cfbb328 100644 --- a/src/lib/mlxcel-xla/src/emitter/molmo2_vision.rs +++ b/src/lib/mlxcel-xla/src/emitter/molmo2_vision.rs @@ -138,7 +138,10 @@ fn silu(builder: &mut Builder, value: &Val) -> Val { let negative = builder.negate(value); let exponential = builder.exponential(&negative); let denominator = builder.add(&one, &exponential); - builder.divide(value, &denominator) + // Preserve MLX's `x * sigmoid(x)` operation order. Reassociating this as + // `x / (1 + exp(-x))` changes F32 rounding before the wide w2 projection. + let sigmoid = builder.divide(&one, &denominator); + builder.multiply(value, &sigmoid) } fn softmax_last(builder: &mut Builder, scores: &Val) -> Val { @@ -562,4 +565,48 @@ mod tests { assert!(diagnostics.contains("tensor<2xf32>")); assert!(diagnostics.contains("tensor<2x10xf32>")); } + + #[test] + fn projector_silu_preserves_mlx_sigmoid_multiply_rounding_order() { + let mut builder = Builder::new(); + let input = Builder::arg(0, Ty::f32(vec![3])); + let output = silu(&mut builder, &input); + let body = builder.body(); + let lines = body.lines().collect::>(); + let broadcast = lines + .iter() + .find(|line| line.contains("stablehlo.broadcast_in_dim")) + .expect("SiLU must broadcast one"); + let one = broadcast + .split_once(" = ") + .map(|(name, _)| name.trim()) + .expect("broadcast result"); + let divide = lines + .iter() + .find(|line| line.contains("stablehlo.divide")) + .expect("SiLU must compute sigmoid"); + assert!( + divide.contains(&format!("stablehlo.divide {one}, ")), + "SiLU must divide one by the denominator instead of reassociating x / denominator: {divide}" + ); + let sigmoid = divide + .split_once(" = ") + .map(|(name, _)| name.trim()) + .expect("sigmoid result"); + let multiply = lines + .iter() + .find(|line| line.contains("stablehlo.multiply")) + .expect("SiLU must multiply x by sigmoid"); + assert!( + multiply.contains(&format!("stablehlo.multiply %arg0, {sigmoid}")), + "SiLU must preserve MLX x * sigmoid(x) order: {multiply}" + ); + assert_eq!( + output.name, + multiply + .split_once(" = ") + .map(|(name, _)| name.trim()) + .expect("SiLU output") + ); + } } From 905375495204de0dafa4eaaca7fa9a198f8529af Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Sun, 26 Jul 2026 23:50:18 +0900 Subject: [PATCH 10/20] test(xla): report Molmo2 gate progress The ignored actual-checkpoint reference check only emitted comparison output after both eager MLX and IREE work completed, so long checkpoint loading, projection, compilation, or invocation could appear stalled without any flushed progress. Emit explicitly flushed start/completion messages and 60-second heartbeats around MLX vision-only loading and projection plus IREE diagnostic compile/load and invocation. Flush every stage comparison before enforcing the existing limits, and add a synthetic regression for output flushing and the sub-five-minute heartbeat interval. No numerical threshold, model operation, or CSV changes are included. The heavyweight actual-checkpoint gate was not run. Validation: - Molmo2 diagnostics integration tests: 5 passed, 1 ignored - xla-diagnostics-cpu integration check - focused Clippy with existing repository baseline lints allowed - cargo fmt --all -- --check - git diff --check Refs #871 --- tests/molmo2_xla_vision_parity.rs | 152 ++++++++++++++++++++++++++---- 1 file changed, 132 insertions(+), 20 deletions(-) diff --git a/tests/molmo2_xla_vision_parity.rs b/tests/molmo2_xla_vision_parity.rs index 6f23d6be7..6b73c92d3 100644 --- a/tests/molmo2_xla_vision_parity.rs +++ b/tests/molmo2_xla_vision_parity.rs @@ -18,7 +18,13 @@ //! the IREE vision projector. It never loads either text decoder. #[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] -use std::path::PathBuf; +use std::{ + io::{self, Write}, + path::PathBuf, + sync::mpsc::{self, Sender}, + thread::{self, JoinHandle}, + time::{Duration, Instant}, +}; use anyhow::{Result, anyhow}; #[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] @@ -27,6 +33,74 @@ use mlxcel_xla::add_molmo2_projected_features; #[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] use mlxcel_xla::{IreeMolmo2VisionDiagnosticProjector, Molmo2VisionInput}; +#[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] +const PROGRESS_HEARTBEAT_INTERVAL: Duration = Duration::from_secs(60); + +#[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] +fn write_progress(output: &mut impl Write, message: &str) -> io::Result<()> { + writeln!(output, "[molmo2-reference] {message}")?; + output.flush() +} + +#[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] +fn emit_progress(message: &str) { + let stderr = io::stderr(); + let mut stderr = stderr.lock(); + let _ = write_progress(&mut stderr, message); +} + +#[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] +struct ProgressHeartbeat { + stop: Sender<()>, + worker: Option>, +} + +#[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] +impl Drop for ProgressHeartbeat { + fn drop(&mut self) { + let _ = self.stop.send(()); + if let Some(worker) = self.worker.take() { + let _ = worker.join(); + } + } +} + +#[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] +fn with_progress( + label: &'static str, + operation: impl FnOnce() -> Result, +) -> Result { + emit_progress(&format!("{label}: started")); + let started = Instant::now(); + let (stop, receiver) = mpsc::channel(); + let worker = thread::spawn(move || { + while let Err(mpsc::RecvTimeoutError::Timeout) = + receiver.recv_timeout(PROGRESS_HEARTBEAT_INTERVAL) + { + emit_progress(&format!( + "{label}: still running (elapsed={}s)", + started.elapsed().as_secs() + )); + } + }); + let progress = ProgressHeartbeat { + stop, + worker: Some(worker), + }; + let result = operation(); + drop(progress); + let outcome = if result.is_ok() { + "completed" + } else { + "failed" + }; + emit_progress(&format!( + "{label}: {outcome} (elapsed={}s)", + started.elapsed().as_secs() + )); + result +} + #[derive(Debug, Clone, Copy)] struct Comparison { max_abs: f32, @@ -77,6 +151,10 @@ fn assert_within( rms_limit: f32, ) -> Result<()> { let comparison = compare(actual, expected)?; + emit_progress(&format!( + "{label}: max_abs={} at {}, rms={}", + comparison.max_abs, comparison.max_index, comparison.rms + )); if comparison.max_abs > max_abs_limit || comparison.rms > rms_limit { return Err(anyhow!( "{label} parity failed: max_abs={} at {}, rms={}, limits=({}, {})", @@ -87,10 +165,6 @@ fn assert_within( rms_limit )); } - eprintln!( - "{label}: max_abs={} at {}, rms={}", - comparison.max_abs, comparison.max_index, comparison.rms - ); Ok(()) } @@ -193,6 +267,37 @@ fn synthetic_negative_controls_detect_layer_denominator_and_clamped_index_drift( assert_ne!(masked_gather, unmasked_clamped_gather); } +#[test] +#[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] +fn synthetic_progress_flushes_and_heartbeats_before_five_minutes() { + #[derive(Default)] + struct FlushProbe { + bytes: Vec, + flushes: usize, + } + + impl Write for FlushProbe { + fn write(&mut self, buffer: &[u8]) -> io::Result { + self.bytes.extend_from_slice(buffer); + Ok(buffer.len()) + } + + fn flush(&mut self) -> io::Result<()> { + self.flushes += 1; + Ok(()) + } + } + + let mut probe = FlushProbe::default(); + write_progress(&mut probe, "MLX eager diagnostic projection: started").unwrap(); + assert_eq!(probe.flushes, 1); + assert_eq!( + String::from_utf8(probe.bytes).unwrap(), + "[molmo2-reference] MLX eager diagnostic projection: started\n" + ); + assert!(PROGRESS_HEARTBEAT_INTERVAL < Duration::from_secs(5 * 60)); +} + #[test] #[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] #[ignore = "requires a Molmo2 checkpoint plus configured MLX and IREE runtimes"] @@ -211,10 +316,14 @@ fn real_checkpoint_mlx_iree_vision_and_scatter_parity() -> Result<()> { let rms_limit = tolerance("MLXCEL_MOLMO2_RMS", 0.01)?; let _runtime = initialize_runtime(); - let reference = load_molmo2_xla_vision_reference(&model)?; + let reference = with_progress("MLX vision-only checkpoint load", || { + load_molmo2_xla_vision_reference(&model) + })?; let image = image::open(&image_path) .map_err(|error| anyhow!("open {}: {error}", image_path.display()))?; - let eager = reference.project(&image)?; + let eager = with_progress("MLX eager diagnostic projection", || { + reference.project(&image) + })?; let processed = &eager.processed; let pooling_shape = processed .image_token_pooling_shape @@ -232,24 +341,27 @@ fn real_checkpoint_mlx_iree_vision_and_scatter_parity() -> Result<()> { )); } - let mut projector = - IreeMolmo2VisionDiagnosticProjector::load(&model, &device).map_err(anyhow::Error::msg)?; + let mut projector = with_progress("IREE diagnostic compile/load", || { + IreeMolmo2VisionDiagnosticProjector::load(&model, &device).map_err(anyhow::Error::msg) + })?; if projector.image_patch_id() != reference.image_patch_id() || projector.text_hidden_size() != reference.text_hidden_size() { return Err(anyhow!("MLX and IREE Molmo2 metadata disagree")); } - let iree = projector - .project(Molmo2VisionInput { - patches: &processed.pixel_values, - patches_shape: processed.pixel_values_shape.map(|value| value as usize), - image_token_pooling: &processed.image_token_pooling, - pooling_shape, - image_grid: processed.image_grid, - image_num_crops: processed.image_num_crops as usize, - prompt_image_patch_count: independently_active.len(), - }) - .map_err(anyhow::Error::msg)?; + let iree = with_progress("IREE diagnostic invocation", || { + projector + .project(Molmo2VisionInput { + patches: &processed.pixel_values, + patches_shape: processed.pixel_values_shape.map(|value| value as usize), + image_token_pooling: &processed.image_token_pooling, + pooling_shape, + image_grid: processed.image_grid, + image_num_crops: processed.image_num_crops as usize, + prompt_image_patch_count: independently_active.len(), + }) + .map_err(anyhow::Error::msg) + })?; if iree.active_groups != independently_active || iree.projected_shape != eager.shape { return Err(anyhow!( "active-row mismatch: processor={independently_active:?}, IREE={:?}, MLX shape={:?}, IREE shape={:?}", From 162c22fe7706b365db6e3d0dcf612c9a0449e47f Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 00:28:23 +0900 Subject: [PATCH 11/20] test(xla): localize Molmo2 projector divergence The reciprocal-before-multiply SiLU change reached the actual local-task VMFB and survived IREE lowering, but the real-checkpoint parity gate still failed at the same final projector boundary. Capture the projector w1, SiLU, w3, and product intermediates in both eager MLX and IREE diagnostics so the next checkpoint run identifies the first divergent operation without another speculative fix. Refs #871 --- .../mlxcel-xla/src/emitter/molmo2_vision.rs | 24 ++++-- .../mlxcel-xla/src/molmo2_vision_runtime.rs | 22 +++++- src/vision/encoders/molmo2.rs | 78 ++++++++++++++++--- 3 files changed, 108 insertions(+), 16 deletions(-) diff --git a/src/lib/mlxcel-xla/src/emitter/molmo2_vision.rs b/src/lib/mlxcel-xla/src/emitter/molmo2_vision.rs index 05cfbb328..c0a6f93fe 100644 --- a/src/lib/mlxcel-xla/src/emitter/molmo2_vision.rs +++ b/src/lib/mlxcel-xla/src/emitter/molmo2_vision.rs @@ -433,11 +433,11 @@ fn emit_molmo2_vision_inner(config: &Molmo2VisionConfig, diagnostics: bool) -> S let w1 = args.take(); let w2 = args.take(); let w3 = args.take(); - let gate = linear(&mut builder, &pool.pooled, &w1, None); - let gate = silu(&mut builder, &gate); - let up = linear(&mut builder, &pool.pooled, &w3, None); - let projected = builder.multiply(&gate, &up); - let projected = linear(&mut builder, &projected, &w2, None); + let gate_linear = linear(&mut builder, &pool.pooled, &w1, None); + let gate_activation = silu(&mut builder, &gate_linear); + let up_linear = linear(&mut builder, &pool.pooled, &w3, None); + let projector_product = builder.multiply(&gate_activation, &up_linear); + let projected = linear(&mut builder, &projector_product, &w2, None); assert_eq!( args.cursor, specs.len(), @@ -457,6 +457,10 @@ fn emit_molmo2_vision_inner(config: &Molmo2VisionConfig, diagnostics: bool) -> S pool.counts, pool.query, pool.pooled, + gate_linear, + gate_activation, + up_linear, + projector_product, projected, ]); outputs @@ -564,6 +568,16 @@ mod tests { ); assert!(diagnostics.contains("tensor<2xf32>")); assert!(diagnostics.contains("tensor<2x10xf32>")); + let return_values = diagnostics + .lines() + .find_map(|line| line.trim().strip_prefix("return ")) + .and_then(|line| line.split_once(" : ").map(|(values, _)| values)) + .expect("Molmo2 diagnostic graph must return named stage values"); + assert_eq!( + return_values.split(',').count(), + 16, + "diagnostics must include w1, SiLU, w3, and product before projector output" + ); } #[test] diff --git a/src/lib/mlxcel-xla/src/molmo2_vision_runtime.rs b/src/lib/mlxcel-xla/src/molmo2_vision_runtime.rs index c4f892f07..9ad6ed33b 100644 --- a/src/lib/mlxcel-xla/src/molmo2_vision_runtime.rs +++ b/src/lib/mlxcel-xla/src/molmo2_vision_runtime.rs @@ -615,6 +615,26 @@ fn diagnostic_stage_specs( static_shape: vec![config.static_pool_groups, config.pool_hidden], active_shape: vec![groups, config.pool_hidden], }, + Molmo2DiagnosticStageSpec { + name: "projector.w1".to_string(), + static_shape: vec![config.static_pool_groups, config.projector_intermediate], + active_shape: vec![groups, config.projector_intermediate], + }, + Molmo2DiagnosticStageSpec { + name: "projector.silu".to_string(), + static_shape: vec![config.static_pool_groups, config.projector_intermediate], + active_shape: vec![groups, config.projector_intermediate], + }, + Molmo2DiagnosticStageSpec { + name: "projector.w3".to_string(), + static_shape: vec![config.static_pool_groups, config.projector_intermediate], + active_shape: vec![groups, config.projector_intermediate], + }, + Molmo2DiagnosticStageSpec { + name: "projector.product".to_string(), + static_shape: vec![config.static_pool_groups, config.projector_intermediate], + active_shape: vec![groups, config.projector_intermediate], + }, Molmo2DiagnosticStageSpec { name: "projector.output_all".to_string(), static_shape: vec![config.static_pool_groups, config.text_hidden], @@ -641,7 +661,7 @@ impl IreeMolmo2VisionDiagnosticProjector { &config, &mlir, "molmo2-vision-diagnostics", - Some("first-divergence-v1"), + Some("first-divergence-v2-projector-stages"), )?; Ok(Self { module, config }) } diff --git a/src/vision/encoders/molmo2.rs b/src/vision/encoders/molmo2.rs index 520ef4eb9..3a6aa74ed 100644 --- a/src/vision/encoders/molmo2.rs +++ b/src/vision/encoders/molmo2.rs @@ -417,6 +417,15 @@ pub(crate) struct ImageProjectorMLP { w3: Linear, } +#[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] +struct ImageProjectorDiagnostics { + gate_linear: UniquePtr, + gate_activation: UniquePtr, + up_linear: UniquePtr, + product: UniquePtr, + output: UniquePtr, +} + impl ImageProjectorMLP { pub(crate) fn forward(&self, x: &MlxArray) -> UniquePtr { // silu(w1(x)) * w3(x) → w2(...) @@ -427,6 +436,37 @@ impl ImageProjectorMLP { self.w2.forward(&h) } + #[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] + fn forward_diagnostics( + &self, + x: &MlxArray, + ) -> (UniquePtr, ImageProjectorDiagnostics) { + let gate_linear = self.w1.forward(x); + let gate_activation = mlxcel_core::silu(&gate_linear); + let up_linear = self.w3.forward(x); + let product = mlxcel_core::multiply(&gate_activation, &up_linear); + let output = self.w2.forward(&product); + let projector_width = mlxcel_core::array_shape(&gate_linear) + .last() + .copied() + .expect("Molmo2 projector gate must have a feature dimension"); + let output_width = mlxcel_core::array_shape(&output) + .last() + .copied() + .expect("Molmo2 projector output must have a feature dimension"); + let captured_output = mlxcel_core::reshape(&output, &[-1, output_width]); + ( + output, + ImageProjectorDiagnostics { + gate_linear: mlxcel_core::reshape(&gate_linear, &[-1, projector_width]), + gate_activation: mlxcel_core::reshape(&gate_activation, &[-1, projector_width]), + up_linear: mlxcel_core::reshape(&up_linear, &[-1, projector_width]), + product: mlxcel_core::reshape(&product, &[-1, projector_width]), + output: captured_output, + }, + ) + } + pub(crate) fn from_weights(weights: &WeightMap, prefix: &str) -> Result { let w1 = Linear::from_weights(weights, &format!("{}.w1", prefix))?; let w2 = Linear::from_weights(weights, &format!("{}.w2", prefix))?; @@ -678,21 +718,21 @@ impl Molmo2VisionModel { #[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] let pooling_output = CAPTURE.then(|| mlxcel_core::reshape(&pooled, &[-1, pooled_dim])); - // Project through SwiGLU MLP + // Project through SwiGLU MLP. + #[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] + let (projected, projector_capture) = if CAPTURE { + let (projected, diagnostics) = self.image_projector.forward_diagnostics(&pooled); + (projected, Some(diagnostics)) + } else { + (self.image_projector.forward(&pooled), None) + }; + #[cfg(not(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu")))] let projected = self.image_projector.forward(&pooled); // Flatten to [total_valid_tokens, output_dim] let proj_shape = mlxcel_core::array_shape(&projected); let out_dim = proj_shape[proj_shape.len() - 1]; let projected = mlxcel_core::reshape(&projected, &[-1, out_dim]); - #[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] - let projector_output = CAPTURE.then(|| { - mlxcel_core::copy( - projected - .as_ref() - .expect("Molmo2 projected features must be materialized"), - ) - }); // Filter valid tokens: valid_token = any(valid, axis=-1) // sum valid along pool_size axis, then check > 0 @@ -725,6 +765,8 @@ impl Molmo2VisionModel { #[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] if CAPTURE { let encode = _encode_capture.expect("Molmo2 encode capture must exist"); + let projector = + projector_capture.expect("Molmo2 projector diagnostics must be captured"); let mut stages = vec![ Molmo2VisionDiagnosticTensor { name: "vit.patch_embedding".to_string(), @@ -770,9 +812,25 @@ impl Molmo2VisionModel { name: "pool.output".to_string(), tensor: pooling_output.expect("Molmo2 pool output capture must exist"), }, + Molmo2VisionDiagnosticTensor { + name: "projector.w1".to_string(), + tensor: projector.gate_linear, + }, + Molmo2VisionDiagnosticTensor { + name: "projector.silu".to_string(), + tensor: projector.gate_activation, + }, + Molmo2VisionDiagnosticTensor { + name: "projector.w3".to_string(), + tensor: projector.up_linear, + }, + Molmo2VisionDiagnosticTensor { + name: "projector.product".to_string(), + tensor: projector.product, + }, Molmo2VisionDiagnosticTensor { name: "projector.output_all".to_string(), - tensor: projector_output.expect("Molmo2 projector capture must exist"), + tensor: projector.output, }, ]); return (active_projected, Some(Molmo2VisionDiagnostics { stages })); From af3ad9f1763321e2b57b0b9ae2c2510868a3642a Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 01:55:58 +0900 Subject: [PATCH 12/20] test(xla): probe Molmo2 layer 24 row Compare layer 18 before layer 24, then capture only the known failing row across the seven layer-24 block boundaries. Keep the production graph, strict thresholds, and checkpoint execution unchanged. --- .../mlxcel-xla/src/emitter/molmo2_vision.rs | 93 ++++++++- .../mlxcel-xla/src/molmo2_vision_runtime.rs | 46 ++++- src/vision/encoders/molmo2.rs | 179 +++++++++++++++++- 3 files changed, 310 insertions(+), 8 deletions(-) diff --git a/src/lib/mlxcel-xla/src/emitter/molmo2_vision.rs b/src/lib/mlxcel-xla/src/emitter/molmo2_vision.rs index c0a6f93fe..ad4977729 100644 --- a/src/lib/mlxcel-xla/src/emitter/molmo2_vision.rs +++ b/src/lib/mlxcel-xla/src/emitter/molmo2_vision.rs @@ -17,6 +17,10 @@ use super::builder::{Builder, Ty, Val}; use super::molmo2_config::{Molmo2VisionConfig, Molmo2VisionWeightSpec}; +const MOLMO2_VIT_PROBE_LAYER: usize = 24; +// 591490 = 513 * checkpoint hidden width 1152 + component 514. +const MOLMO2_VIT_PROBE_FLAT_ROW: usize = 513; + struct Args { values: Vec, declarations: Vec, @@ -173,6 +177,16 @@ fn selected_slot(selected_layers: &[usize], layer: usize) -> Option { .position(|&selected| selected == layer) } +fn diagnostic_probe_row(builder: &mut Builder, value: &Val, config: &Molmo2VisionConfig) -> Val { + let crop = MOLMO2_VIT_PROBE_FLAT_ROW / config.patches_per_crop; + let token = MOLMO2_VIT_PROBE_FLAT_ROW % config.patches_per_crop; + let row = builder.slice( + value, + &[(crop, crop + 1), (token, token + 1), (0, config.hidden)], + ); + builder.reshape(&row, vec![1, config.hidden]) +} + fn self_attention( builder: &mut Builder, hidden: &Val, @@ -358,6 +372,7 @@ fn emit_molmo2_vision_inner(config: &Molmo2VisionConfig, diagnostics: bool) -> S let positioned_embedding = diagnostics.then(|| hidden.clone()); let mut early_block = None; let mut selected = vec![None::; config.selected_layers.len()]; + let mut probe_rows = Vec::new(); for layer in 0..config.emitted_layers { // Norm weights follow attention projection weights in the persisted // schema. Pull them before emitting the attention and pass normalized @@ -365,6 +380,12 @@ fn emit_molmo2_vision_inner(config: &Molmo2VisionConfig, diagnostics: bool) -> S let block_start = args.cursor; let norm_weight = args.values[block_start + 8].clone(); let norm_bias = args.values[block_start + 9].clone(); + let capture_probe = diagnostics + && layer == MOLMO2_VIT_PROBE_LAYER + && config.static_crops * config.patches_per_crop > MOLMO2_VIT_PROBE_FLAT_ROW; + if capture_probe { + probe_rows.push(diagnostic_probe_row(&mut builder, &hidden, config)); + } let normalized = layer_norm( &mut builder, &hidden, @@ -372,7 +393,13 @@ fn emit_molmo2_vision_inner(config: &Molmo2VisionConfig, diagnostics: bool) -> S &norm_bias, config.layer_norm_eps, ); + if capture_probe { + probe_rows.push(diagnostic_probe_row(&mut builder, &normalized, config)); + } let attention = self_attention(&mut builder, &normalized, &mut args, config); + if capture_probe { + probe_rows.push(diagnostic_probe_row(&mut builder, &attention, config)); + } // encoder_block consumes the already-taken attention schema, so finish // this block explicitly. let _attention_norm_weight = args.take(); @@ -380,6 +407,9 @@ fn emit_molmo2_vision_inner(config: &Molmo2VisionConfig, diagnostics: bool) -> S let ffn_norm_weight = args.take(); let ffn_norm_bias = args.take(); let residual = builder.add(&hidden, &attention); + if capture_probe { + probe_rows.push(diagnostic_probe_row(&mut builder, &residual, config)); + } let normalized = layer_norm( &mut builder, &residual, @@ -387,6 +417,9 @@ fn emit_molmo2_vision_inner(config: &Molmo2VisionConfig, diagnostics: bool) -> S &ffn_norm_bias, config.layer_norm_eps, ); + if capture_probe { + probe_rows.push(diagnostic_probe_row(&mut builder, &normalized, config)); + } let w1 = args.take(); let b1 = args.take(); let w2 = args.take(); @@ -394,7 +427,13 @@ fn emit_molmo2_vision_inner(config: &Molmo2VisionConfig, diagnostics: bool) -> S let mlp = linear(&mut builder, &normalized, &w1, Some(&b1)); let mlp = tanh_gelu(&mut builder, &mlp); let mlp = linear(&mut builder, &mlp, &w2, Some(&b2)); + if capture_probe { + probe_rows.push(diagnostic_probe_row(&mut builder, &mlp, config)); + } hidden = builder.add(&residual, &mlp); + if capture_probe { + probe_rows.push(diagnostic_probe_row(&mut builder, &hidden, config)); + } if diagnostics && layer == 0 { early_block = Some(hidden.clone()); } @@ -450,7 +489,26 @@ fn emit_molmo2_vision_inner(config: &Molmo2VisionConfig, diagnostics: bool) -> S positioned_embedding.expect("Molmo2 diagnostic positioned embedding"), early_block.expect("Molmo2 diagnostics require an early block"), ]; - outputs.extend(selected); + let mut diagnostic_selected = config + .selected_layers + .iter() + .copied() + .zip(selected.iter().cloned()) + .collect::>(); + diagnostic_selected.sort_by_key(|(layer, _)| *layer); + let probe_split = + diagnostic_selected.partition_point(|(layer, _)| *layer < MOLMO2_VIT_PROBE_LAYER); + outputs.extend( + diagnostic_selected[..probe_split] + .iter() + .map(|(_, value)| value.clone()), + ); + outputs.extend(probe_rows); + outputs.extend( + diagnostic_selected[probe_split..] + .iter() + .map(|(_, value)| value.clone()), + ); outputs.extend([ selected_features, pool.gathered_masked, @@ -580,6 +638,39 @@ mod tests { ); } + #[test] + fn diagnostic_graph_probes_only_the_known_layer24_failure_row() { + let mut config = test_config(true); + config.layers = 25; + config.emitted_layers = 25; + config.selected_layers = vec![24, 18]; + config.static_crops = 1; + config.patches_per_crop = MOLMO2_VIT_PROBE_FLAT_ROW + 1; + config.position_count = config.patches_per_crop; + + let diagnostics = emit_molmo2_vision_diagnostics(&config); + assert_eq!( + diagnostics.matches("stablehlo.slice").count(), + 7, + "input, two norms, attention, residual, MLP, and output must each expose one row" + ); + assert_eq!( + diagnostics.matches("[0:1, 513:514, 0:8]").count(), + 7, + "the probe must stay bounded to the row containing flat failure index 591490" + ); + let return_values = diagnostics + .lines() + .find_map(|line| line.trim().strip_prefix("return ")) + .and_then(|line| line.split_once(" : ").map(|(values, _)| values)) + .expect("Molmo2 diagnostic graph must return named stage values"); + assert_eq!( + return_values.split(',').count(), + 23, + "the seven row probes must be the only new diagnostic transfers" + ); + } + #[test] fn projector_silu_preserves_mlx_sigmoid_multiply_rounding_order() { let mut builder = Builder::new(); diff --git a/src/lib/mlxcel-xla/src/molmo2_vision_runtime.rs b/src/lib/mlxcel-xla/src/molmo2_vision_runtime.rs index 9ad6ed33b..8823efa74 100644 --- a/src/lib/mlxcel-xla/src/molmo2_vision_runtime.rs +++ b/src/lib/mlxcel-xla/src/molmo2_vision_runtime.rs @@ -36,6 +36,11 @@ use crate::molmo2::Molmo2SafePooling; use crate::weights::{bf16_to_f32, f16_to_f32, f32_le_to_f32}; const ENTRY_NAME: &str = "molmo2_vision.main"; +#[cfg(feature = "diagnostics")] +const MOLMO2_VIT_PROBE_LAYER: usize = 24; +#[cfg(feature = "diagnostics")] +// 591490 = 513 * checkpoint hidden width 1152 + component 514. +const MOLMO2_VIT_PROBE_FLAT_ROW: usize = 513; #[derive(Debug, Clone, PartialEq)] pub struct Molmo2VisionProjection { @@ -572,9 +577,44 @@ fn diagnostic_stage_specs( active_shape: hidden_active.clone(), }, ]; + let mut selected_layers = config.selected_layers.clone(); + selected_layers.sort_unstable(); + let probe_split = selected_layers.partition_point(|layer| *layer < MOLMO2_VIT_PROBE_LAYER); + specs.extend( + selected_layers[..probe_split] + .iter() + .map(|layer| Molmo2DiagnosticStageSpec { + name: format!("vit.selected.{layer}"), + static_shape: hidden_static.clone(), + active_shape: hidden_active.clone(), + }), + ); + if config.emitted_layers > MOLMO2_VIT_PROBE_LAYER + && config.static_crops * config.patches_per_crop > MOLMO2_VIT_PROBE_FLAT_ROW + { + specs.extend( + [ + "input", + "attention_norm", + "attention", + "post_attention_residual", + "ffn_norm", + "mlp", + "output", + ] + .into_iter() + .map(|stage| Molmo2DiagnosticStageSpec { + name: format!( + "vit.probe.{}.row.{}.{}", + MOLMO2_VIT_PROBE_LAYER, MOLMO2_VIT_PROBE_FLAT_ROW, stage + ), + static_shape: vec![1, config.hidden], + active_shape: vec![1, config.hidden], + }), + ); + } specs.extend( - config - .selected_layers + selected_layers[probe_split..] .iter() .map(|layer| Molmo2DiagnosticStageSpec { name: format!("vit.selected.{layer}"), @@ -661,7 +701,7 @@ impl IreeMolmo2VisionDiagnosticProjector { &config, &mlir, "molmo2-vision-diagnostics", - Some("first-divergence-v2-projector-stages"), + Some("first-divergence-v3-layer24-row513"), )?; Ok(Self { module, config }) } diff --git a/src/vision/encoders/molmo2.rs b/src/vision/encoders/molmo2.rs index 3a6aa74ed..7dbd91961 100644 --- a/src/vision/encoders/molmo2.rs +++ b/src/vision/encoders/molmo2.rs @@ -27,6 +27,64 @@ use mlxcel_core::layers::{LayerNorm, Linear}; use mlxcel_core::weights::WeightMap; use mlxcel_core::{MlxArray, UniquePtr}; +#[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] +const MOLMO2_VIT_PROBE_LAYER: usize = 24; +#[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] +// The pinned actual failure was flat index 591490 at hidden width 1152: +// 591490 = 513 * 1152 + 514. Snapshot the whole row at producer boundaries. +const MOLMO2_VIT_PROBE_FLAT_ROW: usize = 513; + +#[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] +fn diagnostic_flat_row_snapshot( + value: &MlxArray, + tokens_per_crop: usize, + flat_row: usize, +) -> UniquePtr { + let shape = mlxcel_core::array_shape(value); + assert_eq!( + shape.len(), + 3, + "Molmo2 ViT probe requires [crop, token, hidden]" + ); + let crop = flat_row / tokens_per_crop; + let token = flat_row % tokens_per_crop; + assert!( + crop < usize::try_from(shape[0]).expect("non-negative Molmo2 crop count"), + "Molmo2 ViT probe row is outside the active crops" + ); + let hidden = shape[2]; + let row = mlxcel_core::slice( + value, + &[crop as i32, token as i32, 0], + &[crop as i32 + 1, token as i32 + 1, hidden], + ); + let row = mlxcel_core::reshape(&row, &[1, hidden]); + let row = mlxcel_core::astype(&row, mlxcel_core::dtype::FLOAT32); + mlxcel_core::eval(&row); + let raw = mlxcel_core::array_to_raw_bytes(&row); + let values = raw + .chunks_exact(std::mem::size_of::()) + .map(|chunk| f32::from_ne_bytes(chunk.try_into().expect("four-byte f32 probe value"))) + .collect::>(); + assert_eq!( + values.len(), + usize::try_from(hidden).expect("non-negative Molmo2 hidden width") + ); + mlxcel_core::from_slice_f32(&values, &[1, hidden]) +} + +#[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] +fn diagnostic_probe_stage( + stage: &'static str, + value: &MlxArray, + tokens_per_crop: usize, +) -> (&'static str, UniquePtr) { + ( + stage, + diagnostic_flat_row_snapshot(value, tokens_per_crop, MOLMO2_VIT_PROBE_FLAT_ROW), + ) +} + /// Hugging Face's `gelu_pytorch_tanh`, evaluated in F32 so the eager Molmo2 /// reference matches both the checkpoint's declared activation and StableHLO. fn gelu_pytorch_tanh(x: &MlxArray) -> UniquePtr { @@ -202,6 +260,43 @@ impl Molmo2VisionBlock { mlxcel_core::add(&h, &mlp_out) } + #[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] + fn forward_probe( + &self, + x: &MlxArray, + tokens_per_crop: usize, + ) -> ( + UniquePtr, + Vec<(&'static str, UniquePtr)>, + ) { + let mut stages = vec![diagnostic_probe_stage("input", x, tokens_per_crop)]; + let normed = self.attention_norm.forward(x); + stages.push(diagnostic_probe_stage( + "attention_norm", + &normed, + tokens_per_crop, + )); + let attn_out = self.attention.forward(&normed, None, None); + stages.push(diagnostic_probe_stage( + "attention", + &attn_out, + tokens_per_crop, + )); + let residual = mlxcel_core::add(x, &attn_out); + stages.push(diagnostic_probe_stage( + "post_attention_residual", + &residual, + tokens_per_crop, + )); + let normed = self.ffn_norm.forward(&residual); + stages.push(diagnostic_probe_stage("ffn_norm", &normed, tokens_per_crop)); + let mlp_out = self.feed_forward.forward(&normed); + stages.push(diagnostic_probe_stage("mlp", &mlp_out, tokens_per_crop)); + let output = mlxcel_core::add(&residual, &mlp_out); + stages.push(diagnostic_probe_stage("output", &output, tokens_per_crop)); + (output, stages) + } + fn from_weights( weights: &WeightMap, prefix: &str, @@ -259,6 +354,7 @@ struct Molmo2VitDiagnostics { patch_embedding: UniquePtr, position_embedding: UniquePtr, positioned_embedding: UniquePtr, + probe_rows: Vec<(&'static str, UniquePtr)>, } #[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] @@ -325,7 +421,7 @@ impl Molmo2VisionTransformer { let position_embedding = self.position_embedding(&patch_embedding, patch_h, patch_w); let mut x = mlxcel_core::add(&patch_embedding, &position_embedding); #[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] - let capture = CAPTURE.then(|| Molmo2VitDiagnostics { + let mut capture = CAPTURE.then(|| Molmo2VitDiagnostics { patch_embedding: mlxcel_core::copy( patch_embedding .as_ref() @@ -349,11 +445,37 @@ impl Molmo2VisionTransformer { x.as_ref() .expect("Molmo2 positioned embedding must be materialized"), ), + probe_rows: Vec::new(), }); let mut hidden_states = Vec::with_capacity(self.blocks.len()); - for block in &self.blocks { - x = block.forward(&x); + #[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] + let tokens_per_crop = + usize::try_from(patch_h * patch_w).expect("positive Molmo2 patch grid"); + for (layer, block) in self.blocks.iter().enumerate() { + #[cfg(not(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu")))] + let _ = layer; + #[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] + if CAPTURE + && layer == MOLMO2_VIT_PROBE_LAYER + && usize::try_from(mlxcel_core::array_shape(&x)[0]) + .expect("non-negative Molmo2 crop count") + * tokens_per_crop + > MOLMO2_VIT_PROBE_FLAT_ROW + { + let (output, probe_rows) = block.forward_probe(&x, tokens_per_crop); + x = output; + capture + .as_mut() + .expect("Molmo2 probe requires diagnostics capture") + .probe_rows = probe_rows; + } else { + x = block.forward(&x); + } + #[cfg(not(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu")))] + { + x = block.forward(&x); + } hidden_states.push(mlxcel_core::copy(&x)); } #[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] @@ -785,7 +907,26 @@ impl Molmo2VisionModel { tensor: encode.early_block, }, ]; - stages.extend(encode.selected_layers.into_iter().map(|(layer, tensor)| { + let mut selected_layers = encode.selected_layers; + selected_layers.sort_by_key(|(layer, _)| *layer); + let probe_split = + selected_layers.partition_point(|(layer, _)| *layer < MOLMO2_VIT_PROBE_LAYER); + stages.extend(selected_layers.drain(..probe_split).map(|(layer, tensor)| { + Molmo2VisionDiagnosticTensor { + name: format!("vit.selected.{layer}"), + tensor, + } + })); + stages.extend(encode.vit.probe_rows.into_iter().map(|(stage, tensor)| { + Molmo2VisionDiagnosticTensor { + name: format!( + "vit.probe.{}.row.{}.{}", + MOLMO2_VIT_PROBE_LAYER, MOLMO2_VIT_PROBE_FLAT_ROW, stage + ), + tensor, + } + })); + stages.extend(selected_layers.into_iter().map(|(layer, tensor)| { Molmo2VisionDiagnosticTensor { name: format!("vit.selected.{layer}"), tensor, @@ -971,6 +1112,8 @@ fn get_weight_copy(weights: &WeightMap, name: &str) -> Result>(); + let input = + mlxcel_core::from_slice_f32(&values, &[1, (MOLMO2_VIT_PROBE_FLAT_ROW + 1) as i32, 2]); + let row = diagnostic_flat_row_snapshot( + &input, + MOLMO2_VIT_PROBE_FLAT_ROW + 1, + MOLMO2_VIT_PROBE_FLAT_ROW, + ); + assert_eq!(mlxcel_core::array_shape(&row), vec![1, 2]); + let bytes = mlxcel_core::array_to_raw_bytes(&row); + let actual = bytes + .chunks_exact(4) + .map(|chunk| f32::from_ne_bytes(chunk.try_into().expect("four-byte f32"))) + .collect::>(); + assert_eq!( + actual, + vec![ + (MOLMO2_VIT_PROBE_FLAT_ROW * 2) as f32, + (MOLMO2_VIT_PROBE_FLAT_ROW * 2 + 1) as f32, + ] + ); + } } From aae31e77113513cc41dda75b00bb7558bab2a3a4 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 02:08:40 +0900 Subject: [PATCH 13/20] test(xla): probe Molmo2 selected layer 18 --- .../mlxcel-xla/src/emitter/molmo2_vision.rs | 8 ++- .../mlxcel-xla/src/molmo2_vision_runtime.rs | 4 +- .../src/molmo2_vision_runtime_tests.rs | 54 +++++++++++++++++++ src/vision/encoders/molmo2.rs | 2 +- 4 files changed, 63 insertions(+), 5 deletions(-) diff --git a/src/lib/mlxcel-xla/src/emitter/molmo2_vision.rs b/src/lib/mlxcel-xla/src/emitter/molmo2_vision.rs index ad4977729..846e3c756 100644 --- a/src/lib/mlxcel-xla/src/emitter/molmo2_vision.rs +++ b/src/lib/mlxcel-xla/src/emitter/molmo2_vision.rs @@ -17,7 +17,7 @@ use super::builder::{Builder, Ty, Val}; use super::molmo2_config::{Molmo2VisionConfig, Molmo2VisionWeightSpec}; -const MOLMO2_VIT_PROBE_LAYER: usize = 24; +const MOLMO2_VIT_PROBE_LAYER: usize = 18; // 591490 = 513 * checkpoint hidden width 1152 + component 514. const MOLMO2_VIT_PROBE_FLAT_ROW: usize = 513; @@ -639,7 +639,7 @@ mod tests { } #[test] - fn diagnostic_graph_probes_only_the_known_layer24_failure_row() { + fn diagnostic_graph_probes_the_first_selected_layer_failure_row() { let mut config = test_config(true); config.layers = 25; config.emitted_layers = 25; @@ -649,6 +649,10 @@ mod tests { config.position_count = config.patches_per_crop; let diagnostics = emit_molmo2_vision_diagnostics(&config); + assert_eq!( + MOLMO2_VIT_PROBE_LAYER, 18, + "the row probes must precede the first failing selected-layer comparison" + ); assert_eq!( diagnostics.matches("stablehlo.slice").count(), 7, diff --git a/src/lib/mlxcel-xla/src/molmo2_vision_runtime.rs b/src/lib/mlxcel-xla/src/molmo2_vision_runtime.rs index 8823efa74..17990173c 100644 --- a/src/lib/mlxcel-xla/src/molmo2_vision_runtime.rs +++ b/src/lib/mlxcel-xla/src/molmo2_vision_runtime.rs @@ -37,7 +37,7 @@ use crate::weights::{bf16_to_f32, f16_to_f32, f32_le_to_f32}; const ENTRY_NAME: &str = "molmo2_vision.main"; #[cfg(feature = "diagnostics")] -const MOLMO2_VIT_PROBE_LAYER: usize = 24; +const MOLMO2_VIT_PROBE_LAYER: usize = 18; #[cfg(feature = "diagnostics")] // 591490 = 513 * checkpoint hidden width 1152 + component 514. const MOLMO2_VIT_PROBE_FLAT_ROW: usize = 513; @@ -701,7 +701,7 @@ impl IreeMolmo2VisionDiagnosticProjector { &config, &mlir, "molmo2-vision-diagnostics", - Some("first-divergence-v3-layer24-row513"), + Some("first-divergence-v4-layer18-row513"), )?; Ok(Self { module, config }) } diff --git a/src/lib/mlxcel-xla/src/molmo2_vision_runtime_tests.rs b/src/lib/mlxcel-xla/src/molmo2_vision_runtime_tests.rs index 999ddf7df..99f637882 100644 --- a/src/lib/mlxcel-xla/src/molmo2_vision_runtime_tests.rs +++ b/src/lib/mlxcel-xla/src/molmo2_vision_runtime_tests.rs @@ -31,3 +31,57 @@ fn compiler_identity_changes_when_same_path_and_version_bytes_change() { std::fs::remove_file(path).ok(); assert_ne!(first, second); } + +#[cfg(feature = "diagnostics")] +#[test] +fn layer18_row_probes_precede_the_first_failing_selected_stage() { + let mut config = Molmo2VisionConfig::from_json_strs( + &serde_json::json!({ + "model_type":"molmo2","image_patch_id":151938, + "vit_config":{"hidden_size":8,"intermediate_size":16,"num_attention_heads":2, + "head_dim":4,"num_hidden_layers":2,"image_default_input_size":[28,28], + "image_patch_size":14,"image_num_pos":4,"layer_norm_eps":1e-6, + "hidden_act":"gelu_pytorch_tanh"}, + "adapter_config":{"hidden_size":8,"intermediate_size":12,"text_hidden_size":10, + "num_attention_heads":2,"head_dim":4,"vit_layers":[0,1], + "pooling_attention_mask":true} + }) + .to_string(), + &serde_json::json!({"patch_size":14,"max_crops":1,"overlap_margins":[0,0], + "pooling_size":[2,2],"size":{"height":28,"width":28}}) + .to_string(), + ) + .unwrap(); + config.layers = 25; + config.emitted_layers = 25; + config.selected_layers = vec![24, 18]; + config.static_crops = 1; + config.patches_per_crop = MOLMO2_VIT_PROBE_FLAT_ROW + 1; + config.position_count = config.patches_per_crop; + + let names = diagnostic_stage_specs(&config, 1, 1) + .into_iter() + .map(|spec| spec.name) + .collect::>(); + let first_probe = names + .iter() + .position(|name| name == "vit.probe.18.row.513.input") + .expect("layer 18 input probe"); + let selected18 = names + .iter() + .position(|name| name == "vit.selected.18") + .expect("selected layer 18"); + assert_eq!( + &names[first_probe..selected18], + &[ + "vit.probe.18.row.513.input", + "vit.probe.18.row.513.attention_norm", + "vit.probe.18.row.513.attention", + "vit.probe.18.row.513.post_attention_residual", + "vit.probe.18.row.513.ffn_norm", + "vit.probe.18.row.513.mlp", + "vit.probe.18.row.513.output", + ], + "all row-local producer boundaries must be compared before fail-fast reaches selected18" + ); +} diff --git a/src/vision/encoders/molmo2.rs b/src/vision/encoders/molmo2.rs index 7dbd91961..765b0567d 100644 --- a/src/vision/encoders/molmo2.rs +++ b/src/vision/encoders/molmo2.rs @@ -28,7 +28,7 @@ use mlxcel_core::weights::WeightMap; use mlxcel_core::{MlxArray, UniquePtr}; #[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] -const MOLMO2_VIT_PROBE_LAYER: usize = 24; +const MOLMO2_VIT_PROBE_LAYER: usize = 18; #[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] // The pinned actual failure was flat index 591490 at hidden width 1152: // 591490 = 513 * 1152 + 514. Snapshot the whole row at producer boundaries. From 0e846fc43a31c1febb68c104ff447d7104e28f9d Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 09:21:12 +0900 Subject: [PATCH 14/20] fix(xla): mark Molmo2 artifact unqualified Keep the rebased Molmo2 runtime on the explicit legacy artifact path until the selected-layer numeric divergence is resolved.\n\nRefs #871 --- src/lib/mlxcel-xla/src/molmo2_vision_runtime.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/lib/mlxcel-xla/src/molmo2_vision_runtime.rs b/src/lib/mlxcel-xla/src/molmo2_vision_runtime.rs index 17990173c..df99063ef 100644 --- a/src/lib/mlxcel-xla/src/molmo2_vision_runtime.rs +++ b/src/lib/mlxcel-xla/src/molmo2_vision_runtime.rs @@ -444,7 +444,7 @@ fn compile_vision_module( let graph_identity = diagnostic_identity .map(|identity| format!("{};diagnostics={identity}", config.fingerprint())) .unwrap_or_else(|| config.fingerprint()); - let contract = AuxiliaryArtifactContract::new( + let contract = AuxiliaryArtifactContract::new_legacy_unqualified( ENTRY_NAME, format!( "{graph_identity};checkpoint_schema_sha256={}", From ce1891fcae36276ef14fb3ac85196aad39dc00e8 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 13:12:23 +0900 Subject: [PATCH 15/20] refactor(xla): share Molmo2 numeric decompositions Use the merged contract-sensitive numeric helpers for Molmo2 tanh GELU, SiLU, and stable softmax so production emission and bounded probes cannot drift. Preserve the existing StableHLO operation order, including projector x * sigmoid(x); keep the rank-3 LayerNorm local because the shared helper is intentionally row-wise. Validation: Molmo2 diagnostics unit tests (21 passed), shared numeric helper tests (13 passed), synthetic Molmo2 parity tests (5 passed), rustfmt check, and git diff check. Refs #871 --- src/lib/mlxcel-xla/src/emitter/mod.rs | 2 +- .../mlxcel-xla/src/emitter/molmo2_vision.rs | 55 +------------------ 2 files changed, 3 insertions(+), 54 deletions(-) diff --git a/src/lib/mlxcel-xla/src/emitter/mod.rs b/src/lib/mlxcel-xla/src/emitter/mod.rs index 65d203ba9..115cb6095 100644 --- a/src/lib/mlxcel-xla/src/emitter/mod.rs +++ b/src/lib/mlxcel-xla/src/emitter/mod.rs @@ -57,9 +57,9 @@ mod gemma3n_schema; mod gemma3n_weights; mod model; mod moe; -pub(crate) mod numeric_ops; mod molmo2_config; mod molmo2_vision; +pub(crate) mod numeric_ops; mod phi4_audio; mod qwen2_vl; mod rope; diff --git a/src/lib/mlxcel-xla/src/emitter/molmo2_vision.rs b/src/lib/mlxcel-xla/src/emitter/molmo2_vision.rs index 846e3c756..f7e5ffbfc 100644 --- a/src/lib/mlxcel-xla/src/emitter/molmo2_vision.rs +++ b/src/lib/mlxcel-xla/src/emitter/molmo2_vision.rs @@ -16,6 +16,7 @@ use super::builder::{Builder, Ty, Val}; use super::molmo2_config::{Molmo2VisionConfig, Molmo2VisionWeightSpec}; +use super::numeric_ops::{silu, stable_softmax, tanh_gelu}; const MOLMO2_VIT_PROBE_LAYER: usize = 18; // 591490 = 513 * checkpoint hidden width 1152 + component 514. @@ -114,61 +115,9 @@ fn layer_norm(builder: &mut Builder, value: &Val, weight: &Val, bias: &Val, epsi builder.add(&normalized, &bias) } -fn tanh_gelu(builder: &mut Builder, value: &Val) -> Val { - let shape = value.ty.shape.clone(); - let half = builder.const_f32(0.5); - let half = builder.broadcast(&half, &[], shape.clone()); - let one = builder.const_f32(1.0); - let one = builder.broadcast(&one, &[], shape.clone()); - let coefficient = builder.const_f32(0.044_715); - let coefficient = builder.broadcast(&coefficient, &[], shape.clone()); - let scale = builder.const_f32(0.797_884_6); - let scale = builder.broadcast(&scale, &[], shape); - let squared = builder.multiply(value, value); - let cubed = builder.multiply(&squared, value); - let nonlinear = builder.multiply(&coefficient, &cubed); - let inner = builder.add(value, &nonlinear); - let scaled = builder.multiply(&scale, &inner); - let tanh = builder.tanh(&scaled); - let cdf = builder.add(&one, &tanh); - let half_value = builder.multiply(value, &half); - builder.multiply(&half_value, &cdf) -} - -fn silu(builder: &mut Builder, value: &Val) -> Val { - let shape = value.ty.shape.clone(); - let one = builder.const_f32(1.0); - let one = builder.broadcast(&one, &[], shape); - let negative = builder.negate(value); - let exponential = builder.exponential(&negative); - let denominator = builder.add(&one, &exponential); - // Preserve MLX's `x * sigmoid(x)` operation order. Reassociating this as - // `x / (1 + exp(-x))` changes F32 rounding before the wide w2 projection. - let sigmoid = builder.divide(&one, &denominator); - builder.multiply(value, &sigmoid) -} - fn softmax_last(builder: &mut Builder, scores: &Val) -> Val { let axis = scores.ty.shape.len() - 1; - let leading = scores.ty.shape[..axis].to_vec(); - let negative_infinity = builder.const_f32(f32::NEG_INFINITY); - let maximum = builder.reduce_max(scores, axis, &negative_infinity); - let maximum = builder.broadcast( - &maximum, - &(0..axis).collect::>(), - scores.ty.shape.clone(), - ); - let shifted = builder.subtract(scores, &maximum); - let exponentials = builder.exponential(&shifted); - let zero = builder.const_f32(0.0); - let denominator = builder.reduce_add(&exponentials, axis, &zero); - let denominator = builder.broadcast( - &denominator, - &(0..axis).collect::>(), - scores.ty.shape.clone(), - ); - debug_assert_eq!(denominator.ty.shape[..axis], leading); - builder.divide(&exponentials, &denominator) + stable_softmax(builder, scores, axis) } fn selected_slot(selected_layers: &[usize], layer: usize) -> Option { From 2df457f2c846268624ef99b639885aa2763f4192 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Tue, 28 Jul 2026 02:32:28 +0900 Subject: [PATCH 16/20] fix(xla): configure local-task threads before Molmo2 IREE creation The Molmo2 vision parity gate created its IREE diagnostic projector without first applying the diagnostics-only local-task thread configuration added in #945. `configure_diagnostic_local_task_threads` had no callers on this branch, so the gate aborted before any comparison ran: xla_aux_create failed (status 13): iree/base/threading/thread_pthreads.c:159: INTERNAL; thread creation failed with 22 IREE parses its process-global flag registry when the first instance is created, so the call has to precede the projector load to have any effect. This is the same defect fixed for the Molmo v1 gate. Refs #871 --- tests/molmo2_xla_vision_parity.rs | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/tests/molmo2_xla_vision_parity.rs b/tests/molmo2_xla_vision_parity.rs index 6b73c92d3..5ab99f3d8 100644 --- a/tests/molmo2_xla_vision_parity.rs +++ b/tests/molmo2_xla_vision_parity.rs @@ -341,6 +341,14 @@ fn real_checkpoint_mlx_iree_vision_and_scatter_parity() -> Result<()> { )); } + // Bound the local-task worker topology and stack before the first IREE + // instance exists. IREE parses its process-global flag registry on that + // first creation, so this has to run ahead of the projector load to take + // effect. Without it, `local-task` worker creation fails on this host with + // `thread creation failed with 22` out of `thread_pthreads.c`, because + // IREE's exact PTHREAD_STACK_MIN request is rejected. Diagnostics-only: + // the production runtime never applies these flags. + mlxcel_xla::configure_diagnostic_local_task_threads().map_err(anyhow::Error::msg)?; let mut projector = with_progress("IREE diagnostic compile/load", || { IreeMolmo2VisionDiagnosticProjector::load(&model, &device).map_err(anyhow::Error::msg) })?; From ed881bd75c8c74806106ffb7b26b828e2b2ed3f2 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 21 Aug 2026 16:48:09 +0900 Subject: [PATCH 17/20] fix(xla): gate the diagnostics-only weight-loader re-export This branch promotes `load_weights_from_dir_with_filter` to `pub(crate)` and re-exports it from `src/models/mod.rs` so the Molmo2 eager vision reference can load filtered vision tensors. That reference, `load_molmo2_xla_vision_reference` in `src/loading/vlm_special.rs`, is gated on `any(test, feature = "xla-diagnostics", feature = "xla-diagnostics-cpu")`, so an unconditional re-export is dead in a default build and `-D warnings` rejects it. Gate the re-export with the same predicate its only consumer carries, so the two move together. --- src/models/mod.rs | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/models/mod.rs b/src/models/mod.rs index 3c4f014ed..26de87db8 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -277,9 +277,13 @@ pub use rwkv7::Rwkv7; pub(crate) use sanitize::{ Gemma4WeightBacking, load_gemma4_text_weights_with_backing, load_gemma4_unified_weights_with_backing, load_gemma4_vlm_weights_with_backing, - load_weights_from_dir_with_filter, sanitize_gemma4_nvfp4_weights, - strip_gemma4_kv_shared_weights, + sanitize_gemma4_nvfp4_weights, strip_gemma4_kv_shared_weights, }; +// The only consumer outside `sanitize` is the diagnostics-gated Molmo2 vision +// reference loader, so an unconditional re-export is dead in a default build +// and `-D warnings` rejects it. +#[cfg(any(test, feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] +pub(crate) use sanitize::load_weights_from_dir_with_filter; pub use sanitize::{ convert_bf16_weights, convert_bf16_weights_with_keep, gemma3n_language_mlp_bf16_key, load_and_sanitize_weights, load_text_weights, sanitize_config_json, sanitize_tied_embeddings, From 272fbc27af836f5c4ca59fa3d99089d0b831bd2c Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 21 Aug 2026 16:48:22 +0900 Subject: [PATCH 18/20] fix(xla): reject a context capacity that can never admit one image The OpenXLA context capacity is a static graph shape fixed when the engine is built, and it defaults to 256. A Molmo2 image expands to far more than that: 424 logical tokens for a square image and up to 1834 for a tall one on the pinned 4B checkpoint. Serving images from that default therefore ran the whole vision tower on every request and only then failed at admission, wasting the work and reporting the problem per request rather than once at startup. `xla_image_context_floor` derives the worst case from `config.json` and `preprocessor_config.json` alone, with no weights, processor, or vision tower loaded, because the capacity has to be chosen before any preprocessor exists. It walks every admissible tiling rather than assuming the squarest or largest one: a column token per pooled row makes tall tilings cost more than wide ones of equal area, so `(max_crops, 1)` is the maximum. `ensure_xla_image_context_capacity` turns that into a startup check on both the server and CLI paths. A capacity the operator pinned is honored verbatim, since text-only serving from a VLM checkpoint is a real workload and the capacity is also the length every decode step attends over, so a smaller graph may be a deliberate throughput choice. Measured on this checkpoint: 3.18 tok/s at 256, 2.17 at 1024, 1.41 at 2048. The rejection message carries the derived requirement, that cost, and both remedies. A worker that logged and returned left the process alive with a model that never loads: `loaded` stays false, nothing observes the dead thread, and the server parks forever. `exit_on_worker_startup_failure` ends the process instead, extending the fail-fast posture this file already documents for panics to startup failures that can never resolve. It exits rather than aborts, because a misconfiguration is not a broken invariant. `HostPreprocessorError` also stops calling every checkpoint LLaVA; Qwen2-VL, Muse Glimmer, and Molmo2 all share these variants now. --- src/backend/xla.rs | 18 ++- src/lib.rs | 5 +- src/multimodal/host_preprocessor.rs | 133 +++++++++++++++++++++- src/multimodal/host_preprocessor_tests.rs | 109 +++++++++++++++++- src/server/model_worker.rs | 39 +++++-- src/worker_failfast.rs | 24 ++++ 6 files changed, 308 insertions(+), 20 deletions(-) diff --git a/src/backend/xla.rs b/src/backend/xla.rs index 4d6f473f5..9d4f1da08 100644 --- a/src/backend/xla.rs +++ b/src/backend/xla.rs @@ -109,8 +109,22 @@ impl ComputeBackend for XlaBackend { // session owns its own KV and samples greedily on-device. let image_preprocessor = crate::load_xla_image_preprocessor(model_path) .map_err(|error| anyhow::anyhow!("OpenXLA image preprocessor load failed: {error}"))?; - let session = - XlaInferenceSession::load(model_path, num_layers).map_err(|e| anyhow::anyhow!(e))?; + // Match the server path: reject a graph that could never admit one + // image before generation starts, instead of after the vision tower has + // already run on the first image request. + let context_capacity = + mlxcel_xla::context_capacity_from_env().map_err(|e| anyhow::anyhow!(e))?; + crate::ensure_xla_image_context_capacity( + model_path, + context_capacity, + std::env::var_os(mlxcel_xla::CONTEXT_CAPACITY_ENV).is_some(), + )?; + let session = XlaInferenceSession::load_with_context_capacity( + model_path, + num_layers, + context_capacity, + ) + .map_err(|e| anyhow::anyhow!(e))?; Ok(Session::xla(session, image_preprocessor)) } diff --git a/src/lib.rs b/src/lib.rs index f4f3a35f1..ccb2445fc 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -80,8 +80,9 @@ pub use multimodal::host_preprocessor::LlavaHostReferenceCapture; #[cfg(feature = "xla-iree")] pub use multimodal::host_preprocessor::LlavaIreeHostPreprocessor; pub use multimodal::host_preprocessor::{ - FakeHostMultimodalPreprocessor, HostMultimodalPreprocessor, HostPreprocessorError, - LlavaHostPreprocessor, XlaVisionBackend, load_xla_image_preprocessor, + CONTEXT_CAPACITY_ENV, FakeHostMultimodalPreprocessor, HostMultimodalPreprocessor, + HostPreprocessorError, LlavaHostPreprocessor, XlaVisionBackend, + ensure_xla_image_context_capacity, load_xla_image_preprocessor, xla_image_context_floor, }; pub use multimodal::{ falcon_ocr_prompt, internvl_prompt, kimi_vl_prompt, locateanything_prompt, minicpmo_prompt, diff --git a/src/multimodal/host_preprocessor.rs b/src/multimodal/host_preprocessor.rs index 1aa9f3caa..e1e0ffe73 100644 --- a/src/multimodal/host_preprocessor.rs +++ b/src/multimodal/host_preprocessor.rs @@ -258,6 +258,135 @@ pub fn load_xla_image_preprocessor( } } +/// Worst-case logical prompt length one image contributes, from config alone. +/// +/// The static StableHLO context shape is fixed when the engine is built, which +/// on the server path happens before any preprocessor exists, so this answers +/// the question from `config.json` and `preprocessor_config.json` without +/// loading weights, a processor, or a vision tower. +/// +/// `None` means "not derived for this checkpoint", not "no images": a family +/// without a formula here, or a config this cannot read, falls through without +/// a capacity guard rather than inventing a number. +/// +/// # Molmo2 +/// +/// The prompt carries one pooled token per pooled patch, plus one column token +/// per pooled row, for a low-resolution crop and a high-resolution tiling: +/// +/// ```text +/// tokens = lo_h * (lo_w + 1) + hi_h * (hi_w + 1) + 4 +/// ``` +/// +/// The low-resolution crop is always one tile. The high-resolution tiling is +/// whichever `(rows, cols)` with `rows * cols <= max_crops` the processor picks +/// for the image, so the worst case is the admissible tiling that maximizes the +/// token count. That is the tallest one, `(max_crops, 1)`, because every extra +/// row also adds a column token. On the pinned 4B checkpoint (`max_crops = 8`, +/// 378px crops, 14px patches, 2x2 pooling) a square image expands to 424 tokens +/// while a tall one reaches 1834, so sizing on the observed square case would +/// still reject ordinary photographs. +#[must_use] +pub fn xla_image_context_floor(model_path: &Path) -> Option { + match crate::models::get_model_type(model_path) { + Ok(crate::models::ModelType::Molmo2VLM) => molmo2_image_context_floor(model_path), + // Other qualified families expand images too, but their worst case is + // not derived here yet. Reporting `None` keeps this guard silent for + // them instead of guarding with a Molmo2-shaped guess. + _ => None, + } +} + +/// Name of the variable that selects the static OpenXLA context shape. +/// +/// Duplicated from `mlxcel_xla` so this guard, and its tests, still compile in +/// a build without the `xla-backend` crate. The two are kept in step by +/// `the_capacity_variable_name_matches_the_xla_crate`. +pub const CONTEXT_CAPACITY_ENV: &str = "MLXCEL_XLA_CONTEXT_CAPACITY"; + +/// Reject a graph too small to ever admit one image from this checkpoint. +/// +/// Returns `Ok(())` when no floor is derived, when the graph already fits, or +/// when the operator pinned the capacity themselves. Pinning is treated as a +/// decision, not a mistake: text-only serving from a VLM checkpoint is a real +/// workload, and the capacity is also what every decode step attends over, so +/// a larger graph is a throughput cost the operator may be declining on +/// purpose. What this stops is the silent case, where an unset default admits +/// no image at all and every image request runs its vision tower and only then +/// fails at admission. +/// +/// # Errors +/// +/// Returns a message naming the derived requirement and the variable to set. +pub fn ensure_xla_image_context_capacity( + model_path: &Path, + context_capacity: usize, + operator_pinned: bool, +) -> Result<(), HostPreprocessorError> { + if operator_pinned { + return Ok(()); + } + let Some(floor) = xla_image_context_floor(model_path) else { + return Ok(()); + }; + if context_capacity >= floor { + return Ok(()); + } + Err(HostPreprocessorError::InvalidConfig(format!( + "this checkpoint expands one image into up to {floor} tokens, which the default OpenXLA \ + context capacity of {context_capacity} cannot admit, so every image request would run \ + its vision tower and then fail. Set {env}={floor} or higher to serve images (the \ + capacity is also the length every decode step attends over, so a larger graph costs \ + throughput), or set {env}={context_capacity} explicitly to keep this graph for \ + text-only serving.", + env = CONTEXT_CAPACITY_ENV, + ))) +} + +/// Worst-case Molmo2 image expansion, read from `preprocessor_config.json`. +fn molmo2_image_context_floor(model_path: &Path) -> Option { + let raw = std::fs::read_to_string(model_path.join("preprocessor_config.json")).ok()?; + let config: serde_json::Value = serde_json::from_str(&raw).ok()?; + + let max_crops = usize::try_from(config.get("max_crops")?.as_u64()?).ok()?; + let patch_size = usize::try_from(config.get("patch_size")?.as_u64()?).ok()?; + let size = config.get("size")?; + let crop_height = usize::try_from(size.get("height")?.as_u64()?).ok()?; + let crop_width = usize::try_from(size.get("width")?.as_u64()?).ok()?; + let pooling = config.get("pooling_size")?.as_array()?; + let pool_h = usize::try_from(pooling.first()?.as_u64()?).ok()?; + let pool_w = usize::try_from(pooling.get(1)?.as_u64()?).ok()?; + if max_crops == 0 || patch_size == 0 || pool_h == 0 || pool_w == 0 { + return None; + } + + let patch_rows = crop_height.checked_div(patch_size)?; + let patch_cols = crop_width.checked_div(patch_size)?; + let pooled = |patches: usize, pool: usize| patches.div_ceil(pool); + + let lo_h = pooled(patch_rows, pool_h); + let lo_w = pooled(patch_cols, pool_w); + let low_res = lo_h.checked_mul(lo_w.checked_add(1)?)?; + + // Every admissible tiling, because the maximum is not always the squarest + // or the largest-area one: a column token per row makes tall tilings cost + // more than wide ones of the same area. + let mut high_res = 0usize; + for rows in 1..=max_crops { + for cols in 1..=max_crops { + if rows.checked_mul(cols)? > max_crops { + continue; + } + let hi_h = pooled(patch_rows.checked_mul(rows)?, pool_h); + let hi_w = pooled(patch_cols.checked_mul(cols)?, pool_w); + high_res = high_res.max(hi_h.checked_mul(hi_w.checked_add(1)?)?); + } + } + + // The four framing tokens are the low-res start/end and high-res start/end. + low_res.checked_add(high_res)?.checked_add(4) +} + fn load_llava_host_preprocessor_boxed( model_path: &Path, ) -> Result>, HostPreprocessorError> { @@ -1142,9 +1271,9 @@ pub enum HostPreprocessorError { Placeholder(#[from] ImageTokenBlockError), #[error("incompatible multimodal family: expected LLaVA, got {actual}")] FamilyMismatch { actual: String }, - #[error("invalid LLaVA host-preprocessor config: {0}")] + #[error("invalid host-preprocessor config: {0}")] InvalidConfig(String), - #[error("failed to load LLaVA host-preprocessor weights: {0}")] + #[error("failed to load host-preprocessor weights: {0}")] WeightLoad(String), #[error("IREE vision backend failed: {0}")] Iree(String), diff --git a/src/multimodal/host_preprocessor_tests.rs b/src/multimodal/host_preprocessor_tests.rs index 1b6120d6f..968353e9c 100644 --- a/src/multimodal/host_preprocessor_tests.rs +++ b/src/multimodal/host_preprocessor_tests.rs @@ -18,9 +18,11 @@ use image::DynamicImage; use mlxcel_core::dtype; use super::{ - FakeHostMultimodalPreprocessor, HostMultimodalPreprocessor, HostPreprocessorError, - XlaVisionBackend, XlaVisionBackendPolicy, export_llava_prefill, export_mlx_tensor, + CONTEXT_CAPACITY_ENV, FakeHostMultimodalPreprocessor, HostMultimodalPreprocessor, + HostPreprocessorError, XlaVisionBackend, XlaVisionBackendPolicy, + ensure_xla_image_context_capacity, export_llava_prefill, export_mlx_tensor, export_qwen2_vl_prefill, load_xla_image_preprocessor, validate_processor_shape, + xla_image_context_floor, }; use crate::multimodal::vlm_prompt::ImageTokenBlockError; use crate::vision::merge::merge_llava; @@ -350,3 +352,106 @@ fn qwen2_vl_export_rejects_video_and_cross_image_run_drift() { HostPreprocessorError::InvalidConfig(_) )); } + +/// The pinned 4B checkpoint's geometry, so the expected numbers below are the +/// ones a real Molmo2 load would produce rather than invented ones. +fn molmo2_checkpoint_dir() -> tempfile::TempDir { + let dir = tempfile::tempdir().unwrap(); + std::fs::write( + dir.path().join("config.json"), + r#"{"model_type":"molmo2","image_patch_id":151938}"#, + ) + .unwrap(); + std::fs::write( + dir.path().join("preprocessor_config.json"), + r#"{"max_crops":8,"patch_size":14,"pooling_size":[2,2],"size":{"height":378,"width":378}}"#, + ) + .unwrap(); + dir +} + +#[test] +fn molmo2_image_floor_is_the_tallest_tiling_not_the_square_one() { + let dir = molmo2_checkpoint_dir(); + let floor = xla_image_context_floor(dir.path()).expect("Molmo2 geometry must derive a floor"); + // 8x1 tiling: low-res 14*(14+1)=210, high-res 108*(14+1)=1620, framing 4. + assert_eq!(floor, 1834); + // A square image reaches only 424 (210 + 210 + 4), which is what the + // 224x224 fixture produces. Sizing on that would reject tall photographs, + // so the floor must be strictly larger. + assert!( + floor > 424, + "floor {floor} must exceed the square-image case" + ); +} + +#[test] +fn a_checkpoint_without_a_derived_formula_has_no_floor() { + let dir = tempfile::tempdir().unwrap(); + std::fs::write( + dir.path().join("config.json"), + r#"{"model_type":"qwen2_vl"}"#, + ) + .unwrap(); + assert_eq!(xla_image_context_floor(dir.path()), None); + // No floor means no guard, not a rejection. + assert!(ensure_xla_image_context_capacity(dir.path(), 256, false).is_ok()); +} + +#[test] +fn a_missing_preprocessor_config_reports_no_floor_instead_of_guessing() { + let dir = tempfile::tempdir().unwrap(); + std::fs::write( + dir.path().join("config.json"), + r#"{"model_type":"molmo2","image_patch_id":151938}"#, + ) + .unwrap(); + assert_eq!(xla_image_context_floor(dir.path()), None); +} + +#[test] +fn the_default_capacity_is_rejected_with_the_derived_requirement() { + let dir = molmo2_checkpoint_dir(); + let error = ensure_xla_image_context_capacity(dir.path(), 256, false) + .err() + .expect("a graph that cannot admit any image must fail at startup"); + let HostPreprocessorError::InvalidConfig(message) = &error else { + panic!("expected a configuration error, got {error:?}"); + }; + assert!( + message.contains("1834"), + "the message must carry the derived requirement: {message}" + ); + assert!( + message.contains(CONTEXT_CAPACITY_ENV), + "the message must name the variable to set: {message}" + ); + assert!( + message.contains("256"), + "the message must show the capacity that was rejected: {message}" + ); +} + +#[test] +fn an_operator_pinned_capacity_is_never_second_guessed() { + let dir = molmo2_checkpoint_dir(); + // Text-only serving from a VLM checkpoint is a real workload, and capacity + // is what every decode step attends over, so a deliberately small graph + // must start. + assert!(ensure_xla_image_context_capacity(dir.path(), 256, true).is_ok()); +} + +#[test] +fn a_graph_that_fits_the_worst_case_image_starts_without_pinning() { + let dir = molmo2_checkpoint_dir(); + assert!(ensure_xla_image_context_capacity(dir.path(), 1834, false).is_ok()); + assert!(ensure_xla_image_context_capacity(dir.path(), 1833, false).is_err()); +} + +#[cfg(feature = "xla-backend")] +#[test] +fn the_capacity_variable_name_matches_the_xla_crate() { + // The name is duplicated so this guard compiles without the xla crate; this + // catches a rename on either side. + assert_eq!(CONTEXT_CAPACITY_ENV, mlxcel_xla::CONTEXT_CAPACITY_ENV); +} diff --git a/src/server/model_worker.rs b/src/server/model_worker.rs index 49534d944..2afbb779a 100644 --- a/src/server/model_worker.rs +++ b/src/server/model_worker.rs @@ -1106,12 +1106,27 @@ pub(crate) fn spawn_xla_model_worker( .unwrap_or_else(|_| mlxcel_xla::default_device().to_string()); let context_capacity = match mlxcel_xla::context_capacity_from_env() { Ok(value) => value, - Err(err) => { - tracing::error!("Failed to configure the OpenXLA engine: {err}"); - return; - } + Err(err) => crate::worker_failfast::exit_on_worker_startup_failure( + "model-worker-xla", + &format!("failed to configure the OpenXLA engine: {err}"), + ), }; + // The capacity is a static graph shape, so this has to be decided + // before the engine is built and cannot be revisited per request. + let operator_pinned = std::env::var_os(mlxcel_xla::CONTEXT_CAPACITY_ENV).is_some(); + if let Err(err) = crate::ensure_xla_image_context_capacity( + &model_path, + context_capacity, + operator_pinned, + ) { + crate::worker_failfast::exit_on_worker_startup_failure( + "model-worker-xla", + &err.to_string(), + ); + } tracing::info!( + context_capacity, + capacity_pinned_by_operator = operator_pinned, "Model worker thread starting (OpenXLA continuous batching, B_max={b_max}, \ context_capacity={context_capacity}, device={device}), loading model..." ); @@ -1131,10 +1146,10 @@ pub(crate) fn spawn_xla_model_worker( let load_start = Instant::now(); let tokenizer = match crate::tokenizer::load_tokenizer(&model_path) { Ok(t) => t, - Err(err) => { - tracing::error!("Failed to load tokenizer for the OpenXLA backend: {err}"); - return; - } + Err(err) => crate::worker_failfast::exit_on_worker_startup_failure( + "model-worker-xla", + &format!("failed to load the tokenizer for the OpenXLA backend: {err}"), + ), }; let engine = match mlxcel_xla::XlaBatchEngine::load_with_context_capacity( &model_path, @@ -1143,10 +1158,10 @@ pub(crate) fn spawn_xla_model_worker( context_capacity, ) { Ok(engine) => engine, - Err(err) => { - tracing::error!("Failed to load the OpenXLA engine: {err}"); - return; - } + Err(err) => crate::worker_failfast::exit_on_worker_startup_failure( + "model-worker-xla", + &format!("failed to load the OpenXLA engine: {err}"), + ), }; let mut worker = match crate::server::batch::XlaServeWorker::new( engine, diff --git a/src/worker_failfast.rs b/src/worker_failfast.rs index 7729cb113..fe76be078 100644 --- a/src/worker_failfast.rs +++ b/src/worker_failfast.rs @@ -61,3 +61,27 @@ where } } } + +/// Terminate the process on a worker startup failure that can never resolve. +/// +/// A core worker thread that logs and returns leaves the process alive with a +/// model that never loads: `loaded` is never set, nothing observes the dead +/// thread, and the server parks forever answering no requests. For a startup +/// failure that is always a misconfiguration (an unusable context capacity, an +/// unloadable engine or tokenizer), parking is the worst outcome, because an +/// operator watching for readiness sees a hang rather than the reason. +/// +/// This exits rather than [`std::process::abort`]: a configuration error is not +/// a broken invariant, so it should not look like a crash or leave a core dump. +/// The message is already logged by the caller with its own context; this adds +/// the thread label and the exit. +// Only the OpenXLA worker takes this path today, so an unconditional +// definition is dead in a default build and `-D warnings` rejects it. +#[cfg(feature = "xla-backend")] +pub(crate) fn exit_on_worker_startup_failure(label: &str, reason: &str) -> ! { + tracing::error!( + target: "mlxcel::worker", + "core worker thread '{label}' cannot start: {reason}" + ); + std::process::exit(1); +} From 22a753c922d0449268e59a40e0fef413f172c926 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 21 Aug 2026 16:48:33 +0900 Subject: [PATCH 19/20] test(xla): gate Molmo2 vision parity on a scale-aware tolerance The gate ran one 224x224 solid-color fixture, whose `1x1` tiling uses a single crop and never reaches the high-resolution tiling, pooling offsets, or multi-crop merge. Four ordinary photographs covering `2x2`, `4x1`, `1x4`, and `2x3` tilings all failed, every one of them at `projector.output_all` and nowhere earlier. The failures were the metric, not the math. That tensor carries values around 2e4, where one f32 ULP is already about 0.004, so the fixed `max_abs = 0.05` bound asked a 2000-term dot product to match to within tens of ULPs regardless of accumulation order. The four images agreed to within 0.001% relatively, put 4 to 12 elements out of ~2e6 past the bound, and held rms at 5 to 9% of its budget. The fixture passed at 0.0498 against 0.05, which was luck: it is the smallest and least textured input available. The contract is now `|a - b| <= atol + rtol * |b|` with `atol` unchanged at 0.05 and `rtol` at 1e-5, so the absolute bound still floors the comparison near zero and nothing that was strict before becomes loose. All five images now pass using at most 41% of the tolerance, rather than one passing at 99.6% and four failing. Blunting is checked, not assumed: swapped selected layers are rejected at 1158x the tolerance and a 2% error on a large value at 1693x, against 0.41x for the worst real measurement. The comparison also reports the values behind the worst gap, its relative size, and how many elements exceed the bound, because a single absolute maximum cannot separate one outlier from broad drift. --- tests/molmo2_xla_vision_parity.rs | 214 ++++++++++++++++++++++++++++-- 1 file changed, 203 insertions(+), 11 deletions(-) diff --git a/tests/molmo2_xla_vision_parity.rs b/tests/molmo2_xla_vision_parity.rs index 5ab99f3d8..c2bcbef91 100644 --- a/tests/molmo2_xla_vision_parity.rs +++ b/tests/molmo2_xla_vision_parity.rs @@ -106,9 +106,79 @@ struct Comparison { max_abs: f32, max_index: usize, rms: f32, + /// The two values behind `max_abs`, so an absolute gap can be read against + /// the magnitude it sits on. + actual_at_max: f32, + expected_at_max: f32, + /// How many elements exceed the combined tolerance. One element out of a + /// million is an outlier; a broad fraction is a systematic mismatch, and + /// `max_abs` alone cannot tell those apart. + over_limit: usize, + len: usize, + /// Worst `difference / (atol + rtol * |expected|)` seen. At most 1.0 the + /// tensor is inside the contract; above 1.0 it is not. + worst_ratio: f32, + ratio_index: usize, + ratio_actual: f32, + ratio_expected: f32, +} + +impl Comparison { + /// Relative size of the worst gap against the value it sits on. + fn relative_at_max(&self) -> f32 { + let scale = self.expected_at_max.abs().max(self.actual_at_max.abs()); + if scale > 0.0 { + self.max_abs / scale + } else { + 0.0 + } + } + + fn detail(&self) -> String { + format!( + "max_abs={} at {} (actual={}, expected={}, relative={:.4}%), rms={}, \ + worst_tolerance_ratio={:.3} at {} (actual={}, expected={}), over_limit={}/{} ({:.4}%)", + self.max_abs, + self.max_index, + self.actual_at_max, + self.expected_at_max, + self.relative_at_max() * 100.0, + self.rms, + self.worst_ratio, + self.ratio_index, + self.ratio_actual, + self.ratio_expected, + self.over_limit, + self.len, + self.over_limit as f64 * 100.0 / self.len as f64 + ) + } } fn compare(actual: &[f32], expected: &[f32]) -> Result { + compare_with_limit(actual, expected, f32::INFINITY, 0.0) +} + +/// Compare against the combined tolerance `atol + rtol * |expected|`. +/// +/// A single absolute limit cannot express this contract across the pipeline. +/// The ViT stages carry values around 1, where 0.05 is a strict bound, but the +/// projector output reaches magnitudes of 2e4, where one f32 ULP is already +/// about 0.004 and a 2000-term dot product accumulated in a different order +/// lands tens of ULPs away for reasons that have nothing to do with the +/// emitter. Measured on the pinned checkpoint, four ordinary photographs put +/// 4 to 12 elements out of ~2e6 past 0.05 while every one of them agreed to +/// within 0.001% relatively. +/// +/// `atol` still floors the comparison near zero, so nothing that was strict +/// before becomes loose: the bound only widens where the values themselves are +/// large enough to make an absolute bound meaningless. +fn compare_with_limit( + actual: &[f32], + expected: &[f32], + max_abs_limit: f32, + relative_limit: f32, +) -> Result { if actual.len() != expected.len() { return Err(anyhow!( "comparison length mismatch: actual={}, expected={}", @@ -122,6 +192,13 @@ fn compare(actual: &[f32], expected: &[f32]) -> Result { let mut max_abs = 0.0f32; let mut max_index = 0usize; let mut squared = 0.0f64; + let mut over_limit = 0usize; + let mut actual_at_max = 0.0f32; + let mut expected_at_max = 0.0f32; + let mut worst_ratio = 0.0f32; + let mut ratio_index = 0usize; + let mut ratio_actual = 0.0f32; + let mut ratio_expected = 0.0f32; for (index, (&actual, &expected)) in actual.iter().zip(expected).enumerate() { if !actual.is_finite() || !expected.is_finite() { return Err(anyhow!( @@ -132,13 +209,43 @@ fn compare(actual: &[f32], expected: &[f32]) -> Result { if difference > max_abs { max_abs = difference; max_index = index; + actual_at_max = actual; + expected_at_max = expected; + } + let allowed = max_abs_limit + relative_limit * expected.abs(); + if difference > allowed { + over_limit += 1; + } + // A zero allowance can only be met exactly, so report an exceeded + // ratio rather than dividing by zero. + let ratio = if allowed > 0.0 { + difference / allowed + } else if difference > 0.0 { + f32::INFINITY + } else { + 0.0 + }; + if ratio > worst_ratio { + worst_ratio = ratio; + ratio_index = index; + ratio_actual = actual; + ratio_expected = expected; } squared += f64::from(difference) * f64::from(difference); } + let len = actual.len(); Ok(Comparison { max_abs, max_index, - rms: (squared / actual.len() as f64).sqrt() as f32, + rms: (squared / len as f64).sqrt() as f32, + actual_at_max, + expected_at_max, + over_limit, + len, + worst_ratio, + ratio_index, + ratio_actual, + ratio_expected, }) } @@ -148,20 +255,17 @@ fn assert_within( actual: &[f32], expected: &[f32], max_abs_limit: f32, + relative_limit: f32, rms_limit: f32, ) -> Result<()> { - let comparison = compare(actual, expected)?; - emit_progress(&format!( - "{label}: max_abs={} at {}, rms={}", - comparison.max_abs, comparison.max_index, comparison.rms - )); - if comparison.max_abs > max_abs_limit || comparison.rms > rms_limit { + let comparison = compare_with_limit(actual, expected, max_abs_limit, relative_limit)?; + emit_progress(&format!("{label}: {}", comparison.detail())); + if comparison.worst_ratio > 1.0 || comparison.rms > rms_limit { return Err(anyhow!( - "{label} parity failed: max_abs={} at {}, rms={}, limits=({}, {})", - comparison.max_abs, - comparison.max_index, - comparison.rms, + "{label} parity failed: {}, limits=(atol={}, rtol={}, rms={})", + comparison.detail(), max_abs_limit, + relative_limit, rms_limit )); } @@ -247,6 +351,53 @@ fn synthetic_comparison_reports_max_and_rms() { assert_eq!(comparison.max_abs, 0.5); assert_eq!(comparison.max_index, 1); assert!((comparison.rms - (0.25f32 / 3.0).sqrt()).abs() < 1e-7); + // The worst gap is reported against the magnitude it sits on, so an + // absolute limit can be read as a relative one. + assert_eq!(comparison.actual_at_max, 2.5); + assert_eq!(comparison.expected_at_max, 2.0); + assert!((comparison.relative_at_max() - 0.2).abs() < 1e-6); +} + +#[test] +fn synthetic_comparison_separates_one_outlier_from_broad_drift() { + // One element over the limit out of four: an outlier. + let outlier = + compare_with_limit(&[1.0, 1.0, 1.0, 9.0], &[1.0, 1.0, 1.0, 1.0], 0.5, 0.0).unwrap(); + assert_eq!(outlier.over_limit, 1); + assert_eq!(outlier.len, 4); + + // Every element over the limit at the same max: same max_abs, different story. + let drift = compare_with_limit(&[9.0, 9.0, 9.0, 9.0], &[1.0, 1.0, 1.0, 1.0], 0.5, 0.0).unwrap(); + assert_eq!(drift.max_abs, outlier.max_abs); + assert_eq!(drift.over_limit, 4); +} + +#[test] +fn the_relative_half_only_widens_the_bound_where_values_are_large() { + let atol = 0.05f32; + let rtol = 1e-5f32; + + // Near zero the absolute floor still rules: a gap just past atol fails, so + // nothing that was strict before this change became loose. + assert!( + assert_within("small", &[0.06], &[0.0], atol, rtol, 1.0).is_err(), + "atol must still bound values around zero" + ); + assert!(assert_within("small-ok", &[0.04], &[0.0], atol, rtol, 1.0).is_ok()); + + // At the projector output's magnitude the same absolute gap is a fraction + // of one f32 ULP's worth of relative error and must pass. + assert!( + assert_within("large", &[25507.98], &[25507.887], atol, rtol, 1.0).is_ok(), + "a 0.09 gap on a 2.5e4 value is reduction-order noise, not a defect" + ); + + // A genuine defect changes leading digits, which the relative half still + // catches no matter how large the value is. + assert!( + assert_within("large-wrong", &[25507.98], &[25000.0], atol, rtol, 1.0).is_err(), + "a 2% error on a large value must still fail" + ); } #[test] @@ -254,6 +405,21 @@ fn synthetic_negative_controls_detect_layer_denominator_and_clamped_index_drift( let canonical_layers = [24.0, 18.0, 240.0, 180.0]; let wrong_layers = [18.0, 24.0, 180.0, 240.0]; assert!(compare(&canonical_layers, &wrong_layers).unwrap().max_abs > 0.0); + // The combined tolerance must not blunt this: swapping the selected layers + // is a defect at any magnitude, so it has to fail under the same atol/rtol + // the real gate runs with. + assert!( + assert_within( + "wrong-layers", + &canonical_layers, + &wrong_layers, + 0.05, + 1e-5, + 0.01 + ) + .is_err(), + "selected-layer drift must still be rejected" + ); let masked_values = [2.0, 6.0, 0.0, 0.0]; let valid_mean = masked_values.iter().sum::() / 2.0; @@ -298,6 +464,25 @@ fn synthetic_progress_flushes_and_heartbeats_before_five_minutes() { assert!(PROGRESS_HEARTBEAT_INTERVAL < Duration::from_secs(5 * 60)); } +/// Compare the whole vision path against eager MLX on a real checkpoint. +/// +/// `MLXCEL_MOLMO2_IMAGE` selects the image, which decides the crop tiling and +/// therefore which parts of the path run at all. The bundled fixture is a +/// 224x224 solid square, whose `1x1` tiling exercises a single crop and skips +/// the high-resolution tiling, pooling offsets, and multi-crop merge entirely. +/// Validated tilings, all against the pinned 4B checkpoint on CUDA: +/// +/// | image | size | tiling | crops | prompt tokens | +/// | ------------ | -------- | ------ | ----- | ------------- | +/// | fixture | 224x224 | 1x1 | 2 | 424 | +/// | square-large | 756x756 | 2x2 | 5 | 970 | +/// | tall | 378x1512 | 4x1 | 5 | 1024 | +/// | wide | 1512x378 | 1x4 | 5 | 984 | +/// | photo-ish | 1024x768 | 2x3 | 7 | 1348 | +/// +/// The non-fixture images are generated rather than committed; any textured +/// image of those dimensions reproduces the tiling. A flat color will not: it +/// hides pooling and interpolation differences. #[test] #[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] #[ignore = "requires a Molmo2 checkpoint plus configured MLX and IREE runtimes"] @@ -313,6 +498,10 @@ fn real_checkpoint_mlx_iree_vision_and_scatter_parity() -> Result<()> { }); let device = std::env::var("MLXCEL_XLA_DEVICE").unwrap_or_else(|_| "local-task".to_string()); let max_abs_limit = tolerance("MLXCEL_MOLMO2_MAX_ABS", 0.05)?; + // Relative half of the contract. 1e-5 is roughly four f32 ULPs at the + // projector output's magnitude, so it admits reduction-order differences + // while still rejecting anything that changes a value's leading digits. + let relative_limit = tolerance("MLXCEL_MOLMO2_RTOL", 1e-5)?; let rms_limit = tolerance("MLXCEL_MOLMO2_RMS", 0.01)?; let _runtime = initialize_runtime(); @@ -400,6 +589,7 @@ fn real_checkpoint_mlx_iree_vision_and_scatter_parity() -> Result<()> { &iree_stage.values, &eager_stage.values, max_abs_limit, + relative_limit, rms_limit, )?; } @@ -408,6 +598,7 @@ fn real_checkpoint_mlx_iree_vision_and_scatter_parity() -> Result<()> { &iree.projected_values, &eager.values, max_abs_limit, + relative_limit, rms_limit, )?; @@ -454,6 +645,7 @@ fn real_checkpoint_mlx_iree_vision_and_scatter_parity() -> Result<()> { &production_iree_merged, &eager_merged, max_abs_limit, + relative_limit, rms_limit, ) } From ba25807a000c64bb36e94fad4b61038f5db3ba4d Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 21 Aug 2026 20:10:53 +0900 Subject: [PATCH 20/20] test(xla): check the tolerance shape in every build of the gate The tolerance-shape and negative-control assertions called `assert_within`, which is gated on the diagnostics features along with the progress reporting it emits and the `std::io` import behind that. A `--features cuda,xla-iree` build therefore failed to compile the test target with E0425, because the helper was configured out from under callers that are not. Rather than ungating the helper and its progress plumbing, extract `within`, the same admission predicate without the reporting, and have the synthetic tests use it. The assertions now run in every build of this target instead of only the diagnostics ones that can reach a checkpoint. --- tests/molmo2_xla_vision_parity.rs | 30 ++++++++++++++++-------------- 1 file changed, 16 insertions(+), 14 deletions(-) diff --git a/tests/molmo2_xla_vision_parity.rs b/tests/molmo2_xla_vision_parity.rs index c2bcbef91..1e53e66f9 100644 --- a/tests/molmo2_xla_vision_parity.rs +++ b/tests/molmo2_xla_vision_parity.rs @@ -272,6 +272,15 @@ fn assert_within( Ok(()) } +/// The admission predicate `assert_within` applies, without its progress +/// reporting. Kept ungated so the tolerance shape is verified in every build of +/// this target, not only the diagnostics ones that can reach a checkpoint. +fn within(actual: &[f32], expected: &[f32], max_abs_limit: f32, relative_limit: f32) -> bool { + compare_with_limit(actual, expected, max_abs_limit, relative_limit) + .map(|comparison| comparison.worst_ratio <= 1.0) + .unwrap_or(false) +} + fn active_groups(pooling: &[i32], groups: usize, group_size: usize) -> Result> { if group_size == 0 || pooling.len() != groups * group_size { return Err(anyhow!( @@ -380,22 +389,23 @@ fn the_relative_half_only_widens_the_bound_where_values_are_large() { // Near zero the absolute floor still rules: a gap just past atol fails, so // nothing that was strict before this change became loose. assert!( - assert_within("small", &[0.06], &[0.0], atol, rtol, 1.0).is_err(), + !within(&[0.06], &[0.0], atol, rtol), "atol must still bound values around zero" ); - assert!(assert_within("small-ok", &[0.04], &[0.0], atol, rtol, 1.0).is_ok()); + assert!(within(&[0.04], &[0.0], atol, rtol)); // At the projector output's magnitude the same absolute gap is a fraction - // of one f32 ULP's worth of relative error and must pass. + // of one f32 ULP's worth of relative error and must pass. These are the + // measured `tall` values from the real gate. assert!( - assert_within("large", &[25507.98], &[25507.887], atol, rtol, 1.0).is_ok(), + within(&[25507.98], &[25507.887], atol, rtol), "a 0.09 gap on a 2.5e4 value is reduction-order noise, not a defect" ); // A genuine defect changes leading digits, which the relative half still // catches no matter how large the value is. assert!( - assert_within("large-wrong", &[25507.98], &[25000.0], atol, rtol, 1.0).is_err(), + !within(&[25507.98], &[25000.0], atol, rtol), "a 2% error on a large value must still fail" ); } @@ -409,15 +419,7 @@ fn synthetic_negative_controls_detect_layer_denominator_and_clamped_index_drift( // is a defect at any magnitude, so it has to fail under the same atol/rtol // the real gate runs with. assert!( - assert_within( - "wrong-layers", - &canonical_layers, - &wrong_layers, - 0.05, - 1e-5, - 0.01 - ) - .is_err(), + !within(&canonical_layers, &wrong_layers, 0.05, 1e-5), "selected-layer drift must still be rejected" );