diff --git a/recipe/cua_s1/native.md b/recipe/cua_s1/native.md index 3a89e9d9..c7094390 100644 --- a/recipe/cua_s1/native.md +++ b/recipe/cua_s1/native.md @@ -104,7 +104,7 @@ not substitute for the full 139-request legacy upgrade and unmerged-reference comparisons. See the [four-case diagnosis and retest scope](../../docs/benchmarks/cua-s1-text-parity-20261013/README.md#legacy-merged-upgrade-regression). For the screenshot worker, [the vision Graph recipe](native_vision_graph.md) -documents `CUA_S1_VISION_GRAPH`, retained four-grid FIFO scratch, memory/lifetime +documents `CUA_S1_VISION_GRAPH`, bounded exact-grid scratch, memory/lifetime limits and the real CUDA capture-failure wrapper build and test commands. The ABI-8 parent synchronization preserves the unmerged precision profile. diff --git a/recipe/cua_s1/native_vision_graph.md b/recipe/cua_s1/native_vision_graph.md index ac999256..5518f8ca 100644 --- a/recipe/cua_s1/native_vision_graph.md +++ b/recipe/cua_s1/native_vision_graph.md @@ -7,23 +7,60 @@ unset or set it to `0` for eager execution. This switch is independent of `CUA_S1_GRAPH`, which controls text and multimodal language execution. Enable each switch separately; keep both unset for eager execution. -Both modes use the shared Qwen vision executor and retain at most four exact -[T,H,W] scratch/Graph pairs in FIFO order. Equal patch counts with different -heights or widths remain distinct. Before eviction the worker synchronizes, -destroys the old executable, and releases its buffers. Pixel values are uploaded -on every request, including replay; geometry tables belong to their exact grid. -Each retained grid has at most 4608 patches. This bounds the number of retained -activations, not their total bytes, loaded weights or the GEMM plan cache. - -The first ordinary forward for a grid finishes eagerly, then captures operators +Both modes use the shared Qwen vision executor and retain scratch by the exact `[T,H,W]` patch grid. The default entry +limit is 1; set `CUA_S1_VISION_CACHE_ENTRIES=4` to retain four geometries across +size changes. The entry setting is capped at 16; zero disables retention. Equal +patch counts with different heights or widths are distinct keys. The active +scratch is most recently used; a hit restores its existing buffers and Graph, +while a miss synchronizes and evicts least recently used entries before allocating +the new shape. Each executable is destroyed before its referenced buffers. +Pixel values are uploaded on every request, including replay. Geometry tables +remain valid only for that exact grid. + +`CUA_S1_VISION_CACHE_BYTES` bounds total requested bytes in retained scratch, +defaulting to 268435456 (256 MiB). The fixed 4B layout requests 64096 bytes per +patch, including pixels, position/rotary tables, LoRA work, activations and output. +The four tested grids contain 2304 patches in total and retain 147677184 bytes +(140.84 MiB). This budget excludes weights, GEMM workspace/plans, Graph metadata +and allocator/runtime overhead. Trace messages report retained entries/bytes; +measure actual device memory separately. + +Both limits are read once by the Cua 4B wrapper when the model loads. Other shared +vision consumers keep the existing four-grid FIFO policy unless they explicitly +select the bounded cache constructor. JEMM's reference execution retains its +separate four-entry ordered multi-image FIFO and reference cuDNN/attention +workspaces. Multi-image forwards require the generic constructor. Reference numerics also +cannot be combined with the bounded constructor: +ABI 8 does not expose the opaque cuDNN patch-plan workspace size, so it cannot +be charged exactly to this budget. Unset values use their defaults; +an explicitly invalid unsigned integer (including trailing whitespace or a +negative value) logs a warning to stderr and uses that limit's default. The +warning names the setting, supplied value and effective fallback, independently +of `CUA_S1_GRAPH_TRACE`. Zero for either limit disables retention. + +A valid geometry larger than the byte budget, or any geometry with retention +disabled, executes eagerly using transient scratch. It preserves existing +cache entries, synchronizes on success or error and releases the transient +buffers after a host result is read. Resident output remains model-owned while +its device view is borrowed; the worker releases transient storage after both +vision and language streams drain. Direct diagnostic users may call +`release_transient` after consuming the view; the next valid forward or model +drop also retires it. The default byte budget therefore prevents +retention even for a single oversized grid (the processor permits up to 4608 +patches); increase it deliberately if that workload needs retention. Temporary +scratch required to answer such a request is additional peak memory, so this is +a retained-activation budget rather than a whole-device memory limit. + +The first ordinary forward for a retained 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. +downloads and preserves an already matching Graph for later ordinary forwards; +it does not capture a new Graph itself. -For optional capture/replay messages, set `CUA_S1_GRAPH_TRACE=1`. Failure messages +For optional cache and 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. @@ -45,6 +82,8 @@ 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_CACHE_ENTRIES=4 \ +CUA_S1_VISION_CACHE_BYTES=268435456 \ CUA_S1_VISION_GRAPH=1 target/release/omni-cua-s1-vision ``` @@ -137,6 +176,19 @@ CUA_S1_TEST_COPY_FAILURE_EVIDENCE=/tmp/fresh-copy-drain.txt \ --lib device_copy_failure_tests -- --ignored --test-threads=1 --nocapture ``` +For the zero-retention resident path, use the same wrapper with two fresh evidence +paths. This production-entry probe checks both stream drain and release of the +matching transient output before the failed request returns, then checks recovery: + +```sh +CUA_S1_CUDA_LIB=/tmp/cua-fail-device-copy.so CUA_S1_TEST_COPY_FAILURE=1 \ +CUA_S1_VISION_CACHE_ENTRIES=0 \ +CUA_S1_TEST_COPY_FAILURE_EVIDENCE=/tmp/fresh-transient-drain.txt \ +CUA_S1_TEST_COPY_FAILURE_RELEASE_EVIDENCE=/tmp/fresh-transient-release.txt \ + cargo test --release --locked -p omni-cua-s1-native \ + --lib device_transient_failure_tests -- --ignored --test-threads=1 --nocapture +``` + ## Regressions and native measurements Export the same `CUA_S1_BASE`, `CUA_S1_VISION_ADAPTER`, `CUA_S1_MODEL`, and @@ -145,6 +197,13 @@ Export the same `CUA_S1_BASE`, `CUA_S1_VISION_ADAPTER`, `CUA_S1_MODEL`, and ```sh cargo test --release --locked -p omni-qwen3-5-native --lib vision::model::vision_tests::vision_graph_ \ -- --ignored --test-threads=1 --nocapture +CUA_S1_VISION_CACHE_ENTRIES=4 cargo test --release --locked \ + -p omni-qwen3-5-native --lib vision_cache_ \ + -- --ignored --test-threads=1 --nocapture +CUA_S1_VISION_CACHE_ENTRIES=4 CUA_S1_GRAPH_TRACE=1 \ + cargo test --release --locked -p omni-qwen3-5-native --lib vision_cache_aba_probe \ + -- --ignored --test-threads=1 --nocapture > /tmp/vision-cache-aba.log 2>&1 +python3 tests/cua_s1/vision_cache_trace_probe.py /tmp/vision-cache-aba.log CUA_S1_GRAPH=0 CUA_S1_VISION_GRAPH=0 cargo run --release --locked -p omni-cua-s1-native \ --example vision_graph_bench -- /tmp/vision-eager.json CUA_S1_GRAPH=0 CUA_S1_VISION_GRAPH=1 cargo run --release --locked -p omni-cua-s1-native \ @@ -162,11 +221,13 @@ 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-qwen3-5-native \ +CUA_S1_VISION_CACHE_ENTRIES=4 CUA_S1_GRAPH_TRACE=0 \ + cargo test --release --locked -p omni-qwen3-5-native \ --lib vision::model::vision_tests::vision_capture_failure_ -- --ignored --test-threads=1 --nocapture ``` -The root GPU tests cover changed pixels, exact-grid residency and eviction, retirement, +The root GPU tests cover changed pixels, exact-grid identity, A/B/A reuse, +count/byte LRU eviction, transient oversize and zero retention, 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. diff --git a/src/backends/cuda/qwen3_5/README.md b/src/backends/cuda/qwen3_5/README.md index 56df461f..96490530 100644 --- a/src/backends/cuda/qwen3_5/README.md +++ b/src/backends/cuda/qwen3_5/README.md @@ -96,14 +96,17 @@ Torch versions; the original approximate GDN entry points remain available. ## Shared vision cache and operator checks The generic shared vision constructor keeps a four-entry exact-grid FIFO for -compatibility. Its cap counts entries, not device bytes; four large grids may +compatibility. JEMM also retains a separate four-entry FIFO keyed by the ordered +multi-image grids. Their caps count entries, not device bytes; four large grids may retain substantially more memory than four small grids. A borrowed resident output prevents mutation or eviction until its consumer is finished. [Cua-S1's bounded-cache integration in PR #127](https://github.com/ThinkFlowLab/system1-omni/pull/127) uses a separate explicit policy: count/byte LRU limits, default one entry and 256 MiB, with zero retention available. The Cua wrapper owns environment parsing; -generic Qwen/JEMM users retain their existing policy. Retention-disabled or +generic Qwen/JEMM users retain their existing policy. Multi-image and JEMM +reference execution require the generic constructor: the bounded policy cannot +account for opaque cuDNN patch-plan allocations through ABI 8. Retention-disabled or oversized Cua output stays model-owned through both producer and consumer stream drains. Graph enablement and scratch retention are separate policies: turning Graph off alone does not release retained scratch. diff --git a/src/models/cua_s1/native/src/vision/model.rs b/src/models/cua_s1/native/src/vision/model.rs index 14532a28..c5d1a61e 100644 --- a/src/models/cua_s1/native/src/vision/model.rs +++ b/src/models/cua_s1/native/src/vision/model.rs @@ -3,6 +3,7 @@ use super::VisionCheckpoint; use crate::{cuda::DeviceBf16Matrix, image_preprocess::ProcessedImage}; use anyhow::Result; use half::bf16; +use omni_qwen3_5_native::vision::CacheLimits; use std::path::Path; pub struct VisionModel(omni_qwen3_5_native::vision::VisionModel); impl VisionModel { @@ -25,7 +26,13 @@ impl VisionModel { .map(|n| Ok((n.to_owned(), checkpoint.adapter_tensor(n)?))) .collect::>>()?; Ok(Self( - omni_qwen3_5_native::vision::VisionModel::from_tensors(config, base, lora, library)?, + omni_qwen3_5_native::vision::VisionModel::from_tensors_with_cache( + config, + base, + lora, + library, + cache_limits_from_env(), + )?, )) } pub fn synchronize(&self) -> Result<()> { @@ -38,6 +45,9 @@ impl VisionModel { pub fn forward_device(&mut self, image: &ProcessedImage) -> Result> { self.0.forward_device(image) } + pub fn release_transient(&mut self) -> Result<()> { + self.0.release_transient() + } pub fn forward_with_trace( &mut self, image: &ProcessedImage, @@ -46,3 +56,26 @@ impl VisionModel { self.0.forward_with_trace(image, callback) } } + +fn cache_limits_from_env() -> CacheLimits { + let value = |name: &str, default| match std::env::var_os(name) { + None => default, + Some(raw) => match raw.to_str().and_then(|v| v.parse::().ok()) { + Some(value) => value, + None => { + eprintln!( + "Invalid {name}={raw:?}; expected an unsigned integer, using default {default}" + ); + default + } + }, + }; + CacheLimits { + entries: value("CUA_S1_VISION_CACHE_ENTRIES", 1).min(16), + bytes: value("CUA_S1_VISION_CACHE_BYTES", 256 << 20), + } +} + +#[cfg(test)] +#[path = "../../../../../../tests/cua_s1/vision_cache_env.rs"] +mod cache_tests; diff --git a/src/models/cua_s1/native/src/vision_engine.rs b/src/models/cua_s1/native/src/vision_engine.rs index 69ff501e..a529f2ea 100644 --- a/src/models/cua_s1/native/src/vision_engine.rs +++ b/src/models/cua_s1/native/src/vision_engine.rs @@ -246,6 +246,7 @@ impl VisionEngine { let language_sync = self.language.synchronize(); vision_sync?; language_sync?; + self.vision.release_transient()?; result } @@ -320,3 +321,7 @@ mod device_feature_tests; #[cfg(test)] #[path = "../../../../../tests/cua_s1/device_copy_failure_gpu.rs"] mod device_copy_failure_tests; + +#[cfg(test)] +#[path = "../../../../../tests/cua_s1/device_transient_failure_gpu.rs"] +mod device_transient_failure_tests; diff --git a/src/models/qwen3_5/native/src/vision/mod.rs b/src/models/qwen3_5/native/src/vision/mod.rs index c0e09f84..b815369f 100644 --- a/src/models/qwen3_5/native/src/vision/mod.rs +++ b/src/models/qwen3_5/native/src/vision/mod.rs @@ -318,4 +318,4 @@ impl<'de, T: Deserialize<'de>> Deserialize<'de> for UniqueMap { mod geometry; mod model; pub use geometry::VisionGeometry; -pub use model::VisionModel; +pub use model::{CacheLimits, VisionModel}; diff --git a/src/models/qwen3_5/native/src/vision/model.rs b/src/models/qwen3_5/native/src/vision/model.rs index c17b2d06..e16865c9 100644 --- a/src/models/qwen3_5/native/src/vision/model.rs +++ b/src/models/qwen3_5/native/src/vision/model.rs @@ -39,6 +39,10 @@ pub struct VisionModel { config: VisionConfig, // Captures and their buffers retire before weights, GEMM workspace and stream. scratch: VecDeque, + // Uncached device output remains owned until its borrow is consumed. + transient: Option, + // None preserves the legacy four-grid FIFO policy for generic consumers. + cache_limits: Option, batch_scratch: VecDeque, graph_enabled: bool, base: BTreeMap, @@ -47,6 +51,11 @@ pub struct VisionModel { reference: Option<&'static cuda::Reference>, stream: OwnedStream, } +#[derive(Clone, Copy)] +enum ScratchLocation { + Retained(usize), + Transient, +} impl VisionModel { pub fn load(dir: impl AsRef, library: &Path) -> Result { let checkpoint = VisionCheckpoint::load(dir)?; @@ -63,6 +72,27 @@ impl VisionModel { base: Vec<(String, TensorView<'_>)>, lora: Vec<(String, TensorView<'_>)>, library: &Path, + ) -> Result { + Self::build(config, base, lora, library, None) + } + /// Explicit retained-byte/count LRU policy; environment settings belong to callers. + /// Reference numerics requires the generic FIFO constructor because the ABI does + /// not expose opaque cuDNN plan allocations for exact budget accounting. + pub fn from_tensors_with_cache( + config: VisionConfig, + base: Vec<(String, TensorView<'_>)>, + lora: Vec<(String, TensorView<'_>)>, + library: &Path, + cache_limits: CacheLimits, + ) -> Result { + Self::build(config, base, lora, library, Some(cache_limits)) + } + fn build( + config: VisionConfig, + base: Vec<(String, TensorView<'_>)>, + lora: Vec<(String, TensorView<'_>)>, + library: &Path, + cache_limits: Option, ) -> Result { let config = VisionConfig::from_value(serde_json::to_value(config)?)?; let expected = config.inventory(); @@ -132,6 +162,8 @@ impl VisionModel { let mut model = Self { config, scratch: VecDeque::new(), + transient: None, + cache_limits, batch_scratch: VecDeque::new(), graph_enabled: std::env::var("CUA_S1_VISION_GRAPH").as_deref() == Ok("1"), base: BTreeMap::new(), @@ -167,9 +199,13 @@ impl VisionModel { "reference vision requires unadapted 27B vision" ); ensure!( - self.scratch.is_empty() && self.batch_scratch.is_empty(), + self.scratch.is_empty() && self.batch_scratch.is_empty() && self.transient.is_none(), "select reference vision before the first forward" ); + ensure!( + self.cache_limits.is_none(), + "reference vision uses the generic FIFO cache; bounded cache cannot account for opaque cuDNN workspace" + ); if self.reference.is_some() { return Ok(()); } @@ -355,6 +391,10 @@ impl VisionModel { self.lora.is_empty(), "batched vision does not support vision adapters" ); + ensure!( + self.cache_limits.is_none(), + "batched vision requires the generic FIFO cache" + ); let grids: Vec<_> = images.iter().map(|image| image.image_grid_thw).collect(); let lengths = BatchGeometry::lengths(&grids, &self.config)?; let n: usize = lengths.iter().sum(); @@ -449,10 +489,10 @@ impl VisionModel { } /// Synchronized resident merger output; the borrow prevents cache mutation/retirement. pub fn forward_device(&mut self, image: &ProcessedImage) -> Result> { - let result: Result<(usize, usize)> = (|| { + let result: Result<(usize, ScratchLocation)> = (|| { let (n, index, capture) = self.enqueue_image(image, &mut None)?; self.synchronize()?; - if capture { + if let (true, ScratchLocation::Retained(index)) = (capture, index) { self.capture_current(index); } Ok((n, index)) @@ -460,9 +500,10 @@ impl VisionModel { if result.is_err() { // Drain while every referenced allocation and Graph is still owned. self.synchronize()?; + self.release_transient()?; } let (n, index) = result?; - self.scratch[index] + self.scratch_at(index) .buffers .out .bf16_matrix(n / 4, self.config.out_hidden_size) @@ -471,12 +512,12 @@ impl VisionModel { let result = (|| { let (n, index, capture) = self.enqueue_image(image, &mut callback)?; let output = self.read( - &self.scratch[index].buffers.out, + &self.scratch_at(index).buffers.out, n / 4 * self.config.out_hidden_size, )?; if let Some(callback) = callback.as_mut() { callback("merger.output", &output)?; - } else if capture { + } else if let (true, ScratchLocation::Retained(index)) = (capture, index) { self.capture_current(index); } Ok(output) @@ -484,8 +525,32 @@ impl VisionModel { if result.is_err() { self.synchronize()?; } + self.release_transient()?; result } + /// Release uncached output after every consuming stream has completed. + /// A live DeviceBf16Matrix borrow prevents calling this method. + /// ```compile_fail + /// use omni_qwen3_5_native::{vision::VisionModel, image_preprocess::ProcessedImage}; + /// fn retire_live_output(model: &mut VisionModel, image: &ProcessedImage) { + /// let output = model.forward_device(image).unwrap(); + /// model.release_transient().unwrap(); + /// let _ = output.download(); + /// } + /// ``` + pub fn release_transient(&mut self) -> Result<()> { + if self.transient.is_some() { + self.synchronize()?; + self.transient = None; + } + Ok(()) + } + fn scratch_at(&self, location: ScratchLocation) -> &Scratch { + match location { + ScratchLocation::Retained(index) => &self.scratch[index], + ScratchLocation::Transient => self.transient.as_ref().unwrap(), + } + } fn capture_current(&mut self, index: usize) { let scratch = &self.scratch[index]; // The caller completes warmup before recording; capture never reruns math. @@ -511,7 +576,7 @@ impl VisionModel { &mut self, image: &ProcessedImage, callback: &mut Trace<'_>, - ) -> Result<(usize, usize, bool)> { + ) -> Result<(usize, ScratchLocation, bool)> { cuda::set_device(0)?; let cached = self .scratch @@ -547,24 +612,58 @@ impl VisionModel { .iter() .flat_map(|v| bf16::from_f32(*v).to_bits().to_le_bytes()) .collect(); - let index = if let Some(index) = cached { - index - } else { + self.release_transient()?; + let bytes = + Scratch::bytes_for_grid(&self.config, !self.lora.is_empty(), image.image_grid_thw); + let index = if self + .cache_limits + .is_some_and(|limits| !limits.retains(bytes)) + { self.synchronize()?; - // FIFO is bounded to four exact grids; graph drops before its buffers. - if self.scratch.len() == 4 { - self.scratch.pop_front(); - } - let scratch = Scratch::new(self, image.image_grid_thw, geometry.unwrap())?; - self.scratch.push_back(scratch); - self.scratch.len() - 1 + let geometry = match geometry { + Some(geometry) => geometry, + None => VisionGeometry::new(image.image_grid_thw, &self.config)?, + }; + self.transient = Some(Scratch::new(self, image.image_grid_thw, geometry)?); + self.cache_trace("transient", image.image_grid_thw); + ScratchLocation::Transient + } else { + let index = if let Some(index) = cached { + if self.cache_limits.is_some() { + // Restore the exact grid as MRU, keeping inactive entries oldest first. + let scratch = self.scratch.remove(index).unwrap(); + self.scratch.push_back(scratch); + self.cache_trace("hit", image.image_grid_thw); + self.scratch.len() - 1 + } else { + index + } + } else { + self.synchronize()?; + if let Some(limits) = self.cache_limits { + while limits.needs_eviction(self.scratch.len(), self.retained_bytes(), bytes) { + let retired = self.scratch.pop_front().unwrap(); + let grid = retired.grid; + drop(retired); // Graph field retires before every captured buffer. + self.cache_trace("evicted", grid); + } + } else if self.scratch.len() == 4 { + self.scratch.pop_front(); + } + let scratch = Scratch::new(self, image.image_grid_thw, geometry.unwrap())?; + self.scratch.push_back(scratch); + self.cache_trace("miss", image.image_grid_thw); + self.scratch.len() - 1 + }; + ScratchLocation::Retained(index) }; - let scratch = &self.scratch[index]; + let scratch = self.scratch_at(index); // SAFETY: the current image has the validated exact resident buffer size. unsafe { cuda::upload(scratch.buffers.pixels.at(0), &pixels, self.stream.0)?; } if callback.is_none() + && matches!(index, ScratchLocation::Retained(_)) && self.graph_enabled && let Some(graph) = &scratch.graph { @@ -573,7 +672,32 @@ impl VisionModel { return Ok((n, index, false)); } self.enqueue(&scratch.buffers, n, &[n], callback)?; - Ok((n, index, callback.is_none() && self.graph_enabled)) + Ok(( + n, + index, + callback.is_none() + && self.graph_enabled + && matches!(index, ScratchLocation::Retained(_)), + )) + } + fn retained_bytes(&self) -> usize { + self.scratch + .iter() + .map(|s| Scratch::bytes_for_grid(&self.config, !self.lora.is_empty(), s.grid)) + .sum() + } + fn cache_trace(&self, action: &str, grid: [usize; 3]) { + if let Some(limits) = self.cache_limits + && std::env::var("CUA_S1_GRAPH_TRACE").as_deref() == Ok("1") + { + eprintln!( + "Vision cache {action} grid={grid:?} retained_entries={} retained_bytes={} entry_limit={} byte_limit={}", + self.scratch.len(), + self.retained_bytes(), + limits.entries, + limits.bytes + ); + } } fn enqueue( &self, @@ -842,7 +966,58 @@ struct Buffers { mlp: DeviceBuffer, out: DeviceBuffer, } +/// Budgets cover every retained Scratch allocation, excluding shared weights/GEMM workspace. +#[derive(Clone, Copy)] +pub struct CacheLimits { + pub entries: usize, + pub bytes: usize, +} +impl CacheLimits { + fn retains(self, bytes: usize) -> bool { + self.entries > 0 && bytes <= self.bytes + } + fn needs_eviction(self, entries: usize, retained_bytes: usize, incoming_bytes: usize) -> bool { + entries >= self.entries || incoming_bytes > self.bytes.saturating_sub(retained_bytes) + } +} + impl Scratch { + fn layout(config: &VisionConfig, lora: bool, grid: [usize; 3]) -> [usize; 18] { + let n = grid[1] * grid[2]; + let h = config.hidden_size; + let width = config + .intermediate_size + .max(h * 4) + .max(config.out_hidden_size); + let work_rows = if lora { n } else { 0 }; + [ + n * 1536 * 2, + n * 4 * 4, + n * 4 * 4, + n * (config.head_dim() / 2) * 4, + n * (config.head_dim() / 2) * 4, + work_rows * width * 4, + work_rows * 16 * 4, + work_rows * width * 4, + if h == 1024 { + 0 + } else { + n * config.num_heads * 80 * 4 * 2 + }, + n * h * 2, + n * h * 2, + n * h * 3 * 2, + n * h * 2, + n * h * 2, + n * h * 2, + n * h * 2, + n * width * 2, + n / 4 * config.out_hidden_size * 2, + ] + } + fn bytes_for_grid(config: &VisionConfig, lora: bool, grid: [usize; 3]) -> usize { + Self::layout(config, lora, grid).iter().sum() + } fn new(model: &VisionModel, grid: [usize; 3], geo: VisionGeometry) -> Result { let n = grid[1] * grid[2]; Ok(Self { @@ -859,12 +1034,14 @@ impl Buffers { attention_rows: usize, geo: VisionGeometry, ) -> Result { - let h = model.config.hidden_size; - let width = model - .config - .intermediate_size - .max(h * 4) - .max(model.config.out_hidden_size); + let mut sizes = Scratch::layout(&model.config, !model.lora.is_empty(), [1, 1, n]); + ensure!( + geo.indices.len() * 4 == sizes[1] + && geo.weights.len() * 4 == sizes[2] + && geo.cos.len() * 4 == sizes[3] + && geo.sin.len() * 4 == sizes[4], + "vision geometry storage does not match layout" + ); let patch = if let Some(reference) = model.reference { // SAFETY: descriptors/workspace are allocated outside capture for validated rows. let plan = PatchPlan(unsafe { (reference.patch_create)(n as i32, model.stream.0) }); @@ -876,7 +1053,7 @@ impl Buffers { } else { None }; - let attention_bytes = if let Some(reference) = model.reference { + sizes[8] = if let Some(reference) = model.reference { let floats = unsafe { (reference.vision_attention_workspace_floats)( attention_rows as i32, @@ -888,14 +1065,14 @@ impl Buffers { "cannot size reference vision attention workspace" ); attention_rows * model.config.num_heads * 80 * 4 * 2 + floats * 4 - } else if h == 1024 { + } else if model.config.hidden_size == 1024 { 0 } else { attention_rows * model.config.num_heads * 80 * 4 * 2 }; Ok(Self { patch, - pixels: DeviceBuffer::new(n * 1536 * 2)?, + pixels: DeviceBuffer::new(sizes[0])?, indices: model.upload( &geo.indices .iter() @@ -920,17 +1097,17 @@ impl Buffers { .flat_map(|v| v.to_le_bytes()) .collect::>(), )?, - w: Work::new(if model.lora.is_empty() { 0 } else { n }, width)?, - attention_workspace: DeviceBuffer::new(attention_bytes)?, - x: DeviceBuffer::new(n * h * 2)?, - norm: DeviceBuffer::new(n * h * 2)?, - qkv: DeviceBuffer::new(n * h * 3 * 2)?, - q: DeviceBuffer::new(n * h * 2)?, - k: DeviceBuffer::new(n * h * 2)?, - attn: DeviceBuffer::new(n * h * 2)?, - delta: DeviceBuffer::new(n * h * 2)?, - mlp: DeviceBuffer::new(n * width * 2)?, - out: DeviceBuffer::new(n / 4 * model.config.out_hidden_size * 2)?, + w: Work::new([sizes[5], sizes[6], sizes[7]])?, + attention_workspace: DeviceBuffer::new(sizes[8])?, + x: DeviceBuffer::new(sizes[9])?, + norm: DeviceBuffer::new(sizes[10])?, + qkv: DeviceBuffer::new(sizes[11])?, + q: DeviceBuffer::new(sizes[12])?, + k: DeviceBuffer::new(sizes[13])?, + attn: DeviceBuffer::new(sizes[14])?, + delta: DeviceBuffer::new(sizes[15])?, + mlp: DeviceBuffer::new(sizes[16])?, + out: DeviceBuffer::new(sizes[17])?, }) } } @@ -941,11 +1118,11 @@ struct Work { delta: DeviceBuffer, } impl Work { - fn new(n: usize, width: usize) -> Result { + fn new(sizes: [usize; 3]) -> Result { Ok(Self { - float_input: DeviceBuffer::new(n * width * 4)?, - rank: DeviceBuffer::new(n * 16 * 4)?, - delta: DeviceBuffer::new(n * width * 4)?, + float_input: DeviceBuffer::new(sizes[0])?, + rank: DeviceBuffer::new(sizes[1])?, + delta: DeviceBuffer::new(sizes[2])?, }) } } diff --git a/tests/cua_s1/device_transient_failure_gpu.rs b/tests/cua_s1/device_transient_failure_gpu.rs new file mode 100644 index 00000000..5861cd94 --- /dev/null +++ b/tests/cua_s1/device_transient_failure_gpu.rs @@ -0,0 +1,43 @@ +use super::device_feature_tests::{load, questions}; + +#[test] +#[ignore = "requires test-only queued-copy failure wrapper, zero-entry cache and CUA_S1_* weights"] +fn device_features_production_transient_failure_drains_and_releases_before_return() { + assert_eq!( + std::env::var("CUA_S1_TEST_COPY_FAILURE").as_deref(), + Ok("1") + ); + assert_eq!( + std::env::var("CUA_S1_VISION_CACHE_ENTRIES").as_deref(), + Ok("0") + ); + let evidence = + std::path::PathBuf::from(std::env::var_os("CUA_S1_TEST_COPY_FAILURE_EVIDENCE").unwrap()); + let release = std::path::PathBuf::from( + std::env::var_os("CUA_S1_TEST_COPY_FAILURE_RELEASE_EVIDENCE").unwrap(), + ); + assert!( + !evidence.exists() && !release.exists(), + "use fresh evidence paths" + ); + let mut engine = load(); + let first = engine + .prepare(256, 256, &vec![80; 256 * 256 * 3], &questions(1)) + .unwrap(); + let changed = engine + .prepare(512, 256, &vec![120; 512 * 256 * 3], &questions(2)) + .unwrap(); + engine.device_features = false; + let expected_first = engine.predict_prepared(&first).unwrap(); + let expected_changed = engine.predict_prepared(&changed).unwrap(); + engine.device_features = true; + let error = engine + .predict_prepared(&first) + .expect_err("queued D2D failure must surface"); + assert!(error.to_string().contains("device copy"), "{error:#}"); + // Check before recovery: the next enqueue would otherwise mask a missed release. + assert_eq!(std::fs::read_to_string(&evidence).unwrap(), "drained\n"); + assert_eq!(std::fs::read_to_string(&release).unwrap(), "released\n"); + assert_eq!(engine.predict_prepared(&changed).unwrap(), expected_changed); + assert_eq!(engine.predict_prepared(&first).unwrap(), expected_first); +} diff --git a/tests/cua_s1/fail_device_copy.cpp b/tests/cua_s1/fail_device_copy.cpp index a4c0376c..80ac4107 100644 --- a/tests/cua_s1/fail_device_copy.cpp +++ b/tests/cua_s1/fail_device_copy.cpp @@ -8,6 +8,16 @@ namespace { std::atomic injected{false}; std::atomic pending{false}; +std::atomic drained{false}; +std::atomic injected_source{nullptr}; +std::atomic injected_stream{nullptr}; +void record(const char* path, const char* message) { + if (!path) std::abort(); + FILE* evidence = std::fopen(path, "w"); + if (!evidence) std::abort(); + std::fputs(message, evidence); + if (std::fclose(evidence) != 0) std::abort(); +} void* real_library() { const char* path = std::getenv("CUA_S1_REAL_CUDA_LIB"); if (!path) std::abort(); @@ -23,6 +33,8 @@ extern "C" int cs1_copy_dd(void* dst, const void* src, size_t bytes, void* strea if (!copy || copy == &cs1_copy_dd) std::abort(); int result = copy(dst, src, bytes, stream); if (result == 0 && !injected.exchange(true)) { + injected_source.store(src); + injected_stream.store(stream); pending.store(true); std::fputs("Injected device-copy failure after a real queued D2D copy\n", stderr); result = 2; // cudaErrorMemoryAllocation, after the real copy was enqueued. @@ -36,13 +48,28 @@ extern "C" int cs1_stream_sync(void* stream) { auto sync = reinterpret_cast(dlsym(library, "cs1_stream_sync")); if (!sync || sync == &cs1_stream_sync) std::abort(); int result = sync(stream); - if (result == 0 && pending.exchange(false)) { - const char* path = std::getenv("CUA_S1_TEST_COPY_FAILURE_EVIDENCE"); - if (!path) std::abort(); - FILE* evidence = std::fopen(path, "w"); - if (!evidence) std::abort(); - std::fputs("drained\n", evidence); - if (std::fclose(evidence) != 0) std::abort(); + if (result == 0 && stream == injected_stream.load() && pending.exchange(false)) { + drained.store(true); + record(std::getenv("CUA_S1_TEST_COPY_FAILURE_EVIDENCE"), "drained\n"); + } + dlclose(library); + return result; +} + +// Optional production-path evidence: the producer allocation must retire only +// after the exact consumer stream drained, and before predict_prepared returns. +extern "C" int cs1_free(void* pointer) { + void* library = real_library(); + using Free = int (*)(void*); + auto free_device = reinterpret_cast(dlsym(library, "cs1_free")); + if (!free_device || free_device == &cs1_free) std::abort(); + const char* release = std::getenv("CUA_S1_TEST_COPY_FAILURE_RELEASE_EVIDENCE"); + const bool source = release && pointer == injected_source.load(); + if (source && !drained.load()) std::abort(); + int result = free_device(pointer); + if (source && result == 0) { + record(release, "released\n"); + injected_source.store(nullptr); } dlclose(library); return result; diff --git a/tests/cua_s1/vision_cache_env.rs b/tests/cua_s1/vision_cache_env.rs new file mode 100644 index 00000000..ba1f0e91 --- /dev/null +++ b/tests/cua_s1/vision_cache_env.rs @@ -0,0 +1,74 @@ +use super::*; + +#[test] +fn vision_cache_limits_from_env() { + // Each probe gets a private environment; parallel tests never mutate global env. + if std::env::var_os("CUA_S1_TEST_CACHE_LIMITS_PROBE").is_some() { + let limits = cache_limits_from_env(); + eprintln!("CACHE_LIMITS {} {}", limits.entries, limits.bytes); + return; + } + let test = format!( + "{}::vision_cache_limits_from_env", + module_path!().split_once("::").unwrap().1 + ); + let probe = |entries: Option<&str>, + bytes: Option<&str>, + expected: (usize, usize), + warning: Option<&str>| { + let mut command = std::process::Command::new(std::env::current_exe().unwrap()); + command + .args(["--exact", &test, "--nocapture"]) + .env("CUA_S1_TEST_CACHE_LIMITS_PROBE", "1") + .env_remove("CUA_S1_VISION_CACHE_ENTRIES") + .env_remove("CUA_S1_VISION_CACHE_BYTES"); + if let Some(value) = entries { + command.env("CUA_S1_VISION_CACHE_ENTRIES", value); + } + if let Some(value) = bytes { + command.env("CUA_S1_VISION_CACHE_BYTES", value); + } + let output = command.output().unwrap(); + let stderr = String::from_utf8(output.stderr).unwrap(); + assert!(output.status.success(), "{stderr}"); + assert!( + stderr.contains(&format!("CACHE_LIMITS {} {}", expected.0, expected.1)), + "{stderr}" + ); + assert_eq!( + stderr.matches("Invalid CUA_S1_VISION_CACHE_").count(), + usize::from(warning.is_some()), + "{stderr}" + ); + if let Some(name) = warning { + let default = if name.ends_with("ENTRIES") { + 1 + } else { + 256 << 20 + }; + assert!(stderr.contains(&format!("Invalid {name}=")), "{stderr}"); + assert!( + stderr.contains(&format!("using default {default}")), + "{stderr}" + ); + } + }; + probe(None, None, (1, 256 << 20), None); + probe(Some("0"), Some("0"), (0, 0), None); + probe(Some("4"), Some("1024"), (4, 1024), None); + probe(Some("17"), None, (16, 256 << 20), None); + for invalid in ["4 ", "-1", "four"] { + probe( + Some(invalid), + None, + (1, 256 << 20), + Some("CUA_S1_VISION_CACHE_ENTRIES"), + ); + probe( + None, + Some(invalid), + (1, 256 << 20), + Some("CUA_S1_VISION_CACHE_BYTES"), + ); + } +} diff --git a/tests/cua_s1/vision_cache_trace_probe.py b/tests/cua_s1/vision_cache_trace_probe.py new file mode 100644 index 00000000..303c1441 --- /dev/null +++ b/tests/cua_s1/vision_cache_trace_probe.py @@ -0,0 +1,20 @@ +#!/usr/bin/env python3 +"""Check live A/B/A CUDA logs after the Rust feature-parity probe. + +Run vision_cache_aba_probe with CUA_S1_VISION_CACHE_ENTRIES=4 and +CUA_S1_GRAPH_TRACE=1, redirect both stdout and stderr, then pass that log here. +A one-entry cache must fail this four-entry reuse assertion. +""" +import pathlib +import sys + +log = pathlib.Path(sys.argv[1]).read_text() +assert "test result: ok. 1 passed" in log, "the GPU feature-parity probe must pass" +assert log.count("VISION_CACHE_ABA_BEGIN") == log.count("VISION_CACHE_ABA_END") == 1 +window = log.split("VISION_CACHE_ABA_BEGIN")[1].split("VISION_CACHE_ABA_END")[0] +captures = window.count("Vision CUDA Graph captured") +replays = window.count("Vision CUDA Graph replayed") +assert (captures, replays) == (2, 1), ( + f"A/B/A expected 2 captures and 1 replay, observed {captures} captures and {replays} replays" +) +print(f"A/B/A retained graph: {captures} captures, {replays} replay; feature parity passed") diff --git a/tests/qwen3_5/vision_batch.rs b/tests/qwen3_5/vision_batch.rs index 2f6f5a4f..a157bf39 100644 --- a/tests/qwen3_5/vision_batch.rs +++ b/tests/qwen3_5/vision_batch.rs @@ -62,6 +62,12 @@ fn cpu_batch_uses_aggregate_gemms_and_image_local_attention() { "cs1_gemm_create", "cs1_gemm_destroy", "cs1_vision_linear", + "cs1_vision_position", + "cs1_vision_rope", + "cs1_vision_attention", + "cs1_vision_to_float", + "cs1_gemm_f32", + "cs1_vision_lora_add", "cs1_vision_norm", "cs1_vision_bias", "cs1_vision_gelu", @@ -180,6 +186,21 @@ fn cpu_batch_uses_aggregate_gemms_and_image_local_attention() { .get:: i32>(b"batch_live\0") .unwrap() }; + let live_bytes = unsafe { + *metrics + .get:: usize>(b"batch_live_bytes\0") + .unwrap() + }; + let unsafe_frees = unsafe { + *metrics + .get:: i32>(b"batch_unsafe_frees\0") + .unwrap() + }; + let fail_download = unsafe { + *metrics + .get::(b"batch_fail_download\0") + .unwrap() + }; let config = config(); let mut base = BTreeMap::new(); for name in config.inventory().keys() { @@ -192,6 +213,8 @@ fn cpu_batch_uses_aggregate_gemms_and_image_local_attention() { let mut model = VisionModel { config, scratch: VecDeque::new(), + transient: None, + cache_limits: None, batch_scratch: VecDeque::new(), graph_enabled: false, reference: None, @@ -297,6 +320,126 @@ fn cpu_batch_uses_aggregate_gemms_and_image_local_attention() { model.forward_images(&images).unwrap(); assert_eq!(unsafe { captures() }, capture_count); + // Cua's explicit policy uses the same Buffers owner as the generic/batch path. + // Retire generic state before measuring exact requested activation bytes. + model.synchronize().unwrap(); + model.scratch.clear(); + model.batch_scratch.clear(); + assert_eq!(unsafe { live_bytes() }, 0); + model.cache_limits = Some(CacheLimits { + entries: 2, + bytes: 32 * 47_168, + }); + model.graph_enabled = true; + let a = image([1, 2, 4]); + let b = image([1, 4, 2]); + let c = image([1, 2, 8]); + for input in [&a, &b, &a] { + model.forward(input).unwrap(); + assert_eq!(unsafe { live_bytes() }, model.retained_bytes()); + } + let a_pixels = model.scratch.back().unwrap().buffers.pixels.at(0); + model.forward(&c).unwrap(); + assert_eq!(model.scratch.len(), 2); + assert_eq!( + model.scratch.front().unwrap().buffers.pixels.at(0), + a_pixels + ); + assert_eq!(model.scratch.back().unwrap().grid, c.image_grid_thw); + assert_eq!(unsafe { live_bytes() }, 24 * 47_168); + assert_eq!(model.retained_bytes(), unsafe { live_bytes() }); + // This image exceeds the budget: retained captures/addresses survive it. + let large = image([1, 4, 16]); + model.forward(&large).unwrap(); + assert!(model.transient.is_none()); + assert_eq!(unsafe { live_bytes() }, model.retained_bytes()); + { + let output = model.forward_device(&large).unwrap(); + assert_eq!(output.download().unwrap().len(), 16 * 5120); + } + assert!(model.transient.as_ref().unwrap().graph.is_none()); + assert_eq!( + unsafe { live_bytes() }, + model.retained_bytes() + 64 * 47_168 + ); + model.release_transient().unwrap(); + assert_eq!(unsafe { live_bytes() }, model.retained_bytes()); + assert!( + model + .forward_with_trace(&large, |_, _| anyhow::bail!("trace failure")) + .is_err() + ); + assert!(model.transient.is_none()); + assert_eq!(unsafe { live_bytes() }, model.retained_bytes()); + // A failed download can leave work queued: synchronization must precede release. + unsafe { + fail_download(1); + } + assert!(model.forward(&large).is_err()); + assert!(model.transient.is_none()); + assert_eq!(unsafe { live_bytes() }, model.retained_bytes()); + unsafe { + fail_download(0); + } + model.forward(&a).unwrap(); + assert_eq!(model.scratch.back().unwrap().buffers.pixels.at(0), a_pixels); + model.synchronize().unwrap(); + model.scratch.clear(); + model.cache_limits = Some(CacheLimits { + entries: 0, + bytes: usize::MAX, + }); + model.forward_device(&a).unwrap(); + assert!( + model + .forward_images(&[image([1, 2, 4]), image([1, 4, 2])]) + .is_err() + ); + assert!(model.batch_scratch.is_empty()); + assert!(model.scratch.is_empty()); + assert!(model.transient.is_some()); + // Reference numerics cannot be enabled over an existing transient allocation. + assert!( + model + .enable_reference_numerics() + .unwrap_err() + .to_string() + .contains("before the first forward") + ); + model.release_transient().unwrap(); + // Opaque cuDNN workspaces cannot satisfy this explicit byte budget. + assert!( + model + .enable_reference_numerics() + .unwrap_err() + .to_string() + .contains("bounded cache") + ); + assert_eq!(unsafe { live_bytes() }, 0); + // The Cua 4B adapted layout additionally retains all three FP32 LoRA work buffers. + let mut cua_config = serde_json::to_value(&model.config).unwrap(); + cua_config["depth"] = 24.into(); + cua_config["hidden_size"] = 1024.into(); + cua_config["intermediate_size"] = 4096.into(); + cua_config["out_hidden_size"] = 2560.into(); + model.config = VisionConfig::from_value(cua_config).unwrap(); + for suffix in ["lora_A.weight", "lora_B.weight"] { + model.lora.insert( + format!("blocks.0.mlp.linear_fc1.{suffix}"), + DeviceBuffer::new(0).unwrap(), + ); + } + model.cache_limits = Some(CacheLimits { + entries: 1, + bytes: 8 * 64_096, + }); + model.forward(&a).unwrap(); + assert_eq!(model.retained_bytes(), 8 * 64_096); + assert_eq!(unsafe { live_bytes() }, model.retained_bytes()); + model.forward(&large).unwrap(); + assert!(model.transient.is_none()); + assert_eq!(unsafe { live_bytes() }, model.retained_bytes()); + assert_eq!(unsafe { unsafe_frees() }, 0); drop(model); assert_eq!(unsafe { live() }, 0); drop(metrics); diff --git a/tests/qwen3_5/vision_batch_fake.c b/tests/qwen3_5/vision_batch_fake.c index cf5d0aff..1540930d 100644 --- a/tests/qwen3_5/vision_batch_fake.c +++ b/tests/qwen3_5/vision_batch_fake.c @@ -1,6 +1,10 @@ #include +#include #include #include +static size_t live_bytes; +static int pending, in_capture, unsafe_frees, fail_download; +struct Allocation { size_t bytes; max_align_t alignment; }; static int live, linear_count, attention_count, captures, launches, destroyed, fail_capture; static int rows[16384], outs[16384], lens[16384]; static uintptr_t qkv_base; @@ -9,17 +13,18 @@ unsigned cs1_abi_version(void){return 9;} const char *cs1_error_string(int c){(void)c;return "fake CUDA error";} int cs1_get_device(int *d){if(!d)return 1;*d=0;return 0;} int cs1_set_device(int d){(void)d;return 0;} -int cs1_malloc(void **p,size_t n){*p=calloc(1,n);if(!*p)return 1;live++;return 0;} -int cs1_free(void *p){free(p);live--;return 0;} +int cs1_malloc(void **p,size_t n){struct Allocation *a=calloc(1,sizeof(*a)+n);if(!a)return 1;a->bytes=n;*p=a+1;live++;live_bytes+=n;return 0;} +int cs1_free(void *p){struct Allocation *a=(struct Allocation*)p-1;if(pending)unsafe_frees++;live_bytes-=a->bytes;free(a);live--;return 0;} int cs1_stream_create(void **p){*p=(void*)1;return 0;} -int cs1_stream_sync(void *s){(void)s;return 0;} +int cs1_stream_sync(void *s){(void)s;pending=0;return 0;} int cs1_stream_destroy(void *s){(void)s;return 0;} -int cs1_upload(void *d,const void *s,size_t n,void *stream){(void)stream;memcpy(d,s,n);return 0;} -int cs1_download(void *d,const void *s,size_t n,void *stream){(void)stream;memcpy(d,s,n);return 0;} +int cs1_upload(void *d,const void *s,size_t n,void *stream){(void)stream;memcpy(d,s,n);pending=0;return 0;} +int cs1_download(void *d,const void *s,size_t n,void *stream){if(fail_download){pending=1;return 1;}(void)stream;memcpy(d,s,n);pending=0;return 0;} void *cs1_gemm_create(size_t n){(void)n;return (void*)1;} void cs1_gemm_destroy(void *g){(void)g;} int cs1_vision_linear(void *g,const void *x,const void *w,const void *bias,void *y,int m,int n,int k,void *s){ (void)g;(void)x;(void)w;(void)bias;(void)y;(void)k;(void)s; + if(!in_capture)pending=1; rows[linear_count]=m;outs[linear_count]=n;linear_count++;return 0; } int cs1_vision_norm(const void*x,const void*w,const void*b,void*y,int n,int h,void*s){(void)x;(void)w;(void)b;(void)y;(void)n;(void)h;(void)s;return 0;} @@ -32,9 +37,9 @@ int cs1_vision_attention_v2(const void*q,const void*k,const void*v,void*o,int n, (void)h;(void)d;(void)w;(void)s; lens[attention_count]=n;qs[attention_count]=(uintptr_t)q;ks[attention_count]=(uintptr_t)k;vs[attention_count]=(uintptr_t)v;os[attention_count]=(uintptr_t)o;attention_count++;return 0; } -int cs1_graph_begin(void*s){(void)s;return 0;} -int cs1_graph_end(void*s,void**p){(void)s;if(fail_capture){*p=NULL;return 1;}*p=malloc(1);captures++;return 0;} -int cs1_graph_launch(void*g,void*s){(void)g;(void)s;launches++;return 0;} +int cs1_graph_begin(void*s){(void)s;in_capture=1;return 0;} +int cs1_graph_end(void*s,void**p){(void)s;in_capture=0;if(fail_capture){*p=NULL;return 1;}*p=malloc(1);captures++;return 0;} +int cs1_graph_launch(void*g,void*s){(void)g;(void)s;pending=1;launches++;return 0;} int cs1_graph_destroy(void*g){free(g);destroyed++;return 0;} int batch_live(void){return live;} int batch_linear_count(void){return linear_count;} @@ -53,3 +58,14 @@ int batch_destroyed(void){return destroyed;} uintptr_t batch_qkv_base(void){return qkv_base;} void batch_fail_capture(int value){fail_capture=value;} + +size_t batch_live_bytes(void){return live_bytes;} +int batch_unsafe_frees(void){return unsafe_frees;} +void batch_fail_download(int value){fail_download=value;} + +int cs1_vision_position(void*x,const void*t,const int*i,const float*w,int n,void*s){(void)x;(void)t;(void)i;(void)w;(void)n;(void)s;return 0;} +int cs1_vision_rope(const void*x,const float*c,const float*si,void*q,void*k,int n,void*s){(void)x;(void)c;(void)si;(void)q;(void)k;(void)n;(void)s;return 0;} +int cs1_vision_attention(const void*q,const void*k,const void*v,void*o,int n,void*s){(void)q;(void)k;(void)v;(void)o;(void)n;(void)s;return 0;} +int cs1_vision_to_float(const void*x,float*y,size_t n,void*s){(void)x;(void)y;(void)n;(void)s;return 0;} +int cs1_gemm_f32(void*g,const float*x,const float*w,float*y,int m,int n,int k,void*s){(void)g;(void)x;(void)w;(void)y;(void)m;(void)n;(void)k;(void)s;return 0;} +int cs1_vision_lora_add(void*x,const float*d,size_t n,float scale,void*s){(void)x;(void)d;(void)n;(void)scale;(void)s;return 0;} diff --git a/tests/qwen3_5/vision_graph.rs b/tests/qwen3_5/vision_graph.rs index 89933718..1726c5a5 100644 --- a/tests/qwen3_5/vision_graph.rs +++ b/tests/qwen3_5/vision_graph.rs @@ -2,6 +2,22 @@ use super::*; fn load() -> VisionModel { + load_with_cache(None) +} +fn load_cached() -> VisionModel { + load_with_cache(Some(CacheLimits { + entries: std::env::var("CUA_S1_VISION_CACHE_ENTRIES") + .ok() + .and_then(|v| v.parse::().ok()) + .unwrap_or(1) + .min(16), + bytes: std::env::var("CUA_S1_VISION_CACHE_BYTES") + .ok() + .and_then(|v| v.parse::().ok()) + .unwrap_or(256 << 20), + })) +} +fn load_with_cache(cache: Option) -> VisionModel { use memmap2::Mmap; use safetensors::SafeTensors; use std::{collections::BTreeSet, fs::File}; @@ -52,7 +68,17 @@ fn load() -> VisionModel { .filter(|name| name.starts_with("base_model.model.model.visual.")) .map(|name| ((*name).to_owned(), adapter.tensor(name).unwrap())) .collect(); - VisionModel::from_tensors(config, tensors, lora, &path("CUA_S1_CUDA_LIB")).unwrap() + match cache { + Some(limits) => VisionModel::from_tensors_with_cache( + config, + tensors, + lora, + &path("CUA_S1_CUDA_LIB"), + limits, + ), + None => VisionModel::from_tensors(config, tensors, lora, &path("CUA_S1_CUDA_LIB")), + } + .unwrap() } fn image(grid: [usize; 3], seed: usize) -> ProcessedImage { let n = grid[1] * grid[2] * 1536; @@ -181,12 +207,19 @@ fn vision_graph_equal_patch_count_geometry_and_retirement() { #[test] #[ignore = "requires the test-only failing CUDA wrapper; see native_vision_graph.md"] fn vision_capture_failure_keeps_eager_result_and_disables_replay() { + check_capture_failure(load()); +} +#[test] +#[ignore = "requires the test-only failing CUDA wrapper; see native_vision_graph.md"] +fn vision_capture_failure_cached_keeps_eager_result_and_disables_replay() { + check_capture_failure(load_cached()); +} +fn check_capture_failure(mut model: VisionModel) { 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); @@ -204,6 +237,216 @@ fn vision_capture_failure_keeps_eager_result_and_disables_replay() { assert_eq!(bits(&model.forward(&input).unwrap()), expected); } +#[test] +#[ignore = "requires CUDA and pinned CUA_S1_BASE/CUA_S1_VISION_ADAPTER/CUA_S1_CUDA_LIB"] +fn vision_cache_aba_probe() { + let mut model = load_cached(); + model.graph_enabled = false; + let inputs = [image([1, 8, 32], 21), image([1, 16, 16], 22)]; + let expected: Vec<_> = inputs + .iter() + .map(|i| bits(&model.forward(i).unwrap())) + .collect(); + model.graph_enabled = true; + eprintln!("VISION_CACHE_ABA_BEGIN"); + for index in [0, 1, 0] { + assert_eq!( + bits(&model.forward(&inputs[index]).unwrap()), + expected[index] + ); + } + eprintln!("VISION_CACHE_ABA_END"); +} + +fn config(hidden: usize) -> VisionConfig { + VisionConfig::from_value(serde_json::json!({"depth": if hidden==1024 {24} else {27}, + "hidden_size":hidden,"intermediate_size":if hidden==1024 {4096} else {4304}, + "num_heads":16,"num_position_embeddings":2304,"out_hidden_size":if hidden==1024 {2560} else {5120}, + "in_channels":3,"patch_size":16,"temporal_patch_size":2,"spatial_merge_size":2, + "hidden_act":"gelu_pytorch_tanh","deepstack_visual_indexes":[],"model_type":"qwen3_5"})).unwrap() +} + +#[test] +fn vision_cache_dimension_storage() { + for (hidden, lora, per_patch) in [ + (1024, true, 64_096), + (1024, false, 31_264), + (1152, false, 47_168), + ] { + let config = config(hidden); + for grid in [[1, 2, 2], [1, 8, 32], [1, 32, 8]] { + let geometry = VisionGeometry::new(grid, &config).unwrap(); + let sizes = Scratch::layout(&config, lora, grid); + assert_eq!( + [ + geometry.indices.len() * 4, + geometry.weights.len() * 4, + geometry.cos.len() * 4, + geometry.sin.len() * 4 + ], + sizes[1..5] + ); + assert_eq!( + Scratch::bytes_for_grid(&config, lora, grid), + grid[1] * grid[2] * per_patch + ); + } + } +} + +#[test] +fn vision_cache_budget_policy() { + let patch_bytes = Scratch::bytes_for_grid(&config(1024), true, [1, 2, 2]) / 4; + assert_eq!(patch_bytes, 64_096); + let geometry = VisionGeometry::new([1, 2, 2], &config(1024)).unwrap(); + let geometry_bytes = + (geometry.indices.len() + geometry.weights.len() + geometry.cos.len() + geometry.sin.len()) + * 4; + assert_eq!(geometry_bytes / 4, 288); + assert_eq!( + Scratch::bytes_for_grid(&config(1024), true, [1, 32, 72]), + 147_677_184 + ); + let limits = CacheLimits { + entries: 2, + bytes: 100, + }; + assert!(limits.retains(100)); + assert!(!limits.retains(101)); + assert!( + !CacheLimits { + entries: 0, + bytes: 100 + } + .retains(1) + ); + assert!(!limits.needs_eviction(1, 60, 40)); + assert!(limits.needs_eviction(2, 0, 1)); + assert!(limits.needs_eviction(1, 61, 40)); +} + +fn clear_scratch(model: &mut VisionModel) { + model.synchronize().unwrap(); + model.scratch.clear(); +} + +#[test] +#[ignore = "requires CUDA and pinned CUA_S1_BASE/CUA_S1_VISION_ADAPTER/CUA_S1_CUDA_LIB"] +fn vision_cache_capacity_lru_and_invalid_input() { + let mut model = load_cached(); + model.cache_limits = Some(CacheLimits { + entries: 2, + bytes: 256 << 20, + }); + model.graph_enabled = false; + let inputs = [ + image([1, 8, 32], 30), + image([1, 16, 16], 31), + image([1, 32, 8], 32), + ]; + let expected: Vec<_> = inputs + .iter() + .map(|i| bits(&model.forward(i).unwrap())) + .collect(); + clear_scratch(&mut model); + model.graph_enabled = true; + for index in [0, 1, 0] { + assert_eq!( + bits(&model.forward(&inputs[index]).unwrap()), + expected[index] + ); + } + let a = model.scratch.back().unwrap().buffers.pixels.at(0); + assert!(model.scratch.back().unwrap().graph.is_some()); + assert_eq!( + model.scratch.front().unwrap().grid, + inputs[1].image_grid_thw + ); + let b = model.scratch.front().unwrap().buffers.pixels.at(0); + let mut invalid = image([1, 16, 16], 31); + invalid.pixel_values.pop(); + assert!(model.forward(&invalid).is_err()); + invalid = image([2, 8, 32], 30); + assert!(model.forward(&invalid).is_err()); + assert_eq!(model.scratch.back().unwrap().buffers.pixels.at(0), a); + assert_eq!(model.scratch.front().unwrap().buffers.pixels.at(0), b); + assert_eq!(bits(&model.forward(&inputs[2]).unwrap()), expected[2]); + assert_eq!(model.scratch.len(), 2); + assert_eq!( + model.scratch.front().unwrap().grid, + inputs[0].image_grid_thw, + "B must be LRU victim after A/B/A/C" + ); + assert_eq!(model.scratch.front().unwrap().buffers.pixels.at(0), a); + assert_eq!(model.retained_bytes(), 512 * 64_096); + assert_eq!(bits(&model.forward(&inputs[0]).unwrap()), expected[0]); + assert_eq!(model.scratch.back().unwrap().buffers.pixels.at(0), a); + drop(model); +} + +#[test] +#[ignore = "requires CUDA and pinned CUA_S1_BASE/CUA_S1_VISION_ADAPTER/CUA_S1_CUDA_LIB"] +fn vision_cache_byte_eviction_and_transient_oversize() { + let mut model = load_cached(); + model.cache_limits = Some(CacheLimits { + entries: 4, + bytes: 512 * 64_096, + }); + model.graph_enabled = false; + let inputs = [ + image([1, 8, 32], 40), + image([1, 16, 32], 41), + image([1, 32, 8], 42), + image([1, 32, 32], 43), + ]; + let expected: Vec<_> = inputs + .iter() + .map(|i| bits(&model.forward(i).unwrap())) + .collect(); + clear_scratch(&mut model); + model.graph_enabled = true; + for index in [0, 1, 2, 0] { + assert_eq!( + bits(&model.forward(&inputs[index]).unwrap()), + expected[index] + ); + assert!(model.retained_bytes() <= model.cache_limits.as_ref().unwrap().bytes); + } + assert_eq!(model.scratch.len(), 2); + assert_eq!( + model.scratch.front().unwrap().grid, + inputs[2].image_grid_thw + ); + let a = model.scratch.back().unwrap().buffers.pixels.at(0); + let c = model.scratch.front().unwrap().buffers.pixels.at(0); + for _ in 0..2 { + assert_eq!(bits(&model.forward(&inputs[3]).unwrap()), expected[3]); + assert_eq!(model.scratch.back().unwrap().buffers.pixels.at(0), a); + assert_eq!(model.scratch.front().unwrap().buffers.pixels.at(0), c); + assert_eq!( + model.retained_bytes(), + model.cache_limits.as_ref().unwrap().bytes + ); + assert!(model.graph_enabled); + } + assert!( + model + .forward_with_trace(&inputs[3], |_, _| anyhow::bail!("test trace failure")) + .is_err() + ); + assert_eq!(bits(&model.forward(&inputs[0]).unwrap()), expected[0]); + model.cache_limits.as_mut().unwrap().bytes = 1; + assert_eq!(bits(&model.forward(&inputs[0]).unwrap()), expected[0]); + // Internal test-only budget changes may encounter an already resident same grid. + assert_eq!(model.scratch.back().unwrap().buffers.pixels.at(0), a); + clear_scratch(&mut model); + model.cache_limits.as_mut().unwrap().entries = 0; + assert_eq!(bits(&model.forward(&inputs[0]).unwrap()), expected[0]); + assert_eq!(bits(&model.forward(&inputs[0]).unwrap()), expected[0]); + assert_eq!(model.retained_bytes(), 0); + assert!(model.scratch.is_empty()); +} + #[test] #[ignore = "requires CUDA and pinned CUA_S1_BASE/CUA_S1_VISION_ADAPTER/CUA_S1_CUDA_LIB"] fn device_vision_graph_replays_changed_pixels_and_replaces_geometry() { @@ -268,3 +511,57 @@ fn device_vision_capture_failure_preserves_resident_eager_result_and_recovery() expected ); } + +#[test] +#[ignore = "requires CUDA and pinned CUA_S1_BASE/CUA_S1_VISION_ADAPTER/CUA_S1_CUDA_LIB"] +fn device_cache_transient_output_keeps_retained_entries_and_releases_after_consumption() { + let mut model = load_cached(); + model.cache_limits = Some(CacheLimits { + entries: 1, + bytes: 256 * 64_096, + }); + model.graph_enabled = false; + let small = image([1, 16, 16], 61); + let large = image([1, 16, 32], 62); + let expected_small = bits(&model.forward(&small).unwrap()); + let expected_large = bits(&model.forward(&large).unwrap()); + assert!( + model.transient.is_none(), + "host transient must be released before return" + ); + let retained = model.scratch[0].buffers.pixels.at(0); + model.graph_enabled = true; + for input in [&small, &large, &large, &small] { + let observed = { + let output = model.forward_device(input).unwrap(); + bits(&output.download().unwrap()) + }; + assert_eq!( + observed, + if input.image_grid_thw == small.image_grid_thw { + &expected_small + } else { + &expected_large + } + .clone() + ); + assert_eq!(model.scratch[0].buffers.pixels.at(0), retained); + assert_eq!(model.retained_bytes(), 256 * 64_096); + if input.image_grid_thw == large.image_grid_thw { + assert!(model.transient.as_ref().unwrap().graph.is_none()); + } + model.release_transient().unwrap(); + assert!(model.transient.is_none()); + } + clear_scratch(&mut model); + model.cache_limits.as_mut().unwrap().entries = 0; + let observed = { + let output = model.forward_device(&small).unwrap(); + bits(&output.download().unwrap()) + }; + assert_eq!(observed, expected_small); + assert!(model.scratch.is_empty()); + assert!(model.transient.is_some()); + model.release_transient().unwrap(); + assert!(model.transient.is_none()); +}