feat(xla): support Molmo2 indexed attention pooling

Cargo.toml

@@ -104,6 +104,10 @@ xla-diagnostics = ["cuda", "xla-iree", "mlxcel-xla/diagnostics"]
 # CPU-capable bounded MLX/IREE operator-oracle harness. CUDA-specific production
 # probes may add `cuda`, but the shared report/comparison layer does not require it.
 xla-micro-oracle = ["xla-iree", "mlxcel-xla/micro-oracle"]
+# CPU-only companion for bounded host-reference checks on machines where the
+# CUDA driver is unavailable. This deliberately does not replace the qualified
+# `xla-diagnostics` production-target gate.
+xla-diagnostics-cpu = ["xla-iree", "mlxcel-xla/diagnostics"]
 # Axis A "weight-load surgery" framework.
 #
 # On by default so production `mlxcel` and `mlxcel-server` binaries expose
@@ -324,6 +328,10 @@ required-features = ["test-utils"]
 name = "xla_prepared_prefill"
 required-features = ["xla-iree"]
 
+[[test]]
+name = "molmo2_xla_vision_parity"
+required-features = ["xla-iree"]
+
 # Reference-equivalence + throughput harness for the OpenXLA/IREE continuous
 # batching engine (#449 M3 Stage 2b). Needs real IREE execution, so it builds
 # only under `xla-iree` (which carries the runtime link recipe); default/CI

src/backend/xla.rs

@@ -109,8 +109,22 @@ impl ComputeBackend for XlaBackend {
         // session owns its own KV and samples greedily on-device.
         let image_preprocessor = crate::load_xla_image_preprocessor(model_path)
             .map_err(|error| anyhow::anyhow!("OpenXLA image preprocessor load failed: {error}"))?;
-        let session =
-            XlaInferenceSession::load(model_path, num_layers).map_err(|e| anyhow::anyhow!(e))?;
+        // Match the server path: reject a graph that could never admit one
+        // image before generation starts, instead of after the vision tower has
+        // already run on the first image request.
+        let context_capacity =
+            mlxcel_xla::context_capacity_from_env().map_err(|e| anyhow::anyhow!(e))?;
+        crate::ensure_xla_image_context_capacity(
+            model_path,
+            context_capacity,
+            std::env::var_os(mlxcel_xla::CONTEXT_CAPACITY_ENV).is_some(),
+        )?;
+        let session = XlaInferenceSession::load_with_context_capacity(
+            model_path,
+            num_layers,
+            context_capacity,
+        )
+        .map_err(|e| anyhow::anyhow!(e))?;
         Ok(Session::xla(session, image_preprocessor))
     }
 

src/lib.rs

