From 75a80cbdf95f281eb9e29054091f7bdd7c3d2898 Mon Sep 17 00:00:00 2001 From: levius <2114377220@qq.com> Date: Thu, 1 Oct 2026 09:59:24 +0800 Subject: [PATCH 1/5] cua_s1: accept adapted image rows and three-axis language positions --- recipe/cua_s1/native.md | 112 +++++- .../native/examples/multimodal_boundary.rs | 184 ++++++++++ src/models/cua_s1/native/src/inputs.rs | 213 ++++++++++++ src/models/cua_s1/native/src/lib.rs | 1 + src/models/cua_s1/native/src/model.rs | 328 ++++++++++++++---- src/models/cua_s1/native/tests/multimodal.rs | 90 +++++ 6 files changed, 867 insertions(+), 61 deletions(-) create mode 100644 src/models/cua_s1/native/examples/multimodal_boundary.rs create mode 100644 src/models/cua_s1/native/src/inputs.rs create mode 100644 src/models/cua_s1/native/tests/multimodal.rs diff --git a/recipe/cua_s1/native.md b/recipe/cua_s1/native.md index 018feff..35dccb4 100644 --- a/recipe/cua_s1/native.md +++ b/recipe/cua_s1/native.md @@ -24,9 +24,11 @@ CUA_S1_MODEL=weights/cua-s1-4b-0.2-text-merged target/release/omni-cua-s1-native ``` For the local CUDA Graph experiment, also set `CUA_S1_GRAPH=1`. The first use of -each exact prompt length warms the GEMM plans and captures the forward pass; -later requests replay it with freshly uploaded token ids. At most eight lengths -are cached. Growing the scratch allocation clears the captures before freeing +each exact prompt length warms the GEMM plans and captures the forward pass. +The first call returns the eager result; later requests replay it after fresh +token embedding. The graph contains the language layers, which update residuals +in place, so a cache miss must not replay those layers over its eager result. +At most eight lengths are cached. Growing the scratch allocation clears the captures before freeing their buffers. Capture adds first-use latency; leave the variable unset to use the eager control. Rebuild both the worker and CUDA library together (ABI 3). If capture fails, the worker returns the completed eager result and disables @@ -41,3 +43,107 @@ cargo test -p omni-cua-s1-native CUA_S1_CUDA_LIB=$PWD/target/release/libqwen3_5_cuda.so \ cargo test --release -p omni-cua-s1-native --test kernels -- --ignored ``` + +## Multimodal language boundary + +`Model::forward_multimodal` consumes token IDs, adapted BF16 image features +`[image_tokens, hidden_size]`, their sorted placeholder indices, and the T/H/W +slices of int64 `position_ids [3, 1, sequence]`. It returns the last position's +final-normalized hidden state, like `Model::forward`. + +Inputs are one unpadded prompt. Every image placeholder must have exactly one +feature row; all other rows come from the token embedding table. The caller +calculates positions and runs the image processor, vision tower and vision LoRA. +Positions must be nonnegative and below `max_position_embeddings`. The language +path uses Qwen3.5's interleaved MRoPE sections, not three contiguous rotary blocks. +Text calls and their captured graphs use immutable text-position tables; +multimodal calls use separate device tables, so returning to text requires no +host table rebuild or restoration copy. Multimodal calls execute eagerly even +when `CUA_S1_GRAPH=1`. The extra tables use +`4 * scratch_capacity * rotary_half` bytes (2 MiB at 16,384 rows). + +This is a Rust model API for integrating a vision producer. The HTTP worker +above continues to serve the text adapter. A native vision encoder, image HTTP +requests, padding, video and batching are not implemented by this API. + +### Prepare a matching language checkpoint + +The language weights must contain the **multimodal** adapter, not the `text` +adapter. In the pinned reference environment, with upstream-verified weights, +export just the merged language model (about 7.5 GB) to a new directory: + +```sh +PYTHONPATH=src HF_HUB_OFFLINE=1 .venv/bin/python - <<'PY' +import json +from pathlib import Path +from models.cua_s1.multimodal.model import ( + ADAPTER_REVISION, BASE_REVISION, MultimodalEngine, +) + +out = Path("weights/cua-s1-4b-0.2-multimodal-language-merged") +if out.exists(): + raise FileExistsError(out) +engine = MultimodalEngine( + "weights/Qwen3.5-4B", "weights/cua-s1-4b-0.2/multimodal" +) +merged = engine.model.merge_and_unload() +merged.model.language_model.save_pretrained(out, max_shard_size="5GB") +# Preserve the root image_token_id and text_config for the native input contract. +merged.config.to_json_file(out / "config.json") +(out / "cua_s1_language_export.json").write_text(json.dumps({ + "format": "cua-s1-multimodal-language-merged/1", + "base_revision": BASE_REVISION, + "adapter_revision": ADAPTER_REVISION, +})) +PY +``` + +No `cua_s1_export.json` text-worker marker is created. The low-level `Model` API +does not verify checkpoint provenance; retain the export metadata and use the +matching adapter for the supplied features. Standalone language safetensors +names and the existing full-model prefixes are supported. + +### Replay a reference boundary + +This optional example consumes the `cua-s1-multimodal-reference-v1` format +from [#53](https://github.com/ThinkFlowLab/system1-omni/pull/53), which is still +open. The exporter and checksum verifier are not yet available on `main`. +Use a separate checkout of exporter revision +`1b64fa2ceb0a82b6a66a69ecdc9bc5cc1b1a0b66` to generate and verify the bundle; +the Rust model API itself does not depend on that PR being merged. +Its eight questions include different image grids and question lengths, JPEG, +structured/non-ASCII text, and 1/3/26 candidates. Then run: + +```sh +CUA_S1_CUDA_LIB=$PWD/target/release/libqwen3_5_cuda.so \ + cargo run --release --locked -p omni-cua-s1-native \ + --example multimodal_boundary -- \ + weights/cua-s1-4b-0.2-multimodal-language-merged \ + /path/to/verified-reference-bundle /tmp/native-language.json + +CUA_S1_MODEL=$PWD/weights/cua-s1-4b-0.2-multimodal-language-merged \ +CUA_S1_CUDA_LIB=$PWD/target/release/libqwen3_5_cuda.so \ + cargo test --release --locked -p omni-cua-s1-native \ + --test multimodal -- --ignored + +# Compare graph misses, hits, eviction and scratch growth with eager hidden states. +CUA_S1_MODEL=$PWD/weights/cua-s1-4b-0.2-multimodal-language-merged \ +CUA_S1_CUDA_LIB=$PWD/target/release/libqwen3_5_cuda.so \ + cargo test --release --locked -p omni-cua-s1-native \ + --lib graph_tests::graph_misses_hits_eviction_growth_and_multimodal_match_eager -- --ignored +``` + +The example checks repeated native hidden-state equality and writes last hidden +states, candidate logits and probabilities. It uses the text engine's FP32 +letter-row readout with FP64 accumulation. Output must be a new file. Verify +bundle integrity before invoking the example; it checks tensor shapes and input +contracts but is not the bundle checksum verifier. + +For accuracy validation, compare against an unmerged FP32 **language** control +with TF32 disabled, feeding the same fixed exported embeddings and positions. +Use the [declared native tolerance](../../src/models/cua_s1/README.md#validation): +maximum probability error over the set must be at most twice the BF16 reference +error plus 0.01, and the top option must match for FP32 margins at least 0.05. +The FP32 control starts after the BF16-exported vision boundary; it does not +validate a full FP32 vision pipeline. Native kernel and LoRA-merge rounding can +change hidden states and logits; bitwise equality to Transformers is not claimed. diff --git a/src/models/cua_s1/native/examples/multimodal_boundary.rs b/src/models/cua_s1/native/examples/multimodal_boundary.rs new file mode 100644 index 0000000..7c93fb4 --- /dev/null +++ b/src/models/cua_s1/native/examples/multimodal_boundary.rs @@ -0,0 +1,184 @@ +//! Replay #53's exported language boundary through the native model. +//! The v1 bundle producer/verifier are on open PR #53 at exporter revision +//! 1b64fa2ceb0a82b6a66a69ecdc9bc5cc1b1a0b66; see recipe/cua_s1/native.md. +use std::path::{Path, PathBuf}; + +use anyhow::{Context, Result, ensure}; +use half::bf16; +use omni_cua_s1_native::{inputs::MultimodalInput, model::Model}; +use safetensors::{Dtype, SafeTensors}; +use serde_json::{Value, json}; + +fn integers(st: &SafeTensors<'_>, name: &str, shape: &[usize]) -> Result> { + let v = st.tensor(name)?; + ensure!( + v.dtype() == Dtype::I64 && v.shape() == shape, + "{name}: expected I64 {shape:?}" + ); + Ok(v.data() + .as_chunks::<8>() + .0 + .iter() + .map(|b| i64::from_le_bytes(*b)) + .collect()) +} + +fn features(st: &SafeTensors<'_>, shape: &[usize]) -> Result> { + let v = st.tensor("image_features")?; + ensure!( + v.dtype() == Dtype::BF16 && v.shape() == shape, + "image_features: expected BF16 {shape:?}" + ); + Ok(v.data() + .as_chunks::<2>() + .0 + .iter() + .map(|b| bf16::from_le_bytes(*b)) + .collect()) +} + +fn embedding_file(dir: &Path) -> Result<(PathBuf, String)> { + let names = [ + "model.language_model.embed_tokens.weight", + "model.embed_tokens.weight", + "embed_tokens.weight", + ]; + if dir.join("model.safetensors.index.json").exists() { + let index: Value = + serde_json::from_slice(&std::fs::read(dir.join("model.safetensors.index.json"))?)?; + for name in names { + if let Some(file) = index["weight_map"][name].as_str() { + return Ok((dir.join(file), name.into())); + } + } + } else { + let path = dir.join("model.safetensors"); + let bytes = std::fs::read(&path)?; + let st = SafeTensors::deserialize(&bytes)?; + for name in names { + if st.tensor(name).is_ok() { + return Ok((path, name.into())); + } + } + } + anyhow::bail!("missing embedding weight") +} + +fn main() -> Result<()> { + let args: Vec<_> = std::env::args_os().skip(1).collect(); + ensure!( + args.len() == 3, + "usage: multimodal_boundary MODEL_DIR REFERENCE_BUNDLE OUTPUT_JSON" + ); + let (dir, bundle, out) = ( + Path::new(&args[0]), + Path::new(&args[1]), + Path::new(&args[2]), + ); + ensure!(!out.exists(), "output already exists"); + let library = PathBuf::from(std::env::var_os("CUA_S1_CUDA_LIB").context("CUA_S1_CUDA_LIB")?); + let manifest: Value = serde_json::from_slice(&std::fs::read(bundle.join("manifest.json"))?)?; + ensure!( + manifest["schema"] == "cua-s1-multimodal-reference-v1", + "unsupported reference schema" + ); + let mut model = Model::load(dir, &library)?; + let (file, name) = embedding_file(dir)?; + let file = std::fs::File::open(file)?; + // SAFETY: the checkpoint is immutable while the example runs. + let map = unsafe { memmap2::Mmap::map(&file)? }; + let weights = SafeTensors::deserialize(&map)?; + let embed = weights.tensor(&name)?; + ensure!( + embed.dtype() == Dtype::BF16 + && embed.shape().len() == 2 + && embed.shape()[1] == model.cfg.hidden, + "embedding shape/dtype" + ); + let mut rows = Vec::new(); + for entry in manifest["questions"].as_array().context("questions")? { + let relative = Path::new(entry["tensors_file"].as_str().context("tensors_file")?); + ensure!( + relative + .components() + .all(|c| matches!(c, std::path::Component::Normal(_))), + "unsafe tensor path" + ); + let bytes = std::fs::read(bundle.join(relative))?; + let st = SafeTensors::deserialize(&bytes)?; + let ids = st.tensor("input_ids")?; + ensure!( + ids.shape().len() == 2 && ids.shape()[0] == 1, + "expected batch one" + ); + let t = ids.shape()[1]; + let ids: Vec = integers(&st, "input_ids", &[1, t])? + .into_iter() + .map(u32::try_from) + .collect::>()?; + let indices = st.tensor("image_token_indices")?; + ensure!( + indices.shape().len() == 1, + "image indices must be one-dimensional" + ); + let count = indices.shape()[0]; + let indices: Vec = integers(&st, "image_token_indices", &[count])? + .into_iter() + .map(usize::try_from) + .collect::>()?; + let features = features(&st, &[count, model.cfg.hidden])?; + let positions = integers(&st, "position_ids", &[3, 1, t])?; + let input = MultimodalInput { + token_ids: &ids, + image_token_indices: &indices, + image_embeddings: &features, + position_ids: [&positions[..t], &positions[t..2 * t], &positions[2 * t..]], + }; + let last = model.forward_multimodal(&input)?; + ensure!( + last.iter().all(|x| x.is_finite()), + "non-finite hidden state" + ); + // Repeating the same boundary in the same model checks buffer reuse. + ensure!( + last == model.forward_multimodal(&input)?, + "repeat changed native hidden state" + ); + let n = entry["option_keys"] + .as_array() + .context("option_keys")? + .len(); + ensure!((1..=26).contains(&n), "candidate count"); + let candidates = integers(&st, "candidate_token_ids", &[n])?; + let mut logits = Vec::new(); + for id in candidates { + let id = usize::try_from(id)?; + ensure!(id < embed.shape()[0], "candidate outside vocabulary"); + let row = &embed.data()[id * last.len() * 2..(id + 1) * last.len() * 2]; + let dot: f64 = row + .as_chunks::<2>() + .0 + .iter() + .zip(&last) + .map(|(b, &h)| bf16::from_le_bytes(*b).to_f64() * h as f64) + .sum(); + logits.push(dot as f32); + } + let max = logits.iter().copied().fold(f32::NEG_INFINITY, f32::max) as f64; + let exps: Vec = logits.iter().map(|&l| (l as f64 - max).exp()).collect(); + let total: f64 = exps.iter().sum(); + let probabilities: Vec = exps.iter().map(|e| (e / total) as f32).collect(); + ensure!( + probabilities.iter().all(|x| x.is_finite()), + "non-finite readout" + ); + rows.push(json!({"case": entry["case"], "question": entry["question"], "sequence": t, "image_tokens": count, "last_hidden_state": last, "candidate_logits": logits, "probabilities": probabilities, "repeat_equal": true})); + } + std::fs::write( + out, + serde_json::to_vec_pretty( + &json!({"schema": "cua-s1-native-language-boundary-v1", "questions": rows}), + )?, + )?; + Ok(()) +} diff --git a/src/models/cua_s1/native/src/inputs.rs b/src/models/cua_s1/native/src/inputs.rs new file mode 100644 index 0000000..1f0c980 --- /dev/null +++ b/src/models/cua_s1/native/src/inputs.rs @@ -0,0 +1,213 @@ +//! Batch-one, unpadded inputs at the adapted-vision / language-model boundary. + +use anyhow::{Result, ensure}; +use half::bf16; + +/// Image rows are already adapted to the language hidden size, in placeholder +/// order. Positions are the T/H/W slices of an int64 `[3, 1, sequence]` tensor. +/// The caller owns preprocessing, vision execution and the position calculation. +pub struct MultimodalInput<'a> { + pub token_ids: &'a [u32], + pub image_token_indices: &'a [usize], + pub image_embeddings: &'a [bf16], + pub position_ids: [&'a [i64]; 3], +} + +impl MultimodalInput<'_> { + /// Check the entire boundary before allocating buffers or launching CUDA. + pub fn validate( + &self, + hidden: usize, + vocab: usize, + image_token: u32, + max_position: usize, + ) -> Result<()> { + let t = self.token_ids.len(); + ensure!(t > 0 && t <= max_position, "empty or oversized prompt"); + ensure!( + self.token_ids.iter().all(|&id| (id as usize) < vocab), + "token id outside the vocabulary" + ); + let expected: Vec = self + .token_ids + .iter() + .enumerate() + .filter_map(|(i, &id)| (id == image_token).then_some(i)) + .collect(); + ensure!( + self.image_token_indices == expected, + "image indices must exactly match the ordered placeholders" + ); + ensure!( + Some(self.image_embeddings.len()) == expected.len().checked_mul(hidden), + "image embedding shape mismatch" + ); + ensure!( + self.image_embeddings.iter().all(|x| x.is_finite()), + "non-finite image embedding" + ); + ensure!( + self.position_ids.iter().all(|axis| axis.len() == t), + "position_ids must have shape [3, 1, sequence]" + ); + ensure!( + self.position_ids + .iter() + .flat_map(|axis| axis.iter()) + .all(|&p| p >= 0 && (p as u64) < max_position as u64), + "position outside the configured range" + ); + Ok(()) + } +} + +/// Qwen3.5's interleaved recomposition: overwrite H at 1::3 and W at 2::3 up +/// to section[axis] * 3, retaining T elsewhere. The second rotary half repeats +/// these frequencies, which the existing attention-prep kernel handles. +/// Float32 inverse frequencies/products and host float64 trig preserve the +/// original native text table rounding when all three axes are equal. +pub(crate) fn rotary_tables( + positions: [&[i64]; 3], + half: usize, + theta: f64, + sections: [usize; 3], +) -> (Vec, Vec) { + let inv: Vec = (0..half) + .map(|i| 1f32 / (theta as f32).powf((2 * i) as f32 / (2 * half) as f32)) + .collect(); + let mut cos = Vec::with_capacity(positions[0].len() * half * 2); + let mut sin = Vec::with_capacity(cos.capacity()); + for (t, _) in positions[0].iter().enumerate() { + for (i, &f) in inv.iter().enumerate() { + let axis = if i % 3 == 1 && i < sections[1] * 3 { + 1 + } else if i % 3 == 2 && i < sections[2] * 3 { + 2 + } else { + 0 + }; + let angle = (f * positions[axis][t] as f32) as f64; + cos.extend(bf16::from_f32(angle.cos() as f32).to_le_bytes()); + sin.extend(bf16::from_f32(angle.sin() as f32).to_le_bytes()); + } + } + (cos, sin) +} + +#[cfg(test)] +mod tests { + use super::*; + use half::bf16; + + #[test] + fn valid_input_and_three_distinct_axes() { + let input = MultimodalInput { + token_ids: &[1, 99, 99, 2], + image_token_indices: &[1, 2], + image_embeddings: &[bf16::ONE; 8], + position_ids: [&[0, 1, 1, 3], &[0, 1, 2, 3], &[0, 2, 1, 3]], + }; + input.validate(4, 100, 99, 100).unwrap(); + } + + #[test] + fn rejects_bad_placeholder_inventory_and_feature_rows() { + for indices in [vec![2, 1], vec![1, 1], vec![1], vec![0, 1], vec![1, 4]] { + let input = MultimodalInput { + token_ids: &[1, 99, 99, 2], + image_token_indices: &indices, + image_embeddings: &[bf16::ONE; 8], + position_ids: [&[0, 1, 1, 3]; 3], + }; + assert!(input.validate(4, 100, 99, 100).is_err(), "{indices:?}"); + } + for features in [vec![bf16::ONE; 7], vec![bf16::ONE; 9], vec![bf16::NAN; 8]] { + let input = MultimodalInput { + token_ids: &[1, 99, 99, 2], + image_token_indices: &[1, 2], + image_embeddings: &features, + position_ids: [&[0, 1, 1, 3]; 3], + }; + assert!(input.validate(4, 100, 99, 100).is_err()); + } + } + + #[test] + fn rejects_bad_tokens_positions_and_empty_sequence() { + let features = [bf16::ONE; 4]; + for (ids, positions) in [ + (vec![100, 99], vec![0, 1]), + (vec![1, 99], vec![0]), + (vec![1, 99], vec![0, -1]), + (vec![1, 99], vec![0, 100]), + (vec![], vec![]), + ] { + let input = MultimodalInput { + token_ids: &ids, + image_token_indices: &[1], + image_embeddings: &features, + position_ids: [&positions; 3], + }; + assert!(input.validate(4, 100, 99, 100).is_err()); + } + } + + #[test] + fn image_free_explicit_positions_are_valid() { + MultimodalInput { + token_ids: &[1, 2], + image_token_indices: &[], + image_embeddings: &[], + position_ids: [&[7, 8]; 3], + } + .validate(4, 100, 99, 100) + .unwrap(); + } + + #[test] + fn rotary_interleaves_height_width_and_leaves_temporal_tail() { + // theta=1 makes every inverse frequency 1. Axis values differ so a plain + // text table or a contiguous-section implementation fails this check. + for (sections, tail) in [([12, 10, 10], [0, 0]), ([11, 11, 10], [0, 1])] { + let (cos, sin) = rotary_tables([&[0], &[1], &[2]], 32, 1.0, sections); + // Ten T/H/W triples, then T/T for the synthetic layout or T/H for + // the real checkpoint. In particular, frequency 31 must use H. + let axes = [ + 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, + 1, 2, tail[0], tail[1], + ]; + for (i, axis) in axes.into_iter().enumerate() { + let angle = axis as f32; + assert_eq!( + &cos[i * 2..i * 2 + 2], + &bf16::from_f32(angle.cos()).to_le_bytes() + ); + assert_eq!( + &sin[i * 2..i * 2 + 2], + &bf16::from_f32(angle.sin()).to_le_bytes() + ); + } + } + } + + #[test] + fn equal_axes_reproduce_the_existing_text_tables() { + let pos: Vec = (0..257).collect(); + let (cos, sin) = rotary_tables([&pos; 3], 32, 10_000_000.0, [11, 11, 10]); + for (t, &p) in pos.iter().enumerate() { + for i in 0..32 { + let inv = 1f32 / 10_000_000f32.powf((2 * i) as f32 / 64.0); + let angle = (inv * p as f32) as f64; + let offset = (t * 32 + i) * 2; + assert_eq!( + &cos[offset..offset + 2], + &bf16::from_f32(angle.cos() as f32).to_le_bytes() + ); + assert_eq!( + &sin[offset..offset + 2], + &bf16::from_f32(angle.sin() as f32).to_le_bytes() + ); + } + } + } +} diff --git a/src/models/cua_s1/native/src/lib.rs b/src/models/cua_s1/native/src/lib.rs index 0d3aa5f..3d09d9c 100644 --- a/src/models/cua_s1/native/src/lib.rs +++ b/src/models/cua_s1/native/src/lib.rs @@ -5,5 +5,6 @@ pub mod contract; pub mod cuda; pub mod engine; +pub mod inputs; pub mod json; pub mod model; diff --git a/src/models/cua_s1/native/src/model.rs b/src/models/cua_s1/native/src/model.rs index ae14d55..d6d6597 100644 --- a/src/models/cua_s1/native/src/model.rs +++ b/src/models/cua_s1/native/src/model.rs @@ -6,8 +6,8 @@ //! The order of operations follows `modeling_qwen3_5.py`, and so do the points where //! it rounds to bfloat16, except inside attention and the Gated DeltaNet prefill (see //! their kernels). Text prompts use one position per -//! token, so the multimodal rotary sections all get the same position and the -//! rotary embedding is the plain one. +//! token. The explicit multimodal boundary inserts adapted image rows and supplies +//! the interleaved temporal/height/width rotary positions to the same layer loop. use std::collections::{HashMap, VecDeque}; use std::ffi::c_void; @@ -17,6 +17,7 @@ use anyhow::{Context, Result, bail, ensure}; use serde_json::Value as Json; use crate::cuda::{self, DeviceBuffer, Stream, check}; +use crate::inputs::{MultimodalInput, rotary_tables}; const ALIGN: usize = 256; const BF16: usize = 2; @@ -35,6 +36,9 @@ pub struct Config { /// Half the number of rotary dims (rotate_half pairs dim i with dim i + half). pub rotary_half: usize, pub rope_theta: f64, + pub mrope_section: [usize; 3], + pub max_positions: usize, + pub image_token_id: Option, pub lin_k_heads: usize, pub lin_v_heads: usize, pub lin_k_dim: usize, @@ -69,6 +73,31 @@ impl Config { other => bail!("unknown layer type {other:?}"), }) .collect::>>()?; + let sections = rope + .get("mrope_section") + .cloned() + .unwrap_or(serde_json::json!([11, 11, 10])); + let sections = sections + .as_array() + .context("mrope_section must be an array")? + .iter() + .map(|v| { + v.as_u64() + .and_then(|n| usize::try_from(n).ok()) + .context("mrope_section must contain integers") + }) + .collect::>>()?; + let mrope_section: [usize; 3] = sections + .try_into() + .map_err(|_| anyhow::anyhow!("mrope_section must have three entries"))?; + let image_token_id = root + .get("image_token_id") + .map(|v| { + v.as_u64() + .and_then(|id| u32::try_from(id).ok()) + .context("image_token_id must be a u32") + }) + .transpose()?; let cfg = Config { hidden: int("hidden_size")?, intermediate: int("intermediate_size")?, @@ -81,6 +110,9 @@ impl Config { .as_f64() .or(c["rope_theta"].as_f64()) .context("rope_theta")?, + mrope_section, + max_positions: int("max_position_embeddings")?, + image_token_id, lin_k_heads: int("linear_num_key_heads")?, lin_v_heads: int("linear_num_value_heads")?, lin_k_dim: int("linear_key_head_dim")?, @@ -115,6 +147,25 @@ impl Config { cfg.lin_v_dim ); ensure!(cfg.rotary_half == 32, "{} rotary dims", 2 * cfg.rotary_half); + ensure!( + cfg.rope_theta.is_finite() && cfg.rope_theta > 0.0, + "invalid rope_theta" + ); + ensure!( + rope["mrope_interleaved"].as_bool() != Some(false), + "non-interleaved mrope is unsupported" + ); + ensure!( + cfg.mrope_section + .iter() + .try_fold(0usize, |sum, &x| sum.checked_add(x)) + == Some(cfg.rotary_half), + "mrope sections must sum to rotary_half" + ); + ensure!( + cfg.max_positions > 0 && cfg.max_positions <= (1 << 24), + "unsupported max_position_embeddings" + ); ensure!( cfg.kv_heads > 0 && cfg.heads.is_multiple_of(cfg.kv_heads), "attention heads" @@ -218,7 +269,7 @@ impl Weights { .enumerate() .flat_map(|(i, st)| st.names().into_iter().map(move |n| (i, n.to_string()))) .collect(); - let prefix = ["model.language_model.", "model."] + let prefix = ["model.language_model.", "model.", ""] .into_iter() .find(|p| { names @@ -385,6 +436,8 @@ struct Scratch { act: usize, cos: usize, sin: usize, + custom_cos: usize, + custom_sin: usize, } impl Scratch { @@ -423,6 +476,8 @@ impl Scratch { take(cap * cfg.intermediate * BF16), take(cap * cfg.rotary_half * BF16), take(cap * cfg.rotary_half * BF16), + take(cap * cfg.rotary_half * BF16), + take(cap * cfg.rotary_half * BF16), ]; let buf = DeviceBuffer::new(next)?; let [ @@ -448,24 +503,16 @@ impl Scratch { act, cos, sin, + custom_cos, + custom_sin, ] = offsets; - // Rotary tables close to how Qwen3_5TextRotaryEmbedding builds them: inv_freq and - // freqs = inv_freq * position in float32, cos and sin rounded to bfloat16. Here - // cos and sin are taken in float64 on the host rather than in float32 on the - // GPU, so a few of the rounded values can differ by one bfloat16 step. - let half = cfg.rotary_half; - let inv: Vec = (0..half) - .map(|i| 1.0f32 / (cfg.rope_theta as f32).powf((2 * i) as f32 / (2 * half) as f32)) - .collect(); - let mut cos_t = Vec::with_capacity(cap * half * BF16); - let mut sin_t = Vec::with_capacity(cap * half * BF16); - for pos in 0..cap { - for &f in &inv { - let freq = (f * pos as f32) as f64; - cos_t.extend(half::bf16::from_f32(freq.cos() as f32).to_le_bytes()); - sin_t.extend(half::bf16::from_f32(freq.sin() as f32).to_le_bytes()); - } - } + let positions: Vec = (0..cap as i64).collect(); + let (cos_t, sin_t) = rotary_tables( + [&positions; 3], + cfg.rotary_half, + cfg.rope_theta, + cfg.mrope_section, + ); // SAFETY: both tables were laid out for cap * rotary_half bfloat16 values. unsafe { cuda::upload(buf.at(cos), &cos_t, stream)?; @@ -496,6 +543,8 @@ impl Scratch { act, cos, sin, + custom_cos, + custom_sin, }) } @@ -632,15 +681,7 @@ impl Model { ) } - /// The final-norm hidden state at the last position, as float32. - pub fn forward(&mut self, ids: &[u32]) -> Result> { - let t = ids.len(); - ensure!(t > 0, "empty prompt"); - let (vocab, h) = (self.embed.shape[0], self.cfg.hidden); - ensure!( - ids.iter().all(|&i| (i as usize) < vocab), - "token id outside the vocabulary" - ); + fn prepare_scratch(&mut self, t: usize) -> Result<()> { cuda::set_device(0)?; if self.scratch.as_ref().is_none_or(|s| t > s.cap) { self.graphs.clear(); @@ -651,16 +692,33 @@ impl Model { self.stream, )?); } + Ok(()) + } + + /// The final-norm hidden state at the last position, as float32. + pub fn forward(&mut self, ids: &[u32]) -> Result> { + let t = ids.len(); + ensure!( + t > 0 && t <= self.cfg.max_positions, + "empty or oversized prompt" + ); + ensure!( + ids.iter().all(|&i| (i as usize) < self.embed.shape[0]), + "token id outside the vocabulary" + ); + self.prepare_scratch(t)?; let s = self.scratch.as_ref().unwrap(); - let ids32: Vec = ids.iter().flat_map(|&i| (i as i32).to_le_bytes()).collect(); - // SAFETY: the ids buffer holds at least t int32 values. - unsafe { cuda::upload(s.at(s.ids), &ids32, self.stream)? }; + self.embed_tokens(s, ids)?; if self.graph_enabled { - if !self.graphs.iter().any(|(length, _)| *length == t) { - // Initialize every cuBLASLt plan before stream capture. - self.run(s, t)?; + if let Some((_, graph)) = self.graphs.iter().find(|(length, _)| *length == t) { + graph.launch(self.stream)?; + } else { + // Warm GEMM plans and keep this eager result for the cache miss. + // run() advances s.res in place and no longer embeds tokens, so + // launching the new graph here would advance the residual twice. + self.run(s, t, false)?; cuda::synchronize(self.stream)?; - match cuda::Graph::capture(self.stream, || self.run(s, t)) { + match cuda::Graph::capture(self.stream, || self.run(s, t, false)) { Ok(graph) => { if self.graphs.len() == 8 { self.graphs.pop_front(); @@ -669,27 +727,111 @@ impl Model { } Err(error) => { // Capture records without executing: the eager result is valid. - // Disable graphs for this worker rather than retrying failures. eprintln!("CUDA Graph capture failed; using eager execution: {error:#}"); self.graph_enabled = false; self.graphs.clear(); } } } - if self.graph_enabled { - self.graphs - .iter() - .find(|(length, _)| *length == t) - .unwrap() - .1 - .launch(self.stream)?; - } } else { - self.run(s, t)?; + self.run(s, t, false)?; + } + self.last_hidden(s, t) + } + + /// Prefill one unpadded prompt with already-adapted BF16 image embeddings and + /// explicit `[3, 1, sequence]` T/H/W positions. No vision tower runs here. + /// Load a checkpoint with the matching multimodal language adapter merged. + pub fn forward_multimodal(&mut self, input: &MultimodalInput<'_>) -> Result> { + let image_token = self + .cfg + .image_token_id + .context("checkpoint has no image_token_id")?; + input.validate( + self.cfg.hidden, + self.embed.shape[0], + image_token, + self.cfg.max_positions, + )?; + let t = input.token_ids.len(); + self.prepare_scratch(t)?; + let s = self.scratch.as_ref().unwrap(); + self.upload_positions(s, input.position_ids)?; + self.embed_tokens(s, input.token_ids)?; + let bytes: Vec = input + .image_embeddings + .iter() + .flat_map(|x| x.to_le_bytes()) + .collect(); + // Coalesce adjacent placeholders. Text rows remain those of embed_tokens. + let indices = input.image_token_indices; + let mut begin = 0; + while begin < indices.len() { + let mut end = begin + 1; + while end < indices.len() && indices[end] == indices[end - 1] + 1 { + end += 1; + } + let row_bytes = self.cfg.hidden * BF16; + // SAFETY: validated indices lie in the t-row residual buffer, and + // features contain exactly one hidden-size BF16 row per placeholder. + unsafe { + cuda::upload( + s.at(s.res + indices[begin] * row_bytes), + &bytes[begin * row_bytes..end * row_bytes], + self.stream, + )?; + } + begin = end; + } + self.run(s, t, true)?; + self.last_hidden(s, t) + } + + fn upload_positions(&self, s: &Scratch, positions: [&[i64]; 3]) -> Result<()> { + let (cos, sin) = rotary_tables( + positions, + self.cfg.rotary_half, + self.cfg.rope_theta, + self.cfg.mrope_section, + ); + // SAFETY: tables contain at most s.cap rows of rotary_half BF16 values. + unsafe { + cuda::upload(s.at(s.custom_cos), &cos, self.stream)?; + cuda::upload(s.at(s.custom_sin), &sin, self.stream)?; + } + Ok(()) + } + + fn embed_tokens(&self, s: &Scratch, ids: &[u32]) -> Result<()> { + let ids32: Vec = ids.iter().flat_map(|&i| (i as i32).to_le_bytes()).collect(); + // SAFETY: IDs were checked against the vocabulary; scratch holds t rows. + unsafe { + cuda::upload(s.at(s.ids), &ids32, self.stream)?; + check( + (cuda::api().cs1_embed)( + s.at(s.ids).cast(), + self.embed.ptr, + s.at(s.res), + ids.len() as i32, + self.cfg.hidden as i32, + self.stream, + ), + "embed", + )?; } - let mut last = vec![0u8; h * BF16]; + Ok(()) + } + + fn last_hidden(&self, s: &Scratch, t: usize) -> Result> { + let mut last = vec![0u8; self.cfg.hidden * BF16]; // SAFETY: x holds at least t rows of the hidden size. - unsafe { cuda::download(&mut last, s.at(s.x + (t - 1) * h * BF16), self.stream)? }; + unsafe { + cuda::download( + &mut last, + s.at(s.x + (t - 1) * self.cfg.hidden * BF16), + self.stream, + )?; + } let (pairs, _) = last.as_chunks::<2>(); Ok(pairs .iter() @@ -697,9 +839,10 @@ impl Model { .collect()) } - /// Queue one forward pass over the first `t` ids in `s`. The final-norm hidden - /// states end up in `s.x`. - fn run(&self, s: &Scratch, t: usize) -> Result<()> { + /// Queue language layers over prepared embeddings in `s.res`, with rotary + /// tables in immutable text buffers or separate explicit-position buffers. + /// Final-norm hidden states end up in `s.x`. + fn run(&self, s: &Scratch, t: usize, custom_positions: bool) -> Result<()> { let cfg = &self.cfg; let st = self.stream; let (ti, hi, eps) = (t as i32, cfg.hidden as i32, cfg.eps); @@ -707,13 +850,14 @@ impl Model { let (hq, hk, hd) = (cfg.heads as i32, cfg.kv_heads as i32, cfg.head_dim as i32); let w = Widths::of(cfg); let p = |off: usize| s.at(off); - // SAFETY (every kernel call below): the pointers are weights in the arena or + let (cos, sin) = if custom_positions { + (s.custom_cos, s.custom_sin) + } else { + (s.cos, s.sin) + }; + // SAFETY (every kernel call below): pointers are weights in the arena or // scratch buffers laid out for at least t tokens with the widths used here. unsafe { - check( - (cuda::api().cs1_embed)(p(s.ids).cast(), self.embed.ptr, p(s.res), ti, hi, st), - "embed", - )?; check( (cuda::api().cs1_rms_norm)( p(s.res), @@ -814,8 +958,8 @@ impl Model { ld, fa.q_norm.ptr, fa.k_norm.ptr, - p(s.cos), - p(s.sin), + p(cos), + p(sin), p(s.aq), p(s.agate), p(s.ak), @@ -911,3 +1055,71 @@ impl Model { Ok(()) } } + +#[cfg(test)] +mod graph_tests { + use super::*; + + #[test] + #[ignore = "needs CUA_S1_MODEL and ABI-3 CUA_S1_CUDA_LIB on a GPU"] + fn graph_misses_hits_eviction_growth_and_multimodal_match_eager() { + let dir = std::path::PathBuf::from(std::env::var_os("CUA_S1_MODEL").unwrap()); + let lib = std::path::PathBuf::from(std::env::var_os("CUA_S1_CUDA_LIB").unwrap()); + let prompts: Vec> = (4..=13) + .chain([1025]) + .map(|t| vec![32 + (t % 3) as u32; t]) + .collect(); + let changed_ids = vec![35; 4]; + let mut eager = Model::load(&dir, &lib).unwrap(); + eager.graph_enabled = false; + let expected: Vec<_> = prompts + .iter() + .map(|ids| eager.forward(ids).unwrap()) + .collect(); + let changed_expected = eager.forward(&changed_ids).unwrap(); + let image_ids = [32, eager.cfg.image_token_id.unwrap(), 33, 34]; + let features = vec![half::bf16::ONE; eager.cfg.hidden]; + let boundary = MultimodalInput { + token_ids: &image_ids, + image_token_indices: &[1], + image_embeddings: &features, + position_ids: [&[0, 1, 2, 3], &[0, 7, 8, 9], &[0, 3, 4, 5]], + }; + let multimodal_expected = eager.forward_multimodal(&boundary).unwrap(); + drop(eager); + + let mut model = Model::load(&dir, &lib).unwrap(); + model.graph_enabled = true; + for (ids, expected) in prompts[..10].iter().zip(&expected) { + // First use returns the eager result, then the graph is a cache hit. + assert_eq!(expected, &model.forward(ids).unwrap()); + assert!( + model.graph_enabled, + "capture unexpectedly fell back to eager" + ); + assert!(model.graphs.iter().any(|(t, _)| *t == ids.len())); + assert_eq!(expected, &model.forward(ids).unwrap()); + } + assert_eq!(model.graphs.len(), 8); + assert!(!model.graphs.iter().any(|(t, _)| *t == 4)); + // Evicted length is another miss; then change token IDs on a warm hit. + assert_eq!(expected[0], model.forward(&prompts[0]).unwrap()); + assert_eq!(changed_expected, model.forward(&changed_ids).unwrap()); + + // The same cached length must still use eager multimodal execution: + // reusing the text graph would read text positions instead of T/H/W. + assert!(model.graphs.iter().any(|(t, _)| *t == image_ids.len())); + assert_eq!( + multimodal_expected, + model.forward_multimodal(&boundary).unwrap() + ); + assert_eq!(expected[0], model.forward(&prompts[0]).unwrap()); + + // Growth invalidates all captures before freeing their device buffers. + assert_eq!(expected[10], model.forward(&prompts[10]).unwrap()); + assert_eq!(model.graphs.len(), 1); + assert_eq!(model.graphs[0].0, 1025); + assert_eq!(expected[10], model.forward(&prompts[10]).unwrap()); + assert_eq!(expected[0], model.forward(&prompts[0]).unwrap()); + } +} diff --git a/src/models/cua_s1/native/tests/multimodal.rs b/src/models/cua_s1/native/tests/multimodal.rs new file mode 100644 index 0000000..20afda5 --- /dev/null +++ b/src/models/cua_s1/native/tests/multimodal.rs @@ -0,0 +1,90 @@ +//! Real CUDA regression: explicit positions must not contaminate later text calls. +use std::path::PathBuf; + +use half::bf16; +use omni_cua_s1_native::{inputs::MultimodalInput, model::Model}; + +#[test] +#[ignore = "needs CUA_S1_MODEL and CUA_S1_CUDA_LIB on a GPU"] +fn text_multimodal_text_keeps_text_positions_and_overwrites_image_rows() { + let dir = PathBuf::from(std::env::var_os("CUA_S1_MODEL").expect("CUA_S1_MODEL")); + let lib = PathBuf::from(std::env::var_os("CUA_S1_CUDA_LIB").expect("CUA_S1_CUDA_LIB")); + let mut model = Model::load(&dir, &lib).unwrap(); + let ids = [32, 33, 34, 35]; + let positions = [0, 1, 2, 3]; + let baseline = model.forward(&ids).unwrap(); + let explicit = model + .forward_multimodal(&MultimodalInput { + token_ids: &ids, + image_token_indices: &[], + image_embeddings: &[], + position_ids: [&positions; 3], + }) + .unwrap(); + assert_eq!(baseline, explicit); + + let image_token = model.cfg.image_token_id.unwrap(); + let image_ids = [32, image_token, image_token, 35]; + let features = vec![bf16::ONE; 2 * model.cfg.hidden]; + let different = model + .forward_multimodal(&MultimodalInput { + token_ids: &image_ids, + image_token_indices: &[1, 2], + image_embeddings: &features, + position_ids: [&[0, 1, 1, 2], &[0, 1, 2, 2], &[0, 2, 1, 2]], + }) + .unwrap(); + assert!(different.iter().all(|x| x.is_finite())); + assert_ne!(baseline, different); + assert_eq!(baseline, model.forward(&ids).unwrap()); + + // Reusing the same layout with changed features must not retain old rows. + let features = vec![bf16::from_f32(-1.0); 2 * model.cfg.hidden]; + let changed = model + .forward_multimodal(&MultimodalInput { + token_ids: &image_ids, + image_token_indices: &[1, 2], + image_embeddings: &features, + position_ids: [&[0, 1, 1, 2], &[0, 1, 2, 2], &[0, 2, 1, 2]], + }) + .unwrap(); + assert_ne!(different, changed); + assert_eq!(baseline, model.forward(&ids).unwrap()); + + // Non-adjacent image spans exercise separate uploads and untouched text rows. + let disjoint_ids = [image_token, 33, image_token, 35]; + let disjoint = MultimodalInput { + token_ids: &disjoint_ids, + image_token_indices: &[0, 2], + image_embeddings: &features, + position_ids: [&positions; 3], + }; + let last = model.forward_multimodal(&disjoint).unwrap(); + assert_eq!(last, model.forward_multimodal(&disjoint).unwrap()); + + // A rejected boundary must not interfere with restoring the next text call. + assert!( + model + .forward_multimodal(&MultimodalInput { + image_token_indices: &[2, 0], + ..disjoint + }) + .is_err() + ); + assert_eq!(baseline, model.forward(&ids).unwrap()); + + // Cross the 1024-row allocation boundary, then use a shorter custom layout + // before reusing the immutable text tables in the larger scratch allocation. + let long_ids = vec![32; 1025]; + let long_text = model.forward(&long_ids).unwrap(); + model + .forward_multimodal(&MultimodalInput { + token_ids: &image_ids, + image_token_indices: &[1, 2], + image_embeddings: &features, + position_ids: [&[0, 1, 1, 2], &[0, 1, 2, 2], &[0, 2, 1, 2]], + }) + .unwrap(); + assert_eq!(long_text, model.forward(&long_ids).unwrap()); + assert_eq!(baseline, model.forward(&ids).unwrap()); +} From de46f070dd3ebee398b59778c017119a07e6c5e6 Mon Sep 17 00:00:00 2001 From: levius <2114377220@qq.com> Date: Thu, 1 Oct 2026 20:01:23 +0800 Subject: [PATCH 2/5] cua_s1: load and validate native vision checkpoints on CPU --- recipe/cua_s1/native.md | 3 + recipe/cua_s1/native_vision.md | 64 + src/models/cua_s1/native/Cargo.toml | 6 +- .../cua_s1/native/examples/inspect_vision.rs | 22 + src/models/cua_s1/native/src/lib.rs | 2 + src/models/cua_s1/native/src/vision/mod.rs | 477 +++++ tests/cua_s1/fixtures/vision/README.md | 31 + .../fixtures/vision/adapter-tensors.json | 702 ++++++ .../fixtures/vision/adapter_config.json | 48 + .../cua_s1/fixtures/vision/base-tensors.json | 1887 +++++++++++++++++ tests/cua_s1/fixtures/vision/config.json | 104 + tests/cua_s1/vision_loader.rs | 466 ++++ 12 files changed, 3811 insertions(+), 1 deletion(-) create mode 100644 recipe/cua_s1/native_vision.md create mode 100644 src/models/cua_s1/native/examples/inspect_vision.rs create mode 100644 src/models/cua_s1/native/src/vision/mod.rs create mode 100644 tests/cua_s1/fixtures/vision/README.md create mode 100644 tests/cua_s1/fixtures/vision/adapter-tensors.json create mode 100644 tests/cua_s1/fixtures/vision/adapter_config.json create mode 100644 tests/cua_s1/fixtures/vision/base-tensors.json create mode 100644 tests/cua_s1/fixtures/vision/config.json create mode 100644 tests/cua_s1/vision_loader.rs diff --git a/recipe/cua_s1/native.md b/recipe/cua_s1/native.md index 018feff..ba27138 100644 --- a/recipe/cua_s1/native.md +++ b/recipe/cua_s1/native.md @@ -41,3 +41,6 @@ cargo test -p omni-cua-s1-native CUA_S1_CUDA_LIB=$PWD/target/release/libqwen3_5_cuda.so \ cargo test --release -p omni-cua-s1-native --test kernels -- --ignored ``` + +For CPU-only structural inspection of the separate vision base weights and +multimodal adapter, see [native vision checkpoint inspection](native_vision.md). diff --git a/recipe/cua_s1/native_vision.md b/recipe/cua_s1/native_vision.md new file mode 100644 index 0000000..2a51b83 --- /dev/null +++ b/recipe/cua_s1/native_vision.md @@ -0,0 +1,64 @@ +# Cua-S1 native vision checkpoint inspection + +The model-owned Rust module `omni_cua_s1_native::vision` loads and structurally +validates the Qwen3.5-4B vision weights and Cua-S1 4B 0.2 **multimodal** LoRA on +CPU. It keeps the 297 BF16 base tensors and 100 FP32 adapter tensors separate. +The adapter contains 50 A/B pairs, rank 16, alpha 32, and scale 2. No weights are +merged or converted. This increment does not preprocess images, run the vision +encoder, call CUDA, or enable image requests in the native text worker. + +Prepare checkpoints from these pinned upstream revisions: + +| Checkpoint | Revision | Required files | +| --- | --- | --- | +| `Qwen/Qwen3.5-4B` | `851bf6e806efd8d0a36b00ddf55e13ccb7b8cd0a` | `config.json` and `model.safetensors`, or `model.safetensors.index.json` and its vision-bearing shards | +| `cua-ai/cua-s1-4b-0.2` multimodal adapter | `16818868b0cc7813808aae4e87b417657046ab79` | `adapter_config.json` and `adapter_model.safetensors` from the multimodal directory | + +Verify downloaded files against trusted upstream hashes before loading them. +Structural validation checks names, shapes, dtypes, configuration, and shard +mapping; it does **not** establish cryptographic identity or validate tensor +values. A bare base checkpoint or the `text` adapter is insufficient. + +From the repository root, pass exactly the base and multimodal adapter directories: + +```sh +cargo run --locked -p omni-cua-s1-native --example inspect_vision -- \ + weights/Qwen3.5-4B weights/cua-s1-4b-0.2/multimodal +``` + +The example reports the vision configuration, LoRA rank and scale, and tensor +counts. It needs neither a GPU nor a CUDA shared library. The index loader opens +only shards assigned visual tensors; unrelated language tensors in those shards +are ignored. Unexpected visual tensors, incompatible configuration, invalid +safetensors, missing tensors, mismatched index entries, and paths escaping the +checkpoint directory (including symlinks) are rejected. + +`VisionCheckpoint::load(base_dir, adapter_dir)` owns the memory maps and caches +validated tensor metadata. `config()` and `adapter()` expose the configuration; +`base_names()` / `adapter_names()` enumerate visual tensor names, and +`base_tensor(name)` / `adapter_tensor(name)` return borrowed safetensors +`TensorView`s with the original bytes. Views cannot outlive the checkpoint. +Accessors do not parse the file headers again. + +**Do not modify, replace in place, or truncate checkpoint files while a loaded +checkpoint or any of its views exists.** As with the existing native text +loader, callers must ensure that the memory-mapped files remain immutable. + +The supported layout is the pinned 4B configuration: 24 vision blocks, hidden +size 1024, intermediate size 4096, 16 heads, 2304 positions, output width 2560, +3 input channels, patch size 16, temporal patch size 2, spatial merge size 2, +`gelu_pytorch_tanh`, and no DeepStack. Text hidden size must equal vision output +width. LoRA targets both MLP matrices in every vision block and both merger +matrices. Math-changing PEFT options such as DoRA, rsLoRA, biases, target +exclusions, and custom rank/alpha patterns are unsupported. + +Run the CPU integration tests with: + +```sh +cargo test --locked -p omni-cua-s1-native --test vision_loader +``` + +Tests use sparse safetensors with real tensor shapes and small sentinel values; +no weight download or multi-gigabyte in-memory tensor allocation is required. +Their independent metadata oracle and provenance live in +[`tests/cua_s1/fixtures/vision/`](../../tests/cua_s1/fixtures/vision/). diff --git a/src/models/cua_s1/native/Cargo.toml b/src/models/cua_s1/native/Cargo.toml index 2e8caf6..e6903b7 100644 --- a/src/models/cua_s1/native/Cargo.toml +++ b/src/models/cua_s1/native/Cargo.toml @@ -17,8 +17,12 @@ half = "2.7.1" libloading = "0.8" memmap2 = "0.9.9" safetensors = "0.8.0" -serde = "1" +serde = { version = "1", features = ["derive"] } serde_json = { version = "1.0.149", features = ["float_roundtrip", "preserve_order"] } # the onig regex backend, as in the Python tokenizers wheel tokenizers = { version = "=0.22.2", default-features = false, features = ["onig"] } tokio = { version = "1.49.0", features = ["macros", "net", "rt-multi-thread", "sync"] } + +[[test]] +name = "vision_loader" +path = "../../../../tests/cua_s1/vision_loader.rs" diff --git a/src/models/cua_s1/native/examples/inspect_vision.rs b/src/models/cua_s1/native/examples/inspect_vision.rs new file mode 100644 index 0000000..afffb19 --- /dev/null +++ b/src/models/cua_s1/native/examples/inspect_vision.rs @@ -0,0 +1,22 @@ +use anyhow::{Result, ensure}; +use omni_cua_s1_native::vision::VisionCheckpoint; + +fn main() -> Result<()> { + let args: Vec<_> = std::env::args_os().skip(1).collect(); + ensure!( + args.len() == 2, + "usage: inspect_vision BASE_DIR MULTIMODAL_ADAPTER_DIR" + ); + let checkpoint = VisionCheckpoint::load(&args[0], &args[1])?; + println!("Vision: {:#?}", checkpoint.config()); + println!("Base: {} BF16 tensors", checkpoint.base_names().count()); + println!( + "Adapter: {} FP32 tensors ({} pairs), rank {}, alpha {}, scale {}", + checkpoint.adapter_names().count(), + checkpoint.adapter_names().count() / 2, + checkpoint.adapter().rank, + checkpoint.adapter().alpha, + checkpoint.adapter().scale() + ); + Ok(()) +} diff --git a/src/models/cua_s1/native/src/lib.rs b/src/models/cua_s1/native/src/lib.rs index 0d3aa5f..547e704 100644 --- a/src/models/cua_s1/native/src/lib.rs +++ b/src/models/cua_s1/native/src/lib.rs @@ -7,3 +7,5 @@ pub mod cuda; pub mod engine; pub mod json; pub mod model; + +pub mod vision; diff --git a/src/models/cua_s1/native/src/vision/mod.rs b/src/models/cua_s1/native/src/vision/mod.rs new file mode 100644 index 0000000..16cc059 --- /dev/null +++ b/src/models/cua_s1/native/src/vision/mod.rs @@ -0,0 +1,477 @@ +//! CPU-only, structurally validated vision checkpoint storage. No LoRA merge or execution. +//! +//! Callers must keep checkpoint files immutable (including no truncation) for the +//! lifetime of the checkpoint and its borrowed views. Structural checks do not +//! verify upstream hashes or tensor values. + +use anyhow::{Context, Result, ensure}; +use memmap2::Mmap; +use safetensors::{ + Dtype, SafeTensors, + tensor::{TensorInfo, TensorView}, +}; +use serde::{ + Deserialize, + de::{self, MapAccess, Visitor}, +}; +use serde_json::Value; +use std::{ + collections::{BTreeMap, BTreeSet}, + fmt, + fs::{self, File}, + path::{Component, Path, PathBuf}, +}; + +const BASE: &str = "model.visual."; +const ADAPTER: &str = "base_model.model.model.visual."; +type Inventory = BTreeMap>; + +/// The supported Qwen3.5-4B vision architecture; all fields are validated at load. +#[derive(Debug, Deserialize)] +pub struct VisionConfig { + pub depth: usize, + pub hidden_size: usize, + pub intermediate_size: usize, + pub num_heads: usize, + pub num_position_embeddings: usize, + pub out_hidden_size: usize, + pub in_channels: usize, + pub patch_size: usize, + pub temporal_patch_size: usize, + pub spatial_merge_size: usize, + pub hidden_act: String, + pub deepstack_visual_indexes: Vec, + pub model_type: String, +} +impl VisionConfig { + fn load(dir: &Path) -> Result { + let config: Value = serde_json::from_slice(&fs::read(inside(dir, "config.json")?)?)?; + let vision: Self = + serde_json::from_value(config["vision_config"].clone()).context("vision config")?; + let sizes = [ + vision.depth, + vision.hidden_size, + vision.intermediate_size, + vision.num_heads, + vision.num_position_embeddings, + vision.out_hidden_size, + vision.in_channels, + vision.patch_size, + vision.temporal_patch_size, + vision.spatial_merge_size, + ]; + ensure!( + sizes == [24, 1024, 4096, 16, 2304, 2560, 3, 16, 2, 2] + && vision.hidden_act == "gelu_pytorch_tanh" + && vision.deepstack_visual_indexes.is_empty() + && vision.model_type == "qwen3_5" + && config["model_type"] == "qwen3_5" + && config["text_config"]["hidden_size"].as_u64() + == Some(vision.out_hidden_size as u64), + "unsupported vision/text config: expected pinned Qwen3.5-4B layout" + ); + Ok(vision) + } + fn inventory(&self) -> Inventory { + let mut tensors = Inventory::new(); + let h = self.hidden_size; + let i = self.intermediate_size; + let mut linear = |name: String, output: usize, input: Option| { + tensors.insert( + format!("{BASE}{name}.weight"), + input.map_or_else(|| vec![output], |input| vec![output, input]), + ); + tensors.insert(format!("{BASE}{name}.bias"), vec![output]); + }; + for block in 0..self.depth { + for norm in ["norm1", "norm2"] { + linear(format!("blocks.{block}.{norm}"), h, None); + } + for (name, output, input) in [ + ("attn.qkv", 3 * h, h), + ("attn.proj", h, h), + ("mlp.linear_fc1", i, h), + ("mlp.linear_fc2", h, i), + ] { + linear(format!("blocks.{block}.{name}"), output, Some(input)); + } + } + let merged = h * self.spatial_merge_size * self.spatial_merge_size; + linear("merger.norm".into(), h, None); + linear("merger.linear_fc1".into(), merged, Some(merged)); + linear( + "merger.linear_fc2".into(), + self.out_hidden_size, + Some(merged), + ); + tensors.insert( + format!("{BASE}patch_embed.proj.weight"), + vec![ + h, + self.in_channels, + self.temporal_patch_size, + self.patch_size, + self.patch_size, + ], + ); + tensors.insert(format!("{BASE}patch_embed.proj.bias"), vec![h]); + tensors.insert( + format!("{BASE}pos_embed.weight"), + vec![self.num_position_embeddings, h], + ); + tensors + } +} + +/// Inference LoRA parameters; base and adapter bytes remain separate. +#[derive(Debug)] +pub struct AdapterConfig { + pub rank: usize, + pub alpha: usize, +} +impl AdapterConfig { + pub fn scale(&self) -> f64 { + self.alpha as f64 / self.rank as f64 + } + fn load(dir: &Path) -> Result { + let config: Value = + serde_json::from_slice(&fs::read(inside(dir, "adapter_config.json")?)?)?; + let c = config + .as_object() + .context("adapter config must be an object")?; + for (key, value) in c { + let supported = match key.as_str() { + "r" => value == 16, + "lora_alpha" => value == 32, + "peft_type" => value == "LORA", + "bias" => value == "none", + "base_model_name_or_path" => value == "Qwen/Qwen3.5-4B", + "task_type" => value == "CAUSAL_LM", + "lora_bias" + | "use_dora" + | "use_rslora" + | "use_qalora" + | "fan_in_fan_out" + | "ensure_weight_tying" => value == false, + "rank_pattern" | "alpha_pattern" | "loftq_config" => { + value.as_object().is_some_and(|v| v.is_empty()) + } + "exclude_modules" + | "modules_to_save" + | "layers_to_transform" + | "layers_pattern" + | "layer_replication" + | "target_parameters" + | "trainable_token_indices" + | "alora_invocation_tokens" + | "arrow_config" + | "corda_config" + | "eva_config" + | "megatron_config" => value.is_null(), + // Training/serialization metadata does not change ordinary inference LoRA. + "auto_mapping" | "inference_mode" | "init_lora_weights" | "lora_dropout" + | "megatron_core" | "peft_version" | "qalora_group_size" | "revision" + | "target_modules" => true, + _ => false, + }; + ensure!( + supported, + "unsupported adapter config option {key}: {value}" + ); + } + for (key, value) in [ + ("r", Value::from(16)), + ("lora_alpha", Value::from(32)), + ("peft_type", Value::from("LORA")), + ("bias", Value::from("none")), + ] { + ensure!(config[key] == value, "unsupported adapter config {key}"); + } + let targets = config["target_modules"] + .as_array() + .context("adapter config target_modules must be an array")?; + let expected = BTreeSet::from([ + "up_proj", + "k_proj", + "linear_fc1", + "q_proj", + "linear_fc2", + "down_proj", + "gate_proj", + "o_proj", + "v_proj", + ]); + let actual: BTreeSet<_> = targets.iter().filter_map(Value::as_str).collect(); + ensure!( + actual == expected && targets.len() == expected.len(), + "adapter config requires the full multimodal target_modules" + ); + Ok(Self { + rank: 16, + alpha: 32, + }) + } + fn inventory(&self, base: &Inventory) -> Inventory { + let mut tensors = Inventory::new(); + for (name, shape) in base { + if name.ends_with(".weight") + && (name.contains(".linear_fc1.") || name.contains(".linear_fc2.")) + { + let module = name + .strip_prefix(BASE) + .unwrap() + .strip_suffix(".weight") + .unwrap(); + tensors.insert( + format!("{ADAPTER}{module}.lora_A.weight"), + vec![self.rank, shape[1]], + ); + tensors.insert( + format!("{ADAPTER}{module}.lora_B.weight"), + vec![shape[0], self.rank], + ); + } + } + tensors + } +} + +/// Immutable mmap storage for one base checkpoint and its multimodal adapter. +pub struct VisionCheckpoint { + config: VisionConfig, + adapter: AdapterConfig, + base: TensorStore, + lora: TensorStore, +} +impl VisionCheckpoint { + /// Loads and validates headers on CPU. Keep the files immutable while mapped. + pub fn load(base_dir: impl AsRef, adapter_dir: impl AsRef) -> Result { + let base_dir = fs::canonicalize(base_dir).context("base checkpoint directory")?; + let adapter_dir = fs::canonicalize(adapter_dir).context("adapter checkpoint directory")?; + let config = VisionConfig::load(&base_dir).context("base config")?; + let adapter = AdapterConfig::load(&adapter_dir).context("adapter config")?; + let expected = config.inventory(); + let lora_expected = adapter.inventory(&expected); + let index_path = base_dir.join("model.safetensors.index.json"); + let index = if index_path.try_exists()? { + #[derive(Deserialize)] + struct Index { + weight_map: UniqueMap, + } + let index: Index = serde_json::from_slice(&fs::read(inside( + &base_dir, + "model.safetensors.index.json", + )?)?) + .context("safetensors index")?; + let map = index.weight_map.0; + for name in map.keys().filter(|n| is_visual(n)) { + ensure!( + expected.contains_key(name), + "unexpected visual tensor in index: {name}" + ); + } + for name in expected.keys() { + ensure!( + map.contains_key(name), + "missing visual tensor in index: {name}" + ); + } + Some(map) + } else { + None + }; + let files: BTreeSet = match &index { + Some(index) => expected.keys().map(|n| index[n].clone()).collect(), + None => BTreeSet::from(["model.safetensors".into()]), + }; + let base = TensorStore::load(&base_dir, files, &expected, Dtype::BF16, index.as_ref())?; + let lora = TensorStore::load( + &adapter_dir, + BTreeSet::from(["adapter_model.safetensors".into()]), + &lora_expected, + Dtype::F32, + None, + )?; + Ok(Self { + config, + adapter, + base, + lora, + }) + } + pub fn config(&self) -> &VisionConfig { + &self.config + } + pub fn adapter(&self) -> &AdapterConfig { + &self.adapter + } + pub fn base_names(&self) -> impl Iterator { + self.base.tensors.keys().map(String::as_str) + } + pub fn adapter_names(&self) -> impl Iterator { + self.lora.tensors.keys().map(String::as_str) + } + pub fn base_tensor(&self, name: &str) -> Result> { + self.base.tensor(name) + } + pub fn adapter_tensor(&self, name: &str) -> Result> { + self.lora.tensor(name) + } +} + +struct TensorStore { + maps: Vec, + tensors: BTreeMap, +} +impl TensorStore { + fn load( + dir: &Path, + files: BTreeSet, + expected: &Inventory, + dtype: Dtype, + index: Option<&BTreeMap>, + ) -> Result { + let mut store = Self { + maps: Vec::new(), + tensors: BTreeMap::new(), + }; + for filename in files { + let path = inside(dir, &filename)?; + let file = File::open(&path).with_context(|| format!("open {}", path.display()))?; + // SAFETY: callers must not mutate or truncate checkpoint files while mapped. + let map = unsafe { Mmap::map(&file) } + .with_context(|| format!("mmap safetensors {}", path.display()))?; + validate_header(&map) + .with_context(|| format!("safetensors header {}", path.display()))?; + let (header_len, metadata) = SafeTensors::read_metadata(&map) + .with_context(|| format!("safetensors {}", path.display()))?; + for (name, info) in metadata.tensors() { + if !is_visual(&name) { + continue; + } + let shape = expected + .get(&name) + .with_context(|| format!("unexpected visual tensor {name} in {filename}"))?; + if let Some(index) = index { + ensure!( + index.get(&name) == Some(&filename), + "index mismatch for {name} in {filename}" + ); + } + ensure!( + &info.shape == shape, + "shape mismatch for {name}: {:?}, expected {shape:?}", + info.shape + ); + ensure!( + info.dtype == dtype, + "dtype mismatch for {name}: {:?}, expected {dtype:?}", + info.dtype + ); + let mut info = info.clone(); + info.data_offsets.0 += 8 + header_len; + info.data_offsets.1 += 8 + header_len; + ensure!( + store + .tensors + .insert(name.clone(), (store.maps.len(), info)) + .is_none(), + "duplicate visual tensor {name}" + ); + } + store.maps.push(map); + } + for name in expected.keys() { + ensure!( + store.tensors.contains_key(name), + "missing visual tensor {name}" + ); + } + Ok(store) + } + fn tensor(&self, name: &str) -> Result> { + let (shard, info) = self + .tensors + .get(name) + .with_context(|| format!("unknown visual tensor {name}"))?; + Ok(TensorView::new( + info.dtype, + info.shape.clone(), + &self.maps[*shard][info.data_offsets.0..info.data_offsets.1], + )?) + } +} +// Bound all offsets before safetensors 0.8 adds payload size to header size: +// its final length check uses unchecked addition, even for ignored language tensors. +fn validate_header(bytes: &[u8]) -> Result<()> { + let length_bytes = bytes.get(..8).context("missing header length")?; + let header_len = usize::try_from(u64::from_le_bytes(length_bytes.try_into()?))?; + // Match safetensors 0.8's header allocation limit. + ensure!(header_len <= 100_000_000, "header too large"); + let data_start = header_len + .checked_add(8) + .context("header length overflow")?; + let header = bytes.get(8..data_start).context("truncated header")?; + let payload_len = bytes.len() - data_start; + let entries: UniqueMap = serde_json::from_slice(header)?; + for (name, entry) in entries.0 { + if name == "__metadata__" { + continue; + } + let (start, end): (usize, usize) = serde_json::from_value(entry["data_offsets"].clone()) + .with_context(|| format!("invalid offsets for {name}"))?; + ensure!( + start <= end && end <= payload_len, + "tensor offsets exceed payload for {name}" + ); + } + Ok(()) +} +fn is_visual(name: &str) -> bool { + name.split('.').any(|part| part == "visual") +} +fn inside(dir: &Path, name: &str) -> Result { + ensure!( + !name.is_empty() + && Path::new(name) + .components() + .all(|c| matches!(c, Component::Normal(_))), + "path must remain inside checkpoint directory: {name}" + ); + let resolved = + fs::canonicalize(dir.join(name)).with_context(|| format!("checkpoint file {name}"))?; + ensure!( + resolved.starts_with(dir), + "path escapes checkpoint directory: {name}" + ); + Ok(resolved) +} + +// serde_json's ordinary maps overwrite repeated keys. Reject ambiguous headers/indexes. +struct UniqueMap(BTreeMap); +impl<'de, T: Deserialize<'de>> Deserialize<'de> for UniqueMap { + fn deserialize>( + deserializer: D, + ) -> std::result::Result { + struct UniqueVisitor(std::marker::PhantomData); + impl<'de, T: Deserialize<'de>> Visitor<'de> for UniqueVisitor { + type Value = UniqueMap; + fn expecting(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str("an object with unique names") + } + fn visit_map>( + self, + mut map: M, + ) -> std::result::Result { + let mut entries = BTreeMap::new(); + while let Some((key, value)) = map.next_entry::()? { + if entries.insert(key.clone(), value).is_some() { + return Err(de::Error::custom(format!("duplicate JSON key: {key}"))); + } + } + Ok(UniqueMap(entries)) + } + } + deserializer.deserialize_map(UniqueVisitor(std::marker::PhantomData)) + } +} diff --git a/tests/cua_s1/fixtures/vision/README.md b/tests/cua_s1/fixtures/vision/README.md new file mode 100644 index 0000000..21a8ca6 --- /dev/null +++ b/tests/cua_s1/fixtures/vision/README.md @@ -0,0 +1,31 @@ +# Pinned vision metadata fixtures + +These fixtures are an independent test oracle captured from upstream metadata, +not generated by the production Rust inventory builder: + +- Base: `Qwen/Qwen3.5-4B`, revision + `851bf6e806efd8d0a36b00ddf55e13ccb7b8cd0a`. +- Adapter: `cua-ai/cua-s1-4b-0.2`, revision + `16818868b0cc7813808aae4e87b417657046ab79`, multimodal adapter. +- `config.json` and `adapter_config.json`: upstream configuration files with LF line endings. +- `base-tensors.json`: visual names, shapes, and dtypes read from the pinned + base safetensors headers (prefix `model.visual.`): 297 BF16 tensors. +- `adapter-tensors.json`: visual names, shapes, and dtypes read from the pinned + multimodal adapter safetensors header (prefix + `base_model.model.model.visual.`): 100 FP32 tensors / 50 LoRA pairs. + +No weight payloads are included. Tests construct sparse files with these shapes +and sentinel values. Production code never reads these fixtures. Fixture +matching proves structural compatibility, not checkpoint identity: verify real +checkpoint downloads against trusted upstream hashes separately. + +Exact pinned sources: + +- [Base configuration](https://huggingface.co/Qwen/Qwen3.5-4B/resolve/851bf6e806efd8d0a36b00ddf55e13ccb7b8cd0a/config.json) +- [Multimodal adapter configuration](https://huggingface.co/cua-ai/cua-s1-4b-0.2/resolve/16818868b0cc7813808aae4e87b417657046ab79/multimodal/adapter_config.json) +- [Base second shard, containing the vision tensors](https://huggingface.co/Qwen/Qwen3.5-4B/resolve/851bf6e806efd8d0a36b00ddf55e13ccb7b8cd0a/model.safetensors-00002-of-00002.safetensors) +- [Multimodal adapter safetensors](https://huggingface.co/cua-ai/cua-s1-4b-0.2/resolve/16818868b0cc7813808aae4e87b417657046ab79/multimodal/adapter_model.safetensors) + +The tensor inventories retain only the `shape` and `dtype` fields of the visual +entries in those safetensors headers; byte offsets and language tensors are not +part of the oracle. diff --git a/tests/cua_s1/fixtures/vision/adapter-tensors.json b/tests/cua_s1/fixtures/vision/adapter-tensors.json new file mode 100644 index 0000000..ac90604 --- /dev/null +++ b/tests/cua_s1/fixtures/vision/adapter-tensors.json @@ -0,0 +1,702 @@ +{ + "base_model.model.model.visual.blocks.0.mlp.linear_fc1.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 1024 + ] + }, + "base_model.model.model.visual.blocks.0.mlp.linear_fc1.lora_B.weight": { + "dtype": "F32", + "shape": [ + 4096, + 16 + ] + }, + "base_model.model.model.visual.blocks.0.mlp.linear_fc2.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 4096 + ] + }, + "base_model.model.model.visual.blocks.0.mlp.linear_fc2.lora_B.weight": { + "dtype": "F32", + "shape": [ + 1024, + 16 + ] + }, + "base_model.model.model.visual.blocks.1.mlp.linear_fc1.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 1024 + ] + }, + "base_model.model.model.visual.blocks.1.mlp.linear_fc1.lora_B.weight": { + "dtype": "F32", + "shape": [ + 4096, + 16 + ] + }, + "base_model.model.model.visual.blocks.1.mlp.linear_fc2.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 4096 + ] + }, + "base_model.model.model.visual.blocks.1.mlp.linear_fc2.lora_B.weight": { + "dtype": "F32", + "shape": [ + 1024, + 16 + ] + }, + "base_model.model.model.visual.blocks.10.mlp.linear_fc1.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 1024 + ] + }, + "base_model.model.model.visual.blocks.10.mlp.linear_fc1.lora_B.weight": { + "dtype": "F32", + "shape": [ + 4096, + 16 + ] + }, + "base_model.model.model.visual.blocks.10.mlp.linear_fc2.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 4096 + ] + }, + "base_model.model.model.visual.blocks.10.mlp.linear_fc2.lora_B.weight": { + "dtype": "F32", + "shape": [ + 1024, + 16 + ] + }, + "base_model.model.model.visual.blocks.11.mlp.linear_fc1.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 1024 + ] + }, + "base_model.model.model.visual.blocks.11.mlp.linear_fc1.lora_B.weight": { + "dtype": "F32", + "shape": [ + 4096, + 16 + ] + }, + "base_model.model.model.visual.blocks.11.mlp.linear_fc2.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 4096 + ] + }, + "base_model.model.model.visual.blocks.11.mlp.linear_fc2.lora_B.weight": { + "dtype": "F32", + "shape": [ + 1024, + 16 + ] + }, + "base_model.model.model.visual.blocks.12.mlp.linear_fc1.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 1024 + ] + }, + "base_model.model.model.visual.blocks.12.mlp.linear_fc1.lora_B.weight": { + "dtype": "F32", + "shape": [ + 4096, + 16 + ] + }, + "base_model.model.model.visual.blocks.12.mlp.linear_fc2.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 4096 + ] + }, + "base_model.model.model.visual.blocks.12.mlp.linear_fc2.lora_B.weight": { + "dtype": "F32", + "shape": [ + 1024, + 16 + ] + }, + "base_model.model.model.visual.blocks.13.mlp.linear_fc1.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 1024 + ] + }, + "base_model.model.model.visual.blocks.13.mlp.linear_fc1.lora_B.weight": { + "dtype": "F32", + "shape": [ + 4096, + 16 + ] + }, + "base_model.model.model.visual.blocks.13.mlp.linear_fc2.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 4096 + ] + }, + "base_model.model.model.visual.blocks.13.mlp.linear_fc2.lora_B.weight": { + "dtype": "F32", + "shape": [ + 1024, + 16 + ] + }, + "base_model.model.model.visual.blocks.14.mlp.linear_fc1.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 1024 + ] + }, + "base_model.model.model.visual.blocks.14.mlp.linear_fc1.lora_B.weight": { + "dtype": "F32", + "shape": [ + 4096, + 16 + ] + }, + "base_model.model.model.visual.blocks.14.mlp.linear_fc2.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 4096 + ] + }, + "base_model.model.model.visual.blocks.14.mlp.linear_fc2.lora_B.weight": { + "dtype": "F32", + "shape": [ + 1024, + 16 + ] + }, + "base_model.model.model.visual.blocks.15.mlp.linear_fc1.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 1024 + ] + }, + "base_model.model.model.visual.blocks.15.mlp.linear_fc1.lora_B.weight": { + "dtype": "F32", + "shape": [ + 4096, + 16 + ] + }, + "base_model.model.model.visual.blocks.15.mlp.linear_fc2.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 4096 + ] + }, + "base_model.model.model.visual.blocks.15.mlp.linear_fc2.lora_B.weight": { + "dtype": "F32", + "shape": [ + 1024, + 16 + ] + }, + "base_model.model.model.visual.blocks.16.mlp.linear_fc1.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 1024 + ] + }, + "base_model.model.model.visual.blocks.16.mlp.linear_fc1.lora_B.weight": { + "dtype": "F32", + "shape": [ + 4096, + 16 + ] + }, + "base_model.model.model.visual.blocks.16.mlp.linear_fc2.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 4096 + ] + }, + "base_model.model.model.visual.blocks.16.mlp.linear_fc2.lora_B.weight": { + "dtype": "F32", + "shape": [ + 1024, + 16 + ] + }, + "base_model.model.model.visual.blocks.17.mlp.linear_fc1.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 1024 + ] + }, + "base_model.model.model.visual.blocks.17.mlp.linear_fc1.lora_B.weight": { + "dtype": "F32", + "shape": [ + 4096, + 16 + ] + }, + "base_model.model.model.visual.blocks.17.mlp.linear_fc2.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 4096 + ] + }, + "base_model.model.model.visual.blocks.17.mlp.linear_fc2.lora_B.weight": { + "dtype": "F32", + "shape": [ + 1024, + 16 + ] + }, + "base_model.model.model.visual.blocks.18.mlp.linear_fc1.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 1024 + ] + }, + "base_model.model.model.visual.blocks.18.mlp.linear_fc1.lora_B.weight": { + "dtype": "F32", + "shape": [ + 4096, + 16 + ] + }, + "base_model.model.model.visual.blocks.18.mlp.linear_fc2.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 4096 + ] + }, + "base_model.model.model.visual.blocks.18.mlp.linear_fc2.lora_B.weight": { + "dtype": "F32", + "shape": [ + 1024, + 16 + ] + }, + "base_model.model.model.visual.blocks.19.mlp.linear_fc1.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 1024 + ] + }, + "base_model.model.model.visual.blocks.19.mlp.linear_fc1.lora_B.weight": { + "dtype": "F32", + "shape": [ + 4096, + 16 + ] + }, + "base_model.model.model.visual.blocks.19.mlp.linear_fc2.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 4096 + ] + }, + "base_model.model.model.visual.blocks.19.mlp.linear_fc2.lora_B.weight": { + "dtype": "F32", + "shape": [ + 1024, + 16 + ] + }, + "base_model.model.model.visual.blocks.2.mlp.linear_fc1.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 1024 + ] + }, + "base_model.model.model.visual.blocks.2.mlp.linear_fc1.lora_B.weight": { + "dtype": "F32", + "shape": [ + 4096, + 16 + ] + }, + "base_model.model.model.visual.blocks.2.mlp.linear_fc2.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 4096 + ] + }, + "base_model.model.model.visual.blocks.2.mlp.linear_fc2.lora_B.weight": { + "dtype": "F32", + "shape": [ + 1024, + 16 + ] + }, + "base_model.model.model.visual.blocks.20.mlp.linear_fc1.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 1024 + ] + }, + "base_model.model.model.visual.blocks.20.mlp.linear_fc1.lora_B.weight": { + "dtype": "F32", + "shape": [ + 4096, + 16 + ] + }, + "base_model.model.model.visual.blocks.20.mlp.linear_fc2.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 4096 + ] + }, + "base_model.model.model.visual.blocks.20.mlp.linear_fc2.lora_B.weight": { + "dtype": "F32", + "shape": [ + 1024, + 16 + ] + }, + "base_model.model.model.visual.blocks.21.mlp.linear_fc1.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 1024 + ] + }, + "base_model.model.model.visual.blocks.21.mlp.linear_fc1.lora_B.weight": { + "dtype": "F32", + "shape": [ + 4096, + 16 + ] + }, + "base_model.model.model.visual.blocks.21.mlp.linear_fc2.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 4096 + ] + }, + "base_model.model.model.visual.blocks.21.mlp.linear_fc2.lora_B.weight": { + "dtype": "F32", + "shape": [ + 1024, + 16 + ] + }, + "base_model.model.model.visual.blocks.22.mlp.linear_fc1.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 1024 + ] + }, + "base_model.model.model.visual.blocks.22.mlp.linear_fc1.lora_B.weight": { + "dtype": "F32", + "shape": [ + 4096, + 16 + ] + }, + "base_model.model.model.visual.blocks.22.mlp.linear_fc2.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 4096 + ] + }, + "base_model.model.model.visual.blocks.22.mlp.linear_fc2.lora_B.weight": { + "dtype": "F32", + "shape": [ + 1024, + 16 + ] + }, + "base_model.model.model.visual.blocks.23.mlp.linear_fc1.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 1024 + ] + }, + "base_model.model.model.visual.blocks.23.mlp.linear_fc1.lora_B.weight": { + "dtype": "F32", + "shape": [ + 4096, + 16 + ] + }, + "base_model.model.model.visual.blocks.23.mlp.linear_fc2.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 4096 + ] + }, + "base_model.model.model.visual.blocks.23.mlp.linear_fc2.lora_B.weight": { + "dtype": "F32", + "shape": [ + 1024, + 16 + ] + }, + "base_model.model.model.visual.blocks.3.mlp.linear_fc1.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 1024 + ] + }, + "base_model.model.model.visual.blocks.3.mlp.linear_fc1.lora_B.weight": { + "dtype": "F32", + "shape": [ + 4096, + 16 + ] + }, + "base_model.model.model.visual.blocks.3.mlp.linear_fc2.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 4096 + ] + }, + "base_model.model.model.visual.blocks.3.mlp.linear_fc2.lora_B.weight": { + "dtype": "F32", + "shape": [ + 1024, + 16 + ] + }, + "base_model.model.model.visual.blocks.4.mlp.linear_fc1.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 1024 + ] + }, + "base_model.model.model.visual.blocks.4.mlp.linear_fc1.lora_B.weight": { + "dtype": "F32", + "shape": [ + 4096, + 16 + ] + }, + "base_model.model.model.visual.blocks.4.mlp.linear_fc2.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 4096 + ] + }, + "base_model.model.model.visual.blocks.4.mlp.linear_fc2.lora_B.weight": { + "dtype": "F32", + "shape": [ + 1024, + 16 + ] + }, + "base_model.model.model.visual.blocks.5.mlp.linear_fc1.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 1024 + ] + }, + "base_model.model.model.visual.blocks.5.mlp.linear_fc1.lora_B.weight": { + "dtype": "F32", + "shape": [ + 4096, + 16 + ] + }, + "base_model.model.model.visual.blocks.5.mlp.linear_fc2.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 4096 + ] + }, + "base_model.model.model.visual.blocks.5.mlp.linear_fc2.lora_B.weight": { + "dtype": "F32", + "shape": [ + 1024, + 16 + ] + }, + "base_model.model.model.visual.blocks.6.mlp.linear_fc1.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 1024 + ] + }, + "base_model.model.model.visual.blocks.6.mlp.linear_fc1.lora_B.weight": { + "dtype": "F32", + "shape": [ + 4096, + 16 + ] + }, + "base_model.model.model.visual.blocks.6.mlp.linear_fc2.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 4096 + ] + }, + "base_model.model.model.visual.blocks.6.mlp.linear_fc2.lora_B.weight": { + "dtype": "F32", + "shape": [ + 1024, + 16 + ] + }, + "base_model.model.model.visual.blocks.7.mlp.linear_fc1.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 1024 + ] + }, + "base_model.model.model.visual.blocks.7.mlp.linear_fc1.lora_B.weight": { + "dtype": "F32", + "shape": [ + 4096, + 16 + ] + }, + "base_model.model.model.visual.blocks.7.mlp.linear_fc2.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 4096 + ] + }, + "base_model.model.model.visual.blocks.7.mlp.linear_fc2.lora_B.weight": { + "dtype": "F32", + "shape": [ + 1024, + 16 + ] + }, + "base_model.model.model.visual.blocks.8.mlp.linear_fc1.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 1024 + ] + }, + "base_model.model.model.visual.blocks.8.mlp.linear_fc1.lora_B.weight": { + "dtype": "F32", + "shape": [ + 4096, + 16 + ] + }, + "base_model.model.model.visual.blocks.8.mlp.linear_fc2.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 4096 + ] + }, + "base_model.model.model.visual.blocks.8.mlp.linear_fc2.lora_B.weight": { + "dtype": "F32", + "shape": [ + 1024, + 16 + ] + }, + "base_model.model.model.visual.blocks.9.mlp.linear_fc1.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 1024 + ] + }, + "base_model.model.model.visual.blocks.9.mlp.linear_fc1.lora_B.weight": { + "dtype": "F32", + "shape": [ + 4096, + 16 + ] + }, + "base_model.model.model.visual.blocks.9.mlp.linear_fc2.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 4096 + ] + }, + "base_model.model.model.visual.blocks.9.mlp.linear_fc2.lora_B.weight": { + "dtype": "F32", + "shape": [ + 1024, + 16 + ] + }, + "base_model.model.model.visual.merger.linear_fc1.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 4096 + ] + }, + "base_model.model.model.visual.merger.linear_fc1.lora_B.weight": { + "dtype": "F32", + "shape": [ + 4096, + 16 + ] + }, + "base_model.model.model.visual.merger.linear_fc2.lora_A.weight": { + "dtype": "F32", + "shape": [ + 16, + 4096 + ] + }, + "base_model.model.model.visual.merger.linear_fc2.lora_B.weight": { + "dtype": "F32", + "shape": [ + 2560, + 16 + ] + } +} diff --git a/tests/cua_s1/fixtures/vision/adapter_config.json b/tests/cua_s1/fixtures/vision/adapter_config.json new file mode 100644 index 0000000..52b492f --- /dev/null +++ b/tests/cua_s1/fixtures/vision/adapter_config.json @@ -0,0 +1,48 @@ +{ + "alora_invocation_tokens": null, + "alpha_pattern": {}, + "arrow_config": null, + "auto_mapping": null, + "base_model_name_or_path": "Qwen/Qwen3.5-4B", + "bias": "none", + "corda_config": null, + "ensure_weight_tying": false, + "eva_config": null, + "exclude_modules": null, + "fan_in_fan_out": false, + "inference_mode": true, + "init_lora_weights": true, + "layer_replication": null, + "layers_pattern": null, + "layers_to_transform": null, + "loftq_config": {}, + "lora_alpha": 32, + "lora_bias": false, + "lora_dropout": 0.05, + "megatron_config": null, + "megatron_core": "megatron.core", + "modules_to_save": null, + "peft_type": "LORA", + "peft_version": "0.18.1", + "qalora_group_size": 16, + "r": 16, + "rank_pattern": {}, + "revision": null, + "target_modules": [ + "up_proj", + "k_proj", + "linear_fc1", + "q_proj", + "linear_fc2", + "down_proj", + "gate_proj", + "o_proj", + "v_proj" + ], + "target_parameters": null, + "task_type": "CAUSAL_LM", + "trainable_token_indices": null, + "use_dora": false, + "use_qalora": false, + "use_rslora": false +} diff --git a/tests/cua_s1/fixtures/vision/base-tensors.json b/tests/cua_s1/fixtures/vision/base-tensors.json new file mode 100644 index 0000000..0c2d0b9 --- /dev/null +++ b/tests/cua_s1/fixtures/vision/base-tensors.json @@ -0,0 +1,1887 @@ +{ + "model.visual.blocks.0.attn.proj.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.0.attn.proj.weight": { + "dtype": "BF16", + "shape": [ + 1024, + 1024 + ] + }, + "model.visual.blocks.0.attn.qkv.bias": { + "dtype": "BF16", + "shape": [ + 3072 + ] + }, + "model.visual.blocks.0.attn.qkv.weight": { + "dtype": "BF16", + "shape": [ + 3072, + 1024 + ] + }, + "model.visual.blocks.0.mlp.linear_fc1.bias": { + "dtype": "BF16", + "shape": [ + 4096 + ] + }, + "model.visual.blocks.0.mlp.linear_fc1.weight": { + "dtype": "BF16", + "shape": [ + 4096, + 1024 + ] + }, + "model.visual.blocks.0.mlp.linear_fc2.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.0.mlp.linear_fc2.weight": { + "dtype": "BF16", + "shape": [ + 1024, + 4096 + ] + }, + "model.visual.blocks.0.norm1.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.0.norm1.weight": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.0.norm2.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.0.norm2.weight": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.1.attn.proj.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.1.attn.proj.weight": { + "dtype": "BF16", + "shape": [ + 1024, + 1024 + ] + }, + "model.visual.blocks.1.attn.qkv.bias": { + "dtype": "BF16", + "shape": [ + 3072 + ] + }, + "model.visual.blocks.1.attn.qkv.weight": { + "dtype": "BF16", + "shape": [ + 3072, + 1024 + ] + }, + "model.visual.blocks.1.mlp.linear_fc1.bias": { + "dtype": "BF16", + "shape": [ + 4096 + ] + }, + "model.visual.blocks.1.mlp.linear_fc1.weight": { + "dtype": "BF16", + "shape": [ + 4096, + 1024 + ] + }, + "model.visual.blocks.1.mlp.linear_fc2.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.1.mlp.linear_fc2.weight": { + "dtype": "BF16", + "shape": [ + 1024, + 4096 + ] + }, + "model.visual.blocks.1.norm1.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.1.norm1.weight": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.1.norm2.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.1.norm2.weight": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.10.attn.proj.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.10.attn.proj.weight": { + "dtype": "BF16", + "shape": [ + 1024, + 1024 + ] + }, + "model.visual.blocks.10.attn.qkv.bias": { + "dtype": "BF16", + "shape": [ + 3072 + ] + }, + "model.visual.blocks.10.attn.qkv.weight": { + "dtype": "BF16", + "shape": [ + 3072, + 1024 + ] + }, + "model.visual.blocks.10.mlp.linear_fc1.bias": { + "dtype": "BF16", + "shape": [ + 4096 + ] + }, + "model.visual.blocks.10.mlp.linear_fc1.weight": { + "dtype": "BF16", + "shape": [ + 4096, + 1024 + ] + }, + "model.visual.blocks.10.mlp.linear_fc2.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.10.mlp.linear_fc2.weight": { + "dtype": "BF16", + "shape": [ + 1024, + 4096 + ] + }, + "model.visual.blocks.10.norm1.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.10.norm1.weight": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.10.norm2.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.10.norm2.weight": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.11.attn.proj.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.11.attn.proj.weight": { + "dtype": "BF16", + "shape": [ + 1024, + 1024 + ] + }, + "model.visual.blocks.11.attn.qkv.bias": { + "dtype": "BF16", + "shape": [ + 3072 + ] + }, + "model.visual.blocks.11.attn.qkv.weight": { + "dtype": "BF16", + "shape": [ + 3072, + 1024 + ] + }, + "model.visual.blocks.11.mlp.linear_fc1.bias": { + "dtype": "BF16", + "shape": [ + 4096 + ] + }, + "model.visual.blocks.11.mlp.linear_fc1.weight": { + "dtype": "BF16", + "shape": [ + 4096, + 1024 + ] + }, + "model.visual.blocks.11.mlp.linear_fc2.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.11.mlp.linear_fc2.weight": { + "dtype": "BF16", + "shape": [ + 1024, + 4096 + ] + }, + "model.visual.blocks.11.norm1.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.11.norm1.weight": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.11.norm2.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.11.norm2.weight": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.12.attn.proj.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.12.attn.proj.weight": { + "dtype": "BF16", + "shape": [ + 1024, + 1024 + ] + }, + "model.visual.blocks.12.attn.qkv.bias": { + "dtype": "BF16", + "shape": [ + 3072 + ] + }, + "model.visual.blocks.12.attn.qkv.weight": { + "dtype": "BF16", + "shape": [ + 3072, + 1024 + ] + }, + "model.visual.blocks.12.mlp.linear_fc1.bias": { + "dtype": "BF16", + "shape": [ + 4096 + ] + }, + "model.visual.blocks.12.mlp.linear_fc1.weight": { + "dtype": "BF16", + "shape": [ + 4096, + 1024 + ] + }, + "model.visual.blocks.12.mlp.linear_fc2.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.12.mlp.linear_fc2.weight": { + "dtype": "BF16", + "shape": [ + 1024, + 4096 + ] + }, + "model.visual.blocks.12.norm1.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.12.norm1.weight": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.12.norm2.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.12.norm2.weight": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.13.attn.proj.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.13.attn.proj.weight": { + "dtype": "BF16", + "shape": [ + 1024, + 1024 + ] + }, + "model.visual.blocks.13.attn.qkv.bias": { + "dtype": "BF16", + "shape": [ + 3072 + ] + }, + "model.visual.blocks.13.attn.qkv.weight": { + "dtype": "BF16", + "shape": [ + 3072, + 1024 + ] + }, + "model.visual.blocks.13.mlp.linear_fc1.bias": { + "dtype": "BF16", + "shape": [ + 4096 + ] + }, + "model.visual.blocks.13.mlp.linear_fc1.weight": { + "dtype": "BF16", + "shape": [ + 4096, + 1024 + ] + }, + "model.visual.blocks.13.mlp.linear_fc2.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.13.mlp.linear_fc2.weight": { + "dtype": "BF16", + "shape": [ + 1024, + 4096 + ] + }, + "model.visual.blocks.13.norm1.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.13.norm1.weight": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.13.norm2.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.13.norm2.weight": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.14.attn.proj.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.14.attn.proj.weight": { + "dtype": "BF16", + "shape": [ + 1024, + 1024 + ] + }, + "model.visual.blocks.14.attn.qkv.bias": { + "dtype": "BF16", + "shape": [ + 3072 + ] + }, + "model.visual.blocks.14.attn.qkv.weight": { + "dtype": "BF16", + "shape": [ + 3072, + 1024 + ] + }, + "model.visual.blocks.14.mlp.linear_fc1.bias": { + "dtype": "BF16", + "shape": [ + 4096 + ] + }, + "model.visual.blocks.14.mlp.linear_fc1.weight": { + "dtype": "BF16", + "shape": [ + 4096, + 1024 + ] + }, + "model.visual.blocks.14.mlp.linear_fc2.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.14.mlp.linear_fc2.weight": { + "dtype": "BF16", + "shape": [ + 1024, + 4096 + ] + }, + "model.visual.blocks.14.norm1.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.14.norm1.weight": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.14.norm2.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.14.norm2.weight": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.15.attn.proj.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.15.attn.proj.weight": { + "dtype": "BF16", + "shape": [ + 1024, + 1024 + ] + }, + "model.visual.blocks.15.attn.qkv.bias": { + "dtype": "BF16", + "shape": [ + 3072 + ] + }, + "model.visual.blocks.15.attn.qkv.weight": { + "dtype": "BF16", + "shape": [ + 3072, + 1024 + ] + }, + "model.visual.blocks.15.mlp.linear_fc1.bias": { + "dtype": "BF16", + "shape": [ + 4096 + ] + }, + "model.visual.blocks.15.mlp.linear_fc1.weight": { + "dtype": "BF16", + "shape": [ + 4096, + 1024 + ] + }, + "model.visual.blocks.15.mlp.linear_fc2.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.15.mlp.linear_fc2.weight": { + "dtype": "BF16", + "shape": [ + 1024, + 4096 + ] + }, + "model.visual.blocks.15.norm1.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.15.norm1.weight": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.15.norm2.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.15.norm2.weight": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.16.attn.proj.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.16.attn.proj.weight": { + "dtype": "BF16", + "shape": [ + 1024, + 1024 + ] + }, + "model.visual.blocks.16.attn.qkv.bias": { + "dtype": "BF16", + "shape": [ + 3072 + ] + }, + "model.visual.blocks.16.attn.qkv.weight": { + "dtype": "BF16", + "shape": [ + 3072, + 1024 + ] + }, + "model.visual.blocks.16.mlp.linear_fc1.bias": { + "dtype": "BF16", + "shape": [ + 4096 + ] + }, + "model.visual.blocks.16.mlp.linear_fc1.weight": { + "dtype": "BF16", + "shape": [ + 4096, + 1024 + ] + }, + "model.visual.blocks.16.mlp.linear_fc2.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.16.mlp.linear_fc2.weight": { + "dtype": "BF16", + "shape": [ + 1024, + 4096 + ] + }, + "model.visual.blocks.16.norm1.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.16.norm1.weight": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.16.norm2.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.16.norm2.weight": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.17.attn.proj.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.17.attn.proj.weight": { + "dtype": "BF16", + "shape": [ + 1024, + 1024 + ] + }, + "model.visual.blocks.17.attn.qkv.bias": { + "dtype": "BF16", + "shape": [ + 3072 + ] + }, + "model.visual.blocks.17.attn.qkv.weight": { + "dtype": "BF16", + "shape": [ + 3072, + 1024 + ] + }, + "model.visual.blocks.17.mlp.linear_fc1.bias": { + "dtype": "BF16", + "shape": [ + 4096 + ] + }, + "model.visual.blocks.17.mlp.linear_fc1.weight": { + "dtype": "BF16", + "shape": [ + 4096, + 1024 + ] + }, + "model.visual.blocks.17.mlp.linear_fc2.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.17.mlp.linear_fc2.weight": { + "dtype": "BF16", + "shape": [ + 1024, + 4096 + ] + }, + "model.visual.blocks.17.norm1.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.17.norm1.weight": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.17.norm2.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.17.norm2.weight": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.18.attn.proj.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.18.attn.proj.weight": { + "dtype": "BF16", + "shape": [ + 1024, + 1024 + ] + }, + "model.visual.blocks.18.attn.qkv.bias": { + "dtype": "BF16", + "shape": [ + 3072 + ] + }, + "model.visual.blocks.18.attn.qkv.weight": { + "dtype": "BF16", + "shape": [ + 3072, + 1024 + ] + }, + "model.visual.blocks.18.mlp.linear_fc1.bias": { + "dtype": "BF16", + "shape": [ + 4096 + ] + }, + "model.visual.blocks.18.mlp.linear_fc1.weight": { + "dtype": "BF16", + "shape": [ + 4096, + 1024 + ] + }, + "model.visual.blocks.18.mlp.linear_fc2.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.18.mlp.linear_fc2.weight": { + "dtype": "BF16", + "shape": [ + 1024, + 4096 + ] + }, + "model.visual.blocks.18.norm1.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.18.norm1.weight": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.18.norm2.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.18.norm2.weight": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.19.attn.proj.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.19.attn.proj.weight": { + "dtype": "BF16", + "shape": [ + 1024, + 1024 + ] + }, + "model.visual.blocks.19.attn.qkv.bias": { + "dtype": "BF16", + "shape": [ + 3072 + ] + }, + "model.visual.blocks.19.attn.qkv.weight": { + "dtype": "BF16", + "shape": [ + 3072, + 1024 + ] + }, + "model.visual.blocks.19.mlp.linear_fc1.bias": { + "dtype": "BF16", + "shape": [ + 4096 + ] + }, + "model.visual.blocks.19.mlp.linear_fc1.weight": { + "dtype": "BF16", + "shape": [ + 4096, + 1024 + ] + }, + "model.visual.blocks.19.mlp.linear_fc2.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.19.mlp.linear_fc2.weight": { + "dtype": "BF16", + "shape": [ + 1024, + 4096 + ] + }, + "model.visual.blocks.19.norm1.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.19.norm1.weight": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.19.norm2.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.19.norm2.weight": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.2.attn.proj.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.2.attn.proj.weight": { + "dtype": "BF16", + "shape": [ + 1024, + 1024 + ] + }, + "model.visual.blocks.2.attn.qkv.bias": { + "dtype": "BF16", + "shape": [ + 3072 + ] + }, + "model.visual.blocks.2.attn.qkv.weight": { + "dtype": "BF16", + "shape": [ + 3072, + 1024 + ] + }, + "model.visual.blocks.2.mlp.linear_fc1.bias": { + "dtype": "BF16", + "shape": [ + 4096 + ] + }, + "model.visual.blocks.2.mlp.linear_fc1.weight": { + "dtype": "BF16", + "shape": [ + 4096, + 1024 + ] + }, + "model.visual.blocks.2.mlp.linear_fc2.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.2.mlp.linear_fc2.weight": { + "dtype": "BF16", + "shape": [ + 1024, + 4096 + ] + }, + "model.visual.blocks.2.norm1.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.2.norm1.weight": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.2.norm2.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.2.norm2.weight": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.20.attn.proj.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.20.attn.proj.weight": { + "dtype": "BF16", + "shape": [ + 1024, + 1024 + ] + }, + "model.visual.blocks.20.attn.qkv.bias": { + "dtype": "BF16", + "shape": [ + 3072 + ] + }, + "model.visual.blocks.20.attn.qkv.weight": { + "dtype": "BF16", + "shape": [ + 3072, + 1024 + ] + }, + "model.visual.blocks.20.mlp.linear_fc1.bias": { + "dtype": "BF16", + "shape": [ + 4096 + ] + }, + "model.visual.blocks.20.mlp.linear_fc1.weight": { + "dtype": "BF16", + "shape": [ + 4096, + 1024 + ] + }, + "model.visual.blocks.20.mlp.linear_fc2.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.20.mlp.linear_fc2.weight": { + "dtype": "BF16", + "shape": [ + 1024, + 4096 + ] + }, + "model.visual.blocks.20.norm1.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.20.norm1.weight": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.20.norm2.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.20.norm2.weight": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.21.attn.proj.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.21.attn.proj.weight": { + "dtype": "BF16", + "shape": [ + 1024, + 1024 + ] + }, + "model.visual.blocks.21.attn.qkv.bias": { + "dtype": "BF16", + "shape": [ + 3072 + ] + }, + "model.visual.blocks.21.attn.qkv.weight": { + "dtype": "BF16", + "shape": [ + 3072, + 1024 + ] + }, + "model.visual.blocks.21.mlp.linear_fc1.bias": { + "dtype": "BF16", + "shape": [ + 4096 + ] + }, + "model.visual.blocks.21.mlp.linear_fc1.weight": { + "dtype": "BF16", + "shape": [ + 4096, + 1024 + ] + }, + "model.visual.blocks.21.mlp.linear_fc2.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.21.mlp.linear_fc2.weight": { + "dtype": "BF16", + "shape": [ + 1024, + 4096 + ] + }, + "model.visual.blocks.21.norm1.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.21.norm1.weight": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.21.norm2.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.21.norm2.weight": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.22.attn.proj.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.22.attn.proj.weight": { + "dtype": "BF16", + "shape": [ + 1024, + 1024 + ] + }, + "model.visual.blocks.22.attn.qkv.bias": { + "dtype": "BF16", + "shape": [ + 3072 + ] + }, + "model.visual.blocks.22.attn.qkv.weight": { + "dtype": "BF16", + "shape": [ + 3072, + 1024 + ] + }, + "model.visual.blocks.22.mlp.linear_fc1.bias": { + "dtype": "BF16", + "shape": [ + 4096 + ] + }, + "model.visual.blocks.22.mlp.linear_fc1.weight": { + "dtype": "BF16", + "shape": [ + 4096, + 1024 + ] + }, + "model.visual.blocks.22.mlp.linear_fc2.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.22.mlp.linear_fc2.weight": { + "dtype": "BF16", + "shape": [ + 1024, + 4096 + ] + }, + "model.visual.blocks.22.norm1.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.22.norm1.weight": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.22.norm2.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.22.norm2.weight": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.23.attn.proj.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.23.attn.proj.weight": { + "dtype": "BF16", + "shape": [ + 1024, + 1024 + ] + }, + "model.visual.blocks.23.attn.qkv.bias": { + "dtype": "BF16", + "shape": [ + 3072 + ] + }, + "model.visual.blocks.23.attn.qkv.weight": { + "dtype": "BF16", + "shape": [ + 3072, + 1024 + ] + }, + "model.visual.blocks.23.mlp.linear_fc1.bias": { + "dtype": "BF16", + "shape": [ + 4096 + ] + }, + "model.visual.blocks.23.mlp.linear_fc1.weight": { + "dtype": "BF16", + "shape": [ + 4096, + 1024 + ] + }, + "model.visual.blocks.23.mlp.linear_fc2.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.23.mlp.linear_fc2.weight": { + "dtype": "BF16", + "shape": [ + 1024, + 4096 + ] + }, + "model.visual.blocks.23.norm1.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.23.norm1.weight": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.23.norm2.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.23.norm2.weight": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.3.attn.proj.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.3.attn.proj.weight": { + "dtype": "BF16", + "shape": [ + 1024, + 1024 + ] + }, + "model.visual.blocks.3.attn.qkv.bias": { + "dtype": "BF16", + "shape": [ + 3072 + ] + }, + "model.visual.blocks.3.attn.qkv.weight": { + "dtype": "BF16", + "shape": [ + 3072, + 1024 + ] + }, + "model.visual.blocks.3.mlp.linear_fc1.bias": { + "dtype": "BF16", + "shape": [ + 4096 + ] + }, + "model.visual.blocks.3.mlp.linear_fc1.weight": { + "dtype": "BF16", + "shape": [ + 4096, + 1024 + ] + }, + "model.visual.blocks.3.mlp.linear_fc2.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.3.mlp.linear_fc2.weight": { + "dtype": "BF16", + "shape": [ + 1024, + 4096 + ] + }, + "model.visual.blocks.3.norm1.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.3.norm1.weight": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.3.norm2.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.3.norm2.weight": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.4.attn.proj.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.4.attn.proj.weight": { + "dtype": "BF16", + "shape": [ + 1024, + 1024 + ] + }, + "model.visual.blocks.4.attn.qkv.bias": { + "dtype": "BF16", + "shape": [ + 3072 + ] + }, + "model.visual.blocks.4.attn.qkv.weight": { + "dtype": "BF16", + "shape": [ + 3072, + 1024 + ] + }, + "model.visual.blocks.4.mlp.linear_fc1.bias": { + "dtype": "BF16", + "shape": [ + 4096 + ] + }, + "model.visual.blocks.4.mlp.linear_fc1.weight": { + "dtype": "BF16", + "shape": [ + 4096, + 1024 + ] + }, + "model.visual.blocks.4.mlp.linear_fc2.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.4.mlp.linear_fc2.weight": { + "dtype": "BF16", + "shape": [ + 1024, + 4096 + ] + }, + "model.visual.blocks.4.norm1.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.4.norm1.weight": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.4.norm2.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.4.norm2.weight": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.5.attn.proj.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.5.attn.proj.weight": { + "dtype": "BF16", + "shape": [ + 1024, + 1024 + ] + }, + "model.visual.blocks.5.attn.qkv.bias": { + "dtype": "BF16", + "shape": [ + 3072 + ] + }, + "model.visual.blocks.5.attn.qkv.weight": { + "dtype": "BF16", + "shape": [ + 3072, + 1024 + ] + }, + "model.visual.blocks.5.mlp.linear_fc1.bias": { + "dtype": "BF16", + "shape": [ + 4096 + ] + }, + "model.visual.blocks.5.mlp.linear_fc1.weight": { + "dtype": "BF16", + "shape": [ + 4096, + 1024 + ] + }, + "model.visual.blocks.5.mlp.linear_fc2.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.5.mlp.linear_fc2.weight": { + "dtype": "BF16", + "shape": [ + 1024, + 4096 + ] + }, + "model.visual.blocks.5.norm1.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.5.norm1.weight": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.5.norm2.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.5.norm2.weight": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.6.attn.proj.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.6.attn.proj.weight": { + "dtype": "BF16", + "shape": [ + 1024, + 1024 + ] + }, + "model.visual.blocks.6.attn.qkv.bias": { + "dtype": "BF16", + "shape": [ + 3072 + ] + }, + "model.visual.blocks.6.attn.qkv.weight": { + "dtype": "BF16", + "shape": [ + 3072, + 1024 + ] + }, + "model.visual.blocks.6.mlp.linear_fc1.bias": { + "dtype": "BF16", + "shape": [ + 4096 + ] + }, + "model.visual.blocks.6.mlp.linear_fc1.weight": { + "dtype": "BF16", + "shape": [ + 4096, + 1024 + ] + }, + "model.visual.blocks.6.mlp.linear_fc2.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.6.mlp.linear_fc2.weight": { + "dtype": "BF16", + "shape": [ + 1024, + 4096 + ] + }, + "model.visual.blocks.6.norm1.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.6.norm1.weight": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.6.norm2.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.6.norm2.weight": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.7.attn.proj.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.7.attn.proj.weight": { + "dtype": "BF16", + "shape": [ + 1024, + 1024 + ] + }, + "model.visual.blocks.7.attn.qkv.bias": { + "dtype": "BF16", + "shape": [ + 3072 + ] + }, + "model.visual.blocks.7.attn.qkv.weight": { + "dtype": "BF16", + "shape": [ + 3072, + 1024 + ] + }, + "model.visual.blocks.7.mlp.linear_fc1.bias": { + "dtype": "BF16", + "shape": [ + 4096 + ] + }, + "model.visual.blocks.7.mlp.linear_fc1.weight": { + "dtype": "BF16", + "shape": [ + 4096, + 1024 + ] + }, + "model.visual.blocks.7.mlp.linear_fc2.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.7.mlp.linear_fc2.weight": { + "dtype": "BF16", + "shape": [ + 1024, + 4096 + ] + }, + "model.visual.blocks.7.norm1.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.7.norm1.weight": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.7.norm2.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.7.norm2.weight": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.8.attn.proj.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.8.attn.proj.weight": { + "dtype": "BF16", + "shape": [ + 1024, + 1024 + ] + }, + "model.visual.blocks.8.attn.qkv.bias": { + "dtype": "BF16", + "shape": [ + 3072 + ] + }, + "model.visual.blocks.8.attn.qkv.weight": { + "dtype": "BF16", + "shape": [ + 3072, + 1024 + ] + }, + "model.visual.blocks.8.mlp.linear_fc1.bias": { + "dtype": "BF16", + "shape": [ + 4096 + ] + }, + "model.visual.blocks.8.mlp.linear_fc1.weight": { + "dtype": "BF16", + "shape": [ + 4096, + 1024 + ] + }, + "model.visual.blocks.8.mlp.linear_fc2.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.8.mlp.linear_fc2.weight": { + "dtype": "BF16", + "shape": [ + 1024, + 4096 + ] + }, + "model.visual.blocks.8.norm1.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.8.norm1.weight": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.8.norm2.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.8.norm2.weight": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.9.attn.proj.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.9.attn.proj.weight": { + "dtype": "BF16", + "shape": [ + 1024, + 1024 + ] + }, + "model.visual.blocks.9.attn.qkv.bias": { + "dtype": "BF16", + "shape": [ + 3072 + ] + }, + "model.visual.blocks.9.attn.qkv.weight": { + "dtype": "BF16", + "shape": [ + 3072, + 1024 + ] + }, + "model.visual.blocks.9.mlp.linear_fc1.bias": { + "dtype": "BF16", + "shape": [ + 4096 + ] + }, + "model.visual.blocks.9.mlp.linear_fc1.weight": { + "dtype": "BF16", + "shape": [ + 4096, + 1024 + ] + }, + "model.visual.blocks.9.mlp.linear_fc2.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.9.mlp.linear_fc2.weight": { + "dtype": "BF16", + "shape": [ + 1024, + 4096 + ] + }, + "model.visual.blocks.9.norm1.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.9.norm1.weight": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.9.norm2.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.blocks.9.norm2.weight": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.merger.linear_fc1.bias": { + "dtype": "BF16", + "shape": [ + 4096 + ] + }, + "model.visual.merger.linear_fc1.weight": { + "dtype": "BF16", + "shape": [ + 4096, + 4096 + ] + }, + "model.visual.merger.linear_fc2.bias": { + "dtype": "BF16", + "shape": [ + 2560 + ] + }, + "model.visual.merger.linear_fc2.weight": { + "dtype": "BF16", + "shape": [ + 2560, + 4096 + ] + }, + "model.visual.merger.norm.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.merger.norm.weight": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.patch_embed.proj.bias": { + "dtype": "BF16", + "shape": [ + 1024 + ] + }, + "model.visual.patch_embed.proj.weight": { + "dtype": "BF16", + "shape": [ + 1024, + 3, + 2, + 16, + 16 + ] + }, + "model.visual.pos_embed.weight": { + "dtype": "BF16", + "shape": [ + 2304, + 1024 + ] + } +} diff --git a/tests/cua_s1/fixtures/vision/config.json b/tests/cua_s1/fixtures/vision/config.json new file mode 100644 index 0000000..cca4041 --- /dev/null +++ b/tests/cua_s1/fixtures/vision/config.json @@ -0,0 +1,104 @@ +{ + "architectures": [ + "Qwen3_5ForConditionalGeneration" + ], + "image_token_id": 248056, + "model_type": "qwen3_5", + "text_config": { + "attention_bias": false, + "attention_dropout": 0.0, + "attn_output_gate": true, + "dtype": "bfloat16", + "eos_token_id": 248044, + "full_attention_interval": 4, + "head_dim": 256, + "hidden_act": "silu", + "hidden_size": 2560, + "initializer_range": 0.02, + "intermediate_size": 9216, + "layer_types": [ + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention" + ], + "linear_conv_kernel_dim": 4, + "linear_key_head_dim": 128, + "linear_num_key_heads": 16, + "linear_num_value_heads": 32, + "linear_value_head_dim": 128, + "max_position_embeddings": 262144, + "mlp_only_layers": [], + "model_type": "qwen3_5_text", + "mtp_num_hidden_layers": 1, + "mtp_use_dedicated_embeddings": false, + "num_attention_heads": 16, + "num_hidden_layers": 32, + "num_key_value_heads": 4, + "rms_norm_eps": 1e-06, + "tie_word_embeddings": true, + "use_cache": true, + "vocab_size": 248320, + "mamba_ssm_dtype": "float32", + "rope_parameters": { + "mrope_interleaved": true, + "mrope_section": [ + 11, + 11, + 10 + ], + "rope_type": "default", + "rope_theta": 10000000, + "partial_rotary_factor": 0.25 + } + }, + "tie_word_embeddings": true, + "transformers_version": "4.57.0.dev0", + "video_token_id": 248057, + "vision_config": { + "deepstack_visual_indexes": [], + "depth": 24, + "hidden_act": "gelu_pytorch_tanh", + "hidden_size": 1024, + "in_channels": 3, + "initializer_range": 0.02, + "intermediate_size": 4096, + "model_type": "qwen3_5", + "num_heads": 16, + "num_position_embeddings": 2304, + "out_hidden_size": 2560, + "patch_size": 16, + "spatial_merge_size": 2, + "temporal_patch_size": 2 + }, + "vision_end_token_id": 248054, + "vision_start_token_id": 248053 +} diff --git a/tests/cua_s1/vision_loader.rs b/tests/cua_s1/vision_loader.rs new file mode 100644 index 0000000..8de73e2 --- /dev/null +++ b/tests/cua_s1/vision_loader.rs @@ -0,0 +1,466 @@ +use omni_cua_s1_native::vision::VisionCheckpoint; +use safetensors::Dtype; +use serde_json::{Value, json}; +use std::{ + collections::BTreeMap, + fs::{self, File}, + io::{Seek, SeekFrom, Write}, + path::{Path, PathBuf}, + sync::atomic::{AtomicUsize, Ordering}, +}; + +const BASE: &str = "model.visual."; +const ADAPTER: &str = "base_model.model.model.visual."; +type Inventory = BTreeMap; +fn inventory(adapter: bool) -> Inventory { + serde_json::from_str(if adapter { + include_str!("fixtures/vision/adapter-tensors.json") + } else { + include_str!("fixtures/vision/base-tensors.json") + }) + .unwrap() +} +fn sparse(path: &Path, tensors: &Inventory) { + let mut offset = 0u64; + let mut header = serde_json::Map::new(); + for (name, tensor) in tensors { + let bytes = tensor["shape"] + .as_array() + .unwrap() + .iter() + .map(|n| n.as_u64().unwrap()) + .product::() + * if tensor["dtype"] == "BF16" { 2 } else { 4 }; + header.insert(name.clone(), json!({"dtype":tensor["dtype"],"shape":tensor["shape"],"data_offsets":[offset, offset+bytes]})); + offset += bytes; + } + let header = serde_json::to_vec(&header).unwrap(); + let mut file = File::create(path).unwrap(); + file.write_all(&(header.len() as u64).to_le_bytes()) + .unwrap(); + file.write_all(&header).unwrap(); + file.set_len(8 + header.len() as u64 + offset).unwrap(); + for tensor in serde_json::from_slice::(&header) + .unwrap() + .as_object() + .unwrap() + .values() + { + file.seek(SeekFrom::Start( + 8 + header.len() as u64 + tensor["data_offsets"][0].as_u64().unwrap(), + )) + .unwrap(); + file.write_all(if tensor["dtype"] == "BF16" { + &[0x80, 0x3f] + } else { + &[0, 0, 0x80, 0x3f] + }) + .unwrap(); + } +} +struct Fixture { + root: PathBuf, + base: PathBuf, + adapter: PathBuf, +} +impl Fixture { + fn new() -> Self { + static NEXT: AtomicUsize = AtomicUsize::new(0); + let root = std::env::temp_dir().join(format!( + "cua-vision-{}-{}", + std::process::id(), + NEXT.fetch_add(1, Ordering::Relaxed) + )); + let base = root.join("base"); + let adapter = root.join("adapter"); + fs::create_dir_all(&base).unwrap(); + fs::create_dir_all(&adapter).unwrap(); + fs::write( + base.join("config.json"), + include_str!("fixtures/vision/config.json"), + ) + .unwrap(); + fs::write( + adapter.join("adapter_config.json"), + include_str!("fixtures/vision/adapter_config.json"), + ) + .unwrap(); + sparse(&base.join("model.safetensors"), &inventory(false)); + sparse(&adapter.join("adapter_model.safetensors"), &inventory(true)); + Self { + root, + base, + adapter, + } + } + fn load(&self) -> anyhow::Result { + VisionCheckpoint::load(&self.base, &self.adapter) + } + fn reject(&self, expected: &str) { + let error = match self.load() { + Ok(_) => panic!("accepted invalid checkpoint: {expected}"), + Err(error) => format!("{error:#}"), + }; + assert!( + error.contains(expected), + "expected {expected:?}, got {error}" + ); + } + fn config(&self, adapter: bool, key: &str, value: Value) { + let path = if adapter { + self.adapter.join("adapter_config.json") + } else { + self.base.join("config.json") + }; + let mut config: Value = serde_json::from_str(if adapter { + include_str!("fixtures/vision/adapter_config.json") + } else { + include_str!("fixtures/vision/config.json") + }) + .unwrap(); + if adapter { + config[key] = value; + } else if key == "text_hidden" { + config["text_config"]["hidden_size"] = value; + } else { + config["vision_config"][key] = value; + } + fs::write(path, serde_json::to_vec(&config).unwrap()).unwrap(); + } + fn shards(&self) -> BTreeMap { + let (a, b): (Inventory, Inventory) = inventory(false) + .into_iter() + .partition(|(n, _)| n.contains("blocks.")); + sparse(&self.base.join("a.safetensors"), &a); + sparse(&self.base.join("b.safetensors"), &b); + let mut map: BTreeMap = a + .keys() + .map(|n| (n.clone(), "a.safetensors".into())) + .chain(b.keys().map(|n| (n.clone(), "b.safetensors".into()))) + .collect(); + map.insert( + "model.language_model.weight".into(), + "absent-language.safetensors".into(), + ); + self.index(&map); + fs::remove_file(self.base.join("model.safetensors")).unwrap(); + map + } + fn index(&self, map: &BTreeMap) { + fs::write( + self.base.join("model.safetensors.index.json"), + serde_json::to_vec(&json!({"weight_map":map})).unwrap(), + ) + .unwrap(); + } +} +impl Drop for Fixture { + fn drop(&mut self) { + let _ = fs::remove_dir_all(&self.root); + } +} +fn check(checkpoint: &VisionCheckpoint) { + assert_eq!(checkpoint.config().depth, 24); + assert_eq!(checkpoint.config().hidden_size, 1024); + assert_eq!(checkpoint.adapter().rank, 16); + assert_eq!(checkpoint.adapter().alpha, 32); + assert_eq!(checkpoint.adapter().scale(), 2.0); + for adapter in [false, true] { + let expected = inventory(adapter); + let names: Vec<_> = if adapter { + checkpoint.adapter_names().collect() + } else { + checkpoint.base_names().collect() + }; + assert_eq!(names.len(), if adapter { 100 } else { 297 }); + for (name, spec) in expected { + let tensor = if adapter { + checkpoint.adapter_tensor(&name) + } else { + checkpoint.base_tensor(&name) + } + .unwrap(); + assert_eq!( + tensor.shape(), + serde_json::from_value::>(spec["shape"].clone()).unwrap() + ); + assert_eq!( + tensor.dtype(), + if adapter { Dtype::F32 } else { Dtype::BF16 } + ); + assert_eq!( + &tensor.data()[..if adapter { 4 } else { 2 }], + if adapter { + &[0, 0, 0x80, 0x3f][..] + } else { + &[0x80, 0x3f][..] + } + ); + } + } + assert!(checkpoint.base_tensor("missing").is_err()); +} +#[test] +fn valid_single_matches_independent_inventory() { + let f = Fixture::new(); + check(&f.load().unwrap()); +} +#[test] +fn valid_shards_skip_language_only_files() { + let f = Fixture::new(); + f.shards(); + check(&f.load().unwrap()); +} +#[test] +fn ignores_language_tensors() { + let f = Fixture::new(); + for adapter in [false, true] { + let mut inv = inventory(adapter); + inv.insert( + "unrelated.language.weight".into(), + json!({"shape":[1],"dtype":"F32"}), + ); + sparse( + &if adapter { + f.adapter.join("adapter_model.safetensors") + } else { + f.base.join("model.safetensors") + }, + &inv, + ); + } + check(&f.load().unwrap()); +} +#[test] +fn rejects_missing_wrong_and_extra_tensors() { + let f = Fixture::new(); + for adapter in [false, true] { + let path = if adapter { + f.adapter.join("adapter_model.safetensors") + } else { + f.base.join("model.safetensors") + }; + let original = inventory(adapter); + let name = original.keys().next().unwrap().clone(); + let mut inv = original.clone(); + inv.remove(&name); + sparse(&path, &inv); + f.reject("missing"); + let mut inv = original.clone(); + inv.get_mut(&name).unwrap()["shape"] = json!([1]); + sparse(&path, &inv); + f.reject("shape"); + let mut inv = original.clone(); + inv.get_mut(&name).unwrap()["dtype"] = json!(if adapter { "BF16" } else { "F32" }); + sparse(&path, &inv); + f.reject("dtype"); + let mut inv = original.clone(); + inv.insert( + format!("{}unexpected.weight", if adapter { ADAPTER } else { BASE }), + json!({"shape":[1],"dtype":"F32"}), + ); + sparse(&path, &inv); + f.reject("unexpected"); + sparse(&path, &original); + } +} +#[test] +fn rejects_truncated_and_malformed_safetensors() { + let f = Fixture::new(); + for adapter in [false, true] { + let path = if adapter { + f.adapter.join("adapter_model.safetensors") + } else { + f.base.join("model.safetensors") + }; + let len = fs::metadata(&path).unwrap().len(); + File::options() + .write(true) + .open(&path) + .unwrap() + .set_len(len - 1) + .unwrap(); + f.reject("safetensors"); + fs::write(&path, b"invalid").unwrap(); + f.reject("safetensors"); + sparse(&path, &inventory(adapter)); + } +} +#[test] +fn rejects_index_missing_extra_misrouting_and_duplicates() { + let f = Fixture::new(); + let original = f.shards(); + let name = inventory(false).keys().next().unwrap().clone(); + let mut map = original.clone(); + map.remove(&name); + f.index(&map); + f.reject("missing"); + let mut map = original.clone(); + map.insert(format!("{BASE}unexpected"), "a.safetensors".into()); + f.index(&map); + f.reject("unexpected"); + let mut map = original.clone(); + map.insert(name.clone(), "b.safetensors".into()); + f.index(&map); + f.reject("index"); + f.index(&original); + let mut b: Inventory = inventory(false) + .into_iter() + .filter(|(n, _)| !n.contains("blocks.")) + .collect(); + b.insert(name.clone(), inventory(false)[&name].clone()); + sparse(&f.base.join("b.safetensors"), &b); + f.reject("index"); + fs::write( + f.base.join("model.safetensors.index.json"), + format!("{{\"weight_map\":{{\"{name}\":\"a.safetensors\",\"{name}\":\"b.safetensors\"}}}}"), + ) + .unwrap(); + f.reject("duplicate"); +} +#[test] +fn rejects_index_paths_outside_directory() { + let f = Fixture::new(); + let original = f.shards(); + let name = inventory(false).keys().next().unwrap().clone(); + for path in ["../outside.safetensors", "/tmp/outside.safetensors"] { + let mut map = original.clone(); + map.insert(name.clone(), path.into()); + f.index(&map); + f.reject("checkpoint directory"); + } +} +#[cfg(unix)] +#[test] +fn rejects_symlink_escape() { + use std::os::unix::fs::symlink; + let f = Fixture::new(); + fs::rename( + f.base.join("model.safetensors"), + f.root.join("outside.safetensors"), + ) + .unwrap(); + symlink( + f.root.join("outside.safetensors"), + f.base.join("model.safetensors"), + ) + .unwrap(); + f.reject("checkpoint directory"); +} +#[test] +fn validates_full_vision_configuration() { + let f = Fixture::new(); + for (key, value) in [ + ("depth", json!(23)), + ("hidden_size", json!(512)), + ("intermediate_size", json!(2048)), + ("num_heads", json!(8)), + ("num_position_embeddings", json!(1024)), + ("out_hidden_size", json!(2048)), + ("in_channels", json!(1)), + ("patch_size", json!(14)), + ("temporal_patch_size", json!(1)), + ("spatial_merge_size", json!(1)), + ("hidden_act", json!("gelu")), + ("deepstack_visual_indexes", json!([1])), + ("text_hidden", json!(2048)), + ] { + f.config(false, key, value); + f.reject("config"); + } +} +#[test] +fn rejects_adapter_math_and_target_changes() { + let f = Fixture::new(); + for (key, value) in [ + ("r", json!(8)), + ("lora_alpha", json!(16)), + ("bias", json!("all")), + ("lora_bias", json!(true)), + ("use_dora", json!(true)), + ("use_rslora", json!(true)), + ("use_qalora", json!(true)), + ("fan_in_fan_out", json!(true)), + ("rank_pattern", json!({"linear_fc1":8})), + ("alpha_pattern", json!({"linear_fc1":16})), + ("exclude_modules", json!(["linear_fc1"])), + ("target_modules", json!(["up_proj", "down_proj"])), + ("modules_to_save", json!(["merger"])), + ("layers_to_transform", json!([1])), + ("target_parameters", json!(["weight"])), + ("layer_replication", json!([[0, 1]])), + ("alora_invocation_tokens", json!([1])), + ("peft_type", json!("IA3")), + ] { + f.config(true, key, value); + f.reject("adapter config"); + } +} +#[test] +fn requires_multimodal_adapter() { + let f = Fixture::new(); + fs::remove_file(f.adapter.join("adapter_model.safetensors")).unwrap(); + f.reject("adapter_model.safetensors"); + sparse( + &f.adapter.join("adapter_model.safetensors"), + &BTreeMap::from([( + "base_model.model.model.language_model.weight".into(), + json!({"shape":[1],"dtype":"F32"}), + )]), + ); + f.reject("missing"); +} + +#[test] +fn rejects_duplicate_visual_keys_in_safetensors_header() { + use std::io::Read; + let f = Fixture::new(); + let path = f.base.join("model.safetensors"); + let mut file = File::open(&path).unwrap(); + let total = file.metadata().unwrap().len(); + let mut size = [0; 8]; + file.read_exact(&mut size).unwrap(); + let old_len = u64::from_le_bytes(size); + let mut header = vec![0; old_len as usize]; + file.read_exact(&mut header).unwrap(); + drop(file); + let entries: BTreeMap = serde_json::from_slice(&header).unwrap(); + let (name, info) = entries.first_key_value().unwrap(); + let mut duplicate = String::from_utf8(header).unwrap(); + duplicate.pop(); + duplicate.push_str(&format!( + ",{}:{info}}}", + serde_json::to_string(name).unwrap() + )); + let mut file = File::create(&path).unwrap(); + file.write_all(&(duplicate.len() as u64).to_le_bytes()) + .unwrap(); + file.write_all(duplicate.as_bytes()).unwrap(); + file.set_len(total - old_len + duplicate.len() as u64) + .unwrap(); + drop(file); + f.reject("duplicate"); +} + +#[test] +fn rejects_overflowing_payload_offsets_without_panicking() { + let f = Fixture::new(); + let elements = usize::MAX / 8; + let mut header = serde_json::Map::new(); + for i in 0..8 { + header.insert( + format!("unrelated.language.{i}.weight"), + json!({ + "dtype": "U8", "shape": [elements], + "data_offsets": [i * elements, (i + 1) * elements] + }), + ); + } + let header = serde_json::to_vec(&header).unwrap(); + let mut file = File::create(f.base.join("model.safetensors")).unwrap(); + file.write_all(&(header.len() as u64).to_le_bytes()) + .unwrap(); + file.write_all(&header).unwrap(); + drop(file); + f.reject("safetensors"); +} From 591097e46c7dcbe2f4d7c8b4dc1a81ac55324ade Mon Sep 17 00:00:00 2001 From: levius <2114377220@qq.com> Date: Fri, 2 Oct 2026 12:02:44 +0800 Subject: [PATCH 3/5] cua_s1: preprocess RGB images for native vision on CPU --- Cargo.lock | 1 + recipe/cua_s1/native.md | 3 + recipe/cua_s1/native_image_preprocess.md | 78 +++++ src/models/cua_s1/native/Cargo.toml | 7 + .../cua_s1/native/THIRD_PARTY_NOTICES.md | 24 ++ .../native/examples/preprocess_image.rs | 75 +++++ src/models/cua_s1/native/licenses/APACHE-2.0 | 203 ++++++++++++ .../cua_s1/native/licenses/PILLOW-LICENSE | 30 ++ .../cua_s1/native/licenses/PYTORCH-LICENSE | 84 +++++ .../cua_s1/native/src/image_preprocess.rs | 246 ++++++++++++++ src/models/cua_s1/native/src/lib.rs | 1 + .../fixtures/image_preprocess/README.md | 30 ++ .../fixtures/image_preprocess/generate.py | 119 +++++++ .../fixtures/image_preprocess/manifest.json | 311 ++++++++++++++++++ .../image_preprocess/preprocessor_config.json | 21 ++ tests/cua_s1/image_preprocess.rs | 223 +++++++++++++ 16 files changed, 1456 insertions(+) create mode 100644 recipe/cua_s1/native_image_preprocess.md create mode 100644 src/models/cua_s1/native/THIRD_PARTY_NOTICES.md create mode 100644 src/models/cua_s1/native/examples/preprocess_image.rs create mode 100644 src/models/cua_s1/native/licenses/APACHE-2.0 create mode 100644 src/models/cua_s1/native/licenses/PILLOW-LICENSE create mode 100644 src/models/cua_s1/native/licenses/PYTORCH-LICENSE create mode 100644 src/models/cua_s1/native/src/image_preprocess.rs create mode 100644 tests/cua_s1/fixtures/image_preprocess/README.md create mode 100644 tests/cua_s1/fixtures/image_preprocess/generate.py create mode 100644 tests/cua_s1/fixtures/image_preprocess/manifest.json create mode 100644 tests/cua_s1/fixtures/image_preprocess/preprocessor_config.json create mode 100644 tests/cua_s1/image_preprocess.rs diff --git a/Cargo.lock b/Cargo.lock index 4b9adda..fad6575 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -954,6 +954,7 @@ dependencies = [ "safetensors 0.8.0", "serde", "serde_json", + "sha2", "tokenizers", "tokio", ] diff --git a/recipe/cua_s1/native.md b/recipe/cua_s1/native.md index 018feff..8b1d839 100644 --- a/recipe/cua_s1/native.md +++ b/recipe/cua_s1/native.md @@ -41,3 +41,6 @@ cargo test -p omni-cua-s1-native CUA_S1_CUDA_LIB=$PWD/target/release/libqwen3_5_cuda.so \ cargo test --release -p omni-cua-s1-native --test kernels -- --ignored ``` + +For the separate decoded-RGB8 CPU preprocessing API and example, see +[native image preprocessing](native_image_preprocess.md). diff --git a/recipe/cua_s1/native_image_preprocess.md b/recipe/cua_s1/native_image_preprocess.md new file mode 100644 index 0000000..eb63043 --- /dev/null +++ b/recipe/cua_s1/native_image_preprocess.md @@ -0,0 +1,78 @@ +# Native CPU image preprocessing + +The native crate exposes `image_preprocess::preprocess_rgb8(width, height, rgb)` +for **already decoded, interleaved RGB8** data. It prepares the image tensor for +the fixed Qwen3.5-4B / Cua-S1 4B processor. It does not decode PNG/JPEG, fetch +URLs, handle HTTP requests, run the vision encoder, or use a GPU. The existing +native text worker remains separate. + +The input must contain exactly `width * height * 3` bytes, in row-major RGB +order. Both dimensions must be nonzero and at most 2048; the area must be at +most 1,048,576 pixels and the aspect ratio at most 200. The library checks +geometry, lengths, and allocation arithmetic before creating image buffers. +These are input limits; smart resize can produce a side longer than 2048 for +very narrow inputs. + +The fixed processor uses a factor of 32, minimum area 65,536 and maximum area +16,777,216. Smart resize follows Python ties-to-even rounding and floating-point +square-root scaling with floor/ceil. Resampling matches the CPU torchvision +uint8 bicubic antialias path: Keys cubic coefficient `a = -0.5`, float64 weights, +per-axis int16 fixed-point coefficients, horizontal then vertical passes, and +rounding/clamping to uint8 after each pass. Unchanged axes bypass resampling. +The maximum-area downscale branch is retained for parity with the processor, +although the smaller input cap makes it unreachable through this API. + +`ProcessedImage` contains: + +- `pixel_values: Vec`, contiguous `[patches, 1536]` values normalized as + `(pixel - 127.5) / 127.5` using float32 operations. +- `image_grid_thw: [usize; 3]`, equal to `[1, resized_height / 16, resized_width / 16]`. +- `resized_width` and `resized_height`. +- `image_tokens()`, the patch count divided by four for the 2×2 spatial merge. + +Packing order is `block_y, block_x, merge_y (2), merge_x (2), channel (3), +temporal repeat (2), patch_y (16), patch_x (16)`. Each single image is repeated +across the two temporal positions. There is no video input support. + +## Run without a GPU + +From the repository root, provide a raw RGB8 file and its dimensions: + +```sh +cargo run --locked -p omni-cua-s1-native --example preprocess_image -- \ + 256 256 image.rgb pixel_values.f32 +``` + +The example prints the tensor shape, grid, resized dimensions and image token +count. The optional fourth argument writes every output float as little-endian +float32, with no header. It validates argument count, decimal dimensions and +exact file length, and bounds the input read. No model weights, Python, CUDA +library or image decoder are needed for this command. + +## Reference and validation + +Reference hashes are produced by the actual Hugging Face `AutoImageProcessor` +on CPU, with Python 3.12 and these exact package pins: Transformers 5.17.0, +PyTorch 2.14.0, torchvision 0.29.0, NumPy 2.5.3 and Pillow 11.3.0. The processor +configuration is the Qwen3.5-4B file at revision +`851bf6e806efd8d0a36b00ddf55e13ccb7b8cd0a` (see the fixture manifest for its URL +and SHA-256). Fixtures and regeneration instructions live in +[`tests/cua_s1/fixtures/image_preprocess/`](../../tests/cua_s1/fixtures/image_preprocess/). + +```sh +cargo test --locked -p omni-cua-s1-native --test image_preprocess +cargo test --release --locked -p omni-cua-s1-native --test image_preprocess +``` + +The tests compare input hashes, shape/grid metadata and SHA-256 of **every +little-endian float32 output byte** for 14 deterministic images. Cases include +tiny images, noise, ramps, checkerboards, ties-to-even dimensions, unchanged +axes, one- and two-axis resizes, extreme aspect ratios and the input area cap. +Additional checks cover all RGB byte values, channel/temporal/patch order, +constant images, inclusive limits and malformed geometry/buffers. This +validates the pinned CPU preprocessing behavior; it does not establish CUDA +preprocessing, image decoding, vision inference, or end-to-end model parity. + +The implementation adapts upstream algorithms; retained attributions and +license texts are in +[`THIRD_PARTY_NOTICES.md`](../../src/models/cua_s1/native/THIRD_PARTY_NOTICES.md). diff --git a/src/models/cua_s1/native/Cargo.toml b/src/models/cua_s1/native/Cargo.toml index 2e8caf6..0861b42 100644 --- a/src/models/cua_s1/native/Cargo.toml +++ b/src/models/cua_s1/native/Cargo.toml @@ -22,3 +22,10 @@ serde_json = { version = "1.0.149", features = ["float_roundtrip", "preserve_ord # the onig regex backend, as in the Python tokenizers wheel tokenizers = { version = "=0.22.2", default-features = false, features = ["onig"] } tokio = { version = "1.49.0", features = ["macros", "net", "rt-multi-thread", "sync"] } + +[[test]] +name = "image_preprocess" +path = "../../../../tests/cua_s1/image_preprocess.rs" + +[dev-dependencies] +sha2 = "0.10" diff --git a/src/models/cua_s1/native/THIRD_PARTY_NOTICES.md b/src/models/cua_s1/native/THIRD_PARTY_NOTICES.md new file mode 100644 index 0000000..54b788e --- /dev/null +++ b/src/models/cua_s1/native/THIRD_PARTY_NOTICES.md @@ -0,0 +1,24 @@ +# Image preprocessing attributions + +`src/image_preprocess.rs` is a Rust adaptation of the following algorithms. +It is modified for decoded interleaved RGB8 input, fixed Cua-S1 4B settings, +bounded dimensions, standard-library buffers, and a standalone CPU API. + +- PyTorch 2.14.0, `aten/src/ATen/native/cpu/UpSampleKernel.cpp` + (`_compute_indices_min_size_weights_aa`, `_compute_index_ranges_int16_weights`, + and the separable uint8 horizontal/vertical loops), plus the cubic polynomial + helpers in `aten/src/ATen/native/UpSample.h`. See [PyTorch license](licenses/PYTORCH-LICENSE) + for the retained copyright notices, redistribution conditions, and disclaimer. +- PyTorch's bicubic filter credits Pillow's `src/libImaging/Resample.c`. + The retained PIL/Pillow notice is in [Pillow license](licenses/PILLOW-LICENSE). +- Transformers 5.17.0, + `src/transformers/models/qwen2_vl/image_processing_qwen2_vl.py` (smart resize + and patch ordering) and `src/transformers/image_processing_backends.py` + (fused normalization). Copyright 2024 The Qwen team, Alibaba Group and the + HuggingFace Inc. team. All rights reserved. The backend file is + Copyright 2025 The HuggingFace Inc. team. Licensed under the + [Apache License, Version 2.0](licenses/APACHE-2.0). + +No upstream runtime or image decoder is linked by this module. The above +notices and license texts must accompany redistributed adaptations as required +by their respective licenses. diff --git a/src/models/cua_s1/native/examples/preprocess_image.rs b/src/models/cua_s1/native/examples/preprocess_image.rs new file mode 100644 index 0000000..310c661 --- /dev/null +++ b/src/models/cua_s1/native/examples/preprocess_image.rs @@ -0,0 +1,75 @@ +//! CPU-only RGB8 preprocessing; run from the repository root with: +//! cargo run -p omni-cua-s1-native --example preprocess_image -- 256 256 image.rgb + +use std::{ + env, + fs::File, + io::{BufWriter, Read, Write}, +}; + +use anyhow::{Context, Result, ensure}; +use omni_cua_s1_native::image_preprocess::preprocess_rgb8; + +fn main() -> Result<()> { + let args: Vec<_> = env::args_os().skip(1).collect(); + ensure!( + args.len() == 3 || args.len() == 4, + "usage: preprocess_image WIDTH HEIGHT RAW_RGB_PATH [OUTPUT_F32_PATH]" + ); + let dimension = |index: usize| -> Result { + let text = args[index] + .to_str() + .context("dimensions must be UTF-8 decimal integers")?; + ensure!( + !text.is_empty() && text.bytes().all(|byte| byte.is_ascii_digit()), + "dimensions must be unsigned decimal integers" + ); + text.parse().context("dimension is too large") + }; + let width = dimension(0)?; + let height = dimension(1)?; + // Bound the file read before allocating. The library validates the complete + // contract too; these checks keep malformed CLI inputs cheap to reject. + ensure!( + width > 0 && height > 0 && width <= 2048 && height <= 2048, + "dimensions must be in 1..=2048" + ); + let area = width.checked_mul(height).context("image area overflow")?; + ensure!( + area <= 1_048_576, + "image area must not exceed 1048576 pixels" + ); + ensure!( + width.max(height) <= width.min(height) * 200, + "image aspect ratio must not exceed 200" + ); + let expected = area.checked_mul(3).context("RGB length overflow")?; + let mut rgb = Vec::with_capacity(expected + 1); + File::open(&args[2]) + .context("opening RGB input")? + .take((expected + 1) as u64) + .read_to_end(&mut rgb) + .context("reading RGB input")?; + ensure!( + rgb.len() == expected, + "RGB input must contain exactly {expected} bytes" + ); + let image = preprocess_rgb8(width, height, &rgb)?; + println!( + "pixel_values shape: [{}, 1536]", + image.pixel_values.len() / 1536 + ); + println!("image_grid_thw: {:?}", image.image_grid_thw); + println!("resized: {}x{}", image.resized_width, image.resized_height); + println!("image_tokens: {}", image.image_tokens()); + if let Some(path) = args.get(3) { + let mut output = BufWriter::new(File::create(path).context("creating float32 output")?); + for value in image.pixel_values { + output + .write_all(&value.to_le_bytes()) + .context("writing float32 output")?; + } + output.flush().context("flushing float32 output")?; + } + Ok(()) +} diff --git a/src/models/cua_s1/native/licenses/APACHE-2.0 b/src/models/cua_s1/native/licenses/APACHE-2.0 new file mode 100644 index 0000000..68b7d66 --- /dev/null +++ b/src/models/cua_s1/native/licenses/APACHE-2.0 @@ -0,0 +1,203 @@ +Copyright 2018- The Hugging Face team. All rights reserved. + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + 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. diff --git a/src/models/cua_s1/native/licenses/PILLOW-LICENSE b/src/models/cua_s1/native/licenses/PILLOW-LICENSE new file mode 100644 index 0000000..10dd42d --- /dev/null +++ b/src/models/cua_s1/native/licenses/PILLOW-LICENSE @@ -0,0 +1,30 @@ +The Python Imaging Library (PIL) is + + Copyright © 1997-2011 by Secret Labs AB + Copyright © 1995-2011 by Fredrik Lundh and contributors + +Pillow is the friendly PIL fork. It is + + Copyright © 2010 by Jeffrey A. Clark and contributors + +Like PIL, Pillow is licensed under the open source MIT-CMU License: + +By obtaining, using, and/or copying this software and/or its associated +documentation, you agree that you have read, understood, and will comply +with the following terms and conditions: + +Permission to use, copy, modify and distribute this software and its +documentation for any purpose and without fee is hereby granted, +provided that the above copyright notice appears in all copies, and that +both that copyright notice and this permission notice appear in supporting +documentation, and that the name of Secret Labs AB or the author not be +used in advertising or publicity pertaining to distribution of the software +without specific, written prior permission. + +SECRET LABS AB AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH REGARD TO THIS +SOFTWARE, INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS. +IN NO EVENT SHALL SECRET LABS AB OR THE AUTHOR BE LIABLE FOR ANY SPECIAL, +INDIRECT OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM +LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE +OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR +PERFORMANCE OF THIS SOFTWARE. diff --git a/src/models/cua_s1/native/licenses/PYTORCH-LICENSE b/src/models/cua_s1/native/licenses/PYTORCH-LICENSE new file mode 100644 index 0000000..c23172f --- /dev/null +++ b/src/models/cua_s1/native/licenses/PYTORCH-LICENSE @@ -0,0 +1,84 @@ +From PyTorch: + +Copyright (c) 2016- Facebook, Inc (Adam Paszke) +Copyright (c) 2014- Facebook, Inc (Soumith Chintala) +Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert) +Copyright (c) 2012-2014 Deepmind Technologies (Koray Kavukcuoglu) +Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu) +Copyright (c) 2011-2013 NYU (Clement Farabet) +Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston) +Copyright (c) 2006 Idiap Research Institute (Samy Bengio) +Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz) + +From Caffe2: + +Copyright (c) 2016-present, Facebook Inc. All rights reserved. + +All contributions by Facebook: +Copyright (c) 2016 Facebook Inc. + +All contributions by Google: +Copyright (c) 2015 Google Inc. +All rights reserved. + +All contributions by Yangqing Jia: +Copyright (c) 2015 Yangqing Jia +All rights reserved. + +All contributions by Kakao Brain: +Copyright 2019-2020 Kakao Brain + +All contributions by Cruise LLC: +Copyright (c) 2022 Cruise LLC. +All rights reserved. + +All contributions by Tri Dao: +Copyright (c) 2024 Tri Dao. +All rights reserved. + +All contributions by Arm: +Copyright (c) 2021, 2023-2025 Arm Limited and/or its affiliates + +All contributions from Caffe: +Copyright(c) 2013, 2014, 2015, the respective contributors +All rights reserved. + +All other contributions: +Copyright(c) 2015, 2016 the respective contributors +All rights reserved. + +Caffe2 uses a copyright model similar to Caffe: each contributor holds +copyright over their contributions to Caffe2. The project versioning records +all such contribution and copyright details. If a contributor wants to further +mark their specific copyright on a particular contribution, they should +indicate their copyright solely in the commit message of the change when it is +committed. + +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright + notice, this list of conditions and the following disclaimer in the + documentation and/or other materials provided with the distribution. + +3. Neither the names of Facebook, Deepmind Technologies, NYU, NEC Laboratories America + and IDIAP Research Institute nor the names of its contributors may be + used to endorse or promote products derived from this software without + specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE +LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE +POSSIBILITY OF SUCH DAMAGE. diff --git a/src/models/cua_s1/native/src/image_preprocess.rs b/src/models/cua_s1/native/src/image_preprocess.rs new file mode 100644 index 0000000..9009fb2 --- /dev/null +++ b/src/models/cua_s1/native/src/image_preprocess.rs @@ -0,0 +1,246 @@ +//! CPU preprocessing for the fixed Qwen3.5-4B / Cua-S1 image processor. +//! +//! Input is already decoded, interleaved RGB8. No image codecs, GPU, model +//! weights, or HTTP handling are involved. The output is row-major +//! `[patches, 1536]`, ready for a separate vision encoder. +//! +//! The resize algorithm is a Rust adaptation of PyTorch's CPU uint8 bicubic +//! antialias implementation (which credits Pillow). Smart resize and packing +//! follow Transformers' Qwen2VLImageProcessor. See `../THIRD_PARTY_NOTICES.md`. + +use std::borrow::Cow; + +use anyhow::{Result, ensure}; + +const PATCH_SIZE: usize = 16; +const MERGE_SIZE: usize = 2; +const FACTOR: usize = PATCH_SIZE * MERGE_SIZE; +const PATCH_VALUES: usize = 3 * 2 * PATCH_SIZE * PATCH_SIZE; +const MIN_PIXELS: usize = 65_536; +const MAX_PIXELS: usize = 16_777_216; + +/// Normalized image patches and the spatial metadata used by the vision model. +#[derive(Debug)] +pub struct ProcessedImage { + /// Contiguous row-major `[patches, 1536]` float32 values. + pub pixel_values: Vec, + /// `[1, resized_height / 16, resized_width / 16]`. + pub image_grid_thw: [usize; 3], + pub resized_width: usize, + pub resized_height: usize, +} + +impl ProcessedImage { + /// Number of image tokens after the model's 2-by-2 spatial merge. + pub fn image_tokens(&self) -> usize { + self.pixel_values.len() / PATCH_VALUES / (MERGE_SIZE * MERGE_SIZE) + } +} + +/// Preprocess a decoded RGB8 image using the fixed 4B processor settings. +/// +/// Rejects empty dimensions, sides over 2048, area over 1,048,576 pixels, +/// aspect ratios over 200, and buffers whose length is not `width * height * 3`. +/// Geometry and buffer arithmetic are checked before allocating image buffers. +pub fn preprocess_rgb8(width: usize, height: usize, rgb: &[u8]) -> Result { + ensure!(width > 0 && height > 0, "image dimensions must be nonzero"); + ensure!( + width <= 2048 && height <= 2048, + "image sides must not exceed 2048" + ); + let area = width + .checked_mul(height) + .ok_or_else(|| anyhow::anyhow!("image area overflow"))?; + ensure!( + area <= 1_048_576, + "image area must not exceed 1048576 pixels" + ); + ensure!( + width.max(height) <= width.min(height) * 200, + "image aspect ratio must not exceed 200" + ); + let input_len = area + .checked_mul(3) + .ok_or_else(|| anyhow::anyhow!("RGB buffer length overflow"))?; + ensure!( + rgb.len() == input_len, + "RGB buffer length must be {input_len}, got {}", + rgb.len() + ); + + let (resized_width, resized_height) = smart_resize(width, height); + let resized_area = resized_width + .checked_mul(resized_height) + .ok_or_else(|| anyhow::anyhow!("resized area overflow"))?; + let resized_len = resized_area + .checked_mul(3) + .ok_or_else(|| anyhow::anyhow!("resized buffer length overflow"))?; + let horizontal_len = resized_width + .checked_mul(height) + .and_then(|area| area.checked_mul(3)) + .ok_or_else(|| anyhow::anyhow!("horizontal buffer length overflow"))?; + let output_len = resized_area + .checked_mul(6) + .ok_or_else(|| anyhow::anyhow!("patch buffer length overflow"))?; + output_len + .checked_mul(std::mem::size_of::()) + .ok_or_else(|| anyhow::anyhow!("patch buffer byte length overflow"))?; + + let mut resized = Cow::Borrowed(rgb); + if resized_width != width { + let axis = AxisWeights::new(width, resized_width); + let mut horizontal = vec![0; horizontal_len]; + for y in 0..height { + for (x, kernel) in axis.kernels.iter().enumerate() { + for channel in 0..3 { + horizontal[(y * resized_width + x) * 3 + channel] = + axis.apply(kernel, |source_x| rgb[(y * width + source_x) * 3 + channel]); + } + } + } + resized = Cow::Owned(horizontal); + } + if resized_height != height { + let axis = AxisWeights::new(height, resized_height); + let mut vertical = vec![0; resized_len]; + for (y, kernel) in axis.kernels.iter().enumerate() { + for x in 0..resized_width { + for channel in 0..3 { + vertical[(y * resized_width + x) * 3 + channel] = axis + .apply(kernel, |source_y| { + resized[(source_y * resized_width + x) * 3 + channel] + }); + } + } + } + resized = Cow::Owned(vertical); + } + + let mut pixel_values = Vec::with_capacity(output_len); + for block_y in 0..resized_height / FACTOR { + for block_x in 0..resized_width / FACTOR { + for merge_y in 0..MERGE_SIZE { + for merge_x in 0..MERGE_SIZE { + for channel in 0..3 { + for _temporal in 0..2 { + for patch_y in 0..PATCH_SIZE { + for patch_x in 0..PATCH_SIZE { + let y = block_y * FACTOR + merge_y * PATCH_SIZE + patch_y; + let x = block_x * FACTOR + merge_x * PATCH_SIZE + patch_x; + let pixel = resized[(y * resized_width + x) * 3 + channel]; + // Match the fused float32 torchvision normalization, + // including its operation order (no reciprocal multiply). + pixel_values.push((f32::from(pixel) - 127.5) / 127.5); + } + } + } + } + } + } + } + } + Ok(ProcessedImage { + pixel_values, + image_grid_thw: [1, resized_height / PATCH_SIZE, resized_width / PATCH_SIZE], + resized_width, + resized_height, + }) +} + +fn smart_resize(width: usize, height: usize) -> (usize, usize) { + // Python round uses ties-to-even; Rust's ordinary round does not. + let mut w = (width as f64 / FACTOR as f64).round_ties_even() as usize * FACTOR; + let mut h = (height as f64 / FACTOR as f64).round_ties_even() as usize * FACTOR; + if w * h > MAX_PIXELS { + let beta = ((width * height) as f64 / MAX_PIXELS as f64).sqrt(); + w = ((width as f64 / beta / FACTOR as f64).floor() as usize * FACTOR).max(FACTOR); + h = ((height as f64 / beta / FACTOR as f64).floor() as usize * FACTOR).max(FACTOR); + } else if w * h < MIN_PIXELS { + let beta = (MIN_PIXELS as f64 / (width * height) as f64).sqrt(); + w = (width as f64 * beta / FACTOR as f64).ceil() as usize * FACTOR; + h = (height as f64 * beta / FACTOR as f64).ceil() as usize * FACTOR; + } + (w, h) +} + +struct Kernel { + start: usize, + weights: Vec, +} + +struct AxisWeights { + kernels: Vec, + precision: u32, +} + +impl AxisWeights { + fn new(input: usize, output: usize) -> Self { + let scale = input as f64 / output as f64; + let support = 2.0 * scale.max(1.0); + let invscale = if scale >= 1.0 { 1.0 / scale } else { 1.0 }; + let max_size = support.ceil() as usize * 2 + 1; + let mut maximum = 0.0_f64; + let mut floating = Vec::with_capacity(output); + for index in 0..output { + let center = scale * (index as f64 + 0.5); + // C++ conversion truncates toward zero before clamping the bounds. + let start = ((center - support + 0.5) as isize).max(0) as usize; + let end = ((center + support + 0.5) as usize).min(input); + let count = end.saturating_sub(start).min(max_size); + let mut weights: Vec = (0..count) + .map(|j| cubic((j as f64 + start as f64 - center + 0.5) * invscale)) + .collect(); + let total: f64 = weights.iter().sum(); + if total != 0.0 { + for weight in &mut weights { + *weight /= total; + maximum = maximum.max(*weight); + } + } + floating.push((start, weights)); + } + // One precision for the whole axis, as in PyTorch's int16 path. + let mut precision = 0; + while precision < 22 { + if (0.5 + maximum * f64::from(1 << (precision + 1))) as i32 >= (1 << 15) { + break; + } + precision += 1; + } + let multiplier = f64::from(1 << precision); + let kernels = floating + .into_iter() + .map(|(start, weights)| Kernel { + start, + weights: weights + .into_iter() + .map(|weight| { + let value = weight * multiplier; + (value + if value < 0.0 { -0.5 } else { 0.5 }) as i16 + }) + .collect(), + }) + .collect(); + Self { kernels, precision } + } + + fn apply(&self, kernel: &Kernel, pixel: impl Fn(usize) -> u8) -> u8 { + let mut accumulator = 1_i32 << (self.precision - 1); + for (offset, &weight) in kernel.weights.iter().enumerate() { + accumulator += i32::from(pixel(kernel.start + offset)) * i32::from(weight); + } + (accumulator >> self.precision).clamp(0, 255) as u8 + } +} + +fn cubic(x: f64) -> f64 { + let x = x.abs(); + const A: f64 = -0.5; + if x < 1.0 { + ((A + 2.0) * x - (A + 3.0)) * x * x + 1.0 + } else if x < 2.0 { + ((A * x - 5.0 * A) * x + 8.0 * A) * x - 4.0 * A + } else { + 0.0 + } +} diff --git a/src/models/cua_s1/native/src/lib.rs b/src/models/cua_s1/native/src/lib.rs index 0d3aa5f..ff12872 100644 --- a/src/models/cua_s1/native/src/lib.rs +++ b/src/models/cua_s1/native/src/lib.rs @@ -5,5 +5,6 @@ pub mod contract; pub mod cuda; pub mod engine; +pub mod image_preprocess; pub mod json; pub mod model; diff --git a/tests/cua_s1/fixtures/image_preprocess/README.md b/tests/cua_s1/fixtures/image_preprocess/README.md new file mode 100644 index 0000000..1eff2a5 --- /dev/null +++ b/tests/cua_s1/fixtures/image_preprocess/README.md @@ -0,0 +1,30 @@ +# Native RGB preprocessing reference fixtures + +`preprocessor_config.json` is from pinned Qwen3.5-4B revision +`851bf6e806efd8d0a36b00ddf55e13ccb7b8cd0a`: +[upstream configuration](https://huggingface.co/Qwen/Qwen3.5-4B/resolve/851bf6e806efd8d0a36b00ddf55e13ccb7b8cd0a/preprocessor_config.json). + +`manifest.json` contains complete FP32 output byte hashes generated by the actual +Transformers `AutoImageProcessor` with this configuration on CPU. It records +package versions, platform, configuration fingerprint, deterministic RGB input +hashes, resized dimensions, grid, patch shape and image token count. No weights +or generated output tensors are committed. Rust tests duplicate only the input +generator and compare every output byte through SHA-256. + +The input generator covers constant RGB, spatial/channel ramps, checkerboards +and seeded xorshift32 noise. Cases include no resize, upsampling, downsampling, +Python ties-to-even dimensions, portrait/wide inputs, the 200:1 aspect boundary +and the request pixel cap. These establish CPU preprocessing parity on the +recorded reference environment; they are not an encoder or GPU accuracy test. + +Regenerate with Python 3.12 and the pinned packages (no GPU or model downloads): + +```sh +python -m pip install torch==2.14.0 torchvision==0.29.0 transformers==5.17.0 \ + numpy==2.5.3 Pillow==11.3.0 +python tests/cua_s1/fixtures/image_preprocess/generate.py +``` + +The generator rejects mismatched package versions. Keep new generated hashes +reviewable against the recorded source; do not update expected values merely +to make a native mismatch disappear. diff --git a/tests/cua_s1/fixtures/image_preprocess/generate.py b/tests/cua_s1/fixtures/image_preprocess/generate.py new file mode 100644 index 0000000..70efc1b --- /dev/null +++ b/tests/cua_s1/fixtures/image_preprocess/generate.py @@ -0,0 +1,119 @@ +"""Regenerate CPU parity hashes using the actual pinned Hugging Face processor. + +Run in a Python 3.12 environment with torch==2.14.0, torchvision==0.29.0, +transformers==5.17.0, numpy==2.5.3 and Pillow==11.3.0. No model weights or GPU. +""" + +import hashlib +import importlib.metadata +import json +import platform +from pathlib import Path + +import numpy as np +import torch +from PIL import Image +from transformers import AutoImageProcessor + +PACKAGES = { + "torch": "2.14.0", + "torchvision": "0.29.0", + "transformers": "5.17.0", + "numpy": "2.5.3", + "Pillow": "11.3.0", +} +CASES = [ + ("single_pixel", 1, 1, "constant", 1), + ("tiny_noise", 17, 19, "noise", 17), + ("aligned_noise", 256, 256, "noise", 123), + ("odd_noise", 319, 241, "noise", 42), + ("tie_down", 272, 256, "noise", 999), + ("tie_up", 304, 256, "noise", 101), + ("both_down", 271, 271, "checker", 1), + ("one_axis_down", 256, 257, "ramp", 1), + ("wide", 640, 320, "ramp", 1), + ("portrait", 319, 641, "noise", 23), + ("aspect_limit", 200, 1, "noise", 47), + ("tall_aspect_limit", 7, 1400, "checker", 1), + ("input_pixel_limit", 2048, 512, "noise", 55), + ("large_round_up", 1600, 600, "noise", 77), +] + + +def pixels(width, height, pattern, seed): + """Only input generation is duplicated in Rust; outputs come from HF.""" + if pattern == "noise": + out = bytearray(width * height * 3) + state = seed + for i in range(len(out)): + state ^= (state << 13) & 0xFFFFFFFF + state ^= state >> 17 + state ^= (state << 5) & 0xFFFFFFFF + out[i] = state & 255 + return bytes(out) + return bytes( + ( + [0, 128, 255][c] + if pattern == "constant" + else (255 if (x + y + c) % 2 else 0) + if pattern == "checker" + else (x * 13 + y * 7 + c * 83) % 256 + ) + for y in range(height) + for x in range(width) + for c in range(3) + ) + + +def evaluate(processor, name, width, height, pattern, seed): + raw = pixels(width, height, pattern, seed) + array = np.frombuffer(raw, dtype=np.uint8).reshape(height, width, 3) + result = processor(images=Image.fromarray(array), return_tensors="pt", device="cpu") + tensor = result["pixel_values"].contiguous() + grid = result["image_grid_thw"][0].tolist() + assert tensor.dtype == torch.float32 and tensor.device.type == "cpu" + return { + "name": name, + "width": width, + "height": height, + "pattern": pattern, + "seed": seed, + "input_sha256": hashlib.sha256(raw).hexdigest(), + "grid": grid, + "shape": list(tensor.shape), + "image_tokens": grid[1] * grid[2] // 4, + "resized_width": grid[2] * 16, + "resized_height": grid[1] * 16, + "output_sha256": hashlib.sha256( + tensor.numpy().astype(" f32 { + (f32::from(pixel) - 127.5) / 127.5 +} + +#[test] +fn tiny_rgb_is_upscaled_and_temporally_repeated() { + let image = preprocess_rgb8(1, 1, &[0, 127, 255]).unwrap(); + assert_eq!((image.resized_width, image.resized_height), (256, 256)); + assert_eq!(image.image_grid_thw, [1, 16, 16]); + assert_eq!(image.image_tokens(), 64); + assert_eq!(image.pixel_values.len(), 256 * 1536); + for patch in image.pixel_values.as_chunks::<1536>().0 { + for (channel, expected) in [0, 127, 255].into_iter().enumerate() { + assert!( + patch[channel * 512..(channel + 1) * 512] + .iter() + .all(|&value| value.to_bits() == normalized(expected).to_bits()) + ); + } + } +} + +#[test] +fn identity_resize_preserves_every_byte_and_patch_merge_order() { + let width = 256; + let height = 256; + let rgb: Vec = (0..height) + .flat_map(|y| (0..width).flat_map(move |x| [x as u8, y as u8, (x ^ y) as u8])) + .collect(); + let image = preprocess_rgb8(width, height, &rgb).unwrap(); + for block_y in 0..8 { + for block_x in 0..8 { + for merge_y in 0..2 { + for merge_x in 0..2 { + let patch = ((block_y * 8 + block_x) * 2 + merge_y) * 2 + merge_x; + for channel in 0..3 { + for temporal in 0..2 { + for py in 0..16 { + for px in 0..16 { + let x = block_x * 32 + merge_x * 16 + px; + let y = block_y * 32 + merge_y * 16 + py; + let index = patch * 1536 + + channel * 512 + + temporal * 256 + + py * 16 + + px; + assert_eq!( + image.pixel_values[index].to_bits(), + normalized(rgb[(y * width + x) * 3 + channel]).to_bits() + ); + } + } + } + } + } + } + } + } +} + +#[test] +fn smart_resize_uses_python_ties_even_rounding() { + for (side, expected) in [(272, 256), (304, 320)] { + let image = preprocess_rgb8(side, side, &vec![128; side * side * 3]).unwrap(); + assert_eq!( + (image.resized_width, image.resized_height), + (expected, expected) + ); + assert!( + image + .pixel_values + .iter() + .all(|&value| value == normalized(128)) + ); + } +} + +#[test] +fn rejects_invalid_geometry_before_buffer_length_validation() { + for (width, height) in [(0, 1), (1, 0), (0, 0)] { + assert_eq!( + preprocess_rgb8(width, height, &[]).unwrap_err().to_string(), + "image dimensions must be nonzero", + "{width}x{height}" + ); + } + for (width, height) in [(usize::MAX, 1), (1, usize::MAX), (usize::MAX, usize::MAX)] { + assert_eq!( + preprocess_rgb8(width, height, &[]).unwrap_err().to_string(), + "image sides must not exceed 2048", + "{width}x{height}" + ); + } + for (width, height, expected_error) in [ + (2049, 32, "image sides must not exceed 2048"), + (32, 2049, "image sides must not exceed 2048"), + (1025, 1024, "image area must not exceed 1048576 pixels"), + (201, 1, "image aspect ratio must not exceed 200"), + (1, 201, "image aspect ratio must not exceed 200"), + ] { + // A valid byte length ensures the geometry check itself rejects this + // image, rather than accidentally passing due to a truncated buffer. + let rgb = vec![0; width * height * 3]; + assert_eq!( + preprocess_rgb8(width, height, &rgb) + .unwrap_err() + .to_string(), + expected_error, + "{width}x{height}" + ); + } +} + +#[test] +fn rejects_incorrect_rgb_buffer_lengths() { + for rgb in [&[][..], &[1, 2][..], &[1, 2, 3, 4][..]] { + assert_eq!( + preprocess_rgb8(1, 1, rgb).unwrap_err().to_string(), + format!("RGB buffer length must be 3, got {}", rgb.len()) + ); + } +} + +#[test] +fn accepts_input_limits_inclusively() { + for (width, height) in [(200, 1), (1, 200), (2048, 512), (512, 2048)] { + let image = preprocess_rgb8(width, height, &vec![255; width * height * 3]).unwrap(); + assert_eq!(image.resized_width % 32, 0); + assert_eq!(image.resized_height % 32, 0); + assert!(image.pixel_values.iter().all(|&value| value == 1.0)); + } +} + +#[test] +fn matches_pinned_processor_full_output_hashes() { + use sha2::{Digest, Sha256}; + let manifest: serde_json::Value = + serde_json::from_str(include_str!("fixtures/image_preprocess/manifest.json")).unwrap(); + assert_eq!( + format!( + "{:x}", + Sha256::digest(include_bytes!( + "fixtures/image_preprocess/preprocessor_config.json" + )) + ), + manifest["config_sha256"].as_str().unwrap(), + "processor configuration must match the reference manifest" + ); + for case in manifest["cases"].as_array().unwrap() { + let name = case["name"].as_str().unwrap(); + let width = case["width"].as_u64().unwrap() as usize; + let height = case["height"].as_u64().unwrap() as usize; + let pattern = case["pattern"].as_str().unwrap(); + let mut state = case["seed"].as_u64().unwrap() as u32; + let mut rgb = Vec::with_capacity(width * height * 3); + for y in 0..height { + for x in 0..width { + for channel in 0..3 { + rgb.push(match pattern { + "noise" => { + state ^= state << 13; + state ^= state >> 17; + state ^= state << 5; + state as u8 + } + "constant" => [0, 128, 255][channel], + "checker" => { + if (x + y + channel) % 2 == 0 { + 0 + } else { + 255 + } + } + "ramp" => ((x * 13 + y * 7 + channel * 83) % 256) as u8, + _ => panic!("unknown fixture pattern {pattern}"), + }); + } + } + } + assert_eq!( + format!("{:x}", Sha256::digest(&rgb)), + case["input_sha256"].as_str().unwrap(), + "{name} input" + ); + let image = preprocess_rgb8(width, height, &rgb).unwrap(); + assert_eq!( + serde_json::json!(image.image_grid_thw), + case["grid"], + "{name} grid" + ); + assert_eq!( + serde_json::json!([image.pixel_values.len() / 1536, 1536]), + case["shape"], + "{name} shape" + ); + assert_eq!( + image.image_tokens(), + case["image_tokens"].as_u64().unwrap() as usize, + "{name} tokens" + ); + assert_eq!( + image.resized_width, + case["resized_width"].as_u64().unwrap() as usize, + "{name} width" + ); + assert_eq!( + image.resized_height, + case["resized_height"].as_u64().unwrap() as usize, + "{name} height" + ); + let mut hash = Sha256::new(); + for value in image.pixel_values { + hash.update(value.to_le_bytes()); + } + assert_eq!( + format!("{:x}", hash.finalize()), + case["output_sha256"].as_str().unwrap(), + "{name} output" + ); + } +} From 932e645b49947c6bf51c97b2f823314a9dea1db9 Mon Sep 17 00:00:00 2001 From: levius <2114377220@qq.com> Date: Fri, 2 Oct 2026 12:57:21 +0800 Subject: [PATCH 4/5] feat(cua_s1): execute native vision and serve screenshot decisions --- Cargo.lock | 155 +++++++ README.md | 3 +- .../plans/2026-10-02-native-vision.md | 13 + recipe/cua_s1/check_native_vision_http.py | 131 ++++++ .../experiments/native-vision/README.md | 63 +++ .../experiments/native-vision/boundary.json | 86 ++++ .../native-vision/source-manifest.json | 55 +++ .../experiments/native-vision/standard.json | 65 +++ recipe/cua_s1/export_multimodal_language.py | 57 +++ recipe/cua_s1/native.md | 32 +- recipe/cua_s1/native_multimodal.md | 128 ++++++ recipe/cua_s1/native_vision.md | 4 +- recipe/cua_s1/native_vision_cases.py | 119 ++++++ .../cua_s1/requirements-native-validation.txt | 4 + recipe/cua_s1/validate_native_vision.py | 197 +++++++++ recipe/cua_s1/verify_native_vision.py | 89 ++++ src/backends/cuda/qwen3_5/README.md | 5 + src/backends/cuda/qwen3_5/gemm.cu | 58 ++- src/backends/cuda/qwen3_5/ops.h | 18 +- src/backends/cuda/qwen3_5/runtime.cu | 2 + src/backends/cuda/qwen3_5/vision.cu | 258 +++++++++++ src/models/cua_s1/README.md | 4 +- src/models/cua_s1/native/Cargo.toml | 9 +- .../cua_s1/native/THIRD_PARTY_NOTICES.md | 6 + .../cua_s1/native/examples/native_vision.rs | 138 ++++++ .../cua_s1/native/examples/vision_stage.rs | 28 ++ src/models/cua_s1/native/src/contract.rs | 1 + src/models/cua_s1/native/src/cuda.rs | 14 +- src/models/cua_s1/native/src/image_request.rs | 266 ++++++++++++ src/models/cua_s1/native/src/lib.rs | 8 +- src/models/cua_s1/native/src/multimodal.rs | 217 ++++++++++ src/models/cua_s1/native/src/provenance.rs | 128 ++++++ .../cua_s1/native/src/vision/geometry.rs | 95 +++++ src/models/cua_s1/native/src/vision/mod.rs | 7 +- src/models/cua_s1/native/src/vision/model.rs | 402 ++++++++++++++++++ src/models/cua_s1/native/src/vision_engine.rs | 228 ++++++++++ src/models/cua_s1/native/src/vision_main.rs | 115 +++++ .../cua_s1/fixtures/native_decode/generate.py | 33 ++ .../fixtures/native_decode/l16-trns.png | Bin 0 -> 82 bytes tests/cua_s1/fixtures/native_decode/l16.png | Bin 0 -> 68 bytes tests/cua_s1/fixtures/native_decode/la16.png | Bin 0 -> 70 bytes .../fixtures/native_decode/rgb16-trns.png | Bin 0 -> 90 bytes tests/cua_s1/fixtures/native_decode/rgb16.png | Bin 0 -> 72 bytes .../cua_s1/fixtures/native_decode/rgba16.png | Bin 0 -> 74 bytes tests/cua_s1/test_native_vision_cuda.py | 114 +++++ 45 files changed, 3306 insertions(+), 49 deletions(-) create mode 100644 docs/superpowers/plans/2026-10-02-native-vision.md create mode 100644 recipe/cua_s1/check_native_vision_http.py create mode 100644 recipe/cua_s1/experiments/native-vision/README.md create mode 100644 recipe/cua_s1/experiments/native-vision/boundary.json create mode 100644 recipe/cua_s1/experiments/native-vision/source-manifest.json create mode 100644 recipe/cua_s1/experiments/native-vision/standard.json create mode 100644 recipe/cua_s1/export_multimodal_language.py create mode 100644 recipe/cua_s1/native_multimodal.md create mode 100644 recipe/cua_s1/native_vision_cases.py create mode 100644 recipe/cua_s1/requirements-native-validation.txt create mode 100644 recipe/cua_s1/validate_native_vision.py create mode 100644 recipe/cua_s1/verify_native_vision.py create mode 100644 src/backends/cuda/qwen3_5/vision.cu create mode 100644 src/models/cua_s1/native/examples/native_vision.rs create mode 100644 src/models/cua_s1/native/examples/vision_stage.rs create mode 100644 src/models/cua_s1/native/src/image_request.rs create mode 100644 src/models/cua_s1/native/src/multimodal.rs create mode 100644 src/models/cua_s1/native/src/provenance.rs create mode 100644 src/models/cua_s1/native/src/vision/geometry.rs create mode 100644 src/models/cua_s1/native/src/vision/model.rs create mode 100644 src/models/cua_s1/native/src/vision_engine.rs create mode 100644 src/models/cua_s1/native/src/vision_main.rs create mode 100644 tests/cua_s1/fixtures/native_decode/generate.py create mode 100644 tests/cua_s1/fixtures/native_decode/l16-trns.png create mode 100644 tests/cua_s1/fixtures/native_decode/l16.png create mode 100644 tests/cua_s1/fixtures/native_decode/la16.png create mode 100644 tests/cua_s1/fixtures/native_decode/rgb16-trns.png create mode 100644 tests/cua_s1/fixtures/native_decode/rgb16.png create mode 100644 tests/cua_s1/fixtures/native_decode/rgba16.png create mode 100644 tests/cua_s1/test_native_vision_cuda.py diff --git a/Cargo.lock b/Cargo.lock index fad6575..55f2244 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,6 +2,12 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "adler2" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" + [[package]] name = "ahash" version = "0.8.12" @@ -43,6 +49,12 @@ version = "1.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0" +[[package]] +name = "autocfg" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" + [[package]] name = "axum" version = "0.8.9" @@ -134,6 +146,18 @@ version = "3.20.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" +[[package]] +name = "bytemuck" +version = "1.25.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "95832e849adfb21180ccb6826a99da14e5d266ae5c2e668e1602cf234f153797" + +[[package]] +name = "byteorder-lite" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f1fe948ff07f4bd06c30984e69f5b4899c516a3ef74f34df92a2df2ab535495" + [[package]] name = "bytes" version = "1.12.1" @@ -215,6 +239,15 @@ dependencies = [ "libc", ] +[[package]] +name = "crc32fast" +version = "1.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "01a7799fd6b852db0e61728dde9a204c423b44d689dbd432522543614b490e78" +dependencies = [ + "cfg-if", +] + [[package]] name = "crossbeam-deque" version = "0.8.8" @@ -386,12 +419,32 @@ version = "2.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "da7c62ceae207dd37ea5b845da6a0696c799f85e97da1ab5b7910be3c1c80223" +[[package]] +name = "fdeflate" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e6853b52649d4ac5c0bd02320cddc5ba956bdb407c4b75a2c6b75bf51500f8c" +dependencies = [ + "simd-adler32", +] + [[package]] name = "find-msvc-tools" version = "0.1.13" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ef25905e51abafe4dcea6c15fec58c57b601cdbd0ee53d22ea1d3016c587d39b" +[[package]] +name = "flate2" +version = "1.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e634e2e0ebac1ee034020da1ca582e17ffe4e0f5e985823721e168928136dcb" +dependencies = [ + "crc32fast", + "miniz_oxide 0.9.1", + "zlib-rs", +] + [[package]] name = "fnv" version = "1.0.7" @@ -768,6 +821,21 @@ dependencies = [ "icu_properties", ] +[[package]] +name = "image" +version = "0.25.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85ab80394333c02fe689eaf900ab500fbd0c2213da414687ebf995a65d5a6104" +dependencies = [ + "bytemuck", + "byteorder-lite", + "moxcms", + "num-traits", + "png", + "zune-core", + "zune-jpeg", +] + [[package]] name = "indexmap" version = "2.14.2" @@ -899,6 +967,26 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" +[[package]] +name = "miniz_oxide" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fa76a2c86f704bdb222d66965fb3d63269ce38518b83cb0575fca855ebb6316" +dependencies = [ + "adler2", + "simd-adler32", +] + +[[package]] +name = "miniz_oxide" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b63fbc4a50860e98e7b2aa7804ded1db5cbc3aff9193adaff57a6931bf7c4b4c" +dependencies = [ + "adler2", + "simd-adler32", +] + [[package]] name = "mio" version = "1.2.3" @@ -932,6 +1020,16 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "moxcms" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bb85c154ba489f01b25c0d36ae69a87e4a1c73a72631fc6c0eb6dde34a73e44b" +dependencies = [ + "num-traits", + "pxfm", +] + [[package]] name = "nom" version = "7.1.3" @@ -942,13 +1040,24 @@ dependencies = [ "minimal-lexical", ] +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", +] + [[package]] name = "omni-cua-s1-native" version = "0.1.0" dependencies = [ "anyhow", "axum", + "base64 0.22.1", "half", + "image", "libloading", "memmap2", "safetensors 0.8.0", @@ -1041,6 +1150,19 @@ version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f6b464fbc74e149a392436b17d523f769e057cb6877f6a5c4618bc6f11800548" +[[package]] +name = "png" +version = "0.18.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60769b8b31b2a9f263dae2776c37b1b28ae246943cf719eb6946a1db05128a61" +dependencies = [ + "bitflags", + "crc32fast", + "fdeflate", + "flate2", + "miniz_oxide 0.8.9", +] + [[package]] name = "potential_utf" version = "0.1.6" @@ -1068,6 +1190,12 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "pxfm" +version = "0.1.30" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d55d956fa96f5ec02be2e13af0e20391a5aa83d6a074e3ad368959d0fab299ea" + [[package]] name = "quinn" version = "0.11.12" @@ -1498,6 +1626,12 @@ dependencies = [ "libc", ] +[[package]] +name = "simd-adler32" +version = "0.3.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a219298ac11a56ea9a6d2120044824d6f01aeb034955e7af7bc16858527deea" + [[package]] name = "slab" version = "0.4.12" @@ -2190,8 +2324,29 @@ dependencies = [ "syn 3.0.6", ] +[[package]] +name = "zlib-rs" +version = "0.6.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b268e58e7c693d7c271f93ffc4ba3b380412554231c85bf61ca7af91042a4112" + [[package]] name = "zmij" version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" + +[[package]] +name = "zune-core" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d56377fd46368984a170bc5aac5567e52ca5da874caa60bea39fcbca78fb658b" + +[[package]] +name = "zune-jpeg" +version = "0.5.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "27bc9d5b815bc103f142aa054f561d9187d191692ec7c2d1e2b4737f8dbd7296" +dependencies = [ + "zune-core", +] diff --git a/README.md b/README.md index 35bc8b9..c157738 100644 --- a/README.md +++ b/README.md @@ -4,7 +4,7 @@ Documentation: A community-maintained inference engine for prefill-only System1-Omni models, designed around a Rust frontend, model-owned execution, and high-performance CUDA and Metal backends. -The Rust frontend forwards requests to a separately running model worker. The Cua-S1 4B 0.2 `text` adapter has a native worker with CUDA kernels in this repository; other in-repository model engines and GPU backends are not implemented yet. +The Rust frontend forwards requests to a separately running model worker. Cua-S1 4B 0.2 has native text and screenshot workers with CUDA kernels in this repository; other in-repository model engines and GPU backends are not implemented yet. ## Run the frontend @@ -59,6 +59,7 @@ LAYA can run as an external Python worker for text requests; its in-repository m | --- | --- | | LAYA | [External worker](recipe/laya/README.md); [CPU checkpoint reader](src/models/laya/README.md); model execution planned | | Cua-S1 4B 0.2 (`text` adapter) | [Python worker](recipe/cua_s1/text.md); [native worker](recipe/cua_s1/native.md), CUDA, run on sm_89 | +| Cua-S1 4B 0.2 (`multimodal` adapter) | [Native screenshot worker](recipe/cua_s1/native_multimodal.md), complete RGB/vision/language CUDA path, validated on RTX 4090 | CUDA and Metal coverage will be documented per model as implementations are added and validated. diff --git a/docs/superpowers/plans/2026-10-02-native-vision.md b/docs/superpowers/plans/2026-10-02-native-vision.md new file mode 100644 index 0000000..1e5fd9b --- /dev/null +++ b/docs/superpowers/plans/2026-10-02-native-vision.md @@ -0,0 +1,13 @@ +# Cua-S1 native vision implementation plan + +Goal: finish native RGB-to-decision inference and verify against the pinned unmerged reference on RTX 4090. + +Architecture: combine #59 checkpoint validation, #63 CPU RGB preparation and #56 language boundary on an isolated integration branch. Keep original BF16 base and FP32 vision LoRA branches separate. Execute patch projection, interpolated positions, 24 bidirectional vision transformer blocks and merger on CUDA. Build single-image token positions in Rust and expose an end-to-end RGB API/example. Retain the text API. Reject invalid geometry before GPU execution. + +- [x] Inspect dependencies and GPU environment; preserve original branches. +- [x] Implement CUDA vision execution with CPU geometry tests and numerical GPU checks against Transformers stage outputs. +- [x] Implement native prompt image insertion, three-axis positions, RGB orchestration and candidate scoring; test invalid inputs and position oracles. +- [x] Export deterministic multi-geometry reference cases using pinned weights, unmerged BF16 and full FP32 controls, then run native GPU end-to-end comparison. Acceptance is the repository probability criterion: max error <= 2*BF16 reference error + 0.01, matching top choices at FP32 margin >= 0.05. Record intermediate errors without inventing tolerances after measurements. +- [x] Run workspace fmt, strict Clippy, tests, release build, CUDA regressions and independent spec/code review. Resolve actionable findings. Preserve commands and evidence in recipe documentation. + +The first implementation is eager; no acceleration claim. GPU validation must exercise actual native vision, never substitute exported reference image features. Large attention inputs require bounded-memory attention. Verification uses fresh source hashes and exact tool/environment versions. diff --git a/recipe/cua_s1/check_native_vision_http.py b/recipe/cua_s1/check_native_vision_http.py new file mode 100644 index 0000000..615e66f --- /dev/null +++ b/recipe/cua_s1/check_native_vision_http.py @@ -0,0 +1,131 @@ +"""Check live native screenshot HTTP answers against native replay outputs.""" + +import argparse +import json +import math +import urllib.error +import urllib.request +from pathlib import Path + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--url", default="http://127.0.0.1:28003") + parser.add_argument("--controls", required=True, type=Path) + parser.add_argument("--native", required=True, type=Path) + parser.add_argument("--out", required=True, type=Path) + args = parser.parse_args() + manifest = json.loads((args.controls / "manifest.json").read_text()) + native = json.loads(args.native.read_text()) + if not manifest["cases"] or len(manifest["cases"]) != len(native["cases"]): + raise ValueError("case coverage mismatch") + records = [] + + def post(raw): + request = urllib.request.Request( + args.url + "/v1/systemone", + data=raw, + headers={"Content-Type": "application/json"}, + ) + try: + with urllib.request.urlopen(request, timeout=180) as response: + return response.status, json.load(response) + except urllib.error.HTTPError as error: + return error.code, json.load(error) + + with urllib.request.urlopen(args.url + "/health") as response: + health = json.load(response) + assert ( + health["status"] == "ready" + and health["model"] + == "cua-ai/cua-s1-4b-0.2@16818868b0cc7813808aae4e87b417657046ab79:multimodal" + ) + for case, result in zip(manifest["cases"], native["cases"]): + assert case["case"] == result["case"] + assert len(case["questions"]) == len(result["questions"]) + assert [q["name"] for q in case["questions"]] == [ + q["name"] for q in result["questions"] + ] + raw = (args.controls / case["request_file"]).read_bytes() + body = json.loads(raw) + assert list(body["questions"]) == [q["name"] for q in case["questions"]] + status, answer = post(raw) + assert status == 200, answer + assert answer["model"] == health["model"] + assert answer["usage"]["input_tokens"] == sum( + len(q["input_ids"]) for q in case["questions"] + ) + assert answer["usage"]["output_tokens"] == 0 + assert list(answer["answers"]) == list(body["questions"]) + for q in result["questions"]: + probs = list(answer["answers"][q["name"]]["probabilities"].values()) + keys = list(body["questions"][q["name"]]["criteria"]) + response = answer["answers"][q["name"]] + assert list(response["probabilities"]) == keys + assert response["type"] == "choice" + assert ( + response["choice"] + == keys[max(range(len(probs)), key=probs.__getitem__)] + ) + confidence = ( + 1.0 + if len(probs) == 1 + else max( + 0.0, + 1.0 + + sum(p * math.log(p) for p in probs if p) / math.log(len(probs)), + ) + ) + assert abs(response["confidence"] - confidence) < 1e-7 + assert len(probs) == len(q["probabilities"]) + assert max(abs(a - b) for a, b in zip(probs, q["probabilities"])) < 1e-7 + records.append( + { + "case": case["case"], + "status": status, + "usage": answer["usage"], + "answers": answer["answers"], + } + ) + # Run input failures after successful requests to exercise the live loaded worker. + valid = json.loads( + (args.controls / manifest["cases"][0]["request_file"]).read_bytes() + ) + bad = [] + bad.append((b'{"a":1,"a":2}', 400)) + bad.append((b'{"x":NaN}', 400)) + bad.append((b"[]", 400)) + bad.append((b"{}", 422)) + invalid = json.loads(json.dumps(valid)) + invalid["state"]["image"] = "https://example.invalid/a.png" + bad.append((json.dumps(invalid).encode(), 422)) + invalid = json.loads(json.dumps(valid)) + next(iter(invalid["questions"].values()))["instructions"] = "<|image_pad|>" + bad.append((json.dumps(invalid).encode(), 422)) + invalid = json.loads(json.dumps(valid)) + invalid["questions"] = { + str(i): next(iter(valid["questions"].values())) for i in range(9) + } + bad.append((json.dumps(invalid).encode(), 422)) + invalid = json.loads(json.dumps(valid)) + invalid["extra"] = True + bad.append((json.dumps(invalid).encode(), 422)) + bad.append((b" " * (8 * 1024 * 1024 + 1), 413)) + for raw, expected in bad: + status, response = post(raw) + assert status == expected and isinstance(response.get("detail"), str), ( + status, + response, + ) + args.out.write_text( + json.dumps( + {"passed": True, "requests": records, "invalid_inputs_checked": len(bad)}, + indent=2, + ) + + "\n" + ) + print(f"PASS: {len(records)} screenshot requests and {len(bad)} input failures") + + +if __name__ == "__main__": + main() diff --git a/recipe/cua_s1/experiments/native-vision/README.md b/recipe/cua_s1/experiments/native-vision/README.md new file mode 100644 index 0000000..d05083a --- /dev/null +++ b/recipe/cua_s1/experiments/native-vision/README.md @@ -0,0 +1,63 @@ +# Native vision validation — 2026-10-02 + +Validated on NVIDIA GeForce RTX 4090 (sm_89, 24 GiB), driver 595.71.05, +CUDA toolkit 13.2; Python 3.12.3, Torch 2.14.0+cu130, Transformers 5.17.0, +PEFT 0.21.0, Pillow 11.3.0. Reference controls use unmerged adapters in BF16 +and full FP32 with TF32 disabled. Weights match the pinned upstream lock. +Native execution uses unmerged FP32 vision LoRA and merged BF16 language LoRA. + +| Set | Questions | Maximum native probability error | BF16 reference error | Allowed error | Matching choices | +| --- | ---: | ---: | ---: | ---: | ---: | +| [Standard](standard.json) | 8 | 0.00331324 | 0.00526386 | 0.02052773 | 8/8 | +| [Boundary](boundary.json) | 11 | 0.07950398 | 0.09646034 | 0.20292068 | 11/11 | + +The sets are evaluated separately so adding a difficult boundary case does not +relax the standard set's threshold. Every FP32 margin exceeds 0.05. Token IDs, +image grids and all three position axes match exactly. Repeated language calls +and vision reruns after language execution produce identical native outputs. +Standard PNG RGB pixels match Pillow; JPEG differs by up to 3 intensity levels. + +The standard set includes square, landscape, portrait, JPEG, 1/26 choices, +structured Chinese instructions and two questions. Boundary images include +1024×1024, 200×1, 383×257 and 1×1, the latter with eight questions. These are a +finite synthetic validation corpus, not a general accuracy benchmark. + +Vision stage checks identified and fixed Conv3d patch projection rounding: +FP32 accumulation is rounded to BF16 before the patch bias. On the standard +small image, patch mean absolute error is 2.33e-8; final feature RMS error is +0.005306 and maximum error 0.632813. Feature outputs are not bitwise identical. +Stage downloads are diagnostic and are excluded from any performance claim. +No speedup claim is made. + +## Verification + +- Workspace format, strict all-target Clippy, Rust tests and release build. +- 63 CPU Rust tests passed; 7 tests requiring external GPU/checkpoint data are + skipped by ordinary CPU tests. Five Cua-S1 GPU regressions were executed + separately, including graph misses/hits, changed tokens, eviction, allocation + growth and multimodal/text mixing, kernel oracles and image insertion. +- CUDA primitives compared with PyTorch: patch Conv3d, biased projection, + LayerNorm, FP32 LoRA GEMMs/addition, rotary, both GELUs and bidirectional attention. +- Live HTTP: 11 valid requests, 19 question outputs equal to the direct native + replay (tolerance 1e-7); 9 malformed/unsupported/body-limit cases checked on + each set. Model identity, usage and `detail` envelopes checked. +- Six 16-bit PNG modes independently checked against Pillow, including L16 and + RGB16 with transparency. Same-size weight corruption and unlisted overrides + are rejected by provenance tests. + +See [reproduction instructions](../../native_multimodal.md). The case generator, +reference exporter, native replay, HTTP checker and acceptance verifier are all +included in this branch. Raw controls, per-stage BF16 tensors and logs are retained +in `/root/cua-native-vision-20261002/` on the authorized GPU host and the local +`artifacts/cua-native-vision-20261002/` evidence directory. The language checkpoint +reuses the previously verified multimodal export without changing its weights; +its tokenizer and SHA-256 manifest are saved in a new task directory. + +## Review + +Independent review covered specification and implementation. Five findings were +fixed and re-reviewed: pinned source/export provenance, pre-decode PNG allocation +limits, 16-bit PNG conversions, public question invariants, and CUDA device +selection when moving vision execution between threads. A final review also fixed +HTTP checker coverage truncation and added choice/confidence assertions. No unresolved actionable +finding remained in that review. Contributor and maintainer review remain separate. diff --git a/recipe/cua_s1/experiments/native-vision/boundary.json b/recipe/cua_s1/experiments/native-vision/boundary.json new file mode 100644 index 0000000..4a26afc --- /dev/null +++ b/recipe/cua_s1/experiments/native-vision/boundary.json @@ -0,0 +1,86 @@ +{ + "passed": true, + "questions": 11, + "max_native_error": 0.07950398325920105, + "max_bf16_reference_error": 0.09646034240722656, + "allowance": 0.20292068481445313, + "choices": [ + { + "case": "maximum", + "question": "q0", + "fp32_margin": 0.19149160385131836, + "same_choice": true, + "max_probability_error": 0.01613593101501465 + }, + { + "case": "narrow", + "question": "q0", + "fp32_margin": 0.9099561609327793, + "same_choice": true, + "max_probability_error": 0.0005906363949179649 + }, + { + "case": "noise", + "question": "q0", + "fp32_margin": 0.22005164623260498, + "same_choice": true, + "max_probability_error": 0.07950398325920105 + }, + { + "case": "tiny-eight", + "question": "q0", + "fp32_margin": 0.7194288372993469, + "same_choice": true, + "max_probability_error": 0.02142578363418579 + }, + { + "case": "tiny-eight", + "question": "q1", + "fp32_margin": 0.8431737795472145, + "same_choice": true, + "max_probability_error": 0.01713651418685913 + }, + { + "case": "tiny-eight", + "question": "q2", + "fp32_margin": 0.7816882357001305, + "same_choice": true, + "max_probability_error": 0.018895387649536133 + }, + { + "case": "tiny-eight", + "question": "q3", + "fp32_margin": 0.7832479327917099, + "same_choice": true, + "max_probability_error": 0.03977161645889282 + }, + { + "case": "tiny-eight", + "question": "q4", + "fp32_margin": 0.7988598495721817, + "same_choice": true, + "max_probability_error": 0.027261078357696533 + }, + { + "case": "tiny-eight", + "question": "q5", + "fp32_margin": 0.7969894707202911, + "same_choice": true, + "max_probability_error": 0.0319133996963501 + }, + { + "case": "tiny-eight", + "question": "q6", + "fp32_margin": 0.7931957244873047, + "same_choice": true, + "max_probability_error": 0.026585102081298828 + }, + { + "case": "tiny-eight", + "question": "q7", + "fp32_margin": 0.7984120920300484, + "same_choice": true, + "max_probability_error": 0.03071284294128418 + } + ] +} diff --git a/recipe/cua_s1/experiments/native-vision/source-manifest.json b/recipe/cua_s1/experiments/native-vision/source-manifest.json new file mode 100644 index 0000000..f185586 --- /dev/null +++ b/recipe/cua_s1/experiments/native-vision/source-manifest.json @@ -0,0 +1,55 @@ +{ + "source_base": "6c40bba", + "files_sha256": { + "src/backends/cuda/qwen3_5/README.md": "c56d6505658914c68aa0a5e44695a8cb523bb8909641f1dd8cbeb51db2b28d21", + "src/backends/cuda/qwen3_5/attention.cu": "a8c19aacf46d39fca5461a0a43d99a1d224e6a45e8b1eb62ec7db6fe7f753054", + "src/backends/cuda/qwen3_5/build.sh": "178d94396fdb46ee929f0da919f120683990ed7af8fd079796106ab5995b3f0f", + "src/backends/cuda/qwen3_5/common.cuh": "ea924f7c903b6d14b0e1f10c8c33a1a9ca23806df093628a178661a303d54d55", + "src/backends/cuda/qwen3_5/elementwise.cu": "679bbce2ed940eb41620c4044c5e360293b5df077e9821001a9ca6144cf106b9", + "src/backends/cuda/qwen3_5/gdn_prefill.cu": "1c9027b6f2d3677381d8d19611ecdf7e138718459cd61f341802cc2047a582be", + "src/backends/cuda/qwen3_5/gemm.cu": "521282cb1f2fd9236301cf32d3af38aaf50c943e6e8c0d13c3e9f5dfc9c86596", + "src/backends/cuda/qwen3_5/mma.cuh": "64654f8b937585805daa72b692c243ae35d34b746e2cb6b0b55c255748495764", + "src/backends/cuda/qwen3_5/norm.cu": "7b0361f93933fddcb1615531217d77a3cdee005146218599089db06283e18c54", + "src/backends/cuda/qwen3_5/ops.h": "985f98191329023fd15195c5f93aafda24718ba8dbdab278575cf576a8d144fa", + "src/backends/cuda/qwen3_5/runtime.cu": "993ead3c2120c9b594a479f59c1666cf6bf59c046fd3abe4af04264b887a8f5d", + "src/backends/cuda/qwen3_5/vision.cu": "1e5d643e8da15da001641e40272a84dacc876be29d66aeb3e4950a52930abc75", + "src/models/cua_s1/multimodal/graph_runtime.py": "c0572ceb03dcb7f14364b4bdc697d0435c3cb739d387b72711501d6478e68e17", + "src/models/cua_s1/multimodal/model.py": "4cfe9d69d4083f502cf1b41971fdeb2115f5f1b4a00316e1f5a89d9e5a36ece0", + "src/models/cua_s1/multimodal/protocol.py": "eefc421dc485ab0a23d475a3a733b5d46f33cbe9febd6295e5d794e7fdfcec5f", + "src/models/cua_s1/native/Cargo.toml": "4fa7afd109ecb3feab6941510f757dd2984ae25ad804ca9a8a6603ca5fccbd9e", + "src/models/cua_s1/native/THIRD_PARTY_NOTICES.md": "b697db90fec79f2b2843445f526b8b21e7b6e93b91d78dfe4d9632509434cc08", + "src/models/cua_s1/native/examples/inspect_vision.rs": "2db1c38e78ac13a09fb6293b5f3dae86b0c8d2d1a44d80e76c36a3b1cafa78a0", + "src/models/cua_s1/native/examples/multimodal_boundary.rs": "c2503b83f69bf6ab5383d7058f0444d0d8786c63a6d80d9254e9d81c090c62ac", + "src/models/cua_s1/native/examples/native_vision.rs": "cd2b9f7408ff155db117bd556dbd2e032fb97eff36d49ac6d05e5c0c88cc403a", + "src/models/cua_s1/native/examples/preprocess_image.rs": "a401f4d8bc1aae4656f9de7ae671f16e8721a9fa3acb4bfeef22a4ba2d862480", + "src/models/cua_s1/native/examples/vision_stage.rs": "c0c70338f63fb97e4ab7409c875fa2b5d676bff88438252dc1a4b67227d9a5ae", + "src/models/cua_s1/native/licenses/APACHE-2.0": "77fd4710def9ec3c0f6225800e0235f15a425abd4a8b03559127fcd782612049", + "src/models/cua_s1/native/licenses/PILLOW-LICENSE": "17f240ae101143707e5e7303a5d800450d9ccc7475b463cedb555cefdb3c6ece", + "src/models/cua_s1/native/licenses/PYTORCH-LICENSE": "bd018feef8825e88181c84eb7e3aa4eafb8f08a20d9fd6ef948569610c4a3e43", + "src/models/cua_s1/native/src/contract.rs": "4278a6eaf2bfd8e76aa2287e1af9c5887e7092fd55750ff74dd0f0f23eb7af95", + "src/models/cua_s1/native/src/cuda.rs": "fe3c7011bf9cafa21eec6463eec5a633767336e87b797f1227a91f62799125af", + "src/models/cua_s1/native/src/engine.rs": "46d20775d6bd691d5375e50fb5f0aa19d5f4b6b87d59605413b393603084ed1b", + "src/models/cua_s1/native/src/image_preprocess.rs": "cda77fae0931a186158856316392b35abbcb3d71a4db387cc3e8d7135f615408", + "src/models/cua_s1/native/src/image_request.rs": "cf3957050c53b6faf131cae49721fbb734bf059e01a209dfb74a458814103142", + "src/models/cua_s1/native/src/inputs.rs": "e3c376366cdd1f9e941b6f9e4fa32010448048a6282ba8d8a99c6c4c99a09565", + "src/models/cua_s1/native/src/json.rs": "5b78af18f737ec830d3acd2e052fb7c70829aec6a7a6603a87b731635e6a9610", + "src/models/cua_s1/native/src/lib.rs": "fe1c85c646b4332359802b4f5f352daa1c7e17c49021edb61f3b20e01c313fb3", + "src/models/cua_s1/native/src/main.rs": "2066930417c9b047d40771dd8a7949d266503e09dc6b1b76e9908760b7f2bf34", + "src/models/cua_s1/native/src/model.rs": "d2308b86fac452d07b755bf56e710503964dbd3fd57c58eacf20804696cc9267", + "src/models/cua_s1/native/src/multimodal.rs": "ef56d6e7b293ca4ed4979c129a203d0ed9a34bea7b361a43d9e249b5d4f5f83b", + "src/models/cua_s1/native/src/provenance.rs": "d60f06da537b2be62acb9b612fba718ff33cea730d6f9c9f53113760c4abe65d", + "src/models/cua_s1/native/src/vision/geometry.rs": "02d1f68c53a4cb07a8d3bf7d3adff081093c2504a721a994277bc52bc6b4a0b9", + "src/models/cua_s1/native/src/vision/mod.rs": "057accf927a3c480a3b662c08ad1645f7f54e6b07ce370ecf8785252046f6417", + "src/models/cua_s1/native/src/vision/model.rs": "806a9815b891bf2c2365d3a435b42c22084a5fe6ac0d0919ee15bc9504db1c80", + "src/models/cua_s1/native/src/vision_engine.rs": "ad24024aaf7483ba5e482a45e34d6280c401ae42d171ef0a12508d4b4e31f487", + "src/models/cua_s1/native/src/vision_main.rs": "3ee30185429b2cb9f008481d635eabb2d83b20f16d6c651e3d050a8eec3f034a", + "src/models/cua_s1/native/tests/kernels.rs": "c585b28ba48f1951f540e02d3a8b3200ec89e8604081acfe46c7667289be0bb3", + "src/models/cua_s1/native/tests/multimodal.rs": "202262ae6d19b0e32e399d360812666fe68865c80598f26a620839a08e1d2bd0" + }, + "executed_binaries_sha256": { + "/root/cua-native-vision-20261002/build/libqwen3_5_cuda.so": "0f5f458f96b14b4a7ed0e06aa949cc833a1044e4e0ee54b03371e11e07e05b13", + "/root/pr52-review/code/target/release/examples/native_vision": "302e86a31c7ec11546e538725e619535e076d74809ee227a107744133befe092", + "/root/pr52-review/code/target/release/omni-cua-s1-vision": "5a5c565b5bd746e49162dd56723f8944130ec1568a966c1a34ff89524043bbb3" + }, + "reference": "full BF16 and FP32, pinned model.py plus Transformers 5.17.0, source and raw outputs retained in evidence" +} diff --git a/recipe/cua_s1/experiments/native-vision/standard.json b/recipe/cua_s1/experiments/native-vision/standard.json new file mode 100644 index 0000000..a18cfe7 --- /dev/null +++ b/recipe/cua_s1/experiments/native-vision/standard.json @@ -0,0 +1,65 @@ +{ + "passed": true, + "questions": 8, + "max_native_error": 0.0033132433891296387, + "max_bf16_reference_error": 0.005263864994049072, + "allowance": 0.020527729988098146, + "choices": [ + { + "case": "26-options", + "question": "next", + "fp32_margin": 0.4796072989702225, + "same_choice": true, + "max_probability_error": 0.0033132433891296387 + }, + { + "case": "jpeg", + "question": "next", + "fp32_margin": 0.9943915181793272, + "same_choice": true, + "max_probability_error": 0.0004887116374447942 + }, + { + "case": "portrait", + "question": "next", + "fp32_margin": 0.9951829832280055, + "same_choice": true, + "max_probability_error": 0.0002958049299195409 + }, + { + "case": "single-option", + "question": "next", + "fp32_margin": 1.0, + "same_choice": true, + "max_probability_error": 0.0 + }, + { + "case": "small", + "question": "next", + "fp32_margin": 0.994548179442063, + "same_choice": true, + "max_probability_error": 0.00012633949518203735 + }, + { + "case": "two-questions", + "question": "next", + "fp32_margin": 0.994548179442063, + "same_choice": true, + "max_probability_error": 0.00012633949518203735 + }, + { + "case": "two-questions", + "question": "second", + "fp32_margin": 0.997789517743513, + "same_choice": true, + "max_probability_error": 5.620718002319336e-05 + }, + { + "case": "wide", + "question": "next", + "fp32_margin": 0.9933466790243983, + "same_choice": true, + "max_probability_error": 7.712841033935547e-05 + } + ] +} diff --git a/recipe/cua_s1/export_multimodal_language.py b/recipe/cua_s1/export_multimodal_language.py new file mode 100644 index 0000000..1866003 --- /dev/null +++ b/recipe/cua_s1/export_multimodal_language.py @@ -0,0 +1,57 @@ +"""Export the pinned BF16 language model with the multimodal LoRA merged. + +Vision execution keeps its original BF16 base and FP32 adapters separately. +This one-time export requires the pinned Transformers/PEFT reference environment. +""" + +import argparse +import hashlib +import json +from pathlib import Path + +from models.cua_s1.multimodal.model import ( + ADAPTER_REVISION, + BASE_REVISION, + MultimodalEngine, +) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base", required=True) + parser.add_argument("--adapter", required=True) + parser.add_argument("--out", required=True, type=Path) + args = parser.parse_args() + if args.out.exists(): + parser.error("output already exists") + engine = MultimodalEngine(args.base, args.adapter) + model = engine.model.merge_and_unload() + model.model.language_model.save_pretrained(args.out, max_shard_size="5GB") + model.config.to_json_file(args.out / "config.json") + engine.tokenizer.save_pretrained(args.out) + (args.out / "cua_s1_language_export.json").write_text( + json.dumps( + { + "format": "cua-s1-multimodal-language-merged/1", + "base_revision": BASE_REVISION, + "adapter_revision": ADAPTER_REVISION, + "vision": "separate unmerged base and adapter", + "files": { + p.name: { + "size": p.stat().st_size, + "sha256": hashlib.file_digest( + p.open("rb"), "sha256" + ).hexdigest(), + } + for p in args.out.iterdir() + if p.is_file() + }, + }, + indent=2, + ) + + "\n" + ) + + +if __name__ == "__main__": + main() diff --git a/recipe/cua_s1/native.md b/recipe/cua_s1/native.md index 33c1863..1c93b01 100644 --- a/recipe/cua_s1/native.md +++ b/recipe/cua_s1/native.md @@ -30,7 +30,7 @@ token embedding. The graph contains the language layers, which update residuals in place, so a cache miss must not replay those layers over its eager result. At most eight lengths are cached. Growing the scratch allocation clears the captures before freeing their buffers. Capture adds first-use latency; leave the variable unset to use -the eager control. Rebuild both the worker and CUDA library together (ABI 3). +the eager control. Rebuild both the worker and CUDA library together (ABI 4). If capture fails, the worker returns the completed eager result and disables Graph capture/replay for its remaining lifetime, logging the failure to stderr. @@ -68,8 +68,8 @@ when `CUA_S1_GRAPH=1`. The extra tables use `4 * scratch_capacity * rotary_half` bytes (2 MiB at 16,384 rows). This is a Rust model API for integrating a vision producer. The HTTP worker -above continues to serve the text adapter. A native vision encoder, image HTTP -requests, padding, video and batching are not implemented by this API. +above continues to serve the text adapter. For native vision and image HTTP requests use the separate +[screenshot worker](native_multimodal.md). Padding, video and batching remain unsupported. ### Prepare a matching language checkpoint @@ -78,29 +78,9 @@ adapter. In the pinned reference environment, with upstream-verified weights, export just the merged language model (about 7.5 GB) to a new directory: ```sh -PYTHONPATH=src HF_HUB_OFFLINE=1 .venv/bin/python - <<'PY' -import json -from pathlib import Path -from models.cua_s1.multimodal.model import ( - ADAPTER_REVISION, BASE_REVISION, MultimodalEngine, -) - -out = Path("weights/cua-s1-4b-0.2-multimodal-language-merged") -if out.exists(): - raise FileExistsError(out) -engine = MultimodalEngine( - "weights/Qwen3.5-4B", "weights/cua-s1-4b-0.2/multimodal" -) -merged = engine.model.merge_and_unload() -merged.model.language_model.save_pretrained(out, max_shard_size="5GB") -# Preserve the root image_token_id and text_config for the native input contract. -merged.config.to_json_file(out / "config.json") -(out / "cua_s1_language_export.json").write_text(json.dumps({ - "format": "cua-s1-multimodal-language-merged/1", - "base_revision": BASE_REVISION, - "adapter_revision": ADAPTER_REVISION, -})) -PY +PYTHONPATH=src .venv/bin/python recipe/cua_s1/export_multimodal_language.py \ + --base weights/Qwen3.5-4B --adapter weights/cua-s1-4b-0.2/multimodal \ + --out weights/cua-s1-4b-0.2-multimodal-language-merged ``` No `cua_s1_export.json` text-worker marker is created. The low-level `Model` API diff --git a/recipe/cua_s1/native_multimodal.md b/recipe/cua_s1/native_multimodal.md new file mode 100644 index 0000000..424cd3a --- /dev/null +++ b/recipe/cua_s1/native_multimodal.md @@ -0,0 +1,128 @@ +# Cua-S1 native screenshot worker + +`omni-cua-s1-vision` serves the same screenshot `/v1/systemone` contract as the +[Python multimodal model](../../src/models/cua_s1/multimodal/model.py). Runtime request handling, PNG/JPEG decoding, +RGB preprocessing, the 24-block vision encoder, 2×2 merger, image insertion, +three-axis positions, language forward and candidate scoring execute in Rust/CUDA. +Python is used only for the one-time language export and reference validation. + +The worker combines the checkpoint loader (#59), RGB processor (#63) and language +input boundary (#56). These dependencies are included in the integration branch; +merging this branch does not require users to manually combine working trees. +The existing `omni-cua-s1-native` binary continues to serve the text adapter. + +## Build and load + +Download the pinned base and adapter revisions using the commands in the +[text recipe](text.md). Use its Python 3.12 environment with the additional image +packages and save the upstream lock beside the base directory: + +```sh +.venv/bin/python -m pip install -r recipe/cua_s1/requirements-native-validation.txt +curl --fail -L https://raw.githubusercontent.com/trycua/cua/0e75660ce4c2edda519e0c795fa3ad98abf4e76f/libs/cua-s1/ci/weights.lock.json \ + -o weights/weights.lock.json +``` + +The exporter and native worker verify this lock's exact trusted SHA-256. Export +the language weights once to a new directory: + +```sh +PYTHONPATH=src .venv/bin/python recipe/cua_s1/export_multimodal_language.py \ + --base weights/Qwen3.5-4B \ + --adapter weights/cua-s1-4b-0.2/multimodal \ + --out weights/cua-s1-multimodal-language +src/backends/cuda/qwen3_5/build.sh target/release 89 +cargo build --release --locked -p omni-cua-s1-native --bins --examples +CUA_S1_BASE=weights/Qwen3.5-4B \ +CUA_S1_VISION_ADAPTER=weights/cua-s1-4b-0.2/multimodal \ +CUA_S1_MODEL=weights/cua-s1-multimodal-language \ +CUA_S1_CUDA_LIB=$PWD/target/release/libqwen3_5_cuda.so \ + target/release/omni-cua-s1-vision +``` + +`CUA_S1_HOST` and `CUA_S1_PORT` default to `127.0.0.1:8000`. `/health` reports ready +after loading. The library and worker must both use **CUDA ABI 4**; rebuild both. +RTX 4090 / sm_89 is validated. Tensor-core operations require sm_80 or newer. + +At startup the worker verifies the upstream lock's trusted SHA-256, all pinned +base/adapter files, and the language export's per-file hashes before CUDA loading. +The export manifest is locally generated, trusted provenance; it is not a signature +of an externally supplied checkpoint. Keep all checkpoint and manifest files +immutable during inference. Low-level `VisionCheckpoint` and `VisionModel` APIs +perform structural validation; callers of those APIs own artifact provenance. + +The original vision tensors stay BF16 and the 50 rank-16 LoRA pairs stay FP32. +Each LoRA branch uses FP32 GEMMs, scaling and addition before BF16 rounding. +Language LoRA uses the existing merged BF16 path. The encoder uses bounded-memory, +bidirectional attention; it never allocates a full image-token-square score matrix. +Multimodal execution remains eager. This is a correctness implementation, with no +throughput or speedup claim. + +## Request and API behavior + +The screenshot contract accepts one inline PNG/JPEG, up to eight choice questions, +1–26 options each, a body up to 8 MiB and compressed image up to 4 MiB. Images are +bounded to 2048 per side, 1,048,576 pixels, and aspect ratio 200:1; animated PNGs +are rejected. PNG metadata allocations are bounded before decoding. Six 16-bit PNG +modes, including grayscale transparency, have Pillow-compatible RGB conversions. + +All questions and tokenized prompt lengths (at most 4096) are checked before +vision execution. A request preprocesses and encodes its image once, then scores +its questions independently. Features are request-local. Changed images, prompt +lengths and language/vision scratch reuse are included in validation. Malformed +JSON returns 400; invalid screenshot/question inputs return 422, body-limit +violations 413, and inference failures 500, all with a `detail` envelope. + +`VisionEngine::prepare` and `predict` accept decoded RGB requests. The lower-level +`VisionModel::forward(&ProcessedImage)` returns row-major BF16 `[image_tokens,2560]`. +`forward_with_trace` exposes synchronized intermediate downloads for diagnostics. +`prepare_prompt` constructs the token sequence and the T/H/W positions. Video, +padding, batching and CUDA Graph capture of vision are outside this worker. + +JPEG decoding uses the Rust decoder; it is not byte-identical to Pillow/libjpeg. +The validation JPEG differed by at most 3 intensity levels per channel and its +complete native probabilities passed the same acceptance test. PNG fixture pixels +were identical. CPU preprocessing parity for identical decoded RGB remains the +separate [RGB processor guarantee](native_image_preprocess.md). + +## Reproduce GPU alignment + +Generate the standard and independently added boundary requests with Pillow 11.3.0: + +```sh +.venv/bin/python recipe/cua_s1/native_vision_cases.py /tmp/cua-standard +.venv/bin/python recipe/cua_s1/native_vision_cases.py /tmp/cua-boundary --boundary +PYTHONPATH=src .venv/bin/python recipe/cua_s1/validate_native_vision.py \ + --base weights/Qwen3.5-4B --adapter weights/cua-s1-4b-0.2/multimodal \ + --requests /tmp/cua-standard --out /tmp/cua-standard-controls +CUA_S1_CUDA_LIB=$PWD/target/release/libqwen3_5_cuda.so \ + target/release/examples/native_vision \ + weights/Qwen3.5-4B weights/cua-s1-4b-0.2/multimodal \ + weights/cua-s1-multimodal-language /tmp/cua-standard-controls /tmp/cua-native.json +python recipe/cua_s1/verify_native_vision.py \ + /tmp/cua-standard-controls/manifest.json /tmp/cua-native.json +``` + +Repeat with `/tmp/cua-boundary` and separate output paths. Output directories/files +must not exist. Reference controls execute the **full vision and language path** +with unmerged BF16 and FP32 models, TF32 disabled; they do not feed BF16 reference +image embeddings into the native language model. The native replay starts from +request PNG/JPEG data and compares independently produced token IDs, image grids, +positions, candidate probabilities and repeated forwards. Stage feature errors +are reported separately. The acceptance rule is unchanged: maximum native +probability error over a set ≤ twice the maximum BF16 reference error plus 0.01; +top choices match where the FP32 top-two margin is at least 0.05. + +Additional CUDA regressions: + +```sh +.venv/bin/python tests/cua_s1/test_native_vision_cuda.py \ + target/release/libqwen3_5_cuda.so +CUA_S1_MODEL=weights/cua-s1-multimodal-language \ +CUA_S1_CUDA_LIB=$PWD/target/release/libqwen3_5_cuda.so \ + cargo test --release --locked -p omni-cua-s1-native -- --ignored +``` + +The optional `vision_stage` example emits each vision block's BF16 tensors. These +are diagnostics, not a bitwise-equivalence claim. See the recorded results in +[native vision validation](experiments/native-vision/README.md). diff --git a/recipe/cua_s1/native_vision.md b/recipe/cua_s1/native_vision.md index 2a51b83..f6ffd39 100644 --- a/recipe/cua_s1/native_vision.md +++ b/recipe/cua_s1/native_vision.md @@ -4,8 +4,8 @@ The model-owned Rust module `omni_cua_s1_native::vision` loads and structurally validates the Qwen3.5-4B vision weights and Cua-S1 4B 0.2 **multimodal** LoRA on CPU. It keeps the 297 BF16 base tensors and 100 FP32 adapter tensors separate. The adapter contains 50 A/B pairs, rank 16, alpha 32, and scale 2. No weights are -merged or converted. This increment does not preprocess images, run the vision -encoder, call CUDA, or enable image requests in the native text worker. +merged or converted. This CPU inspection API is separate from the +[native CUDA vision encoder and screenshot worker](native_multimodal.md). Prepare checkpoints from these pinned upstream revisions: diff --git a/recipe/cua_s1/native_vision_cases.py b/recipe/cua_s1/native_vision_cases.py new file mode 100644 index 0000000..d0099ce --- /dev/null +++ b/recipe/cua_s1/native_vision_cases.py @@ -0,0 +1,119 @@ +"""Deterministic screenshot requests for native end-to-end validation. + +Standard cases preserve the #53 reference suite; boundary cases add maximum +area, a 200:1 image, odd dimensions, and eight questions sharing one tiny image. +""" + +import argparse +import base64 +import io +import json +from pathlib import Path + +from PIL import Image, ImageDraw + + +def write_json(path, value): + path.write_text(json.dumps(value, ensure_ascii=False, indent=2) + "\n") + + +def make_cases(folder): + folder.mkdir(parents=True) + cases = [] + for name, size, fmt in [ + ("small", (320, 240), "PNG"), + ("wide", (640, 320), "PNG"), + ("portrait", (320, 640), "PNG"), + ("jpeg", (640, 480), "JPEG"), + ("single-option", (320, 240), "PNG"), + ("26-options", (256, 256), "PNG"), + ("two-questions", (320, 240), "PNG"), + ]: + image = Image.new("RGB", size, "#f4f6f8") + draw = ImageDraw.Draw(image) + width, height = size + draw.rectangle( + (16, 16, width - 16, height - 16), fill="white", outline="#8899aa" + ) + draw.text((24, 24), "Account settings", fill="black") + draw.text((24, 48), "Display name: Alice", fill="black") + draw.rectangle((24, height // 2, width // 2, height // 2 + 32), fill="#1460b4") + draw.text((28, height // 2 + 8), "Save", fill="white") + draw.text((width // 2 + 16, height // 2 + 8), "Cancel", fill="black") + image_path = folder / (name + (".jpg" if fmt == "JPEG" else ".png")) + image.save(image_path, format=fmt) + criteria = {"save": "Click Save", "cancel": "Click Cancel", "wait": "Wait"} + if name == "single-option": + criteria = {"save": "Click Save"} + elif name == "26-options": + criteria = {f"option-{i}": f"Choose action {i}" for i in range(26)} + questions = { + "next": { + "type": "choice", + "instructions": "Save the changed display name.", + "criteria": criteria, + } + } + if name == "two-questions": + questions["second"] = { + "type": "choice", + "instructions": {"goal": "保存名称"}, + "criteria": {"continue": {"label": "Save"}, "cancel": None}, + } + mime = "jpeg" if fmt == "JPEG" else "png" + request = { + "model": "cua-s1-4b-0.2", + "state": { + "image": f"data:image/{mime};base64," + + base64.b64encode(image_path.read_bytes()).decode() + }, + "questions": questions, + } + write_json(folder / f"{name}.json", request) + cases.append({"name": name, "image": image_path.name, "request": request}) + return cases + + +def make_boundaries(folder): + folder.mkdir(parents=True, exist_ok=False) + for name, w, h, count in [ + ("maximum", 1024, 1024, 1), + ("narrow", 200, 1, 1), + ("noise", 383, 257, 1), + ("tiny-eight", 1, 1, 8), + ]: + rgb = bytes((i * 73 + (i // 3) * 17) % 256 for i in range(w * h * 3)) + image = Image.frombytes("RGB", (w, h), rgb) + buffer = io.BytesIO() + image.save(buffer, format="PNG") + questions = { + f"q{i}": { + "type": "choice", + "instructions": f"Choose a useful next action. Case {i}.", + "criteria": { + "save": "Click Save", + "cancel": "Click Cancel", + "wait": "Wait", + }, + } + for i in range(count) + } + write_json( + folder / (name + ".json"), + { + "model": "cua-s1-4b-0.2", + "state": { + "image": "data:image/png;base64," + + base64.b64encode(buffer.getvalue()).decode() + }, + "questions": questions, + }, + ) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("output", type=Path) + parser.add_argument("--boundary", action="store_true") + args = parser.parse_args() + (make_boundaries if args.boundary else make_cases)(args.output) diff --git a/recipe/cua_s1/requirements-native-validation.txt b/recipe/cua_s1/requirements-native-validation.txt new file mode 100644 index 0000000..5609fec --- /dev/null +++ b/recipe/cua_s1/requirements-native-validation.txt @@ -0,0 +1,4 @@ +-r requirements-text.txt +torchvision==0.29.0 +numpy==2.5.3 +Pillow==11.3.0 diff --git a/recipe/cua_s1/validate_native_vision.py b/recipe/cua_s1/validate_native_vision.py new file mode 100644 index 0000000..e173689 --- /dev/null +++ b/recipe/cua_s1/validate_native_vision.py @@ -0,0 +1,197 @@ +"""Export fresh, full-vision unmerged BF16/FP32 controls for native RGB replay. + +Run with PYTHONPATH=src in the pinned reference environment. Output contains +RGB inputs, token/position oracles, vision traces and full-model probabilities. +""" + +import argparse +import gc +import hashlib +import json +import os +from pathlib import Path + +os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8") +os.environ.setdefault("HF_HUB_OFFLINE", "1") + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base", required=True) + parser.add_argument("--adapter", required=True) + parser.add_argument( + "--requests", required=True, help="directory of screenshot request JSON files" + ) + parser.add_argument("--out", required=True) + args = parser.parse_args() + import torch + import transformers + from safetensors.torch import save_file + + from models.cua_s1.multimodal.model import ( + ADAPTER_REVISION, + BASE_REVISION, + MultimodalEngine, + letter_ids, + ) + from models.cua_s1.multimodal.protocol import decode_request, parse_request + + out = Path(args.out) + out.mkdir(parents=True, exist_ok=False) + torch.set_num_threads(8) + torch.manual_seed(0) + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + cases = [] + for path in sorted(Path(args.requests).glob("*.json")): + raw = path.read_bytes() + request = parse_request(decode_request(raw)) + rgb = request.image.tobytes() + name = path.stem + (out / f"{name}.rgb").write_bytes(rgb) + (out / f"{name}.json").write_bytes(raw) + cases.append( + ( + name, + request, + { + "case": name, + "width": request.image.width, + "height": request.image.height, + "rgb_file": f"{name}.rgb", + "request_file": f"{name}.json", + "rgb_sha256": hashlib.sha256(rgb).hexdigest(), + "questions": [], + }, + ) + ) + for dtype in ["bfloat16", "float32"]: + engine = MultimodalEngine(args.base, args.adapter, dtype=dtype) + core = engine.model.get_base_model().model + for name, request, record in cases: + prepared = engine.prepare_reused(request.image, request.questions) + trace = {} + handles = [] + + def hook(key, storage=trace): + def keep(module, ins, output): + if hasattr(output, "pooler_output"): + output = output.pooler_output + storage[key] = output.detach().cpu().contiguous().clone() + + return keep + + if name == "small" and dtype == "bfloat16": + handles.append( + core.visual.patch_embed.register_forward_hook(hook("patch_embed")) + ) + + def keep_position(module, ins, kwargs, storage=trace): + x = ins[0] if ins else kwargs["hidden_states"] + storage["position"] = x.detach().cpu().contiguous().clone() + + handles.append( + core.visual.blocks[0].register_forward_pre_hook( + keep_position, with_kwargs=True + ) + ) + for i, block in enumerate(core.visual.blocks): + handles.append(block.register_forward_hook(hook(f"blocks.{i}"))) + handles.append( + core.visual.merger.norm.register_forward_hook(hook("merger.norm")) + ) + handles.append( + core.visual.merger.linear_fc1.register_forward_hook( + hook("merger.fc1") + ) + ) + with torch.no_grad(): + features = engine.encode_image(prepared[0]) + for handle in handles: + handle.remove() + tensors = {"image_features": features.detach().cpu().contiguous()} + if trace: + trace["merger.output"] = tensors["image_features"].clone() + save_file(trace, str(out / "vision-trace.safetensors")) + for idx, (q, inputs) in enumerate(zip(request.questions, prepared)): + with torch.no_grad(): + device_inputs = { + k: v.cuda() for k, v in inputs.items() if k != "pixel_values" + } + ids = device_inputs["input_ids"] + pos, _ = core.get_rope_index( + input_ids=ids, + mm_token_type_ids=device_inputs["mm_token_type_ids"], + image_grid_thw=device_inputs["image_grid_thw"], + attention_mask=device_inputs["attention_mask"], + ) + embeds = core.get_input_embeddings()(ids) + mask = ids == engine.model.config.image_token_id + embeds[mask] = features.to(embeds.dtype) + hidden = core.language_model( + inputs_embeds=embeds, + position_ids=pos, + attention_mask=device_inputs["attention_mask"], + use_cache=False, + ).last_hidden_state[0, -1] + candidate_ids = letter_ids(engine.tokenizer, len(q.keys)) + rows = engine.model.get_output_embeddings().weight[candidate_ids] + logits = torch.nn.functional.linear(hidden, rows) + probs = logits.float().softmax(-1).tolist() + if dtype == "bfloat16": + record["grid"] = inputs["image_grid_thw"][0].tolist() + record["questions"].append( + { + "name": q.name, + "input_ids": ids[0].tolist(), + "position_ids": pos[:, 0].tolist(), + "candidate_ids": candidate_ids, + "bf16_probabilities": probs, + } + ) + tensors["pixel_values"] = inputs["pixel_values"].contiguous() + else: + record["questions"][idx]["fp32_probabilities"] = probs + tensors[f"hidden.{idx}"] = hidden.detach().cpu().contiguous() + print(dtype, name, q.name, probs, flush=True) + save_file(tensors, str(out / f"{name}.{dtype}.safetensors")) + del ( + features, + hidden, + embeds, + rows, + tensors, + logits, + trace, + inputs, + device_inputs, + ids, + pos, + ) + del engine, core + gc.collect() + torch.cuda.empty_cache() + manifest = { + "schema": "cua-s1-native-vision-controls-v1", + "base_revision": BASE_REVISION, + "adapter_revision": ADAPTER_REVISION, + "torch": str(torch.__version__), + "transformers": transformers.__version__, + "gpu": torch.cuda.get_device_name(), + "tf32": False, + "control": "full-vision unmerged adapter, candidate projection at last position", + "acceptance": "max native probability error <= 2 * max BF16 probability error + 0.01; same top choice for FP32 margin >= 0.05", + "cases": [r for _, _, r in cases], + } + manifest["files"] = { + p.name: hashlib.sha256(p.read_bytes()).hexdigest() + for p in out.iterdir() + if p.is_file() + } + (out / "manifest.json").write_text( + json.dumps(manifest, ensure_ascii=False, indent=2) + "\n" + ) + + +if __name__ == "__main__": + main() diff --git a/recipe/cua_s1/verify_native_vision.py b/recipe/cua_s1/verify_native_vision.py new file mode 100644 index 0000000..1c83842 --- /dev/null +++ b/recipe/cua_s1/verify_native_vision.py @@ -0,0 +1,89 @@ +"""Fail closed on missing coverage or violations of the repository parity rule.""" + +import argparse +import json +import math +from pathlib import Path + + +def verify(control, native): + if ( + control["schema"] != "cua-s1-native-vision-controls-v1" + or native["schema"] != "cua-s1-native-vision-results-v1" + ): + raise ValueError("unsupported schema") + refs, results = control["cases"], native["cases"] + if not refs or len(refs) != len(results): + raise ValueError("case coverage mismatch") + errors, bf_errors, choices = [], [], [] + for ref, result in zip(refs, results): + if ref["case"] != result["case"] or ref["grid"] != result["grid"]: + raise ValueError("case/grid mismatch") + if len(ref["questions"]) != len(result["questions"]): + raise ValueError("question coverage mismatch") + for r, n in zip(ref["questions"], result["questions"]): + if r["name"] != n["name"] or not all( + n[k] is True + for k in ["token_ids_equal", "position_ids_equal", "repeat_equal"] + ): + raise ValueError("CPU boundary or replay mismatch") + fp, bf, actual = ( + r["fp32_probabilities"], + r["bf16_probabilities"], + n["probabilities"], + ) + if not 1 <= len(fp) <= 26 or len(fp) != len(bf) or len(fp) != len(actual): + raise ValueError("candidate count mismatch") + for p in [fp, bf, actual]: + if ( + not all(math.isfinite(x) and 0 <= x <= 1 for x in p) + or abs(sum(p) - 1) > 1e-5 + ): + raise ValueError("invalid probability distribution") + error = max(abs(x - y) for x, y in zip(fp, actual)) + errors.append(error) + bf_errors.append(max(abs(x - y) for x, y in zip(fp, bf))) + order = sorted(range(len(fp)), key=lambda i: -fp[i]) + margin = fp[order[0]] - fp[order[1]] if len(fp) > 1 else 1.0 + chosen = max(range(len(actual)), key=actual.__getitem__) + if margin >= 0.05 and chosen != order[0]: + raise ValueError(f"choice mismatch: {ref['case']}/{r['name']}") + choices.append( + { + "case": ref["case"], + "question": r["name"], + "fp32_margin": margin, + "same_choice": chosen == order[0], + "max_probability_error": error, + } + ) + allowance = 2 * max(bf_errors) + 0.01 + if max(errors) > allowance: + raise ValueError(f"probability error {max(errors)} exceeds {allowance}") + return { + "passed": True, + "questions": len(errors), + "max_native_error": max(errors), + "max_bf16_reference_error": max(bf_errors), + "allowance": allowance, + "choices": choices, + } + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("control", type=Path) + parser.add_argument("native", type=Path) + parser.add_argument("--out", type=Path) + args = parser.parse_args() + result = verify( + json.loads(args.control.read_text()), json.loads(args.native.read_text()) + ) + text = json.dumps(result, indent=2) + "\n" + if args.out: + args.out.write_text(text) + print(text, end="") + + +if __name__ == "__main__": + main() diff --git a/src/backends/cuda/qwen3_5/README.md b/src/backends/cuda/qwen3_5/README.md index b84f50b..176c949 100644 --- a/src/backends/cuda/qwen3_5/README.md +++ b/src/backends/cuda/qwen3_5/README.md @@ -7,3 +7,8 @@ src/backends/cuda/qwen3_5/build.sh [compute capability, default 89] ``` The norm, elementwise and q/k preparation kernels round to bfloat16 where Transformers (`modeling_qwen3_5.py`) does. Attention (FlashAttention-2 style, on tensor cores) and the chunked gated delta rule keep some intermediate results in bfloat16, as FlashAttention and flash-linear-attention do. GEMMs go through cuBLASLt with its first heuristic choice. Tensor-core kernels need sm_80 or newer; only sm_89 has been run. + +ABI 4 adds native vision LayerNorm, patch/biased projection, spatial rotary, +bidirectional attention, GELU, residual operations and FP32 LoRA GEMMs. Vision +GEMMs disable split-K BF16 reduction; existing language GEMM dispatch is preserved. +Rebuild the library and native workers together after ABI changes. diff --git a/src/backends/cuda/qwen3_5/gemm.cu b/src/backends/cuda/qwen3_5/gemm.cu index 7de47f5..154ad69 100644 --- a/src/backends/cuda/qwen3_5/gemm.cu +++ b/src/backends/cuda/qwen3_5/gemm.cu @@ -22,7 +22,7 @@ struct Plan { cublasLtMatmulAlgo_t algo{}; }; -using Key = std::tuple; // M, N, K, ldy +using Key = std::tuple; // M, N, K, ldy, FP32, bias struct Gemm { cublasLtHandle_t handle = nullptr; @@ -41,26 +41,32 @@ void destroy(Plan& p) { p = Plan{}; } -int describe(int M, int N, int K, int ldy, Plan& p) { +int describe(int M, int N, int K, int ldy, Plan& p, bool fp32, bool bias) { cublasStatus_t s = cublasLtMatmulDescCreate(&p.op, CUBLAS_COMPUTE_32F, CUDA_R_32F); if (s != CUBLAS_STATUS_SUCCESS) return status(s); const cublasOperation_t ta = CUBLAS_OP_T, tb = CUBLAS_OP_N; cublasLtMatmulDescSetAttribute(p.op, CUBLASLT_MATMUL_DESC_TRANSA, &ta, sizeof(ta)); cublasLtMatmulDescSetAttribute(p.op, CUBLASLT_MATMUL_DESC_TRANSB, &tb, sizeof(tb)); - if ((s = cublasLtMatrixLayoutCreate(&p.a, CUDA_R_16BF, K, N, K)) != CUBLAS_STATUS_SUCCESS) return status(s); - if ((s = cublasLtMatrixLayoutCreate(&p.b, CUDA_R_16BF, K, M, K)) != CUBLAS_STATUS_SUCCESS) return status(s); - if ((s = cublasLtMatrixLayoutCreate(&p.c, CUDA_R_16BF, N, M, ldy)) != CUBLAS_STATUS_SUCCESS) return status(s); + if (bias) { + cublasLtEpilogue_t epilogue = CUBLASLT_EPILOGUE_BIAS; + if ((s = cublasLtMatmulDescSetAttribute(p.op, CUBLASLT_MATMUL_DESC_EPILOGUE, &epilogue, sizeof(epilogue))) != CUBLAS_STATUS_SUCCESS) return status(s); + } + const cudaDataType_t dtype = fp32 ? CUDA_R_32F : CUDA_R_16BF; + if ((s = cublasLtMatrixLayoutCreate(&p.a, dtype, K, N, K)) != CUBLAS_STATUS_SUCCESS) return status(s); + if ((s = cublasLtMatrixLayoutCreate(&p.b, dtype, K, M, K)) != CUBLAS_STATUS_SUCCESS) return status(s); + if ((s = cublasLtMatrixLayoutCreate(&p.c, dtype, N, M, ldy)) != CUBLAS_STATUS_SUCCESS) return status(s); return 0; } -// The heuristic's first choice, without in-place split-K reductions. -int first_choice(Gemm& g, Plan& p) { +// The heuristic's first choice. Vision disables all split-K to avoid BF16 +// intermediate reductions; existing language GEMMs exclude only in-place reductions. +int first_choice(Gemm& g, Plan& p, bool vision) { cublasLtMatmulPreference_t pref; cublasStatus_t s = cublasLtMatmulPreferenceCreate(&pref); if (s != CUBLAS_STATUS_SUCCESS) return status(s); cublasLtMatmulPreferenceSetAttribute(pref, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &g.workspace_bytes, sizeof(g.workspace_bytes)); - const uint32_t schemes = CUBLASLT_REDUCTION_SCHEME_MASK & ~CUBLASLT_REDUCTION_SCHEME_INPLACE; + const uint32_t schemes = vision ? CUBLASLT_REDUCTION_SCHEME_NONE : (CUBLASLT_REDUCTION_SCHEME_MASK & ~CUBLASLT_REDUCTION_SCHEME_INPLACE); cublasLtMatmulPreferenceSetAttribute(pref, CUBLASLT_MATMUL_PREF_REDUCTION_SCHEME_MASK, &schemes, sizeof(schemes)); cublasLtMatmulHeuristicResult_t r{}; @@ -74,16 +80,16 @@ int first_choice(Gemm& g, Plan& p) { } // The plan for a shape, created on first use. -int plan_for(Gemm& g, int M, int N, int K, int ldy, Plan*& out) { - const Key key{M, N, K, ldy}; +int plan_for(Gemm& g, int M, int N, int K, int ldy, Plan*& out, bool fp32 = false, bool bias = false) { + const Key key{M, N, K, ldy, fp32, bias}; auto it = g.plans.find(key); if (it != g.plans.end()) { out = &it->second; return 0; } Plan p; - int rc = describe(M, N, K, ldy, p); - if (rc == 0) rc = first_choice(g, p); + int rc = describe(M, N, K, ldy, p, fp32, bias); + if (rc == 0) rc = first_choice(g, p, fp32 || bias); if (rc != 0) { destroy(p); return rc; @@ -127,3 +133,31 @@ extern "C" int cs1_gemm(void* gemm, const void* x, const void* w, void* y, int M return status(cublasLtMatmul(g->handle, p->op, &alpha, w, p->a, x, p->b, &beta, y, p->c, y, p->c, &p->algo, g->workspace, g->workspace_bytes, static_cast(stream))); } + +// Vision's biased BF16 linears round only after adding bias. The patch +// convolution passes zero bias here and applies its bias after BF16 rounding. +extern "C" int cs1_vision_linear(void* gemm, const void* x, const void* w, const void* bias, + void* y, int M, int N, int K, void* stream) { + Gemm* g = static_cast(gemm); + if (!g || !bias || M <= 0 || N <= 0 || K <= 0) return cudaErrorInvalidValue; + Plan* p = nullptr; + int rc = plan_for(*g, M, N, K, N, p, false, true); + if (rc) return rc; + auto s = cublasLtMatmulDescSetAttribute(p->op, CUBLASLT_MATMUL_DESC_BIAS_POINTER, &bias, sizeof(bias)); + if (s != CUBLAS_STATUS_SUCCESS) return status(s); + const float alpha = 1.f, beta = 0.f; + return status(cublasLtMatmul(g->handle, p->op, &alpha, w, p->a, x, p->b, &beta, y, p->c, y, p->c, &p->algo, + g->workspace, g->workspace_bytes, static_cast(stream))); +} +// Separate unmerged LoRA matrices and intermediates stay FP32. No TF32 fast compute. +extern "C" int cs1_gemm_f32(void* gemm, const float* x, const float* w, float* y, + int M, int N, int K, void* stream) { + Gemm* g = static_cast(gemm); + if (!g || M <= 0 || N <= 0 || K <= 0) return cudaErrorInvalidValue; + Plan* p = nullptr; + int rc = plan_for(*g, M, N, K, N, p, true, false); + if (rc) return rc; + const float alpha = 1.f, beta = 0.f; + return status(cublasLtMatmul(g->handle, p->op, &alpha, w, p->a, x, p->b, &beta, y, p->c, y, p->c, &p->algo, + g->workspace, g->workspace_bytes, static_cast(stream))); +} diff --git a/src/backends/cuda/qwen3_5/ops.h b/src/backends/cuda/qwen3_5/ops.h index ec0dcd6..d4de0fe 100644 --- a/src/backends/cuda/qwen3_5/ops.h +++ b/src/backends/cuda/qwen3_5/ops.h @@ -13,7 +13,7 @@ #include // Bumped whenever a signature below changes. -#define CS1_ABI_VERSION 3 +#define CS1_ABI_VERSION 4 #ifdef __cplusplus extern "C" { @@ -28,6 +28,7 @@ int cs1_malloc(void** ptr, size_t bytes); int cs1_free(void* ptr); int cs1_stream_create(void** stream); int cs1_stream_sync(void* stream); +int cs1_stream_destroy(void* stream); int cs1_graph_begin(void* stream); int cs1_graph_end(void* stream, void** exec); int cs1_graph_launch(void* exec, void* stream); @@ -94,6 +95,21 @@ void cs1_gemm_destroy(void* gemm); int cs1_gemm(void* gemm, const void* x, const void* w, void* y, int M, int N, int K, int ldy, void* stream); +// ---- vision: single image, 1024 hidden, 16 heads of 64; all BF16 except explicit float pointers ---- +int cs1_vision_linear(void* gemm, const void* x, const void* w, const void* bias, void* y, int M, int N, int K, void* stream); +int cs1_gemm_f32(void* gemm, const float* x, const float* w, float* y, int M, int N, int K, void* stream); +int cs1_vision_norm(const void* x, const void* w, const void* b, void* y, int rows, int d, void* stream); +// indices/weights [N,4], 48x48 learned table; FP32 rotary cos/sin [N,32]. +int cs1_vision_position(void* x, const void* table, const int* indices, const float* weights, int n, void* stream); +int cs1_vision_rope(const void* qkv, const float* co, const float* si, void* q, void* k, int n, void* stream); +// q/k [N,1024], V is a slice in qkv [N,3072]. No causal mask, O(N) memory. +int cs1_vision_attention(const void* q, const void* k, const void* v, void* out, int n, void* stream); +int cs1_vision_bias(void* x, const void* bias, size_t n, int d, void* stream); +int cs1_vision_gelu(void* x, size_t n, int exact, void* stream); +int cs1_vision_add(void* x, const void* delta, size_t n, void* stream); +int cs1_vision_to_float(const void* x, float* out, size_t n, void* stream); +int cs1_vision_lora_add(void* x, const float* delta, size_t n, float scale, void* stream); + #ifdef __cplusplus } #endif diff --git a/src/backends/cuda/qwen3_5/runtime.cu b/src/backends/cuda/qwen3_5/runtime.cu index 02a7db6..697c9db 100644 --- a/src/backends/cuda/qwen3_5/runtime.cu +++ b/src/backends/cuda/qwen3_5/runtime.cu @@ -20,6 +20,8 @@ int cs1_stream_create(void** stream) { return cudaStreamCreateWithFlags(reinterpret_cast(stream), cudaStreamNonBlocking); } +int cs1_stream_destroy(void* stream) { return cudaStreamDestroy(static_cast(stream)); } + int cs1_stream_sync(void* stream) { return cudaStreamSynchronize(static_cast(stream)); } int cs1_upload(void* dst, const void* src, size_t bytes, void* stream) { diff --git a/src/backends/cuda/qwen3_5/vision.cu b/src/backends/cuda/qwen3_5/vision.cu new file mode 100644 index 0000000..d77701a --- /dev/null +++ b/src/backends/cuda/qwen3_5/vision.cu @@ -0,0 +1,258 @@ +// Native Qwen3.5 vision kernels. BF16 activations, FP32 normalization/rotary/LoRA. +#include "common.cuh" +#include "mma.cuh" +#include "ops.h" +namespace cs1 { namespace vision { +namespace flash { + +constexpr int D = 64, BM = 64, BN = 32, THREADS = 128; +constexpr int LDS = D + 8; // shared row stride in elements: 528 bytes keeps ldmatrix conflict-free +constexpr int SMEM_BYTES = (BM + 2 * BN) * LDS * 2; + +__global__ void __launch_bounds__(THREADS) + flash_kernel(const bf16* __restrict__ q, const bf16* __restrict__ k, const bf16* __restrict__ v, int ldv, + bf16* __restrict__ out, int T, int Hq, int Hk, float scale_log2) { + extern __shared__ __align__(16) unsigned char smem[]; + bf16* qs = reinterpret_cast(smem); + bf16* ks = qs + BM * LDS; + bf16* vs = ks + BN * LDS; + const int h = blockIdx.y, hk = h / (Hq / Hk); + const int q0 = (gridDim.x - 1 - blockIdx.x) * BM; // the longest blocks first + const int tid = threadIdx.x, warp = tid / 32, lane = tid % 32; + const int g = lane / 4, t = lane % 4; + const int row0 = q0 + warp * 16; // this warp's first query + + for (int c = tid; c < BM * (D / 8); c += THREADS) { + const int r = c / (D / 8), col = (c % (D / 8)) * 8, row = q0 + r; + cp_async16(qs + r * LDS + col, q + ((size_t)min(row, T - 1) * Hq + h) * D + col, row < T); + } + cp_async_commit(); + + float o[D / 8][4]; +#pragma unroll + for (int n = 0; n < D / 8; n++) o[n][0] = o[n][1] = o[n][2] = o[n][3] = 0.f; + float m[2] = {-INFINITY, -INFINITY}, l[2] = {0.f, 0.f}; + + const int kv_end = T; + for (int k0 = 0; k0 < kv_end; k0 += BN) { + for (int c = tid; c < BN * (D / 8); c += THREADS) { + const int r = c / (D / 8), col = (c % (D / 8)) * 8, s = k0 + r; + cp_async16(ks + r * LDS + col, k + ((size_t)min(s, T - 1) * Hk + hk) * D + col, s < T); + } + cp_async_commit(); + for (int c = tid; c < BN * (D / 8); c += THREADS) { + const int r = c / (D / 8), col = (c % (D / 8)) * 8, s = k0 + r; + cp_async16(vs + r * LDS + col, v + (size_t)min(s, T - 1) * ldv + (size_t)hk * D + col, s < T); + } + cp_async_commit(); + cp_async_wait<1>(); // Q and K + __syncthreads(); + + // keys past every query of this warp contribute nothing + const bool active = true; + float sc[BN / 8][4]; +#pragma unroll + for (int n = 0; n < BN / 8; n++) sc[n][0] = sc[n][1] = sc[n][2] = sc[n][3] = 0.f; + if (active) { +#pragma unroll + for (int kk = 0; kk < D; kk += 16) { + uint32_t a[4]; + load_a(a, qs, LDS, warp * 16, kk, lane); +#pragma unroll + for (int n = 0; n < BN / 8; n += 2) { + uint32_t b[4]; + load_b_nk(b, ks, LDS, kk, n * 8, lane); + mma16816(sc[n], a, b[0], b[1]); + mma16816(sc[n + 1], a, b[2], b[3]); + } + } + } + uint32_t p[BN / 16][4]; + if (active) { + // causal and length mask, then the online softmax in base 2 + float mx[2] = {-INFINITY, -INFINITY}; +#pragma unroll + for (int n = 0; n < BN / 8; n++) { +#pragma unroll + for (int e = 0; e < 4; e++) { + const int key = k0 + n * 8 + 2 * t + (e & 1); + sc[n][e] = (key < T) ? sc[n][e] * scale_log2 : -INFINITY; + mx[e >> 1] = fmaxf(mx[e >> 1], sc[n][e]); + } + } + float alpha[2], base[2]; +#pragma unroll + for (int r = 0; r < 2; r++) { + mx[r] = fmaxf(mx[r], __shfl_xor_sync(0xffffffffu, mx[r], 1)); + mx[r] = fmaxf(mx[r], __shfl_xor_sync(0xffffffffu, mx[r], 2)); + const float mn = fmaxf(m[r], mx[r]); + base[r] = mn == -INFINITY ? 0.f : mn; + alpha[r] = exp2f(m[r] - base[r]); + m[r] = mn; + l[r] *= alpha[r]; + } +#pragma unroll + for (int n = 0; n < BN / 8; n++) { +#pragma unroll + for (int e = 0; e < 4; e++) { + sc[n][e] = exp2f(sc[n][e] - base[e >> 1]); + l[e >> 1] += sc[n][e]; + } + } +#pragma unroll + for (int n = 0; n < D / 8; n++) { + o[n][0] *= alpha[0]; + o[n][1] *= alpha[0]; + o[n][2] *= alpha[1]; + o[n][3] *= alpha[1]; + } + // the score accumulators, two 8-key tiles at a time, are the A fragments of P*V +#pragma unroll + for (int j = 0; j < BN / 16; j++) { + p[j][0] = pack_bf16(sc[2 * j][0], sc[2 * j][1]); + p[j][1] = pack_bf16(sc[2 * j][2], sc[2 * j][3]); + p[j][2] = pack_bf16(sc[2 * j + 1][0], sc[2 * j + 1][1]); + p[j][3] = pack_bf16(sc[2 * j + 1][2], sc[2 * j + 1][3]); + } + } + cp_async_wait<0>(); // V + __syncthreads(); + if (active) { +#pragma unroll + for (int j = 0; j < BN / 16; j++) { +#pragma unroll + for (int n = 0; n < D / 8; n += 2) { + uint32_t b[4]; + load_b_kn(b, vs, LDS, j * 16, n * 8, lane); + mma16816(o[n], p[j], b[0], b[1]); + mma16816(o[n + 1], p[j], b[2], b[3]); + } + } + } + __syncthreads(); // before the next tile overwrites K and V + } + + // the four lanes of a row each summed a quarter of its keys +#pragma unroll + for (int r = 0; r < 2; r++) { + l[r] += __shfl_xor_sync(0xffffffffu, l[r], 1); + l[r] += __shfl_xor_sync(0xffffffffu, l[r], 2); + } + const float inv[2] = {1.f / l[0], 1.f / l[1]}; +#pragma unroll + for (int r = 0; r < 2; r++) { + const int row = row0 + g + r * 8; + if (row >= T) continue; + bf16* dst = out + ((size_t)row * Hq + h) * D + 2 * t; +#pragma unroll + for (int n = 0; n < D / 8; n++) + *reinterpret_cast(dst + n * 8) = pack_bf16(o[n][2 * r] * inv[r], o[n][2 * r + 1] * inv[r]); + } +} + +} // namespace flash +} } +namespace cs1 { namespace vision { +__global__ void norm_kernel(const bf16* x, const bf16* w, const bf16* b, bf16* y, int d) { + __shared__ float scratch[32]; + const size_t off = (size_t)blockIdx.x*d; + float sum = 0.f; + for (int i=threadIdx.x;i=n) return; + const int t=i/1024,d=i%1024; + float pos=0.f; + // Separate FP32 multiply and sum, then position BF16 rounding before the residual add. + for(int j=0;j<4;j++) pos=__fadd_rn(pos,__fmul_rn(f32(table[indices[t*4+j]*1024+d]),weights[t*4+j])); + x[i]=to_bf16(f32(x[i])+round_bf16(pos)); +} +__global__ void rope_kernel(const bf16* qkv,const float* co,const float* si,bf16* q,bf16* k,size_t n) { + size_t i=(size_t)blockIdx.x*blockDim.x+threadIdx.x; + if(i>=n) return; + int t=i/1024,d=i%64,channel=i%1024; + const float c=co[t*32+d%32],s=si[t*32+d%32]; + int partner=channel+(d<32?32:-32); + float sign=d<32?-1.f:1.f; + // PyTorch materializes both products in float32 (not a fused multiply-add). + q[i]=to_bf16(__fadd_rn(__fmul_rn(f32(qkv[t*3072+channel]),c),__fmul_rn(sign*f32(qkv[t*3072+partner]),s))); + k[i]=to_bf16(__fadd_rn(__fmul_rn(f32(qkv[t*3072+1024+channel]),c),__fmul_rn(sign*f32(qkv[t*3072+1024+partner]),s))); +} +__global__ void gelu_kernel(bf16* x,size_t n,int exact) { + size_t i=(size_t)blockIdx.x*blockDim.x+threadIdx.x; + if(i>=n) return; + float a=f32(x[i]); + float v=exact?0.5f*a*(1.f+erff(a*0.7071067811865475244f)): + 0.5f*a*(1.f+tanhf(0.7978845608028654f*(a+0.044715f*a*a*a))); + x[i]=to_bf16(v); +} +__global__ void add_kernel(bf16* x,const bf16* delta,size_t n) { + size_t i=(size_t)blockIdx.x*blockDim.x+threadIdx.x; + if(i>>((const bf16*)x,(const bf16*)w,(const bf16*)b,(bf16*)y,d); + return cudaGetLastError(); +} +extern "C" int cs1_vision_position(void* x,const void* table,const int* indices,const float* weights,int n,void* stream) { + if(n<=0) return cudaErrorInvalidValue; + vision::position_kernel<<<(n*1024+255)/256,256,0,(cudaStream_t)stream>>>((bf16*)x,(const bf16*)table,indices,weights,(size_t)n*1024); + return cudaGetLastError(); +} +extern "C" int cs1_vision_rope(const void* qkv,const float* co,const float* si,void* q,void* k,int n,void* stream) { + if(n<=0) return cudaErrorInvalidValue; + vision::rope_kernel<<<(n*1024+255)/256,256,0,(cudaStream_t)stream>>>((const bf16*)qkv,co,si,(bf16*)q,(bf16*)k,(size_t)n*1024); + return cudaGetLastError(); +} +extern "C" int cs1_vision_attention(const void* q,const void* k,const void* v,void* out,int n,void* stream) { + if(n<=0) return cudaErrorInvalidValue; + namespace f=vision::flash; + f::flash_kernel<<>>( + (const bf16*)q,(const bf16*)k,(const bf16*)v,3072,(bf16*)out,n,16,16,0.125f*1.4426950408889634f); + return cudaGetLastError(); +} +extern "C" int cs1_vision_gelu(void* x,size_t n,int exact,void* stream) { + if(n==0) return cudaSuccess; + vision::gelu_kernel<<<(n+255)/256,256,0,(cudaStream_t)stream>>>((bf16*)x,n,exact); return cudaGetLastError(); +} +extern "C" int cs1_vision_add(void* x,const void* delta,size_t n,void* stream) { + if(n==0) return cudaSuccess; + vision::add_kernel<<<(n+255)/256,256,0,(cudaStream_t)stream>>>((bf16*)x,(const bf16*)delta,n); return cudaGetLastError(); +} +extern "C" int cs1_vision_to_float(const void* x,float* out,size_t n,void* stream) { + if(n==0) return cudaSuccess; + vision::to_float_kernel<<<(n+255)/256,256,0,(cudaStream_t)stream>>>((const bf16*)x,out,n); return cudaGetLastError(); +} +extern "C" int cs1_vision_lora_add(void* x,const float* delta,size_t n,float scale,void* stream) { + if(n==0) return cudaSuccess; + vision::lora_add_kernel<<<(n+255)/256,256,0,(cudaStream_t)stream>>>((bf16*)x,delta,n,scale); return cudaGetLastError(); +} + +// cuDNN Conv3d rounds its convolution output before its separate bias addition. +extern "C" int cs1_vision_bias(void* x,const void* bias,size_t n,int d,void* stream) { + if(d<=0) return cudaErrorInvalidValue; + if(n==0) return cudaSuccess; + vision::bias_kernel<<<(n+255)/256,256,0,(cudaStream_t)stream>>>((bf16*)x,(const bf16*)bias,n,d); return cudaGetLastError(); +} diff --git a/src/models/cua_s1/README.md b/src/models/cua_s1/README.md index 999ae0a..18f47fd 100644 --- a/src/models/cua_s1/README.md +++ b/src/models/cua_s1/README.md @@ -2,7 +2,7 @@ This directory owns Cua-S1 4B 0.2 ([#10](https://github.com/ThinkFlowLab/system1-omni/issues/10)): request mapping, prompt construction, adapter selection, execution, and the answer-letter readout. This page records the pinned upstream revisions, the inference contract an implementation must match, and how its outputs will be compared with the upstream reference. -Status: a reference worker for the `text` adapter loads the model through Hugging Face Transformers and PEFT: [`text/`](text/), served by [`src/frontend/cua_s1_text.py`](../../frontend/cua_s1_text.py), with setup in [`recipe/cua_s1/text.md`](../../../recipe/cua_s1/text.md). It is the correctness reference for the native worker in [`native/`](native/): Rust, with the Qwen3.5 forward pass on the CUDA kernels in [`src/backends/cuda/qwen3_5/`](../../backends/cuda/qwen3_5/), set up as in [`recipe/cua_s1/native.md`](../../../recipe/cua_s1/native.md). The `multimodal` adapter is deferred; see [Not covered yet](#not-covered-yet). +Status: a reference worker for the `text` adapter loads the model through Hugging Face Transformers and PEFT: [`text/`](text/), served by [`src/frontend/cua_s1_text.py`](../../frontend/cua_s1_text.py), with setup in [`recipe/cua_s1/text.md`](../../../recipe/cua_s1/text.md). It is the correctness reference for the native worker in [`native/`](native/): Rust, with the Qwen3.5 forward pass on the CUDA kernels in [`src/backends/cuda/qwen3_5/`](../../backends/cuda/qwen3_5/), set up as in [`recipe/cua_s1/native.md`](../../../recipe/cua_s1/native.md). The `multimodal` adapter has a [native screenshot worker](../../../recipe/cua_s1/native_multimodal.md), including Rust RGB preprocessing, CUDA vision and language execution, and RTX 4090 output validation. ## Pinned revisions @@ -105,7 +105,7 @@ The bfloat16 worker's own difference from the fp32 worker is reported next to ea ## Not covered yet -- The `multimodal` adapter: image preprocessing, the vision tower and the vision LoRA. This is tracked in [#10](https://github.com/ThinkFlowLab/system1-omni/issues/10). +- Video, padded or batched multimodal inputs, and vision CUDA Graph capture. - `score` and `noul` questions. - More than 26 options per question. - The Metal backend. diff --git a/src/models/cua_s1/native/Cargo.toml b/src/models/cua_s1/native/Cargo.toml index 6615b89..b24e878 100644 --- a/src/models/cua_s1/native/Cargo.toml +++ b/src/models/cua_s1/native/Cargo.toml @@ -11,8 +11,11 @@ path = "src/main.rs" [dependencies] anyhow = "1.0.100" +sha2 = "0.10" axum = "0.8.8" +base64 = "0.22.1" half = "2.7.1" +image = { version = "0.25", default-features = false, features = ["png", "jpeg"] } # the CUDA kernels live in libqwen3_5_cuda.so, loaded at run time libloading = "0.8" memmap2 = "0.9.9" @@ -31,5 +34,7 @@ path = "../../../../tests/cua_s1/vision_loader.rs" name = "image_preprocess" path = "../../../../tests/cua_s1/image_preprocess.rs" -[dev-dependencies] -sha2 = "0.10" + +[[bin]] +name = "omni-cua-s1-vision" +path = "src/vision_main.rs" diff --git a/src/models/cua_s1/native/THIRD_PARTY_NOTICES.md b/src/models/cua_s1/native/THIRD_PARTY_NOTICES.md index 54b788e..7d0ec38 100644 --- a/src/models/cua_s1/native/THIRD_PARTY_NOTICES.md +++ b/src/models/cua_s1/native/THIRD_PARTY_NOTICES.md @@ -22,3 +22,9 @@ bounded dimensions, standard-library buffers, and a standalone CPU API. No upstream runtime or image decoder is linked by this module. The above notices and license texts must accompany redistributed adaptations as required by their respective licenses. + +The native vision geometry, rotary, block and merger execution in +`src/vision/` follows Transformers 5.17.0 `modeling_qwen3_5.py`, Copyright +2025 The Qwen team, Alibaba Group and the HuggingFace Inc. team, under the +[Apache License, Version 2.0](licenses/APACHE-2.0). The implementation is +adapted to Rust/CUDA and separate FP32 visual LoRA execution. diff --git a/src/models/cua_s1/native/examples/native_vision.rs b/src/models/cua_s1/native/examples/native_vision.rs new file mode 100644 index 0000000..b40367b --- /dev/null +++ b/src/models/cua_s1/native/examples/native_vision.rs @@ -0,0 +1,138 @@ +//! Replay full screenshot requests and independently compare all CPU boundaries. +use anyhow::{Context, Result, ensure}; +use half::bf16; +use omni_cua_s1_native::{contract, image_request::parse_image_body, vision_engine::VisionEngine}; +use safetensors::SafeTensors; +use serde_json::{Value, json}; +use std::path::{Path, PathBuf}; + +fn main() -> Result<()> { + let args: Vec<_> = std::env::args_os().skip(1).collect(); + ensure!( + args.len() == 5, + "usage: native_vision BASE ADAPTER LANGUAGE CONTROL_DIR OUTPUT_JSON" + ); + let library = PathBuf::from(std::env::var_os("CUA_S1_CUDA_LIB").context("CUA_S1_CUDA_LIB")?); + let bundle = Path::new(&args[3]); + let out = Path::new(&args[4]); + ensure!(!out.exists(), "output already exists"); + let manifest: Value = serde_json::from_slice(&std::fs::read(bundle.join("manifest.json"))?)?; + ensure!( + manifest["schema"] == "cua-s1-native-vision-controls-v1", + "unsupported control schema" + ); + let mut engine = VisionEngine::load( + Path::new(&args[0]), + Path::new(&args[1]), + Path::new(&args[2]), + &library, + )?; + let mut records = Vec::new(); + for case in manifest["cases"].as_array().context("cases")? { + let name = case["case"].as_str().context("case name")?; + let read = |field: &str| -> Result> { + let p = Path::new(case[field].as_str().context("file name")?); + ensure!( + p.components() + .all(|c| matches!(c, std::path::Component::Normal(_))), + "invalid relative path" + ); + Ok(std::fs::read(bundle.join(p))?) + }; + let body = + contract::parse_body(&read("request_file")?).map_err(|e| anyhow::anyhow!(e.message))?; + let request = parse_image_body(&body)?; + let reference_rgb = read("rgb_file")?; + ensure!( + request.rgb.len() == reference_rgb.len(), + "decoded RGB shape mismatch" + ); + let decode_max_abs = request + .rgb + .iter() + .zip(&reference_rgb) + .map(|(&a, &b)| a.abs_diff(b)) + .max() + .unwrap_or(0); + let prepared = engine.prepare( + request.width, + request.height, + &request.rgb, + &request.questions, + )?; + let grid: Vec = serde_json::from_value(case["grid"].clone())?; + ensure!( + prepared.image.image_grid_thw.as_slice() == grid, + "grid mismatch" + ); + let start = std::time::Instant::now(); + let features = engine.vision.forward(&prepared.image)?; + let vision_ms = start.elapsed().as_secs_f64() * 1000.; + let reference = std::fs::read(bundle.join(format!("{name}.bfloat16.safetensors")))?; + let st = SafeTensors::deserialize(&reference)?; + let f = st.tensor("image_features")?; + ensure!( + f.dtype() == safetensors::Dtype::BF16 && f.data().len() == features.len() * 2, + "feature shape/dtype mismatch" + ); + let reference_features: Vec = f + .data() + .as_chunks::<2>() + .0 + .iter() + .map(|b| bf16::from_le_bytes(*b).to_f32()) + .collect(); + let errors: Vec = features + .iter() + .zip(&reference_features) + .map(|(a, &b)| a.to_f64() - b as f64) + .collect(); + let feature_max_abs = errors.iter().copied().map(f64::abs).fold(0., f64::max); + let feature_rmse = (errors.iter().map(|e| e * e).sum::() / errors.len() as f64).sqrt(); + let mut questions = Vec::new(); + for ((q, prompt), control) in request + .questions + .iter() + .zip(&prepared.prompts) + .zip(case["questions"].as_array().context("questions")?) + { + let ids: Vec = serde_json::from_value(control["input_ids"].clone())?; + let positions: [Vec; 3] = serde_json::from_value(control["position_ids"].clone())?; + ensure!( + prompt.token_ids == ids, + "{name}: tokenization differs from reference" + ); + ensure!( + prompt.position_ids == positions, + "{name}: positions differ from reference" + ); + let score = engine.score(prompt, &features, q.keys.len())?; + let repeated = engine.score(prompt, &features, q.keys.len())?; + ensure!( + score.hidden == repeated.hidden, + "repeated native language execution changed" + ); + questions.push( + json!({"name":q.name, "sequence":ids.len(), "token_ids_equal":true, + "position_ids_equal":true, "repeat_equal":true, "probabilities":score.probabilities, + "logits":score.logits, "hidden":score.hidden}), + ); + } + // Replay the same image after LM calls to expose cross-model scratch corruption. + ensure!( + features == engine.vision.forward(&prepared.image)?, + "vision replay changed features" + ); + records.push(json!({"case":name,"decode_max_abs":decode_max_abs,"grid":grid, + "image_tokens":features.len()/2560,"feature_max_abs":feature_max_abs,"feature_rmse":feature_rmse, + "vision_ms_first_call":vision_ms,"questions":questions})); + eprintln!("completed {name}: feature max {feature_max_abs}, RMS {feature_rmse}"); + } + std::fs::write( + out, + serde_json::to_vec_pretty( + &json!({"schema":"cua-s1-native-vision-results-v1","cases":records}), + )?, + )?; + Ok(()) +} diff --git a/src/models/cua_s1/native/examples/vision_stage.rs b/src/models/cua_s1/native/examples/vision_stage.rs new file mode 100644 index 0000000..1d7aa62 --- /dev/null +++ b/src/models/cua_s1/native/examples/vision_stage.rs @@ -0,0 +1,28 @@ +//! GPU parity stage runner. Raw RGB8 input; emits BF16 stage files and shapes JSON. +use anyhow::{Result, ensure}; +use omni_cua_s1_native::{image_preprocess::preprocess_rgb8, vision::VisionModel}; +use std::{env, fs, path::Path}; +fn main() -> Result<()> { + let args: Vec = env::args().collect(); + ensure!( + args.len() == 8, + "usage: vision_stage BASE ADAPTER LIB RGB WIDTH HEIGHT OUTDIR" + ); + let image = preprocess_rgb8(args[5].parse()?, args[6].parse()?, &fs::read(&args[4])?)?; + fs::create_dir_all(&args[7])?; + let mut model = VisionModel::load(&args[1], &args[2], Path::new(&args[3]))?; + let start = std::time::Instant::now(); + let output = model.forward_with_trace(&image, |name, values| { + let bytes: Vec = values + .iter() + .flat_map(|v| v.to_bits().to_le_bytes()) + .collect(); + fs::write(Path::new(&args[7]).join(format!("{name}.bf16")), bytes)?; + Ok(()) + })?; + println!( + "{}", + serde_json::json!({"grid":image.image_grid_thw,"output_shape":[image.image_tokens(),2560],"values":output.len(),"elapsed_seconds":start.elapsed().as_secs_f64()}) + ); + Ok(()) +} diff --git a/src/models/cua_s1/native/src/contract.rs b/src/models/cua_s1/native/src/contract.rs index 0ca10f2..23b91e6 100644 --- a/src/models/cua_s1/native/src/contract.rs +++ b/src/models/cua_s1/native/src/contract.rs @@ -53,6 +53,7 @@ fn error(status: u16, message: impl Into) -> RequestError { } } +#[derive(Clone)] pub struct Question { pub name: String, pub goal: String, diff --git a/src/models/cua_s1/native/src/cuda.rs b/src/models/cua_s1/native/src/cuda.rs index 719dc45..73c6725 100644 --- a/src/models/cua_s1/native/src/cuda.rs +++ b/src/models/cua_s1/native/src/cuda.rs @@ -9,7 +9,7 @@ use std::sync::OnceLock; use anyhow::{Context, Result, bail, ensure}; /// `CS1_ABI_VERSION` in ops.h. -const ABI_VERSION: u32 = 3; +const ABI_VERSION: u32 = 4; pub const LIBRARY: &str = "libqwen3_5_cuda.so"; /// A `cudaStream_t`. @@ -48,6 +48,17 @@ macro_rules! api { } api! { + cs1_vision_linear(gemm: *mut c_void, x: *const c_void, w: *const c_void, bias: *const c_void, y: *mut c_void, m: c_int, n: c_int, k: c_int, stream: Stream) -> c_int; + cs1_gemm_f32(gemm: *mut c_void, x: *const f32, w: *const f32, y: *mut f32, m: c_int, n: c_int, k: c_int, stream: Stream) -> c_int; + cs1_vision_norm(x: *const c_void, w: *const c_void, b: *const c_void, y: *mut c_void, rows: c_int, d: c_int, stream: Stream) -> c_int; + cs1_vision_position(x: *mut c_void, table: *const c_void, indices: *const i32, weights: *const f32, n: c_int, stream: Stream) -> c_int; + cs1_vision_rope(qkv: *const c_void, co: *const f32, si: *const f32, q: *mut c_void, k: *mut c_void, n: c_int, stream: Stream) -> c_int; + cs1_vision_attention(q: *const c_void, k: *const c_void, v: *const c_void, out: *mut c_void, n: c_int, stream: Stream) -> c_int; + cs1_vision_bias(x: *mut c_void, bias: *const c_void, n: usize, d: c_int, stream: Stream) -> c_int; + cs1_vision_gelu(x: *mut c_void, n: usize, exact: c_int, stream: Stream) -> c_int; + cs1_vision_add(x: *mut c_void, delta: *const c_void, n: usize, stream: Stream) -> c_int; + cs1_vision_to_float(x: *const c_void, out: *mut f32, n: usize, stream: Stream) -> c_int; + cs1_vision_lora_add(x: *mut c_void, delta: *const f32, n: usize, scale: f32, stream: Stream) -> c_int; cs1_abi_version() -> u32; cs1_error_string(code: c_int) -> *const c_char; cs1_set_device(device: c_int) -> c_int; @@ -55,6 +66,7 @@ api! { cs1_free(ptr: *mut c_void) -> c_int; cs1_stream_create(stream: *mut Stream) -> c_int; cs1_stream_sync(stream: Stream) -> c_int; + cs1_stream_destroy(stream: Stream) -> c_int; cs1_graph_begin(stream: Stream) -> c_int; cs1_graph_end(stream: Stream, exec: *mut *mut c_void) -> c_int; cs1_graph_launch(exec: *mut c_void, stream: Stream) -> c_int; diff --git a/src/models/cua_s1/native/src/image_request.rs b/src/models/cua_s1/native/src/image_request.rs new file mode 100644 index 0000000..3cf4315 --- /dev/null +++ b/src/models/cua_s1/native/src/image_request.rs @@ -0,0 +1,266 @@ +//! Bounded PNG/JPEG screenshot request mapping for the native worker. +use crate::contract::Question; +use anyhow::{Context, Result, ensure}; +use base64::Engine; +use serde_json::{Map, Value}; +use std::io::Cursor; + +pub const MAX_BODY: usize = 8 * 1024 * 1024; +pub const MAX_IMAGE_BYTES: usize = 4 * 1024 * 1024; +pub const MODEL_ID: &str = + "cua-ai/cua-s1-4b-0.2@16818868b0cc7813808aae4e87b417657046ab79:multimodal"; + +pub struct ImageRequest { + pub width: usize, + pub height: usize, + pub rgb: Vec, + pub questions: Vec, +} + +pub fn parse_image_body(body: &Map) -> Result { + ensure!( + body.len() == 3 + && ["model", "state", "questions"] + .iter() + .all(|k| body.contains_key(*k)), + "request must contain model, state and questions only" + ); + let questions = body + .get("questions") + .and_then(Value::as_object) + .context("questions must be an object")?; + ensure!( + (1..=8).contains(&questions.len()), + "questions must contain 1 to 8 questions" + ); + for (name, q) in questions { + bounded_key(name)?; + let q = q.as_object().context("question must be an object")?; + ensure!( + q.keys() + .all(|k| ["type", "instructions", "criteria"].contains(&k.as_str())), + "unsupported question fields" + ); + let goal_len = match q.get("instructions") { + Some(Value::Null) => 0, + Some(v) => checked_text(v)?.chars().count(), + None => anyhow::bail!("instructions is required"), + }; + let criteria = q + .get("criteria") + .and_then(Value::as_object) + .context("criteria must be an object")?; + let mut length = goal_len; + for (key, value) in criteria { + bounded_key(key)?; + let label = if value.is_null() { + checked_text(&Value::String(key.clone()))? + } else { + checked_text(value)? + }; + // Match Python's escaped-label character budget. + length += crate::json::quote(&label).chars().count() - 2; + } + ensure!( + length <= 16384, + "combined question text exceeds 16384 characters" + ); + } + let mut mapped = body.clone(); + mapped.insert("state".into(), Value::String("image".into())); + let (_, questions) = + crate::contract::map_request(&mapped).map_err(|e| anyhow::anyhow!(e.message))?; + let state = body["state"] + .as_object() + .context("state must contain exactly one image")?; + ensure!(state.len() == 1, "state must contain exactly one image"); + let url = state + .get("image") + .and_then(Value::as_str) + .context("state.image must be a data URL")?; + let (prefix, encoded) = url.split_once(',').context("invalid image data URL")?; + let format = match prefix { + "data:image/png;base64" => image::ImageFormat::Png, + "data:image/jpeg;base64" => image::ImageFormat::Jpeg, + _ => anyhow::bail!("only inline PNG/JPEG images are supported"), + }; + ensure!( + encoded.len() <= 4 * MAX_IMAGE_BYTES.div_ceil(3), + "encoded image exceeds 4 MiB" + ); + let raw = base64::engine::general_purpose::STANDARD + .decode(encoded) + .context("invalid base64 image")?; + ensure!(raw.len() <= MAX_IMAGE_BYTES, "image exceeds 4 MiB"); + ensure!( + image::guess_format(&raw)? == format, + "image format does not match MIME type" + ); + let mut limits = image::Limits::default(); + limits.max_image_width = Some(2048); + limits.max_image_height = Some(2048); + limits.max_alloc = Some(32 * 1024 * 1024); + let mut header = image::ImageReader::with_format(Cursor::new(&raw), format); + header.limits(limits.clone()); + let (width, height) = header.into_dimensions()?; + let (w, h) = (width as usize, height as usize); + ensure!( + w > 0 && h > 0 && w <= 2048 && h <= 2048 && w * h <= 1048576 && w.max(h) <= 200 * w.min(h), + "image dimensions exceed supported limits" + ); + let decoded = if format == image::ImageFormat::Png { + let decoder = image::codecs::png::PngDecoder::with_limits(Cursor::new(&raw), limits)?; + ensure!(!decoder.is_apng()?, "image must be single-frame"); + image::DynamicImage::from_decoder(decoder)? + } else { + let mut reader = image::ImageReader::with_format(Cursor::new(&raw), format); + reader.limits(limits); + reader.decode()? + }; + Ok(ImageRequest { + width: w, + height: h, + rgb: pillow_rgb( + decoded, + format == image::ImageFormat::Png && raw.get(25) == Some(&0), + ), + questions, + }) +} + +fn pillow_rgb(image: image::DynamicImage, png_grayscale: bool) -> Vec { + use image::DynamicImage::*; + match image { + ImageLuma16(p) => p.pixels().flat_map(|v| [v[0].min(255) as u8; 3]).collect(), + ImageLumaA16(p) if png_grayscale => { + p.pixels().flat_map(|v| [v[0].min(255) as u8; 3]).collect() + } + ImageLumaA16(p) => p.pixels().flat_map(|v| [(v[0] >> 8) as u8; 3]).collect(), + ImageRgb16(p) => p + .pixels() + .flat_map(|v| [(v[0] >> 8) as u8, (v[1] >> 8) as u8, (v[2] >> 8) as u8]) + .collect(), + ImageRgba16(p) => p + .pixels() + .flat_map(|v| [(v[0] >> 8) as u8, (v[1] >> 8) as u8, (v[2] >> 8) as u8]) + .collect(), + other => other.to_rgb8().into_raw(), + } +} + +fn bounded_key(key: &str) -> Result<()> { + ensure!( + (1..=256).contains(&key.chars().count()), + "names and option keys must contain 1 to 256 characters" + ); + Ok(()) +} + +fn checked_text(value: &Value) -> Result { + let text = match value { + Value::String(s) => s.clone(), + Value::Object(_) | Value::Array(_) => crate::json::dumps(value), + _ => anyhow::bail!("text must be a string, object or array"), + }; + ensure!( + text.chars().count() <= 16384, + "text exceeds 16384 characters" + ); + ensure!( + ![ + "<|image_pad|>", + "<|video_pad|>", + "<|vision_start|>", + "<|vision_end|>" + ] + .iter() + .any(|t| text.contains(t)), + "unsupported media control token" + ); + Ok(text) +} + +#[cfg(test)] +mod tests { + use super::*; + use base64::Engine; + use serde_json::json; + fn body() -> Map { + let mut png = std::io::Cursor::new(Vec::new()); + image::RgbImage::from_pixel(32, 32, image::Rgb([17, 83, 191])) + .write_to(&mut png, image::ImageFormat::Png) + .unwrap(); + json!({"model":"cua-s1-4b-0.2", "state":{"image":format!("data:image/png;base64,{}", base64::engine::general_purpose::STANDARD.encode(png.into_inner()))}, "questions":{"q":{"type":"choice","instructions":{"goal":"保存"},"criteria":{"save":"Click Save", "cancel":null}}}}).as_object().unwrap().clone() + } + #[test] + fn sixteen_bit_modes_match_pinned_pillow_oracles() { + let dir = std::path::Path::new(env!("CARGO_MANIFEST_DIR")) + .join("../../../../tests/cua_s1/fixtures/native_decode"); + for (name, expected) in [ + ("l16", [255, 255, 255]), + ("l16-trns", [255, 255, 255]), + ("la16", [3, 3, 3]), + ("rgb16", [3, 0, 128]), + ("rgb16-trns", [3, 0, 128]), + ("rgba16", [3, 0, 128]), + ] { + let png = std::fs::read(dir.join(format!("{name}.png"))).unwrap(); + let mut b = body(); + b["state"]["image"] = json!(format!( + "data:image/png;base64,{}", + base64::engine::general_purpose::STANDARD.encode(png) + )); + assert_eq!(parse_image_body(&b).unwrap().rgb, expected, "{name}"); + } + } + #[test] + fn sixteen_bit_grayscale_matches_pillow_clamping() { + let mut png = std::io::Cursor::new(Vec::new()); + let pixels = image::ImageBuffer::, Vec>::from_raw( + 3, + 1, + vec![100, 1000, 65535], + ) + .unwrap(); + image::DynamicImage::ImageLuma16(pixels) + .write_to(&mut png, image::ImageFormat::Png) + .unwrap(); + let mut b = body(); + b["state"]["image"] = json!(format!( + "data:image/png;base64,{}", + base64::engine::general_purpose::STANDARD.encode(png.into_inner()) + )); + assert_eq!( + parse_image_body(&b).unwrap().rgb, + [100, 100, 100, 255, 255, 255, 255, 255, 255] + ); + } + #[test] + fn decodes_png_and_matches_question_mapping() { + let r = parse_image_body(&body()).unwrap(); + assert_eq!((r.width, r.height), (32, 32)); + assert_eq!(r.rgb, [17, 83, 191].repeat(32 * 32)); + assert_eq!(r.questions[0].goal, "{\"goal\": \"保存\"}"); + assert_eq!(r.questions[0].keys, ["save", "cancel"]); + } + #[test] + fn rejects_media_controls_unknown_fields_and_mime_mismatch() { + let mut b = body(); + b["questions"]["q"]["instructions"] = json!("<|image_pad|>"); + assert!(parse_image_body(&b).is_err()); + let mut b = body(); + b.insert("extra".into(), Value::Null); + assert!(parse_image_body(&b).is_err()); + let mut b = body(); + b["state"]["image"] = json!( + b["state"]["image"] + .as_str() + .unwrap() + .replace("image/png", "image/jpeg") + ); + assert!(parse_image_body(&b).is_err()); + let mut b = body(); + b["questions"]["q"]["criteria"]["save"] = json!("x".repeat(16385)); + assert!(parse_image_body(&b).is_err()); + } +} diff --git a/src/models/cua_s1/native/src/lib.rs b/src/models/cua_s1/native/src/lib.rs index d34eb9c..56e6e3f 100644 --- a/src/models/cua_s1/native/src/lib.rs +++ b/src/models/cua_s1/native/src/lib.rs @@ -1,4 +1,4 @@ -//! A native `/v1/systemone` worker for Cua-S1 4B 0.2 (`text` adapter): request +//! A native `/v1/systemone` worker for Cua-S1 4B 0.2 (text and multimodal adapters): request //! handling, tokenization and scoring in Rust, the Qwen3.5 forward pass on the CUDA //! kernels of `src/backends/cuda/qwen3_5`, loaded at run time. @@ -6,8 +6,14 @@ pub mod contract; pub mod cuda; pub mod engine; pub mod image_preprocess; +pub mod image_request; pub mod inputs; pub mod json; pub mod model; +pub mod multimodal; pub mod vision; + +pub mod vision_engine; + +mod provenance; diff --git a/src/models/cua_s1/native/src/multimodal.rs b/src/models/cua_s1/native/src/multimodal.rs new file mode 100644 index 0000000..f6e2876 --- /dev/null +++ b/src/models/cua_s1/native/src/multimodal.rs @@ -0,0 +1,217 @@ +//! Single-image native prompt preparation and end-to-end execution. + +use crate::contract::{LETTERS, Question, SYSTEM_PROMPT}; +use anyhow::{Result, ensure}; + +pub const MAX_TOKENS: usize = 4096; + +pub fn chat_image(question: &Question, image_tokens: usize) -> Result { + ensure!( + (1..MAX_TOKENS).contains(&image_tokens), + "invalid image token count" + ); + ensure!( + (1..=26).contains(&question.labels.len()), + "invalid option count" + ); + let goal = if question.goal.is_empty() { + String::new() + } else { + format!("Goal: {}\n\n", question.goal) + }; + let options = LETTERS + .chars() + .zip(&question.labels) + .map(|(letter, label)| format!("{letter}. Decision \"{label}\" -> select")) + .collect::>() + .join("\n"); + Ok(format!( + "<|im_start|>system\n{SYSTEM_PROMPT}<|im_end|>\n<|im_start|>user\n<|vision_start|>{}<|vision_end|>{goal}App: Cua Driver\nTask family: closed-candidate decision\n\nThe current screenshot is attached.\n\nOptions:\n{options}\n\nAnswer with a single letter.<|im_end|>\n<|im_start|>assistant\n\n", + "<|image_pad|>".repeat(image_tokens) + )) +} + +pub fn image_positions(ids: &[u32], image_token: u32, grid: [usize; 3]) -> Result<[Vec; 3]> { + let [t, h, w] = grid; + ensure!( + t == 1 && h > 0 && w > 0 && h.is_multiple_of(2) && w.is_multiple_of(2), + "expected one image with even, nonzero spatial grid" + ); + ensure!( + !ids.is_empty() && ids.len() <= MAX_TOKENS, + "processed prompt exceeds {MAX_TOKENS} tokens or is empty" + ); + let count = (h / 2) + .checked_mul(w / 2) + .ok_or_else(|| anyhow::anyhow!("grid overflow"))?; + let start = ids + .iter() + .position(|&id| id == image_token) + .ok_or_else(|| anyhow::anyhow!("missing image placeholders"))?; + let end = start + .checked_add(count) + .ok_or_else(|| anyhow::anyhow!("grid overflow"))?; + ensure!( + end <= ids.len() + && ids[start..end].iter().all(|&id| id == image_token) + && ids[end..].iter().all(|&id| id != image_token), + "image placeholders must be one contiguous span matching the grid" + ); + let mut positions: [Vec; 3] = std::array::from_fn(|_| (0..start as i64).collect()); + for y in 0..h / 2 { + for x in 0..w / 2 { + positions[0].push(start as i64); + positions[1].push((start + y) as i64); + positions[2].push((start + x) as i64); + } + } + let next = start + h.max(w) / 2; + for axis in &mut positions { + axis.extend((next..next + ids.len() - end).map(|p| p as i64)); + } + Ok(positions) +} + +/// A prepared single-image prompt. Every prompt in a request is prepared before +/// running the shared image encoder. +pub struct ImagePrompt { + pub token_ids: Vec, + pub image_token_indices: Vec, + pub position_ids: [Vec; 3], +} + +pub fn prepare_prompt( + tokenizer: &tokenizers::Tokenizer, + question: &Question, + grid: [usize; 3], + image_token: u32, +) -> Result { + let count = grid[1] + .checked_mul(grid[2]) + .ok_or_else(|| anyhow::anyhow!("grid overflow"))? + / 4; + let encoded = tokenizer + .encode(chat_image(question, count)?, false) + .map_err(anyhow::Error::msg)?; + let token_ids = encoded.get_ids().to_vec(); + let position_ids = image_positions(&token_ids, image_token, grid)?; + let image_token_indices = token_ids + .iter() + .enumerate() + .filter_map(|(i, &id)| (id == image_token).then_some(i)) + .collect(); + Ok(ImagePrompt { + token_ids, + image_token_indices, + position_ids, + }) +} + +/// Validate public RGB callers as well as the JSON mapper before GPU work. +pub(crate) fn validate_questions(questions: &[Question]) -> Result<()> { + ensure!( + (1..=8).contains(&questions.len()), + "expected 1 to 8 questions" + ); + let mut names = std::collections::HashSet::new(); + for q in questions { + ensure!( + (1..=256).contains(&q.name.chars().count()) && names.insert(&q.name), + "invalid or duplicate question name" + ); + ensure!( + (1..=26).contains(&q.keys.len()) && q.keys.len() == q.labels.len(), + "option key/label counts must agree and be 1 to 26" + ); + let mut keys = std::collections::HashSet::new(); + ensure!( + q.keys + .iter() + .all(|k| (1..=256).contains(&k.chars().count()) && keys.insert(k)), + "invalid or duplicate option key" + ); + ensure!( + q.goal.chars().count() + q.labels.iter().map(|s| s.chars().count()).sum::() + <= 16384, + "combined question text exceeds 16384 characters" + ); + for text in std::iter::once(&q.goal).chain(&q.labels) { + ensure!( + ![ + "<|image_pad|>", + "<|video_pad|>", + "<|vision_start|>", + "<|vision_end|>" + ] + .iter() + .any(|t| text.contains(t)), + "unsupported media control token" + ); + } + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn public_questions_reject_inconsistent_options_and_media_controls() { + let mut q = Question { + name: "q".into(), + goal: "Save".into(), + keys: vec!["a".into()], + labels: vec!["A".into(), "B".into()], + }; + assert!(validate_questions(&[q.clone()]).is_err()); + q.labels.pop(); + validate_questions(&[q.clone()]).unwrap(); + q.goal = "<|image_pad|>".into(); + assert!(validate_questions(&[q]).is_err()); + } + #[test] + fn rectangular_grid_uses_three_axes_and_compressed_text_suffix() { + let ids = [7, 8, 99, 99, 99, 99, 99, 99, 9, 10]; + let p = image_positions(&ids, 99, [1, 4, 6]).unwrap(); + assert_eq!(p[0], [0, 1, 2, 2, 2, 2, 2, 2, 5, 6]); + assert_eq!(p[1], [0, 1, 2, 2, 2, 3, 3, 3, 5, 6]); + assert_eq!(p[2], [0, 1, 2, 3, 4, 2, 3, 4, 5, 6]); + } + + #[test] + fn rejects_invalid_grid_and_placeholder_counts() { + for grid in [ + [0, 2, 2], + [2, 2, 2], + [1, 0, 2], + [1, 3, 2], + [1, usize::MAX, 2], + ] { + assert!(image_positions(&[1, 99, 2], 99, grid).is_err()); + } + for ids in [vec![1, 2], vec![1, 99, 99, 2], vec![99, 1, 99]] { + assert!(image_positions(&ids, 99, [1, 2, 2]).is_err()); + } + } + + #[test] + fn prompt_expands_image_without_extra_separator() { + let q = Question { + name: "q".into(), + goal: "Save".into(), + keys: vec!["save".into()], + labels: vec!["Click Save".into()], + }; + let p = chat_image(&q, 2).unwrap(); + assert!(p.contains( + "user\n<|vision_start|><|image_pad|><|image_pad|><|vision_end|>Goal: Save\n\nApp:" + )); + assert!(p.contains( + "The current screenshot is attached.\n\nOptions:\nA. Decision \"Click Save\" -> select" + )); + assert!(p.ends_with("<|im_start|>assistant\n\n")); + assert!(chat_image(&q, 0).is_err()); + assert!(chat_image(&q, MAX_TOKENS).is_err()); + } +} diff --git a/src/models/cua_s1/native/src/provenance.rs b/src/models/cua_s1/native/src/provenance.rs new file mode 100644 index 0000000..6a42f12 --- /dev/null +++ b/src/models/cua_s1/native/src/provenance.rs @@ -0,0 +1,128 @@ +//! Hash checks before assigning the pinned multimodal model identity. +use anyhow::{Context, Result, ensure}; +use serde_json::Value; +use sha2::{Digest, Sha256}; +use std::{ + fs::{self, File}, + io::Read, + path::Path, +}; +const LOCK_SHA: &str = "9820bd232c5762f114e19680c0f8203d7e1faaf8a60c196cfe01964d6d8a6c09"; + +fn hash_file(path: &Path) -> Result { + let mut file = File::open(path)?; + let mut hash = Sha256::new(); + let mut buffer = vec![0; 1024 * 1024]; + loop { + let n = file.read(&mut buffer)?; + if n == 0 { + break; + } + hash.update(&buffer[..n]); + } + Ok(format!("{:x}", hash.finalize())) +} + +fn verify_files( + root: &Path, + files: &serde_json::Map, + extra: Option<&str>, +) -> Result<()> { + ensure!(!files.is_empty(), "empty artifact hash manifest"); + for (name, expected) in files { + ensure!( + Path::new(name).components().count() == 1 + && matches!( + Path::new(name).components().next(), + Some(std::path::Component::Normal(_)) + ), + "invalid artifact name" + ); + let path = root.join(name); + ensure!( + Some(fs::metadata(&path)?.len()) == expected["size"].as_u64(), + "{} size mismatch", + path.display() + ); + ensure!( + Some(hash_file(&path)?.as_str()) == expected["sha256"].as_str(), + "{} checksum mismatch", + path.display() + ); + } + for entry in fs::read_dir(root)? { + let entry = entry?; + let name = entry.file_name(); + let name = name.to_str().context("non-UTF8 artifact")?; + ensure!( + name == ".cache" || Some(name) == extra || files.contains_key(name), + "unlisted artifact: {name}" + ); + } + Ok(()) +} + +pub fn verify_sources(base: &Path, adapter: &Path) -> Result<()> { + let lock = fs::read( + base.parent() + .context("base has no parent")? + .join("weights.lock.json"), + )?; + ensure!( + format!("{:x}", Sha256::digest(&lock)) == LOCK_SHA, + "upstream weights manifest checksum mismatch" + ); + let lock: Value = serde_json::from_slice(&lock)?; + for artifact in lock["artifacts"].as_array().context("artifacts")? { + let mut selected = serde_json::Map::new(); + let is_adapter = artifact["role"] == "adapter"; + for (name, value) in artifact["files"].as_object().context("files")? { + if is_adapter { + if let Some(name) = name.strip_prefix("multimodal/") { + selected.insert(name.into(), value.clone()); + } + } else { + selected.insert(name.clone(), value.clone()); + } + } + verify_files(if is_adapter { adapter } else { base }, &selected, None)?; + } + Ok(()) +} + +pub fn verify_export(dir: &Path, marker: &Value) -> Result<()> { + let files = marker["files"] + .as_object() + .context("language export has no file hashes; rerun export_multimodal_language.py")?; + for name in ["config.json", "tokenizer.json"] { + ensure!(files.contains_key(name), "export must hash {name}"); + } + ensure!( + files.keys().any(|n| n.ends_with(".safetensors")), + "export must hash model weights" + ); + verify_files(dir, files, Some("cua_s1_language_export.json")) +} + +#[cfg(test)] +mod tests { + use super::*; + #[test] + fn rejects_missing_export_hashes_before_loading_cuda() { + assert!(verify_export(Path::new("unused"), &serde_json::json!({})).is_err()); + } + #[test] + fn detects_same_size_mutation_and_unlisted_override() { + let dir = std::env::temp_dir().join(format!("cua-provenance-{}", std::process::id())); + fs::create_dir_all(&dir).unwrap(); + fs::write(dir.join("tensor"), b"abc").unwrap(); + let files = serde_json::json!({"tensor":{"size":3, "sha256":format!("{:x}",Sha256::digest(b"abc"))}}); + verify_files(&dir, files.as_object().unwrap(), None).unwrap(); + fs::write(dir.join("tensor"), b"abd").unwrap(); + assert!(verify_files(&dir, files.as_object().unwrap(), None).is_err()); + fs::write(dir.join("tensor"), b"abc").unwrap(); + fs::write(dir.join("override"), b"x").unwrap(); + assert!(verify_files(&dir, files.as_object().unwrap(), None).is_err()); + fs::remove_dir_all(dir).unwrap(); + } +} diff --git a/src/models/cua_s1/native/src/vision/geometry.rs b/src/models/cua_s1/native/src/vision/geometry.rs new file mode 100644 index 0000000..11842aa --- /dev/null +++ b/src/models/cua_s1/native/src/vision/geometry.rs @@ -0,0 +1,95 @@ +//! Geometry in the processor's 2×2 block-major patch order. +use anyhow::{Result, ensure}; +pub(super) struct Geometry { + pub indices: Vec, + pub weights: Vec, + pub cos: Vec, + pub sin: Vec, +} +impl Geometry { + pub fn new([t, h, w]: [usize; 3]) -> Result { + ensure!( + t == 1 && h > 0 && w > 0 && h % 2 == 0 && w % 2 == 0, + "expected one image with an even, nonzero patch grid" + ); + let n = h + .checked_mul(w) + .ok_or_else(|| anyhow::anyhow!("vision grid overflow"))?; + // Processor bounds allow rounding up a 1,048,576-pixel source and narrow upscaled images. + ensure!( + n <= 4608 && h <= 512 && w <= 512, + "vision grid exceeds processor bounds" + ); + let mut g = Self { + indices: Vec::with_capacity(n * 4), + weights: Vec::with_capacity(n * 4), + cos: Vec::with_capacity(n * 32), + sin: Vec::with_capacity(n * 32), + }; + for br in 0..h / 2 { + for bc in 0..w / 2 { + for ir in 0..2 { + for ic in 0..2 { + let row = br * 2 + ir; + let col = bc * 2 + ic; + let y = (row as f32 * 47.) / (h - 1) as f32; + let x = (col as f32 * 47.) / (w - 1) as f32; + let yl = y.floor() as usize; + let xl = x.floor() as usize; + let fy = y - yl as f32; + let fx = x - xl as f32; + for (yy, wy) in [(yl, 1. - fy), ((yl + 1).min(47), fy)] { + for (xx, wx) in [(xl, 1. - fx), ((xl + 1).min(47), fx)] { + g.indices.push((yy * 48 + xx) as i32); + g.weights.push(wy * wx); + } + } + for pos in [row, col] { + for i in 0..16 { + let angle = pos as f32 / 10000f32.powf(i as f32 / 16.); + g.cos.push(angle.cos()); + g.sin.push(angle.sin()); + } + } + } + } + } + } + Ok(g) + } +} +#[cfg(test)] +mod tests { + use super::*; + #[test] + fn interpolation_corners_and_merge_order() { + let g = Geometry::new([1, 2, 4]).unwrap(); + // Sequence: (0,0), (0,1), (1,0), (1,1), (0,2), (0,3), (1,2), (1,3). + assert_eq!(&g.indices[..4], &[0, 1, 48, 49]); + assert_eq!(&g.weights[..4], &[1., 0., 0., 0.]); + assert_eq!(g.indices[2 * 4], 47 * 48); + assert_eq!(g.indices[7 * 4], 2303); + assert!((g.weights[4] - 1. / 3.).abs() < 2e-6); + assert_eq!(g.cos[0], 1.); + assert_eq!(g.sin[0], 0.); + assert!((g.sin[32 + 16] - 1f32.sin()).abs() < 1e-6); + assert!((g.sin[2 * 32] - 1f32.sin()).abs() < 1e-6); + for w in g.weights.as_chunks::<4>().0 { + assert!((w.iter().sum::() - 1.).abs() < 1e-6); + } + } + #[test] + fn geometry_rejects_video_odd_empty_and_oversize_grids() { + for grid in [ + [2, 16, 16], + [1, 0, 16], + [1, 15, 16], + [1, 128, 128], + [1, usize::MAX, 2], + ] { + assert!(Geometry::new(grid).is_err(), "accepted {grid:?}"); + } + assert!(Geometry::new([1, 64, 64]).is_ok()); + assert!(Geometry::new([1, 2, 224]).is_ok()); + } +} diff --git a/src/models/cua_s1/native/src/vision/mod.rs b/src/models/cua_s1/native/src/vision/mod.rs index 16cc059..940e4fb 100644 --- a/src/models/cua_s1/native/src/vision/mod.rs +++ b/src/models/cua_s1/native/src/vision/mod.rs @@ -1,4 +1,5 @@ -//! CPU-only, structurally validated vision checkpoint storage. No LoRA merge or execution. +//! Structurally validated vision checkpoints and native CUDA vision execution. +//! `VisionCheckpoint` loads on CPU; `VisionModel` uploads separate base/LoRA tensors. //! //! Callers must keep checkpoint files immutable (including no truncation) for the //! lifetime of the checkpoint and its borrowed views. Structural checks do not @@ -475,3 +476,7 @@ impl<'de, T: Deserialize<'de>> Deserialize<'de> for UniqueMap { deserializer.deserialize_map(UniqueVisitor(std::marker::PhantomData)) } } + +mod geometry; +mod model; +pub use model::VisionModel; diff --git a/src/models/cua_s1/native/src/vision/model.rs b/src/models/cua_s1/native/src/vision/model.rs new file mode 100644 index 0000000..293db48 --- /dev/null +++ b/src/models/cua_s1/native/src/vision/model.rs @@ -0,0 +1,402 @@ +//! GPU vision forward for the structurally validated, unmerged checkpoint. +use super::{VisionCheckpoint, geometry::Geometry}; +use crate::{ + cuda::{self, DeviceBuffer, Stream, api, check}, + image_preprocess::ProcessedImage, +}; +use anyhow::{Result, ensure}; +use half::bf16; +use std::{collections::BTreeMap, ffi::c_void, path::Path}; + +type Trace<'a> = Option<&'a mut dyn FnMut(&str, &[bf16]) -> Result<()>>; + +struct OwnedStream(Stream); +impl Drop for OwnedStream { + fn drop(&mut self) { + unsafe { (api().cs1_stream_destroy)(self.0) }; + } +} +struct Gemm(*mut c_void); +// SAFETY: the model uses its GEMM handle and stream serially. +unsafe impl Send for Gemm {} +impl Drop for Gemm { + fn drop(&mut self) { + unsafe { (api().cs1_gemm_destroy)(self.0) }; + } +} + +/// The base remains BF16; all 50 LoRA A/B pairs remain FP32, applied at scale 2. +/// A model is used by one request at a time and returns row-major [image_tokens,2560]. +pub struct VisionModel { + base: BTreeMap, + lora: BTreeMap, + stream: OwnedStream, + gemm: Gemm, +} +impl VisionModel { + pub fn load(base: impl AsRef, adapter: impl AsRef, library: &Path) -> Result { + let checkpoint = VisionCheckpoint::load(base, adapter)?; + cuda::load(library)?; + cuda::set_device(0)?; + let stream = OwnedStream(cuda::new_stream()?); + let gemm = Gemm(unsafe { (api().cs1_gemm_create)(32 << 20) }); + ensure!(!gemm.0.is_null(), "cannot create vision cuBLAS handle"); + let mut model = Self { + base: BTreeMap::new(), + lora: BTreeMap::new(), + stream, + gemm, + }; + for name in checkpoint.base_names() { + let tensor = checkpoint.base_tensor(name)?; + model.base.insert( + name.strip_prefix("model.visual.").unwrap().into(), + model.upload(tensor.data())?, + ); + } + for name in checkpoint.adapter_names() { + let tensor = checkpoint.adapter_tensor(name)?; + model.lora.insert( + name.strip_prefix("base_model.model.model.visual.") + .unwrap() + .into(), + model.upload(tensor.data())?, + ); + } + model + .base + .insert("__zero_bias".into(), model.upload(&[0; 2048])?); + Ok(model) + } + fn upload(&self, bytes: &[u8]) -> Result { + let buffer = DeviceBuffer::new(bytes.len())?; + unsafe { + cuda::upload(buffer.at(0), bytes, self.stream.0)?; + } + Ok(buffer) + } + fn weight(&self, name: &str) -> *const c_void { + self.base[name].at(0) + } + fn norm(&self, name: &str, x: &DeviceBuffer, y: &DeviceBuffer, rows: usize) -> Result<()> { + unsafe { + check( + (api().cs1_vision_norm)( + x.at(0), + self.weight(&format!("{name}.weight")), + self.weight(&format!("{name}.bias")), + y.at(0), + rows as i32, + 1024, + self.stream.0, + ), + name, + ) + } + } + #[allow(clippy::too_many_arguments)] // Mirrors the fixed-shape GEMM operation. + fn linear( + &self, + name: &str, + x: &DeviceBuffer, + y: &DeviceBuffer, + rows: usize, + output: usize, + input: usize, + work: &Work, + ) -> Result<()> { + // SAFETY: each caller provides rows*input/output sized allocations. All dimensions + // derive from validated fixed architecture and bounded image geometry. + let bias_name = if name == "patch_embed.proj" { + "__zero_bias".into() + } else { + format!("{name}.bias") + }; + unsafe { + check( + (api().cs1_vision_linear)( + self.gemm.0, + x.at(0), + self.weight(&format!("{name}.weight")), + self.weight(&bias_name), + y.at(0), + rows as i32, + output as i32, + input as i32, + self.stream.0, + ), + name, + )?; + if name == "patch_embed.proj" { + check( + (api().cs1_vision_bias)( + y.at(0), + self.weight("patch_embed.proj.bias"), + rows * output, + output as i32, + self.stream.0, + ), + "vision patch bias", + )?; + } + if let Some(a) = self.lora.get(&format!("{name}.lora_A.weight")) { + let b = &self.lora[&format!("{name}.lora_B.weight")]; + check( + (api().cs1_vision_to_float)( + x.at(0), + work.float_input.at(0).cast(), + rows * input, + self.stream.0, + ), + "vision LoRA input", + )?; + check( + (api().cs1_gemm_f32)( + self.gemm.0, + work.float_input.at(0).cast(), + a.at(0).cast(), + work.rank.at(0).cast(), + rows as i32, + 16, + input as i32, + self.stream.0, + ), + "vision LoRA A", + )?; + check( + (api().cs1_gemm_f32)( + self.gemm.0, + work.rank.at(0).cast(), + b.at(0).cast(), + work.delta.at(0).cast(), + rows as i32, + output as i32, + 16, + self.stream.0, + ), + "vision LoRA B", + )?; + check( + (api().cs1_vision_lora_add)( + y.at(0), + work.delta.at(0).cast(), + rows * output, + 2., + self.stream.0, + ), + "vision LoRA add", + )?; + } + } + Ok(()) + } + fn read(&self, x: &DeviceBuffer, n: usize) -> Result> { + let mut bytes = vec![0u8; n * 2]; + unsafe { + cuda::download(&mut bytes, x.at(0), self.stream.0)?; + } + Ok(bytes + .as_chunks::<2>() + .0 + .iter() + .map(|b| bf16::from_bits(u16::from_le_bytes([b[0], b[1]]))) + .collect()) + } + fn trace( + &self, + callback: &mut Trace<'_>, + name: &str, + x: &DeviceBuffer, + n: usize, + ) -> Result<()> { + if let Some(callback) = callback.as_mut() { + callback(name, &self.read(x, n)?)?; + } + Ok(()) + } + pub fn forward(&mut self, image: &ProcessedImage) -> Result> { + self.run(image, None) + } + /// Optional synchronized stage downloads for parity diagnosis; ordinary forward skips them. + pub fn forward_with_trace( + &mut self, + image: &ProcessedImage, + mut callback: impl FnMut(&str, &[bf16]) -> Result<()>, + ) -> Result> { + self.run(image, Some(&mut callback)) + } + fn run(&mut self, image: &ProcessedImage, mut callback: Trace<'_>) -> Result> { + cuda::set_device(0)?; + let geo = Geometry::new(image.image_grid_thw)?; + let n = image.image_grid_thw[1] * image.image_grid_thw[2]; + ensure!( + image.pixel_values.len() == n * 1536, + "vision pixel_values length does not match grid" + ); + ensure!( + image.resized_height == image.image_grid_thw[1] * 16 + && image.resized_width == image.image_grid_thw[2] * 16, + "vision resized geometry does not match grid" + ); + ensure!( + image.pixel_values.iter().all(|v| v.is_finite()), + "vision pixels must be finite" + ); + let pixels: Vec = image + .pixel_values + .iter() + .flat_map(|v| bf16::from_f32(*v).to_bits().to_le_bytes()) + .collect(); + let pixels = self.upload(&pixels)?; + let indices = self.upload( + &geo.indices + .iter() + .flat_map(|v| v.to_le_bytes()) + .collect::>(), + )?; + let weights = self.upload( + &geo.weights + .iter() + .flat_map(|v| v.to_le_bytes()) + .collect::>(), + )?; + let co = self.upload( + &geo.cos + .iter() + .flat_map(|v| v.to_le_bytes()) + .collect::>(), + )?; + let si = self.upload( + &geo.sin + .iter() + .flat_map(|v| v.to_le_bytes()) + .collect::>(), + )?; + let w = Work::new(n)?; + let x = DeviceBuffer::new(n * 1024 * 2)?; + let norm = DeviceBuffer::new(n * 1024 * 2)?; + let qkv = DeviceBuffer::new(n * 3072 * 2)?; + let q = DeviceBuffer::new(n * 1024 * 2)?; + let k = DeviceBuffer::new(n * 1024 * 2)?; + let attn = DeviceBuffer::new(n * 1024 * 2)?; + let delta = DeviceBuffer::new(n * 1024 * 2)?; + let mlp = DeviceBuffer::new(n * 4096 * 2)?; + self.linear("patch_embed.proj", &pixels, &x, n, 1024, 1536, &w)?; + self.trace(&mut callback, "patch_embed", &x, n * 1024)?; + unsafe { + check( + (api().cs1_vision_position)( + x.at(0), + self.weight("pos_embed.weight"), + indices.at(0).cast(), + weights.at(0).cast(), + n as i32, + self.stream.0, + ), + "vision learned positions", + )?; + } + self.trace(&mut callback, "position", &x, n * 1024)?; + for i in 0..24 { + let p = format!("blocks.{i}"); + self.norm(&format!("{p}.norm1"), &x, &norm, n)?; + self.linear(&format!("{p}.attn.qkv"), &norm, &qkv, n, 3072, 1024, &w)?; + unsafe { + check( + (api().cs1_vision_rope)( + qkv.at(0), + co.at(0).cast(), + si.at(0).cast(), + q.at(0), + k.at(0), + n as i32, + self.stream.0, + ), + "vision rotary", + )?; + check( + (api().cs1_vision_attention)( + q.at(0), + k.at(0), + qkv.at(2048 * 2), + attn.at(0), + n as i32, + self.stream.0, + ), + "vision attention", + )?; + } + self.linear(&format!("{p}.attn.proj"), &attn, &delta, n, 1024, 1024, &w)?; + unsafe { + check( + (api().cs1_vision_add)(x.at(0), delta.at(0), n * 1024, self.stream.0), + "vision attention residual", + )?; + } + self.norm(&format!("{p}.norm2"), &x, &norm, n)?; + self.linear( + &format!("{p}.mlp.linear_fc1"), + &norm, + &mlp, + n, + 4096, + 1024, + &w, + )?; + unsafe { + check( + (api().cs1_vision_gelu)(mlp.at(0), n * 4096, 0, self.stream.0), + "vision tanh GELU", + )?; + } + self.linear( + &format!("{p}.mlp.linear_fc2"), + &mlp, + &delta, + n, + 1024, + 4096, + &w, + )?; + unsafe { + check( + (api().cs1_vision_add)(x.at(0), delta.at(0), n * 1024, self.stream.0), + "vision MLP residual", + )?; + } + self.trace(&mut callback, &p, &x, n * 1024)?; + } + self.norm("merger.norm", &x, &norm, n)?; + self.trace(&mut callback, "merger.norm", &norm, n * 1024)?; + // Consecutive groups of four patches already have the required 2x2 merge order. + self.linear("merger.linear_fc1", &norm, &mlp, n / 4, 4096, 4096, &w)?; + self.trace(&mut callback, "merger.linear_fc1", &mlp, n * 1024)?; + unsafe { + check( + (api().cs1_vision_gelu)(mlp.at(0), n * 1024, 1, self.stream.0), + "vision exact GELU", + )?; + } + let out = DeviceBuffer::new(n / 4 * 2560 * 2)?; + self.linear("merger.linear_fc2", &mlp, &out, n / 4, 2560, 4096, &w)?; + let result = self.read(&out, n / 4 * 2560)?; + if let Some(callback) = callback.as_mut() { + callback("merger.output", &result)?; + } + Ok(result) + } +} +struct Work { + float_input: DeviceBuffer, + rank: DeviceBuffer, + delta: DeviceBuffer, +} +impl Work { + fn new(n: usize) -> Result { + Ok(Self { + float_input: DeviceBuffer::new(n * 4096 * 4)?, + rank: DeviceBuffer::new(n * 16 * 4)?, + delta: DeviceBuffer::new(n * 4096 * 4)?, + }) + } +} diff --git a/src/models/cua_s1/native/src/vision_engine.rs b/src/models/cua_s1/native/src/vision_engine.rs new file mode 100644 index 0000000..5d5e94a --- /dev/null +++ b/src/models/cua_s1/native/src/vision_engine.rs @@ -0,0 +1,228 @@ +//! Native screenshot-to-decision orchestration, with request-local vision reuse. +use crate::{ + contract::{self, LETTERS, Question}, + image_preprocess::{ProcessedImage, preprocess_rgb8}, + image_request::{ImageRequest, MODEL_ID}, + inputs::MultimodalInput, + model::Model, + multimodal::{ImagePrompt, prepare_prompt}, + vision::VisionModel, +}; +use anyhow::{Context, Result, ensure}; +use half::bf16; +use safetensors::{Dtype, SafeTensors}; +use serde_json::{Value, json}; +use std::{fs::File, path::Path}; +use tokenizers::Tokenizer; + +pub struct VisionEngine { + tokenizer: Tokenizer, + pub vision: VisionModel, + language: Model, + letters: Vec, +} + +pub struct PreparedRequest { + pub image: ProcessedImage, + pub prompts: Vec, + questions: Vec, +} + +pub struct Readout { + pub hidden: Vec, + pub logits: Vec, + pub probabilities: Vec, +} + +impl VisionEngine { + pub fn load(base: &Path, adapter: &Path, language: &Path, library: &Path) -> Result { + let marker: Value = serde_json::from_slice( + &std::fs::read(language.join("cua_s1_language_export.json")) + .context("export the multimodal language checkpoint first")?, + )?; + ensure!( + marker["format"] == "cua-s1-multimodal-language-merged/1" + && marker["base_revision"] == "851bf6e806efd8d0a36b00ddf55e13ccb7b8cd0a" + && marker["adapter_revision"] == "16818868b0cc7813808aae4e87b417657046ab79", + "expected pinned multimodal language export" + ); + crate::provenance::verify_sources(base, adapter)?; + crate::provenance::verify_export(language, &marker)?; + let tokenizer = + Tokenizer::from_file(language.join("tokenizer.json")).map_err(anyhow::Error::msg)?; + let ids = LETTERS + .chars() + .map(|c| { + let enc = tokenizer + .encode(c.to_string(), false) + .map_err(anyhow::Error::msg)?; + ensure!(enc.len() == 1, "each candidate must be one token"); + Ok(enc.get_ids()[0]) + }) + .collect::>>()?; + let model = Model::load(language, library)?; + ensure!( + tokenizer.token_to_id("<|image_pad|>") == model.cfg.image_token_id + && model.cfg.image_token_id.is_some(), + "tokenizer/image-token mismatch" + ); + let letters = letter_rows(language, &ids, model.cfg.hidden)?; + let vision = VisionModel::load(base, adapter, library)?; + Ok(Self { + tokenizer, + vision, + language: model, + letters, + }) + } + + /// Validate all tokenized prompts before the first CUDA forward of a request. + pub fn prepare( + &self, + width: usize, + height: usize, + rgb: &[u8], + questions: &[Question], + ) -> Result { + crate::multimodal::validate_questions(questions)?; + let image = preprocess_rgb8(width, height, rgb)?; + let prompts = questions + .iter() + .map(|q| { + prepare_prompt( + &self.tokenizer, + q, + image.image_grid_thw, + self.language.cfg.image_token_id.unwrap(), + ) + }) + .collect::>()?; + Ok(PreparedRequest { + image, + prompts, + questions: questions.to_vec(), + }) + } + + pub fn score( + &mut self, + prompt: &ImagePrompt, + features: &[bf16], + options: usize, + ) -> Result { + ensure!((1..=26).contains(&options), "expected 1 to 26 options"); + let hidden = self.language.forward_multimodal(&MultimodalInput { + token_ids: &prompt.token_ids, + image_token_indices: &prompt.image_token_indices, + image_embeddings: features, + position_ids: [ + &prompt.position_ids[0], + &prompt.position_ids[1], + &prompt.position_ids[2], + ], + })?; + let logits: Vec = self + .letters + .chunks_exact(hidden.len()) + .take(options) + .map(|w| { + w.iter() + .zip(&hidden) + .map(|(&a, &b)| a as f64 * b as f64) + .sum::() as f32 + }) + .collect(); + ensure!( + logits.iter().all(|x| x.is_finite()), + "non-finite candidate logits" + ); + let max = logits.iter().copied().fold(f32::NEG_INFINITY, f32::max) as f64; + let exps: Vec = logits.iter().map(|&x| (x as f64 - max).exp()).collect(); + let total: f64 = exps.iter().sum(); + let probabilities = exps.iter().map(|x| (x / total) as f32).collect(); + Ok(Readout { + hidden, + logits, + probabilities, + }) + } + + pub fn predict(&mut self, request: &ImageRequest) -> Result { + let prepared = self.prepare( + request.width, + request.height, + &request.rgb, + &request.questions, + )?; + self.predict_prepared(&prepared) + } + + pub fn predict_prepared(&mut self, prepared: &PreparedRequest) -> Result { + let questions = &prepared.questions; + ensure!( + prepared.prompts.len() == questions.len() && !questions.is_empty(), + "question/prompt count mismatch" + ); + let features = self.vision.forward(&prepared.image)?; + let mut answers = serde_json::Map::new(); + for (q, p) in questions.iter().zip(&prepared.prompts) { + answers.insert( + q.name.clone(), + contract::answer(q, &self.score(p, &features, q.keys.len())?.probabilities), + ); + } + let tokens: usize = prepared.prompts.iter().map(|p| p.token_ids.len()).sum(); + Ok( + json!({"model": MODEL_ID, "answers": answers, "usage":{"input_tokens":tokens,"output_tokens":0}}), + ) + } +} + +fn letter_rows(dir: &Path, ids: &[u32], hidden: usize) -> Result> { + let names = [ + "embed_tokens.weight", + "model.embed_tokens.weight", + "model.language_model.embed_tokens.weight", + ]; + let (path, requested) = if dir.join("model.safetensors.index.json").exists() { + let index: Value = + serde_json::from_slice(&std::fs::read(dir.join("model.safetensors.index.json"))?)?; + names + .iter() + .find_map(|name| { + index["weight_map"][name] + .as_str() + .map(|file| (dir.join(file), Some(*name))) + }) + .context("embedding missing from index")? + } else { + (dir.join("model.safetensors"), None) + }; + let file = File::open(path)?; + // SAFETY: exported checkpoint files must remain immutable during inference. + let map = unsafe { memmap2::Mmap::map(&file)? }; + let tensors = SafeTensors::deserialize(&map)?; + let name = requested + .or_else(|| names.iter().copied().find(|n| tensors.tensor(n).is_ok())) + .context("missing embedding")?; + let tensor = tensors.tensor(name)?; + ensure!( + tensor.dtype() == Dtype::BF16 && tensor.shape().len() == 2 && tensor.shape()[1] == hidden, + "embedding shape/dtype mismatch" + ); + let mut result = Vec::with_capacity(ids.len() * hidden); + for &id in ids { + ensure!( + (id as usize) < tensor.shape()[0], + "candidate outside vocabulary" + ); + let row = &tensor.data()[id as usize * hidden * 2..(id as usize + 1) * hidden * 2]; + result.extend( + row.as_chunks::<2>() + .0 + .iter() + .map(|b| bf16::from_le_bytes(*b).to_f32()), + ); + } + Ok(result) +} diff --git a/src/models/cua_s1/native/src/vision_main.rs b/src/models/cua_s1/native/src/vision_main.rs new file mode 100644 index 0000000..b6eb53b --- /dev/null +++ b/src/models/cua_s1/native/src/vision_main.rs @@ -0,0 +1,115 @@ +//! Screenshot `/v1/systemone` worker using native RGB, vision and language stages. +use anyhow::{Context, Result}; +use axum::{ + Json, Router, + body::Bytes, + extract::{DefaultBodyLimit, State, rejection::BytesRejection}, + http::StatusCode, + response::{IntoResponse, Response}, + routing::{get, post}, +}; +use omni_cua_s1_native::{ + contract, cuda, + image_request::{self, MODEL_ID}, + vision_engine::VisionEngine, +}; +use serde_json::{Value, json}; +use std::{ + path::PathBuf, + sync::{Arc, Mutex}, +}; + +type Shared = Arc>; +fn reply(status: StatusCode, value: Value) -> Response { + (status, Json(value)).into_response() +} + +async fn decide(State(engine): State, body: Result) -> Response { + let raw = match body { + Ok(b) => b, + Err(e) => return reply(e.status(), json!({"detail": e.body_text()})), + }; + let body = match contract::parse_body(&raw) { + Ok(b) => b, + Err(e) => { + return reply( + StatusCode::from_u16(e.status).unwrap(), + json!({"detail":e.message}), + ); + } + }; + // Decode and inference both run off the async executor. The mutex serializes + // the CUDA models; prepared image features live only for this request. + match tokio::task::spawn_blocking(move || -> Result<(StatusCode, Value)> { + let request = match image_request::parse_image_body(&body) { + Ok(r) => r, + Err(e) => { + return Ok(( + StatusCode::UNPROCESSABLE_ENTITY, + json!({"detail":e.to_string()}), + )); + } + }; + let mut engine = engine + .lock() + .map_err(|_| anyhow::anyhow!("poisoned engine"))?; + let prepared = match engine.prepare( + request.width, + request.height, + &request.rgb, + &request.questions, + ) { + Ok(p) => p, + Err(e) => { + return Ok(( + StatusCode::UNPROCESSABLE_ENTITY, + json!({"detail": e.to_string()}), + )); + } + }; + Ok((StatusCode::OK, engine.predict_prepared(&prepared)?)) + }) + .await + { + Ok(Ok((status, body))) => reply(status, body), + err => { + eprintln!("native image inference failed: {err:?}"); + reply( + StatusCode::INTERNAL_SERVER_ERROR, + json!({"detail":"inference failed"}), + ) + } + } +} + +#[tokio::main] +async fn main() -> Result<()> { + let path = |name| { + std::env::var_os(name) + .map(PathBuf::from) + .with_context(|| format!("set {name}")) + }; + let library = std::env::var_os("CUA_S1_CUDA_LIB") + .map(PathBuf::from) + .map_or_else(cuda::default_library, Ok)?; + let engine = VisionEngine::load( + &path("CUA_S1_BASE")?, + &path("CUA_S1_VISION_ADAPTER")?, + &path("CUA_S1_MODEL")?, + &library, + )?; + let host = std::env::var("CUA_S1_HOST").unwrap_or_else(|_| "127.0.0.1".into()); + let port: u16 = std::env::var("CUA_S1_PORT").map_or(Ok(8000), |p| p.parse())?; + let app = Router::new() + .route( + "/health", + get(|| async { Json(json!({"status":"ready", "model": MODEL_ID})) }), + ) + .route("/v1/systemone", post(decide)) + .layer(DefaultBodyLimit::max(image_request::MAX_BODY)) + .with_state(Arc::new(Mutex::new(engine))); + let listener = tokio::net::TcpListener::bind((host.as_str(), port)).await?; + println!("native vision worker listening on {host}:{port}"); + axum::serve(listener, app).await?; + Ok(()) +} diff --git a/tests/cua_s1/fixtures/native_decode/generate.py b/tests/cua_s1/fixtures/native_decode/generate.py new file mode 100644 index 0000000..cf04db8 --- /dev/null +++ b/tests/cua_s1/fixtures/native_decode/generate.py @@ -0,0 +1,33 @@ +"""Tiny lossless PNG oracles; expectations checked with pinned Pillow 11.3.0.""" + +import struct +import zlib +from pathlib import Path + + +def chunk(kind, data): + return ( + struct.pack(">I", len(data)) + + kind + + data + + struct.pack(">I", zlib.crc32(kind + data)) + ) + + +for name, kind, pixels, transparent in [ + ("l16", 0, [1000], None), + ("l16-trns", 0, [1000], struct.pack(">H", 1000)), + ("la16", 4, [1000, 0], None), + ("rgb16", 2, [1000, 255, 32768], None), + ("rgb16-trns", 2, [1000, 255, 32768], struct.pack(">HHH", 1000, 255, 32768)), + ("rgba16", 6, [1000, 255, 32768, 0], None), +]: + data = b"\x89PNG\r\n\x1a\n" + chunk( + b"IHDR", struct.pack(">IIBBBBB", 1, 1, 16, kind, 0, 0, 0) + ) + if transparent is not None: + data += chunk(b"tRNS", transparent) + data += chunk( + b"IDAT", zlib.compress(b"\0" + struct.pack(">" + "H" * len(pixels), *pixels)) + ) + chunk(b"IEND", b"") + Path(__file__).with_name(name + ".png").write_bytes(data) diff --git a/tests/cua_s1/fixtures/native_decode/l16-trns.png b/tests/cua_s1/fixtures/native_decode/l16-trns.png new file mode 100644 index 0000000000000000000000000000000000000000..fc7f7b88afbf2196cfe4e4943b57846d48deb1b6 GIT binary patch literal 82 zcmeAS@N?(olHy`uVBq!ia0vp^j36ukBp9;Zxr+fQrjj7PVCEN}^S^8a^0+-+978JR bBq!7X6@FxRqx0Ha52V=B)z4*}Q$iB}shtzf literal 0 HcmV?d00001 diff --git a/tests/cua_s1/fixtures/native_decode/l16.png b/tests/cua_s1/fixtures/native_decode/l16.png new file mode 100644 index 0000000000000000000000000000000000000000..cd96bb28f5aeba7c94c3904c8f7a4c7c362a40a9 GIT binary patch literal 68 zcmeAS@N?(olHy`uVBq!ia0vp^j36ukBp9;Zxr+fQZci7-kcv6U33WiJj|^{gUYqNI Oq&;2zT-G@yGywpVmJV|O literal 0 HcmV?d00001 diff --git a/tests/cua_s1/fixtures/native_decode/la16.png b/tests/cua_s1/fixtures/native_decode/la16.png new file mode 100644 index 0000000000000000000000000000000000000000..e9a3c0c7d355b5a6bfa2a753e6dfa3d0a1e9d600 GIT binary patch literal 70 zcmeAS@N?(olHy`uVBq!ia0vp^j36w)0wkaITyO+SxHyJX%t=nDJ9vPBf$1c}8=n0a Rw*n;?JYD@<);T3K0RZr&5p@6n literal 0 HcmV?d00001 diff --git a/tests/cua_s1/fixtures/native_decode/rgb16-trns.png b/tests/cua_s1/fixtures/native_decode/rgb16-trns.png new file mode 100644 index 0000000000000000000000000000000000000000..f9f707e065e7ce127f0f6db6b4701ad1d092564f GIT binary patch literal 90 zcmeAS@N?(olHy`uVBq!ia0vp^j36w)1SAhU@1F~#*h+%@f|*}1{BK|=?G0@Qitu~7 lIEGZrNlvIc_+ft|BLl}grffU=ug`%}44$rjF6*2UngE#p7*GHJ literal 0 HcmV?d00001 diff --git a/tests/cua_s1/fixtures/native_decode/rgb16.png b/tests/cua_s1/fixtures/native_decode/rgb16.png new file mode 100644 index 0000000000000000000000000000000000000000..2f36ac7ea0e831a0264633e0e45f75efe14a033c GIT binary patch literal 72 zcmeAS@N?(olHy`uVBq!ia0vp^j36w)1SAhU@1F~#_&r@5Ln`JZC)6GMu)mR!fny$1 Uww?Xg=Rhe2Pgg&ebxsLQ04akKn*aa+ literal 0 HcmV?d00001 diff --git a/tests/cua_s1/fixtures/native_decode/rgba16.png b/tests/cua_s1/fixtures/native_decode/rgba16.png new file mode 100644 index 0000000000000000000000000000000000000000..1477c0f0c3354bddb3cee98d905a5c8f2f30fd1b GIT binary patch literal 74 zcmeAS@N?(olHy`uVBq!ia0vp^j36w)1|~CyrWMtqi WVan!{Ju3o~WbkzLb6Mw<&;$S!%Mla+ literal 0 HcmV?d00001 diff --git a/tests/cua_s1/test_native_vision_cuda.py b/tests/cua_s1/test_native_vision_cuda.py new file mode 100644 index 0000000..526291e --- /dev/null +++ b/tests/cua_s1/test_native_vision_cuda.py @@ -0,0 +1,114 @@ +"""GPU regression tests: python test_native_vision_cuda.py /path/libqwen3_5_cuda.so.""" + + +def main(): + import ctypes as C + import sys + + import torch + import torch.nn.functional as F + + lib = C.CDLL(sys.argv[1]) + P, I, Z = C.c_void_p, C.c_int, C.c_size_t + signatures = { + "cs1_vision_linear": [P, P, P, P, P, I, I, I, P], + "cs1_gemm_f32": [P, P, P, P, I, I, I, P], + "cs1_vision_bias": [P, P, Z, I, P], + "cs1_vision_norm": [P, P, P, P, I, I, P], + "cs1_vision_rope": [P, P, P, P, P, I, P], + "cs1_vision_attention": [P, P, P, P, I, P], + "cs1_vision_gelu": [P, Z, I, P], + "cs1_vision_lora_add": [P, P, Z, C.c_float, P], + "cs1_gemm_create": [Z], + } + for name, args in signatures.items(): + fn = getattr(lib, name) + fn.argtypes = args + fn.restype = P if name == "cs1_gemm_create" else I + + def run(name, *args): + rc = getattr(lib, name)( + *[a.data_ptr() if isinstance(a, torch.Tensor) else a for a in args] + ) + assert rc == 0, (name, rc) + + def close(got, want, atol=0.02, rtol=0.02): + torch.cuda.synchronize() + torch.testing.assert_close(got, want, atol=atol, rtol=rtol) + + torch.manual_seed(93) + s = torch.cuda.current_stream().cuda_stream + g = lib.cs1_gemm_create(32 << 20) + assert g + x = torch.randn(68, 1024, device="cuda", dtype=torch.bfloat16) + w = torch.randn(1024, 1024, device="cuda", dtype=torch.bfloat16) / 32 + b = torch.randn(1024, device="cuda", dtype=torch.bfloat16) + y = torch.empty_like(x) + run("cs1_vision_linear", g, x, w, b, y, 68, 1024, 1024, s) + close(y, F.linear(x, w, b), atol=0.015) + run("cs1_vision_norm", x, w[0], b, y, 68, 1024, s) + close(y, F.layer_norm(x, (1024,), w[0], b, 1e-6), atol=0.008) + for exact in (0, 1): + y.copy_(x) + run("cs1_vision_gelu", y, y.numel(), exact, s) + close( + y, + F.gelu(x, approximate="none" if exact else "tanh"), + atol=0.0001, + rtol=0.001, + ) + a = torch.randn(16, 1024, device="cuda") / 32 + xf = x.float() + r = torch.empty(68, 16, device="cuda") + run("cs1_gemm_f32", g, xf, a, r, 68, 16, 1024, s) + close(r, F.linear(xf, a), atol=2e-6, rtol=2e-5) + lora = torch.randn_like(xf) * 0.01 + y.copy_(x) + run("cs1_vision_lora_add", y, lora, y.numel(), 2.0, s) + close(y, (x.float() + 2 * lora).bfloat16(), atol=0, rtol=0) + for n in (4, 68, 256): + qkv = torch.randn(n, 3, 16, 64, device="cuda", dtype=torch.bfloat16) + angles = torch.randn(n, 32, device="cuda") + co, si = angles.cos(), angles.sin() + q, k = ( + torch.empty(n, 16, 64, device="cuda", dtype=torch.bfloat16), + torch.empty(n, 16, 64, device="cuda", dtype=torch.bfloat16), + ) + run("cs1_vision_rope", qkv, co, si, q, k, n, s) + co2, si2 = co.repeat(1, 2)[:, None], si.repeat(1, 2)[:, None] + + def rope(z, co2=co2, si2=si2): + z = z.float() + return ( + z * co2 + torch.cat((-z[..., 32:], z[..., :32]), -1) * si2 + ).bfloat16() + + close(q, rope(qkv[:, 0]), atol=0, rtol=0) + close(k, rope(qkv[:, 1]), atol=0, rtol=0) + out = torch.empty_like(q) + run("cs1_vision_attention", q, k, qkv[:, 2].data_ptr(), out, n, s) + want = F.scaled_dot_product_attention( + q.transpose(0, 1), k.transpose(0, 1), qkv[:, 2].transpose(0, 1) + ).transpose(0, 1) + close(out, want, atol=0.008, rtol=0.015) + # Conv3d rounds the convolution before its separately applied BF16 bias. + xp = torch.randn(4, 1536, device="cuda", dtype=torch.bfloat16) + wp = torch.randn(1024, 1536, device="cuda", dtype=torch.bfloat16) / 32 + bp = torch.randn(1024, device="cuda", dtype=torch.bfloat16) + z = torch.zeros_like(bp) + yp = torch.empty(4, 1024, device="cuda", dtype=torch.bfloat16) + run("cs1_vision_linear", g, xp, wp, z, yp, 4, 1024, 1536, s) + run("cs1_vision_bias", yp, bp, yp.numel(), 1024, s) + want = F.conv3d( + xp.reshape(4, 3, 2, 16, 16), + wp.reshape(1024, 3, 2, 16, 16), + bp, + stride=(2, 16, 16), + ).flatten(1) + close(yp, want, atol=0.008, rtol=0.001) + assert (yp.float() - want.float()).abs().mean().item() < 1e-5 + print("native vision CUDA primitives: PASS") + + +if __name__ == "__main__": + main() From 63f974f88d68a970a8fcdf744b06c7f4f733ca6b Mon Sep 17 00:00:00 2001 From: levius <2114377220@qq.com> Date: Fri, 2 Oct 2026 13:18:20 +0800 Subject: [PATCH 5/5] chore(cua_s1): limit PR diff to core implementation --- README.md | 3 +- .../plans/2026-10-02-native-vision.md | 13 - recipe/cua_s1/check_native_vision_http.py | 131 -- .../experiments/native-vision/README.md | 63 - .../experiments/native-vision/boundary.json | 86 - .../native-vision/source-manifest.json | 55 - .../experiments/native-vision/standard.json | 65 - recipe/cua_s1/native.md | 99 +- recipe/cua_s1/native_image_preprocess.md | 78 - recipe/cua_s1/native_multimodal.md | 128 -- recipe/cua_s1/native_vision.md | 64 - recipe/cua_s1/native_vision_cases.py | 119 -- .../cua_s1/requirements-native-validation.txt | 4 - recipe/cua_s1/validate_native_vision.py | 197 -- recipe/cua_s1/verify_native_vision.py | 89 - src/backends/cuda/qwen3_5/README.md | 5 - src/models/cua_s1/README.md | 4 +- src/models/cua_s1/native/Cargo.toml | 9 - .../cua_s1/native/examples/inspect_vision.rs | 22 - .../native/examples/multimodal_boundary.rs | 184 -- .../cua_s1/native/examples/native_vision.rs | 138 -- .../native/examples/preprocess_image.rs | 75 - .../cua_s1/native/examples/vision_stage.rs | 28 - src/models/cua_s1/native/src/image_request.rs | 85 - src/models/cua_s1/native/src/inputs.rs | 118 -- src/models/cua_s1/native/src/model.rs | 68 - src/models/cua_s1/native/src/multimodal.rs | 64 - src/models/cua_s1/native/src/provenance.rs | 23 - .../cua_s1/native/src/vision/geometry.rs | 35 - src/models/cua_s1/native/tests/multimodal.rs | 90 - .../fixtures/image_preprocess/README.md | 30 - .../fixtures/image_preprocess/generate.py | 119 -- .../fixtures/image_preprocess/manifest.json | 311 --- .../image_preprocess/preprocessor_config.json | 21 - .../cua_s1/fixtures/native_decode/generate.py | 33 - .../fixtures/native_decode/l16-trns.png | Bin 82 -> 0 bytes tests/cua_s1/fixtures/native_decode/l16.png | Bin 68 -> 0 bytes tests/cua_s1/fixtures/native_decode/la16.png | Bin 70 -> 0 bytes .../fixtures/native_decode/rgb16-trns.png | Bin 90 -> 0 bytes tests/cua_s1/fixtures/native_decode/rgb16.png | Bin 72 -> 0 bytes .../cua_s1/fixtures/native_decode/rgba16.png | Bin 74 -> 0 bytes tests/cua_s1/fixtures/vision/README.md | 31 - .../fixtures/vision/adapter-tensors.json | 702 ------ .../fixtures/vision/adapter_config.json | 48 - .../cua_s1/fixtures/vision/base-tensors.json | 1887 ----------------- tests/cua_s1/fixtures/vision/config.json | 104 - tests/cua_s1/image_preprocess.rs | 223 -- tests/cua_s1/test_native_vision_cuda.py | 114 - tests/cua_s1/vision_loader.rs | 466 ---- 49 files changed, 7 insertions(+), 6224 deletions(-) delete mode 100644 docs/superpowers/plans/2026-10-02-native-vision.md delete mode 100644 recipe/cua_s1/check_native_vision_http.py delete mode 100644 recipe/cua_s1/experiments/native-vision/README.md delete mode 100644 recipe/cua_s1/experiments/native-vision/boundary.json delete mode 100644 recipe/cua_s1/experiments/native-vision/source-manifest.json delete mode 100644 recipe/cua_s1/experiments/native-vision/standard.json delete mode 100644 recipe/cua_s1/native_image_preprocess.md delete mode 100644 recipe/cua_s1/native_multimodal.md delete mode 100644 recipe/cua_s1/native_vision.md delete mode 100644 recipe/cua_s1/native_vision_cases.py delete mode 100644 recipe/cua_s1/requirements-native-validation.txt delete mode 100644 recipe/cua_s1/validate_native_vision.py delete mode 100644 recipe/cua_s1/verify_native_vision.py delete mode 100644 src/models/cua_s1/native/examples/inspect_vision.rs delete mode 100644 src/models/cua_s1/native/examples/multimodal_boundary.rs delete mode 100644 src/models/cua_s1/native/examples/native_vision.rs delete mode 100644 src/models/cua_s1/native/examples/preprocess_image.rs delete mode 100644 src/models/cua_s1/native/examples/vision_stage.rs delete mode 100644 src/models/cua_s1/native/tests/multimodal.rs delete mode 100644 tests/cua_s1/fixtures/image_preprocess/README.md delete mode 100644 tests/cua_s1/fixtures/image_preprocess/generate.py delete mode 100644 tests/cua_s1/fixtures/image_preprocess/manifest.json delete mode 100644 tests/cua_s1/fixtures/image_preprocess/preprocessor_config.json delete mode 100644 tests/cua_s1/fixtures/native_decode/generate.py delete mode 100644 tests/cua_s1/fixtures/native_decode/l16-trns.png delete mode 100644 tests/cua_s1/fixtures/native_decode/l16.png delete mode 100644 tests/cua_s1/fixtures/native_decode/la16.png delete mode 100644 tests/cua_s1/fixtures/native_decode/rgb16-trns.png delete mode 100644 tests/cua_s1/fixtures/native_decode/rgb16.png delete mode 100644 tests/cua_s1/fixtures/native_decode/rgba16.png delete mode 100644 tests/cua_s1/fixtures/vision/README.md delete mode 100644 tests/cua_s1/fixtures/vision/adapter-tensors.json delete mode 100644 tests/cua_s1/fixtures/vision/adapter_config.json delete mode 100644 tests/cua_s1/fixtures/vision/base-tensors.json delete mode 100644 tests/cua_s1/fixtures/vision/config.json delete mode 100644 tests/cua_s1/image_preprocess.rs delete mode 100644 tests/cua_s1/test_native_vision_cuda.py delete mode 100644 tests/cua_s1/vision_loader.rs diff --git a/README.md b/README.md index c157738..35bc8b9 100644 --- a/README.md +++ b/README.md @@ -4,7 +4,7 @@ Documentation: A community-maintained inference engine for prefill-only System1-Omni models, designed around a Rust frontend, model-owned execution, and high-performance CUDA and Metal backends. -The Rust frontend forwards requests to a separately running model worker. Cua-S1 4B 0.2 has native text and screenshot workers with CUDA kernels in this repository; other in-repository model engines and GPU backends are not implemented yet. +The Rust frontend forwards requests to a separately running model worker. The Cua-S1 4B 0.2 `text` adapter has a native worker with CUDA kernels in this repository; other in-repository model engines and GPU backends are not implemented yet. ## Run the frontend @@ -59,7 +59,6 @@ LAYA can run as an external Python worker for text requests; its in-repository m | --- | --- | | LAYA | [External worker](recipe/laya/README.md); [CPU checkpoint reader](src/models/laya/README.md); model execution planned | | Cua-S1 4B 0.2 (`text` adapter) | [Python worker](recipe/cua_s1/text.md); [native worker](recipe/cua_s1/native.md), CUDA, run on sm_89 | -| Cua-S1 4B 0.2 (`multimodal` adapter) | [Native screenshot worker](recipe/cua_s1/native_multimodal.md), complete RGB/vision/language CUDA path, validated on RTX 4090 | CUDA and Metal coverage will be documented per model as implementations are added and validated. diff --git a/docs/superpowers/plans/2026-10-02-native-vision.md b/docs/superpowers/plans/2026-10-02-native-vision.md deleted file mode 100644 index 1e5fd9b..0000000 --- a/docs/superpowers/plans/2026-10-02-native-vision.md +++ /dev/null @@ -1,13 +0,0 @@ -# Cua-S1 native vision implementation plan - -Goal: finish native RGB-to-decision inference and verify against the pinned unmerged reference on RTX 4090. - -Architecture: combine #59 checkpoint validation, #63 CPU RGB preparation and #56 language boundary on an isolated integration branch. Keep original BF16 base and FP32 vision LoRA branches separate. Execute patch projection, interpolated positions, 24 bidirectional vision transformer blocks and merger on CUDA. Build single-image token positions in Rust and expose an end-to-end RGB API/example. Retain the text API. Reject invalid geometry before GPU execution. - -- [x] Inspect dependencies and GPU environment; preserve original branches. -- [x] Implement CUDA vision execution with CPU geometry tests and numerical GPU checks against Transformers stage outputs. -- [x] Implement native prompt image insertion, three-axis positions, RGB orchestration and candidate scoring; test invalid inputs and position oracles. -- [x] Export deterministic multi-geometry reference cases using pinned weights, unmerged BF16 and full FP32 controls, then run native GPU end-to-end comparison. Acceptance is the repository probability criterion: max error <= 2*BF16 reference error + 0.01, matching top choices at FP32 margin >= 0.05. Record intermediate errors without inventing tolerances after measurements. -- [x] Run workspace fmt, strict Clippy, tests, release build, CUDA regressions and independent spec/code review. Resolve actionable findings. Preserve commands and evidence in recipe documentation. - -The first implementation is eager; no acceleration claim. GPU validation must exercise actual native vision, never substitute exported reference image features. Large attention inputs require bounded-memory attention. Verification uses fresh source hashes and exact tool/environment versions. diff --git a/recipe/cua_s1/check_native_vision_http.py b/recipe/cua_s1/check_native_vision_http.py deleted file mode 100644 index 615e66f..0000000 --- a/recipe/cua_s1/check_native_vision_http.py +++ /dev/null @@ -1,131 +0,0 @@ -"""Check live native screenshot HTTP answers against native replay outputs.""" - -import argparse -import json -import math -import urllib.error -import urllib.request -from pathlib import Path - - -def main(): - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--url", default="http://127.0.0.1:28003") - parser.add_argument("--controls", required=True, type=Path) - parser.add_argument("--native", required=True, type=Path) - parser.add_argument("--out", required=True, type=Path) - args = parser.parse_args() - manifest = json.loads((args.controls / "manifest.json").read_text()) - native = json.loads(args.native.read_text()) - if not manifest["cases"] or len(manifest["cases"]) != len(native["cases"]): - raise ValueError("case coverage mismatch") - records = [] - - def post(raw): - request = urllib.request.Request( - args.url + "/v1/systemone", - data=raw, - headers={"Content-Type": "application/json"}, - ) - try: - with urllib.request.urlopen(request, timeout=180) as response: - return response.status, json.load(response) - except urllib.error.HTTPError as error: - return error.code, json.load(error) - - with urllib.request.urlopen(args.url + "/health") as response: - health = json.load(response) - assert ( - health["status"] == "ready" - and health["model"] - == "cua-ai/cua-s1-4b-0.2@16818868b0cc7813808aae4e87b417657046ab79:multimodal" - ) - for case, result in zip(manifest["cases"], native["cases"]): - assert case["case"] == result["case"] - assert len(case["questions"]) == len(result["questions"]) - assert [q["name"] for q in case["questions"]] == [ - q["name"] for q in result["questions"] - ] - raw = (args.controls / case["request_file"]).read_bytes() - body = json.loads(raw) - assert list(body["questions"]) == [q["name"] for q in case["questions"]] - status, answer = post(raw) - assert status == 200, answer - assert answer["model"] == health["model"] - assert answer["usage"]["input_tokens"] == sum( - len(q["input_ids"]) for q in case["questions"] - ) - assert answer["usage"]["output_tokens"] == 0 - assert list(answer["answers"]) == list(body["questions"]) - for q in result["questions"]: - probs = list(answer["answers"][q["name"]]["probabilities"].values()) - keys = list(body["questions"][q["name"]]["criteria"]) - response = answer["answers"][q["name"]] - assert list(response["probabilities"]) == keys - assert response["type"] == "choice" - assert ( - response["choice"] - == keys[max(range(len(probs)), key=probs.__getitem__)] - ) - confidence = ( - 1.0 - if len(probs) == 1 - else max( - 0.0, - 1.0 - + sum(p * math.log(p) for p in probs if p) / math.log(len(probs)), - ) - ) - assert abs(response["confidence"] - confidence) < 1e-7 - assert len(probs) == len(q["probabilities"]) - assert max(abs(a - b) for a, b in zip(probs, q["probabilities"])) < 1e-7 - records.append( - { - "case": case["case"], - "status": status, - "usage": answer["usage"], - "answers": answer["answers"], - } - ) - # Run input failures after successful requests to exercise the live loaded worker. - valid = json.loads( - (args.controls / manifest["cases"][0]["request_file"]).read_bytes() - ) - bad = [] - bad.append((b'{"a":1,"a":2}', 400)) - bad.append((b'{"x":NaN}', 400)) - bad.append((b"[]", 400)) - bad.append((b"{}", 422)) - invalid = json.loads(json.dumps(valid)) - invalid["state"]["image"] = "https://example.invalid/a.png" - bad.append((json.dumps(invalid).encode(), 422)) - invalid = json.loads(json.dumps(valid)) - next(iter(invalid["questions"].values()))["instructions"] = "<|image_pad|>" - bad.append((json.dumps(invalid).encode(), 422)) - invalid = json.loads(json.dumps(valid)) - invalid["questions"] = { - str(i): next(iter(valid["questions"].values())) for i in range(9) - } - bad.append((json.dumps(invalid).encode(), 422)) - invalid = json.loads(json.dumps(valid)) - invalid["extra"] = True - bad.append((json.dumps(invalid).encode(), 422)) - bad.append((b" " * (8 * 1024 * 1024 + 1), 413)) - for raw, expected in bad: - status, response = post(raw) - assert status == expected and isinstance(response.get("detail"), str), ( - status, - response, - ) - args.out.write_text( - json.dumps( - {"passed": True, "requests": records, "invalid_inputs_checked": len(bad)}, - indent=2, - ) - + "\n" - ) - print(f"PASS: {len(records)} screenshot requests and {len(bad)} input failures") - - -if __name__ == "__main__": - main() diff --git a/recipe/cua_s1/experiments/native-vision/README.md b/recipe/cua_s1/experiments/native-vision/README.md deleted file mode 100644 index d05083a..0000000 --- a/recipe/cua_s1/experiments/native-vision/README.md +++ /dev/null @@ -1,63 +0,0 @@ -# Native vision validation — 2026-10-02 - -Validated on NVIDIA GeForce RTX 4090 (sm_89, 24 GiB), driver 595.71.05, -CUDA toolkit 13.2; Python 3.12.3, Torch 2.14.0+cu130, Transformers 5.17.0, -PEFT 0.21.0, Pillow 11.3.0. Reference controls use unmerged adapters in BF16 -and full FP32 with TF32 disabled. Weights match the pinned upstream lock. -Native execution uses unmerged FP32 vision LoRA and merged BF16 language LoRA. - -| Set | Questions | Maximum native probability error | BF16 reference error | Allowed error | Matching choices | -| --- | ---: | ---: | ---: | ---: | ---: | -| [Standard](standard.json) | 8 | 0.00331324 | 0.00526386 | 0.02052773 | 8/8 | -| [Boundary](boundary.json) | 11 | 0.07950398 | 0.09646034 | 0.20292068 | 11/11 | - -The sets are evaluated separately so adding a difficult boundary case does not -relax the standard set's threshold. Every FP32 margin exceeds 0.05. Token IDs, -image grids and all three position axes match exactly. Repeated language calls -and vision reruns after language execution produce identical native outputs. -Standard PNG RGB pixels match Pillow; JPEG differs by up to 3 intensity levels. - -The standard set includes square, landscape, portrait, JPEG, 1/26 choices, -structured Chinese instructions and two questions. Boundary images include -1024×1024, 200×1, 383×257 and 1×1, the latter with eight questions. These are a -finite synthetic validation corpus, not a general accuracy benchmark. - -Vision stage checks identified and fixed Conv3d patch projection rounding: -FP32 accumulation is rounded to BF16 before the patch bias. On the standard -small image, patch mean absolute error is 2.33e-8; final feature RMS error is -0.005306 and maximum error 0.632813. Feature outputs are not bitwise identical. -Stage downloads are diagnostic and are excluded from any performance claim. -No speedup claim is made. - -## Verification - -- Workspace format, strict all-target Clippy, Rust tests and release build. -- 63 CPU Rust tests passed; 7 tests requiring external GPU/checkpoint data are - skipped by ordinary CPU tests. Five Cua-S1 GPU regressions were executed - separately, including graph misses/hits, changed tokens, eviction, allocation - growth and multimodal/text mixing, kernel oracles and image insertion. -- CUDA primitives compared with PyTorch: patch Conv3d, biased projection, - LayerNorm, FP32 LoRA GEMMs/addition, rotary, both GELUs and bidirectional attention. -- Live HTTP: 11 valid requests, 19 question outputs equal to the direct native - replay (tolerance 1e-7); 9 malformed/unsupported/body-limit cases checked on - each set. Model identity, usage and `detail` envelopes checked. -- Six 16-bit PNG modes independently checked against Pillow, including L16 and - RGB16 with transparency. Same-size weight corruption and unlisted overrides - are rejected by provenance tests. - -See [reproduction instructions](../../native_multimodal.md). The case generator, -reference exporter, native replay, HTTP checker and acceptance verifier are all -included in this branch. Raw controls, per-stage BF16 tensors and logs are retained -in `/root/cua-native-vision-20261002/` on the authorized GPU host and the local -`artifacts/cua-native-vision-20261002/` evidence directory. The language checkpoint -reuses the previously verified multimodal export without changing its weights; -its tokenizer and SHA-256 manifest are saved in a new task directory. - -## Review - -Independent review covered specification and implementation. Five findings were -fixed and re-reviewed: pinned source/export provenance, pre-decode PNG allocation -limits, 16-bit PNG conversions, public question invariants, and CUDA device -selection when moving vision execution between threads. A final review also fixed -HTTP checker coverage truncation and added choice/confidence assertions. No unresolved actionable -finding remained in that review. Contributor and maintainer review remain separate. diff --git a/recipe/cua_s1/experiments/native-vision/boundary.json b/recipe/cua_s1/experiments/native-vision/boundary.json deleted file mode 100644 index 4a26afc..0000000 --- a/recipe/cua_s1/experiments/native-vision/boundary.json +++ /dev/null @@ -1,86 +0,0 @@ -{ - "passed": true, - "questions": 11, - "max_native_error": 0.07950398325920105, - "max_bf16_reference_error": 0.09646034240722656, - "allowance": 0.20292068481445313, - "choices": [ - { - "case": "maximum", - "question": "q0", - "fp32_margin": 0.19149160385131836, - "same_choice": true, - "max_probability_error": 0.01613593101501465 - }, - { - "case": "narrow", - "question": "q0", - "fp32_margin": 0.9099561609327793, - "same_choice": true, - "max_probability_error": 0.0005906363949179649 - }, - { - "case": "noise", - "question": "q0", - "fp32_margin": 0.22005164623260498, - "same_choice": true, - "max_probability_error": 0.07950398325920105 - }, - { - "case": "tiny-eight", - "question": "q0", - "fp32_margin": 0.7194288372993469, - "same_choice": true, - "max_probability_error": 0.02142578363418579 - }, - { - "case": "tiny-eight", - "question": "q1", - "fp32_margin": 0.8431737795472145, - "same_choice": true, - "max_probability_error": 0.01713651418685913 - }, - { - "case": "tiny-eight", - "question": "q2", - "fp32_margin": 0.7816882357001305, - "same_choice": true, - "max_probability_error": 0.018895387649536133 - }, - { - "case": "tiny-eight", - "question": "q3", - "fp32_margin": 0.7832479327917099, - "same_choice": true, - "max_probability_error": 0.03977161645889282 - }, - { - "case": "tiny-eight", - "question": "q4", - "fp32_margin": 0.7988598495721817, - "same_choice": true, - "max_probability_error": 0.027261078357696533 - }, - { - "case": "tiny-eight", - "question": "q5", - "fp32_margin": 0.7969894707202911, - "same_choice": true, - "max_probability_error": 0.0319133996963501 - }, - { - "case": "tiny-eight", - "question": "q6", - "fp32_margin": 0.7931957244873047, - "same_choice": true, - "max_probability_error": 0.026585102081298828 - }, - { - "case": "tiny-eight", - "question": "q7", - "fp32_margin": 0.7984120920300484, - "same_choice": true, - "max_probability_error": 0.03071284294128418 - } - ] -} diff --git a/recipe/cua_s1/experiments/native-vision/source-manifest.json b/recipe/cua_s1/experiments/native-vision/source-manifest.json deleted file mode 100644 index f185586..0000000 --- a/recipe/cua_s1/experiments/native-vision/source-manifest.json +++ /dev/null @@ -1,55 +0,0 @@ -{ - "source_base": "6c40bba", - "files_sha256": { - "src/backends/cuda/qwen3_5/README.md": "c56d6505658914c68aa0a5e44695a8cb523bb8909641f1dd8cbeb51db2b28d21", - "src/backends/cuda/qwen3_5/attention.cu": "a8c19aacf46d39fca5461a0a43d99a1d224e6a45e8b1eb62ec7db6fe7f753054", - "src/backends/cuda/qwen3_5/build.sh": "178d94396fdb46ee929f0da919f120683990ed7af8fd079796106ab5995b3f0f", - "src/backends/cuda/qwen3_5/common.cuh": "ea924f7c903b6d14b0e1f10c8c33a1a9ca23806df093628a178661a303d54d55", - "src/backends/cuda/qwen3_5/elementwise.cu": "679bbce2ed940eb41620c4044c5e360293b5df077e9821001a9ca6144cf106b9", - "src/backends/cuda/qwen3_5/gdn_prefill.cu": "1c9027b6f2d3677381d8d19611ecdf7e138718459cd61f341802cc2047a582be", - "src/backends/cuda/qwen3_5/gemm.cu": "521282cb1f2fd9236301cf32d3af38aaf50c943e6e8c0d13c3e9f5dfc9c86596", - "src/backends/cuda/qwen3_5/mma.cuh": "64654f8b937585805daa72b692c243ae35d34b746e2cb6b0b55c255748495764", - "src/backends/cuda/qwen3_5/norm.cu": "7b0361f93933fddcb1615531217d77a3cdee005146218599089db06283e18c54", - "src/backends/cuda/qwen3_5/ops.h": "985f98191329023fd15195c5f93aafda24718ba8dbdab278575cf576a8d144fa", - "src/backends/cuda/qwen3_5/runtime.cu": "993ead3c2120c9b594a479f59c1666cf6bf59c046fd3abe4af04264b887a8f5d", - "src/backends/cuda/qwen3_5/vision.cu": "1e5d643e8da15da001641e40272a84dacc876be29d66aeb3e4950a52930abc75", - "src/models/cua_s1/multimodal/graph_runtime.py": "c0572ceb03dcb7f14364b4bdc697d0435c3cb739d387b72711501d6478e68e17", - "src/models/cua_s1/multimodal/model.py": "4cfe9d69d4083f502cf1b41971fdeb2115f5f1b4a00316e1f5a89d9e5a36ece0", - "src/models/cua_s1/multimodal/protocol.py": "eefc421dc485ab0a23d475a3a733b5d46f33cbe9febd6295e5d794e7fdfcec5f", - "src/models/cua_s1/native/Cargo.toml": "4fa7afd109ecb3feab6941510f757dd2984ae25ad804ca9a8a6603ca5fccbd9e", - "src/models/cua_s1/native/THIRD_PARTY_NOTICES.md": "b697db90fec79f2b2843445f526b8b21e7b6e93b91d78dfe4d9632509434cc08", - "src/models/cua_s1/native/examples/inspect_vision.rs": "2db1c38e78ac13a09fb6293b5f3dae86b0c8d2d1a44d80e76c36a3b1cafa78a0", - "src/models/cua_s1/native/examples/multimodal_boundary.rs": "c2503b83f69bf6ab5383d7058f0444d0d8786c63a6d80d9254e9d81c090c62ac", - "src/models/cua_s1/native/examples/native_vision.rs": "cd2b9f7408ff155db117bd556dbd2e032fb97eff36d49ac6d05e5c0c88cc403a", - "src/models/cua_s1/native/examples/preprocess_image.rs": "a401f4d8bc1aae4656f9de7ae671f16e8721a9fa3acb4bfeef22a4ba2d862480", - "src/models/cua_s1/native/examples/vision_stage.rs": "c0c70338f63fb97e4ab7409c875fa2b5d676bff88438252dc1a4b67227d9a5ae", - "src/models/cua_s1/native/licenses/APACHE-2.0": "77fd4710def9ec3c0f6225800e0235f15a425abd4a8b03559127fcd782612049", - "src/models/cua_s1/native/licenses/PILLOW-LICENSE": "17f240ae101143707e5e7303a5d800450d9ccc7475b463cedb555cefdb3c6ece", - "src/models/cua_s1/native/licenses/PYTORCH-LICENSE": "bd018feef8825e88181c84eb7e3aa4eafb8f08a20d9fd6ef948569610c4a3e43", - "src/models/cua_s1/native/src/contract.rs": "4278a6eaf2bfd8e76aa2287e1af9c5887e7092fd55750ff74dd0f0f23eb7af95", - "src/models/cua_s1/native/src/cuda.rs": "fe3c7011bf9cafa21eec6463eec5a633767336e87b797f1227a91f62799125af", - "src/models/cua_s1/native/src/engine.rs": "46d20775d6bd691d5375e50fb5f0aa19d5f4b6b87d59605413b393603084ed1b", - "src/models/cua_s1/native/src/image_preprocess.rs": "cda77fae0931a186158856316392b35abbcb3d71a4db387cc3e8d7135f615408", - "src/models/cua_s1/native/src/image_request.rs": "cf3957050c53b6faf131cae49721fbb734bf059e01a209dfb74a458814103142", - "src/models/cua_s1/native/src/inputs.rs": "e3c376366cdd1f9e941b6f9e4fa32010448048a6282ba8d8a99c6c4c99a09565", - "src/models/cua_s1/native/src/json.rs": "5b78af18f737ec830d3acd2e052fb7c70829aec6a7a6603a87b731635e6a9610", - "src/models/cua_s1/native/src/lib.rs": "fe1c85c646b4332359802b4f5f352daa1c7e17c49021edb61f3b20e01c313fb3", - "src/models/cua_s1/native/src/main.rs": "2066930417c9b047d40771dd8a7949d266503e09dc6b1b76e9908760b7f2bf34", - "src/models/cua_s1/native/src/model.rs": "d2308b86fac452d07b755bf56e710503964dbd3fd57c58eacf20804696cc9267", - "src/models/cua_s1/native/src/multimodal.rs": "ef56d6e7b293ca4ed4979c129a203d0ed9a34bea7b361a43d9e249b5d4f5f83b", - "src/models/cua_s1/native/src/provenance.rs": "d60f06da537b2be62acb9b612fba718ff33cea730d6f9c9f53113760c4abe65d", - "src/models/cua_s1/native/src/vision/geometry.rs": "02d1f68c53a4cb07a8d3bf7d3adff081093c2504a721a994277bc52bc6b4a0b9", - "src/models/cua_s1/native/src/vision/mod.rs": "057accf927a3c480a3b662c08ad1645f7f54e6b07ce370ecf8785252046f6417", - "src/models/cua_s1/native/src/vision/model.rs": "806a9815b891bf2c2365d3a435b42c22084a5fe6ac0d0919ee15bc9504db1c80", - "src/models/cua_s1/native/src/vision_engine.rs": "ad24024aaf7483ba5e482a45e34d6280c401ae42d171ef0a12508d4b4e31f487", - "src/models/cua_s1/native/src/vision_main.rs": "3ee30185429b2cb9f008481d635eabb2d83b20f16d6c651e3d050a8eec3f034a", - "src/models/cua_s1/native/tests/kernels.rs": "c585b28ba48f1951f540e02d3a8b3200ec89e8604081acfe46c7667289be0bb3", - "src/models/cua_s1/native/tests/multimodal.rs": "202262ae6d19b0e32e399d360812666fe68865c80598f26a620839a08e1d2bd0" - }, - "executed_binaries_sha256": { - "/root/cua-native-vision-20261002/build/libqwen3_5_cuda.so": "0f5f458f96b14b4a7ed0e06aa949cc833a1044e4e0ee54b03371e11e07e05b13", - "/root/pr52-review/code/target/release/examples/native_vision": "302e86a31c7ec11546e538725e619535e076d74809ee227a107744133befe092", - "/root/pr52-review/code/target/release/omni-cua-s1-vision": "5a5c565b5bd746e49162dd56723f8944130ec1568a966c1a34ff89524043bbb3" - }, - "reference": "full BF16 and FP32, pinned model.py plus Transformers 5.17.0, source and raw outputs retained in evidence" -} diff --git a/recipe/cua_s1/experiments/native-vision/standard.json b/recipe/cua_s1/experiments/native-vision/standard.json deleted file mode 100644 index a18cfe7..0000000 --- a/recipe/cua_s1/experiments/native-vision/standard.json +++ /dev/null @@ -1,65 +0,0 @@ -{ - "passed": true, - "questions": 8, - "max_native_error": 0.0033132433891296387, - "max_bf16_reference_error": 0.005263864994049072, - "allowance": 0.020527729988098146, - "choices": [ - { - "case": "26-options", - "question": "next", - "fp32_margin": 0.4796072989702225, - "same_choice": true, - "max_probability_error": 0.0033132433891296387 - }, - { - "case": "jpeg", - "question": "next", - "fp32_margin": 0.9943915181793272, - "same_choice": true, - "max_probability_error": 0.0004887116374447942 - }, - { - "case": "portrait", - "question": "next", - "fp32_margin": 0.9951829832280055, - "same_choice": true, - "max_probability_error": 0.0002958049299195409 - }, - { - "case": "single-option", - "question": "next", - "fp32_margin": 1.0, - "same_choice": true, - "max_probability_error": 0.0 - }, - { - "case": "small", - "question": "next", - "fp32_margin": 0.994548179442063, - "same_choice": true, - "max_probability_error": 0.00012633949518203735 - }, - { - "case": "two-questions", - "question": "next", - "fp32_margin": 0.994548179442063, - "same_choice": true, - "max_probability_error": 0.00012633949518203735 - }, - { - "case": "two-questions", - "question": "second", - "fp32_margin": 0.997789517743513, - "same_choice": true, - "max_probability_error": 5.620718002319336e-05 - }, - { - "case": "wide", - "question": "next", - "fp32_margin": 0.9933466790243983, - "same_choice": true, - "max_probability_error": 7.712841033935547e-05 - } - ] -} diff --git a/recipe/cua_s1/native.md b/recipe/cua_s1/native.md index 1c93b01..018feff 100644 --- a/recipe/cua_s1/native.md +++ b/recipe/cua_s1/native.md @@ -24,13 +24,11 @@ CUA_S1_MODEL=weights/cua-s1-4b-0.2-text-merged target/release/omni-cua-s1-native ``` For the local CUDA Graph experiment, also set `CUA_S1_GRAPH=1`. The first use of -each exact prompt length warms the GEMM plans and captures the forward pass. -The first call returns the eager result; later requests replay it after fresh -token embedding. The graph contains the language layers, which update residuals -in place, so a cache miss must not replay those layers over its eager result. -At most eight lengths are cached. Growing the scratch allocation clears the captures before freeing +each exact prompt length warms the GEMM plans and captures the forward pass; +later requests replay it with freshly uploaded token ids. At most eight lengths +are cached. Growing the scratch allocation clears the captures before freeing their buffers. Capture adds first-use latency; leave the variable unset to use -the eager control. Rebuild both the worker and CUDA library together (ABI 4). +the eager control. Rebuild both the worker and CUDA library together (ABI 3). If capture fails, the worker returns the completed eager result and disables Graph capture/replay for its remaining lifetime, logging the failure to stderr. @@ -43,92 +41,3 @@ cargo test -p omni-cua-s1-native CUA_S1_CUDA_LIB=$PWD/target/release/libqwen3_5_cuda.so \ cargo test --release -p omni-cua-s1-native --test kernels -- --ignored ``` - -For CPU-only structural inspection of the separate vision base weights and -multimodal adapter, see [native vision checkpoint inspection](native_vision.md). - -For the separate decoded-RGB8 CPU preprocessing API and example, see -[native image preprocessing](native_image_preprocess.md). -## Multimodal language boundary - -`Model::forward_multimodal` consumes token IDs, adapted BF16 image features -`[image_tokens, hidden_size]`, their sorted placeholder indices, and the T/H/W -slices of int64 `position_ids [3, 1, sequence]`. It returns the last position's -final-normalized hidden state, like `Model::forward`. - -Inputs are one unpadded prompt. Every image placeholder must have exactly one -feature row; all other rows come from the token embedding table. The caller -calculates positions and runs the image processor, vision tower and vision LoRA. -Positions must be nonnegative and below `max_position_embeddings`. The language -path uses Qwen3.5's interleaved MRoPE sections, not three contiguous rotary blocks. -Text calls and their captured graphs use immutable text-position tables; -multimodal calls use separate device tables, so returning to text requires no -host table rebuild or restoration copy. Multimodal calls execute eagerly even -when `CUA_S1_GRAPH=1`. The extra tables use -`4 * scratch_capacity * rotary_half` bytes (2 MiB at 16,384 rows). - -This is a Rust model API for integrating a vision producer. The HTTP worker -above continues to serve the text adapter. For native vision and image HTTP requests use the separate -[screenshot worker](native_multimodal.md). Padding, video and batching remain unsupported. - -### Prepare a matching language checkpoint - -The language weights must contain the **multimodal** adapter, not the `text` -adapter. In the pinned reference environment, with upstream-verified weights, -export just the merged language model (about 7.5 GB) to a new directory: - -```sh -PYTHONPATH=src .venv/bin/python recipe/cua_s1/export_multimodal_language.py \ - --base weights/Qwen3.5-4B --adapter weights/cua-s1-4b-0.2/multimodal \ - --out weights/cua-s1-4b-0.2-multimodal-language-merged -``` - -No `cua_s1_export.json` text-worker marker is created. The low-level `Model` API -does not verify checkpoint provenance; retain the export metadata and use the -matching adapter for the supplied features. Standalone language safetensors -names and the existing full-model prefixes are supported. - -### Replay a reference boundary - -This optional example consumes the `cua-s1-multimodal-reference-v1` format -from [#53](https://github.com/ThinkFlowLab/system1-omni/pull/53), which is still -open. The exporter and checksum verifier are not yet available on `main`. -Use a separate checkout of exporter revision -`1b64fa2ceb0a82b6a66a69ecdc9bc5cc1b1a0b66` to generate and verify the bundle; -the Rust model API itself does not depend on that PR being merged. -Its eight questions include different image grids and question lengths, JPEG, -structured/non-ASCII text, and 1/3/26 candidates. Then run: - -```sh -CUA_S1_CUDA_LIB=$PWD/target/release/libqwen3_5_cuda.so \ - cargo run --release --locked -p omni-cua-s1-native \ - --example multimodal_boundary -- \ - weights/cua-s1-4b-0.2-multimodal-language-merged \ - /path/to/verified-reference-bundle /tmp/native-language.json - -CUA_S1_MODEL=$PWD/weights/cua-s1-4b-0.2-multimodal-language-merged \ -CUA_S1_CUDA_LIB=$PWD/target/release/libqwen3_5_cuda.so \ - cargo test --release --locked -p omni-cua-s1-native \ - --test multimodal -- --ignored - -# Compare graph misses, hits, eviction and scratch growth with eager hidden states. -CUA_S1_MODEL=$PWD/weights/cua-s1-4b-0.2-multimodal-language-merged \ -CUA_S1_CUDA_LIB=$PWD/target/release/libqwen3_5_cuda.so \ - cargo test --release --locked -p omni-cua-s1-native \ - --lib graph_tests::graph_misses_hits_eviction_growth_and_multimodal_match_eager -- --ignored -``` - -The example checks repeated native hidden-state equality and writes last hidden -states, candidate logits and probabilities. It uses the text engine's FP32 -letter-row readout with FP64 accumulation. Output must be a new file. Verify -bundle integrity before invoking the example; it checks tensor shapes and input -contracts but is not the bundle checksum verifier. - -For accuracy validation, compare against an unmerged FP32 **language** control -with TF32 disabled, feeding the same fixed exported embeddings and positions. -Use the [declared native tolerance](../../src/models/cua_s1/README.md#validation): -maximum probability error over the set must be at most twice the BF16 reference -error plus 0.01, and the top option must match for FP32 margins at least 0.05. -The FP32 control starts after the BF16-exported vision boundary; it does not -validate a full FP32 vision pipeline. Native kernel and LoRA-merge rounding can -change hidden states and logits; bitwise equality to Transformers is not claimed. diff --git a/recipe/cua_s1/native_image_preprocess.md b/recipe/cua_s1/native_image_preprocess.md deleted file mode 100644 index eb63043..0000000 --- a/recipe/cua_s1/native_image_preprocess.md +++ /dev/null @@ -1,78 +0,0 @@ -# Native CPU image preprocessing - -The native crate exposes `image_preprocess::preprocess_rgb8(width, height, rgb)` -for **already decoded, interleaved RGB8** data. It prepares the image tensor for -the fixed Qwen3.5-4B / Cua-S1 4B processor. It does not decode PNG/JPEG, fetch -URLs, handle HTTP requests, run the vision encoder, or use a GPU. The existing -native text worker remains separate. - -The input must contain exactly `width * height * 3` bytes, in row-major RGB -order. Both dimensions must be nonzero and at most 2048; the area must be at -most 1,048,576 pixels and the aspect ratio at most 200. The library checks -geometry, lengths, and allocation arithmetic before creating image buffers. -These are input limits; smart resize can produce a side longer than 2048 for -very narrow inputs. - -The fixed processor uses a factor of 32, minimum area 65,536 and maximum area -16,777,216. Smart resize follows Python ties-to-even rounding and floating-point -square-root scaling with floor/ceil. Resampling matches the CPU torchvision -uint8 bicubic antialias path: Keys cubic coefficient `a = -0.5`, float64 weights, -per-axis int16 fixed-point coefficients, horizontal then vertical passes, and -rounding/clamping to uint8 after each pass. Unchanged axes bypass resampling. -The maximum-area downscale branch is retained for parity with the processor, -although the smaller input cap makes it unreachable through this API. - -`ProcessedImage` contains: - -- `pixel_values: Vec`, contiguous `[patches, 1536]` values normalized as - `(pixel - 127.5) / 127.5` using float32 operations. -- `image_grid_thw: [usize; 3]`, equal to `[1, resized_height / 16, resized_width / 16]`. -- `resized_width` and `resized_height`. -- `image_tokens()`, the patch count divided by four for the 2×2 spatial merge. - -Packing order is `block_y, block_x, merge_y (2), merge_x (2), channel (3), -temporal repeat (2), patch_y (16), patch_x (16)`. Each single image is repeated -across the two temporal positions. There is no video input support. - -## Run without a GPU - -From the repository root, provide a raw RGB8 file and its dimensions: - -```sh -cargo run --locked -p omni-cua-s1-native --example preprocess_image -- \ - 256 256 image.rgb pixel_values.f32 -``` - -The example prints the tensor shape, grid, resized dimensions and image token -count. The optional fourth argument writes every output float as little-endian -float32, with no header. It validates argument count, decimal dimensions and -exact file length, and bounds the input read. No model weights, Python, CUDA -library or image decoder are needed for this command. - -## Reference and validation - -Reference hashes are produced by the actual Hugging Face `AutoImageProcessor` -on CPU, with Python 3.12 and these exact package pins: Transformers 5.17.0, -PyTorch 2.14.0, torchvision 0.29.0, NumPy 2.5.3 and Pillow 11.3.0. The processor -configuration is the Qwen3.5-4B file at revision -`851bf6e806efd8d0a36b00ddf55e13ccb7b8cd0a` (see the fixture manifest for its URL -and SHA-256). Fixtures and regeneration instructions live in -[`tests/cua_s1/fixtures/image_preprocess/`](../../tests/cua_s1/fixtures/image_preprocess/). - -```sh -cargo test --locked -p omni-cua-s1-native --test image_preprocess -cargo test --release --locked -p omni-cua-s1-native --test image_preprocess -``` - -The tests compare input hashes, shape/grid metadata and SHA-256 of **every -little-endian float32 output byte** for 14 deterministic images. Cases include -tiny images, noise, ramps, checkerboards, ties-to-even dimensions, unchanged -axes, one- and two-axis resizes, extreme aspect ratios and the input area cap. -Additional checks cover all RGB byte values, channel/temporal/patch order, -constant images, inclusive limits and malformed geometry/buffers. This -validates the pinned CPU preprocessing behavior; it does not establish CUDA -preprocessing, image decoding, vision inference, or end-to-end model parity. - -The implementation adapts upstream algorithms; retained attributions and -license texts are in -[`THIRD_PARTY_NOTICES.md`](../../src/models/cua_s1/native/THIRD_PARTY_NOTICES.md). diff --git a/recipe/cua_s1/native_multimodal.md b/recipe/cua_s1/native_multimodal.md deleted file mode 100644 index 424cd3a..0000000 --- a/recipe/cua_s1/native_multimodal.md +++ /dev/null @@ -1,128 +0,0 @@ -# Cua-S1 native screenshot worker - -`omni-cua-s1-vision` serves the same screenshot `/v1/systemone` contract as the -[Python multimodal model](../../src/models/cua_s1/multimodal/model.py). Runtime request handling, PNG/JPEG decoding, -RGB preprocessing, the 24-block vision encoder, 2×2 merger, image insertion, -three-axis positions, language forward and candidate scoring execute in Rust/CUDA. -Python is used only for the one-time language export and reference validation. - -The worker combines the checkpoint loader (#59), RGB processor (#63) and language -input boundary (#56). These dependencies are included in the integration branch; -merging this branch does not require users to manually combine working trees. -The existing `omni-cua-s1-native` binary continues to serve the text adapter. - -## Build and load - -Download the pinned base and adapter revisions using the commands in the -[text recipe](text.md). Use its Python 3.12 environment with the additional image -packages and save the upstream lock beside the base directory: - -```sh -.venv/bin/python -m pip install -r recipe/cua_s1/requirements-native-validation.txt -curl --fail -L https://raw.githubusercontent.com/trycua/cua/0e75660ce4c2edda519e0c795fa3ad98abf4e76f/libs/cua-s1/ci/weights.lock.json \ - -o weights/weights.lock.json -``` - -The exporter and native worker verify this lock's exact trusted SHA-256. Export -the language weights once to a new directory: - -```sh -PYTHONPATH=src .venv/bin/python recipe/cua_s1/export_multimodal_language.py \ - --base weights/Qwen3.5-4B \ - --adapter weights/cua-s1-4b-0.2/multimodal \ - --out weights/cua-s1-multimodal-language -src/backends/cuda/qwen3_5/build.sh target/release 89 -cargo build --release --locked -p omni-cua-s1-native --bins --examples -CUA_S1_BASE=weights/Qwen3.5-4B \ -CUA_S1_VISION_ADAPTER=weights/cua-s1-4b-0.2/multimodal \ -CUA_S1_MODEL=weights/cua-s1-multimodal-language \ -CUA_S1_CUDA_LIB=$PWD/target/release/libqwen3_5_cuda.so \ - target/release/omni-cua-s1-vision -``` - -`CUA_S1_HOST` and `CUA_S1_PORT` default to `127.0.0.1:8000`. `/health` reports ready -after loading. The library and worker must both use **CUDA ABI 4**; rebuild both. -RTX 4090 / sm_89 is validated. Tensor-core operations require sm_80 or newer. - -At startup the worker verifies the upstream lock's trusted SHA-256, all pinned -base/adapter files, and the language export's per-file hashes before CUDA loading. -The export manifest is locally generated, trusted provenance; it is not a signature -of an externally supplied checkpoint. Keep all checkpoint and manifest files -immutable during inference. Low-level `VisionCheckpoint` and `VisionModel` APIs -perform structural validation; callers of those APIs own artifact provenance. - -The original vision tensors stay BF16 and the 50 rank-16 LoRA pairs stay FP32. -Each LoRA branch uses FP32 GEMMs, scaling and addition before BF16 rounding. -Language LoRA uses the existing merged BF16 path. The encoder uses bounded-memory, -bidirectional attention; it never allocates a full image-token-square score matrix. -Multimodal execution remains eager. This is a correctness implementation, with no -throughput or speedup claim. - -## Request and API behavior - -The screenshot contract accepts one inline PNG/JPEG, up to eight choice questions, -1–26 options each, a body up to 8 MiB and compressed image up to 4 MiB. Images are -bounded to 2048 per side, 1,048,576 pixels, and aspect ratio 200:1; animated PNGs -are rejected. PNG metadata allocations are bounded before decoding. Six 16-bit PNG -modes, including grayscale transparency, have Pillow-compatible RGB conversions. - -All questions and tokenized prompt lengths (at most 4096) are checked before -vision execution. A request preprocesses and encodes its image once, then scores -its questions independently. Features are request-local. Changed images, prompt -lengths and language/vision scratch reuse are included in validation. Malformed -JSON returns 400; invalid screenshot/question inputs return 422, body-limit -violations 413, and inference failures 500, all with a `detail` envelope. - -`VisionEngine::prepare` and `predict` accept decoded RGB requests. The lower-level -`VisionModel::forward(&ProcessedImage)` returns row-major BF16 `[image_tokens,2560]`. -`forward_with_trace` exposes synchronized intermediate downloads for diagnostics. -`prepare_prompt` constructs the token sequence and the T/H/W positions. Video, -padding, batching and CUDA Graph capture of vision are outside this worker. - -JPEG decoding uses the Rust decoder; it is not byte-identical to Pillow/libjpeg. -The validation JPEG differed by at most 3 intensity levels per channel and its -complete native probabilities passed the same acceptance test. PNG fixture pixels -were identical. CPU preprocessing parity for identical decoded RGB remains the -separate [RGB processor guarantee](native_image_preprocess.md). - -## Reproduce GPU alignment - -Generate the standard and independently added boundary requests with Pillow 11.3.0: - -```sh -.venv/bin/python recipe/cua_s1/native_vision_cases.py /tmp/cua-standard -.venv/bin/python recipe/cua_s1/native_vision_cases.py /tmp/cua-boundary --boundary -PYTHONPATH=src .venv/bin/python recipe/cua_s1/validate_native_vision.py \ - --base weights/Qwen3.5-4B --adapter weights/cua-s1-4b-0.2/multimodal \ - --requests /tmp/cua-standard --out /tmp/cua-standard-controls -CUA_S1_CUDA_LIB=$PWD/target/release/libqwen3_5_cuda.so \ - target/release/examples/native_vision \ - weights/Qwen3.5-4B weights/cua-s1-4b-0.2/multimodal \ - weights/cua-s1-multimodal-language /tmp/cua-standard-controls /tmp/cua-native.json -python recipe/cua_s1/verify_native_vision.py \ - /tmp/cua-standard-controls/manifest.json /tmp/cua-native.json -``` - -Repeat with `/tmp/cua-boundary` and separate output paths. Output directories/files -must not exist. Reference controls execute the **full vision and language path** -with unmerged BF16 and FP32 models, TF32 disabled; they do not feed BF16 reference -image embeddings into the native language model. The native replay starts from -request PNG/JPEG data and compares independently produced token IDs, image grids, -positions, candidate probabilities and repeated forwards. Stage feature errors -are reported separately. The acceptance rule is unchanged: maximum native -probability error over a set ≤ twice the maximum BF16 reference error plus 0.01; -top choices match where the FP32 top-two margin is at least 0.05. - -Additional CUDA regressions: - -```sh -.venv/bin/python tests/cua_s1/test_native_vision_cuda.py \ - target/release/libqwen3_5_cuda.so -CUA_S1_MODEL=weights/cua-s1-multimodal-language \ -CUA_S1_CUDA_LIB=$PWD/target/release/libqwen3_5_cuda.so \ - cargo test --release --locked -p omni-cua-s1-native -- --ignored -``` - -The optional `vision_stage` example emits each vision block's BF16 tensors. These -are diagnostics, not a bitwise-equivalence claim. See the recorded results in -[native vision validation](experiments/native-vision/README.md). diff --git a/recipe/cua_s1/native_vision.md b/recipe/cua_s1/native_vision.md deleted file mode 100644 index f6ffd39..0000000 --- a/recipe/cua_s1/native_vision.md +++ /dev/null @@ -1,64 +0,0 @@ -# Cua-S1 native vision checkpoint inspection - -The model-owned Rust module `omni_cua_s1_native::vision` loads and structurally -validates the Qwen3.5-4B vision weights and Cua-S1 4B 0.2 **multimodal** LoRA on -CPU. It keeps the 297 BF16 base tensors and 100 FP32 adapter tensors separate. -The adapter contains 50 A/B pairs, rank 16, alpha 32, and scale 2. No weights are -merged or converted. This CPU inspection API is separate from the -[native CUDA vision encoder and screenshot worker](native_multimodal.md). - -Prepare checkpoints from these pinned upstream revisions: - -| Checkpoint | Revision | Required files | -| --- | --- | --- | -| `Qwen/Qwen3.5-4B` | `851bf6e806efd8d0a36b00ddf55e13ccb7b8cd0a` | `config.json` and `model.safetensors`, or `model.safetensors.index.json` and its vision-bearing shards | -| `cua-ai/cua-s1-4b-0.2` multimodal adapter | `16818868b0cc7813808aae4e87b417657046ab79` | `adapter_config.json` and `adapter_model.safetensors` from the multimodal directory | - -Verify downloaded files against trusted upstream hashes before loading them. -Structural validation checks names, shapes, dtypes, configuration, and shard -mapping; it does **not** establish cryptographic identity or validate tensor -values. A bare base checkpoint or the `text` adapter is insufficient. - -From the repository root, pass exactly the base and multimodal adapter directories: - -```sh -cargo run --locked -p omni-cua-s1-native --example inspect_vision -- \ - weights/Qwen3.5-4B weights/cua-s1-4b-0.2/multimodal -``` - -The example reports the vision configuration, LoRA rank and scale, and tensor -counts. It needs neither a GPU nor a CUDA shared library. The index loader opens -only shards assigned visual tensors; unrelated language tensors in those shards -are ignored. Unexpected visual tensors, incompatible configuration, invalid -safetensors, missing tensors, mismatched index entries, and paths escaping the -checkpoint directory (including symlinks) are rejected. - -`VisionCheckpoint::load(base_dir, adapter_dir)` owns the memory maps and caches -validated tensor metadata. `config()` and `adapter()` expose the configuration; -`base_names()` / `adapter_names()` enumerate visual tensor names, and -`base_tensor(name)` / `adapter_tensor(name)` return borrowed safetensors -`TensorView`s with the original bytes. Views cannot outlive the checkpoint. -Accessors do not parse the file headers again. - -**Do not modify, replace in place, or truncate checkpoint files while a loaded -checkpoint or any of its views exists.** As with the existing native text -loader, callers must ensure that the memory-mapped files remain immutable. - -The supported layout is the pinned 4B configuration: 24 vision blocks, hidden -size 1024, intermediate size 4096, 16 heads, 2304 positions, output width 2560, -3 input channels, patch size 16, temporal patch size 2, spatial merge size 2, -`gelu_pytorch_tanh`, and no DeepStack. Text hidden size must equal vision output -width. LoRA targets both MLP matrices in every vision block and both merger -matrices. Math-changing PEFT options such as DoRA, rsLoRA, biases, target -exclusions, and custom rank/alpha patterns are unsupported. - -Run the CPU integration tests with: - -```sh -cargo test --locked -p omni-cua-s1-native --test vision_loader -``` - -Tests use sparse safetensors with real tensor shapes and small sentinel values; -no weight download or multi-gigabyte in-memory tensor allocation is required. -Their independent metadata oracle and provenance live in -[`tests/cua_s1/fixtures/vision/`](../../tests/cua_s1/fixtures/vision/). diff --git a/recipe/cua_s1/native_vision_cases.py b/recipe/cua_s1/native_vision_cases.py deleted file mode 100644 index d0099ce..0000000 --- a/recipe/cua_s1/native_vision_cases.py +++ /dev/null @@ -1,119 +0,0 @@ -"""Deterministic screenshot requests for native end-to-end validation. - -Standard cases preserve the #53 reference suite; boundary cases add maximum -area, a 200:1 image, odd dimensions, and eight questions sharing one tiny image. -""" - -import argparse -import base64 -import io -import json -from pathlib import Path - -from PIL import Image, ImageDraw - - -def write_json(path, value): - path.write_text(json.dumps(value, ensure_ascii=False, indent=2) + "\n") - - -def make_cases(folder): - folder.mkdir(parents=True) - cases = [] - for name, size, fmt in [ - ("small", (320, 240), "PNG"), - ("wide", (640, 320), "PNG"), - ("portrait", (320, 640), "PNG"), - ("jpeg", (640, 480), "JPEG"), - ("single-option", (320, 240), "PNG"), - ("26-options", (256, 256), "PNG"), - ("two-questions", (320, 240), "PNG"), - ]: - image = Image.new("RGB", size, "#f4f6f8") - draw = ImageDraw.Draw(image) - width, height = size - draw.rectangle( - (16, 16, width - 16, height - 16), fill="white", outline="#8899aa" - ) - draw.text((24, 24), "Account settings", fill="black") - draw.text((24, 48), "Display name: Alice", fill="black") - draw.rectangle((24, height // 2, width // 2, height // 2 + 32), fill="#1460b4") - draw.text((28, height // 2 + 8), "Save", fill="white") - draw.text((width // 2 + 16, height // 2 + 8), "Cancel", fill="black") - image_path = folder / (name + (".jpg" if fmt == "JPEG" else ".png")) - image.save(image_path, format=fmt) - criteria = {"save": "Click Save", "cancel": "Click Cancel", "wait": "Wait"} - if name == "single-option": - criteria = {"save": "Click Save"} - elif name == "26-options": - criteria = {f"option-{i}": f"Choose action {i}" for i in range(26)} - questions = { - "next": { - "type": "choice", - "instructions": "Save the changed display name.", - "criteria": criteria, - } - } - if name == "two-questions": - questions["second"] = { - "type": "choice", - "instructions": {"goal": "保存名称"}, - "criteria": {"continue": {"label": "Save"}, "cancel": None}, - } - mime = "jpeg" if fmt == "JPEG" else "png" - request = { - "model": "cua-s1-4b-0.2", - "state": { - "image": f"data:image/{mime};base64," - + base64.b64encode(image_path.read_bytes()).decode() - }, - "questions": questions, - } - write_json(folder / f"{name}.json", request) - cases.append({"name": name, "image": image_path.name, "request": request}) - return cases - - -def make_boundaries(folder): - folder.mkdir(parents=True, exist_ok=False) - for name, w, h, count in [ - ("maximum", 1024, 1024, 1), - ("narrow", 200, 1, 1), - ("noise", 383, 257, 1), - ("tiny-eight", 1, 1, 8), - ]: - rgb = bytes((i * 73 + (i // 3) * 17) % 256 for i in range(w * h * 3)) - image = Image.frombytes("RGB", (w, h), rgb) - buffer = io.BytesIO() - image.save(buffer, format="PNG") - questions = { - f"q{i}": { - "type": "choice", - "instructions": f"Choose a useful next action. Case {i}.", - "criteria": { - "save": "Click Save", - "cancel": "Click Cancel", - "wait": "Wait", - }, - } - for i in range(count) - } - write_json( - folder / (name + ".json"), - { - "model": "cua-s1-4b-0.2", - "state": { - "image": "data:image/png;base64," - + base64.b64encode(buffer.getvalue()).decode() - }, - "questions": questions, - }, - ) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("output", type=Path) - parser.add_argument("--boundary", action="store_true") - args = parser.parse_args() - (make_boundaries if args.boundary else make_cases)(args.output) diff --git a/recipe/cua_s1/requirements-native-validation.txt b/recipe/cua_s1/requirements-native-validation.txt deleted file mode 100644 index 5609fec..0000000 --- a/recipe/cua_s1/requirements-native-validation.txt +++ /dev/null @@ -1,4 +0,0 @@ --r requirements-text.txt -torchvision==0.29.0 -numpy==2.5.3 -Pillow==11.3.0 diff --git a/recipe/cua_s1/validate_native_vision.py b/recipe/cua_s1/validate_native_vision.py deleted file mode 100644 index e173689..0000000 --- a/recipe/cua_s1/validate_native_vision.py +++ /dev/null @@ -1,197 +0,0 @@ -"""Export fresh, full-vision unmerged BF16/FP32 controls for native RGB replay. - -Run with PYTHONPATH=src in the pinned reference environment. Output contains -RGB inputs, token/position oracles, vision traces and full-model probabilities. -""" - -import argparse -import gc -import hashlib -import json -import os -from pathlib import Path - -os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8") -os.environ.setdefault("HF_HUB_OFFLINE", "1") - - -def main(): - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--base", required=True) - parser.add_argument("--adapter", required=True) - parser.add_argument( - "--requests", required=True, help="directory of screenshot request JSON files" - ) - parser.add_argument("--out", required=True) - args = parser.parse_args() - import torch - import transformers - from safetensors.torch import save_file - - from models.cua_s1.multimodal.model import ( - ADAPTER_REVISION, - BASE_REVISION, - MultimodalEngine, - letter_ids, - ) - from models.cua_s1.multimodal.protocol import decode_request, parse_request - - out = Path(args.out) - out.mkdir(parents=True, exist_ok=False) - torch.set_num_threads(8) - torch.manual_seed(0) - torch.backends.cuda.matmul.allow_tf32 = False - torch.backends.cudnn.allow_tf32 = False - cases = [] - for path in sorted(Path(args.requests).glob("*.json")): - raw = path.read_bytes() - request = parse_request(decode_request(raw)) - rgb = request.image.tobytes() - name = path.stem - (out / f"{name}.rgb").write_bytes(rgb) - (out / f"{name}.json").write_bytes(raw) - cases.append( - ( - name, - request, - { - "case": name, - "width": request.image.width, - "height": request.image.height, - "rgb_file": f"{name}.rgb", - "request_file": f"{name}.json", - "rgb_sha256": hashlib.sha256(rgb).hexdigest(), - "questions": [], - }, - ) - ) - for dtype in ["bfloat16", "float32"]: - engine = MultimodalEngine(args.base, args.adapter, dtype=dtype) - core = engine.model.get_base_model().model - for name, request, record in cases: - prepared = engine.prepare_reused(request.image, request.questions) - trace = {} - handles = [] - - def hook(key, storage=trace): - def keep(module, ins, output): - if hasattr(output, "pooler_output"): - output = output.pooler_output - storage[key] = output.detach().cpu().contiguous().clone() - - return keep - - if name == "small" and dtype == "bfloat16": - handles.append( - core.visual.patch_embed.register_forward_hook(hook("patch_embed")) - ) - - def keep_position(module, ins, kwargs, storage=trace): - x = ins[0] if ins else kwargs["hidden_states"] - storage["position"] = x.detach().cpu().contiguous().clone() - - handles.append( - core.visual.blocks[0].register_forward_pre_hook( - keep_position, with_kwargs=True - ) - ) - for i, block in enumerate(core.visual.blocks): - handles.append(block.register_forward_hook(hook(f"blocks.{i}"))) - handles.append( - core.visual.merger.norm.register_forward_hook(hook("merger.norm")) - ) - handles.append( - core.visual.merger.linear_fc1.register_forward_hook( - hook("merger.fc1") - ) - ) - with torch.no_grad(): - features = engine.encode_image(prepared[0]) - for handle in handles: - handle.remove() - tensors = {"image_features": features.detach().cpu().contiguous()} - if trace: - trace["merger.output"] = tensors["image_features"].clone() - save_file(trace, str(out / "vision-trace.safetensors")) - for idx, (q, inputs) in enumerate(zip(request.questions, prepared)): - with torch.no_grad(): - device_inputs = { - k: v.cuda() for k, v in inputs.items() if k != "pixel_values" - } - ids = device_inputs["input_ids"] - pos, _ = core.get_rope_index( - input_ids=ids, - mm_token_type_ids=device_inputs["mm_token_type_ids"], - image_grid_thw=device_inputs["image_grid_thw"], - attention_mask=device_inputs["attention_mask"], - ) - embeds = core.get_input_embeddings()(ids) - mask = ids == engine.model.config.image_token_id - embeds[mask] = features.to(embeds.dtype) - hidden = core.language_model( - inputs_embeds=embeds, - position_ids=pos, - attention_mask=device_inputs["attention_mask"], - use_cache=False, - ).last_hidden_state[0, -1] - candidate_ids = letter_ids(engine.tokenizer, len(q.keys)) - rows = engine.model.get_output_embeddings().weight[candidate_ids] - logits = torch.nn.functional.linear(hidden, rows) - probs = logits.float().softmax(-1).tolist() - if dtype == "bfloat16": - record["grid"] = inputs["image_grid_thw"][0].tolist() - record["questions"].append( - { - "name": q.name, - "input_ids": ids[0].tolist(), - "position_ids": pos[:, 0].tolist(), - "candidate_ids": candidate_ids, - "bf16_probabilities": probs, - } - ) - tensors["pixel_values"] = inputs["pixel_values"].contiguous() - else: - record["questions"][idx]["fp32_probabilities"] = probs - tensors[f"hidden.{idx}"] = hidden.detach().cpu().contiguous() - print(dtype, name, q.name, probs, flush=True) - save_file(tensors, str(out / f"{name}.{dtype}.safetensors")) - del ( - features, - hidden, - embeds, - rows, - tensors, - logits, - trace, - inputs, - device_inputs, - ids, - pos, - ) - del engine, core - gc.collect() - torch.cuda.empty_cache() - manifest = { - "schema": "cua-s1-native-vision-controls-v1", - "base_revision": BASE_REVISION, - "adapter_revision": ADAPTER_REVISION, - "torch": str(torch.__version__), - "transformers": transformers.__version__, - "gpu": torch.cuda.get_device_name(), - "tf32": False, - "control": "full-vision unmerged adapter, candidate projection at last position", - "acceptance": "max native probability error <= 2 * max BF16 probability error + 0.01; same top choice for FP32 margin >= 0.05", - "cases": [r for _, _, r in cases], - } - manifest["files"] = { - p.name: hashlib.sha256(p.read_bytes()).hexdigest() - for p in out.iterdir() - if p.is_file() - } - (out / "manifest.json").write_text( - json.dumps(manifest, ensure_ascii=False, indent=2) + "\n" - ) - - -if __name__ == "__main__": - main() diff --git a/recipe/cua_s1/verify_native_vision.py b/recipe/cua_s1/verify_native_vision.py deleted file mode 100644 index 1c83842..0000000 --- a/recipe/cua_s1/verify_native_vision.py +++ /dev/null @@ -1,89 +0,0 @@ -"""Fail closed on missing coverage or violations of the repository parity rule.""" - -import argparse -import json -import math -from pathlib import Path - - -def verify(control, native): - if ( - control["schema"] != "cua-s1-native-vision-controls-v1" - or native["schema"] != "cua-s1-native-vision-results-v1" - ): - raise ValueError("unsupported schema") - refs, results = control["cases"], native["cases"] - if not refs or len(refs) != len(results): - raise ValueError("case coverage mismatch") - errors, bf_errors, choices = [], [], [] - for ref, result in zip(refs, results): - if ref["case"] != result["case"] or ref["grid"] != result["grid"]: - raise ValueError("case/grid mismatch") - if len(ref["questions"]) != len(result["questions"]): - raise ValueError("question coverage mismatch") - for r, n in zip(ref["questions"], result["questions"]): - if r["name"] != n["name"] or not all( - n[k] is True - for k in ["token_ids_equal", "position_ids_equal", "repeat_equal"] - ): - raise ValueError("CPU boundary or replay mismatch") - fp, bf, actual = ( - r["fp32_probabilities"], - r["bf16_probabilities"], - n["probabilities"], - ) - if not 1 <= len(fp) <= 26 or len(fp) != len(bf) or len(fp) != len(actual): - raise ValueError("candidate count mismatch") - for p in [fp, bf, actual]: - if ( - not all(math.isfinite(x) and 0 <= x <= 1 for x in p) - or abs(sum(p) - 1) > 1e-5 - ): - raise ValueError("invalid probability distribution") - error = max(abs(x - y) for x, y in zip(fp, actual)) - errors.append(error) - bf_errors.append(max(abs(x - y) for x, y in zip(fp, bf))) - order = sorted(range(len(fp)), key=lambda i: -fp[i]) - margin = fp[order[0]] - fp[order[1]] if len(fp) > 1 else 1.0 - chosen = max(range(len(actual)), key=actual.__getitem__) - if margin >= 0.05 and chosen != order[0]: - raise ValueError(f"choice mismatch: {ref['case']}/{r['name']}") - choices.append( - { - "case": ref["case"], - "question": r["name"], - "fp32_margin": margin, - "same_choice": chosen == order[0], - "max_probability_error": error, - } - ) - allowance = 2 * max(bf_errors) + 0.01 - if max(errors) > allowance: - raise ValueError(f"probability error {max(errors)} exceeds {allowance}") - return { - "passed": True, - "questions": len(errors), - "max_native_error": max(errors), - "max_bf16_reference_error": max(bf_errors), - "allowance": allowance, - "choices": choices, - } - - -def main(): - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("control", type=Path) - parser.add_argument("native", type=Path) - parser.add_argument("--out", type=Path) - args = parser.parse_args() - result = verify( - json.loads(args.control.read_text()), json.loads(args.native.read_text()) - ) - text = json.dumps(result, indent=2) + "\n" - if args.out: - args.out.write_text(text) - print(text, end="") - - -if __name__ == "__main__": - main() diff --git a/src/backends/cuda/qwen3_5/README.md b/src/backends/cuda/qwen3_5/README.md index 176c949..b84f50b 100644 --- a/src/backends/cuda/qwen3_5/README.md +++ b/src/backends/cuda/qwen3_5/README.md @@ -7,8 +7,3 @@ src/backends/cuda/qwen3_5/build.sh [compute capability, default 89] ``` The norm, elementwise and q/k preparation kernels round to bfloat16 where Transformers (`modeling_qwen3_5.py`) does. Attention (FlashAttention-2 style, on tensor cores) and the chunked gated delta rule keep some intermediate results in bfloat16, as FlashAttention and flash-linear-attention do. GEMMs go through cuBLASLt with its first heuristic choice. Tensor-core kernels need sm_80 or newer; only sm_89 has been run. - -ABI 4 adds native vision LayerNorm, patch/biased projection, spatial rotary, -bidirectional attention, GELU, residual operations and FP32 LoRA GEMMs. Vision -GEMMs disable split-K BF16 reduction; existing language GEMM dispatch is preserved. -Rebuild the library and native workers together after ABI changes. diff --git a/src/models/cua_s1/README.md b/src/models/cua_s1/README.md index 18f47fd..999ae0a 100644 --- a/src/models/cua_s1/README.md +++ b/src/models/cua_s1/README.md @@ -2,7 +2,7 @@ This directory owns Cua-S1 4B 0.2 ([#10](https://github.com/ThinkFlowLab/system1-omni/issues/10)): request mapping, prompt construction, adapter selection, execution, and the answer-letter readout. This page records the pinned upstream revisions, the inference contract an implementation must match, and how its outputs will be compared with the upstream reference. -Status: a reference worker for the `text` adapter loads the model through Hugging Face Transformers and PEFT: [`text/`](text/), served by [`src/frontend/cua_s1_text.py`](../../frontend/cua_s1_text.py), with setup in [`recipe/cua_s1/text.md`](../../../recipe/cua_s1/text.md). It is the correctness reference for the native worker in [`native/`](native/): Rust, with the Qwen3.5 forward pass on the CUDA kernels in [`src/backends/cuda/qwen3_5/`](../../backends/cuda/qwen3_5/), set up as in [`recipe/cua_s1/native.md`](../../../recipe/cua_s1/native.md). The `multimodal` adapter has a [native screenshot worker](../../../recipe/cua_s1/native_multimodal.md), including Rust RGB preprocessing, CUDA vision and language execution, and RTX 4090 output validation. +Status: a reference worker for the `text` adapter loads the model through Hugging Face Transformers and PEFT: [`text/`](text/), served by [`src/frontend/cua_s1_text.py`](../../frontend/cua_s1_text.py), with setup in [`recipe/cua_s1/text.md`](../../../recipe/cua_s1/text.md). It is the correctness reference for the native worker in [`native/`](native/): Rust, with the Qwen3.5 forward pass on the CUDA kernels in [`src/backends/cuda/qwen3_5/`](../../backends/cuda/qwen3_5/), set up as in [`recipe/cua_s1/native.md`](../../../recipe/cua_s1/native.md). The `multimodal` adapter is deferred; see [Not covered yet](#not-covered-yet). ## Pinned revisions @@ -105,7 +105,7 @@ The bfloat16 worker's own difference from the fp32 worker is reported next to ea ## Not covered yet -- Video, padded or batched multimodal inputs, and vision CUDA Graph capture. +- The `multimodal` adapter: image preprocessing, the vision tower and the vision LoRA. This is tracked in [#10](https://github.com/ThinkFlowLab/system1-omni/issues/10). - `score` and `noul` questions. - More than 26 options per question. - The Metal backend. diff --git a/src/models/cua_s1/native/Cargo.toml b/src/models/cua_s1/native/Cargo.toml index b24e878..31275dd 100644 --- a/src/models/cua_s1/native/Cargo.toml +++ b/src/models/cua_s1/native/Cargo.toml @@ -26,15 +26,6 @@ serde_json = { version = "1.0.149", features = ["float_roundtrip", "preserve_ord tokenizers = { version = "=0.22.2", default-features = false, features = ["onig"] } tokio = { version = "1.49.0", features = ["macros", "net", "rt-multi-thread", "sync"] } -[[test]] -name = "vision_loader" -path = "../../../../tests/cua_s1/vision_loader.rs" - -[[test]] -name = "image_preprocess" -path = "../../../../tests/cua_s1/image_preprocess.rs" - - [[bin]] name = "omni-cua-s1-vision" path = "src/vision_main.rs" diff --git a/src/models/cua_s1/native/examples/inspect_vision.rs b/src/models/cua_s1/native/examples/inspect_vision.rs deleted file mode 100644 index afffb19..0000000 --- a/src/models/cua_s1/native/examples/inspect_vision.rs +++ /dev/null @@ -1,22 +0,0 @@ -use anyhow::{Result, ensure}; -use omni_cua_s1_native::vision::VisionCheckpoint; - -fn main() -> Result<()> { - let args: Vec<_> = std::env::args_os().skip(1).collect(); - ensure!( - args.len() == 2, - "usage: inspect_vision BASE_DIR MULTIMODAL_ADAPTER_DIR" - ); - let checkpoint = VisionCheckpoint::load(&args[0], &args[1])?; - println!("Vision: {:#?}", checkpoint.config()); - println!("Base: {} BF16 tensors", checkpoint.base_names().count()); - println!( - "Adapter: {} FP32 tensors ({} pairs), rank {}, alpha {}, scale {}", - checkpoint.adapter_names().count(), - checkpoint.adapter_names().count() / 2, - checkpoint.adapter().rank, - checkpoint.adapter().alpha, - checkpoint.adapter().scale() - ); - Ok(()) -} diff --git a/src/models/cua_s1/native/examples/multimodal_boundary.rs b/src/models/cua_s1/native/examples/multimodal_boundary.rs deleted file mode 100644 index 7c93fb4..0000000 --- a/src/models/cua_s1/native/examples/multimodal_boundary.rs +++ /dev/null @@ -1,184 +0,0 @@ -//! Replay #53's exported language boundary through the native model. -//! The v1 bundle producer/verifier are on open PR #53 at exporter revision -//! 1b64fa2ceb0a82b6a66a69ecdc9bc5cc1b1a0b66; see recipe/cua_s1/native.md. -use std::path::{Path, PathBuf}; - -use anyhow::{Context, Result, ensure}; -use half::bf16; -use omni_cua_s1_native::{inputs::MultimodalInput, model::Model}; -use safetensors::{Dtype, SafeTensors}; -use serde_json::{Value, json}; - -fn integers(st: &SafeTensors<'_>, name: &str, shape: &[usize]) -> Result> { - let v = st.tensor(name)?; - ensure!( - v.dtype() == Dtype::I64 && v.shape() == shape, - "{name}: expected I64 {shape:?}" - ); - Ok(v.data() - .as_chunks::<8>() - .0 - .iter() - .map(|b| i64::from_le_bytes(*b)) - .collect()) -} - -fn features(st: &SafeTensors<'_>, shape: &[usize]) -> Result> { - let v = st.tensor("image_features")?; - ensure!( - v.dtype() == Dtype::BF16 && v.shape() == shape, - "image_features: expected BF16 {shape:?}" - ); - Ok(v.data() - .as_chunks::<2>() - .0 - .iter() - .map(|b| bf16::from_le_bytes(*b)) - .collect()) -} - -fn embedding_file(dir: &Path) -> Result<(PathBuf, String)> { - let names = [ - "model.language_model.embed_tokens.weight", - "model.embed_tokens.weight", - "embed_tokens.weight", - ]; - if dir.join("model.safetensors.index.json").exists() { - let index: Value = - serde_json::from_slice(&std::fs::read(dir.join("model.safetensors.index.json"))?)?; - for name in names { - if let Some(file) = index["weight_map"][name].as_str() { - return Ok((dir.join(file), name.into())); - } - } - } else { - let path = dir.join("model.safetensors"); - let bytes = std::fs::read(&path)?; - let st = SafeTensors::deserialize(&bytes)?; - for name in names { - if st.tensor(name).is_ok() { - return Ok((path, name.into())); - } - } - } - anyhow::bail!("missing embedding weight") -} - -fn main() -> Result<()> { - let args: Vec<_> = std::env::args_os().skip(1).collect(); - ensure!( - args.len() == 3, - "usage: multimodal_boundary MODEL_DIR REFERENCE_BUNDLE OUTPUT_JSON" - ); - let (dir, bundle, out) = ( - Path::new(&args[0]), - Path::new(&args[1]), - Path::new(&args[2]), - ); - ensure!(!out.exists(), "output already exists"); - let library = PathBuf::from(std::env::var_os("CUA_S1_CUDA_LIB").context("CUA_S1_CUDA_LIB")?); - let manifest: Value = serde_json::from_slice(&std::fs::read(bundle.join("manifest.json"))?)?; - ensure!( - manifest["schema"] == "cua-s1-multimodal-reference-v1", - "unsupported reference schema" - ); - let mut model = Model::load(dir, &library)?; - let (file, name) = embedding_file(dir)?; - let file = std::fs::File::open(file)?; - // SAFETY: the checkpoint is immutable while the example runs. - let map = unsafe { memmap2::Mmap::map(&file)? }; - let weights = SafeTensors::deserialize(&map)?; - let embed = weights.tensor(&name)?; - ensure!( - embed.dtype() == Dtype::BF16 - && embed.shape().len() == 2 - && embed.shape()[1] == model.cfg.hidden, - "embedding shape/dtype" - ); - let mut rows = Vec::new(); - for entry in manifest["questions"].as_array().context("questions")? { - let relative = Path::new(entry["tensors_file"].as_str().context("tensors_file")?); - ensure!( - relative - .components() - .all(|c| matches!(c, std::path::Component::Normal(_))), - "unsafe tensor path" - ); - let bytes = std::fs::read(bundle.join(relative))?; - let st = SafeTensors::deserialize(&bytes)?; - let ids = st.tensor("input_ids")?; - ensure!( - ids.shape().len() == 2 && ids.shape()[0] == 1, - "expected batch one" - ); - let t = ids.shape()[1]; - let ids: Vec = integers(&st, "input_ids", &[1, t])? - .into_iter() - .map(u32::try_from) - .collect::>()?; - let indices = st.tensor("image_token_indices")?; - ensure!( - indices.shape().len() == 1, - "image indices must be one-dimensional" - ); - let count = indices.shape()[0]; - let indices: Vec = integers(&st, "image_token_indices", &[count])? - .into_iter() - .map(usize::try_from) - .collect::>()?; - let features = features(&st, &[count, model.cfg.hidden])?; - let positions = integers(&st, "position_ids", &[3, 1, t])?; - let input = MultimodalInput { - token_ids: &ids, - image_token_indices: &indices, - image_embeddings: &features, - position_ids: [&positions[..t], &positions[t..2 * t], &positions[2 * t..]], - }; - let last = model.forward_multimodal(&input)?; - ensure!( - last.iter().all(|x| x.is_finite()), - "non-finite hidden state" - ); - // Repeating the same boundary in the same model checks buffer reuse. - ensure!( - last == model.forward_multimodal(&input)?, - "repeat changed native hidden state" - ); - let n = entry["option_keys"] - .as_array() - .context("option_keys")? - .len(); - ensure!((1..=26).contains(&n), "candidate count"); - let candidates = integers(&st, "candidate_token_ids", &[n])?; - let mut logits = Vec::new(); - for id in candidates { - let id = usize::try_from(id)?; - ensure!(id < embed.shape()[0], "candidate outside vocabulary"); - let row = &embed.data()[id * last.len() * 2..(id + 1) * last.len() * 2]; - let dot: f64 = row - .as_chunks::<2>() - .0 - .iter() - .zip(&last) - .map(|(b, &h)| bf16::from_le_bytes(*b).to_f64() * h as f64) - .sum(); - logits.push(dot as f32); - } - let max = logits.iter().copied().fold(f32::NEG_INFINITY, f32::max) as f64; - let exps: Vec = logits.iter().map(|&l| (l as f64 - max).exp()).collect(); - let total: f64 = exps.iter().sum(); - let probabilities: Vec = exps.iter().map(|e| (e / total) as f32).collect(); - ensure!( - probabilities.iter().all(|x| x.is_finite()), - "non-finite readout" - ); - rows.push(json!({"case": entry["case"], "question": entry["question"], "sequence": t, "image_tokens": count, "last_hidden_state": last, "candidate_logits": logits, "probabilities": probabilities, "repeat_equal": true})); - } - std::fs::write( - out, - serde_json::to_vec_pretty( - &json!({"schema": "cua-s1-native-language-boundary-v1", "questions": rows}), - )?, - )?; - Ok(()) -} diff --git a/src/models/cua_s1/native/examples/native_vision.rs b/src/models/cua_s1/native/examples/native_vision.rs deleted file mode 100644 index b40367b..0000000 --- a/src/models/cua_s1/native/examples/native_vision.rs +++ /dev/null @@ -1,138 +0,0 @@ -//! Replay full screenshot requests and independently compare all CPU boundaries. -use anyhow::{Context, Result, ensure}; -use half::bf16; -use omni_cua_s1_native::{contract, image_request::parse_image_body, vision_engine::VisionEngine}; -use safetensors::SafeTensors; -use serde_json::{Value, json}; -use std::path::{Path, PathBuf}; - -fn main() -> Result<()> { - let args: Vec<_> = std::env::args_os().skip(1).collect(); - ensure!( - args.len() == 5, - "usage: native_vision BASE ADAPTER LANGUAGE CONTROL_DIR OUTPUT_JSON" - ); - let library = PathBuf::from(std::env::var_os("CUA_S1_CUDA_LIB").context("CUA_S1_CUDA_LIB")?); - let bundle = Path::new(&args[3]); - let out = Path::new(&args[4]); - ensure!(!out.exists(), "output already exists"); - let manifest: Value = serde_json::from_slice(&std::fs::read(bundle.join("manifest.json"))?)?; - ensure!( - manifest["schema"] == "cua-s1-native-vision-controls-v1", - "unsupported control schema" - ); - let mut engine = VisionEngine::load( - Path::new(&args[0]), - Path::new(&args[1]), - Path::new(&args[2]), - &library, - )?; - let mut records = Vec::new(); - for case in manifest["cases"].as_array().context("cases")? { - let name = case["case"].as_str().context("case name")?; - let read = |field: &str| -> Result> { - let p = Path::new(case[field].as_str().context("file name")?); - ensure!( - p.components() - .all(|c| matches!(c, std::path::Component::Normal(_))), - "invalid relative path" - ); - Ok(std::fs::read(bundle.join(p))?) - }; - let body = - contract::parse_body(&read("request_file")?).map_err(|e| anyhow::anyhow!(e.message))?; - let request = parse_image_body(&body)?; - let reference_rgb = read("rgb_file")?; - ensure!( - request.rgb.len() == reference_rgb.len(), - "decoded RGB shape mismatch" - ); - let decode_max_abs = request - .rgb - .iter() - .zip(&reference_rgb) - .map(|(&a, &b)| a.abs_diff(b)) - .max() - .unwrap_or(0); - let prepared = engine.prepare( - request.width, - request.height, - &request.rgb, - &request.questions, - )?; - let grid: Vec = serde_json::from_value(case["grid"].clone())?; - ensure!( - prepared.image.image_grid_thw.as_slice() == grid, - "grid mismatch" - ); - let start = std::time::Instant::now(); - let features = engine.vision.forward(&prepared.image)?; - let vision_ms = start.elapsed().as_secs_f64() * 1000.; - let reference = std::fs::read(bundle.join(format!("{name}.bfloat16.safetensors")))?; - let st = SafeTensors::deserialize(&reference)?; - let f = st.tensor("image_features")?; - ensure!( - f.dtype() == safetensors::Dtype::BF16 && f.data().len() == features.len() * 2, - "feature shape/dtype mismatch" - ); - let reference_features: Vec = f - .data() - .as_chunks::<2>() - .0 - .iter() - .map(|b| bf16::from_le_bytes(*b).to_f32()) - .collect(); - let errors: Vec = features - .iter() - .zip(&reference_features) - .map(|(a, &b)| a.to_f64() - b as f64) - .collect(); - let feature_max_abs = errors.iter().copied().map(f64::abs).fold(0., f64::max); - let feature_rmse = (errors.iter().map(|e| e * e).sum::() / errors.len() as f64).sqrt(); - let mut questions = Vec::new(); - for ((q, prompt), control) in request - .questions - .iter() - .zip(&prepared.prompts) - .zip(case["questions"].as_array().context("questions")?) - { - let ids: Vec = serde_json::from_value(control["input_ids"].clone())?; - let positions: [Vec; 3] = serde_json::from_value(control["position_ids"].clone())?; - ensure!( - prompt.token_ids == ids, - "{name}: tokenization differs from reference" - ); - ensure!( - prompt.position_ids == positions, - "{name}: positions differ from reference" - ); - let score = engine.score(prompt, &features, q.keys.len())?; - let repeated = engine.score(prompt, &features, q.keys.len())?; - ensure!( - score.hidden == repeated.hidden, - "repeated native language execution changed" - ); - questions.push( - json!({"name":q.name, "sequence":ids.len(), "token_ids_equal":true, - "position_ids_equal":true, "repeat_equal":true, "probabilities":score.probabilities, - "logits":score.logits, "hidden":score.hidden}), - ); - } - // Replay the same image after LM calls to expose cross-model scratch corruption. - ensure!( - features == engine.vision.forward(&prepared.image)?, - "vision replay changed features" - ); - records.push(json!({"case":name,"decode_max_abs":decode_max_abs,"grid":grid, - "image_tokens":features.len()/2560,"feature_max_abs":feature_max_abs,"feature_rmse":feature_rmse, - "vision_ms_first_call":vision_ms,"questions":questions})); - eprintln!("completed {name}: feature max {feature_max_abs}, RMS {feature_rmse}"); - } - std::fs::write( - out, - serde_json::to_vec_pretty( - &json!({"schema":"cua-s1-native-vision-results-v1","cases":records}), - )?, - )?; - Ok(()) -} diff --git a/src/models/cua_s1/native/examples/preprocess_image.rs b/src/models/cua_s1/native/examples/preprocess_image.rs deleted file mode 100644 index 310c661..0000000 --- a/src/models/cua_s1/native/examples/preprocess_image.rs +++ /dev/null @@ -1,75 +0,0 @@ -//! CPU-only RGB8 preprocessing; run from the repository root with: -//! cargo run -p omni-cua-s1-native --example preprocess_image -- 256 256 image.rgb - -use std::{ - env, - fs::File, - io::{BufWriter, Read, Write}, -}; - -use anyhow::{Context, Result, ensure}; -use omni_cua_s1_native::image_preprocess::preprocess_rgb8; - -fn main() -> Result<()> { - let args: Vec<_> = env::args_os().skip(1).collect(); - ensure!( - args.len() == 3 || args.len() == 4, - "usage: preprocess_image WIDTH HEIGHT RAW_RGB_PATH [OUTPUT_F32_PATH]" - ); - let dimension = |index: usize| -> Result { - let text = args[index] - .to_str() - .context("dimensions must be UTF-8 decimal integers")?; - ensure!( - !text.is_empty() && text.bytes().all(|byte| byte.is_ascii_digit()), - "dimensions must be unsigned decimal integers" - ); - text.parse().context("dimension is too large") - }; - let width = dimension(0)?; - let height = dimension(1)?; - // Bound the file read before allocating. The library validates the complete - // contract too; these checks keep malformed CLI inputs cheap to reject. - ensure!( - width > 0 && height > 0 && width <= 2048 && height <= 2048, - "dimensions must be in 1..=2048" - ); - let area = width.checked_mul(height).context("image area overflow")?; - ensure!( - area <= 1_048_576, - "image area must not exceed 1048576 pixels" - ); - ensure!( - width.max(height) <= width.min(height) * 200, - "image aspect ratio must not exceed 200" - ); - let expected = area.checked_mul(3).context("RGB length overflow")?; - let mut rgb = Vec::with_capacity(expected + 1); - File::open(&args[2]) - .context("opening RGB input")? - .take((expected + 1) as u64) - .read_to_end(&mut rgb) - .context("reading RGB input")?; - ensure!( - rgb.len() == expected, - "RGB input must contain exactly {expected} bytes" - ); - let image = preprocess_rgb8(width, height, &rgb)?; - println!( - "pixel_values shape: [{}, 1536]", - image.pixel_values.len() / 1536 - ); - println!("image_grid_thw: {:?}", image.image_grid_thw); - println!("resized: {}x{}", image.resized_width, image.resized_height); - println!("image_tokens: {}", image.image_tokens()); - if let Some(path) = args.get(3) { - let mut output = BufWriter::new(File::create(path).context("creating float32 output")?); - for value in image.pixel_values { - output - .write_all(&value.to_le_bytes()) - .context("writing float32 output")?; - } - output.flush().context("flushing float32 output")?; - } - Ok(()) -} diff --git a/src/models/cua_s1/native/examples/vision_stage.rs b/src/models/cua_s1/native/examples/vision_stage.rs deleted file mode 100644 index 1d7aa62..0000000 --- a/src/models/cua_s1/native/examples/vision_stage.rs +++ /dev/null @@ -1,28 +0,0 @@ -//! GPU parity stage runner. Raw RGB8 input; emits BF16 stage files and shapes JSON. -use anyhow::{Result, ensure}; -use omni_cua_s1_native::{image_preprocess::preprocess_rgb8, vision::VisionModel}; -use std::{env, fs, path::Path}; -fn main() -> Result<()> { - let args: Vec = env::args().collect(); - ensure!( - args.len() == 8, - "usage: vision_stage BASE ADAPTER LIB RGB WIDTH HEIGHT OUTDIR" - ); - let image = preprocess_rgb8(args[5].parse()?, args[6].parse()?, &fs::read(&args[4])?)?; - fs::create_dir_all(&args[7])?; - let mut model = VisionModel::load(&args[1], &args[2], Path::new(&args[3]))?; - let start = std::time::Instant::now(); - let output = model.forward_with_trace(&image, |name, values| { - let bytes: Vec = values - .iter() - .flat_map(|v| v.to_bits().to_le_bytes()) - .collect(); - fs::write(Path::new(&args[7]).join(format!("{name}.bf16")), bytes)?; - Ok(()) - })?; - println!( - "{}", - serde_json::json!({"grid":image.image_grid_thw,"output_shape":[image.image_tokens(),2560],"values":output.len(),"elapsed_seconds":start.elapsed().as_secs_f64()}) - ); - Ok(()) -} diff --git a/src/models/cua_s1/native/src/image_request.rs b/src/models/cua_s1/native/src/image_request.rs index 3cf4315..25b3c2b 100644 --- a/src/models/cua_s1/native/src/image_request.rs +++ b/src/models/cua_s1/native/src/image_request.rs @@ -179,88 +179,3 @@ fn checked_text(value: &Value) -> Result { ); Ok(text) } - -#[cfg(test)] -mod tests { - use super::*; - use base64::Engine; - use serde_json::json; - fn body() -> Map { - let mut png = std::io::Cursor::new(Vec::new()); - image::RgbImage::from_pixel(32, 32, image::Rgb([17, 83, 191])) - .write_to(&mut png, image::ImageFormat::Png) - .unwrap(); - json!({"model":"cua-s1-4b-0.2", "state":{"image":format!("data:image/png;base64,{}", base64::engine::general_purpose::STANDARD.encode(png.into_inner()))}, "questions":{"q":{"type":"choice","instructions":{"goal":"保存"},"criteria":{"save":"Click Save", "cancel":null}}}}).as_object().unwrap().clone() - } - #[test] - fn sixteen_bit_modes_match_pinned_pillow_oracles() { - let dir = std::path::Path::new(env!("CARGO_MANIFEST_DIR")) - .join("../../../../tests/cua_s1/fixtures/native_decode"); - for (name, expected) in [ - ("l16", [255, 255, 255]), - ("l16-trns", [255, 255, 255]), - ("la16", [3, 3, 3]), - ("rgb16", [3, 0, 128]), - ("rgb16-trns", [3, 0, 128]), - ("rgba16", [3, 0, 128]), - ] { - let png = std::fs::read(dir.join(format!("{name}.png"))).unwrap(); - let mut b = body(); - b["state"]["image"] = json!(format!( - "data:image/png;base64,{}", - base64::engine::general_purpose::STANDARD.encode(png) - )); - assert_eq!(parse_image_body(&b).unwrap().rgb, expected, "{name}"); - } - } - #[test] - fn sixteen_bit_grayscale_matches_pillow_clamping() { - let mut png = std::io::Cursor::new(Vec::new()); - let pixels = image::ImageBuffer::, Vec>::from_raw( - 3, - 1, - vec![100, 1000, 65535], - ) - .unwrap(); - image::DynamicImage::ImageLuma16(pixels) - .write_to(&mut png, image::ImageFormat::Png) - .unwrap(); - let mut b = body(); - b["state"]["image"] = json!(format!( - "data:image/png;base64,{}", - base64::engine::general_purpose::STANDARD.encode(png.into_inner()) - )); - assert_eq!( - parse_image_body(&b).unwrap().rgb, - [100, 100, 100, 255, 255, 255, 255, 255, 255] - ); - } - #[test] - fn decodes_png_and_matches_question_mapping() { - let r = parse_image_body(&body()).unwrap(); - assert_eq!((r.width, r.height), (32, 32)); - assert_eq!(r.rgb, [17, 83, 191].repeat(32 * 32)); - assert_eq!(r.questions[0].goal, "{\"goal\": \"保存\"}"); - assert_eq!(r.questions[0].keys, ["save", "cancel"]); - } - #[test] - fn rejects_media_controls_unknown_fields_and_mime_mismatch() { - let mut b = body(); - b["questions"]["q"]["instructions"] = json!("<|image_pad|>"); - assert!(parse_image_body(&b).is_err()); - let mut b = body(); - b.insert("extra".into(), Value::Null); - assert!(parse_image_body(&b).is_err()); - let mut b = body(); - b["state"]["image"] = json!( - b["state"]["image"] - .as_str() - .unwrap() - .replace("image/png", "image/jpeg") - ); - assert!(parse_image_body(&b).is_err()); - let mut b = body(); - b["questions"]["q"]["criteria"]["save"] = json!("x".repeat(16385)); - assert!(parse_image_body(&b).is_err()); - } -} diff --git a/src/models/cua_s1/native/src/inputs.rs b/src/models/cua_s1/native/src/inputs.rs index 1f0c980..bf719ec 100644 --- a/src/models/cua_s1/native/src/inputs.rs +++ b/src/models/cua_s1/native/src/inputs.rs @@ -93,121 +93,3 @@ pub(crate) fn rotary_tables( } (cos, sin) } - -#[cfg(test)] -mod tests { - use super::*; - use half::bf16; - - #[test] - fn valid_input_and_three_distinct_axes() { - let input = MultimodalInput { - token_ids: &[1, 99, 99, 2], - image_token_indices: &[1, 2], - image_embeddings: &[bf16::ONE; 8], - position_ids: [&[0, 1, 1, 3], &[0, 1, 2, 3], &[0, 2, 1, 3]], - }; - input.validate(4, 100, 99, 100).unwrap(); - } - - #[test] - fn rejects_bad_placeholder_inventory_and_feature_rows() { - for indices in [vec![2, 1], vec![1, 1], vec![1], vec![0, 1], vec![1, 4]] { - let input = MultimodalInput { - token_ids: &[1, 99, 99, 2], - image_token_indices: &indices, - image_embeddings: &[bf16::ONE; 8], - position_ids: [&[0, 1, 1, 3]; 3], - }; - assert!(input.validate(4, 100, 99, 100).is_err(), "{indices:?}"); - } - for features in [vec![bf16::ONE; 7], vec![bf16::ONE; 9], vec![bf16::NAN; 8]] { - let input = MultimodalInput { - token_ids: &[1, 99, 99, 2], - image_token_indices: &[1, 2], - image_embeddings: &features, - position_ids: [&[0, 1, 1, 3]; 3], - }; - assert!(input.validate(4, 100, 99, 100).is_err()); - } - } - - #[test] - fn rejects_bad_tokens_positions_and_empty_sequence() { - let features = [bf16::ONE; 4]; - for (ids, positions) in [ - (vec![100, 99], vec![0, 1]), - (vec![1, 99], vec![0]), - (vec![1, 99], vec![0, -1]), - (vec![1, 99], vec![0, 100]), - (vec![], vec![]), - ] { - let input = MultimodalInput { - token_ids: &ids, - image_token_indices: &[1], - image_embeddings: &features, - position_ids: [&positions; 3], - }; - assert!(input.validate(4, 100, 99, 100).is_err()); - } - } - - #[test] - fn image_free_explicit_positions_are_valid() { - MultimodalInput { - token_ids: &[1, 2], - image_token_indices: &[], - image_embeddings: &[], - position_ids: [&[7, 8]; 3], - } - .validate(4, 100, 99, 100) - .unwrap(); - } - - #[test] - fn rotary_interleaves_height_width_and_leaves_temporal_tail() { - // theta=1 makes every inverse frequency 1. Axis values differ so a plain - // text table or a contiguous-section implementation fails this check. - for (sections, tail) in [([12, 10, 10], [0, 0]), ([11, 11, 10], [0, 1])] { - let (cos, sin) = rotary_tables([&[0], &[1], &[2]], 32, 1.0, sections); - // Ten T/H/W triples, then T/T for the synthetic layout or T/H for - // the real checkpoint. In particular, frequency 31 must use H. - let axes = [ - 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, - 1, 2, tail[0], tail[1], - ]; - for (i, axis) in axes.into_iter().enumerate() { - let angle = axis as f32; - assert_eq!( - &cos[i * 2..i * 2 + 2], - &bf16::from_f32(angle.cos()).to_le_bytes() - ); - assert_eq!( - &sin[i * 2..i * 2 + 2], - &bf16::from_f32(angle.sin()).to_le_bytes() - ); - } - } - } - - #[test] - fn equal_axes_reproduce_the_existing_text_tables() { - let pos: Vec = (0..257).collect(); - let (cos, sin) = rotary_tables([&pos; 3], 32, 10_000_000.0, [11, 11, 10]); - for (t, &p) in pos.iter().enumerate() { - for i in 0..32 { - let inv = 1f32 / 10_000_000f32.powf((2 * i) as f32 / 64.0); - let angle = (inv * p as f32) as f64; - let offset = (t * 32 + i) * 2; - assert_eq!( - &cos[offset..offset + 2], - &bf16::from_f32(angle.cos() as f32).to_le_bytes() - ); - assert_eq!( - &sin[offset..offset + 2], - &bf16::from_f32(angle.sin() as f32).to_le_bytes() - ); - } - } - } -} diff --git a/src/models/cua_s1/native/src/model.rs b/src/models/cua_s1/native/src/model.rs index d6d6597..ab4439a 100644 --- a/src/models/cua_s1/native/src/model.rs +++ b/src/models/cua_s1/native/src/model.rs @@ -1055,71 +1055,3 @@ impl Model { Ok(()) } } - -#[cfg(test)] -mod graph_tests { - use super::*; - - #[test] - #[ignore = "needs CUA_S1_MODEL and ABI-3 CUA_S1_CUDA_LIB on a GPU"] - fn graph_misses_hits_eviction_growth_and_multimodal_match_eager() { - let dir = std::path::PathBuf::from(std::env::var_os("CUA_S1_MODEL").unwrap()); - let lib = std::path::PathBuf::from(std::env::var_os("CUA_S1_CUDA_LIB").unwrap()); - let prompts: Vec> = (4..=13) - .chain([1025]) - .map(|t| vec![32 + (t % 3) as u32; t]) - .collect(); - let changed_ids = vec![35; 4]; - let mut eager = Model::load(&dir, &lib).unwrap(); - eager.graph_enabled = false; - let expected: Vec<_> = prompts - .iter() - .map(|ids| eager.forward(ids).unwrap()) - .collect(); - let changed_expected = eager.forward(&changed_ids).unwrap(); - let image_ids = [32, eager.cfg.image_token_id.unwrap(), 33, 34]; - let features = vec![half::bf16::ONE; eager.cfg.hidden]; - let boundary = MultimodalInput { - token_ids: &image_ids, - image_token_indices: &[1], - image_embeddings: &features, - position_ids: [&[0, 1, 2, 3], &[0, 7, 8, 9], &[0, 3, 4, 5]], - }; - let multimodal_expected = eager.forward_multimodal(&boundary).unwrap(); - drop(eager); - - let mut model = Model::load(&dir, &lib).unwrap(); - model.graph_enabled = true; - for (ids, expected) in prompts[..10].iter().zip(&expected) { - // First use returns the eager result, then the graph is a cache hit. - assert_eq!(expected, &model.forward(ids).unwrap()); - assert!( - model.graph_enabled, - "capture unexpectedly fell back to eager" - ); - assert!(model.graphs.iter().any(|(t, _)| *t == ids.len())); - assert_eq!(expected, &model.forward(ids).unwrap()); - } - assert_eq!(model.graphs.len(), 8); - assert!(!model.graphs.iter().any(|(t, _)| *t == 4)); - // Evicted length is another miss; then change token IDs on a warm hit. - assert_eq!(expected[0], model.forward(&prompts[0]).unwrap()); - assert_eq!(changed_expected, model.forward(&changed_ids).unwrap()); - - // The same cached length must still use eager multimodal execution: - // reusing the text graph would read text positions instead of T/H/W. - assert!(model.graphs.iter().any(|(t, _)| *t == image_ids.len())); - assert_eq!( - multimodal_expected, - model.forward_multimodal(&boundary).unwrap() - ); - assert_eq!(expected[0], model.forward(&prompts[0]).unwrap()); - - // Growth invalidates all captures before freeing their device buffers. - assert_eq!(expected[10], model.forward(&prompts[10]).unwrap()); - assert_eq!(model.graphs.len(), 1); - assert_eq!(model.graphs[0].0, 1025); - assert_eq!(expected[10], model.forward(&prompts[10]).unwrap()); - assert_eq!(expected[0], model.forward(&prompts[0]).unwrap()); - } -} diff --git a/src/models/cua_s1/native/src/multimodal.rs b/src/models/cua_s1/native/src/multimodal.rs index f6e2876..555eda8 100644 --- a/src/models/cua_s1/native/src/multimodal.rs +++ b/src/models/cua_s1/native/src/multimodal.rs @@ -151,67 +151,3 @@ pub(crate) fn validate_questions(questions: &[Question]) -> Result<()> { } Ok(()) } - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn public_questions_reject_inconsistent_options_and_media_controls() { - let mut q = Question { - name: "q".into(), - goal: "Save".into(), - keys: vec!["a".into()], - labels: vec!["A".into(), "B".into()], - }; - assert!(validate_questions(&[q.clone()]).is_err()); - q.labels.pop(); - validate_questions(&[q.clone()]).unwrap(); - q.goal = "<|image_pad|>".into(); - assert!(validate_questions(&[q]).is_err()); - } - #[test] - fn rectangular_grid_uses_three_axes_and_compressed_text_suffix() { - let ids = [7, 8, 99, 99, 99, 99, 99, 99, 9, 10]; - let p = image_positions(&ids, 99, [1, 4, 6]).unwrap(); - assert_eq!(p[0], [0, 1, 2, 2, 2, 2, 2, 2, 5, 6]); - assert_eq!(p[1], [0, 1, 2, 2, 2, 3, 3, 3, 5, 6]); - assert_eq!(p[2], [0, 1, 2, 3, 4, 2, 3, 4, 5, 6]); - } - - #[test] - fn rejects_invalid_grid_and_placeholder_counts() { - for grid in [ - [0, 2, 2], - [2, 2, 2], - [1, 0, 2], - [1, 3, 2], - [1, usize::MAX, 2], - ] { - assert!(image_positions(&[1, 99, 2], 99, grid).is_err()); - } - for ids in [vec![1, 2], vec![1, 99, 99, 2], vec![99, 1, 99]] { - assert!(image_positions(&ids, 99, [1, 2, 2]).is_err()); - } - } - - #[test] - fn prompt_expands_image_without_extra_separator() { - let q = Question { - name: "q".into(), - goal: "Save".into(), - keys: vec!["save".into()], - labels: vec!["Click Save".into()], - }; - let p = chat_image(&q, 2).unwrap(); - assert!(p.contains( - "user\n<|vision_start|><|image_pad|><|image_pad|><|vision_end|>Goal: Save\n\nApp:" - )); - assert!(p.contains( - "The current screenshot is attached.\n\nOptions:\nA. Decision \"Click Save\" -> select" - )); - assert!(p.ends_with("<|im_start|>assistant\n\n")); - assert!(chat_image(&q, 0).is_err()); - assert!(chat_image(&q, MAX_TOKENS).is_err()); - } -} diff --git a/src/models/cua_s1/native/src/provenance.rs b/src/models/cua_s1/native/src/provenance.rs index 6a42f12..4ea295b 100644 --- a/src/models/cua_s1/native/src/provenance.rs +++ b/src/models/cua_s1/native/src/provenance.rs @@ -103,26 +103,3 @@ pub fn verify_export(dir: &Path, marker: &Value) -> Result<()> { ); verify_files(dir, files, Some("cua_s1_language_export.json")) } - -#[cfg(test)] -mod tests { - use super::*; - #[test] - fn rejects_missing_export_hashes_before_loading_cuda() { - assert!(verify_export(Path::new("unused"), &serde_json::json!({})).is_err()); - } - #[test] - fn detects_same_size_mutation_and_unlisted_override() { - let dir = std::env::temp_dir().join(format!("cua-provenance-{}", std::process::id())); - fs::create_dir_all(&dir).unwrap(); - fs::write(dir.join("tensor"), b"abc").unwrap(); - let files = serde_json::json!({"tensor":{"size":3, "sha256":format!("{:x}",Sha256::digest(b"abc"))}}); - verify_files(&dir, files.as_object().unwrap(), None).unwrap(); - fs::write(dir.join("tensor"), b"abd").unwrap(); - assert!(verify_files(&dir, files.as_object().unwrap(), None).is_err()); - fs::write(dir.join("tensor"), b"abc").unwrap(); - fs::write(dir.join("override"), b"x").unwrap(); - assert!(verify_files(&dir, files.as_object().unwrap(), None).is_err()); - fs::remove_dir_all(dir).unwrap(); - } -} diff --git a/src/models/cua_s1/native/src/vision/geometry.rs b/src/models/cua_s1/native/src/vision/geometry.rs index 11842aa..6cd9f4a 100644 --- a/src/models/cua_s1/native/src/vision/geometry.rs +++ b/src/models/cua_s1/native/src/vision/geometry.rs @@ -58,38 +58,3 @@ impl Geometry { Ok(g) } } -#[cfg(test)] -mod tests { - use super::*; - #[test] - fn interpolation_corners_and_merge_order() { - let g = Geometry::new([1, 2, 4]).unwrap(); - // Sequence: (0,0), (0,1), (1,0), (1,1), (0,2), (0,3), (1,2), (1,3). - assert_eq!(&g.indices[..4], &[0, 1, 48, 49]); - assert_eq!(&g.weights[..4], &[1., 0., 0., 0.]); - assert_eq!(g.indices[2 * 4], 47 * 48); - assert_eq!(g.indices[7 * 4], 2303); - assert!((g.weights[4] - 1. / 3.).abs() < 2e-6); - assert_eq!(g.cos[0], 1.); - assert_eq!(g.sin[0], 0.); - assert!((g.sin[32 + 16] - 1f32.sin()).abs() < 1e-6); - assert!((g.sin[2 * 32] - 1f32.sin()).abs() < 1e-6); - for w in g.weights.as_chunks::<4>().0 { - assert!((w.iter().sum::() - 1.).abs() < 1e-6); - } - } - #[test] - fn geometry_rejects_video_odd_empty_and_oversize_grids() { - for grid in [ - [2, 16, 16], - [1, 0, 16], - [1, 15, 16], - [1, 128, 128], - [1, usize::MAX, 2], - ] { - assert!(Geometry::new(grid).is_err(), "accepted {grid:?}"); - } - assert!(Geometry::new([1, 64, 64]).is_ok()); - assert!(Geometry::new([1, 2, 224]).is_ok()); - } -} diff --git a/src/models/cua_s1/native/tests/multimodal.rs b/src/models/cua_s1/native/tests/multimodal.rs deleted file mode 100644 index 20afda5..0000000 --- a/src/models/cua_s1/native/tests/multimodal.rs +++ /dev/null @@ -1,90 +0,0 @@ -//! Real CUDA regression: explicit positions must not contaminate later text calls. -use std::path::PathBuf; - -use half::bf16; -use omni_cua_s1_native::{inputs::MultimodalInput, model::Model}; - -#[test] -#[ignore = "needs CUA_S1_MODEL and CUA_S1_CUDA_LIB on a GPU"] -fn text_multimodal_text_keeps_text_positions_and_overwrites_image_rows() { - let dir = PathBuf::from(std::env::var_os("CUA_S1_MODEL").expect("CUA_S1_MODEL")); - let lib = PathBuf::from(std::env::var_os("CUA_S1_CUDA_LIB").expect("CUA_S1_CUDA_LIB")); - let mut model = Model::load(&dir, &lib).unwrap(); - let ids = [32, 33, 34, 35]; - let positions = [0, 1, 2, 3]; - let baseline = model.forward(&ids).unwrap(); - let explicit = model - .forward_multimodal(&MultimodalInput { - token_ids: &ids, - image_token_indices: &[], - image_embeddings: &[], - position_ids: [&positions; 3], - }) - .unwrap(); - assert_eq!(baseline, explicit); - - let image_token = model.cfg.image_token_id.unwrap(); - let image_ids = [32, image_token, image_token, 35]; - let features = vec![bf16::ONE; 2 * model.cfg.hidden]; - let different = model - .forward_multimodal(&MultimodalInput { - token_ids: &image_ids, - image_token_indices: &[1, 2], - image_embeddings: &features, - position_ids: [&[0, 1, 1, 2], &[0, 1, 2, 2], &[0, 2, 1, 2]], - }) - .unwrap(); - assert!(different.iter().all(|x| x.is_finite())); - assert_ne!(baseline, different); - assert_eq!(baseline, model.forward(&ids).unwrap()); - - // Reusing the same layout with changed features must not retain old rows. - let features = vec![bf16::from_f32(-1.0); 2 * model.cfg.hidden]; - let changed = model - .forward_multimodal(&MultimodalInput { - token_ids: &image_ids, - image_token_indices: &[1, 2], - image_embeddings: &features, - position_ids: [&[0, 1, 1, 2], &[0, 1, 2, 2], &[0, 2, 1, 2]], - }) - .unwrap(); - assert_ne!(different, changed); - assert_eq!(baseline, model.forward(&ids).unwrap()); - - // Non-adjacent image spans exercise separate uploads and untouched text rows. - let disjoint_ids = [image_token, 33, image_token, 35]; - let disjoint = MultimodalInput { - token_ids: &disjoint_ids, - image_token_indices: &[0, 2], - image_embeddings: &features, - position_ids: [&positions; 3], - }; - let last = model.forward_multimodal(&disjoint).unwrap(); - assert_eq!(last, model.forward_multimodal(&disjoint).unwrap()); - - // A rejected boundary must not interfere with restoring the next text call. - assert!( - model - .forward_multimodal(&MultimodalInput { - image_token_indices: &[2, 0], - ..disjoint - }) - .is_err() - ); - assert_eq!(baseline, model.forward(&ids).unwrap()); - - // Cross the 1024-row allocation boundary, then use a shorter custom layout - // before reusing the immutable text tables in the larger scratch allocation. - let long_ids = vec![32; 1025]; - let long_text = model.forward(&long_ids).unwrap(); - model - .forward_multimodal(&MultimodalInput { - token_ids: &image_ids, - image_token_indices: &[1, 2], - image_embeddings: &features, - position_ids: [&[0, 1, 1, 2], &[0, 1, 2, 2], &[0, 2, 1, 2]], - }) - .unwrap(); - assert_eq!(long_text, model.forward(&long_ids).unwrap()); - assert_eq!(baseline, model.forward(&ids).unwrap()); -} diff --git a/tests/cua_s1/fixtures/image_preprocess/README.md b/tests/cua_s1/fixtures/image_preprocess/README.md deleted file mode 100644 index 1eff2a5..0000000 --- a/tests/cua_s1/fixtures/image_preprocess/README.md +++ /dev/null @@ -1,30 +0,0 @@ -# Native RGB preprocessing reference fixtures - -`preprocessor_config.json` is from pinned Qwen3.5-4B revision -`851bf6e806efd8d0a36b00ddf55e13ccb7b8cd0a`: -[upstream configuration](https://huggingface.co/Qwen/Qwen3.5-4B/resolve/851bf6e806efd8d0a36b00ddf55e13ccb7b8cd0a/preprocessor_config.json). - -`manifest.json` contains complete FP32 output byte hashes generated by the actual -Transformers `AutoImageProcessor` with this configuration on CPU. It records -package versions, platform, configuration fingerprint, deterministic RGB input -hashes, resized dimensions, grid, patch shape and image token count. No weights -or generated output tensors are committed. Rust tests duplicate only the input -generator and compare every output byte through SHA-256. - -The input generator covers constant RGB, spatial/channel ramps, checkerboards -and seeded xorshift32 noise. Cases include no resize, upsampling, downsampling, -Python ties-to-even dimensions, portrait/wide inputs, the 200:1 aspect boundary -and the request pixel cap. These establish CPU preprocessing parity on the -recorded reference environment; they are not an encoder or GPU accuracy test. - -Regenerate with Python 3.12 and the pinned packages (no GPU or model downloads): - -```sh -python -m pip install torch==2.14.0 torchvision==0.29.0 transformers==5.17.0 \ - numpy==2.5.3 Pillow==11.3.0 -python tests/cua_s1/fixtures/image_preprocess/generate.py -``` - -The generator rejects mismatched package versions. Keep new generated hashes -reviewable against the recorded source; do not update expected values merely -to make a native mismatch disappear. diff --git a/tests/cua_s1/fixtures/image_preprocess/generate.py b/tests/cua_s1/fixtures/image_preprocess/generate.py deleted file mode 100644 index 70efc1b..0000000 --- a/tests/cua_s1/fixtures/image_preprocess/generate.py +++ /dev/null @@ -1,119 +0,0 @@ -"""Regenerate CPU parity hashes using the actual pinned Hugging Face processor. - -Run in a Python 3.12 environment with torch==2.14.0, torchvision==0.29.0, -transformers==5.17.0, numpy==2.5.3 and Pillow==11.3.0. No model weights or GPU. -""" - -import hashlib -import importlib.metadata -import json -import platform -from pathlib import Path - -import numpy as np -import torch -from PIL import Image -from transformers import AutoImageProcessor - -PACKAGES = { - "torch": "2.14.0", - "torchvision": "0.29.0", - "transformers": "5.17.0", - "numpy": "2.5.3", - "Pillow": "11.3.0", -} -CASES = [ - ("single_pixel", 1, 1, "constant", 1), - ("tiny_noise", 17, 19, "noise", 17), - ("aligned_noise", 256, 256, "noise", 123), - ("odd_noise", 319, 241, "noise", 42), - ("tie_down", 272, 256, "noise", 999), - ("tie_up", 304, 256, "noise", 101), - ("both_down", 271, 271, "checker", 1), - ("one_axis_down", 256, 257, "ramp", 1), - ("wide", 640, 320, "ramp", 1), - ("portrait", 319, 641, "noise", 23), - ("aspect_limit", 200, 1, "noise", 47), - ("tall_aspect_limit", 7, 1400, "checker", 1), - ("input_pixel_limit", 2048, 512, "noise", 55), - ("large_round_up", 1600, 600, "noise", 77), -] - - -def pixels(width, height, pattern, seed): - """Only input generation is duplicated in Rust; outputs come from HF.""" - if pattern == "noise": - out = bytearray(width * height * 3) - state = seed - for i in range(len(out)): - state ^= (state << 13) & 0xFFFFFFFF - state ^= state >> 17 - state ^= (state << 5) & 0xFFFFFFFF - out[i] = state & 255 - return bytes(out) - return bytes( - ( - [0, 128, 255][c] - if pattern == "constant" - else (255 if (x + y + c) % 2 else 0) - if pattern == "checker" - else (x * 13 + y * 7 + c * 83) % 256 - ) - for y in range(height) - for x in range(width) - for c in range(3) - ) - - -def evaluate(processor, name, width, height, pattern, seed): - raw = pixels(width, height, pattern, seed) - array = np.frombuffer(raw, dtype=np.uint8).reshape(height, width, 3) - result = processor(images=Image.fromarray(array), return_tensors="pt", device="cpu") - tensor = result["pixel_values"].contiguous() - grid = result["image_grid_thw"][0].tolist() - assert tensor.dtype == torch.float32 and tensor.device.type == "cpu" - return { - "name": name, - "width": width, - "height": height, - "pattern": pattern, - "seed": seed, - "input_sha256": hashlib.sha256(raw).hexdigest(), - "grid": grid, - "shape": list(tensor.shape), - "image_tokens": grid[1] * grid[2] // 4, - "resized_width": grid[2] * 16, - "resized_height": grid[1] * 16, - "output_sha256": hashlib.sha256( - tensor.numpy().astype("I", len(data)) - + kind - + data - + struct.pack(">I", zlib.crc32(kind + data)) - ) - - -for name, kind, pixels, transparent in [ - ("l16", 0, [1000], None), - ("l16-trns", 0, [1000], struct.pack(">H", 1000)), - ("la16", 4, [1000, 0], None), - ("rgb16", 2, [1000, 255, 32768], None), - ("rgb16-trns", 2, [1000, 255, 32768], struct.pack(">HHH", 1000, 255, 32768)), - ("rgba16", 6, [1000, 255, 32768, 0], None), -]: - data = b"\x89PNG\r\n\x1a\n" + chunk( - b"IHDR", struct.pack(">IIBBBBB", 1, 1, 16, kind, 0, 0, 0) - ) - if transparent is not None: - data += chunk(b"tRNS", transparent) - data += chunk( - b"IDAT", zlib.compress(b"\0" + struct.pack(">" + "H" * len(pixels), *pixels)) - ) + chunk(b"IEND", b"") - Path(__file__).with_name(name + ".png").write_bytes(data) diff --git a/tests/cua_s1/fixtures/native_decode/l16-trns.png b/tests/cua_s1/fixtures/native_decode/l16-trns.png deleted file mode 100644 index fc7f7b88afbf2196cfe4e4943b57846d48deb1b6..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 82 zcmeAS@N?(olHy`uVBq!ia0vp^j36ukBp9;Zxr+fQrjj7PVCEN}^S^8a^0+-+978JR bBq!7X6@FxRqx0Ha52V=B)z4*}Q$iB}shtzf diff --git a/tests/cua_s1/fixtures/native_decode/l16.png b/tests/cua_s1/fixtures/native_decode/l16.png deleted file mode 100644 index cd96bb28f5aeba7c94c3904c8f7a4c7c362a40a9..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 68 zcmeAS@N?(olHy`uVBq!ia0vp^j36ukBp9;Zxr+fQZci7-kcv6U33WiJj|^{gUYqNI Oq&;2zT-G@yGywpVmJV|O diff --git a/tests/cua_s1/fixtures/native_decode/la16.png b/tests/cua_s1/fixtures/native_decode/la16.png deleted file mode 100644 index e9a3c0c7d355b5a6bfa2a753e6dfa3d0a1e9d600..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 70 zcmeAS@N?(olHy`uVBq!ia0vp^j36w)0wkaITyO+SxHyJX%t=nDJ9vPBf$1c}8=n0a Rw*n;?JYD@<);T3K0RZr&5p@6n diff --git a/tests/cua_s1/fixtures/native_decode/rgb16-trns.png b/tests/cua_s1/fixtures/native_decode/rgb16-trns.png deleted file mode 100644 index f9f707e065e7ce127f0f6db6b4701ad1d092564f..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 90 zcmeAS@N?(olHy`uVBq!ia0vp^j36w)1SAhU@1F~#*h+%@f|*}1{BK|=?G0@Qitu~7 lIEGZrNlvIc_+ft|BLl}grffU=ug`%}44$rjF6*2UngE#p7*GHJ diff --git a/tests/cua_s1/fixtures/native_decode/rgb16.png b/tests/cua_s1/fixtures/native_decode/rgb16.png deleted file mode 100644 index 2f36ac7ea0e831a0264633e0e45f75efe14a033c..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 72 zcmeAS@N?(olHy`uVBq!ia0vp^j36w)1SAhU@1F~#_&r@5Ln`JZC)6GMu)mR!fny$1 Uww?Xg=Rhe2Pgg&ebxsLQ04akKn*aa+ diff --git a/tests/cua_s1/fixtures/native_decode/rgba16.png b/tests/cua_s1/fixtures/native_decode/rgba16.png deleted file mode 100644 index 1477c0f0c3354bddb3cee98d905a5c8f2f30fd1b..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 74 zcmeAS@N?(olHy`uVBq!ia0vp^j36w)1|~CyrWMtqi WVan!{Ju3o~WbkzLb6Mw<&;$S!%Mla+ diff --git a/tests/cua_s1/fixtures/vision/README.md b/tests/cua_s1/fixtures/vision/README.md deleted file mode 100644 index 21a8ca6..0000000 --- a/tests/cua_s1/fixtures/vision/README.md +++ /dev/null @@ -1,31 +0,0 @@ -# Pinned vision metadata fixtures - -These fixtures are an independent test oracle captured from upstream metadata, -not generated by the production Rust inventory builder: - -- Base: `Qwen/Qwen3.5-4B`, revision - `851bf6e806efd8d0a36b00ddf55e13ccb7b8cd0a`. -- Adapter: `cua-ai/cua-s1-4b-0.2`, revision - `16818868b0cc7813808aae4e87b417657046ab79`, multimodal adapter. -- `config.json` and `adapter_config.json`: upstream configuration files with LF line endings. -- `base-tensors.json`: visual names, shapes, and dtypes read from the pinned - base safetensors headers (prefix `model.visual.`): 297 BF16 tensors. -- `adapter-tensors.json`: visual names, shapes, and dtypes read from the pinned - multimodal adapter safetensors header (prefix - `base_model.model.model.visual.`): 100 FP32 tensors / 50 LoRA pairs. - -No weight payloads are included. Tests construct sparse files with these shapes -and sentinel values. Production code never reads these fixtures. Fixture -matching proves structural compatibility, not checkpoint identity: verify real -checkpoint downloads against trusted upstream hashes separately. - -Exact pinned sources: - -- [Base configuration](https://huggingface.co/Qwen/Qwen3.5-4B/resolve/851bf6e806efd8d0a36b00ddf55e13ccb7b8cd0a/config.json) -- [Multimodal adapter configuration](https://huggingface.co/cua-ai/cua-s1-4b-0.2/resolve/16818868b0cc7813808aae4e87b417657046ab79/multimodal/adapter_config.json) -- [Base second shard, containing the vision tensors](https://huggingface.co/Qwen/Qwen3.5-4B/resolve/851bf6e806efd8d0a36b00ddf55e13ccb7b8cd0a/model.safetensors-00002-of-00002.safetensors) -- [Multimodal adapter safetensors](https://huggingface.co/cua-ai/cua-s1-4b-0.2/resolve/16818868b0cc7813808aae4e87b417657046ab79/multimodal/adapter_model.safetensors) - -The tensor inventories retain only the `shape` and `dtype` fields of the visual -entries in those safetensors headers; byte offsets and language tensors are not -part of the oracle. diff --git a/tests/cua_s1/fixtures/vision/adapter-tensors.json b/tests/cua_s1/fixtures/vision/adapter-tensors.json deleted file mode 100644 index ac90604..0000000 --- a/tests/cua_s1/fixtures/vision/adapter-tensors.json +++ /dev/null @@ -1,702 +0,0 @@ -{ - "base_model.model.model.visual.blocks.0.mlp.linear_fc1.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 1024 - ] - }, - "base_model.model.model.visual.blocks.0.mlp.linear_fc1.lora_B.weight": { - "dtype": "F32", - "shape": [ - 4096, - 16 - ] - }, - "base_model.model.model.visual.blocks.0.mlp.linear_fc2.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 4096 - ] - }, - "base_model.model.model.visual.blocks.0.mlp.linear_fc2.lora_B.weight": { - "dtype": "F32", - "shape": [ - 1024, - 16 - ] - }, - "base_model.model.model.visual.blocks.1.mlp.linear_fc1.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 1024 - ] - }, - "base_model.model.model.visual.blocks.1.mlp.linear_fc1.lora_B.weight": { - "dtype": "F32", - "shape": [ - 4096, - 16 - ] - }, - "base_model.model.model.visual.blocks.1.mlp.linear_fc2.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 4096 - ] - }, - "base_model.model.model.visual.blocks.1.mlp.linear_fc2.lora_B.weight": { - "dtype": "F32", - "shape": [ - 1024, - 16 - ] - }, - "base_model.model.model.visual.blocks.10.mlp.linear_fc1.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 1024 - ] - }, - "base_model.model.model.visual.blocks.10.mlp.linear_fc1.lora_B.weight": { - "dtype": "F32", - "shape": [ - 4096, - 16 - ] - }, - "base_model.model.model.visual.blocks.10.mlp.linear_fc2.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 4096 - ] - }, - "base_model.model.model.visual.blocks.10.mlp.linear_fc2.lora_B.weight": { - "dtype": "F32", - "shape": [ - 1024, - 16 - ] - }, - "base_model.model.model.visual.blocks.11.mlp.linear_fc1.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 1024 - ] - }, - "base_model.model.model.visual.blocks.11.mlp.linear_fc1.lora_B.weight": { - "dtype": "F32", - "shape": [ - 4096, - 16 - ] - }, - "base_model.model.model.visual.blocks.11.mlp.linear_fc2.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 4096 - ] - }, - "base_model.model.model.visual.blocks.11.mlp.linear_fc2.lora_B.weight": { - "dtype": "F32", - "shape": [ - 1024, - 16 - ] - }, - "base_model.model.model.visual.blocks.12.mlp.linear_fc1.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 1024 - ] - }, - "base_model.model.model.visual.blocks.12.mlp.linear_fc1.lora_B.weight": { - "dtype": "F32", - "shape": [ - 4096, - 16 - ] - }, - "base_model.model.model.visual.blocks.12.mlp.linear_fc2.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 4096 - ] - }, - "base_model.model.model.visual.blocks.12.mlp.linear_fc2.lora_B.weight": { - "dtype": "F32", - "shape": [ - 1024, - 16 - ] - }, - "base_model.model.model.visual.blocks.13.mlp.linear_fc1.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 1024 - ] - }, - "base_model.model.model.visual.blocks.13.mlp.linear_fc1.lora_B.weight": { - "dtype": "F32", - "shape": [ - 4096, - 16 - ] - }, - "base_model.model.model.visual.blocks.13.mlp.linear_fc2.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 4096 - ] - }, - "base_model.model.model.visual.blocks.13.mlp.linear_fc2.lora_B.weight": { - "dtype": "F32", - "shape": [ - 1024, - 16 - ] - }, - "base_model.model.model.visual.blocks.14.mlp.linear_fc1.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 1024 - ] - }, - "base_model.model.model.visual.blocks.14.mlp.linear_fc1.lora_B.weight": { - "dtype": "F32", - "shape": [ - 4096, - 16 - ] - }, - "base_model.model.model.visual.blocks.14.mlp.linear_fc2.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 4096 - ] - }, - "base_model.model.model.visual.blocks.14.mlp.linear_fc2.lora_B.weight": { - "dtype": "F32", - "shape": [ - 1024, - 16 - ] - }, - "base_model.model.model.visual.blocks.15.mlp.linear_fc1.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 1024 - ] - }, - "base_model.model.model.visual.blocks.15.mlp.linear_fc1.lora_B.weight": { - "dtype": "F32", - "shape": [ - 4096, - 16 - ] - }, - "base_model.model.model.visual.blocks.15.mlp.linear_fc2.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 4096 - ] - }, - "base_model.model.model.visual.blocks.15.mlp.linear_fc2.lora_B.weight": { - "dtype": "F32", - "shape": [ - 1024, - 16 - ] - }, - "base_model.model.model.visual.blocks.16.mlp.linear_fc1.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 1024 - ] - }, - "base_model.model.model.visual.blocks.16.mlp.linear_fc1.lora_B.weight": { - "dtype": "F32", - "shape": [ - 4096, - 16 - ] - }, - "base_model.model.model.visual.blocks.16.mlp.linear_fc2.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 4096 - ] - }, - "base_model.model.model.visual.blocks.16.mlp.linear_fc2.lora_B.weight": { - "dtype": "F32", - "shape": [ - 1024, - 16 - ] - }, - "base_model.model.model.visual.blocks.17.mlp.linear_fc1.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 1024 - ] - }, - "base_model.model.model.visual.blocks.17.mlp.linear_fc1.lora_B.weight": { - "dtype": "F32", - "shape": [ - 4096, - 16 - ] - }, - "base_model.model.model.visual.blocks.17.mlp.linear_fc2.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 4096 - ] - }, - "base_model.model.model.visual.blocks.17.mlp.linear_fc2.lora_B.weight": { - "dtype": "F32", - "shape": [ - 1024, - 16 - ] - }, - "base_model.model.model.visual.blocks.18.mlp.linear_fc1.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 1024 - ] - }, - "base_model.model.model.visual.blocks.18.mlp.linear_fc1.lora_B.weight": { - "dtype": "F32", - "shape": [ - 4096, - 16 - ] - }, - "base_model.model.model.visual.blocks.18.mlp.linear_fc2.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 4096 - ] - }, - "base_model.model.model.visual.blocks.18.mlp.linear_fc2.lora_B.weight": { - "dtype": "F32", - "shape": [ - 1024, - 16 - ] - }, - "base_model.model.model.visual.blocks.19.mlp.linear_fc1.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 1024 - ] - }, - "base_model.model.model.visual.blocks.19.mlp.linear_fc1.lora_B.weight": { - "dtype": "F32", - "shape": [ - 4096, - 16 - ] - }, - "base_model.model.model.visual.blocks.19.mlp.linear_fc2.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 4096 - ] - }, - "base_model.model.model.visual.blocks.19.mlp.linear_fc2.lora_B.weight": { - "dtype": "F32", - "shape": [ - 1024, - 16 - ] - }, - "base_model.model.model.visual.blocks.2.mlp.linear_fc1.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 1024 - ] - }, - "base_model.model.model.visual.blocks.2.mlp.linear_fc1.lora_B.weight": { - "dtype": "F32", - "shape": [ - 4096, - 16 - ] - }, - "base_model.model.model.visual.blocks.2.mlp.linear_fc2.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 4096 - ] - }, - "base_model.model.model.visual.blocks.2.mlp.linear_fc2.lora_B.weight": { - "dtype": "F32", - "shape": [ - 1024, - 16 - ] - }, - "base_model.model.model.visual.blocks.20.mlp.linear_fc1.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 1024 - ] - }, - "base_model.model.model.visual.blocks.20.mlp.linear_fc1.lora_B.weight": { - "dtype": "F32", - "shape": [ - 4096, - 16 - ] - }, - "base_model.model.model.visual.blocks.20.mlp.linear_fc2.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 4096 - ] - }, - "base_model.model.model.visual.blocks.20.mlp.linear_fc2.lora_B.weight": { - "dtype": "F32", - "shape": [ - 1024, - 16 - ] - }, - "base_model.model.model.visual.blocks.21.mlp.linear_fc1.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 1024 - ] - }, - "base_model.model.model.visual.blocks.21.mlp.linear_fc1.lora_B.weight": { - "dtype": "F32", - "shape": [ - 4096, - 16 - ] - }, - "base_model.model.model.visual.blocks.21.mlp.linear_fc2.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 4096 - ] - }, - "base_model.model.model.visual.blocks.21.mlp.linear_fc2.lora_B.weight": { - "dtype": "F32", - "shape": [ - 1024, - 16 - ] - }, - "base_model.model.model.visual.blocks.22.mlp.linear_fc1.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 1024 - ] - }, - "base_model.model.model.visual.blocks.22.mlp.linear_fc1.lora_B.weight": { - "dtype": "F32", - "shape": [ - 4096, - 16 - ] - }, - "base_model.model.model.visual.blocks.22.mlp.linear_fc2.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 4096 - ] - }, - "base_model.model.model.visual.blocks.22.mlp.linear_fc2.lora_B.weight": { - "dtype": "F32", - "shape": [ - 1024, - 16 - ] - }, - "base_model.model.model.visual.blocks.23.mlp.linear_fc1.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 1024 - ] - }, - "base_model.model.model.visual.blocks.23.mlp.linear_fc1.lora_B.weight": { - "dtype": "F32", - "shape": [ - 4096, - 16 - ] - }, - "base_model.model.model.visual.blocks.23.mlp.linear_fc2.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 4096 - ] - }, - "base_model.model.model.visual.blocks.23.mlp.linear_fc2.lora_B.weight": { - "dtype": "F32", - "shape": [ - 1024, - 16 - ] - }, - "base_model.model.model.visual.blocks.3.mlp.linear_fc1.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 1024 - ] - }, - "base_model.model.model.visual.blocks.3.mlp.linear_fc1.lora_B.weight": { - "dtype": "F32", - "shape": [ - 4096, - 16 - ] - }, - "base_model.model.model.visual.blocks.3.mlp.linear_fc2.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 4096 - ] - }, - "base_model.model.model.visual.blocks.3.mlp.linear_fc2.lora_B.weight": { - "dtype": "F32", - "shape": [ - 1024, - 16 - ] - }, - "base_model.model.model.visual.blocks.4.mlp.linear_fc1.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 1024 - ] - }, - "base_model.model.model.visual.blocks.4.mlp.linear_fc1.lora_B.weight": { - "dtype": "F32", - "shape": [ - 4096, - 16 - ] - }, - "base_model.model.model.visual.blocks.4.mlp.linear_fc2.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 4096 - ] - }, - "base_model.model.model.visual.blocks.4.mlp.linear_fc2.lora_B.weight": { - "dtype": "F32", - "shape": [ - 1024, - 16 - ] - }, - "base_model.model.model.visual.blocks.5.mlp.linear_fc1.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 1024 - ] - }, - "base_model.model.model.visual.blocks.5.mlp.linear_fc1.lora_B.weight": { - "dtype": "F32", - "shape": [ - 4096, - 16 - ] - }, - "base_model.model.model.visual.blocks.5.mlp.linear_fc2.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 4096 - ] - }, - "base_model.model.model.visual.blocks.5.mlp.linear_fc2.lora_B.weight": { - "dtype": "F32", - "shape": [ - 1024, - 16 - ] - }, - "base_model.model.model.visual.blocks.6.mlp.linear_fc1.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 1024 - ] - }, - "base_model.model.model.visual.blocks.6.mlp.linear_fc1.lora_B.weight": { - "dtype": "F32", - "shape": [ - 4096, - 16 - ] - }, - "base_model.model.model.visual.blocks.6.mlp.linear_fc2.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 4096 - ] - }, - "base_model.model.model.visual.blocks.6.mlp.linear_fc2.lora_B.weight": { - "dtype": "F32", - "shape": [ - 1024, - 16 - ] - }, - "base_model.model.model.visual.blocks.7.mlp.linear_fc1.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 1024 - ] - }, - "base_model.model.model.visual.blocks.7.mlp.linear_fc1.lora_B.weight": { - "dtype": "F32", - "shape": [ - 4096, - 16 - ] - }, - "base_model.model.model.visual.blocks.7.mlp.linear_fc2.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 4096 - ] - }, - "base_model.model.model.visual.blocks.7.mlp.linear_fc2.lora_B.weight": { - "dtype": "F32", - "shape": [ - 1024, - 16 - ] - }, - "base_model.model.model.visual.blocks.8.mlp.linear_fc1.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 1024 - ] - }, - "base_model.model.model.visual.blocks.8.mlp.linear_fc1.lora_B.weight": { - "dtype": "F32", - "shape": [ - 4096, - 16 - ] - }, - "base_model.model.model.visual.blocks.8.mlp.linear_fc2.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 4096 - ] - }, - "base_model.model.model.visual.blocks.8.mlp.linear_fc2.lora_B.weight": { - "dtype": "F32", - "shape": [ - 1024, - 16 - ] - }, - "base_model.model.model.visual.blocks.9.mlp.linear_fc1.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 1024 - ] - }, - "base_model.model.model.visual.blocks.9.mlp.linear_fc1.lora_B.weight": { - "dtype": "F32", - "shape": [ - 4096, - 16 - ] - }, - "base_model.model.model.visual.blocks.9.mlp.linear_fc2.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 4096 - ] - }, - "base_model.model.model.visual.blocks.9.mlp.linear_fc2.lora_B.weight": { - "dtype": "F32", - "shape": [ - 1024, - 16 - ] - }, - "base_model.model.model.visual.merger.linear_fc1.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 4096 - ] - }, - "base_model.model.model.visual.merger.linear_fc1.lora_B.weight": { - "dtype": "F32", - "shape": [ - 4096, - 16 - ] - }, - "base_model.model.model.visual.merger.linear_fc2.lora_A.weight": { - "dtype": "F32", - "shape": [ - 16, - 4096 - ] - }, - "base_model.model.model.visual.merger.linear_fc2.lora_B.weight": { - "dtype": "F32", - "shape": [ - 2560, - 16 - ] - } -} diff --git a/tests/cua_s1/fixtures/vision/adapter_config.json b/tests/cua_s1/fixtures/vision/adapter_config.json deleted file mode 100644 index 52b492f..0000000 --- a/tests/cua_s1/fixtures/vision/adapter_config.json +++ /dev/null @@ -1,48 +0,0 @@ -{ - "alora_invocation_tokens": null, - "alpha_pattern": {}, - "arrow_config": null, - "auto_mapping": null, - "base_model_name_or_path": "Qwen/Qwen3.5-4B", - "bias": "none", - "corda_config": null, - "ensure_weight_tying": false, - "eva_config": null, - "exclude_modules": null, - "fan_in_fan_out": false, - "inference_mode": true, - "init_lora_weights": true, - "layer_replication": null, - "layers_pattern": null, - "layers_to_transform": null, - "loftq_config": {}, - "lora_alpha": 32, - "lora_bias": false, - "lora_dropout": 0.05, - "megatron_config": null, - "megatron_core": "megatron.core", - "modules_to_save": null, - "peft_type": "LORA", - "peft_version": "0.18.1", - "qalora_group_size": 16, - "r": 16, - "rank_pattern": {}, - "revision": null, - "target_modules": [ - "up_proj", - "k_proj", - "linear_fc1", - "q_proj", - "linear_fc2", - "down_proj", - "gate_proj", - "o_proj", - "v_proj" - ], - "target_parameters": null, - "task_type": "CAUSAL_LM", - "trainable_token_indices": null, - "use_dora": false, - "use_qalora": false, - "use_rslora": false -} diff --git a/tests/cua_s1/fixtures/vision/base-tensors.json b/tests/cua_s1/fixtures/vision/base-tensors.json deleted file mode 100644 index 0c2d0b9..0000000 --- a/tests/cua_s1/fixtures/vision/base-tensors.json +++ /dev/null @@ -1,1887 +0,0 @@ -{ - "model.visual.blocks.0.attn.proj.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.0.attn.proj.weight": { - "dtype": "BF16", - "shape": [ - 1024, - 1024 - ] - }, - "model.visual.blocks.0.attn.qkv.bias": { - "dtype": "BF16", - "shape": [ - 3072 - ] - }, - "model.visual.blocks.0.attn.qkv.weight": { - "dtype": "BF16", - "shape": [ - 3072, - 1024 - ] - }, - "model.visual.blocks.0.mlp.linear_fc1.bias": { - "dtype": "BF16", - "shape": [ - 4096 - ] - }, - "model.visual.blocks.0.mlp.linear_fc1.weight": { - "dtype": "BF16", - "shape": [ - 4096, - 1024 - ] - }, - "model.visual.blocks.0.mlp.linear_fc2.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.0.mlp.linear_fc2.weight": { - "dtype": "BF16", - "shape": [ - 1024, - 4096 - ] - }, - "model.visual.blocks.0.norm1.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.0.norm1.weight": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.0.norm2.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.0.norm2.weight": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.1.attn.proj.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.1.attn.proj.weight": { - "dtype": "BF16", - "shape": [ - 1024, - 1024 - ] - }, - "model.visual.blocks.1.attn.qkv.bias": { - "dtype": "BF16", - "shape": [ - 3072 - ] - }, - "model.visual.blocks.1.attn.qkv.weight": { - "dtype": "BF16", - "shape": [ - 3072, - 1024 - ] - }, - "model.visual.blocks.1.mlp.linear_fc1.bias": { - "dtype": "BF16", - "shape": [ - 4096 - ] - }, - "model.visual.blocks.1.mlp.linear_fc1.weight": { - "dtype": "BF16", - "shape": [ - 4096, - 1024 - ] - }, - "model.visual.blocks.1.mlp.linear_fc2.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.1.mlp.linear_fc2.weight": { - "dtype": "BF16", - "shape": [ - 1024, - 4096 - ] - }, - "model.visual.blocks.1.norm1.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.1.norm1.weight": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.1.norm2.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.1.norm2.weight": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.10.attn.proj.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.10.attn.proj.weight": { - "dtype": "BF16", - "shape": [ - 1024, - 1024 - ] - }, - "model.visual.blocks.10.attn.qkv.bias": { - "dtype": "BF16", - "shape": [ - 3072 - ] - }, - "model.visual.blocks.10.attn.qkv.weight": { - "dtype": "BF16", - "shape": [ - 3072, - 1024 - ] - }, - "model.visual.blocks.10.mlp.linear_fc1.bias": { - "dtype": "BF16", - "shape": [ - 4096 - ] - }, - "model.visual.blocks.10.mlp.linear_fc1.weight": { - "dtype": "BF16", - "shape": [ - 4096, - 1024 - ] - }, - "model.visual.blocks.10.mlp.linear_fc2.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.10.mlp.linear_fc2.weight": { - "dtype": "BF16", - "shape": [ - 1024, - 4096 - ] - }, - "model.visual.blocks.10.norm1.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.10.norm1.weight": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.10.norm2.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.10.norm2.weight": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.11.attn.proj.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.11.attn.proj.weight": { - "dtype": "BF16", - "shape": [ - 1024, - 1024 - ] - }, - "model.visual.blocks.11.attn.qkv.bias": { - "dtype": "BF16", - "shape": [ - 3072 - ] - }, - "model.visual.blocks.11.attn.qkv.weight": { - "dtype": "BF16", - "shape": [ - 3072, - 1024 - ] - }, - "model.visual.blocks.11.mlp.linear_fc1.bias": { - "dtype": "BF16", - "shape": [ - 4096 - ] - }, - "model.visual.blocks.11.mlp.linear_fc1.weight": { - "dtype": "BF16", - "shape": [ - 4096, - 1024 - ] - }, - "model.visual.blocks.11.mlp.linear_fc2.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.11.mlp.linear_fc2.weight": { - "dtype": "BF16", - "shape": [ - 1024, - 4096 - ] - }, - "model.visual.blocks.11.norm1.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.11.norm1.weight": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.11.norm2.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.11.norm2.weight": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.12.attn.proj.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.12.attn.proj.weight": { - "dtype": "BF16", - "shape": [ - 1024, - 1024 - ] - }, - "model.visual.blocks.12.attn.qkv.bias": { - "dtype": "BF16", - "shape": [ - 3072 - ] - }, - "model.visual.blocks.12.attn.qkv.weight": { - "dtype": "BF16", - "shape": [ - 3072, - 1024 - ] - }, - "model.visual.blocks.12.mlp.linear_fc1.bias": { - "dtype": "BF16", - "shape": [ - 4096 - ] - }, - "model.visual.blocks.12.mlp.linear_fc1.weight": { - "dtype": "BF16", - "shape": [ - 4096, - 1024 - ] - }, - "model.visual.blocks.12.mlp.linear_fc2.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.12.mlp.linear_fc2.weight": { - "dtype": "BF16", - "shape": [ - 1024, - 4096 - ] - }, - "model.visual.blocks.12.norm1.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.12.norm1.weight": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.12.norm2.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.12.norm2.weight": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.13.attn.proj.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.13.attn.proj.weight": { - "dtype": "BF16", - "shape": [ - 1024, - 1024 - ] - }, - "model.visual.blocks.13.attn.qkv.bias": { - "dtype": "BF16", - "shape": [ - 3072 - ] - }, - "model.visual.blocks.13.attn.qkv.weight": { - "dtype": "BF16", - "shape": [ - 3072, - 1024 - ] - }, - "model.visual.blocks.13.mlp.linear_fc1.bias": { - "dtype": "BF16", - "shape": [ - 4096 - ] - }, - "model.visual.blocks.13.mlp.linear_fc1.weight": { - "dtype": "BF16", - "shape": [ - 4096, - 1024 - ] - }, - "model.visual.blocks.13.mlp.linear_fc2.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.13.mlp.linear_fc2.weight": { - "dtype": "BF16", - "shape": [ - 1024, - 4096 - ] - }, - "model.visual.blocks.13.norm1.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.13.norm1.weight": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.13.norm2.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.13.norm2.weight": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.14.attn.proj.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.14.attn.proj.weight": { - "dtype": "BF16", - "shape": [ - 1024, - 1024 - ] - }, - "model.visual.blocks.14.attn.qkv.bias": { - "dtype": "BF16", - "shape": [ - 3072 - ] - }, - "model.visual.blocks.14.attn.qkv.weight": { - "dtype": "BF16", - "shape": [ - 3072, - 1024 - ] - }, - "model.visual.blocks.14.mlp.linear_fc1.bias": { - "dtype": "BF16", - "shape": [ - 4096 - ] - }, - "model.visual.blocks.14.mlp.linear_fc1.weight": { - "dtype": "BF16", - "shape": [ - 4096, - 1024 - ] - }, - "model.visual.blocks.14.mlp.linear_fc2.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.14.mlp.linear_fc2.weight": { - "dtype": "BF16", - "shape": [ - 1024, - 4096 - ] - }, - "model.visual.blocks.14.norm1.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.14.norm1.weight": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.14.norm2.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.14.norm2.weight": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.15.attn.proj.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.15.attn.proj.weight": { - "dtype": "BF16", - "shape": [ - 1024, - 1024 - ] - }, - "model.visual.blocks.15.attn.qkv.bias": { - "dtype": "BF16", - "shape": [ - 3072 - ] - }, - "model.visual.blocks.15.attn.qkv.weight": { - "dtype": "BF16", - "shape": [ - 3072, - 1024 - ] - }, - "model.visual.blocks.15.mlp.linear_fc1.bias": { - "dtype": "BF16", - "shape": [ - 4096 - ] - }, - "model.visual.blocks.15.mlp.linear_fc1.weight": { - "dtype": "BF16", - "shape": [ - 4096, - 1024 - ] - }, - "model.visual.blocks.15.mlp.linear_fc2.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.15.mlp.linear_fc2.weight": { - "dtype": "BF16", - "shape": [ - 1024, - 4096 - ] - }, - "model.visual.blocks.15.norm1.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.15.norm1.weight": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.15.norm2.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.15.norm2.weight": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.16.attn.proj.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.16.attn.proj.weight": { - "dtype": "BF16", - "shape": [ - 1024, - 1024 - ] - }, - "model.visual.blocks.16.attn.qkv.bias": { - "dtype": "BF16", - "shape": [ - 3072 - ] - }, - "model.visual.blocks.16.attn.qkv.weight": { - "dtype": "BF16", - "shape": [ - 3072, - 1024 - ] - }, - "model.visual.blocks.16.mlp.linear_fc1.bias": { - "dtype": "BF16", - "shape": [ - 4096 - ] - }, - "model.visual.blocks.16.mlp.linear_fc1.weight": { - "dtype": "BF16", - "shape": [ - 4096, - 1024 - ] - }, - "model.visual.blocks.16.mlp.linear_fc2.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.16.mlp.linear_fc2.weight": { - "dtype": "BF16", - "shape": [ - 1024, - 4096 - ] - }, - "model.visual.blocks.16.norm1.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.16.norm1.weight": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.16.norm2.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.16.norm2.weight": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.17.attn.proj.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.17.attn.proj.weight": { - "dtype": "BF16", - "shape": [ - 1024, - 1024 - ] - }, - "model.visual.blocks.17.attn.qkv.bias": { - "dtype": "BF16", - "shape": [ - 3072 - ] - }, - "model.visual.blocks.17.attn.qkv.weight": { - "dtype": "BF16", - "shape": [ - 3072, - 1024 - ] - }, - "model.visual.blocks.17.mlp.linear_fc1.bias": { - "dtype": "BF16", - "shape": [ - 4096 - ] - }, - "model.visual.blocks.17.mlp.linear_fc1.weight": { - "dtype": "BF16", - "shape": [ - 4096, - 1024 - ] - }, - "model.visual.blocks.17.mlp.linear_fc2.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.17.mlp.linear_fc2.weight": { - "dtype": "BF16", - "shape": [ - 1024, - 4096 - ] - }, - "model.visual.blocks.17.norm1.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.17.norm1.weight": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.17.norm2.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.17.norm2.weight": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.18.attn.proj.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.18.attn.proj.weight": { - "dtype": "BF16", - "shape": [ - 1024, - 1024 - ] - }, - "model.visual.blocks.18.attn.qkv.bias": { - "dtype": "BF16", - "shape": [ - 3072 - ] - }, - "model.visual.blocks.18.attn.qkv.weight": { - "dtype": "BF16", - "shape": [ - 3072, - 1024 - ] - }, - "model.visual.blocks.18.mlp.linear_fc1.bias": { - "dtype": "BF16", - "shape": [ - 4096 - ] - }, - "model.visual.blocks.18.mlp.linear_fc1.weight": { - "dtype": "BF16", - "shape": [ - 4096, - 1024 - ] - }, - "model.visual.blocks.18.mlp.linear_fc2.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.18.mlp.linear_fc2.weight": { - "dtype": "BF16", - "shape": [ - 1024, - 4096 - ] - }, - "model.visual.blocks.18.norm1.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.18.norm1.weight": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.18.norm2.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.18.norm2.weight": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.19.attn.proj.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.19.attn.proj.weight": { - "dtype": "BF16", - "shape": [ - 1024, - 1024 - ] - }, - "model.visual.blocks.19.attn.qkv.bias": { - "dtype": "BF16", - "shape": [ - 3072 - ] - }, - "model.visual.blocks.19.attn.qkv.weight": { - "dtype": "BF16", - "shape": [ - 3072, - 1024 - ] - }, - "model.visual.blocks.19.mlp.linear_fc1.bias": { - "dtype": "BF16", - "shape": [ - 4096 - ] - }, - "model.visual.blocks.19.mlp.linear_fc1.weight": { - "dtype": "BF16", - "shape": [ - 4096, - 1024 - ] - }, - "model.visual.blocks.19.mlp.linear_fc2.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.19.mlp.linear_fc2.weight": { - "dtype": "BF16", - "shape": [ - 1024, - 4096 - ] - }, - "model.visual.blocks.19.norm1.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.19.norm1.weight": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.19.norm2.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.19.norm2.weight": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.2.attn.proj.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.2.attn.proj.weight": { - "dtype": "BF16", - "shape": [ - 1024, - 1024 - ] - }, - "model.visual.blocks.2.attn.qkv.bias": { - "dtype": "BF16", - "shape": [ - 3072 - ] - }, - "model.visual.blocks.2.attn.qkv.weight": { - "dtype": "BF16", - "shape": [ - 3072, - 1024 - ] - }, - "model.visual.blocks.2.mlp.linear_fc1.bias": { - "dtype": "BF16", - "shape": [ - 4096 - ] - }, - "model.visual.blocks.2.mlp.linear_fc1.weight": { - "dtype": "BF16", - "shape": [ - 4096, - 1024 - ] - }, - "model.visual.blocks.2.mlp.linear_fc2.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.2.mlp.linear_fc2.weight": { - "dtype": "BF16", - "shape": [ - 1024, - 4096 - ] - }, - "model.visual.blocks.2.norm1.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.2.norm1.weight": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.2.norm2.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.2.norm2.weight": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.20.attn.proj.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.20.attn.proj.weight": { - "dtype": "BF16", - "shape": [ - 1024, - 1024 - ] - }, - "model.visual.blocks.20.attn.qkv.bias": { - "dtype": "BF16", - "shape": [ - 3072 - ] - }, - "model.visual.blocks.20.attn.qkv.weight": { - "dtype": "BF16", - "shape": [ - 3072, - 1024 - ] - }, - "model.visual.blocks.20.mlp.linear_fc1.bias": { - "dtype": "BF16", - "shape": [ - 4096 - ] - }, - "model.visual.blocks.20.mlp.linear_fc1.weight": { - "dtype": "BF16", - "shape": [ - 4096, - 1024 - ] - }, - "model.visual.blocks.20.mlp.linear_fc2.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.20.mlp.linear_fc2.weight": { - "dtype": "BF16", - "shape": [ - 1024, - 4096 - ] - }, - "model.visual.blocks.20.norm1.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.20.norm1.weight": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.20.norm2.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.20.norm2.weight": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.21.attn.proj.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.21.attn.proj.weight": { - "dtype": "BF16", - "shape": [ - 1024, - 1024 - ] - }, - "model.visual.blocks.21.attn.qkv.bias": { - "dtype": "BF16", - "shape": [ - 3072 - ] - }, - "model.visual.blocks.21.attn.qkv.weight": { - "dtype": "BF16", - "shape": [ - 3072, - 1024 - ] - }, - "model.visual.blocks.21.mlp.linear_fc1.bias": { - "dtype": "BF16", - "shape": [ - 4096 - ] - }, - "model.visual.blocks.21.mlp.linear_fc1.weight": { - "dtype": "BF16", - "shape": [ - 4096, - 1024 - ] - }, - "model.visual.blocks.21.mlp.linear_fc2.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.21.mlp.linear_fc2.weight": { - "dtype": "BF16", - "shape": [ - 1024, - 4096 - ] - }, - "model.visual.blocks.21.norm1.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.21.norm1.weight": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.21.norm2.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.21.norm2.weight": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.22.attn.proj.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.22.attn.proj.weight": { - "dtype": "BF16", - "shape": [ - 1024, - 1024 - ] - }, - "model.visual.blocks.22.attn.qkv.bias": { - "dtype": "BF16", - "shape": [ - 3072 - ] - }, - "model.visual.blocks.22.attn.qkv.weight": { - "dtype": "BF16", - "shape": [ - 3072, - 1024 - ] - }, - "model.visual.blocks.22.mlp.linear_fc1.bias": { - "dtype": "BF16", - "shape": [ - 4096 - ] - }, - "model.visual.blocks.22.mlp.linear_fc1.weight": { - "dtype": "BF16", - "shape": [ - 4096, - 1024 - ] - }, - "model.visual.blocks.22.mlp.linear_fc2.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.22.mlp.linear_fc2.weight": { - "dtype": "BF16", - "shape": [ - 1024, - 4096 - ] - }, - "model.visual.blocks.22.norm1.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.22.norm1.weight": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.22.norm2.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.22.norm2.weight": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.23.attn.proj.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.23.attn.proj.weight": { - "dtype": "BF16", - "shape": [ - 1024, - 1024 - ] - }, - "model.visual.blocks.23.attn.qkv.bias": { - "dtype": "BF16", - "shape": [ - 3072 - ] - }, - "model.visual.blocks.23.attn.qkv.weight": { - "dtype": "BF16", - "shape": [ - 3072, - 1024 - ] - }, - "model.visual.blocks.23.mlp.linear_fc1.bias": { - "dtype": "BF16", - "shape": [ - 4096 - ] - }, - "model.visual.blocks.23.mlp.linear_fc1.weight": { - "dtype": "BF16", - "shape": [ - 4096, - 1024 - ] - }, - "model.visual.blocks.23.mlp.linear_fc2.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.23.mlp.linear_fc2.weight": { - "dtype": "BF16", - "shape": [ - 1024, - 4096 - ] - }, - "model.visual.blocks.23.norm1.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.23.norm1.weight": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.23.norm2.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.23.norm2.weight": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.3.attn.proj.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.3.attn.proj.weight": { - "dtype": "BF16", - "shape": [ - 1024, - 1024 - ] - }, - "model.visual.blocks.3.attn.qkv.bias": { - "dtype": "BF16", - "shape": [ - 3072 - ] - }, - "model.visual.blocks.3.attn.qkv.weight": { - "dtype": "BF16", - "shape": [ - 3072, - 1024 - ] - }, - "model.visual.blocks.3.mlp.linear_fc1.bias": { - "dtype": "BF16", - "shape": [ - 4096 - ] - }, - "model.visual.blocks.3.mlp.linear_fc1.weight": { - "dtype": "BF16", - "shape": [ - 4096, - 1024 - ] - }, - "model.visual.blocks.3.mlp.linear_fc2.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.3.mlp.linear_fc2.weight": { - "dtype": "BF16", - "shape": [ - 1024, - 4096 - ] - }, - "model.visual.blocks.3.norm1.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.3.norm1.weight": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.3.norm2.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.3.norm2.weight": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.4.attn.proj.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.4.attn.proj.weight": { - "dtype": "BF16", - "shape": [ - 1024, - 1024 - ] - }, - "model.visual.blocks.4.attn.qkv.bias": { - "dtype": "BF16", - "shape": [ - 3072 - ] - }, - "model.visual.blocks.4.attn.qkv.weight": { - "dtype": "BF16", - "shape": [ - 3072, - 1024 - ] - }, - "model.visual.blocks.4.mlp.linear_fc1.bias": { - "dtype": "BF16", - "shape": [ - 4096 - ] - }, - "model.visual.blocks.4.mlp.linear_fc1.weight": { - "dtype": "BF16", - "shape": [ - 4096, - 1024 - ] - }, - "model.visual.blocks.4.mlp.linear_fc2.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.4.mlp.linear_fc2.weight": { - "dtype": "BF16", - "shape": [ - 1024, - 4096 - ] - }, - "model.visual.blocks.4.norm1.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.4.norm1.weight": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.4.norm2.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.4.norm2.weight": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.5.attn.proj.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.5.attn.proj.weight": { - "dtype": "BF16", - "shape": [ - 1024, - 1024 - ] - }, - "model.visual.blocks.5.attn.qkv.bias": { - "dtype": "BF16", - "shape": [ - 3072 - ] - }, - "model.visual.blocks.5.attn.qkv.weight": { - "dtype": "BF16", - "shape": [ - 3072, - 1024 - ] - }, - "model.visual.blocks.5.mlp.linear_fc1.bias": { - "dtype": "BF16", - "shape": [ - 4096 - ] - }, - "model.visual.blocks.5.mlp.linear_fc1.weight": { - "dtype": "BF16", - "shape": [ - 4096, - 1024 - ] - }, - "model.visual.blocks.5.mlp.linear_fc2.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.5.mlp.linear_fc2.weight": { - "dtype": "BF16", - "shape": [ - 1024, - 4096 - ] - }, - "model.visual.blocks.5.norm1.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.5.norm1.weight": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.5.norm2.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.5.norm2.weight": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.6.attn.proj.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.6.attn.proj.weight": { - "dtype": "BF16", - "shape": [ - 1024, - 1024 - ] - }, - "model.visual.blocks.6.attn.qkv.bias": { - "dtype": "BF16", - "shape": [ - 3072 - ] - }, - "model.visual.blocks.6.attn.qkv.weight": { - "dtype": "BF16", - "shape": [ - 3072, - 1024 - ] - }, - "model.visual.blocks.6.mlp.linear_fc1.bias": { - "dtype": "BF16", - "shape": [ - 4096 - ] - }, - "model.visual.blocks.6.mlp.linear_fc1.weight": { - "dtype": "BF16", - "shape": [ - 4096, - 1024 - ] - }, - "model.visual.blocks.6.mlp.linear_fc2.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.6.mlp.linear_fc2.weight": { - "dtype": "BF16", - "shape": [ - 1024, - 4096 - ] - }, - "model.visual.blocks.6.norm1.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.6.norm1.weight": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.6.norm2.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.6.norm2.weight": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.7.attn.proj.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.7.attn.proj.weight": { - "dtype": "BF16", - "shape": [ - 1024, - 1024 - ] - }, - "model.visual.blocks.7.attn.qkv.bias": { - "dtype": "BF16", - "shape": [ - 3072 - ] - }, - "model.visual.blocks.7.attn.qkv.weight": { - "dtype": "BF16", - "shape": [ - 3072, - 1024 - ] - }, - "model.visual.blocks.7.mlp.linear_fc1.bias": { - "dtype": "BF16", - "shape": [ - 4096 - ] - }, - "model.visual.blocks.7.mlp.linear_fc1.weight": { - "dtype": "BF16", - "shape": [ - 4096, - 1024 - ] - }, - "model.visual.blocks.7.mlp.linear_fc2.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.7.mlp.linear_fc2.weight": { - "dtype": "BF16", - "shape": [ - 1024, - 4096 - ] - }, - "model.visual.blocks.7.norm1.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.7.norm1.weight": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.7.norm2.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.7.norm2.weight": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.8.attn.proj.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.8.attn.proj.weight": { - "dtype": "BF16", - "shape": [ - 1024, - 1024 - ] - }, - "model.visual.blocks.8.attn.qkv.bias": { - "dtype": "BF16", - "shape": [ - 3072 - ] - }, - "model.visual.blocks.8.attn.qkv.weight": { - "dtype": "BF16", - "shape": [ - 3072, - 1024 - ] - }, - "model.visual.blocks.8.mlp.linear_fc1.bias": { - "dtype": "BF16", - "shape": [ - 4096 - ] - }, - "model.visual.blocks.8.mlp.linear_fc1.weight": { - "dtype": "BF16", - "shape": [ - 4096, - 1024 - ] - }, - "model.visual.blocks.8.mlp.linear_fc2.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.8.mlp.linear_fc2.weight": { - "dtype": "BF16", - "shape": [ - 1024, - 4096 - ] - }, - "model.visual.blocks.8.norm1.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.8.norm1.weight": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.8.norm2.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.8.norm2.weight": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.9.attn.proj.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.9.attn.proj.weight": { - "dtype": "BF16", - "shape": [ - 1024, - 1024 - ] - }, - "model.visual.blocks.9.attn.qkv.bias": { - "dtype": "BF16", - "shape": [ - 3072 - ] - }, - "model.visual.blocks.9.attn.qkv.weight": { - "dtype": "BF16", - "shape": [ - 3072, - 1024 - ] - }, - "model.visual.blocks.9.mlp.linear_fc1.bias": { - "dtype": "BF16", - "shape": [ - 4096 - ] - }, - "model.visual.blocks.9.mlp.linear_fc1.weight": { - "dtype": "BF16", - "shape": [ - 4096, - 1024 - ] - }, - "model.visual.blocks.9.mlp.linear_fc2.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.9.mlp.linear_fc2.weight": { - "dtype": "BF16", - "shape": [ - 1024, - 4096 - ] - }, - "model.visual.blocks.9.norm1.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.9.norm1.weight": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.9.norm2.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.blocks.9.norm2.weight": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.merger.linear_fc1.bias": { - "dtype": "BF16", - "shape": [ - 4096 - ] - }, - "model.visual.merger.linear_fc1.weight": { - "dtype": "BF16", - "shape": [ - 4096, - 4096 - ] - }, - "model.visual.merger.linear_fc2.bias": { - "dtype": "BF16", - "shape": [ - 2560 - ] - }, - "model.visual.merger.linear_fc2.weight": { - "dtype": "BF16", - "shape": [ - 2560, - 4096 - ] - }, - "model.visual.merger.norm.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.merger.norm.weight": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.patch_embed.proj.bias": { - "dtype": "BF16", - "shape": [ - 1024 - ] - }, - "model.visual.patch_embed.proj.weight": { - "dtype": "BF16", - "shape": [ - 1024, - 3, - 2, - 16, - 16 - ] - }, - "model.visual.pos_embed.weight": { - "dtype": "BF16", - "shape": [ - 2304, - 1024 - ] - } -} diff --git a/tests/cua_s1/fixtures/vision/config.json b/tests/cua_s1/fixtures/vision/config.json deleted file mode 100644 index cca4041..0000000 --- a/tests/cua_s1/fixtures/vision/config.json +++ /dev/null @@ -1,104 +0,0 @@ -{ - "architectures": [ - "Qwen3_5ForConditionalGeneration" - ], - "image_token_id": 248056, - "model_type": "qwen3_5", - "text_config": { - "attention_bias": false, - "attention_dropout": 0.0, - "attn_output_gate": true, - "dtype": "bfloat16", - "eos_token_id": 248044, - "full_attention_interval": 4, - "head_dim": 256, - "hidden_act": "silu", - "hidden_size": 2560, - "initializer_range": 0.02, - "intermediate_size": 9216, - "layer_types": [ - "linear_attention", - "linear_attention", - "linear_attention", - "full_attention", - "linear_attention", - "linear_attention", - "linear_attention", - "full_attention", - "linear_attention", - "linear_attention", - "linear_attention", - "full_attention", - "linear_attention", - "linear_attention", - "linear_attention", - "full_attention", - "linear_attention", - "linear_attention", - "linear_attention", - "full_attention", - "linear_attention", - "linear_attention", - "linear_attention", - "full_attention", - "linear_attention", - "linear_attention", - "linear_attention", - "full_attention", - "linear_attention", - "linear_attention", - "linear_attention", - "full_attention" - ], - "linear_conv_kernel_dim": 4, - "linear_key_head_dim": 128, - "linear_num_key_heads": 16, - "linear_num_value_heads": 32, - "linear_value_head_dim": 128, - "max_position_embeddings": 262144, - "mlp_only_layers": [], - "model_type": "qwen3_5_text", - "mtp_num_hidden_layers": 1, - "mtp_use_dedicated_embeddings": false, - "num_attention_heads": 16, - "num_hidden_layers": 32, - "num_key_value_heads": 4, - "rms_norm_eps": 1e-06, - "tie_word_embeddings": true, - "use_cache": true, - "vocab_size": 248320, - "mamba_ssm_dtype": "float32", - "rope_parameters": { - "mrope_interleaved": true, - "mrope_section": [ - 11, - 11, - 10 - ], - "rope_type": "default", - "rope_theta": 10000000, - "partial_rotary_factor": 0.25 - } - }, - "tie_word_embeddings": true, - "transformers_version": "4.57.0.dev0", - "video_token_id": 248057, - "vision_config": { - "deepstack_visual_indexes": [], - "depth": 24, - "hidden_act": "gelu_pytorch_tanh", - "hidden_size": 1024, - "in_channels": 3, - "initializer_range": 0.02, - "intermediate_size": 4096, - "model_type": "qwen3_5", - "num_heads": 16, - "num_position_embeddings": 2304, - "out_hidden_size": 2560, - "patch_size": 16, - "spatial_merge_size": 2, - "temporal_patch_size": 2 - }, - "vision_end_token_id": 248054, - "vision_start_token_id": 248053 -} diff --git a/tests/cua_s1/image_preprocess.rs b/tests/cua_s1/image_preprocess.rs deleted file mode 100644 index d2bf8ef..0000000 --- a/tests/cua_s1/image_preprocess.rs +++ /dev/null @@ -1,223 +0,0 @@ -use omni_cua_s1_native::image_preprocess::preprocess_rgb8; - -fn normalized(pixel: u8) -> f32 { - (f32::from(pixel) - 127.5) / 127.5 -} - -#[test] -fn tiny_rgb_is_upscaled_and_temporally_repeated() { - let image = preprocess_rgb8(1, 1, &[0, 127, 255]).unwrap(); - assert_eq!((image.resized_width, image.resized_height), (256, 256)); - assert_eq!(image.image_grid_thw, [1, 16, 16]); - assert_eq!(image.image_tokens(), 64); - assert_eq!(image.pixel_values.len(), 256 * 1536); - for patch in image.pixel_values.as_chunks::<1536>().0 { - for (channel, expected) in [0, 127, 255].into_iter().enumerate() { - assert!( - patch[channel * 512..(channel + 1) * 512] - .iter() - .all(|&value| value.to_bits() == normalized(expected).to_bits()) - ); - } - } -} - -#[test] -fn identity_resize_preserves_every_byte_and_patch_merge_order() { - let width = 256; - let height = 256; - let rgb: Vec = (0..height) - .flat_map(|y| (0..width).flat_map(move |x| [x as u8, y as u8, (x ^ y) as u8])) - .collect(); - let image = preprocess_rgb8(width, height, &rgb).unwrap(); - for block_y in 0..8 { - for block_x in 0..8 { - for merge_y in 0..2 { - for merge_x in 0..2 { - let patch = ((block_y * 8 + block_x) * 2 + merge_y) * 2 + merge_x; - for channel in 0..3 { - for temporal in 0..2 { - for py in 0..16 { - for px in 0..16 { - let x = block_x * 32 + merge_x * 16 + px; - let y = block_y * 32 + merge_y * 16 + py; - let index = patch * 1536 - + channel * 512 - + temporal * 256 - + py * 16 - + px; - assert_eq!( - image.pixel_values[index].to_bits(), - normalized(rgb[(y * width + x) * 3 + channel]).to_bits() - ); - } - } - } - } - } - } - } - } -} - -#[test] -fn smart_resize_uses_python_ties_even_rounding() { - for (side, expected) in [(272, 256), (304, 320)] { - let image = preprocess_rgb8(side, side, &vec![128; side * side * 3]).unwrap(); - assert_eq!( - (image.resized_width, image.resized_height), - (expected, expected) - ); - assert!( - image - .pixel_values - .iter() - .all(|&value| value == normalized(128)) - ); - } -} - -#[test] -fn rejects_invalid_geometry_before_buffer_length_validation() { - for (width, height) in [(0, 1), (1, 0), (0, 0)] { - assert_eq!( - preprocess_rgb8(width, height, &[]).unwrap_err().to_string(), - "image dimensions must be nonzero", - "{width}x{height}" - ); - } - for (width, height) in [(usize::MAX, 1), (1, usize::MAX), (usize::MAX, usize::MAX)] { - assert_eq!( - preprocess_rgb8(width, height, &[]).unwrap_err().to_string(), - "image sides must not exceed 2048", - "{width}x{height}" - ); - } - for (width, height, expected_error) in [ - (2049, 32, "image sides must not exceed 2048"), - (32, 2049, "image sides must not exceed 2048"), - (1025, 1024, "image area must not exceed 1048576 pixels"), - (201, 1, "image aspect ratio must not exceed 200"), - (1, 201, "image aspect ratio must not exceed 200"), - ] { - // A valid byte length ensures the geometry check itself rejects this - // image, rather than accidentally passing due to a truncated buffer. - let rgb = vec![0; width * height * 3]; - assert_eq!( - preprocess_rgb8(width, height, &rgb) - .unwrap_err() - .to_string(), - expected_error, - "{width}x{height}" - ); - } -} - -#[test] -fn rejects_incorrect_rgb_buffer_lengths() { - for rgb in [&[][..], &[1, 2][..], &[1, 2, 3, 4][..]] { - assert_eq!( - preprocess_rgb8(1, 1, rgb).unwrap_err().to_string(), - format!("RGB buffer length must be 3, got {}", rgb.len()) - ); - } -} - -#[test] -fn accepts_input_limits_inclusively() { - for (width, height) in [(200, 1), (1, 200), (2048, 512), (512, 2048)] { - let image = preprocess_rgb8(width, height, &vec![255; width * height * 3]).unwrap(); - assert_eq!(image.resized_width % 32, 0); - assert_eq!(image.resized_height % 32, 0); - assert!(image.pixel_values.iter().all(|&value| value == 1.0)); - } -} - -#[test] -fn matches_pinned_processor_full_output_hashes() { - use sha2::{Digest, Sha256}; - let manifest: serde_json::Value = - serde_json::from_str(include_str!("fixtures/image_preprocess/manifest.json")).unwrap(); - assert_eq!( - format!( - "{:x}", - Sha256::digest(include_bytes!( - "fixtures/image_preprocess/preprocessor_config.json" - )) - ), - manifest["config_sha256"].as_str().unwrap(), - "processor configuration must match the reference manifest" - ); - for case in manifest["cases"].as_array().unwrap() { - let name = case["name"].as_str().unwrap(); - let width = case["width"].as_u64().unwrap() as usize; - let height = case["height"].as_u64().unwrap() as usize; - let pattern = case["pattern"].as_str().unwrap(); - let mut state = case["seed"].as_u64().unwrap() as u32; - let mut rgb = Vec::with_capacity(width * height * 3); - for y in 0..height { - for x in 0..width { - for channel in 0..3 { - rgb.push(match pattern { - "noise" => { - state ^= state << 13; - state ^= state >> 17; - state ^= state << 5; - state as u8 - } - "constant" => [0, 128, 255][channel], - "checker" => { - if (x + y + channel) % 2 == 0 { - 0 - } else { - 255 - } - } - "ramp" => ((x * 13 + y * 7 + channel * 83) % 256) as u8, - _ => panic!("unknown fixture pattern {pattern}"), - }); - } - } - } - assert_eq!( - format!("{:x}", Sha256::digest(&rgb)), - case["input_sha256"].as_str().unwrap(), - "{name} input" - ); - let image = preprocess_rgb8(width, height, &rgb).unwrap(); - assert_eq!( - serde_json::json!(image.image_grid_thw), - case["grid"], - "{name} grid" - ); - assert_eq!( - serde_json::json!([image.pixel_values.len() / 1536, 1536]), - case["shape"], - "{name} shape" - ); - assert_eq!( - image.image_tokens(), - case["image_tokens"].as_u64().unwrap() as usize, - "{name} tokens" - ); - assert_eq!( - image.resized_width, - case["resized_width"].as_u64().unwrap() as usize, - "{name} width" - ); - assert_eq!( - image.resized_height, - case["resized_height"].as_u64().unwrap() as usize, - "{name} height" - ); - let mut hash = Sha256::new(); - for value in image.pixel_values { - hash.update(value.to_le_bytes()); - } - assert_eq!( - format!("{:x}", hash.finalize()), - case["output_sha256"].as_str().unwrap(), - "{name} output" - ); - } -} diff --git a/tests/cua_s1/test_native_vision_cuda.py b/tests/cua_s1/test_native_vision_cuda.py deleted file mode 100644 index 526291e..0000000 --- a/tests/cua_s1/test_native_vision_cuda.py +++ /dev/null @@ -1,114 +0,0 @@ -"""GPU regression tests: python test_native_vision_cuda.py /path/libqwen3_5_cuda.so.""" - - -def main(): - import ctypes as C - import sys - - import torch - import torch.nn.functional as F - - lib = C.CDLL(sys.argv[1]) - P, I, Z = C.c_void_p, C.c_int, C.c_size_t - signatures = { - "cs1_vision_linear": [P, P, P, P, P, I, I, I, P], - "cs1_gemm_f32": [P, P, P, P, I, I, I, P], - "cs1_vision_bias": [P, P, Z, I, P], - "cs1_vision_norm": [P, P, P, P, I, I, P], - "cs1_vision_rope": [P, P, P, P, P, I, P], - "cs1_vision_attention": [P, P, P, P, I, P], - "cs1_vision_gelu": [P, Z, I, P], - "cs1_vision_lora_add": [P, P, Z, C.c_float, P], - "cs1_gemm_create": [Z], - } - for name, args in signatures.items(): - fn = getattr(lib, name) - fn.argtypes = args - fn.restype = P if name == "cs1_gemm_create" else I - - def run(name, *args): - rc = getattr(lib, name)( - *[a.data_ptr() if isinstance(a, torch.Tensor) else a for a in args] - ) - assert rc == 0, (name, rc) - - def close(got, want, atol=0.02, rtol=0.02): - torch.cuda.synchronize() - torch.testing.assert_close(got, want, atol=atol, rtol=rtol) - - torch.manual_seed(93) - s = torch.cuda.current_stream().cuda_stream - g = lib.cs1_gemm_create(32 << 20) - assert g - x = torch.randn(68, 1024, device="cuda", dtype=torch.bfloat16) - w = torch.randn(1024, 1024, device="cuda", dtype=torch.bfloat16) / 32 - b = torch.randn(1024, device="cuda", dtype=torch.bfloat16) - y = torch.empty_like(x) - run("cs1_vision_linear", g, x, w, b, y, 68, 1024, 1024, s) - close(y, F.linear(x, w, b), atol=0.015) - run("cs1_vision_norm", x, w[0], b, y, 68, 1024, s) - close(y, F.layer_norm(x, (1024,), w[0], b, 1e-6), atol=0.008) - for exact in (0, 1): - y.copy_(x) - run("cs1_vision_gelu", y, y.numel(), exact, s) - close( - y, - F.gelu(x, approximate="none" if exact else "tanh"), - atol=0.0001, - rtol=0.001, - ) - a = torch.randn(16, 1024, device="cuda") / 32 - xf = x.float() - r = torch.empty(68, 16, device="cuda") - run("cs1_gemm_f32", g, xf, a, r, 68, 16, 1024, s) - close(r, F.linear(xf, a), atol=2e-6, rtol=2e-5) - lora = torch.randn_like(xf) * 0.01 - y.copy_(x) - run("cs1_vision_lora_add", y, lora, y.numel(), 2.0, s) - close(y, (x.float() + 2 * lora).bfloat16(), atol=0, rtol=0) - for n in (4, 68, 256): - qkv = torch.randn(n, 3, 16, 64, device="cuda", dtype=torch.bfloat16) - angles = torch.randn(n, 32, device="cuda") - co, si = angles.cos(), angles.sin() - q, k = ( - torch.empty(n, 16, 64, device="cuda", dtype=torch.bfloat16), - torch.empty(n, 16, 64, device="cuda", dtype=torch.bfloat16), - ) - run("cs1_vision_rope", qkv, co, si, q, k, n, s) - co2, si2 = co.repeat(1, 2)[:, None], si.repeat(1, 2)[:, None] - - def rope(z, co2=co2, si2=si2): - z = z.float() - return ( - z * co2 + torch.cat((-z[..., 32:], z[..., :32]), -1) * si2 - ).bfloat16() - - close(q, rope(qkv[:, 0]), atol=0, rtol=0) - close(k, rope(qkv[:, 1]), atol=0, rtol=0) - out = torch.empty_like(q) - run("cs1_vision_attention", q, k, qkv[:, 2].data_ptr(), out, n, s) - want = F.scaled_dot_product_attention( - q.transpose(0, 1), k.transpose(0, 1), qkv[:, 2].transpose(0, 1) - ).transpose(0, 1) - close(out, want, atol=0.008, rtol=0.015) - # Conv3d rounds the convolution before its separately applied BF16 bias. - xp = torch.randn(4, 1536, device="cuda", dtype=torch.bfloat16) - wp = torch.randn(1024, 1536, device="cuda", dtype=torch.bfloat16) / 32 - bp = torch.randn(1024, device="cuda", dtype=torch.bfloat16) - z = torch.zeros_like(bp) - yp = torch.empty(4, 1024, device="cuda", dtype=torch.bfloat16) - run("cs1_vision_linear", g, xp, wp, z, yp, 4, 1024, 1536, s) - run("cs1_vision_bias", yp, bp, yp.numel(), 1024, s) - want = F.conv3d( - xp.reshape(4, 3, 2, 16, 16), - wp.reshape(1024, 3, 2, 16, 16), - bp, - stride=(2, 16, 16), - ).flatten(1) - close(yp, want, atol=0.008, rtol=0.001) - assert (yp.float() - want.float()).abs().mean().item() < 1e-5 - print("native vision CUDA primitives: PASS") - - -if __name__ == "__main__": - main() diff --git a/tests/cua_s1/vision_loader.rs b/tests/cua_s1/vision_loader.rs deleted file mode 100644 index 8de73e2..0000000 --- a/tests/cua_s1/vision_loader.rs +++ /dev/null @@ -1,466 +0,0 @@ -use omni_cua_s1_native::vision::VisionCheckpoint; -use safetensors::Dtype; -use serde_json::{Value, json}; -use std::{ - collections::BTreeMap, - fs::{self, File}, - io::{Seek, SeekFrom, Write}, - path::{Path, PathBuf}, - sync::atomic::{AtomicUsize, Ordering}, -}; - -const BASE: &str = "model.visual."; -const ADAPTER: &str = "base_model.model.model.visual."; -type Inventory = BTreeMap; -fn inventory(adapter: bool) -> Inventory { - serde_json::from_str(if adapter { - include_str!("fixtures/vision/adapter-tensors.json") - } else { - include_str!("fixtures/vision/base-tensors.json") - }) - .unwrap() -} -fn sparse(path: &Path, tensors: &Inventory) { - let mut offset = 0u64; - let mut header = serde_json::Map::new(); - for (name, tensor) in tensors { - let bytes = tensor["shape"] - .as_array() - .unwrap() - .iter() - .map(|n| n.as_u64().unwrap()) - .product::() - * if tensor["dtype"] == "BF16" { 2 } else { 4 }; - header.insert(name.clone(), json!({"dtype":tensor["dtype"],"shape":tensor["shape"],"data_offsets":[offset, offset+bytes]})); - offset += bytes; - } - let header = serde_json::to_vec(&header).unwrap(); - let mut file = File::create(path).unwrap(); - file.write_all(&(header.len() as u64).to_le_bytes()) - .unwrap(); - file.write_all(&header).unwrap(); - file.set_len(8 + header.len() as u64 + offset).unwrap(); - for tensor in serde_json::from_slice::(&header) - .unwrap() - .as_object() - .unwrap() - .values() - { - file.seek(SeekFrom::Start( - 8 + header.len() as u64 + tensor["data_offsets"][0].as_u64().unwrap(), - )) - .unwrap(); - file.write_all(if tensor["dtype"] == "BF16" { - &[0x80, 0x3f] - } else { - &[0, 0, 0x80, 0x3f] - }) - .unwrap(); - } -} -struct Fixture { - root: PathBuf, - base: PathBuf, - adapter: PathBuf, -} -impl Fixture { - fn new() -> Self { - static NEXT: AtomicUsize = AtomicUsize::new(0); - let root = std::env::temp_dir().join(format!( - "cua-vision-{}-{}", - std::process::id(), - NEXT.fetch_add(1, Ordering::Relaxed) - )); - let base = root.join("base"); - let adapter = root.join("adapter"); - fs::create_dir_all(&base).unwrap(); - fs::create_dir_all(&adapter).unwrap(); - fs::write( - base.join("config.json"), - include_str!("fixtures/vision/config.json"), - ) - .unwrap(); - fs::write( - adapter.join("adapter_config.json"), - include_str!("fixtures/vision/adapter_config.json"), - ) - .unwrap(); - sparse(&base.join("model.safetensors"), &inventory(false)); - sparse(&adapter.join("adapter_model.safetensors"), &inventory(true)); - Self { - root, - base, - adapter, - } - } - fn load(&self) -> anyhow::Result { - VisionCheckpoint::load(&self.base, &self.adapter) - } - fn reject(&self, expected: &str) { - let error = match self.load() { - Ok(_) => panic!("accepted invalid checkpoint: {expected}"), - Err(error) => format!("{error:#}"), - }; - assert!( - error.contains(expected), - "expected {expected:?}, got {error}" - ); - } - fn config(&self, adapter: bool, key: &str, value: Value) { - let path = if adapter { - self.adapter.join("adapter_config.json") - } else { - self.base.join("config.json") - }; - let mut config: Value = serde_json::from_str(if adapter { - include_str!("fixtures/vision/adapter_config.json") - } else { - include_str!("fixtures/vision/config.json") - }) - .unwrap(); - if adapter { - config[key] = value; - } else if key == "text_hidden" { - config["text_config"]["hidden_size"] = value; - } else { - config["vision_config"][key] = value; - } - fs::write(path, serde_json::to_vec(&config).unwrap()).unwrap(); - } - fn shards(&self) -> BTreeMap { - let (a, b): (Inventory, Inventory) = inventory(false) - .into_iter() - .partition(|(n, _)| n.contains("blocks.")); - sparse(&self.base.join("a.safetensors"), &a); - sparse(&self.base.join("b.safetensors"), &b); - let mut map: BTreeMap = a - .keys() - .map(|n| (n.clone(), "a.safetensors".into())) - .chain(b.keys().map(|n| (n.clone(), "b.safetensors".into()))) - .collect(); - map.insert( - "model.language_model.weight".into(), - "absent-language.safetensors".into(), - ); - self.index(&map); - fs::remove_file(self.base.join("model.safetensors")).unwrap(); - map - } - fn index(&self, map: &BTreeMap) { - fs::write( - self.base.join("model.safetensors.index.json"), - serde_json::to_vec(&json!({"weight_map":map})).unwrap(), - ) - .unwrap(); - } -} -impl Drop for Fixture { - fn drop(&mut self) { - let _ = fs::remove_dir_all(&self.root); - } -} -fn check(checkpoint: &VisionCheckpoint) { - assert_eq!(checkpoint.config().depth, 24); - assert_eq!(checkpoint.config().hidden_size, 1024); - assert_eq!(checkpoint.adapter().rank, 16); - assert_eq!(checkpoint.adapter().alpha, 32); - assert_eq!(checkpoint.adapter().scale(), 2.0); - for adapter in [false, true] { - let expected = inventory(adapter); - let names: Vec<_> = if adapter { - checkpoint.adapter_names().collect() - } else { - checkpoint.base_names().collect() - }; - assert_eq!(names.len(), if adapter { 100 } else { 297 }); - for (name, spec) in expected { - let tensor = if adapter { - checkpoint.adapter_tensor(&name) - } else { - checkpoint.base_tensor(&name) - } - .unwrap(); - assert_eq!( - tensor.shape(), - serde_json::from_value::>(spec["shape"].clone()).unwrap() - ); - assert_eq!( - tensor.dtype(), - if adapter { Dtype::F32 } else { Dtype::BF16 } - ); - assert_eq!( - &tensor.data()[..if adapter { 4 } else { 2 }], - if adapter { - &[0, 0, 0x80, 0x3f][..] - } else { - &[0x80, 0x3f][..] - } - ); - } - } - assert!(checkpoint.base_tensor("missing").is_err()); -} -#[test] -fn valid_single_matches_independent_inventory() { - let f = Fixture::new(); - check(&f.load().unwrap()); -} -#[test] -fn valid_shards_skip_language_only_files() { - let f = Fixture::new(); - f.shards(); - check(&f.load().unwrap()); -} -#[test] -fn ignores_language_tensors() { - let f = Fixture::new(); - for adapter in [false, true] { - let mut inv = inventory(adapter); - inv.insert( - "unrelated.language.weight".into(), - json!({"shape":[1],"dtype":"F32"}), - ); - sparse( - &if adapter { - f.adapter.join("adapter_model.safetensors") - } else { - f.base.join("model.safetensors") - }, - &inv, - ); - } - check(&f.load().unwrap()); -} -#[test] -fn rejects_missing_wrong_and_extra_tensors() { - let f = Fixture::new(); - for adapter in [false, true] { - let path = if adapter { - f.adapter.join("adapter_model.safetensors") - } else { - f.base.join("model.safetensors") - }; - let original = inventory(adapter); - let name = original.keys().next().unwrap().clone(); - let mut inv = original.clone(); - inv.remove(&name); - sparse(&path, &inv); - f.reject("missing"); - let mut inv = original.clone(); - inv.get_mut(&name).unwrap()["shape"] = json!([1]); - sparse(&path, &inv); - f.reject("shape"); - let mut inv = original.clone(); - inv.get_mut(&name).unwrap()["dtype"] = json!(if adapter { "BF16" } else { "F32" }); - sparse(&path, &inv); - f.reject("dtype"); - let mut inv = original.clone(); - inv.insert( - format!("{}unexpected.weight", if adapter { ADAPTER } else { BASE }), - json!({"shape":[1],"dtype":"F32"}), - ); - sparse(&path, &inv); - f.reject("unexpected"); - sparse(&path, &original); - } -} -#[test] -fn rejects_truncated_and_malformed_safetensors() { - let f = Fixture::new(); - for adapter in [false, true] { - let path = if adapter { - f.adapter.join("adapter_model.safetensors") - } else { - f.base.join("model.safetensors") - }; - let len = fs::metadata(&path).unwrap().len(); - File::options() - .write(true) - .open(&path) - .unwrap() - .set_len(len - 1) - .unwrap(); - f.reject("safetensors"); - fs::write(&path, b"invalid").unwrap(); - f.reject("safetensors"); - sparse(&path, &inventory(adapter)); - } -} -#[test] -fn rejects_index_missing_extra_misrouting_and_duplicates() { - let f = Fixture::new(); - let original = f.shards(); - let name = inventory(false).keys().next().unwrap().clone(); - let mut map = original.clone(); - map.remove(&name); - f.index(&map); - f.reject("missing"); - let mut map = original.clone(); - map.insert(format!("{BASE}unexpected"), "a.safetensors".into()); - f.index(&map); - f.reject("unexpected"); - let mut map = original.clone(); - map.insert(name.clone(), "b.safetensors".into()); - f.index(&map); - f.reject("index"); - f.index(&original); - let mut b: Inventory = inventory(false) - .into_iter() - .filter(|(n, _)| !n.contains("blocks.")) - .collect(); - b.insert(name.clone(), inventory(false)[&name].clone()); - sparse(&f.base.join("b.safetensors"), &b); - f.reject("index"); - fs::write( - f.base.join("model.safetensors.index.json"), - format!("{{\"weight_map\":{{\"{name}\":\"a.safetensors\",\"{name}\":\"b.safetensors\"}}}}"), - ) - .unwrap(); - f.reject("duplicate"); -} -#[test] -fn rejects_index_paths_outside_directory() { - let f = Fixture::new(); - let original = f.shards(); - let name = inventory(false).keys().next().unwrap().clone(); - for path in ["../outside.safetensors", "/tmp/outside.safetensors"] { - let mut map = original.clone(); - map.insert(name.clone(), path.into()); - f.index(&map); - f.reject("checkpoint directory"); - } -} -#[cfg(unix)] -#[test] -fn rejects_symlink_escape() { - use std::os::unix::fs::symlink; - let f = Fixture::new(); - fs::rename( - f.base.join("model.safetensors"), - f.root.join("outside.safetensors"), - ) - .unwrap(); - symlink( - f.root.join("outside.safetensors"), - f.base.join("model.safetensors"), - ) - .unwrap(); - f.reject("checkpoint directory"); -} -#[test] -fn validates_full_vision_configuration() { - let f = Fixture::new(); - for (key, value) in [ - ("depth", json!(23)), - ("hidden_size", json!(512)), - ("intermediate_size", json!(2048)), - ("num_heads", json!(8)), - ("num_position_embeddings", json!(1024)), - ("out_hidden_size", json!(2048)), - ("in_channels", json!(1)), - ("patch_size", json!(14)), - ("temporal_patch_size", json!(1)), - ("spatial_merge_size", json!(1)), - ("hidden_act", json!("gelu")), - ("deepstack_visual_indexes", json!([1])), - ("text_hidden", json!(2048)), - ] { - f.config(false, key, value); - f.reject("config"); - } -} -#[test] -fn rejects_adapter_math_and_target_changes() { - let f = Fixture::new(); - for (key, value) in [ - ("r", json!(8)), - ("lora_alpha", json!(16)), - ("bias", json!("all")), - ("lora_bias", json!(true)), - ("use_dora", json!(true)), - ("use_rslora", json!(true)), - ("use_qalora", json!(true)), - ("fan_in_fan_out", json!(true)), - ("rank_pattern", json!({"linear_fc1":8})), - ("alpha_pattern", json!({"linear_fc1":16})), - ("exclude_modules", json!(["linear_fc1"])), - ("target_modules", json!(["up_proj", "down_proj"])), - ("modules_to_save", json!(["merger"])), - ("layers_to_transform", json!([1])), - ("target_parameters", json!(["weight"])), - ("layer_replication", json!([[0, 1]])), - ("alora_invocation_tokens", json!([1])), - ("peft_type", json!("IA3")), - ] { - f.config(true, key, value); - f.reject("adapter config"); - } -} -#[test] -fn requires_multimodal_adapter() { - let f = Fixture::new(); - fs::remove_file(f.adapter.join("adapter_model.safetensors")).unwrap(); - f.reject("adapter_model.safetensors"); - sparse( - &f.adapter.join("adapter_model.safetensors"), - &BTreeMap::from([( - "base_model.model.model.language_model.weight".into(), - json!({"shape":[1],"dtype":"F32"}), - )]), - ); - f.reject("missing"); -} - -#[test] -fn rejects_duplicate_visual_keys_in_safetensors_header() { - use std::io::Read; - let f = Fixture::new(); - let path = f.base.join("model.safetensors"); - let mut file = File::open(&path).unwrap(); - let total = file.metadata().unwrap().len(); - let mut size = [0; 8]; - file.read_exact(&mut size).unwrap(); - let old_len = u64::from_le_bytes(size); - let mut header = vec![0; old_len as usize]; - file.read_exact(&mut header).unwrap(); - drop(file); - let entries: BTreeMap = serde_json::from_slice(&header).unwrap(); - let (name, info) = entries.first_key_value().unwrap(); - let mut duplicate = String::from_utf8(header).unwrap(); - duplicate.pop(); - duplicate.push_str(&format!( - ",{}:{info}}}", - serde_json::to_string(name).unwrap() - )); - let mut file = File::create(&path).unwrap(); - file.write_all(&(duplicate.len() as u64).to_le_bytes()) - .unwrap(); - file.write_all(duplicate.as_bytes()).unwrap(); - file.set_len(total - old_len + duplicate.len() as u64) - .unwrap(); - drop(file); - f.reject("duplicate"); -} - -#[test] -fn rejects_overflowing_payload_offsets_without_panicking() { - let f = Fixture::new(); - let elements = usize::MAX / 8; - let mut header = serde_json::Map::new(); - for i in 0..8 { - header.insert( - format!("unrelated.language.{i}.weight"), - json!({ - "dtype": "U8", "shape": [elements], - "data_offsets": [i * elements, (i + 1) * elements] - }), - ); - } - let header = serde_json::to_vec(&header).unwrap(); - let mut file = File::create(f.base.join("model.safetensors")).unwrap(); - file.write_all(&(header.len() as u64).to_le_bytes()) - .unwrap(); - file.write_all(&header).unwrap(); - drop(file); - f.reject("safetensors"); -}