diff --git a/Cargo.toml b/Cargo.toml index ab506bce0..c4248540b 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", "mlxcel-xla/diagnostics"] # Axis A "weight-load surgery" framework. # # On by default so production `mlxcel` and `mlxcel-server` binaries expose @@ -290,6 +294,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..f5674f1b4 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -121,6 +121,11 @@ pub fn decode_image_payloads_with_limits( // Re-export split modules 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, Molmo2XlaVisionReferenceStage, + 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/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/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..91bc2a31f 100644 --- a/src/lib/mlxcel-xla/src/emitter/mod.rs +++ b/src/lib/mlxcel-xla/src/emitter/mod.rs @@ -57,6 +57,8 @@ mod gemma3n_schema; mod gemma3n_weights; mod model; mod moe; +mod molmo2_config; +mod molmo2_vision; pub(crate) mod numeric_ops; mod phi4_audio; mod qwen2_vl; @@ -152,6 +154,12 @@ 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; +#[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"))] 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..b1ca2ea01 --- /dev/null +++ b/src/lib/mlxcel-xla/src/emitter/molmo2_config.rs @@ -0,0 +1,473 @@ +// 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")?; + 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 + .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")?; + 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 { + 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) + .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;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, + 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..41ca55469 --- /dev/null +++ b/src/lib/mlxcel-xla/src/emitter/molmo2_config_tests.rs @@ -0,0 +1,132 @@ +// 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, + "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": [-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, 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("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")); +} + +#[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] +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!([-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() + .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 new file mode 100644 index 000000000..f7e5ffbfc --- /dev/null +++ b/src/lib/mlxcel-xla/src/emitter/molmo2_vision.rs @@ -0,0 +1,670 @@ +// 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}; +use super::numeric_ops::{silu, stable_softmax, tanh_gelu}; + +const MOLMO2_VIT_PROBE_LAYER: usize = 18; +// 591490 = 513 * checkpoint hidden width 1152 + component 514. +const MOLMO2_VIT_PROBE_FLAT_ROW: usize = 513; + +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 softmax_last(builder: &mut Builder, scores: &Val) -> Val { + let axis = scores.ty.shape.len() - 1; + stable_softmax(builder, scores, axis) +} + +fn selected_slot(selected_layers: &[usize], layer: usize) -> Option { + selected_layers + .iter() + .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, + 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)) +} + +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, +) -> IndexedPoolValues { + 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 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); + + 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)); + IndexedPoolValues { + gathered_masked: gathered, + query, + pooled, + counts, + } +} + +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(); + 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 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()]; + 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 + // 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 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, + &norm_weight, + &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(); + let _attention_norm_bias = args.take(); + 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, + &ffn_norm_weight, + &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(); + 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)); + 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()); + } + if let Some(slot) = selected_slot(&config.selected_layers, layer) { + selected[slot] = Some(hidden.clone()); + } + } + 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, + vec![ + config.static_crops * config.patches_per_crop, + config.selected_width(), + ], + ); + let pool = 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 = 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(), + "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"), + ]; + 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, + pool.counts, + pool.query, + pool.pooled, + gate_linear, + gate_activation, + up_linear, + projector_product, + 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}) -> {result_type} {{\n{body} return {result_values} : {return_types}\n }}\n}}\n", + signature = args.declarations.join(", "), + body = builder.body(), + 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::*; + + 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, + "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":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() + } + + #[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")); + 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 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(&[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); + } + + #[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>")); + 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] + fn diagnostic_graph_probes_the_first_selected_layer_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!( + MOLMO2_VIT_PROBE_LAYER, 18, + "the row probes must precede the first failing selected-layer comparison" + ); + 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(); + 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") + ); + } +} 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..6ee3aa39b 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,18 @@ 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 = "diagnostics")] +pub use molmo2_vision_runtime::{ + IreeMolmo2VisionDiagnosticProjector, Molmo2VisionDiagnosticStage, Molmo2VisionDiagnostics, +}; +#[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..24b7d0bb1 --- /dev/null +++ b/src/lib/mlxcel-xla/src/molmo2.rs @@ -0,0 +1,482 @@ +// 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, + }, + ActiveTokenCount { + active_groups: usize, + grid_groups: usize, + prompt_positions: usize, + all_invalid_groups: Vec, + }, + 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::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, + "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 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, 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() + .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) + } +} + +/// 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().is_multiple_of(hidden_size) { + 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(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) + .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 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]; + 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..df99063ef --- /dev/null +++ b/src/lib/mlxcel-xla/src/molmo2_vision_runtime.rs @@ -0,0 +1,814 @@ +// 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}; +#[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; +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 = 18; +#[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 { + 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, + 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 { + 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) +} + +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_legacy_unqualified( + 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, +} + +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 module = + compile_vision_module(model_dir, device, &config, &mlir, "molmo2-vision", None)?; + 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 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 = [ + 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 [AuxiliaryOutput { + bytes: &mut output, + dtype: AuxiliaryTensorDType::Float32, + shape: &output_shape, + }], + )?; + let all_values = decode_output(&output)?; + 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]); + } + Ok(Molmo2VisionProjection { + shape: [ + values.len() / self.config.text_hidden, + self.config.text_hidden, + ], + values, + 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(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(), + }, + ]; + 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( + selected_layers[probe_split..] + .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.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], + 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-v4-layer18-row513"), + )?; + 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/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..99f637882 --- /dev/null +++ b/src/lib/mlxcel-xla/src/molmo2_vision_runtime_tests.rs @@ -0,0 +1,87 @@ +// 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); +} + +#[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/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..1dd6ebfa0 100644 --- a/src/loading/mod.rs +++ b/src/loading/mod.rs @@ -62,9 +62,15 @@ 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(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] +pub use self::vlm::{ + Molmo2XlaVisionReference, Molmo2XlaVisionReferenceProjection, Molmo2XlaVisionReferenceStage, + load_molmo2_xla_vision_reference, +}; #[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..c1ae91c56 100644 --- a/src/loading/vlm.rs +++ b/src/loading/vlm.rs @@ -125,9 +125,15 @@ 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(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] +pub use special::{ + Molmo2XlaVisionReference, Molmo2XlaVisionReferenceProjection, Molmo2XlaVisionReferenceStage, + load_molmo2_xla_vision_reference, +}; #[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..8bea03711 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 @@ -1724,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.") { @@ -1942,31 +2031,54 @@ fn read_clip_triple(config: Option<&Value>, key: &str) -> Option<[f32; 3]> { }) } -pub(crate) fn load_molmo2_vlm(model_path: &Path) -> Result { +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, + 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 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) + .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`" + )); + } - 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()) @@ -2011,7 +2123,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()) @@ -2033,24 +2145,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, @@ -2068,19 +2164,242 @@ 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))?; + .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", feature = "xla-diagnostics-cpu"))] +#[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", feature = "xla-diagnostics-cpu"))] +#[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, + 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"))] +#[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]); + #[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}"))?; + 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, + stages, + }) + } +} + +/// Load only Molmo2's eager vision encoder/projector for 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)?; + 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 mut raw_weights = models::load_weights_from_dir_with_filter( + model_path, + |name| name.starts_with("vision_tower.") || name.starts_with("model.vision_backbone."), + false, + ) + .map_err(anyhow::Error::msg)?; + super::finish_vlm_weights_common(model_path, &mut raw_weights, None, false)?; + 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/src/loading/vlm_special_tests.rs b/src/loading/vlm_special_tests.rs index f76ee5288..b5c85dfca 100644 --- a/src/loading/vlm_special_tests.rs +++ b/src/loading/vlm_special_tests.rs @@ -16,10 +16,11 @@ 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, + 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, @@ -766,6 +767,40 @@ 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 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!( diff --git a/src/models/mod.rs b/src/models/mod.rs index 12acd5957..142e29b5f 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -229,7 +229,8 @@ pub use rwkv7::Rwkv7; pub(crate) use sanitize::{ Gemma4WeightBacking, load_gemma4_text_weights_with_backing, load_gemma4_unified_weights_with_backing, load_gemma4_vlm_weights_with_backing, - sanitize_gemma4_nvfp4_weights, strip_gemma4_kv_shared_weights, + load_weights_from_dir_with_filter, sanitize_gemma4_nvfp4_weights, + strip_gemma4_kv_shared_weights, }; pub use sanitize::{ convert_bf16_weights, convert_bf16_weights_with_keep, gemma3n_language_mlp_bf16_key, diff --git a/src/models/sanitize.rs b/src/models/sanitize.rs index fbb0f1316..51d917d22 100644 --- a/src/models/sanitize.rs +++ b/src/models/sanitize.rs @@ -1252,7 +1252,14 @@ where Ok(()) } -fn load_weights_from_dir_with_filter( +/// 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/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..e4b63f7e5 --- /dev/null +++ b/src/multimodal/molmo2_xla_preprocessor.rs @@ -0,0 +1,352 @@ +// 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)?; + 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, + &[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)?, + prompt_image_patch_count, + }) + .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() + ); + } +} diff --git a/src/vision/encoders/molmo2.rs b/src/vision/encoders/molmo2.rs index b3daae083..765b0567d 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,86 @@ 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 = 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. +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 { + 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 +116,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) } @@ -179,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, @@ -231,13 +349,21 @@ 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, + probe_rows: Vec<(&'static str, 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 +385,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 +393,99 @@ 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 mut 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"), + ), + 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)); } - 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( @@ -333,6 +539,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(...) @@ -343,6 +558,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))?; @@ -360,9 +606,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 +648,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 +679,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 +770,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 +794,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 +823,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,8 +837,18 @@ 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]); - - // Project through SwiGLU MLP + #[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. + #[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] @@ -502,12 +876,114 @@ 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 projector = + projector_capture.expect("Molmo2 projector diagnostics must be captured"); + 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, + }, + ]; + 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, + } + })); + 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.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, + }, + ]); + 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 @@ -574,9 +1050,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 +1091,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, }) } @@ -636,3 +1108,53 @@ 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, + ] + ); + } +} diff --git a/tests/molmo2_xla_vision_parity.rs b/tests/molmo2_xla_vision_parity.rs new file mode 100644 index 000000000..5ab99f3d8 --- /dev/null +++ b/tests/molmo2_xla_vision_parity.rs @@ -0,0 +1,459 @@ +// 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(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] +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"))] +use mlxcel::{initialize_runtime, load_molmo2_xla_vision_reference}; +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, + 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(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] +fn assert_within( + label: &str, + actual: &[f32], + expected: &[f32], + max_abs_limit: f32, + rms_limit: f32, +) -> Result<()> { + let comparison = compare(actual, expected)?; + emit_progress(&format!( + "{label}: max_abs={} at {}, rms={}", + comparison.max_abs, comparison.max_index, comparison.rms + )); + if comparison.max_abs > max_abs_limit || comparison.rms > rms_limit { + return Err(anyhow!( + "{label} parity failed: max_abs={} at {}, rms={}, limits=({}, {})", + comparison.max_abs, + comparison.max_index, + comparison.rms, + max_abs_limit, + rms_limit + )); + } + 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(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] +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] +fn synthetic_negative_controls_detect_layer_denominator_and_clamped_index_drift() { + let canonical_layers = [24.0, 18.0, 240.0, 180.0]; + let wrong_layers = [18.0, 24.0, 180.0, 240.0]; + assert!(compare(&canonical_layers, &wrong_layers).unwrap().max_abs > 0.0); + + 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"))] +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"] +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 = 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 = with_progress("MLX eager diagnostic projection", || { + 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 + )); + } + + // 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) + })?; + 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 = 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={:?}", + iree.active_groups, + eager.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.projected_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.projected_values, + )?; + let mut production_iree_merged = base; + add_molmo2_projected_features( + &tokens, + reference.image_patch_id(), + &mut production_iree_merged, + hidden, + &iree.projected_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, + ) +}