diff --git a/Cargo.toml b/Cargo.toml index ab506bce0..63e20f08c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -71,11 +71,14 @@ xla-backend = ["dep:mlxcel-xla"] # StableHLO prefill / decode_step graphs. Needs IREE_DIST at build time (the # extracted iree dist), so it is a local / opt-in build, not a CI default. xla-iree = ["xla-backend", "mlxcel-xla/iree"] +# Device-neutral eager MLX/IREE reference capture for CPU `local-task` +# qualification without requiring an IREE CUDA source/build tree. +xla-reference-diagnostics = ["xla-iree", "mlxcel-xla/diagnostics"] # Explicit test-only forwarding for Gemma3n intermediate oracle validation. # Its pinned reference is the production CUDA runtime, so enabling diagnostics # also enables the root MLX CUDA backend. This is not part of `xla-iree`; normal # production bundles remain unchanged. -xla-diagnostics = ["cuda", "xla-iree", "mlxcel-xla/diagnostics"] +xla-diagnostics = ["cuda", "xla-reference-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"] diff --git a/examples/xla_molmo_reference_check.rs b/examples/xla_molmo_reference_check.rs new file mode 100644 index 000000000..b6417c461 --- /dev/null +++ b/examples/xla_molmo_reference_check.rs @@ -0,0 +1,24 @@ +// 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. + +//! Dedicated Molmo v1 actual-checkpoint reference gate. +//! +//! This executable keeps eager MLX on CUDA while the IREE vision and decoder +//! references remain on `local-task` or `local-sync`. It deliberately avoids +//! the root crate's full libtest harness and emits flushed progress plus +//! periodic heartbeats around every potentially long eager materialization. + +fn main() { + mlxcel::multimodal::host_preprocessor::run_pinned_molmo_eager_mlx_iree_boundaries(); +} diff --git a/src/lib/mlxcel-xla/src/emitter/config.rs b/src/lib/mlxcel-xla/src/emitter/config.rs index e5c450eaa..7d3488a88 100644 --- a/src/lib/mlxcel-xla/src/emitter/config.rs +++ b/src/lib/mlxcel-xla/src/emitter/config.rs @@ -132,6 +132,9 @@ 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, + /// Molmo v1's OLMo-style decoder layout under `language_model.model`, with + /// fused `att_proj` QKV and fused `ff_proj` gate/up projections. + Molmo, } /// MLX affine weight quantization (`config.json` `quantization`). The linear / @@ -509,10 +512,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("molmo") => WeightScheme::Molmo, + _ => WeightScheme::Llama, }; // Interleaved (GPT-J-style) RoPE reaches the supported families only through @@ -870,6 +873,35 @@ impl Config { tie_default = false; rotary_dim = partial_rotary(1.0); } + Some("molmo") => { + if !matches!( + v.get("rope_impl").and_then(serde_json::Value::as_str), + None | Some("interleave") + ) { + return Err("Molmo v1 XLA requires rope_impl = \"interleave\"".to_string()); + } + if v.get("qkv_bias").and_then(serde_json::Value::as_bool) == Some(false) { + return Err("Molmo v1 XLA requires qkv_bias = true".to_string()); + } + if !matches!( + v.get("layer_norm_type").and_then(serde_json::Value::as_str), + None | Some("rms") + ) { + return Err("Molmo v1 XLA requires layer_norm_type = \"rms\"".to_string()); + } + if v.get("tie_word_embeddings") + .and_then(serde_json::Value::as_bool) + == Some(true) + || v.get("weight_tying").and_then(serde_json::Value::as_bool) == Some(true) + { + return Err("Molmo v1 XLA requires an untied output head".to_string()); + } + qkv_bias = true; + tie_default = false; + fused_qkv = true; + fused_gate_up = true; + rope_interleaved = true; + } Some("stablelm") => { // LayerNorm with bias, partial RoPE, optional q/k/v bias, untied. layernorm = true; @@ -1003,7 +1035,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, Molmo, 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)" @@ -1343,10 +1375,22 @@ impl Config { None }; + let intermediate = u("intermediate_size")?; + let intermediate = if model_type == Some("molmo") { + if !intermediate.is_multiple_of(2) { + return Err(format!( + "Molmo fused intermediate_size={intermediate} must be even" + )); + } + intermediate / 2 + } else { + intermediate + }; + Ok(Config { context_capacity: crate::DEFAULT_CONTEXT_CAPACITY, hidden, - inter: u("intermediate_size")?, + inter: intermediate, n_layers, n_q, n_kv: u("num_key_value_heads")?, diff --git a/src/lib/mlxcel-xla/src/emitter/mod.rs b/src/lib/mlxcel-xla/src/emitter/mod.rs index 0065385c7..f8fb9d365 100644 --- a/src/lib/mlxcel-xla/src/emitter/mod.rs +++ b/src/lib/mlxcel-xla/src/emitter/mod.rs @@ -57,6 +57,9 @@ mod gemma3n_schema; mod gemma3n_weights; mod model; mod moe; +mod molmo_vision; +mod molmo_vision_config; +mod molmo_vision_ops; pub(crate) mod numeric_ops; mod phi4_audio; mod qwen2_vl; @@ -152,6 +155,14 @@ pub(crate) use model::{ validate_prefill_embeddings_metadata, }; #[allow(unused_imports)] +pub(crate) use molmo_vision::emit_molmo_vision; +#[cfg(feature = "diagnostics")] +pub(crate) use molmo_vision::emit_molmo_vision_diagnostics; +#[allow(unused_imports)] +pub(crate) use molmo_vision_config::{ + MolmoVisionConfig, MolmoVisionWeightDType, MolmoVisionWeightSpec, +}; +#[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/molmo_vision.rs b/src/lib/mlxcel-xla/src/emitter/molmo_vision.rs new file mode 100644 index 000000000..5eafe09af --- /dev/null +++ b/src/lib/mlxcel-xla/src/emitter/molmo_vision.rs @@ -0,0 +1,485 @@ +// 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 Molmo v1's ViT, attention pool, and SwiGLU projector. + +use super::builder::{Builder, Ty, Val}; +use super::molmo_vision_config::{ + MolmoVisionConfig, MolmoVisionWeightDType, MolmoVisionWeightSpec, +}; +use super::molmo_vision_ops::{scalar_broadcast, silu, softmax_last, tanh_gelu}; + +struct Args { + values: Vec, + declarations: Vec, + cursor: usize, +} + +impl Args { + fn new(specs: &[MolmoVisionWeightSpec]) -> Self { + let mut values = Vec::with_capacity(specs.len() + 2); + let mut declarations = Vec::with_capacity(specs.len() + 2); + for (index, spec) in specs.iter().enumerate() { + let element = match spec.dtype { + MolmoVisionWeightDType::Float32 => "f32", + MolmoVisionWeightDType::Float16 => "f16", + MolmoVisionWeightDType::Uint32 => "ui32", + }; + let ty = Ty::new(spec.shape.clone(), element); + 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 take_quant(&mut self, builder: &mut Builder, config: &MolmoVisionConfig) -> Val { + let packed = self.take(); + let scales = self.take(); + let biases = self.take(); + builder.dequant_affine(&packed, &scales, &biases, config.bits, config.group_size) + } + + fn push_input(&mut self, ty: Ty, name: &str) -> Val { + let index = self.values.len(); + self.declarations + .push(format!("%arg{index}: {} loc(\"{name}\")", ty.render())); + let value = Builder::arg(index, ty); + self.values.push(value.clone()); + value + } +} + +fn bias_rows(builder: &mut Builder, value: &Val, bias: &Val) -> Val { + let bias = builder.broadcast(bias, &[1], value.ty.shape.clone()); + builder.add(value, &bias) +} + +fn quant_linear( + builder: &mut Builder, + args: &mut Args, + config: &MolmoVisionConfig, + value: &Val, + bias: bool, +) -> Val { + let weight = args.take_quant(builder, config); + let projected = builder.linear_seq(value, &weight); + if bias { + bias_rows(builder, &projected, &args.take()) + } else { + projected + } +} + +fn layer_norm(builder: &mut Builder, value: &Val, weight: &Val, bias: &Val, epsilon: f32) -> Val { + let rows = value.ty.shape[0]; + let width = value.ty.shape[1]; + let zero = builder.const_f32(0.0); + let width_scalar = builder.const_f32(width as f32); + let width_rows = builder.broadcast(&width_scalar, &[], vec![rows]); + let sum = builder.reduce_add(value, 1, &zero); + let mean = builder.divide(&sum, &width_rows); + let mean = builder.broadcast(&mean, &[0], vec![rows, width]); + let centered = builder.subtract(value, &mean); + let squared = builder.multiply(¢ered, ¢ered); + let squared_sum = builder.reduce_add(&squared, 1, &zero); + let variance = builder.divide(&squared_sum, &width_rows); + let epsilon = builder.const_f32(epsilon); + let epsilon = builder.broadcast(&epsilon, &[], vec![rows]); + let variance = builder.add(&variance, &epsilon); + let inv_std = builder.rsqrt(&variance); + let inv_std = builder.broadcast(&inv_std, &[0], vec![rows, width]); + let normalized = builder.multiply(¢ered, &inv_std); + let weight = builder.broadcast(weight, &[1], vec![rows, width]); + let bias = builder.broadcast(bias, &[1], vec![rows, width]); + let normalized = builder.multiply(&normalized, &weight); + builder.add(&normalized, &bias) +} + +fn self_attention( + builder: &mut Builder, + args: &mut Args, + config: &MolmoVisionConfig, + hidden: &Val, +) -> Val { + let crops = config.max_crops; + let tokens = config.positions; + let rows = crops * tokens; + let q = quant_linear(builder, args, config, hidden, true); + let k = quant_linear(builder, args, config, hidden, true); + let v = quant_linear(builder, args, config, hidden, true); + 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 = scalar_broadcast( + builder, + (config.head_dim as f32).powf(-0.5), + 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![rows, config.hidden]); + quant_linear(builder, args, config, &context, true) +} + +fn encoder_layer( + builder: &mut Builder, + args: &mut Args, + config: &MolmoVisionConfig, + hidden: &Val, +) -> Val { + let epsilon = f32::from_bits(config.layer_norm_eps_bits); + let norm = layer_norm(builder, hidden, &args.take(), &args.take(), epsilon); + let attention = self_attention(builder, args, config, &norm); + let residual = builder.add(hidden, &attention); + let norm = layer_norm(builder, &residual, &args.take(), &args.take(), epsilon); + let expanded = quant_linear(builder, args, config, &norm, true); + let activated = tanh_gelu(builder, &expanded); + let contracted = quant_linear(builder, args, config, &activated, true); + builder.add(&residual, &contracted) +} + +fn mask_padding_crops( + builder: &mut Builder, + pixels: &Val, + features: &Val, + config: &MolmoVisionConfig, +) -> Val { + let negative_one = scalar_broadcast(builder, -1.0, pixels.ty.shape.clone()); + let is_padding = builder.compare("EQ", pixels, &negative_one, "FLOAT"); + let is_padding = builder.convert(&is_padding, "f32"); + let zero = builder.const_f32(0.0); + let patch_sums = builder.reduce_add(&is_padding, 2, &zero); + let crop_sums = builder.reduce_add(&patch_sums, 1, &zero); + let expected = scalar_broadcast( + builder, + (config.patches_per_crop * config.patch_width) as f32, + vec![config.max_crops], + ); + let all_padding = builder.compare("EQ", &crop_sums, &expected, "FLOAT"); + let all_padding = builder.broadcast( + &all_padding, + &[0], + vec![ + config.max_crops, + config.patches_per_crop, + features.ty.shape[2], + ], + ); + let zero = scalar_broadcast(builder, 0.0, features.ty.shape.clone()); + builder.select(&all_padding, &zero, features) +} + +fn apply_pad_embed(builder: &mut Builder, features: &Val, masks: &Val, pad_embed: &Val) -> Val { + let width = features.ty.shape[2]; + let zero_mask = scalar_broadcast(builder, 0.0, masks.ty.shape.clone()); + let one_mask = scalar_broadcast(builder, 1.0, masks.ty.shape.clone()); + let all_padding = builder.compare("EQ", masks, &zero_mask, "FLOAT"); + let below_one = builder.compare("LT", masks, &one_mask, "FLOAT"); + let not_all_padding = builder.compare("NE", masks, &zero_mask, "FLOAT"); + let partial_padding = builder.and(&below_one, ¬_all_padding); + let all_padding = builder.convert(&all_padding, "f32"); + let partial_padding = builder.convert(&partial_padding, "f32"); + let all_padding = builder.reshape( + &all_padding, + vec![features.ty.shape[0], features.ty.shape[1], 1], + ); + let partial_padding = builder.reshape( + &partial_padding, + vec![features.ty.shape[0], features.ty.shape[1], 1], + ); + let all_padding = builder.broadcast(&all_padding, &[0, 1, 2], features.ty.shape.clone()); + let partial_padding = + builder.broadcast(&partial_padding, &[0, 1, 2], features.ty.shape.clone()); + let complete = builder.slice(pad_embed, &[(0, 1), (0, width)]); + let complete = builder.reshape(&complete, vec![width]); + let complete = builder.broadcast(&complete, &[2], features.ty.shape.clone()); + let partial = builder.slice(pad_embed, &[(1, 2), (0, width)]); + let partial = builder.reshape(&partial, vec![width]); + let partial = builder.broadcast(&partial, &[2], features.ty.shape.clone()); + let complete = builder.multiply(&complete, &all_padding); + let partial = builder.multiply(&partial, &partial_padding); + let features = builder.add(features, &complete); + builder.add(&features, &partial) +} + +fn attention_pool( + builder: &mut Builder, + args: &mut Args, + config: &MolmoVisionConfig, + features: &Val, +) -> Val { + let patch_side = usize::try_from((config.patches_per_crop as f64).sqrt() as u64) + .expect("patch side fits usize"); + let block_h = patch_side / config.pool_h; + let block_w = patch_side / config.pool_w; + let width = features.ty.shape[2]; + let blocks = builder.reshape( + features, + vec![ + config.max_crops, + block_h, + config.pool_h, + block_w, + config.pool_w, + width, + ], + ); + let blocks = builder.transpose(&blocks, &[0, 1, 3, 2, 4, 5]); + let batch = config.max_crops * block_h * block_w; + let pool_size = config.pool_h * config.pool_w; + let blocks = builder.reshape(&blocks, vec![batch, pool_size, width]); + let flattened = builder.reshape(&blocks, vec![batch * pool_size, width]); + let zero = builder.const_f32(0.0); + let sum = builder.reduce_add(&blocks, 1, &zero); + let divisor = scalar_broadcast(builder, pool_size as f32, vec![batch]); + let divisor = builder.broadcast(&divisor, &[0], vec![batch, width]); + let query = builder.divide(&sum, &divisor); + let q = quant_linear(builder, args, config, &query, true); + let k = quant_linear(builder, args, config, &flattened, true); + let v = quant_linear(builder, args, config, &flattened, true); + let q = builder.reshape(&q, vec![batch, config.heads, 1, config.head_dim]); + let k = builder.reshape(&k, vec![batch, pool_size, config.heads, config.head_dim]); + let k = builder.transpose(&k, &[0, 2, 1, 3]); + let v = builder.reshape(&v, vec![batch, pool_size, 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![batch, config.heads, 1, pool_size], + ); + let scale = scalar_broadcast( + builder, + (config.head_dim as f32).powf(-0.5), + 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![batch, config.heads, 1, config.head_dim], + ); + let context = builder.reshape(&context, vec![batch, config.hidden]); + let output = quant_linear(builder, args, config, &context, true); + builder.reshape( + &output, + vec![ + config.max_crops, + config.projected_rows_per_crop(), + config.hidden, + ], + ) +} + +fn emit_molmo_vision_impl(config: &MolmoVisionConfig, diagnostics: bool) -> String { + let specs = config.weight_specs(); + let mut args = Args::new(&specs); + let mut builder = Builder::new(); + let class_embedding = args.take(); + let patch_embedding = args.take(); + let position_embedding = args.take(); + let pre_ln_weight = args.take(); + let pre_ln_bias = args.take(); + let pixels = args.push_input( + Ty::f32(vec![ + config.max_crops, + config.patches_per_crop, + config.patch_width, + ]), + "molmo.pixel_values", + ); + let masks = args.push_input( + Ty::f32(vec![config.max_crops, config.patches_per_crop]), + "molmo.image_masks", + ); + let rows = config.max_crops * config.patches_per_crop; + let pixels_2d = builder.reshape(&pixels, vec![rows, config.patch_width]); + let patches = builder.linear_seq(&pixels_2d, &patch_embedding); + let patches = builder.reshape( + &patches, + vec![config.max_crops, config.patches_per_crop, config.hidden], + ); + let diagnostic_patches = diagnostics.then(|| patches.clone()); + let class_embedding = builder.reshape(&class_embedding, vec![1, 1, config.hidden]); + let class_embedding = builder.broadcast( + &class_embedding, + &[0, 1, 2], + vec![config.max_crops, 1, config.hidden], + ); + let hidden = builder.concatenate(&class_embedding, &patches, 1); + let position_embedding = builder.broadcast( + &position_embedding, + &[1, 2], + vec![config.max_crops, config.positions, config.hidden], + ); + let hidden = builder.add(&hidden, &position_embedding); + let hidden = builder.reshape( + &hidden, + vec![config.max_crops * config.positions, config.hidden], + ); + let mut hidden = layer_norm( + &mut builder, + &hidden, + &pre_ln_weight, + &pre_ln_bias, + f32::from_bits(config.layer_norm_eps_bits), + ); + let mut selected = vec![None; config.selected_layers.len()]; + for layer in 0..config.emitted_layers() { + hidden = encoder_layer(&mut builder, &mut args, config, &hidden); + for (slot, &selected_layer) in config.selected_layers.iter().enumerate() { + if layer == selected_layer { + selected[slot] = Some(hidden.clone()); + } + } + } + let diagnostic_selected = diagnostics.then(|| { + selected + .iter() + .map(|value| { + value + .clone() + .expect("selected Molmo diagnostic layer was emitted") + }) + .collect::>() + }); + let mut selected_values = Vec::with_capacity(selected.len()); + for value in selected { + let value = value.expect("selected layer was emitted"); + let value = builder.reshape( + &value, + vec![config.max_crops, config.positions, config.hidden], + ); + selected_values.push(builder.slice( + &value, + &[ + (0, config.max_crops), + (1, config.positions), + (0, config.hidden), + ], + )); + } + let mut selected = selected_values.remove(0); + for value in selected_values { + selected = builder.concatenate(&selected, &value, 2); + } + let selected = mask_padding_crops(&mut builder, &pixels, &selected, config); + let pad_embed = args.take(); + let selected = apply_pad_embed(&mut builder, &selected, &masks, &pad_embed); + let pooled = attention_pool(&mut builder, &mut args, config, &selected); + let pooled_rows = config.max_crops * config.projected_rows_per_crop(); + let pooled = builder.reshape(&pooled, vec![pooled_rows, config.hidden]); + let gate = quant_linear(&mut builder, &mut args, config, &pooled, false); + let gate = silu(&mut builder, &gate); + let up = quant_linear(&mut builder, &mut args, config, &pooled, false); + let activated = builder.multiply(&gate, &up); + let projected = quant_linear(&mut builder, &mut args, config, &activated, false); + let projected = builder.reshape( + &projected, + vec![ + config.max_crops, + config.projected_rows_per_crop(), + config.text_hidden, + ], + ); + assert_eq!(args.cursor, specs.len(), "Molmo vision schema drifted"); + let outputs = + if let (Some(patches), Some(mut selected)) = (diagnostic_patches, diagnostic_selected) { + let mut outputs = Vec::with_capacity(selected.len() + 2); + outputs.push(patches); + outputs.append(&mut selected); + outputs.push(projected); + outputs + } else { + vec![projected] + }; + let output_types = outputs + .iter() + .map(|value| value.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(|value| value.name.as_str()) + .collect::>() + .join(", "); + format!( + "module @molmo_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_molmo_vision(config: &MolmoVisionConfig) -> String { + emit_molmo_vision_impl(config, false) +} + +#[cfg(any(test, feature = "diagnostics"))] +pub(crate) fn emit_molmo_vision_diagnostics(config: &MolmoVisionConfig) -> String { + emit_molmo_vision_impl(config, true) +} + +#[cfg(test)] +#[path = "molmo_vision_tests.rs"] +mod tests; diff --git a/src/lib/mlxcel-xla/src/emitter/molmo_vision_config.rs b/src/lib/mlxcel-xla/src/emitter/molmo_vision_config.rs new file mode 100644 index 000000000..005481bfd --- /dev/null +++ b/src/lib/mlxcel-xla/src/emitter/molmo_vision_config.rs @@ -0,0 +1,547 @@ +// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Static Molmo v1 native vision/pool/projector artifact contract. + +use std::collections::BTreeSet; +use std::path::Path; + +use serde_json::Value; + +use crate::MOLMO_V1_MERGE_MODE; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum MolmoVisionWeightDType { + Float32, + Float16, + Uint32, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct MolmoVisionWeightSpec { + pub(crate) name: String, + pub(crate) shape: Vec, + pub(crate) dtype: MolmoVisionWeightDType, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct MolmoVisionConfig { + pub(crate) max_crops: usize, + pub(crate) patches_per_crop: usize, + pub(crate) patch_width: usize, + pub(crate) hidden: usize, + pub(crate) intermediate: usize, + pub(crate) heads: usize, + pub(crate) head_dim: usize, + pub(crate) positions: usize, + pub(crate) layers: usize, + pub(crate) selected_layers: Vec, + pub(crate) pool_h: usize, + pub(crate) pool_w: usize, + pub(crate) text_hidden: usize, + pub(crate) projector_hidden: usize, + pub(crate) group_size: usize, + pub(crate) bits: usize, + pub(crate) layer_norm_eps_bits: u32, + pub(crate) text_architecture: String, +} + +fn positive(value: Option, default: usize, name: &str) -> Result { + let value = value.unwrap_or(default as u64); + let value = usize::try_from(value).map_err(|_| format!("{name} does not fit usize"))?; + if value == 0 { + return Err(format!("{name} must be greater than zero")); + } + Ok(value) +} + +fn checked_product(values: &[usize], name: &str) -> Result { + values.iter().try_fold(1usize, |product, value| { + product + .checked_mul(*value) + .ok_or_else(|| format!("{name} overflows usize")) + }) +} + +impl MolmoVisionConfig { + pub(crate) fn from_model_dir(model_dir: &Path) -> Result { + let config_path = model_dir.join("config.json"); + let config = std::fs::read_to_string(&config_path) + .map_err(|error| format!("{}: {error}", config_path.display()))?; + let preprocessor_path = model_dir.join("preprocessor_config.json"); + let preprocessor = std::fs::read_to_string(&preprocessor_path) + .map_err(|error| format!("{}: {error}", preprocessor_path.display()))?; + Self::from_json_strs(&config, &preprocessor) + } + + pub(crate) fn from_json_strs(config: &str, preprocessor: &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("molmo") { + return Err("native Molmo vision requires model_type=molmo".to_string()); + } + let vision = root.get("vision_config").unwrap_or(&Value::Null); + let preprocess: Value = serde_json::from_str(preprocessor) + .map_err(|error| format!("parse preprocessor_config.json: {error}"))?; + + let patch_size = positive( + vision.get("image_patch_size").and_then(Value::as_u64), + 14, + "vision_config.image_patch_size", + )?; + let input_size = vision + .get("image_default_input_size") + .and_then(Value::as_array) + .map(|size| { + let height = positive( + size.first().and_then(Value::as_u64), + 336, + "image_default_input_size height", + )?; + let width = positive( + size.get(1).and_then(Value::as_u64), + 336, + "image_default_input_size width", + )?; + Ok::<_, String>((height, width)) + }) + .transpose()? + .unwrap_or((336, 336)); + if input_size.0 % patch_size != 0 || input_size.1 % patch_size != 0 { + return Err("Molmo crop size must be divisible by image_patch_size".to_string()); + } + let patch_h = input_size.0 / patch_size; + let patch_w = input_size.1 / patch_size; + let patches_per_crop = checked_product(&[patch_h, patch_w], "Molmo patches per crop")?; + let patch_width = checked_product(&[patch_size, patch_size, 3], "Molmo patch width")?; + let positions = patches_per_crop + .checked_add(1) + .ok_or_else(|| "Molmo vision position count overflows usize".to_string())?; + let pool_h = positive( + vision.get("image_pooling_h").and_then(Value::as_u64), + 2, + "vision_config.image_pooling_h", + )?; + let pool_w = positive( + vision.get("image_pooling_w").and_then(Value::as_u64), + 2, + "vision_config.image_pooling_w", + )?; + if !patch_h.is_multiple_of(pool_h) || !patch_w.is_multiple_of(pool_w) { + return Err("Molmo patch grid must tile evenly into pooling windows".to_string()); + } + + let layers = positive( + vision.get("image_num_layers").and_then(Value::as_u64), + 23, + "vision_config.image_num_layers", + )?; + let raw_layers = root + .get("vit_layers") + .or_else(|| vision.get("vit_layers")) + .and_then(Value::as_array) + .map(|layers| layers.iter().filter_map(Value::as_i64).collect::>()) + .unwrap_or_else(|| vec![-2, -9]); + if raw_layers.is_empty() { + return Err("Molmo vit_layers must not be empty".to_string()); + } + let layers_i64 = + i64::try_from(layers).map_err(|_| "Molmo layer count does not fit i64".to_string())?; + let mut selected_layers = Vec::with_capacity(raw_layers.len()); + let mut unique_layers = BTreeSet::new(); + for layer in raw_layers { + let resolved = if layer < 0 { layers_i64 + layer } else { layer }; + if resolved < 0 || resolved >= layers_i64 { + return Err(format!("Molmo vit_layers entry {layer} is out of range")); + } + let resolved = usize::try_from(resolved) + .map_err(|_| "Molmo selected layer does not fit usize".to_string())?; + if !unique_layers.insert(resolved) { + return Err("Molmo vit_layers must be unique".to_string()); + } + selected_layers.push(resolved); + } + + let hidden = positive( + vision.get("image_emb_dim").and_then(Value::as_u64), + 1024, + "vision_config.image_emb_dim", + )?; + let heads = positive( + vision.get("image_num_heads").and_then(Value::as_u64), + 16, + "vision_config.image_num_heads", + )?; + let head_dim = positive( + vision.get("image_head_dim").and_then(Value::as_u64), + 64, + "vision_config.image_head_dim", + )?; + if heads.checked_mul(head_dim) != Some(hidden) { + return Err( + "Molmo image_num_heads * image_head_dim must equal image_emb_dim".to_string(), + ); + } + let bits = positive( + root.get("quantization") + .and_then(|quant| quant.get("bits")) + .and_then(Value::as_u64), + 4, + "quantization.bits", + )?; + if bits != 4 && bits != 8 { + return Err(format!( + "Molmo native vision supports 4/8-bit affine weights, got {bits}" + )); + } + let group_size = positive( + root.get("quantization") + .and_then(|quant| quant.get("group_size")) + .and_then(Value::as_u64), + 64, + "quantization.group_size", + )?; + let text_hidden = positive( + root.get("hidden_size").and_then(Value::as_u64), + 3584, + "hidden_size", + )?; + let text_layers = positive( + root.get("num_hidden_layers").and_then(Value::as_u64), + 28, + "num_hidden_layers", + )?; + let text_heads = positive( + root.get("num_attention_heads").and_then(Value::as_u64), + 28, + "num_attention_heads", + )?; + let text_kv_heads = positive( + root.get("num_key_value_heads").and_then(Value::as_u64), + 4, + "num_key_value_heads", + )?; + if !text_hidden.is_multiple_of(text_heads) { + return Err("Molmo hidden_size must be divisible by num_attention_heads".to_string()); + } + if text_heads % text_kv_heads != 0 { + return Err( + "Molmo num_attention_heads must be divisible by num_key_value_heads".to_string(), + ); + } + let fused_intermediate = positive( + root.get("intermediate_size").and_then(Value::as_u64), + 37_888, + "intermediate_size", + )?; + if fused_intermediate % 2 != 0 { + return Err("Molmo intermediate_size must contain equal gate/up halves".to_string()); + } + if root.get("qkv_bias").and_then(Value::as_bool) == Some(false) { + return Err("pinned Molmo v1 native path requires qkv_bias=true".to_string()); + } + let rope_theta = root + .get("rope_theta") + .and_then(Value::as_f64) + .unwrap_or(1_000_000.0); + if !rope_theta.is_finite() || rope_theta <= 0.0 { + return Err("Molmo rope_theta must be finite and positive".to_string()); + } + let selected_width = checked_product( + &[hidden, selected_layers.len()], + "Molmo selected vision width", + )?; + let packed_group = 32usize + .checked_div(bits) + .ok_or_else(|| "Molmo quantization bits must divide 32".to_string())?; + for (name, width) in [ + ("vision hidden", hidden), + ("vision intermediate", 4096), + ("selected vision width", selected_width), + ("projector hidden", fused_intermediate / 2), + ] { + if !width.is_multiple_of(packed_group) { + return Err(format!( + "Molmo {name}={width} must be divisible by packed group {packed_group}" + )); + } + if !width.is_multiple_of(group_size) { + return Err(format!( + "Molmo {name}={width} must be divisible by quantization group_size={group_size}" + )); + } + } + let max_high_res_crops = positive( + preprocess.get("max_crops").and_then(Value::as_u64), + 12, + "preprocessor_config.max_crops", + )?; + let layer_norm_eps = vision + .get("image_norm_eps") + .and_then(Value::as_f64) + .unwrap_or(1e-5); + if !layer_norm_eps.is_finite() || layer_norm_eps <= 0.0 { + return Err("vision_config.image_norm_eps must be finite and positive".to_string()); + } + let max_crops = max_high_res_crops + .checked_add(1) + .ok_or_else(|| "Molmo max crop count overflows usize".to_string())?; + let text_layer_norm_eps = root + .get("layer_norm_eps") + .and_then(Value::as_f64) + .unwrap_or(1e-6); + if !text_layer_norm_eps.is_finite() || text_layer_norm_eps <= 0.0 { + return Err("Molmo layer_norm_eps must be finite and positive".to_string()); + } + let text_architecture = format!( + "olmo=hidden:{text_hidden},layers:{text_layers},heads:{text_heads},kv_heads:{text_kv_heads},\ + fused_intermediate:{fused_intermediate},rope_theta:{:016x},rope:interleave,\ + qkv_bias:true,norm:rms:{:016x},tied:false", + rope_theta.to_bits(), + text_layer_norm_eps.to_bits(), + ); + Ok(Self { + max_crops, + patches_per_crop, + patch_width, + hidden, + intermediate: 4096, + heads, + head_dim, + positions, + layers, + selected_layers, + pool_h, + pool_w, + text_hidden, + projector_hidden: fused_intermediate / 2, + group_size, + bits, + layer_norm_eps_bits: (layer_norm_eps as f32).to_bits(), + text_architecture, + }) + } + + pub(crate) fn emitted_layers(&self) -> usize { + self.selected_layers.iter().copied().max().unwrap_or(0) + 1 + } + + pub(crate) fn projected_rows_per_crop(&self) -> usize { + (self.patches_per_crop / self.pool_h) / self.pool_w + } + + pub(crate) fn fingerprint(&self) -> String { + format!( + "family=molmo-v1;crop={}x{}x{};max_crops={};vit={}x{}x{};layers={}/{};\ + selected={:?};pool={}x{};projector={}->{}->{};quant={}/{};merge={};eps={:08x};{}", + self.patches_per_crop, + self.patch_width, + self.hidden, + self.max_crops, + self.hidden, + self.heads, + self.head_dim, + self.emitted_layers(), + self.layers, + self.selected_layers, + self.pool_h, + self.pool_w, + self.hidden, + self.projector_hidden, + self.text_hidden, + self.bits, + self.group_size, + MOLMO_V1_MERGE_MODE, + self.layer_norm_eps_bits, + self.text_architecture, + ) + } + + pub(crate) fn weight_specs(&self) -> Vec { + let mut specs = Vec::new(); + push_f32( + &mut specs, + "vision_tower.image_vit.class_embedding", + &[self.hidden], + ); + push_f32( + &mut specs, + "vision_tower.image_vit.patch_embedding.weight", + &[self.hidden, self.patch_width], + ); + push_f32( + &mut specs, + "vision_tower.image_vit.positional_embedding", + &[self.positions, self.hidden], + ); + push_f32( + &mut specs, + "vision_tower.image_vit.pre_ln.weight", + &[self.hidden], + ); + push_f32( + &mut specs, + "vision_tower.image_vit.pre_ln.bias", + &[self.hidden], + ); + for layer in 0..self.emitted_layers() { + let prefix = format!("vision_tower.image_vit.transformer.resblocks.{layer}"); + push_f32( + &mut specs, + &format!("{prefix}.attention_norm.weight"), + &[self.hidden], + ); + push_f32( + &mut specs, + &format!("{prefix}.attention_norm.bias"), + &[self.hidden], + ); + for projection in ["wq", "wk", "wv", "wo"] { + push_quant( + &mut specs, + &format!("{prefix}.attention.{projection}"), + self.hidden, + self.hidden, + self.bits, + self.group_size, + ); + push_f32( + &mut specs, + &format!("{prefix}.attention.{projection}.bias"), + &[self.hidden], + ); + } + push_f32( + &mut specs, + &format!("{prefix}.ffn_norm.weight"), + &[self.hidden], + ); + push_f32( + &mut specs, + &format!("{prefix}.ffn_norm.bias"), + &[self.hidden], + ); + push_quant( + &mut specs, + &format!("{prefix}.feed_forward.w1"), + self.intermediate, + self.hidden, + self.bits, + self.group_size, + ); + push_f32( + &mut specs, + &format!("{prefix}.feed_forward.w1.bias"), + &[self.intermediate], + ); + push_quant( + &mut specs, + &format!("{prefix}.feed_forward.w2"), + self.hidden, + self.intermediate, + self.bits, + self.group_size, + ); + push_f32( + &mut specs, + &format!("{prefix}.feed_forward.w2.bias"), + &[self.hidden], + ); + } + let selected_width = self + .hidden + .checked_mul(self.selected_layers.len()) + .expect("validated selected vision width"); + push_f32(&mut specs, "vision_tower.pad_embed", &[2, selected_width]); + for projection in ["wq", "wk", "wv"] { + push_quant( + &mut specs, + &format!("vision_tower.image_pooling_2d.{projection}"), + self.hidden, + selected_width, + self.bits, + self.group_size, + ); + push_f32( + &mut specs, + &format!("vision_tower.image_pooling_2d.{projection}.bias"), + &[self.hidden], + ); + } + push_quant( + &mut specs, + "vision_tower.image_pooling_2d.wo", + self.hidden, + self.hidden, + self.bits, + self.group_size, + ); + push_f32( + &mut specs, + "vision_tower.image_pooling_2d.wo.bias", + &[self.hidden], + ); + for (name, output, input) in [ + ("w1", self.projector_hidden, self.hidden), + ("w3", self.projector_hidden, self.hidden), + ("w2", self.text_hidden, self.projector_hidden), + ] { + push_quant( + &mut specs, + &format!("vision_tower.image_projector.{name}"), + output, + input, + self.bits, + self.group_size, + ); + } + specs + } +} + +fn push_f32(specs: &mut Vec, name: &str, shape: &[usize]) { + specs.push(MolmoVisionWeightSpec { + name: name.to_string(), + shape: shape.to_vec(), + dtype: MolmoVisionWeightDType::Float32, + }); +} + +fn push_quant( + specs: &mut Vec, + prefix: &str, + output: usize, + input: usize, + bits: usize, + group_size: usize, +) { + let packed_input = input / (32 / bits); + let groups = input / group_size; + specs.push(MolmoVisionWeightSpec { + name: format!("{prefix}.weight"), + shape: vec![output, packed_input], + dtype: MolmoVisionWeightDType::Uint32, + }); + for suffix in ["scales", "biases"] { + specs.push(MolmoVisionWeightSpec { + name: format!("{prefix}.{suffix}"), + shape: vec![output, groups], + dtype: MolmoVisionWeightDType::Float16, + }); + } +} + +#[cfg(test)] +#[path = "molmo_vision_config_tests.rs"] +mod tests; diff --git a/src/lib/mlxcel-xla/src/emitter/molmo_vision_config_tests.rs b/src/lib/mlxcel-xla/src/emitter/molmo_vision_config_tests.rs new file mode 100644 index 000000000..104b3a1ef --- /dev/null +++ b/src/lib/mlxcel-xla/src/emitter/molmo_vision_config_tests.rs @@ -0,0 +1,89 @@ +use super::*; + +fn config() -> MolmoVisionConfig { + MolmoVisionConfig::from_json_strs( + r#"{ + "model_type":"molmo", + "hidden_size":3584, + "intermediate_size":37888, + "quantization":{"bits":4,"group_size":64}, + "vision_config":{} + }"#, + r#"{"max_crops":12}"#, + ) + .unwrap() +} + +#[test] +fn pinned_defaults_bind_sparse_add_artifact_identity() { + let config = config(); + assert_eq!(config.max_crops, 13); + assert_eq!(config.patches_per_crop, 576); + assert_eq!(config.patch_width, 588); + assert_eq!(config.selected_layers, vec![21, 14]); + assert_eq!(config.emitted_layers(), 22); + assert_eq!(config.projected_rows_per_crop(), 144); + assert!(config.fingerprint().contains(MOLMO_V1_MERGE_MODE)); + assert!(config.fingerprint().contains("olmo=hidden:3584,layers:28")); +} + +#[test] +fn schema_keeps_packed_quantized_weights_resident() { + let specs = config().weight_specs(); + assert_eq!(specs.len(), 647); + assert_eq!( + specs[1], + MolmoVisionWeightSpec { + name: "vision_tower.image_vit.patch_embedding.weight".to_string(), + shape: vec![1024, 588], + dtype: MolmoVisionWeightDType::Float32, + } + ); + assert!(specs.iter().any(|spec| { + spec.name == "vision_tower.image_projector.w2.weight" + && spec.shape == [3584, 2368] + && spec.dtype == MolmoVisionWeightDType::Uint32 + })); +} + +#[test] +fn rejects_duplicate_or_out_of_range_selected_layers() { + let duplicate = r#"{ + "model_type":"molmo", + "vision_config":{"image_num_layers":3}, + "vit_layers":[-1,2] + }"#; + assert!( + MolmoVisionConfig::from_json_strs(duplicate, r#"{"max_crops":1}"#) + .unwrap_err() + .contains("unique") + ); + let invalid = duplicate.replace("[-1,2]", "[-4]"); + assert!( + MolmoVisionConfig::from_json_strs(&invalid, r#"{"max_crops":1}"#) + .unwrap_err() + .contains("out of range") + ); + let non_adjacent = duplicate.replace("[-1,2]", "[-1,0,2]"); + assert!( + MolmoVisionConfig::from_json_strs(&non_adjacent, r#"{"max_crops":1}"#) + .unwrap_err() + .contains("unique") + ); +} + +#[test] +fn rejects_quantized_widths_that_would_truncate_weight_layouts() { + let invalid = r#"{ + "model_type":"molmo", + "hidden_size":3584, + "intermediate_size":37890, + "quantization":{"bits":4,"group_size":64}, + "vision_config":{} + }"#; + assert!( + MolmoVisionConfig::from_json_strs(invalid, r#"{"max_crops":12}"#) + .unwrap_err() + .contains("projector hidden") + ); +} diff --git a/src/lib/mlxcel-xla/src/emitter/molmo_vision_ops.rs b/src/lib/mlxcel-xla/src/emitter/molmo_vision_ops.rs new file mode 100644 index 000000000..bdfb2c07a --- /dev/null +++ b/src/lib/mlxcel-xla/src/emitter/molmo_vision_ops.rs @@ -0,0 +1,50 @@ +use super::builder::{Builder, Val}; + +pub(super) fn scalar_broadcast(builder: &mut Builder, value: f32, shape: Vec) -> Val { + let scalar = builder.const_f32(value); + builder.broadcast(&scalar, &[], shape) +} + +pub(super) fn tanh_gelu(builder: &mut Builder, value: &Val) -> Val { + let shape = value.ty.shape.clone(); + let half = scalar_broadcast(builder, 0.5, shape.clone()); + let one = scalar_broadcast(builder, 1.0, shape.clone()); + let coefficient = scalar_broadcast(builder, 0.044_715, shape.clone()); + let scale = scalar_broadcast(builder, 0.797_884_6, shape); + let squared = builder.multiply(value, value); + let cubed = builder.multiply(&squared, value); + let nonlinear = builder.multiply(&coefficient, &cubed); + let inner = builder.add(value, &nonlinear); + let scaled = builder.multiply(&scale, &inner); + let tanh = builder.tanh(&scaled); + let cdf = builder.add(&one, &tanh); + let half_value = builder.multiply(value, &half); + builder.multiply(&half_value, &cdf) +} + +pub(super) fn softmax_last(builder: &mut Builder, scores: &Val) -> Val { + let last = scores.ty.shape.len() - 1; + let negative_infinity = builder.const_f32(f32::NEG_INFINITY); + let maximum = builder.reduce_max(scores, last, &negative_infinity); + let mut broadcast_shape = maximum.ty.shape.clone(); + broadcast_shape.push(scores.ty.shape[last]); + let dimensions = (0..maximum.ty.shape.len()).collect::>(); + let maximum = builder.broadcast(&maximum, &dimensions, broadcast_shape); + let shifted = builder.subtract(scores, &maximum); + let exponentials = builder.exponential(&shifted); + let zero = builder.const_f32(0.0); + let denominator = builder.reduce_add(&exponentials, last, &zero); + let mut broadcast_shape = denominator.ty.shape.clone(); + broadcast_shape.push(scores.ty.shape[last]); + let dimensions = (0..denominator.ty.shape.len()).collect::>(); + let denominator = builder.broadcast(&denominator, &dimensions, broadcast_shape); + builder.divide(&exponentials, &denominator) +} + +pub(super) fn silu(builder: &mut Builder, value: &Val) -> Val { + let negated = builder.negate(value); + let exponential = builder.exponential(&negated); + let one = scalar_broadcast(builder, 1.0, value.ty.shape.clone()); + let denominator = builder.add(&one, &exponential); + builder.divide(value, &denominator) +} diff --git a/src/lib/mlxcel-xla/src/emitter/molmo_vision_tests.rs b/src/lib/mlxcel-xla/src/emitter/molmo_vision_tests.rs new file mode 100644 index 000000000..d8ae01e37 --- /dev/null +++ b/src/lib/mlxcel-xla/src/emitter/molmo_vision_tests.rs @@ -0,0 +1,75 @@ +use super::*; + +fn config() -> MolmoVisionConfig { + MolmoVisionConfig::from_json_strs( + r#"{ + "model_type":"molmo", + "hidden_size":3584, + "intermediate_size":37888, + "quantization":{"bits":4,"group_size":64}, + "vision_config":{} + }"#, + r#"{"max_crops":12}"#, + ) + .unwrap() +} + +#[test] +fn pinned_graph_contains_vit_pool_projector_and_static_output() { + let mlir = emit_molmo_vision(&config()); + assert!(mlir.contains("stablehlo.dot_general")); + assert!(mlir.contains("stablehlo.reduce")); + assert!(mlir.contains("stablehlo.exponential")); + assert!(mlir.contains("stablehlo.compare")); + assert!(mlir.contains("tensor<13x144x3584xf32>")); + assert!(mlir.contains("molmo.image_masks")); +} + +#[test] +fn diagnostics_return_patch_selected_layer_and_projector_boundaries() { + let diagnostic = emit_molmo_vision_diagnostics(&config()); + assert!(diagnostic.contains( + "-> (tensor<13x576x1024xf32>, tensor<7501x1024xf32>, \ + tensor<7501x1024xf32>, tensor<13x144x3584xf32>)" + )); + let production = emit_molmo_vision(&config()); + assert!(!production.contains("tensor<7501x1024xf32>, tensor<7501x1024xf32>")); +} + +#[cfg(feature = "iree")] +#[test] +fn pinned_graph_compiles_for_cpu() { + let mlir = emit_molmo_vision(&config()); + let compiler = crate::iree::iree_compile_bin().unwrap(); + let cache = std::env::temp_dir().join("mlxcel-xla-molmo-vision-emitter-test"); + std::fs::create_dir_all(&cache).unwrap(); + let vmfb = crate::iree::compile_one( + &compiler, + &mlir, + crate::iree::target_flags("local-task").unwrap(), + &cache, + "pinned-molmo-v1-vision", + 0, + ) + .unwrap(); + assert!(vmfb.metadata().unwrap().len() > 0); +} + +#[cfg(feature = "iree")] +#[test] +fn diagnostic_graph_compiles_for_cpu() { + let mlir = emit_molmo_vision_diagnostics(&config()); + let compiler = crate::iree::iree_compile_bin().unwrap(); + let cache = std::env::temp_dir().join("mlxcel-xla-molmo-vision-emitter-test"); + std::fs::create_dir_all(&cache).unwrap(); + let vmfb = crate::iree::compile_one( + &compiler, + &mlir, + crate::iree::target_flags("local-task").unwrap(), + &cache, + "pinned-molmo-v1-vision-diagnostics", + 0, + ) + .unwrap(); + assert!(vmfb.metadata().unwrap().len() > 0); +} diff --git a/src/lib/mlxcel-xla/src/float.rs b/src/lib/mlxcel-xla/src/float.rs new file mode 100644 index 000000000..4c71fe937 --- /dev/null +++ b/src/lib/mlxcel-xla/src/float.rs @@ -0,0 +1,82 @@ +// 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. + +//! Scalar floating-point conversions shared by feature-neutral host contracts. + +/// Round one f32 value to BF16 using round-to-nearest, ties-to-even, then widen +/// it back to an f32 carrier. +#[inline] +pub(crate) fn round_bf16_f32(value: f32) -> f32 { + let bits = value.to_bits(); + if bits & 0x7f80_0000 == 0x7f80_0000 { + return value; + } + let rounded = bits.wrapping_add(0x7fff + ((bits >> 16) & 1)); + f32::from_bits(rounded & 0xffff_0000) +} + +/// Widen one IEEE 754 half bit pattern to f32. +pub(crate) fn half_to_f32(half: u16) -> f32 { + let sign = if half >> 15 == 1 { -1.0 } else { 1.0 }; + let exponent = (half >> 10) & 0x1f; + let mantissa = (half & 0x3ff) as f32; + match exponent { + 0 => sign * mantissa * 2f32.powi(-24), + 0x1f if mantissa == 0.0 => sign * f32::INFINITY, + 0x1f => f32::NAN, + _ => sign * (1.0 + mantissa / 1024.0) * 2f32.powi(exponent as i32 - 15), + } +} + +/// Convert one f32 to an IEEE 754 half bit pattern using round-to-nearest, +/// ties-to-even. +pub(crate) fn f32_to_f16_bits(value: f32) -> u16 { + let bits = value.to_bits(); + let sign = ((bits >> 16) & 0x8000) as u16; + let absolute = bits & 0x7fff_ffff; + + if absolute >= 0x7f80_0000 { + return sign + | if absolute > 0x7f80_0000 { + 0x7e00 + } else { + 0x7c00 + }; + } + + let exponent = (absolute >> 23) as i32 - 127 + 15; + if exponent >= 0x1f { + return sign | 0x7c00; + } + + if exponent <= 0 { + if exponent < -10 { + return sign; + } + let mantissa = (absolute & 0x007f_ffff) | 0x0080_0000; + let shift = (14 - exponent) as u32; + let quotient = mantissa >> shift; + let remainder = mantissa & ((1 << shift) - 1); + let halfway = 1u32 << (shift - 1); + let round = u32::from(remainder > halfway || (remainder == halfway && quotient & 1 == 1)); + return sign | (quotient + round) as u16; + } + + let mantissa = absolute & 0x007f_ffff; + let base = ((exponent as u32) << 10) | (mantissa >> 13); + let remainder = mantissa & 0x1fff; + let halfway = 0x1000u32; + let round = u32::from(remainder > halfway || (remainder == halfway && base & 1 == 1)); + sign | (base + round) as u16 +} diff --git a/src/lib/mlxcel-xla/src/iree.rs b/src/lib/mlxcel-xla/src/iree.rs index e31100190..7a0490783 100644 --- a/src/lib/mlxcel-xla/src/iree.rs +++ b/src/lib/mlxcel-xla/src/iree.rs @@ -1658,6 +1658,27 @@ fn load_weights( checked_ffi_i64(*end - *start, &format!("weight {name} sliced rows"))?; dims[i * 4 + 1] = checked_ffi_i64(shape[1], &format!("weight {name} dim 1"))?; } + WeightSpec::Elements { start, end, .. } => { + if shape.len() != 1 { + return Err(format!( + "element-slice weight {name} is rank {} (expected 1)", + shape.len() + )); + } + let sliced = data.get(*start..*end).ok_or_else(|| { + format!( + "element-slice {name}: range [{start},{end}) is outside length {}", + shape[0] + ) + })?; + bufs[i] = WeightBuf::F32(sliced.to_vec()); + ranks[i] = 1; + let length = end.checked_sub(*start).ok_or_else(|| { + format!("element-slice {name}: start {start} exceeds end {end}") + })?; + dims[i * 4] = + checked_ffi_i64(length, &format!("weight {name} sliced elements"))?; + } WeightSpec::QuantRaw { .. } => { unreachable!("QuantRaw is handled by the early-continue above") } diff --git a/src/lib/mlxcel-xla/src/lib.rs b/src/lib/mlxcel-xla/src/lib.rs index 2b1328267..062c667df 100644 --- a/src/lib/mlxcel-xla/src/lib.rs +++ b/src/lib/mlxcel-xla/src/lib.rs @@ -52,6 +52,7 @@ use mlxcel_core::session::{InferenceSession, PreparedPrefill, SessionCapabilitie mod context; #[cfg(any(feature = "diagnostics", test))] mod diagnostic_flags; +mod float; #[cfg(any(feature = "iree", test))] #[allow(dead_code)] mod numeric_dtype_contract; @@ -67,6 +68,7 @@ mod operator_numeric_contract; mod prepared; mod prepared_deepstack; mod prepared_gemma3n; +mod prepared_molmo; #[cfg(feature = "iree")] mod aux; @@ -77,6 +79,8 @@ mod aux_smoke; #[cfg(feature = "iree")] mod iree; #[cfg(feature = "iree")] +mod molmo_vision_runtime; +#[cfg(feature = "iree")] mod phi4_audio; #[cfg(feature = "iree")] mod qwen2_vl_runtime; @@ -168,6 +172,10 @@ pub use emitter::{ }; #[cfg(feature = "diagnostics")] pub use iree::PreparedPrefillDiagnostics; +#[cfg(feature = "diagnostics")] +pub use molmo_vision_runtime::{IreeMolmoVisionDiagnosticProjector, MolmoVisionDiagnostics}; +#[cfg(feature = "iree")] +pub use molmo_vision_runtime::{IreeMolmoVisionProjector, MolmoVisionProjection}; #[cfg(feature = "iree")] pub use phi4_audio::{ PHI4MM_AUDIO_CHECKPOINT_REVISION, PHI4MM_AUDIO_FRAME_BUCKETS, Phi4AudioOutput, @@ -176,6 +184,10 @@ pub use phi4_audio::{ #[cfg(feature = "iree")] #[doc(hidden)] pub use phi4_audio::{Phi4AudioCheckpoint, Phi4AudioDiagnosticRuntime, Phi4AudioDiagnostics}; +pub use prepared_molmo::{ + MOLMO_V1_MERGE_MODE, MolmoEmbeddingDType, MolmoPreparedImage, MolmoSparseAddError, + MolmoSparseAddPair, MolmoSparseAddPlan, +}; #[cfg(feature = "iree")] pub use qwen2_vl_runtime::{ IreeQwen2VlProjector, Qwen2VlVisionExecutionMetrics, Qwen2VlVisionProjection, diff --git a/src/lib/mlxcel-xla/src/molmo_vision_diagnostics.rs b/src/lib/mlxcel-xla/src/molmo_vision_diagnostics.rs new file mode 100644 index 000000000..5466ff8ab --- /dev/null +++ b/src/lib/mlxcel-xla/src/molmo_vision_diagnostics.rs @@ -0,0 +1,191 @@ +// 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::*; +use crate::emitter::emit_molmo_vision_diagnostics; + +/// Diagnostics-only outputs from the resident Molmo v1 vision graph. +#[derive(Debug, Clone, PartialEq)] +pub struct MolmoVisionDiagnostics { + pub patch_embeddings: Vec, + pub patch_shape: [usize; 3], + pub selected_hidden_states: Vec>, + pub selected_shape: [usize; 3], + pub selected_layers: Vec, + pub projected_features: Vec, + pub projected_shape: [usize; 3], + pub elapsed_seconds: f64, + pub upload_bytes: usize, + pub transfer_bytes: usize, +} + +/// Diagnostics-only Molmo graph. Production continues to load the one-output +/// [`IreeMolmoVisionProjector`] artifact. +pub struct IreeMolmoVisionDiagnosticProjector { + module: IreeAuxiliaryModule, + config: MolmoVisionConfig, +} + +impl IreeMolmoVisionDiagnosticProjector { + pub fn load(model_dir: &Path, device: &str) -> Result { + let config = MolmoVisionConfig::from_model_dir(model_dir)?; + let mlir = emit_molmo_vision_diagnostics(&config); + let module = compile_and_load( + model_dir, + device, + &config, + &mlir, + "molmo-v1-vision-diagnostics", + )?; + Ok(Self { module, config }) + } + + #[must_use] + pub fn artifact_fingerprint(&self) -> u64 { + self.module.fingerprint() + } + + pub fn project( + &mut self, + pixels: &[f32], + masks: &[f32], + crop_count: usize, + ) -> Result { + let input = validate_and_pad_input(&self.config, pixels, masks, crop_count)?; + let patch_shape = [ + self.config.max_crops, + self.config.patches_per_crop, + self.config.hidden, + ]; + let selected_shape = [ + self.config.max_crops * self.config.positions, + self.config.hidden, + ]; + let projected_shape = [ + self.config.max_crops, + self.config.projected_rows_per_crop(), + self.config.text_hidden, + ]; + let mut shapes = Vec::with_capacity(self.config.selected_layers.len() + 2); + shapes.push(patch_shape.to_vec()); + shapes.extend((0..self.config.selected_layers.len()).map(|_| selected_shape.to_vec())); + shapes.push(projected_shape.to_vec()); + let mut buffers = shapes + .iter() + .map(|shape| { + checked_product(shape, "diagnostic output").and_then(|elements| { + elements + .checked_mul(std::mem::size_of::()) + .map(|bytes| vec![0u8; bytes]) + .ok_or_else(|| { + "Molmo diagnostic output byte count overflows usize".to_string() + }) + }) + }) + .collect::, String>>()?; + let transfer_bytes = buffers.iter().map(Vec::len).sum(); + let mut outputs = buffers + .iter_mut() + .zip(&shapes) + .map(|(bytes, shape)| AuxiliaryOutput { + bytes, + dtype: AuxiliaryTensorDType::Float32, + shape, + }) + .collect::>(); + let started = Instant::now(); + self.module.invoke( + &[ + AuxiliaryInput { + bytes: f32_as_bytes(&input.padded_pixels), + dtype: AuxiliaryTensorDType::Float32, + shape: &input.pixel_shape, + }, + AuxiliaryInput { + bytes: f32_as_bytes(&input.padded_masks), + dtype: AuxiliaryTensorDType::Float32, + shape: &input.mask_shape, + }, + ], + &mut outputs, + )?; + let elapsed_seconds = started.elapsed().as_secs_f64(); + drop(outputs); + let values = buffers + .into_iter() + .enumerate() + .map(|(index, bytes)| { + checked_f32_output(&format!("IREE Molmo diagnostic output {index}"), bytes) + }) + .collect::, String>>()?; + let mut values = values.into_iter(); + let mut patch_embeddings = values + .next() + .expect("Molmo diagnostics include patch embeddings"); + let active_patch_count = checked_product( + &[ + input.crop_count, + self.config.patches_per_crop, + self.config.hidden, + ], + "active diagnostic patch output", + )?; + patch_embeddings.truncate(active_patch_count); + let active_selected_count = checked_product( + &[input.crop_count, self.config.positions, self.config.hidden], + "active diagnostic selected output", + )?; + let selected_hidden_states = values + .by_ref() + .take(self.config.selected_layers.len()) + .map(|mut state| { + state.truncate(active_selected_count); + state + }) + .collect::>(); + let mut projected_features = values + .next() + .expect("Molmo diagnostics include projected features"); + let active_projected_count = checked_product( + &[ + input.crop_count, + self.config.projected_rows_per_crop(), + self.config.text_hidden, + ], + "active diagnostic projected output", + )?; + projected_features.truncate(active_projected_count); + debug_assert!(values.next().is_none()); + Ok(MolmoVisionDiagnostics { + patch_embeddings, + patch_shape: [ + input.crop_count, + self.config.patches_per_crop, + self.config.hidden, + ], + selected_hidden_states, + selected_shape: [input.crop_count, self.config.positions, self.config.hidden], + selected_layers: self.config.selected_layers.clone(), + projected_features, + projected_shape: [ + input.crop_count, + self.config.projected_rows_per_crop(), + self.config.text_hidden, + ], + elapsed_seconds, + upload_bytes: input.upload_bytes, + transfer_bytes, + }) + } +} diff --git a/src/lib/mlxcel-xla/src/molmo_vision_runtime.rs b/src/lib/mlxcel-xla/src/molmo_vision_runtime.rs new file mode 100644 index 000000000..ff45818bf --- /dev/null +++ b/src/lib/mlxcel-xla/src/molmo_vision_runtime.rs @@ -0,0 +1,421 @@ +// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Resident packed-weight IREE runtime for Molmo v1 vision preprocessing. + +use std::collections::{BTreeMap, BTreeSet}; +use std::fs::File; +use std::path::{Path, PathBuf}; +use std::time::Instant; + +use memmap2::Mmap; +use safetensors::{Dtype, SafeTensors}; + +use crate::aux::{ + AuxiliaryInput, AuxiliaryOutput, AuxiliaryTensorDType, AuxiliaryWeight, AuxiliaryWeightDType, + IreeAuxiliaryModule, +}; +use crate::aux_manifest::{AuxiliaryArtifactContract, ensure_qualified_auxiliary_artifact}; +use crate::emitter::{ + MolmoVisionConfig, MolmoVisionWeightDType, MolmoVisionWeightSpec, emit_molmo_vision, +}; +use crate::iree::{cached_vmfb_path, compile_one_to, iree_compile_bin, target_flags}; +use crate::vision_runtime::{ + checked_f32_output, compiler_generation_identity, f32_as_bytes, model_shards, native_f32_bytes, + sha256_hex, validate_finite_values, +}; +use crate::weights::{bf16_to_f32, f16_to_f32, f32_le_to_f32}; + +const ENTRY_NAME: &str = "molmo_vision.main"; + +fn checked_product(dimensions: &[usize], label: &str) -> Result { + dimensions.iter().try_fold(1usize, |count, dimension| { + count + .checked_mul(*dimension) + .ok_or_else(|| format!("Molmo {label} shape overflows usize")) + }) +} + +#[derive(Debug, Clone, PartialEq)] +pub struct MolmoVisionProjection { + pub values: Vec, + pub shape: [usize; 3], + pub elapsed_seconds: f64, + pub upload_bytes: usize, + pub transfer_bytes: usize, +} + +fn resolve_weight_shards( + model_dir: &Path, + specs: &[MolmoVisionWeightSpec], +) -> 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: this read-only map lives through the header scan. + let mmap = unsafe { Mmap::map(&file) } + .map_err(|error| format!("mmap {}: {error}", shard.display()))?; + let tensors = SafeTensors::deserialize(&mmap) + .map_err(|error| format!("parse {}: {error}", shard.display()))?; + for name in tensors.names() { + if required.contains(name) + && let Some(previous) = locations.insert(name.to_string(), shard.clone()) + { + return Err(format!( + "Molmo vision tensor {name:?} is duplicated 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 {} Molmo vision tensor(s): {}", + missing.len(), + missing + .iter() + .take(8) + .copied() + .collect::>() + .join(", ") + )); + } + Ok(specs + .iter() + .map(|spec| locations[&spec.name].clone()) + .collect()) +} + +fn load_weights( + model_dir: &Path, + specs: &[MolmoVisionWeightSpec], +) -> Result<(Vec, String), String> { + let shards = resolve_weight_shards(model_dir, specs)?; + let mut by_shard = BTreeMap::<&Path, Vec>::new(); + for (index, shard) in shards.iter().enumerate() { + by_shard.entry(shard.as_path()).or_default().push(index); + } + let mut loaded = (0..specs.len()) + .map(|_| None) + .collect::>>(); + let mut source_schema = vec![String::new(); specs.len()]; + for (shard, indices) in by_shard { + let file = + File::open(shard).map_err(|error| format!("open {}: {error}", shard.display()))?; + // Safety: this map lives while every selected tensor is 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!("resolved Molmo vision tensor {}: {error}", spec.name))?; + if tensor.shape() != spec.shape { + return Err(format!( + "Molmo vision tensor {} has shape {:?}, expected {:?}", + spec.name, + tensor.shape(), + spec.shape + )); + } + let (bytes, dtype) = match spec.dtype { + MolmoVisionWeightDType::Float32 => { + 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()), + other => { + return Err(format!( + "Molmo tensor {} has {other:?}, expected a floating dtype", + spec.name + )); + } + }; + validate_finite_values(&spec.name, &values)?; + (native_f32_bytes(values), AuxiliaryWeightDType::Float32) + } + MolmoVisionWeightDType::Float16 => { + if tensor.dtype() != Dtype::F16 { + return Err(format!( + "Molmo tensor {} has {:?}, expected F16", + spec.name, + tensor.dtype() + )); + } + (tensor.data().to_vec(), AuxiliaryWeightDType::Float16) + } + MolmoVisionWeightDType::Uint32 => { + if tensor.dtype() != Dtype::U32 { + return Err(format!( + "Molmo tensor {} has {:?}, expected U32", + spec.name, + tensor.dtype() + )); + } + (tensor.data().to_vec(), AuxiliaryWeightDType::Uint32) + } + }; + source_schema[index] = + format!("{}:{:?}:{:?}", spec.name, tensor.dtype(), tensor.shape()); + loaded[index] = Some(AuxiliaryWeight { + name: spec.name.clone(), + bytes, + dtype, + shape: spec.shape.clone(), + }); + } + } + Ok(( + loaded + .into_iter() + .map(|weight| weight.expect("all resolved tensors loaded")) + .collect(), + source_schema.join("\n"), + )) +} + +fn compile_and_load( + model_dir: &Path, + device: &str, + config: &MolmoVisionConfig, + mlir: &str, + tag: &str, +) -> Result { + let compiler = iree_compile_bin()?; + let flags = target_flags(device)?; + let cache = std::env::temp_dir().join("mlxcel-xla-molmo-vision-vmfb"); + std::fs::create_dir_all(&cache) + .map_err(|error| format!("mkdir {}: {error}", cache.display()))?; + let specs = config.weight_specs(); + let (weights, checkpoint_schema) = load_weights(model_dir, &specs)?; + let processor_path = model_dir.join("preprocessor_config.json"); + let processor = std::fs::read(&processor_path) + .map_err(|error| format!("{}: {error}", processor_path.display()))?; + let contract = AuxiliaryArtifactContract::new_legacy_unqualified( + ENTRY_NAME, + format!( + "{};processor_sha256={};checkpoint_schema_sha256={}", + config.fingerprint(), + sha256_hex(&processor), + sha256_hex(checkpoint_schema.as_bytes()) + ), + compiler_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) +} + +struct ValidatedMolmoVisionInput { + padded_pixels: Vec, + padded_masks: Vec, + pixel_shape: [usize; 3], + mask_shape: [usize; 2], + crop_count: usize, + upload_bytes: usize, +} + +fn validate_and_pad_input( + config: &MolmoVisionConfig, + pixels: &[f32], + masks: &[f32], + crop_count: usize, +) -> Result { + if crop_count == 0 || crop_count > config.max_crops { + return Err(format!( + "Molmo crop count {crop_count} is outside 1..={}", + config.max_crops + )); + } + let pixel_count = crop_count + .checked_mul(config.patches_per_crop) + .and_then(|count| count.checked_mul(config.patch_width)) + .ok_or_else(|| "Molmo pixel shape overflowed".to_string())?; + let mask_count = crop_count + .checked_mul(config.patches_per_crop) + .ok_or_else(|| "Molmo mask shape overflowed".to_string())?; + if pixels.len() != pixel_count || masks.len() != mask_count { + return Err(format!( + "Molmo processor counts pixels={}/{} masks={}/{}", + pixels.len(), + pixel_count, + masks.len(), + mask_count + )); + } + validate_finite_values("Molmo pixel_values", pixels)?; + validate_finite_values("Molmo image_masks", masks)?; + if let Some((index, value)) = masks + .iter() + .enumerate() + .find(|(_, value)| !(**value >= -1.0 && **value <= 1.0)) + { + return Err(format!( + "Molmo image_masks[{index}]={value} is outside [-1,1]" + )); + } + let pixel_shape = [ + config.max_crops, + config.patches_per_crop, + config.patch_width, + ]; + let mask_shape = [config.max_crops, config.patches_per_crop]; + let static_pixels = checked_product(&pixel_shape, "static pixel")?; + let static_masks = checked_product(&mask_shape, "static mask")?; + let mut padded_pixels = vec![-1.0f32; static_pixels]; + padded_pixels[..pixels.len()].copy_from_slice(pixels); + let mut padded_masks = vec![-1.0f32; static_masks]; + padded_masks[..masks.len()].copy_from_slice(masks); + let upload_bytes = static_pixels + .checked_add(static_masks) + .and_then(|count| count.checked_mul(std::mem::size_of::())) + .ok_or_else(|| "Molmo invocation upload byte count overflows usize".to_string())?; + Ok(ValidatedMolmoVisionInput { + padded_pixels, + padded_masks, + pixel_shape, + mask_shape, + crop_count, + upload_bytes, + }) +} + +pub struct IreeMolmoVisionProjector { + module: IreeAuxiliaryModule, + config: MolmoVisionConfig, +} + +impl IreeMolmoVisionProjector { + pub fn load(model_dir: &Path, device: &str) -> Result { + let config = MolmoVisionConfig::from_model_dir(model_dir)?; + let mlir = emit_molmo_vision(&config); + let module = compile_and_load(model_dir, device, &config, &mlir, "molmo-v1-vision")?; + Ok(Self { module, config }) + } + + #[must_use] + pub fn max_crops(&self) -> usize { + self.config.max_crops + } + + #[must_use] + pub fn patches_per_crop(&self) -> usize { + self.config.patches_per_crop + } + + #[must_use] + pub fn patch_width(&self) -> usize { + self.config.patch_width + } + + #[must_use] + pub fn projected_rows_per_crop(&self) -> usize { + self.config.projected_rows_per_crop() + } + + #[must_use] + pub fn text_hidden(&self) -> usize { + self.config.text_hidden + } + + #[must_use] + pub fn artifact_fingerprint(&self) -> u64 { + self.module.fingerprint() + } + + pub fn project( + &mut self, + pixels: &[f32], + masks: &[f32], + crop_count: usize, + ) -> Result { + let input = validate_and_pad_input(&self.config, pixels, masks, crop_count)?; + let output_shape = [ + self.config.max_crops, + self.config.projected_rows_per_crop(), + self.config.text_hidden, + ]; + let output_elements = checked_product(&output_shape, "projected output")?; + let output_bytes = output_elements + .checked_mul(std::mem::size_of::()) + .ok_or_else(|| "Molmo projected output byte count overflows usize".to_string())?; + let mut output = vec![0u8; output_bytes]; + let started = Instant::now(); + self.module.invoke( + &[ + AuxiliaryInput { + bytes: f32_as_bytes(&input.padded_pixels), + dtype: AuxiliaryTensorDType::Float32, + shape: &input.pixel_shape, + }, + AuxiliaryInput { + bytes: f32_as_bytes(&input.padded_masks), + dtype: AuxiliaryTensorDType::Float32, + shape: &input.mask_shape, + }, + ], + &mut [AuxiliaryOutput { + bytes: &mut output, + dtype: AuxiliaryTensorDType::Float32, + shape: &output_shape, + }], + )?; + let elapsed_seconds = started.elapsed().as_secs_f64(); + let mut values = checked_f32_output("IREE Molmo projected output", output)?; + let active_count = checked_product( + &[ + input.crop_count, + self.config.projected_rows_per_crop(), + self.config.text_hidden, + ], + "active projected output", + )?; + values.truncate(active_count); + let transfer_bytes = active_count + .checked_mul(std::mem::size_of::()) + .ok_or_else(|| "Molmo invocation transfer byte count overflows usize".to_string())?; + Ok(MolmoVisionProjection { + values, + shape: [ + input.crop_count, + self.config.projected_rows_per_crop(), + self.config.text_hidden, + ], + elapsed_seconds, + upload_bytes: input.upload_bytes, + transfer_bytes, + }) + } +} + +#[cfg(feature = "diagnostics")] +#[path = "molmo_vision_diagnostics.rs"] +mod diagnostics; +#[cfg(feature = "diagnostics")] +pub use diagnostics::{IreeMolmoVisionDiagnosticProjector, MolmoVisionDiagnostics}; diff --git a/src/lib/mlxcel-xla/src/prepared_molmo.rs b/src/lib/mlxcel-xla/src/prepared_molmo.rs new file mode 100644 index 000000000..fedec4c64 --- /dev/null +++ b/src/lib/mlxcel-xla/src/prepared_molmo.rs @@ -0,0 +1,470 @@ +// 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. + +//! Owned Molmo v1 image input and sparse additive embedding merge contract. +//! +//! Molmo v1 is deliberately distinct from token-scanned VLM merges. Every +//! `image_input_idx` entry corresponds to the projected feature row at the same +//! flat index. Negative entries are padding sentinels: removing one must not +//! renumber later feature rows. Valid target positions receive an addition to +//! the post-scale OLMo text embedding, never a replacement. + +use std::collections::HashSet; +use std::fmt; + +/// Artifact-identity spelling for the only qualified Molmo v1 merge. +pub const MOLMO_V1_MERGE_MODE: &str = "processor-indexed-sparse-add-v1"; + +/// One active processor-supplied sparse merge pair. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct MolmoSparseAddPair { + /// Row in the projected feature tensor before sentinel removal. + pub feature_row: usize, + /// Position in the expanded OLMo text sequence. + pub target_position: usize, +} + +/// Owned processor output retained across the host/compiler boundary. +#[derive(Debug, Clone, PartialEq)] +pub struct MolmoPreparedImage { + pub pixel_values: Vec, + pub crop_count: usize, + pub patches_per_crop: usize, + pub patch_width: usize, + pub image_masks: Vec, + /// Rows emitted by attention pooling for each crop. + pub projected_rows_per_crop: usize, + /// Authoritative mapping, including negative sentinels. + pub image_input_idx: Vec, +} + +/// A validated sparse-add plan whose row mapping cannot drift after filtering. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct MolmoSparseAddPlan { + pairs: Vec, + feature_rows: usize, + text_len: usize, + hidden_size: usize, +} + +/// Floating-point storage dtype of the post-scale OLMo text embeddings. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum MolmoEmbeddingDType { + Float16, + BFloat16, + Float32, +} + +impl MolmoEmbeddingDType { + fn round(self, value: f32) -> f32 { + match self { + Self::Float16 => crate::float::half_to_f32(crate::float::f32_to_f16_bits(value)), + Self::BFloat16 => crate::float::round_bf16_f32(value), + Self::Float32 => value, + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum MolmoSparseAddError { + ZeroDimension(&'static str), + Capacity { + dimension: &'static str, + actual: usize, + maximum: usize, + }, + ShapeOverflow, + PixelCount { + expected: usize, + actual: usize, + }, + MaskCount { + expected: usize, + actual: usize, + }, + MaskValue { + flat_index: usize, + value_bits: u32, + }, + IndexCount { + expected: usize, + actual: usize, + }, + TargetOutOfRange { + feature_row: usize, + target_position: i32, + text_len: usize, + }, + DuplicateTarget { + first_feature_row: usize, + feature_row: usize, + target_position: usize, + }, + TensorCount { + tensor: &'static str, + expected: usize, + actual: usize, + }, + NonFinite { + tensor: &'static str, + flat_index: usize, + }, +} + +impl fmt::Display for MolmoSparseAddError { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::ZeroDimension(name) => write!(formatter, "Molmo {name} must be non-zero"), + Self::Capacity { + dimension, + actual, + maximum, + } => write!( + formatter, + "Molmo {dimension} {actual} exceeds static maximum {maximum}" + ), + Self::ShapeOverflow => formatter.write_str("Molmo tensor shape overflowed"), + Self::PixelCount { expected, actual } => write!( + formatter, + "Molmo pixel_values has {actual} elements, expected {expected}" + ), + Self::MaskCount { expected, actual } => write!( + formatter, + "Molmo image_masks has {actual} elements, expected {expected}" + ), + Self::MaskValue { + flat_index, + value_bits, + } => write!( + formatter, + "Molmo image_masks[{flat_index}]={} is outside [-1,1]", + f32::from_bits(*value_bits) + ), + Self::IndexCount { expected, actual } => write!( + formatter, + "Molmo image_input_idx has {actual} rows, expected {expected}" + ), + Self::TargetOutOfRange { + feature_row, + target_position, + text_len, + } => write!( + formatter, + "Molmo image_input_idx feature row {feature_row} targets {target_position}, outside [0,{text_len})" + ), + Self::DuplicateTarget { + first_feature_row, + feature_row, + target_position, + } => write!( + formatter, + "Molmo image_input_idx rows {first_feature_row} and {feature_row} both target {target_position}" + ), + Self::TensorCount { + tensor, + expected, + actual, + } => write!( + formatter, + "Molmo {tensor} has {actual} elements, expected {expected}" + ), + Self::NonFinite { tensor, flat_index } => { + write!(formatter, "Molmo {tensor}[{flat_index}] is not finite") + } + } + } +} + +impl std::error::Error for MolmoSparseAddError {} + +fn checked_product(dimensions: &[usize]) -> Result { + dimensions.iter().try_fold(1usize, |size, dimension| { + size.checked_mul(*dimension) + .ok_or(MolmoSparseAddError::ShapeOverflow) + }) +} + +fn validate_capacity( + dimension: &'static str, + actual: usize, + maximum: usize, +) -> Result<(), MolmoSparseAddError> { + if actual > maximum { + return Err(MolmoSparseAddError::Capacity { + dimension, + actual, + maximum, + }); + } + Ok(()) +} + +impl MolmoPreparedImage { + /// Validate all processor shapes and static buckets before native execution. + pub fn validate( + &self, + max_crops: usize, + max_patches_per_crop: usize, + max_feature_rows: usize, + ) -> Result<(), MolmoSparseAddError> { + for (name, value) in [ + ("crop count", self.crop_count), + ("patches per crop", self.patches_per_crop), + ("patch width", self.patch_width), + ("projected rows per crop", self.projected_rows_per_crop), + ] { + if value == 0 { + return Err(MolmoSparseAddError::ZeroDimension(name)); + } + } + validate_capacity("crop count", self.crop_count, max_crops)?; + validate_capacity( + "patches per crop", + self.patches_per_crop, + max_patches_per_crop, + )?; + let patch_rows = checked_product(&[self.crop_count, self.patches_per_crop])?; + let feature_rows = checked_product(&[self.crop_count, self.projected_rows_per_crop])?; + validate_capacity("feature rows", feature_rows, max_feature_rows)?; + let expected_pixels = + checked_product(&[self.crop_count, self.patches_per_crop, self.patch_width])?; + if self.pixel_values.len() != expected_pixels { + return Err(MolmoSparseAddError::PixelCount { + expected: expected_pixels, + actual: self.pixel_values.len(), + }); + } + if self.image_masks.len() != patch_rows { + return Err(MolmoSparseAddError::MaskCount { + expected: patch_rows, + actual: self.image_masks.len(), + }); + } + if self.image_input_idx.len() != feature_rows { + return Err(MolmoSparseAddError::IndexCount { + expected: feature_rows, + actual: self.image_input_idx.len(), + }); + } + for (flat_index, value) in self.pixel_values.iter().enumerate() { + if !value.is_finite() { + return Err(MolmoSparseAddError::NonFinite { + tensor: "pixel_values", + flat_index, + }); + } + } + for (flat_index, value) in self.image_masks.iter().enumerate() { + if !value.is_finite() { + return Err(MolmoSparseAddError::NonFinite { + tensor: "image_masks", + flat_index, + }); + } + if !(-1.0..=1.0).contains(value) { + return Err(MolmoSparseAddError::MaskValue { + flat_index, + value_bits: value.to_bits(), + }); + } + } + Ok(()) + } + + /// Validate the processor payload and build its authoritative sparse plan. + #[allow(clippy::too_many_arguments)] + pub fn sparse_add_plan( + &self, + text_len: usize, + hidden_size: usize, + max_crops: usize, + max_patches_per_crop: usize, + max_feature_rows: usize, + max_text_len: usize, + ) -> Result { + self.validate(max_crops, max_patches_per_crop, max_feature_rows)?; + MolmoSparseAddPlan::from_image_input_idx( + &self.image_input_idx, + text_len, + hidden_size, + max_feature_rows, + max_text_len, + ) + } +} + +impl MolmoSparseAddPlan { + /// Build the authoritative mapping from the processor output. + /// + /// Sentinel rows are skipped, but their indices remain consumed so a later + /// active entry keeps its original projected feature row. + pub fn from_image_input_idx( + image_input_idx: &[i32], + text_len: usize, + hidden_size: usize, + max_feature_rows: usize, + max_text_len: usize, + ) -> Result { + if text_len == 0 { + return Err(MolmoSparseAddError::ZeroDimension("text length")); + } + if hidden_size == 0 { + return Err(MolmoSparseAddError::ZeroDimension("hidden size")); + } + validate_capacity("feature rows", image_input_idx.len(), max_feature_rows)?; + validate_capacity("expanded text length", text_len, max_text_len)?; + checked_product(&[text_len, hidden_size])?; + checked_product(&[image_input_idx.len(), hidden_size])?; + + let mut pairs = Vec::new(); + let mut targets = std::collections::HashMap::new(); + for (feature_row, &target) in image_input_idx.iter().enumerate() { + if target < 0 { + continue; + } + let target_position = + usize::try_from(target).map_err(|_| MolmoSparseAddError::TargetOutOfRange { + feature_row, + target_position: target, + text_len, + })?; + if target_position >= text_len { + return Err(MolmoSparseAddError::TargetOutOfRange { + feature_row, + target_position: target, + text_len, + }); + } + if let Some(first_feature_row) = targets.insert(target_position, feature_row) { + return Err(MolmoSparseAddError::DuplicateTarget { + first_feature_row, + feature_row, + target_position, + }); + } + pairs.push(MolmoSparseAddPair { + feature_row, + target_position, + }); + } + debug_assert_eq!( + pairs + .iter() + .map(|pair| pair.target_position) + .collect::>() + .len(), + pairs.len() + ); + Ok(Self { + pairs, + feature_rows: image_input_idx.len(), + text_len, + hidden_size, + }) + } + + #[must_use] + pub fn pairs(&self) -> &[MolmoSparseAddPair] { + &self.pairs + } + + /// Add F32 projected rows to already post-scaled F32 OLMo embeddings. + pub fn apply( + &self, + text_embeddings: &mut [f32], + projected_features: &[f32], + ) -> Result<(), MolmoSparseAddError> { + self.apply_in_dtype( + text_embeddings, + projected_features, + MolmoEmbeddingDType::Float32, + ) + } + + /// Cast projected rows to the text-embedding dtype, add in that dtype, and + /// widen the rounded result to F32 for the prepared-prefill boundary. + pub fn apply_in_dtype( + &self, + text_embeddings: &mut [f32], + projected_features: &[f32], + dtype: MolmoEmbeddingDType, + ) -> Result<(), MolmoSparseAddError> { + let expected_text = checked_product(&[self.text_len, self.hidden_size])?; + let expected_features = checked_product(&[self.feature_rows, self.hidden_size])?; + for (tensor, expected, actual) in [ + ("text embeddings", expected_text, text_embeddings.len()), + ( + "projected features", + expected_features, + projected_features.len(), + ), + ] { + if expected != actual { + return Err(MolmoSparseAddError::TensorCount { + tensor, + expected, + actual, + }); + } + } + if let Some((flat_index, _)) = text_embeddings + .iter() + .enumerate() + .find(|(_, value)| !value.is_finite()) + { + return Err(MolmoSparseAddError::NonFinite { + tensor: "text embeddings", + flat_index, + }); + } + if let Some((flat_index, _)) = projected_features + .iter() + .enumerate() + .find(|(_, value)| !value.is_finite()) + { + return Err(MolmoSparseAddError::NonFinite { + tensor: "projected features", + flat_index, + }); + } + for pair in &self.pairs { + let text_start = pair.target_position * self.hidden_size; + let feature_start = pair.feature_row * self.hidden_size; + for offset in 0..self.hidden_size { + let text = dtype.round(text_embeddings[text_start + offset]); + let feature = dtype.round(projected_features[feature_start + offset]); + if !dtype.round(text + feature).is_finite() { + return Err(MolmoSparseAddError::NonFinite { + tensor: "merged embeddings", + flat_index: text_start + offset, + }); + } + } + } + for pair in &self.pairs { + let text_start = pair.target_position * self.hidden_size; + let feature_start = pair.feature_row * self.hidden_size; + for offset in 0..self.hidden_size { + let text = dtype.round(text_embeddings[text_start + offset]); + let feature = dtype.round(projected_features[feature_start + offset]); + text_embeddings[text_start + offset] = dtype.round(text + feature); + } + } + Ok(()) + } +} + +#[cfg(test)] +#[path = "prepared_molmo_tests.rs"] +mod tests; diff --git a/src/lib/mlxcel-xla/src/prepared_molmo_tests.rs b/src/lib/mlxcel-xla/src/prepared_molmo_tests.rs new file mode 100644 index 000000000..42c88d027 --- /dev/null +++ b/src/lib/mlxcel-xla/src/prepared_molmo_tests.rs @@ -0,0 +1,175 @@ +// 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 sentinels_preserve_original_feature_rows_and_non_contiguous_targets() { + let plan = MolmoSparseAddPlan::from_image_input_idx(&[-100, 3, -1, 0], 5, 2, 8, 8).unwrap(); + assert_eq!( + plan.pairs(), + [ + MolmoSparseAddPair { + feature_row: 1, + target_position: 3, + }, + MolmoSparseAddPair { + feature_row: 3, + target_position: 0, + }, + ] + ); + let mut text = vec![10.0; 10]; + let features = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]; + plan.apply(&mut text, &features).unwrap(); + assert_eq!( + text, + [17.0, 18.0, 10.0, 10.0, 10.0, 10.0, 13.0, 14.0, 10.0, 10.0] + ); +} + +#[test] +fn rejects_duplicate_and_out_of_range_targets_before_merge() { + assert!(matches!( + MolmoSparseAddPlan::from_image_input_idx(&[1, -100, 1], 3, 2, 4, 4), + Err(MolmoSparseAddError::DuplicateTarget { .. }) + )); + assert!(matches!( + MolmoSparseAddPlan::from_image_input_idx(&[-1, 3], 3, 2, 4, 4), + Err(MolmoSparseAddError::TargetOutOfRange { .. }) + )); +} + +#[test] +fn empty_active_features_are_valid_and_leave_text_unchanged() { + let plan = MolmoSparseAddPlan::from_image_input_idx(&[-100, -1], 2, 2, 4, 4).unwrap(); + let mut text = vec![1.0, 2.0, 3.0, 4.0]; + plan.apply(&mut text, &[5.0, 6.0, 7.0, 8.0]).unwrap(); + assert_eq!(text, [1.0, 2.0, 3.0, 4.0]); +} + +#[test] +fn replacement_negative_fixture_is_caught_by_additive_oracle() { + fn oracle_logits(embeddings: &[f32]) -> [f32; 2] { + let active = &embeddings[2..4]; + [active[0] + 2.0 * active[1], 3.0 * active[0] - active[1]] + } + + let plan = MolmoSparseAddPlan::from_image_input_idx(&[1], 2, 2, 1, 2).unwrap(); + let original = vec![2.0, 3.0, 5.0, 7.0]; + let features = vec![11.0, 13.0]; + let mut additive = original.clone(); + plan.apply(&mut additive, &features).unwrap(); + + let mut replacement = original; + replacement[2..4].copy_from_slice(&features); + assert_eq!(additive, [2.0, 3.0, 16.0, 20.0]); + assert_ne!(replacement, additive); + assert_eq!(oracle_logits(&additive), [56.0, 28.0]); + assert_ne!( + oracle_logits(&replacement), + oracle_logits(&additive), + "the downstream logit oracle must reject replacement semantics" + ); +} + +#[test] +fn sparse_add_rounds_projected_and_sum_in_text_embedding_dtype() { + let plan = MolmoSparseAddPlan::from_image_input_idx(&[-100, 0], 1, 1, 2, 1).unwrap(); + let projected = [100.0, 0.0006]; + + let mut f16_text = [1.0]; + plan.apply_in_dtype(&mut f16_text, &projected, MolmoEmbeddingDType::Float16) + .unwrap(); + assert_eq!(f16_text, [1.000_976_6]); + assert_ne!(f16_text[0], 1.0 + projected[1]); + + let mut bf16_text = [1.0]; + plan.apply_in_dtype( + &mut bf16_text, + &[100.0, 0.005], + MolmoEmbeddingDType::BFloat16, + ) + .unwrap(); + assert_eq!(bf16_text, [1.007_812_5]); + assert_ne!(bf16_text[0], 1.0 + 0.005); +} + +#[test] +fn capacity_shape_and_nonfinite_fail_closed() { + assert!(matches!( + MolmoSparseAddPlan::from_image_input_idx(&[0, 1], 2, 2, 1, 2), + Err(MolmoSparseAddError::Capacity { + dimension: "feature rows", + .. + }) + )); + let plan = MolmoSparseAddPlan::from_image_input_idx(&[0], 1, 2, 1, 1).unwrap(); + assert!(matches!( + plan.apply(&mut [0.0, 0.0], &[f32::NAN, 1.0]), + Err(MolmoSparseAddError::NonFinite { + tensor: "projected features", + .. + }) + )); +} + +#[test] +fn processor_patch_rows_and_projected_rows_have_distinct_shapes() { + let prepared = prepared_image(1.0); + let plan = prepared.sparse_add_plan(4, 2, 2, 4, 2, 4).unwrap(); + assert_eq!( + plan.pairs(), + [MolmoSparseAddPair { + feature_row: 1, + target_position: 2, + }] + ); +} + +#[test] +fn invalid_mask_domain_and_late_overflow_are_atomic_failures() { + let mut invalid_mask = prepared_image(1.5); + assert!(matches!( + invalid_mask.validate(2, 4, 2), + Err(MolmoSparseAddError::MaskValue { .. }) + )); + invalid_mask.image_masks.fill(-1.0); + invalid_mask.validate(2, 4, 2).unwrap(); + + let plan = MolmoSparseAddPlan::from_image_input_idx(&[0, 1], 2, 1, 2, 2).unwrap(); + let original = vec![1.0, f32::MAX]; + let mut text = original.clone(); + let error = plan.apply(&mut text, &[2.0, f32::MAX]).unwrap_err(); + assert!(matches!( + error, + MolmoSparseAddError::NonFinite { + tensor: "merged embeddings", + flat_index: 1, + } + )); + assert_eq!(text, original, "failed merge must not modify any row"); +} + +fn prepared_image(mask: f32) -> MolmoPreparedImage { + MolmoPreparedImage { + pixel_values: vec![0.0; 2 * 4 * 3], + crop_count: 2, + patches_per_crop: 4, + patch_width: 3, + image_masks: vec![mask; 2 * 4], + projected_rows_per_crop: 1, + image_input_idx: vec![-100, 2], + } +} diff --git a/src/lib/mlxcel-xla/src/vision_runtime.rs b/src/lib/mlxcel-xla/src/vision_runtime.rs index b16bc9d62..e0e00483f 100644 --- a/src/lib/mlxcel-xla/src/vision_runtime.rs +++ b/src/lib/mlxcel-xla/src/vision_runtime.rs @@ -77,7 +77,7 @@ pub struct VisionDiagnosticProjection { pub metrics: VisionExecutionMetrics, } -fn sha256_hex(bytes: &[u8]) -> String { +pub(crate) fn sha256_hex(bytes: &[u8]) -> String { hex_bytes(&Sha256::digest(bytes)) } @@ -90,7 +90,7 @@ fn hex_bytes(bytes: &[u8]) -> String { output } -fn compiler_generation_identity( +pub(crate) fn compiler_generation_identity( compiler: &Path, flags: &[&str], mlir: &str, @@ -344,7 +344,7 @@ impl VisionProcessorContract { } } -fn model_shards(model_dir: &Path) -> Result, String> { +pub(crate) 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| { @@ -426,7 +426,7 @@ fn resolve_weight_shards( .collect()) } -fn native_f32_bytes(values: Vec) -> Vec { +pub(crate) fn native_f32_bytes(values: Vec) -> Vec { let mut bytes = Vec::with_capacity(values.len() * std::mem::size_of::()); for value in values { bytes.extend_from_slice(&value.to_ne_bytes()); @@ -597,7 +597,7 @@ fn compile_and_load( IreeAuxiliaryModule::load(device, &vmfb, &contract, weights) } -fn f32_as_bytes(values: &[f32]) -> &[u8] { +pub(crate) fn f32_as_bytes(values: &[f32]) -> &[u8] { // Safety: f32 has no invalid bit patterns and the byte slice cannot outlive // the immutable source slice. unsafe { @@ -605,7 +605,7 @@ fn f32_as_bytes(values: &[f32]) -> &[u8] { } } -fn validate_finite_values(label: &str, values: &[f32]) -> Result<(), String> { +pub(crate) fn validate_finite_values(label: &str, values: &[f32]) -> Result<(), String> { if let Some((index, value)) = values .iter() .enumerate() @@ -618,7 +618,7 @@ fn validate_finite_values(label: &str, values: &[f32]) -> Result<(), String> { Ok(()) } -fn checked_f32_output(label: &str, bytes: Vec) -> Result, String> { +pub(crate) fn checked_f32_output(label: &str, bytes: Vec) -> Result, String> { if !bytes.len().is_multiple_of(std::mem::size_of::()) { return Err(format!( "{label} returned {} bytes, which is not a whole number of f32 values", diff --git a/src/lib/mlxcel-xla/src/weight_names.rs b/src/lib/mlxcel-xla/src/weight_names.rs index ab7b20dc6..8bdbcb516 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::Molmo => SchemeNames { + embed: "language_model.model.wte.embedding", + final_norm: "language_model.model.ln_f.weight", + lm_head: "language_model.model.ff_out.weight", + layer_stem: "language_model.model.blocks.", + down: "ff_out.weight", + gate: "ff_proj.weight", + input_layernorm: "attn_norm.weight", + post_attention_layernorm: "ff_norm.weight", + up: "ff_proj.weight", + k_proj: "att_proj.weight", + o_proj: "attn_out.weight", + q_proj: "att_proj.weight", + v_proj: "att_proj.weight", + k_bias: "att_proj.bias", + q_bias: "att_proj.bias", + v_bias: "att_proj.bias", + q_norm: "unused.q_norm.weight", + k_norm: "unused.k_norm.weight", + pre_ff_norm: "unused.pre_ff_norm.weight", + post_ff_norm: "unused.post_ff_norm.weight", + }, } } diff --git a/src/lib/mlxcel-xla/src/weights.rs b/src/lib/mlxcel-xla/src/weights.rs index 22338ab7e..5778ad0e6 100644 --- a/src/lib/mlxcel-xla/src/weights.rs +++ b/src/lib/mlxcel-xla/src/weights.rs @@ -49,6 +49,13 @@ pub(crate) enum WeightSpec { start: usize, end: usize, }, + /// Load elements `[start, end)` of a rank-1 tensor. Molmo v1 uses this for + /// the bias paired with its fused `[Q|K|V]` attention projection. + Elements { + name: String, + start: usize, + end: usize, + }, /// Upload one part of an MLX affine-quantized projection as RAW bytes without /// dequantizing (issue #516 packed path): the packed `[out, in_packed]` U32 /// weight, or its `[out, in/group_size]` f16 `scales` / `biases`. The graph @@ -74,6 +81,7 @@ impl WeightSpec { WeightSpec::Whole(n) => n, WeightSpec::Proj(n) => n, WeightSpec::Rows { name, .. } => name, + WeightSpec::Elements { name, .. } => name, WeightSpec::QuantRaw { name, .. } => name, } } @@ -180,10 +188,22 @@ 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 fused = if cfg.weight_scheme == crate::emitter::WeightScheme::Molmo { + format!("{p}{}", s.gate) + } else { + format!("{p}mlp.gate_up_proj.weight") + }; + let (start, end) = if cfg.weight_scheme == crate::emitter::WeightScheme::Molmo { + // Molmo stores ff_proj as [value | gate], while the + // shared emitter applies SiLU to `gate`. + (inter, 2 * inter) + } else { + (0, inter) + }; 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, fused), + start, + end, }); } else { push_proj( @@ -208,10 +228,22 @@ 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 fused = if cfg.weight_scheme == crate::emitter::WeightScheme::Molmo { + format!("{p}{}", s.up) + } else { + format!("{p}mlp.gate_up_proj.weight") + }; + let (start, end) = if cfg.weight_scheme == crate::emitter::WeightScheme::Molmo { + // The first Molmo half is the multiplicative value + // (`x_part` in the MLX reference), not the SiLU gate. + (0, inter) + } else { + (inter, 2 * inter) + }; 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, fused), + start, + end, }); } else if cfg.dense_mlp { out.push(WeightSpec::Whole(format!("{p}mlp.c_fc.weight"))); @@ -225,7 +257,12 @@ 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 fused = if cfg.weight_scheme == crate::emitter::WeightScheme::Molmo { + format!("{p}{}", s.q_proj) + } else { + format!("{p}self_attn.qkv_proj.weight") + }; + let qkv = phi4_base_projection(cfg, fused); out.push(WeightSpec::Rows { name: qkv.clone(), start: nq, @@ -269,9 +306,28 @@ fn weight_specs_q(cfg: &Config, quant: bool) -> Vec { } // q/k/v projection biases (k, q, v order). if cfg.qkv_bias { - out.push(WeightSpec::Whole(format!("{p}{}", s.k_bias))); - out.push(WeightSpec::Whole(format!("{p}{}", s.q_bias))); - out.push(WeightSpec::Whole(format!("{p}{}", s.v_bias))); + if cfg.fused_qkv && cfg.weight_scheme == crate::emitter::WeightScheme::Molmo { + let bias = format!("{p}{}", s.q_bias); + out.push(WeightSpec::Elements { + name: bias.clone(), + start: nq, + end: nq + nkv, + }); + out.push(WeightSpec::Elements { + name: bias.clone(), + start: 0, + end: nq, + }); + out.push(WeightSpec::Elements { + name: bias, + start: nq + nkv, + end: nq + 2 * nkv, + }); + } else { + out.push(WeightSpec::Whole(format!("{p}{}", s.k_bias))); + out.push(WeightSpec::Whole(format!("{p}{}", s.q_bias))); + out.push(WeightSpec::Whole(format!("{p}{}", s.v_bias))); + } } // q/k norms (Qwen3 / Gemma3 per-head, OLMo2/3 flat), q then k. if cfg.qk_norm.is_some() { @@ -393,32 +449,7 @@ pub(crate) fn bf16_to_f32(bytes: &[u8]) -> Vec { .collect() } -/// Round one f32 value to BF16 using round-to-nearest, ties-to-even, then widen -/// it back to an f32 carrier. -#[inline] -pub(crate) fn round_bf16_f32(value: f32) -> f32 { - let bits = value.to_bits(); - if bits & 0x7f80_0000 == 0x7f80_0000 { - return value; - } - let rounded = bits.wrapping_add(0x7fff + ((bits >> 16) & 1)); - f32::from_bits(rounded & 0xffff_0000) -} - -/// One IEEE 754 half (f16) -> f32. The arithmetic forms are exact: a normal's -/// `1 + mant/1024` is a dyadic with denominator 2^10 and the `2^(exp-15)` / `2^-24` -/// scales are exact powers of two, so the widening is bit-for-bit. -pub(crate) fn half_to_f32(h: u16) -> f32 { - let sign = if h >> 15 == 1 { -1.0 } else { 1.0 }; - let exp = (h >> 10) & 0x1f; - let mant = (h & 0x3ff) as f32; - match exp { - 0 => sign * mant * 2f32.powi(-24), // zero / subnormal - 0x1f if mant == 0.0 => sign * f32::INFINITY, // +/- inf - 0x1f => f32::NAN, // nan - _ => sign * (1.0 + mant / 1024.0) * 2f32.powi(exp as i32 - 15), // normal - } -} +pub(crate) use crate::float::{f32_to_f16_bits, half_to_f32, round_bf16_f32}; /// f16 little-endian bytes -> f32, via a 65536-entry `u16 -> f32` lookup table. /// The table is built once (every f16 bit pattern, exact) and then each element @@ -441,57 +472,6 @@ pub(crate) fn f32_le_to_f32(bytes: &[u8]) -> Vec { .collect() } -/// One f32 -> IEEE 754 half (f16) bit pattern, round-to-nearest, ties-to-even (the -/// IEEE default), matching a `stablehlo.convert` f32 -> f16. So a projection weight -/// packed here (issue #572, f16-resident) is bit-identical to demoting the same f32 -/// weight inside the graph, and the contraction sees the same f16 operand and stays -/// token-exact. It is the exact inverse of [`half_to_f32`]: -/// `f32_to_f16_bits(half_to_f32(h)) == h` for every finite, non-NaN f16 `h`. -pub(crate) fn f32_to_f16_bits(x: f32) -> u16 { - let bits = x.to_bits(); - let sign = ((bits >> 16) & 0x8000) as u16; - let abs = bits & 0x7fff_ffff; - - // NaN / Inf (f32 exponent all ones): NaN -> a canonical quiet f16 NaN, Inf -> f16 Inf. - if abs >= 0x7f80_0000 { - return sign | if abs > 0x7f80_0000 { 0x7e00 } else { 0x7c00 }; - } - - // f16 biased exponent = (f32 biased exponent - 127) + 15. - let e = (abs >> 23) as i32 - 127 + 15; - - if e >= 0x1f { - return sign | 0x7c00; // overflow -> Inf - } - - if e <= 0 { - // Subnormal f16, or underflow to a signed zero. - if e < -10 { - return sign; // below half the smallest subnormal -> +/- 0 - } - // 24-bit significand (implicit leading 1), shifted into the subnormal range - // and rounded to nearest, ties to even. - let mant = (abs & 0x007f_ffff) | 0x0080_0000; - let shift = (14 - e) as u32; // e in [-10, 0] -> shift in [14, 24] - let q = mant >> shift; - let rem = mant & ((1 << shift) - 1); - let half = 1u32 << (shift - 1); - let round = u32::from(rem > half || (rem == half && q & 1 == 1)); - // q + round may reach 0x400, which is exactly the smallest normal (correct). - return sign | (q + round) as u16; - } - - // Normal f16: keep the top 10 mantissa bits, round to nearest even on bit 12. A - // mantissa carry rolls into the exponent, an exponent carry into 0x7c00 (Inf) -- - // both the correct results. - let mant = abs & 0x007f_ffff; - let base = ((e as u32) << 10) | (mant >> 13); - let rem = mant & 0x1fff; // the 13 dropped low bits - let half = 0x1000u32; // 1 << 12 - let round = u32::from(rem > half || (rem == half && base & 1 == 1)); - sign | (base + round) as u16 -} - /// Pack a row-major f32 weight to its little-endian f16 bit pattern for an /// f16-resident device upload (issue #572), via [`f32_to_f16_bits`] (RNE, matching /// the in-graph demotion). The `u16` values are native-endian; the shim copies the @@ -739,7 +719,83 @@ pub(crate) fn dequantize_affine_stacked( #[cfg(test)] mod tests { use super::*; - use crate::emitter::Phi4LoraConfig; + use crate::emitter::{Phi4LoraConfig, WeightScheme}; + + #[test] + fn molmo_v1_maps_fused_olmo_weights_without_decoder_duplication() { + let config = Config::from_json_str( + r#"{ + "model_type":"molmo", + "hidden_size":8, + "intermediate_size":24, + "num_hidden_layers":1, + "num_attention_heads":2, + "num_key_value_heads":1, + "layer_norm_eps":0.00001, + "rope_theta":1000000.0, + "vocab_size":32, + "qkv_bias":true, + "tie_word_embeddings":false, + "quantization":{"bits":4,"group_size":4} + }"#, + ) + .unwrap(); + assert_eq!(config.weight_scheme, WeightScheme::Molmo); + assert_eq!(config.inter, 12, "Molmo config stores fused gate+up width"); + assert!(config.fused_qkv && config.fused_gate_up); + assert!(config.qkv_bias && config.rope_interleaved); + + let specs = weight_specs_q(&config, false); + assert_eq!( + specs[0], + WeightSpec::Whole("language_model.model.wte.embedding".to_string()) + ); + assert_eq!( + specs[2], + WeightSpec::Whole("language_model.model.ff_out.weight".to_string()) + ); + let rows = |start, end| WeightSpec::Rows { + name: "language_model.model.blocks.0.att_proj.weight".to_string(), + start, + end, + }; + assert!(specs.contains(&rows(8, 12))); + assert!(specs.contains(&rows(0, 8))); + assert!(specs.contains(&rows(12, 16))); + let fused_mlp_slices = specs + .iter() + .filter(|spec| spec.tensor_name() == "language_model.model.blocks.0.ff_proj.weight") + .collect::>(); + assert_eq!( + fused_mlp_slices, + [ + &WeightSpec::Rows { + name: "language_model.model.blocks.0.ff_proj.weight".to_string(), + start: 12, + end: 24, + }, + &WeightSpec::Rows { + name: "language_model.model.blocks.0.ff_proj.weight".to_string(), + start: 0, + end: 12, + }, + ], + "the shared emitter takes gate then value, while Molmo stores value then gate" + ); + let value = 1.0f32; + let gate = 4.0f32; + let expected = gate / (1.0 + (-gate).exp()) * value; + let reversed = value / (1.0 + (-value).exp()) * gate; + assert_ne!( + expected, reversed, + "the fixture must detect applying SiLU to the wrong fused half" + ); + assert!(specs.contains(&WeightSpec::Elements { + name: "language_model.model.blocks.0.att_proj.bias".to_string(), + start: 8, + end: 12, + })); + } fn phi4_config() -> Config { Config { diff --git a/src/loading/mod.rs b/src/loading/mod.rs index fb2cab9d4..3c1da6d9d 100644 --- a/src/loading/mod.rs +++ b/src/loading/mod.rs @@ -59,6 +59,10 @@ use self::vlm::*; pub(crate) use self::vlm::load_llava_host_preprocessor; #[cfg(feature = "xla-iree")] pub(crate) use self::vlm::load_llava_iree_host_preprocessor; +#[cfg(feature = "xla-backend")] +pub(crate) use self::vlm::load_molmo_host_preprocessor; +#[cfg(feature = "xla-iree")] +pub(crate) use self::vlm::load_molmo_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; diff --git a/src/loading/vlm.rs b/src/loading/vlm.rs index b87c6c0a2..b9642affe 100644 --- a/src/loading/vlm.rs +++ b/src/loading/vlm.rs @@ -71,6 +71,9 @@ mod llava; mod minimax_m3_vl; #[path = "vlm_mllama.rs"] mod mllama; +#[cfg(feature = "xla-backend")] +#[path = "vlm_molmo_xla.rs"] +mod molmo_xla; #[path = "vlm_nemotron_h_nano_omni.rs"] mod nemotron_h_nano_omni; #[path = "vlm_paddleocr.rs"] @@ -111,6 +114,10 @@ pub(crate) use llava::load_llava_iree_host_preprocessor; pub(crate) use llava::{load_llava_bunny_vlm, load_llava_host_preprocessor, load_llava_vlm}; pub(crate) use minimax_m3_vl::load_minimax_m3_vl; pub(crate) use mllama::load_mllama_vlm; +#[cfg(feature = "xla-backend")] +pub(crate) use molmo_xla::load_molmo_host_preprocessor; +#[cfg(feature = "xla-iree")] +pub(crate) use molmo_xla::load_molmo_iree_host_preprocessor; pub(crate) use nemotron_h_nano_omni::load_nemotron_h_nano_omni_vlm; pub(crate) use paddleocr::load_paddleocr_vl; pub(crate) use pixtral::{load_mistral3_vlm, load_pixtral_vlm}; diff --git a/src/loading/vlm_molmo_xla.rs b/src/loading/vlm_molmo_xla.rs new file mode 100644 index 000000000..0b93a2447 --- /dev/null +++ b/src/loading/vlm_molmo_xla.rs @@ -0,0 +1,317 @@ +// 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. + +//! Filtered Molmo v1 host producer loader for XLA prepared prefills. + +use std::path::Path; + +use mlxcel_core::weights::WeightMap; +use serde_json::Value; + +use crate::models; +use crate::multimodal::host_preprocessor::{HostPreprocessorError, MolmoHostPreprocessor}; +use crate::vision::encoders::molmo::{MolmoVisionConfig, MolmoVisionModel}; +use crate::vision::processors::molmo::{MolmoImageTokens, MolmoProcessor}; + +use super::special::{ + inherit_quantization_if_missing, molmo_vision_i32, read_clip_triple, rewrite_molmo_weight_key, +}; +use super::{ + load_vlm_weights_common_filtered_canonical, read_optional_model_json, read_sanitized_vlm_config, +}; + +/// Load only Molmo's dual WTE, vision stack, tokenizer, and processor. +pub(crate) fn load_molmo_host_preprocessor( + model_path: &Path, +) -> Result { + let (_config_str, full_config) = read_sanitized_vlm_config(model_path) + .map_err(|error| HostPreprocessorError::InvalidConfig(error.to_string()))?; + if full_config.get("model_type").and_then(Value::as_str) != Some("molmo") { + return Err(HostPreprocessorError::FamilyMismatch { + actual: full_config + .get("model_type") + .and_then(Value::as_str) + .unwrap_or("unknown") + .to_string(), + }); + } + + 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) + .map_err(|error| HostPreprocessorError::InvalidConfig(error.to_string()))?; + let text_config: models::molmo::MolmoTextConfig = serde_json::from_value(text_config_value) + .map_err(|error| { + HostPreprocessorError::InvalidConfig(format!( + "failed to parse Molmo v1 text config: {error}" + )) + })?; + + let empty = Value::Object(Default::default()); + let vision_config = full_config.get("vision_config").unwrap_or(&empty); + let mut vit_layers = full_config + .get("vit_layers") + .or_else(|| vision_config.get("vit_layers")) + .and_then(Value::as_array) + .map(|layers| { + layers + .iter() + .filter_map(|layer| layer.as_i64().map(|layer| layer as i32)) + .collect::>() + }) + .unwrap_or_else(|| vec![-2, -9]); + if vit_layers.is_empty() { + vit_layers = vec![-2, -9]; + } + let image_patch_size = molmo_vision_i32(vision_config, "image_patch_size", 14) as usize; + let (input_h, input_w) = vision_config + .get("image_default_input_size") + .and_then(Value::as_array) + .and_then(|dimensions| { + Some(( + dimensions.first()?.as_i64()? as i32, + dimensions.get(1)?.as_i64()? as i32, + )) + }) + .unwrap_or((336, 336)); + let patch_h = input_h / image_patch_size as i32; + let patch_w = input_w / image_patch_size as i32; + let pool_h = molmo_vision_i32(vision_config, "image_pooling_h", 2); + let pool_w = molmo_vision_i32(vision_config, "image_pooling_w", 2); + if patch_h <= 0 || patch_w <= 0 || pool_h <= 0 || pool_w <= 0 { + return Err(HostPreprocessorError::InvalidConfig( + "Molmo v1 patch and pooling dimensions must be positive".to_string(), + )); + } + let vision_cfg = MolmoVisionConfig { + image_num_layers: molmo_vision_i32(vision_config, "image_num_layers", 23) as usize, + image_emb_dim: molmo_vision_i32(vision_config, "image_emb_dim", 1024), + image_num_heads: molmo_vision_i32(vision_config, "image_num_heads", 16), + image_num_kv_heads: molmo_vision_i32(vision_config, "image_num_key_value_heads", 16), + image_head_dim: molmo_vision_i32(vision_config, "image_head_dim", 64), + image_num_pos: molmo_vision_i32(vision_config, "image_num_pos", 577) as usize, + image_norm_eps: vision_config + .get("image_norm_eps") + .and_then(Value::as_f64) + .unwrap_or(1e-5) as f32, + image_num_patch: (patch_h, patch_w), + image_pooling_h: pool_h, + image_pooling_w: pool_w, + vit_layers, + group_size: text_config.group_size(), + bits: text_config.bits(), + }; + + let raw_weights = load_vlm_weights_common_filtered_canonical(model_path, |name| { + name.starts_with("vision_tower.") + || name.starts_with("model.vision_backbone.") + || name.starts_with("language_model.model.wte.") + || name.starts_with("model.transformer.wte.") + }) + .map_err(|error| HostPreprocessorError::WeightLoad(error.to_string()))?; + let mut weights = WeightMap::new(); + for (name, value) in raw_weights { + weights.insert(rewrite_molmo_weight_key(&name), value); + } + let text_embeddings = + models::molmo::Molmo2Embedding::from_weights(&weights, "language_model.model.wte") + .map_err(|error| HostPreprocessorError::WeightLoad(error.to_string()))?; + let vision_tower = MolmoVisionModel::from_weights(&weights, "vision_tower", vision_cfg) + .map_err(HostPreprocessorError::WeightLoad)?; + let tokenizer = crate::tokenizer::load_tokenizer(model_path) + .map_err(|error| HostPreprocessorError::InvalidConfig(error.to_string()))?; + + let preprocessor_config = read_optional_model_json(model_path, "preprocessor_config.json"); + let max_high_res_crops = preprocessor_config + .as_ref() + .and_then(|config| config.get("max_crops")) + .and_then(Value::as_u64) + .unwrap_or(12) as usize; + let overlap = preprocessor_config + .as_ref() + .and_then(|config| config.get("overlap_margins")) + .and_then(Value::as_array) + .and_then(|values| { + Some(( + values.first()?.as_u64()? as usize, + values.get(1)?.as_u64()? as usize, + )) + }); + let base_size = preprocessor_config + .as_ref() + .and_then(|config| config.get("base_image_input_size")) + .and_then(Value::as_array) + .and_then(|values| { + Some(( + values.first()?.as_u64()? as usize, + values.get(1)?.as_u64()? as usize, + )) + }); + let token_len = preprocessor_config.as_ref().and_then(|config| { + Some(( + config.get("image_token_length_h")?.as_u64()? as usize, + config.get("image_token_length_w")?.as_u64()? as usize, + )) + }); + let processor = MolmoProcessor::new( + max_high_res_crops, + overlap, + Some(image_patch_size), + base_size, + token_len, + read_clip_triple(preprocessor_config.as_ref(), "image_mean"), + read_clip_triple(preprocessor_config.as_ref(), "image_std"), + MolmoImageTokens::default(), + ); + let max_crops = max_high_res_crops + .checked_add(1) + .ok_or(HostPreprocessorError::ShapeOverflow)?; + let patches_per_crop = usize::try_from(patch_h) + .ok() + .and_then(|height| usize::try_from(patch_w).ok().map(|width| height * width)) + .ok_or(HostPreprocessorError::ShapeOverflow)?; + let projected_rows_per_crop = usize::try_from(patch_h / pool_h) + .ok() + .and_then(|height| { + usize::try_from(patch_w / pool_w) + .ok() + .and_then(|width| height.checked_mul(width)) + }) + .ok_or(HostPreprocessorError::ShapeOverflow)?; + let max_sequence_len = full_config + .get("max_position_embeddings") + .and_then(Value::as_u64) + .unwrap_or(4096) as usize; + + MolmoHostPreprocessor::from_parts( + processor, + vision_tower, + text_embeddings, + tokenizer, + text_config.hidden_size, + max_sequence_len, + max_crops, + patches_per_crop, + projected_rows_per_crop, + ) +} + +#[cfg(feature = "xla-iree")] +pub(crate) fn load_molmo_iree_host_preprocessor( + model_path: &Path, + device: &str, +) -> Result { + let (_config_str, full_config) = read_sanitized_vlm_config(model_path) + .map_err(|error| HostPreprocessorError::InvalidConfig(error.to_string()))?; + if full_config.get("model_type").and_then(Value::as_str) != Some("molmo") { + return Err(HostPreprocessorError::FamilyMismatch { + actual: full_config + .get("model_type") + .and_then(Value::as_str) + .unwrap_or("unknown") + .to_string(), + }); + } + 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) + .map_err(|error| HostPreprocessorError::InvalidConfig(error.to_string()))?; + let text_config: models::molmo::MolmoTextConfig = serde_json::from_value(text_config_value) + .map_err(|error| { + HostPreprocessorError::InvalidConfig(format!( + "failed to parse Molmo v1 text config: {error}" + )) + })?; + let raw_weights = load_vlm_weights_common_filtered_canonical(model_path, |name| { + name.starts_with("language_model.model.wte.") || name.starts_with("model.transformer.wte.") + }) + .map_err(|error| HostPreprocessorError::WeightLoad(error.to_string()))?; + let mut weights = WeightMap::new(); + for (name, value) in raw_weights { + weights.insert(rewrite_molmo_weight_key(&name), value); + } + let text_embeddings = + models::molmo::Molmo2Embedding::from_weights(&weights, "language_model.model.wte") + .map_err(HostPreprocessorError::WeightLoad)?; + let tokenizer = crate::tokenizer::load_tokenizer(model_path) + .map_err(|error| HostPreprocessorError::InvalidConfig(error.to_string()))?; + let preprocessor_config = read_optional_model_json(model_path, "preprocessor_config.json"); + let max_crops = preprocessor_config + .as_ref() + .and_then(|config| config.get("max_crops")) + .and_then(Value::as_u64) + .unwrap_or(12) as usize; + let overlap = preprocessor_config + .as_ref() + .and_then(|config| config.get("overlap_margins")) + .and_then(Value::as_array) + .and_then(|values| { + Some(( + values.first()?.as_u64()? as usize, + values.get(1)?.as_u64()? as usize, + )) + }); + let base_size = preprocessor_config + .as_ref() + .and_then(|config| config.get("base_image_input_size")) + .and_then(Value::as_array) + .and_then(|values| { + Some(( + values.first()?.as_u64()? as usize, + values.get(1)?.as_u64()? as usize, + )) + }); + let token_len = preprocessor_config.as_ref().and_then(|config| { + Some(( + config.get("image_token_length_h")?.as_u64()? as usize, + config.get("image_token_length_w")?.as_u64()? as usize, + )) + }); + let vision_config = full_config.get("vision_config").unwrap_or(&full_config); + let image_patch_size = molmo_vision_i32(vision_config, "image_patch_size", 14) as usize; + let processor = MolmoProcessor::new( + max_crops, + overlap, + Some(image_patch_size), + base_size, + token_len, + read_clip_triple(preprocessor_config.as_ref(), "image_mean"), + read_clip_triple(preprocessor_config.as_ref(), "image_std"), + MolmoImageTokens::default(), + ); + let max_sequence_len = full_config + .get("max_position_embeddings") + .and_then(Value::as_u64) + .unwrap_or(4096) as usize; + let projector = mlxcel_xla::IreeMolmoVisionProjector::load(model_path, device) + .map_err(HostPreprocessorError::Iree)?; + if projector.text_hidden() != text_config.hidden_size { + return Err(HostPreprocessorError::InvalidConfig(format!( + "Molmo IREE text hidden {} disagrees with text config {}", + projector.text_hidden(), + text_config.hidden_size + ))); + } + MolmoHostPreprocessor::from_iree_parts( + processor, + text_embeddings, + tokenizer, + max_sequence_len, + projector, + ) +} diff --git a/src/loading/vlm_special.rs b/src/loading/vlm_special.rs index ea8353ddb..fd2621643 100644 --- a/src/loading/vlm_special.rs +++ b/src/loading/vlm_special.rs @@ -1779,7 +1779,7 @@ fn remap_molmo_weights(raw_weights: WeightMap) -> WeightMap { /// Read a vision config field with a reference default. Molmo-7B ships an almost /// empty `vision_config`, so most fields fall back to the dataclass defaults. -fn molmo_vision_i32(vision_config: &Value, key: &str, default: i32) -> i32 { +pub(super) fn molmo_vision_i32(vision_config: &Value, key: &str, default: i32) -> i32 { vision_config .get(key) .and_then(|v| v.as_i64()) @@ -1930,7 +1930,7 @@ pub(crate) fn load_molmo_vlm(model_path: &Path) -> Result { } /// Read a 3-element CLIP normalization triple from the preprocessor config. -fn read_clip_triple(config: Option<&Value>, key: &str) -> Option<[f32; 3]> { +pub(super) fn read_clip_triple(config: Option<&Value>, key: &str) -> Option<[f32; 3]> { config .and_then(|c| c.get(key)) .and_then(|v| v.as_array()) diff --git a/src/multimodal/host_preprocessor.rs b/src/multimodal/host_preprocessor.rs index bfad9cb9d..905b400ab 100644 --- a/src/multimodal/host_preprocessor.rs +++ b/src/multimodal/host_preprocessor.rs @@ -49,6 +49,15 @@ use export::{ validate_sequence_capacity, }; +#[cfg(feature = "xla-backend")] +#[path = "host_preprocessor_molmo.rs"] +mod molmo; +#[cfg(feature = "xla-backend")] +pub use molmo::MolmoHostPreprocessor; +#[cfg(feature = "xla-reference-diagnostics")] +#[doc(hidden)] +pub use molmo::run_pinned_molmo_eager_mlx_iree_boundaries; + /// Vision implementation selected for OpenXLA multimodal preprocessing. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum XlaVisionBackend { @@ -175,12 +184,60 @@ pub fn load_xla_image_preprocessor( )); } } - if model_type != crate::models::ModelType::LlavaVLM { - return Ok(None); + let policy = XlaVisionBackendPolicy::from_env()?; + match model_type { + crate::models::ModelType::LlavaVLM => load_llava_image_preprocessor(model_path, policy), + #[cfg(feature = "xla-backend")] + crate::models::ModelType::MolmoVLM => load_molmo_image_preprocessor(model_path, policy), + _ => Ok(None), } +} - let policy = XlaVisionBackendPolicy::from_env()?; - load_llava_image_preprocessor(model_path, policy) +#[cfg(feature = "xla-backend")] +fn load_molmo_image_preprocessor( + model_path: &Path, + policy: XlaVisionBackendPolicy, +) -> Result>, HostPreprocessorError> { + #[cfg(feature = "xla-iree")] + if policy != XlaVisionBackendPolicy::Host { + let device = std::env::var("MLXCEL_XLA_DEVICE") + .unwrap_or_else(|_| mlxcel_xla::default_device().to_string()); + match MolmoHostPreprocessor::load_iree(model_path, &device) { + Ok(preprocessor) => { + tracing::info!( + vision_backend = "iree", + vision_device = %device, + vision_backend_policy = ?policy, + merge_mode = mlxcel_xla::MOLMO_V1_MERGE_MODE, + "OpenXLA Molmo v1 native vision backend selected" + ); + return Ok(Some(Box::new(preprocessor))); + } + Err(error) if policy == XlaVisionBackendPolicy::Auto => { + tracing::warn!( + vision_backend = "host", + vision_device = %device, + fallback_reason = %error, + "Molmo v1 native vision unavailable; using observable host fallback" + ); + } + Err(error) => return Err(error), + } + } + #[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(), + )); + } + let preprocessor = MolmoHostPreprocessor::load(model_path)?; + tracing::info!( + vision_backend = "host", + vision_backend_policy = ?policy, + merge_mode = mlxcel_xla::MOLMO_V1_MERGE_MODE, + "OpenXLA Molmo v1 sparse-add preprocessor selected" + ); + Ok(Some(Box::new(preprocessor))) } fn load_llava_host_preprocessor_boxed( diff --git a/src/multimodal/host_preprocessor_molmo.rs b/src/multimodal/host_preprocessor_molmo.rs new file mode 100644 index 000000000..509cbbd71 --- /dev/null +++ b/src/multimodal/host_preprocessor_molmo.rs @@ -0,0 +1,414 @@ +// 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. + +//! Molmo v1 prepared-prefill producer. +//! +//! The producer owns only the tokenizer, dual embedding table, processor, and +//! vision tower. The OLMo decoder and LM head remain exclusively in the XLA +//! session. Processor-provided `image_input_idx` values are the authoritative +//! sparse-add map; negative entries never consume a target position. + +#[cfg(feature = "xla-iree")] +use std::cell::RefCell; +use std::path::Path; + +use image::DynamicImage; +use mlxcel_core::session::{OwnedTensor, PreparedPrefill, PreparedTensorDType}; + +use crate::models::molmo::Molmo2Embedding; +use crate::tokenizer::MlxcelTokenizer; +use crate::vision::encoders::molmo::MolmoVisionModel; +use crate::vision::processors::molmo::MolmoProcessor; + +use super::export::{ + build_prepared_prefill, export_mlx_tensor, usize_to_i32, validate_embedding_shape, + validate_sequence_capacity, +}; +use super::{HostMultimodalPreprocessor, HostPreprocessorError}; + +const MOLMO_V1_BOS_TOKEN_ID: i32 = 151_643; + +/// Host components that build Molmo v1's owned XLA prepared-prefill payload. +pub struct MolmoHostPreprocessor { + processor: MolmoProcessor, + vision: MolmoVision, + text_embeddings: Molmo2Embedding, + tokenizer: MlxcelTokenizer, + hidden_size: usize, + max_sequence_len: usize, + max_crops: usize, + patches_per_crop: usize, + projected_rows_per_crop: usize, +} + +enum MolmoVision { + Host(MolmoVisionModel), + #[cfg(feature = "xla-iree")] + Iree(RefCell), +} + +impl MolmoHostPreprocessor { + /// Load only Molmo v1's prefill-side components; no decoder is retained. + pub fn load(model_path: &Path) -> Result { + crate::loading::load_molmo_host_preprocessor(model_path) + } + + #[cfg(feature = "xla-iree")] + pub fn load_iree(model_path: &Path, device: &str) -> Result { + crate::loading::load_molmo_iree_host_preprocessor(model_path, device) + } + + #[allow(clippy::too_many_arguments)] + pub(crate) fn from_parts( + processor: MolmoProcessor, + vision_tower: MolmoVisionModel, + text_embeddings: Molmo2Embedding, + tokenizer: MlxcelTokenizer, + hidden_size: usize, + max_sequence_len: usize, + max_crops: usize, + patches_per_crop: usize, + projected_rows_per_crop: usize, + ) -> Result { + for (name, value) in [ + ("hidden size", hidden_size), + ("maximum sequence length", max_sequence_len), + ("maximum crop count", max_crops), + ("patches per crop", patches_per_crop), + ("projected rows per crop", projected_rows_per_crop), + ] { + if value == 0 { + return Err(HostPreprocessorError::InvalidConfig(format!( + "Molmo v1 {name} must be non-zero" + ))); + } + } + Ok(Self { + processor, + vision: MolmoVision::Host(vision_tower), + text_embeddings, + tokenizer, + hidden_size, + max_sequence_len, + max_crops, + patches_per_crop, + projected_rows_per_crop, + }) + } + + #[cfg(feature = "xla-iree")] + pub(crate) fn from_iree_parts( + processor: MolmoProcessor, + text_embeddings: Molmo2Embedding, + tokenizer: MlxcelTokenizer, + max_sequence_len: usize, + projector: mlxcel_xla::IreeMolmoVisionProjector, + ) -> Result { + let hidden_size = projector.text_hidden(); + let max_crops = projector.max_crops(); + let patches_per_crop = projector.patches_per_crop(); + let projected_rows_per_crop = projector.projected_rows_per_crop(); + if projector.patch_width() == 0 || max_sequence_len == 0 { + return Err(HostPreprocessorError::InvalidConfig( + "Molmo IREE static dimensions must be non-zero".to_string(), + )); + } + Ok(Self { + processor, + vision: MolmoVision::Iree(RefCell::new(projector)), + text_embeddings, + tokenizer, + hidden_size, + max_sequence_len, + max_crops, + patches_per_crop, + projected_rows_per_crop, + }) + } + + fn input_prompt(&self, token_ids: &[i32]) -> Result { + let ids = token_ids + .iter() + .map(|&token| { + u32::try_from(token).map_err(|_| { + HostPreprocessorError::InvalidConfig(format!( + "Molmo v1 prompt token id {token} is negative" + )) + }) + }) + .collect::, _>>()?; + self.tokenizer + .decode(&ids, true) + .map_err(|error| HostPreprocessorError::InvalidConfig(error.to_string())) + } + + fn logical_tokens( + &self, + token_ids: &[i32], + image_tokens: &[i32], + ) -> Result, HostPreprocessorError> { + let prompt = format_prompt(&self.input_prompt(token_ids)?); + let prompt_ids = self + .tokenizer + .encode(&prompt, false) + .map_err(|error| HostPreprocessorError::InvalidConfig(error.to_string()))?; + let logical_len = 1usize + .checked_add(image_tokens.len()) + .and_then(|length| length.checked_add(prompt_ids.len())) + .ok_or(HostPreprocessorError::ShapeOverflow)?; + validate_sequence_capacity(logical_len, self.max_sequence_len)?; + let prompt_ids = prompt_ids + .into_iter() + .map(|token| { + i32::try_from(token).map_err(|_| { + HostPreprocessorError::InvalidConfig(format!( + "Molmo v1 tokenizer id {token} does not fit i32" + )) + }) + }) + .collect::, _>>()?; + let mut logical_tokens = Vec::with_capacity(logical_len); + logical_tokens.push(MOLMO_V1_BOS_TOKEN_ID); + logical_tokens.extend_from_slice(image_tokens); + logical_tokens.extend(prompt_ids); + Ok(logical_tokens) + } + + fn prepare_image( + &self, + token_ids: &[i32], + image: &DynamicImage, + ) -> Result { + let output = self.processor.preprocess_image(image); + let logical_tokens = self.logical_tokens(token_ids, &output.image_token_ids)?; + let crop_count = usize::try_from(output.pixel_values_shape[0]) + .map_err(|_| HostPreprocessorError::ShapeOverflow)?; + let patches_per_crop = usize::try_from(output.pixel_values_shape[1]) + .map_err(|_| HostPreprocessorError::ShapeOverflow)?; + let patch_width = usize::try_from(output.pixel_values_shape[2]) + .map_err(|_| HostPreprocessorError::ShapeOverflow)?; + let shifted_indices = output + .image_input_idx + .iter() + .map(|&index| { + if index < 0 { + Ok(index) + } else { + index + .checked_add(1) + .ok_or(HostPreprocessorError::ShapeOverflow) + } + }) + .collect::, _>>()?; + let prepared_image = mlxcel_xla::MolmoPreparedImage { + pixel_values: output.pixel_values, + crop_count, + patches_per_crop, + patch_width, + image_masks: output.image_masks, + projected_rows_per_crop: self.projected_rows_per_crop, + image_input_idx: shifted_indices, + }; + let max_feature_rows = self + .max_crops + .checked_mul(self.projected_rows_per_crop) + .ok_or(HostPreprocessorError::ShapeOverflow)?; + let plan = prepared_image + .sparse_add_plan( + logical_tokens.len(), + self.hidden_size, + self.max_crops, + self.patches_per_crop, + max_feature_rows, + self.max_sequence_len, + ) + .map_err(|error| HostPreprocessorError::InvalidConfig(error.to_string()))?; + + let input_ids = mlxcel_core::from_slice_i32( + &logical_tokens, + &[1, usize_to_i32(logical_tokens.len(), "sequence length")?], + ); + let text = self.text_embeddings.forward(&input_ids); + let embedding_dtype = match mlxcel_core::array_dtype(&text) { + mlxcel_core::dtype::FLOAT16 => mlxcel_xla::MolmoEmbeddingDType::Float16, + mlxcel_core::dtype::BFLOAT16 => mlxcel_xla::MolmoEmbeddingDType::BFloat16, + mlxcel_core::dtype::FLOAT32 => mlxcel_xla::MolmoEmbeddingDType::Float32, + other => return Err(HostPreprocessorError::UnsupportedDType(other)), + }; + validate_embedding_shape( + &mlxcel_core::array_shape(&text), + logical_tokens.len(), + self.hidden_size, + "Molmo text embedding table", + )?; + + let text = mlxcel_core::astype(&text, mlxcel_core::dtype::FLOAT32); + let mut text_values = tensor_f32(export_mlx_tensor(&text, "Molmo text embeddings")?)?; + let projected_values = match &self.vision { + MolmoVision::Host(vision_tower) => { + let pixels = mlxcel_core::from_slice_f32( + &prepared_image.pixel_values, + &[ + 1, + usize_to_i32(crop_count, "crop count")?, + usize_to_i32(patches_per_crop, "patches per crop")?, + usize_to_i32(patch_width, "patch width")?, + ], + ); + let masks = mlxcel_core::from_slice_f32( + &prepared_image.image_masks, + &[ + 1, + usize_to_i32(crop_count, "crop count")?, + usize_to_i32(patches_per_crop, "patches per crop")?, + ], + ); + let projected = mlxcel_core::astype( + &vision_tower.forward(&pixels, &masks), + mlxcel_core::dtype::FLOAT32, + ); + let expected = [ + 1, + usize_to_i32(crop_count, "crop count")?, + usize_to_i32(self.projected_rows_per_crop, "projected rows per crop")?, + usize_to_i32(self.hidden_size, "hidden size")?, + ]; + if mlxcel_core::array_shape(&projected) != expected { + return Err(HostPreprocessorError::InvalidConfig(format!( + "Molmo projected image shape {:?} does not match {expected:?}", + mlxcel_core::array_shape(&projected) + ))); + } + tensor_f32(export_mlx_tensor(&projected, "Molmo projected features")?)? + } + #[cfg(feature = "xla-iree")] + MolmoVision::Iree(projector) => { + let projection = projector + .try_borrow_mut() + .map_err(|_| { + HostPreprocessorError::Iree( + "concurrent/re-entrant Molmo IREE vision invocation is unsupported" + .to_string(), + ) + })? + .project( + &prepared_image.pixel_values, + &prepared_image.image_masks, + crop_count, + ) + .map_err(HostPreprocessorError::Iree)?; + tracing::info!( + vision_backend = "iree", + crop_count, + upload_bytes = projection.upload_bytes, + transfer_bytes = projection.transfer_bytes, + iree_vision_seconds = projection.elapsed_seconds, + "Molmo v1 native vision projection completed" + ); + projection.values + } + }; + plan.apply_in_dtype(&mut text_values, &projected_values, embedding_dtype) + .map_err(|error| HostPreprocessorError::InvalidConfig(error.to_string()))?; + let mut bytes = Vec::with_capacity( + text_values + .len() + .checked_mul(std::mem::size_of::()) + .ok_or(HostPreprocessorError::ShapeOverflow)?, + ); + for value in text_values { + bytes.extend_from_slice(&value.to_ne_bytes()); + } + let embeddings = OwnedTensor::new( + bytes, + PreparedTensorDType::Float32, + vec![1, logical_tokens.len(), self.hidden_size], + )?; + build_prepared_prefill( + logical_tokens, + embeddings, + 1, + output.image_token_ids.len(), + "molmo-v1", + ) + } +} + +impl HostMultimodalPreprocessor for MolmoHostPreprocessor { + fn backend(&self) -> super::XlaVisionBackend { + match &self.vision { + MolmoVision::Host(_) => super::XlaVisionBackend::Host, + #[cfg(feature = "xla-iree")] + MolmoVision::Iree(_) => super::XlaVisionBackend::Iree, + } + } + + fn prepare( + &self, + token_ids: &[i32], + images: &[DynamicImage], + ) -> Result { + match images { + [image] => self.prepare_image(token_ids, image), + [] => Err(HostPreprocessorError::InvalidConfig( + "Molmo v1 image preprocessor requires exactly one image".to_string(), + )), + _ => Err(HostPreprocessorError::InvalidConfig( + "Molmo v1 currently supports exactly one image per request".to_string(), + )), + } + } +} + +fn tensor_f32(tensor: OwnedTensor) -> Result, HostPreprocessorError> { + if tensor.dtype != PreparedTensorDType::Float32 { + return Err(HostPreprocessorError::InvalidConfig(format!( + "Molmo sparse add requires float32 tensors, got {:?}", + tensor.dtype + ))); + } + Ok(tensor + .bytes + .chunks_exact(std::mem::size_of::()) + .map(|bytes| { + f32::from_ne_bytes( + bytes + .try_into() + .expect("chunks_exact yields one native f32"), + ) + }) + .collect()) +} + +fn format_prompt(prompt: &str) -> String { + let without_image_placeholder = prompt.replace("<|image|>", ""); + let prompt = without_image_placeholder.trim(); + if prompt.starts_with("User:") && prompt.ends_with("Assistant:") { + format!(" {prompt}") + } else { + format!(" User: {prompt} Assistant:") + } +} + +#[cfg(any( + all(test, feature = "xla-diagnostics"), + feature = "xla-reference-diagnostics" +))] +#[path = "host_preprocessor_molmo_diagnostics_tests.rs"] +mod diagnostics_tests; +#[cfg(feature = "xla-reference-diagnostics")] +pub use diagnostics_tests::run_pinned_molmo_eager_mlx_iree_boundaries; +#[cfg(test)] +#[path = "host_preprocessor_molmo_tests.rs"] +mod tests; diff --git a/src/multimodal/host_preprocessor_molmo_diagnostics_support.rs b/src/multimodal/host_preprocessor_molmo_diagnostics_support.rs new file mode 100644 index 000000000..ec830339c --- /dev/null +++ b/src/multimodal/host_preprocessor_molmo_diagnostics_support.rs @@ -0,0 +1,246 @@ +use std::fs; +use std::io::{self, Write}; +use std::path::Path; + +#[cfg(test)] +use image::DynamicImage; +use mlxcel_core::session::{PreparedPrefill, PreparedTensorDType}; +use sha2::{Digest, Sha256}; + +#[cfg(test)] +use crate::vision::processors::molmo::{MolmoImageTokens, MolmoProcessor, MolmoProcessorOutput}; + +pub(super) const PINNED_REVISION: &str = "5c04b3a418979597b1968e41414ad799c87533e8"; +pub(super) const PINNED_CONFIG_SHA256: &str = + "4b27cd3177a990224d4f31aa941bcf53b3bb99b90b097d0dd7084a99c423224c"; +pub(super) const PINNED_PREPROCESSOR_SHA256: &str = + "01a54ee26803f8a4088030ce2b6be93cf87d402fe8e8e4dbf56a2adbbb7ab3d7"; +pub(super) const PINNED_PROCESSOR_SHA256: &str = + "9c40772b47f74c061e4829611bb581e5498ac0f90ac0851874a907462a6a026e"; +pub(super) const PINNED_IMAGE_SHA256: &str = + "5e7d54e8a7d21802378c87d2d70cf551e29739fe27599ddf129ebccdad1e6261"; +pub(super) const PINNED_PIXEL_VALUES_SHA256: &str = + "0a1bef503b209de9b4454f636f03c55f4d84309cb7fd6bbf573750dd499799aa"; +pub(super) const PINNED_IMAGE_MASKS_SHA256: &str = + "5a0de017af320127aa8685a4671487960f2aa498b541d5445a0962410bf3b92a"; +#[cfg(test)] +pub(super) const PINNED_IMAGE_TOKEN_IDS_SHA256: &str = + "23a0878a6a2818f113460cc9a116f3d68f4d3d4b1a34b1f34dfe7805395562fa"; +pub(super) const PINNED_IMAGE_INPUT_IDX_SHA256: &str = + "243da1bc98ebb61f5d98696623b41ade1d643994db6fcfb764392452d3f9c1b4"; +#[cfg(test)] +pub(super) const PINNED_GRID_SHA256: &str = + "29c898799e5d25bb6fa02a097a3878d280419aeab84e720cab55d889c6821bd5"; + +#[derive(Debug, Clone, Copy)] +pub(super) struct Tolerance { + atol: f64, + rtol: f64, +} + +pub(super) const EXACT: Tolerance = Tolerance { + atol: 0.0, + rtol: 0.0, +}; +// The pinned host path retains F16 checkpoint rounding while the local-task +// StableHLO graph accumulates in F32. This is the existing #870 prepared-visual +// envelope, now applied at the first observable internal boundary. +pub(super) const VISION: Tolerance = Tolerance { + atol: 0.25, + rtol: 0.05, +}; + +pub(super) fn progress(stage: &str) { + eprintln!("[molmo-v1-boundary] {stage}"); + io::stderr().flush().expect("flush diagnostic progress"); +} + +pub(super) fn compare_stage( + stage: &str, + observed: &[f32], + reference: &[f32], + tolerance: Tolerance, + first_divergence: &mut Option, +) { + if observed.len() != reference.len() { + let detail = format!("{stage}: length {} != {}", observed.len(), reference.len()); + first_divergence.get_or_insert(detail.clone()); + eprintln!("[molmo-v1-boundary] stage={stage} status=FAIL {detail}"); + return; + } + let mut max_abs = 0.0f64; + let mut max_rel = 0.0f64; + let mut failures = 0usize; + let mut first_failure = None; + for (index, (&observed, &reference)) in observed.iter().zip(reference).enumerate() { + if !observed.is_finite() || !reference.is_finite() { + failures += 1; + first_failure.get_or_insert(index); + continue; + } + let absolute = f64::from((observed - reference).abs()); + let relative = absolute / f64::from(reference.abs()).max(f64::MIN_POSITIVE); + max_abs = max_abs.max(absolute); + max_rel = max_rel.max(relative); + if absolute > tolerance.atol + tolerance.rtol * f64::from(reference.abs()) { + failures += 1; + first_failure.get_or_insert(index); + } + } + let status = if failures == 0 { "PASS" } else { "FAIL" }; + eprintln!( + "[molmo-v1-boundary] stage={stage} status={status} elements={} \ + atol={:.3e} rtol={:.3e} max_abs={max_abs:.6e} max_rel={max_rel:.6e} \ + failures={failures} first_failure={first_failure:?}", + observed.len(), + tolerance.atol, + tolerance.rtol, + ); + io::stderr().flush().expect("flush diagnostic stage"); + if failures != 0 { + first_divergence.get_or_insert_with(|| { + format!( + "{stage} at flat index {}", + first_failure.expect("failed comparison has an index") + ) + }); + } +} + +pub(super) fn compare_exact_i32( + stage: &str, + observed: &[i32], + reference: &[i32], + first_divergence: &mut Option, +) { + let first = observed + .iter() + .zip(reference) + .position(|(observed, reference)| observed != reference); + let failed = observed.len() != reference.len() || first.is_some(); + eprintln!( + "[molmo-v1-boundary] stage={stage} status={} elements={} first_failure={first:?}", + if failed { "FAIL" } else { "PASS" }, + observed.len() + ); + if failed { + first_divergence.get_or_insert_with(|| { + format!( + "{stage} at {}", + first.map_or_else( + || format!("length {} != {}", observed.len(), reference.len()), + |index| format!("flat index {index}") + ) + ) + }); + } +} + +pub(super) fn sha256(path: &Path) -> String { + format!( + "{:x}", + Sha256::digest( + fs::read(path).unwrap_or_else(|error| panic!("read {}: {error}", path.display())) + ) + ) +} + +pub(super) fn assert_sha256(path: &Path, expected: &str, label: &str) { + assert_eq!( + sha256(path), + expected, + "{label} differs from the pinned #870 fixture: {}", + path.display() + ); +} + +pub(super) fn digest_f32(values: &[f32]) -> String { + let mut digest = Sha256::new(); + for value in values { + digest.update(value.to_bits().to_le_bytes()); + } + format!("{:x}", digest.finalize()) +} + +pub(super) fn digest_i32(values: &[i32]) -> String { + let mut digest = Sha256::new(); + for value in values { + digest.update(value.to_le_bytes()); + } + format!("{:x}", digest.finalize()) +} + +#[cfg(test)] +pub(super) fn digest_grid(output: &MolmoProcessorOutput) -> String { + let mut values = Vec::new(); + values.extend(output.pixel_values_shape); + values.extend(output.image_masks_shape); + values.push(output.image_input_idx_len); + values.push(i32::try_from(output.image_token_ids.len()).expect("token count fits i32")); + digest_i32(&values) +} + +pub(super) fn pinned_revision(model: &Path) -> String { + if let Ok(revision) = std::env::var("MLXCEL_MOLMO_REVISION") { + return revision; + } + let metadata = model.join(".cache/huggingface/download/config.json.metadata"); + fs::read_to_string(&metadata) + .unwrap_or_else(|error| { + panic!( + "read pinned revision from {} ({error}); set MLXCEL_MOLMO_REVISION", + metadata.display() + ) + }) + .lines() + .next() + .expect("Hugging Face metadata contains a revision") + .to_string() +} + +#[cfg(test)] +fn fixture_image() -> DynamicImage { + image::open("tests/fixtures/test_image.png").expect("load pinned Molmo image fixture") +} + +#[cfg(test)] +pub(super) fn pinned_processor_output() -> MolmoProcessorOutput { + MolmoProcessor::new( + 12, + Some((4, 4)), + Some(14), + Some((336, 336)), + Some((12, 12)), + None, + None, + MolmoImageTokens::default(), + ) + .preprocess_image(&fixture_image()) +} + +pub(super) fn mlx_f32(array: &mlxcel_core::MlxArray) -> Vec { + let widened = mlxcel_core::astype(array, mlxcel_core::dtype::FLOAT32); + mlxcel_core::try_array_to_raw_bytes(&widened) + .expect("export MLX diagnostic tensor") + .chunks_exact(4) + .map(|bytes| f32::from_ne_bytes(bytes.try_into().expect("four-byte f32 chunk"))) + .collect() +} + +pub(super) fn prepared_f32(prepared: &PreparedPrefill) -> Vec { + assert_eq!(prepared.embeddings.dtype, PreparedTensorDType::Float32); + prepared + .embeddings + .bytes + .chunks_exact(4) + .map(|bytes| f32::from_ne_bytes(bytes.try_into().expect("four-byte f32 chunk"))) + .collect() +} + +pub(super) fn argmax(values: &[f32]) -> usize { + values + .iter() + .enumerate() + .max_by(|left, right| left.1.total_cmp(right.1)) + .map_or(0, |(index, _)| index) +} diff --git a/src/multimodal/host_preprocessor_molmo_diagnostics_tests.rs b/src/multimodal/host_preprocessor_molmo_diagnostics_tests.rs new file mode 100644 index 000000000..f0ad28009 --- /dev/null +++ b/src/multimodal/host_preprocessor_molmo_diagnostics_tests.rs @@ -0,0 +1,471 @@ +// 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::BTreeSet; +use std::path::PathBuf; +use std::sync::mpsc::{self, RecvTimeoutError, Sender}; +use std::thread::{self, JoinHandle}; +use std::time::{Duration, Instant}; + +use super::*; +#[cfg(test)] +use crate::vision::processors::molmo::MolmoImageTokens; + +#[path = "host_preprocessor_molmo_diagnostics_support.rs"] +mod support; +use support::*; + +#[test] +fn pinned_processor_fixture_preserves_patch_grid_and_sparse_mapping() { + let output = pinned_processor_output(); + assert_eq!(digest_f32(&output.pixel_values), PINNED_PIXEL_VALUES_SHA256); + assert_eq!(digest_f32(&output.image_masks), PINNED_IMAGE_MASKS_SHA256); + assert_eq!( + digest_i32(&output.image_token_ids), + PINNED_IMAGE_TOKEN_IDS_SHA256 + ); + assert_eq!( + digest_i32(&output.image_input_idx), + PINNED_IMAGE_INPUT_IDX_SHA256 + ); + assert_eq!(digest_grid(&output), PINNED_GRID_SHA256); +} + +#[test] +fn strict_oracle_rejects_molmo2_token_scan_and_replacement() { + let canonical = + mlxcel_xla::MolmoSparseAddPlan::from_image_input_idx(&[-1, 4, 1], 6, 2, 3, 6).unwrap(); + let patch_id = MolmoImageTokens::default().image_patch_id; + let tokens = [7, patch_id, patch_id, 8, patch_id, 9]; + let token_scanned = tokens + .iter() + .enumerate() + .filter(|(_, token)| **token == patch_id) + .enumerate() + .map( + |(feature_row, (target_position, _))| mlxcel_xla::MolmoSparseAddPair { + feature_row, + target_position, + }, + ) + .collect::>(); + assert_ne!( + canonical.pairs(), + token_scanned, + "Molmo2 token scanning must not satisfy the Molmo v1 mapping oracle" + ); + + let original = vec![2.0, 3.0, 5.0, 7.0]; + let features = vec![11.0, 13.0]; + let plan = mlxcel_xla::MolmoSparseAddPlan::from_image_input_idx(&[1], 2, 2, 1, 2).unwrap(); + let mut additive = original.clone(); + plan.apply(&mut additive, &features).unwrap(); + let mut replacement = original; + replacement[2..].copy_from_slice(&features); + assert_ne!( + additive, replacement, + "replacement must not satisfy the Molmo v1 additive oracle" + ); +} + +/// Strict actual-checkpoint boundary gate for #870. +/// +/// Run from the repository root only with the pinned 5.3 GB checkpoint: +/// +/// ```text +/// IREE_DIST=/home/inureyes/Development/mlxcel/spike/iree-ffi/iree-dist \ +/// cargo build --release --example xla_molmo_reference_check \ +/// --features cuda,xla-reference-diagnostics +/// +/// MLXCEL_DEVICE=gpu \ +/// MLXCEL_MOLMO_FIXTURE=/home/inureyes/models/molmo-7b \ +/// MLXCEL_MOLMO_REVISION=5c04b3a418979597b1968e41414ad799c87533e8 \ +/// MLXCEL_MOLMO_IMAGE=tests/fixtures/test_image.png \ +/// MLXCEL_MOLMO_IREE_DEVICE=local-task \ +/// MLXCEL_MOLMO_HEARTBEAT_SECS=30 \ +/// ./target/release/examples/xla_molmo_reference_check +/// ``` +/// +/// The gate uses one resident IREE decoder for both prepared payloads. It +/// captures K/V and logits propagation but does not claim an independent +/// decoder oracle. +pub fn run_pinned_molmo_eager_mlx_iree_boundaries() { + assert!( + cfg!(feature = "cuda"), + "the Molmo actual-checkpoint reference binary must be built with --features cuda" + ); + mlxcel_core::set_default_device(true); + assert!( + mlxcel_core::is_gpu_available(), + "the Molmo eager reference must run on the MLX GPU device" + ); + let model = PathBuf::from( + std::env::var_os("MLXCEL_MOLMO_FIXTURE") + .expect("MLXCEL_MOLMO_FIXTURE must name the pinned checkpoint"), + ); + let image_path = std::env::var_os("MLXCEL_MOLMO_IMAGE") + .map(PathBuf::from) + .unwrap_or_else(|| PathBuf::from("tests/fixtures/test_image.png")); + let device = + std::env::var("MLXCEL_MOLMO_IREE_DEVICE").unwrap_or_else(|_| "local-task".to_string()); + assert!( + matches!(device.as_str(), "local-task" | "local-sync"), + "the #870 reference gate requires IREE local-task or local-sync" + ); + progress(&format!( + "runtime split ready: eager MLX=CUDA, IREE={device}" + )); + progress("validate pinned checkpoint and image"); + assert_eq!(pinned_revision(&model), PINNED_REVISION); + assert_sha256( + &model.join("config.json"), + PINNED_CONFIG_SHA256, + "checkpoint config", + ); + assert_sha256( + &model.join("preprocessor_config.json"), + PINNED_PREPROCESSOR_SHA256, + "checkpoint preprocessor", + ); + assert_sha256( + &model.join("processor_config.json"), + PINNED_PROCESSOR_SHA256, + "checkpoint processor", + ); + assert_sha256(&image_path, PINNED_IMAGE_SHA256, "fixture image"); + let image = image::open(&image_path).expect("open pinned Molmo image"); + let mut first_divergence = None; + + progress("load filtered eager MLX and production IREE preprocessors"); + let host = MolmoHostPreprocessor::load(&model).expect("load Molmo host preprocessor"); + // Bound the local-task worker topology and stack before the first IREE + // instance exists. IREE's flag registry is process-global and is parsed on + // that first creation, so this has to run ahead of `load_iree` 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() + .expect("configure diagnostics-only IREE local-task threads"); + let production = MolmoHostPreprocessor::load_iree(&model, &device) + .expect("load Molmo production IREE preprocessor"); + let host_processor = host.processor.preprocess_image(&image); + let iree_processor = production.processor.preprocess_image(&image); + assert_eq!( + digest_f32(&host_processor.pixel_values), + PINNED_PIXEL_VALUES_SHA256 + ); + assert_eq!( + digest_f32(&host_processor.image_masks), + PINNED_IMAGE_MASKS_SHA256 + ); + assert_eq!( + digest_i32(&host_processor.image_input_idx), + PINNED_IMAGE_INPUT_IDX_SHA256 + ); + compare_stage( + "processor.patches", + &iree_processor.pixel_values, + &host_processor.pixel_values, + EXACT, + &mut first_divergence, + ); + compare_stage( + "processor.patch_masks", + &iree_processor.image_masks, + &host_processor.image_masks, + EXACT, + &mut first_divergence, + ); + compare_exact_i32( + "processor.image_token_ids", + &iree_processor.image_token_ids, + &host_processor.image_token_ids, + &mut first_divergence, + ); + compare_exact_i32( + "processor.image_input_idx", + &iree_processor.image_input_idx, + &host_processor.image_input_idx, + &mut first_divergence, + ); + compare_exact_i32( + "processor.grid", + &[ + iree_processor.pixel_values_shape[0], + iree_processor.pixel_values_shape[1], + iree_processor.image_input_idx_len, + ], + &[ + host_processor.pixel_values_shape[0], + host_processor.pixel_values_shape[1], + host_processor.image_input_idx_len, + ], + &mut first_divergence, + ); + + progress("capture eager MLX patch, selected-layer, and projector states"); + let [crops, patches, width] = host_processor.pixel_values_shape; + let pixels = + mlxcel_core::from_slice_f32(&host_processor.pixel_values, &[1, crops, patches, width]); + let masks = mlxcel_core::from_slice_f32(&host_processor.image_masks, &[1, crops, patches]); + let MolmoVision::Host(vision) = &host.vision else { + panic!("host preprocessor must retain the eager MLX vision tower"); + }; + let mut heartbeat = DiagnosticHeartbeat::from_env(); + let (_, host_diagnostics) = + vision.forward_with_diagnostics(&pixels, &masks, &mut |stage| match stage { + crate::vision::encoders::molmo::MolmoHostVisionDiagnosticStage::PatchEmbeddingStarted => { + progress("eager MLX patch embedding materialization started"); + heartbeat.start("eager MLX patch embedding"); + } + crate::vision::encoders::molmo::MolmoHostVisionDiagnosticStage::PatchEmbeddingCompleted => { + heartbeat.stop(); + progress("eager MLX patch embedding materialized"); + } + crate::vision::encoders::molmo::MolmoHostVisionDiagnosticStage::SelectedLayerStarted(layer) => { + progress(&format!("eager MLX selected ViT layer {layer} materialization started")); + heartbeat.start(format!("eager MLX selected ViT layer {layer}")); + } + crate::vision::encoders::molmo::MolmoHostVisionDiagnosticStage::SelectedLayerCompleted(layer) => { + heartbeat.stop(); + progress(&format!("eager MLX selected ViT layer {layer} materialized")); + } + crate::vision::encoders::molmo::MolmoHostVisionDiagnosticStage::ProjectorStarted => { + progress("eager MLX pool/projector materialization started"); + heartbeat.start("eager MLX pool/projector"); + } + crate::vision::encoders::molmo::MolmoHostVisionDiagnosticStage::ProjectorCompleted => { + heartbeat.stop(); + progress("eager MLX pool/projector materialized"); + } + }); + heartbeat.stop(); + let host_patch = mlx_f32(host_diagnostics.patch_embeddings.as_ref().unwrap()); + let host_selected = host_diagnostics + .selected_hidden_states + .iter() + .map(|state| mlx_f32(state.as_ref().unwrap())) + .collect::>(); + let host_projected = mlx_f32(host_diagnostics.projected_features.as_ref().unwrap()); + + progress("capture resident IREE patch, selected-layer, and projector states"); + let mut diagnostic = mlxcel_xla::IreeMolmoVisionDiagnosticProjector::load(&model, &device) + .expect("load Molmo IREE diagnostic projector"); + let iree_diagnostics = diagnostic + .project( + &host_processor.pixel_values, + &host_processor.image_masks, + usize::try_from(crops).expect("positive crop count"), + ) + .expect("run Molmo IREE diagnostic projector"); + compare_stage( + "vision.patch_embedding", + &iree_diagnostics.patch_embeddings, + &host_patch, + VISION, + &mut first_divergence, + ); + assert_eq!( + iree_diagnostics.selected_hidden_states.len(), + host_selected.len() + ); + for ((layer, iree), host) in iree_diagnostics + .selected_layers + .iter() + .zip(&iree_diagnostics.selected_hidden_states) + .zip(&host_selected) + { + compare_stage( + &format!("vision.selected_layer_{layer}"), + iree, + host, + VISION, + &mut first_divergence, + ); + } + compare_stage( + "vision.projector", + &iree_diagnostics.projected_features, + &host_projected, + VISION, + &mut first_divergence, + ); + + progress("compare production prepared sparse-add merge"); + let tokenizer = crate::tokenizer::load_tokenizer(&model).expect("load Molmo tokenizer"); + let token_ids = tokenizer + .encode("Describe the image.", true) + .expect("tokenize Molmo prompt") + .into_iter() + .map(|token| i32::try_from(token).expect("token id fits i32")) + .collect::>(); + let host_prepared = host + .prepare(&token_ids, std::slice::from_ref(&image)) + .expect("prepare eager MLX Molmo prefill"); + let iree_prepared = production + .prepare(&token_ids, &[image]) + .expect("prepare production IREE Molmo prefill"); + assert_eq!(iree_prepared.token_ids, host_prepared.token_ids); + assert_eq!(iree_prepared.positions, host_prepared.positions); + assert_eq!(iree_prepared.attention_bias, host_prepared.attention_bias); + assert_eq!(iree_prepared.modalities, host_prepared.modalities); + let visual_rows = host_processor + .image_input_idx + .iter() + .filter_map(|&position| usize::try_from(position).ok()) + .map(|position| position + 1) + .collect::>(); + let hidden = host_prepared.embeddings.shape[2]; + let host_values = prepared_f32(&host_prepared); + let iree_values = prepared_f32(&iree_prepared); + let mut host_non_visual = Vec::new(); + let mut iree_non_visual = Vec::new(); + let mut host_visual = Vec::new(); + let mut iree_visual = Vec::new(); + for row in 0..host_prepared.sequence_len { + let range = row * hidden..(row + 1) * hidden; + if visual_rows.contains(&row) { + host_visual.extend_from_slice(&host_values[range.clone()]); + iree_visual.extend_from_slice(&iree_values[range]); + } else { + host_non_visual.extend_from_slice(&host_values[range.clone()]); + iree_non_visual.extend_from_slice(&iree_values[range]); + } + } + compare_stage( + "prepared.non_visual_rows", + &iree_non_visual, + &host_non_visual, + EXACT, + &mut first_divergence, + ); + compare_stage( + "prepared.sparse_add_visual_rows", + &iree_visual, + &host_visual, + VISION, + &mut first_divergence, + ); + + progress("compare one resident decoder's layer-0 KV, all KV, and logits"); + let mut engine = mlxcel_xla::LlavaReferenceDiagnosticEngine::load( + &model, + &device, + host_prepared.sequence_len, + ) + .expect("load one Molmo IREE decoder diagnostic engine"); + let host_decoder = engine + .capture(&host_prepared, 64, 1) + .expect("capture eager-prepared decoder boundary"); + let iree_decoder = engine + .capture(&iree_prepared, 64, 1) + .expect("capture IREE-prepared decoder boundary"); + assert_eq!(host_decoder.prefill.kv_width, 64); + assert_eq!(iree_decoder.prefill.kv_width, 64); + assert!(host_decoder.prefill.kv.len() >= 128); + assert!(iree_decoder.prefill.kv.len() >= 128); + compare_stage( + "decoder.layer0_kv", + &iree_decoder.prefill.kv[..128], + &host_decoder.prefill.kv[..128], + VISION, + &mut first_divergence, + ); + compare_stage( + "decoder.all_layer_kv", + &iree_decoder.prefill.kv, + &host_decoder.prefill.kv, + VISION, + &mut first_divergence, + ); + compare_stage( + "decoder.prefill_logits", + &iree_decoder.prefill.logits, + &host_decoder.prefill.logits, + VISION, + &mut first_divergence, + ); + assert_eq!( + argmax(&iree_decoder.prefill.logits), + argmax(&host_decoder.prefill.logits), + "prepared-path top-1 token diverged" + ); + if let Some(first_divergence) = first_divergence { + panic!("Molmo v1 eager MLX/IREE first divergence: {first_divergence}"); + } +} + +struct DiagnosticHeartbeat { + interval: Duration, + stop_sender: Option>, + worker: Option>, +} + +impl DiagnosticHeartbeat { + fn from_env() -> Self { + let seconds = std::env::var("MLXCEL_MOLMO_HEARTBEAT_SECS") + .ok() + .and_then(|value| value.parse::().ok()) + .filter(|seconds| *seconds > 0) + .unwrap_or(30); + Self { + interval: Duration::from_secs(seconds), + stop_sender: None, + worker: None, + } + } + + fn start(&mut self, stage: impl Into) { + self.stop(); + let stage = stage.into(); + let interval = self.interval; + let (stop_sender, stop_receiver) = mpsc::channel(); + self.stop_sender = Some(stop_sender); + self.worker = Some(thread::spawn(move || { + let started = Instant::now(); + loop { + match stop_receiver.recv_timeout(interval) { + Ok(()) | Err(RecvTimeoutError::Disconnected) => break, + Err(RecvTimeoutError::Timeout) => progress(&format!( + "heartbeat stage={stage} elapsed={}s", + started.elapsed().as_secs() + )), + } + } + })); + } + + fn stop(&mut self) { + if let Some(stop_sender) = self.stop_sender.take() { + let _ = stop_sender.send(()); + } + if let Some(worker) = self.worker.take() { + worker.join().expect("Molmo diagnostic heartbeat panicked"); + } + } +} + +impl Drop for DiagnosticHeartbeat { + fn drop(&mut self) { + self.stop(); + } +} + +#[cfg(test)] +#[test] +#[ignore = "requires pinned Molmo checkpoint, MLX CUDA, and IREE local-task"] +fn pinned_molmo_eager_mlx_matches_iree_boundaries() { + run_pinned_molmo_eager_mlx_iree_boundaries(); +} diff --git a/src/multimodal/host_preprocessor_molmo_tests.rs b/src/multimodal/host_preprocessor_molmo_tests.rs new file mode 100644 index 000000000..700a2447b --- /dev/null +++ b/src/multimodal/host_preprocessor_molmo_tests.rs @@ -0,0 +1,242 @@ +// 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::BTreeSet; + +use super::*; + +const VISUAL_MAX_ABS: f32 = 0.25; +const VISUAL_MIN_COSINE: f64 = 0.99; +const VISUAL_MAX_NORMALIZED_RMSE: f64 = 0.05; + +fn real_inputs() -> (std::path::PathBuf, Vec, DynamicImage) { + let model_path = std::env::var_os("MLXCEL_TEST_MOLMO_MODEL") + .map(std::path::PathBuf::from) + .expect("MLXCEL_TEST_MOLMO_MODEL is required"); + let tokenizer = crate::tokenizer::load_tokenizer(&model_path).unwrap(); + let token_ids = tokenizer + .encode("Describe the image.", true) + .unwrap() + .into_iter() + .map(|token| i32::try_from(token).unwrap()) + .collect::>(); + let image = image::open("tests/fixtures/test_image.png").unwrap(); + (model_path, token_ids, image) +} + +fn assert_real_prepared(prepared: &PreparedPrefill) { + assert_eq!(prepared.embeddings.dtype, PreparedTensorDType::Float32); + assert_eq!(prepared.embeddings.shape[0], 1); + assert_eq!(prepared.embeddings.shape[1], prepared.token_ids.len()); + assert_eq!(prepared.modalities[0].family, "molmo-v1"); + assert_eq!(prepared.modalities[0].item_count, 1); +} + +fn real_parity_inputs() -> (std::path::PathBuf, Vec, DynamicImage, String) { + let model_path = std::env::var_os("MLXCEL_TEST_MOLMO_MODEL") + .map(std::path::PathBuf::from) + .expect("MLXCEL_TEST_MOLMO_MODEL is required"); + let image_path = std::env::var_os("MLXCEL_TEST_MOLMO_IMAGE") + .map(std::path::PathBuf::from) + .expect("MLXCEL_TEST_MOLMO_IMAGE is required"); + let device = std::env::var("MLXCEL_TEST_MOLMO_IREE_DEVICE") + .expect("MLXCEL_TEST_MOLMO_IREE_DEVICE is required"); + let tokenizer = crate::tokenizer::load_tokenizer(&model_path).unwrap(); + let token_ids = tokenizer + .encode("Describe the image.", true) + .unwrap() + .into_iter() + .map(|token| i32::try_from(token).unwrap()) + .collect::>(); + let image = image::open(&image_path).unwrap_or_else(|error| { + panic!( + "failed to load MLXCEL_TEST_MOLMO_IMAGE {}: {error}", + image_path.display() + ) + }); + (model_path, token_ids, image, device) +} + +fn prepared_f32(prepared: &PreparedPrefill) -> Vec { + assert_eq!(prepared.embeddings.dtype, PreparedTensorDType::Float32); + prepared + .embeddings + .bytes + .chunks_exact(std::mem::size_of::()) + .map(|bytes| { + f32::from_ne_bytes( + bytes + .try_into() + .expect("chunks_exact yields one native f32"), + ) + }) + .collect() +} + +#[test] +fn prompt_format_matches_molmo_v1_processor_contract() { + assert_eq!( + format_prompt("<|image|> Describe the image."), + " User: Describe the image. Assistant:" + ); + assert_eq!( + format_prompt("User: Describe the image. Assistant:"), + " User: Describe the image. Assistant:" + ); +} + +#[test] +#[ignore = "requires MLXCEL_TEST_MOLMO_MODEL and the checked-in image fixture"] +fn real_checkpoint_builds_owned_sparse_add_prefill() { + mlxcel_core::set_default_device(false); + let (model_path, token_ids, image) = real_inputs(); + let preprocessor = MolmoHostPreprocessor::load(&model_path).unwrap(); + let prepared = preprocessor.prepare(&token_ids, &[image]).unwrap(); + assert_real_prepared(&prepared); +} + +#[cfg(feature = "xla-iree")] +#[test] +#[ignore = "requires MLXCEL_TEST_MOLMO_MODEL and a real IREE runtime"] +fn real_checkpoint_builds_native_iree_sparse_add_prefill() { + mlxcel_core::set_default_device(false); + let (model_path, token_ids, image) = real_inputs(); + let device = + std::env::var("MLXCEL_TEST_MOLMO_IREE_DEVICE").unwrap_or_else(|_| "cuda".to_string()); + let preprocessor = MolmoHostPreprocessor::load_iree(&model_path, &device).unwrap(); + assert_eq!(preprocessor.backend(), super::super::XlaVisionBackend::Iree); + let prepared = preprocessor.prepare(&token_ids, &[image]).unwrap(); + assert_real_prepared(&prepared); +} + +#[cfg(feature = "xla-iree")] +#[test] +#[ignore = "requires MLXCEL_TEST_MOLMO_MODEL, MLXCEL_TEST_MOLMO_IMAGE, and MLXCEL_TEST_MOLMO_IREE_DEVICE"] +fn real_checkpoint_host_and_iree_prepared_prefill_match() { + mlxcel_core::set_default_device(false); + let (model_path, token_ids, image, device) = real_parity_inputs(); + let host = MolmoHostPreprocessor::load(&model_path).unwrap(); + + // `image_input_idx` is relative to the image-token block. Prepared + // prefills prepend BOS, so active visual rows are shifted by one. + let processor_output = host.processor.preprocess_image(&image); + let active_count = processor_output + .image_input_idx + .iter() + .filter(|&&position| position >= 0) + .count(); + let visual_rows = processor_output + .image_input_idx + .iter() + .filter_map(|&position| { + usize::try_from(position) + .ok() + .and_then(|position| position.checked_add(1)) + }) + .collect::>(); + assert!(!visual_rows.is_empty(), "fixture must contain visual rows"); + assert_eq!( + visual_rows.len(), + active_count, + "fixture must retain unique authoritative visual targets" + ); + + let host_prepared = host + .prepare(&token_ids, std::slice::from_ref(&image)) + .unwrap(); + let iree = MolmoHostPreprocessor::load_iree(&model_path, &device).unwrap(); + let iree_prepared = iree.prepare(&token_ids, &[image]).unwrap(); + assert_real_prepared(&host_prepared); + assert_real_prepared(&iree_prepared); + + assert_eq!(iree_prepared.token_ids, host_prepared.token_ids); + assert_eq!( + iree_prepared.embeddings.shape, + host_prepared.embeddings.shape + ); + assert_eq!( + iree_prepared.embeddings.dtype, + host_prepared.embeddings.dtype + ); + assert_eq!(iree_prepared.positions, host_prepared.positions); + assert_eq!(iree_prepared.attention_bias, host_prepared.attention_bias); + assert_eq!(iree_prepared.modalities, host_prepared.modalities); + assert_eq!(iree_prepared.adapter_mode, host_prepared.adapter_mode); + + let host_values = prepared_f32(&host_prepared); + let iree_values = prepared_f32(&iree_prepared); + assert_eq!(iree_values.len(), host_values.len()); + let hidden_size = host_prepared.embeddings.shape[2]; + for row in 0..host_prepared.sequence_len { + assert!( + row < host_prepared.token_ids.len(), + "prepared row must have a logical token" + ); + if !visual_rows.contains(&row) { + let start = row * hidden_size; + let end = start + hidden_size; + assert_eq!( + &iree_values[start..end], + &host_values[start..end], + "non-visual prepared row {row} must be bit-exact" + ); + } + } + + let mut max_abs = 0.0f32; + let mut squared_error = 0.0f64; + let mut host_square = 0.0f64; + let mut iree_square = 0.0f64; + let mut dot = 0.0f64; + for &row in &visual_rows { + assert!( + row < host_prepared.sequence_len, + "visual target {row} must fit the prepared sequence" + ); + let start = row * hidden_size; + let end = start + hidden_size; + for (&host_value, &iree_value) in + host_values[start..end].iter().zip(&iree_values[start..end]) + { + assert!(host_value.is_finite() && iree_value.is_finite()); + let difference = f64::from(iree_value) - f64::from(host_value); + max_abs = max_abs.max((iree_value - host_value).abs()); + squared_error += difference * difference; + host_square += f64::from(host_value) * f64::from(host_value); + iree_square += f64::from(iree_value) * f64::from(iree_value); + dot += f64::from(host_value) * f64::from(iree_value); + } + } + assert!(host_square > 0.0 && iree_square > 0.0); + let normalized_rmse = (squared_error / host_square).sqrt(); + let cosine = dot / (host_square.sqrt() * iree_square.sqrt()); + eprintln!( + "Molmo host/IREE visual prepared parity: rows={} max_abs={max_abs:.8} \ + cosine={cosine:.8} normalized_rmse={normalized_rmse:.8}", + visual_rows.len() + ); + + // The host reference retains checkpoint BF16/F16 rounding while IREE + // may use F32 contractions on local-task. The combined envelope remains + // tight enough to reject projected-row reordering or a materially different + // vision result without demanding bit equality across the 22-layer ViT, + // attention pooler, and SwiGLU projector. + assert!( + max_abs <= VISUAL_MAX_ABS + && cosine >= VISUAL_MIN_COSINE + && normalized_rmse <= VISUAL_MAX_NORMALIZED_RMSE, + "Molmo visual prepared parity exceeded the documented envelope: \ + max_abs={max_abs}, cosine={cosine}, normalized_rmse={normalized_rmse}" + ); +} diff --git a/src/vision/encoders/molmo.rs b/src/vision/encoders/molmo.rs index 42c25e099..31f5d2441 100644 --- a/src/vision/encoders/molmo.rs +++ b/src/vision/encoders/molmo.rs @@ -38,6 +38,17 @@ use mlxcel_core::layers::{LayerNorm, Linear, UnifiedLinear}; use mlxcel_core::weights::WeightMap; use mlxcel_core::{MlxArray, UniquePtr}; +/// Materialization boundaries reported by the pinned Molmo host diagnostic. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum MolmoHostVisionDiagnosticStage { + PatchEmbeddingStarted, + PatchEmbeddingCompleted, + SelectedLayerStarted(usize), + SelectedLayerCompleted(usize), + ProjectorStarted, + ProjectorCompleted, +} + /// Static configuration for the Molmo v1 vision tower. #[derive(Debug, Clone)] pub struct MolmoVisionConfig { @@ -291,7 +302,12 @@ impl MolmoVisionTransformer { mlxcel_core::add(x, &pe) } - fn forward(&self, x: &MlxArray) -> Vec> { + fn forward_internal( + &self, + x: &MlxArray, + selected_layers: &[usize], + mut diagnostic_progress: Option<&mut dyn FnMut(MolmoHostVisionDiagnosticStage)>, + ) -> (Vec>, Option>) { // x: [B, num_patch, n_pixels]. Pad last dim up to intermediate_size if // the processor produced a narrower patch (quantization padding). let shape = mlxcel_core::array_shape(x); @@ -313,6 +329,17 @@ impl MolmoVisionTransformer { }; let x = self.patch_embedding.forward(&x); + let patch_embedding = diagnostic_progress.as_mut().map(|progress| { + progress(MolmoHostVisionDiagnosticStage::PatchEmbeddingStarted); + let patch_embedding = mlxcel_core::copy(x.as_ref().unwrap()); + mlxcel_core::eval( + patch_embedding + .as_ref() + .expect("diagnostic patch embedding must not be null"), + ); + progress(MolmoHostVisionDiagnosticStage::PatchEmbeddingCompleted); + patch_embedding + }); let bsz = mlxcel_core::array_shape(&x)[0]; // Prepend class embedding: broadcast [emb] -> [B, 1, emb]. @@ -325,12 +352,79 @@ impl MolmoVisionTransformer { let x = self.add_pos_emb(&x); let mut x = self.pre_ln.forward(&x); - let mut hidden_states = Vec::with_capacity(self.blocks.len()); - for block in &self.blocks { + let mut selected_hidden_states = (0..selected_layers.len()) + .map(|_| None) + .collect::>>>(); + if let (Some(progress), Some(first_layer)) = ( + diagnostic_progress.as_mut(), + selected_layers.iter().copied().min(), + ) { + progress(MolmoHostVisionDiagnosticStage::SelectedLayerStarted( + first_layer, + )); + } + for (layer, block) in self.blocks.iter().enumerate() { x = block.forward(&x); - hidden_states.push(mlxcel_core::copy(&x)); + if selected_layers.contains(&layer) { + if let Some(progress) = diagnostic_progress.as_mut() { + mlxcel_core::eval(&x); + progress(MolmoHostVisionDiagnosticStage::SelectedLayerCompleted( + layer, + )); + if let Some(next_layer) = selected_layers + .iter() + .copied() + .filter(|&selected_layer| selected_layer > layer) + .min() + { + progress(MolmoHostVisionDiagnosticStage::SelectedLayerStarted( + next_layer, + )); + } + } + for (slot, &selected_layer) in selected_layers.iter().enumerate() { + if selected_layer == layer { + selected_hidden_states[slot] = Some(mlxcel_core::copy(&x)); + } + } + } } - hidden_states + let selected_hidden_states = selected_hidden_states + .into_iter() + .zip(selected_layers) + .map(|(state, layer)| { + state.unwrap_or_else(|| { + panic!( + "selected Molmo ViT layer {layer} was not evaluated; loaded depth is {}", + self.blocks.len() + ) + }) + }) + .collect(); + (selected_hidden_states, patch_embedding) + } + + fn forward_selected( + &self, + x: &MlxArray, + selected_layers: &[usize], + ) -> Vec> { + self.forward_internal(x, selected_layers, None).0 + } + + #[cfg(any(feature = "xla-diagnostics", feature = "xla-reference-diagnostics"))] + fn forward_selected_with_diagnostics( + &self, + x: &MlxArray, + selected_layers: &[usize], + diagnostic_progress: &mut dyn FnMut(MolmoHostVisionDiagnosticStage), + ) -> (Vec>, UniquePtr) { + let (hidden_states, patch_embedding) = + self.forward_internal(x, selected_layers, Some(diagnostic_progress)); + ( + hidden_states, + patch_embedding.expect("diagnostic patch embedding was requested"), + ) } fn from_weights( @@ -413,11 +507,27 @@ pub struct MolmoVisionModel { num_prefix_tokens: usize, } +#[allow(dead_code)] +pub(crate) struct MolmoHostVisionDiagnostics { + pub(crate) patch_embeddings: UniquePtr, + pub(crate) selected_hidden_states: Vec>, + pub(crate) projected_features: UniquePtr, +} + impl MolmoVisionModel { /// Run the ViT over `[B, T, N, D]` crops, select+concat `vit_layers`, strip /// the cls token, and mask all-padding crops. Returns `[B, T, N_patch, /// feat_dim]` where `feat_dim = image_emb_dim * len(vit_layers)`. - fn encode_image(&self, images: &MlxArray) -> UniquePtr { + fn encode_image_internal( + &self, + images: &MlxArray, + diagnostic_progress: &mut Option<&mut dyn FnMut(MolmoHostVisionDiagnosticStage)>, + ) -> ( + UniquePtr, + Option>, + Vec>, + ) { + let capture_diagnostics = diagnostic_progress.is_some(); let shape = mlxcel_core::array_shape(images); let b = shape[0]; let t = shape[1]; @@ -441,17 +551,46 @@ impl MolmoVisionModel { mlxcel_core::subtract(&one, &all_pad_i) // 1 if valid, 0 if padding }; - let hidden_states = self.image_vit.forward(&flat); + #[cfg(any(feature = "xla-diagnostics", feature = "xla-reference-diagnostics"))] + let (hidden_states, patch_embeddings) = if capture_diagnostics { + let (hidden_states, patch_embeddings) = + self.image_vit.forward_selected_with_diagnostics( + &flat, + &self.vit_layers, + diagnostic_progress + .as_mut() + .map(|progress| &mut **progress) + .expect("diagnostic progress callback must be present"), + ); + (hidden_states, Some(patch_embeddings)) + } else { + ( + self.image_vit.forward_selected(&flat, &self.vit_layers), + None, + ) + }; + #[cfg(not(any(feature = "xla-diagnostics", feature = "xla-reference-diagnostics")))] + let (hidden_states, patch_embeddings) = { + debug_assert!(!capture_diagnostics); + ( + self.image_vit.forward_selected(&flat, &self.vit_layers), + None, + ) + }; + let selected_hidden_states = if capture_diagnostics { + hidden_states + .iter() + .map(|state| mlxcel_core::copy(state.as_ref().unwrap())) + .collect() + } else { + Vec::new() + }; // Select and concat the requested ViT layers along the feature dim. - let mut image_features = - mlxcel_core::copy(hidden_states[self.vit_layers[0]].as_ref().unwrap()); - for &layer in &self.vit_layers[1..] { - image_features = mlxcel_core::concatenate( - &image_features, - hidden_states[layer].as_ref().unwrap(), - -1, - ); + let mut image_features = mlxcel_core::copy(hidden_states[0].as_ref().unwrap()); + for hidden_state in &hidden_states[1..] { + image_features = + mlxcel_core::concatenate(&image_features, hidden_state.as_ref().unwrap(), -1); } // Strip cls token (num_prefix_tokens == 1). @@ -473,7 +612,11 @@ impl MolmoVisionModel { let not_pad_f = mlxcel_core::reshape(¬_pad_f, &[b * t, 1, 1]); let image_features = mlxcel_core::multiply(&image_features, ¬_pad_f); - mlxcel_core::reshape(&image_features, &[b, t, n, feat_dim]) + ( + mlxcel_core::reshape(&image_features, &[b, t, n, feat_dim]), + patch_embeddings, + selected_hidden_states, + ) } /// Apply `pad_and_partial_pad` learned pad embeddings using `image_masks`. @@ -519,13 +662,20 @@ impl MolmoVisionModel { /// Full forward: encode -> pad_embed -> spatial 2x2 pool (attention-meanq) /// -> projector. Returns `[B, num_image, h*w, text_hidden]`. - pub fn forward(&self, images: &MlxArray, image_masks: &MlxArray) -> UniquePtr { + fn forward_internal( + &self, + images: &MlxArray, + image_masks: &MlxArray, + mut diagnostic_progress: Option<&mut dyn FnMut(MolmoHostVisionDiagnosticStage)>, + ) -> (UniquePtr, Option) { let shape = mlxcel_core::array_shape(images); let batch_size = shape[0]; let num_image = shape[1]; + let capture_diagnostics = diagnostic_progress.is_some(); let cfg = &self.config; - let mut image_features = self.encode_image(images); + let (mut image_features, patch_embeddings, selected_hidden_states) = + self.encode_image_internal(images, &mut diagnostic_progress); image_features = self.apply_pad_embed(image_features, image_masks); let feat_dim = mlxcel_core::array_shape(&image_features); @@ -585,7 +735,38 @@ impl MolmoVisionModel { let pooled = mlxcel_core::reshape(&pooled, &[batch_size, num_image, lh * lw, pooled_dim]); // Project to the text hidden dim. - self.image_projector.forward(&pooled) + let projected = self.image_projector.forward(&pooled); + if let Some(progress) = diagnostic_progress.as_mut() { + progress(MolmoHostVisionDiagnosticStage::ProjectorStarted); + mlxcel_core::eval(&projected); + progress(MolmoHostVisionDiagnosticStage::ProjectorCompleted); + } + let diagnostics = capture_diagnostics.then(|| MolmoHostVisionDiagnostics { + patch_embeddings: patch_embeddings.expect("diagnostic patch embedding was requested"), + selected_hidden_states, + projected_features: mlxcel_core::copy(projected.as_ref().unwrap()), + }); + (projected, diagnostics) + } + + pub fn forward(&self, images: &MlxArray, image_masks: &MlxArray) -> UniquePtr { + self.forward_internal(images, image_masks, None).0 + } + + #[cfg(any(feature = "xla-diagnostics", feature = "xla-reference-diagnostics"))] + #[allow(dead_code)] + pub(crate) fn forward_with_diagnostics( + &self, + images: &MlxArray, + image_masks: &MlxArray, + diagnostic_progress: &mut dyn FnMut(MolmoHostVisionDiagnosticStage), + ) -> (UniquePtr, MolmoHostVisionDiagnostics) { + let (projected, diagnostics) = + self.forward_internal(images, image_masks, Some(diagnostic_progress)); + ( + projected, + diagnostics.expect("Molmo host diagnostics were requested"), + ) } pub fn from_weights(