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..3d85b9a6a --- /dev/null +++ b/examples/xla_youtu_vl_reference_check.rs @@ -0,0 +1,700 @@ +// 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, + configure_diagnostic_local_task_threads, +}; +use serde_json::{Map, Value, json}; +use sha2::{Digest, Sha256}; + +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 = + "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 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; +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 has_flag(flag: &str) -> bool { + std::env::args().any(|value| value == flag) +} + +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 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); + 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); + 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]; + 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; + } + } + } + } + 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(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, [(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"); + 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, 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"); + 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(), + MERGED_TOKENS + ); + + 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, PATCH_GRID * PATCH_SIZE, PATCH_GRID * PATCH_SIZE], + ); + 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", + &[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"), + &[MERGED_TOKENS], + ); + write_i32( + &mut arrays, + &out, + "reverse_group_index", + &usize_i32(&vision.reverse_group_index, "reverse group index"), + &[MERGED_TOKENS], + ); + 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..22a3d7141 --- /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[:] = [[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"}: + 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..4cd08ba9a --- /dev/null +++ b/spike/openxla/youtu_vl_reference_oracle.py @@ -0,0 +1,897 @@ +#!/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-v2" +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 +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 + +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: + 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(6, 0, 2, 4, 1, 3, 5) + .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=MAX_IMAGE_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() != [ + [PATCH_GRID, PATCH_GRID] + ]: + raise ContractError( + "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 + 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, PATCH_GRID, PATCH_GRID), + "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") + # 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(), + "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, + RESIZED_IMAGE_SIDE, + RESIZED_IMAGE_SIDE, + ], + "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() != [[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"]) + 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..56209a70d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -77,12 +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; 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, 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/config.rs b/src/lib/mlxcel-xla/src/emitter/config.rs index e5c450eaa..71e301152 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,38 @@ impl Config { one_plus: false, }); } + 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")?, + 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 +1434,7 @@ impl Config { sliding_pattern, use_rope_layers, mrope, + mla, deepstack, moe, weight_scheme, @@ -1426,6 +1490,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 +1523,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..10edf44d6 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,16 @@ 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; +#[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, +}; #[cfg(test)] mod tests { @@ -1401,6 +1413,85 @@ mod tests { assert!(e.contains("MiniCPM3") && e.contains("MLA"), "got: {e}"); } + #[test] + fn youtu_vl_selects_dense_mla_and_interleaved_rope() { + 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}"#; + 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); + 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']")); + + 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] + #[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..cd66af416 --- /dev/null +++ b/src/lib/mlxcel-xla/src/emitter/youtu_vl.rs @@ -0,0 +1,442 @@ +// 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::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}; + +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 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); + 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 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 value = rms_norm_2d(builder, value, weight, epsilon); + round_bf16(builder, &value) +} + +fn gelu_tanh(builder: &mut Builder, value: &Val) -> Val { + 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 { + 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); + let value = builder.add(&direct, &rotated); + round_bf16(builder, &value) +} + +fn softmax(builder: &mut Builder, value: &Val) -> Val { + let probabilities = stable_softmax(builder, value, 2); + round_bf16(builder, &probabilities) +} + +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 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, + &v, + &[0], + &[0], + &[2], + &[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()) +} + +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 residual = round_bf16(builder, &residual); + 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()); + let residual = builder.add(&residual, &down); + round_bf16(builder, &residual) +} + +struct YoutuVlGraph { + args: Args, + builder: Builder, + patch_projection: Val, + layer_outputs: Vec, + post_layernorm: Val, + 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" + ); + 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 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); + 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); + layer_outputs.push(hidden.clone()); + } + hidden = layer_norm( + &mut builder, + &hidden, + &args.take(), + &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]); + let merged = linear(&mut builder, &merged, &args.take(), &args.take()); + 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 { + args, + builder, + patch_projection, + layer_outputs, + post_layernorm, + 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 = graph.args.declarations.join(", "), + result_type = graph.projected.ty.render(), + body = graph.builder.body(), + result = graph.projected.name, + ) +} + +#[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); + 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()) + .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::*; + + 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( + 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>")); + 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] + 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<64x16xf32>, tensor<16x12xf32>" + )); + assert!(ir.contains("Youtu-VL diagnostic checkpoints")); + assert_eq!( + config.diagnostic_contract_identity(), + "youtu-vl-vision-diagnostics-v3: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-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"); + 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 new file mode 100644 index 000000000..460c91366 --- /dev/null +++ b/src/lib/mlxcel-xla/src/emitter/youtu_vl_plan.rs @@ -0,0 +1,675 @@ +// 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]; +#[cfg(any(test, feature = "diagnostics"))] +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; +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, +} + +/// 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, + 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, + ) + } + + #[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!( + "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()); + } + // 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 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(&patch_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_processor_merge_groups() { + 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]; + assert_eq!( + group_rows, + &[group * 4, group * 4 + 1, group * 4 + 2, group * 4 + 3,] + ); + } + } +} 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, diff --git a/src/lib/mlxcel-xla/src/lib.rs b/src/lib/mlxcel-xla/src/lib.rs index 2b1328267..4b8aa0fc8 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::{ @@ -145,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, @@ -156,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, @@ -184,6 +189,15 @@ 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, + YoutuVlVisionDiagnosticStageOutput, +}; +#[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..0602a82be --- /dev/null +++ b/src/lib/mlxcel-xla/src/youtu_vl_runtime.rs @@ -0,0 +1,601 @@ +// 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}; +#[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, +}; +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"; +#[cfg(feature = "diagnostics")] +const DIAGNOSTIC_ENTRY_NAME: &str = "youtu_vl_vision_diagnostics.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, +} + +#[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 { + /// 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], + /// 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 { + 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_legacy_unqualified( + 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) +} + +#[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_legacy_unqualified( + DIAGNOSTIC_ENTRY_NAME, + format!( + "{};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-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) + })?; + 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, + }, + }) + } +} + +#[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 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() + .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) + .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 values = buffers + .into_iter() + .enumerate() + .map(|(index, bytes)| { + checked_f32_output(&format!("Youtu-VL IREE diagnostic output {index}"), bytes) + }) + .collect::, _>>()?; + let actual_tokens = plan.reverse_group_index.len(); + let actual_values = actual_tokens * self.config.text_hidden; + 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; + } else { + values.truncate(actual_patches * shape[1]); + shape[0] = actual_patches; + } + 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], + }); + 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, + }) + } +} 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/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/multimodal/host_preprocessor.rs b/src/multimodal/host_preprocessor.rs index bfad9cb9d..67883921d 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,127 @@ 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 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 token_ids.contains(&video_token_id) { + // Placeholder compatibility only: the request boundary remains + // image-only until a video preprocessing oracle is qualified. + 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, + 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 +609,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..b16fe833d 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,98 @@ 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, IMAGE, IMAGE, 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] +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"] +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.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 diff --git a/src/vision/encoders/youtu_vl_window.rs b/src/vision/encoders/youtu_vl_window.rs index bb2e00d14..1f224f8a9 100644 --- a/src/vision/encoders/youtu_vl_window.rs +++ b/src/vision/encoders/youtu_vl_window.rs @@ -152,8 +152,36 @@ 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() { + // 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_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", + ] + ); + } +} diff --git a/src/vision/processors/youtu_vl.rs b/src/vision/processors/youtu_vl.rs index e4d8fd8dd..7c9b620c1 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; @@ -280,34 +299,61 @@ 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 { + 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, 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] + ); +}