From fcd990434a328290021c72b0446ea91988833506 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 24 Jul 2026 18:41:21 +0900 Subject: [PATCH 01/16] 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 | 96 ++-- src/multimodal/molmo2_xla_preprocessor.rs | 347 +++++++++++++ 17 files changed, 2716 insertions(+), 66 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 0065385c7..9743e88c4 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; #[cfg(any(test, feature = "diagnostics"))] pub(crate) use vision::emit_vision_diagnostics; 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 2b1328267..9115c973a 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 fb2cab9d4..b8d2150ef 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, }; /// Resolve model path: if a file is given, use its parent directory. diff --git a/src/loading/vlm.rs b/src/loading/vlm.rs index b87c6c0a2..bdf2eb6eb 100644 --- a/src/loading/vlm.rs +++ b/src/loading/vlm.rs @@ -127,7 +127,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 ea8353ddb..7ae60aa04 100644 --- a/src/loading/vlm_special.rs +++ b/src/loading/vlm_special.rs @@ -1362,6 +1362,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 bfad9cb9d..05fd99739 100644 --- a/src/multimodal/host_preprocessor.rs +++ b/src/multimodal/host_preprocessor.rs @@ -42,12 +42,17 @@ 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, export_qwen2_vl_prefill, usize_to_i32, validate_embedding_shape, validate_processor_shape, validate_projected_shape, validate_sequence_capacity, }; +#[cfg(feature = "xla-iree")] +pub use molmo2_xla::Molmo2IreeHostPreprocessor; /// Vision implementation selected for OpenXLA multimodal preprocessing. #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -147,40 +152,71 @@ pub fn load_xla_image_preprocessor( model_path.display() )) })?; - 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))); - } - #[cfg(not(feature = "xla-iree"))] - { - return Err(HostPreprocessorError::InvalidConfig( - "Qwen2-VL XLA image execution requires the xla-iree feature".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))) + } + #[cfg(not(feature = "xla-iree"))] + { + Err(HostPreprocessorError::InvalidConfig( + "Qwen2-VL XLA image execution requires the xla-iree feature".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 ef024c63c83085c52f7426224c1ee88aaeb46cb0 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 24 Jul 2026 18:52:06 +0900 Subject: [PATCH 02/16] 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 88b13b994582a9dc815579ef76164b691514aef1 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 24 Jul 2026 20:11:45 +0900 Subject: [PATCH 03/16] 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 44d58ce5e468505aa30ff48769cad04ae6f38c04 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 24 Jul 2026 20:55:19 +0900 Subject: [PATCH 04/16] 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 | 31 +++++++-- src/vision/encoders/molmo2.rs | 24 +++---- 6 files changed, 147 insertions(+), 37 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 7ae60aa04..44942851b 100644 --- a/src/loading/vlm_special.rs +++ b/src/loading/vlm_special.rs @@ -1765,6 +1765,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.") { @@ -2002,12 +2050,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()) @@ -2109,7 +2163,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 f76ee5288..bfff86a3c 100644 --- a/src/loading/vlm_special_tests.rs +++ b/src/loading/vlm_special_tests.rs @@ -16,12 +16,13 @@ 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, 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, + 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; @@ -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 262585ac817314d22fc4c8bf30c67fba05fa35fb Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 24 Jul 2026 21:11:41 +0900 Subject: [PATCH 05/16] 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 ab506bce0..9a00388d4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -290,6 +290,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 d9520a12f..3a88fb0f7 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 b8d2150ef..fba2c72a2 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 bdf2eb6eb..d9af1e037 100644 --- a/src/loading/vlm.rs +++ b/src/loading/vlm.rs @@ -125,6 +125,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 44942851b..447143501 100644 --- a/src/loading/vlm_special.rs +++ b/src/loading/vlm_special.rs @@ -2031,22 +2031,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); @@ -2106,7 +2098,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()) @@ -2128,24 +2120,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, @@ -2166,16 +2142,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 418a4abd430c90b1e1dda79a33d170489325802d Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Sun, 26 Jul 2026 22:08:02 +0900 Subject: [PATCH 06/16] 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 9a00388d4..76d882f68 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -79,6 +79,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 3a88fb0f7..2f16e8661 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 72ac0d1aa..74d7fd80d 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; @@ -550,6 +553,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()))?; @@ -559,7 +581,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); @@ -900,6 +925,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 fba2c72a2..e730da0f0 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 d9af1e037..fb97423cb 100644 --- a/src/loading/vlm.rs +++ b/src/loading/vlm.rs @@ -125,7 +125,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 447143501..0d7300585 100644 --- a/src/loading/vlm_special.rs +++ b/src/loading/vlm_special.rs @@ -2160,7 +2160,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, @@ -2170,7 +2170,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, @@ -2179,7 +2179,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 { @@ -2252,7 +2252,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)?; @@ -2268,9 +2268,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 8dbdceab6b5a5308a068684b0a290e4aef0ebac5 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Sun, 26 Jul 2026 22:41:36 +0900 Subject: [PATCH 07/16] 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 76d882f68..c4248540b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -82,7 +82,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 2f16e8661..f5674f1b4 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 9743e88c4..4f46744c7 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; #[cfg(any(test, feature = "diagnostics"))] 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 9115c973a..6ee3aa39b 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 e730da0f0..1dd6ebfa0 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 fb97423cb..c1ae91c56 100644 --- a/src/loading/vlm.rs +++ b/src/loading/vlm.rs @@ -127,7 +127,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 0d7300585..0d37d678a 100644 --- a/src/loading/vlm_special.rs +++ b/src/loading/vlm_special.rs @@ -2041,6 +2041,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. @@ -2177,6 +2186,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"))] @@ -2207,7 +2225,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}"))?; @@ -2247,6 +2311,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 eb5a5101c5c7ddda6a582c0621d9961c6faafcb5 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Sun, 26 Jul 2026 22:59:09 +0900 Subject: [PATCH 08/16] 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 | 30 +++++++++++++++++++++++------- 2 files changed, 42 insertions(+), 10 deletions(-) diff --git a/src/loading/vlm_special.rs b/src/loading/vlm_special.rs index 0d37d678a..8bea03711 100644 --- a/src/loading/vlm_special.rs +++ b/src/loading/vlm_special.rs @@ -2031,6 +2031,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, @@ -2038,9 +2056,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 bfff86a3c..b5c85dfca 100644 --- a/src/loading/vlm_special_tests.rs +++ b/src/loading/vlm_special_tests.rs @@ -16,13 +16,13 @@ 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, 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, + 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; @@ -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 1d071b994fbf48fbbf92488d9e4bd9645937f40d Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Sun, 26 Jul 2026 23:28:29 +0900 Subject: [PATCH 09/16] 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 7261e4f530632e10130171480ca584eca7e4347f Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Sun, 26 Jul 2026 23:50:18 +0900 Subject: [PATCH 10/16] 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 aa0e4831235ecc0c57da66b08f67da453f907425 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 00:28:23 +0900 Subject: [PATCH 11/16] 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 6a6652944848cbe145152b046aac8c59987773d9 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 01:55:58 +0900 Subject: [PATCH 12/16] 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 e32c595e12d86e8c945e9deb597951d42d17d144 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 02:08:40 +0900 Subject: [PATCH 13/16] 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 5a5ab482e6e3cc6eb665fcb23b21a9e77e29172b Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 09:21:12 +0900 Subject: [PATCH 14/16] 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 684c6ee1d37497b17b5baaf5fc2ff25cf980c536 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 13:12:23 +0900 Subject: [PATCH 15/16] 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 4f46744c7..91bc2a31f 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 329fbd996e3fec92be0b1a3f4a4dd4b7d8cc139b Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Tue, 28 Jul 2026 02:32:28 +0900 Subject: [PATCH 16/16] 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) })?;