diff --git a/mkdocs.yml b/mkdocs.yml index ef244fd2..0fb2b855 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -65,6 +65,7 @@ nav: - Laya on Apple Silicon: recipe/laya/apple-silicon.md - Cua-S1 text worker: recipe/cua_s1/text.md - Cua-S1 native text worker: recipe/cua_s1/native.md + - Cua-S1 vision CUDA Graph: recipe/cua_s1/native_vision_graph.md - Open-Jev native text worker: recipe/open_jev/native.md - Open-Jev-27B H200 validation: recipe/open_jev/validation.md - Open-Jev-9B validation: recipe/open_jev/validation-9b.md diff --git a/recipe/cua_s1/native.md b/recipe/cua_s1/native.md index 83322b25..35336704 100644 --- a/recipe/cua_s1/native.md +++ b/recipe/cua_s1/native.md @@ -24,9 +24,8 @@ 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 ordered prompt-length vector warms the GEMM plans and captures the forward pass; +later requests replay it with freshly uploaded token ids. The text worker keeps a 64-entry FIFO cache. 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 6). If capture fails, the worker returns the completed eager result and disables @@ -41,3 +40,7 @@ cargo test -p omni-cua-s1-native CUA_S1_CUDA_LIB=$PWD/target/release/libqwen3_5_cuda.so \ cargo test --release -p omni-qwen3-5-native --test kernels -- --ignored ``` + +For the screenshot worker, [the vision Graph recipe](native_vision_graph.md) +documents `CUA_S1_VISION_GRAPH`, retained single-grid scratch, memory/lifetime +limits and the real CUDA capture-failure wrapper build and test commands. diff --git a/recipe/cua_s1/native_vision_graph.md b/recipe/cua_s1/native_vision_graph.md new file mode 100644 index 00000000..e6fcc71c --- /dev/null +++ b/recipe/cua_s1/native_vision_graph.md @@ -0,0 +1,96 @@ +# Native Cua-S1 vision CUDA Graph + +The `omni-cua-s1-vision` screenshot worker can capture the complete vision +encoder, including all 24 blocks, unmerged FP32 LoRA branches and the merger. +Set `CUA_S1_VISION_GRAPH=1` before loading the worker to enable it. Leave it +unset or set it to `0` for eager execution. This switch is independent of +`CUA_S1_GRAPH`, which controls language execution; this change does not enable +multimodal language Graph replay. + +Both modes retain one scratch allocation for the latest exact `[T,H,W]` patch +grid. Equal patch counts with different heights or widths replace the scratch +and Graph. The worker synchronizes, destroys the old executable, and releases +its buffers before allocating the new shape. Pixel values are uploaded on every +request, including replay. Geometry tables remain valid only for that exact grid. +This bounds activation storage to one grid (at most 4608 patches); it does not +bound loaded weights or the GEMM plan cache. + +The first ordinary forward for a grid finishes eagerly, then captures operators +for subsequent requests and returns the completed eager result. Capture time is +part of that first request. Uploads, downloads and CPU processing stay outside +the Graph. Capture failure always logs to stderr and disables vision Graph for +that loaded model; the completed eager result is preserved. Replay failures +propagate as inference errors. `forward_with_trace` runs eagerly with all stage +downloads and leaves a matching Graph available for later ordinary forwards. + +For optional capture/replay messages, set `CUA_S1_GRAPH_TRACE=1`. Failure messages +are unconditional. Requests execute serially through the existing worker runtime; +this feature does not add batching or concurrent forwards. + +## Build and launch + +Use the pinned base and multimodal adapter revisions listed in the +[model contract](../../src/models/cua_s1/README.md#pinned-revisions), and the +verified language export produced by `export_multimodal_language.py`. Keep all +checkpoint files immutable during execution. Build the worker and the current +CUDA library together; this change adds no CUDA ABI symbols. + +```sh +src/backends/cuda/qwen3_5/build.sh target/release 89 +cargo build --release --locked -p omni-cua-s1-native + +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 \ +CUA_S1_VISION_GRAPH=1 target/release/omni-cua-s1-vision +``` + +Set `CUA_S1_PORT` to choose the listening port (default 8000). Only the existing +single-image native CUDA path is covered. Python workers and Metal are unchanged. + +## Regressions and native measurements + +Export the same `CUA_S1_BASE`, `CUA_S1_VISION_ADAPTER`, `CUA_S1_MODEL`, and +`CUA_S1_CUDA_LIB` variables for these commands: + +```sh +cargo test --release --locked -p omni-cua-s1-native --lib vision_graph_ \ + -- --ignored --test-threads=1 --nocapture +CUA_S1_VISION_GRAPH=0 cargo run --release --locked -p omni-cua-s1-native \ + --example vision_graph_bench -- /tmp/vision-eager.json +CUA_S1_VISION_GRAPH=1 cargo run --release --locked -p omni-cua-s1-native \ + --example vision_graph_bench -- /tmp/vision-graph.json +``` + +The capture-failure test uses a separate process and a test-only wrapper library. +It delegates all math to the real CUDA library, ends the real capture, destroys +the executable, then reports an instantiation failure. No production backend is +modified. Set an absolute path for the real library: + +```sh +export CUA_S1_REAL_CUDA_LIB="$CUA_S1_CUDA_LIB" +g++ -std=c++17 -shared -fPIC -Wall -Wextra tests/cua_s1/fail_graph_capture.cpp \ + -Wl,--no-as-needed "$CUA_S1_REAL_CUDA_LIB" -ldl \ + -o /tmp/cua-fail-capture.so +CUA_S1_CUDA_LIB=/tmp/cua-fail-capture.so CUA_S1_TEST_CAPTURE_FAILURE=1 \ +CUA_S1_GRAPH_TRACE=0 cargo test --release --locked -p omni-cua-s1-native \ + --lib vision_capture_failure_ -- --ignored --test-threads=1 --nocapture +``` + +The root GPU tests cover changed pixels, full-grid replacement, retirement, +trace bypass, invalid input and capture-failure fallback. Ordinary CI compiles +these tests but skips device execution. Run them explicitly on a reserved GPU. + +The benchmark prepares ten deterministic RGB screenshots for each of three +sizes (256x256, 512x256, 512x512), with two choice questions each. It reports model +loading and first vision/capture time separately, excludes three warmup decisions +per shape, and retains every timed vision call and complete native response. +`vision_ms` includes validation, BF16 conversion, upload, forward, synchronized +download and output conversion. `native_decision_ms` includes vision, both eager +language forwards, letter projections and response construction; it excludes CPU +preparation, HTTP and model loading. Feature SHA256 and complete responses must +match between modes. These serial synthetic measurements do not establish +labelled task quality, production throughput, HTTP latency or a universal speedup. +Preserve two repetitions with reversed mode order, raw outputs, binary/source +hashes and failures in external evidence rather than committing generated runs. diff --git a/src/models/cua_s1/native/Cargo.toml b/src/models/cua_s1/native/Cargo.toml index 80339951..b06b6dc4 100644 --- a/src/models/cua_s1/native/Cargo.toml +++ b/src/models/cua_s1/native/Cargo.toml @@ -30,3 +30,7 @@ tokio = { version = "1.49.0", features = ["macros", "net", "rt-multi-thread", "s [[bin]] name = "omni-cua-s1-vision" path = "src/vision_main.rs" + +[[example]] +name = "vision_graph_bench" +path = "../../../../tests/cua_s1/vision_graph_bench.rs" diff --git a/src/models/cua_s1/native/src/vision/model.rs b/src/models/cua_s1/native/src/vision/model.rs index d89833ef..e2ba8fe3 100644 --- a/src/models/cua_s1/native/src/vision/model.rs +++ b/src/models/cua_s1/native/src/vision/model.rs @@ -28,10 +28,13 @@ impl Drop for Gemm { /// 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 { + // Field order retires graph/scratch before weights, GEMM workspace and stream. + scratch: Option, base: BTreeMap, lora: BTreeMap, - stream: OwnedStream, gemm: Gemm, + stream: OwnedStream, + graph_enabled: bool, } impl VisionModel { pub fn load(base: impl AsRef, adapter: impl AsRef, library: &Path) -> Result { @@ -42,6 +45,8 @@ impl VisionModel { let gemm = Gemm(unsafe { (api().cs1_gemm_create)(32 << 20) }); ensure!(!gemm.0.is_null(), "cannot create vision cuBLAS handle"); let mut model = Self { + scratch: None, + graph_enabled: std::env::var("CUA_S1_VISION_GRAPH").as_deref() == Ok("1"), base: BTreeMap::new(), lora: BTreeMap::new(), stream, @@ -232,7 +237,15 @@ impl VisionModel { } fn run(&mut self, image: &ProcessedImage, mut callback: Trace<'_>) -> Result> { cuda::set_device(0)?; - let geo = Geometry::new(image.image_grid_thw)?; + let geo = if self + .scratch + .as_ref() + .is_some_and(|s| s.grid == image.image_grid_thw) + { + None + } else { + Some(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, @@ -252,42 +265,75 @@ impl VisionModel { .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)?; + if let Some(geo) = geo { + self.synchronize()?; + // Destroy the executable before releasing any address it references. + self.scratch = None; + self.scratch = Some(Scratch::new(self, image.image_grid_thw, geo)?); + } + let scratch = self.scratch.as_ref().unwrap(); + // SAFETY: pixels have the validated exact size of the resident allocation. + unsafe { + cuda::upload(scratch.pixels.at(0), &pixels, self.stream.0)?; + } + if callback.is_none() + && self.graph_enabled + && let Some(graph) = &scratch.graph + { + graph.launch(self.stream.0)?; + graph_trace("replayed", scratch.grid); + return self.read(&scratch.out, n / 4 * 2560); + } + self.enqueue(scratch, &mut callback)?; + let result = self.read(&scratch.out, n / 4 * 2560)?; + if let Some(callback) = callback.as_mut() { + callback("merger.output", &result)?; + } else if self.graph_enabled { + // The eager read completed warmup. Capture does not execute the operators; + // return the already completed result even if recording fails. + let captured = cuda::Graph::capture(self.stream.0, || self.enqueue(scratch, &mut None)); + self.finish_capture(captured); + } + Ok(result) + } + fn finish_capture(&mut self, captured: Result) { + match captured { + Ok(graph) => { + let scratch = self.scratch.as_mut().unwrap(); + scratch.graph = Some(graph); + graph_trace("captured", scratch.grid); + } + Err(error) => { + eprintln!("Vision CUDA Graph capture failed; using eager execution: {error:#}"); + self.graph_enabled = false; + if let Some(scratch) = self.scratch.as_mut() { + scratch.graph = None; + } + } + } + } + fn enqueue(&self, scratch: &Scratch, callback: &mut Trace<'_>) -> Result<()> { + let n = scratch.grid[1] * scratch.grid[2]; + let Scratch { + pixels, + indices, + weights, + co, + si, + w, + x, + norm, + qkv, + q, + k, + attn, + delta, + mlp, + out, + .. + } = scratch; + self.linear("patch_embed.proj", pixels, x, n, 1024, 1536, w)?; + self.trace(callback, "patch_embed", x, n * 1024)?; unsafe { check( (api().cs1_vision_position)( @@ -301,11 +347,11 @@ impl VisionModel { "vision learned positions", )?; } - self.trace(&mut callback, "position", &x, n * 1024)?; + self.trace(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)?; + 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)( @@ -331,64 +377,118 @@ impl VisionModel { "vision attention", )?; } - self.linear(&format!("{p}.attn.proj"), &attn, &delta, n, 1024, 1024, &w)?; + 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, - )?; + 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, - )?; + 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.trace(callback, &p, x, n * 1024)?; } - self.norm("merger.norm", &x, &norm, n)?; - self.trace(&mut callback, "merger.norm", &norm, n * 1024)?; + self.norm("merger.norm", x, norm, n)?; + self.trace(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)?; + self.linear("merger.linear_fc1", norm, mlp, n / 4, 4096, 4096, w)?; + self.trace(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) + self.linear("merger.linear_fc2", mlp, out, n / 4, 2560, 4096, w)?; + Ok(()) + } +} +impl Drop for VisionModel { + fn drop(&mut self) { + // Error paths can leave work queued; all referenced resources are still alive here. + let _ = self.synchronize(); + } +} +fn graph_trace(action: &str, grid: [usize; 3]) { + if std::env::var("CUA_S1_GRAPH_TRACE").as_deref() == Ok("1") { + eprintln!("Vision CUDA Graph {action} grid={grid:?}"); + } +} +struct Scratch { + // Executable is dropped first. One exact grid bounds resident activation memory. + graph: Option, + grid: [usize; 3], + pixels: DeviceBuffer, + indices: DeviceBuffer, + weights: DeviceBuffer, + co: DeviceBuffer, + si: DeviceBuffer, + w: Work, + x: DeviceBuffer, + norm: DeviceBuffer, + qkv: DeviceBuffer, + q: DeviceBuffer, + k: DeviceBuffer, + attn: DeviceBuffer, + delta: DeviceBuffer, + mlp: DeviceBuffer, + out: DeviceBuffer, +} +impl Scratch { + fn new(model: &VisionModel, grid: [usize; 3], geo: Geometry) -> Result { + let n = grid[1] * grid[2]; + Ok(Self { + graph: None, + grid, + pixels: DeviceBuffer::new(n * 1536 * 2)?, + indices: model.upload( + &geo.indices + .iter() + .flat_map(|v| v.to_le_bytes()) + .collect::>(), + )?, + weights: model.upload( + &geo.weights + .iter() + .flat_map(|v| v.to_le_bytes()) + .collect::>(), + )?, + co: model.upload( + &geo.cos + .iter() + .flat_map(|v| v.to_le_bytes()) + .collect::>(), + )?, + si: model.upload( + &geo.sin + .iter() + .flat_map(|v| v.to_le_bytes()) + .collect::>(), + )?, + w: Work::new(n)?, + x: DeviceBuffer::new(n * 1024 * 2)?, + norm: DeviceBuffer::new(n * 1024 * 2)?, + qkv: DeviceBuffer::new(n * 3072 * 2)?, + q: DeviceBuffer::new(n * 1024 * 2)?, + k: DeviceBuffer::new(n * 1024 * 2)?, + attn: DeviceBuffer::new(n * 1024 * 2)?, + delta: DeviceBuffer::new(n * 1024 * 2)?, + mlp: DeviceBuffer::new(n * 4096 * 2)?, + out: DeviceBuffer::new(n / 4 * 2560 * 2)?, + }) } } struct Work { @@ -405,3 +505,7 @@ impl Work { }) } } + +#[cfg(test)] +#[path = "../../../../../../tests/cua_s1/vision_graph.rs"] +mod graph_tests; diff --git a/tests/cua_s1/fail_graph_capture.cpp b/tests/cua_s1/fail_graph_capture.cpp new file mode 100644 index 00000000..5b9fac03 --- /dev/null +++ b/tests/cua_s1/fail_graph_capture.cpp @@ -0,0 +1,26 @@ +// Test-only loader wrapper. All other cs1 symbols come from the linked real library. +// Fail instantiation after real capture, destroy the executable, and leave math intact. +#include +#include +#include + +extern "C" int cs1_graph_end(void* stream, void** exec) { + const char* path = std::getenv("CUA_S1_REAL_CUDA_LIB"); + if (!path) std::abort(); + void* library = dlopen(path, RTLD_NOW | RTLD_LOCAL); + if (!library) std::abort(); + using End = int (*)(void*, void**); + using Destroy = int (*)(void*); + auto end = reinterpret_cast(dlsym(library, "cs1_graph_end")); + auto destroy = reinterpret_cast(dlsym(library, "cs1_graph_destroy")); + if (!end || !destroy || end == &cs1_graph_end) std::abort(); + int result = end(stream, exec); + if (result == 0) { + if (*exec && destroy(*exec) != 0) std::abort(); + *exec = nullptr; + std::fputs("Injected vision Graph instantiation failure\n", stderr); + result = 2; // cudaErrorMemoryAllocation; no CUDA operation was changed. + } + dlclose(library); + return result; +} diff --git a/tests/cua_s1/vision_graph.rs b/tests/cua_s1/vision_graph.rs new file mode 100644 index 00000000..ae3e737a --- /dev/null +++ b/tests/cua_s1/vision_graph.rs @@ -0,0 +1,133 @@ +//! GPU regression bodies live outside production source. Run serially with pinned weights. +use super::*; + +fn load() -> VisionModel { + let path = |name| std::path::PathBuf::from(std::env::var_os(name).expect(name)); + VisionModel::load( + path("CUA_S1_BASE"), + path("CUA_S1_VISION_ADAPTER"), + &path("CUA_S1_CUDA_LIB"), + ) + .unwrap() +} +fn image(grid: [usize; 3], seed: usize) -> ProcessedImage { + let n = grid[1] * grid[2] * 1536; + ProcessedImage { + image_grid_thw: grid, + resized_height: grid[1] * 16, + resized_width: grid[2] * 16, + pixel_values: (0..n) + .map(|i| ((i * 17 + seed * 31) % 251) as f32 / 125. - 1.) + .collect(), + } +} +fn bits(values: &[bf16]) -> Vec { + values.iter().map(|v| v.to_bits()).collect() +} + +#[test] +#[ignore = "requires CUDA and pinned CUA_S1_BASE/CUA_S1_VISION_ADAPTER/CUA_S1_CUDA_LIB"] +fn vision_graph_changed_pixels_and_trace() { + let mut model = load(); + model.graph_enabled = false; + let inputs = [ + image([1, 16, 16], 1), + image([1, 16, 16], 2), + image([1, 16, 16], 3), + ]; + let expected: Vec<_> = inputs + .iter() + .map(|i| bits(&model.forward(i).unwrap())) + .collect(); + assert_ne!( + expected[0], expected[1], + "changed pixels must change the feature oracle" + ); + model.graph_enabled = true; + assert_eq!(bits(&model.forward(&inputs[0]).unwrap()), expected[0]); + assert!( + model.scratch.as_ref().unwrap().graph.is_some(), + "missing vision capture" + ); + for (input, reference) in inputs.iter().zip(&expected).cycle().take(9) { + assert_eq!(bits(&model.forward(input).unwrap()), *reference); + } + let mut stages = Vec::new(); + let traced = model + .forward_with_trace(&inputs[1], |name, _| { + stages.push(name.to_string()); + Ok(()) + }) + .unwrap(); + assert_eq!(bits(&traced), expected[1]); + assert_eq!(stages.len(), 29); + assert!(model.scratch.as_ref().unwrap().graph.is_some()); + assert_eq!(bits(&model.forward(&inputs[2]).unwrap()), expected[2]); + let mut invalid = image([1, 16, 16], 4); + invalid.pixel_values[0] = f32::NAN; + assert!(model.forward(&invalid).is_err()); + assert!(model.scratch.as_ref().unwrap().graph.is_some()); + assert_eq!(bits(&model.forward(&inputs[0]).unwrap()), expected[0]); +} + +#[test] +#[ignore = "requires CUDA and pinned CUA_S1_BASE/CUA_S1_VISION_ADAPTER/CUA_S1_CUDA_LIB"] +fn vision_graph_equal_patch_count_geometry_and_retirement() { + let mut model = load(); + model.graph_enabled = false; + let inputs = [ + image([1, 8, 32], 7), + image([1, 16, 16], 7), + image([1, 32, 8], 7), + image([1, 16, 32], 8), + ]; + let expected: Vec<_> = inputs + .iter() + .map(|i| bits(&model.forward(i).unwrap())) + .collect(); + assert_ne!( + expected[0], expected[1], + "geometry must affect the feature oracle" + ); + model.graph_enabled = true; + for index in [0, 0, 1, 1, 2, 2, 3, 3, 0, 0] { + assert_eq!( + bits(&model.forward(&inputs[index]).unwrap()), + expected[index] + ); + let scratch = model.scratch.as_ref().unwrap(); + assert_eq!(scratch.grid, inputs[index].image_grid_thw); + assert!(scratch.graph.is_some()); + } + drop(model); // graph must retire before scratch/weights/GEMM/stream + let mut fresh = load(); + fresh.graph_enabled = true; + assert_eq!(bits(&fresh.forward(&inputs[1]).unwrap()), expected[1]); + assert_eq!(bits(&fresh.forward(&inputs[1]).unwrap()), expected[1]); +} + +#[test] +#[ignore = "requires the test-only failing CUDA wrapper; see native_vision_graph.md"] +fn vision_capture_failure_keeps_eager_result_and_disables_replay() { + assert_eq!( + std::env::var("CUA_S1_TEST_CAPTURE_FAILURE").as_deref(), + Ok("1"), + "use the test-only failing CUDA wrapper" + ); + let mut model = load(); + model.graph_enabled = false; + let input = image([1, 16, 16], 9); + let changed = image([1, 8, 32], 10); + let expected = bits(&model.forward(&input).unwrap()); + let changed_expected = bits(&model.forward(&changed).unwrap()); + model.graph_enabled = true; + // The wrapper fails graph instantiation after the real eager forward and + // capture have finished. This must return that same forward's eager result. + assert_eq!(bits(&model.forward(&input).unwrap()), expected); + assert!(!model.graph_enabled); + assert!(model.scratch.as_ref().unwrap().graph.is_none()); + assert_eq!(bits(&model.forward(&changed).unwrap()), changed_expected); + assert!(!model.graph_enabled); + assert!(model.scratch.as_ref().unwrap().graph.is_none()); + assert_eq!(bits(&model.forward(&input).unwrap()), expected); +} diff --git a/tests/cua_s1/vision_graph_bench.rs b/tests/cua_s1/vision_graph_bench.rs new file mode 100644 index 00000000..58cefacc --- /dev/null +++ b/tests/cua_s1/vision_graph_bench.rs @@ -0,0 +1,90 @@ +//! Maintained native benchmark: upload + forward + download, and complete decisions. +use anyhow::{Context, Result, ensure}; +use omni_cua_s1_native::{contract::Question, vision_engine::VisionEngine}; +use serde_json::json; +use sha2::{Digest, Sha256}; +use std::{path::PathBuf, time::Instant}; + +fn main() -> Result<()> { + let args: Vec<_> = std::env::args_os().skip(1).collect(); + ensure!( + args.len() == 1, + "usage: vision_graph_bench OUTPUT_JSON (CUA_S1_* environment)" + ); + let output = PathBuf::from(&args[0]); + ensure!(!output.exists(), "refusing to overwrite evidence"); + let path = |name| { + std::env::var_os(name) + .map(PathBuf::from) + .with_context(|| format!("set {name}")) + }; + let started = Instant::now(); + let mut engine = VisionEngine::load( + &path("CUA_S1_BASE")?, + &path("CUA_S1_VISION_ADAPTER")?, + &path("CUA_S1_MODEL")?, + &path("CUA_S1_CUDA_LIB")?, + )?; + let load_ms = started.elapsed().as_secs_f64() * 1000.; + let mut cold = Vec::new(); + let mut samples = Vec::new(); + for (width, height) in [(256, 256), (512, 256), (512, 512)] { + let mut prepared = Vec::new(); + for sample in 0..10usize { + let rgb: Vec = (0..width * height * 3) + .map(|i| ((i * 17 + sample * 31 + (i / (width * 3)) * 7) % 251) as u8) + .collect(); + let questions = [ + Question { + name: "move".into(), + goal: format!("Select the next move for screenshot {sample}."), + keys: vec![ + "left".into(), + "right".into(), + "continue".into(), + "stop".into(), + ], + labels: vec![ + "Move left".into(), + "Move right".into(), + "Continue".into(), + "Stop".into(), + ], + }, + Question { + name: "visible".into(), + goal: "Is the scene visible?".into(), + keys: vec!["yes".into(), "no".into()], + labels: vec!["Yes".into(), "No".into()], + }, + ]; + prepared.push(engine.prepare(width, height, &rgb, &questions)?); + } + let start = Instant::now(); + engine.vision.forward(&prepared[0].image)?; + cold.push(json!({"width":width,"height":height,"grid":prepared[0].image.image_grid_thw,"vision_first_ms":start.elapsed().as_secs_f64()*1000.})); + for request in prepared.iter().take(3) { + engine.predict_prepared(request)?; + } + for (sample, request) in prepared.iter().enumerate() { + let start = Instant::now(); + let features = engine.vision.forward(&request.image)?; + let vision_ms = start.elapsed().as_secs_f64() * 1000.; + let bytes: Vec = features + .iter() + .flat_map(|v| v.to_bits().to_le_bytes()) + .collect(); + let start = Instant::now(); + let response = engine.predict_prepared(request)?; + let decision_ms = start.elapsed().as_secs_f64() * 1000.; + samples.push(json!({"width":width,"height":height,"sample":sample,"grid":request.image.image_grid_thw,"features_sha256":format!("{:x}",Sha256::digest(bytes)),"vision_ms":vision_ms,"native_decision_ms":decision_ms,"response":response})); + } + } + std::fs::write( + output, + serde_json::to_vec_pretty( + &json!({"schema":"cua-s1-vision-graph-native-v1","load_ms":load_ms,"cold":cold,"samples":samples}), + )?, + )?; + Ok(()) +}