@@ -80,8 +80,9 @@ pub use multimodal::host_preprocessor::LlavaHostReferenceCapture;
 #[cfg(feature = "xla-iree")]
 pub use multimodal::host_preprocessor::LlavaIreeHostPreprocessor;
 pub use multimodal::host_preprocessor::{
-    FakeHostMultimodalPreprocessor, HostMultimodalPreprocessor, HostPreprocessorError,
-    LlavaHostPreprocessor, XlaVisionBackend, load_xla_image_preprocessor,
+    CONTEXT_CAPACITY_ENV, FakeHostMultimodalPreprocessor, HostMultimodalPreprocessor,
+    HostPreprocessorError, LlavaHostPreprocessor, XlaVisionBackend,
+    ensure_xla_image_context_capacity, load_xla_image_preprocessor, xla_image_context_floor,
 };
 pub use multimodal::{
     falcon_ocr_prompt, internvl_prompt, kimi_vl_prompt, locateanything_prompt, minicpmo_prompt,
@@ -121,6 +122,11 @@ pub fn decode_image_payloads_with_limits(
 // Re-export split modules
 pub use loaded_model::LoadedModel;
 pub use loaded_model_capabilities::VlmRuntimeRef;
+#[cfg(any(feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))]
+pub use loading::{
+    Molmo2XlaVisionReference, Molmo2XlaVisionReferenceProjection, Molmo2XlaVisionReferenceStage,
+    load_molmo2_xla_vision_reference,
+};
 pub use loading::{
     context_window_from_config, load_model, load_model_with_adapter,
     load_model_with_tensor_parallel, load_qwen3_omni_speech, read_eos_token_ids,

src/lib/mlxcel-core/src/weights.rs

@@ -182,7 +182,7 @@ fn safetensors_dtype_itemsize(dtype: &str) -> Option<u64> {
 ///
 /// Returns `None` when the file does not exist or cannot be parsed; callers
 /// should fall back to the analytical estimate in that case.
-fn read_safetensors_header_bytes(path: &Path) -> Option<u64> {
+fn read_safetensors_header(path: &Path) -> Option<serde_json::Map<String, serde_json::Value>> {
     use std::io::Read;
 
     let mut f = std::fs::File::open(path).ok()?;
@@ -201,10 +201,13 @@ fn read_safetensors_header_bytes(path: &Path) -> Option<u64> {
     let mut header_bytes = vec![0u8; header_len as usize];
     f.read_exact(&mut header_bytes).ok()?;
     let header_json = serde_json::from_slice::<serde_json::Value>(&header_bytes).ok()?;
+    header_json.as_object().cloned()
+}
 
-    let obj = header_json.as_object()?;
+fn read_safetensors_header_bytes(path: &Path) -> Option<u64> {
+    let obj = read_safetensors_header(path)?;
     let mut total: u64 = 0;
-    for (key, meta) in obj {
+    for (key, meta) in &obj {
         // The special "__metadata__" key is not a tensor entry.
         if key == "__metadata__" {
             continue;
@@ -551,6 +554,25 @@ where
     F: FnMut(&str) -> bool,
 {
     let path = path.as_ref();
+    // A stale index can force the directory loader to inspect every local shard.
+    // Read the lightweight safetensors header first so a filtered sub-stack does
+    // not enter MLX's native loader for a shard that cannot contain any retained
+    // tensor. Besides avoiding unnecessary mmap state, this is important for
+    // repackaged quantized shards whose unrelated tensor dtypes/layouts may not be
+    // loadable by the current MLX build.
+    let retained_names = read_safetensors_header(path).map(|header| {
+        header
+            .into_iter()
+            .map(|(name, _)| name)
+            .filter(|name| name != "__metadata__" && keep(name))
+            .collect::<std::collections::HashSet<_>>()
+    });
+    if retained_names
+        .as_ref()
+        .is_some_and(|names| names.is_empty())
+    {
+        return Ok(HashMap::new());
+    }
     let path_str = path
         .to_str()
         .ok_or_else(|| format!("Non-UTF8 path: {}", path.display()))?;
@@ -560,7 +582,10 @@ where
     let mut weights = HashMap::with_capacity(len);
     for i in 0..len {
         let name = ffi::loaded_weights_name(&loaded, i);
-        if !keep(&name) {
+        let retain = retained_names
+            .as_ref()
+            .map_or_else(|| keep(&name), |names| names.contains(&name));
+        if !retain {
             continue;
         }
         let array = ffi::loaded_weights_take(loaded.pin_mut(), i);
@@ -901,6 +926,21 @@ mod tests {
         f.write_all(header_bytes).unwrap();
     }
 
+    #[test]
+    fn filtered_loader_skips_nonmatching_shard_before_native_load() {
+        let dir = tempfile::tempdir().unwrap();
+        let file = dir.path().join("model-00001-of-00002.safetensors");
+        write_safetensors_stub(&file, "U32", &[4, 4]);
+        let mut inspected = Vec::new();
+        let weights = load_safetensors_filtered(&file, |name| {
+            inspected.push(name.to_string());
+            name.starts_with("vision_tower.")
+        })
+        .expect("a shard with no retained header names is skipped before native loading");
+        assert!(weights.is_empty());
+        assert_eq!(inspected, ["test_tensor"]);
+    }
+
     #[test]
     fn test_parse_shard_index_with_total_size_present() {
         let dir = tempfile::tempdir().unwrap();

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

@@ -132,6 +132,8 @@ pub enum WeightScheme {
     /// ExaOne 3.x GPT-2-style names (`transformer.h.{i}...`, gated MLP `c_fc_0` /
     /// `c_fc_1` / `c_proj`, `out_proj` attention output).
     Exaone,
+    /// Molmo2's OLMo-style names under a VLM `language_model.*` namespace.
+    Molmo2,
 }
 
 /// MLX affine weight quantization (`config.json` `quantization`). The linear /
@@ -482,14 +484,17 @@ impl Config {
             serde_json::from_str(s).map_err(|e| format!("parse config.json: {e}"))?;
         let wrapper_model_type = root.get("model_type").and_then(serde_json::Value::as_str);
         let is_phi4mm = wrapper_model_type == Some("phi4mm");
-        let v = if matches!(wrapper_model_type, Some("llava" | "llava_next")) {
+        let is_molmo2 = wrapper_model_type == Some("molmo2");
+        let v = if matches!(wrapper_model_type, Some("llava" | "llava_next" | "molmo2")) {
             let mut text = root
                 .get("text_config")
                 .and_then(serde_json::Value::as_object)
                 .cloned()
                 .ok_or_else(|| {
-                    "LLaVA config.json missing object `text_config` for the XLA text graph"
-                        .to_string()
+                    format!(
+                        "{} config.json missing object `text_config` for the XLA text graph",
+                        wrapper_model_type.unwrap_or("VLM")
+                    )
                 })?;
             // mlx-community quantized VLMs commonly keep the affine scheme at
             // the wrapper level even though the tensors belong to the nested
@@ -509,10 +514,10 @@ impl Config {
         // ExaOne 3.x keeps GPT-2-style tensor names; every other supported family
         // uses the standard HF Llama layout. Loader-only (see [`WeightScheme`]), so
         // it never changes the emitted graph.
-        let weight_scheme = if model_type == Some("exaone") {
-            WeightScheme::Exaone
-        } else {
-            WeightScheme::Llama
+        let weight_scheme = match model_type {
+            Some("exaone") => WeightScheme::Exaone,
+            Some("molmo2_text") if is_molmo2 => WeightScheme::Molmo2,
+            _ => WeightScheme::Llama,
         };
 
         // Interleaved (GPT-J-style) RoPE reaches the supported families only through
@@ -870,6 +875,19 @@ impl Config {
                 tie_default = false;
                 rotary_dim = partial_rotary(1.0);
             }
+            Some("molmo2_text") if is_molmo2 => {
+                // Molmo2 reuses the shared dense text graph: plain pre-norm,
+                // fused QKV, per-head raw q/k RMSNorm, and fused SwiGLU. Its
+                // checkpoint names and gate/up half ordering are loader-only
+                // differences captured by WeightScheme::Molmo2.
+                fused_qkv = true;
+                fused_gate_up = true;
+                tie_default = false;
+                qk_norm = Some(QkNorm {
+                    per_head: true,
+                    one_plus: false,
+                });
+            }
             Some("stablelm") => {
                 // LayerNorm with bias, partial RoPE, optional q/k/v bias, untied.
                 layernorm = true;
@@ -1003,7 +1021,7 @@ impl Config {
                 return Err(format!(
                     "the OpenXLA emitter supports the dense architectures Llama, Qwen2, \
                      Qwen3, Gemma1/2/3, SmolLM3, OLMo2/3, Seed-OSS, MiMo, InternLM3, ExaOne, \
-                     Cohere, Cohere2, Phi3, Phi4MM, StableLM, StarCoder2, Granite, and MiniCPM, plus \
+                     Cohere, Cohere2, Phi3, Phi4MM, Molmo2, StableLM, StarCoder2, Granite, and MiniCPM, plus \
                      the Mixtral, Qwen2-MoE, Qwen3-MoE, and OLMoE mixture-of-experts \
                      architectures; config.json model_type = {other:?} (other MoE / MLA / \
                      novel-activation variants are follow-ups)"

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

@@ -57,6 +57,8 @@ mod gemma3n_schema;
 mod gemma3n_weights;
 mod model;
 mod moe;
+mod molmo2_config;
+mod molmo2_vision;
 pub(crate) mod numeric_ops;
 mod phi4_audio;
 mod qwen2_vl;
@@ -152,6 +154,12 @@ pub(crate) use model::{
     validate_prefill_embeddings_metadata,
 };
 #[allow(unused_imports)]
+pub(crate) use molmo2_config::{Molmo2VisionConfig, Molmo2VisionWeightSpec};
+#[allow(unused_imports)]
+pub(crate) use molmo2_vision::emit_molmo2_vision;
+#[cfg(feature = "diagnostics")]
+pub(crate) use molmo2_vision::emit_molmo2_vision_diagnostics;
+#[allow(unused_imports)]
 pub(crate) use vision::emit_vision;
 // The only consumer is `IreeVisionDiagnosticProjector` in `vision_runtime`, which
 // is gated on `diagnostics` alone. The `test` arm is kept deliberately: it is the