diff --git a/Cargo.toml b/Cargo.toml index ab506bce0..e887e38b0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -71,6 +71,10 @@ 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 reference capture for native IREE diagnostics. This keeps the +# eager MLX side on its default CPU backend and qualifies IREE `local-task` +# 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 @@ -344,3 +348,9 @@ required-features = ["xla-micro-oracle"] [[example]] name = "xla_phi4_audio_check" required-features = ["xla-iree"] + +# Pinned #869 eager MLX CUDA versus IREE local-task Gemma3 boundary gate. +# This is a standalone executable so running it does not link libtest. +[[example]] +name = "xla_gemma3_reference_check" +required-features = ["cuda", "xla-reference-diagnostics"] diff --git a/examples/xla_gemma3_reference_check.rs b/examples/xla_gemma3_reference_check.rs new file mode 100644 index 000000000..5367c70fc --- /dev/null +++ b/examples/xla_gemma3_reference_check.rs @@ -0,0 +1,58 @@ +// 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. + +//! Standalone pinned Gemma3 eager-MLX-CUDA versus IREE-local-task boundary gate. +//! +//! This executable intentionally avoids the Rust libtest harness. It accepts +//! command-line arguments first and falls back to the historical environment +//! variables used by the ignored test. +//! +//! ```text +//! IREE_DIST=/path/to/iree-dist \ +//! cargo run --example xla_gemma3_reference_check \ +//! --features cuda,xla-reference-diagnostics -- \ +//! --model /path/to/gemma-3-4b-it-4bit \ +//! --image tests/fixtures/test_image.png \ +//! --device local-task +//! ``` + +use std::path::PathBuf; + +fn argument(flag: &str) -> Option { + let args = std::env::args().collect::>(); + args.iter() + .position(|argument| argument == flag) + .and_then(|index| args.get(index + 1)) + .cloned() +} + +fn required_path(flag: &str, variable: &str) -> PathBuf { + argument(flag) + .or_else(|| std::env::var(variable).ok()) + .map(PathBuf::from) + .unwrap_or_else(|| panic!("missing required {flag} or {variable}")) +} + +fn main() { + let model = required_path("--model", "MLXCEL_GEMMA3_FIXTURE"); + let image = argument("--image") + .or_else(|| std::env::var("MLXCEL_GEMMA3_IMAGE").ok()) + .map(PathBuf::from) + .unwrap_or_else(|| PathBuf::from("tests/fixtures/test_image.png")); + let device = argument("--device") + .or_else(|| std::env::var("MLXCEL_XLA_DEVICE").ok()) + .unwrap_or_else(|| "local-task".to_string()); + + mlxcel::run_gemma3_eager_mlx_iree_prepared_boundary(&model, &image, &device); +} diff --git a/spike/openxla/gemma3_vlm_mask_check.py b/spike/openxla/gemma3_vlm_mask_check.py new file mode 100644 index 000000000..f58c3403d --- /dev/null +++ b/spike/openxla/gemma3_vlm_mask_check.py @@ -0,0 +1,432 @@ +#!/usr/bin/env python3 +"""Independent Gemma3 VLM embeddings/mask execution check for issue #869. + +The production VLM path hands `prefill_embeddings.main` post-Gemma-scale text +rows, unscaled projected image rows, and one authoritative additive f32 mask. +This fixture compares that graph with an independent Hugging Face Gemma3 eager +model using the same tiny deterministic checkpoint. The HF side receives the +corresponding pre-scale rows because Gemma3 applies `sqrt(hidden_size)` inside +its model forward. + +The canonical case compares last-token logits, every layer's K/V cache, and the +greedy token. Four negative fixtures must then diverge: accidental causal +masking, a multiplicative 0/1 mask, double Gemma scaling, and a missing image +pre-divide. The tiny CPU graph keeps the check suitable for a short, bounded +local-task gate; pinned-checkpoint vision and server gates remain separate. + +Run from the repository root with the shared OpenXLA spike environment: + + spike/openxla/.venv/bin/python spike/openxla/gemma3_vlm_mask_check.py + +Exit 0 means the independent canonical oracle matched and every negative +fixture was detected. +""" + +from __future__ import annotations + +import json +import os +import subprocess +import sys +import tempfile +from pathlib import Path + +import numpy as np +import torch +from iree.compiler.tools import compile_file +from iree.runtime import load_vm_flatbuffer_file +from transformers import Gemma3ForCausalLM, Gemma3TextConfig + +REPO_ROOT = Path(__file__).resolve().parents[2] +CARGO = os.environ.get("CARGO", "cargo") +PREFILL_LP = 256 +REAL_LEN = 12 +HIDDEN = 16 +INTERMEDIATE = 32 +N_LAYERS = 4 +N_Q = 4 +N_KV = 2 +HEAD_DIM = 4 +VOCAB = 64 +EPS = 1e-6 +ROPE_THETA = 10_000.0 +ROPE_LOCAL_BASE = 100.0 +SLIDING_WINDOW = 4 +SLIDING_PATTERN = 3 +MASKED = np.float32(np.finfo(np.float32).min) +LOGIT_ATOL = 2e-2 +KV_MAX_ABS = 2.5e-1 +KV_RMS = 1e-1 +KV_MIN_COSINE = 9.9e-1 +NEGATIVE_MIN_DIFF = 1e-4 +IMAGE_POSITIONS = (3, 4) + + +def dimensions() -> dict[str, object]: + return { + "hidden_size": HIDDEN, + "num_attention_heads": N_Q, + "num_key_value_heads": N_KV, + "head_dim": HEAD_DIM, + "intermediate_size": INTERMEDIATE, + "num_hidden_layers": N_LAYERS, + "vocab_size": VOCAB, + "rms_norm_eps": EPS, + "rope_theta": ROPE_THETA, + "max_position_embeddings": 512, + "attention_bias": False, + } + + +def hf_config() -> Gemma3TextConfig: + return Gemma3TextConfig( + **dimensions(), + tie_word_embeddings=True, + query_pre_attn_scalar=HEAD_DIM, + rope_local_base_freq=ROPE_LOCAL_BASE, + sliding_window=SLIDING_WINDOW, + sliding_window_pattern=SLIDING_PATTERN, + attn_logit_softcapping=None, + final_logit_softcapping=None, + hidden_activation="gelu_pytorch_tanh", + ) + + +def emitter_config() -> dict[str, object]: + config = dimensions() + config.pop("max_position_embeddings") + config.update( + model_type="gemma3_text", + tie_word_embeddings=True, + query_pre_attn_scalar=HEAD_DIM, + rope_local_base_freq=ROPE_LOCAL_BASE, + sliding_window=SLIDING_WINDOW, + sliding_window_pattern=SLIDING_PATTERN, + attn_logit_softcapping=None, + final_logit_softcapping=None, + hidden_activation="gelu_pytorch_tanh", + ) + return config + + +def argument_names() -> list[str]: + names = ["model.embed_tokens.weight", "model.norm.weight"] + for index in range(N_LAYERS): + prefix = f"model.layers.{index}." + names.extend( + [ + prefix + "mlp.down_proj.weight", + prefix + "mlp.gate_proj.weight", + prefix + "input_layernorm.weight", + prefix + "post_attention_layernorm.weight", + prefix + "mlp.up_proj.weight", + prefix + "self_attn.k_proj.weight", + prefix + "self_attn.o_proj.weight", + prefix + "self_attn.q_proj.weight", + prefix + "self_attn.v_proj.weight", + prefix + "self_attn.q_norm.weight", + prefix + "self_attn.k_norm.weight", + prefix + "pre_feedforward_layernorm.weight", + prefix + "post_feedforward_layernorm.weight", + ] + ) + return names + + +def build_checkpoint() -> tuple[Gemma3ForCausalLM, list[np.ndarray]]: + torch.manual_seed(869) + model = Gemma3ForCausalLM(hf_config()).eval().float() + model.config._attn_implementation = "eager" + with torch.no_grad(): + for _, parameter in model.named_parameters(): + if parameter.dim() == 1: + parameter.copy_(torch.randn_like(parameter) * 0.1) + state = model.state_dict() + names = argument_names() + missing = [name for name in names if name not in state] + if missing: + raise RuntimeError(f"HF checkpoint is missing emitter weights: {missing[:4]}") + weights = [ + np.ascontiguousarray(state[name].detach().numpy(), dtype=np.float32) + for name in names + ] + return model, weights + + +def emit_and_compile() -> object: + work = Path(tempfile.mkdtemp(prefix="gemma3_vlm_mask_")) + config_path = work / "config.json" + config_path.write_text(json.dumps(emitter_config()), encoding="utf-8") + + print("[emit] Gemma3 embeddings prefill StableHLO", flush=True) + subprocess.run( + [ + CARGO, + "test", + "-p", + "mlxcel-xla", + "--lib", + "emitter::tests::dump_prefill_embeddings_parity_graphs", + "--", + "--ignored", + "--nocapture", + ], + cwd=REPO_ROOT, + env={ + **os.environ, + "MLXCEL_DUMP_CONFIG": str(config_path), + "MLXCEL_DUMP_DIR": str(work), + }, + check=True, + ) + + source = work / "prefill_embeddings_logits.mlir" + output = work / "prefill_embeddings_logits.vmfb" + print("[compile] Gemma3 embeddings prefill (llvm-cpu)", flush=True) + compile_file( + str(source), + output_file=str(output), + input_type="stablehlo", + target_backends=["llvm-cpu"], + ) + return load_vm_flatbuffer_file(str(output), driver="local-task") + + +def canonical_inputs( + embedding_table: np.ndarray, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + rng = np.random.default_rng(869) + tokens = rng.integers(1, VOCAB, size=REAL_LEN, dtype=np.int32) + pre_scale = np.ascontiguousarray(embedding_table[tokens], dtype=np.float32) + projected = np.ascontiguousarray( + rng.normal(0.0, 0.2, (len(IMAGE_POSITIONS), HIDDEN)), + dtype=np.float32, + ) + normalizer = np.float32(np.sqrt(HIDDEN)) + for row, position in enumerate(IMAGE_POSITIONS): + pre_scale[position] = projected[row] / normalizer + + post_scale = np.zeros((PREFILL_LP, HIDDEN), dtype=np.float32) + post_scale[:REAL_LEN] = pre_scale * normalizer + + mask = np.full((PREFILL_LP, PREFILL_LP), MASKED, dtype=np.float32) + mask[:REAL_LEN, :REAL_LEN] = 0.0 + return pre_scale, np.ascontiguousarray(post_scale), np.ascontiguousarray(mask) + + +def to_host(value: object) -> np.ndarray: + host = value.to_host() if hasattr(value, "to_host") else value + return np.asarray(host, dtype=np.float32) + + +def run_iree( + module: object, + weights: list[np.ndarray], + embeddings: np.ndarray, + mask: np.ndarray, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + positions = np.arange(PREFILL_LP, dtype=np.int32) + real_len = np.asarray(REAL_LEN, dtype=np.int32) + outputs = module.main(*weights, embeddings, positions, real_len, mask) + return tuple(to_host(value) for value in outputs) + + +def legacy_cache(cache: object) -> tuple[tuple[torch.Tensor, torch.Tensor], ...]: + if hasattr(cache, "to_legacy_cache"): + return cache.to_legacy_cache() + if isinstance(cache, (tuple, list)): + return tuple(cache) + layers = getattr(cache, "layers", None) + if layers is None: + raise RuntimeError(f"unsupported HF cache type: {type(cache).__name__}") + return tuple((layer.keys, layer.values) for layer in layers) + + +def run_hf( + model: Gemma3ForCausalLM, + pre_scale: np.ndarray, + mask: np.ndarray, +) -> tuple[np.ndarray, list[np.ndarray], list[np.ndarray]]: + real_mask = np.ascontiguousarray(mask[:REAL_LEN, :REAL_LEN]) + with torch.no_grad(): + outputs = model( + inputs_embeds=torch.from_numpy(pre_scale[None, ...]), + attention_mask=torch.from_numpy(real_mask[None, None, ...]), + position_ids=torch.arange(REAL_LEN, dtype=torch.long)[None, :], + use_cache=True, + return_dict=True, + ) + logits = outputs.logits[0, REAL_LEN - 1].detach().numpy().astype(np.float32) + cache = legacy_cache(outputs.past_key_values) + keys = [ + layer[0][0].detach().numpy().transpose(1, 0, 2).astype(np.float32) + for layer in cache + ] + values = [ + layer[1][0].detach().numpy().transpose(1, 0, 2).astype(np.float32) + for layer in cache + ] + return logits, keys, values + + +def compare_canonical( + iree: tuple[np.ndarray, np.ndarray, np.ndarray], + hf: tuple[np.ndarray, list[np.ndarray], list[np.ndarray]], +) -> bool: + logits_close = np.allclose(iree[0], hf[0], rtol=0.0, atol=LOGIT_ATOL) + logits_diff = float(np.max(np.abs(iree[0] - hf[0]))) + print( + f"[canonical/logits] shape={iree[0].shape}/{hf[0].shape} " + f"max|diff|={logits_diff:.3e} -> {'PASS' if logits_close else 'FAIL'}", + flush=True, + ) + ok = bool(logits_close) + for name, actual, expected_layers in ( + ("kcache", iree[1], hf[1]), + ("vcache", iree[2], hf[2]), + ): + for layer, expected in enumerate(expected_layers): + cache_len = expected.shape[0] + actual_slice = actual[layer, REAL_LEN - cache_len : REAL_LEN] + same_shape = actual_slice.shape == expected.shape + max_diff = ( + float(np.max(np.abs(actual_slice - expected))) + if same_shape + else float("inf") + ) + rms_diff = ( + float(np.sqrt(np.mean(np.square(actual_slice - expected)))) + if same_shape + else float("inf") + ) + denominator = ( + float(np.linalg.norm(actual_slice) * np.linalg.norm(expected)) + if same_shape + else 0.0 + ) + cosine = ( + float(np.vdot(actual_slice, expected) / denominator) + if denominator > 0.0 + else 1.0 + ) + close = ( + same_shape + and max_diff <= KV_MAX_ABS + and rms_diff <= KV_RMS + and cosine >= KV_MIN_COSINE + ) + ok = ok and close + best = "" + if not close and cache_len < REAL_LEN: + candidates = [ + ( + start, + float( + np.max( + np.abs( + actual[layer, start : start + cache_len] - expected + ) + ) + ), + ) + for start in range(REAL_LEN - cache_len + 1) + ] + best_start, best_diff = min(candidates, key=lambda item: item[1]) + best = f" best_slice={best_start}:{best_start + cache_len}/{best_diff:.3e}" + print( + f"[canonical/{name}/layer{layer}] " + f"shape={actual_slice.shape}/{expected.shape} " + f"max|diff|={max_diff:.3e} rms={rms_diff:.3e} " + f"cos={cosine:.6f}{best} " + f"-> {'PASS' if close else 'FAIL'}", + flush=True, + ) + iree_token = int(np.argmax(iree[0])) + hf_token = int(np.argmax(hf[0])) + token_ok = iree_token == hf_token + print( + f"[canonical/token] iree={iree_token} hf={hf_token} " + f"-> {'PASS' if token_ok else 'FAIL'}", + flush=True, + ) + return ok and token_ok + + +def negative_detected( + label: str, + canonical_logits: np.ndarray, + negative_logits: np.ndarray, +) -> bool: + difference = float(np.max(np.abs(canonical_logits - negative_logits))) + detected = difference > NEGATIVE_MIN_DIFF + print( + f"[negative/{label}] max|logit diff|={difference:.3e} " + f"-> {'DETECTED' if detected else 'MISSED'}", + flush=True, + ) + return detected + + +def main() -> int: + model, weights = build_checkpoint() + module = emit_and_compile() + pre_scale, embeddings, mask = canonical_inputs(weights[0]) + + print("[run] canonical independent HF and IREE paths", flush=True) + canonical = run_iree(module, weights, embeddings, mask) + reference = run_hf(model, pre_scale, mask) + checks = [compare_canonical(canonical, reference)] + + causal = np.full_like(mask, MASKED) + valid = np.arange(REAL_LEN) + causal[valid[:, None], valid[None, :]] = np.where( + valid[None, :] <= valid[:, None], 0.0, MASKED + ) + checks.append( + negative_detected( + "causal-mask", + canonical[0], + run_iree(module, weights, embeddings, causal)[0], + ) + ) + + multiplicative = np.zeros_like(mask) + multiplicative[:REAL_LEN, :REAL_LEN] = 1.0 + checks.append( + negative_detected( + "multiplicative-mask", + canonical[0], + run_iree(module, weights, embeddings, multiplicative)[0], + ) + ) + + checks.append( + negative_detected( + "double-scale", + canonical[0], + run_iree(module, weights, embeddings * np.float32(np.sqrt(HIDDEN)), mask)[0], + ) + ) + + missing_image_prescale = embeddings.copy() + missing_image_prescale[list(IMAGE_POSITIONS)] *= np.float32(np.sqrt(HIDDEN)) + checks.append( + negative_detected( + "missing-image-prescale", + canonical[0], + run_iree(module, weights, missing_image_prescale, mask)[0], + ) + ) + + ok = all(checks) + print( + f"RESULT: {'PASS' if ok else 'FAIL'} " + f"(logit_atol={LOGIT_ATOL:g}, kv_max={KV_MAX_ABS:g}, " + f"kv_rms={KV_RMS:g}, kv_cos={KV_MIN_COSINE:g}, local-task)", + flush=True, + ) + return 0 if ok else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/lib.rs b/src/lib.rs index d9520a12f..daec017d6 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -75,6 +75,8 @@ pub use mlxcel_core::generate::{ SamplingConfig, }; pub use mlxcel_core::speculative::SpeculativeGenerator; +#[cfg(feature = "xla-iree")] +pub use multimodal::host_preprocessor::Gemma3IreeHostPreprocessor; #[cfg(feature = "xla-diagnostics")] pub use multimodal::host_preprocessor::LlavaHostReferenceCapture; #[cfg(feature = "xla-iree")] @@ -121,6 +123,8 @@ pub fn decode_image_payloads_with_limits( // Re-export split modules pub use loaded_model::LoadedModel; pub use loaded_model_capabilities::VlmRuntimeRef; +#[cfg(feature = "xla-reference-diagnostics")] +pub use loading::run_gemma3_eager_mlx_iree_prepared_boundary; pub use loading::{ context_window_from_config, load_model, load_model_with_adapter, load_model_with_tensor_parallel, load_qwen3_omni_speech, read_eos_token_ids, diff --git a/src/lib/mlxcel-xla/src/emitter/config.rs b/src/lib/mlxcel-xla/src/emitter/config.rs index e5c450eaa..08dbd8878 100644 --- a/src/lib/mlxcel-xla/src/emitter/config.rs +++ b/src/lib/mlxcel-xla/src/emitter/config.rs @@ -313,6 +313,11 @@ pub struct Config { /// uses 4; Cohere2 uses `sliding_window_pattern` (4). Only meaningful when /// `sliding_window` is `Some`. pub sliding_pattern: usize, + /// The embeddings-prefill caller supplies the complete per-layer attention + /// policy. Gemma3 VLM uses one bidirectional padding mask for every layer, + /// so the emitter must not intersect it with the text model's ordinary + /// sliding-window schedule. Token prefill remains causal/windowed. + pub embeddings_prefill_uses_authoritative_mask: bool, /// Per-layer NoPE mask (SmolLM3): `use_rope_layers[li] == false` skips RoPE on /// that layer (`no_rope_layers`). `None` applies RoPE on every layer. pub use_rope_layers: Option>, @@ -439,6 +444,7 @@ impl Config { final_logit_softcap: None, sliding_window: None, sliding_pattern: 2, + embeddings_prefill_uses_authoritative_mask: false, use_rope_layers: None, mrope: None, deepstack: None, @@ -482,15 +488,44 @@ 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_gemma3_vlm_wrapper = wrapper_model_type == Some("gemma3") + && root + .get("vision_config") + .is_some_and(serde_json::Value::is_object); + let v = if matches!(wrapper_model_type, Some("llava" | "llava_next")) + || is_gemma3_vlm_wrapper + { 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!( + "{wrapper_model_type:?} config.json missing object `text_config` for the XLA text graph" + ) })?; + if is_gemma3_vlm_wrapper { + // mlx-vlm's Gemma3 TextConfig supplies these architecture + // defaults. mlx-community conversions commonly omit them from + // the nested object, so resolve the same explicit contract here + // instead of deriving head_dim from hidden_size (Gemma3's q + // projection width is intentionally smaller than hidden_size). + let defaults = [ + ("num_attention_heads", serde_json::json!(8)), + ("num_key_value_heads", serde_json::json!(4)), + ("head_dim", serde_json::json!(256)), + ("rms_norm_eps", serde_json::json!(1.0e-6)), + ("vocab_size", serde_json::json!(262_208)), + ("rope_theta", serde_json::json!(1_000_000.0)), + ("rope_local_base_freq", serde_json::json!(10_000.0)), + ("query_pre_attn_scalar", serde_json::json!(256.0)), + ("sliding_window", serde_json::json!(1024)), + ("sliding_window_pattern", serde_json::json!(6)), + ]; + for (name, value) in defaults { + text.entry(name.to_string()).or_insert(value); + } + } // mlx-community quantized VLMs commonly keep the affine scheme at // the wrapper level even though the tensors belong to the nested // language model. @@ -653,6 +688,7 @@ impl Config { let mut final_logit_softcap: Option = None; let mut sliding_window: Option = None; let mut sliding_pattern = 2usize; + let mut embeddings_prefill_uses_authoritative_mask = false; let mut use_rope_layers: Option> = None; // issue #498 dense arch pack flags. let mut layernorm = false; @@ -765,6 +801,7 @@ impl Config { }); sliding_pattern = ou("sliding_window_pattern").unwrap_or(6).max(1); sliding_window = Some(ou("sliding_window").unwrap_or(4096)); + embeddings_prefill_uses_authoritative_mask = true; rope_local_base = Some(of("rope_local_base_freq").unwrap_or(10000.0)); read_gemma_common( &mut query_pre_attn_scalar, @@ -1072,6 +1109,12 @@ impl Config { // type while carrying the M-RoPE section table; the table is // the explicit schema signal in that representation. Some("default") | Some("dynamic") | Some("mrope") => RopeScaling::Plain, + // The qualified Gemma3 MLX implementation and upstream + // mlx-vlm Gemma3 model select distinct local/global bases + // but do not apply the conversion metadata's legacy linear + // scaling block. Preserve that reference contract instead + // of silently applying a different HF-only position scale. + Some("linear") if model_type == Some("gemma3_text") => RopeScaling::Plain, None if has_mrope_section => RopeScaling::Plain, Some("llama3") => { let sf = |k: &str| -> Result { @@ -1369,6 +1412,7 @@ impl Config { final_logit_softcap, sliding_window, sliding_pattern, + embeddings_prefill_uses_authoritative_mask, use_rope_layers, mrope, deepstack, @@ -1545,6 +1589,59 @@ impl Config { mod tests { use super::*; + #[test] + fn gemma3_embeddings_prefill_owns_the_complete_external_mask() { + let config = Config::from_json_str( + r#"{"model_type":"gemma3_text","hidden_size":8,"num_attention_heads":2, + "num_key_value_heads":1,"head_dim":4,"intermediate_size":16, + "num_hidden_layers":4,"rms_norm_eps":1e-6,"rope_theta":1000000, + "rope_local_base_freq":10000,"sliding_window":2, + "sliding_window_pattern":3,"vocab_size":12, + "hidden_activation":"gelu_pytorch_tanh"}"#, + ) + .unwrap(); + assert!(config.embeddings_prefill_uses_authoritative_mask); + assert_eq!(config.sliding_window, Some(2)); + assert!(config.is_sliding_layer(0), "token prefill remains windowed"); + } + + #[test] + fn gemma3_wrapper_uses_nested_text_config_and_wrapper_quantization() { + let config = Config::from_json_str( + r#"{"model_type":"gemma3","quantization":{"bits":4,"group_size":64}, + "vision_config":{}, + "text_config":{"model_type":"gemma3_text","hidden_size":2560, + "intermediate_size":10240,"num_hidden_layers":34,"sliding_window":1024, + "rope_scaling":{"factor":8.0,"rope_type":"linear"}}}"#, + ) + .unwrap(); + assert!(config.embeddings_prefill_uses_authoritative_mask); + assert_eq!(config.quantization.unwrap().bits, 4); + assert_eq!(config.hidden, 2560); + assert_eq!(config.n_q, 8); + assert_eq!(config.n_kv, 4); + assert_eq!(config.head_dim, 256); + assert_eq!(config.rope, RopeScaling::Plain); + } + + #[test] + fn gemma3_root_text_config_remains_direct_without_vision() { + let config = Config::from_json_str( + r#"{"model_type":"gemma3","hidden_size":8,"num_attention_heads":2, + "num_key_value_heads":1,"head_dim":4,"intermediate_size":16, + "num_hidden_layers":4,"rms_norm_eps":1e-6,"rope_theta":1000000, + "rope_local_base_freq":10000,"sliding_window":2, + "sliding_window_pattern":3,"vocab_size":12, + "hidden_activation":"gelu_pytorch_tanh"}"#, + ) + .unwrap(); + assert_eq!(config.hidden, 8); + assert_eq!(config.n_q, 2); + assert_eq!(config.n_kv, 1); + assert_eq!(config.head_dim, 4); + assert!(config.embeddings_prefill_uses_authoritative_mask); + } + /// ERNIE-4.5 is rejected with a message naming its interleaved (GPT-J-style) /// RoPE: it looks like a plain-RoPE Llama in config.json but its `rotate_half` /// rotates the (2i, 2i+1) pairs, not the (i, i+d/2) halves the Llama emit uses, diff --git a/src/lib/mlxcel-xla/src/emitter/mod.rs b/src/lib/mlxcel-xla/src/emitter/mod.rs index 0065385c7..07e5bdd16 100644 --- a/src/lib/mlxcel-xla/src/emitter/mod.rs +++ b/src/lib/mlxcel-xla/src/emitter/mod.rs @@ -156,7 +156,9 @@ pub(crate) use vision::emit_vision; #[cfg(any(test, feature = "diagnostics"))] pub(crate) use vision::emit_vision_diagnostics; #[allow(unused_imports)] -pub(crate) use vision_config::{LlavaVisionConfig, VisionActivation, VisionWeightSpec}; +pub(crate) use vision_config::{ + LlavaVisionConfig, VisionActivation, VisionProjector, VisionWeightSpec, +}; #[cfg(test)] mod tests { diff --git a/src/lib/mlxcel-xla/src/emitter/model.rs b/src/lib/mlxcel-xla/src/emitter/model.rs index 16a748cc6..b3e17a5a1 100644 --- a/src/lib/mlxcel-xla/src/emitter/model.rs +++ b/src/lib/mlxcel-xla/src/emitter/model.rs @@ -157,9 +157,9 @@ pub(crate) fn validate_prefill_embeddings_metadata( Ok(()) } -/// Validate a canonical additive attention-bias payload. Only `0.0` (allowed) -/// and [`PREFILL_EMBEDDINGS_MASKED_VALUE`] (masked) are accepted; NaN, infinity, -/// and intermediate additive values are rejected to keep polarity unambiguous. +/// Validate a canonical additive attention-bias payload. Generic families use +/// [`PREFILL_EMBEDDINGS_MASKED_VALUE`]; Gemma3's authoritative VLM mask uses +/// `f32::MIN`. NaN, infinity, and intermediate values are rejected. pub(crate) fn validate_prefill_embeddings_attention_bias( c: &Config, bias: &[f32], @@ -179,15 +179,20 @@ pub(crate) fn validate_prefill_embeddings_attention_bias( bias.len() )); } + let masked_value = if c.embeddings_prefill_uses_authoritative_mask { + f32::MIN + } else { + PREFILL_EMBEDDINGS_MASKED_VALUE + }; if let Some((index, value)) = bias .iter() .copied() .enumerate() - .find(|(_, value)| *value != 0.0 && *value != PREFILL_EMBEDDINGS_MASKED_VALUE) + .find(|(_, value)| *value != 0.0 && *value != masked_value) { return Err(format!( "prefill attention bias at flat index {index} is {value}; expected only 0 or {}", - PREFILL_EMBEDDINGS_MASKED_VALUE + masked_value )); } Ok(()) @@ -3304,6 +3309,14 @@ fn apply_deepstack_after_layer( b.add(hidden, &delta) } +fn embeddings_prefill_local_window(c: &Config) -> Option { + if c.embeddings_prefill_uses_authoritative_mask { + None + } else { + c.sliding_window + } +} + fn emit_prefill_module( c: &Config, sample: bool, @@ -3387,7 +3400,7 @@ fn emit_prefill_module( } PrefillInput::Embeddings { attention_bias, .. } => { let cmask = attention_bias.clone(); - let cmask_local = c.sliding_window.map(|w| { + let cmask_local = embeddings_prefill_local_window(c).map(|w| { let irow = b.iota(lp); let row = b.broadcast(&irow, &[0], vec![lp, lp]); let jcol = b.iota(lp); @@ -3403,7 +3416,7 @@ fn emit_prefill_module( } PrefillInput::DeepStack(deepstack) => { let cmask = deepstack.attention_bias.clone(); - let cmask_local = c.sliding_window.map(|w| { + let cmask_local = embeddings_prefill_local_window(c).map(|w| { let irow = b.iota(lp); let row = b.broadcast(&irow, &[0], vec![lp, lp]); let jcol = b.iota(lp); @@ -3517,6 +3530,24 @@ fn emit_prefill_module( } } +#[cfg(test)] +mod gemma3_vlm_mask_tests { + use super::*; + + #[test] + fn authoritative_embeddings_mask_bypasses_only_the_external_local_intersection() { + let mut config = Config::llama_3_2_1b(); + config.sliding_window = Some(4096); + assert_eq!(embeddings_prefill_local_window(&config), Some(4096)); + config.embeddings_prefill_uses_authoritative_mask = true; + assert_eq!(embeddings_prefill_local_window(&config), None); + assert!( + config.is_sliding_layer(0), + "the token-prefill/decode layer schedule must remain windowed" + ); + } +} + // =========================================================================== // MoE FFN block probe (issue #500): a standalone module for the execution check // =========================================================================== diff --git a/src/lib/mlxcel-xla/src/emitter/vision.rs b/src/lib/mlxcel-xla/src/emitter/vision.rs index 43f7f713e..b936596fb 100644 --- a/src/lib/mlxcel-xla/src/emitter/vision.rs +++ b/src/lib/mlxcel-xla/src/emitter/vision.rs @@ -23,7 +23,9 @@ use super::builder::{Builder, Ty, Val}; use super::numeric_ops::{exact_gelu, layer_norm_2d as layer_norm, stable_softmax, tanh_gelu}; -use super::vision_config::{LlavaVisionConfig, VisionActivation, VisionWeightSpec}; +use super::vision_config::{ + LlavaVisionConfig, VisionActivation, VisionProjector, VisionWeightSpec, +}; struct Args { values: Vec, @@ -58,6 +60,36 @@ mod tests { assert!(mlir.contains("return ")); } + #[test] + fn gemma3_graph_avg_pools_norms_and_projects_without_llava_mlp() { + let config = LlavaVisionConfig::from_json_str( + &serde_json::json!({ + "model_type": "gemma3", + "image_token_index": 99, + "mm_tokens_per_image": 1, + "vision_config": { + "model_type": "siglip_vision_model", + "image_size": 28, + "patch_size": 14, + "hidden_size": 8, + "intermediate_size": 16, + "num_hidden_layers": 1, + "num_attention_heads": 2 + }, + "text_config": {"hidden_size": 12} + }) + .to_string(), + ) + .unwrap(); + let mlir = emit_vision(&config); + assert!(mlir.contains("vision_tower.vision_model.post_layernorm.weight")); + assert!(mlir.contains("vision_tower.vision_model.post_layernorm.bias")); + assert!(mlir.contains("multi_modal_projector.mm_soft_emb_norm.weight")); + assert!(mlir.contains("multi_modal_projector.mm_input_projection_weight")); + assert!(!mlir.contains("multi_modal_projector.linear_1")); + assert!(mlir.contains("-> tensor<1x12xf32>")); + } + #[cfg(feature = "iree")] #[test] fn pinned_siglip_graph_compiles_for_cpu() { @@ -129,6 +161,73 @@ fn linear_2d(builder: &mut Builder, value: &Val, weight: &Val, bias: &Val) -> Va bias_2d(builder, &value, bias) } +fn gemma_rms_norm(builder: &mut Builder, value: &Val, weight: &Val, epsilon: f32) -> Val { + let rows = value.ty.shape[0]; + let width = value.ty.shape[1]; + let zero = builder.const_f32(0.0); + let squared = builder.multiply(value, value); + let squared_sum = builder.reduce_add(&squared, 1, &zero); + let width_scalar = builder.const_f32(width as f32); + let width_rows = builder.broadcast(&width_scalar, &[], vec![rows]); + let mean_square = builder.divide(&squared_sum, &width_rows); + let epsilon = builder.const_f32(epsilon); + let epsilon = builder.broadcast(&epsilon, &[], vec![rows]); + let mean_square = builder.add(&mean_square, &epsilon); + let inv_rms = builder.rsqrt(&mean_square); + let inv_rms = builder.broadcast(&inv_rms, &[0], vec![rows, width]); + let normalized = builder.multiply(value, &inv_rms); + let one = builder.const_f32(1.0); + let one = builder.broadcast(&one, &[], vec![width]); + let weight = builder.add(weight, &one); + let weight = builder.broadcast(&weight, &[1], vec![rows, width]); + builder.multiply(&normalized, &weight) +} + +fn gemma3_project( + builder: &mut Builder, + hidden: &Val, + args: &mut Args, + config: &LlavaVisionConfig, + tokens_per_side: usize, + kernel_size: usize, +) -> Val { + let pooled = builder.reshape( + hidden, + vec![ + tokens_per_side, + kernel_size, + tokens_per_side, + kernel_size, + config.hidden, + ], + ); + let zero = builder.const_f32(0.0); + let pooled = builder.reduce_add(&pooled, 3, &zero); + let pooled = builder.reduce_add(&pooled, 1, &zero); + let divisor = builder.const_f32((kernel_size * kernel_size) as f32); + let divisor = builder.broadcast( + &divisor, + &[], + vec![tokens_per_side, tokens_per_side, config.hidden], + ); + let pooled = builder.divide(&pooled, &divisor); + let pooled = builder.reshape( + &pooled, + vec![tokens_per_side * tokens_per_side, config.hidden], + ); + let normalized = gemma_rms_norm(builder, &pooled, &args.take(), config.layer_norm_eps); + let projection = args.take(); + builder.dot_general( + &normalized, + &projection, + &[], + &[], + &[1], + &[0], + vec![tokens_per_side * tokens_per_side, config.text_hidden], + ) +} + fn activate(builder: &mut Builder, value: &Val, activation: VisionActivation) -> Val { match activation { VisionActivation::ExactGelu => exact_gelu(builder, value), @@ -314,12 +413,37 @@ fn emit_vision_impl(config: &LlavaVisionConfig, diagnostics: bool) -> String { outputs.push(hidden.clone()); } } + if matches!(config.projector, VisionProjector::Gemma3AvgPool { .. }) { + hidden = layer_norm( + &mut builder, + &hidden, + &args.take(), + &args.take(), + config.layer_norm_eps, + ); + } if config.drop_first_token { hidden = builder.slice(&hidden, &[(1, config.position_count()), (0, config.hidden)]); } - let projected = linear_2d(&mut builder, &hidden, &args.take(), &args.take()); - let projected = exact_gelu(&mut builder, &projected); - let projected = linear_2d(&mut builder, &projected, &args.take(), &args.take()); + let projected = match config.projector { + VisionProjector::LlavaMlp => { + let projected = linear_2d(&mut builder, &hidden, &args.take(), &args.take()); + let projected = exact_gelu(&mut builder, &projected); + linear_2d(&mut builder, &projected, &args.take(), &args.take()) + } + VisionProjector::Gemma3AvgPool { + tokens_per_side, + kernel_size, + .. + } => gemma3_project( + &mut builder, + &hidden, + &mut args, + config, + tokens_per_side, + kernel_size, + ), + }; assert_eq!(args.cursor, specs.len(), "vision weight schema drifted"); let outputs = if let Some(mut values) = diagnostic_values { values.push(hidden); diff --git a/src/lib/mlxcel-xla/src/emitter/vision_config.rs b/src/lib/mlxcel-xla/src/emitter/vision_config.rs index aee70e746..87a659983 100644 --- a/src/lib/mlxcel-xla/src/emitter/vision_config.rs +++ b/src/lib/mlxcel-xla/src/emitter/vision_config.rs @@ -24,6 +24,51 @@ pub(crate) enum VisionActivation { GeluPytorchTanh, } +impl VisionActivation { + const fn stable_identity(self) -> &'static str { + match self { + Self::ExactGelu => "ExactGelu", + Self::GeluPytorchTanh => "GeluPytorchTanh", + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub(crate) enum VisionProjector { + LlavaMlp, + Gemma3AvgPool { + tokens_per_side: usize, + kernel_size: usize, + image_token_id: i32, + pad_token_id: i32, + boi_token_id: i32, + eoi_token_id: i32, + newline_token_id: i32, + }, +} + +impl VisionProjector { + fn stable_identity(self) -> String { + match self { + Self::LlavaMlp => "LlavaMlp".to_string(), + Self::Gemma3AvgPool { + tokens_per_side, + kernel_size, + image_token_id, + pad_token_id, + boi_token_id, + eoi_token_id, + newline_token_id, + } => format!( + "Gemma3AvgPool {{ tokens_per_side: {tokens_per_side}, kernel_size: {kernel_size}, \ + image_token_id: {image_token_id}, pad_token_id: {pad_token_id}, \ + boi_token_id: {boi_token_id}, eoi_token_id: {eoi_token_id}, \ + newline_token_id: {newline_token_id} }}" + ), + } + } +} + #[cfg(test)] mod tests { use super::*; @@ -171,6 +216,65 @@ mod tests { } ); } + + #[test] + fn gemma3_avg_pool_projector_binds_pool_and_prompt_contract() { + let config = LlavaVisionConfig::from_json_str( + &serde_json::json!({ + "model_type": "gemma3", + "image_token_index": 99, + "pad_token_id": 0, + "boi_token_index": 97, + "eoi_token_index": 98, + "mm_tokens_per_image": 1, + "vision_config": { + "model_type": "siglip_vision_model", + "image_size": 28, + "patch_size": 14, + "num_channels": 3, + "hidden_size": 8, + "intermediate_size": 16, + "num_hidden_layers": 2, + "num_attention_heads": 2 + }, + "text_config": {"hidden_size": 12} + }) + .to_string(), + ) + .unwrap(); + assert_eq!(config.feature_layer, 1); + assert!(!config.drop_first_token); + assert_eq!(config.image_tokens(), 1); + assert_eq!( + config.projector, + VisionProjector::Gemma3AvgPool { + tokens_per_side: 1, + kernel_size: 2, + image_token_id: 99, + pad_token_id: 0, + boi_token_id: 97, + eoi_token_id: 98, + newline_token_id: 108, + } + ); + assert!(config.fingerprint().contains("image_token_id: 99")); + assert!(config.fingerprint().contains("newline_token_id: 108")); + assert!(config.fingerprint().starts_with("iree-vision-v3:")); + let specs = config.weight_specs(); + assert_eq!( + specs[specs.len() - 4].name, + "vision_tower.vision_model.post_layernorm.weight" + ); + assert_eq!( + specs[specs.len() - 3].name, + "vision_tower.vision_model.post_layernorm.bias" + ); + assert_eq!( + specs[specs.len() - 2].name, + "multi_modal_projector.mm_soft_emb_norm.weight" + ); + assert_eq!(specs.last().unwrap().shape, [8, 12]); + } } #[derive(Debug, Clone, PartialEq)] @@ -188,6 +292,7 @@ pub(crate) struct LlavaVisionConfig { pub(crate) feature_layer: usize, pub(crate) drop_first_token: bool, pub(crate) text_hidden: usize, + pub(crate) projector: VisionProjector, } #[derive(Debug, Clone, PartialEq, Eq)] @@ -232,8 +337,14 @@ impl LlavaVisionConfig { pub(crate) fn from_json_str(text: &str) -> Result { let root: Value = serde_json::from_str(text).map_err(|error| format!("parse config.json: {error}"))?; - if root.get("model_type").and_then(Value::as_str) != Some("llava") { - return Err("IREE vision currently supports only model_type=llava".to_string()); + let model_type = root + .get("model_type") + .and_then(Value::as_str) + .ok_or_else(|| "config.json model_type is required".to_string())?; + if !matches!(model_type, "llava" | "gemma3") { + return Err(format!( + "IREE vision currently supports model_type=llava or gemma3, got {model_type:?}" + )); } let vision = object(&root, "vision_config")?; let vision_model_type = vision @@ -292,12 +403,14 @@ impl LlavaVisionConfig { )); } }; - match root.get("projector_hidden_act").and_then(Value::as_str) { - None | Some("gelu") => {} - other => { - return Err(format!( - "unsupported projector_hidden_act {other:?}; only exact GELU is qualified" - )); + if model_type == "llava" { + match root.get("projector_hidden_act").and_then(Value::as_str) { + None | Some("gelu") => {} + other => { + return Err(format!( + "unsupported projector_hidden_act {other:?}; only exact GELU is qualified" + )); + } } } let text_hidden = object(&root, "text_config")? @@ -311,10 +424,13 @@ impl LlavaVisionConfig { if text_hidden == 0 { return Err("text_config.hidden_size must be greater than zero".to_string()); } - let requested_layer = root - .get("vision_feature_layer") - .and_then(Value::as_i64) - .unwrap_or(-2); + let requested_layer = if model_type == "gemma3" { + -1 + } else { + root.get("vision_feature_layer") + .and_then(Value::as_i64) + .unwrap_or(-2) + }; let resolved = if requested_layer < 0 { i64::try_from(layers).map_err(|_| "vision layer count does not fit i64".to_string())? + requested_layer @@ -326,18 +442,69 @@ impl LlavaVisionConfig { "vision_feature_layer={requested_layer} resolves outside {layers} encoder layers" )); } - let strategy = root - .get("vision_feature_select_strategy") - .and_then(Value::as_str) - .unwrap_or("default"); - let drop_first_token = match strategy { - "default" => true, - "full" => false, - other => { + let drop_first_token = if model_type == "gemma3" { + false + } else { + match root + .get("vision_feature_select_strategy") + .and_then(Value::as_str) + .unwrap_or("default") + { + "default" => true, + "full" => false, + other => { + return Err(format!( + "unsupported vision_feature_select_strategy={other:?}" + )); + } + } + }; + let projector = if model_type == "gemma3" { + if class_token { + return Err("Gemma3 IREE vision requires the class-token-free SigLIP tower".into()); + } + let image_tokens = root + .get("mm_tokens_per_image") + .and_then(Value::as_u64) + .map_or(Ok(256usize), |value| { + usize::try_from(value) + .map_err(|_| "mm_tokens_per_image does not fit usize".to_string()) + })?; + let tokens_per_side = (image_tokens as f64).sqrt() as usize; + if tokens_per_side == 0 + || tokens_per_side + .checked_mul(tokens_per_side) + .ok_or_else(|| "Gemma3 image-token grid overflowed".to_string())? + != image_tokens + { return Err(format!( - "unsupported vision_feature_select_strategy={other:?}" + "Gemma3 mm_tokens_per_image={image_tokens} must be a non-zero square" )); } + let patch_grid = image_size / patch_size; + if image_size % patch_size != 0 || patch_grid % tokens_per_side != 0 { + return Err(format!( + "Gemma3 patch grid {patch_grid} must divide exactly into {tokens_per_side} pooled tokens per side" + )); + } + let token_id = |name: &str, default: i32| -> Result { + match root.get(name).and_then(Value::as_i64) { + Some(value) => i32::try_from(value) + .map_err(|_| format!("config.json {name} does not fit i32")), + None => Ok(default), + } + }; + VisionProjector::Gemma3AvgPool { + tokens_per_side, + kernel_size: patch_grid / tokens_per_side, + image_token_id: token_id("image_token_index", 262_144)?, + pad_token_id: token_id("pad_token_id", 0)?, + boi_token_id: token_id("boi_token_index", 255_999)?, + eoi_token_id: token_id("eoi_token_index", 256_000)?, + newline_token_id: 108, + } + } else { + VisionProjector::LlavaMlp }; Ok(Self { image_size, @@ -353,6 +520,7 @@ impl LlavaVisionConfig { feature_layer: resolved as usize, drop_first_token, text_hidden, + projector, }) } @@ -368,14 +536,24 @@ impl LlavaVisionConfig { #[must_use] pub(crate) fn image_tokens(&self) -> usize { - self.position_count() - usize::from(self.drop_first_token) + match self.projector { + VisionProjector::LlavaMlp => self.position_count() - usize::from(self.drop_first_token), + VisionProjector::Gemma3AvgPool { + tokens_per_side, .. + } => tokens_per_side * tokens_per_side, + } } #[must_use] pub(crate) fn fingerprint(&self) -> String { + let schema = if matches!(self.projector, VisionProjector::Gemma3AvgPool { .. }) { + "iree-vision-v3" + } else { + "iree-vision-v2" + }; format!( - "llava-vision-v1:image={}:patch={}:channels={}:hidden={}:intermediate={}:layers={}:\ - heads={}:eps={:08x}:activation={:?}:class={}:feature={}:drop_first={}:text={}", + "{schema}:image={}:patch={}:channels={}:hidden={}:intermediate={}:layers={}:\ + heads={}:eps={:08x}:activation={}:class={}:feature={}:drop_first={}:text={}:projector={}", self.image_size, self.patch_size, self.channels, @@ -384,11 +562,12 @@ impl LlavaVisionConfig { self.layers, self.heads, self.layer_norm_eps.to_bits(), - self.activation, + self.activation.stable_identity(), self.class_token, self.feature_layer, self.drop_first_token, - self.text_hidden + self.text_hidden, + self.projector.stable_identity(), ) } @@ -448,18 +627,42 @@ impl LlavaVisionConfig { }); } } - specs.extend([ - self.spec( - "multi_modal_projector.linear_1.weight", - [self.text_hidden, self.hidden], - ), - self.spec("multi_modal_projector.linear_1.bias", [self.text_hidden]), - self.spec( - "multi_modal_projector.linear_2.weight", - [self.text_hidden, self.text_hidden], - ), - self.spec("multi_modal_projector.linear_2.bias", [self.text_hidden]), - ]); + if matches!(self.projector, VisionProjector::Gemma3AvgPool { .. }) { + specs.extend([ + self.spec( + "vision_tower.vision_model.post_layernorm.weight", + [self.hidden], + ), + self.spec( + "vision_tower.vision_model.post_layernorm.bias", + [self.hidden], + ), + ]); + } + match self.projector { + VisionProjector::LlavaMlp => specs.extend([ + self.spec( + "multi_modal_projector.linear_1.weight", + [self.text_hidden, self.hidden], + ), + self.spec("multi_modal_projector.linear_1.bias", [self.text_hidden]), + self.spec( + "multi_modal_projector.linear_2.weight", + [self.text_hidden, self.text_hidden], + ), + self.spec("multi_modal_projector.linear_2.bias", [self.text_hidden]), + ]), + VisionProjector::Gemma3AvgPool { .. } => specs.extend([ + self.spec( + "multi_modal_projector.mm_soft_emb_norm.weight", + [self.hidden], + ), + self.spec( + "multi_modal_projector.mm_input_projection_weight", + [self.hidden, self.text_hidden], + ), + ]), + } specs } diff --git a/src/lib/mlxcel-xla/src/iree.rs b/src/lib/mlxcel-xla/src/iree.rs index e31100190..22be32c3b 100644 --- a/src/lib/mlxcel-xla/src/iree.rs +++ b/src/lib/mlxcel-xla/src/iree.rs @@ -58,9 +58,9 @@ use mlxcel_core::session::PreparedPrefill; use safetensors::{Dtype, SafeTensors}; use crate::emitter::{ - Config, DeepStackConfig, Gemma3nConfig, Gemma3nWeightSpec, Precision, QuantConfig, - check_packed_supported, emit_decode_ragged_with, emit_decode_with, - emit_gemma3n_decode_ragged_with_qmv, emit_gemma3n_decode_with_qmv, + Config, DeepStackConfig, Gemma3nConfig, Gemma3nWeightSpec, LlavaVisionConfig, Precision, + QuantConfig, VisionProjector, check_packed_supported, emit_decode_ragged_with, + emit_decode_with, emit_gemma3n_decode_ragged_with_qmv, emit_gemma3n_decode_with_qmv, emit_gemma3n_prefill_embeddings_ple_with_qmv, emit_gemma3n_prefill_with_qmv, emit_prefill_embeddings_deepstack_with, emit_prefill_embeddings_with, emit_prefill_with, gemma3n_qmv_artifact_identity, gemma3n_qmv_is_available, gemma3n_weight_specs, quant_in_graph, @@ -76,6 +76,10 @@ use crate::prepared::{ validate_slot, }; use crate::prepared_deepstack::PreparedDeepStack; +use crate::prepared_gemma3::{ + GEMMA3_VLM_NEWLINE_WRAPPER_TOKEN_ID, Gemma3VlmCompatibilityContract, + bind_gemma3_vlm_compatibility, +}; use crate::{DeepStackFeatures, DeepStackPreparedPrefill, Gemma3nDensePle, Gemma3nPreparedPrefill}; // The loader reads the per-architecture checkpoint-weight order from // `weights::weight_specs`, which sources its names from `weight_names::scheme_names` @@ -582,6 +586,177 @@ impl RuntimeConfig { } } +fn gemma3_vlm_compatibility_from_json( + config_json: &str, + cfg: &RuntimeConfig, +) -> Result, String> { + let root: serde_json::Value = + serde_json::from_str(config_json).map_err(|error| format!("parse config.json: {error}"))?; + let is_vlm = root.get("model_type").and_then(serde_json::Value::as_str) == Some("gemma3") + && root + .get("vision_config") + .is_some_and(serde_json::Value::is_object); + if !is_vlm { + return Ok(None); + } + let RuntimeConfig::Dense(text) = cfg else { + return Err("Gemma3 VLM requires the dense Gemma3 text runtime".to_string()); + }; + if !text.embed_scale || !text.embeddings_prefill_uses_authoritative_mask { + return Err( + "Gemma3 VLM requires sqrt(hidden) token scaling and an authoritative embeddings-prefill mask" + .to_string(), + ); + } + let vision = LlavaVisionConfig::from_json_str(config_json)?; + let VisionProjector::Gemma3AvgPool { + image_token_id, + pad_token_id, + boi_token_id, + eoi_token_id, + newline_token_id, + .. + } = vision.projector + else { + return Err("Gemma3 VLM requires the average-pool vision projector".to_string()); + }; + if newline_token_id != GEMMA3_VLM_NEWLINE_WRAPPER_TOKEN_ID { + return Err(format!( + "Gemma3 VLM newline wrapper token {newline_token_id} does not match the prepared-prefill contract {}", + GEMMA3_VLM_NEWLINE_WRAPPER_TOKEN_ID + )); + } + if vision.text_hidden != text.hidden { + return Err(format!( + "Gemma3 VLM projector hidden size {} does not match language hidden size {}", + vision.text_hidden, text.hidden + )); + } + Ok(Some(Gemma3VlmCompatibilityContract::new( + image_token_id, + pad_token_id, + boi_token_id, + eoi_token_id, + vision.image_tokens(), + text.hidden, + text.sliding_window, + text.sliding_pattern, + text.rope_theta, + text.rope_local_base, + vision.fingerprint(), + ))) +} + +fn gemma3_vlm_compatibility( + model_dir: &Path, + cfg: &RuntimeConfig, +) -> Result, String> { + let path = model_dir.join("config.json"); + let config_json = std::fs::read_to_string(&path) + .map_err(|error| format!("read {}: {error}", path.display()))?; + gemma3_vlm_compatibility_from_json(&config_json, cfg) + .map_err(|error| format!("{}: {error}", path.display())) +} + +#[cfg(test)] +mod gemma3_vlm_compatibility_tests { + use super::*; + + fn gemma3_text_config() -> serde_json::Value { + serde_json::json!({ + "model_type": "gemma3_text", + "hidden_size": 12, + "intermediate_size": 24, + "num_hidden_layers": 2, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "head_dim": 6, + "rms_norm_eps": 1e-6, + "rope_theta": 1_000_000.0, + "rope_local_base_freq": 10_000.0, + "sliding_window": 4, + "sliding_window_pattern": 2, + "vocab_size": 128, + "hidden_activation": "gelu_pytorch_tanh" + }) + } + + fn runtime(config_json: &str) -> RuntimeConfig { + RuntimeConfig::Dense(Box::new( + Config::from_json_str(config_json) + .unwrap() + .with_context_capacity(8) + .unwrap(), + )) + } + + #[test] + fn gemma3_wrapper_binds_prepared_and_vision_contracts() { + let config_json = serde_json::json!({ + "model_type": "gemma3", + "image_token_index": 99, + "pad_token_id": 0, + "boi_token_index": 97, + "eoi_token_index": 98, + "mm_tokens_per_image": 1, + "vision_config": { + "model_type": "siglip_vision_model", + "image_size": 28, + "patch_size": 14, + "num_channels": 3, + "hidden_size": 8, + "intermediate_size": 16, + "num_hidden_layers": 2, + "num_attention_heads": 2 + }, + "text_config": gemma3_text_config() + }) + .to_string(); + let runtime = runtime(&config_json); + let contract = gemma3_vlm_compatibility_from_json(&config_json, &runtime) + .unwrap() + .unwrap(); + let identity = contract.stable_identity(); + + for component in [ + "prepared_prefill=gemma3-vlm-owned-f32-embeddings", + "mask_mode=gemma3-vlm-bidirectional-padding-f32-min-v1", + "post_scale_policy=gemma3-vlm-post-scale-text-sqrt-hidden-image-identity-pad-zero-v1", + "newline_wrapper_token_id=108", + "image_token_id=99", + "pad_token_id=0", + "boi_token_id=97", + "eoi_token_id=98", + "sliding_window=4", + "sliding_pattern=2", + "vision_identity=iree-vision-v3:", + ] { + assert!( + identity.contains(component), + "missing {component}: {identity}" + ); + } + assert_ne!( + bind_gemma3_vlm_compatibility(runtime.artifact_identity(), Some(&contract)), + runtime.artifact_identity() + ); + } + + #[test] + fn direct_root_gemma3_text_keeps_existing_runtime_identity() { + let config_json = gemma3_text_config().to_string(); + let runtime = runtime(&config_json); + let contract = gemma3_vlm_compatibility_from_json(&config_json, &runtime).unwrap(); + assert!(contract.is_none()); + + let existing = runtime.artifact_identity(); + assert_eq!( + bind_gemma3_vlm_compatibility(existing.clone(), contract.as_ref()), + existing + ); + } +} + fn effective_precision(device: &str, cfg: &RuntimeConfig) -> Result { match cfg { RuntimeConfig::Dense(_) => resolve_precision_checked(device), @@ -882,6 +1057,7 @@ struct CompiledBundle { fn compile_bundle( device: &str, cfg: &RuntimeConfig, + gemma3_vlm_contract: Option<&Gemma3VlmCompatibilityContract>, gemma3n_qmv: bool, prefill_mlir: &str, prefill_tag: &str, @@ -942,7 +1118,8 @@ fn compile_bundle( )?; let mut fingerprint = std::collections::hash_map::DefaultHasher::new(); "mlxcel-xla-runtime-bundle-v2-explicit-position-mode".hash(&mut fingerprint); - cfg.artifact_identity().hash(&mut fingerprint); + bind_gemma3_vlm_compatibility(cfg.artifact_identity(), gemma3_vlm_contract) + .hash(&mut fingerprint); format!("{:?}", cfg.weight_specs()).hash(&mut fingerprint); format!("{:?}", effective_precision(device, cfg)?).hash(&mut fingerprint); if gemma3n_qmv { @@ -965,7 +1142,11 @@ fn compile_bundle( } /// Emit and compile the argmax prefill bundle for a single-sequence engine. -fn compile_vmfbs(device: &str, cfg: &RuntimeConfig) -> Result { +fn compile_vmfbs( + device: &str, + cfg: &RuntimeConfig, + gemma3_vlm_contract: Option<&Gemma3VlmCompatibilityContract>, +) -> Result { let precision = effective_precision(device, cfg)?; let native_qmv = match cfg { RuntimeConfig::Dense(_) => false, @@ -994,6 +1175,7 @@ fn compile_vmfbs(device: &str, cfg: &RuntimeConfig) -> Result Result { let cfg = RuntimeConfig::from_json(model_dir, context_capacity)?; runtime_ffi_dimensions(&cfg, cfg.weight_specs().len())?; - let bundle = compile_vmfbs(device, &cfg)?; + let gemma3_vlm_contract = gemma3_vlm_compatibility(model_dir, &cfg)?; + let bundle = compile_vmfbs(device, &cfg, gemma3_vlm_contract.as_ref())?; let ctx = create_ctx(model_dir, &cfg, device, &bundle)?; Ok(Self { ctx, @@ -2265,6 +2448,7 @@ impl IreeRaggedLlama { let bundle = compile_bundle( device, &cfg, + gemma3_vlm_compatibility(model_dir, &cfg)?.as_ref(), native_qmv, &prefill_mlir, "prefill_logits", diff --git a/src/lib/mlxcel-xla/src/lib.rs b/src/lib/mlxcel-xla/src/lib.rs index 2b1328267..04c8f73a5 100644 --- a/src/lib/mlxcel-xla/src/lib.rs +++ b/src/lib/mlxcel-xla/src/lib.rs @@ -66,6 +66,7 @@ mod operator_numeric_contract; #[cfg_attr(not(feature = "iree"), allow(dead_code))] mod prepared; mod prepared_deepstack; +mod prepared_gemma3; mod prepared_gemma3n; #[cfg(feature = "iree")] @@ -210,6 +211,11 @@ pub use emitter::{Gemma3nDiagnosticLayout, Gemma3nDiagnosticSegment}; #[cfg(feature = "iree")] pub use prepared::PreparedInputError; pub use prepared_deepstack::{DeepStackFeatures, DeepStackInputError, DeepStackPreparedPrefill}; +pub use prepared_gemma3::{ + GEMMA3_VLM_MASK_MODE, GEMMA3_VLM_MASKED_VALUE, GEMMA3_VLM_NEWLINE_WRAPPER_TOKEN_ID, + GEMMA3_VLM_POST_SCALE_POLICY, GEMMA3_VLM_PREPARED_PREFILL_CONTRACT, Gemma3VlmPreparedError, + prepare_gemma3_vlm_prefill, +}; pub use prepared_gemma3n::{Gemma3nDensePle, Gemma3nDensePleError, Gemma3nPreparedPrefill}; #[cfg(feature = "iree")] pub use sampler::SampleParams; diff --git a/src/lib/mlxcel-xla/src/prepared.rs b/src/lib/mlxcel-xla/src/prepared.rs index bc1b4eaad..67cbc382d 100644 --- a/src/lib/mlxcel-xla/src/prepared.rs +++ b/src/lib/mlxcel-xla/src/prepared.rs @@ -213,7 +213,7 @@ impl fmt::Display for PreparedInputError { "prepared {tensor} byte count mismatch: expected {expected}, got {actual}" ), Self::InvalidFloat { tensor } => { - write!(f, "prepared {tensor} contains NaN or an invalid positive mask value") + write!(f, "prepared {tensor} contains an invalid floating-point or mask value") } Self::ShapeOverflow => f.write_str("prepared IREE static shape overflowed"), } @@ -560,10 +560,30 @@ impl PreparedIreePrefill { bias_tensor.shape.clone(), )); } + let gemma3_vlm_mask = value + .modalities + .iter() + .any(|modality| modality.family == "gemma3"); + if gemma3_vlm_mask + && (value.attention_bias.causal + || !matrix_bias + || compact_bias + .iter() + .any(|value| *value != 0.0 && *value != f32::MIN)) + { + return Err(PreparedInputError::InvalidFloat { + tensor: "Gemma3 VLM attention bias", + }); + } + let masked_value = if gemma3_vlm_mask { + f32::MIN + } else { + MASKED_VALUE + }; let bias_count = context_capacity .checked_mul(context_capacity) .ok_or(PreparedInputError::ShapeOverflow)?; - let mut attention_bias = vec![MASKED_VALUE; bias_count]; + let mut attention_bias = vec![masked_value; bias_count]; for query in 0..effective_len { for key in 0..effective_len { let base = if key_bias { @@ -573,7 +593,7 @@ impl PreparedIreePrefill { }; attention_bias[query * context_capacity + key] = if value.attention_bias.causal && key > query { - MASKED_VALUE + masked_value } else { base }; diff --git a/src/lib/mlxcel-xla/src/prepared_gemma3.rs b/src/lib/mlxcel-xla/src/prepared_gemma3.rs new file mode 100644 index 000000000..cd2691b5f --- /dev/null +++ b/src/lib/mlxcel-xla/src/prepared_gemma3.rs @@ -0,0 +1,636 @@ +// 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. + +//! Canonical Gemma3 VLM merge at the XLA embeddings-prefill boundary. +//! +//! The ordinary Gemma3 token graph gathers embeddings and multiplies them by +//! `sqrt(hidden_size)`. The embeddings entry deliberately does neither, so this +//! producer exports text rows after that scale while keeping projected image +//! rows at their original magnitude. It also owns the reference's bidirectional +//! padding mask instead of substituting a causal or sliding-window mask. + +use std::fmt; + +use mlxcel_core::session::{ + OwnedTensor, PreparedAttentionBias, PreparedModality, PreparedPositions, PreparedPrefill, + PreparedPrefillError, PreparedTensorDType, +}; + +/// Exact masked value emitted by `vision::merge::prepare_inputs_for_multimodal`. +pub const GEMMA3_VLM_MASKED_VALUE: f32 = f32::MIN; + +/// Stable identity component for the qualified external-mask behavior. +pub const GEMMA3_VLM_MASK_MODE: &str = "gemma3-vlm-bidirectional-padding-f32-min-v1"; + +/// Stable identity for the embedding values handed to embeddings-prefill. +/// +/// Text rows have already received Gemma's `sqrt(hidden_size)` scale, projected +/// image rows stay at projector magnitude, and padding rows are zero. +pub const GEMMA3_VLM_POST_SCALE_POLICY: &str = + "gemma3-vlm-post-scale-text-sqrt-hidden-image-identity-pad-zero-v1"; + +/// Gemma3Processor wraps every expanded image block with this newline token. +pub const GEMMA3_VLM_NEWLINE_WRAPPER_TOKEN_ID: i32 = 108; + +/// Stable schema for the owned payload consumed by the authoritative +/// embeddings-prefill entry. +pub const GEMMA3_VLM_PREPARED_PREFILL_CONTRACT: &str = + "gemma3-vlm-owned-f32-embeddings-sequential-1d-additive-f32-1x1xlxl-noncausal-v1"; + +/// Cross-runtime compatibility contract for a Gemma3 VLM bundle. +/// +/// The language graphs, resident SigLIP/projector module, and host producer are +/// separate runtime components. This explicit identity prevents their semantic +/// boundary from depending on derived `Debug` output. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct Gemma3VlmCompatibilityContract { + pub(crate) prepared_prefill: String, + pub(crate) mask_mode: String, + pub(crate) masked_value_bits: u32, + pub(crate) post_scale_policy: String, + pub(crate) newline_wrapper_token_id: i32, + pub(crate) image_token_id: i32, + pub(crate) pad_token_id: i32, + pub(crate) boi_token_id: i32, + pub(crate) eoi_token_id: i32, + pub(crate) tokens_per_image: usize, + pub(crate) hidden_size: usize, + pub(crate) sliding_window: Option, + pub(crate) sliding_pattern: usize, + pub(crate) global_rope_theta_bits: u64, + pub(crate) local_rope_base_bits: Option, + pub(crate) vision_identity: String, +} + +impl Gemma3VlmCompatibilityContract { + #[allow(clippy::too_many_arguments)] + pub(crate) fn new( + image_token_id: i32, + pad_token_id: i32, + boi_token_id: i32, + eoi_token_id: i32, + tokens_per_image: usize, + hidden_size: usize, + sliding_window: Option, + sliding_pattern: usize, + global_rope_theta: f64, + local_rope_base: Option, + vision_identity: String, + ) -> Self { + Self { + prepared_prefill: GEMMA3_VLM_PREPARED_PREFILL_CONTRACT.to_string(), + mask_mode: GEMMA3_VLM_MASK_MODE.to_string(), + masked_value_bits: GEMMA3_VLM_MASKED_VALUE.to_bits(), + post_scale_policy: GEMMA3_VLM_POST_SCALE_POLICY.to_string(), + newline_wrapper_token_id: GEMMA3_VLM_NEWLINE_WRAPPER_TOKEN_ID, + image_token_id, + pad_token_id, + boi_token_id, + eoi_token_id, + tokens_per_image, + hidden_size, + sliding_window, + sliding_pattern, + global_rope_theta_bits: global_rope_theta.to_bits(), + local_rope_base_bits: local_rope_base.map(f64::to_bits), + vision_identity, + } + } + + /// Canonical, field-labelled identity. Integer bit patterns avoid + /// locale/toolchain-dependent float formatting. + pub(crate) fn stable_identity(&self) -> String { + let sliding_window = self + .sliding_window + .map_or_else(|| "none".to_string(), |value| value.to_string()); + let local_rope_base = self + .local_rope_base_bits + .map_or_else(|| "none".to_string(), |bits| format!("{bits:016x}")); + format!( + "gemma3-vlm-runtime-bundle-v1\ + ;prepared_prefill={}\ + ;mask_mode={}\ + ;masked_value_bits={:08x}\ + ;post_scale_policy={}\ + ;newline_wrapper_token_id={}\ + ;image_token_id={}\ + ;pad_token_id={}\ + ;boi_token_id={}\ + ;eoi_token_id={}\ + ;tokens_per_image={}\ + ;hidden_size={}\ + ;sliding_window={}\ + ;sliding_pattern={}\ + ;global_rope_theta_bits={:016x}\ + ;local_rope_base_bits={}\ + ;vision_identity={}", + self.prepared_prefill, + self.mask_mode, + self.masked_value_bits, + self.post_scale_policy, + self.newline_wrapper_token_id, + self.image_token_id, + self.pad_token_id, + self.boi_token_id, + self.eoi_token_id, + self.tokens_per_image, + self.hidden_size, + sliding_window, + self.sliding_pattern, + self.global_rope_theta_bits, + local_rope_base, + self.vision_identity, + ) + } +} + +/// Add the prepared-prefill identity only for a qualified multimodal runtime. +/// +/// Returning the input verbatim for `None` is intentional: ordinary text-only +/// bundle identities must not move when this VLM-only contract evolves. +pub(crate) fn bind_gemma3_vlm_compatibility( + runtime_identity: String, + contract: Option<&Gemma3VlmCompatibilityContract>, +) -> String { + match contract { + Some(contract) => format!( + "{runtime_identity};prepared_prefill_contract={}", + contract.stable_identity() + ), + None => runtime_identity, + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum Gemma3VlmPreparedError { + Empty, + ZeroHiddenSize, + Capacity { + sequence_len: usize, + context_capacity: usize, + }, + ShapeOverflow, + EmbeddingCount { + expected: usize, + actual: usize, + }, + ProjectedShape { + values: usize, + hidden_size: usize, + }, + AttentionMaskCount { + expected: usize, + actual: usize, + }, + InvalidAttentionMask { + index: usize, + value: i32, + }, + PaddingMaskMismatch { + index: usize, + token_id: i32, + mask: i32, + }, + PlaceholderCount { + positions: usize, + projected_tokens: usize, + }, + ImageCount { + images: usize, + projected_tokens: usize, + }, + NonFinite { + tensor: &'static str, + index: usize, + }, + InvalidTokenConfiguration, + Prepared(PreparedPrefillError), +} + +impl fmt::Display for Gemma3VlmPreparedError { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Empty => formatter.write_str("Gemma3 VLM prefill requires at least one token"), + Self::ZeroHiddenSize => { + formatter.write_str("Gemma3 VLM hidden size must be greater than zero") + } + Self::Capacity { + sequence_len, + context_capacity, + } => write!( + formatter, + "Gemma3 VLM sequence length {sequence_len} exceeds context capacity {context_capacity}" + ), + Self::ShapeOverflow => formatter.write_str("Gemma3 VLM tensor shape overflowed"), + Self::EmbeddingCount { expected, actual } => write!( + formatter, + "Gemma3 VLM raw text embeddings have {actual} values, expected {expected}" + ), + Self::ProjectedShape { + values, + hidden_size, + } => write!( + formatter, + "Gemma3 VLM projected image features have {values} values, which is not divisible by hidden size {hidden_size}" + ), + Self::AttentionMaskCount { expected, actual } => write!( + formatter, + "Gemma3 VLM attention mask has {actual} values, expected {expected}" + ), + Self::InvalidAttentionMask { index, value } => write!( + formatter, + "Gemma3 VLM attention mask[{index}] is {value}; expected 0 or 1" + ), + Self::PaddingMaskMismatch { + index, + token_id, + mask, + } => write!( + formatter, + "Gemma3 VLM token {index} id {token_id} disagrees with padding mask {mask}" + ), + Self::PlaceholderCount { + positions, + projected_tokens, + } => write!( + formatter, + "Gemma3 VLM prompt has {positions} image positions but projector returned {projected_tokens} tokens" + ), + Self::ImageCount { + images, + projected_tokens, + } => write!( + formatter, + "Gemma3 VLM declares {images} image(s) but projector returned {projected_tokens} token(s)" + ), + Self::NonFinite { tensor, index } => write!( + formatter, + "Gemma3 VLM {tensor} contains a non-finite value at flat index {index}" + ), + Self::InvalidTokenConfiguration => { + formatter.write_str("Gemma3 VLM pad_token_id and image_token_id must be distinct") + } + Self::Prepared(error) => write!(formatter, "Gemma3 VLM prepared prefill: {error}"), + } + } +} + +impl std::error::Error for Gemma3VlmPreparedError {} + +impl From for Gemma3VlmPreparedError { + fn from(value: PreparedPrefillError) -> Self { + Self::Prepared(value) + } +} + +fn f32_bytes(values: &[f32]) -> Vec { + values + .iter() + .flat_map(|value| value.to_le_bytes()) + .collect() +} + +fn first_non_finite(values: &[f32]) -> Option { + values.iter().position(|value| !value.is_finite()) +} + +/// Build the exact post-scale embedding and additive-mask payload for one +/// Gemma3 VLM request. +#[allow(clippy::too_many_arguments)] +pub fn prepare_gemma3_vlm_prefill( + token_ids: Vec, + raw_text_embeddings: &[f32], + projected_image_features: &[f32], + attention_mask: &[i32], + hidden_size: usize, + context_capacity: usize, + pad_token_id: i32, + image_token_id: i32, + image_count: usize, +) -> Result { + let sequence_len = token_ids.len(); + if sequence_len == 0 { + return Err(Gemma3VlmPreparedError::Empty); + } + if hidden_size == 0 { + return Err(Gemma3VlmPreparedError::ZeroHiddenSize); + } + if sequence_len > context_capacity { + return Err(Gemma3VlmPreparedError::Capacity { + sequence_len, + context_capacity, + }); + } + if pad_token_id == image_token_id { + return Err(Gemma3VlmPreparedError::InvalidTokenConfiguration); + } + + let expected_embeddings = sequence_len + .checked_mul(hidden_size) + .ok_or(Gemma3VlmPreparedError::ShapeOverflow)?; + if raw_text_embeddings.len() != expected_embeddings { + return Err(Gemma3VlmPreparedError::EmbeddingCount { + expected: expected_embeddings, + actual: raw_text_embeddings.len(), + }); + } + if attention_mask.len() != sequence_len { + return Err(Gemma3VlmPreparedError::AttentionMaskCount { + expected: sequence_len, + actual: attention_mask.len(), + }); + } + if let Some(index) = first_non_finite(raw_text_embeddings) { + return Err(Gemma3VlmPreparedError::NonFinite { + tensor: "raw text embeddings", + index, + }); + } + if let Some(index) = first_non_finite(projected_image_features) { + return Err(Gemma3VlmPreparedError::NonFinite { + tensor: "projected image features", + index, + }); + } + if projected_image_features.len() % hidden_size != 0 { + return Err(Gemma3VlmPreparedError::ProjectedShape { + values: projected_image_features.len(), + hidden_size, + }); + } + + let mut image_positions = Vec::new(); + for (index, (&token_id, &mask)) in token_ids.iter().zip(attention_mask).enumerate() { + if mask != 0 && mask != 1 { + return Err(Gemma3VlmPreparedError::InvalidAttentionMask { index, value: mask }); + } + let is_padding = token_id == pad_token_id; + if is_padding != (mask == 0) { + return Err(Gemma3VlmPreparedError::PaddingMaskMismatch { + index, + token_id, + mask, + }); + } + if token_id == image_token_id { + image_positions.push(index); + } + } + let projected_tokens = projected_image_features.len() / hidden_size; + if image_positions.len() != projected_tokens { + return Err(Gemma3VlmPreparedError::PlaceholderCount { + positions: image_positions.len(), + projected_tokens, + }); + } + if (image_count == 0) != (projected_tokens == 0) + || (image_count > 0 && !projected_tokens.is_multiple_of(image_count)) + { + return Err(Gemma3VlmPreparedError::ImageCount { + images: image_count, + projected_tokens, + }); + } + + let normalizer = (hidden_size as f64).sqrt() as f32; + let mut merged = vec![0.0f32; expected_embeddings]; + let mut image_row = 0usize; + for (position, &token_id) in token_ids.iter().enumerate() { + let destination = position * hidden_size; + if token_id == pad_token_id { + continue; + } + if token_id == image_token_id { + let source = image_row * hidden_size; + merged[destination..destination + hidden_size] + .copy_from_slice(&projected_image_features[source..source + hidden_size]); + image_row += 1; + continue; + } + for offset in 0..hidden_size { + let value = raw_text_embeddings[destination + offset] * normalizer; + if !value.is_finite() { + return Err(Gemma3VlmPreparedError::NonFinite { + tensor: "post-scale text embeddings", + index: destination + offset, + }); + } + merged[destination + offset] = value; + } + } + + let bias_count = sequence_len + .checked_mul(sequence_len) + .ok_or(Gemma3VlmPreparedError::ShapeOverflow)?; + let mut attention_bias = vec![GEMMA3_VLM_MASKED_VALUE; bias_count]; + for query in 0..sequence_len { + for key in 0..sequence_len { + if attention_mask[query] == 1 && attention_mask[key] == 1 { + attention_bias[query * sequence_len + key] = 0.0; + } + } + } + + let embeddings = OwnedTensor::new( + f32_bytes(&merged), + PreparedTensorDType::Float32, + vec![1, sequence_len, hidden_size], + )?; + let attention_bias = PreparedAttentionBias { + tensor: OwnedTensor::new( + f32_bytes(&attention_bias), + PreparedTensorDType::Float32, + vec![1, 1, sequence_len, sequence_len], + )?, + causal: false, + }; + let modalities = (image_count > 0).then(|| PreparedModality { + family: "gemma3".to_string(), + item_count: image_count, + token_count: image_positions.len(), + }); + PreparedPrefill::new( + token_ids, + embeddings, + PreparedPositions::Sequential { + start: 0, + length: sequence_len, + }, + attention_bias, + modalities.into_iter().collect(), + ) + .map_err(Into::into) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn read_f32(tensor: &OwnedTensor) -> Vec { + tensor + .bytes + .chunks_exact(4) + .map(|chunk| f32::from_le_bytes(chunk.try_into().unwrap())) + .collect() + } + + fn compatibility_contract() -> Gemma3VlmCompatibilityContract { + Gemma3VlmCompatibilityContract::new( + 262_144, + 0, + 255_999, + 256_000, + 256, + 2_560, + Some(1_024), + 6, + 1_000_000.0, + Some(10_000.0), + "iree-vision-v3:fixture".to_string(), + ) + } + + #[test] + fn compatibility_identity_binds_every_cross_runtime_field() { + let base = compatibility_contract(); + let identity = base.stable_identity(); + let mut variants = Vec::new(); + + macro_rules! changed { + ($field:ident, $value:expr) => {{ + let mut variant = base.clone(); + variant.$field = $value; + variants.push(variant); + }}; + } + + changed!(prepared_prefill, "prepared-v2".to_string()); + changed!(mask_mode, "causal-mask".to_string()); + changed!(masked_value_bits, 0); + changed!(post_scale_policy, "double-scale".to_string()); + changed!(newline_wrapper_token_id, 109); + changed!(image_token_id, 262_145); + changed!(pad_token_id, 1); + changed!(boi_token_id, 255_998); + changed!(eoi_token_id, 256_001); + changed!(tokens_per_image, 255); + changed!(hidden_size, 2_561); + changed!(sliding_window, Some(2_048)); + changed!(sliding_pattern, 5); + changed!(global_rope_theta_bits, 1_000_001.0f64.to_bits()); + changed!(local_rope_base_bits, Some(10_001.0f64.to_bits())); + changed!(vision_identity, "iree-vision-v3:changed".to_string()); + + for variant in variants { + assert_ne!(variant.stable_identity(), identity); + } + } + + #[test] + fn text_only_runtime_identity_is_byte_for_byte_unchanged() { + let identity = "dense:ordinary-text-config".to_string(); + assert_eq!( + bind_gemma3_vlm_compatibility(identity.clone(), None), + identity + ); + assert_ne!( + bind_gemma3_vlm_compatibility(identity.clone(), Some(&compatibility_contract())), + identity + ); + } + + #[test] + fn exports_post_scale_text_unscaled_images_and_exact_bidirectional_mask() { + let prepared = prepare_gemma3_vlm_prefill( + vec![0, 9, 7, 0], + &[1.0, 2.0, 100.0, 200.0, 3.0, 4.0, 5.0, 6.0], + &[11.0, 13.0], + &[0, 1, 1, 0], + 2, + 8, + 0, + 9, + 1, + ) + .unwrap(); + + let sqrt_two = 2.0f32.sqrt(); + assert_eq!( + read_f32(&prepared.embeddings), + [ + 0.0, + 0.0, + 11.0, + 13.0, + 3.0 * sqrt_two, + 4.0 * sqrt_two, + 0.0, + 0.0 + ] + ); + let bias = read_f32(&prepared.attention_bias.tensor); + assert_eq!(bias.len(), 16); + assert_eq!(bias[1 * 4 + 1], 0.0); + assert_eq!(bias[1 * 4 + 2], 0.0, "valid future keys stay bidirectional"); + assert_eq!(bias[2 * 4 + 1], 0.0); + assert_eq!(bias[0], f32::MIN); + assert_eq!(bias[3 * 4 + 2], f32::MIN); + assert!(!prepared.attention_bias.causal); + assert_eq!(prepared.modalities[0].token_count, 1); + + let static_payload = crate::prepared::PreparedIreePrefill::prepare(&prepared, 2, 8) + .expect("Gemma3 mask materializes into the static XLA bucket"); + assert_eq!(static_payload.attention_bias.len(), 64); + assert_eq!(static_payload.attention_bias[1 * 8 + 2], 0.0); + assert_eq!(static_payload.attention_bias[1 * 8 + 7], f32::MIN); + assert_eq!(static_payload.attention_bias[7 * 8 + 1], f32::MIN); + } + + #[test] + fn rejects_padding_disagreement_placeholder_mismatch_and_double_scale_overflow() { + assert!(matches!( + prepare_gemma3_vlm_prefill(vec![0], &[1.0], &[], &[1], 1, 1, 0, 9, 0), + Err(Gemma3VlmPreparedError::PaddingMaskMismatch { .. }) + )); + assert!(matches!( + prepare_gemma3_vlm_prefill(vec![9], &[1.0], &[], &[1], 1, 1, 0, 9, 1), + Err(Gemma3VlmPreparedError::PlaceholderCount { .. }) + )); + assert!(matches!( + prepare_gemma3_vlm_prefill( + vec![7], + &[f32::MAX, 0.0, 0.0, 0.0], + &[], + &[1], + 4, + 1, + 0, + 9, + 0, + ), + Err(Gemma3VlmPreparedError::NonFinite { + tensor: "post-scale text embeddings", + .. + }) + )); + assert!(matches!( + prepare_gemma3_vlm_prefill(vec![9], &[1.0, 2.0], &[3.0], &[1], 2, 1, 0, 9, 1,), + Err(Gemma3VlmPreparedError::ProjectedShape { .. }) + )); + assert!(matches!( + prepare_gemma3_vlm_prefill(vec![9], &[1.0], &[3.0], &[1], 1, 1, 0, 9, 0,), + Err(Gemma3VlmPreparedError::ImageCount { .. }) + )); + } +} diff --git a/src/lib/mlxcel-xla/src/vision_runtime.rs b/src/lib/mlxcel-xla/src/vision_runtime.rs index b16bc9d62..eef7a8b77 100644 --- a/src/lib/mlxcel-xla/src/vision_runtime.rs +++ b/src/lib/mlxcel-xla/src/vision_runtime.rs @@ -38,7 +38,7 @@ use crate::aux::{ use crate::aux_manifest::{AuxiliaryArtifactContract, ensure_qualified_auxiliary_artifact}; #[cfg(feature = "diagnostics")] use crate::emitter::emit_vision_diagnostics; -use crate::emitter::{LlavaVisionConfig, VisionWeightSpec, emit_vision}; +use crate::emitter::{LlavaVisionConfig, VisionProjector, VisionWeightSpec, emit_vision}; use crate::iree::{cached_vmfb_path, compile_one_to, iree_compile_bin, target_flags}; use crate::weights::{bf16_to_f32, f16_to_f32, f32_le_to_f32}; @@ -265,8 +265,30 @@ impl VisionProcessorContract { config.image_size )); } - if object.get("resample").and_then(serde_json::Value::as_u64) != Some(3) { - return Err("preprocessor_config.json resample must be PIL bicubic (3)".to_string()); + let expected_resample = match config.projector { + VisionProjector::LlavaMlp => 3, + VisionProjector::Gemma3AvgPool { .. } => 2, + }; + if matches!(config.projector, VisionProjector::Gemma3AvgPool { .. }) + && object + .get("image_processor_type") + .and_then(serde_json::Value::as_str) + != Some("Gemma3ImageProcessor") + { + return Err( + "preprocessor_config.json image_processor_type must be Gemma3ImageProcessor" + .to_string(), + ); + } + if object.get("resample").and_then(serde_json::Value::as_u64) != Some(expected_resample) { + return Err(format!( + "preprocessor_config.json resample must be {} ({expected_resample})", + if expected_resample == 3 { + "PIL bicubic" + } else { + "PIL bilinear" + } + )); } let factor = finite_number(object, "rescale_factor")?; if factor <= 0.0 { @@ -296,44 +318,66 @@ impl VisionProcessorContract { let processor_object = processor .as_object() .ok_or_else(|| format!("{} must be a JSON object", processor_path.display()))?; - if processor_object - .get("patch_size") - .and_then(serde_json::Value::as_u64) - != Some(config.patch_size as u64) - { - return Err(format!( - "processor patch_size disagrees with vision patch_size={}", - config.patch_size - )); - } - let expected_strategy = if config.drop_first_token { - "default" + if matches!(config.projector, VisionProjector::LlavaMlp) { + if processor_object + .get("patch_size") + .and_then(serde_json::Value::as_u64) + != Some(config.patch_size as u64) + { + return Err(format!( + "processor patch_size disagrees with vision patch_size={}", + config.patch_size + )); + } + let expected_strategy = if config.drop_first_token { + "default" + } else { + "full" + }; + if processor_object + .get("vision_feature_select_strategy") + .and_then(serde_json::Value::as_str) + != Some(expected_strategy) + { + return Err(format!( + "processor vision_feature_select_strategy disagrees with config ({expected_strategy})" + )); + } + let expected_additional_tokens = usize::from(config.class_token) as u64; + if processor_object + .get("num_additional_image_tokens") + .and_then(serde_json::Value::as_u64) + != Some(expected_additional_tokens) + { + return Err(format!( + "processor num_additional_image_tokens disagrees with class-token contract ({expected_additional_tokens})" + )); + } } else { - "full" - }; - if processor_object - .get("vision_feature_select_strategy") - .and_then(serde_json::Value::as_str) - != Some(expected_strategy) - { - return Err(format!( - "processor vision_feature_select_strategy disagrees with config ({expected_strategy})" - )); - } - let expected_additional_tokens = usize::from(config.class_token) as u64; - if processor_object - .get("num_additional_image_tokens") - .and_then(serde_json::Value::as_u64) - != Some(expected_additional_tokens) - { - return Err(format!( - "processor num_additional_image_tokens disagrees with class-token contract ({expected_additional_tokens})" - )); + if processor_object + .get("processor_class") + .and_then(serde_json::Value::as_str) + != Some("Gemma3Processor") + { + return Err( + "processor_config.json processor_class must be Gemma3Processor".to_string(), + ); + } + if processor_object + .get("image_seq_length") + .and_then(serde_json::Value::as_u64) + != Some(config.image_tokens() as u64) + { + return Err(format!( + "processor image_seq_length disagrees with Gemma3 projector output {}", + config.image_tokens() + )); + } } Ok(Self { identity: format!( - "preprocessor={};processor={};resolved=size:{};crop:{};resample:bicubic;\ - rescale:{:016x};mean:{mean:?};std:{std:?}", + "preprocessor={};processor={};resolved=size:{};crop:{};\ + resample:{expected_resample};rescale:{:016x};mean:{mean:?};std:{std:?}", preprocessor, processor, config.image_size, @@ -663,7 +707,11 @@ impl IreeVisionProjector { pub fn load(model_dir: &Path, device: &str) -> Result { let config = LlavaVisionConfig::from_model_dir(model_dir)?; let mlir = emit_vision(&config); - let module = compile_and_load(model_dir, device, &config, &mlir, "llava-vision")?; + let cache_tag = match config.projector { + VisionProjector::LlavaMlp => "llava-vision", + VisionProjector::Gemma3AvgPool { .. } => "gemma3-vision", + }; + let module = compile_and_load(model_dir, device, &config, &mlir, cache_tag)?; Ok(Self { module, config }) } diff --git a/src/loading/mod.rs b/src/loading/mod.rs index fb2cab9d4..02c76e205 100644 --- a/src/loading/mod.rs +++ b/src/loading/mod.rs @@ -56,12 +56,16 @@ use self::special::try_load_special_model_from_weights; use self::vlm::*; // Re-exported at the crate root for the CLI's lazy `--output-audio` load. +#[cfg(feature = "xla-iree")] +pub(crate) use self::vlm::load_gemma3_iree_host_preprocessor; 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; pub use self::vlm::load_qwen3_omni_speech; +#[cfg(feature = "xla-reference-diagnostics")] +pub use self::vlm::run_gemma3_eager_mlx_iree_prepared_boundary; #[cfg(feature = "xla-iree")] pub(crate) use self::vlm::{ Phi4MMXlaVisionComponents, load_phi4mm_xla_media_components, load_phi4mm_xla_text_embeddings, diff --git a/src/loading/vlm.rs b/src/loading/vlm.rs index b87c6c0a2..acc186817 100644 --- a/src/loading/vlm.rs +++ b/src/loading/vlm.rs @@ -97,6 +97,10 @@ pub(crate) use deepseekocr::{ pub(crate) use dots_ocr::load_dots_ocr_vl; pub(crate) use ernie4_5_vl::load_ernie4_5_moe_vlm; pub(crate) use fastvlm::load_fastvlm_vlm; +#[cfg(feature = "xla-iree")] +pub(crate) use gemma::load_gemma3_iree_host_preprocessor; +#[cfg(feature = "xla-reference-diagnostics")] +pub use gemma::run_gemma3_eager_mlx_iree_prepared_boundary; pub(crate) use gemma::{load_gemma3_vlm, load_gemma3n_vlm, load_gemma4_vlm}; pub(crate) use gemma_unified::load_gemma4_unified; pub(crate) use granite_vision::load_granite_vision_vlm; diff --git a/src/loading/vlm_gemma.rs b/src/loading/vlm_gemma.rs index 91b3528f4..f64c0ea64 100644 --- a/src/loading/vlm_gemma.rs +++ b/src/loading/vlm_gemma.rs @@ -22,17 +22,22 @@ //! wrapper assembly out of the generic VLM router. use anyhow::Result; +#[cfg(feature = "xla-iree")] +use mlxcel_core::layers::UnifiedEmbedding; use mlxcel_core::weights::WeightMap; use serde_json::Value; use std::path::Path; use crate::LoadedModel; use crate::models; +#[cfg(feature = "xla-iree")] +use crate::multimodal::host_preprocessor::{Gemma3IreeHostPreprocessor, HostPreprocessorError}; use crate::vision; use super::{ - load_vlm_weights_common, parse_required_vlm_subconfig, parse_vlm_config, - read_optional_model_json, read_sanitized_vlm_config, strip_language_model_prefix, + load_vlm_weights_common, load_vlm_weights_common_filtered_canonical, + parse_required_vlm_subconfig, parse_vlm_config, read_optional_model_json, + read_sanitized_vlm_config, strip_language_model_prefix, }; struct Gemma3nMetadata { @@ -224,6 +229,113 @@ pub(crate) fn load_gemma3_vlm(model_path: &Path) -> Result { Ok(LoadedModel::Gemma3VLM(vlm)) } +#[cfg(feature = "xla-iree")] +fn is_gemma3_text_embedding_weight(name: &str) -> bool { + name.strip_prefix("language_model.") + .unwrap_or(name) + .starts_with("model.embed_tokens.") +} + +/// Load the bounded Gemma3 host producer paired with the resident IREE +/// SigLIP/averaging-projector module. Only the text embedding table is retained +/// in MLX; no language layer or host vision tower is constructed. +#[cfg(feature = "xla-iree")] +pub(crate) fn load_gemma3_iree_host_preprocessor( + model_path: &Path, + device: &str, +) -> Result { + use vision::config::VLMConfig; + use vision::processors::siglip::SigLipProcessor; + + let (_config_str, full_config) = read_sanitized_vlm_config(model_path) + .map_err(|error| HostPreprocessorError::InvalidConfig(error.to_string()))?; + if !full_config + .get("vision_config") + .is_some_and(Value::is_object) + || !full_config.get("text_config").is_some_and(Value::is_object) + { + return Err(HostPreprocessorError::FamilyMismatch { + actual: "unqualified gemma3 config without vision_config/text_config".to_string(), + }); + } + let vlm_config: VLMConfig = serde_json::from_value(full_config.clone()) + .map_err(|error| HostPreprocessorError::InvalidConfig(error.to_string()))?; + if vlm_config.model_type != "gemma3" + || vlm_config + .text_config + .get("model_type") + .and_then(Value::as_str) + != Some("gemma3_text") + { + return Err(HostPreprocessorError::FamilyMismatch { + actual: format!( + "model_type={:?}, text_model_type={:?}", + vlm_config.model_type, + vlm_config + .text_config + .get("model_type") + .and_then(Value::as_str) + ), + }); + } + let text_config: models::gemma3::ModelArgs = + serde_json::from_value(vlm_config.text_config.clone()) + .map_err(|error| HostPreprocessorError::InvalidConfig(error.to_string()))?; + + let weights = load_vlm_weights_common_filtered_canonical(model_path, |name| { + is_gemma3_text_embedding_weight(name) + }) + .map(strip_language_model_prefix) + .map_err(|error| HostPreprocessorError::WeightLoad(error.to_string()))?; + let quant_group_size = full_config + .get("quantization") + .and_then(|value| value.get("group_size")) + .and_then(Value::as_i64) + .unwrap_or(64) as i32; + let quant_bits = full_config + .get("quantization") + .and_then(|value| value.get("bits")) + .and_then(Value::as_i64) + .unwrap_or(4) as i32; + let text_embeddings = UnifiedEmbedding::from_weights( + &weights, + "model.embed_tokens", + quant_group_size, + quant_bits, + ) + .map_err(|error| { + HostPreprocessorError::WeightLoad(format!( + "missing or invalid Gemma3 text embedding table: {error}" + )) + })?; + let tokens_per_image = vlm_config.get_mm_tokens_per_image(); + let processor = SigLipProcessor::new(vlm_config.vision_config.image_size); + let projector = mlxcel_xla::IreeVisionProjector::load(model_path, device) + .map_err(HostPreprocessorError::Iree)?; + + Gemma3IreeHostPreprocessor::from_parts( + Box::new(processor), + text_embeddings, + projector, + vlm_config.image_token_index, + vlm_config.pad_token_id, + vlm_config.boi_token_index, + vlm_config.eoi_token_index, + tokens_per_image, + text_config.hidden_size, + vlm_config.vision_config.image_size, + text_config.max_position_embeddings, + device.to_string(), + ) +} + +#[cfg(all(feature = "xla-iree", any(test, feature = "xla-reference-diagnostics")))] +#[path = "vlm_gemma_xla_tests.rs"] +mod xla_tests; + +#[cfg(feature = "xla-reference-diagnostics")] +pub use xla_tests::reference_boundary::run_gemma3_eager_mlx_iree_prepared_boundary; + /// Load a Gemma3n VLM model. pub(crate) fn load_gemma3n_vlm(model_path: &Path) -> Result { use vision::encoders::gemma3n::load_gemma3n_vision; diff --git a/src/loading/vlm_gemma_xla_tests.rs b/src/loading/vlm_gemma_xla_tests.rs new file mode 100644 index 000000000..e0c4b4d9e --- /dev/null +++ b/src/loading/vlm_gemma_xla_tests.rs @@ -0,0 +1,1153 @@ +// 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 pinned_gemma3_projector_loads_and_returns_finite_features() { + let Ok(model) = std::env::var("MLXCEL_GEMMA3_FIXTURE") else { + return; + }; + let device = std::env::var("MLXCEL_XLA_DEVICE").unwrap_or_else(|_| "cuda".to_string()); + let mut projector = mlxcel_xla::IreeVisionProjector::load(Path::new(&model), &device) + .expect("load pinned Gemma3 IREE vision projector"); + assert_eq!(projector.input_shape(), [1, 3, 896, 896]); + assert_eq!(projector.output_shape(), [256, 2560]); + let pixels = vec![0.0; projector.input_shape().into_iter().product()]; + let projection = projector + .project(&pixels) + .expect("execute pinned Gemma3 IREE vision projector"); + assert_eq!(projection.shape, [256, 2560]); + assert!(projection.values.iter().all(|value| value.is_finite())); +} + +#[cfg(feature = "xla-reference-diagnostics")] +pub mod reference_boundary { + use std::fs; + use std::io::{self, Write}; + use std::path::Path; + #[cfg(test)] + use std::path::PathBuf; + use std::sync::mpsc::{self, RecvTimeoutError, Sender}; + use std::thread::{self, JoinHandle}; + use std::time::{Duration, Instant}; + + use mlxcel_core::layers::UnifiedEmbedding; + use mlxcel_core::session::{OwnedTensor, PreparedPrefill, PreparedTensorDType}; + use sha2::{Digest, Sha256}; + + use super::*; + use crate::multimodal::host_preprocessor::{ + Gemma3IreeHostPreprocessor, HostMultimodalPreprocessor, + }; + use crate::multimodal::vlm_prompt::{ImageTokenBlockInfo, apply_image_token_blocks}; + use crate::vision::config::VLMConfig; + use crate::vision::connectors::MultiModalConnector; + use crate::vision::connectors::avg_pool::AvgPoolProjector; + use crate::vision::encoders::VisionEncoder; + use crate::vision::encoders::siglip::SigLipVisionModel; + use crate::vision::processors::ImageProcessor; + use crate::vision::processors::siglip::SigLipProcessor; + + const PINNED_REVISION: &str = "93724907d4ed1745d2fe50baadf3b0b01a65abf2"; + const PINNED_CONFIG_SHA256: &str = + "5ccdde91da736e6e6f8f138268c620adcbf1219c973b884240b719a54122465b"; + const PINNED_PREPROCESSOR_SHA256: &str = + "f688d6bb20c5017601c4011de7ca656da8485b540b05013efdaf986c0fcc918d"; + const PINNED_PROCESSOR_SHA256: &str = + "3ffd5f11778dc73e2b69b3c00535e4121e1badf7018136263cd17b5b34fbaa53"; + const PINNED_IMAGE_SHA256: &str = + "5e7d54e8a7d21802378c87d2d70cf551e29739fe27599ddf129ebccdad1e6261"; + const BLOCK0_STAGES: [&str; 12] = [ + "siglip.block0.layer_norm1", + "siglip.block0.q_proj", + "siglip.block0.k_proj", + "siglip.block0.v_proj", + "siglip.block0.attention_context", + "siglip.block0.attention_output", + "siglip.block0.attention_residual", + "siglip.block0.layer_norm2", + "siglip.block0.mlp_fc1", + "siglip.block0.mlp_activation", + "siglip.block0.mlp_fc2", + "siglip.block0.output", + ]; + const PINNED_SIGLIP_BLOCK_COUNT: usize = 27; + + #[derive(Debug, Clone, Copy, PartialEq, Eq)] + struct SiglipHiddenCheckpoint { + label: &'static str, + block_index: usize, + hidden_state_index: usize, + } + + // Block 6 already closes the first balanced bisection interval. Capture + // blocks 1..=5 immediately before it so the next run identifies the exact + // first failing block inside the known 1..=6 range. + const SIGLIP_HIDDEN_REFINEMENT_CHECKPOINTS: [SiglipHiddenCheckpoint; 5] = [ + SiglipHiddenCheckpoint { + label: "siglip.hidden.block1.output", + block_index: 1, + hidden_state_index: 2, + }, + SiglipHiddenCheckpoint { + label: "siglip.hidden.block2.output", + block_index: 2, + hidden_state_index: 3, + }, + SiglipHiddenCheckpoint { + label: "siglip.hidden.block3.output", + block_index: 3, + hidden_state_index: 4, + }, + SiglipHiddenCheckpoint { + label: "siglip.hidden.block4.output", + block_index: 4, + hidden_state_index: 5, + }, + SiglipHiddenCheckpoint { + label: "siglip.hidden.block5.output", + block_index: 5, + hidden_state_index: 6, + }, + ]; + + // Block 0 already has full sub-stage coverage. These ordered checkpoints + // split blocks 1..=26 into balanced inclusive ranges 1..=6, 7..=13, + // 14..=20, and 21..=26 while reusing the existing final-state comparison. + const SIGLIP_HIDDEN_BISECTION_CHECKPOINTS: [SiglipHiddenCheckpoint; 4] = [ + SiglipHiddenCheckpoint { + label: "siglip.hidden.block6.output", + block_index: 6, + hidden_state_index: 7, + }, + SiglipHiddenCheckpoint { + label: "siglip.hidden.block13.output", + block_index: 13, + hidden_state_index: 14, + }, + SiglipHiddenCheckpoint { + label: "siglip.hidden.block20.output", + block_index: 20, + hidden_state_index: 21, + }, + SiglipHiddenCheckpoint { + label: "siglip.hidden.last_pre_layernorm", + block_index: 26, + hidden_state_index: 27, + }, + ]; + + fn siglip_hidden_checkpoints() -> impl Iterator { + SIGLIP_HIDDEN_REFINEMENT_CHECKPOINTS + .iter() + .chain(SIGLIP_HIDDEN_BISECTION_CHECKPOINTS.iter()) + .copied() + } + + #[derive(Debug, Clone, Copy)] + struct Tolerance { + atol: f64, + rtol: f64, + } + + const PIXEL_TOLERANCE: Tolerance = Tolerance { + atol: 1e-6, + rtol: 1e-6, + }; + // The pinned MLX checkpoint executes the eager vision path in BF16 while + // the qualified IREE graph widens immutable checkpoint weights to F32. + // These are the same BF16 stage envelopes used by the #863 vision gate. + const VISION_TOLERANCE: Tolerance = Tolerance { + atol: 8e-2, + rtol: 4e-2, + }; + const IREE_REPLAY_TOLERANCE: Tolerance = Tolerance { + atol: 1e-6, + rtol: 1e-6, + }; + const EXACT_TOLERANCE: Tolerance = Tolerance { + atol: 0.0, + rtol: 0.0, + }; + + #[derive(Debug, Clone, Copy, PartialEq)] + struct ComparisonStats { + max_absolute: f64, + max_relative: f64, + failures: usize, + non_finite_count: usize, + first_failure: Option, + } + + fn comparison_stats( + observed: &[f32], + reference: &[f32], + tolerance: Tolerance, + ) -> ComparisonStats { + assert_eq!(observed.len(), reference.len(), "comparison lengths differ"); + let mut stats = ComparisonStats { + max_absolute: 0.0, + max_relative: 0.0, + failures: 0, + non_finite_count: 0, + first_failure: None, + }; + for (index, (&observed, &reference)) in observed.iter().zip(reference).enumerate() { + if !observed.is_finite() || !reference.is_finite() { + stats.failures += 1; + stats.non_finite_count += 1; + stats.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); + stats.max_absolute = stats.max_absolute.max(absolute); + stats.max_relative = stats.max_relative.max(relative); + if absolute > tolerance.atol + tolerance.rtol * f64::from(reference.abs()) { + stats.failures += 1; + stats.first_failure.get_or_insert(index); + } + } + stats + } + + fn progress(stage: &str) { + eprintln!("[gemma3-vlm-boundary] {stage}"); + io::stderr().flush().expect("flush diagnostic progress"); + } + + fn create_after_diagnostic_iree_configuration( + configure: impl FnOnce() -> Result<(), String>, + create: impl FnOnce() -> Result, + ) -> Result { + configure()?; + create() + } + + struct ProgressHeartbeat { + stop: Option>, + worker: Option>, + } + + impl ProgressHeartbeat { + fn start(stage: &str) -> Self { + progress(stage); + let (stop, receiver) = mpsc::channel(); + let stage = stage.to_string(); + let started = Instant::now(); + let worker = thread::spawn(move || { + while let Err(RecvTimeoutError::Timeout) = + receiver.recv_timeout(Duration::from_secs(60)) + { + progress(&format!( + "heartbeat stage={stage} elapsed={}s", + started.elapsed().as_secs() + )); + } + }); + Self { + stop: Some(stop), + worker: Some(worker), + } + } + } + + impl Drop for ProgressHeartbeat { + fn drop(&mut self) { + if let Some(stop) = self.stop.take() { + let _ = stop.send(()); + } + if let Some(worker) = self.worker.take() { + let _ = worker.join(); + } + } + } + + fn compare_stage( + stage: &str, + observed: &[f32], + reference: &[f32], + tolerance: Tolerance, + first_divergence: &mut Option, + ) -> bool { + if observed.len() != reference.len() { + let detail = format!("{stage}: length {} != {}", observed.len(), reference.len()); + first_divergence.get_or_insert(detail.clone()); + eprintln!("[gemma3-vlm-boundary] stage={stage} status=FAIL {detail}"); + return false; + } + let stats = comparison_stats(observed, reference, tolerance); + let passed = stats.failures == 0; + let status = if passed { "PASS" } else { "FAIL" }; + eprintln!( + "[gemma3-vlm-boundary] stage={stage} status={status} elements={} \ + atol={:.3e} rtol={:.3e} max_abs={:.6e} max_rel={:.6e} \ + failures={} non_finite={} first_failure={:?}", + observed.len(), + tolerance.atol, + tolerance.rtol, + stats.max_absolute, + stats.max_relative, + stats.failures, + stats.non_finite_count, + stats.first_failure, + ); + io::stderr().flush().expect("flush diagnostic stage"); + if stats.failures != 0 { + first_divergence.get_or_insert_with(|| { + format!( + "{stage} at flat index {}", + stats + .first_failure + .expect("a failed comparison has an index") + ) + }); + } + passed + } + + fn first_failing_siglip_block_interval( + checkpoint_results: &[(SiglipHiddenCheckpoint, bool)], + ) -> Option<(usize, usize)> { + checkpoint_results + .iter() + .position(|(_, passed)| !passed) + .map(|failed| { + let first = failed + .checked_sub(1) + .and_then(|previous| checkpoint_results.get(previous)) + .map_or(1, |(checkpoint, _)| checkpoint.block_index + 1); + (first, checkpoint_results[failed].0.block_index) + }) + } + + fn first_failing_siglip_block( + checkpoint_results: &[(SiglipHiddenCheckpoint, bool)], + ) -> Option { + checkpoint_results + .iter() + .find_map(|(checkpoint, passed)| (!passed).then_some(checkpoint.block_index)) + } + + fn checkpoint_statuses(checkpoint_results: &[(SiglipHiddenCheckpoint, bool)]) -> String { + checkpoint_results + .iter() + .map(|(checkpoint, passed)| { + format!( + "{}:{}", + checkpoint.label, + if *passed { "PASS" } else { "FAIL" } + ) + }) + .collect::>() + .join(",") + } + + fn compare_siglip_hidden_checkpoints( + observed_hidden_states: &[Vec], + reference_checkpoints: &[Vec], + first_divergence: &mut Option, + ) { + assert_eq!( + observed_hidden_states.len(), + PINNED_SIGLIP_BLOCK_COUNT + 1, + "IREE SigLIP diagnostics must contain embedding plus every pinned block output" + ); + assert_eq!( + reference_checkpoints.len(), + SIGLIP_HIDDEN_REFINEMENT_CHECKPOINTS.len() + SIGLIP_HIDDEN_BISECTION_CHECKPOINTS.len(), + "eager SigLIP ordered checkpoint count drifted" + ); + let results = siglip_hidden_checkpoints() + .zip(reference_checkpoints) + .map(|(checkpoint, reference)| { + let observed = observed_hidden_states + .get(checkpoint.hidden_state_index) + .unwrap_or_else(|| { + panic!( + "IREE SigLIP diagnostics are missing {} at hidden-state index {}", + checkpoint.label, checkpoint.hidden_state_index + ) + }); + let passed = compare_stage( + checkpoint.label, + observed, + reference, + VISION_TOLERANCE, + first_divergence, + ); + (checkpoint, passed) + }) + .collect::>(); + let refinement_end = SIGLIP_HIDDEN_REFINEMENT_CHECKPOINTS.len() + 1; + let refinement_results = &results[..refinement_end]; + let refinement_statuses = checkpoint_statuses(refinement_results); + if let Some(block) = first_failing_siglip_block(refinement_results) { + eprintln!( + "[gemma3-vlm-boundary] stage=siglip.hidden.block1_to_6 status=FAIL \ + exact_first_failing_block={block} checkpoints=[{refinement_statuses}]" + ); + } else { + eprintln!( + "[gemma3-vlm-boundary] stage=siglip.hidden.block1_to_6 status=PASS \ + checked_block_range=1..=6 checkpoints=[{refinement_statuses}]" + ); + } + let bisection_results = &results[SIGLIP_HIDDEN_REFINEMENT_CHECKPOINTS.len()..]; + let bisection_statuses = checkpoint_statuses(bisection_results); + if let Some((first, last)) = first_failing_siglip_block_interval(bisection_results) { + eprintln!( + "[gemma3-vlm-boundary] stage=siglip.hidden.bisection status=FAIL \ + first_failing_block_range={first}..={last} checkpoints=[{bisection_statuses}]" + ); + } else { + eprintln!( + "[gemma3-vlm-boundary] stage=siglip.hidden.bisection status=PASS \ + checked_block_range=1..={} checkpoints=[{bisection_statuses}]", + PINNED_SIGLIP_BLOCK_COUNT - 1 + ); + } + io::stderr() + .flush() + .expect("flush SigLIP hidden checkpoint report"); + } + + fn sha256(path: &Path) -> String { + let bytes = + fs::read(path).unwrap_or_else(|error| panic!("read {}: {error}", path.display())); + format!("{:x}", Sha256::digest(bytes)) + } + + fn assert_sha256(path: &Path, expected: &str, label: &str) { + assert_eq!( + sha256(path), + expected, + "{label} differs from the pinned #869 fixture: {}", + path.display() + ); + } + + fn pinned_revision(model: &Path) -> String { + if let Ok(revision) = std::env::var("MLXCEL_GEMMA3_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_GEMMA3_REVISION", + metadata.display() + ) + }) + .lines() + .next() + .expect("Hugging Face metadata contains a revision") + .to_string() + } + + fn tensor_f32(tensor: &OwnedTensor, label: &str) -> Vec { + assert_eq!( + tensor.dtype, + PreparedTensorDType::Float32, + "{label} must be float32" + ); + tensor + .bytes + .chunks_exact(4) + .map(|bytes| f32::from_le_bytes(bytes.try_into().expect("four-byte f32 chunk"))) + .collect() + } + + fn mlx_f32(array: &mlxcel_core::MlxArray, label: &str) -> Vec { + let widened = mlxcel_core::astype(array, mlxcel_core::dtype::FLOAT32); + mlxcel_core::try_array_to_raw_bytes(&widened) + .unwrap_or_else(|error| panic!("export {label}: {error}")) + .chunks_exact(4) + .map(|bytes| f32::from_ne_bytes(bytes.try_into().expect("four-byte f32 chunk"))) + .collect() + } + + fn reference_weight(name: &str) -> bool { + let canonical = name.strip_prefix("language_model.").unwrap_or(name); + canonical.starts_with("vision_tower.") + || canonical.starts_with("multi_modal_projector.") + || canonical.starts_with("model.embed_tokens.") + } + + fn image_block_info(config: &VLMConfig, tokens_per_image: usize) -> ImageTokenBlockInfo { + ImageTokenBlockInfo { + use_boi_eoi: true, + image_token_id: config.image_token_index, + mm_tokens_per_image: tokens_per_image, + boi_token_id: config.boi_token_index, + eoi_token_id: config.eoi_token_index, + has_bos: true, + separator_token_id: None, + suffix_tokens: Vec::new(), + block_prefix_tokens: vec![108], + block_suffix_tokens: vec![108], + } + } + + fn image_rows(prepared: &PreparedPrefill, image_token_id: i32, hidden: usize) -> Vec { + let embeddings = tensor_f32(&prepared.embeddings, "prepared embeddings"); + prepared + .token_ids + .iter() + .enumerate() + .filter(|(_, token)| **token == image_token_id) + .flat_map(|(position, _)| { + embeddings[position * hidden..(position + 1) * hidden] + .iter() + .copied() + }) + .collect() + } + + fn expected_mask(token_ids: &[i32], pad_token_id: i32) -> Vec { + let mut expected = vec![f32::MIN; token_ids.len() * token_ids.len()]; + for (query, query_token) in token_ids.iter().enumerate() { + for (key, key_token) in token_ids.iter().enumerate() { + if *query_token != pad_token_id && *key_token != pad_token_id { + expected[query * token_ids.len() + key] = 0.0; + } + } + } + expected + } + + fn expected_post_scale( + token_ids: &[i32], + raw_text: &[f32], + projected: &[f32], + hidden: usize, + pad_token_id: i32, + image_token_id: i32, + ) -> Vec { + let normalizer = (hidden as f64).sqrt() as f32; + let mut expected = vec![0.0; token_ids.len() * hidden]; + let mut image_row = 0; + for (position, token) in token_ids.iter().enumerate() { + let destination = position * hidden; + if *token == pad_token_id { + continue; + } + if *token == image_token_id { + let source = image_row * hidden; + expected[destination..destination + hidden] + .copy_from_slice(&projected[source..source + hidden]); + image_row += 1; + } else { + for offset in 0..hidden { + expected[destination + offset] = raw_text[destination + offset] * normalizer; + } + } + } + expected + } + + #[cfg(test)] + #[test] + fn comparison_reports_the_first_failed_element() { + let stats = comparison_stats( + &[1.0, 2.2, f32::NAN], + &[1.0, 2.0, 3.0], + Tolerance { + atol: 0.01, + rtol: 0.01, + }, + ); + assert_eq!(stats.failures, 2); + assert_eq!(stats.non_finite_count, 1); + assert_eq!(stats.first_failure, Some(1)); + } + + #[cfg(test)] + #[test] + fn siglip_hidden_refinement_descriptors_are_exact_and_ordered() { + assert_eq!( + SIGLIP_HIDDEN_REFINEMENT_CHECKPOINTS, + [ + SiglipHiddenCheckpoint { + label: "siglip.hidden.block1.output", + block_index: 1, + hidden_state_index: 2, + }, + SiglipHiddenCheckpoint { + label: "siglip.hidden.block2.output", + block_index: 2, + hidden_state_index: 3, + }, + SiglipHiddenCheckpoint { + label: "siglip.hidden.block3.output", + block_index: 3, + hidden_state_index: 4, + }, + SiglipHiddenCheckpoint { + label: "siglip.hidden.block4.output", + block_index: 4, + hidden_state_index: 5, + }, + SiglipHiddenCheckpoint { + label: "siglip.hidden.block5.output", + block_index: 5, + hidden_state_index: 6, + }, + ] + ); + assert_eq!( + siglip_hidden_checkpoints() + .map(|checkpoint| checkpoint.block_index) + .collect::>(), + [1, 2, 3, 4, 5, 6, 13, 20, 26] + ); + } + + #[cfg(test)] + #[test] + fn siglip_hidden_refinement_identifies_each_exact_failed_block() { + let refinement = siglip_hidden_checkpoints() + .take(SIGLIP_HIDDEN_REFINEMENT_CHECKPOINTS.len() + 1) + .collect::>(); + assert_eq!(refinement.len(), 6); + + for failed in 0..refinement.len() { + let mut results = refinement + .iter() + .copied() + .map(|checkpoint| (checkpoint, true)) + .collect::>(); + results[failed].1 = false; + assert_eq!( + first_failing_siglip_block(&results), + Some(refinement[failed].block_index) + ); + } + } + + #[cfg(test)] + #[test] + fn siglip_hidden_bisection_descriptors_are_exact_and_balanced() { + assert_eq!( + SIGLIP_HIDDEN_BISECTION_CHECKPOINTS, + [ + SiglipHiddenCheckpoint { + label: "siglip.hidden.block6.output", + block_index: 6, + hidden_state_index: 7, + }, + SiglipHiddenCheckpoint { + label: "siglip.hidden.block13.output", + block_index: 13, + hidden_state_index: 14, + }, + SiglipHiddenCheckpoint { + label: "siglip.hidden.block20.output", + block_index: 20, + hidden_state_index: 21, + }, + SiglipHiddenCheckpoint { + label: "siglip.hidden.last_pre_layernorm", + block_index: 26, + hidden_state_index: 27, + }, + ] + ); + let mut first = 1; + let widths = SIGLIP_HIDDEN_BISECTION_CHECKPOINTS.map(|checkpoint| { + assert_eq!(checkpoint.hidden_state_index, checkpoint.block_index + 1); + let width = checkpoint.block_index - first + 1; + first = checkpoint.block_index + 1; + width + }); + assert_eq!(widths, [6, 7, 7, 6]); + assert_eq!(first, PINNED_SIGLIP_BLOCK_COUNT); + } + + #[cfg(test)] + #[test] + fn siglip_hidden_bisection_bounds_the_first_failed_checkpoint() { + let results = SIGLIP_HIDDEN_BISECTION_CHECKPOINTS.map(|checkpoint| (checkpoint, true)); + assert_eq!(first_failing_siglip_block_interval(&results), None); + + for (failed, expected) in [(0, (1, 6)), (1, (7, 13)), (2, (14, 20)), (3, (21, 26))] { + let mut results = + SIGLIP_HIDDEN_BISECTION_CHECKPOINTS.map(|checkpoint| (checkpoint, true)); + for (_, passed) in &mut results[failed..] { + *passed = false; + } + assert_eq!( + first_failing_siglip_block_interval(&results), + Some(expected) + ); + } + } + + /// Run the pinned mixed-runtime boundary gate for #869. + /// + /// This entry point loads only the eager MLX SigLIP/projector/text embedding + /// weights and the resident IREE vision projector. The caller selects the + /// MLX runtime at build time; the dedicated example requires CUDA while the + /// IREE side remains pinned to `local-task`. + /// + /// # Panics + /// + /// Panics if the pinned fixture identity or any ordered comparison differs. + pub fn run_gemma3_eager_mlx_iree_prepared_boundary( + model: &Path, + image_path: &Path, + device: &str, + ) { + assert_eq!( + device, "local-task", + "the pinned #869 mixed-runtime gate qualifies IREE local-task only" + ); + + let _heartbeat = ProgressHeartbeat::start("validate pinned checkpoint and image"); + assert_eq!( + pinned_revision(model), + PINNED_REVISION, + "checkpoint revision differs from pinned mlx-community/gemma-3-4b-it-4bit" + ); + assert_sha256( + &model.join("config.json"), + PINNED_CONFIG_SHA256, + "checkpoint config", + ); + assert_sha256( + &model.join("preprocessor_config.json"), + PINNED_PREPROCESSOR_SHA256, + "checkpoint image preprocessor", + ); + assert_sha256( + &model.join("processor_config.json"), + PINNED_PROCESSOR_SHA256, + "checkpoint processor", + ); + assert_sha256(image_path, PINNED_IMAGE_SHA256, "image"); + drop(_heartbeat); + + let _heartbeat = + ProgressHeartbeat::start("load filtered MLX SigLIP/projector/text embedding weights"); + let (_config_str, full_config) = + read_sanitized_vlm_config(model).expect("read pinned Gemma3 config"); + let config: VLMConfig = + serde_json::from_value(full_config.clone()).expect("parse pinned Gemma3 VLM config"); + let text_config: models::gemma3::ModelArgs = + serde_json::from_value(config.text_config.clone()) + .expect("parse pinned Gemma3 text config"); + let weights = load_vlm_weights_common_filtered_canonical(model, reference_weight) + .map(strip_language_model_prefix) + .expect("load only Gemma3 vision/projector/embedding weights"); + let quant_group_size = full_config + .pointer("/quantization/group_size") + .and_then(Value::as_i64) + .unwrap_or(64) as i32; + let quant_bits = full_config + .pointer("/quantization/bits") + .and_then(Value::as_i64) + .unwrap_or(4) as i32; + let text_embeddings = UnifiedEmbedding::from_weights( + &weights, + "model.embed_tokens", + quant_group_size, + quant_bits, + ) + .expect("load filtered Gemma3 text embedding table"); + let encoder = SigLipVisionModel::from_weights( + &weights, + &config.vision_config, + "vision_tower.vision_model", + ) + .expect("load filtered Gemma3 SigLIP tower"); + let tokens_per_image = config.get_mm_tokens_per_image(); + let connector = AvgPoolProjector::from_weights( + &weights, + "multi_modal_projector", + config.vision_config.hidden_size, + config.vision_config.image_size, + config.vision_config.patch_size, + tokens_per_image, + config.vision_config.layer_norm_eps, + ) + .expect("load filtered Gemma3 average-pool projector"); + let image = image::open(image_path) + .unwrap_or_else(|error| panic!("decode {}: {error}", image_path.display())); + let images = vec![image]; + drop(_heartbeat); + + progress("compare independently constructed processor pixels"); + let mlx_processor = SigLipProcessor::new(config.vision_config.image_size); + let iree_processor = SigLipProcessor::new(config.vision_config.image_size); + let mlx_pixels = mlx_processor.preprocess(&images); + let iree_pixels = iree_processor.preprocess(&images); + let mlx_pixel_values = mlx_f32(&mlx_pixels, "MLX processor pixels"); + let iree_pixel_values = mlx_f32(&iree_pixels, "IREE host processor pixels"); + let mut first_divergence = None; + compare_stage( + "processor.pixel_values", + &iree_pixel_values, + &mlx_pixel_values, + PIXEL_TOLERANCE, + &mut first_divergence, + ); + + progress("expand pinned padded image-token fixture and embed text"); + let mut logical_tokens = vec![ + config.pad_token_id, + 2, + config.boi_token_index, + 1, + config.pad_token_id, + ]; + apply_image_token_blocks( + &mut logical_tokens, + image_block_info(&config, tokens_per_image), + images.len(), + ) + .expect("expand Gemma3 image-token block"); + let input_ids = mlxcel_core::from_slice_i32( + &logical_tokens, + &[ + 1, + i32::try_from(logical_tokens.len()).expect("fixture sequence length fits i32"), + ], + ); + let raw_text_array = text_embeddings.forward(&input_ids); + let embed_dtype = mlxcel_core::array_dtype(&raw_text_array); + let raw_text = mlx_f32(&raw_text_array, "raw text embeddings"); + + let _heartbeat = + ProgressHeartbeat::start("capture eager MLX SigLIP hidden and projector stages"); + let mlx_vision_input = mlxcel_core::astype( + &mlxcel_core::transpose_axes(&mlx_pixels, &[0, 2, 3, 1]), + embed_dtype, + ); + let (mlx_selected, mlx_hidden, mlx_block0) = + encoder.forward_with_hidden_state_diagnostics(&mlx_vision_input); + let mlx_selected_values = + mlx_f32(&mlx_selected.hidden_states, "MLX selected vision features"); + let mlx_projected = connector.forward(&mlx_selected.hidden_states); + let mlx_projected_values = mlx_f32(&mlx_projected, "MLX projected image features"); + let mlx_hidden0 = mlx_f32(&mlx_hidden[0], "MLX SigLIP embedding output"); + assert_eq!( + mlx_hidden.len(), + PINNED_SIGLIP_BLOCK_COUNT + 1, + "eager SigLIP diagnostics must contain embedding plus every pinned block output" + ); + let mlx_hidden_checkpoints = siglip_hidden_checkpoints() + .map(|checkpoint| { + assert_eq!( + checkpoint.hidden_state_index, + checkpoint.block_index + 1, + "{} descriptor must map block output to embedding-prefixed hidden-state index", + checkpoint.label + ); + mlx_f32( + mlx_hidden + .get(checkpoint.hidden_state_index) + .unwrap_or_else(|| { + panic!( + "eager SigLIP diagnostics are missing {} at hidden-state index {}", + checkpoint.label, checkpoint.hidden_state_index + ) + }), + checkpoint.label, + ) + }) + .collect::>(); + let mlx_block0_values = mlx_block0 + .iter() + .map(|stage| mlx_f32(stage, "MLX SigLIP block 0 stage")) + .collect::>(); + assert_eq!(mlx_block0_values.len(), BLOCK0_STAGES.len()); + drop(_heartbeat); + + let _heartbeat = + ProgressHeartbeat::start("run IREE diagnostic SigLIP and average-pool projector"); + let mut diagnostic = create_after_diagnostic_iree_configuration( + mlxcel_xla::configure_diagnostic_local_task_threads, + || mlxcel_xla::IreeVisionDiagnosticProjector::load(model, device), + ) + .expect("configure and load Gemma3 IREE diagnostic projector"); + let iree = diagnostic + .project(&iree_pixel_values) + .expect("execute Gemma3 IREE diagnostic projector"); + compare_stage( + "siglip.hidden.embedding", + &iree.hidden_states[0], + &mlx_hidden0, + VISION_TOLERANCE, + &mut first_divergence, + ); + assert_eq!(iree.block0_states.len(), BLOCK0_STAGES.len()); + for ((stage, observed), reference) in BLOCK0_STAGES + .iter() + .zip(&iree.block0_states) + .zip(&mlx_block0_values) + { + compare_stage( + stage, + observed, + reference, + VISION_TOLERANCE, + &mut first_divergence, + ); + } + compare_siglip_hidden_checkpoints( + &iree.hidden_states, + &mlx_hidden_checkpoints, + &mut first_divergence, + ); + compare_stage( + "siglip.selected.post_layernorm", + &iree.selected_vision_features, + &mlx_selected_values, + VISION_TOLERANCE, + &mut first_divergence, + ); + compare_stage( + "projector.avg_pool_projection", + &iree.projected_image_features, + &mlx_projected_values, + VISION_TOLERANCE, + &mut first_divergence, + ); + drop(_heartbeat); + + let _heartbeat = ProgressHeartbeat::start( + "construct MLX-reference and production IREE prepared prefills", + ); + let attention_mask = logical_tokens + .iter() + .map(|token| i32::from(*token != config.pad_token_id)) + .collect::>(); + let mlx_prepared = mlxcel_xla::prepare_gemma3_vlm_prefill( + logical_tokens.clone(), + &raw_text, + &mlx_projected_values, + &attention_mask, + text_config.hidden_size, + text_config.max_position_embeddings, + config.pad_token_id, + config.image_token_index, + images.len(), + ) + .expect("construct eager MLX reference prepared prefill"); + let production = Gemma3IreeHostPreprocessor::load(model, device) + .expect("load production Gemma3 IREE host preprocessor") + .prepare( + &[ + config.pad_token_id, + 2, + config.boi_token_index, + 1, + config.pad_token_id, + ], + &images, + ) + .expect("construct production Gemma3 IREE prepared prefill"); + assert_eq!(production.token_ids, logical_tokens); + assert_eq!(production.positions, mlx_prepared.positions); + assert_eq!(production.modalities, mlx_prepared.modalities); + drop(_heartbeat); + + progress("compare final projected image rows and one-time scaling"); + let mlx_image_rows = image_rows( + &mlx_prepared, + config.image_token_index, + text_config.hidden_size, + ); + let iree_image_rows = image_rows( + &production, + config.image_token_index, + text_config.hidden_size, + ); + compare_stage( + "prepared.mlx_projected_image_rows", + &mlx_image_rows, + &mlx_projected_values, + EXACT_TOLERANCE, + &mut first_divergence, + ); + compare_stage( + "prepared.iree_projected_image_rows", + &iree_image_rows, + &iree.projected_image_features, + IREE_REPLAY_TOLERANCE, + &mut first_divergence, + ); + compare_stage( + "prepared.mlx_vs_iree_image_rows", + &iree_image_rows, + &mlx_image_rows, + VISION_TOLERANCE, + &mut first_divergence, + ); + let expected_mlx = expected_post_scale( + &logical_tokens, + &raw_text, + &mlx_projected_values, + text_config.hidden_size, + config.pad_token_id, + config.image_token_index, + ); + let mlx_embeddings = tensor_f32(&mlx_prepared.embeddings, "MLX prepared embeddings"); + compare_stage( + "prepared.text_sqrt_hidden_and_image_identity", + &mlx_embeddings, + &expected_mlx, + EXACT_TOLERANCE, + &mut first_divergence, + ); + let expected_iree = expected_post_scale( + &logical_tokens, + &raw_text, + &iree.projected_image_features, + text_config.hidden_size, + config.pad_token_id, + config.image_token_index, + ); + let production_embeddings = tensor_f32( + &production.embeddings, + "IREE production prepared embeddings", + ); + compare_stage( + "prepared.iree_text_sqrt_hidden_and_image_identity", + &production_embeddings, + &expected_iree, + IREE_REPLAY_TOLERANCE, + &mut first_divergence, + ); + + progress("compare exact additive bidirectional padding mask"); + assert!(!production.attention_bias.causal); + assert!(!mlx_prepared.attention_bias.causal); + assert_eq!( + production.attention_bias.tensor.shape, + vec![1, 1, logical_tokens.len(), logical_tokens.len()] + ); + assert_eq!( + production.attention_bias.tensor.dtype, + PreparedTensorDType::Float32 + ); + let expected_attention = expected_mask(&logical_tokens, config.pad_token_id); + let mlx_attention = tensor_f32(&mlx_prepared.attention_bias.tensor, "MLX attention bias"); + let iree_attention = tensor_f32(&production.attention_bias.tensor, "IREE attention bias"); + compare_stage( + "prepared.mask.mlx_cell_exact", + &mlx_attention, + &expected_attention, + EXACT_TOLERANCE, + &mut first_divergence, + ); + compare_stage( + "prepared.mask.iree_cell_exact", + &iree_attention, + &expected_attention, + EXACT_TOLERANCE, + &mut first_divergence, + ); + + if let Some(first_divergence) = first_divergence { + panic!("Gemma3 eager MLX/IREE first divergence: {first_divergence}"); + } + progress("PASS all pinned Gemma3 eager MLX/IREE boundary stages"); + } + + /// Retained libtest wrapper for developers who already use the historical + /// ignored gate. The dedicated example avoids linking the libtest harness. + #[cfg(test)] + #[test] + #[ignore = "requires pinned Gemma3 checkpoint, image, MLX CUDA, and IREE local-task"] + fn pinned_gemma3_eager_mlx_matches_iree_prepared_boundary() { + let model = PathBuf::from( + std::env::var("MLXCEL_GEMMA3_FIXTURE") + .expect("MLXCEL_GEMMA3_FIXTURE must name the pinned checkpoint"), + ); + let image_path = std::env::var("MLXCEL_GEMMA3_IMAGE") + .map(PathBuf::from) + .unwrap_or_else(|_| PathBuf::from("tests/fixtures/test_image.png")); + let device = + std::env::var("MLXCEL_XLA_DEVICE").unwrap_or_else(|_| "local-task".to_string()); + run_gemma3_eager_mlx_iree_prepared_boundary(&model, &image_path, &device); + } + + #[cfg(test)] + #[test] + fn diagnostic_runner_configures_threads_before_creating_iree() { + use std::cell::Cell; + + let configured = Cell::new(false); + let result = create_after_diagnostic_iree_configuration( + || { + configured.set(true); + Ok(()) + }, + || { + assert!(configured.get()); + Ok("created") + }, + ); + assert_eq!(result.as_deref(), Ok("created")); + } + + #[cfg(test)] + #[test] + fn diagnostic_runner_propagates_configuration_failure_before_creation() { + let created = std::cell::Cell::new(false); + let result = create_after_diagnostic_iree_configuration( + || Err("configuration failed".to_string()), + || { + created.set(true); + Ok(()) + }, + ); + assert_eq!(result, Err("configuration failed".to_string())); + assert!(!created.get()); + } + + #[cfg(test)] + #[test] + fn diagnostics_thread_overrides_stay_out_of_production_iree_paths() { + let production_rust = include_str!("../lib/mlxcel-xla/src/aux.rs"); + let production_aux_c = include_str!("../lib/mlxcel-xla/csrc/xla_aux.c"); + let production_iree_c = include_str!("../lib/mlxcel-xla/csrc/xla_iree.c"); + for source in [production_rust, production_aux_c, production_iree_c] { + assert!(!source.contains("task_topology_group_count")); + assert!(!source.contains("task_worker_stack_size")); + assert!(!source.contains("configure_diagnostic_local_task_threads")); + } + } + + #[cfg(test)] + #[test] + fn diagnostic_iree_thread_flags_pin_group_and_use_host_stack_default() { + let diagnostic_c = include_str!("../lib/mlxcel-xla/csrc/xla_diagnostic_flags.c"); + assert!(diagnostic_c.contains("--task_topology_group_count=1")); + assert!(diagnostic_c.contains("--task_worker_stack_size=0")); + assert!(diagnostic_c.contains("if (argc != 1)")); + } + + #[cfg(test)] + #[test] + fn native_diagnostic_thread_configuration_is_reusable() { + mlxcel_xla::configure_diagnostic_local_task_threads() + .expect("configure diagnostics-only IREE task threads"); + assert!(mlxcel_xla::diagnostic_local_task_threads_are_configured()); + mlxcel_xla::configure_diagnostic_local_task_threads() + .expect("reuse diagnostics-only IREE task thread configuration"); + } +} diff --git a/src/loading/vlm_qwen.rs b/src/loading/vlm_qwen.rs index 52971240c..1902a3b49 100644 --- a/src/loading/vlm_qwen.rs +++ b/src/loading/vlm_qwen.rs @@ -124,6 +124,14 @@ pub(crate) fn load_qwen2_vl_iree_host_preprocessor( .to_string(), }); } + if !full_config + .get("vision_config") + .is_some_and(serde_json::Value::is_object) + { + return Err(HostPreprocessorError::FamilyMismatch { + actual: "unqualified qwen2_vl config without vision_config".to_string(), + }); + } let mut vision_config: Qwen2VLVisionConfig = parse_required_vlm_subconfig(&full_config, "vision_config", "Qwen2VL vision config") .map_err(|error| HostPreprocessorError::InvalidConfig(error.to_string()))?; diff --git a/src/multimodal/host_preprocessor.rs b/src/multimodal/host_preprocessor.rs index bfad9cb9d..a18cefe3f 100644 --- a/src/multimodal/host_preprocessor.rs +++ b/src/multimodal/host_preprocessor.rs @@ -130,13 +130,13 @@ pub trait HostMultimodalPreprocessor { /// /// `Ok(None)` is the conservative result for text-only checkpoints and VLM /// families whose processor/position contract has not been qualified for XLA -/// yet. Once a checkpoint is identified as the supported LLaVA family, missing -/// or malformed processor/projector weights are startup errors rather than a -/// capability downgrade. +/// yet. Once a checkpoint is identified as a supported multimodal family, +/// missing or malformed processor/projector weights are startup errors rather +/// than a capability downgrade. /// /// # Errors /// -/// Returns a typed configuration or weight-loading error for a supported LLaVA +/// Returns a typed configuration or weight-loading error for a supported /// checkpoint that cannot construct its complete host preprocessor. pub fn load_xla_image_preprocessor( model_path: &Path, @@ -147,8 +147,8 @@ pub fn load_xla_image_preprocessor( model_path.display() )) })?; + let policy = XlaVisionBackendPolicy::from_env()?; if model_type == crate::models::ModelType::Qwen2VL { - let policy = XlaVisionBackendPolicy::from_env()?; if policy == XlaVisionBackendPolicy::Host { return Err(HostPreprocessorError::InvalidConfig( "Qwen2-VL XLA vision has no MLX fallback; MLXCEL_XLA_VISION_BACKEND=host is unsupported" @@ -159,7 +159,11 @@ pub fn load_xla_image_preprocessor( { let device = std::env::var("MLXCEL_XLA_DEVICE") .unwrap_or_else(|_| mlxcel_xla::default_device().to_string()); - let preprocessor = Qwen2VlIreeHostPreprocessor::load(model_path, &device)?; + let preprocessor = match Qwen2VlIreeHostPreprocessor::load(model_path, &device) { + Ok(preprocessor) => preprocessor, + Err(HostPreprocessorError::FamilyMismatch { .. }) => return Ok(None), + Err(error) => return Err(error), + }; tracing::info!( vision_backend = "iree", vision_device = %device, @@ -175,14 +179,54 @@ pub fn load_xla_image_preprocessor( )); } } + if model_type == crate::models::ModelType::Gemma3VLM { + return load_gemma3_image_preprocessor(model_path, policy); + } if model_type != crate::models::ModelType::LlavaVLM { return Ok(None); } - let policy = XlaVisionBackendPolicy::from_env()?; load_llava_image_preprocessor(model_path, policy) } +fn load_gemma3_image_preprocessor( + model_path: &Path, + policy: XlaVisionBackendPolicy, +) -> Result>, HostPreprocessorError> { + if policy == XlaVisionBackendPolicy::Host { + return Err(HostPreprocessorError::InvalidConfig( + "Gemma3 VLM requires MLXCEL_XLA_VISION_BACKEND=iree; the host LLaVA merge does not implement Gemma3 scaling or additive-mask semantics" + .to_string(), + )); + } + + #[cfg(feature = "xla-iree")] + { + let device = std::env::var("MLXCEL_XLA_DEVICE") + .unwrap_or_else(|_| mlxcel_xla::default_device().to_string()); + let preprocessor = match Gemma3IreeHostPreprocessor::load(model_path, &device) { + Ok(preprocessor) => preprocessor, + Err(HostPreprocessorError::FamilyMismatch { .. }) => return Ok(None), + Err(error) => return Err(error), + }; + tracing::info!( + vision_backend = "iree", + vision_device = %device, + vision_backend_policy = ?policy, + "OpenXLA Gemma3 multimodal vision backend selected" + ); + return Ok(Some(Box::new(preprocessor))); + } + + #[cfg(not(feature = "xla-iree"))] + { + let _ = model_path; + Err(HostPreprocessorError::InvalidConfig( + "Gemma3 VLM image input requires the xla-iree feature".to_string(), + )) + } +} + fn load_llava_host_preprocessor_boxed( model_path: &Path, ) -> Result>, HostPreprocessorError> { @@ -458,6 +502,251 @@ impl HostMultimodalPreprocessor for Qwen2VlIreeHostPreprocessor { } } +/// Gemma3 producer that retains only host image processing and the filtered +/// text embedding table. SigLIP and the averaging projector execute in IREE; +/// this owner then constructs Gemma3's post-scale embeddings and exact +/// bidirectional additive prefill mask. +#[cfg(feature = "xla-iree")] +pub struct Gemma3IreeHostPreprocessor { + processor: Box, + text_embeddings: UnifiedEmbedding, + projector: RefCell, + image_token_id: i32, + pad_token_id: i32, + boi_token_id: i32, + eoi_token_id: i32, + tokens_per_image: usize, + hidden_size: usize, + image_size: usize, + max_sequence_len: usize, + device: String, +} + +#[cfg(feature = "xla-iree")] +impl Gemma3IreeHostPreprocessor { + pub fn load(model_path: &Path, device: &str) -> Result { + crate::loading::load_gemma3_iree_host_preprocessor(model_path, device) + } + + #[allow(clippy::too_many_arguments)] + pub(crate) fn from_parts( + processor: Box, + text_embeddings: UnifiedEmbedding, + projector: mlxcel_xla::IreeVisionProjector, + image_token_id: i32, + pad_token_id: i32, + boi_token_id: i32, + eoi_token_id: i32, + tokens_per_image: usize, + hidden_size: usize, + image_size: usize, + max_sequence_len: usize, + device: String, + ) -> Result { + validate_image_preprocessor_dimensions( + tokens_per_image, + hidden_size, + image_size, + max_sequence_len, + )?; + if [image_token_id, pad_token_id, boi_token_id, eoi_token_id] + .windows(2) + .any(|pair| pair[0] == pair[1]) + || image_token_id == boi_token_id + || image_token_id == eoi_token_id + || pad_token_id == eoi_token_id + { + return Err(HostPreprocessorError::InvalidConfig( + "Gemma3 image, padding, BOI, and EOI token IDs must be distinct".to_string(), + )); + } + let expected_input = [1, 3, image_size, image_size]; + if projector.input_shape() != expected_input { + return Err(HostPreprocessorError::InvalidConfig(format!( + "Gemma3 IREE vision input shape {:?} does not match processor shape {expected_input:?}", + projector.input_shape() + ))); + } + let expected_output = [tokens_per_image, hidden_size]; + if projector.output_shape() != expected_output { + return Err(HostPreprocessorError::InvalidConfig(format!( + "Gemma3 IREE vision output shape {:?} does not match prepared-prefill shape {expected_output:?}", + projector.output_shape() + ))); + } + Ok(Self { + processor, + text_embeddings, + projector: RefCell::new(projector), + image_token_id, + pad_token_id, + boi_token_id, + eoi_token_id, + tokens_per_image, + hidden_size, + image_size, + max_sequence_len, + device, + }) + } + + fn token_block_info(&self) -> ImageTokenBlockInfo { + ImageTokenBlockInfo { + use_boi_eoi: true, + image_token_id: self.image_token_id, + mm_tokens_per_image: self.tokens_per_image, + boi_token_id: self.boi_token_id, + eoi_token_id: self.eoi_token_id, + has_bos: true, + separator_token_id: None, + suffix_tokens: Vec::new(), + block_prefix_tokens: vec![mlxcel_xla::GEMMA3_VLM_NEWLINE_WRAPPER_TOKEN_ID], + block_suffix_tokens: vec![mlxcel_xla::GEMMA3_VLM_NEWLINE_WRAPPER_TOKEN_ID], + } + } + + fn prepare_iree( + &self, + token_ids: &[i32], + images: &[DynamicImage], + ) -> Result { + let mut logical_tokens = token_ids.to_vec(); + apply_image_token_blocks(&mut logical_tokens, self.token_block_info(), images.len())?; + validate_sequence_capacity(logical_tokens.len(), self.max_sequence_len)?; + let attention_mask = logical_tokens + .iter() + .map(|token| i32::from(*token != self.pad_token_id)) + .collect::>(); + let input_ids = mlxcel_core::from_slice_i32( + &logical_tokens, + &[1, usize_to_i32(logical_tokens.len(), "sequence length")?], + ); + let text_embeddings = mlxcel_core::astype( + &self.text_embeddings.forward(&input_ids), + mlxcel_core::dtype::FLOAT32, + ); + validate_embedding_shape( + &mlxcel_core::array_shape(&text_embeddings), + logical_tokens.len(), + self.hidden_size, + "Gemma3 text embedding table", + )?; + let raw_text_embeddings = export_mlx_tensor(&text_embeddings, "Gemma3 text embeddings")?; + let raw_text_embeddings = raw_text_embeddings + .bytes + .chunks_exact(std::mem::size_of::()) + .map(|bytes| f32::from_ne_bytes(bytes.try_into().expect("four-byte f32 chunk"))) + .collect::>(); + + let mut projected_values = Vec::new(); + if !images.is_empty() { + let pixels = self.processor.preprocess(images); + validate_processor_shape( + &mlxcel_core::array_shape(&pixels), + images.len(), + self.image_size, + )?; + let pixels = export_mlx_tensor(&pixels, "Gemma3 processor pixel_values")?; + if pixels.dtype != PreparedTensorDType::Float32 { + return Err(HostPreprocessorError::InvalidConfig(format!( + "Gemma3 IREE vision requires float32 processor output, got {:?}", + pixels.dtype + ))); + } + let pixel_values = pixels + .bytes + .chunks_exact(std::mem::size_of::()) + .map(|bytes| f32::from_ne_bytes(bytes.try_into().expect("four-byte f32 chunk"))) + .collect::>(); + let pixels_per_image = 3usize + .checked_mul(self.image_size) + .and_then(|count| count.checked_mul(self.image_size)) + .ok_or(HostPreprocessorError::ShapeOverflow)?; + let projected_per_image = self + .tokens_per_image + .checked_mul(self.hidden_size) + .ok_or(HostPreprocessorError::ShapeOverflow)?; + projected_values.reserve( + images + .len() + .checked_mul(projected_per_image) + .ok_or(HostPreprocessorError::ShapeOverflow)?, + ); + let mut elapsed_seconds = 0.0; + let mut upload_bytes = 0usize; + let mut transfer_bytes = 0usize; + let mut projector = self.projector.try_borrow_mut().map_err(|_| { + HostPreprocessorError::Iree( + "concurrent/re-entrant Gemma3 IREE vision invocation is unsupported" + .to_string(), + ) + })?; + for image_pixels in pixel_values.chunks_exact(pixels_per_image) { + let projection = projector + .project(image_pixels) + .map_err(HostPreprocessorError::Iree)?; + if projection.shape != [self.tokens_per_image, self.hidden_size] { + return Err(HostPreprocessorError::ProjectedShape { + actual: projection + .shape + .into_iter() + .map(|dimension| i32::try_from(dimension).unwrap_or(i32::MAX)) + .collect(), + image_count: 1, + tokens_per_image: self.tokens_per_image, + hidden_size: self.hidden_size, + }); + } + elapsed_seconds += projection.metrics.elapsed_seconds; + upload_bytes = upload_bytes + .checked_add(projection.metrics.pixel_upload_bytes) + .ok_or(HostPreprocessorError::ShapeOverflow)?; + transfer_bytes = transfer_bytes + .checked_add(projection.metrics.projected_transfer_bytes) + .ok_or(HostPreprocessorError::ShapeOverflow)?; + projected_values.extend(projection.values); + } + tracing::info!( + vision_backend = "iree", + vision_device = %self.device, + image_count = images.len(), + pixel_upload_bytes = upload_bytes, + projected_transfer_bytes = transfer_bytes, + iree_vision_seconds = elapsed_seconds, + "OpenXLA Gemma3 vision projection completed" + ); + } + + mlxcel_xla::prepare_gemma3_vlm_prefill( + logical_tokens, + &raw_text_embeddings, + &projected_values, + &attention_mask, + self.hidden_size, + self.max_sequence_len, + self.pad_token_id, + self.image_token_id, + images.len(), + ) + .map_err(|error| HostPreprocessorError::Gemma3(error.to_string())) + } +} + +#[cfg(feature = "xla-iree")] +impl HostMultimodalPreprocessor for Gemma3IreeHostPreprocessor { + fn backend(&self) -> XlaVisionBackend { + XlaVisionBackend::Iree + } + + fn prepare( + &self, + token_ids: &[i32], + images: &[DynamicImage], + ) -> Result { + self.prepare_iree(token_ids, images) + } +} + #[cfg(feature = "xla-iree")] impl LlavaIreeHostPreprocessor { /// Load the host processor/text embedding table and resident IREE vision @@ -478,7 +767,7 @@ impl LlavaIreeHostPreprocessor { max_sequence_len: usize, device: String, ) -> Result { - validate_llava_preprocessor_dimensions( + validate_image_preprocessor_dimensions( tokens_per_image, hidden_size, image_size, @@ -740,7 +1029,7 @@ impl LlavaHostPreprocessor { image_size: usize, max_sequence_len: usize, ) -> Result { - validate_llava_preprocessor_dimensions( + validate_image_preprocessor_dimensions( tokens_per_image, hidden_size, image_size, @@ -939,7 +1228,7 @@ impl HostMultimodalPreprocessor for LlavaHostPreprocessor { } } -fn validate_llava_preprocessor_dimensions( +fn validate_image_preprocessor_dimensions( tokens_per_image: usize, hidden_size: usize, image_size: usize, @@ -947,22 +1236,22 @@ fn validate_llava_preprocessor_dimensions( ) -> Result<(), HostPreprocessorError> { if tokens_per_image == 0 { return Err(HostPreprocessorError::InvalidConfig( - "LLaVA mm_tokens_per_image must be greater than zero".to_string(), + "multimodal tokens_per_image must be greater than zero".to_string(), )); } if hidden_size == 0 { return Err(HostPreprocessorError::InvalidConfig( - "LLaVA text hidden_size must be greater than zero".to_string(), + "multimodal text hidden_size must be greater than zero".to_string(), )); } if image_size == 0 { return Err(HostPreprocessorError::InvalidConfig( - "LLaVA processor image_size must be greater than zero".to_string(), + "multimodal processor image_size must be greater than zero".to_string(), )); } if max_sequence_len == 0 { return Err(HostPreprocessorError::InvalidConfig( - "LLaVA max sequence length must be greater than zero".to_string(), + "multimodal max sequence length must be greater than zero".to_string(), )); } Ok(()) @@ -1065,14 +1354,16 @@ impl HostMultimodalPreprocessor for FakeHostMultimodalPreprocessor { pub enum HostPreprocessorError { #[error(transparent)] Placeholder(#[from] ImageTokenBlockError), - #[error("incompatible multimodal family: expected LLaVA, got {actual}")] + #[error("incompatible multimodal family: {actual}")] FamilyMismatch { actual: String }, - #[error("invalid LLaVA host-preprocessor config: {0}")] + #[error("invalid multimodal host-preprocessor config: {0}")] InvalidConfig(String), - #[error("failed to load LLaVA host-preprocessor weights: {0}")] + #[error("failed to load multimodal host-preprocessor weights: {0}")] WeightLoad(String), #[error("IREE vision backend failed: {0}")] Iree(String), + #[error("Gemma3 multimodal prefill failed: {0}")] + Gemma3(String), #[error( "processor output shape {actual:?} does not match decoded RGB batch [{image_count}, 3, {image_size}, {image_size}]" )] diff --git a/src/multimodal/host_preprocessor_tests.rs b/src/multimodal/host_preprocessor_tests.rs index a83e0f032..27a6648b6 100644 --- a/src/multimodal/host_preprocessor_tests.rs +++ b/src/multimodal/host_preprocessor_tests.rs @@ -67,7 +67,7 @@ fn iree_vision_contract_policy_is_explicit_and_strict() { #[test] fn xla_loader_keeps_text_and_unqualified_vlm_image_capability_false() { - for model_type in ["llama", "qwen2_vl"] { + for model_type in ["llama", "qwen2_vl", "gemma3"] { let model_dir = tempfile::tempdir().unwrap(); std::fs::write( model_dir.path().join("config.json"), diff --git a/src/multimodal/vlm_prompt.rs b/src/multimodal/vlm_prompt.rs index 84f5ecf15..247864499 100644 --- a/src/multimodal/vlm_prompt.rs +++ b/src/multimodal/vlm_prompt.rs @@ -168,6 +168,8 @@ pub fn apply_image_token_blocks( let block_len = info .mm_tokens_per_image .checked_add(wrapper_tokens) + .and_then(|tokens| tokens.checked_add(info.block_prefix_tokens.len())) + .and_then(|tokens| tokens.checked_add(info.block_suffix_tokens.len())) .ok_or(ImageTokenBlockError::CapacityOverflow)?; let image_token_capacity = block_len .checked_mul(num_images) @@ -182,6 +184,7 @@ pub fn apply_image_token_blocks( let mut image_tokens = Vec::with_capacity(image_token_capacity); for _ in 0..num_images { + image_tokens.extend_from_slice(&info.block_prefix_tokens); if info.use_boi_eoi { image_tokens.push(info.boi_token_id); } @@ -191,6 +194,7 @@ pub fn apply_image_token_blocks( if info.use_boi_eoi { image_tokens.push(info.eoi_token_id); } + image_tokens.extend_from_slice(&info.block_suffix_tokens); } if info.has_bos { diff --git a/src/multimodal/vlm_prompt_tests.rs b/src/multimodal/vlm_prompt_tests.rs index 65eb67c97..5434c2adc 100644 --- a/src/multimodal/vlm_prompt_tests.rs +++ b/src/multimodal/vlm_prompt_tests.rs @@ -171,6 +171,37 @@ fn apply_image_token_blocks_expands_with_block_prefix_and_suffix() { assert_eq!(prompt_tokens, vec![1, 108, 10, 99, 99, 99, 11, 108, 2]); } +#[test] +fn apply_image_token_blocks_inserts_with_block_prefix_and_suffix() { + // Gemma3Processor expands a bare image into + // "\n\n" + BOI + image tokens + EOI + "\n\n" even when the caller did not + // provide an existing BOI placeholder. + let info = ImageTokenBlockInfo { + use_boi_eoi: true, + image_token_id: 99, + mm_tokens_per_image: 3, + boi_token_id: 10, + eoi_token_id: 11, + has_bos: true, + separator_token_id: None, + suffix_tokens: Vec::new(), + block_prefix_tokens: vec![108], + block_suffix_tokens: vec![108], + }; + let mut prompt_tokens = vec![1, 2]; + + let stats = apply_image_token_blocks(&mut prompt_tokens, info, 1).unwrap(); + + assert_eq!( + stats, + Some(ImageTokenBlockStats { + action: ImageTokenBlockAction::Inserted { image_blocks: 1 }, + tokens_per_image: 3, + }) + ); + assert_eq!(prompt_tokens, vec![1, 108, 10, 99, 99, 99, 11, 108, 2]); +} + #[test] fn apply_image_token_blocks_rejects_media_cardinality_mismatch() { let info = ImageTokenBlockInfo { diff --git a/src/vision/encoders/mod.rs b/src/vision/encoders/mod.rs index a2dab9ff7..bd507f598 100644 --- a/src/vision/encoders/mod.rs +++ b/src/vision/encoders/mod.rs @@ -65,7 +65,7 @@ pub trait VisionEncoder { /// Run the exact encoder path while retaining ordered hidden states for a /// reference-oracle first-divergence report. - #[cfg(feature = "xla-diagnostics")] + #[cfg(any(feature = "xla-diagnostics", feature = "xla-reference-diagnostics"))] fn forward_with_hidden_state_diagnostics( &self, pixel_values: &MlxArray, diff --git a/src/vision/encoders/siglip.rs b/src/vision/encoders/siglip.rs index 86a6a4078..c16800365 100644 --- a/src/vision/encoders/siglip.rs +++ b/src/vision/encoders/siglip.rs @@ -644,7 +644,7 @@ impl VisionEncoder for SigLipVisionModel { self.forward_impl(pixel_values, false).0 } - #[cfg(feature = "xla-diagnostics")] + #[cfg(any(feature = "xla-diagnostics", feature = "xla-reference-diagnostics"))] fn forward_with_hidden_state_diagnostics( &self, pixel_values: &MlxArray,