cua_s1: retain bounded vision cache across screenshot grids

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.

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.
 

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.

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::<Result<Vec<_>>>()?;
         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<DeviceBf16Matrix<'_>> {
         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::<usize>().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;

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;

src/models/qwen3_5/native/src/vision/mod.rs

@@ -318,4 +318,4 @@ impl<'de, T: Deserialize<'de>> Deserialize<'de> for UniqueMap<T> {
 mod geometry;
 mod model;
 pub use geometry::VisionGeometry;
-pub use model::VisionModel;
+pub use model::{CacheLimits, VisionModel};