From 05cc645bcb97fd6e1b03b31e8e9a07f77f55a29f Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 24 Jul 2026 19:32:40 +0900 Subject: [PATCH 01/14] feat(xla): add Youtu-VL vision and MLA inference Add checked static patch buckets, window/full attention isolation, 2D vision RoPE, the built-in merger, and resident IREE checkpoint loading for Youtu-VL. Map the dense DeepSeek-style MLA weights and compressed KV cache into token prefill, embedding prefill, single decode, and ragged continuous batching while retaining traditional interleaved RoPE. Wire filtered host preprocessing, placeholder selection, sequence export, and strict no-MLX-fallback capability routing; fix window-order restoration and preserve pre-allocation patch caps. Validated with xla-iree check, focused unit tests, clippy, pinned MLA IREE graph compilation, and one pinned checkpoint IREE vision execution. Closes #872 --- src/lib/mlxcel-xla/src/emitter/config.rs | 73 ++- src/lib/mlxcel-xla/src/emitter/mod.rs | 77 +++ src/lib/mlxcel-xla/src/emitter/model.rs | 590 +++++++++++++++-- src/lib/mlxcel-xla/src/emitter/youtu_vl.rs | 321 +++++++++ .../mlxcel-xla/src/emitter/youtu_vl_plan.rs | 615 ++++++++++++++++++ src/lib/mlxcel-xla/src/lib.rs | 6 + src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs | 16 +- src/lib/mlxcel-xla/src/weights.rs | 21 +- src/lib/mlxcel-xla/src/youtu_vl_runtime.rs | 356 ++++++++++ src/loading/mod.rs | 2 + src/loading/vlm.rs | 2 + src/loading/vlm_youtu_vl.rs | 120 +++- src/multimodal/host_preprocessor.rs | 244 ++++++- src/multimodal/host_preprocessor_export.rs | 61 ++ src/multimodal/host_preprocessor_tests.rs | 56 ++ src/vision/encoders/youtu_vl_window.rs | 22 +- src/vision/processors/youtu_vl.rs | 44 +- 17 files changed, 2566 insertions(+), 60 deletions(-) create mode 100644 src/lib/mlxcel-xla/src/emitter/youtu_vl.rs create mode 100644 src/lib/mlxcel-xla/src/emitter/youtu_vl_plan.rs create mode 100644 src/lib/mlxcel-xla/src/youtu_vl_runtime.rs diff --git a/src/lib/mlxcel-xla/src/emitter/config.rs b/src/lib/mlxcel-xla/src/emitter/config.rs index e5c450eaa..ff37d8c8d 100644 --- a/src/lib/mlxcel-xla/src/emitter/config.rs +++ b/src/lib/mlxcel-xla/src/emitter/config.rs @@ -113,6 +113,32 @@ pub struct MropeConfig { pub layout: MropeLayout, } +/// Multi-head latent-attention dimensions used by Youtu-VL's DeepSeek-V3 +/// language backbone. +/// +/// MLA caches one normalized KV latent plus the decoupled rotary key instead of +/// materializing one full K/V pair per attention head. These dimensions affect +/// graph signatures and resident cache shapes, so they are part of [`Config`] +/// and therefore the compiled-artifact identity. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub struct MlaConfig { + pub q_lora_rank: usize, + pub kv_lora_rank: usize, + pub qk_nope_head_dim: usize, + pub qk_rope_head_dim: usize, + pub v_head_dim: usize, +} + +impl MlaConfig { + pub fn q_head_dim(self) -> usize { + self.qk_nope_head_dim + self.qk_rope_head_dim + } + + pub fn cache_dim(self) -> usize { + self.kv_lora_rank + self.qk_rope_head_dim + } +} + /// The checkpoint tensor-naming scheme (issue #499). Almost every Llama-family /// checkpoint uses the standard HF layout /// (`model.layers.{i}.self_attn.q_proj.weight`, `model.embed_tokens.weight`, @@ -319,6 +345,9 @@ pub struct Config { /// Three-axis multimodal RoPE. `None` preserves the exact one-dimensional /// prefill/decode schemas and emitted modules. pub mrope: Option, + /// DeepSeek-V3 style dense MLA attention. `None` retains the ordinary + /// per-head K/V attention and cache contract. + pub mla: Option, /// Optional sparse, prefill-only residual additions injected immediately /// after selected transformer layers. Ordinary embeddings prefill and decode /// ignore this declaration and keep their existing schemas. @@ -441,6 +470,7 @@ impl Config { sliding_pattern: 2, use_rope_layers: None, mrope: None, + mla: None, deepstack: None, moe: None, weight_scheme: WeightScheme::Llama, @@ -577,7 +607,7 @@ impl Config { let hidden = u("hidden_size")?; let n_q = u("num_attention_heads")?; // head_dim is explicit in recent configs; otherwise it is hidden / heads. - let head_dim = ou("head_dim").unwrap_or(hidden / n_q.max(1)); + let mut head_dim = ou("head_dim").unwrap_or(hidden / n_q.max(1)); // ExaOne 3.x uses `num_layers` in place of `num_hidden_layers`. let n_layers = u_any(&["num_hidden_layers", "num_layers"])?; let deepstack = match ( @@ -671,6 +701,7 @@ impl Config { let mut logit_div: Option = None; let mut fused_qkv = false; let mut fused_gate_up = false; + let mut mla: Option = None; // Read a Gemma family's soft-caps + query scale (shared by gemma2 / gemma3). // Gemma3's soft-caps are null, so this yields `None` there. @@ -731,6 +762,32 @@ impl Config { one_plus: false, }); } + Some("youtu_vl") => { + let parsed = MlaConfig { + q_lora_rank: u("q_lora_rank")?, + kv_lora_rank: u("kv_lora_rank")?, + qk_nope_head_dim: u("qk_nope_head_dim")?, + qk_rope_head_dim: u("qk_rope_head_dim")?, + v_head_dim: u("v_head_dim")?, + }; + if parsed.q_lora_rank == 0 + || parsed.kv_lora_rank == 0 + || parsed.qk_nope_head_dim == 0 + || parsed.qk_rope_head_dim == 0 + || parsed.v_head_dim == 0 + { + return Err("Youtu-VL MLA dimensions must all be positive".to_string()); + } + if parsed.qk_rope_head_dim % 2 != 0 { + return Err( + "Youtu-VL qk_rope_head_dim must be even for interleaved RoPE".to_string(), + ); + } + mla = Some(parsed); + head_dim = parsed.q_head_dim(); + rope_interleaved = ob("rope_interleave").unwrap_or(true); + rotary_dim = Some(parsed.qk_rope_head_dim); + } Some("gemma") => { // Gemma1: Llama-shaped norm placement, but embedding scale, `(1+w)` // RMSNorm, and a GeGLU MLP. No soft-caps, sliding, or q/k norm. @@ -1371,6 +1428,7 @@ impl Config { sliding_pattern, use_rope_layers, mrope, + mla, deepstack, moe, weight_scheme, @@ -1426,6 +1484,7 @@ impl Config { && !self.fused_gate_up && !self.dense_mlp && self.moe.is_none() + && self.mla.is_none() } /// Whether the issue #572 f16-resident-weight path can apply: the standard @@ -1458,6 +1517,18 @@ impl Config { } } + /// Number of logical cache heads. MLA owns one compressed latent row per + /// token; ordinary attention retains the checkpoint's KV-head count. + pub fn cache_heads(&self) -> usize { + if self.mla.is_some() { 1 } else { self.n_kv } + } + + /// Elements in each logical cache row. MLA packs + /// `[kv_lora_rank, qk_rope_head_dim]`; ordinary attention stores one head. + pub fn cache_dim(&self) -> usize { + self.mla.map_or(self.head_dim, MlaConfig::cache_dim) + } + /// The RoPE rotation width: `rotary_dim` for a partial-RoPE arch (StableLM), /// else the full `head_dim` (Llama family). Always even. pub fn rotary_width(&self) -> usize { diff --git a/src/lib/mlxcel-xla/src/emitter/mod.rs b/src/lib/mlxcel-xla/src/emitter/mod.rs index 0065385c7..463d77d12 100644 --- a/src/lib/mlxcel-xla/src/emitter/mod.rs +++ b/src/lib/mlxcel-xla/src/emitter/mod.rs @@ -63,6 +63,8 @@ mod qwen2_vl; mod rope; mod vision; mod vision_config; +mod youtu_vl; +mod youtu_vl_plan; #[cfg(test)] mod context_tests; @@ -157,6 +159,12 @@ pub(crate) use vision::emit_vision; pub(crate) use vision::emit_vision_diagnostics; #[allow(unused_imports)] pub(crate) use vision_config::{LlavaVisionConfig, VisionActivation, VisionWeightSpec}; +#[allow(unused_imports)] +pub(crate) use youtu_vl::emit_youtu_vl; +#[allow(unused_imports)] +pub(crate) use youtu_vl_plan::{ + YoutuVlHostInputs, YoutuVlVisionConfig, prepare_youtu_vl_host_inputs, +}; #[cfg(test)] mod tests { @@ -1401,6 +1409,75 @@ mod tests { assert!(e.contains("MiniCPM3") && e.contains("MLA"), "got: {e}"); } + #[test] + fn youtu_vl_selects_dense_mla_and_interleaved_rope() { + let config = Config::from_json_str( + r#"{"model_type":"youtu_vl","hidden_size":64,"intermediate_size":128, + "num_hidden_layers":2,"num_attention_heads":4,"num_key_value_heads":4, + "vocab_size":128,"rms_norm_eps":1e-6,"rope_theta":500000, + "kv_lora_rank":16,"q_lora_rank":24,"qk_nope_head_dim":8, + "qk_rope_head_dim":4,"v_head_dim":8,"rope_interleave":true, + "tie_word_embeddings":true,"attention_bias":false,"mlp_bias":false}"#, + ) + .unwrap() + .with_context_capacity(16) + .unwrap(); + let mla = config.mla.expect("Youtu-VL must select MLA"); + assert_eq!(config.head_dim, 12); + assert_eq!(config.rotary_width(), 4); + assert_eq!(config.cache_heads(), 1); + assert_eq!(config.cache_dim(), 20); + assert!(config.rope_interleaved); + assert_eq!(mla.q_head_dim(), 12); + let prefill = emit_prefill(&config, false); + assert!(prefill.contains("['mla_q_a']")); + assert!(prefill.contains("['mla_kv_b']")); + assert!(!prefill.contains("['wq']")); + } + + #[test] + #[ignore = "requires YOUTU_VL_MODEL_DIR and IREE_CUDA_COMPILE"] + fn pinned_youtu_vl_mla_prefill_and_decode_compile_with_iree() { + let model_dir = + std::env::var("YOUTU_VL_MODEL_DIR").expect("set YOUTU_VL_MODEL_DIR to the checkpoint"); + let compiler = + std::env::var("IREE_CUDA_COMPILE").expect("set IREE_CUDA_COMPILE to iree-compile"); + let config = Config::from_json(std::path::Path::new(&model_dir)) + .unwrap() + .with_context_capacity(16) + .unwrap(); + let directory = std::env::temp_dir().join(format!( + "mlxcel-youtu-vl-mla-compile-{}", + std::process::id() + )); + std::fs::create_dir_all(&directory).unwrap(); + for (name, graph) in [ + ("prefill", emit_prefill(&config, false)), + ("decode", emit_decode(&config, false)), + ("ragged", emit_decode_ragged(&config, 2, false)), + ] { + let mlir_path = directory.join(format!("{name}.mlir")); + let vmfb_path = directory.join(format!("{name}.vmfb")); + std::fs::write(&mlir_path, graph).unwrap(); + let output = std::process::Command::new(&compiler) + .arg(&mlir_path) + .args([ + "--iree-hal-target-device=local", + "--iree-hal-local-target-device-backends=llvm-cpu", + ]) + .arg("-o") + .arg(&vmfb_path) + .output() + .unwrap(); + assert!( + output.status.success(), + "{name} iree-compile failed:\n{}", + String::from_utf8_lossy(&output.stderr) + ); + } + let _ = std::fs::remove_dir_all(directory); + } + /// The parallel-block archs (Cohere/Cohere2) carry a single `input_layernorm` /// per layer and NO `post_attention_layernorm`; the sequential archs carry both. /// Holds across every shared-core graph kind. diff --git a/src/lib/mlxcel-xla/src/emitter/model.rs b/src/lib/mlxcel-xla/src/emitter/model.rs index 16a748cc6..83115104b 100644 --- a/src/lib/mlxcel-xla/src/emitter/model.rs +++ b/src/lib/mlxcel-xla/src/emitter/model.rs @@ -221,10 +221,13 @@ struct LayerW { /// MLP up projection. `None` on a MoE layer (issue #500); `Some` for every dense /// layer, so the dense-MLP op sequence is unchanged. up: Option, - wk: Val, - wo: Val, - wq: Val, - wv: Val, + wk: Option, + wo: Option, + wq: Option, + wv: Option, + /// Youtu-VL's dense DeepSeek-V3 MLA projections. Mutually exclusive with + /// the ordinary q/k/v projection handles above. + mla: Option, bk: Option, bq: Option, bv: Option, @@ -257,6 +260,16 @@ struct LayerW { phi4_lora: Option, } +struct MlaLayerW { + q_a: Val, + q_a_norm: Val, + q_b: Val, + kv_a: Val, + kv_a_norm: Val, + kv_b: Val, + o: Val, +} + #[derive(Clone)] struct LoraPairW { a: Val, @@ -463,13 +476,58 @@ fn take_layer_weights( .then(|| take_arg(decls, idx, Ty::f32(vec![h]), p("in_ln"))); let post_ln = has_post.then(|| take_arg(decls, idx, Ty::f32(vec![h]), p("post_ln"))); let up = (!moe_layer).then(|| take_weight(b, decls, idx, c, inter, h, p("up"))); - let wk = take_weight(b, decls, idx, c, kv, h, p("wk")); - // o_proj maps `[n_q*head_dim]` -> `[hidden]`, so its weight is `[h, qd]` (HF's - // `[out, in]`). For Llama / Qwen2 `qd == h`, so this renders the same square - // type as before (byte-identical); Gemma is genuinely non-square. - let wo = take_weight(b, decls, idx, c, h, qd, p("wo")); - let wq = take_weight(b, decls, idx, c, qd, h, p("wq")); - let wv = take_weight(b, decls, idx, c, kv, h, p("wv")); + let (wk, wo, wq, wv, mla) = if let Some(m) = c.mla { + let q_a = take_weight(b, decls, idx, c, m.q_lora_rank, h, p("mla_q_a")); + let q_a_norm = take_arg(decls, idx, Ty::f32(vec![m.q_lora_rank]), p("mla_q_a_norm")); + let q_b = take_weight( + b, + decls, + idx, + c, + c.n_q * m.q_head_dim(), + m.q_lora_rank, + p("mla_q_b"), + ); + let kv_a = take_weight(b, decls, idx, c, m.cache_dim(), h, p("mla_kv_a")); + let kv_a_norm = take_arg( + decls, + idx, + Ty::f32(vec![m.kv_lora_rank]), + p("mla_kv_a_norm"), + ); + let kv_b = take_weight( + b, + decls, + idx, + c, + c.n_q * (m.qk_nope_head_dim + m.v_head_dim), + m.kv_lora_rank, + p("mla_kv_b"), + ); + let o = take_weight(b, decls, idx, c, h, c.n_q * m.v_head_dim, p("mla_o")); + ( + None, + None, + None, + None, + Some(MlaLayerW { + q_a, + q_a_norm, + q_b, + kv_a, + kv_a_norm, + kv_b, + o, + }), + ) + } else { + let wk = take_weight(b, decls, idx, c, kv, h, p("wk")); + // o_proj maps `[n_q*head_dim]` -> `[hidden]`, so its weight is `[h, qd]`. + let wo = take_weight(b, decls, idx, c, h, qd, p("wo")); + let wq = take_weight(b, decls, idx, c, qd, h, p("wq")); + let wv = take_weight(b, decls, idx, c, kv, h, p("wv")); + (Some(wk), Some(wo), Some(wq), Some(wv), None) + }; let (bk, bq, bv) = if c.qkv_bias { let bk = take_arg(decls, idx, Ty::f32(vec![kv]), p("bk")); let bq = take_arg(decls, idx, Ty::f32(vec![qd]), p("bq")); @@ -592,6 +650,7 @@ fn take_layer_weights( wo, wq, wv, + mla, bk, bq, bv, @@ -727,13 +786,23 @@ fn build_arg_schema(b: &mut Builder, c: &Config) -> (Vec, Args) { let kcache = take_arg( &mut decls, &mut idx, - Ty::f32(vec![c.n_layers, c.context_capacity, c.n_kv, c.head_dim]), + Ty::f32(vec![ + c.n_layers, + c.context_capacity, + c.cache_heads(), + c.cache_dim(), + ]), "kcache".into(), ); let vcache = take_arg( &mut decls, &mut idx, - Ty::f32(vec![c.n_layers, c.context_capacity, c.n_kv, c.head_dim]), + Ty::f32(vec![ + c.n_layers, + c.context_capacity, + c.cache_heads(), + c.cache_dim(), + ]), "vcache".into(), ); @@ -1530,7 +1599,7 @@ impl AttnLayout { .phi4_lora .as_ref() .map(|lora| self.modal_lora_delta(b, hn, &lora.qkv)); - let q = b.linear(hn, &lw.wq); + let q = b.linear(hn, lw.wq.as_ref().expect("ordinary q projection")); let q = match &lora { Some(delta) => { let delta = b.slice(delta, &[(0, nq * d)]); @@ -1540,7 +1609,7 @@ impl AttnLayout { }; let q = add_proj_bias(b, q, &lw.bq); let q = b.reshape(&q, vec![nq, d]); - let kk = b.linear(hn, &lw.wk); + let kk = b.linear(hn, lw.wk.as_ref().expect("ordinary k projection")); let kk = match &lora { Some(delta) => { let delta = b.slice(delta, &[(nq * d, nq * d + nkv * d)]); @@ -1550,7 +1619,7 @@ impl AttnLayout { }; let kk = add_proj_bias(b, kk, &lw.bk); let kk = b.reshape(&kk, vec![nkv, d]); - let vv = b.linear(hn, &lw.wv); + let vv = b.linear(hn, lw.wv.as_ref().expect("ordinary v projection")); let vv = match &lora { Some(delta) => { let delta = b.slice(delta, &[(nq * d + nkv * d, nq * d + 2 * nkv * d)]); @@ -1582,7 +1651,7 @@ impl AttnLayout { .phi4_lora .as_ref() .map(|lora| self.modal_lora_delta(b, hn, &lora.qkv)); - let q = b.linear_seq(hn, &lw.wq); + let q = b.linear_seq(hn, lw.wq.as_ref().expect("ordinary q projection")); let q = match &lora { Some(delta) => { let delta = b.slice(delta, &[(0, n), (0, nq * d)]); @@ -1592,7 +1661,7 @@ impl AttnLayout { }; let q = add_proj_bias_seq(b, q, &lw.bq, n, nq * d); let q = b.reshape(&q, vec![n, nq, d]); - let kk = b.linear_seq(hn, &lw.wk); + let kk = b.linear_seq(hn, lw.wk.as_ref().expect("ordinary k projection")); let kk = match &lora { Some(delta) => { let delta = b.slice(delta, &[(0, n), (nq * d, nq * d + nkv * d)]); @@ -1602,7 +1671,7 @@ impl AttnLayout { }; let kk = add_proj_bias_seq(b, kk, &lw.bk, n, nkv * d); let kk = b.reshape(&kk, vec![n, nkv, d]); - let vv = b.linear_seq(hn, &lw.wv); + let vv = b.linear_seq(hn, lw.wv.as_ref().expect("ordinary v projection")); let vv = match &lora { Some(delta) => { let delta = b.slice(delta, &[(0, n), (nq * d + nkv * d, nq * d + 2 * nkv * d)]); @@ -1876,8 +1945,10 @@ impl AttnLayout { /// family, byte-identical). fn o_proj(&self, b: &mut Builder, o: &Val, lw: &LayerW) -> Val { let out = match self { - AttnLayout::Single { .. } => b.linear(o, &lw.wo), - _ => b.linear_seq(o, &lw.wo), + AttnLayout::Single { .. } => { + b.linear(o, lw.wo.as_ref().expect("ordinary output projection")) + } + _ => b.linear_seq(o, lw.wo.as_ref().expect("ordinary output projection")), }; let out = match &lw.phi4_lora { Some(lora) => { @@ -1933,6 +2004,440 @@ fn apply_scale_and_softcap(b: &mut Builder, c: &Config, k: &Consts, scores: Val) } } +fn mla_rms_norm( + b: &mut Builder, + x: &Val, + weight: &Val, + eps: &Val, + rows: Option, + width: usize, +) -> Val { + let width_f = b.const_f32(width as f32); + let zero = b.const_f32(0.0); + let square = b.multiply(x, x); + match rows { + None => { + let sum = b.reduce_add(&square, 0, &zero); + let mean = b.divide(&sum, &width_f); + let variance = b.add(&mean, eps); + let scale = b.rsqrt(&variance); + let scale = b.broadcast(&scale, &[], vec![width]); + let normalized = b.multiply(x, &scale); + b.multiply(&normalized, weight) + } + Some(rows) => { + let sum = b.reduce_add(&square, 1, &zero); + let width_f = b.broadcast(&width_f, &[], vec![rows]); + let mean = b.divide(&sum, &width_f); + let eps = b.broadcast(eps, &[], vec![rows]); + let variance = b.add(&mean, &eps); + let scale = b.rsqrt(&variance); + let scale = b.broadcast(&scale, &[0], vec![rows, width]); + let weight = b.broadcast(weight, &[1], vec![rows, width]); + let normalized = b.multiply(x, &scale); + b.multiply(&normalized, &weight) + } + } +} + +fn mla_expand_latent( + b: &mut Builder, + latent: &Val, + kv_b: &Val, + rows: Option, + c: &Config, +) -> (Val, Val) { + let m = c.mla.expect("MLA config"); + let weights = b.reshape( + kv_b, + vec![c.n_q, m.qk_nope_head_dim + m.v_head_dim, m.kv_lora_rank], + ); + let key_weights = b.slice( + &weights, + &[(0, c.n_q), (0, m.qk_nope_head_dim), (0, m.kv_lora_rank)], + ); + let value_weights = b.slice( + &weights, + &[ + (0, c.n_q), + (m.qk_nope_head_dim, m.qk_nope_head_dim + m.v_head_dim), + (0, m.kv_lora_rank), + ], + ); + match rows { + None => { + let key = b.dot_general( + latent, + &key_weights, + &[], + &[], + &[0], + &[2], + vec![c.n_q, m.qk_nope_head_dim], + ); + let value = b.dot_general( + latent, + &value_weights, + &[], + &[], + &[0], + &[2], + vec![c.n_q, m.v_head_dim], + ); + (key, value) + } + Some(rows) => { + let key = b.dot_general( + latent, + &key_weights, + &[], + &[], + &[latent.ty.shape.len() - 1], + &[2], + if latent.ty.shape.len() == 2 { + vec![rows, c.n_q, m.qk_nope_head_dim] + } else { + vec![latent.ty.shape[0], rows, c.n_q, m.qk_nope_head_dim] + }, + ); + let value = b.dot_general( + latent, + &value_weights, + &[], + &[], + &[latent.ty.shape.len() - 1], + &[2], + if latent.ty.shape.len() == 2 { + vec![rows, c.n_q, m.v_head_dim] + } else { + vec![latent.ty.shape[0], rows, c.n_q, m.v_head_dim] + }, + ); + (key, value) + } + } +} + +#[allow(clippy::too_many_arguments)] +fn emit_mla_attention( + b: &mut Builder, + c: &Config, + k: &Consts, + lw: &LayerW, + li: usize, + x: &Val, + layout: &AttnLayout, + kcache: &mut Val, + vcache: &mut Val, +) -> (Val, Val) { + let m = c.mla.expect("MLA config"); + let mw = lw.mla.as_ref().expect("MLA layer weights"); + let hn = arch_norm(b, c, k, layout, x, &lw.in_ln, lw.in_ln_bias.as_ref()); + let rows = match layout { + AttnLayout::Single { .. } => None, + AttnLayout::Ragged { bsz, .. } => Some(*bsz), + AttnLayout::Prefill { lp, .. } => Some(*lp), + }; + let q_latent = match rows { + None => b.linear(&hn, &mw.q_a), + Some(_) => b.linear_seq(&hn, &mw.q_a), + }; + let q_latent = mla_rms_norm(b, &q_latent, &mw.q_a_norm, &k.eps, rows, m.q_lora_rank); + let q = match rows { + None => b.linear(&q_latent, &mw.q_b), + Some(_) => b.linear_seq(&q_latent, &mw.q_b), + }; + let q = match rows { + None => b.reshape(&q, vec![c.n_q, m.q_head_dim()]), + Some(n) => b.reshape(&q, vec![n, c.n_q, m.q_head_dim()]), + }; + let (q_nope, q_pe) = match rows { + None => ( + b.slice(&q, &[(0, c.n_q), (0, m.qk_nope_head_dim)]), + b.slice(&q, &[(0, c.n_q), (m.qk_nope_head_dim, m.q_head_dim())]), + ), + Some(n) => ( + b.slice(&q, &[(0, n), (0, c.n_q), (0, m.qk_nope_head_dim)]), + b.slice( + &q, + &[(0, n), (0, c.n_q), (m.qk_nope_head_dim, m.q_head_dim())], + ), + ), + }; + let kv = match rows { + None => b.linear(&hn, &mw.kv_a), + Some(_) => b.linear_seq(&hn, &mw.kv_a), + }; + let (latent, k_pe) = match rows { + None => ( + b.slice(&kv, &[(0, m.kv_lora_rank)]), + b.slice(&kv, &[(m.kv_lora_rank, m.cache_dim())]), + ), + Some(n) => ( + b.slice(&kv, &[(0, n), (0, m.kv_lora_rank)]), + b.slice(&kv, &[(0, n), (m.kv_lora_rank, m.cache_dim())]), + ), + }; + let latent = mla_rms_norm(b, &latent, &mw.kv_a_norm, &k.eps, rows, m.kv_lora_rank); + let (q_pe, k_pe) = match layout { + AttnLayout::Single { cos, sin, .. } => { + let q_pe = rotate(b, c, &q_pe, cos, sin, c.n_q, m.qk_rope_head_dim); + let k_pe = b.reshape(&k_pe, vec![1, m.qk_rope_head_dim]); + let k_pe = rotate(b, c, &k_pe, cos, sin, 1, m.qk_rope_head_dim); + (q_pe, b.reshape(&k_pe, vec![m.qk_rope_head_dim])) + } + AttnLayout::Ragged { bsz, cos, sin, .. } => { + let q_pe = rotate_seq(b, c, &q_pe, cos, sin, *bsz, c.n_q, m.qk_rope_head_dim); + let k_pe = b.reshape(&k_pe, vec![*bsz, 1, m.qk_rope_head_dim]); + let k_pe = rotate_seq(b, c, &k_pe, cos, sin, *bsz, 1, m.qk_rope_head_dim); + (q_pe, b.reshape(&k_pe, vec![*bsz, m.qk_rope_head_dim])) + } + AttnLayout::Prefill { lp, cos, sin, .. } => { + let q_pe = rotate_seq(b, c, &q_pe, cos, sin, *lp, c.n_q, m.qk_rope_head_dim); + let k_pe = b.reshape(&k_pe, vec![*lp, 1, m.qk_rope_head_dim]); + let k_pe = rotate_seq(b, c, &k_pe, cos, sin, *lp, 1, m.qk_rope_head_dim); + (q_pe, b.reshape(&k_pe, vec![*lp, m.qk_rope_head_dim])) + } + }; + + let cache_width = m.cache_dim(); + let (latent_slab, pe_slab) = match layout { + AttnLayout::Single { cache_len, .. } => { + let zeros_pe = b.broadcast(&k.zero, &[], vec![m.qk_rope_head_dim]); + let latent_row = b.concatenate(&latent, &zeros_pe, 0); + let zeros_latent = b.broadcast(&k.zero, &[], vec![m.kv_lora_rank]); + let pe_row = b.concatenate(&zeros_latent, &k_pe, 0); + let latent_update = b.reshape(&latent_row, vec![1, 1, 1, cache_width]); + *kcache = b.dynamic_update_slice( + kcache, + &latent_update, + &[&k.layer_idx[li], cache_len, &k.c0, &k.c0], + ); + let pe_update = b.reshape(&pe_row, vec![1, 1, 1, cache_width]); + *vcache = b.dynamic_update_slice( + vcache, + &pe_update, + &[&k.layer_idx[li], cache_len, &k.c0, &k.c0], + ); + let latent_slab = b.slice( + kcache, + &[ + (li, li + 1), + (0, c.context_capacity), + (0, 1), + (0, m.kv_lora_rank), + ], + ); + let pe_slab = b.slice( + vcache, + &[ + (li, li + 1), + (0, c.context_capacity), + (0, 1), + (m.kv_lora_rank, cache_width), + ], + ); + ( + b.reshape(&latent_slab, vec![c.context_capacity, m.kv_lora_rank]), + b.reshape(&pe_slab, vec![c.context_capacity, m.qk_rope_head_dim]), + ) + } + AttnLayout::Ragged { + bsz, pos, row_idx, .. + } => { + let zeros_pe = b.broadcast(&k.zero, &[], vec![*bsz, m.qk_rope_head_dim]); + let latent_rows = b.concatenate(&latent, &zeros_pe, 1); + let zeros_latent = b.broadcast(&k.zero, &[], vec![*bsz, m.kv_lora_rank]); + let pe_rows = b.concatenate(&zeros_latent, &k_pe, 1); + for r in 0..*bsz { + let pos_slice = b.slice(pos, &[(r, r + 1)]); + let pos_r = b.reshape(&pos_slice, vec![]); + let latent_row = b.slice(&latent_rows, &[(r, r + 1), (0, cache_width)]); + let latent_update = b.reshape(&latent_row, vec![1, 1, 1, 1, cache_width]); + *kcache = b.dynamic_update_slice( + kcache, + &latent_update, + &[&row_idx[r], &k.layer_idx[li], &pos_r, &k.c0, &k.c0], + ); + let pe_row = b.slice(&pe_rows, &[(r, r + 1), (0, cache_width)]); + let pe_update = b.reshape(&pe_row, vec![1, 1, 1, 1, cache_width]); + *vcache = b.dynamic_update_slice( + vcache, + &pe_update, + &[&row_idx[r], &k.layer_idx[li], &pos_r, &k.c0, &k.c0], + ); + } + let latent_slab = b.slice( + kcache, + &[ + (0, *bsz), + (li, li + 1), + (0, c.context_capacity), + (0, 1), + (0, m.kv_lora_rank), + ], + ); + let pe_slab = b.slice( + vcache, + &[ + (0, *bsz), + (li, li + 1), + (0, c.context_capacity), + (0, 1), + (m.kv_lora_rank, cache_width), + ], + ); + ( + b.reshape(&latent_slab, vec![*bsz, c.context_capacity, m.kv_lora_rank]), + b.reshape(&pe_slab, vec![*bsz, c.context_capacity, m.qk_rope_head_dim]), + ) + } + AttnLayout::Prefill { lp, .. } => { + let zeros_pe = b.broadcast(&k.zero, &[], vec![*lp, m.qk_rope_head_dim]); + let latent_rows = b.concatenate(&latent, &zeros_pe, 1); + let zeros_latent = b.broadcast(&k.zero, &[], vec![*lp, m.kv_lora_rank]); + let pe_rows = b.concatenate(&zeros_latent, &k_pe, 1); + let latent_update = b.reshape(&latent_rows, vec![1, *lp, 1, cache_width]); + *kcache = b.dynamic_update_slice( + kcache, + &latent_update, + &[&k.layer_idx[li], &k.c0, &k.c0, &k.c0], + ); + let pe_update = b.reshape(&pe_rows, vec![1, *lp, 1, cache_width]); + *vcache = b.dynamic_update_slice( + vcache, + &pe_update, + &[&k.layer_idx[li], &k.c0, &k.c0, &k.c0], + ); + (latent.clone(), k_pe.clone()) + } + }; + let sequence = match layout { + AttnLayout::Single { .. } | AttnLayout::Ragged { .. } => c.context_capacity, + AttnLayout::Prefill { lp, .. } => *lp, + }; + let (keys, values) = mla_expand_latent(b, &latent_slab, &mw.kv_b, Some(sequence), c); + let scores = match layout { + AttnLayout::Single { .. } => { + let keys = b.transpose(&keys, &[1, 0, 2]); + let nope = b.dot_general( + &q_nope, + &keys, + &[0], + &[0], + &[1], + &[2], + vec![c.n_q, sequence], + ); + let pe = b.transpose(&pe_slab, &[1, 0]); + let rotary = b.dot_general(&q_pe, &pe, &[], &[], &[1], &[0], vec![c.n_q, sequence]); + b.add(&nope, &rotary) + } + AttnLayout::Ragged { bsz, .. } => { + let keys = b.transpose(&keys, &[0, 2, 1, 3]); + let q_nope = b.reshape(&q_nope, vec![*bsz, c.n_q, 1, m.qk_nope_head_dim]); + let nope = b.dot_general( + &q_nope, + &keys, + &[0, 1], + &[0, 1], + &[3], + &[3], + vec![*bsz, c.n_q, 1, sequence], + ); + let nope = b.reshape(&nope, vec![*bsz, c.n_q, sequence]); + let pe = b.transpose(&pe_slab, &[0, 2, 1]); + let q_pe = b.reshape(&q_pe, vec![*bsz, c.n_q, 1, m.qk_rope_head_dim]); + let pe = b.broadcast( + &pe, + &[0, 2, 3], + vec![*bsz, c.n_q, m.qk_rope_head_dim, sequence], + ); + let rotary = b.dot_general( + &q_pe, + &pe, + &[0, 1], + &[0, 1], + &[3], + &[2], + vec![*bsz, c.n_q, 1, sequence], + ); + let rotary = b.reshape(&rotary, vec![*bsz, c.n_q, sequence]); + b.add(&nope, &rotary) + } + AttnLayout::Prefill { lp, .. } => { + let q_nope = b.transpose(&q_nope, &[1, 0, 2]); + let keys = b.transpose(&keys, &[1, 0, 2]); + let nope = b.dot_general( + &q_nope, + &keys, + &[0], + &[0], + &[2], + &[2], + vec![c.n_q, *lp, *lp], + ); + let q_pe = b.transpose(&q_pe, &[1, 0, 2]); + let pe = b.transpose(&pe_slab, &[1, 0]); + let rotary = b.dot_general(&q_pe, &pe, &[], &[], &[2], &[0], vec![c.n_q, *lp, *lp]); + let scores = b.add(&nope, &rotary); + b.reshape(&scores, vec![c.n_q, *lp, 1, *lp]) + } + }; + let scores = apply_scale_and_softcap(b, c, k, scores); + let scores = layout.add_mask(b, c, li, &scores); + let attention = attn_softmax(b, k, &scores, layout.score_axis()); + let context = match layout { + AttnLayout::Single { .. } => { + let values = b.transpose(&values, &[1, 0, 2]); + let context = b.dot_general( + &attention, + &values, + &[0], + &[0], + &[1], + &[1], + vec![c.n_q, m.v_head_dim], + ); + b.reshape(&context, vec![c.n_q * m.v_head_dim]) + } + AttnLayout::Ragged { bsz, .. } => { + let values = b.transpose(&values, &[0, 2, 1, 3]); + let context = b.dot_general( + &attention, + &values, + &[0, 1], + &[0, 1], + &[2], + &[2], + vec![*bsz, c.n_q, m.v_head_dim], + ); + b.reshape(&context, vec![*bsz, c.n_q * m.v_head_dim]) + } + AttnLayout::Prefill { lp, .. } => { + let attention = b.reshape(&attention, vec![c.n_q, *lp, *lp]); + let values = b.transpose(&values, &[1, 0, 2]); + let context = b.dot_general( + &attention, + &values, + &[0], + &[0], + &[2], + &[1], + vec![c.n_q, *lp, m.v_head_dim], + ); + let context = b.transpose(&context, &[1, 0, 2]); + b.reshape(&context, vec![*lp, c.n_q * m.v_head_dim]) + } + }; + let output = match layout { + AttnLayout::Single { .. } => b.linear(&context, &mw.o), + _ => b.linear_seq(&context, &mw.o), + }; + let output = post_attn_norm(b, c, k, layout, output, lw); + (hn, output) +} + /// The input RMSNorm applied at a layout's rank: the Gemma `(1 + w)` weight offset /// (a no-op handle-copy for the non-Gemma families) followed by the layout's /// rank-appropriate RMSNorm. The OLMo reordered post-norm has NO input norm @@ -2002,6 +2507,9 @@ fn emit_attention( kcache: &mut Val, vcache: &mut Val, ) -> (Val, Val) { + if c.mla.is_some() { + return emit_mla_attention(b, c, k, lw, li, x, layout, kcache, vcache); + } let hn = arch_norm(b, c, k, layout, x, &lw.in_ln, lw.in_ln_bias.as_ref()); let (q, kk, vv) = layout.project_qkv(b, c, &hn, lw); // q/k-norm hook (issue #494): Qwen3 / Gemma3 norm each head over head_dim, @@ -2179,7 +2687,7 @@ pub fn emit_decode_with(c: &Config, sample: bool, precision: Precision) -> Strin }; let sig = render_signature(&decls); - let cache_ty = Ty::f32(vec![c.n_layers, seq, c.n_kv, c.head_dim]).render(); + let cache_ty = Ty::f32(vec![c.n_layers, seq, c.cache_heads(), c.cache_dim()]).render(); format!( "module @decode_step {{\n func.func public @main({sig}) -> ({out_ty}, {cache_ty}, {cache_ty}) {{\n{body} return {l}, {kc}, {vc} : {out_ty}, {cache_ty}, {cache_ty}\n }}\n}}\n", sig = sig, @@ -2279,8 +2787,8 @@ fn build_batched_arg_schema( bsz, c.n_layers, c.context_capacity, - c.n_kv, - c.head_dim, + c.cache_heads(), + c.cache_dim(), ]), "kcache".into(), ); @@ -2291,8 +2799,8 @@ fn build_batched_arg_schema( bsz, c.n_layers, c.context_capacity, - c.n_kv, - c.head_dim, + c.cache_heads(), + c.cache_dim(), ]), "vcache".into(), ); @@ -2351,6 +2859,10 @@ pub fn emit_decode_batched_with( sample: bool, precision: Precision, ) -> String { + assert!( + c.mla.is_none(), + "uniform lockstep decode does not support MLA; use ragged decode" + ); assert!( c.phi4_lora.is_none(), "Phi4MM requires per-row ragged decode; uniform batched decode cannot carry request-scoped adapter modes" @@ -2405,13 +2917,13 @@ pub fn emit_decode_batched_with( .as_ref() .expect("uniform-B batched decode is emitted only for pre-norm archs"); let hn = rms_norm_seq(&mut b, &x, in_ln, &k, bsz, h); // [B, H] - let q = b.linear_seq(&hn, &lw.wq); // [B, qd] + let q = b.linear_seq(&hn, lw.wq.as_ref().expect("ordinary q projection")); // [B, qd] let q = add_proj_bias_seq(&mut b, q, &lw.bq, bsz, nq * d); let q = b.reshape(&q, vec![bsz, nq, d]); - let kk = b.linear_seq(&hn, &lw.wk); // [B, kv] + let kk = b.linear_seq(&hn, lw.wk.as_ref().expect("ordinary k projection")); // [B, kv] let kk = add_proj_bias_seq(&mut b, kk, &lw.bk, bsz, nkv * d); let kk = b.reshape(&kk, vec![bsz, nkv, d]); - let vv = b.linear_seq(&hn, &lw.wv); // [B, kv] + let vv = b.linear_seq(&hn, lw.wv.as_ref().expect("ordinary v projection")); // [B, kv] let vv = add_proj_bias_seq(&mut b, vv, &lw.bv, bsz, nkv * d); let vv = b.reshape(&vv, vec![bsz, nkv, d]); @@ -2484,7 +2996,7 @@ pub fn emit_decode_batched_with( ); let o = b.reshape(&o, vec![bsz, nq, d]); let o = b.reshape(&o, vec![bsz, nq * d]); - let attn_out = b.linear_seq(&o, &lw.wo); // [B, H] + let attn_out = b.linear_seq(&o, lw.wo.as_ref().expect("ordinary output projection")); // [B, H] x = b.add(&x, &attn_out); // FFN: the MoE block (issue #500) on a MoE layer, else the dense SwiGLU MLP @@ -2541,7 +3053,7 @@ pub fn emit_decode_batched_with( }; let sig = render_signature(&decls); - let cache_ty = Ty::f32(vec![bsz, c.n_layers, seq, c.n_kv, c.head_dim]).render(); + let cache_ty = Ty::f32(vec![bsz, c.n_layers, seq, c.cache_heads(), c.cache_dim()]).render(); format!( "module @decode_step {{\n func.func public @main({sig}) -> ({out_ty}, {cache_ty}, {cache_ty}) {{\n{body} return {l}, {kc}, {vc} : {out_ty}, {cache_ty}, {cache_ty}\n }}\n}}\n", sig = sig, @@ -2649,8 +3161,8 @@ fn build_ragged_arg_schema(b: &mut Builder, c: &Config, bsz: usize) -> (Vec (Vec ({out_ty}, {cache_ty}, {cache_ty}) {{\n{body} return {l}, {kc}, {vc} : {out_ty}, {cache_ty}, {cache_ty}\n }}\n}}\n", sig = sig, @@ -3321,8 +3833,8 @@ fn emit_prefill_module( let k = emit_consts(&mut b, c); let h = c.hidden; - let d = c.head_dim; - let nkv = c.n_kv; + let d = c.cache_dim(); + let nkv = c.cache_heads(); // --- input head --- // Token prefill gathers + applies the architecture's embedding scale. The @@ -3488,7 +4000,7 @@ fn emit_prefill_module( }; let sig = render_signature(&decls); - let cache_ty = Ty::f32(vec![c.n_layers, lp, c.n_kv, c.head_dim]).render(); + let cache_ty = Ty::f32(vec![c.n_layers, lp, c.cache_heads(), c.cache_dim()]).render(); let module_name = match (input_kind, capture_deepstack_states) { (PrefillInputKind::DeepStack, true) => "prefill_embeddings_deepstack_diagnostics", (PrefillInputKind::Tokens, false) => "prefill", diff --git a/src/lib/mlxcel-xla/src/emitter/youtu_vl.rs b/src/lib/mlxcel-xla/src/emitter/youtu_vl.rs new file mode 100644 index 000000000..8f1ae5407 --- /dev/null +++ b/src/lib/mlxcel-xla/src/emitter/youtu_vl.rs @@ -0,0 +1,321 @@ +// 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. + +//! StableHLO Youtu-VL windowed vision tower and built-in patch merger. + +use super::builder::{Builder, Ty, Val}; +use super::youtu_vl_plan::{YOUTU_VL_PATCH_BUCKETS, YoutuVlVisionConfig, YoutuVlWeightSpec}; + +struct Args { + values: Vec, + declarations: Vec, + cursor: usize, +} + +impl Args { + fn new(specs: &[YoutuVlWeightSpec]) -> Self { + let mut values = Vec::with_capacity(specs.len() + 4); + let mut declarations = Vec::with_capacity(specs.len() + 4); + for (index, spec) in specs.iter().enumerate() { + let ty = Ty::f32(spec.shape.clone()); + declarations.push(format!( + "%arg{index}: {} loc(\"{}\")", + ty.render(), + spec.name + )); + values.push(Builder::arg(index, ty)); + } + Self { + values, + declarations, + cursor: 0, + } + } + + fn take(&mut self) -> Val { + let value = self.values[self.cursor].clone(); + self.cursor += 1; + value + } + + fn input(&mut self, ty: Ty, name: &str) -> Val { + let index = self.values.len(); + self.declarations + .push(format!("%arg{index}: {} loc(\"{name}\")", ty.render())); + let value = Builder::arg(index, ty); + self.values.push(value.clone()); + value + } +} + +fn add_bias(builder: &mut Builder, value: &Val, bias: &Val) -> Val { + let shape = value.ty.shape.clone(); + let bias = builder.broadcast(bias, &[1], shape); + builder.add(value, &bias) +} + +fn linear(builder: &mut Builder, value: &Val, weight: &Val, bias: &Val) -> Val { + let value = builder.linear_seq(value, weight); + add_bias(builder, &value, bias) +} + +fn layer_norm(builder: &mut Builder, value: &Val, weight: &Val, bias: &Val, epsilon: f32) -> Val { + let rows = value.ty.shape[0]; + let width = value.ty.shape[1]; + let zero = builder.const_f32(0.0); + let width_scalar = builder.const_f32(width as f32); + let width_rows = builder.broadcast(&width_scalar, &[], vec![rows]); + let sum = builder.reduce_add(value, 1, &zero); + let mean = builder.divide(&sum, &width_rows); + let mean = builder.broadcast(&mean, &[0], vec![rows, width]); + let centered = builder.subtract(value, &mean); + let squared = builder.multiply(¢ered, ¢ered); + let sum = builder.reduce_add(&squared, 1, &zero); + let variance = builder.divide(&sum, &width_rows); + let epsilon = builder.const_f32(epsilon); + let epsilon = builder.broadcast(&epsilon, &[], vec![rows]); + let variance = builder.add(&variance, &epsilon); + let inverse = builder.rsqrt(&variance); + let inverse = builder.broadcast(&inverse, &[0], vec![rows, width]); + let normalized = builder.multiply(¢ered, &inverse); + let weight = builder.broadcast(weight, &[1], vec![rows, width]); + let bias = builder.broadcast(bias, &[1], vec![rows, width]); + let normalized = builder.multiply(&normalized, &weight); + builder.add(&normalized, &bias) +} + +fn 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 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 = builder.divide(&sum, &width_rows); + let epsilon = builder.const_f32(epsilon); + let epsilon = builder.broadcast(&epsilon, &[], vec![rows]); + let mean = builder.add(&mean, &epsilon); + let scale = builder.rsqrt(&mean); + let scale = builder.broadcast(&scale, &[0], vec![rows, width]); + let weight = builder.broadcast(weight, &[1], vec![rows, width]); + let value = builder.multiply(value, &scale); + builder.multiply(&value, &weight) +} + +fn gelu_tanh(builder: &mut Builder, value: &Val) -> Val { + let shape = value.ty.shape.clone(); + let half_scalar = builder.const_f32(0.5); + let half = builder.broadcast(&half_scalar, &[], shape.clone()); + let one_scalar = builder.const_f32(1.0); + let one = builder.broadcast(&one_scalar, &[], shape.clone()); + let coefficient_scalar = builder.const_f32(0.044_715); + let coefficient = builder.broadcast(&coefficient_scalar, &[], shape.clone()); + let scale_scalar = builder.const_f32(0.797_884_6); + let scale = builder.broadcast(&scale_scalar, &[], shape); + let squared = builder.multiply(value, value); + let cubed = builder.multiply(&squared, value); + let nonlinear = builder.multiply(&coefficient, &cubed); + let inner = builder.add(value, &nonlinear); + let scaled = builder.multiply(&scale, &inner); + let tanh = builder.tanh(&scaled); + let cdf = builder.add(&one, &tanh); + let half_value = builder.multiply(value, &half); + builder.multiply(&half_value, &cdf) +} + +fn rotate_half(builder: &mut Builder, value: &Val) -> Val { + let (tokens, heads, width) = (value.ty.shape[0], value.ty.shape[1], value.ty.shape[2]); + let half = width / 2; + let first = builder.slice(value, &[(0, tokens), (0, heads), (0, half)]); + let second = builder.slice(value, &[(0, tokens), (0, heads), (half, width)]); + let second = builder.negate(&second); + builder.concatenate(&second, &first, 2) +} + +fn apply_rope(builder: &mut Builder, value: &Val, freqs: &Val) -> Val { + let tokens = value.ty.shape[0]; + let heads = value.ty.shape[1]; + let half = freqs.ty.shape[1]; + let cos = builder.cosine(freqs); + let sin = builder.sine(freqs); + let cos = builder.concatenate(&cos, &cos, 1); + let sin = builder.concatenate(&sin, &sin, 1); + let cos = builder.broadcast(&cos, &[0, 2], vec![tokens, heads, half * 2]); + let sin = builder.broadcast(&sin, &[0, 2], vec![tokens, heads, half * 2]); + let direct = builder.multiply(value, &cos); + let rotated = rotate_half(builder, value); + let rotated = builder.multiply(&rotated, &sin); + builder.add(&direct, &rotated) +} + +fn softmax(builder: &mut Builder, value: &Val) -> Val { + let shape = value.ty.shape.clone(); + let negative = builder.const_f32(f32::NEG_INFINITY); + let maximum = builder.reduce_max(value, 2, &negative); + let maximum = builder.broadcast(&maximum, &[0, 1], shape.clone()); + let shifted = builder.subtract(value, &maximum); + let exponentials = builder.exponential(&shifted); + let zero = builder.const_f32(0.0); + let denominator = builder.reduce_add(&exponentials, 2, &zero); + let denominator = builder.broadcast(&denominator, &[0, 1], shape); + builder.divide(&exponentials, &denominator) +} + +fn attention( + builder: &mut Builder, + hidden: &Val, + args: &mut Args, + config: &YoutuVlVisionConfig, + freqs: &Val, + bias: &Val, +) -> Val { + let tokens = hidden.ty.shape[0]; + let head_dim = config.hidden / config.heads; + let q = linear(builder, hidden, &args.take(), &args.take()); + let k = linear(builder, hidden, &args.take(), &args.take()); + let v = linear(builder, hidden, &args.take(), &args.take()); + let q = builder.reshape(&q, vec![tokens, config.heads, head_dim]); + let k = builder.reshape(&k, vec![tokens, config.heads, head_dim]); + let v = builder.reshape(&v, vec![tokens, config.heads, head_dim]); + let q = apply_rope(builder, &q, freqs); + let k = apply_rope(builder, &k, freqs); + let q = builder.transpose(&q, &[1, 0, 2]); + let k = builder.transpose(&k, &[1, 0, 2]); + let v = builder.transpose(&v, &[1, 0, 2]); + let scores = builder.dot_general( + &q, + &k, + &[0], + &[0], + &[2], + &[2], + vec![config.heads, tokens, tokens], + ); + let scale = builder.const_f32((head_dim as f32).powf(-0.5)); + let scale = builder.broadcast(&scale, &[], vec![config.heads, tokens, tokens]); + let scores = builder.multiply(&scores, &scale); + let bias = builder.broadcast(bias, &[1, 2], vec![config.heads, tokens, tokens]); + let scores = builder.add(&scores, &bias); + let probabilities = softmax(builder, &scores); + let context = builder.dot_general( + &probabilities, + &v, + &[0], + &[0], + &[2], + &[1], + vec![config.heads, tokens, head_dim], + ); + let context = builder.transpose(&context, &[1, 0, 2]); + let context = builder.reshape(&context, vec![tokens, config.hidden]); + linear(builder, &context, &args.take(), &args.take()) +} + +fn layer( + builder: &mut Builder, + hidden: &Val, + args: &mut Args, + config: &YoutuVlVisionConfig, + freqs: &Val, + bias: &Val, +) -> Val { + let normalized = layer_norm( + builder, + hidden, + &args.take(), + &args.take(), + config.layer_norm_eps, + ); + let attention = attention(builder, &normalized, args, config, freqs, bias); + let residual = builder.add(hidden, &attention); + let normalized = layer_norm( + builder, + &residual, + &args.take(), + &args.take(), + config.layer_norm_eps, + ); + let up = linear(builder, &normalized, &args.take(), &args.take()); + let up = gelu_tanh(builder, &up); + let down = linear(builder, &up, &args.take(), &args.take()); + builder.add(&residual, &down) +} + +pub(crate) fn emit_youtu_vl(config: &YoutuVlVisionConfig, patch_bucket: usize) -> String { + assert!( + YOUTU_VL_PATCH_BUCKETS.contains(&patch_bucket), + "unqualified Youtu-VL patch bucket" + ); + let specs = config.weight_specs(); + let mut args = Args::new(&specs); + let mut builder = Builder::new(); + let patch_width = config.channels * config.patch_size * config.patch_size; + let head_dim = config.hidden / config.heads; + let patches = args.input( + Ty::f32(vec![patch_bucket, patch_width]), + "patches.window_order", + ); + let freqs = args.input( + Ty::f32(vec![patch_bucket, head_dim / 2]), + "vision_rope.freqs", + ); + let window_bias = args.input( + Ty::f32(vec![patch_bucket, patch_bucket]), + "window_attention.bias", + ); + let full_bias = args.input( + Ty::f32(vec![patch_bucket, patch_bucket]), + "full_attention.bias", + ); + let mut hidden = linear(&mut builder, &patches, &args.take(), &args.take()); + for layer_index in 0..config.depth { + let bias = if config.full_attention_layers.contains(&layer_index) { + &full_bias + } else { + &window_bias + }; + hidden = layer(&mut builder, &hidden, &mut args, config, &freqs, bias); + } + hidden = layer_norm( + &mut builder, + &hidden, + &args.take(), + &args.take(), + config.layer_norm_eps, + ); + hidden = rms_norm(&mut builder, &hidden, &args.take(), 1e-6); + let merged_width = config.hidden * config.spatial_merge_size * config.spatial_merge_size; + let merged = builder.reshape(&hidden, vec![patch_bucket / 4, merged_width]); + let merged = linear(&mut builder, &merged, &args.take(), &args.take()); + let merged = gelu_tanh(&mut builder, &merged); + let projected = linear(&mut builder, &merged, &args.take(), &args.take()); + assert_eq!(args.cursor, specs.len(), "Youtu-VL weight schema drifted"); + format!( + "module @youtu_vl_vision {{\n func.func public @main({signature}) -> {result_type} {{\n{body} return {result} : {result_type}\n }}\n}}\n", + signature = args.declarations.join(", "), + result_type = projected.ty.render(), + body = builder.body(), + result = projected.name, + ) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn emits_window_and_full_attention_inputs() { + let config = YoutuVlVisionConfig::from_json_str( + r#"{"model_type":"youtu_vl","hidden_size":12,"vision_config":{"num_hidden_layers":2,"hidden_size":8,"intermediate_size":16,"num_attention_heads":2,"num_channels":3,"patch_size":2,"spatial_merge_size":2,"window_size":8,"fullatt_block_indexes":[1],"out_hidden_size":12,"num_patches":256}}"#, + ) + .unwrap(); + let ir = emit_youtu_vl(&config, 16); + assert!(ir.contains("window_attention.bias")); + assert!(ir.contains("full_attention.bias")); + assert!(ir.contains("merger.mlp.2.weight")); + assert!(ir.contains("tensor<4x12xf32>")); + } +} diff --git a/src/lib/mlxcel-xla/src/emitter/youtu_vl_plan.rs b/src/lib/mlxcel-xla/src/emitter/youtu_vl_plan.rs new file mode 100644 index 000000000..f9919b461 --- /dev/null +++ b/src/lib/mlxcel-xla/src/emitter/youtu_vl_plan.rs @@ -0,0 +1,615 @@ +// 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 + +//! Static image-grid and window-attention contract for Youtu-VL. + +use std::collections::HashSet; +use std::path::Path; + +use serde_json::Value; + +pub(crate) const YOUTU_VL_PATCH_BUCKETS: [usize; 3] = [16, 64, 256]; +const MAX_IMAGES: usize = 4; +const DEFAULT_IMAGE_TOKEN_ID: i32 = 128_264; +const DEFAULT_VIDEO_TOKEN_ID: i32 = 128_265; +const DEFAULT_VISION_START_TOKEN_ID: i32 = 128_262; +const DEFAULT_VISION_END_TOKEN_ID: i32 = 128_263; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct YoutuVlWeightSpec { + pub(crate) name: String, + pub(crate) shape: Vec, +} + +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct YoutuVlVisionConfig { + pub(crate) depth: usize, + pub(crate) hidden: usize, + pub(crate) intermediate: usize, + pub(crate) heads: usize, + pub(crate) channels: usize, + pub(crate) patch_size: usize, + pub(crate) spatial_merge_size: usize, + pub(crate) window_size: usize, + pub(crate) full_attention_layers: Vec, + pub(crate) text_hidden: usize, + pub(crate) layer_norm_eps: f32, + pub(crate) max_patches_per_image: usize, + pub(crate) resample: u8, + pub(crate) image_token_id: i32, + pub(crate) video_token_id: i32, + pub(crate) vision_start_token_id: i32, + pub(crate) vision_end_token_id: i32, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct YoutuVlGridPlan { + pub(crate) patch_bucket: usize, + pub(crate) actual_patches: usize, + pub(crate) merged_tokens_per_image: Vec, + pub(crate) window_group_index: Vec, + pub(crate) reverse_group_index: Vec, + pub(crate) window_cu_seqlens: Vec, + pub(crate) full_cu_seqlens: Vec, +} + +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct YoutuVlHostInputs { + pub(crate) plan: YoutuVlGridPlan, + pub(crate) patches: Vec, + pub(crate) rope_freqs: Vec, + pub(crate) window_attention_bias: Vec, + pub(crate) full_attention_bias: Vec, +} + +fn positive(object: &serde_json::Map, field: &str) -> Result { + object + .get(field) + .and_then(Value::as_u64) + .and_then(|value| usize::try_from(value).ok()) + .filter(|&value| value > 0) + .ok_or_else(|| format!("Youtu-VL config field {field} must be a positive integer")) +} + +impl YoutuVlVisionConfig { + pub(crate) fn from_model_dir(model_dir: &Path) -> Result { + let config_path = model_dir.join("config.json"); + let text = std::fs::read_to_string(&config_path) + .map_err(|error| format!("{}: {error}", config_path.display()))?; + let mut config = Self::from_json_str(&text)?; + let processor_path = model_dir.join("preprocessor_config.json"); + let processor_text = std::fs::read_to_string(&processor_path) + .map_err(|error| format!("{}: {error}", processor_path.display()))?; + let processor: Value = serde_json::from_str(&processor_text) + .map_err(|error| format!("parse {}: {error}", processor_path.display()))?; + config.max_patches_per_image = processor + .get("max_num_patches") + .and_then(Value::as_u64) + .and_then(|value| usize::try_from(value).ok()) + .filter(|&value| value > 0) + .ok_or_else(|| { + "Youtu-VL preprocessor_config.json requires positive max_num_patches".to_string() + })?; + config.resample = processor + .get("resample") + .and_then(Value::as_u64) + .and_then(|value| u8::try_from(value).ok()) + .ok_or_else(|| { + "Youtu-VL preprocessor_config.json requires byte-sized resample".to_string() + })?; + if config.max_patches_per_image > YOUTU_VL_PATCH_BUCKETS[2] { + return Err(format!( + "Youtu-VL processor cap {} exceeds qualified XLA capacity {}", + config.max_patches_per_image, YOUTU_VL_PATCH_BUCKETS[2] + )); + } + Ok(config) + } + + 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("youtu_vl") { + return Err("Youtu-VL IREE vision requires model_type=youtu_vl".to_string()); + } + if root + .get("quantization") + .is_some_and(|value| !value.is_null()) + { + return Err("Youtu-VL IREE vision currently requires BF16/F16/F32 weights".to_string()); + } + let vision = root + .get("vision_config") + .and_then(Value::as_object) + .ok_or_else(|| "Youtu-VL config.json vision_config must be an object".to_string())?; + let depth = positive(vision, "num_hidden_layers")?; + let hidden = positive(vision, "hidden_size")?; + let heads = positive(vision, "num_attention_heads")?; + if !hidden.is_multiple_of(heads) || !(hidden / heads).is_multiple_of(4) { + return Err("Youtu-VL vision head dimension must be divisible by four".to_string()); + } + let spatial_merge_size = vision + .get("spatial_merge_size") + .and_then(Value::as_u64) + .map_or(2usize, |value| value as usize); + if spatial_merge_size != 2 { + return Err(format!( + "Youtu-VL IREE vision qualifies spatial_merge_size=2, got {spatial_merge_size}" + )); + } + let patch_size = positive(vision, "patch_size")?; + let window_size = positive(vision, "window_size")?; + if !window_size.is_multiple_of(patch_size * spatial_merge_size) { + return Err( + "Youtu-VL window_size must be divisible by patch_size * spatial_merge_size" + .to_string(), + ); + } + let full_attention_layers = vision + .get("fullatt_block_indexes") + .and_then(Value::as_array) + .ok_or_else(|| "Youtu-VL fullatt_block_indexes must be an array".to_string())? + .iter() + .map(|value| { + value + .as_u64() + .and_then(|value| usize::try_from(value).ok()) + .filter(|&value| value < depth) + .ok_or_else(|| "invalid Youtu-VL full-attention layer".to_string()) + }) + .collect::, _>>()?; + if full_attention_layers.is_empty() { + return Err("Youtu-VL requires at least one full-attention layer".to_string()); + } + Ok(Self { + depth, + hidden, + intermediate: positive(vision, "intermediate_size")?, + heads, + channels: positive(vision, "num_channels")?, + patch_size, + spatial_merge_size, + window_size, + full_attention_layers, + text_hidden: positive(vision, "out_hidden_size")?, + layer_norm_eps: vision + .get("layer_norm_eps") + .and_then(Value::as_f64) + .unwrap_or(1e-6) as f32, + max_patches_per_image: positive(vision, "num_patches")?, + resample: 2, + image_token_id: token_id(&root, "image_token_id", DEFAULT_IMAGE_TOKEN_ID)?, + video_token_id: token_id(&root, "video_token_id", DEFAULT_VIDEO_TOKEN_ID)?, + vision_start_token_id: token_id( + &root, + "vision_start_token_id", + DEFAULT_VISION_START_TOKEN_ID, + )?, + vision_end_token_id: token_id( + &root, + "vision_end_token_id", + DEFAULT_VISION_END_TOKEN_ID, + )?, + }) + } + + pub(crate) fn fingerprint(&self, bucket: usize) -> String { + format!( + "youtu-vl-vision-v1:bucket={bucket}:patch={}:merge={}:window={}:channels={}:hidden={}:intermediate={}:heads={}:depth={}:full={:?}:text={}:eps={}:max_patches={}:resample={}:image_token={}:video_token={}:vision_start={}:vision_end={}:dtype=f32", + self.patch_size, + self.spatial_merge_size, + self.window_size, + self.channels, + self.hidden, + self.intermediate, + self.heads, + self.depth, + self.full_attention_layers, + self.text_hidden, + self.layer_norm_eps, + self.max_patches_per_image, + self.resample, + self.image_token_id, + self.video_token_id, + self.vision_start_token_id, + self.vision_end_token_id, + ) + } + + pub(crate) fn plan(&self, shapes: &[(i32, i32)]) -> Result { + if shapes.is_empty() || shapes.len() > MAX_IMAGES { + return Err(format!( + "Youtu-VL XLA requires 1..={MAX_IMAGES} images, got {}", + shapes.len() + )); + } + let merge = self.spatial_merge_size; + let window = self.window_size / self.patch_size / merge; + let merge_unit = merge * merge; + let mut total_patches = 0usize; + let mut merged_tokens_per_image = Vec::with_capacity(shapes.len()); + let mut window_group_index = Vec::new(); + let mut window_cu_seqlens = vec![0usize]; + let mut full_cu_seqlens = vec![0usize]; + let mut group_offset = 0usize; + for (index, &(height, width)) in shapes.iter().enumerate() { + let height = usize::try_from(height) + .map_err(|_| format!("Youtu-VL image {index} has invalid height {height}"))?; + let width = usize::try_from(width) + .map_err(|_| format!("Youtu-VL image {index} has invalid width {width}"))?; + if height == 0 + || width == 0 + || !height.is_multiple_of(merge) + || !width.is_multiple_of(merge) + { + return Err(format!( + "Youtu-VL image {index} grid {height}x{width} must be positive and divisible by merge={merge}" + )); + } + let patches = height + .checked_mul(width) + .ok_or_else(|| "Youtu-VL image patch count overflowed".to_string())?; + if patches > self.max_patches_per_image { + return Err(format!( + "Youtu-VL image {index} has {patches} patches, exceeding cap {}", + self.max_patches_per_image + )); + } + total_patches = total_patches + .checked_add(patches) + .ok_or_else(|| "Youtu-VL packed patch count overflowed".to_string())?; + let groups_h = height / merge; + let groups_w = width / merge; + let groups = groups_h * groups_w; + merged_tokens_per_image.push(groups); + full_cu_seqlens.push(total_patches); + for window_h in (0..groups_h).step_by(window) { + for window_w in (0..groups_w).step_by(window) { + let before = window_group_index.len(); + for h in window_h..(window_h + window).min(groups_h) { + for w in window_w..(window_w + window).min(groups_w) { + window_group_index.push(group_offset + h * groups_w + w); + } + } + let group_count = window_group_index.len() - before; + window_cu_seqlens.push( + window_cu_seqlens + .last() + .copied() + .unwrap() + .checked_add(group_count * merge_unit) + .ok_or_else(|| "Youtu-VL window boundary overflowed".to_string())?, + ); + } + } + group_offset += groups; + } + let patch_bucket = YOUTU_VL_PATCH_BUCKETS + .into_iter() + .find(|&bucket| total_patches <= bucket) + .ok_or_else(|| { + format!( + "Youtu-VL packed grid has {total_patches} patches, exceeding qualified capacity {}", + YOUTU_VL_PATCH_BUCKETS[2] + ) + })?; + let mut reverse_group_index = vec![0usize; window_group_index.len()]; + for (window_position, &original_position) in window_group_index.iter().enumerate() { + reverse_group_index[original_position] = window_position; + } + Ok(YoutuVlGridPlan { + patch_bucket, + actual_patches: total_patches, + merged_tokens_per_image, + window_group_index, + reverse_group_index, + window_cu_seqlens, + full_cu_seqlens, + }) + } + + pub(crate) fn weight_specs(&self) -> Vec { + let patch_width = self.channels * self.patch_size * self.patch_size; + let mut specs = vec![ + self.spec( + "siglip2.vision_model.embeddings.patch_embedding.weight", + [self.hidden, patch_width], + ), + self.spec( + "siglip2.vision_model.embeddings.patch_embedding.bias", + [self.hidden], + ), + ]; + for layer in 0..self.depth { + let prefix = format!("siglip2.vision_model.encoder.layers.{layer}"); + for (suffix, shape) in [ + ("layer_norm1.weight", vec![self.hidden]), + ("layer_norm1.bias", vec![self.hidden]), + ("self_attn.q_proj.weight", vec![self.hidden, self.hidden]), + ("self_attn.q_proj.bias", vec![self.hidden]), + ("self_attn.k_proj.weight", vec![self.hidden, self.hidden]), + ("self_attn.k_proj.bias", vec![self.hidden]), + ("self_attn.v_proj.weight", vec![self.hidden, self.hidden]), + ("self_attn.v_proj.bias", vec![self.hidden]), + ("self_attn.out_proj.weight", vec![self.hidden, self.hidden]), + ("self_attn.out_proj.bias", vec![self.hidden]), + ("layer_norm2.weight", vec![self.hidden]), + ("layer_norm2.bias", vec![self.hidden]), + ("mlp.fc1.weight", vec![self.intermediate, self.hidden]), + ("mlp.fc1.bias", vec![self.intermediate]), + ("mlp.fc2.weight", vec![self.hidden, self.intermediate]), + ("mlp.fc2.bias", vec![self.hidden]), + ] { + specs.push(YoutuVlWeightSpec { + name: format!("{prefix}.{suffix}"), + shape, + }); + } + } + let merged = self.hidden * 4; + specs.extend([ + self.spec("siglip2.vision_model.post_layernorm.weight", [self.hidden]), + self.spec("siglip2.vision_model.post_layernorm.bias", [self.hidden]), + self.spec("merger.ln_q.weight", [self.hidden]), + self.spec("merger.mlp.0.weight", [merged, merged]), + self.spec("merger.mlp.0.bias", [merged]), + self.spec("merger.mlp.2.weight", [self.text_hidden, merged]), + self.spec("merger.mlp.2.bias", [self.text_hidden]), + ]); + specs + } + + fn spec(&self, name: &str, shape: [usize; N]) -> YoutuVlWeightSpec { + YoutuVlWeightSpec { + name: name.to_string(), + shape: shape.into(), + } + } +} + +fn token_id(root: &Value, field: &str, default: i32) -> Result { + match root.get(field) { + None | Some(Value::Null) => Ok(default), + Some(value) => value + .as_i64() + .and_then(|value| i32::try_from(value).ok()) + .ok_or_else(|| format!("Youtu-VL config field {field} must fit i32")), + } +} + +fn attention_bias(bucket: usize, boundaries: &[usize]) -> Vec { + let mut bias = vec![f32::NEG_INFINITY; bucket * bucket]; + for segment in boundaries.windows(2) { + for row in segment[0]..segment[1] { + for column in segment[0]..segment[1] { + bias[row * bucket + column] = 0.0; + } + } + } + for index in boundaries.last().copied().unwrap_or(0)..bucket { + bias[index * bucket + index] = 0.0; + } + bias +} + +pub(crate) fn prepare_youtu_vl_host_inputs( + config: &YoutuVlVisionConfig, + shapes: &[(i32, i32)], + patch_rows: &[f32], +) -> Result { + let plan = config.plan(shapes)?; + let patch_width = config.channels * config.patch_size * config.patch_size; + if patch_rows.len() != plan.actual_patches * patch_width { + return Err(format!( + "Youtu-VL processor produced {} values, expected {}", + patch_rows.len(), + plan.actual_patches * patch_width + )); + } + if patch_rows.iter().any(|value| !value.is_finite()) { + return Err("Youtu-VL processor produced a non-finite patch value".to_string()); + } + let merge_unit = config.spatial_merge_size * config.spatial_merge_size; + let mut grouped_rows = Vec::with_capacity(patch_rows.len()); + let mut image_patch_offset = 0usize; + for &(height, width) in shapes { + let (height, width) = (height as usize, width as usize); + for block_h in 0..height / config.spatial_merge_size { + for block_w in 0..width / config.spatial_merge_size { + for inner_h in 0..config.spatial_merge_size { + for inner_w in 0..config.spatial_merge_size { + let patch = image_patch_offset + + (block_h * config.spatial_merge_size + inner_h) * width + + block_w * config.spatial_merge_size + + inner_w; + let start = patch * patch_width; + grouped_rows.extend_from_slice(&patch_rows[start..start + patch_width]); + } + } + } + } + image_patch_offset += height * width; + } + let mut patches = Vec::with_capacity(plan.patch_bucket * patch_width); + for &group in &plan.window_group_index { + let start = group * merge_unit * patch_width; + patches.extend_from_slice(&grouped_rows[start..start + merge_unit * patch_width]); + } + patches.resize(plan.patch_bucket * patch_width, 0.0); + + let head_dim = config.hidden / config.heads; + let quarter = head_dim / 4; + let inv = (0..quarter) + .map(|index| 1.0 / 10_000.0f32.powf((2 * index) as f32 / (head_dim / 2) as f32)) + .collect::>(); + let mut grouped_freqs = Vec::with_capacity(plan.actual_patches * head_dim / 2); + for &(height, width) in shapes { + let (height, width) = (height as usize, width as usize); + for block_h in 0..height / 2 { + for block_w in 0..width / 2 { + for inner_h in 0..2 { + for inner_w in 0..2 { + let h = (block_h * 2 + inner_h) as f32; + let w = (block_w * 2 + inner_w) as f32; + grouped_freqs.extend(inv.iter().map(|frequency| h * frequency)); + grouped_freqs.extend(inv.iter().map(|frequency| w * frequency)); + } + } + } + } + } + let mut rope_freqs = Vec::with_capacity(plan.patch_bucket * head_dim / 2); + for &group in &plan.window_group_index { + let start = group * merge_unit * head_dim / 2; + rope_freqs.extend_from_slice(&grouped_freqs[start..start + merge_unit * head_dim / 2]); + } + rope_freqs.resize(plan.patch_bucket * head_dim / 2, 0.0); + let window_attention_bias = attention_bias(plan.patch_bucket, &plan.window_cu_seqlens); + let full_attention_bias = attention_bias(plan.patch_bucket, &plan.full_cu_seqlens); + debug_assert_eq!( + plan.window_group_index + .iter() + .copied() + .collect::>() + .len(), + plan.window_group_index.len() + ); + Ok(YoutuVlHostInputs { + plan, + patches, + rope_freqs, + window_attention_bias, + full_attention_bias, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn config() -> YoutuVlVisionConfig { + YoutuVlVisionConfig { + depth: 2, + hidden: 16, + intermediate: 32, + heads: 4, + channels: 3, + patch_size: 2, + spatial_merge_size: 2, + window_size: 8, + full_attention_layers: vec![1], + text_hidden: 12, + layer_norm_eps: 1e-6, + max_patches_per_image: 256, + resample: 2, + image_token_id: DEFAULT_IMAGE_TOKEN_ID, + video_token_id: DEFAULT_VIDEO_TOKEN_ID, + vision_start_token_id: DEFAULT_VISION_START_TOKEN_ID, + vision_end_token_id: DEFAULT_VISION_END_TOKEN_ID, + } + } + + #[test] + fn plans_aspect_ratios_and_restores_window_group_order() { + let plan = config().plan(&[(4, 8), (8, 4)]).unwrap(); + assert_eq!(plan.patch_bucket, 64); + assert_eq!(plan.actual_patches, 64); + assert_eq!(plan.merged_tokens_per_image, vec![8, 8]); + assert_eq!(plan.full_cu_seqlens, vec![0, 32, 64]); + assert_ne!( + plan.window_group_index, + (0..plan.window_group_index.len()).collect::>() + ); + let window_values = plan + .window_group_index + .iter() + .map(|&original| format!("group-{original}")) + .collect::>(); + let restored = plan + .reverse_group_index + .iter() + .map(|&window_position| window_values[window_position].as_str()) + .collect::>(); + let expected = (0..16) + .map(|original| format!("group-{original}")) + .collect::>(); + assert_eq!( + restored, + expected.iter().map(String::as_str).collect::>() + ); + } + + #[test] + fn attention_bias_keeps_images_and_windows_isolated() { + let width = config().channels * config().patch_size * config().patch_size; + let inputs = + prepare_youtu_vl_host_inputs(&config(), &[(4, 8), (8, 4)], &vec![0.0; 64 * width]) + .unwrap(); + let bucket = inputs.plan.patch_bucket; + assert_eq!(inputs.full_attention_bias[0 * bucket + 31], 0.0); + assert!(inputs.full_attention_bias[0 * bucket + 32].is_infinite()); + assert!(inputs.window_attention_bias[0 * bucket + 16].is_infinite()); + assert_eq!( + inputs.window_attention_bias[(bucket - 1) * bucket + bucket - 1], + 0.0 + ); + } + + #[test] + fn rejects_per_image_and_packed_bucket_overflow() { + let mut capped = config(); + capped.max_patches_per_image = 15; + assert!( + capped + .plan(&[(4, 4)]) + .unwrap_err() + .contains("exceeding cap") + ); + + let error = config().plan(&[(16, 16), (4, 4)]).unwrap_err(); + assert!(error.contains("packed grid")); + assert_eq!(config().plan(&[(16, 16)]).unwrap().patch_bucket, 256); + } + + #[test] + fn artifact_identity_binds_placeholder_ids() { + let baseline = config(); + let mut changed = baseline.clone(); + changed.video_token_id += 1; + assert_ne!(baseline.fingerprint(16), changed.fingerprint(16)); + assert!(baseline.fingerprint(16).contains("image_token=128264")); + } + + #[test] + fn host_patch_reorder_preserves_merge_group_membership() { + let config = config(); + let width = config.channels * config.patch_size * config.patch_size; + let mut rows = vec![0.0; 32 * width]; + for patch in 0..32 { + rows[patch * width..(patch + 1) * width].fill(patch as f32); + } + let inputs = prepare_youtu_vl_host_inputs(&config, &[(4, 8)], &rows).unwrap(); + let observed = (0..inputs.plan.actual_patches) + .map(|row| inputs.patches[row * width] as usize) + .collect::>(); + for (window_position, &group) in inputs.plan.window_group_index.iter().enumerate() { + let group_rows = &observed[window_position * 4..window_position * 4 + 4]; + let group_h = group / 4; + let group_w = group % 4; + assert_eq!( + group_rows, + &[ + group_h * 16 + group_w * 2, + group_h * 16 + group_w * 2 + 1, + group_h * 16 + group_w * 2 + 8, + group_h * 16 + group_w * 2 + 9, + ] + ); + } + } +} diff --git a/src/lib/mlxcel-xla/src/lib.rs b/src/lib/mlxcel-xla/src/lib.rs index 2b1328267..f4f8fd4a3 100644 --- a/src/lib/mlxcel-xla/src/lib.rs +++ b/src/lib/mlxcel-xla/src/lib.rs @@ -82,6 +82,8 @@ mod phi4_audio; mod qwen2_vl_runtime; #[cfg(feature = "iree")] mod vision_runtime; +#[cfg(feature = "iree")] +mod youtu_vl_runtime; #[cfg(feature = "micro-oracle")] pub use numeric_oracle::{ @@ -184,6 +186,10 @@ pub use qwen2_vl_runtime::{ pub use vision_runtime::{IreeVisionDiagnosticProjector, VisionDiagnosticProjection}; #[cfg(feature = "iree")] pub use vision_runtime::{IreeVisionProjector, VisionExecutionMetrics, VisionProjection}; +#[cfg(feature = "iree")] +pub use youtu_vl_runtime::{ + IreeYoutuVlProjector, YoutuVlVisionExecutionMetrics, YoutuVlVisionProjection, +}; #[cfg(any(test, feature = "diagnostics"))] #[must_use] pub fn llava_diagnostic_device_memory_note(device: &str) -> &'static str { diff --git a/src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs b/src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs index 2634f24b3..e86312a0b 100644 --- a/src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs +++ b/src/lib/mlxcel-xla/src/qwen2_vl_runtime.rs @@ -81,7 +81,7 @@ fn hex_bytes(bytes: &[u8]) -> String { output } -fn sha256_hex(bytes: &[u8]) -> String { +pub(crate) fn sha256_hex(bytes: &[u8]) -> String { hex_bytes(&Sha256::digest(bytes)) } @@ -101,7 +101,7 @@ fn sha256_file(path: &Path) -> Result { Ok(hex_bytes(&digest.finalize())) } -fn compiler_generation_identity( +pub(crate) fn compiler_generation_identity( compiler: &Path, flags: &[&str], mlir: &str, @@ -179,7 +179,7 @@ fn processor_identity(model_dir: &Path, config: &Qwen2VlConfig) -> Result Result, String> { +pub(crate) fn model_shards(model_dir: &Path) -> Result, String> { let mut shards = std::fs::read_dir(model_dir) .map_err(|error| format!("read {}: {error}", model_dir.display()))? .map(|entry| { @@ -230,7 +230,7 @@ fn tensor_locations(model_dir: &Path) -> Result, Strin Ok(locations) } -fn native_f32_bytes(values: Vec) -> Vec { +pub(crate) fn native_f32_bytes(values: Vec) -> Vec { let mut bytes = Vec::with_capacity(values.len() * std::mem::size_of::()); for value in values { bytes.extend_from_slice(&value.to_ne_bytes()); @@ -238,7 +238,7 @@ fn native_f32_bytes(values: Vec) -> Vec { bytes } -fn validate_finite(label: &str, values: &[f32]) -> Result<(), String> { +pub(crate) fn validate_finite(label: &str, values: &[f32]) -> Result<(), String> { if let Some((index, value)) = values .iter() .enumerate() @@ -278,7 +278,7 @@ fn transpose_patch_t_h_w_c_to_t_c_h_w(values: Vec, config: &Qwen2VlConfig) transposed } -fn decode_direct_f32( +pub(crate) fn decode_direct_f32( label: &str, dtype: Dtype, data: &[u8], @@ -518,14 +518,14 @@ fn load_weights( )) } -fn f32_as_bytes(values: &[f32]) -> &[u8] { +pub(crate) fn f32_as_bytes(values: &[f32]) -> &[u8] { // Safety: f32 has no invalid bit patterns and the view cannot outlive input. unsafe { std::slice::from_raw_parts(values.as_ptr().cast::(), std::mem::size_of_val(values)) } } -fn checked_f32_output(label: &str, bytes: Vec) -> Result, String> { +pub(crate) fn checked_f32_output(label: &str, bytes: Vec) -> Result, String> { if !bytes.len().is_multiple_of(std::mem::size_of::()) { return Err(format!( "{label} returned {} bytes, not a whole number of f32 values", diff --git a/src/lib/mlxcel-xla/src/weights.rs b/src/lib/mlxcel-xla/src/weights.rs index 22338ab7e..433eb8fd7 100644 --- a/src/lib/mlxcel-xla/src/weights.rs +++ b/src/lib/mlxcel-xla/src/weights.rs @@ -223,8 +223,25 @@ fn weight_specs_q(cfg: &Config, quant: bool) -> Vec { ); } } - // wk, wo, wq, wv (JAX-alphabetical; a fused Phi3 qkv_proj is [Q|K|V] rows). - if cfg.fused_qkv { + // Attention projections. Youtu-VL uses the checkpoint's native MLA + // factors; ordinary families keep wk/wo/wq/wv in their historical order. + if cfg.mla.is_some() { + push_proj(&mut out, format!("{p}self_attn.q_a_proj.weight"), quant); + out.push(WeightSpec::Whole(format!( + "{p}self_attn.q_a_layernorm.weight" + ))); + push_proj(&mut out, format!("{p}self_attn.q_b_proj.weight"), quant); + push_proj( + &mut out, + format!("{p}self_attn.kv_a_proj_with_mqa.weight"), + quant, + ); + out.push(WeightSpec::Whole(format!( + "{p}self_attn.kv_a_layernorm.weight" + ))); + push_proj(&mut out, format!("{p}self_attn.kv_b_proj.weight"), quant); + push_proj(&mut out, format!("{p}self_attn.o_proj.weight"), quant); + } else if cfg.fused_qkv { let qkv = phi4_base_projection(cfg, format!("{p}self_attn.qkv_proj.weight")); out.push(WeightSpec::Rows { name: qkv.clone(), diff --git a/src/lib/mlxcel-xla/src/youtu_vl_runtime.rs b/src/lib/mlxcel-xla/src/youtu_vl_runtime.rs new file mode 100644 index 000000000..2edbf3d87 --- /dev/null +++ b/src/lib/mlxcel-xla/src/youtu_vl_runtime.rs @@ -0,0 +1,356 @@ +// 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. + +//! Resident IREE execution for Youtu-VL's vision tower and built-in merger. + +use std::collections::{BTreeMap, BTreeSet}; +use std::fs::File; +use std::path::{Path, PathBuf}; +use std::time::Instant; + +use memmap2::Mmap; +use safetensors::SafeTensors; + +use crate::aux::{ + AuxiliaryInput, AuxiliaryOutput, AuxiliaryTensorDType, AuxiliaryWeight, AuxiliaryWeightDType, + IreeAuxiliaryModule, +}; +use crate::aux_manifest::{AuxiliaryArtifactContract, ensure_qualified_auxiliary_artifact}; +use crate::emitter::{ + YoutuVlHostInputs, YoutuVlVisionConfig, emit_youtu_vl, prepare_youtu_vl_host_inputs, +}; +use crate::iree::{cached_vmfb_path, compile_one_to, iree_compile_bin, target_flags}; +use crate::qwen2_vl_runtime::{ + checked_f32_output, compiler_generation_identity, decode_direct_f32, f32_as_bytes, + model_shards, native_f32_bytes, sha256_hex, validate_finite, +}; + +const ENTRY_NAME: &str = "youtu_vl_vision.main"; + +#[derive(Debug, Clone, PartialEq)] +pub struct YoutuVlVisionExecutionMetrics { + pub patch_upload_bytes: usize, + pub metadata_upload_bytes: usize, + pub projected_transfer_bytes: usize, + pub elapsed_seconds: f64, +} + +#[derive(Debug, Clone, PartialEq)] +pub struct YoutuVlVisionProjection { + pub values: Vec, + pub shape: [usize; 2], + pub merged_tokens_per_image: Vec, + pub window_cu_seqlens: Vec, + pub full_cu_seqlens: Vec, + pub metrics: YoutuVlVisionExecutionMetrics, +} + +struct ActiveModule { + patch_bucket: usize, + module: IreeAuxiliaryModule, +} + +pub struct IreeYoutuVlProjector { + model_dir: PathBuf, + device: String, + config: YoutuVlVisionConfig, + processor_identity: String, + active: Option, +} + +fn processor_identity(model_dir: &Path, config: &YoutuVlVisionConfig) -> Result { + let path = model_dir.join("preprocessor_config.json"); + let bytes = std::fs::read(&path).map_err(|error| format!("{}: {error}", path.display()))?; + let value: serde_json::Value = serde_json::from_slice(&bytes) + .map_err(|error| format!("parse {}: {error}", path.display()))?; + let object = value + .as_object() + .ok_or_else(|| format!("{} must contain an object", path.display()))?; + for field in ["do_resize", "do_rescale", "do_normalize"] { + if object.get(field).and_then(serde_json::Value::as_bool) != Some(true) { + return Err(format!( + "Youtu-VL IREE vision requires preprocessor {field}=true" + )); + } + } + if object.get("patch_size").and_then(serde_json::Value::as_u64) + != Some(config.patch_size as u64) + { + return Err("Youtu-VL processor patch_size disagrees with vision config".to_string()); + } + Ok(format!( + "youtu-vl-processor-v1:sha256={}:max_patches={}:resample={}", + sha256_hex(&bytes), + config.max_patches_per_image, + config.resample, + )) +} + +fn tensor_locations( + model_dir: &Path, + required: &BTreeSet<&str>, +) -> Result, String> { + let mut locations = BTreeMap::new(); + for shard in model_shards(model_dir)? { + let file = + File::open(&shard).map_err(|error| format!("open {}: {error}", shard.display()))?; + // Safety: the read-only mapping lives through the header scan. + let mmap = unsafe { Mmap::map(&file) } + .map_err(|error| format!("mmap {}: {error}", shard.display()))?; + let tensors = SafeTensors::deserialize(&mmap) + .map_err(|error| format!("parse {}: {error}", shard.display()))?; + for name in tensors.names() { + if !required.contains(name.as_str()) { + continue; + } + if let Some(previous) = locations.insert(name.to_string(), shard.clone()) { + return Err(format!( + "Youtu-VL tensor {name} is duplicated in {} and {}", + previous.display(), + shard.display() + )); + } + } + } + Ok(locations) +} + +fn load_weights( + model_dir: &Path, + config: &YoutuVlVisionConfig, +) -> Result<(Vec, String), String> { + let specs = config.weight_specs(); + let required = specs + .iter() + .map(|spec| spec.name.as_str()) + .collect::>(); + let locations = tensor_locations(model_dir, &required)?; + let missing = required + .iter() + .filter(|name| !locations.contains_key(**name)) + .copied() + .collect::>(); + if !missing.is_empty() { + return Err(format!( + "checkpoint is missing {} Youtu-VL vision tensor(s): {}", + missing.len(), + missing + .iter() + .take(8) + .copied() + .collect::>() + .join(", ") + )); + } + let mut by_shard = BTreeMap::<&Path, Vec>::new(); + for (index, spec) in specs.iter().enumerate() { + by_shard + .entry( + locations + .get(&spec.name) + .expect("missing specs rejected") + .as_path(), + ) + .or_default() + .push(index); + } + let mut loaded = (0..specs.len()) + .map(|_| None) + .collect::>>(); + let mut schema = vec![String::new(); specs.len()]; + for (shard, indices) in by_shard { + let file = + File::open(shard).map_err(|error| format!("open {}: {error}", shard.display()))?; + // Safety: the read-only mapping lives while selected tensors copy. + let mmap = unsafe { Mmap::map(&file) } + .map_err(|error| format!("mmap {}: {error}", shard.display()))?; + let tensors = SafeTensors::deserialize(&mmap) + .map_err(|error| format!("parse {}: {error}", shard.display()))?; + for index in indices { + let spec = &specs[index]; + let tensor = tensors.tensor(&spec.name).map_err(|error| { + format!( + "resolved Youtu-VL tensor {} in {}: {error}", + spec.name, + shard.display() + ) + })?; + let dtype = tensor.dtype(); + let source_shape = tensor.shape().to_vec(); + let values = + decode_direct_f32(&spec.name, dtype, tensor.data(), &spec.shape, &source_shape)?; + validate_finite(&spec.name, &values)?; + loaded[index] = Some(AuxiliaryWeight { + name: spec.name.clone(), + bytes: native_f32_bytes(values), + dtype: AuxiliaryWeightDType::Float32, + shape: spec.shape.clone(), + }); + schema[index] = format!("{}:{dtype:?}:{source_shape:?}:identity->f32", spec.name); + } + } + Ok(( + loaded + .into_iter() + .map(|weight| weight.expect("all Youtu-VL specs loaded")) + .collect(), + schema.join("\n"), + )) +} + +fn compile_and_load( + model_dir: &Path, + device: &str, + config: &YoutuVlVisionConfig, + processor_identity: &str, + patch_bucket: usize, +) -> Result { + let mlir = emit_youtu_vl(config, patch_bucket); + let compiler = iree_compile_bin()?; + if !compiler.is_file() { + return Err(format!("iree-compile not found at {}", compiler.display())); + } + let flags = target_flags(device)?; + let cache = std::env::temp_dir().join("mlxcel-xla-youtu-vl-vmfb"); + std::fs::create_dir_all(&cache) + .map_err(|error| format!("mkdir {}: {error}", cache.display()))?; + let (weights, checkpoint_schema) = load_weights(model_dir, config)?; + let contract = AuxiliaryArtifactContract::new( + ENTRY_NAME, + format!( + "{};{};checkpoint_schema_sha256={}", + config.fingerprint(patch_bucket), + processor_identity, + sha256_hex(checkpoint_schema.as_bytes()) + ), + compiler_generation_identity(&compiler, flags, &mlir)?, + )?; + let tag = format!("youtu-vl-vision-{patch_bucket}"); + let vmfb = cached_vmfb_path(&compiler, &mlir, flags, &cache, &tag, 0); + ensure_qualified_auxiliary_artifact(&vmfb, &contract, &weights, |temporary| { + compile_one_to(&compiler, &mlir, flags, &cache, &tag, 0, temporary) + })?; + IreeAuxiliaryModule::load(device, &vmfb, &contract, weights) +} + +impl IreeYoutuVlProjector { + pub fn load(model_dir: &Path, device: &str) -> Result { + let config = YoutuVlVisionConfig::from_model_dir(model_dir)?; + let processor_identity = processor_identity(model_dir, &config)?; + Ok(Self { + model_dir: model_dir.to_path_buf(), + device: device.to_string(), + config, + processor_identity, + active: None, + }) + } + + #[must_use] + pub fn text_hidden(&self) -> usize { + self.config.text_hidden + } + + #[must_use] + pub fn active_bucket(&self) -> Option { + self.active.as_ref().map(|active| active.patch_bucket) + } + + fn ensure_bucket(&mut self, bucket: usize) -> Result<&mut IreeAuxiliaryModule, String> { + if self.active.as_ref().map(|active| active.patch_bucket) != Some(bucket) { + let module = compile_and_load( + &self.model_dir, + &self.device, + &self.config, + &self.processor_identity, + bucket, + )?; + self.active = Some(ActiveModule { + patch_bucket: bucket, + module, + }); + } + Ok(&mut self.active.as_mut().expect("active module").module) + } + + pub fn project( + &mut self, + patch_rows: &[f32], + shapes: &[(i32, i32)], + ) -> Result { + let YoutuVlHostInputs { + plan, + patches, + rope_freqs, + window_attention_bias, + full_attention_bias, + } = prepare_youtu_vl_host_inputs(&self.config, shapes, patch_rows)?; + let patch_shape = [ + plan.patch_bucket, + self.config.channels * self.config.patch_size * self.config.patch_size, + ]; + let rope_shape = [ + plan.patch_bucket, + self.config.hidden / self.config.heads / 2, + ]; + let bias_shape = [plan.patch_bucket, plan.patch_bucket]; + let output_shape = [plan.patch_bucket / 4, self.config.text_hidden]; + let mut output = + vec![0u8; output_shape.iter().product::() * std::mem::size_of::()]; + let started = Instant::now(); + self.ensure_bucket(plan.patch_bucket)?.invoke( + &[ + AuxiliaryInput { + bytes: f32_as_bytes(&patches), + dtype: AuxiliaryTensorDType::Float32, + shape: &patch_shape, + }, + AuxiliaryInput { + bytes: f32_as_bytes(&rope_freqs), + dtype: AuxiliaryTensorDType::Float32, + shape: &rope_shape, + }, + AuxiliaryInput { + bytes: f32_as_bytes(&window_attention_bias), + dtype: AuxiliaryTensorDType::Float32, + shape: &bias_shape, + }, + AuxiliaryInput { + bytes: f32_as_bytes(&full_attention_bias), + dtype: AuxiliaryTensorDType::Float32, + shape: &bias_shape, + }, + ], + &mut [AuxiliaryOutput { + bytes: &mut output, + dtype: AuxiliaryTensorDType::Float32, + shape: &output_shape, + }], + )?; + let elapsed_seconds = started.elapsed().as_secs_f64(); + let all_values = checked_f32_output("Youtu-VL IREE projected output", output)?; + let actual_tokens = plan.reverse_group_index.len(); + let mut values = Vec::with_capacity(actual_tokens * self.config.text_hidden); + for &window_position in &plan.reverse_group_index { + let start = window_position * self.config.text_hidden; + values.extend_from_slice(&all_values[start..start + self.config.text_hidden]); + } + Ok(YoutuVlVisionProjection { + values, + shape: [actual_tokens, self.config.text_hidden], + merged_tokens_per_image: plan.merged_tokens_per_image, + window_cu_seqlens: plan.window_cu_seqlens, + full_cu_seqlens: plan.full_cu_seqlens, + metrics: YoutuVlVisionExecutionMetrics { + patch_upload_bytes: std::mem::size_of_val(patches.as_slice()), + metadata_upload_bytes: std::mem::size_of_val(rope_freqs.as_slice()) + + std::mem::size_of_val(window_attention_bias.as_slice()) + + std::mem::size_of_val(full_attention_bias.as_slice()), + projected_transfer_bytes: std::mem::size_of_val(all_values.as_slice()), + elapsed_seconds, + }, + }) + } +} diff --git a/src/loading/mod.rs b/src/loading/mod.rs index fb2cab9d4..4e899527e 100644 --- a/src/loading/mod.rs +++ b/src/loading/mod.rs @@ -63,6 +63,8 @@ pub(crate) use self::vlm::load_llava_iree_host_preprocessor; pub(crate) use self::vlm::load_qwen2_vl_iree_host_preprocessor; pub use self::vlm::load_qwen3_omni_speech; #[cfg(feature = "xla-iree")] +pub(crate) use self::vlm::load_youtu_vl_iree_host_preprocessor; +#[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..97ea75764 100644 --- a/src/loading/vlm.rs +++ b/src/loading/vlm.rs @@ -135,6 +135,8 @@ pub(crate) use special::{ load_phi4mm_vlm, }; pub(crate) use step3p7::load_step3p7_vl; +#[cfg(feature = "xla-iree")] +pub(crate) use youtu_vl_loader::load_youtu_vl_iree_host_preprocessor; pub(crate) use youtu_vl_loader::load_youtu_vl_vlm; fn read_sanitized_vlm_config(model_path: &Path) -> Result<(String, Value)> { diff --git a/src/loading/vlm_youtu_vl.rs b/src/loading/vlm_youtu_vl.rs index a3e7d7a7b..2151f4e3a 100644 --- a/src/loading/vlm_youtu_vl.rs +++ b/src/loading/vlm_youtu_vl.rs @@ -39,13 +39,15 @@ use std::path::Path; use crate::LoadedModel; use crate::models::youtu_vl_lm::{YoutuLanguageModel, YoutuTextConfig, sanitize_text_weights}; +#[cfg(feature = "xla-iree")] +use crate::multimodal::host_preprocessor::{HostPreprocessorError, YoutuVlIreeHostPreprocessor}; use crate::vision::YoutuVLModel; use crate::vision::encoders::youtu_vl::{YoutuVLVisionEncoder, YoutuVisionConfig}; use crate::vision::processors::youtu_vl::YoutuVLProcessor; use super::{ - load_vlm_weights_common, parse_required_vlm_subconfig, parse_vlm_config, - read_sanitized_vlm_config, + load_vlm_weights_common, load_vlm_weights_common_filtered_canonical, + parse_required_vlm_subconfig, parse_vlm_config, read_sanitized_vlm_config, }; /// Default token IDs from upstream `youtu_vl/config.py::ModelConfig`. @@ -149,6 +151,109 @@ pub(crate) fn load_youtu_vl_vlm(model_path: &Path) -> Result { Ok(LoadedModel::YoutuVL(vlm)) } +/// Load only Youtu-VL's checked host processor and text embedding table. The +/// complete vision tower and built-in merger stay resident in IREE. +#[cfg(feature = "xla-iree")] +pub(crate) fn load_youtu_vl_iree_host_preprocessor( + model_path: &Path, + device: &str, +) -> Result { + use mlxcel_core::layers::UnifiedEmbedding; + + let (config_str, full_config) = read_sanitized_vlm_config(model_path) + .map_err(|error| HostPreprocessorError::InvalidConfig(error.to_string()))?; + if full_config + .get("model_type") + .and_then(serde_json::Value::as_str) + != Some("youtu_vl") + { + return Err(HostPreprocessorError::FamilyMismatch { + actual: full_config + .get("model_type") + .and_then(serde_json::Value::as_str) + .unwrap_or("") + .to_string(), + }); + } + let mut text_config: YoutuTextConfig = parse_vlm_config(&config_str, "Youtu-VL text config") + .map_err(|error| HostPreprocessorError::InvalidConfig(error.to_string()))?; + if text_config.quantization.is_none() + && let Some(quantization) = full_config.get("quantization").cloned() + && let Ok(parsed) = + serde_json::from_value::(quantization) + { + text_config.quantization = Some(parsed); + } + let vision_config: YoutuVisionConfig = + parse_required_vlm_subconfig(&full_config, "vision_config", "Youtu-VL vision config") + .map_err(|error| HostPreprocessorError::InvalidConfig(error.to_string()))?; + let raw_embeddings = load_vlm_weights_common_filtered_canonical(model_path, |name| { + name == "model.embed_tokens.weight" + || name.starts_with("model.embed_tokens.") + || name == "language_model.model.embed_tokens.weight" + || name.starts_with("language_model.model.embed_tokens.") + }) + .map_err(|error| HostPreprocessorError::WeightLoad(error.to_string()))?; + let weights = remap_youtu_vl_weights(raw_embeddings); + let (quant_group_size, quant_bits) = text_config + .quantization + .as_ref() + .map(|quantization| (quantization.group_size, quantization.bits)) + .unwrap_or((0, 0)); + let text_embeddings = UnifiedEmbedding::from_weights( + &weights, + "model.embed_tokens", + quant_group_size, + quant_bits, + ) + .map_err(|error| { + HostPreprocessorError::WeightLoad(format!( + "missing or invalid Youtu-VL text embedding table: {error}" + )) + })?; + if text_config.hidden_size == 0 || text_config.max_position_embeddings == 0 { + return Err(HostPreprocessorError::InvalidConfig( + "Youtu-VL hidden_size and max_position_embeddings must be positive".to_string(), + )); + } + let image_token_id = config_token_id(&full_config, "image_token_id", DEFAULT_IMAGE_TOKEN_ID); + let video_token_id = config_token_id(&full_config, "video_token_id", DEFAULT_VIDEO_TOKEN_ID); + let vision_start_token_id = config_token_id( + &full_config, + "vision_start_token_id", + DEFAULT_VISION_START_TOKEN_ID, + ); + let vision_end_token_id = config_token_id( + &full_config, + "vision_end_token_id", + DEFAULT_VISION_END_TOKEN_ID, + ); + let processor = build_processor(model_path, &vision_config); + let projector = mlxcel_xla::IreeYoutuVlProjector::load(model_path, device) + .map_err(HostPreprocessorError::Iree)?; + YoutuVlIreeHostPreprocessor::from_parts( + processor, + text_embeddings, + projector, + image_token_id, + video_token_id, + vision_start_token_id, + vision_end_token_id, + vision_config.spatial_merge_size, + text_config.hidden_size, + text_config.max_position_embeddings, + device.to_string(), + ) +} + +fn config_token_id(config: &serde_json::Value, key: &str, default: i32) -> i32 { + config + .get(key) + .and_then(serde_json::Value::as_i64) + .and_then(|value| i32::try_from(value).ok()) + .unwrap_or(default) +} + fn parse_eos_token_ids(config: &serde_json::Value) -> Vec { match config.get("eos_token_id") { Some(serde_json::Value::Number(n)) => { @@ -269,9 +374,18 @@ fn build_processor(model_path: &Path, vision_config: &YoutuVisionConfig) -> Yout if let (Some(min_p), Some(max_p)) = (min_pixels, max_pixels) { processor = processor.with_pixel_bounds(min_p as usize, max_p as usize); } - if let Some(num_patches) = json.get("num_patches").and_then(|v| v.as_u64()) { + if let Some(num_patches) = json + .get("max_num_patches") + .or_else(|| json.get("num_patches")) + .and_then(|v| v.as_u64()) + { processor = processor.with_max_patches_per_image(num_patches as usize); } + if let Some(resample) = json.get("resample").and_then(|v| v.as_u64()) + && let Ok(resample) = u8::try_from(resample) + { + processor = processor.with_resample(resample); + } processor } diff --git a/src/multimodal/host_preprocessor.rs b/src/multimodal/host_preprocessor.rs index bfad9cb9d..3969adb39 100644 --- a/src/multimodal/host_preprocessor.rs +++ b/src/multimodal/host_preprocessor.rs @@ -42,10 +42,12 @@ use super::vlm_prompt::{ImageTokenBlockError, ImageTokenBlockInfo, apply_image_t mod export; #[cfg(feature = "xla-iree")] use super::qwen_vl::insert_qwen_vl_image_tokens; +#[cfg(feature = "xla-iree")] +use super::youtu_vl_prompt::insert_youtu_vl_image_tokens; use export::export_mlx_tensor; use export::{ - build_prepared_prefill, export_llava_prefill, export_qwen2_vl_prefill, usize_to_i32, - validate_embedding_shape, validate_processor_shape, validate_projected_shape, + build_prepared_prefill, export_llava_prefill, export_qwen2_vl_prefill, export_youtu_vl_prefill, + usize_to_i32, validate_embedding_shape, validate_processor_shape, validate_projected_shape, validate_sequence_capacity, }; @@ -175,6 +177,34 @@ pub fn load_xla_image_preprocessor( )); } } + if model_type == crate::models::ModelType::YoutuVLM { + let policy = XlaVisionBackendPolicy::from_env()?; + if policy == XlaVisionBackendPolicy::Host { + return Err(HostPreprocessorError::InvalidConfig( + "Youtu-VL XLA vision has no MLX fallback; MLXCEL_XLA_VISION_BACKEND=host is unsupported" + .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 = YoutuVlIreeHostPreprocessor::load(model_path, &device)?; + tracing::info!( + vision_backend = "iree", + vision_device = %device, + family = "youtu_vl", + "OpenXLA multimodal vision backend selected" + ); + return Ok(Some(Box::new(preprocessor))); + } + #[cfg(not(feature = "xla-iree"))] + { + return Err(HostPreprocessorError::InvalidConfig( + "Youtu-VL XLA image execution requires the xla-iree feature".to_string(), + )); + } + } if model_type != crate::models::ModelType::LlavaVLM { return Ok(None); } @@ -292,6 +322,59 @@ pub struct Qwen2VlIreeHostPreprocessor { device: String, } +/// Youtu-VL producer with checked flattened-patch preprocessing and the +/// complete windowed vision tower plus built-in merger resident in IREE. +#[cfg(feature = "xla-iree")] +pub struct YoutuVlIreeHostPreprocessor { + processor: crate::vision::processors::youtu_vl::YoutuVLProcessor, + text_embeddings: UnifiedEmbedding, + projector: RefCell, + image_token_id: i32, + video_token_id: i32, + vision_start_token_id: i32, + vision_end_token_id: i32, + spatial_merge_size: usize, + hidden_size: usize, + max_sequence_len: usize, + device: String, +} + +#[cfg(feature = "xla-iree")] +#[allow(clippy::too_many_arguments)] +fn prepare_youtu_vl_prompt( + token_ids: &[i32], + spatial_shapes: &[(i32, i32)], + spatial_merge_size: usize, + image_token_id: i32, + video_token_id: i32, + vision_start_token_id: i32, + vision_end_token_id: i32, +) -> Result<(Vec, i32), HostPreprocessorError> { + let mut logical_tokens = token_ids.to_vec(); + if logical_tokens.contains(&image_token_id) { + return Ok((logical_tokens, image_token_id)); + } + if logical_tokens.contains(&video_token_id) { + // Placeholder compatibility only: the request boundary remains + // image-only until a video preprocessing oracle is qualified. + return Ok((logical_tokens, video_token_id)); + } + insert_youtu_vl_image_tokens( + &mut logical_tokens, + spatial_shapes, + spatial_merge_size, + vision_start_token_id, + vision_end_token_id, + image_token_id, + ) + .ok_or_else(|| { + HostPreprocessorError::InvalidConfig( + "Youtu-VL could not expand image placeholders for the decoded images".to_string(), + ) + })?; + Ok((logical_tokens, image_token_id)) +} + #[cfg(feature = "xla-iree")] impl Qwen2VlIreeHostPreprocessor { pub fn load(model_path: &Path, device: &str) -> Result { @@ -458,6 +541,163 @@ impl HostMultimodalPreprocessor for Qwen2VlIreeHostPreprocessor { } } +#[cfg(feature = "xla-iree")] +impl YoutuVlIreeHostPreprocessor { + pub fn load(model_path: &Path, device: &str) -> Result { + crate::loading::load_youtu_vl_iree_host_preprocessor(model_path, device) + } + + #[allow(clippy::too_many_arguments)] + pub(crate) fn from_parts( + processor: crate::vision::processors::youtu_vl::YoutuVLProcessor, + text_embeddings: UnifiedEmbedding, + projector: mlxcel_xla::IreeYoutuVlProjector, + image_token_id: i32, + video_token_id: i32, + vision_start_token_id: i32, + vision_end_token_id: i32, + spatial_merge_size: usize, + hidden_size: usize, + max_sequence_len: usize, + device: String, + ) -> Result { + if hidden_size == 0 || max_sequence_len == 0 || spatial_merge_size == 0 { + return Err(HostPreprocessorError::InvalidConfig( + "Youtu-VL hidden size, sequence capacity, and spatial merge size must be positive" + .to_string(), + )); + } + if processor.spatial_merge_size != spatial_merge_size { + return Err(HostPreprocessorError::InvalidConfig(format!( + "Youtu-VL processor merge size {} disagrees with model merge size {spatial_merge_size}", + processor.spatial_merge_size + ))); + } + if projector.text_hidden() != hidden_size { + return Err(HostPreprocessorError::InvalidConfig(format!( + "Youtu-VL IREE merger hidden size {} disagrees with text hidden size {hidden_size}", + projector.text_hidden() + ))); + } + Ok(Self { + processor, + text_embeddings, + projector: RefCell::new(projector), + image_token_id, + video_token_id, + vision_start_token_id, + vision_end_token_id, + spatial_merge_size, + hidden_size, + max_sequence_len, + device, + }) + } + + fn prepare_iree( + &self, + token_ids: &[i32], + images: &[DynamicImage], + ) -> Result { + if images.is_empty() { + return Err(HostPreprocessorError::InvalidConfig( + "Youtu-VL image preprocessing requires at least one decoded image; text-only requests use ordinary XLA token prefill" + .to_string(), + )); + } + let (patch_values, spatial_shapes, _patch_width) = self + .processor + .try_preprocess_values_with_spatial(images) + .map_err(|error| HostPreprocessorError::InvalidConfig(error.to_string()))?; + let (logical_tokens, target_token_id) = prepare_youtu_vl_prompt( + token_ids, + &spatial_shapes, + self.spatial_merge_size, + self.image_token_id, + self.video_token_id, + self.vision_start_token_id, + self.vision_end_token_id, + )?; + validate_sequence_capacity(logical_tokens.len(), self.max_sequence_len)?; + 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, + "Youtu-VL text embedding table", + )?; + let mut projector = self.projector.try_borrow_mut().map_err(|_| { + HostPreprocessorError::Iree( + "concurrent/re-entrant Youtu-VL IREE vision invocation is unsupported".to_string(), + ) + })?; + let projection = projector + .project(&patch_values, &spatial_shapes) + .map_err(HostPreprocessorError::Iree)?; + if projection.shape[1] != self.hidden_size { + return Err(HostPreprocessorError::InvalidConfig(format!( + "Youtu-VL IREE merger shape {:?} disagrees with text hidden size {}", + projection.shape, self.hidden_size + ))); + } + tracing::info!( + vision_backend = "iree", + vision_device = %self.device, + image_count = images.len(), + patch_upload_bytes = projection.metrics.patch_upload_bytes, + metadata_upload_bytes = projection.metrics.metadata_upload_bytes, + projected_transfer_bytes = projection.metrics.projected_transfer_bytes, + iree_vision_seconds = projection.metrics.elapsed_seconds, + merged_tokens_per_image = ?projection.merged_tokens_per_image, + "Youtu-VL OpenXLA vision projection completed" + ); + let projected = mlxcel_core::from_slice_f32( + &projection.values, + &[ + usize_to_i32(projection.shape[0], "Youtu-VL projected token count")?, + usize_to_i32(self.hidden_size, "hidden size")?, + ], + ); + let merged = merge_llava(target_token_id, &projected, &text_embeddings, &input_ids); + export_youtu_vl_prefill( + logical_tokens, + InputEmbeddings { + inputs_embeds: mlxcel_core::astype( + &merged.inputs_embeds, + mlxcel_core::dtype::FLOAT32, + ), + attention_mask_4d: merged.attention_mask_4d, + }, + &spatial_shapes, + target_token_id, + self.spatial_merge_size, + self.hidden_size, + ) + } +} + +#[cfg(feature = "xla-iree")] +impl HostMultimodalPreprocessor for YoutuVlIreeHostPreprocessor { + 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 diff --git a/src/multimodal/host_preprocessor_export.rs b/src/multimodal/host_preprocessor_export.rs index e50c288f1..71deb40b8 100644 --- a/src/multimodal/host_preprocessor_export.rs +++ b/src/multimodal/host_preprocessor_export.rs @@ -151,6 +151,67 @@ pub(super) fn export_llava_prefill( ) } +pub(super) fn export_youtu_vl_prefill( + logical_tokens: Vec, + merged: InputEmbeddings, + spatial_shapes: &[(i32, i32)], + target_token_id: i32, + spatial_merge_size: usize, + hidden_size: usize, +) -> Result { + if spatial_merge_size == 0 { + return Err(HostPreprocessorError::InvalidConfig( + "Youtu-VL spatial_merge_size must be positive".to_string(), + )); + } + let merge = i32::try_from(spatial_merge_size).map_err(|_| { + HostPreprocessorError::InvalidConfig( + "Youtu-VL spatial_merge_size exceeds i32::MAX".to_string(), + ) + })?; + let expected_tokens = spatial_shapes.iter().try_fold(0usize, |total, &(h, w)| { + if h <= 0 || w <= 0 || h % merge != 0 || w % merge != 0 { + return Err(HostPreprocessorError::InvalidConfig(format!( + "Youtu-VL spatial shape [{h}, {w}] must be positive and divisible by merge size {merge}" + ))); + } + let image_tokens = usize::try_from((h / merge) * (w / merge)) + .map_err(|_| HostPreprocessorError::ShapeOverflow)?; + total + .checked_add(image_tokens) + .ok_or(HostPreprocessorError::ShapeOverflow) + })?; + let actual_tokens = logical_tokens + .iter() + .filter(|&&token| token == target_token_id) + .count(); + if actual_tokens != expected_tokens { + return Err(HostPreprocessorError::ExpandedLength { + actual: actual_tokens, + expected: expected_tokens, + }); + } + validate_embedding_shape( + &mlxcel_core::array_shape(&merged.inputs_embeds), + logical_tokens.len(), + hidden_size, + "Youtu-VL merged embedding", + )?; + if merged.attention_mask_4d.is_some() { + return Err(HostPreprocessorError::InvalidConfig( + "Youtu-VL uses standard causal masking, not a family-specific 4D mask".to_string(), + )); + } + let embeddings = export_mlx_tensor(&merged.inputs_embeds, "Youtu-VL merged embedding")?; + build_prepared_prefill( + logical_tokens, + embeddings, + spatial_shapes.len(), + expected_tokens, + "youtu_vl", + ) +} + pub(super) fn export_qwen2_vl_prefill( logical_tokens: Vec, merged: InputEmbeddings, diff --git a/src/multimodal/host_preprocessor_tests.rs b/src/multimodal/host_preprocessor_tests.rs index a83e0f032..29f523734 100644 --- a/src/multimodal/host_preprocessor_tests.rs +++ b/src/multimodal/host_preprocessor_tests.rs @@ -17,6 +17,8 @@ use std::sync::Once; use image::DynamicImage; use mlxcel_core::dtype; +#[cfg(feature = "xla-iree")] +use super::prepare_youtu_vl_prompt; use super::{ FakeHostMultimodalPreprocessor, HostMultimodalPreprocessor, HostPreprocessorError, XlaVisionBackend, XlaVisionBackendPolicy, export_llava_prefill, export_mlx_tensor, @@ -43,6 +45,60 @@ fn fake() -> FakeHostMultimodalPreprocessor { } } +#[cfg(feature = "xla-iree")] +#[test] +fn youtu_vl_placeholder_selection_prefers_image_then_video_then_insertion() { + const IMAGE: i32 = 100; + const VIDEO: i32 = 101; + const START: i32 = 102; + const END: i32 = 103; + let shapes = [(4, 4)]; + + let (tokens, target) = + prepare_youtu_vl_prompt(&[1, VIDEO, IMAGE, 2], &shapes, 2, IMAGE, VIDEO, START, END) + .unwrap(); + assert_eq!(target, IMAGE); + assert_eq!(tokens, vec![1, VIDEO, IMAGE, 2]); + + let (tokens, target) = prepare_youtu_vl_prompt( + &[1, VIDEO, VIDEO, VIDEO, VIDEO, 2], + &shapes, + 2, + IMAGE, + VIDEO, + START, + END, + ) + .unwrap(); + assert_eq!(target, VIDEO); + assert!(!tokens.contains(&IMAGE)); + + let (tokens, target) = + prepare_youtu_vl_prompt(&[1, 2], &shapes, 2, IMAGE, VIDEO, START, END).unwrap(); + assert_eq!(target, IMAGE); + assert_eq!(tokens, vec![1, START, IMAGE, IMAGE, IMAGE, IMAGE, END, 2]); +} + +#[cfg(feature = "xla-iree")] +#[test] +#[ignore = "requires YOUTU_VL_MODEL_DIR and the pinned production IREE runtime/compiler"] +fn pinned_youtu_vl_checkpoint_executes_one_iree_vision_bucket() { + let model_dir = + std::env::var("YOUTU_VL_MODEL_DIR").expect("set YOUTU_VL_MODEL_DIR to the checkpoint"); + let device = std::env::var("MLXCEL_XLA_DEVICE").unwrap_or_else(|_| "local-task".to_string()); + let mut projector = + mlxcel_xla::IreeYoutuVlProjector::load(std::path::Path::new(&model_dir), &device).unwrap(); + let patch_width = 3 * 16 * 16; + let patch_rows = (0..16 * patch_width) + .map(|index| (index % 257) as f32 / 128.0 - 1.0) + .collect::>(); + let projection = projector.project(&patch_rows, &[(4, 4)]).unwrap(); + assert_eq!(projection.shape, [4, projector.text_hidden()]); + assert!(projection.values.iter().all(|value| value.is_finite())); + assert_eq!(projection.merged_tokens_per_image, vec![4]); + assert_eq!(projector.active_bucket(), Some(16)); +} + #[test] fn iree_vision_contract_policy_is_explicit_and_strict() { assert_eq!( diff --git a/src/vision/encoders/youtu_vl_window.rs b/src/vision/encoders/youtu_vl_window.rs index bb2e00d14..796fa8e5c 100644 --- a/src/vision/encoders/youtu_vl_window.rs +++ b/src/vision/encoders/youtu_vl_window.rs @@ -152,8 +152,26 @@ pub(super) fn reverse_window_indices(window_index: &[i32]) -> Vec { .collect(); indexed.sort_by_key(|&(v, _)| v); let mut reverse = vec![0i32; window_index.len()]; - for (rank, &(_, orig_idx)) in indexed.iter().enumerate() { - reverse[orig_idx] = rank as i32; + for (original_position, &(_, window_position)) in indexed.iter().enumerate() { + reverse[original_position] = window_position as i32; } reverse } + +#[cfg(test)] +mod tests { + use super::reverse_window_indices; + + #[test] + fn reverse_indices_restore_original_group_order() { + let window_index = vec![0, 1, 4, 5, 2, 3]; + let window_order = vec!["g0", "g1", "g4", "g5", "g2", "g3"]; + let reverse = reverse_window_indices(&window_index); + let restored = reverse + .iter() + .map(|&position| window_order[position as usize]) + .collect::>(); + assert_eq!(reverse, vec![0, 1, 4, 5, 2, 3]); + assert_eq!(restored, vec!["g0", "g1", "g2", "g3", "g4", "g5"]); + } +} diff --git a/src/vision/processors/youtu_vl.rs b/src/vision/processors/youtu_vl.rs index e4d8fd8dd..5a6f138b0 100644 --- a/src/vision/processors/youtu_vl.rs +++ b/src/vision/processors/youtu_vl.rs @@ -79,6 +79,9 @@ pub struct YoutuVLProcessor { pub max_patches_per_image: usize, pub mean: [f32; 3], pub std: [f32; 3], + /// PIL-compatible resize mode. The published Youtu-VL processor uses + /// `resample = 2` (bilinear). + pub resample: u8, } impl YoutuVLProcessor { @@ -95,6 +98,7 @@ impl YoutuVLProcessor { max_patches_per_image: DEFAULT_MAX_PATCHES_PER_IMAGE, mean: [0.5, 0.5, 0.5], std: [0.5, 0.5, 0.5], + resample: 2, } } @@ -115,6 +119,11 @@ impl YoutuVLProcessor { self } + pub fn with_resample(mut self, resample: u8) -> Self { + self.resample = resample; + self + } + fn validate_config(&self) -> Result<(), YoutuVLPreprocessError> { if self.patch_size == 0 { return Err(YoutuVLPreprocessError::InvalidConfig( @@ -157,6 +166,16 @@ impl YoutuVLProcessor { Ok(()) } + fn resize_filter(&self) -> FilterType { + match self.resample { + 0 => FilterType::Nearest, + 3 => FilterType::CatmullRom, + 1 => FilterType::Lanczos3, + // PIL BILINEAR=2. Triangle is image-rs' bilinear kernel. + _ => FilterType::Triangle, + } + } + fn resize_factor(&self) -> u32 { self.patch_size .saturating_mul(self.spatial_merge_size) @@ -218,10 +237,10 @@ impl YoutuVLProcessor { .collect() } - pub fn try_preprocess_with_spatial( + pub fn try_preprocess_values_with_spatial( &self, images: &[DynamicImage], - ) -> Result<(UniquePtr, Vec<(i32, i32)>), YoutuVLPreprocessError> { + ) -> Result<(Vec, Vec<(i32, i32)>, usize), YoutuVLPreprocessError> { self.validate_config()?; let spatial_shapes = self.compute_spatial_shapes(images); @@ -263,7 +282,7 @@ impl YoutuVLProcessor { let target_h = (h_patches as u32) * (self.patch_size as u32); let target_w = (w_patches as u32) * (self.patch_size as u32); - let resized = img.resize_exact(target_w, target_h, FilterType::Lanczos3); + let resized = img.resize_exact(target_w, target_h, self.resize_filter()); let rgb = resized.to_rgb8(); let h = target_h as usize; @@ -308,6 +327,25 @@ impl YoutuVLProcessor { write_offset += total_patches_img; } + i32::try_from(total_patches).map_err(|_| YoutuVLPreprocessError::DimensionTooLarge { + dimension: total_patches, + })?; + i32::try_from(features_per_patch).map_err(|_| { + YoutuVLPreprocessError::DimensionTooLarge { + dimension: features_per_patch, + } + })?; + + Ok((all_patches, spatial_shapes, features_per_patch)) + } + + pub fn try_preprocess_with_spatial( + &self, + images: &[DynamicImage], + ) -> Result<(UniquePtr, Vec<(i32, i32)>), YoutuVLPreprocessError> { + let (all_patches, spatial_shapes, features_per_patch) = + self.try_preprocess_values_with_spatial(images)?; + let total_patches = all_patches.len() / features_per_patch; let total_patches_i32 = i32::try_from(total_patches).map_err(|_| { YoutuVLPreprocessError::DimensionTooLarge { dimension: total_patches, From f203528931c6471ca7959bb42b385c0e12767ea2 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 24 Jul 2026 19:49:48 +0900 Subject: [PATCH 02/14] fix: reject unsupported Youtu-VL attention bias Fail closed when a Youtu-VL MLA checkpoint requests attention projection biases because the emitter does not yet load or apply the q_a, kv_a, and output bias tensors. Strengthen the vision restoration regression with a non-self-inverse window permutation so the prior incorrect permutation cannot satisfy the test accidentally. Validation: cargo fmt --all -- --check; cargo test -p mlxcel-xla --lib youtu_vl_selects_dense_mla_and_interleaved_rope; cargo test --lib reverse_indices_restore_original_group_order. Refs #872 --- src/lib/mlxcel-xla/src/emitter/config.rs | 6 ++++++ src/lib/mlxcel-xla/src/emitter/mod.rs | 24 +++++++++++++++++------- src/vision/encoders/youtu_vl_window.rs | 18 ++++++++++++++---- 3 files changed, 37 insertions(+), 11 deletions(-) diff --git a/src/lib/mlxcel-xla/src/emitter/config.rs b/src/lib/mlxcel-xla/src/emitter/config.rs index ff37d8c8d..71e301152 100644 --- a/src/lib/mlxcel-xla/src/emitter/config.rs +++ b/src/lib/mlxcel-xla/src/emitter/config.rs @@ -763,6 +763,12 @@ impl Config { }); } Some("youtu_vl") => { + if ob("attention_bias") == Some(true) { + return Err("the OpenXLA Youtu-VL MLA emitter does not support \ + attention_bias = true for q_a_proj, kv_a_proj_with_mqa, \ + or o_proj" + .to_string()); + } let parsed = MlaConfig { q_lora_rank: u("q_lora_rank")?, kv_lora_rank: u("kv_lora_rank")?, diff --git a/src/lib/mlxcel-xla/src/emitter/mod.rs b/src/lib/mlxcel-xla/src/emitter/mod.rs index 463d77d12..b769e46a2 100644 --- a/src/lib/mlxcel-xla/src/emitter/mod.rs +++ b/src/lib/mlxcel-xla/src/emitter/mod.rs @@ -1411,17 +1411,16 @@ mod tests { #[test] fn youtu_vl_selects_dense_mla_and_interleaved_rope() { - let config = Config::from_json_str( - r#"{"model_type":"youtu_vl","hidden_size":64,"intermediate_size":128, + let json = r#"{"model_type":"youtu_vl","hidden_size":64,"intermediate_size":128, "num_hidden_layers":2,"num_attention_heads":4,"num_key_value_heads":4, "vocab_size":128,"rms_norm_eps":1e-6,"rope_theta":500000, "kv_lora_rank":16,"q_lora_rank":24,"qk_nope_head_dim":8, "qk_rope_head_dim":4,"v_head_dim":8,"rope_interleave":true, - "tie_word_embeddings":true,"attention_bias":false,"mlp_bias":false}"#, - ) - .unwrap() - .with_context_capacity(16) - .unwrap(); + "tie_word_embeddings":true,"attention_bias":false,"mlp_bias":false}"#; + let config = Config::from_json_str(json) + .unwrap() + .with_context_capacity(16) + .unwrap(); let mla = config.mla.expect("Youtu-VL must select MLA"); assert_eq!(config.head_dim, 12); assert_eq!(config.rotary_width(), 4); @@ -1433,6 +1432,17 @@ mod tests { assert!(prefill.contains("['mla_q_a']")); assert!(prefill.contains("['mla_kv_b']")); assert!(!prefill.contains("['wq']")); + + let biased = json.replace("\"attention_bias\":false", "\"attention_bias\":true"); + let error = Config::from_json_str(&biased) + .expect_err("Youtu-VL attention biases must fail closed until they are emitted"); + assert!( + error.contains("attention_bias") + && error.contains("q_a_proj") + && error.contains("kv_a_proj_with_mqa") + && error.contains("o_proj"), + "unexpected Youtu-VL bias rejection: {error}" + ); } #[test] diff --git a/src/vision/encoders/youtu_vl_window.rs b/src/vision/encoders/youtu_vl_window.rs index 796fa8e5c..1f224f8a9 100644 --- a/src/vision/encoders/youtu_vl_window.rs +++ b/src/vision/encoders/youtu_vl_window.rs @@ -164,14 +164,24 @@ mod tests { #[test] fn reverse_indices_restore_original_group_order() { - let window_index = vec![0, 1, 4, 5, 2, 3]; - let window_order = vec!["g0", "g1", "g4", "g5", "g2", "g3"]; + // A two-row, three-window layout with 2x2 windows. This permutation + // contains a four-cycle, so it cannot accidentally pass as its own inverse. + let window_index = vec![0, 1, 6, 7, 2, 3, 8, 9, 4, 5, 10, 11]; + let window_order = vec![ + "g0", "g1", "g6", "g7", "g2", "g3", "g8", "g9", "g4", "g5", "g10", "g11", + ]; let reverse = reverse_window_indices(&window_index); let restored = reverse .iter() .map(|&position| window_order[position as usize]) .collect::>(); - assert_eq!(reverse, vec![0, 1, 4, 5, 2, 3]); - assert_eq!(restored, vec!["g0", "g1", "g2", "g3", "g4", "g5"]); + assert_ne!(reverse, window_index); + assert_eq!(reverse, vec![0, 1, 4, 5, 8, 9, 2, 3, 6, 7, 10, 11]); + assert_eq!( + restored, + vec![ + "g0", "g1", "g2", "g3", "g4", "g5", "g6", "g7", "g8", "g9", "g10", "g11", + ] + ); } } From 481f7e686b4f9fc514573fd0e3f4b62f4c4102ed Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Sun, 26 Jul 2026 21:49:51 +0900 Subject: [PATCH 03/14] test(xla): add Youtu-VL oracle seams Expose diagnostics-only eager and IREE checkpoints for patch projection, window/full layers, merger window order, and restored output. Add a deterministic two-layer dense MLA fixture for position zero, nonzero positions, replaced image embeddings, and padded IREE cache mapping. Validation: cargo test --lib youtu_vl; cargo test -p mlxcel-xla --lib youtu_vl; cargo clippy --lib --tests --no-deps; cargo check -p mlxcel-xla --features diagnostics. Refs #872 --- src/lib/mlxcel-xla/src/batch.rs | 5 + src/lib/mlxcel-xla/src/emitter/mod.rs | 2 + src/lib/mlxcel-xla/src/emitter/youtu_vl.rs | 88 +++++- src/lib/mlxcel-xla/src/lib.rs | 7 +- src/lib/mlxcel-xla/src/youtu_vl_runtime.rs | 200 ++++++++++++++ src/models/youtu_vl_lm.rs | 8 + src/models/youtu_vl_lm_diagnostics.rs | 306 +++++++++++++++++++++ src/models/youtu_vl_lm_tests.rs | 184 +++++++++++++ src/vision/encoders/youtu_vl.rs | 108 +++++++- src/vision/encoders/youtu_vl_tests.rs | 94 +++++++ 10 files changed, 989 insertions(+), 13 deletions(-) create mode 100644 src/models/youtu_vl_lm_diagnostics.rs diff --git a/src/lib/mlxcel-xla/src/batch.rs b/src/lib/mlxcel-xla/src/batch.rs index b08cb4b2f..fb107420e 100644 --- a/src/lib/mlxcel-xla/src/batch.rs +++ b/src/lib/mlxcel-xla/src/batch.rs @@ -816,6 +816,11 @@ pub struct LlavaReferenceDiagnosticEngine { engine: IreeRaggedLlama, } +/// Youtu-VL uses the same production prepared-embedding diagnostic decoder; +/// the architecture-specific comparison maps its padded MLA K/V rows. +#[cfg(feature = "diagnostics")] +pub type YoutuVlReferenceDiagnosticEngine = LlavaReferenceDiagnosticEngine; + #[cfg(feature = "diagnostics")] impl LlavaReferenceDiagnosticEngine { /// Compile and load the smallest production serve bundle at an explicit diff --git a/src/lib/mlxcel-xla/src/emitter/mod.rs b/src/lib/mlxcel-xla/src/emitter/mod.rs index b769e46a2..67e03f1d6 100644 --- a/src/lib/mlxcel-xla/src/emitter/mod.rs +++ b/src/lib/mlxcel-xla/src/emitter/mod.rs @@ -161,6 +161,8 @@ pub(crate) use vision::emit_vision_diagnostics; pub(crate) use vision_config::{LlavaVisionConfig, VisionActivation, VisionWeightSpec}; #[allow(unused_imports)] pub(crate) use youtu_vl::emit_youtu_vl; +#[cfg(feature = "diagnostics")] +pub(crate) use youtu_vl::emit_youtu_vl_diagnostics; #[allow(unused_imports)] pub(crate) use youtu_vl_plan::{ YoutuVlHostInputs, YoutuVlVisionConfig, prepare_youtu_vl_host_inputs, diff --git a/src/lib/mlxcel-xla/src/emitter/youtu_vl.rs b/src/lib/mlxcel-xla/src/emitter/youtu_vl.rs index 8f1ae5407..e429e3b09 100644 --- a/src/lib/mlxcel-xla/src/emitter/youtu_vl.rs +++ b/src/lib/mlxcel-xla/src/emitter/youtu_vl.rs @@ -244,7 +244,15 @@ fn layer( builder.add(&residual, &down) } -pub(crate) fn emit_youtu_vl(config: &YoutuVlVisionConfig, patch_bucket: usize) -> String { +struct YoutuVlGraph { + args: Args, + builder: Builder, + patch_projection: Val, + layer_outputs: Vec, + projected: Val, +} + +fn build_youtu_vl(config: &YoutuVlVisionConfig, patch_bucket: usize) -> YoutuVlGraph { assert!( YOUTU_VL_PATCH_BUCKETS.contains(&patch_bucket), "unqualified Youtu-VL patch bucket" @@ -271,6 +279,8 @@ pub(crate) fn emit_youtu_vl(config: &YoutuVlVisionConfig, patch_bucket: usize) - "full_attention.bias", ); let mut hidden = linear(&mut builder, &patches, &args.take(), &args.take()); + let patch_projection = hidden.clone(); + let mut layer_outputs = Vec::with_capacity(config.depth); for layer_index in 0..config.depth { let bias = if config.full_attention_layers.contains(&layer_index) { &full_bias @@ -278,6 +288,7 @@ pub(crate) fn emit_youtu_vl(config: &YoutuVlVisionConfig, patch_bucket: usize) - &window_bias }; hidden = layer(&mut builder, &hidden, &mut args, config, &freqs, bias); + layer_outputs.push(hidden.clone()); } hidden = layer_norm( &mut builder, @@ -293,15 +304,67 @@ pub(crate) fn emit_youtu_vl(config: &YoutuVlVisionConfig, patch_bucket: usize) - let merged = gelu_tanh(&mut builder, &merged); let projected = linear(&mut builder, &merged, &args.take(), &args.take()); assert_eq!(args.cursor, specs.len(), "Youtu-VL weight schema drifted"); + YoutuVlGraph { + args, + builder, + patch_projection, + layer_outputs, + projected, + } +} + +pub(crate) fn emit_youtu_vl(config: &YoutuVlVisionConfig, patch_bucket: usize) -> String { + let graph = build_youtu_vl(config, patch_bucket); format!( "module @youtu_vl_vision {{\n func.func public @main({signature}) -> {result_type} {{\n{body} return {result} : {result_type}\n }}\n}}\n", - signature = args.declarations.join(", "), - result_type = projected.ty.render(), - body = builder.body(), - result = projected.name, + signature = graph.args.declarations.join(", "), + result_type = graph.projected.ty.render(), + body = graph.builder.body(), + result = graph.projected.name, ) } +#[cfg(feature = "diagnostics")] +pub(crate) fn emit_youtu_vl_diagnostics( + config: &YoutuVlVisionConfig, + patch_bucket: usize, +) -> Result { + let graph = build_youtu_vl(config, patch_bucket); + if graph.layer_outputs.len() != 2 { + return Err(format!( + "Youtu-VL diagnostic oracle requires exactly two vision layers, got {}", + graph.layer_outputs.len() + )); + } + if config.full_attention_layers != [1] { + return Err(format!( + "Youtu-VL diagnostic oracle requires window layer 0 and full layer 1, got full layers {:?}", + config.full_attention_layers + )); + } + let outputs = [ + &graph.patch_projection, + &graph.layer_outputs[0], + &graph.layer_outputs[1], + &graph.projected, + ]; + let return_values = outputs + .iter() + .map(|value| value.name.as_str()) + .collect::>() + .join(", "); + let return_types = outputs + .iter() + .map(|value| value.ty.render()) + .collect::>() + .join(", "); + Ok(format!( + "module @youtu_vl_vision_diagnostics {{\n func.func public @main({signature}) -> ({return_types}) {{\n{body} return {return_values} : {return_types} loc(\"Youtu-VL diagnostic checkpoints\")\n }}\n}}\n", + signature = graph.args.declarations.join(", "), + body = graph.builder.body(), + )) +} + #[cfg(test)] mod tests { use super::*; @@ -318,4 +381,19 @@ mod tests { assert!(ir.contains("merger.mlp.2.weight")); assert!(ir.contains("tensor<4x12xf32>")); } + + #[cfg(feature = "diagnostics")] + #[test] + fn diagnostic_emitter_exposes_two_layer_oracle_stages() { + let config = YoutuVlVisionConfig::from_json_str( + r#"{"model_type":"youtu_vl","hidden_size":12,"vision_config":{"num_hidden_layers":2,"hidden_size":16,"intermediate_size":32,"num_attention_heads":4,"num_channels":3,"patch_size":2,"spatial_merge_size":2,"window_size":12,"fullatt_block_indexes":[1],"out_hidden_size":12,"num_patches":256}}"#, + ) + .unwrap(); + let ir = emit_youtu_vl_diagnostics(&config, 64).unwrap(); + assert!(ir.contains("module @youtu_vl_vision_diagnostics")); + assert!(ir.contains( + "tensor<64x16xf32>, tensor<64x16xf32>, tensor<64x16xf32>, tensor<16x12xf32>" + )); + assert!(ir.contains("Youtu-VL diagnostic checkpoints")); + } } diff --git a/src/lib/mlxcel-xla/src/lib.rs b/src/lib/mlxcel-xla/src/lib.rs index f4f8fd4a3..d32cd3f83 100644 --- a/src/lib/mlxcel-xla/src/lib.rs +++ b/src/lib/mlxcel-xla/src/lib.rs @@ -147,8 +147,9 @@ pub use batch::{EngineEvent, FinishReason, XlaAdmissionError, XlaBatchEngine, Xl #[cfg(feature = "diagnostics")] pub use batch::{ Gemma3nAllLayerDiagnosticRun, Gemma3nCanonicalDiagnosticRun, Gemma3nPrefixDecodeDiagnosticRun, - LlavaReferenceDiagnosticEngine, LlavaReferenceDiagnosticRun, run_gemma3n_all_layer_diagnostics, - run_gemma3n_canonical_diagnostics, run_gemma3n_prefix_decode_diagnostic, + LlavaReferenceDiagnosticEngine, LlavaReferenceDiagnosticRun, YoutuVlReferenceDiagnosticEngine, + run_gemma3n_all_layer_diagnostics, run_gemma3n_canonical_diagnostics, + run_gemma3n_prefix_decode_diagnostic, }; pub use context::{ CONTEXT_CAPACITY_ENV, ContextCapacityError, DEFAULT_CONTEXT_CAPACITY, @@ -186,6 +187,8 @@ pub use qwen2_vl_runtime::{ pub use vision_runtime::{IreeVisionDiagnosticProjector, VisionDiagnosticProjection}; #[cfg(feature = "iree")] pub use vision_runtime::{IreeVisionProjector, VisionExecutionMetrics, VisionProjection}; +#[cfg(feature = "diagnostics")] +pub use youtu_vl_runtime::{IreeYoutuVlDiagnosticProjector, YoutuVlVisionDiagnosticProjection}; #[cfg(feature = "iree")] pub use youtu_vl_runtime::{ IreeYoutuVlProjector, YoutuVlVisionExecutionMetrics, YoutuVlVisionProjection, diff --git a/src/lib/mlxcel-xla/src/youtu_vl_runtime.rs b/src/lib/mlxcel-xla/src/youtu_vl_runtime.rs index 2edbf3d87..bf139e611 100644 --- a/src/lib/mlxcel-xla/src/youtu_vl_runtime.rs +++ b/src/lib/mlxcel-xla/src/youtu_vl_runtime.rs @@ -18,6 +18,8 @@ use crate::aux::{ IreeAuxiliaryModule, }; use crate::aux_manifest::{AuxiliaryArtifactContract, ensure_qualified_auxiliary_artifact}; +#[cfg(feature = "diagnostics")] +use crate::emitter::emit_youtu_vl_diagnostics; use crate::emitter::{ YoutuVlHostInputs, YoutuVlVisionConfig, emit_youtu_vl, prepare_youtu_vl_host_inputs, }; @@ -28,6 +30,8 @@ use crate::qwen2_vl_runtime::{ }; const ENTRY_NAME: &str = "youtu_vl_vision.main"; +#[cfg(feature = "diagnostics")] +const DIAGNOSTIC_ENTRY_NAME: &str = "youtu_vl_vision_diagnostics.main"; #[derive(Debug, Clone, PartialEq)] pub struct YoutuVlVisionExecutionMetrics { @@ -47,6 +51,25 @@ pub struct YoutuVlVisionProjection { pub metrics: YoutuVlVisionExecutionMetrics, } +#[cfg(feature = "diagnostics")] +#[derive(Debug, Clone, PartialEq)] +pub struct YoutuVlVisionDiagnosticProjection { + /// Linear patch projection in IREE's window order. + pub patch_projection: Vec, + /// Output of window-attention layer 0 in IREE's window order. + pub window_layer0: Vec, + /// Output of full-attention layer 1 in IREE's window order. + pub full_layer1: Vec, + /// Built-in merger output before the host restores original group order. + pub merger_window_order: Vec, + /// Built-in merger output after restoring original per-image group order. + pub restored_output: Vec, + pub patch_shape: [usize; 2], + pub merged_shape: [usize; 2], + pub window_group_index: Vec, + pub reverse_group_index: Vec, +} + struct ActiveModule { patch_bucket: usize, module: IreeAuxiliaryModule, @@ -235,6 +258,42 @@ fn compile_and_load( IreeAuxiliaryModule::load(device, &vmfb, &contract, weights) } +#[cfg(feature = "diagnostics")] +fn compile_and_load_diagnostics( + model_dir: &Path, + device: &str, + config: &YoutuVlVisionConfig, + processor_identity: &str, + patch_bucket: usize, +) -> Result { + let mlir = emit_youtu_vl_diagnostics(config, patch_bucket)?; + let compiler = iree_compile_bin()?; + if !compiler.is_file() { + return Err(format!("iree-compile not found at {}", compiler.display())); + } + let flags = target_flags(device)?; + let cache = std::env::temp_dir().join("mlxcel-xla-youtu-vl-vmfb"); + std::fs::create_dir_all(&cache) + .map_err(|error| format!("mkdir {}: {error}", cache.display()))?; + let (weights, checkpoint_schema) = load_weights(model_dir, config)?; + let contract = AuxiliaryArtifactContract::new( + DIAGNOSTIC_ENTRY_NAME, + format!( + "{};diagnostic_stages=patch,window0,full1,merger;{};checkpoint_schema_sha256={}", + config.fingerprint(patch_bucket), + processor_identity, + sha256_hex(checkpoint_schema.as_bytes()) + ), + compiler_generation_identity(&compiler, flags, &mlir)?, + )?; + let tag = format!("youtu-vl-vision-diagnostics-{patch_bucket}"); + let vmfb = cached_vmfb_path(&compiler, &mlir, flags, &cache, &tag, 0); + ensure_qualified_auxiliary_artifact(&vmfb, &contract, &weights, |temporary| { + compile_one_to(&compiler, &mlir, flags, &cache, &tag, 0, temporary) + })?; + IreeAuxiliaryModule::load(device, &vmfb, &contract, weights) +} + impl IreeYoutuVlProjector { pub fn load(model_dir: &Path, device: &str) -> Result { let config = YoutuVlVisionConfig::from_model_dir(model_dir)?; @@ -354,3 +413,144 @@ impl IreeYoutuVlProjector { }) } } + +#[cfg(feature = "diagnostics")] +pub struct IreeYoutuVlDiagnosticProjector { + model_dir: PathBuf, + device: String, + config: YoutuVlVisionConfig, + processor_identity: String, + active: Option, +} + +#[cfg(feature = "diagnostics")] +impl IreeYoutuVlDiagnosticProjector { + pub fn load(model_dir: &Path, device: &str) -> Result { + let config = YoutuVlVisionConfig::from_model_dir(model_dir)?; + let processor_identity = processor_identity(model_dir, &config)?; + Ok(Self { + model_dir: model_dir.to_path_buf(), + device: device.to_string(), + config, + processor_identity, + active: None, + }) + } + + fn ensure_bucket(&mut self, bucket: usize) -> Result<&mut IreeAuxiliaryModule, String> { + if self.active.as_ref().map(|active| active.patch_bucket) != Some(bucket) { + let module = compile_and_load_diagnostics( + &self.model_dir, + &self.device, + &self.config, + &self.processor_identity, + bucket, + )?; + self.active = Some(ActiveModule { + patch_bucket: bucket, + module, + }); + } + Ok(&mut self.active.as_mut().expect("active module").module) + } + + pub fn capture( + &mut self, + patch_rows: &[f32], + shapes: &[(i32, i32)], + ) -> Result { + let YoutuVlHostInputs { + plan, + patches, + rope_freqs, + window_attention_bias, + full_attention_bias, + } = prepare_youtu_vl_host_inputs(&self.config, shapes, patch_rows)?; + let patch_input_shape = [ + plan.patch_bucket, + self.config.channels * self.config.patch_size * self.config.patch_size, + ]; + let rope_shape = [ + plan.patch_bucket, + self.config.hidden / self.config.heads / 2, + ]; + let bias_shape = [plan.patch_bucket, plan.patch_bucket]; + let patch_shape = [plan.patch_bucket, self.config.hidden]; + let merged_bucket_shape = [plan.patch_bucket / 4, self.config.text_hidden]; + let mut buffers = [ + vec![0u8; patch_shape.iter().product::() * std::mem::size_of::()], + vec![0u8; patch_shape.iter().product::() * std::mem::size_of::()], + vec![0u8; patch_shape.iter().product::() * std::mem::size_of::()], + vec![0u8; merged_bucket_shape.iter().product::() * std::mem::size_of::()], + ]; + let output_shapes = [patch_shape, patch_shape, patch_shape, merged_bucket_shape]; + let mut outputs = buffers + .iter_mut() + .zip(&output_shapes) + .map(|(bytes, shape)| AuxiliaryOutput { + bytes, + dtype: AuxiliaryTensorDType::Float32, + shape, + }) + .collect::>(); + self.ensure_bucket(plan.patch_bucket)?.invoke( + &[ + AuxiliaryInput { + bytes: f32_as_bytes(&patches), + dtype: AuxiliaryTensorDType::Float32, + shape: &patch_input_shape, + }, + AuxiliaryInput { + bytes: f32_as_bytes(&rope_freqs), + dtype: AuxiliaryTensorDType::Float32, + shape: &rope_shape, + }, + AuxiliaryInput { + bytes: f32_as_bytes(&window_attention_bias), + dtype: AuxiliaryTensorDType::Float32, + shape: &bias_shape, + }, + AuxiliaryInput { + bytes: f32_as_bytes(&full_attention_bias), + dtype: AuxiliaryTensorDType::Float32, + shape: &bias_shape, + }, + ], + &mut outputs, + )?; + drop(outputs); + let mut values = buffers + .into_iter() + .enumerate() + .map(|(index, bytes)| { + checked_f32_output(&format!("Youtu-VL IREE diagnostic output {index}"), bytes) + }) + .collect::, _>>()? + .into_iter(); + let patch_projection = values.next().expect("four diagnostic outputs"); + let window_layer0 = values.next().expect("four diagnostic outputs"); + let full_layer1 = values.next().expect("four diagnostic outputs"); + let merger_bucket = values.next().expect("four diagnostic outputs"); + debug_assert!(values.next().is_none()); + let actual_tokens = plan.reverse_group_index.len(); + let actual_values = actual_tokens * self.config.text_hidden; + let merger_window_order = merger_bucket[..actual_values].to_vec(); + let mut restored_output = Vec::with_capacity(actual_values); + for &window_position in &plan.reverse_group_index { + let start = window_position * self.config.text_hidden; + restored_output + .extend_from_slice(&merger_window_order[start..start + self.config.text_hidden]); + } + Ok(YoutuVlVisionDiagnosticProjection { + patch_projection, + window_layer0, + full_layer1, + merger_window_order, + restored_output, + patch_shape, + merged_shape: [actual_tokens, self.config.text_hidden], + window_group_index: plan.window_group_index, + reverse_group_index: plan.reverse_group_index, + }) + } +} diff --git a/src/models/youtu_vl_lm.rs b/src/models/youtu_vl_lm.rs index ccd634eaf..c4d146285 100644 --- a/src/models/youtu_vl_lm.rs +++ b/src/models/youtu_vl_lm.rs @@ -50,10 +50,18 @@ use super::deepseek_v3::DeepSeekV3Attention; #[path = "youtu_vl_lm_config.rs"] mod youtu_vl_lm_config; +#[cfg(any(test, feature = "xla-diagnostics"))] +#[path = "youtu_vl_lm_diagnostics.rs"] +mod youtu_vl_lm_diagnostics; #[path = "youtu_vl_lm_sanitize.rs"] mod youtu_vl_lm_sanitize; pub use youtu_vl_lm_config::{QuantizationConfig, YoutuTextConfig}; +#[cfg(any(test, feature = "xla-diagnostics"))] +pub use youtu_vl_lm_diagnostics::{ + YoutuIreeMlaDiagnosticCapture, YoutuMlaDiagnosticCapture, YoutuMlaParityReport, + compare_youtu_mla_diagnostics, +}; pub use youtu_vl_lm_sanitize::sanitize_text_weights; // Dense SwiGLU MLP — identical layout to other Llama-family MLPs but kept local diff --git a/src/models/youtu_vl_lm_diagnostics.rs b/src/models/youtu_vl_lm_diagnostics.rs new file mode 100644 index 000000000..c63e4b312 --- /dev/null +++ b/src/models/youtu_vl_lm_diagnostics.rs @@ -0,0 +1,306 @@ +// 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. + +//! Diagnostics-only MLX/IREE seams for Youtu-VL's dense MLA decoder. + +use mlxcel_core::{MlxArray, array_shape, array_to_raw_bytes, astype, dtype, eval}; + +use super::YoutuLanguageModel; + +#[derive(Debug, Clone, PartialEq)] +pub struct YoutuMlaDiagnosticCapture { + /// Final-position logits from a production MLX prefill. + pub logits: Vec, + /// Per-layer final-position compressed latent K rows. + pub latent_kv: Vec, + /// Per-layer final-position rotated positional K rows. + pub rotary_kv: Vec, + pub layers: usize, + pub position: usize, + pub kv_lora_rank: usize, + pub qk_rope_head_dim: usize, +} + +#[derive(Debug, Clone, PartialEq)] +pub struct YoutuMlaParityReport { + pub compared_logits: usize, + pub compared_latent_kv: usize, + pub compared_rotary_kv: usize, + pub max_absolute: f32, + pub max_relative: f32, +} + +#[derive(Debug, Clone, PartialEq)] +pub struct YoutuIreeMlaDiagnosticCapture { + pub logits: Vec, + pub kv: Vec, + pub layers: usize, + pub kv_width: usize, +} + +#[cfg(feature = "xla-diagnostics")] +impl From<&mlxcel_xla::PreparedPrefillDiagnostics> for YoutuIreeMlaDiagnosticCapture { + fn from(value: &mlxcel_xla::PreparedPrefillDiagnostics) -> Self { + Self { + logits: value.logits.clone(), + kv: value.kv.clone(), + layers: value.layers, + kv_width: value.kv_width, + } + } +} + +fn array_f32(array: &MlxArray) -> Vec { + let array = astype(array, dtype::FLOAT32); + eval(&array); + array_to_raw_bytes(&array) + .chunks_exact(4) + .map(|bytes| f32::from_ne_bytes(bytes.try_into().expect("four-byte f32 chunk"))) + .collect() +} + +impl YoutuLanguageModel { + /// Capture final-position logits and the native compressed MLA cache seam. + /// + /// A one-token prompt captures position 0; a longer prompt captures a + /// nonzero position. `input_embeddings` may contain image-placeholder + /// replacements produced by the ordinary Youtu-VL host path. + pub fn capture_mla_diagnostics( + &self, + input_ids: &MlxArray, + input_embeddings: Option<&MlxArray>, + ) -> Result { + let input_shape = array_shape(input_ids); + if input_shape.len() != 2 || input_shape[0] != 1 || input_shape[1] <= 0 { + return Err(format!( + "Youtu-VL MLA diagnostics require input_ids [1, sequence], got {input_shape:?}" + )); + } + let sequence = usize::try_from(input_shape[1]) + .map_err(|_| "Youtu-VL diagnostic sequence length does not fit usize".to_string())?; + if let Some(embeddings) = input_embeddings { + let shape = array_shape(embeddings); + if shape != [1, input_shape[1], self.config.hidden_size as i32] { + return Err(format!( + "Youtu-VL diagnostic embeddings must be [1, {sequence}, {}], got {shape:?}", + self.config.hidden_size + )); + } + } + let mut caches = self.make_caches_impl(); + let logits = self.forward_impl(input_ids, input_embeddings, &mut caches, None); + let logits_shape = array_shape(&logits); + let logits_values = array_f32(&logits); + let vocab = self.config.vocab_size; + if logits_shape != [1, input_shape[1], vocab as i32] { + return Err(format!( + "Youtu-VL diagnostic logits must be [1, {sequence}, {vocab}], got {logits_shape:?}" + )); + } + let logits = logits_values[(sequence - 1) * vocab..sequence * vocab].to_vec(); + let mut latent_kv = + Vec::with_capacity(self.config.num_hidden_layers * self.config.kv_lora_rank); + let mut rotary_kv = + Vec::with_capacity(self.config.num_hidden_layers * self.config.qk_rope_head_dim); + for (layer, cache) in caches.iter().enumerate() { + let keys = cache + .keys + .as_deref() + .ok_or_else(|| format!("Youtu-VL diagnostic layer {layer} omitted latent K"))?; + let values = cache + .values + .as_deref() + .ok_or_else(|| format!("Youtu-VL diagnostic layer {layer} omitted rotary K"))?; + let key_shape = array_shape(keys); + let value_shape = array_shape(values); + let valid_keys = key_shape.len() == 4 + && key_shape[0] == 1 + && key_shape[1] == 1 + && key_shape[2] >= input_shape[1] + && key_shape[3] == self.config.kv_lora_rank as i32; + let valid_values = value_shape.len() == 4 + && value_shape[0] == 1 + && value_shape[1] == 1 + && value_shape[2] >= input_shape[1] + && value_shape[3] == self.config.qk_rope_head_dim as i32; + if !valid_keys || !valid_values { + return Err(format!( + "Youtu-VL diagnostic layer {layer} cache shape drifted: K={key_shape:?}, V={value_shape:?}" + )); + } + let keys = array_f32(keys); + let values = array_f32(values); + let key_start = (sequence - 1) * self.config.kv_lora_rank; + let value_start = (sequence - 1) * self.config.qk_rope_head_dim; + latent_kv.extend_from_slice(&keys[key_start..key_start + self.config.kv_lora_rank]); + rotary_kv.extend_from_slice( + &values[value_start..value_start + self.config.qk_rope_head_dim], + ); + } + Ok(YoutuMlaDiagnosticCapture { + logits, + latent_kv, + rotary_kv, + layers: self.config.num_hidden_layers, + position: sequence - 1, + kv_lora_rank: self.config.kv_lora_rank, + qk_rope_head_dim: self.config.qk_rope_head_dim, + }) + } +} + +fn compare_slice( + label: &str, + actual: &[f32], + expected: &[f32], + atol: f32, + rtol: f32, + maxima: &mut (f32, f32), +) -> Result<(), String> { + if actual.len() != expected.len() { + return Err(format!( + "{label} length differs: MLX={} IREE={}", + actual.len(), + expected.len() + )); + } + for (index, (&mlx, &iree)) in actual.iter().zip(expected).enumerate() { + if !mlx.is_finite() || !iree.is_finite() { + return Err(format!( + "{label} contains non-finite value at {index}: MLX={mlx}, IREE={iree}" + )); + } + let absolute = (mlx - iree).abs(); + let relative = absolute / iree.abs().max(f32::MIN_POSITIVE); + maxima.0 = maxima.0.max(absolute); + maxima.1 = maxima.1.max(relative); + if absolute > atol + rtol * iree.abs() { + return Err(format!( + "{label} differs at {index}: MLX={mlx}, IREE={iree}, absolute={absolute}, tolerance={}", + atol + rtol * iree.abs() + )); + } + } + Ok(()) +} + +/// Compare MLX's native `(latent, rotary)` cache against the padded XLA MLA +/// cache layout `[latent, zero-rope]` for K and `[zero-latent, rotary]` for V. +pub fn compare_youtu_mla_diagnostics( + mlx: &YoutuMlaDiagnosticCapture, + iree: &YoutuIreeMlaDiagnosticCapture, + atol: f32, + rtol: f32, +) -> Result { + if atol < 0.0 || rtol < 0.0 || !atol.is_finite() || !rtol.is_finite() { + return Err(format!( + "Youtu-VL diagnostic tolerances must be finite and non-negative, got atol={atol}, rtol={rtol}" + )); + } + if mlx.layers != iree.layers { + return Err(format!( + "Youtu-VL diagnostic layer count differs: MLX={} IREE={}", + mlx.layers, iree.layers + )); + } + let cache_width = mlx.kv_lora_rank + mlx.qk_rope_head_dim; + if iree.kv_width != cache_width { + return Err(format!( + "Youtu-VL IREE diagnostic KV width must be latent+rotary={cache_width}, got {}", + iree.kv_width + )); + } + let expected_iree_kv = iree + .layers + .checked_mul(2) + .and_then(|value| value.checked_mul(cache_width)) + .ok_or_else(|| "Youtu-VL diagnostic IREE KV length overflowed".to_string())?; + if iree.kv.len() != expected_iree_kv { + return Err(format!( + "Youtu-VL diagnostic IREE KV length differs: got {}, expected {expected_iree_kv}", + iree.kv.len() + )); + } + let mut iree_latent = Vec::with_capacity(mlx.latent_kv.len()); + let mut iree_rotary = Vec::with_capacity(mlx.rotary_kv.len()); + for layer in 0..iree.layers { + let layer_start = layer * 2 * cache_width; + let key = &iree.kv[layer_start..layer_start + cache_width]; + let value = &iree.kv[layer_start + cache_width..layer_start + 2 * cache_width]; + if key[mlx.kv_lora_rank..] + .iter() + .chain(value[..mlx.kv_lora_rank].iter()) + .any(|value| *value != 0.0) + { + return Err(format!( + "Youtu-VL IREE diagnostic layer {layer} has nonzero MLA cache padding" + )); + } + iree_latent.extend_from_slice(&key[..mlx.kv_lora_rank]); + iree_rotary.extend_from_slice(&value[mlx.kv_lora_rank..]); + } + let mut maxima = (0.0f32, 0.0f32); + compare_slice( + "Youtu-VL final-position logits", + &mlx.logits, + &iree.logits, + atol, + rtol, + &mut maxima, + )?; + compare_slice( + "Youtu-VL compressed latent K", + &mlx.latent_kv, + &iree_latent, + atol, + rtol, + &mut maxima, + )?; + compare_slice( + "Youtu-VL rotary K", + &mlx.rotary_kv, + &iree_rotary, + atol, + rtol, + &mut maxima, + )?; + Ok(YoutuMlaParityReport { + compared_logits: mlx.logits.len(), + compared_latent_kv: mlx.latent_kv.len(), + compared_rotary_kv: mlx.rotary_kv.len(), + max_absolute: maxima.0, + max_relative: maxima.1, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parity_maps_padded_iree_mla_cache_to_native_mlx_seams() { + let mlx = YoutuMlaDiagnosticCapture { + logits: vec![0.25, -0.5], + latent_kv: vec![1.0, 2.0, 3.0, 4.0], + rotary_kv: vec![5.0, 6.0, 7.0, 8.0], + layers: 2, + position: 1, + kv_lora_rank: 2, + qk_rope_head_dim: 2, + }; + let iree = YoutuIreeMlaDiagnosticCapture { + logits: mlx.logits.clone(), + kv: vec![ + 1.0, 2.0, 0.0, 0.0, 0.0, 0.0, 5.0, 6.0, 3.0, 4.0, 0.0, 0.0, 0.0, 0.0, 7.0, 8.0, + ], + layers: 2, + kv_width: 4, + }; + let report = compare_youtu_mla_diagnostics(&mlx, &iree, 0.0, 0.0).unwrap(); + assert_eq!(report.compared_logits, 2); + assert_eq!(report.compared_latent_kv, 4); + assert_eq!(report.compared_rotary_kv, 4); + } +} diff --git a/src/models/youtu_vl_lm_tests.rs b/src/models/youtu_vl_lm_tests.rs index 7132023ee..68644347e 100644 --- a/src/models/youtu_vl_lm_tests.rs +++ b/src/models/youtu_vl_lm_tests.rs @@ -56,6 +56,144 @@ fn minimal_config() -> YoutuTextConfig { } } +fn tiny_diagnostic_config() -> YoutuTextConfig { + YoutuTextConfig { + hidden_size: 16, + intermediate_size: 32, + num_attention_heads: 2, + num_key_value_heads: Some(2), + kv_lora_rank: 4, + q_lora_rank: 8, + qk_rope_head_dim: 2, + v_head_dim: 2, + qk_nope_head_dim: 2, + ..minimal_config() + } +} + +fn put_small(weights: &mut WeightMap, key: impl Into, shape: &[i32], seed: i32) { + let total = shape + .iter() + .try_fold(1usize, |total, &dimension| { + usize::try_from(dimension) + .ok() + .and_then(|dimension| total.checked_mul(dimension)) + }) + .unwrap(); + let values = (0..total) + .map(|index| ((index as i32 * 7 + seed) % 29 - 14) as f32 * 0.002) + .collect::>(); + let array = mlxcel_core::from_slice_f32(&values, shape); + weights.insert(key.into(), mlxcel_core::copy(&array)); +} + +fn build_tiny_diagnostic_weights(config: &YoutuTextConfig) -> WeightMap { + let mut weights = WeightMap::new(); + put_small( + &mut weights, + "model.embed_tokens.weight", + &[config.vocab_size as i32, config.hidden_size as i32], + 1, + ); + for layer in 0..config.num_hidden_layers { + let prefix = format!("model.layers.{layer}"); + for (name, dimension) in [ + ("input_layernorm.weight", config.hidden_size), + ("post_attention_layernorm.weight", config.hidden_size), + ] { + let values = mlxcel_core::ones(&[dimension as i32], mlxcel_core::dtype::FLOAT32); + weights.insert(format!("{prefix}.{name}"), mlxcel_core::copy(&values)); + } + let attention = format!("{prefix}.self_attn"); + put_small( + &mut weights, + format!("{attention}.q_a_proj.weight"), + &[config.q_lora_rank as i32, config.hidden_size as i32], + 10 + layer as i32, + ); + let norm = mlxcel_core::ones(&[config.q_lora_rank as i32], mlxcel_core::dtype::FLOAT32); + weights.insert( + format!("{attention}.q_a_layernorm.weight"), + mlxcel_core::copy(&norm), + ); + put_small( + &mut weights, + format!("{attention}.q_b_proj.weight"), + &[ + (config.num_attention_heads * (config.qk_nope_head_dim + config.qk_rope_head_dim)) + as i32, + config.q_lora_rank as i32, + ], + 20 + layer as i32, + ); + put_small( + &mut weights, + format!("{attention}.kv_a_proj_with_mqa.weight"), + &[ + (config.kv_lora_rank + config.qk_rope_head_dim) as i32, + config.hidden_size as i32, + ], + 30 + layer as i32, + ); + let norm = mlxcel_core::ones(&[config.kv_lora_rank as i32], mlxcel_core::dtype::FLOAT32); + weights.insert( + format!("{attention}.kv_a_layernorm.weight"), + mlxcel_core::copy(&norm), + ); + put_small( + &mut weights, + format!("{attention}.embed_q.weight"), + &[ + config.num_attention_heads as i32, + config.kv_lora_rank as i32, + config.qk_nope_head_dim as i32, + ], + 40 + layer as i32, + ); + put_small( + &mut weights, + format!("{attention}.unembed_out.weight"), + &[ + config.num_attention_heads as i32, + config.v_head_dim as i32, + config.kv_lora_rank as i32, + ], + 50 + layer as i32, + ); + put_small( + &mut weights, + format!("{attention}.o_proj.weight"), + &[ + config.hidden_size as i32, + (config.num_attention_heads * config.v_head_dim) as i32, + ], + 60 + layer as i32, + ); + let mlp = format!("{prefix}.mlp"); + put_small( + &mut weights, + format!("{mlp}.gate_proj.weight"), + &[config.intermediate_size as i32, config.hidden_size as i32], + 70 + layer as i32, + ); + put_small( + &mut weights, + format!("{mlp}.up_proj.weight"), + &[config.intermediate_size as i32, config.hidden_size as i32], + 80 + layer as i32, + ); + put_small( + &mut weights, + format!("{mlp}.down_proj.weight"), + &[config.hidden_size as i32, config.intermediate_size as i32], + 90 + layer as i32, + ); + } + let norm = mlxcel_core::ones(&[config.hidden_size as i32], mlxcel_core::dtype::FLOAT32); + weights.insert("model.norm.weight".to_string(), mlxcel_core::copy(&norm)); + weights +} + #[test] fn config_defaults_match_upstream() { // Round-trip through serde_json with only the required fields set — @@ -183,3 +321,49 @@ fn kv_b_proj_decompose_returns_error_when_biases_missing() { "error message should identify the layer; got: {msg}" ); } + +#[test] +fn tiny_dense_mla_diagnostics_cover_zero_nonzero_and_replaced_embeddings() { + let config = tiny_diagnostic_config(); + let weights = build_tiny_diagnostic_weights(&config); + let model = YoutuLanguageModel::from_weights(&weights, &config).unwrap(); + + let token0 = mlxcel_core::from_slice_i32(&[1], &[1, 1]); + let position0 = model.capture_mla_diagnostics(&token0, None).unwrap(); + assert_eq!(position0.position, 0); + assert_eq!(position0.layers, 2); + assert_eq!(position0.logits.len(), config.vocab_size); + assert_eq!( + position0.latent_kv.len(), + config.num_hidden_layers * config.kv_lora_rank + ); + assert_eq!( + position0.rotary_kv.len(), + config.num_hidden_layers * config.qk_rope_head_dim + ); + + let tokens = mlxcel_core::from_slice_i32(&[1, 3], &[1, 2]); + let position1 = model.capture_mla_diagnostics(&tokens, None).unwrap(); + assert_eq!(position1.position, 1); + assert_ne!(position0.rotary_kv, position1.rotary_kv); + + let text_embeddings = model.get_embed_tokens(&tokens); + let replacements = mlxcel_core::from_slice_f32( + &(0..config.hidden_size) + .map(|index| 0.25 + index as f32 * 0.01) + .collect::>(), + &[1, 1, config.hidden_size as i32], + ); + let prefix = mlxcel_core::slice( + &text_embeddings, + &[0, 0, 0], + &[1, 1, config.hidden_size as i32], + ); + let replaced = mlxcel_core::concatenate(&prefix, &replacements, 1); + let image_position = model + .capture_mla_diagnostics(&tokens, Some(&replaced)) + .unwrap(); + assert_eq!(image_position.position, 1); + assert_ne!(position1.logits, image_position.logits); + assert_ne!(position1.latent_kv, image_position.latent_kv); +} diff --git a/src/vision/encoders/youtu_vl.rs b/src/vision/encoders/youtu_vl.rs index 9e0276492..7d86a8159 100644 --- a/src/vision/encoders/youtu_vl.rs +++ b/src/vision/encoders/youtu_vl.rs @@ -179,6 +179,28 @@ pub struct YoutuVLVisionEncoder { hidden_size: i32, } +#[cfg_attr(not(any(test, feature = "xla-diagnostics")), allow(dead_code))] +struct YoutuVlVisionForward { + restored_output: UniquePtr, + patch_projection: Option>, + layer_outputs: Vec>, + merger_window_order: Option>, + window_index: Vec, + reverse_indices: Vec, +} + +/// Diagnostics-only eager checkpoints for the deterministic two-layer oracle. +#[cfg(any(test, feature = "xla-diagnostics"))] +pub struct YoutuVlVisionDiagnostics { + pub patch_projection: UniquePtr, + pub window_layer0: UniquePtr, + pub full_layer1: UniquePtr, + pub merger_window_order: UniquePtr, + pub restored_output: UniquePtr, + pub window_index: Vec, + pub reverse_indices: Vec, +} + impl YoutuVLVisionEncoder { pub fn from_weights( weights: &WeightMap, @@ -350,14 +372,12 @@ impl YoutuVLVisionEncoder { ) } - /// Forward pass with explicit `spatial_shapes`. Each entry is `(h, w)` — - /// number of patches along the height and width dimensions for one image - /// (this corresponds to upstream's `spatial_shapes[:, 0]`, `spatial_shapes[:, 1]`). - pub fn forward_with_spatial( + fn forward_internal( &self, pixel_values: &MlxArray, spatial_shapes: &[(i32, i32)], - ) -> VisionEncoderOutput { + capture_diagnostics: bool, + ) -> YoutuVlVisionForward { // Embed → [total_tokens, hidden] let mut h = self.embeddings.forward(pixel_values); @@ -377,6 +397,7 @@ impl YoutuVLVisionEncoder { let win_idx_arr = mlxcel_core::from_slice_i32(&window_index, &[window_index.len() as i32]); let h_reordered = mlxcel_core::take(&h_grouped, &win_idx_arr, 0); h = mlxcel_core::reshape(&h_reordered, &[-1, dim]); + let patch_projection = capture_diagnostics.then(|| mlxcel_core::copy(&h)); // Reorder rotary frequencies the same way. let rope_shape = mlxcel_core::array_shape(&rotary_pos_emb); @@ -405,6 +426,11 @@ impl YoutuVLVisionEncoder { } // Run encoder blocks, swapping cu_seqlens for the windowed/full layers. + let mut layer_outputs = Vec::with_capacity(if capture_diagnostics { + self.layers.len() + } else { + 0 + }); for (layer_num, layer) in self.layers.iter().enumerate() { let cu_seqlens_now = if self.fullatt_block_indexes.contains(&layer_num) { full_cu.as_slice() @@ -412,11 +438,15 @@ impl YoutuVLVisionEncoder { cu_window_seqlens.as_slice() }; h = layer.forward(&h, cu_seqlens_now, &cos, &sin); + if capture_diagnostics { + layer_outputs.push(mlxcel_core::copy(&h)); + } } // Post-LN over the encoder output, then run the patch merger. h = self.post_layernorm.forward(&h); h = self.merger.forward(&h); + let merger_window_order = capture_diagnostics.then(|| mlxcel_core::copy(&h)); // Reverse the window reordering. Equivalent to `argsort(window_index)` // (Python upstream); we compute the inverse permutation directly so we @@ -427,7 +457,73 @@ impl YoutuVLVisionEncoder { mlxcel_core::from_slice_i32(&reverse_indices, &[reverse_indices.len() as i32]); h = mlxcel_core::take(&h, &reverse_arr, 0); - VisionEncoderOutput { hidden_states: h } + YoutuVlVisionForward { + restored_output: h, + patch_projection, + layer_outputs, + merger_window_order, + window_index, + reverse_indices, + } + } + + /// Forward pass with explicit `spatial_shapes`. Each entry is `(h, w)` — + /// number of patches along the height and width dimensions for one image + /// (this corresponds to upstream's `spatial_shapes[:, 0]`, `spatial_shapes[:, 1]`). + pub fn forward_with_spatial( + &self, + pixel_values: &MlxArray, + spatial_shapes: &[(i32, i32)], + ) -> VisionEncoderOutput { + VisionEncoderOutput { + hidden_states: self + .forward_internal(pixel_values, spatial_shapes, false) + .restored_output, + } + } + + /// Capture the eager stages used by the two-layer Youtu-VL MLX/IREE oracle. + /// + /// This path is excluded from production builds. It rejects deeper or + /// differently-routed towers so stage names cannot silently describe the + /// wrong layer semantics. + #[cfg(any(test, feature = "xla-diagnostics"))] + pub fn forward_with_spatial_diagnostics( + &self, + pixel_values: &MlxArray, + spatial_shapes: &[(i32, i32)], + ) -> Result { + if self.layers.len() != 2 || self.fullatt_block_indexes != [1] { + return Err(format!( + "Youtu-VL eager diagnostic oracle requires two layers with fullatt_block_indexes=[1], got depth={} and {:?}", + self.layers.len(), + self.fullatt_block_indexes + )); + } + let forward = self.forward_internal(pixel_values, spatial_shapes, true); + let mut layers = forward.layer_outputs.into_iter(); + let window_layer0 = layers + .next() + .ok_or_else(|| "Youtu-VL diagnostic omitted window layer 0".to_string())?; + let full_layer1 = layers + .next() + .ok_or_else(|| "Youtu-VL diagnostic omitted full layer 1".to_string())?; + if layers.next().is_some() { + return Err("Youtu-VL diagnostic produced unexpected extra layers".to_string()); + } + Ok(YoutuVlVisionDiagnostics { + patch_projection: forward + .patch_projection + .ok_or_else(|| "Youtu-VL diagnostic omitted patch projection".to_string())?, + window_layer0, + full_layer1, + merger_window_order: forward + .merger_window_order + .ok_or_else(|| "Youtu-VL diagnostic omitted merger output".to_string())?, + restored_output: forward.restored_output, + window_index: forward.window_index, + reverse_indices: forward.reverse_indices, + }) } } diff --git a/src/vision/encoders/youtu_vl_tests.rs b/src/vision/encoders/youtu_vl_tests.rs index 481bc3382..acdc330bc 100644 --- a/src/vision/encoders/youtu_vl_tests.rs +++ b/src/vision/encoders/youtu_vl_tests.rs @@ -45,6 +45,25 @@ fn small_config() -> YoutuVisionConfig { } } +fn tiny_oracle_config() -> YoutuVisionConfig { + YoutuVisionConfig { + hidden_size: 16, + out_hidden_size: 12, + intermediate_size: 32, + num_hidden_layers: 2, + num_attention_heads: 4, + num_channels: 3, + num_patches: 64, + patch_size: 2, + spatial_merge_size: 2, + // Three merged groups per edge makes the `(4x8, 8x4)` packed fixture + // exercise a non-self-inverse window permutation. + window_size: 12, + fullatt_block_indexes: vec![1], + ..small_config() + } +} + fn put(weights: &mut WeightMap, key: &str, shape: &[i32]) { let total: i64 = shape.iter().map(|&d| d as i64).product(); let arr = mlxcel_core::arange_f32(0.0, total as f32, 1.0); @@ -126,6 +145,15 @@ fn build_synthetic_weights(prefix: &str, config: &YoutuVisionConfig) -> WeightMa w } +fn array_to_vec_f32(array: &MlxArray) -> Vec { + let array = mlxcel_core::astype(array, mlxcel_core::dtype::FLOAT32); + mlxcel_core::eval(&array); + mlxcel_core::array_to_raw_bytes(&array) + .chunks_exact(4) + .map(|bytes| f32::from_ne_bytes(bytes.try_into().unwrap())) + .collect() +} + #[test] fn config_defaults_match_upstream() { let raw = serde_json::json!({}); @@ -179,6 +207,72 @@ fn window_index_partitions_token_set() { } } +#[test] +fn tiny_two_layer_oracle_exposes_non_vacuous_window_and_restore_stages() { + let config = tiny_oracle_config(); + let prefix = "vision_tower"; + let weights = build_synthetic_weights(prefix, &config); + let encoder = YoutuVLVisionEncoder::from_weights(&weights, &config, prefix).unwrap(); + let patch_width = config.num_channels * config.patch_size * config.patch_size; + let mut values = vec![0.0f32; 64 * patch_width]; + for patch in 0..64 { + for column in 0..patch_width { + values[patch * patch_width + column] = patch as f32 * 0.01 + column as f32 * 0.0001; + } + } + let patches = mlxcel_core::from_slice_f32(&values, &[64, patch_width as i32]); + let capture = encoder + .forward_with_spatial_diagnostics(&patches, &[(4, 8), (8, 4)]) + .unwrap(); + + assert_eq!( + mlxcel_core::array_shape(&capture.patch_projection), + vec![64, 16] + ); + assert_eq!( + mlxcel_core::array_shape(&capture.window_layer0), + vec![64, 16] + ); + assert_eq!(mlxcel_core::array_shape(&capture.full_layer1), vec![64, 16]); + assert_eq!( + mlxcel_core::array_shape(&capture.merger_window_order), + vec![16, 12] + ); + assert_eq!( + mlxcel_core::array_shape(&capture.restored_output), + vec![16, 12] + ); + assert_ne!( + capture.window_index, + (0..capture.window_index.len() as i32).collect::>() + ); + assert_ne!(capture.window_index, capture.reverse_indices); + + let patch_projection = array_to_vec_f32(&capture.patch_projection); + let layer0 = array_to_vec_f32(&capture.window_layer0); + let layer1 = array_to_vec_f32(&capture.full_layer1); + let merger_window = array_to_vec_f32(&capture.merger_window_order); + let restored = array_to_vec_f32(&capture.restored_output); + for stage in [ + &patch_projection, + &layer0, + &layer1, + &merger_window, + &restored, + ] { + assert!(stage.iter().all(|value| value.is_finite())); + } + assert_ne!(patch_projection, layer0); + assert_ne!(layer0, layer1); + assert_ne!(merger_window, restored); + for (original_group, &window_position) in capture.reverse_indices.iter().enumerate() { + assert_eq!( + &restored[original_group * 12..original_group * 12 + 12], + &merger_window[window_position as usize * 12..window_position as usize * 12 + 12] + ); + } +} + #[test] fn window_index_handles_padding() { // h=10, w=8 with merger_window_size=4 forces vertical padding (10 not From e595c73bce58f09efd0c3b177ce257d3909cd2e8 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 00:16:14 +0900 Subject: [PATCH 04/14] test(xla): generalize Youtu-VL diagnostic ABI --- src/lib/mlxcel-xla/src/emitter/mod.rs | 2 + src/lib/mlxcel-xla/src/emitter/youtu_vl.rs | 115 ++++++++++++++---- .../mlxcel-xla/src/emitter/youtu_vl_plan.rs | 84 +++++++++++++ src/lib/mlxcel-xla/src/lib.rs | 7 +- src/lib/mlxcel-xla/src/youtu_vl_runtime.rs | 93 ++++++++------ 5 files changed, 245 insertions(+), 56 deletions(-) diff --git a/src/lib/mlxcel-xla/src/emitter/mod.rs b/src/lib/mlxcel-xla/src/emitter/mod.rs index 67e03f1d6..10edf44d6 100644 --- a/src/lib/mlxcel-xla/src/emitter/mod.rs +++ b/src/lib/mlxcel-xla/src/emitter/mod.rs @@ -163,6 +163,8 @@ pub(crate) use vision_config::{LlavaVisionConfig, VisionActivation, VisionWeight pub(crate) use youtu_vl::emit_youtu_vl; #[cfg(feature = "diagnostics")] pub(crate) use youtu_vl::emit_youtu_vl_diagnostics; +#[cfg(feature = "diagnostics")] +pub use youtu_vl_plan::YOUTU_VL_DIAGNOSTIC_ABI_VERSION; #[allow(unused_imports)] pub(crate) use youtu_vl_plan::{ YoutuVlHostInputs, YoutuVlVisionConfig, prepare_youtu_vl_host_inputs, diff --git a/src/lib/mlxcel-xla/src/emitter/youtu_vl.rs b/src/lib/mlxcel-xla/src/emitter/youtu_vl.rs index e429e3b09..06a32f2e5 100644 --- a/src/lib/mlxcel-xla/src/emitter/youtu_vl.rs +++ b/src/lib/mlxcel-xla/src/emitter/youtu_vl.rs @@ -6,6 +6,8 @@ //! StableHLO Youtu-VL windowed vision tower and built-in patch merger. use super::builder::{Builder, Ty, Val}; +#[cfg(any(test, feature = "diagnostics"))] +use super::youtu_vl_plan::YoutuVlDiagnosticStage; use super::youtu_vl_plan::{YOUTU_VL_PATCH_BUCKETS, YoutuVlVisionConfig, YoutuVlWeightSpec}; struct Args { @@ -249,6 +251,7 @@ struct YoutuVlGraph { builder: Builder, patch_projection: Val, layer_outputs: Vec, + post_layernorm: Val, projected: Val, } @@ -297,6 +300,7 @@ fn build_youtu_vl(config: &YoutuVlVisionConfig, patch_bucket: usize) -> YoutuVlG &args.take(), config.layer_norm_eps, ); + let post_layernorm = hidden.clone(); hidden = rms_norm(&mut builder, &hidden, &args.take(), 1e-6); let merged_width = config.hidden * config.spatial_merge_size * config.spatial_merge_size; let merged = builder.reshape(&hidden, vec![patch_bucket / 4, merged_width]); @@ -309,6 +313,7 @@ fn build_youtu_vl(config: &YoutuVlVisionConfig, patch_bucket: usize) -> YoutuVlG builder, patch_projection, layer_outputs, + post_layernorm, projected, } } @@ -324,30 +329,22 @@ pub(crate) fn emit_youtu_vl(config: &YoutuVlVisionConfig, patch_bucket: usize) - ) } -#[cfg(feature = "diagnostics")] +#[cfg(any(test, feature = "diagnostics"))] pub(crate) fn emit_youtu_vl_diagnostics( config: &YoutuVlVisionConfig, patch_bucket: usize, ) -> Result { let graph = build_youtu_vl(config, patch_bucket); - if graph.layer_outputs.len() != 2 { - return Err(format!( - "Youtu-VL diagnostic oracle requires exactly two vision layers, got {}", - graph.layer_outputs.len() - )); - } - if config.full_attention_layers != [1] { - return Err(format!( - "Youtu-VL diagnostic oracle requires window layer 0 and full layer 1, got full layers {:?}", - config.full_attention_layers - )); - } - let outputs = [ - &graph.patch_projection, - &graph.layer_outputs[0], - &graph.layer_outputs[1], - &graph.projected, - ]; + let stages = config.diagnostic_stages(); + let outputs = stages + .iter() + .map(|stage| match stage { + YoutuVlDiagnosticStage::PatchProjection => &graph.patch_projection, + YoutuVlDiagnosticStage::Layer { index, .. } => &graph.layer_outputs[*index], + YoutuVlDiagnosticStage::PostLayerNorm => &graph.post_layernorm, + YoutuVlDiagnosticStage::MergerWindowOrder => &graph.projected, + }) + .collect::>(); let return_values = outputs .iter() .map(|value| value.name.as_str()) @@ -369,6 +366,13 @@ pub(crate) fn emit_youtu_vl_diagnostics( mod tests { use super::*; + fn actual_schedule_config() -> YoutuVlVisionConfig { + YoutuVlVisionConfig::from_json_str( + r#"{"model_type":"youtu_vl","hidden_size":2560,"vision_config":{"num_hidden_layers":27,"hidden_size":1152,"intermediate_size":4304,"num_attention_heads":16,"num_channels":3,"patch_size":16,"spatial_merge_size":2,"window_size":256,"fullatt_block_indexes":[7,15,23,26],"out_hidden_size":2560,"num_patches":256}}"#, + ) + .unwrap() + } + #[test] fn emits_window_and_full_attention_inputs() { let config = YoutuVlVisionConfig::from_json_str( @@ -382,7 +386,6 @@ mod tests { assert!(ir.contains("tensor<4x12xf32>")); } - #[cfg(feature = "diagnostics")] #[test] fn diagnostic_emitter_exposes_two_layer_oracle_stages() { let config = YoutuVlVisionConfig::from_json_str( @@ -392,8 +395,78 @@ mod tests { let ir = emit_youtu_vl_diagnostics(&config, 64).unwrap(); assert!(ir.contains("module @youtu_vl_vision_diagnostics")); assert!(ir.contains( - "tensor<64x16xf32>, tensor<64x16xf32>, tensor<64x16xf32>, tensor<16x12xf32>" + "tensor<64x16xf32>, tensor<64x16xf32>, tensor<64x16xf32>, tensor<64x16xf32>, tensor<16x12xf32>" )); assert!(ir.contains("Youtu-VL diagnostic checkpoints")); + assert_eq!( + config.diagnostic_contract_identity(), + "youtu-vl-vision-diagnostics-v2:module=[patch_projection,layer.0.window,layer.1.full,post_layernorm,merger.window_order]:host=[merger.restored_order]" + ); + } + + #[test] + fn diagnostic_emitter_accepts_actual_youtu_vl_schedule() { + let config = actual_schedule_config(); + let stages = config + .diagnostic_stages() + .iter() + .map(YoutuVlDiagnosticStage::name) + .collect::>(); + assert_eq!( + stages, + [ + "patch_projection", + "layer.0.window", + "layer.7.full", + "layer.15.full", + "layer.23.full", + "layer.26.full", + "post_layernorm", + "merger.window_order", + ] + ); + let ir = emit_youtu_vl_diagnostics(&config, 16).unwrap(); + let result_signature = stages + .iter() + .take(stages.len() - 1) + .map(|_| "tensor<16x1152xf32>") + .chain(std::iter::once("tensor<4x2560xf32>")) + .collect::>() + .join(", "); + assert!(ir.contains(&result_signature)); + assert!( + config + .diagnostic_contract_identity() + .contains("layer.26.full") + ); + } + + #[test] + #[ignore = "requires MLXCEL_XLA_IREE_COMPILE pointing at the pinned iree-compile"] + fn actual_schedule_diagnostic_graph_compiles_with_iree_cpu() { + let compiler = + std::env::var("MLXCEL_XLA_IREE_COMPILE").expect("set MLXCEL_XLA_IREE_COMPILE"); + let graph = emit_youtu_vl_diagnostics(&actual_schedule_config(), 16).unwrap(); + let stem = format!("mlxcel-youtu-vl-diagnostics-v2-{}", std::process::id()); + let input = std::env::temp_dir().join(format!("{stem}.mlir")); + let output = std::env::temp_dir().join(format!("{stem}.vmfb")); + std::fs::write(&input, graph).expect("write temporary StableHLO"); + let result = std::process::Command::new(compiler) + .arg("--iree-input-type=stablehlo") + .arg("--iree-hal-target-device=local") + .arg("--iree-hal-local-target-device-backends=llvm-cpu") + .arg(&input) + .arg("-o") + .arg(&output) + .output() + .expect("run iree-compile"); + let _ = std::fs::remove_file(&input); + let _ = std::fs::remove_file(&output); + assert!( + result.status.success(), + "iree-compile failed:\nstdout:\n{}\nstderr:\n{}", + String::from_utf8_lossy(&result.stdout), + String::from_utf8_lossy(&result.stderr) + ); } } diff --git a/src/lib/mlxcel-xla/src/emitter/youtu_vl_plan.rs b/src/lib/mlxcel-xla/src/emitter/youtu_vl_plan.rs index f9919b461..e88809979 100644 --- a/src/lib/mlxcel-xla/src/emitter/youtu_vl_plan.rs +++ b/src/lib/mlxcel-xla/src/emitter/youtu_vl_plan.rs @@ -14,6 +14,8 @@ use std::path::Path; use serde_json::Value; pub(crate) const YOUTU_VL_PATCH_BUCKETS: [usize; 3] = [16, 64, 256]; +#[cfg(any(test, feature = "diagnostics"))] +pub const YOUTU_VL_DIAGNOSTIC_ABI_VERSION: u32 = 2; const MAX_IMAGES: usize = 4; const DEFAULT_IMAGE_TOKEN_ID: i32 = 128_264; const DEFAULT_VIDEO_TOKEN_ID: i32 = 128_265; @@ -47,6 +49,52 @@ pub(crate) struct YoutuVlVisionConfig { pub(crate) vision_end_token_id: i32, } +/// Ordered module outputs for the diagnostics-only Youtu-VL vision ABI. +/// +/// ABI v2 selects layers from the checkpoint configuration instead of assuming +/// the original two-layer synthetic fixture. Layer zero is always captured as +/// the first encoder seam; every configured full-attention layer is captured in +/// execution order, without duplicating layer zero when it is itself full. +#[cfg(any(test, feature = "diagnostics"))] +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum YoutuVlDiagnosticStage { + PatchProjection, + Layer { index: usize, full_attention: bool }, + PostLayerNorm, + MergerWindowOrder, +} + +#[cfg(any(test, feature = "diagnostics"))] +impl YoutuVlDiagnosticStage { + pub(crate) fn name(&self) -> String { + match self { + Self::PatchProjection => "patch_projection".to_string(), + Self::Layer { + index, + full_attention, + } => format!( + "layer.{index}.{}", + if *full_attention { "full" } else { "window" } + ), + Self::PostLayerNorm => "post_layernorm".to_string(), + Self::MergerWindowOrder => "merger.window_order".to_string(), + } + } + + pub(crate) fn shape(&self, config: &YoutuVlVisionConfig, patch_bucket: usize) -> [usize; 2] { + match self { + Self::MergerWindowOrder => [patch_bucket / 4, config.text_hidden], + Self::PatchProjection | Self::Layer { .. } | Self::PostLayerNorm => { + [patch_bucket, config.hidden] + } + } + } + + pub(crate) fn is_merger_window_order(&self) -> bool { + matches!(self, Self::MergerWindowOrder) + } +} + #[derive(Debug, Clone, PartialEq, Eq)] pub(crate) struct YoutuVlGridPlan { pub(crate) patch_bucket: usize, @@ -221,6 +269,42 @@ impl YoutuVlVisionConfig { ) } + #[cfg(any(test, feature = "diagnostics"))] + pub(crate) fn diagnostic_stages(&self) -> Vec { + let mut stages = Vec::with_capacity(self.full_attention_layers.len() + 4); + stages.push(YoutuVlDiagnosticStage::PatchProjection); + for index in 0..self.depth { + let full_attention = self.full_attention_layers.contains(&index); + if index == 0 || full_attention { + stages.push(YoutuVlDiagnosticStage::Layer { + index, + full_attention, + }); + } + } + stages.push(YoutuVlDiagnosticStage::PostLayerNorm); + stages.push(YoutuVlDiagnosticStage::MergerWindowOrder); + stages + } + + /// Stable diagnostics-only ABI identity. + /// + /// The final restored merger value is produced by the host from the + /// module's window-order output and the captured reverse permutation, so it + /// participates in the contract even though it is not an IREE result. + #[cfg(any(test, feature = "diagnostics"))] + pub(crate) fn diagnostic_contract_identity(&self) -> String { + let module_stages = self + .diagnostic_stages() + .iter() + .map(YoutuVlDiagnosticStage::name) + .collect::>() + .join(","); + format!( + "youtu-vl-vision-diagnostics-v{YOUTU_VL_DIAGNOSTIC_ABI_VERSION}:module=[{module_stages}]:host=[merger.restored_order]" + ) + } + pub(crate) fn plan(&self, shapes: &[(i32, i32)]) -> Result { if shapes.is_empty() || shapes.len() > MAX_IMAGES { return Err(format!( diff --git a/src/lib/mlxcel-xla/src/lib.rs b/src/lib/mlxcel-xla/src/lib.rs index d32cd3f83..4b8aa0fc8 100644 --- a/src/lib/mlxcel-xla/src/lib.rs +++ b/src/lib/mlxcel-xla/src/lib.rs @@ -159,6 +159,8 @@ pub use context::{ pub use diagnostic_flags::{ configure_diagnostic_local_task_threads, diagnostic_local_task_threads_are_configured, }; +#[cfg(feature = "diagnostics")] +pub use emitter::YOUTU_VL_DIAGNOSTIC_ABI_VERSION; #[cfg(all(feature = "diagnostics", xla_iree_cuda))] pub use emitter::{ run_gemma3n_altup_correct_diagnostic_probe, run_gemma3n_altup_predict_diagnostic_probe, @@ -188,7 +190,10 @@ pub use vision_runtime::{IreeVisionDiagnosticProjector, VisionDiagnosticProjecti #[cfg(feature = "iree")] pub use vision_runtime::{IreeVisionProjector, VisionExecutionMetrics, VisionProjection}; #[cfg(feature = "diagnostics")] -pub use youtu_vl_runtime::{IreeYoutuVlDiagnosticProjector, YoutuVlVisionDiagnosticProjection}; +pub use youtu_vl_runtime::{ + IreeYoutuVlDiagnosticProjector, YoutuVlVisionDiagnosticProjection, + YoutuVlVisionDiagnosticStageOutput, +}; #[cfg(feature = "iree")] pub use youtu_vl_runtime::{ IreeYoutuVlProjector, YoutuVlVisionExecutionMetrics, YoutuVlVisionProjection, diff --git a/src/lib/mlxcel-xla/src/youtu_vl_runtime.rs b/src/lib/mlxcel-xla/src/youtu_vl_runtime.rs index bf139e611..39bde2842 100644 --- a/src/lib/mlxcel-xla/src/youtu_vl_runtime.rs +++ b/src/lib/mlxcel-xla/src/youtu_vl_runtime.rs @@ -19,6 +19,8 @@ use crate::aux::{ }; use crate::aux_manifest::{AuxiliaryArtifactContract, ensure_qualified_auxiliary_artifact}; #[cfg(feature = "diagnostics")] +use crate::emitter::YOUTU_VL_DIAGNOSTIC_ABI_VERSION; +#[cfg(feature = "diagnostics")] use crate::emitter::emit_youtu_vl_diagnostics; use crate::emitter::{ YoutuVlHostInputs, YoutuVlVisionConfig, emit_youtu_vl, prepare_youtu_vl_host_inputs, @@ -51,19 +53,24 @@ pub struct YoutuVlVisionProjection { pub metrics: YoutuVlVisionExecutionMetrics, } +#[cfg(feature = "diagnostics")] +#[derive(Debug, Clone, PartialEq)] +pub struct YoutuVlVisionDiagnosticStageOutput { + /// Stable ABI v2 stage name, such as `layer.7.full`. + pub name: String, + pub values: Vec, + pub shape: [usize; 2], +} + #[cfg(feature = "diagnostics")] #[derive(Debug, Clone, PartialEq)] pub struct YoutuVlVisionDiagnosticProjection { - /// Linear patch projection in IREE's window order. - pub patch_projection: Vec, - /// Output of window-attention layer 0 in IREE's window order. - pub window_layer0: Vec, - /// Output of full-attention layer 1 in IREE's window order. - pub full_layer1: Vec, - /// Built-in merger output before the host restores original group order. - pub merger_window_order: Vec, - /// Built-in merger output after restoring original per-image group order. - pub restored_output: Vec, + /// Explicitly versioned because ABI v1 exposed four hard-coded fields for a + /// two-layer synthetic tower. ABI v2 is an ordered, config-selected stage + /// vector and additionally captures post-layernorm plus restored merger + /// output. + pub abi_version: u32, + pub stages: Vec, pub patch_shape: [usize; 2], pub merged_shape: [usize; 2], pub window_group_index: Vec, @@ -279,14 +286,16 @@ fn compile_and_load_diagnostics( let contract = AuxiliaryArtifactContract::new( DIAGNOSTIC_ENTRY_NAME, format!( - "{};diagnostic_stages=patch,window0,full1,merger;{};checkpoint_schema_sha256={}", + "{};diagnostic_contract={};{};checkpoint_schema_sha256={}", config.fingerprint(patch_bucket), + config.diagnostic_contract_identity(), processor_identity, sha256_hex(checkpoint_schema.as_bytes()) ), compiler_generation_identity(&compiler, flags, &mlir)?, )?; - let tag = format!("youtu-vl-vision-diagnostics-{patch_bucket}"); + let tag = + format!("youtu-vl-vision-diagnostics-v{YOUTU_VL_DIAGNOSTIC_ABI_VERSION}-{patch_bucket}"); let vmfb = cached_vmfb_path(&compiler, &mlir, flags, &cache, &tag, 0); ensure_qualified_auxiliary_artifact(&vmfb, &contract, &weights, |temporary| { compile_one_to(&compiler, &mlir, flags, &cache, &tag, 0, temporary) @@ -476,14 +485,15 @@ impl IreeYoutuVlDiagnosticProjector { ]; let bias_shape = [plan.patch_bucket, plan.patch_bucket]; let patch_shape = [plan.patch_bucket, self.config.hidden]; - let merged_bucket_shape = [plan.patch_bucket / 4, self.config.text_hidden]; - let mut buffers = [ - vec![0u8; patch_shape.iter().product::() * std::mem::size_of::()], - vec![0u8; patch_shape.iter().product::() * std::mem::size_of::()], - vec![0u8; patch_shape.iter().product::() * std::mem::size_of::()], - vec![0u8; merged_bucket_shape.iter().product::() * std::mem::size_of::()], - ]; - let output_shapes = [patch_shape, patch_shape, patch_shape, merged_bucket_shape]; + let module_stages = self.config.diagnostic_stages(); + let output_shapes = module_stages + .iter() + .map(|stage| stage.shape(&self.config, plan.patch_bucket)) + .collect::>(); + let mut buffers = output_shapes + .iter() + .map(|shape| vec![0u8; shape.iter().product::() * std::mem::size_of::()]) + .collect::>(); let mut outputs = buffers .iter_mut() .zip(&output_shapes) @@ -519,34 +529,49 @@ impl IreeYoutuVlDiagnosticProjector { &mut outputs, )?; drop(outputs); - let mut values = buffers + let values = buffers .into_iter() .enumerate() .map(|(index, bytes)| { checked_f32_output(&format!("Youtu-VL IREE diagnostic output {index}"), bytes) }) - .collect::, _>>()? - .into_iter(); - let patch_projection = values.next().expect("four diagnostic outputs"); - let window_layer0 = values.next().expect("four diagnostic outputs"); - let full_layer1 = values.next().expect("four diagnostic outputs"); - let merger_bucket = values.next().expect("four diagnostic outputs"); - debug_assert!(values.next().is_none()); + .collect::, _>>()?; let actual_tokens = plan.reverse_group_index.len(); let actual_values = actual_tokens * self.config.text_hidden; - let merger_window_order = merger_bucket[..actual_values].to_vec(); + let mut stages = module_stages + .iter() + .zip(values) + .zip(output_shapes) + .map(|((stage, mut values), mut shape)| { + if stage.is_merger_window_order() { + values.truncate(actual_values); + shape[0] = actual_tokens; + } + YoutuVlVisionDiagnosticStageOutput { + name: stage.name(), + values, + shape, + } + }) + .collect::>(); + let merger_window_order = &stages + .last() + .expect("diagnostic ABI always ends with merger.window_order") + .values; let mut restored_output = Vec::with_capacity(actual_values); for &window_position in &plan.reverse_group_index { let start = window_position * self.config.text_hidden; restored_output .extend_from_slice(&merger_window_order[start..start + self.config.text_hidden]); } + stages.push(YoutuVlVisionDiagnosticStageOutput { + name: "merger.restored_order".to_string(), + values: restored_output, + shape: [actual_tokens, self.config.text_hidden], + }); Ok(YoutuVlVisionDiagnosticProjection { - patch_projection, - window_layer0, - full_layer1, - merger_window_order, - restored_output, + abi_version: YOUTU_VL_DIAGNOSTIC_ABI_VERSION, + stages, patch_shape, merged_shape: [actual_tokens, self.config.text_hidden], window_group_index: plan.window_group_index, From 9ff3a14e9d8f3f6f3a5e3a407a3774132a730ee3 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 01:03:55 +0900 Subject: [PATCH 05/14] test(xla): add independent Youtu-VL oracle --- Cargo.toml | 7 + examples/xla_youtu_vl_reference_check.rs | 535 +++++++++++ scripts/xla/validate_youtu_vl_reference.sh | 74 ++ .../openxla/test_youtu_vl_reference_oracle.py | 216 +++++ spike/openxla/youtu_vl_reference_oracle.py | 876 ++++++++++++++++++ src/lib.rs | 4 +- src/lib/mlxcel-xla/src/youtu_vl_runtime.rs | 17 + 7 files changed, 1728 insertions(+), 1 deletion(-) create mode 100644 examples/xla_youtu_vl_reference_check.rs create mode 100755 scripts/xla/validate_youtu_vl_reference.sh create mode 100644 spike/openxla/test_youtu_vl_reference_oracle.py create mode 100644 spike/openxla/youtu_vl_reference_oracle.py diff --git a/Cargo.toml b/Cargo.toml index ab506bce0..d31fdf7a4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -329,6 +329,13 @@ required-features = ["xla-diagnostics"] name = "xla_vision_reference_check" required-features = ["xla-diagnostics"] +# Independent pinned HF-eager Youtu-VL oracle runner (#872). Captures the +# production 27-layer vision schedule, prepared embeddings, MLA logits/KV, +# greedy tokens, and reset/reuse lifecycle. +[[example]] +name = "xla_youtu_vl_reference_check" +required-features = ["xla-diagnostics"] + [[example]] name = "xla_aux_smoke" required-features = ["xla-iree"] diff --git a/examples/xla_youtu_vl_reference_check.rs b/examples/xla_youtu_vl_reference_check.rs new file mode 100644 index 000000000..2804c724f --- /dev/null +++ b/examples/xla_youtu_vl_reference_check.rs @@ -0,0 +1,535 @@ +// 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. + +//! Production MLX/IREE capture for the independent pinned Youtu-VL HF oracle. +//! +//! This runner consumes only oracle identity and prompt metadata. It +//! independently tokenizes, preprocesses, executes the actual 27-layer vision +//! graph, prepares multimodal embeddings, captures the 40-layer MLA cache and +//! logits, then repeats the request after an engine reset to prove slot reuse. + +use std::collections::BTreeMap; +use std::fs; +use std::path::{Path, PathBuf}; + +use image::DynamicImage; +use mlxcel::tokenizer::load_tokenizer; +use mlxcel::vision::processors::youtu_vl::YoutuVLProcessor; +use mlxcel::{ + HostMultimodalPreprocessor, PreparedTensorDType, YoutuVlIreeHostPreprocessor, + initialize_runtime, +}; +use mlxcel_xla::{IreeYoutuVlDiagnosticProjector, YoutuVlReferenceDiagnosticEngine}; +use serde_json::{Map, Value, json}; +use sha2::{Digest, Sha256}; + +const CONTRACT: &str = "youtu-vl-hf-eager-oracle-v1"; +const CHECKPOINT_REPO: &str = "tencent/Youtu-VL-4B-Instruct"; +const CHECKPOINT_REVISION: &str = "8d30a0e49662a1d628a472b12df264dbcd768753"; +const CHECKPOINT_ARTIFACT_MANIFEST_SHA256: &str = + "4d67dd2750d1ee8d87e68b0b52f5aa6c5b5b1dd85385df643845e393628812c9"; +const FIXTURE_PATH: &str = "tests/fixtures/test_image.png"; +const FIXTURE_SHA256: &str = "5e7d54e8a7d21802378c87d2d70cf551e29739fe27599ddf129ebccdad1e6261"; +const PROMPT: &str = "<|image_pad|>\nDescribe the image briefly."; +const IMAGE_TOKEN_ID: i32 = 128_264; +const PATCH_SIZE: usize = 16; +const PATCHES: usize = 256; +const PATCH_WIDTH: usize = 3 * PATCH_SIZE * PATCH_SIZE; +const TEXT_HIDDEN: usize = 2_560; +const TEXT_LAYERS: usize = 40; +const KV_WIDTH: usize = 576; +const MAX_NEW_TOKENS: usize = 4; + +fn argument(flag: &str) -> Option { + let args = std::env::args().collect::>(); + args.iter() + .position(|value| value == flag) + .and_then(|index| args.get(index + 1)) + .cloned() +} + +fn required_path(flag: &str) -> PathBuf { + argument(flag) + .map(PathBuf::from) + .unwrap_or_else(|| panic!("missing required {flag}")) +} + +fn sha256_bytes(bytes: &[u8]) -> String { + Sha256::digest(bytes) + .iter() + .map(|byte| format!("{byte:02x}")) + .collect() +} + +fn sha256_file(path: &Path) -> String { + sha256_bytes(&fs::read(path).unwrap_or_else(|error| panic!("read {}: {error}", path.display()))) +} + +fn verify_reference_manifest(root: &Path) -> Value { + let manifest_path = root.join("manifest.json"); + let bytes = fs::read(&manifest_path) + .unwrap_or_else(|error| panic!("read {}: {error}", manifest_path.display())); + let expected = + fs::read_to_string(root.join("manifest.sha256")).expect("read reference manifest.sha256"); + assert_eq!( + sha256_bytes(&bytes), + expected.trim(), + "reference manifest SHA-256 differs" + ); + let manifest: Value = serde_json::from_slice(&bytes).expect("parse reference manifest"); + assert_eq!(manifest["schema"], 1); + assert_eq!(manifest["contract"], CONTRACT); + assert_eq!(manifest["producer"], "hf-transformers-eager"); + assert_eq!(manifest["fixture"]["path"], FIXTURE_PATH); + assert_eq!(manifest["fixture"]["sha256"], FIXTURE_SHA256); + assert_eq!(manifest["case"]["name"], "image_text"); + assert_eq!(manifest["case"]["prompt"], PROMPT); + assert_eq!(manifest["generation"]["mode"], "greedy"); + assert_eq!(manifest["generation"]["max_new_tokens"], MAX_NEW_TOKENS); + assert_eq!(manifest["architecture"]["vision_depth"], 27); + assert_eq!(manifest["architecture"]["text_layers"], TEXT_LAYERS); + assert_eq!(manifest["architecture"]["text_hidden"], TEXT_HIDDEN); + assert_eq!(manifest["architecture"]["kv_lora_rank"], 512); + assert_eq!(manifest["architecture"]["qk_rope_head_dim"], 64); + assert_eq!( + manifest["lifecycle"]["events"], + json!([ + "reset", + "prefill", + "greedy_decode", + "reset_for_reuse", + "prefill_reuse", + "greedy_decode_reuse" + ]) + ); + manifest +} + +fn verify_checkpoint(model: &Path, checkpoint: &Value) { + assert_eq!(checkpoint["repo"], CHECKPOINT_REPO); + assert_eq!(checkpoint["revision"], CHECKPOINT_REVISION); + let artifact_manifest = &checkpoint["artifact_manifest"]; + assert_eq!( + artifact_manifest["canonical_sha256"], + CHECKPOINT_ARTIFACT_MANIFEST_SHA256 + ); + let files = + serde_json::from_value::>(artifact_manifest["files"].clone()) + .expect("checkpoint artifact files must be a string map"); + let canonical = serde_json::to_vec(&files).expect("serialize canonical artifact map"); + assert_eq!( + sha256_bytes(&canonical), + CHECKPOINT_ARTIFACT_MANIFEST_SHA256, + "checkpoint artifact map is not canonical" + ); + for (filename, expected) in &files { + let path = model.join(filename); + assert!(path.is_file(), "checkpoint artifact is missing: {filename}"); + assert_eq!( + sha256_file(&path), + expected.as_str(), + "checkpoint artifact SHA-256 differs: {filename}" + ); + } + let unexpected_weights = fs::read_dir(model) + .expect("read checkpoint directory") + .map(|entry| entry.expect("read checkpoint entry").path()) + .filter(|path| path.extension().and_then(|value| value.to_str()) == Some("safetensors")) + .filter_map(|path| { + let filename = path.file_name()?.to_str()?.to_string(); + (!files.contains_key(&filename)).then_some(filename) + }) + .collect::>(); + assert!( + unexpected_weights.is_empty(), + "checkpoint has unpinned weight artifact(s): {unexpected_weights:?}" + ); +} + +fn f32_bytes(values: &[f32]) -> Vec { + values + .iter() + .flat_map(|value| value.to_le_bytes()) + .collect() +} + +fn i32_bytes(values: &[i32]) -> Vec { + values + .iter() + .flat_map(|value| value.to_le_bytes()) + .collect() +} + +fn write_array(out: &Path, stage: &str, bytes: &[u8], dtype: &str, shape: &[usize]) -> Value { + if dtype == "float32" { + assert!( + bytes + .chunks_exact(4) + .map(|chunk| f32::from_le_bytes(chunk.try_into().expect("four-byte f32"))) + .all(f32::is_finite), + "{stage} contains a non-finite value" + ); + } + assert_eq!( + bytes.len(), + shape.iter().product::() * 4, + "{stage} byte length differs" + ); + let filename = format!("image_text.{stage}.bin"); + let path = out.join(&filename); + fs::write(&path, bytes).unwrap_or_else(|error| panic!("write {}: {error}", path.display())); + json!({ + "file": filename, + "dtype": dtype, + "shape": shape, + "sha256": sha256_file(&path), + }) +} + +fn write_f32( + arrays: &mut Map, + out: &Path, + stage: &str, + values: &[f32], + shape: &[usize], +) { + arrays.insert( + stage.to_string(), + write_array(out, stage, &f32_bytes(values), "float32", shape), + ); +} + +fn write_i32( + arrays: &mut Map, + out: &Path, + stage: &str, + values: &[i32], + shape: &[usize], +) { + arrays.insert( + stage.to_string(), + write_array(out, stage, &i32_bytes(values), "int32", shape), + ); +} + +fn usize_i32(values: &[usize], label: &str) -> Vec { + values + .iter() + .map(|&value| i32::try_from(value).unwrap_or_else(|_| panic!("{label} value exceeds i32"))) + .collect() +} + +fn reconstruct_pixels(patches: &[f32], grid_height: usize, grid_width: usize) -> Vec { + assert_eq!(patches.len(), grid_height * grid_width * PATCH_WIDTH); + let height = grid_height * PATCH_SIZE; + let width = grid_width * PATCH_SIZE; + let mut pixels = vec![0.0f32; 3 * height * width]; + for patch in 0..grid_height * grid_width { + let patch_y = patch / grid_width; + let patch_x = patch % grid_width; + let row = &patches[patch * PATCH_WIDTH..(patch + 1) * PATCH_WIDTH]; + let mut cursor = 0usize; + for channel in 0..3 { + for dy in 0..PATCH_SIZE { + for dx in 0..PATCH_SIZE { + let y = patch_y * PATCH_SIZE + dy; + let x = patch_x * PATCH_SIZE + dx; + pixels[channel * height * width + y * width + x] = row[cursor]; + cursor += 1; + } + } + } + } + pixels +} + +fn main() { + let model = required_path("--model"); + let reference_root = required_path("--reference"); + let image_path = required_path("--image"); + let out = required_path("--out"); + let device = argument("--device").unwrap_or_else(|| "local-task".to_string()); + let context_capacity = argument("--context-capacity") + .map(|value| { + value + .parse::() + .expect("context capacity must be usize") + }) + .unwrap_or(2_048); + assert!( + !out.exists(), + "immutable capture output already exists: {}", + out.display() + ); + fs::create_dir_all(&out).unwrap_or_else(|error| panic!("create {}: {error}", out.display())); + + let reference = verify_reference_manifest(&reference_root); + verify_checkpoint(&model, &reference["checkpoint"]); + assert_eq!( + sha256_file(&image_path), + FIXTURE_SHA256, + "fixture SHA-256 differs" + ); + let image = image::open(&image_path) + .unwrap_or_else(|error| panic!("open {}: {error}", image_path.display())) + .into_rgb8(); + let image = DynamicImage::ImageRgb8(image); + + let tokenizer = load_tokenizer(&model).expect("load pinned Youtu-VL tokenizer"); + let unexpanded = tokenizer + .encode(PROMPT, false) + .expect("tokenize pinned Youtu-VL prompt") + .into_iter() + .map(|value| i32::try_from(value).expect("token id fits i32")) + .collect::>(); + let reference_unexpanded = reference["case"]["unexpanded_input_ids"] + .as_array() + .expect("reference unexpanded_input_ids must be an array") + .iter() + .map(|value| { + value + .as_i64() + .and_then(|value| i32::try_from(value).ok()) + .expect("reference token id fits i32") + }) + .collect::>(); + assert_eq!( + unexpanded, reference_unexpanded, + "independent tokenizer ids differ from HF reference" + ); + + let processor = YoutuVLProcessor::new(PATCH_SIZE, 2) + .with_norm([0.5; 3], [0.5; 3]) + .with_max_patches_per_image(PATCHES) + .with_resample(2); + let (flattened_patches, spatial_shapes, patch_width) = processor + .try_preprocess_values_with_spatial(std::slice::from_ref(&image)) + .expect("run checked Youtu-VL flattened-patch processor"); + assert_eq!(spatial_shapes, [(16, 16)]); + assert_eq!(patch_width, PATCH_WIDTH); + let resized_pixels = reconstruct_pixels(&flattened_patches, 16, 16); + + let _runtime = initialize_runtime(); + let mut diagnostic_projector = IreeYoutuVlDiagnosticProjector::load(&model, &device) + .expect("load Youtu-VL diagnostic vision projector"); + let vision = diagnostic_projector + .capture(&flattened_patches, &spatial_shapes) + .expect("capture actual 27-layer Youtu-VL vision graph"); + assert_eq!(vision.abi_version, 2); + assert_eq!(vision.patch_shape, [PATCHES, 1_152]); + assert_eq!(vision.merged_shape, [64, TEXT_HIDDEN]); + + let production_preprocessor = YoutuVlIreeHostPreprocessor::load(&model, &device) + .expect("load production Youtu-VL IREE host preprocessor"); + let prepared = production_preprocessor + .prepare(&unexpanded, std::slice::from_ref(&image)) + .expect("prepare production Youtu-VL multimodal prefill"); + assert_eq!(prepared.embeddings.dtype, PreparedTensorDType::Float32); + assert_eq!( + prepared.embeddings.shape, + [1, prepared.sequence_len, TEXT_HIDDEN] + ); + assert_eq!( + prepared + .token_ids + .iter() + .filter(|&&token| token == IMAGE_TOKEN_ID) + .count(), + 64 + ); + + let mut engine = YoutuVlReferenceDiagnosticEngine::load(&model, &device, context_capacity) + .expect("load Youtu-VL production diagnostic language engine"); + let fresh = engine + .capture(&prepared, KV_WIDTH, MAX_NEW_TOKENS) + .expect("capture fresh Youtu-VL MLA request"); + let reuse = engine + .capture(&prepared, KV_WIDTH, MAX_NEW_TOKENS) + .expect("reset and capture reused Youtu-VL MLA slot"); + assert_eq!( + fresh.tokens, reuse.tokens, + "slot reuse token stream differs" + ); + assert_eq!( + fresh.prefill, reuse.prefill, + "slot reuse logits or MLA cache differs" + ); + assert_eq!(fresh.prefill.layers, TEXT_LAYERS); + assert_eq!(fresh.prefill.kv_width, KV_WIDTH); + assert_eq!(fresh.prefill.logits.len(), 283_386); + assert_eq!(fresh.prefill.kv.len(), TEXT_LAYERS * 2 * KV_WIDTH); + + let mut arrays = Map::new(); + write_f32( + &mut arrays, + &out, + "resized_normalized_pixels", + &resized_pixels, + &[3, 256, 256], + ); + write_f32( + &mut arrays, + &out, + "flattened_patches", + &flattened_patches, + &[PATCHES, PATCH_WIDTH], + ); + write_f32( + &mut arrays, + &out, + "patches.window_order", + &vision.patches_window_order, + &[PATCHES, PATCH_WIDTH], + ); + write_f32( + &mut arrays, + &out, + "vision_rope.freqs", + &vision.rope_freqs, + &[PATCHES, 36], + ); + for stage in &vision.stages { + write_f32(&mut arrays, &out, &stage.name, &stage.values, &stage.shape); + } + arrays.insert( + "prepared_embeddings".to_string(), + write_array( + &out, + "prepared_embeddings", + &prepared.embeddings.bytes, + "float32", + &prepared.embeddings.shape, + ), + ); + write_f32( + &mut arrays, + &out, + "prefill_logits", + &fresh.prefill.logits, + &[283_386], + ); + write_f32( + &mut arrays, + &out, + "selected_kv", + &fresh.prefill.kv, + &[TEXT_LAYERS, 2, KV_WIDTH], + ); + write_f32( + &mut arrays, + &out, + "reuse_prefill_logits", + &reuse.prefill.logits, + &[283_386], + ); + write_f32( + &mut arrays, + &out, + "reuse_selected_kv", + &reuse.prefill.kv, + &[TEXT_LAYERS, 2, KV_WIDTH], + ); + write_i32( + &mut arrays, + &out, + "expanded_input_ids", + &prepared.token_ids, + &[prepared.sequence_len], + ); + let placeholder_positions = prepared + .token_ids + .iter() + .enumerate() + .filter_map(|(index, &token)| (token == IMAGE_TOKEN_ID).then_some(index as i32)) + .collect::>(); + write_i32( + &mut arrays, + &out, + "placeholder_positions", + &placeholder_positions, + &[placeholder_positions.len()], + ); + write_i32(&mut arrays, &out, "spatial_shapes", &[16, 16], &[1, 2]); + write_i32( + &mut arrays, + &out, + "window_group_index", + &usize_i32(&vision.window_group_index, "window group index"), + &[64], + ); + write_i32( + &mut arrays, + &out, + "reverse_group_index", + &usize_i32(&vision.reverse_group_index, "reverse group index"), + &[64], + ); + write_i32( + &mut arrays, + &out, + "window_cu_seqlens", + &usize_i32(&vision.window_cu_seqlens, "window boundary"), + &[vision.window_cu_seqlens.len()], + ); + write_i32( + &mut arrays, + &out, + "full_cu_seqlens", + &usize_i32(&vision.full_cu_seqlens, "full boundary"), + &[vision.full_cu_seqlens.len()], + ); + write_i32( + &mut arrays, + &out, + "greedy_tokens", + &fresh.tokens, + &[MAX_NEW_TOKENS], + ); + write_i32( + &mut arrays, + &out, + "reuse_greedy_tokens", + &reuse.tokens, + &[MAX_NEW_TOKENS], + ); + + let actual = json!({ + "schema": 1, + "contract": CONTRACT, + "producer": "mlxcel-xla-diagnostics", + "checkpoint": reference["checkpoint"].clone(), + "fixture": reference["fixture"].clone(), + "architecture": reference["architecture"].clone(), + "generation": reference["generation"].clone(), + "lifecycle": reference["lifecycle"].clone(), + "case": { + "name": "image_text", + "prompt": PROMPT, + "unexpanded_input_ids": unexpanded, + "arrays": arrays, + }, + }); + let manifest_path = out.join("manifest.json"); + fs::write( + &manifest_path, + serde_json::to_vec_pretty(&actual).expect("serialize actual manifest"), + ) + .expect("write actual manifest"); + fs::write( + out.join("manifest.sha256"), + format!("{}\n", sha256_file(&manifest_path)), + ) + .expect("write actual manifest hash"); + println!("{}", manifest_path.display()); +} diff --git a/scripts/xla/validate_youtu_vl_reference.sh b/scripts/xla/validate_youtu_vl_reference.sh new file mode 100755 index 000000000..77169388b --- /dev/null +++ b/scripts/xla/validate_youtu_vl_reference.sh @@ -0,0 +1,74 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +PYTHON="${PYTHON:-python3}" +MODE="${1:---contract-only}" + +cd "$ROOT" +"$PYTHON" -m unittest spike/openxla/test_youtu_vl_reference_oracle.py +"$PYTHON" -m py_compile \ + spike/openxla/youtu_vl_reference_oracle.py \ + spike/openxla/test_youtu_vl_reference_oracle.py + +if [[ "$MODE" == "--contract-only" ]]; then + exit 0 +fi + +if [[ "$MODE" != "--actual" || $# -lt 2 || $# -gt 3 ]]; then + echo "usage: $0 [--contract-only | --actual OUTPUT_ROOT [IREE_DEVICE]]" >&2 + exit 2 +fi + +OUTPUT_ROOT="$2" +DEVICE="${3:-local-task}" +if [[ -e "$OUTPUT_ROOT" ]]; then + echo "error: immutable output already exists: $OUTPUT_ROOT" >&2 + exit 2 +fi + +CACHE_ROOT="${MLXCEL_YOUTU_VL_ORACLE_CACHE:-$ROOT/.cache/youtu-vl-oracle}" +REVISION="8d30a0e49662a1d628a472b12df264dbcd768753" +MODEL_DIR="$CACHE_ROOT/Youtu-VL-4B-Instruct-$REVISION" +IMAGE="$ROOT/tests/fixtures/test_image.png" +REFERENCE="$OUTPUT_ROOT/hf-reference" +ACTUAL="$OUTPUT_ROOT/mlxcel-actual" +REPORT="$OUTPUT_ROOT/comparison.json" + +mkdir -p "$CACHE_ROOT" "$OUTPUT_ROOT" +"$PYTHON" - "$MODEL_DIR" <<'PY' +import sys +from pathlib import Path + +from huggingface_hub import snapshot_download + +sys.path.insert(0, "spike/openxla") +import youtu_vl_reference_oracle as oracle + +destination = Path(sys.argv[1]) +snapshot_download( + repo_id=oracle.CHECKPOINT_REPO, + revision=oracle.CHECKPOINT_REVISION, + local_dir=destination, + allow_patterns=sorted(oracle.CHECKPOINT_ARTIFACTS), +) +PY + +"$PYTHON" spike/openxla/youtu_vl_reference_oracle.py capture \ + --model "$MODEL_DIR" \ + --image "$IMAGE" \ + --out "$REFERENCE" + +cargo run --release \ + --features xla-diagnostics \ + --example xla_youtu_vl_reference_check -- \ + --model "$MODEL_DIR" \ + --reference "$REFERENCE" \ + --image "$IMAGE" \ + --out "$ACTUAL" \ + --device "$DEVICE" + +"$PYTHON" spike/openxla/youtu_vl_reference_oracle.py compare \ + --reference "$REFERENCE" \ + --actual "$ACTUAL" \ + --report "$REPORT" diff --git a/spike/openxla/test_youtu_vl_reference_oracle.py b/spike/openxla/test_youtu_vl_reference_oracle.py new file mode 100644 index 000000000..ce342a714 --- /dev/null +++ b/spike/openxla/test_youtu_vl_reference_oracle.py @@ -0,0 +1,216 @@ +#!/usr/bin/env python3 +"""Model-free adversarial tests for the Youtu-VL oracle contract.""" + +from __future__ import annotations + +import copy +import json +import sys +import tempfile +import unittest +from pathlib import Path +from typing import Any + +import numpy as np + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +import youtu_vl_reference_oracle as oracle + + +class YoutuVlOracleContractTests(unittest.TestCase): + def setUp(self) -> None: + self.temporary = tempfile.TemporaryDirectory() + self.root = Path(self.temporary.name) + self.reference = self.root / "reference" + self.actual = self.root / "actual" + self.reference.mkdir() + self.actual.mkdir() + self.reference_manifest = self.make_manifest("hf-transformers-eager") + self.actual_manifest = self.make_manifest("mlxcel-xla-diagnostics") + self.write_capture(self.reference, self.reference_manifest) + self.write_capture(self.actual, self.actual_manifest) + + def tearDown(self) -> None: + self.temporary.cleanup() + + @staticmethod + def make_manifest(producer: str) -> dict[str, Any]: + sequence = oracle.MERGED_TOKENS + 4 + arrays = {} + for stage, shape in oracle.expected_shapes(sequence).items(): + arrays[stage] = { + "file": f"image_text.{stage}.bin", + "dtype": ( + "int32" if stage in oracle.INTEGER_STAGES else "float32" + ), + "shape": shape, + "sha256": "", + } + return { + "schema": oracle.SCHEMA_VERSION, + "contract": oracle.CONTRACT, + "producer": producer, + "checkpoint": { + "repo": oracle.CHECKPOINT_REPO, + "revision": oracle.CHECKPOINT_REVISION, + "artifact_manifest": { + "canonical_sha256": ( + oracle.CHECKPOINT_ARTIFACT_MANIFEST_SHA256 + ), + "files": oracle.CHECKPOINT_ARTIFACTS, + }, + }, + "fixture": { + "path": oracle.FIXTURE_PATH, + "sha256": oracle.FIXTURE_SHA256, + }, + "architecture": { + "vision_depth": oracle.VISION_DEPTH, + "vision_hidden": oracle.VISION_HIDDEN, + "selected_vision_layers": [ + {"index": index, "attention": kind} + for index, kind in oracle.SELECTED_VISION_LAYERS + ], + "text_layers": oracle.TEXT_LAYERS, + "text_hidden": oracle.TEXT_HIDDEN, + "vocab": oracle.VOCAB, + "kv_lora_rank": oracle.KV_LORA_RANK, + "qk_rope_head_dim": oracle.ROPE_WIDTH, + }, + "generation": { + "mode": "greedy", + "max_new_tokens": oracle.MAX_NEW_TOKENS, + }, + "lifecycle": copy.deepcopy(oracle.LIFECYCLE), + "case": { + "name": "image_text", + "prompt": oracle.PROMPT, + "unexpanded_input_ids": [oracle.IMAGE_TOKEN_ID, 42, 43], + "arrays": arrays, + }, + } + + @staticmethod + def array(stage: str, shape: list[int], dtype: str) -> np.ndarray: + value = np.zeros(shape, dtype=dtype) + if stage == "expanded_input_ids": + value[:] = np.arange(value.size, dtype=np.int32) + 1 + value[: oracle.MERGED_TOKENS] = oracle.IMAGE_TOKEN_ID + elif stage == "placeholder_positions": + value[:] = np.arange(oracle.MERGED_TOKENS, dtype=np.int32) + elif stage == "spatial_shapes": + value[:] = [[16, 16]] + elif stage in {"window_group_index", "reverse_group_index"}: + value[:] = np.arange(oracle.MERGED_TOKENS, dtype=np.int32) + elif stage in {"window_cu_seqlens", "full_cu_seqlens"}: + value[:] = [0, oracle.PATCHES] + return value + + def write_capture(self, root: Path, manifest: dict[str, Any]) -> None: + for stage, spec in manifest["case"]["arrays"].items(): + value = self.array(stage, spec["shape"], spec["dtype"]) + path = root / spec["file"] + value.tofile(path) + spec["sha256"] = oracle.sha256(path) + oracle.write_manifest(root, manifest) + + def rewrite(self, root: Path, manifest: dict[str, Any]) -> None: + oracle.write_manifest(root, manifest) + + def assert_contract_error(self, contains: str) -> None: + report = oracle.compare_capture_roots(self.reference, self.actual) + self.assertFalse(report["passed"], report) + self.assertEqual(report["first_divergence"]["stage"], "contract") + self.assertIn(contains, report["error"].lower()) + + def test_valid_synthetic_capture_passes(self) -> None: + report = oracle.compare_capture_roots(self.reference, self.actual) + self.assertTrue(report["passed"], report) + self.assertIsNone(report["first_divergence"]) + + def test_checkpoint_identity_and_canonical_hash_are_pinned(self) -> None: + self.actual_manifest["checkpoint"]["revision"] = "main" + self.rewrite(self.actual, self.actual_manifest) + self.assert_contract_error("checkpoint identity") + self.actual_manifest = self.make_manifest("mlxcel-xla-diagnostics") + self.write_capture(self.actual, self.actual_manifest) + self.actual_manifest["checkpoint"]["artifact_manifest"][ + "canonical_sha256" + ] = "0" * 64 + self.rewrite(self.actual, self.actual_manifest) + self.assert_contract_error("checkpoint identity") + + def test_every_array_is_content_addressed(self) -> None: + path = self.actual / "image_text.prefill_logits.bin" + with path.open("r+b") as stream: + stream.write(np.asarray([1.0], dtype=np.float32).tobytes()) + self.assert_contract_error("artifact hash") + + def test_manifest_hash_is_required(self) -> None: + (self.actual / "manifest.sha256").write_text("0" * 64 + "\n") + self.assert_contract_error("manifest.sha256") + + def test_unlisted_capture_artifact_is_rejected(self) -> None: + (self.actual / "untracked.csv").write_text("not,an,oracle\n") + self.assert_contract_error("file set differs") + + def test_missing_selected_layer_is_rejected(self) -> None: + del self.actual_manifest["case"]["arrays"]["layer.23.full"] + self.rewrite(self.actual, self.actual_manifest) + self.assert_contract_error("stage set differs") + + def test_actual_27_layer_schedule_identity_is_required(self) -> None: + self.actual_manifest["architecture"]["vision_depth"] = 2 + self.rewrite(self.actual, self.actual_manifest) + self.assert_contract_error("architecture contract") + + def test_lifecycle_events_are_closed_and_ordered(self) -> None: + self.actual_manifest["lifecycle"]["events"][3] = "reuse_without_reset" + self.rewrite(self.actual, self.actual_manifest) + self.assert_contract_error("lifecycle contract") + + def test_reuse_must_equal_fresh_capture(self) -> None: + stage = "reuse_greedy_tokens" + spec = self.actual_manifest["case"]["arrays"][stage] + path = self.actual / spec["file"] + value = np.fromfile(path, dtype=np.int32) + value[0] = 99 + value.tofile(path) + spec["sha256"] = oracle.sha256(path) + self.rewrite(self.actual, self.actual_manifest) + self.assert_contract_error("lifecycle reuse differs") + + def test_window_permutation_and_inverse_are_validated(self) -> None: + stage = "window_group_index" + spec = self.actual_manifest["case"]["arrays"][stage] + path = self.actual / spec["file"] + value = np.fromfile(path, dtype=np.int32) + value[-1] = value[-2] + value.tofile(path) + spec["sha256"] = oracle.sha256(path) + self.rewrite(self.actual, self.actual_manifest) + self.assert_contract_error("not a permutation") + + def test_placeholder_selection_is_validated(self) -> None: + stage = "placeholder_positions" + spec = self.actual_manifest["case"]["arrays"][stage] + path = self.actual / spec["file"] + value = np.fromfile(path, dtype=np.int32) + value[-1] += 1 + value.tofile(path) + spec["sha256"] = oracle.sha256(path) + self.rewrite(self.actual, self.actual_manifest) + self.assert_contract_error("placeholder selection") + + def test_capture_output_is_immutable(self) -> None: + with self.assertRaisesRegex(oracle.ContractError, "already exists"): + oracle.ensure_new_output(self.actual) + + def test_generator_has_no_mlxcel_dependency(self) -> None: + source = Path(oracle.__file__).read_text(encoding="utf-8") + self.assertNotIn("import mlxcel", source) + self.assertNotIn("from mlxcel", source) + + +if __name__ == "__main__": + unittest.main() diff --git a/spike/openxla/youtu_vl_reference_oracle.py b/spike/openxla/youtu_vl_reference_oracle.py new file mode 100644 index 000000000..b59826d18 --- /dev/null +++ b/spike/openxla/youtu_vl_reference_oracle.py @@ -0,0 +1,876 @@ +#!/usr/bin/env python3 +"""Immutable HF-eager Youtu-VL capture and MLX/IREE comparator. + +The capture is intentionally independent of mlxcel: it loads only the pinned +Hugging Face checkpoint and its pinned remote code. Generated tensors stay +outside Git. The model-free ``compare`` command verifies every artifact hash, +the complete ordered Youtu-VL diagnostic surface, and the fresh-cache/reuse +lifecycle contract before comparing numeric values. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import math +import sys +from pathlib import Path +from typing import Any + + +SCHEMA_VERSION = 1 +CONTRACT = "youtu-vl-hf-eager-oracle-v1" +CHECKPOINT_REPO = "tencent/Youtu-VL-4B-Instruct" +CHECKPOINT_REVISION = "8d30a0e49662a1d628a472b12df264dbcd768753" +CHECKPOINT_ARTIFACTS = { + "__init__.py": "1ca8666808d0a0fbdaed1dfa81b313f1b3f0ed713d11be2918cce083f90ce11c", + "chat_template.json": "172d9ec2636e63f13326ae4a6da2053a7c031ce829be1e7859a0b2f16b2bd088", + "config.json": "8f6115914aa33b7b48c92efb077c0976bec71407db65c4daa39961793ac4009a", + "configuration_siglip2.py": "dca3068142f1cacf9891fd3e1f896891afa30457e20205d40e643f8bd39892bc", + "configuration_youtu_vl.py": "c69b40aa372cddaaf27feca16205b82650f0db684e74202335d70bfa791df595", + "generation_config.json": "8666c5875737011451ad9874fa297adbb5d6c3de47ef647a17c41ccee5af7e90", + "image_processing_siglip2_fast.py": "56cbb7f814ac582e6ff6abd0f7486ee1e65fef998994c2ebabc4540a079ed904", + "model-00001-of-00003.safetensors": "107939f0c4aaf5fdc7c5d9f4ad741546e4ebaef32fcd93d36a462f1e5ea04d0b", + "model-00002-of-00003.safetensors": "f7898bdeb9c5cbc996c0a2789c7b3495c3faa179c9690cb2c01bad65b615250e", + "model-00003-of-00003.safetensors": "548ebf9f8deea536596d6dffc3fed769a60f15e414e6fb9f735a7bb7fa627340", + "model.safetensors.index.json": "8617dad368bb394a9de5f24f08b5b0ecdd8a68b2b8588d2baf2cc8f2bc5b6b7c", + "modeling_siglip2.py": "af9df4e395d7a1c5dcce6ae2cc87155ec1712c0eb025dc4b14674efb6604ac4f", + "modeling_youtu_vl.py": "6ec2ffee9382f248149462f111a92d58846fca52dfb7f12aa184fd13b8eefda5", + "preprocessor_config.json": "eb41cbafe9513f75c79c6f7c1d64f0d8f7419f98cc6e097e63b9467aee6dd6d7", + "processing_youtu_vl.py": "415c83d7763e5aa97867344f2d71178d5523d310410a08c64b335faf57e4a667", + "special_tokens_map.json": "03d62862d41de30db9e05cb4865de4f36c48ab3032327139fcdfacc5798a4828", + "tokenizer.json": "41998384e9cea31ab97207e2ed59fed66b5481bf0c85fd04f8c7bbd3f7648a6d", + "tokenizer_config.json": "da4d16d9ba9de8afb26c3eb28b474450c2dbd2a647b7159fbbd87d61e78f98a6", +} +# Filled from canonical_artifact_sha(CHECKPOINT_ARTIFACTS). Keeping this +# literal makes accidental edits to the maintained identity fail closed. +CHECKPOINT_ARTIFACT_MANIFEST_SHA256 = ( + "4d67dd2750d1ee8d87e68b0b52f5aa6c5b5b1dd85385df643845e393628812c9" +) + +FIXTURE_PATH = "tests/fixtures/test_image.png" +FIXTURE_SHA256 = "5e7d54e8a7d21802378c87d2d70cf551e29739fe27599ddf129ebccdad1e6261" +PROMPT = "<|image_pad|>\nDescribe the image briefly." +IMAGE_TOKEN_ID = 128_264 +VISION_DEPTH = 27 +VISION_HIDDEN = 1_152 +TEXT_LAYERS = 40 +TEXT_HIDDEN = 2_560 +VOCAB = 283_386 +KV_LORA_RANK = 512 +ROPE_WIDTH = 64 +KV_WIDTH = KV_LORA_RANK + ROPE_WIDTH +PATCH_SIZE = 16 +PATCH_WIDTH = 3 * PATCH_SIZE * PATCH_SIZE +PATCHES = 256 +MERGED_TOKENS = PATCHES // 4 +VISION_ROPE_WIDTH = 36 +MAX_NEW_TOKENS = 4 + +SELECTED_VISION_LAYERS = ( + (0, "window"), + (7, "full"), + (15, "full"), + (23, "full"), + (26, "full"), +) +VISION_STAGE_NAMES = ( + "patch_projection", + *(f"layer.{index}.{kind}" for index, kind in SELECTED_VISION_LAYERS), + "post_layernorm", + "merger.window_order", + "merger.restored_order", +) +FLOAT_STAGES = ( + "resized_normalized_pixels", + "flattened_patches", + "patches.window_order", + "vision_rope.freqs", + *VISION_STAGE_NAMES, + "prepared_embeddings", + "prefill_logits", + "selected_kv", + "reuse_prefill_logits", + "reuse_selected_kv", +) +INTEGER_STAGES = ( + "expanded_input_ids", + "placeholder_positions", + "spatial_shapes", + "window_group_index", + "reverse_group_index", + "window_cu_seqlens", + "full_cu_seqlens", + "greedy_tokens", + "reuse_greedy_tokens", +) +REQUIRED_STAGES = FLOAT_STAGES + INTEGER_STAGES +LIFECYCLE = { + "slot": 0, + "events": [ + "reset", + "prefill", + "greedy_decode", + "reset_for_reuse", + "prefill_reuse", + "greedy_decode_reuse", + ], + "reuse_must_match_fresh": True, +} + + +class ContractError(ValueError): + """The capture does not satisfy the closed oracle contract.""" + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def canonical_artifact_sha(artifacts: dict[str, str]) -> str: + payload = json.dumps( + artifacts, sort_keys=True, separators=(",", ":"), ensure_ascii=True + ) + return hashlib.sha256(payload.encode("utf-8")).hexdigest() + + +def require_checkpoint(model: Path) -> dict[str, Any]: + canonical = canonical_artifact_sha(CHECKPOINT_ARTIFACTS) + if canonical != CHECKPOINT_ARTIFACT_MANIFEST_SHA256: + raise ContractError( + "internal checkpoint artifact manifest hash differs: " + f"{canonical}" + ) + for filename, expected in CHECKPOINT_ARTIFACTS.items(): + path = model / filename + if not path.is_file(): + raise ContractError(f"checkpoint artifact is missing: {filename}") + actual = sha256(path) + if actual != expected: + raise ContractError( + f"checkpoint artifact hash differs for {filename}: " + f"expected {expected}, got {actual}" + ) + unexpected_weights = sorted( + path.name + for path in model.glob("*.safetensors") + if path.name not in CHECKPOINT_ARTIFACTS + ) + if unexpected_weights: + raise ContractError( + f"checkpoint has unpinned weight artifact(s): {unexpected_weights}" + ) + return { + "repo": CHECKPOINT_REPO, + "revision": CHECKPOINT_REVISION, + "artifact_manifest": { + "canonical_sha256": canonical, + "files": CHECKPOINT_ARTIFACTS, + }, + } + + +def require_fixture(image: Path) -> None: + if sha256(image) != FIXTURE_SHA256: + raise ContractError(f"fixture hash differs: {image}") + + +def strict_json_load(path: Path) -> dict[str, Any]: + def reject_duplicate(pairs: list[tuple[str, Any]]) -> dict[str, Any]: + value: dict[str, Any] = {} + for key, item in pairs: + if key in value: + raise ContractError(f"duplicate JSON key {key!r}") + value[key] = item + return value + + try: + parsed = json.loads( + path.read_text(encoding="utf-8"), object_pairs_hook=reject_duplicate + ) + except (OSError, json.JSONDecodeError) as error: + raise ContractError(f"read {path}: {error}") from error + if not isinstance(parsed, dict): + raise ContractError(f"{path} must contain a JSON object") + return parsed + + +def write_manifest(root: Path, manifest: dict[str, Any]) -> None: + path = root / "manifest.json" + path.write_text( + json.dumps(manifest, indent=2, sort_keys=True) + "\n", encoding="utf-8" + ) + (root / "manifest.sha256").write_text(sha256(path) + "\n", encoding="ascii") + + +def ensure_new_output(path: Path) -> None: + try: + path.mkdir(parents=True, exist_ok=False) + except FileExistsError as error: + raise ContractError( + f"immutable capture output already exists: {path}" + ) from error + + +def store_array( + np: Any, + root: Path, + stage: str, + value: Any, + dtype: str, +) -> dict[str, Any]: + array = np.asarray(value, dtype=np.dtype(dtype)) + if dtype.startswith("float") and not np.isfinite(array).all(): + raise ContractError(f"{stage} contains a non-finite value") + filename = f"image_text.{stage}.bin" + path = root / filename + array.tofile(path) + return { + "file": filename, + "dtype": dtype, + "shape": list(array.shape), + "sha256": sha256(path), + } + + +def patch_order(torch: Any, window_index: Any) -> Any: + offsets = torch.arange(4, device=window_index.device) + return (window_index[:, None] * 4 + offsets[None, :]).reshape(-1) + + +def unpatch(torch: Any, patches: Any, height: int, width: int) -> Any: + rows = patches.reshape(height, width, 3, PATCH_SIZE, PATCH_SIZE) + return ( + rows.permute(2, 0, 3, 1, 4) + .contiguous() + .reshape(3, height * PATCH_SIZE, width * PATCH_SIZE) + ) + + +def rotate_half(torch: Any, value: Any) -> Any: + first, second = value.chunk(2, dim=-1) + return torch.cat((-second, first), dim=-1) + + +def interleaved_rotary_k(torch: Any, raw: Any, cos: Any, sin: Any) -> Any: + batch, heads, sequence, width = raw.shape + reordered = ( + raw.view(batch, heads, sequence, width // 2, 2) + .transpose(4, 3) + .reshape(batch, heads, sequence, width) + ) + return reordered * cos.unsqueeze(1) + rotate_half( + torch, reordered + ) * sin.unsqueeze(1) + + +def greedy_language_capture( + torch: Any, + model: Any, + prepared: Any, + max_new_tokens: int, +) -> tuple[Any, Any, Any]: + latent: dict[int, Any] = {} + compressed: dict[int, Any] = {} + handles = [] + for index, layer in enumerate(model.model.layers): + handles.append( + layer.self_attn.kv_a_layernorm.register_forward_hook( + lambda _module, _args, output, index=index: latent.__setitem__( + index, output.detach() + ) + ) + ) + handles.append( + layer.self_attn.kv_a_proj_with_mqa.register_forward_hook( + lambda _module, _args, output, index=index: compressed.__setitem__( + index, output.detach() + ) + ) + ) + try: + positions = torch.arange( + prepared.shape[1], device=prepared.device + ).unsqueeze(0) + cos, sin = model.model.rotary_emb(prepared, positions) + outputs = model.model( + inputs_embeds=prepared, + use_cache=True, + cache_position=positions[0], + ) + logits = model.lm_head(outputs.last_hidden_state[:, -1, :]).float() + if set(latent) != set(range(TEXT_LAYERS)) or set(compressed) != set( + range(TEXT_LAYERS) + ): + raise ContractError("HF eager MLA hooks did not capture every layer") + padded_rows = [] + for index in range(TEXT_LAYERS): + latent_row = latent[index][0, -1, :].float() + raw_rotary = compressed[index][ + :, -1:, KV_LORA_RANK: + ].reshape(1, 1, 1, ROPE_WIDTH) + rotated = interleaved_rotary_k( + torch, raw_rotary, cos[:, -1:, :], sin[:, -1:, :] + )[0, 0, 0].float() + zeros_latent = torch.zeros_like(latent_row) + zeros_rotary = torch.zeros_like(rotated) + padded_rows.append( + torch.stack( + ( + torch.cat((latent_row, zeros_rotary)), + torch.cat((zeros_latent, rotated)), + ) + ) + ) + selected_kv = torch.stack(padded_rows) + tokens = [] + current = logits.argmax(dim=-1, keepdim=True) + tokens.append(int(current.item())) + cache = outputs.past_key_values + while len(tokens) < max_new_tokens: + decoded = model.model( + input_ids=current, + past_key_values=cache, + use_cache=True, + ) + step_logits = model.lm_head( + decoded.last_hidden_state[:, -1, :] + ).float() + current = step_logits.argmax(dim=-1, keepdim=True) + tokens.append(int(current.item())) + cache = decoded.past_key_values + return logits[0], selected_kv, torch.tensor(tokens, dtype=torch.int32) + finally: + for handle in handles: + handle.remove() + + +def capture(args: argparse.Namespace) -> int: + if args.max_new != MAX_NEW_TOKENS: + raise ContractError( + f"pinned capture requires --max-new {MAX_NEW_TOKENS}" + ) + checkpoint = require_checkpoint(args.model) + require_fixture(args.image) + ensure_new_output(args.out) + try: + import numpy as np + import torch + from PIL import Image + from transformers import AutoModelForCausalLM, AutoProcessor + except ImportError as error: + raise ContractError( + "capture requires numpy, torch, Pillow, and transformers" + ) from error + + torch.manual_seed(0) + processor = AutoProcessor.from_pretrained( + args.model, trust_remote_code=True, local_files_only=True + ) + model = AutoModelForCausalLM.from_pretrained( + args.model, + trust_remote_code=True, + local_files_only=True, + torch_dtype=torch.bfloat16, + attn_implementation="eager", + device_map=None, + ) + model.eval() + image = Image.open(args.image).convert("RGB") + unexpanded = processor.tokenizer( + PROMPT, add_special_tokens=False, return_tensors="pt" + )["input_ids"] + inputs = processor( + text=PROMPT, + images=image, + max_image_patches=PATCHES, + add_special_tokens=False, + return_tensors="pt", + ) + input_ids = inputs["input_ids"] + pixel_values = inputs["pixel_values"] + spatial_shapes = inputs["spatial_shapes"] + if list(spatial_shapes.shape) != [1, 2] or spatial_shapes.tolist() != [ + [16, 16] + ]: + raise ContractError( + f"pinned fixture grid must be [[16, 16]], got {spatial_shapes.tolist()}" + ) + flat_patches = pixel_values[0] + encoder = model.siglip2.vision_model.encoder + window_index, window_cu = encoder.get_window_index(spatial_shapes) + window_index = window_index.to(dtype=torch.int64) + reverse_index = torch.argsort(window_index) + ordered_patches = flat_patches[patch_order(torch, window_index)] + rope = encoder.rot_pos_emb(spatial_shapes) + ordered_rope = rope[patch_order(torch, window_index)] + full_cu = torch.nn.functional.pad( + (spatial_shapes[:, 0] * spatial_shapes[:, 1]).cumsum( + dim=0, dtype=torch.int32 + ), + (1, 0), + ) + + selected_outputs: dict[int, Any] = {} + handles = [] + for index, _kind in SELECTED_VISION_LAYERS: + handles.append( + encoder.layers[index].register_forward_hook( + lambda _module, _args, output, index=index: selected_outputs.__setitem__( + index, output[0].detach() + ) + ) + ) + try: + with torch.inference_mode(): + vision = model.siglip2( + pixel_values.to(dtype=model.siglip2.dtype), + inputs.get("pixel_attention_mask"), + spatial_shapes, + ).last_hidden_state + restored_merged = model.merger(vision, spatial_shapes) + finally: + for handle in handles: + handle.remove() + if set(selected_outputs) != {index for index, _ in SELECTED_VISION_LAYERS}: + raise ContractError("HF eager vision hooks omitted a selected layer") + + window_patch_order = patch_order(torch, window_index) + patch_projection = model.siglip2.vision_model.embeddings( + pixel_values.to(dtype=model.siglip2.dtype) + )[window_patch_order] + # The encoder restores group order before returning. LayerNorm is pointwise, + # so reindexing its output recovers the module's window-order seam. + post_layernorm_window = vision[window_patch_order] + merger_window = restored_merged[window_index] + text_embeddings = model.model.embed_tokens(input_ids) + image_mask = (input_ids == IMAGE_TOKEN_ID).unsqueeze(-1).expand_as( + text_embeddings + ) + prepared = text_embeddings.masked_scatter( + image_mask, + restored_merged.to( + device=text_embeddings.device, dtype=text_embeddings.dtype + ).unsqueeze(0), + ) + + with torch.inference_mode(): + fresh_logits, fresh_kv, fresh_tokens = greedy_language_capture( + torch, model, prepared, args.max_new + ) + reuse_logits, reuse_kv, reuse_tokens = greedy_language_capture( + torch, model, prepared, args.max_new + ) + if not torch.equal(fresh_tokens, reuse_tokens): + raise ContractError("fresh-cache and reuse token streams differ") + if not torch.equal(fresh_logits, reuse_logits) or not torch.equal( + fresh_kv, reuse_kv + ): + raise ContractError("fresh-cache and reuse language captures differ") + + arrays: dict[str, Any] = {} + + def store(stage: str, value: Any, dtype: str) -> None: + if hasattr(value, "detach"): + value = value.detach().cpu() + value = value.float() if dtype == "float32" else value.int() + value = value.numpy() + arrays[stage] = store_array(np, args.out, stage, value, dtype) + + store( + "resized_normalized_pixels", + unpatch(torch, flat_patches, 16, 16), + "float32", + ) + store("flattened_patches", flat_patches, "float32") + store("patches.window_order", ordered_patches, "float32") + store("vision_rope.freqs", ordered_rope, "float32") + store("patch_projection", patch_projection, "float32") + for index, kind in SELECTED_VISION_LAYERS: + store(f"layer.{index}.{kind}", selected_outputs[index], "float32") + store("post_layernorm", post_layernorm_window, "float32") + store("merger.window_order", merger_window, "float32") + store("merger.restored_order", restored_merged, "float32") + store("prepared_embeddings", prepared, "float32") + store("prefill_logits", fresh_logits, "float32") + store("selected_kv", fresh_kv, "float32") + store("reuse_prefill_logits", reuse_logits, "float32") + store("reuse_selected_kv", reuse_kv, "float32") + store("expanded_input_ids", input_ids, "int32") + store( + "placeholder_positions", + (input_ids[0] == IMAGE_TOKEN_ID).nonzero().flatten(), + "int32", + ) + store("spatial_shapes", spatial_shapes, "int32") + store("window_group_index", window_index, "int32") + store("reverse_group_index", reverse_index, "int32") + store("window_cu_seqlens", torch.tensor(window_cu), "int32") + store("full_cu_seqlens", full_cu, "int32") + store("greedy_tokens", fresh_tokens, "int32") + store("reuse_greedy_tokens", reuse_tokens, "int32") + + manifest = { + "schema": SCHEMA_VERSION, + "contract": CONTRACT, + "producer": "hf-transformers-eager", + "checkpoint": checkpoint, + "fixture": {"path": FIXTURE_PATH, "sha256": FIXTURE_SHA256}, + "architecture": { + "vision_depth": VISION_DEPTH, + "vision_hidden": VISION_HIDDEN, + "selected_vision_layers": [ + {"index": index, "attention": kind} + for index, kind in SELECTED_VISION_LAYERS + ], + "text_layers": TEXT_LAYERS, + "text_hidden": TEXT_HIDDEN, + "vocab": VOCAB, + "kv_lora_rank": KV_LORA_RANK, + "qk_rope_head_dim": ROPE_WIDTH, + }, + "generation": {"mode": "greedy", "max_new_tokens": MAX_NEW_TOKENS}, + "lifecycle": LIFECYCLE, + "case": { + "name": "image_text", + "prompt": PROMPT, + "unexpanded_input_ids": unexpanded[0].tolist(), + "arrays": arrays, + }, + } + write_manifest(args.out, manifest) + print(args.out / "manifest.json") + return 0 + + +def expected_shapes(sequence: int) -> dict[str, list[int]]: + shapes = { + "resized_normalized_pixels": [3, 256, 256], + "flattened_patches": [PATCHES, PATCH_WIDTH], + "patches.window_order": [PATCHES, PATCH_WIDTH], + "vision_rope.freqs": [PATCHES, VISION_ROPE_WIDTH], + "patch_projection": [PATCHES, VISION_HIDDEN], + "post_layernorm": [PATCHES, VISION_HIDDEN], + "merger.window_order": [MERGED_TOKENS, TEXT_HIDDEN], + "merger.restored_order": [MERGED_TOKENS, TEXT_HIDDEN], + "prepared_embeddings": [1, sequence, TEXT_HIDDEN], + "prefill_logits": [VOCAB], + "selected_kv": [TEXT_LAYERS, 2, KV_WIDTH], + "reuse_prefill_logits": [VOCAB], + "reuse_selected_kv": [TEXT_LAYERS, 2, KV_WIDTH], + "expanded_input_ids": [sequence], + "placeholder_positions": [MERGED_TOKENS], + "spatial_shapes": [1, 2], + "window_group_index": [MERGED_TOKENS], + "reverse_group_index": [MERGED_TOKENS], + "window_cu_seqlens": [2], + "full_cu_seqlens": [2], + "greedy_tokens": [MAX_NEW_TOKENS], + "reuse_greedy_tokens": [MAX_NEW_TOKENS], + } + for index, kind in SELECTED_VISION_LAYERS: + shapes[f"layer.{index}.{kind}"] = [PATCHES, VISION_HIDDEN] + return shapes + + +def validate_checkpoint(value: Any) -> None: + expected = { + "repo": CHECKPOINT_REPO, + "revision": CHECKPOINT_REVISION, + "artifact_manifest": { + "canonical_sha256": CHECKPOINT_ARTIFACT_MANIFEST_SHA256, + "files": CHECKPOINT_ARTIFACTS, + }, + } + if value != expected: + raise ContractError("checkpoint identity or artifact hashes differ") + if ( + canonical_artifact_sha(value["artifact_manifest"]["files"]) + != CHECKPOINT_ARTIFACT_MANIFEST_SHA256 + ): + raise ContractError("checkpoint artifact manifest is not canonical") + + +def validate_manifest_hash(root: Path) -> None: + expected = (root / "manifest.sha256").read_text(encoding="ascii").strip() + if len(expected) != 64 or sha256(root / "manifest.json") != expected: + raise ContractError("manifest.sha256 differs") + + +def validate_capture(root: Path, role: str) -> tuple[dict[str, Any], dict[str, Any]]: + validate_manifest_hash(root) + manifest = strict_json_load(root / "manifest.json") + required_top = { + "schema", + "contract", + "producer", + "checkpoint", + "fixture", + "architecture", + "generation", + "lifecycle", + "case", + } + if set(manifest) != required_top: + raise ContractError(f"{role} manifest has an invalid top-level schema") + if manifest["schema"] != SCHEMA_VERSION or manifest["contract"] != CONTRACT: + raise ContractError(f"{role} schema or contract differs") + expected_producer = ( + "hf-transformers-eager" if role == "reference" else "mlxcel-xla-diagnostics" + ) + if manifest["producer"] != expected_producer: + raise ContractError(f"{role} producer differs") + validate_checkpoint(manifest["checkpoint"]) + if manifest["fixture"] != { + "path": FIXTURE_PATH, + "sha256": FIXTURE_SHA256, + }: + raise ContractError(f"{role} fixture identity differs") + architecture = manifest["architecture"] + if architecture != { + "vision_depth": VISION_DEPTH, + "vision_hidden": VISION_HIDDEN, + "selected_vision_layers": [ + {"index": index, "attention": kind} + for index, kind in SELECTED_VISION_LAYERS + ], + "text_layers": TEXT_LAYERS, + "text_hidden": TEXT_HIDDEN, + "vocab": VOCAB, + "kv_lora_rank": KV_LORA_RANK, + "qk_rope_head_dim": ROPE_WIDTH, + }: + raise ContractError(f"{role} architecture contract differs") + if manifest["generation"] != { + "mode": "greedy", + "max_new_tokens": MAX_NEW_TOKENS, + }: + raise ContractError(f"{role} generation contract differs") + if manifest["lifecycle"] != LIFECYCLE: + raise ContractError(f"{role} lifecycle contract differs") + case = manifest["case"] + if not isinstance(case, dict) or set(case) != { + "name", + "prompt", + "unexpanded_input_ids", + "arrays", + }: + raise ContractError(f"{role} case schema differs") + if case["name"] != "image_text" or case["prompt"] != PROMPT: + raise ContractError(f"{role} case identity differs") + if not isinstance(case["unexpanded_input_ids"], list) or not all( + isinstance(value, int) for value in case["unexpanded_input_ids"] + ): + raise ContractError(f"{role} unexpanded token ids are invalid") + arrays = case["arrays"] + if not isinstance(arrays, dict) or set(arrays) != set(REQUIRED_STAGES): + raise ContractError(f"{role} stage set differs") + expanded = arrays["expanded_input_ids"] + if not isinstance(expanded, dict): + raise ContractError(f"{role} expanded_input_ids spec is invalid") + shape = expanded.get("shape") + if ( + not isinstance(shape, list) + or len(shape) != 1 + or not isinstance(shape[0], int) + or shape[0] <= 0 + ): + raise ContractError(f"{role} expanded sequence shape is invalid") + shapes = expected_shapes(shape[0]) + for stage in REQUIRED_STAGES: + spec = arrays[stage] + expected_dtype = "int32" if stage in INTEGER_STAGES else "float32" + if not isinstance(spec, dict) or set(spec) != { + "file", + "dtype", + "shape", + "sha256", + }: + raise ContractError(f"{role} {stage} array spec schema differs") + if spec["dtype"] != expected_dtype or spec["shape"] != shapes[stage]: + raise ContractError(f"{role} {stage} dtype or shape differs") + filename = f"image_text.{stage}.bin" + if spec["file"] != filename: + raise ContractError(f"{role} {stage} filename differs") + path = root / filename + if not path.is_file() or sha256(path) != spec["sha256"]: + raise ContractError(f"{role} {stage} artifact hash differs") + itemsize = 4 + elements = math.prod(spec["shape"]) + if path.stat().st_size != elements * itemsize: + raise ContractError(f"{role} {stage} byte length differs") + expected_files = { + "manifest.json", + "manifest.sha256", + *(spec["file"] for spec in arrays.values()), + } + actual_files = {path.name for path in root.iterdir() if path.is_file()} + if actual_files != expected_files: + raise ContractError(f"{role} capture file set differs") + return manifest, arrays + + +def load_array(np: Any, root: Path, spec: dict[str, Any]) -> Any: + return np.fromfile(root / spec["file"], dtype=np.dtype(spec["dtype"])).reshape( + spec["shape"] + ) + + +def validate_semantics(np: Any, root: Path, arrays: dict[str, Any], role: str) -> None: + spatial = load_array(np, root, arrays["spatial_shapes"]) + if spatial.tolist() != [[16, 16]]: + raise ContractError(f"{role} spatial grid differs") + expanded = load_array(np, root, arrays["expanded_input_ids"]) + placeholders = load_array(np, root, arrays["placeholder_positions"]) + expected_placeholders = np.flatnonzero(expanded == IMAGE_TOKEN_ID).astype( + np.int32 + ) + if not np.array_equal(placeholders, expected_placeholders): + raise ContractError(f"{role} placeholder selection differs") + if placeholders.size != MERGED_TOKENS: + raise ContractError(f"{role} placeholder count differs") + window = load_array(np, root, arrays["window_group_index"]) + reverse = load_array(np, root, arrays["reverse_group_index"]) + if not np.array_equal(np.sort(window), np.arange(MERGED_TOKENS)): + raise ContractError(f"{role} window_group_index is not a permutation") + if not np.array_equal(reverse[window], np.arange(MERGED_TOKENS)): + raise ContractError(f"{role} reverse_group_index is not the inverse") + for stage in ("window_cu_seqlens", "full_cu_seqlens"): + boundaries = load_array(np, root, arrays[stage]) + if ( + boundaries[0] != 0 + or boundaries[-1] != PATCHES + or np.any(boundaries[1:] <= boundaries[:-1]) + ): + raise ContractError(f"{role} {stage} boundaries differ") + for fresh, reuse in ( + ("prefill_logits", "reuse_prefill_logits"), + ("selected_kv", "reuse_selected_kv"), + ("greedy_tokens", "reuse_greedy_tokens"), + ): + if not np.array_equal( + load_array(np, root, arrays[fresh]), + load_array(np, root, arrays[reuse]), + ): + raise ContractError(f"{role} lifecycle reuse differs for {fresh}") + + +def stage_tolerance(stage: str) -> tuple[float, float]: + if stage in INTEGER_STAGES: + return 0.0, 0.0 + if stage in { + "resized_normalized_pixels", + "flattened_patches", + "patches.window_order", + "vision_rope.freqs", + }: + return 1e-5, 1e-5 + return 2e-2, 2e-2 + + +def compare_capture_roots(reference: Path, actual: Path) -> dict[str, Any]: + try: + import numpy as np + + reference_manifest, reference_arrays = validate_capture( + reference, "reference" + ) + actual_manifest, actual_arrays = validate_capture(actual, "actual") + if ( + reference_manifest["case"]["unexpanded_input_ids"] + != actual_manifest["case"]["unexpanded_input_ids"] + ): + raise ContractError("unexpanded tokenizer ids differ") + validate_semantics(np, reference, reference_arrays, "reference") + validate_semantics(np, actual, actual_arrays, "actual") + except (ContractError, ImportError, OSError) as error: + return { + "passed": False, + "first_divergence": {"stage": "contract"}, + "error": str(error), + } + + stages = [] + first_divergence = None + for stage in REQUIRED_STAGES: + expected = load_array(np, reference, reference_arrays[stage]) + observed = load_array(np, actual, actual_arrays[stage]) + atol, rtol = stage_tolerance(stage) + finite = ( + np.isfinite(expected).all() and np.isfinite(observed).all() + if stage in FLOAT_STAGES + else True + ) + absolute = np.abs( + observed.astype(np.float64) - expected.astype(np.float64) + ) + threshold = atol + rtol * np.abs(expected.astype(np.float64)) + matches = bool(finite and np.all(absolute <= threshold)) + mismatch = ( + int(np.flatnonzero(absolute > threshold)[0]) + if finite and np.any(absolute > threshold) + else None + ) + entry = { + "stage": stage, + "passed": matches, + "max_absolute": float(absolute.max(initial=0.0)), + "atol": atol, + "rtol": rtol, + } + if mismatch is not None: + entry["first_mismatch_flat_index"] = mismatch + stages.append(entry) + if not matches and first_divergence is None: + first_divergence = {"stage": stage, "flat_index": mismatch} + + return { + "passed": first_divergence is None, + "first_divergence": first_divergence, + "stages": stages, + } + + +def compare(args: argparse.Namespace) -> int: + report = compare_capture_roots(args.reference, args.actual) + args.report.parent.mkdir(parents=True, exist_ok=True) + args.report.write_text( + json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8" + ) + print(args.report) + return 0 if report["passed"] else 1 + + +def parser() -> argparse.ArgumentParser: + root = argparse.ArgumentParser() + commands = root.add_subparsers(dest="command", required=True) + capture_parser = commands.add_parser("capture") + capture_parser.add_argument("--model", type=Path, required=True) + capture_parser.add_argument("--image", type=Path, required=True) + capture_parser.add_argument("--out", type=Path, required=True) + capture_parser.add_argument("--max-new", type=int, default=MAX_NEW_TOKENS) + capture_parser.set_defaults(run=capture) + compare_parser = commands.add_parser("compare") + compare_parser.add_argument("--reference", type=Path, required=True) + compare_parser.add_argument("--actual", type=Path, required=True) + compare_parser.add_argument("--report", type=Path, required=True) + compare_parser.set_defaults(run=compare) + return root + + +def main() -> int: + args = parser().parse_args() + try: + return int(args.run(args)) + except ContractError as error: + print(f"error: {error}", file=sys.stderr) + return 2 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/lib.rs b/src/lib.rs index d9520a12f..1b4787099 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -78,7 +78,9 @@ pub use mlxcel_core::speculative::SpeculativeGenerator; #[cfg(feature = "xla-diagnostics")] pub use multimodal::host_preprocessor::LlavaHostReferenceCapture; #[cfg(feature = "xla-iree")] -pub use multimodal::host_preprocessor::LlavaIreeHostPreprocessor; +pub use multimodal::host_preprocessor::{ + LlavaIreeHostPreprocessor, YoutuVlIreeHostPreprocessor, +}; pub use multimodal::host_preprocessor::{ FakeHostMultimodalPreprocessor, HostMultimodalPreprocessor, HostPreprocessorError, LlavaHostPreprocessor, XlaVisionBackend, load_xla_image_preprocessor, diff --git a/src/lib/mlxcel-xla/src/youtu_vl_runtime.rs b/src/lib/mlxcel-xla/src/youtu_vl_runtime.rs index 39bde2842..cbb173ef5 100644 --- a/src/lib/mlxcel-xla/src/youtu_vl_runtime.rs +++ b/src/lib/mlxcel-xla/src/youtu_vl_runtime.rs @@ -73,8 +73,14 @@ pub struct YoutuVlVisionDiagnosticProjection { pub stages: Vec, pub patch_shape: [usize; 2], pub merged_shape: [usize; 2], + /// Processor patch rows after the production window-group permutation. + pub patches_window_order: Vec, + /// Production two-dimensional vision RoPE frequencies in window order. + pub rope_freqs: Vec, pub window_group_index: Vec, pub reverse_group_index: Vec, + pub window_cu_seqlens: Vec, + pub full_cu_seqlens: Vec, } struct ActiveModule { @@ -569,13 +575,24 @@ impl IreeYoutuVlDiagnosticProjector { values: restored_output, shape: [actual_tokens, self.config.text_hidden], }); + let actual_patches = plan.actual_patches; + let patch_width = self.config.channels * self.config.patch_size * self.config.patch_size; + let rope_width = self.config.hidden / self.config.heads / 2; + let mut patches_window_order = patches; + patches_window_order.truncate(actual_patches * patch_width); + let mut rope_freqs = rope_freqs; + rope_freqs.truncate(actual_patches * rope_width); Ok(YoutuVlVisionDiagnosticProjection { abi_version: YOUTU_VL_DIAGNOSTIC_ABI_VERSION, stages, patch_shape, merged_shape: [actual_tokens, self.config.text_hidden], + patches_window_order, + rope_freqs, window_group_index: plan.window_group_index, reverse_group_index: plan.reverse_group_index, + window_cu_seqlens: plan.window_cu_seqlens, + full_cu_seqlens: plan.full_cu_seqlens, }) } } From 7ce28abbc65758daf08c8dc137c30d3d40d811aa Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 01:05:35 +0900 Subject: [PATCH 06/14] style: apply cargo fmt import ordering --- src/lib.rs | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index 1b4787099..56209a70d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -77,14 +77,12 @@ pub use mlxcel_core::generate::{ pub use mlxcel_core::speculative::SpeculativeGenerator; #[cfg(feature = "xla-diagnostics")] pub use multimodal::host_preprocessor::LlavaHostReferenceCapture; -#[cfg(feature = "xla-iree")] -pub use multimodal::host_preprocessor::{ - LlavaIreeHostPreprocessor, YoutuVlIreeHostPreprocessor, -}; pub use multimodal::host_preprocessor::{ FakeHostMultimodalPreprocessor, HostMultimodalPreprocessor, HostPreprocessorError, LlavaHostPreprocessor, XlaVisionBackend, load_xla_image_preprocessor, }; +#[cfg(feature = "xla-iree")] +pub use multimodal::host_preprocessor::{LlavaIreeHostPreprocessor, YoutuVlIreeHostPreprocessor}; pub use multimodal::{ internvl_prompt, kimi_vl_prompt, minicpmo_prompt, moondream2_prompt, moondream3_prompt, phi3v_prompt, phi4_siglip_prompt, phi4mm_prompt, pixtral_prompt, qwen_vl, smolvlm_prompt, From d107d8b397b8fff1d4cc540c6275a09b59fb0af3 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 09:02:41 +0900 Subject: [PATCH 07/14] fix(xla): mark Youtu-VL artifact unqualified Keep the rebased Youtu-VL runtime on the explicit legacy artifact path until the full numeric contract is qualified.\n\nRefs #872 --- src/lib/mlxcel-xla/src/youtu_vl_runtime.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/lib/mlxcel-xla/src/youtu_vl_runtime.rs b/src/lib/mlxcel-xla/src/youtu_vl_runtime.rs index cbb173ef5..8b711079d 100644 --- a/src/lib/mlxcel-xla/src/youtu_vl_runtime.rs +++ b/src/lib/mlxcel-xla/src/youtu_vl_runtime.rs @@ -253,7 +253,7 @@ fn compile_and_load( std::fs::create_dir_all(&cache) .map_err(|error| format!("mkdir {}: {error}", cache.display()))?; let (weights, checkpoint_schema) = load_weights(model_dir, config)?; - let contract = AuxiliaryArtifactContract::new( + let contract = AuxiliaryArtifactContract::new_legacy_unqualified( ENTRY_NAME, format!( "{};{};checkpoint_schema_sha256={}", From 6199aff7870e0cf0cea35e4778f4df0e94ce8eee Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 10:55:34 +0900 Subject: [PATCH 08/14] fix(xla): keep Youtu diagnostics unqualified --- src/lib/mlxcel-xla/src/youtu_vl_runtime.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/lib/mlxcel-xla/src/youtu_vl_runtime.rs b/src/lib/mlxcel-xla/src/youtu_vl_runtime.rs index 8b711079d..2907f43b4 100644 --- a/src/lib/mlxcel-xla/src/youtu_vl_runtime.rs +++ b/src/lib/mlxcel-xla/src/youtu_vl_runtime.rs @@ -289,7 +289,7 @@ fn compile_and_load_diagnostics( std::fs::create_dir_all(&cache) .map_err(|error| format!("mkdir {}: {error}", cache.display()))?; let (weights, checkpoint_schema) = load_weights(model_dir, config)?; - let contract = AuxiliaryArtifactContract::new( + let contract = AuxiliaryArtifactContract::new_legacy_unqualified( DIAGNOSTIC_ENTRY_NAME, format!( "{};diagnostic_contract={};{};checkpoint_schema_sha256={}", From 7378c1b5de9ccfbc255cfc1a30ed16dffdecb479 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 11:22:43 +0900 Subject: [PATCH 09/14] fix(xla): bind Youtu oracle to actual fixture grid --- examples/xla_youtu_vl_reference_check.rs | 33 ++++++++++++------- .../openxla/test_youtu_vl_reference_oracle.py | 2 +- spike/openxla/youtu_vl_reference_oracle.py | 24 +++++++++----- src/lib/mlxcel-xla/src/emitter/youtu_vl.rs | 4 +-- .../mlxcel-xla/src/emitter/youtu_vl_plan.rs | 2 +- src/lib/mlxcel-xla/src/youtu_vl_runtime.rs | 7 ++-- 6 files changed, 46 insertions(+), 26 deletions(-) diff --git a/examples/xla_youtu_vl_reference_check.rs b/examples/xla_youtu_vl_reference_check.rs index 2804c724f..8817c37fd 100644 --- a/examples/xla_youtu_vl_reference_check.rs +++ b/examples/xla_youtu_vl_reference_check.rs @@ -34,7 +34,7 @@ use mlxcel_xla::{IreeYoutuVlDiagnosticProjector, YoutuVlReferenceDiagnosticEngin use serde_json::{Map, Value, json}; use sha2::{Digest, Sha256}; -const CONTRACT: &str = "youtu-vl-hf-eager-oracle-v1"; +const CONTRACT: &str = "youtu-vl-hf-eager-oracle-v2"; const CHECKPOINT_REPO: &str = "tencent/Youtu-VL-4B-Instruct"; const CHECKPOINT_REVISION: &str = "8d30a0e49662a1d628a472b12df264dbcd768753"; const CHECKPOINT_ARTIFACT_MANIFEST_SHA256: &str = @@ -44,7 +44,10 @@ const FIXTURE_SHA256: &str = "5e7d54e8a7d21802378c87d2d70cf551e29739fe27599ddf12 const PROMPT: &str = "<|image_pad|>\nDescribe the image briefly."; const IMAGE_TOKEN_ID: i32 = 128_264; const PATCH_SIZE: usize = 16; -const PATCHES: usize = 256; +const MAX_IMAGE_PATCHES: usize = 256; +const PATCHES: usize = 196; +const PATCH_GRID: usize = 14; +const MERGED_TOKENS: usize = PATCHES / 4; const PATCH_WIDTH: usize = 3 * PATCH_SIZE * PATCH_SIZE; const TEXT_HIDDEN: usize = 2_560; const TEXT_LAYERS: usize = 40; @@ -311,14 +314,14 @@ fn main() { let processor = YoutuVLProcessor::new(PATCH_SIZE, 2) .with_norm([0.5; 3], [0.5; 3]) - .with_max_patches_per_image(PATCHES) + .with_max_patches_per_image(MAX_IMAGE_PATCHES) .with_resample(2); let (flattened_patches, spatial_shapes, patch_width) = processor .try_preprocess_values_with_spatial(std::slice::from_ref(&image)) .expect("run checked Youtu-VL flattened-patch processor"); - assert_eq!(spatial_shapes, [(16, 16)]); + assert_eq!(spatial_shapes, [(PATCH_GRID as i32, PATCH_GRID as i32)]); assert_eq!(patch_width, PATCH_WIDTH); - let resized_pixels = reconstruct_pixels(&flattened_patches, 16, 16); + let resized_pixels = reconstruct_pixels(&flattened_patches, PATCH_GRID, PATCH_GRID); let _runtime = initialize_runtime(); let mut diagnostic_projector = IreeYoutuVlDiagnosticProjector::load(&model, &device) @@ -326,9 +329,9 @@ fn main() { let vision = diagnostic_projector .capture(&flattened_patches, &spatial_shapes) .expect("capture actual 27-layer Youtu-VL vision graph"); - assert_eq!(vision.abi_version, 2); + assert_eq!(vision.abi_version, 3); assert_eq!(vision.patch_shape, [PATCHES, 1_152]); - assert_eq!(vision.merged_shape, [64, TEXT_HIDDEN]); + assert_eq!(vision.merged_shape, [MERGED_TOKENS, TEXT_HIDDEN]); let production_preprocessor = YoutuVlIreeHostPreprocessor::load(&model, &device) .expect("load production Youtu-VL IREE host preprocessor"); @@ -346,7 +349,7 @@ fn main() { .iter() .filter(|&&token| token == IMAGE_TOKEN_ID) .count(), - 64 + MERGED_TOKENS ); let mut engine = YoutuVlReferenceDiagnosticEngine::load(&model, &device, context_capacity) @@ -376,7 +379,7 @@ fn main() { &out, "resized_normalized_pixels", &resized_pixels, - &[3, 256, 256], + &[3, PATCH_GRID * PATCH_SIZE, PATCH_GRID * PATCH_SIZE], ); write_f32( &mut arrays, @@ -460,20 +463,26 @@ fn main() { &placeholder_positions, &[placeholder_positions.len()], ); - write_i32(&mut arrays, &out, "spatial_shapes", &[16, 16], &[1, 2]); + write_i32( + &mut arrays, + &out, + "spatial_shapes", + &[PATCH_GRID as i32, PATCH_GRID as i32], + &[1, 2], + ); write_i32( &mut arrays, &out, "window_group_index", &usize_i32(&vision.window_group_index, "window group index"), - &[64], + &[MERGED_TOKENS], ); write_i32( &mut arrays, &out, "reverse_group_index", &usize_i32(&vision.reverse_group_index, "reverse group index"), - &[64], + &[MERGED_TOKENS], ); write_i32( &mut arrays, diff --git a/spike/openxla/test_youtu_vl_reference_oracle.py b/spike/openxla/test_youtu_vl_reference_oracle.py index ce342a714..22a3d7141 100644 --- a/spike/openxla/test_youtu_vl_reference_oracle.py +++ b/spike/openxla/test_youtu_vl_reference_oracle.py @@ -99,7 +99,7 @@ def array(stage: str, shape: list[int], dtype: str) -> np.ndarray: elif stage == "placeholder_positions": value[:] = np.arange(oracle.MERGED_TOKENS, dtype=np.int32) elif stage == "spatial_shapes": - value[:] = [[16, 16]] + value[:] = [[oracle.PATCH_GRID, oracle.PATCH_GRID]] elif stage in {"window_group_index", "reverse_group_index"}: value[:] = np.arange(oracle.MERGED_TOKENS, dtype=np.int32) elif stage in {"window_cu_seqlens", "full_cu_seqlens"}: diff --git a/spike/openxla/youtu_vl_reference_oracle.py b/spike/openxla/youtu_vl_reference_oracle.py index b59826d18..7ad2ed490 100644 --- a/spike/openxla/youtu_vl_reference_oracle.py +++ b/spike/openxla/youtu_vl_reference_oracle.py @@ -20,7 +20,7 @@ SCHEMA_VERSION = 1 -CONTRACT = "youtu-vl-hf-eager-oracle-v1" +CONTRACT = "youtu-vl-hf-eager-oracle-v2" CHECKPOINT_REPO = "tencent/Youtu-VL-4B-Instruct" CHECKPOINT_REVISION = "8d30a0e49662a1d628a472b12df264dbcd768753" CHECKPOINT_ARTIFACTS = { @@ -63,7 +63,10 @@ KV_WIDTH = KV_LORA_RANK + ROPE_WIDTH PATCH_SIZE = 16 PATCH_WIDTH = 3 * PATCH_SIZE * PATCH_SIZE -PATCHES = 256 +MAX_IMAGE_PATCHES = 256 +PATCHES = 196 +PATCH_GRID = 14 +RESIZED_IMAGE_SIDE = PATCH_GRID * PATCH_SIZE MERGED_TOKENS = PATCHES // 4 VISION_ROPE_WIDTH = 36 MAX_NEW_TOKENS = 4 @@ -388,7 +391,7 @@ def capture(args: argparse.Namespace) -> int: inputs = processor( text=PROMPT, images=image, - max_image_patches=PATCHES, + max_image_patches=MAX_IMAGE_PATCHES, add_special_tokens=False, return_tensors="pt", ) @@ -396,10 +399,11 @@ def capture(args: argparse.Namespace) -> int: pixel_values = inputs["pixel_values"] spatial_shapes = inputs["spatial_shapes"] if list(spatial_shapes.shape) != [1, 2] or spatial_shapes.tolist() != [ - [16, 16] + [PATCH_GRID, PATCH_GRID] ]: raise ContractError( - f"pinned fixture grid must be [[16, 16]], got {spatial_shapes.tolist()}" + "pinned fixture grid must be " + f"[[{PATCH_GRID}, {PATCH_GRID}]], got {spatial_shapes.tolist()}" ) flat_patches = pixel_values[0] encoder = model.siglip2.vision_model.encoder @@ -484,7 +488,7 @@ def store(stage: str, value: Any, dtype: str) -> None: store( "resized_normalized_pixels", - unpatch(torch, flat_patches, 16, 16), + unpatch(torch, flat_patches, PATCH_GRID, PATCH_GRID), "float32", ) store("flattened_patches", flat_patches, "float32") @@ -550,7 +554,11 @@ def store(stage: str, value: Any, dtype: str) -> None: def expected_shapes(sequence: int) -> dict[str, list[int]]: shapes = { - "resized_normalized_pixels": [3, 256, 256], + "resized_normalized_pixels": [ + 3, + RESIZED_IMAGE_SIDE, + RESIZED_IMAGE_SIDE, + ], "flattened_patches": [PATCHES, PATCH_WIDTH], "patches.window_order": [PATCHES, PATCH_WIDTH], "vision_rope.freqs": [PATCHES, VISION_ROPE_WIDTH], @@ -723,7 +731,7 @@ def load_array(np: Any, root: Path, spec: dict[str, Any]) -> Any: def validate_semantics(np: Any, root: Path, arrays: dict[str, Any], role: str) -> None: spatial = load_array(np, root, arrays["spatial_shapes"]) - if spatial.tolist() != [[16, 16]]: + if spatial.tolist() != [[PATCH_GRID, PATCH_GRID]]: raise ContractError(f"{role} spatial grid differs") expanded = load_array(np, root, arrays["expanded_input_ids"]) placeholders = load_array(np, root, arrays["placeholder_positions"]) diff --git a/src/lib/mlxcel-xla/src/emitter/youtu_vl.rs b/src/lib/mlxcel-xla/src/emitter/youtu_vl.rs index 06a32f2e5..bcce4e7e0 100644 --- a/src/lib/mlxcel-xla/src/emitter/youtu_vl.rs +++ b/src/lib/mlxcel-xla/src/emitter/youtu_vl.rs @@ -400,7 +400,7 @@ mod tests { assert!(ir.contains("Youtu-VL diagnostic checkpoints")); assert_eq!( config.diagnostic_contract_identity(), - "youtu-vl-vision-diagnostics-v2:module=[patch_projection,layer.0.window,layer.1.full,post_layernorm,merger.window_order]:host=[merger.restored_order]" + "youtu-vl-vision-diagnostics-v3:module=[patch_projection,layer.0.window,layer.1.full,post_layernorm,merger.window_order]:host=[merger.restored_order]" ); } @@ -447,7 +447,7 @@ mod tests { let compiler = std::env::var("MLXCEL_XLA_IREE_COMPILE").expect("set MLXCEL_XLA_IREE_COMPILE"); let graph = emit_youtu_vl_diagnostics(&actual_schedule_config(), 16).unwrap(); - let stem = format!("mlxcel-youtu-vl-diagnostics-v2-{}", std::process::id()); + let stem = format!("mlxcel-youtu-vl-diagnostics-v3-{}", std::process::id()); let input = std::env::temp_dir().join(format!("{stem}.mlir")); let output = std::env::temp_dir().join(format!("{stem}.vmfb")); std::fs::write(&input, graph).expect("write temporary StableHLO"); diff --git a/src/lib/mlxcel-xla/src/emitter/youtu_vl_plan.rs b/src/lib/mlxcel-xla/src/emitter/youtu_vl_plan.rs index e88809979..02b64bb33 100644 --- a/src/lib/mlxcel-xla/src/emitter/youtu_vl_plan.rs +++ b/src/lib/mlxcel-xla/src/emitter/youtu_vl_plan.rs @@ -15,7 +15,7 @@ use serde_json::Value; pub(crate) const YOUTU_VL_PATCH_BUCKETS: [usize; 3] = [16, 64, 256]; #[cfg(any(test, feature = "diagnostics"))] -pub const YOUTU_VL_DIAGNOSTIC_ABI_VERSION: u32 = 2; +pub const YOUTU_VL_DIAGNOSTIC_ABI_VERSION: u32 = 3; const MAX_IMAGES: usize = 4; const DEFAULT_IMAGE_TOKEN_ID: i32 = 128_264; const DEFAULT_VIDEO_TOKEN_ID: i32 = 128_265; diff --git a/src/lib/mlxcel-xla/src/youtu_vl_runtime.rs b/src/lib/mlxcel-xla/src/youtu_vl_runtime.rs index 2907f43b4..0602a82be 100644 --- a/src/lib/mlxcel-xla/src/youtu_vl_runtime.rs +++ b/src/lib/mlxcel-xla/src/youtu_vl_runtime.rs @@ -490,7 +490,8 @@ impl IreeYoutuVlDiagnosticProjector { self.config.hidden / self.config.heads / 2, ]; let bias_shape = [plan.patch_bucket, plan.patch_bucket]; - let patch_shape = [plan.patch_bucket, self.config.hidden]; + let actual_patches = plan.actual_patches; + let patch_shape = [actual_patches, self.config.hidden]; let module_stages = self.config.diagnostic_stages(); let output_shapes = module_stages .iter() @@ -552,6 +553,9 @@ impl IreeYoutuVlDiagnosticProjector { if stage.is_merger_window_order() { values.truncate(actual_values); shape[0] = actual_tokens; + } else { + values.truncate(actual_patches * shape[1]); + shape[0] = actual_patches; } YoutuVlVisionDiagnosticStageOutput { name: stage.name(), @@ -575,7 +579,6 @@ impl IreeYoutuVlDiagnosticProjector { values: restored_output, shape: [actual_tokens, self.config.text_hidden], }); - let actual_patches = plan.actual_patches; let patch_width = self.config.channels * self.config.patch_size * self.config.patch_size; let rope_width = self.config.hidden / self.config.heads / 2; let mut patches_window_order = patches; From 91714926b31c7ffd3178bf084c57657811915de5 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 11:24:19 +0900 Subject: [PATCH 10/14] fix(xla): bound Youtu diagnostic task threads --- examples/xla_youtu_vl_reference_check.rs | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/examples/xla_youtu_vl_reference_check.rs b/examples/xla_youtu_vl_reference_check.rs index 8817c37fd..7bd1c2ac8 100644 --- a/examples/xla_youtu_vl_reference_check.rs +++ b/examples/xla_youtu_vl_reference_check.rs @@ -30,7 +30,10 @@ use mlxcel::{ HostMultimodalPreprocessor, PreparedTensorDType, YoutuVlIreeHostPreprocessor, initialize_runtime, }; -use mlxcel_xla::{IreeYoutuVlDiagnosticProjector, YoutuVlReferenceDiagnosticEngine}; +use mlxcel_xla::{ + IreeYoutuVlDiagnosticProjector, YoutuVlReferenceDiagnosticEngine, + configure_diagnostic_local_task_threads, +}; use serde_json::{Map, Value, json}; use sha2::{Digest, Sha256}; @@ -323,6 +326,8 @@ fn main() { assert_eq!(patch_width, PATCH_WIDTH); let resized_pixels = reconstruct_pixels(&flattened_patches, PATCH_GRID, PATCH_GRID); + configure_diagnostic_local_task_threads() + .expect("configure diagnostics-only IREE local-task threads"); let _runtime = initialize_runtime(); let mut diagnostic_projector = IreeYoutuVlDiagnosticProjector::load(&model, &device) .expect("load Youtu-VL diagnostic vision projector"); From 4ff8de4f73581ecbadcb4449741c2fbc6decb639 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 12:42:58 +0900 Subject: [PATCH 11/14] fix(xla): expand Youtu image placeholders Expand one logical image or video placeholder per spatial shape to the exact number of merged visual tokens expected by Youtu-VL. Preserve already-expanded prompts and reject ambiguous partial expansion before runtime upload. Refs #872 --- src/multimodal/host_preprocessor.rs | 78 +++++++++++++++++++++-- src/multimodal/host_preprocessor_tests.rs | 40 +++++++++++- 2 files changed, 112 insertions(+), 6 deletions(-) diff --git a/src/multimodal/host_preprocessor.rs b/src/multimodal/host_preprocessor.rs index 3969adb39..67883921d 100644 --- a/src/multimodal/host_preprocessor.rs +++ b/src/multimodal/host_preprocessor.rs @@ -350,15 +350,83 @@ fn prepare_youtu_vl_prompt( vision_start_token_id: i32, vision_end_token_id: i32, ) -> Result<(Vec, i32), HostPreprocessorError> { - let mut logical_tokens = token_ids.to_vec(); - if logical_tokens.contains(&image_token_id) { - return Ok((logical_tokens, image_token_id)); + let merge = i32::try_from(spatial_merge_size).map_err(|_| { + HostPreprocessorError::InvalidConfig( + "Youtu-VL spatial merge size exceeds i32::MAX".to_string(), + ) + })?; + if merge == 0 { + return Err(HostPreprocessorError::InvalidConfig( + "Youtu-VL spatial merge size must be positive".to_string(), + )); + } + let expected_per_image = spatial_shapes + .iter() + .map(|&(height, width)| { + if height <= 0 || width <= 0 || height % merge != 0 || width % merge != 0 { + return Err(HostPreprocessorError::InvalidConfig(format!( + "Youtu-VL spatial shape [{height}, {width}] must be positive and divisible by merge size {merge}" + ))); + } + let tokens = (height / merge) + .checked_mul(width / merge) + .ok_or(HostPreprocessorError::ShapeOverflow)?; + usize::try_from(tokens) + .map_err(|_| HostPreprocessorError::ShapeOverflow) + }) + .collect::, _>>()?; + let expected_total = expected_per_image + .iter() + .try_fold(0usize, |total, &tokens| { + total + .checked_add(tokens) + .ok_or(HostPreprocessorError::ShapeOverflow) + })?; + + let expand_existing = |target_token_id| { + let actual = token_ids + .iter() + .filter(|&&token| token == target_token_id) + .count(); + if actual == expected_total { + return Ok(token_ids.to_vec()); + } + if actual != expected_per_image.len() { + return Err(HostPreprocessorError::ExpandedLength { + actual, + expected: expected_total, + }); + } + let capacity = token_ids + .len() + .checked_sub(actual) + .and_then(|length| length.checked_add(expected_total)) + .ok_or(HostPreprocessorError::ShapeOverflow)?; + let mut expanded = Vec::with_capacity(capacity); + let mut image_index = 0usize; + for &token in token_ids { + if token == target_token_id { + expanded.extend(std::iter::repeat_n( + target_token_id, + expected_per_image[image_index], + )); + image_index += 1; + } else { + expanded.push(token); + } + } + Ok(expanded) + }; + + if token_ids.contains(&image_token_id) { + return expand_existing(image_token_id).map(|tokens| (tokens, image_token_id)); } - if logical_tokens.contains(&video_token_id) { + if token_ids.contains(&video_token_id) { // Placeholder compatibility only: the request boundary remains // image-only until a video preprocessing oracle is qualified. - return Ok((logical_tokens, video_token_id)); + return expand_existing(video_token_id).map(|tokens| (tokens, video_token_id)); } + let mut logical_tokens = token_ids.to_vec(); insert_youtu_vl_image_tokens( &mut logical_tokens, spatial_shapes, diff --git a/src/multimodal/host_preprocessor_tests.rs b/src/multimodal/host_preprocessor_tests.rs index 29f523734..b16fe833d 100644 --- a/src/multimodal/host_preprocessor_tests.rs +++ b/src/multimodal/host_preprocessor_tests.rs @@ -58,7 +58,7 @@ fn youtu_vl_placeholder_selection_prefers_image_then_video_then_insertion() { prepare_youtu_vl_prompt(&[1, VIDEO, IMAGE, 2], &shapes, 2, IMAGE, VIDEO, START, END) .unwrap(); assert_eq!(target, IMAGE); - assert_eq!(tokens, vec![1, VIDEO, IMAGE, 2]); + assert_eq!(tokens, vec![1, VIDEO, IMAGE, IMAGE, IMAGE, IMAGE, 2]); let (tokens, target) = prepare_youtu_vl_prompt( &[1, VIDEO, VIDEO, VIDEO, VIDEO, 2], @@ -79,6 +79,44 @@ fn youtu_vl_placeholder_selection_prefers_image_then_video_then_insertion() { assert_eq!(tokens, vec![1, START, IMAGE, IMAGE, IMAGE, IMAGE, END, 2]); } +#[cfg(feature = "xla-iree")] +#[test] +fn youtu_vl_prompt_expands_one_placeholder_per_image() { + const IMAGE: i32 = 100; + const VIDEO: i32 = 101; + const START: i32 = 102; + const END: i32 = 103; + + let (tokens, target) = prepare_youtu_vl_prompt( + &[1, IMAGE, 2, IMAGE, 3], + &[(4, 4), (2, 2)], + 2, + IMAGE, + VIDEO, + START, + END, + ) + .unwrap(); + + assert_eq!(target, IMAGE); + assert_eq!(tokens, vec![1, IMAGE, IMAGE, IMAGE, IMAGE, 2, IMAGE, 3]); +} + +#[cfg(feature = "xla-iree")] +#[test] +fn youtu_vl_prompt_rejects_ambiguous_partial_expansion() { + const IMAGE: i32 = 100; + let error = prepare_youtu_vl_prompt(&[1, IMAGE, IMAGE, 2], &[(4, 4)], 2, IMAGE, 101, 102, 103) + .unwrap_err(); + assert!(matches!( + error, + HostPreprocessorError::ExpandedLength { + actual: 2, + expected: 4 + } + )); +} + #[cfg(feature = "xla-iree")] #[test] #[ignore = "requires YOUTU_VL_MODEL_DIR and the pinned production IREE runtime/compiler"] From 69b9597eae98b9e5f16a07e909a675a252c46936 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 12:44:00 +0900 Subject: [PATCH 12/14] fix(xla): recover stale safetensors shard hints Resolve dense checkpoint tensors by scanning every shard named by the index instead of trusting an individual stale weight-map entry. Continue rejecting missing, duplicate, and mixed canonical/wrapped namespaces. Refs #872 --- src/lib/mlxcel-xla/src/iree.rs | 219 ++++++++++++++++++++++----------- 1 file changed, 148 insertions(+), 71 deletions(-) diff --git a/src/lib/mlxcel-xla/src/iree.rs b/src/lib/mlxcel-xla/src/iree.rs index e31100190..9b6167151 100644 --- a/src/lib/mlxcel-xla/src/iree.rs +++ b/src/lib/mlxcel-xla/src/iree.rs @@ -1050,42 +1050,6 @@ fn compile_gemma3n_diagnostics( Ok((vmfb, layout)) } -/// Map each needed weight name to the safetensors file that holds it, in the same -/// order as `names`. A single-file checkpoint (`model.safetensors` present) maps -/// every name to that one file; a sharded checkpoint (no `model.safetensors`, only -/// `model-0000k-of-...safetensors` shards) reads `model.safetensors.index.json`'s -/// `weight_map`. The big untied models (e.g. Llama-3.1-8B) ship sharded, so this -/// is what lets them load. A model dir that has BOTH a `model.safetensors` and an -/// index uses the single file (the index then just points back at it). -fn resolve_weight_shards(model_dir: &Path, names: &[String]) -> Result, String> { - let single = model_dir.join("model.safetensors"); - if single.exists() { - return Ok(vec![single; names.len()]); - } - let index = model_dir.join("model.safetensors.index.json"); - let text = std::fs::read_to_string(&index).map_err(|e| { - format!( - "no model.safetensors and no readable model.safetensors.index.json in {}: {e}", - model_dir.display() - ) - })?; - let v: serde_json::Value = - serde_json::from_str(&text).map_err(|e| format!("parse {}: {e}", index.display()))?; - let map = v - .get("weight_map") - .and_then(serde_json::Value::as_object) - .ok_or_else(|| format!("{}: missing object `weight_map`", index.display()))?; - names - .iter() - .map(|name| { - map.get(name) - .and_then(serde_json::Value::as_str) - .map(|f| model_dir.join(f)) - .ok_or_else(|| format!("{}: weight_map has no entry for `{name}`", index.display())) - }) - .collect() -} - fn language_model_tensor_name(name: &str) -> String { format!("language_model.{name}") } @@ -1096,12 +1060,28 @@ enum DenseCheckpointNames { LanguageModelPrefixed, } +fn merge_dense_name_scheme( + current: Option, + observed: DenseCheckpointNames, +) -> Result, String> { + match current { + Some(current) if current != observed => Err( + "dense checkpoint mixes canonical and `language_model.*` tensor namespaces".to_string(), + ), + Some(current) => Ok(Some(current)), + None => Ok(Some(observed)), + } +} + /// Resolve dense-language weights from either a standalone checkpoint or the /// `language_model.*` namespace used by LLaVA wrappers. /// /// A single-file checkpoint defers namespace detection until its safetensors -/// header is open. A sharded checkpoint must choose one complete namespace -/// from the index so a partially mixed wrapper is rejected. +/// header is open. A sharded checkpoint scans every shard named by the index +/// instead of trusting an individual `weight_map` entry: the pinned Youtu-VL +/// index points `model.embed_tokens.weight` at the wrong shard even though the +/// tensor itself is present. The scan still rejects missing, duplicate, or +/// mixed-namespace tensors. fn resolve_dense_weight_sources( model_dir: &Path, names: &[String], @@ -1110,21 +1090,68 @@ fn resolve_dense_weight_sources( if single.exists() { return Ok((vec![single; names.len()], None)); } - if let Ok(paths) = resolve_weight_shards(model_dir, names) { - return Ok((paths, Some(DenseCheckpointNames::Canonical))); + let candidate_shards = checkpoint_candidate_shards(model_dir)?; + let mut resolved: Vec> = vec![None; names.len()]; + let mut scheme = None; + for shard in candidate_shards { + let file = + File::open(&shard).map_err(|error| format!("open {}: {error}", shard.display()))?; + // Safety: the file is read-only for the lifetime of this header scan. + let mmap = unsafe { Mmap::map(&file) } + .map_err(|error| format!("mmap {}: {error}", shard.display()))?; + let tensors = SafeTensors::deserialize(&mmap) + .map_err(|error| format!("parse {}: {error}", shard.display()))?; + for (index, canonical) in names.iter().enumerate() { + let wrapped = language_model_tensor_name(canonical); + let has_canonical = tensors.tensor(canonical).is_ok(); + let has_wrapped = tensors.tensor(&wrapped).is_ok(); + let observed = match (has_canonical, has_wrapped) { + (false, false) => continue, + (true, false) => DenseCheckpointNames::Canonical, + (false, true) => DenseCheckpointNames::LanguageModelPrefixed, + (true, true) => { + return Err(format!( + "dense checkpoint {} contains both `{canonical}` and `{wrapped}`", + shard.display() + )); + } + }; + if let Some(previous) = &resolved[index] { + return Err(format!( + "dense tensor `{canonical}` is duplicated in {} and {}", + previous.display(), + shard.display() + )); + } + resolved[index] = Some(shard.clone()); + scheme = merge_dense_name_scheme(scheme, observed)?; + } } - let wrapped: Vec = names + let missing = resolved .iter() - .map(|name| language_model_tensor_name(name)) - .collect(); - resolve_weight_shards(model_dir, &wrapped) - .map(|paths| (paths, Some(DenseCheckpointNames::LanguageModelPrefixed))) - .map_err(|error| { - format!( - "dense checkpoint has neither a complete canonical nor `language_model.*` \ - weight namespace: {error}" - ) - }) + .zip(names) + .filter_map(|(path, name)| path.is_none().then_some(name.as_str())) + .collect::>(); + if !missing.is_empty() { + return Err(format!( + "dense checkpoint is missing {} graph weight(s): {}", + missing.len(), + missing + .iter() + .take(8) + .copied() + .collect::>() + .join(", ") + )); + } + let scheme = scheme.ok_or_else(|| "dense checkpoint has no graph weights".to_string())?; + Ok(( + resolved + .into_iter() + .map(|path| path.expect("missing paths rejected above")) + .collect(), + Some(scheme), + )) } /// mlx-community Gemma3n quantized repacks retain the upstream index but rewrite @@ -1157,16 +1184,36 @@ fn merge_gemma3n_name_scheme( } } -fn gemma3n_candidate_shards(model_dir: &Path, names: &[String]) -> Result, String> { - let indexed = resolve_weight_shards(model_dir, names); - if let Ok(paths) = &indexed - && paths.iter().all(|path| path.is_file()) - { - let mut unique = paths.clone(); - unique.sort(); - unique.dedup(); - return Ok(unique); - } +fn checkpoint_candidate_shards(model_dir: &Path) -> Result, String> { + let index = model_dir.join("model.safetensors.index.json"); + let indexed = std::fs::read_to_string(&index) + .map_err(|error| format!("read {}: {error}", index.display())) + .and_then(|text| { + let value: serde_json::Value = serde_json::from_str(&text) + .map_err(|error| format!("parse {}: {error}", index.display()))?; + let weight_map = value + .get("weight_map") + .and_then(serde_json::Value::as_object) + .ok_or_else(|| format!("{}: missing object `weight_map`", index.display()))?; + let mut paths = weight_map + .values() + .filter_map(serde_json::Value::as_str) + .map(|file| model_dir.join(file)) + .collect::>(); + paths.sort(); + paths.dedup(); + if paths.is_empty() || paths.iter().any(|path| !path.is_file()) { + return Err(format!( + "{} does not name a complete set of existing safetensors shards", + index.display() + )); + } + Ok(paths) + }); + let indexed_error = match indexed { + Ok(paths) => return Ok(paths), + Err(error) => error, + }; let mut actual: Vec = std::fs::read_dir(model_dir) .map_err(|error| format!("read {}: {error}", model_dir.display()))? .filter_map(Result::ok) @@ -1181,10 +1228,7 @@ fn gemma3n_candidate_shards(model_dir: &Path, names: &[String]) -> Result error, - Ok(_) => format!("no safetensors files found in {}", model_dir.display()), - }); + return Err(indexed_error); } Ok(actual) } @@ -1193,7 +1237,7 @@ fn resolve_gemma3n_weight_sources( model_dir: &Path, names: &[String], ) -> Result<(Vec, Gemma3nCheckpointNames), String> { - let candidate_shards = gemma3n_candidate_shards(model_dir, names)?; + let candidate_shards = checkpoint_candidate_shards(model_dir)?; let mut resolved: Vec> = vec![None; names.len()]; let mut scheme = None; for shard in candidate_shards { @@ -1358,6 +1402,40 @@ mod checkpoint_name_tests { assert!(error.contains("mixes canonical")); } + #[test] + fn candidate_shards_scan_every_file_named_by_a_stale_index() { + let model_dir = std::env::temp_dir().join(format!( + "mlxcel-xla-stale-index-{}-{}", + std::process::id(), + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos() + )); + std::fs::create_dir(&model_dir).unwrap(); + let first = model_dir.join("model-00001-of-00002.safetensors"); + let second = model_dir.join("model-00002-of-00002.safetensors"); + std::fs::write(&first, []).unwrap(); + std::fs::write(&second, []).unwrap(); + std::fs::write( + model_dir.join("model.safetensors.index.json"), + serde_json::json!({ + "weight_map": { + "model.embed_tokens.weight": "model-00001-of-00002.safetensors", + "model.layers.0.input_layernorm.weight": "model-00002-of-00002.safetensors" + } + }) + .to_string(), + ) + .unwrap(); + + assert_eq!( + checkpoint_candidate_shards(&model_dir).unwrap(), + vec![first, second] + ); + std::fs::remove_dir_all(model_dir).unwrap(); + } + #[test] #[ignore = "requires GEMMA3N_MODEL_DIR pointing at a real local checkpoint"] fn real_gemma3n_headers_resolve_every_graph_weight_once() { @@ -1379,11 +1457,10 @@ mod checkpoint_name_tests { /// f32 buffers (kept alive until the shim copies them) plus the flat (ptr, rank, /// dims) arrays the C ABI takes. `cfg` fixes the layer count and weight order. /// -/// Single-file and sharded checkpoints both load: [`resolve_weight_shards`] maps -/// each weight to its file, and the weights are read shard by shard (each shard -/// mmap'd exactly once, its tensors copied out as owned f32) and placed back into -/// the emitter's arg order by index, so the buffers line up with the graph args -/// regardless of how the checkpoint was split. +/// Single-file and sharded checkpoints both load: the architecture-aware +/// resolvers map each weight to its containing shard, then each shard is mmap'd +/// exactly once and its owned buffers are placed back into the emitter's arg +/// order regardless of how the checkpoint was split. #[allow(clippy::type_complexity)] fn load_weights( model_dir: &Path, From bce4739cf25deec3d07f57e7066891e44e8f561c Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 12:44:10 +0900 Subject: [PATCH 13/14] fix(xla): match Youtu HF preprocessing layout Match the pinned HF processor's spatial-merge grouping and channel-fast patch layout, and avoid regrouping those rows again at the IREE boundary. Correct the independent oracle reconstruction and add fail-fast preprocessing and selected vision-stage comparisons with partial diagnostic artifacts. Refs #872 --- examples/xla_youtu_vl_reference_check.rs | 175 ++++++++++++++++-- spike/openxla/youtu_vl_reference_oracle.py | 19 +- .../mlxcel-xla/src/emitter/youtu_vl_plan.rs | 36 +--- src/vision/processors/youtu_vl.rs | 50 ++--- src/vision/processors/youtu_vl_tests.rs | 36 ++++ 5 files changed, 250 insertions(+), 66 deletions(-) diff --git a/examples/xla_youtu_vl_reference_check.rs b/examples/xla_youtu_vl_reference_check.rs index 7bd1c2ac8..3d85b9a6a 100644 --- a/examples/xla_youtu_vl_reference_check.rs +++ b/examples/xla_youtu_vl_reference_check.rs @@ -65,6 +65,10 @@ fn argument(flag: &str) -> Option { .cloned() } +fn has_flag(flag: &str) -> bool { + std::env::args().any(|value| value == flag) +} + fn required_path(flag: &str) -> PathBuf { argument(flag) .map(PathBuf::from) @@ -122,6 +126,75 @@ fn verify_reference_manifest(root: &Path) -> Value { manifest } +fn reference_f32(root: &Path, manifest: &Value, stage: &str) -> Vec { + let spec = &manifest["case"]["arrays"][stage]; + assert_eq!(spec["dtype"], "float32", "{stage} reference dtype differs"); + let filename = spec["file"] + .as_str() + .unwrap_or_else(|| panic!("{stage} reference filename must be a string")); + let path = root.join(filename); + let bytes = fs::read(&path).unwrap_or_else(|error| panic!("read {}: {error}", path.display())); + assert_eq!( + sha256_bytes(&bytes), + spec["sha256"] + .as_str() + .unwrap_or_else(|| panic!("{stage} reference hash must be a string")), + "{stage} reference artifact SHA-256 differs" + ); + assert!(bytes.len().is_multiple_of(4)); + bytes + .chunks_exact(4) + .map(|chunk| f32::from_le_bytes(chunk.try_into().expect("four-byte f32"))) + .collect() +} + +fn assert_reference_close( + root: &Path, + manifest: &Value, + stage: &str, + actual: &[f32], + atol: f32, + rtol: f32, +) { + let expected = reference_f32(root, manifest, stage); + assert_eq!( + actual.len(), + expected.len(), + "{stage} reference element count differs" + ); + let mismatch = + actual + .iter() + .zip(&expected) + .enumerate() + .find_map(|(index, (&observed, &wanted))| { + let threshold = atol + rtol * wanted.abs(); + (!observed.is_finite() || (observed - wanted).abs() > threshold) + .then_some((index, observed, wanted)) + }); + if let Some((index, observed, wanted)) = mismatch { + let mismatch_count = actual + .iter() + .zip(&expected) + .filter(|&(&observed, &wanted)| { + let threshold = atol + rtol * wanted.abs(); + !observed.is_finite() || (observed - wanted).abs() > threshold + }) + .count(); + let max_absolute = actual + .iter() + .zip(&expected) + .map(|(&observed, &wanted)| (observed - wanted).abs()) + .fold(0.0f32, f32::max); + panic!( + "{stage} differs from the pinned HF preprocessing at flat index {index}: \ + actual={observed}, reference={wanted}, absolute={}, \ + max_absolute={max_absolute}, mismatch_count={mismatch_count}", + (observed - wanted).abs() + ); + } +} + fn verify_checkpoint(model: &Path, checkpoint: &Value) { assert_eq!(checkpoint["repo"], CHECKPOINT_REPO); assert_eq!(checkpoint["revision"], CHECKPOINT_REVISION); @@ -238,21 +311,31 @@ fn usize_i32(values: &[usize], label: &str) -> Vec { fn reconstruct_pixels(patches: &[f32], grid_height: usize, grid_width: usize) -> Vec { assert_eq!(patches.len(), grid_height * grid_width * PATCH_WIDTH); + assert!(grid_height.is_multiple_of(2)); + assert!(grid_width.is_multiple_of(2)); let height = grid_height * PATCH_SIZE; let width = grid_width * PATCH_SIZE; let mut pixels = vec![0.0f32; 3 * height * width]; - for patch in 0..grid_height * grid_width { - let patch_y = patch / grid_width; - let patch_x = patch % grid_width; - let row = &patches[patch * PATCH_WIDTH..(patch + 1) * PATCH_WIDTH]; - let mut cursor = 0usize; - for channel in 0..3 { - for dy in 0..PATCH_SIZE { - for dx in 0..PATCH_SIZE { - let y = patch_y * PATCH_SIZE + dy; - let x = patch_x * PATCH_SIZE + dx; - pixels[channel * height * width + y * width + x] = row[cursor]; - cursor += 1; + let mut patch = 0usize; + for block_y in 0..grid_height / 2 { + for block_x in 0..grid_width / 2 { + for inner_y in 0..2 { + for inner_x in 0..2 { + let patch_y = block_y * 2 + inner_y; + let patch_x = block_x * 2 + inner_x; + let row = &patches[patch * PATCH_WIDTH..(patch + 1) * PATCH_WIDTH]; + let mut cursor = 0usize; + for dy in 0..PATCH_SIZE { + for dx in 0..PATCH_SIZE { + let y = patch_y * PATCH_SIZE + dy; + let x = patch_x * PATCH_SIZE + dx; + for channel in 0..3 { + pixels[channel * height * width + y * width + x] = row[cursor]; + cursor += 1; + } + } + } + patch += 1; } } } @@ -325,6 +408,26 @@ fn main() { assert_eq!(spatial_shapes, [(PATCH_GRID as i32, PATCH_GRID as i32)]); assert_eq!(patch_width, PATCH_WIDTH); let resized_pixels = reconstruct_pixels(&flattened_patches, PATCH_GRID, PATCH_GRID); + assert_reference_close( + &reference_root, + &reference, + "resized_normalized_pixels", + &resized_pixels, + 1e-5, + 1e-5, + ); + assert_reference_close( + &reference_root, + &reference, + "flattened_patches", + &flattened_patches, + 1e-5, + 1e-5, + ); + if has_flag("--preprocess-only") { + println!("Youtu-VL preprocessing matches the pinned HF capture"); + return; + } configure_diagnostic_local_task_threads() .expect("configure diagnostics-only IREE local-task threads"); @@ -337,6 +440,54 @@ fn main() { assert_eq!(vision.abi_version, 3); assert_eq!(vision.patch_shape, [PATCHES, 1_152]); assert_eq!(vision.merged_shape, [MERGED_TOKENS, TEXT_HIDDEN]); + for (stage, values) in [ + ("resized_normalized_pixels", resized_pixels.as_slice()), + ("flattened_patches", flattened_patches.as_slice()), + ( + "patches.window_order", + vision.patches_window_order.as_slice(), + ), + ("vision_rope.freqs", vision.rope_freqs.as_slice()), + ] { + fs::write( + out.join(format!("image_text.{stage}.bin")), + f32_bytes(values), + ) + .unwrap_or_else(|error| panic!("write partial {stage}: {error}")); + } + for stage in &vision.stages { + fs::write( + out.join(format!("image_text.{}.bin", stage.name)), + f32_bytes(&stage.values), + ) + .unwrap_or_else(|error| panic!("write partial {}: {error}", stage.name)); + } + assert_reference_close( + &reference_root, + &reference, + "patches.window_order", + &vision.patches_window_order, + 1e-5, + 1e-5, + ); + assert_reference_close( + &reference_root, + &reference, + "vision_rope.freqs", + &vision.rope_freqs, + 1e-5, + 1e-5, + ); + for stage in &vision.stages { + assert_reference_close( + &reference_root, + &reference, + &stage.name, + &stage.values, + 2e-2, + 2e-2, + ); + } let production_preprocessor = YoutuVlIreeHostPreprocessor::load(&model, &device) .expect("load production Youtu-VL IREE host preprocessor"); diff --git a/spike/openxla/youtu_vl_reference_oracle.py b/spike/openxla/youtu_vl_reference_oracle.py index 7ad2ed490..4cd08ba9a 100644 --- a/spike/openxla/youtu_vl_reference_oracle.py +++ b/spike/openxla/youtu_vl_reference_oracle.py @@ -247,9 +247,19 @@ def patch_order(torch: Any, window_index: Any) -> Any: def unpatch(torch: Any, patches: Any, height: int, width: int) -> Any: - rows = patches.reshape(height, width, 3, PATCH_SIZE, PATCH_SIZE) + if height % 2 or width % 2: + raise ContractError("Youtu-VL patch grid must be divisible by merge size 2") + rows = patches.reshape( + height // 2, + width // 2, + 2, + 2, + PATCH_SIZE, + PATCH_SIZE, + 3, + ) return ( - rows.permute(2, 0, 3, 1, 4) + rows.permute(6, 0, 2, 4, 1, 3, 5) .contiguous() .reshape(3, height * PATCH_SIZE, width * PATCH_SIZE) ) @@ -505,7 +515,10 @@ def store(stage: str, value: Any, dtype: str) -> None: store("selected_kv", fresh_kv, "float32") store("reuse_prefill_logits", reuse_logits, "float32") store("reuse_selected_kv", reuse_kv, "float32") - store("expanded_input_ids", input_ids, "int32") + # The oracle contract models one request and records the expanded token + # sequence without a redundant batch dimension. `prepared_embeddings` + # retains `[1, sequence, hidden]` because that is the model input boundary. + store("expanded_input_ids", input_ids[0], "int32") store( "placeholder_positions", (input_ids[0] == IMAGE_TOKEN_ID).nonzero().flatten(), diff --git a/src/lib/mlxcel-xla/src/emitter/youtu_vl_plan.rs b/src/lib/mlxcel-xla/src/emitter/youtu_vl_plan.rs index 02b64bb33..460c91366 100644 --- a/src/lib/mlxcel-xla/src/emitter/youtu_vl_plan.rs +++ b/src/lib/mlxcel-xla/src/emitter/youtu_vl_plan.rs @@ -498,31 +498,14 @@ pub(crate) fn prepare_youtu_vl_host_inputs( if patch_rows.iter().any(|value| !value.is_finite()) { return Err("Youtu-VL processor produced a non-finite patch value".to_string()); } + // The production processor already emits the checkpoint's + // spatial-merge-grouped row order. This boundary only permutes whole + // groups into the attention-window schedule. let merge_unit = config.spatial_merge_size * config.spatial_merge_size; - let mut grouped_rows = Vec::with_capacity(patch_rows.len()); - let mut image_patch_offset = 0usize; - for &(height, width) in shapes { - let (height, width) = (height as usize, width as usize); - for block_h in 0..height / config.spatial_merge_size { - for block_w in 0..width / config.spatial_merge_size { - for inner_h in 0..config.spatial_merge_size { - for inner_w in 0..config.spatial_merge_size { - let patch = image_patch_offset - + (block_h * config.spatial_merge_size + inner_h) * width - + block_w * config.spatial_merge_size - + inner_w; - let start = patch * patch_width; - grouped_rows.extend_from_slice(&patch_rows[start..start + patch_width]); - } - } - } - } - image_patch_offset += height * width; - } let mut patches = Vec::with_capacity(plan.patch_bucket * patch_width); for &group in &plan.window_group_index { let start = group * merge_unit * patch_width; - patches.extend_from_slice(&grouped_rows[start..start + merge_unit * patch_width]); + patches.extend_from_slice(&patch_rows[start..start + merge_unit * patch_width]); } patches.resize(plan.patch_bucket * patch_width, 0.0); @@ -670,7 +653,7 @@ mod tests { } #[test] - fn host_patch_reorder_preserves_merge_group_membership() { + fn host_patch_reorder_preserves_processor_merge_groups() { let config = config(); let width = config.channels * config.patch_size * config.patch_size; let mut rows = vec![0.0; 32 * width]; @@ -683,16 +666,9 @@ mod tests { .collect::>(); for (window_position, &group) in inputs.plan.window_group_index.iter().enumerate() { let group_rows = &observed[window_position * 4..window_position * 4 + 4]; - let group_h = group / 4; - let group_w = group % 4; assert_eq!( group_rows, - &[ - group_h * 16 + group_w * 2, - group_h * 16 + group_w * 2 + 1, - group_h * 16 + group_w * 2 + 8, - group_h * 16 + group_w * 2 + 9, - ] + &[group * 4, group * 4 + 1, group * 4 + 2, group * 4 + 3,] ); } } diff --git a/src/vision/processors/youtu_vl.rs b/src/vision/processors/youtu_vl.rs index 5a6f138b0..7c9b620c1 100644 --- a/src/vision/processors/youtu_vl.rs +++ b/src/vision/processors/youtu_vl.rs @@ -299,32 +299,40 @@ impl YoutuVLProcessor { } } - // Emit one row per spatial patch in the layout - // `[h_patches * w_patches, channels * patch_size * patch_size]`, - // with the inner ordering `(c, dy, dx)` to match how upstream - // unfolds patches via `unfold` over the (C, H, W) image tensor. - let total_patches_img = (h_patches as usize) * (w_patches as usize); - for patch_idx in 0..total_patches_img { - let py = patch_idx / w_patches as usize; - let px = patch_idx % w_patches as usize; - let y_start = py * self.patch_size; - let x_start = px * self.patch_size; - - let row_start = (write_offset + patch_idx) * features_per_patch; - let mut k = 0usize; - for c in 0..in_channels { - for dy in 0..self.patch_size { - for dx in 0..self.patch_size { - let y = y_start + dy; - let x = x_start + dx; - all_patches[row_start + k] = normalized[c * h * w + y * w + x]; - k += 1; + // Match the pinned HF processor's `convert_image_to_patches` + // exactly. Rows are grouped by spatial-merge block, then by the + // patch's position inside that block. Values within each row use + // `(dy, dx, channel)` ordering with channel fastest. + let merge = self.spatial_merge_size; + let mut image_row = 0usize; + for block_y in 0..h_patches as usize / merge { + for block_x in 0..w_patches as usize / merge { + for inner_y in 0..merge { + for inner_x in 0..merge { + let patch_y = block_y * merge + inner_y; + let patch_x = block_x * merge + inner_x; + let y_start = patch_y * self.patch_size; + let x_start = patch_x * self.patch_size; + let row_start = (write_offset + image_row) * features_per_patch; + let mut k = 0usize; + for dy in 0..self.patch_size { + for dx in 0..self.patch_size { + let y = y_start + dy; + let x = x_start + dx; + for c in 0..in_channels { + all_patches[row_start + k] = + normalized[c * h * w + y * w + x]; + k += 1; + } + } + } + image_row += 1; } } } } - write_offset += total_patches_img; + write_offset += image_row; } i32::try_from(total_patches).map_err(|_| YoutuVLPreprocessError::DimensionTooLarge { diff --git a/src/vision/processors/youtu_vl_tests.rs b/src/vision/processors/youtu_vl_tests.rs index 735d88d31..bd676db8b 100644 --- a/src/vision/processors/youtu_vl_tests.rs +++ b/src/vision/processors/youtu_vl_tests.rs @@ -154,3 +154,39 @@ fn normalization_matches_siglip_default() { mlxcel_core::item_f32(&max_abs) ); } + +#[test] +fn flattened_patch_layout_matches_hf_merge_groups_and_channel_order() { + let processor = YoutuVLProcessor::new(1, 2) + .with_pixel_bounds(4 * 4, 4 * 4) + .with_norm([0.0, 0.0, 0.0], [1.0, 1.0, 1.0]); + let mut image = RgbImage::new(4, 4); + for y in 0..4 { + for x in 0..4 { + let base = (y * 4 + x) as u8; + image.put_pixel(x, y, image::Rgb([base, base + 32, base + 64])); + } + } + + let (rows, shapes, width) = processor + .try_preprocess_values_with_spatial(&[DynamicImage::ImageRgb8(image)]) + .unwrap(); + assert_eq!(shapes, vec![(4, 4)]); + assert_eq!(width, 3); + + let observed_pixels = rows + .chunks_exact(3) + .map(|rgb| (rgb[0] * 255.0).round() as u8) + .collect::>(); + assert_eq!( + observed_pixels, + vec![0, 1, 4, 5, 2, 3, 6, 7, 8, 9, 12, 13, 10, 11, 14, 15] + ); + assert_eq!( + rows[..3] + .iter() + .map(|value| (value * 255.0).round() as u8) + .collect::>(), + vec![0, 32, 64] + ); +} From 20c2bffd39bca4707eaa90e50e2f9ad2a0a76ff7 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 12:44:25 +0900 Subject: [PATCH 14/14] fix(xla): preserve Youtu BF16 vision carriers Materialize the pinned checkpoint's BF16 vision boundaries while reusing the shared numeric LayerNorm, RMSNorm, GELU, and softmax decompositions. Keep the vision MLP's tanh GELU distinct from the merger's exact GELU. Refs #872 --- src/lib/mlxcel-xla/src/emitter/youtu_vl.rs | 110 ++++++++------------- 1 file changed, 40 insertions(+), 70 deletions(-) diff --git a/src/lib/mlxcel-xla/src/emitter/youtu_vl.rs b/src/lib/mlxcel-xla/src/emitter/youtu_vl.rs index bcce4e7e0..cd66af416 100644 --- a/src/lib/mlxcel-xla/src/emitter/youtu_vl.rs +++ b/src/lib/mlxcel-xla/src/emitter/youtu_vl.rs @@ -6,6 +6,7 @@ //! StableHLO Youtu-VL windowed vision tower and built-in patch merger. use super::builder::{Builder, Ty, Val}; +use super::numeric_ops::{exact_gelu, layer_norm_2d, rms_norm_2d, stable_softmax, tanh_gelu}; #[cfg(any(test, feature = "diagnostics"))] use super::youtu_vl_plan::YoutuVlDiagnosticStage; use super::youtu_vl_plan::{YOUTU_VL_PATCH_BUCKETS, YoutuVlVisionConfig, YoutuVlWeightSpec}; @@ -58,74 +59,35 @@ fn add_bias(builder: &mut Builder, value: &Val, bias: &Val) -> Val { builder.add(value, &bias) } +fn round_bf16(builder: &mut Builder, value: &Val) -> Val { + let rounded = builder.convert(value, "bf16"); + builder.convert(&rounded, "f32") +} + fn linear(builder: &mut Builder, value: &Val, weight: &Val, bias: &Val) -> Val { let value = builder.linear_seq(value, weight); - add_bias(builder, &value, bias) + let value = add_bias(builder, &value, bias); + round_bf16(builder, &value) } fn layer_norm(builder: &mut Builder, value: &Val, weight: &Val, bias: &Val, epsilon: f32) -> Val { - let rows = value.ty.shape[0]; - let width = value.ty.shape[1]; - let zero = builder.const_f32(0.0); - let width_scalar = builder.const_f32(width as f32); - let width_rows = builder.broadcast(&width_scalar, &[], vec![rows]); - let sum = builder.reduce_add(value, 1, &zero); - let mean = builder.divide(&sum, &width_rows); - let mean = builder.broadcast(&mean, &[0], vec![rows, width]); - let centered = builder.subtract(value, &mean); - let squared = builder.multiply(¢ered, ¢ered); - let sum = builder.reduce_add(&squared, 1, &zero); - let variance = builder.divide(&sum, &width_rows); - let epsilon = builder.const_f32(epsilon); - let epsilon = builder.broadcast(&epsilon, &[], vec![rows]); - let variance = builder.add(&variance, &epsilon); - let inverse = builder.rsqrt(&variance); - let inverse = builder.broadcast(&inverse, &[0], vec![rows, width]); - let normalized = builder.multiply(¢ered, &inverse); - let weight = builder.broadcast(weight, &[1], vec![rows, width]); - let bias = builder.broadcast(bias, &[1], vec![rows, width]); - let normalized = builder.multiply(&normalized, &weight); - builder.add(&normalized, &bias) + let normalized = layer_norm_2d(builder, value, weight, bias, epsilon); + round_bf16(builder, &normalized) } fn 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 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 = builder.divide(&sum, &width_rows); - let epsilon = builder.const_f32(epsilon); - let epsilon = builder.broadcast(&epsilon, &[], vec![rows]); - let mean = builder.add(&mean, &epsilon); - let scale = builder.rsqrt(&mean); - let scale = builder.broadcast(&scale, &[0], vec![rows, width]); - let weight = builder.broadcast(weight, &[1], vec![rows, width]); - let value = builder.multiply(value, &scale); - builder.multiply(&value, &weight) + let value = rms_norm_2d(builder, value, weight, epsilon); + round_bf16(builder, &value) } fn gelu_tanh(builder: &mut Builder, value: &Val) -> Val { - let shape = value.ty.shape.clone(); - let half_scalar = builder.const_f32(0.5); - let half = builder.broadcast(&half_scalar, &[], shape.clone()); - let one_scalar = builder.const_f32(1.0); - let one = builder.broadcast(&one_scalar, &[], shape.clone()); - let coefficient_scalar = builder.const_f32(0.044_715); - let coefficient = builder.broadcast(&coefficient_scalar, &[], shape.clone()); - let scale_scalar = builder.const_f32(0.797_884_6); - let scale = builder.broadcast(&scale_scalar, &[], shape); - let squared = builder.multiply(value, value); - let cubed = builder.multiply(&squared, value); - let nonlinear = builder.multiply(&coefficient, &cubed); - let inner = builder.add(value, &nonlinear); - let scaled = builder.multiply(&scale, &inner); - let tanh = builder.tanh(&scaled); - let cdf = builder.add(&one, &tanh); - let half_value = builder.multiply(value, &half); - builder.multiply(&half_value, &cdf) + let value = tanh_gelu(builder, value); + round_bf16(builder, &value) +} + +fn gelu_exact(builder: &mut Builder, value: &Val) -> Val { + let value = exact_gelu(builder, value); + round_bf16(builder, &value) } fn rotate_half(builder: &mut Builder, value: &Val) -> Val { @@ -150,20 +112,13 @@ fn apply_rope(builder: &mut Builder, value: &Val, freqs: &Val) -> Val { let direct = builder.multiply(value, &cos); let rotated = rotate_half(builder, value); let rotated = builder.multiply(&rotated, &sin); - builder.add(&direct, &rotated) + let value = builder.add(&direct, &rotated); + round_bf16(builder, &value) } fn softmax(builder: &mut Builder, value: &Val) -> Val { - let shape = value.ty.shape.clone(); - let negative = builder.const_f32(f32::NEG_INFINITY); - let maximum = builder.reduce_max(value, 2, &negative); - let maximum = builder.broadcast(&maximum, &[0, 1], shape.clone()); - let shifted = builder.subtract(value, &maximum); - let exponentials = builder.exponential(&shifted); - let zero = builder.const_f32(0.0); - let denominator = builder.reduce_add(&exponentials, 2, &zero); - let denominator = builder.broadcast(&denominator, &[0, 1], shape); - builder.divide(&exponentials, &denominator) + let probabilities = stable_softmax(builder, value, 2); + round_bf16(builder, &probabilities) } fn attention( @@ -196,11 +151,14 @@ fn attention( &[2], vec![config.heads, tokens, tokens], ); + let scores = round_bf16(builder, &scores); let scale = builder.const_f32((head_dim as f32).powf(-0.5)); let scale = builder.broadcast(&scale, &[], vec![config.heads, tokens, tokens]); let scores = builder.multiply(&scores, &scale); + let scores = round_bf16(builder, &scores); let bias = builder.broadcast(bias, &[1, 2], vec![config.heads, tokens, tokens]); let scores = builder.add(&scores, &bias); + let scores = round_bf16(builder, &scores); let probabilities = softmax(builder, &scores); let context = builder.dot_general( &probabilities, @@ -211,6 +169,7 @@ fn attention( &[1], vec![config.heads, tokens, head_dim], ); + let context = round_bf16(builder, &context); let context = builder.transpose(&context, &[1, 0, 2]); let context = builder.reshape(&context, vec![tokens, config.hidden]); linear(builder, &context, &args.take(), &args.take()) @@ -233,6 +192,7 @@ fn layer( ); let attention = attention(builder, &normalized, args, config, freqs, bias); let residual = builder.add(hidden, &attention); + let residual = round_bf16(builder, &residual); let normalized = layer_norm( builder, &residual, @@ -243,7 +203,8 @@ fn layer( let up = linear(builder, &normalized, &args.take(), &args.take()); let up = gelu_tanh(builder, &up); let down = linear(builder, &up, &args.take(), &args.take()); - builder.add(&residual, &down) + let residual = builder.add(&residual, &down); + round_bf16(builder, &residual) } struct YoutuVlGraph { @@ -281,6 +242,7 @@ fn build_youtu_vl(config: &YoutuVlVisionConfig, patch_bucket: usize) -> YoutuVlG Ty::f32(vec![patch_bucket, patch_bucket]), "full_attention.bias", ); + let patches = round_bf16(&mut builder, &patches); let mut hidden = linear(&mut builder, &patches, &args.take(), &args.take()); let patch_projection = hidden.clone(); let mut layer_outputs = Vec::with_capacity(config.depth); @@ -305,7 +267,7 @@ fn build_youtu_vl(config: &YoutuVlVisionConfig, patch_bucket: usize) -> YoutuVlG let merged_width = config.hidden * config.spatial_merge_size * config.spatial_merge_size; let merged = builder.reshape(&hidden, vec![patch_bucket / 4, merged_width]); let merged = linear(&mut builder, &merged, &args.take(), &args.take()); - let merged = gelu_tanh(&mut builder, &merged); + let merged = gelu_exact(&mut builder, &merged); let projected = linear(&mut builder, &merged, &args.take(), &args.take()); assert_eq!(args.cursor, specs.len(), "Youtu-VL weight schema drifted"); YoutuVlGraph { @@ -384,6 +346,14 @@ mod tests { assert!(ir.contains("full_attention.bias")); assert!(ir.contains("merger.mlp.2.weight")); assert!(ir.contains("tensor<4x12xf32>")); + assert!( + ir.contains("stablehlo.convert"), + "the BF16 checkpoint carrier tape must remain explicit" + ); + assert!( + ir.contains("chlo.erf"), + "the merger uses PyTorch's exact GELU, unlike the vision MLP" + ); } #[test]