feat(xla): support Molmo sparse additive visual embeddings

Cargo.toml

@@ -71,11 +71,14 @@ xla-backend = ["dep:mlxcel-xla"]
 # StableHLO prefill / decode_step graphs. Needs IREE_DIST at build time (the
 # extracted iree dist), so it is a local / opt-in build, not a CI default.
 xla-iree = ["xla-backend", "mlxcel-xla/iree"]
+# Device-neutral eager MLX/IREE reference capture for CPU `local-task`
+# qualification without requiring an IREE CUDA source/build tree.
+xla-reference-diagnostics = ["xla-iree", "mlxcel-xla/diagnostics"]
 # Explicit test-only forwarding for Gemma3n intermediate oracle validation.
 # Its pinned reference is the production CUDA runtime, so enabling diagnostics
 # also enables the root MLX CUDA backend. This is not part of `xla-iree`; normal
 # production bundles remain unchanged.
-xla-diagnostics = ["cuda", "xla-iree", "mlxcel-xla/diagnostics"]
+xla-diagnostics = ["cuda", "xla-reference-diagnostics"]
 # CPU-capable bounded MLX/IREE operator-oracle harness. CUDA-specific production
 # probes may add `cuda`, but the shared report/comparison layer does not require it.
 xla-micro-oracle = ["xla-iree", "mlxcel-xla/micro-oracle"]

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();
+}

src/lib/mlxcel-xla/src/emitter/config.rs

@@ -132,6 +132,9 @@ pub enum WeightScheme {
     /// ExaOne 3.x GPT-2-style names (`transformer.h.{i}...`, gated MLP `c_fc_0` /
     /// `c_fc_1` / `c_proj`, `out_proj` attention output).
     Exaone,
+    /// Molmo v1's OLMo-style decoder layout under `language_model.model`, with
+    /// fused `att_proj` QKV and fused `ff_proj` gate/up projections.
+    Molmo,
 }
 
 /// MLX affine weight quantization (`config.json` `quantization`). The linear /
@@ -509,10 +512,10 @@ impl Config {
         // ExaOne 3.x keeps GPT-2-style tensor names; every other supported family
         // uses the standard HF Llama layout. Loader-only (see [`WeightScheme`]), so
         // it never changes the emitted graph.
-        let weight_scheme = if model_type == Some("exaone") {
-            WeightScheme::Exaone
-        } else {
-            WeightScheme::Llama
+        let weight_scheme = match model_type {
+            Some("exaone") => WeightScheme::Exaone,
+            Some("molmo") => WeightScheme::Molmo,
+            _ => WeightScheme::Llama,
         };
 
         // Interleaved (GPT-J-style) RoPE reaches the supported families only through
@@ -870,6 +873,35 @@ impl Config {
                 tie_default = false;
                 rotary_dim = partial_rotary(1.0);
             }
+            Some("molmo") => {
+                if !matches!(
+                    v.get("rope_impl").and_then(serde_json::Value::as_str),
+                    None | Some("interleave")
+                ) {
+                    return Err("Molmo v1 XLA requires rope_impl = \"interleave\"".to_string());
+                }
+                if v.get("qkv_bias").and_then(serde_json::Value::as_bool) == Some(false) {
+                    return Err("Molmo v1 XLA requires qkv_bias = true".to_string());
+                }
+                if !matches!(
+                    v.get("layer_norm_type").and_then(serde_json::Value::as_str),
+                    None | Some("rms")
+                ) {
+                    return Err("Molmo v1 XLA requires layer_norm_type = \"rms\"".to_string());
+                }
+                if v.get("tie_word_embeddings")
+                    .and_then(serde_json::Value::as_bool)
+                    == Some(true)
+                    || v.get("weight_tying").and_then(serde_json::Value::as_bool) == Some(true)
+                {
+                    return Err("Molmo v1 XLA requires an untied output head".to_string());
+                }
+                qkv_bias = true;
+                tie_default = false;
+                fused_qkv = true;
+                fused_gate_up = true;
+                rope_interleaved = true;
+            }
             Some("stablelm") => {
                 // LayerNorm with bias, partial RoPE, optional q/k/v bias, untied.
                 layernorm = true;
@@ -1003,7 +1035,7 @@ impl Config {
                 return Err(format!(
                     "the OpenXLA emitter supports the dense architectures Llama, Qwen2, \
                      Qwen3, Gemma1/2/3, SmolLM3, OLMo2/3, Seed-OSS, MiMo, InternLM3, ExaOne, \
-                     Cohere, Cohere2, Phi3, Phi4MM, StableLM, StarCoder2, Granite, and MiniCPM, plus \
+                     Cohere, Cohere2, Phi3, Phi4MM, Molmo, StableLM, StarCoder2, Granite, and MiniCPM, plus \
                      the Mixtral, Qwen2-MoE, Qwen3-MoE, and OLMoE mixture-of-experts \
                      architectures; config.json model_type = {other:?} (other MoE / MLA / \
                      novel-activation variants are follow-ups)"
@@ -1343,10 +1375,22 @@ impl Config {
             None
         };
 
+        let intermediate = u("intermediate_size")?;
+        let intermediate = if model_type == Some("molmo") {
+            if !intermediate.is_multiple_of(2) {
+                return Err(format!(
+                    "Molmo fused intermediate_size={intermediate} must be even"
+                ));
+            }
+            intermediate / 2
+        } else {
+            intermediate
+        };
+
         Ok(Config {
             context_capacity: crate::DEFAULT_CONTEXT_CAPACITY,
             hidden,
-            inter: u("intermediate_size")?,
+            inter: intermediate,
             n_layers,
             n_q,
             n_kv: u("num_key_value_heads")?,

src/lib/mlxcel-xla/src/emitter/mod.rs

@@ -57,6 +57,9 @@ mod gemma3n_schema;
 mod gemma3n_weights;
 mod model;
 mod moe;
+mod molmo_vision;
+mod molmo_vision_config;
+mod molmo_vision_ops;
 pub(crate) mod numeric_ops;
 mod phi4_audio;
 mod qwen2_vl;
@@ -152,6 +155,14 @@ pub(crate) use model::{
     validate_prefill_embeddings_metadata,
 };
 #[allow(unused_imports)]
+pub(crate) use molmo_vision::emit_molmo_vision;
+#[cfg(feature = "diagnostics")]
+pub(crate) use molmo_vision::emit_molmo_vision_diagnostics;
+#[allow(unused_imports)]
+pub(crate) use molmo_vision_config::{
+    MolmoVisionConfig, MolmoVisionWeightDType, MolmoVisionWeightSpec,
+};
+#[allow(unused_imports)]
 pub(crate) use vision::emit_vision;
 #[cfg(any(test, feature = "diagnostics"))]
 pub(crate) use vision::emit_vision_diagnostics;