From ecc57fc2fafec1cb7a7d49ad584838def5d7b4da Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 24 Jul 2026 18:09:49 +0900 Subject: [PATCH 01/16] feat(xla): define Molmo sparse additive merge contract Preserve processor-supplied image_input_idx row correspondence across negative sentinels and reject duplicate, out-of-range, over-capacity, malformed, or non-finite inputs before native execution. Keep patch-mask rows distinct from pooled feature rows and prove an embedding-replacement negative fixture differs from the required additive result. Refs #870 --- src/lib/mlxcel-xla/src/lib.rs | 5 + src/lib/mlxcel-xla/src/prepared_molmo.rs | 513 +++++++++++++++++++++++ 2 files changed, 518 insertions(+) create mode 100644 src/lib/mlxcel-xla/src/prepared_molmo.rs diff --git a/src/lib/mlxcel-xla/src/lib.rs b/src/lib/mlxcel-xla/src/lib.rs index 2b1328267..87ae07891 100644 --- a/src/lib/mlxcel-xla/src/lib.rs +++ b/src/lib/mlxcel-xla/src/lib.rs @@ -67,6 +67,7 @@ mod operator_numeric_contract; mod prepared; mod prepared_deepstack; mod prepared_gemma3n; +mod prepared_molmo; #[cfg(feature = "iree")] mod aux; @@ -180,6 +181,10 @@ pub use phi4_audio::{Phi4AudioCheckpoint, Phi4AudioDiagnosticRuntime, Phi4AudioD pub use qwen2_vl_runtime::{ IreeQwen2VlProjector, Qwen2VlVisionExecutionMetrics, Qwen2VlVisionProjection, }; +pub use prepared_molmo::{ + MOLMO_V1_MERGE_MODE, MolmoPreparedImage, MolmoSparseAddError, MolmoSparseAddPair, + MolmoSparseAddPlan, +}; #[cfg(feature = "diagnostics")] pub use vision_runtime::{IreeVisionDiagnosticProjector, VisionDiagnosticProjection}; #[cfg(feature = "iree")] 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..eab67eae4 --- /dev/null +++ b/src/lib/mlxcel-xla/src/prepared_molmo.rs @@ -0,0 +1,513 @@ +// 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, +} + +#[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, + }, + 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::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, + }); + } + } + 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> { + 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 { + text_embeddings[text_start + offset] += projected_features[feature_start + offset]; + } + } + if let Some((flat_index, _)) = text_embeddings + .iter() + .enumerate() + .find(|(_, value)| !value.is_finite()) + { + return Err(MolmoSparseAddError::NonFinite { + tensor: "merged embeddings", + flat_index, + }); + } + Ok(()) + } +} + +#[cfg(test)] +mod tests { + 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() { + 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); + } + + #[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 = MolmoPreparedImage { + pixel_values: vec![0.0; 2 * 4 * 3], + crop_count: 2, + patches_per_crop: 4, + patch_width: 3, + image_masks: vec![1.0; 2 * 4], + projected_rows_per_crop: 1, + image_input_idx: vec![-100, 2], + }; + let plan = prepared.sparse_add_plan(4, 2, 2, 4, 2, 4).unwrap(); + assert_eq!( + plan.pairs(), + [MolmoSparseAddPair { + feature_row: 1, + target_position: 2, + }] + ); + } +} From e46bc5d468de194bb8cd29b74695f90b101781ec Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 24 Jul 2026 18:12:58 +0900 Subject: [PATCH 02/16] fix(xla): make Molmo sparse merge atomic Preflight every destination sum so non-finite overflow cannot partially mutate a reusable batch slot, and reject processor mask values outside the qualified [-1,1] coverage-and-sentinel domain. Keep raw patch-mask rows distinct from pooled feature rows and split tests out of the production module to preserve file-size limits. Refs #870 --- src/lib/mlxcel-xla/src/prepared_molmo.rs | 147 +++++------------- .../mlxcel-xla/src/prepared_molmo_tests.rs | 142 +++++++++++++++++ 2 files changed, 177 insertions(+), 112 deletions(-) create mode 100644 src/lib/mlxcel-xla/src/prepared_molmo_tests.rs diff --git a/src/lib/mlxcel-xla/src/prepared_molmo.rs b/src/lib/mlxcel-xla/src/prepared_molmo.rs index eab67eae4..602e5c73f 100644 --- a/src/lib/mlxcel-xla/src/prepared_molmo.rs +++ b/src/lib/mlxcel-xla/src/prepared_molmo.rs @@ -75,6 +75,10 @@ pub enum MolmoSparseAddError { expected: usize, actual: usize, }, + MaskValue { + flat_index: usize, + value_bits: u32, + }, IndexCount { expected: usize, actual: usize, @@ -121,6 +125,14 @@ impl fmt::Display for MolmoSparseAddError { 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}" @@ -242,6 +254,12 @@ impl MolmoPreparedImage { flat_index, }); } + if !(-1.0..=1.0).contains(value) { + return Err(MolmoSparseAddError::MaskValue { + flat_index, + value_bits: value.to_bits(), + }); + } } Ok(()) } @@ -391,123 +409,28 @@ impl MolmoSparseAddPlan { 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 { - text_embeddings[text_start + offset] += projected_features[feature_start + offset]; + if !(text_embeddings[text_start + offset] + + projected_features[feature_start + offset]) + .is_finite() + { + return Err(MolmoSparseAddError::NonFinite { + tensor: "merged embeddings", + flat_index: text_start + offset, + }); + } } } - if let Some((flat_index, _)) = text_embeddings - .iter() - .enumerate() - .find(|(_, value)| !value.is_finite()) - { - return Err(MolmoSparseAddError::NonFinite { - tensor: "merged embeddings", - 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 { + text_embeddings[text_start + offset] += projected_features[feature_start + offset]; + } } Ok(()) } } #[cfg(test)] -mod tests { - 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() { - 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); - } - - #[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 = MolmoPreparedImage { - pixel_values: vec![0.0; 2 * 4 * 3], - crop_count: 2, - patches_per_crop: 4, - patch_width: 3, - image_masks: vec![1.0; 2 * 4], - projected_rows_per_crop: 1, - image_input_idx: vec![-100, 2], - }; - let plan = prepared.sparse_add_plan(4, 2, 2, 4, 2, 4).unwrap(); - assert_eq!( - plan.pairs(), - [MolmoSparseAddPair { - feature_row: 1, - target_position: 2, - }] - ); - } -} +#[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..8bcbc8611 --- /dev/null +++ b/src/lib/mlxcel-xla/src/prepared_molmo_tests.rs @@ -0,0 +1,142 @@ +// 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() { + 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); +} + +#[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], + } +} From 2a50f81a309c7e2ef9904d4465d945997e9f2050 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 24 Jul 2026 18:16:16 +0900 Subject: [PATCH 03/16] feat(xla): map Molmo decoder into the dense emitter Reuse the shared XLA dense language graph for Molmo v1 by binding its OLMo-style checkpoint names, interleaved RoPE, fused QKV bias, and fused SwiGLU projections without constructing a second decoder. Interpret Molmo intermediate_size as the combined gate/up width and test every fused row slice against the checkpoint schema. Refs #870 --- src/lib/mlxcel-xla/src/emitter/config.rs | 40 +++++++-- src/lib/mlxcel-xla/src/weight_names.rs | 22 +++++ src/lib/mlxcel-xla/src/weights.rs | 106 +++++++++++++++++++++-- 3 files changed, 155 insertions(+), 13 deletions(-) diff --git a/src/lib/mlxcel-xla/src/emitter/config.rs b/src/lib/mlxcel-xla/src/emitter/config.rs index e5c450eaa..9f114cad7 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,19 @@ 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()); + } + 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 +1019,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 +1359,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/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..c81f7dabe 100644 --- a/src/lib/mlxcel-xla/src/weights.rs +++ b/src/lib/mlxcel-xla/src/weights.rs @@ -180,8 +180,13 @@ 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") + }; out.push(WeightSpec::Rows { - name: phi4_base_projection(cfg, format!("{p}mlp.gate_up_proj.weight")), + name: phi4_base_projection(cfg, fused), start: 0, end: inter, }); @@ -208,8 +213,13 @@ 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") + }; out.push(WeightSpec::Rows { - name: phi4_base_projection(cfg, format!("{p}mlp.gate_up_proj.weight")), + name: phi4_base_projection(cfg, fused), start: inter, end: 2 * inter, }); @@ -225,7 +235,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 +284,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 { + let bias = format!("{p}{}", s.q_bias); + out.push(WeightSpec::Rows { + name: bias.clone(), + start: nq, + end: nq + nkv, + }); + out.push(WeightSpec::Rows { + name: bias.clone(), + start: 0, + end: nq, + }); + out.push(WeightSpec::Rows { + 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() { @@ -739,7 +773,65 @@ 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))); + assert!(specs.contains(&WeightSpec::Rows { + name: "language_model.model.blocks.0.ff_proj.weight".to_string(), + start: 0, + end: 12, + })); + assert!(specs.contains(&WeightSpec::Rows { + name: "language_model.model.blocks.0.ff_proj.weight".to_string(), + start: 12, + end: 24, + })); + assert!(specs.contains(&WeightSpec::Rows { + name: "language_model.model.blocks.0.att_proj.bias".to_string(), + start: 8, + end: 12, + })); + } fn phi4_config() -> Config { Config { From 30c84932a8a55474d39260c4d839c0d52a13644f Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 24 Jul 2026 18:36:06 +0900 Subject: [PATCH 04/16] feat(xla): wire Molmo sparse prefill producer --- src/lib/mlxcel-xla/src/weights.rs | 2 +- src/loading/mod.rs | 2 + src/loading/vlm.rs | 5 + src/loading/vlm_molmo_xla.rs | 210 +++++++++++++ src/loading/vlm_special.rs | 4 +- src/multimodal/host_preprocessor.rs | 35 ++- src/multimodal/host_preprocessor_molmo.rs | 344 ++++++++++++++++++++++ 7 files changed, 595 insertions(+), 7 deletions(-) create mode 100644 src/loading/vlm_molmo_xla.rs create mode 100644 src/multimodal/host_preprocessor_molmo.rs diff --git a/src/lib/mlxcel-xla/src/weights.rs b/src/lib/mlxcel-xla/src/weights.rs index c81f7dabe..e7a4fdf0b 100644 --- a/src/lib/mlxcel-xla/src/weights.rs +++ b/src/lib/mlxcel-xla/src/weights.rs @@ -284,7 +284,7 @@ fn weight_specs_q(cfg: &Config, quant: bool) -> Vec { } // q/k/v projection biases (k, q, v order). if cfg.qkv_bias { - if cfg.fused_qkv { + if cfg.fused_qkv && cfg.weight_scheme == crate::emitter::WeightScheme::Molmo { let bias = format!("{p}{}", s.q_bias); out.push(WeightSpec::Rows { name: bias.clone(), diff --git a/src/loading/mod.rs b/src/loading/mod.rs index fb2cab9d4..66a6475cb 100644 --- a/src/loading/mod.rs +++ b/src/loading/mod.rs @@ -61,6 +61,8 @@ pub(crate) use self::vlm::load_llava_host_preprocessor; pub(crate) use self::vlm::load_llava_iree_host_preprocessor; #[cfg(feature = "xla-iree")] pub(crate) use self::vlm::load_qwen2_vl_iree_host_preprocessor; +#[cfg(feature = "xla-backend")] +pub(crate) use self::vlm::load_molmo_host_preprocessor; pub use self::vlm::load_qwen3_omni_speech; #[cfg(feature = "xla-iree")] pub(crate) use self::vlm::{ diff --git a/src/loading/vlm.rs b/src/loading/vlm.rs index b87c6c0a2..c43377034 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,8 @@ 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; 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..5681f2971 --- /dev/null +++ b/src/loading/vlm_molmo_xla.rs @@ -0,0 +1,210 @@ +// 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, + ) +} 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..ffef28257 100644 --- a/src/multimodal/host_preprocessor.rs +++ b/src/multimodal/host_preprocessor.rs @@ -49,6 +49,12 @@ use export::{ validate_sequence_capacity, }; +#[cfg(feature = "xla-backend")] +#[path = "host_preprocessor_molmo.rs"] +mod molmo; +#[cfg(feature = "xla-backend")] +pub use molmo::MolmoHostPreprocessor; + /// Vision implementation selected for OpenXLA multimodal preprocessing. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum XlaVisionBackend { @@ -175,12 +181,33 @@ 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> { + if policy == XlaVisionBackendPolicy::Iree { + return Err(HostPreprocessorError::InvalidConfig( + "Molmo v1 native IREE vision is not available in this build".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..3fa7cb458 --- /dev/null +++ b/src/multimodal/host_preprocessor_molmo.rs @@ -0,0 +1,344 @@ +// 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. + +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_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, +} + +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) + } + + #[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_tower, + 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 mut logical_tokens = Vec::with_capacity(1 + image_tokens.len() + prompt_ids.len()); + logical_tokens.push(MOLMO_V1_BOS_TOKEN_ID); + logical_tokens.extend_from_slice(image_tokens); + logical_tokens.extend( + prompt_ids + .into_iter() + .map(|token| i32::try_from(token).unwrap_or(i32::MAX)), + ); + validate_sequence_capacity(logical_tokens.len(), self.max_sequence_len)?; + 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 { index } else { index + 1 }) + .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 = mlxcel_core::astype( + &self.text_embeddings.forward(&input_ids), + mlxcel_core::dtype::FLOAT32, + ); + validate_embedding_shape( + &mlxcel_core::array_shape(&text), + logical_tokens.len(), + self.hidden_size, + "Molmo text embedding table", + )?; + + 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( + &self.vision_tower.forward(&pixels, &masks), + mlxcel_core::dtype::FLOAT32, + ); + let expected_projected = [ + 1, + crop_count as i32, + self.projected_rows_per_crop as i32, + self.hidden_size as i32, + ]; + if mlxcel_core::array_shape(&projected) != expected_projected { + return Err(HostPreprocessorError::InvalidConfig(format!( + "Molmo projected image shape {:?} does not match {expected_projected:?}", + mlxcel_core::array_shape(&projected) + ))); + } + + let mut text_values = tensor_f32(export_mlx_tensor(&text, "Molmo text embeddings")?)?; + let projected_values = + tensor_f32(export_mlx_tensor(&projected, "Molmo projected features")?)?; + plan.apply(&mut text_values, &projected_values) + .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 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(test)] +mod tests { + use super::*; + + #[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 = std::env::var_os("MLXCEL_TEST_MOLMO_MODEL") + .map(std::path::PathBuf::from) + .expect("MLXCEL_TEST_MOLMO_MODEL is required"); + let preprocessor = MolmoHostPreprocessor::load(&model_path).unwrap(); + 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(); + + let prepared = preprocessor.prepare(&token_ids, &[image]).unwrap(); + + 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); + } +} From e4cc57d6207a35342433db0a7d7ccd354cf7f872 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 24 Jul 2026 18:47:49 +0900 Subject: [PATCH 05/16] feat(xla): emit Molmo native vision graph --- src/lib/mlxcel-xla/src/emitter/mod.rs | 8 + .../mlxcel-xla/src/emitter/molmo_vision.rs | 490 ++++++++++++++++++ .../src/emitter/molmo_vision_config.rs | 454 ++++++++++++++++ .../src/emitter/molmo_vision_config_tests.rs | 66 +++ .../src/emitter/molmo_vision_tests.rs | 45 ++ 5 files changed, 1063 insertions(+) create mode 100644 src/lib/mlxcel-xla/src/emitter/molmo_vision.rs create mode 100644 src/lib/mlxcel-xla/src/emitter/molmo_vision_config.rs create mode 100644 src/lib/mlxcel-xla/src/emitter/molmo_vision_config_tests.rs create mode 100644 src/lib/mlxcel-xla/src/emitter/molmo_vision_tests.rs diff --git a/src/lib/mlxcel-xla/src/emitter/mod.rs b/src/lib/mlxcel-xla/src/emitter/mod.rs index 0065385c7..7eacff5fe 100644 --- a/src/lib/mlxcel-xla/src/emitter/mod.rs +++ b/src/lib/mlxcel-xla/src/emitter/mod.rs @@ -58,6 +58,8 @@ mod gemma3n_weights; mod model; mod moe; pub(crate) mod numeric_ops; +mod molmo_vision; +mod molmo_vision_config; mod phi4_audio; mod qwen2_vl; mod rope; @@ -152,6 +154,12 @@ pub(crate) use model::{ validate_prefill_embeddings_metadata, }; #[allow(unused_imports)] +pub(crate) use molmo_vision::emit_molmo_vision; +#[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..f14dc1be9 --- /dev/null +++ b/src/lib/mlxcel-xla/src/emitter/molmo_vision.rs @@ -0,0 +1,490 @@ +// 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, +}; + +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 scalar_broadcast(builder: &mut Builder, value: f32, shape: Vec) -> Val { + let scalar = builder.const_f32(value); + builder.broadcast(&scalar, &[], shape) +} + +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) +} + +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) +} + +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 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) +} + +pub(crate) fn emit_molmo_vision(config: &MolmoVisionConfig) -> 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 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 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"); + format!( + "module @molmo_vision {{\n func.func public @main({signature}) -> {result} {{\n{body} \ + return {value} : {result}\n }}\n}}\n", + signature = args.declarations.join(", "), + result = projected.ty.render(), + body = builder.body(), + value = projected.name, + ) +} + +#[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..e8e7d34e4 --- /dev/null +++ b/src/lib/mlxcel-xla/src/emitter/molmo_vision_config.rs @@ -0,0 +1,454 @@ +// 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::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, +} + +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) +} + +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 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 mut selected_layers = Vec::with_capacity(raw_layers.len()); + for layer in raw_layers { + let resolved = if layer < 0 { + i64::try_from(layers) + .map_err(|_| "Molmo layer count does not fit i64".to_string())? + + layer + } else { + layer + }; + if resolved < 0 || resolved >= layers as i64 { + return Err(format!("Molmo vit_layers entry {layer} is out of range")); + } + selected_layers.push(resolved as usize); + } + if selected_layers.windows(2).any(|pair| pair[0] == pair[1]) { + return Err("Molmo vit_layers must be unique".to_string()); + } + + 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 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()); + } + 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()); + } + Ok(Self { + max_crops: max_high_res_crops + 1, + patches_per_crop: patch_h * patch_w, + patch_width: patch_size * patch_size * 3, + hidden, + intermediate: 4096, + heads, + head_dim, + positions: patch_h * patch_w + 1, + 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(), + }) + } + + 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, + ) + } + + 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 * self.selected_layers.len(); + 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..19c6611cc --- /dev/null +++ b/src/lib/mlxcel-xla/src/emitter/molmo_vision_config_tests.rs @@ -0,0 +1,66 @@ +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)); +} + +#[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") + ); +} 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..0761891e0 --- /dev/null +++ b/src/lib/mlxcel-xla/src/emitter/molmo_vision_tests.rs @@ -0,0 +1,45 @@ +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")); +} + +#[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); +} From 20dfcfdd934f62426d24794d641249351cc08e20 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 24 Jul 2026 19:16:42 +0900 Subject: [PATCH 06/16] feat(xla): execute Molmo vision through resident IREE Load the pinned Molmo v1 vision, pooling, and projector weights into one auxiliary IREE module, retain only the MLX text embedding table for prepared prefill, and route CLI/server preprocessing through an observable backend policy. Bind static dimensions and OLMo architecture into the artifact identity, validate allocation arithmetic before invocation, and cover the replacement-vs-addition logit sentinel. Refs #870 --- .../src/emitter/molmo_vision_config.rs | 129 +++++- .../src/emitter/molmo_vision_config_tests.rs | 23 + src/lib/mlxcel-xla/src/lib.rs | 4 + .../mlxcel-xla/src/molmo_vision_runtime.rs | 406 ++++++++++++++++++ .../mlxcel-xla/src/prepared_molmo_tests.rs | 11 + src/lib/mlxcel-xla/src/vision_runtime.rs | 14 +- src/loading/mod.rs | 2 + src/loading/vlm.rs | 2 + src/loading/vlm_molmo_xla.rs | 107 +++++ src/multimodal/host_preprocessor.rs | 29 +- src/multimodal/host_preprocessor_molmo.rs | 246 ++++++++--- 11 files changed, 883 insertions(+), 90 deletions(-) create mode 100644 src/lib/mlxcel-xla/src/molmo_vision_runtime.rs diff --git a/src/lib/mlxcel-xla/src/emitter/molmo_vision_config.rs b/src/lib/mlxcel-xla/src/emitter/molmo_vision_config.rs index e8e7d34e4..005481bfd 100644 --- a/src/lib/mlxcel-xla/src/emitter/molmo_vision_config.rs +++ b/src/lib/mlxcel-xla/src/emitter/molmo_vision_config.rs @@ -14,6 +14,7 @@ //! Static Molmo v1 native vision/pool/projector artifact contract. +use std::collections::BTreeSet; use std::path::Path; use serde_json::Value; @@ -53,6 +54,7 @@ pub(crate) struct MolmoVisionConfig { 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 { @@ -64,6 +66,14 @@ fn positive(value: Option, default: usize, name: &str) -> Result 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"); @@ -113,6 +123,11 @@ impl MolmoVisionConfig { } 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, @@ -141,22 +156,21 @@ impl MolmoVisionConfig { 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 { - i64::try_from(layers) - .map_err(|_| "Molmo layer count does not fit i64".to_string())? - + layer - } else { - layer - }; - if resolved < 0 || resolved >= layers as i64 { + 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")); } - selected_layers.push(resolved as usize); - } - if selected_layers.windows(2).any(|pair| pair[0] == pair[1]) { - return Err("Molmo vit_layers must be unique".to_string()); + 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( @@ -203,6 +217,29 @@ impl MolmoVisionConfig { 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, @@ -211,6 +248,40 @@ impl MolmoVisionConfig { 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, @@ -223,15 +294,32 @@ impl MolmoVisionConfig { 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: max_high_res_crops + 1, - patches_per_crop: patch_h * patch_w, - patch_width: patch_size * patch_size * 3, + max_crops, + patches_per_crop, + patch_width, hidden, intermediate: 4096, heads, head_dim, - positions: patch_h * patch_w + 1, + positions, layers, selected_layers, pool_h, @@ -241,6 +329,7 @@ impl MolmoVisionConfig { group_size, bits, layer_norm_eps_bits: (layer_norm_eps as f32).to_bits(), + text_architecture, }) } @@ -255,7 +344,7 @@ impl MolmoVisionConfig { 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}", + selected={:?};pool={}x{};projector={}->{}->{};quant={}/{};merge={};eps={:08x};{}", self.patches_per_crop, self.patch_width, self.hidden, @@ -275,6 +364,7 @@ impl MolmoVisionConfig { self.group_size, MOLMO_V1_MERGE_MODE, self.layer_norm_eps_bits, + self.text_architecture, ) } @@ -369,7 +459,10 @@ impl MolmoVisionConfig { &[self.hidden], ); } - let selected_width = self.hidden * self.selected_layers.len(); + 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( 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 index 19c6611cc..104b3a1ef 100644 --- a/src/lib/mlxcel-xla/src/emitter/molmo_vision_config_tests.rs +++ b/src/lib/mlxcel-xla/src/emitter/molmo_vision_config_tests.rs @@ -24,6 +24,7 @@ fn pinned_defaults_bind_sparse_add_artifact_identity() { 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] @@ -63,4 +64,26 @@ fn rejects_duplicate_or_out_of_range_selected_layers() { .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/lib.rs b/src/lib/mlxcel-xla/src/lib.rs index 87ae07891..e7e129d5f 100644 --- a/src/lib/mlxcel-xla/src/lib.rs +++ b/src/lib/mlxcel-xla/src/lib.rs @@ -78,6 +78,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; @@ -170,6 +172,8 @@ pub use emitter::{ #[cfg(feature = "diagnostics")] pub use iree::PreparedPrefillDiagnostics; #[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, Phi4AudioProjectionMode, Phi4AudioRuntime, phi4_audio_bucket_for_frames, 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..025e1c3b9 --- /dev/null +++ b/src/lib/mlxcel-xla/src/molmo_vision_runtime.rs @@ -0,0 +1,406 @@ +// 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, +) -> 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( + 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, "molmo-v1-vision", 0); + ensure_qualified_auxiliary_artifact(&vmfb, &contract, &weights, |temporary| { + compile_one_to( + &compiler, + mlir, + flags, + &cache, + "molmo-v1-vision", + 0, + temporary, + ) + })?; + IreeAuxiliaryModule::load(device, &vmfb, &contract, weights) +} + +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)?; + 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 { + if crop_count == 0 || crop_count > self.config.max_crops { + return Err(format!( + "Molmo crop count {crop_count} is outside 1..={}", + self.config.max_crops + )); + } + let pixel_count = crop_count + .checked_mul(self.config.patches_per_crop) + .and_then(|count| count.checked_mul(self.config.patch_width)) + .ok_or_else(|| "Molmo pixel shape overflowed".to_string())?; + let mask_count = crop_count + .checked_mul(self.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 static_pixels = checked_product( + &[ + self.config.max_crops, + self.config.patches_per_crop, + self.config.patch_width, + ], + "static pixel", + )?; + let static_masks = checked_product( + &[self.config.max_crops, self.config.patches_per_crop], + "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 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 pixel_shape = [ + self.config.max_crops, + self.config.patches_per_crop, + self.config.patch_width, + ]; + let mask_shape = [self.config.max_crops, self.config.patches_per_crop]; + let started = Instant::now(); + self.module.invoke( + &[ + AuxiliaryInput { + bytes: f32_as_bytes(&padded_pixels), + dtype: AuxiliaryTensorDType::Float32, + shape: &pixel_shape, + }, + AuxiliaryInput { + bytes: f32_as_bytes(&padded_masks), + dtype: AuxiliaryTensorDType::Float32, + shape: &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( + &[ + crop_count, + self.config.projected_rows_per_crop(), + self.config.text_hidden, + ], + "active projected output", + )?; + values.truncate(active_count); + 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())?; + 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: [ + crop_count, + self.config.projected_rows_per_crop(), + self.config.text_hidden, + ], + elapsed_seconds, + upload_bytes, + transfer_bytes, + }) + } +} diff --git a/src/lib/mlxcel-xla/src/prepared_molmo_tests.rs b/src/lib/mlxcel-xla/src/prepared_molmo_tests.rs index 8bcbc8611..9bd6582c2 100644 --- a/src/lib/mlxcel-xla/src/prepared_molmo_tests.rs +++ b/src/lib/mlxcel-xla/src/prepared_molmo_tests.rs @@ -61,6 +61,11 @@ fn empty_active_features_are_valid_and_leave_text_unchanged() { #[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]; @@ -71,6 +76,12 @@ fn replacement_negative_fixture_is_caught_by_additive_oracle() { 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] 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/loading/mod.rs b/src/loading/mod.rs index 66a6475cb..638da9336 100644 --- a/src/loading/mod.rs +++ b/src/loading/mod.rs @@ -63,6 +63,8 @@ pub(crate) use self::vlm::load_llava_iree_host_preprocessor; pub(crate) use self::vlm::load_qwen2_vl_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; pub use self::vlm::load_qwen3_omni_speech; #[cfg(feature = "xla-iree")] pub(crate) use self::vlm::{ diff --git a/src/loading/vlm.rs b/src/loading/vlm.rs index c43377034..b9642affe 100644 --- a/src/loading/vlm.rs +++ b/src/loading/vlm.rs @@ -116,6 +116,8 @@ 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 index 5681f2971..0b93a2447 100644 --- a/src/loading/vlm_molmo_xla.rs +++ b/src/loading/vlm_molmo_xla.rs @@ -208,3 +208,110 @@ pub(crate) fn load_molmo_host_preprocessor( 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/multimodal/host_preprocessor.rs b/src/multimodal/host_preprocessor.rs index ffef28257..f276d4a30 100644 --- a/src/multimodal/host_preprocessor.rs +++ b/src/multimodal/host_preprocessor.rs @@ -195,9 +195,36 @@ 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( - "Molmo v1 native IREE vision is not available in this build".to_string(), + "MLXCEL_XLA_VISION_BACKEND=iree requires the xla-iree feature".to_string(), )); } let preprocessor = MolmoHostPreprocessor::load(model_path)?; diff --git a/src/multimodal/host_preprocessor_molmo.rs b/src/multimodal/host_preprocessor_molmo.rs index 3fa7cb458..9a8640715 100644 --- a/src/multimodal/host_preprocessor_molmo.rs +++ b/src/multimodal/host_preprocessor_molmo.rs @@ -19,6 +19,8 @@ //! 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; @@ -40,7 +42,7 @@ 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_tower: MolmoVisionModel, + vision: MolmoVision, text_embeddings: Molmo2Embedding, tokenizer: MlxcelTokenizer, hidden_size: usize, @@ -50,12 +52,23 @@ pub struct MolmoHostPreprocessor { 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, @@ -83,7 +96,37 @@ impl MolmoHostPreprocessor { } Ok(Self { processor, - vision_tower, + 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, @@ -120,15 +163,25 @@ impl MolmoHostPreprocessor { .tokenizer .encode(&prompt, false) .map_err(|error| HostPreprocessorError::InvalidConfig(error.to_string()))?; - let mut logical_tokens = Vec::with_capacity(1 + image_tokens.len() + prompt_ids.len()); + 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 - .into_iter() - .map(|token| i32::try_from(token).unwrap_or(i32::MAX)), - ); - validate_sequence_capacity(logical_tokens.len(), self.max_sequence_len)?; + logical_tokens.extend(prompt_ids); Ok(logical_tokens) } @@ -148,8 +201,16 @@ impl MolmoHostPreprocessor { let shifted_indices = output .image_input_idx .iter() - .map(|&index| if index < 0 { index } else { index + 1 }) - .collect::>(); + .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, @@ -189,43 +250,71 @@ impl MolmoHostPreprocessor { "Molmo text embedding table", )?; - 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( - &self.vision_tower.forward(&pixels, &masks), - mlxcel_core::dtype::FLOAT32, - ); - let expected_projected = [ - 1, - crop_count as i32, - self.projected_rows_per_crop as i32, - self.hidden_size as i32, - ]; - if mlxcel_core::array_shape(&projected) != expected_projected { - return Err(HostPreprocessorError::InvalidConfig(format!( - "Molmo projected image shape {:?} does not match {expected_projected:?}", - mlxcel_core::array_shape(&projected) - ))); - } - let mut text_values = tensor_f32(export_mlx_tensor(&text, "Molmo text embeddings")?)?; - let projected_values = - tensor_f32(export_mlx_tensor(&projected, "Molmo projected features")?)?; + 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(&mut text_values, &projected_values) .map_err(|error| HostPreprocessorError::InvalidConfig(error.to_string()))?; let mut bytes = Vec::with_capacity( @@ -253,6 +342,14 @@ impl MolmoHostPreprocessor { } 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], @@ -304,6 +401,29 @@ fn format_prompt(prompt: &str) -> String { mod tests { use super::*; + 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); + } + #[test] fn prompt_format_matches_molmo_v1_processor_contract() { assert_eq!( @@ -320,25 +440,23 @@ mod tests { #[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 = std::env::var_os("MLXCEL_TEST_MOLMO_MODEL") - .map(std::path::PathBuf::from) - .expect("MLXCEL_TEST_MOLMO_MODEL is required"); + let (model_path, token_ids, image) = real_inputs(); let preprocessor = MolmoHostPreprocessor::load(&model_path).unwrap(); - 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(); - let prepared = preprocessor.prepare(&token_ids, &[image]).unwrap(); + assert_real_prepared(&prepared); + } - 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); + #[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); } } From cefddd67f4e1e0f9da6851c261643cca77c3af3c Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 24 Jul 2026 19:17:10 +0900 Subject: [PATCH 07/16] fix(xla): slice Molmo fused attention bias Represent Molmo's rank-one fused QKV bias with a distinct element-slice schema instead of passing it through the rank-two projection row slicer. Reject incompatible Molmo architecture variants before graph construction so the pinned OLMo layout cannot silently drift. Refs #870 --- src/lib/mlxcel-xla/src/emitter/config.rs | 16 ++++++++++++++++ src/lib/mlxcel-xla/src/iree.rs | 21 +++++++++++++++++++++ src/lib/mlxcel-xla/src/lib.rs | 8 ++++---- src/lib/mlxcel-xla/src/weights.rs | 16 ++++++++++++---- src/loading/mod.rs | 4 ++-- 5 files changed, 55 insertions(+), 10 deletions(-) diff --git a/src/lib/mlxcel-xla/src/emitter/config.rs b/src/lib/mlxcel-xla/src/emitter/config.rs index 9f114cad7..7d3488a88 100644 --- a/src/lib/mlxcel-xla/src/emitter/config.rs +++ b/src/lib/mlxcel-xla/src/emitter/config.rs @@ -880,6 +880,22 @@ impl Config { ) { 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; 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 e7e129d5f..cd5a1cf15 100644 --- a/src/lib/mlxcel-xla/src/lib.rs +++ b/src/lib/mlxcel-xla/src/lib.rs @@ -181,14 +181,14 @@ pub use phi4_audio::{ #[cfg(feature = "iree")] #[doc(hidden)] pub use phi4_audio::{Phi4AudioCheckpoint, Phi4AudioDiagnosticRuntime, Phi4AudioDiagnostics}; -#[cfg(feature = "iree")] -pub use qwen2_vl_runtime::{ - IreeQwen2VlProjector, Qwen2VlVisionExecutionMetrics, Qwen2VlVisionProjection, -}; pub use prepared_molmo::{ MOLMO_V1_MERGE_MODE, MolmoPreparedImage, MolmoSparseAddError, MolmoSparseAddPair, MolmoSparseAddPlan, }; +#[cfg(feature = "iree")] +pub use qwen2_vl_runtime::{ + IreeQwen2VlProjector, Qwen2VlVisionExecutionMetrics, Qwen2VlVisionProjection, +}; #[cfg(feature = "diagnostics")] pub use vision_runtime::{IreeVisionDiagnosticProjector, VisionDiagnosticProjection}; #[cfg(feature = "iree")] diff --git a/src/lib/mlxcel-xla/src/weights.rs b/src/lib/mlxcel-xla/src/weights.rs index e7a4fdf0b..f604833fb 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, } } @@ -286,17 +294,17 @@ fn weight_specs_q(cfg: &Config, quant: bool) -> Vec { if cfg.qkv_bias { if cfg.fused_qkv && cfg.weight_scheme == crate::emitter::WeightScheme::Molmo { let bias = format!("{p}{}", s.q_bias); - out.push(WeightSpec::Rows { + out.push(WeightSpec::Elements { name: bias.clone(), start: nq, end: nq + nkv, }); - out.push(WeightSpec::Rows { + out.push(WeightSpec::Elements { name: bias.clone(), start: 0, end: nq, }); - out.push(WeightSpec::Rows { + out.push(WeightSpec::Elements { name: bias, start: nq + nkv, end: nq + 2 * nkv, @@ -826,7 +834,7 @@ mod tests { start: 12, end: 24, })); - assert!(specs.contains(&WeightSpec::Rows { + assert!(specs.contains(&WeightSpec::Elements { name: "language_model.model.blocks.0.att_proj.bias".to_string(), start: 8, end: 12, diff --git a/src/loading/mod.rs b/src/loading/mod.rs index 638da9336..3c1da6d9d 100644 --- a/src/loading/mod.rs +++ b/src/loading/mod.rs @@ -59,12 +59,12 @@ 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-iree")] -pub(crate) use self::vlm::load_qwen2_vl_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; #[cfg(feature = "xla-iree")] pub(crate) use self::vlm::{ From 234b6c249fa10c912cb5ad209befeb6aaef074db Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 24 Jul 2026 20:04:49 +0900 Subject: [PATCH 08/16] fix(xla): preserve Molmo decoder and merge semantics Map Molmo's fused FFN value and gate halves into the shared emitter's gate and value argument order so SiLU is applied to the correct half. Perform sparse visual addition in the original F16 or BF16 text-embedding dtype before widening the owned prepared payload to F32. Validation: cargo test -p mlxcel-xla molmo --lib; cargo test -p mlxcel-xla phi4 --lib; cargo clippy --fix --features xla-iree --lib --allow-dirty. Refs #870 --- src/lib/mlxcel-xla/src/lib.rs | 4 +- src/lib/mlxcel-xla/src/prepared_molmo.rs | 44 ++++++++++++-- .../mlxcel-xla/src/prepared_molmo_tests.rs | 22 +++++++ src/lib/mlxcel-xla/src/weights.rs | 60 ++++++++++++++----- src/multimodal/host_preprocessor_molmo.rs | 14 +++-- 5 files changed, 118 insertions(+), 26 deletions(-) diff --git a/src/lib/mlxcel-xla/src/lib.rs b/src/lib/mlxcel-xla/src/lib.rs index cd5a1cf15..15937fb61 100644 --- a/src/lib/mlxcel-xla/src/lib.rs +++ b/src/lib/mlxcel-xla/src/lib.rs @@ -182,8 +182,8 @@ pub use phi4_audio::{ #[doc(hidden)] pub use phi4_audio::{Phi4AudioCheckpoint, Phi4AudioDiagnosticRuntime, Phi4AudioDiagnostics}; pub use prepared_molmo::{ - MOLMO_V1_MERGE_MODE, MolmoPreparedImage, MolmoSparseAddError, MolmoSparseAddPair, - MolmoSparseAddPlan, + MOLMO_V1_MERGE_MODE, MolmoEmbeddingDType, MolmoPreparedImage, MolmoSparseAddError, + MolmoSparseAddPair, MolmoSparseAddPlan, }; #[cfg(feature = "iree")] pub use qwen2_vl_runtime::{ diff --git a/src/lib/mlxcel-xla/src/prepared_molmo.rs b/src/lib/mlxcel-xla/src/prepared_molmo.rs index 602e5c73f..b40cd1980 100644 --- a/src/lib/mlxcel-xla/src/prepared_molmo.rs +++ b/src/lib/mlxcel-xla/src/prepared_molmo.rs @@ -58,6 +58,24 @@ pub struct MolmoSparseAddPlan { 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::weights::half_to_f32(crate::weights::f32_to_f16_bits(value)), + Self::BFloat16 => crate::weights::round_bf16_f32(value), + Self::Float32 => value, + } + } +} + #[derive(Debug, Clone, PartialEq, Eq)] pub enum MolmoSparseAddError { ZeroDimension(&'static str), @@ -366,6 +384,21 @@ impl MolmoSparseAddPlan { &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])?; @@ -409,10 +442,9 @@ impl MolmoSparseAddPlan { 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 { - if !(text_embeddings[text_start + offset] - + projected_features[feature_start + offset]) - .is_finite() - { + 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, @@ -424,7 +456,9 @@ impl MolmoSparseAddPlan { 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 { - text_embeddings[text_start + offset] += projected_features[feature_start + offset]; + 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(()) diff --git a/src/lib/mlxcel-xla/src/prepared_molmo_tests.rs b/src/lib/mlxcel-xla/src/prepared_molmo_tests.rs index 9bd6582c2..42c88d027 100644 --- a/src/lib/mlxcel-xla/src/prepared_molmo_tests.rs +++ b/src/lib/mlxcel-xla/src/prepared_molmo_tests.rs @@ -84,6 +84,28 @@ fn replacement_negative_fixture_is_caught_by_additive_oracle() { ); } +#[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!( diff --git a/src/lib/mlxcel-xla/src/weights.rs b/src/lib/mlxcel-xla/src/weights.rs index f604833fb..556acdcb0 100644 --- a/src/lib/mlxcel-xla/src/weights.rs +++ b/src/lib/mlxcel-xla/src/weights.rs @@ -193,10 +193,17 @@ fn weight_specs_q(cfg: &Config, quant: bool) -> Vec { } 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, fused), - start: 0, - end: inter, + start, + end, }); } else { push_proj( @@ -226,10 +233,17 @@ fn weight_specs_q(cfg: &Config, quant: bool) -> Vec { } 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, fused), - start: inter, - end: 2 * inter, + start, + end, }); } else if cfg.dense_mlp { out.push(WeightSpec::Whole(format!("{p}mlp.c_fc.weight"))); @@ -824,16 +838,34 @@ mod tests { assert!(specs.contains(&rows(8, 12))); assert!(specs.contains(&rows(0, 8))); assert!(specs.contains(&rows(12, 16))); - assert!(specs.contains(&WeightSpec::Rows { - name: "language_model.model.blocks.0.ff_proj.weight".to_string(), - start: 0, - end: 12, - })); - assert!(specs.contains(&WeightSpec::Rows { - name: "language_model.model.blocks.0.ff_proj.weight".to_string(), - start: 12, - end: 24, - })); + 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, diff --git a/src/multimodal/host_preprocessor_molmo.rs b/src/multimodal/host_preprocessor_molmo.rs index 9a8640715..0be22716d 100644 --- a/src/multimodal/host_preprocessor_molmo.rs +++ b/src/multimodal/host_preprocessor_molmo.rs @@ -239,10 +239,13 @@ impl MolmoHostPreprocessor { &logical_tokens, &[1, usize_to_i32(logical_tokens.len(), "sequence length")?], ); - let text = mlxcel_core::astype( - &self.text_embeddings.forward(&input_ids), - mlxcel_core::dtype::FLOAT32, - ); + 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(), @@ -250,6 +253,7 @@ impl MolmoHostPreprocessor { "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) => { @@ -315,7 +319,7 @@ impl MolmoHostPreprocessor { projection.values } }; - plan.apply(&mut text_values, &projected_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 From 2939481675040705876954994962af5c0a9f986d Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 24 Jul 2026 21:05:35 +0900 Subject: [PATCH 09/16] test(xla): add Molmo prepared-prefill parity gate --- src/multimodal/host_preprocessor_molmo.rs | 64 +---- .../host_preprocessor_molmo_tests.rs | 242 ++++++++++++++++++ 2 files changed, 244 insertions(+), 62 deletions(-) create mode 100644 src/multimodal/host_preprocessor_molmo_tests.rs diff --git a/src/multimodal/host_preprocessor_molmo.rs b/src/multimodal/host_preprocessor_molmo.rs index 0be22716d..ebc03aa43 100644 --- a/src/multimodal/host_preprocessor_molmo.rs +++ b/src/multimodal/host_preprocessor_molmo.rs @@ -402,65 +402,5 @@ fn format_prompt(prompt: &str) -> String { } #[cfg(test)] -mod tests { - use super::*; - - 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); - } - - #[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); - } -} +#[path = "host_preprocessor_molmo_tests.rs"] +mod tests; 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}" + ); +} From 691d28da02e6e26cc4a1fc455c214bfc8b3ef82a Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Sun, 26 Jul 2026 22:38:18 +0900 Subject: [PATCH 10/16] test(xla): add strict Molmo boundary diagnostics The existing real-checkpoint gate compared only the final PreparedPrefill values, so processor, vision, projector, sparse-merge, and decoder propagation regressions could not be localized. Add a CPU local-task reference feature and a diagnostics-only Molmo graph that captures pinned processor contracts, patch embeddings, selected vision layers, projected features, prepared sparse-add rows, and one resident decoder's KV and logits boundary. Negative controls reject Molmo2 token scanning and replacement semantics without loading a duplicate decoder. Validation: Molmo unit tests 15/15; CPU diagnostic tests 2/2 with the actual-checkpoint test ignored; diagnostic StableHLO compiled for IREE local-task; feature clippy completed with pre-existing unrelated warnings. The pinned 5.3 GB actual-checkpoint gate was not run in this task environment. Refs #870 --- Cargo.toml | 9 +- src/lib/mlxcel-xla/src/emitter/mod.rs | 3 + .../mlxcel-xla/src/emitter/molmo_vision.rs | 103 +++-- .../src/emitter/molmo_vision_ops.rs | 50 +++ .../src/emitter/molmo_vision_tests.rs | 30 ++ src/lib/mlxcel-xla/src/lib.rs | 2 + .../src/molmo_vision_diagnostics.rs | 191 ++++++++++ .../mlxcel-xla/src/molmo_vision_runtime.rs | 169 +++++---- src/multimodal/host_preprocessor_molmo.rs | 6 + ..._preprocessor_molmo_diagnostics_support.rs | 239 ++++++++++++ ...st_preprocessor_molmo_diagnostics_tests.rs | 351 ++++++++++++++++++ src/vision/encoders/molmo.rs | 107 +++++- 12 files changed, 1120 insertions(+), 140 deletions(-) create mode 100644 src/lib/mlxcel-xla/src/emitter/molmo_vision_ops.rs create mode 100644 src/lib/mlxcel-xla/src/molmo_vision_diagnostics.rs create mode 100644 src/multimodal/host_preprocessor_molmo_diagnostics_support.rs create mode 100644 src/multimodal/host_preprocessor_molmo_diagnostics_tests.rs diff --git a/Cargo.toml b/Cargo.toml index ab506bce0..ff6d190ef 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -71,11 +71,18 @@ 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"] +# +# 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"] +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/src/lib/mlxcel-xla/src/emitter/mod.rs b/src/lib/mlxcel-xla/src/emitter/mod.rs index 7eacff5fe..8571a17da 100644 --- a/src/lib/mlxcel-xla/src/emitter/mod.rs +++ b/src/lib/mlxcel-xla/src/emitter/mod.rs @@ -60,6 +60,7 @@ mod moe; pub(crate) mod numeric_ops; mod molmo_vision; mod molmo_vision_config; +mod molmo_vision_ops; mod phi4_audio; mod qwen2_vl; mod rope; @@ -155,6 +156,8 @@ pub(crate) use model::{ }; #[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, diff --git a/src/lib/mlxcel-xla/src/emitter/molmo_vision.rs b/src/lib/mlxcel-xla/src/emitter/molmo_vision.rs index f14dc1be9..5eafe09af 100644 --- a/src/lib/mlxcel-xla/src/emitter/molmo_vision.rs +++ b/src/lib/mlxcel-xla/src/emitter/molmo_vision.rs @@ -18,6 +18,7 @@ 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, @@ -119,47 +120,6 @@ fn layer_norm(builder: &mut Builder, value: &Val, weight: &Val, bias: &Val, epsi builder.add(&normalized, &bias) } -fn scalar_broadcast(builder: &mut Builder, value: f32, shape: Vec) -> Val { - let scalar = builder.const_f32(value); - builder.broadcast(&scalar, &[], shape) -} - -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) -} - -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) -} - fn self_attention( builder: &mut Builder, args: &mut Args, @@ -366,15 +326,7 @@ fn attention_pool( ) } -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) -} - -pub(crate) fn emit_molmo_vision(config: &MolmoVisionConfig) -> String { +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(); @@ -402,6 +354,7 @@ pub(crate) fn emit_molmo_vision(config: &MolmoVisionConfig) -> String { &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, @@ -435,6 +388,16 @@ pub(crate) fn emit_molmo_vision(config: &MolmoVisionConfig) -> String { } } } + 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"); @@ -475,16 +438,48 @@ pub(crate) fn emit_molmo_vision(config: &MolmoVisionConfig) -> String { ], ); 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} {{\n{body} \ - return {value} : {result}\n }}\n}}\n", + "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(", "), - result = projected.ty.render(), body = builder.body(), - value = projected.name, + 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_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 index 0761891e0..d8ae01e37 100644 --- a/src/lib/mlxcel-xla/src/emitter/molmo_vision_tests.rs +++ b/src/lib/mlxcel-xla/src/emitter/molmo_vision_tests.rs @@ -25,6 +25,17 @@ fn pinned_graph_contains_vit_pool_projector_and_static_output() { 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() { @@ -43,3 +54,22 @@ fn pinned_graph_compiles_for_cpu() { .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/lib.rs b/src/lib/mlxcel-xla/src/lib.rs index 15937fb61..ae5b915ec 100644 --- a/src/lib/mlxcel-xla/src/lib.rs +++ b/src/lib/mlxcel-xla/src/lib.rs @@ -171,6 +171,8 @@ 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")] 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 index 025e1c3b9..b927b6d73 100644 --- a/src/lib/mlxcel-xla/src/molmo_vision_runtime.rs +++ b/src/lib/mlxcel-xla/src/molmo_vision_runtime.rs @@ -203,6 +203,7 @@ fn compile_and_load( device: &str, config: &MolmoVisionConfig, mlir: &str, + tag: &str, ) -> Result { let compiler = iree_compile_bin()?; let flags = target_flags(device)?; @@ -224,21 +225,87 @@ fn compile_and_load( ), compiler_generation_identity(&compiler, flags, mlir)?, )?; - let vmfb = cached_vmfb_path(&compiler, mlir, flags, &cache, "molmo-v1-vision", 0); + 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, - "molmo-v1-vision", - 0, - 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, @@ -248,7 +315,7 @@ 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)?; + let module = compile_and_load(model_dir, device, &config, &mlir, "molmo-v1-vision")?; Ok(Self { module, config }) } @@ -288,55 +355,7 @@ impl IreeMolmoVisionProjector { masks: &[f32], crop_count: usize, ) -> Result { - if crop_count == 0 || crop_count > self.config.max_crops { - return Err(format!( - "Molmo crop count {crop_count} is outside 1..={}", - self.config.max_crops - )); - } - let pixel_count = crop_count - .checked_mul(self.config.patches_per_crop) - .and_then(|count| count.checked_mul(self.config.patch_width)) - .ok_or_else(|| "Molmo pixel shape overflowed".to_string())?; - let mask_count = crop_count - .checked_mul(self.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 static_pixels = checked_product( - &[ - self.config.max_crops, - self.config.patches_per_crop, - self.config.patch_width, - ], - "static pixel", - )?; - let static_masks = checked_product( - &[self.config.max_crops, self.config.patches_per_crop], - "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 input = validate_and_pad_input(&self.config, pixels, masks, crop_count)?; let output_shape = [ self.config.max_crops, self.config.projected_rows_per_crop(), @@ -347,24 +366,18 @@ impl IreeMolmoVisionProjector { .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 pixel_shape = [ - self.config.max_crops, - self.config.patches_per_crop, - self.config.patch_width, - ]; - let mask_shape = [self.config.max_crops, self.config.patches_per_crop]; let started = Instant::now(); self.module.invoke( &[ AuxiliaryInput { - bytes: f32_as_bytes(&padded_pixels), + bytes: f32_as_bytes(&input.padded_pixels), dtype: AuxiliaryTensorDType::Float32, - shape: &pixel_shape, + shape: &input.pixel_shape, }, AuxiliaryInput { - bytes: f32_as_bytes(&padded_masks), + bytes: f32_as_bytes(&input.padded_masks), dtype: AuxiliaryTensorDType::Float32, - shape: &mask_shape, + shape: &input.mask_shape, }, ], &mut [AuxiliaryOutput { @@ -377,30 +390,32 @@ impl IreeMolmoVisionProjector { let mut values = checked_f32_output("IREE Molmo projected output", output)?; let active_count = checked_product( &[ - crop_count, + input.crop_count, self.config.projected_rows_per_crop(), self.config.text_hidden, ], "active projected output", )?; values.truncate(active_count); - 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())?; 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: [ - crop_count, + input.crop_count, self.config.projected_rows_per_crop(), self.config.text_hidden, ], elapsed_seconds, - upload_bytes, + 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/multimodal/host_preprocessor_molmo.rs b/src/multimodal/host_preprocessor_molmo.rs index ebc03aa43..f847bd767 100644 --- a/src/multimodal/host_preprocessor_molmo.rs +++ b/src/multimodal/host_preprocessor_molmo.rs @@ -401,6 +401,12 @@ fn format_prompt(prompt: &str) -> String { } } +#[cfg(all( + test, + any(feature = "xla-diagnostics", feature = "xla-reference-diagnostics") +))] +#[path = "host_preprocessor_molmo_diagnostics_tests.rs"] +mod diagnostics_tests; #[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..2fe8c1058 --- /dev/null +++ b/src/multimodal/host_preprocessor_molmo_diagnostics_support.rs @@ -0,0 +1,239 @@ +use std::fs; +use std::io::{self, Write}; +use std::path::Path; + +use image::DynamicImage; +use mlxcel_core::session::{PreparedPrefill, PreparedTensorDType}; +use sha2::{Digest, Sha256}; + +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"; +pub(super) const PINNED_IMAGE_TOKEN_IDS_SHA256: &str = + "23a0878a6a2818f113460cc9a116f3d68f4d3d4b1a34b1f34dfe7805395562fa"; +pub(super) const PINNED_IMAGE_INPUT_IDX_SHA256: &str = + "243da1bc98ebb61f5d98696623b41ade1d643994db6fcfb764392452d3f9c1b4"; +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()) +} + +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() +} + +fn fixture_image() -> DynamicImage { + image::open("tests/fixtures/test_image.png").expect("load pinned Molmo image fixture") +} + +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..7a3baa676 --- /dev/null +++ b/src/multimodal/host_preprocessor_molmo_diagnostics_tests.rs @@ -0,0 +1,351 @@ +// 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 super::*; +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=/path/to/iree-dist \ +/// MLXCEL_MOLMO_FIXTURE=/path/to/Molmo-7B-D-0924-4bit \ +/// cargo test --features xla-reference-diagnostics --lib \ +/// pinned_molmo_eager_mlx_matches_iree_boundaries -- --ignored --nocapture +/// ``` +/// +/// 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. +#[test] +#[ignore = "requires pinned Molmo checkpoint and IREE local-task"] +fn pinned_molmo_eager_mlx_matches_iree_boundaries() { + mlxcel_core::set_default_device(false); + 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_eq!( + device, "local-task", + "the #870 reference gate qualifies IREE local-task only" + ); + 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"); + 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 (_, host_diagnostics) = vision.forward_with_diagnostics(&pixels, &masks); + 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}"); + } +} diff --git a/src/vision/encoders/molmo.rs b/src/vision/encoders/molmo.rs index 42c25e099..60182ca5b 100644 --- a/src/vision/encoders/molmo.rs +++ b/src/vision/encoders/molmo.rs @@ -291,7 +291,11 @@ impl MolmoVisionTransformer { mlxcel_core::add(x, &pe) } - fn forward(&self, x: &MlxArray) -> Vec> { + fn forward_internal( + &self, + x: &MlxArray, + capture_patch_embedding: bool, + ) -> (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 +317,8 @@ impl MolmoVisionTransformer { }; let x = self.patch_embedding.forward(&x); + let patch_embedding = + capture_patch_embedding.then(|| mlxcel_core::copy(x.as_ref().unwrap())); let bsz = mlxcel_core::array_shape(&x)[0]; // Prepend class embedding: broadcast [emb] -> [B, 1, emb]. @@ -330,7 +336,23 @@ impl MolmoVisionTransformer { x = block.forward(&x); hidden_states.push(mlxcel_core::copy(&x)); } - hidden_states + (hidden_states, patch_embedding) + } + + fn forward(&self, x: &MlxArray) -> Vec> { + self.forward_internal(x, false).0 + } + + #[cfg(any(feature = "xla-diagnostics", feature = "xla-reference-diagnostics"))] + fn forward_with_patch_embedding( + &self, + x: &MlxArray, + ) -> (Vec>, UniquePtr) { + let (hidden_states, patch_embedding) = self.forward_internal(x, true); + ( + hidden_states, + patch_embedding.expect("diagnostic patch embedding was requested"), + ) } fn from_weights( @@ -413,11 +435,26 @@ 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, + capture_diagnostics: bool, + ) -> ( + UniquePtr, + Option>, + Vec>, + ) { let shape = mlxcel_core::array_shape(images); let b = shape[0]; let t = shape[1]; @@ -441,7 +478,27 @@ 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_with_patch_embedding(&flat); + (hidden_states, Some(patch_embeddings)) + } else { + (self.image_vit.forward(&flat), None) + }; + #[cfg(not(any(feature = "xla-diagnostics", feature = "xla-reference-diagnostics")))] + let (hidden_states, patch_embeddings) = { + debug_assert!(!capture_diagnostics); + (self.image_vit.forward(&flat), None) + }; + let selected_hidden_states = if capture_diagnostics { + self.vit_layers + .iter() + .map(|&layer| mlxcel_core::copy(hidden_states[layer].as_ref().unwrap())) + .collect() + } else { + Vec::new() + }; // Select and concat the requested ViT layers along the feature dim. let mut image_features = @@ -473,7 +530,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 +580,19 @@ 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, + capture_diagnostics: bool, + ) -> (UniquePtr, Option) { let shape = mlxcel_core::array_shape(images); let batch_size = shape[0]; let num_image = shape[1]; 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, capture_diagnostics); image_features = self.apply_pad_embed(image_features, image_masks); let feat_dim = mlxcel_core::array_shape(&image_features); @@ -585,7 +652,31 @@ 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); + 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, false).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, + ) -> (UniquePtr, MolmoHostVisionDiagnostics) { + let (projected, diagnostics) = self.forward_internal(images, image_masks, true); + ( + projected, + diagnostics.expect("Molmo host diagnostics were requested"), + ) } pub fn from_weights( From 33441139013bebe05239402cf793ee853f701e10 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Sun, 26 Jul 2026 22:57:27 +0900 Subject: [PATCH 11/16] test(xla): stage Molmo eager vision diagnostics Materialize the eager MLX patch embedding, each selected ViT layer, and the pool/projector at explicit diagnostic boundaries so the pinned #870 gate reports progress before and after every potentially long stage. Retain only configured selected hidden states instead of every ViT layer, bounding the lazy graph and allowing the next actual checkpoint run to identify the first stalled or divergent stage under the five-minute no-output policy. Validation: cargo check --features xla-reference-diagnostics --lib; cargo test --features xla-reference-diagnostics --lib strict_oracle_rejects_molmo2_token_scan_and_replacement; cargo fmt --all -- --check; git diff --check. The actual full-checkpoint gate was intentionally not rerun. Refs #870 --- ...st_preprocessor_molmo_diagnostics_tests.rs | 22 ++- src/vision/encoders/molmo.rs | 148 ++++++++++++++---- 2 files changed, 140 insertions(+), 30 deletions(-) diff --git a/src/multimodal/host_preprocessor_molmo_diagnostics_tests.rs b/src/multimodal/host_preprocessor_molmo_diagnostics_tests.rs index 7a3baa676..d9e63a626 100644 --- a/src/multimodal/host_preprocessor_molmo_diagnostics_tests.rs +++ b/src/multimodal/host_preprocessor_molmo_diagnostics_tests.rs @@ -194,7 +194,27 @@ fn pinned_molmo_eager_mlx_matches_iree_boundaries() { let MolmoVision::Host(vision) = &host.vision else { panic!("host preprocessor must retain the eager MLX vision tower"); }; - let (_, host_diagnostics) = vision.forward_with_diagnostics(&pixels, &masks); + 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"); + } + crate::vision::encoders::molmo::MolmoHostVisionDiagnosticStage::PatchEmbeddingCompleted => { + progress("eager MLX patch embedding materialized"); + } + crate::vision::encoders::molmo::MolmoHostVisionDiagnosticStage::SelectedLayerStarted(layer) => { + progress(&format!("eager MLX selected ViT layer {layer} materialization started")); + } + crate::vision::encoders::molmo::MolmoHostVisionDiagnosticStage::SelectedLayerCompleted(layer) => { + progress(&format!("eager MLX selected ViT layer {layer} materialized")); + } + crate::vision::encoders::molmo::MolmoHostVisionDiagnosticStage::ProjectorStarted => { + progress("eager MLX pool/projector materialization started"); + } + crate::vision::encoders::molmo::MolmoHostVisionDiagnosticStage::ProjectorCompleted => { + progress("eager MLX pool/projector materialized"); + } + }); let host_patch = mlx_f32(host_diagnostics.patch_embeddings.as_ref().unwrap()); let host_selected = host_diagnostics .selected_hidden_states diff --git a/src/vision/encoders/molmo.rs b/src/vision/encoders/molmo.rs index 60182ca5b..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 { @@ -294,7 +305,8 @@ impl MolmoVisionTransformer { fn forward_internal( &self, x: &MlxArray, - capture_patch_embedding: bool, + 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). @@ -317,8 +329,17 @@ impl MolmoVisionTransformer { }; let x = self.patch_embedding.forward(&x); - let patch_embedding = - capture_patch_embedding.then(|| mlxcel_core::copy(x.as_ref().unwrap())); + 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]. @@ -331,24 +352,75 @@ 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, patch_embedding) + 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(&self, x: &MlxArray) -> Vec> { - self.forward_internal(x, false).0 + 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_with_patch_embedding( + 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, true); + let (hidden_states, patch_embedding) = + self.forward_internal(x, selected_layers, Some(diagnostic_progress)); ( hidden_states, patch_embedding.expect("diagnostic patch embedding was requested"), @@ -449,12 +521,13 @@ impl MolmoVisionModel { fn encode_image_internal( &self, images: &MlxArray, - capture_diagnostics: bool, + 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]; @@ -481,34 +554,43 @@ impl MolmoVisionModel { #[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_with_patch_embedding(&flat); + 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(&flat), None) + ( + 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(&flat), None) + ( + self.image_vit.forward_selected(&flat, &self.vit_layers), + None, + ) }; let selected_hidden_states = if capture_diagnostics { - self.vit_layers + hidden_states .iter() - .map(|&layer| mlxcel_core::copy(hidden_states[layer].as_ref().unwrap())) + .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). @@ -584,15 +666,16 @@ impl MolmoVisionModel { &self, images: &MlxArray, image_masks: &MlxArray, - capture_diagnostics: bool, + 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, patch_embeddings, selected_hidden_states) = - self.encode_image_internal(images, capture_diagnostics); + 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); @@ -653,6 +736,11 @@ impl MolmoVisionModel { // Project to the text hidden dim. 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, @@ -662,7 +750,7 @@ impl MolmoVisionModel { } pub fn forward(&self, images: &MlxArray, image_masks: &MlxArray) -> UniquePtr { - self.forward_internal(images, image_masks, false).0 + self.forward_internal(images, image_masks, None).0 } #[cfg(any(feature = "xla-diagnostics", feature = "xla-reference-diagnostics"))] @@ -671,8 +759,10 @@ impl MolmoVisionModel { &self, images: &MlxArray, image_masks: &MlxArray, + diagnostic_progress: &mut dyn FnMut(MolmoHostVisionDiagnosticStage), ) -> (UniquePtr, MolmoHostVisionDiagnostics) { - let (projected, diagnostics) = self.forward_internal(images, image_masks, true); + let (projected, diagnostics) = + self.forward_internal(images, image_masks, Some(diagnostic_progress)); ( projected, diagnostics.expect("Molmo host diagnostics were requested"), From 2c5a5a1ddea3c52d55bf0ee403f304b5a3753c4b Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 00:48:39 +0900 Subject: [PATCH 12/16] test(xla): add CUDA Molmo boundary runner Expose the pinned Molmo reference gate as a dedicated example so the eager oracle can run on MLX CUDA while IREE remains on local-task or local-sync. Emit flushed stage progress and periodic heartbeats without changing the existing comparison thresholds. --- examples/xla_molmo_reference_check.rs | 24 ++++ src/multimodal/host_preprocessor.rs | 3 + src/multimodal/host_preprocessor_molmo.rs | 8 +- ..._preprocessor_molmo_diagnostics_support.rs | 7 ++ ...st_preprocessor_molmo_diagnostics_tests.rs | 112 ++++++++++++++++-- 5 files changed, 140 insertions(+), 14 deletions(-) create mode 100644 examples/xla_molmo_reference_check.rs 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/multimodal/host_preprocessor.rs b/src/multimodal/host_preprocessor.rs index f276d4a30..f529d870e 100644 --- a/src/multimodal/host_preprocessor.rs +++ b/src/multimodal/host_preprocessor.rs @@ -52,6 +52,9 @@ use export::{ #[cfg(feature = "xla-backend")] #[path = "host_preprocessor_molmo.rs"] mod molmo; +#[cfg(feature = "xla-reference-diagnostics")] +#[doc(hidden)] +pub use molmo::run_pinned_molmo_eager_mlx_iree_boundaries; #[cfg(feature = "xla-backend")] pub use molmo::MolmoHostPreprocessor; diff --git a/src/multimodal/host_preprocessor_molmo.rs b/src/multimodal/host_preprocessor_molmo.rs index f847bd767..509cbbd71 100644 --- a/src/multimodal/host_preprocessor_molmo.rs +++ b/src/multimodal/host_preprocessor_molmo.rs @@ -401,12 +401,14 @@ fn format_prompt(prompt: &str) -> String { } } -#[cfg(all( - test, - any(feature = "xla-diagnostics", feature = "xla-reference-diagnostics") +#[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 index 2fe8c1058..ec830339c 100644 --- a/src/multimodal/host_preprocessor_molmo_diagnostics_support.rs +++ b/src/multimodal/host_preprocessor_molmo_diagnostics_support.rs @@ -2,10 +2,12 @@ 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"; @@ -21,10 +23,12 @@ 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"; @@ -166,6 +170,7 @@ pub(super) fn digest_i32(values: &[i32]) -> String { 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); @@ -193,10 +198,12 @@ pub(super) fn pinned_revision(model: &Path) -> String { .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, diff --git a/src/multimodal/host_preprocessor_molmo_diagnostics_tests.rs b/src/multimodal/host_preprocessor_molmo_diagnostics_tests.rs index d9e63a626..5413a3672 100644 --- a/src/multimodal/host_preprocessor_molmo_diagnostics_tests.rs +++ b/src/multimodal/host_preprocessor_molmo_diagnostics_tests.rs @@ -14,8 +14,12 @@ 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"] @@ -80,19 +84,31 @@ fn strict_oracle_rejects_molmo2_token_scan_and_replacement() { /// Run from the repository root only with the pinned 5.3 GB checkpoint: /// /// ```text -/// IREE_DIST=/path/to/iree-dist \ -/// MLXCEL_MOLMO_FIXTURE=/path/to/Molmo-7B-D-0924-4bit \ -/// cargo test --features xla-reference-diagnostics --lib \ -/// pinned_molmo_eager_mlx_matches_iree_boundaries -- --ignored --nocapture +/// 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_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. -#[test] -#[ignore = "requires pinned Molmo checkpoint and IREE local-task"] -fn pinned_molmo_eager_mlx_matches_iree_boundaries() { - mlxcel_core::set_default_device(false); +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"), @@ -102,10 +118,13 @@ fn pinned_molmo_eager_mlx_matches_iree_boundaries() { .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_eq!( - device, "local-task", - "the #870 reference gate qualifies IREE local-task only" + 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( @@ -194,27 +213,35 @@ fn pinned_molmo_eager_mlx_matches_iree_boundaries() { 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 @@ -369,3 +396,66 @@ fn pinned_molmo_eager_mlx_matches_iree_boundaries() { 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(); +} From 6803b17490e430608b5d87a224696968d9d284ea Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 00:57:45 +0900 Subject: [PATCH 13/16] style(xla): order Molmo exports --- src/multimodal/host_preprocessor.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/multimodal/host_preprocessor.rs b/src/multimodal/host_preprocessor.rs index f529d870e..905b400ab 100644 --- a/src/multimodal/host_preprocessor.rs +++ b/src/multimodal/host_preprocessor.rs @@ -52,11 +52,11 @@ use export::{ #[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; -#[cfg(feature = "xla-backend")] -pub use molmo::MolmoHostPreprocessor; /// Vision implementation selected for OpenXLA multimodal preprocessing. #[derive(Debug, Clone, Copy, PartialEq, Eq)] From f1c0a18983f1beefa89fe0efac2922b8df56d247 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 09:24:24 +0900 Subject: [PATCH 14/16] fix(xla): mark Molmo artifact unqualified Keep the rebased Molmo runtime on the explicit legacy artifact path until the real-checkpoint boundary oracle completes.\n\nRefs #870 --- src/lib/mlxcel-xla/src/molmo_vision_runtime.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/lib/mlxcel-xla/src/molmo_vision_runtime.rs b/src/lib/mlxcel-xla/src/molmo_vision_runtime.rs index b927b6d73..ff45818bf 100644 --- a/src/lib/mlxcel-xla/src/molmo_vision_runtime.rs +++ b/src/lib/mlxcel-xla/src/molmo_vision_runtime.rs @@ -215,7 +215,7 @@ fn compile_and_load( 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( + let contract = AuxiliaryArtifactContract::new_legacy_unqualified( ENTRY_NAME, format!( "{};processor_sha256={};checkpoint_schema_sha256={}", From a4047b220441b14610ca75ee774fd4b67f833e29 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 13:08:54 +0900 Subject: [PATCH 15/16] fix(xla): restore feature-neutral Molmo rounding Rebasing the Molmo prepared-prefill contract onto the numeric-oracle foundation exposed that its F16 and BF16 merge rounding called through a weight-loader module compiled only for IREE or tests, which broke the default mlxcel-xla library build. Move the three scalar conversion helpers into a feature-neutral module while retaining the weight-loader re-exports and existing bit-contract tests. Preserve both diagnostic feature families during the rebase, and document the pinned revision override required by the locally complete checkpoint fixture. Validation: default mlxcel-xla Clippy, 15 focused Molmo tests, four F16 conversion tests, formatting, and diff checks. Refs #870 --- Cargo.toml | 4 - src/lib/mlxcel-xla/src/emitter/mod.rs | 2 +- src/lib/mlxcel-xla/src/float.rs | 82 +++++++++++++++++++ src/lib/mlxcel-xla/src/lib.rs | 1 + src/lib/mlxcel-xla/src/prepared_molmo.rs | 4 +- src/lib/mlxcel-xla/src/weights.rs | 78 +----------------- ...st_preprocessor_molmo_diagnostics_tests.rs | 1 + 7 files changed, 88 insertions(+), 84 deletions(-) create mode 100644 src/lib/mlxcel-xla/src/float.rs diff --git a/Cargo.toml b/Cargo.toml index ff6d190ef..63e20f08c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -78,10 +78,6 @@ xla-reference-diagnostics = ["xla-iree", "mlxcel-xla/diagnostics"] # 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. -# -# 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"] 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. diff --git a/src/lib/mlxcel-xla/src/emitter/mod.rs b/src/lib/mlxcel-xla/src/emitter/mod.rs index 8571a17da..f8fb9d365 100644 --- a/src/lib/mlxcel-xla/src/emitter/mod.rs +++ b/src/lib/mlxcel-xla/src/emitter/mod.rs @@ -57,10 +57,10 @@ mod gemma3n_schema; mod gemma3n_weights; mod model; mod moe; -pub(crate) mod numeric_ops; mod molmo_vision; mod molmo_vision_config; mod molmo_vision_ops; +pub(crate) mod numeric_ops; mod phi4_audio; mod qwen2_vl; mod rope; 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/lib.rs b/src/lib/mlxcel-xla/src/lib.rs index ae5b915ec..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; diff --git a/src/lib/mlxcel-xla/src/prepared_molmo.rs b/src/lib/mlxcel-xla/src/prepared_molmo.rs index b40cd1980..fedec4c64 100644 --- a/src/lib/mlxcel-xla/src/prepared_molmo.rs +++ b/src/lib/mlxcel-xla/src/prepared_molmo.rs @@ -69,8 +69,8 @@ pub enum MolmoEmbeddingDType { impl MolmoEmbeddingDType { fn round(self, value: f32) -> f32 { match self { - Self::Float16 => crate::weights::half_to_f32(crate::weights::f32_to_f16_bits(value)), - Self::BFloat16 => crate::weights::round_bf16_f32(value), + 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, } } diff --git a/src/lib/mlxcel-xla/src/weights.rs b/src/lib/mlxcel-xla/src/weights.rs index 556acdcb0..5778ad0e6 100644 --- a/src/lib/mlxcel-xla/src/weights.rs +++ b/src/lib/mlxcel-xla/src/weights.rs @@ -449,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 @@ -497,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 diff --git a/src/multimodal/host_preprocessor_molmo_diagnostics_tests.rs b/src/multimodal/host_preprocessor_molmo_diagnostics_tests.rs index 5413a3672..3e04befea 100644 --- a/src/multimodal/host_preprocessor_molmo_diagnostics_tests.rs +++ b/src/multimodal/host_preprocessor_molmo_diagnostics_tests.rs @@ -90,6 +90,7 @@ fn strict_oracle_rejects_molmo2_token_scan_and_replacement() { /// /// 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 \ From a954d214db8abac8aacafaccb62dd79c6d55019f Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 22:31:24 +0900 Subject: [PATCH 16/16] fix(xla): configure local-task threads before Molmo IREE creation The Molmo v1 reference gate created its production IREE preprocessor without first applying the diagnostics-only local-task thread configuration added in #945. `configure_diagnostic_local_task_threads` had no callers on this branch, while the Gemma3 runner already applies it, so the Molmo runner aborted before any comparison ran: xla_aux_create failed (status 13): iree/base/threading/thread_pthreads.c:159: INTERNAL; thread creation failed with 22 IREE parses its process-global flag registry when the first instance is created, so the call has to precede `load_iree` to have any effect. With this, the pinned 5.0 GB checkpoint gate completes end to end for the first time on GB10: processor stages exact, patch embedding, selected layer 14, projector, both prepared sparse-add merge stages, layer-0 KV, all-layer KV and prefill logits all pass, leaving 4 of 1,181,696 values outside tolerance at selected layer 21. Refs #870 --- .../host_preprocessor_molmo_diagnostics_tests.rs | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/src/multimodal/host_preprocessor_molmo_diagnostics_tests.rs b/src/multimodal/host_preprocessor_molmo_diagnostics_tests.rs index 3e04befea..f0ad28009 100644 --- a/src/multimodal/host_preprocessor_molmo_diagnostics_tests.rs +++ b/src/multimodal/host_preprocessor_molmo_diagnostics_tests.rs @@ -149,6 +149,15 @@ pub fn run_pinned_molmo_eager_mlx_iree_boundaries() { 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);