diff --git a/docs/architecture.md b/docs/architecture.md index 276143e3..ad90d1a7 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -71,9 +71,12 @@ not selected. Cua-S1 and Open-Jev compute their decision heads on the CPU after downloading the final hidden state. Cua-S1 screenshots can opt into request-local -aligned token/position prefix reuse with fixed-M64 or full-prompt-matched GEMMs, -host features and language Graph off. Captured contents belong to one request; one engine-owned storage -allocation may retain capacity for later captures. Suffixes execute independently, +aligned token/position prefix reuse with fixed-M64 or full-prompt-matched GEMMs. +Fixed-M64 prefix modes use host features and language Graph off; matched prefixes +support borrowed resident feature windows and bounded language Graph replay. +Current rows are staged outside capture; entries own both recorded prefix +allocations until synchronized retirement. Captured contents belong to one +request; one engine-owned storage allocation may retain capacity for later captures. Suffixes execute independently, and the existing letter head and token accounting are preserved. Decider projects its selected LM-head rows on CUDA. JEMM uses the shared Qwen executor with separate unmerged diff --git a/recipe/cua_s1/native_screenshot.md b/recipe/cua_s1/native_screenshot.md index ef32802c..7b225d8a 100644 --- a/recipe/cua_s1/native_screenshot.md +++ b/recipe/cua_s1/native_screenshot.md @@ -76,11 +76,11 @@ multimodal adapter and reports prompt-token usage with zero generated tokens. ## Execution controls and limits -`CUA_S1_GRAPH=1` enables host-feature language Graph replay; +`CUA_S1_GRAPH=1` enables language Graph replay for host or resident features; `CUA_S1_VISION_GRAPH=1` independently enables vision Graph replay. `CUA_S1_DEVICE_FEATURES=1` keeps vision features on the GPU and copies image rows -directly into language scratch. Its language path remains eager; it can be paired -with either vision Graph setting. All three controls default off. See the +directly into language scratch before capture/replay. It can be paired with +either vision Graph setting. All three controls default off. See the [resident-feature contract](native_vision_graph.md#gpu-resident-vision-features) for strict control parsing, synchronization and reproduction commands. @@ -121,8 +121,9 @@ or less than 64 shared tokens, uses complete fixed passes and never reads a previous snapshot. Retained capacity is additional device memory, separate from the vision-cache budget. -All active prefix modes require host features and language Graph off; vision Graph -remains independently selectable. They keep the existing prompt, CPU letter head, +The `fixed` and `shared` modes require host features and language Graph off. +`matched` supports host or resident features and either language Graph setting; +vision Graph remains independently selectable. They keep the existing prompt, CPU letter head, answer order and usage accounting. Fixed GEMMs can round differently from default execution: compare `shared` with `fixed` for continuation parity, and separately report probability drift/decision changes against `off`. No automatic mode or @@ -185,8 +186,10 @@ and restore the ordinary/fixed-M64 handle selection before returning/unwinding. Descriptor equality and algorithm support do not establish arithmetic invariance across row counts on every GPU. Qualify the exact GPU/cuBLAS/model combination before relying on ordinary-output parity. Keep `off` as the ordinary numerical -control and `fixed`/`shared` as the separate M64 controls. No automatic policy, -language Graph integration, changed prompt or new precision profile is enabled. +control and `fixed`/`shared` as the separate M64 controls. No automatic policy, changed prompt or new precision profile is enabled. +With `CUA_S1_GRAPH=1`, admitted matched prefixes and suffixes use bounded window +Graphs; with `CUA_S1_DEVICE_FEATURES=1`, current image rows are copied directly +from the borrowed vision output into owned language scratch. ```sh cargo test --release --locked -p omni-qwen3-5-native --test kernels \ @@ -239,3 +242,87 @@ and acquisition scripts, failures and final-source mapping are in the [2026-10-26 qualification](https://github.com/Levius-Fubuki/system1-omni/releases/download/cua-mm-qualification-20261026/cua-mm-matched-qualification-20261026.tar.gz). Original real screenshot pixels are omitted. These measurements precede the startup-symbol guard. That guard has host-only missing-symbol regression coverage; no new GPU numerical or timing run is claimed. + +### Resident matched prefixes with window Graphs + +```sh +export CUA_S1_MM_PREFIX=matched CUA_S1_DEVICE_FEATURES=1 +export CUA_S1_GRAPH=1 CUA_S1_VISION_GRAPH=1 +# Keep prefix=off as the ordinary complete-request numerical control. +``` + +The matched implementation is shared with [PR #184](https://github.com/InferMatrixLab/system1-omni/pull/184); this extension combines it +with the resident language-Graph path. Prefix reuse removes repeated computation +across questions; any combined speedup must be attributed to that work as well +as Graph submission savings. A single question or incompatible GEMM plans still +uses ordinary independent execution, preserving the selected Graph/device mode. + +Host and device windows use the same token/T-H-W prefix planner and CPU letter +head. Checked device row slices retain the vision producer borrow. The consumer +copies current rows into its own scratch before capture; no vision pointer is +recorded. Device window exits drain the language stream, including errors. + +Window Graph keys contain absolute begin/end rows, both source and destination +prefix allocation identities, and the selected GEMM handle. Entries retain both +prefix allocations, including the simultaneous continue-and-capture case. This +makes producer release and prefix-storage replacement safe. At most 128 window +Graphs are retained; eviction drains first. Scratch growth and matched-workspace +retirement clear affected captures after synchronization. Recording failure keeps +the warmed eager result, disables all language Graph modes and retires their +caches. Vision Graph is independent. Retained snapshot allocations, Graph objects +and up to 16 matched 32 MiB workspaces add GPU memory beyond the vision budget. + +Warm the largest request and the complete shape working set, then check traces +for new captures or failures in the measured region. Cold capture and cache +churn must be reported separately; all switches remain opt-in. + +```sh +CUA_S1_GRAPH=1 CUA_S1_DEVICE_FEATURES=1 CUA_S1_MM_PREFIX=matched \ + cargo test --release --locked -p omni-cua-s1-native --lib prefix_tests \ + -- --ignored --test-threads=1 --nocapture +# Repeat with vision Graph0 and1. +CUA_S1_TEST_MODEL="$CUA_S1_MODEL" cargo test --release --locked \ + -p omni-qwen3-5-native --lib model::graph_tests \ + -- --ignored --test-threads=1 --nocapture +``` + +### Recorded resident matched qualification + +On RTX 4090 / SM8.9 / CUDA 13.0, the 2026-10-26 frozen complete-request +comparison measured the following. O is original PR183 at `911faa0d`, E is +resident matched prefix eager, and G is resident matched prefix with window +Graphs. All use vision Graph and the same pinned artifacts. Each mode has 64 +synthetic requests and 96 requests from 12 fixed real screenshots, concurrency 1. + +| Workload / mode | Mean ms | p50 ms | p95 ms | Sampled process maximum MiB | +|---|---:|---:|---:|---:| +| Synthetic / O | 130.257 | 117.968 | 269.950 | 10202 | +| Synthetic / E | 105.979 | 102.867 | 169.251 | 10448 | +| Synthetic / G | 104.902 | 101.780 | 172.540 | 10538 | +| Real / O | 248.819 | 258.092 | 263.407 | 10256 | +| Real / E | 150.949 | 151.120 | 157.391 | 10606 | +| Real / G | 147.868 | 148.416 | 154.275 | 10694 | + +G reduces the real mean by 40.6% and p95 by 41.4% versus O, with 438 MiB higher +sampled process maximum. Relative to E, Graph reduces the real mean by 2.0%; the +large combined gain therefore includes prefix computation reuse. Synthetic mean +falls 19.5% versus O, while the smallest two-question case regresses 4.8%; G's +synthetic p95 is 1.9% higher than E. These are scoped measurements, not a default-on +recommendation or a universal speedup. + +Two balanced O/E/G/G/E/O visits per suite each exclude a growth traversal, a +stabilization traversal and five warmups, then time four shuffled traversals. +The 480 timed HTTP samples include response parsing/schema validation; 300 +excluded and 135 feasibility samples remain separate. No timed new capture or +capture failure occurred. Five profiles each produced 192 regular numerical +records: raw logits, full responses and token/T-H-W inputs match original O +exactly, including separate single-question extraction and malformed-input +recovery. Cold/growth timings and startup are retained separately. Real quality +remains 119/192 correlated choice permutations, with existing native-vs-FP32 +reference limitations; this is not official ScreenSpot accuracy. + +The [raw qualification and portable audit](https://github.com/Levius-Fubuki/system1-omni/releases/tag/cua183-resident-20261026) +bind source/build/model pins, corpus reconstruction, cold and measured samples, +per-case regressions, memory, admission traces and GPU lifecycle tests. Original +screenshot pixels are omitted. Other hardware, concurrent serving and general +GUI task quality require separate qualification. diff --git a/recipe/cua_s1/native_vision_graph.md b/recipe/cua_s1/native_vision_graph.md index 7675280f..04b83cc4 100644 --- a/recipe/cua_s1/native_vision_graph.md +++ b/recipe/cua_s1/native_vision_graph.md @@ -65,6 +65,37 @@ For optional cache and capture/replay messages, set `CUA_S1_GRAPH_TRACE=1`. Fail are unconditional. Requests execute serially through the existing worker runtime; this feature does not add batching or concurrent forwards. +## Warming the complete working set + +Language scratch growth invalidates earlier language Graphs. Exercise the largest +input first, then warm every measured sequence length and vision grid. Five +warmups alone need not cover a heterogeneous workload. With tracing enabled, +check for timed replays and no new captures or fallback before describing a run +as stable replay. If the real workload exceeds the cache, include and report +its capture/eviction costs instead. + +Entry and byte limits must both fit the intended working set. The recorded +8-payload synthetic matrix uses six grids: its 1024-square image alone requests +262537216 bytes, leaving too little of a 256 MiB budget for the smallest grid. +The original timed rounds therefore contained 15 vision captures (3/6/6), plus +5 language captures after scratch growth in the first round. The original +aggregate latency comparison is not a pure warmed-replay measurement. + +For that bounded matrix, explicitly setting `CUA_S1_VISION_CACHE_ENTRIES=8` and +`CUA_S1_VISION_CACHE_BYTES=1073741824` eliminated timed vision captures after full +warming; these are workload-specific settings, not new defaults or a total GPU +memory limit. On RTX 4090 / SM8.9, matched resident eager/Graph p50 was +116.891/116.399 ms and p95 was 271.109/269.011 ms (24 timed requests per mode, +3 rounds, concurrency 1). Sampled process maxima were 10138/10202 MiB. This small +sample does not establish a general acceleration: Graphs reduce submission +work, while GEMMs, input staging/checks, CPU preprocessing/readout and stream +synchronization remain. Defaults stay off. + +The [follow-up qualification](https://github.com/Levius-Fubuki/system1-omni/releases/download/cua-mm-qualification-20261025/cua-mm-followup-qualification-20261025.tar.gz) +preserves both cache settings, raw/excluded samples, capture traces, source pins +and the full-warming protocol. This follow-up uses the unchanged runtime from +`a14e1027e7de7507e6ff82b3844a916ff19a26b9`. + ## Build and launch Use the pinned base and multimodal adapter revisions listed in the @@ -104,8 +135,12 @@ export CUA_S1_DEVICE_FEATURES=1 ``` This switch is independent of `CUA_S1_VISION_GRAPH`. Vision can run eagerly or -replay a Graph; the resident-feature language path remains eager in both modes. -The host-feature language path retains the independent `CUA_S1_GRAPH` control. +replay a Graph; `CUA_S1_GRAPH=1` independently enables language capture/replay +for both host and resident features. Current token IDs, T/H/W positions and +image rows are staged before capture/replay into language-owned scratch. Captures +do not retain the borrowed vision allocation; host/device calls with the same +sequence length share the existing multimodal cache. Scratch growth and capture +failure retain the existing invalidation and eager-fallback behavior. The resident path reuses one vision output across all questions in the request and copies contiguous image-placeholder rows directly into language scratch. It avoids the full feature download and per-question feature upload. It still @@ -142,14 +177,24 @@ The finite-scan test needs only `CUA_S1_CUDA_LIB`; the other tests need all four ```sh cargo test --release --locked -p omni-qwen3-5-native --lib device_tests \ -- --ignored --test-threads=1 -CUA_S1_VISION_GRAPH=0 cargo test --release --locked -p omni-cua-s1-native \ +CUA_S1_GRAPH=1 CUA_S1_VISION_GRAPH=0 cargo test --release --locked -p omni-cua-s1-native \ --lib device_feature_tests -- --ignored --test-threads=1 -CUA_S1_VISION_GRAPH=1 cargo test --release --locked -p omni-cua-s1-native \ +CUA_S1_GRAPH=1 CUA_S1_VISION_GRAPH=1 cargo test --release --locked -p omni-cua-s1-native \ --lib device_feature_tests -- --ignored --test-threads=1 cargo test --release --locked -p omni-qwen3-5-native --lib device_vision_graph_ \ -- --ignored --test-threads=1 ``` +The language-Graph regression also runs with the real library: + +```sh +CUA_S1_TEST_MODEL="$CUA_S1_MODEL" cargo test --release --locked \ + -p omni-qwen3-5-native --lib device_multimodal_graph_ \ + -- --ignored --test-threads=1 --nocapture +``` + +It checks changed tokens/features/positions, host/device interleaving after +producer release, scratch growth, and actual recording-error eager fallback. These tests compare feature bits, hidden states, logits, probabilities and complete responses across changed pixels/goals, repeated questions, different grids with equal patch counts, scratch growth and short-request recovery. diff --git a/recipe/omnijev/native.md b/recipe/omnijev/native.md index b6467c99..f8716919 100644 --- a/recipe/omnijev/native.md +++ b/recipe/omnijev/native.md @@ -19,6 +19,16 @@ cargo build --release --locked -p omni-omnijev-native OMNIJEV_MODEL=$PWD/omnijev-export target/release/omni-omnijev-native ``` +`CUA_S1_GRAPH=1` also enables language window Graphs in this worker, through the +shared `Model::readout_window` implementation. It defaults off. Request and +question prefixes and option continuations can capture and replay; current token, +position and image inputs are staged before replay. Graph entries retain their +prefix source and capture destination allocations. The shared 128-entry FIFO and +scratch growth drain dependent work before retiring affected Graphs; recording +failure disables language Graphs and preserves the eager +result. This switch does not enable vision Graphs or Cua's matched-prefix mode. +The Cua-S1 latency measurements do not establish an OmniJev speedup. + `OMNIJEV_HOST` and `OMNIJEV_PORT` default to `127.0.0.1` and `8000`. At start the worker checks every file of the export against its manifest, loads the weights and answers one request with all three question types before it listens, so `/health` @@ -63,6 +73,22 @@ QWEN3_5_CHECKPOINT=$PWD/omnijev-export CUA_S1_CUDA_LIB=$PWD/target/release/libqw cargo test --release --locked -p omni-qwen3-5-native --test readout -- --ignored --test-threads=1 ``` +The direct OmniJev window-Graph regression selects eager and Graph models without +changing process-wide environment variables. It compares raw FP32 head bits and +answers (excluding timing fields), requires actual capture and warm replay, and +covers same-layout new pixels, changed questions/options, geometry and scratch +growth, returning to an earlier request and malformed-input recovery: + +```sh +OMNIJEV_EXPORT=$PWD/omnijev-export OMNIJEV_CUDA_LIB=$PWD/target/release/libqwen3_5_cuda.so \ + cargo test --release --locked -p omni-omnijev-native --lib window_graph_tests \ + -- --ignored --test-threads=1 --nocapture +``` + +Also repeat the `--test worker` command above with `CUA_S1_GRAPH=0` and `1` to +exercise both modes through the production loader. Qualification applies to the +tested model, GPU and CUDA library combination; recheck parity on another stack. + ## Comparison with the reference [`validate.py`](validate.py) runs the reference and the worker on the same images and diff --git a/src/models/cua_s1/README.md b/src/models/cua_s1/README.md index 28a39e20..f27057bc 100644 --- a/src/models/cua_s1/README.md +++ b/src/models/cua_s1/README.md @@ -168,9 +168,11 @@ prompts with the same fixed GEMMs; unset/`off` preserves default execution. The `matched` opt-in instead checks the ordinary full-prompt GEMM algorithms for all questions and supports each prefix/suffix shape before sharing. If plans differ or a window is unsupported, it runs ordinary independent full prompts; -a single question also uses the ordinary path. All active prefix modes reject -language Graph and resident features at startup. Prefix contents -are request-local and branches execute serially. One engine-owned snapshot storage +a single question also uses the ordinary path. `fixed` and `shared` reject +language Graph and resident features at startup. `matched` supports both, using +checked device row windows and bounded Graph capture/replay. Graph entries retain +any recorded source/destination prefix allocations and retire after stream drain. +Prefix contents are request-local and branches execute serially. One engine-owned snapshot storage slot grows as needed; each shared request recaptures it before use. Requests with no shared prefix ignore the retained contents. Prompt construction, the CPU letter head, answer order and token accounting are unchanged. Fixed versus default diff --git a/src/models/cua_s1/native/src/vision_engine.rs b/src/models/cua_s1/native/src/vision_engine.rs index 941a31cb..4cc11f8e 100644 --- a/src/models/cua_s1/native/src/vision_engine.rs +++ b/src/models/cua_s1/native/src/vision_engine.rs @@ -217,28 +217,62 @@ impl VisionEngine { ); if self.prefix_mode == prefix::Mode::Matched { let (lengths, windows) = prefix::matched_windows(&prepared.prompts)?; - let features = self.vision.forward(&prepared.image)?; - let result = self - .language - .with_matched_prefix(&lengths, &windows, |model| { - prefix::execute( - model, - &self.letters, - prepared, - &features, - prefix::Mode::Shared, - &mut self.prefix_storage, - ) - })?; - if let Some(logits) = result { - Ok(logits) + if self.device_features { + let features = self.vision.forward_device(&prepared.image)?; + let result = + self.language + .with_matched_prefix(&lengths, &windows, |model| { + prefix::execute_device( + model, + &self.letters, + prepared, + features, + &mut self.prefix_storage, + ) + })?; + if let Some(logits) = result { + Ok(logits) + } else { + prepared + .prompts + .iter() + .zip(counts) + .map(|(p, n)| { + Self::raw_score_device( + &mut self.language, + &self.letters, + p, + &features, + n, + ) + .map(|(_, logits)| logits) + }) + .collect() + } } else { - prepared - .prompts - .iter() - .zip(counts) - .map(|(p, n)| self.raw_score(p, &features, n).map(|(_, logits)| logits)) - .collect() + let features = self.vision.forward(&prepared.image)?; + let result = + self.language + .with_matched_prefix(&lengths, &windows, |model| { + prefix::execute( + model, + &self.letters, + prepared, + &features, + prefix::Mode::Shared, + &mut self.prefix_storage, + ) + })?; + if let Some(logits) = result { + Ok(logits) + } else { + prepared + .prompts + .iter() + .zip(counts) + .map(|(p, n)| self.raw_score(p, &features, n).map(|(_, logits)| logits)) + .collect() + } } } else if self.prefix_mode != prefix::Mode::Off { let features = self.vision.forward(&prepared.image)?; diff --git a/src/models/cua_s1/native/src/vision_engine/prefix.rs b/src/models/cua_s1/native/src/vision_engine/prefix.rs index f6f0466e..3e27d2cc 100644 --- a/src/models/cua_s1/native/src/vision_engine/prefix.rs +++ b/src/models/cua_s1/native/src/vision_engine/prefix.rs @@ -20,7 +20,7 @@ impl Mode { _ => anyhow::bail!("CUA_S1_MM_PREFIX must be off, fixed, shared, or matched"), }; ensure!( - mode == Self::Off + matches!(mode, Self::Off | Self::Matched) || (!device_features && std::env::var("CUA_S1_GRAPH").as_deref() != Ok("1")), "multimodal prefix modes require Graph and device features off" ); @@ -51,10 +51,16 @@ pub(super) fn common_prefix(prompts: &[ImagePrompt]) -> usize { count / 64 * 64 } +#[derive(Clone, Copy)] +enum Features<'a> { + Host(&'a [bf16]), + Device(DeviceBf16Matrix<'a>), +} + fn window( model: &mut Model, prompt: &ImagePrompt, - features: &[bf16], + features: Features<'_>, range: std::ops::Range, prefix: Option<&PrefixState>, capture: Option<&mut PrefixState>, @@ -78,18 +84,40 @@ fn window( let input = MultimodalInput { token_ids: &prompt.token_ids[range.clone()], image_token_indices: &indices, - image_embeddings: &features[begin * model.cfg.hidden..end * model.cfg.hidden], + image_embeddings: match features { + Features::Host(rows) => &rows[begin * model.cfg.hidden..end * model.cfg.hidden], + Features::Device(_) => &[], + }, position_ids: std::array::from_fn(|axis| &prompt.position_ids[axis][range.clone()]), }; - let mut result = model.readout_window(&input, prefix, capture, &rows, &[])?; + let mut result = if let Features::Device(features) = features { + if let Some(features) = features.slice_rows(begin..end)? { + model.readout_window_device( + &DeviceMultimodalInput { + token_ids: input.token_ids, + image_token_indices: input.image_token_indices, + image_embeddings: features, + position_ids: input.position_ids, + }, + prefix, + capture, + &rows, + &[], + )? + } else { + model.readout_window(&input, prefix, capture, &rows, &[])? + } + } else { + model.readout_window(&input, prefix, capture, &rows, &[])? + }; Ok(result.hidden.pop().unwrap_or_default()) } -pub(super) fn execute( +fn execute_features( model: &mut Model, letters: &[f32], prepared: &PreparedRequest, - features: &[bf16], + features: Features<'_>, mode: Mode, storage: &mut Option, ) -> Result>> { @@ -116,14 +144,20 @@ pub(super) fn execute( .filter_map(|(i, &id)| (id == image_token).then_some(i))), "image indices must exactly match the ordered placeholders" ); - ensure!( - prompt - .image_token_indices - .len() - .checked_mul(model.cfg.hidden) - == Some(features.len()), - "image embedding shape mismatch" - ); + match features { + Features::Host(rows) => ensure!( + prompt + .image_token_indices + .len() + .checked_mul(model.cfg.hidden) + == Some(rows.len()), + "image embedding shape mismatch" + ), + Features::Device(view) => { + view.layout() + .validate(prompt.image_token_indices.len(), model.cfg.hidden, 0)? + } + } } let shared = if mode == Mode::Shared { common_prefix(&prepared.prompts) @@ -190,3 +224,38 @@ pub(super) fn matched_windows(prompts: &[ImagePrompt]) -> Result<(Vec, Ve } Ok((lengths, windows)) } + +pub(super) fn execute( + model: &mut Model, + letters: &[f32], + prepared: &PreparedRequest, + features: &[bf16], + mode: Mode, + storage: &mut Option, +) -> Result>> { + execute_features( + model, + letters, + prepared, + Features::Host(features), + mode, + storage, + ) +} + +pub(super) fn execute_device( + model: &mut Model, + letters: &[f32], + prepared: &PreparedRequest, + features: DeviceBf16Matrix<'_>, + storage: &mut Option, +) -> Result>> { + execute_features( + model, + letters, + prepared, + Features::Device(features), + Mode::Shared, + storage, + ) +} diff --git a/src/models/omnijev/README.md b/src/models/omnijev/README.md index e614f21b..a61a5e6d 100644 --- a/src/models/omnijev/README.md +++ b/src/models/omnijev/README.md @@ -115,6 +115,12 @@ tokens that remain before it, each as a continuation of the state before it (`Model::readout_window`). The passes use the fixed-algorithm GEMMs and split at multiples of 64, so every answer equals that of one plain pass per (question, option) row over the prefix and that row, which `Executor::execute_rows` keeps for the checks. +The shared language executor also honors `CUA_S1_GRAPH=1` for these windows; +it defaults off. Current inputs are staged before replay, and cached Graphs retain +their recorded prefix allocations until synchronized retirement. See the +[native recipe](../../../recipe/omnijev/native.md#tests) for the direct eager/Graph +parity and capture/replay regression. Cua-S1 performance results do not qualify +OmniJev latency. The passes read the final-norm hidden states at `zq` and `u`, which is what the reference's `hidden_states[-1]` holds, and the option-text tokens' log-probabilities: BF16 logits from the tied embedding over the diff --git a/src/models/omnijev/native/src/executor.rs b/src/models/omnijev/native/src/executor.rs index 170034a5..21d09b31 100644 --- a/src/models/omnijev/native/src/executor.rs +++ b/src/models/omnijev/native/src/executor.rs @@ -19,6 +19,10 @@ use crate::export; use crate::heads::{self, Heads}; use crate::processing::{IMAGE_TOKEN, PreparedRequest, Processor, Row}; +#[cfg(test)] +#[path = "../../../../../tests/omnijev/window_graph.rs"] +mod window_graph_tests; + pub struct Executor { vision: VisionModel, language: Model, diff --git a/src/models/qwen3_5/native/src/cuda.rs b/src/models/qwen3_5/native/src/cuda.rs index 612d49da..d1ef3146 100644 --- a/src/models/qwen3_5/native/src/cuda.rs +++ b/src/models/qwen3_5/native/src/cuda.rs @@ -622,6 +622,7 @@ impl DeviceBuffer { /// finish writes before passing the view to a different stream. pub fn bf16_matrix(&self, rows: usize, columns: usize) -> Result> { Ok(DeviceBf16Matrix { + offset: 0, buffer: self, layout: DeviceFeatureLayout::new(rows, columns, self.bytes, self.device)?, }) @@ -674,16 +675,37 @@ impl DeviceFeatureLayout { /// allocation are prevented while this view is alive. #[derive(Clone, Copy)] pub struct DeviceBf16Matrix<'a> { + offset: usize, buffer: &'a DeviceBuffer, layout: DeviceFeatureLayout, } impl DeviceBf16Matrix<'_> { + /// Checked contiguous rows retaining the original producer borrow. Empty + /// ranges have no device view and are used by text-only suffix windows. + pub fn slice_rows(self, range: std::ops::Range) -> Result> { + ensure!( + range.start <= range.end && range.end <= self.layout.rows, + "device row slice outside producer" + ); + if range.is_empty() { + return Ok(None); + } + Ok(Some(Self { + buffer: self.buffer, + offset: self.offset + range.start * self.layout.columns * 2, + layout: DeviceFeatureLayout { + rows: range.end - range.start, + ..self.layout + }, + })) + } + pub fn layout(&self) -> DeviceFeatureLayout { self.layout } pub(crate) fn at(&self, offset: usize) -> *const c_void { - self.buffer.at(offset) + self.buffer.at(self.offset + offset) } pub(crate) fn validate_finite(&self, stream: Stream) -> Result<()> { let flag = DeviceBuffer::new(4)?; diff --git a/src/models/qwen3_5/native/src/model.rs b/src/models/qwen3_5/native/src/model.rs index 05f8c6cb..b29dda10 100644 --- a/src/models/qwen3_5/native/src/model.rs +++ b/src/models/qwen3_5/native/src/model.rs @@ -13,6 +13,7 @@ //! once and every branch from the state after it (see `SharedPrompts`). mod matched; +mod window_graph; use std::collections::{HashMap, VecDeque}; use std::ffi::c_void; @@ -695,14 +696,19 @@ pub struct PrefixState { len: usize, /// Tokens the buffers hold; `forward_shared` reuses a state for shorter prefixes. cap: usize, + buffers: Arc, + /// Total device bytes, for cache-budget accounting by the caller. + bytes: usize, +} + +/// Shared ownership keeps captured device addresses alive until Graph retirement. +struct PrefixBuffers { /// Post-prep K [len, Hk*Dh] and raw V [len, Hk*Dh] rows per full-attention layer. attn_kv: Vec<(DeviceBuffer, DeviceBuffer)>, /// Float32 [H, K, V] recurrent states, one per Gated DeltaNet layer. gdn_state: Vec, /// Three pre-conv projection columns [3, gdn_in width], one per Gated DeltaNet layer. conv_tail: Vec, - /// Total device bytes, for cache-budget accounting by the caller. - bytes: usize, } /// A packed suffix batch uses one freshly captured prefix and original algorithm. @@ -867,6 +873,7 @@ pub struct Model { graph_stats: GraphStats, graphs: VecDeque<(Vec, cuda::Graph)>, mm_graphs: VecDeque<(Vec, cuda::Graph)>, + window_graphs: VecDeque, // The stream retires after every model allocation, including adapter resources. _stream_owner: ModelStream, } @@ -889,6 +896,7 @@ impl Drop for Model { let _ = cuda::synchronize(self.stream); self.graphs.clear(); self.mm_graphs.clear(); + self.window_graphs.clear(); self.scratch = None; self.shared = [None, None]; // Retire the FP32 cuBLAS context while its last stream is still valid. @@ -1043,6 +1051,7 @@ impl Model { graphs: VecDeque::new(), _stream_owner: stream_owner, mm_graphs: VecDeque::new(), + window_graphs: VecDeque::new(), }; Ok(model) } @@ -1050,7 +1059,7 @@ impl Model { pub fn graph_stats(&self) -> GraphStats { GraphStats { enabled: self.graph_enabled, - cached_shapes: self.graphs.len() + self.mm_graphs.len(), + cached_shapes: self.graphs.len() + self.mm_graphs.len() + self.window_graphs.len(), ..self.graph_stats } } @@ -1332,11 +1341,15 @@ impl Model { .next_multiple_of(1024); if self.scratch.as_ref().is_none_or(|s| capacity > s.cap) { cuda::synchronize(self.stream)?; - if !self.graphs.is_empty() || !self.mm_graphs.is_empty() { + if !self.graphs.is_empty() + || !self.mm_graphs.is_empty() + || !self.window_graphs.is_empty() + { self.graph_stats.invalidations += 1; } self.graphs.clear(); self.mm_graphs.clear(); + self.window_graphs.clear(); self.scratch = None; self.scratch = Some(Scratch::new( &self.cfg, @@ -1686,6 +1699,13 @@ impl Model { self.upload_positions(s, input.position_ids)?; self.embed_tokens(s, input.token_ids)?; self.overwrite_image_rows(s, input, 0)?; + self.run_multimodal(t) + } + + /// Inputs are staged in owned scratch before capture; producer addresses + /// never enter the graph, so host and device inputs can share a capture. + fn run_multimodal(&mut self, t: usize) -> Result> { + let s = self.scratch.as_ref().unwrap(); if self.graph_enabled { if let Some((_, graph)) = self .mm_graphs @@ -1753,6 +1773,7 @@ impl Model { self.graph_enabled = false; self.graphs.clear(); self.mm_graphs.clear(); + self.window_graphs.clear(); } } Ok(()) @@ -1794,7 +1815,7 @@ impl Model { targets: &[(usize, u32)], ) -> Result { self.ensure_fixed_profile()?; - let result = self.run_readout_window(window, prefix, capture, rows, targets); + let result = self.run_readout_window(window, prefix, capture, rows, targets, None); if result.is_err() { // Work already queued may still use the buffers. let _ = cuda::synchronize(self.stream); @@ -1802,6 +1823,37 @@ impl Model { result } + /// Window readout from checked borrowed device rows. Current input is copied + /// into model scratch before capture/replay; producer addresses are never + /// recorded. Every exit drains consumer work before the borrow is released. + pub fn readout_window_device( + &mut self, + window: &DeviceMultimodalInput<'_>, + prefix: Option<&PrefixState>, + capture: Option<&mut PrefixState>, + rows: &[usize], + targets: &[(usize, u32)], + ) -> Result { + self.ensure_fixed_profile()?; + let host = MultimodalInput { + token_ids: window.token_ids, + image_token_indices: window.image_token_indices, + image_embeddings: &[], + position_ids: window.position_ids, + }; + let result = self.run_readout_window( + &host, + prefix, + capture, + rows, + targets, + Some(window.image_embeddings), + ); + let synced = self.synchronize(); + synced?; + result + } + fn run_readout_window( &mut self, window: &MultimodalInput<'_>, @@ -1809,13 +1861,31 @@ impl Model { mut capture: Option<&mut PrefixState>, rows: &[usize], targets: &[(usize, u32)], + device: Option>, ) -> Result { let image_token = self .cfg .image_token_id .context("checkpoint has no image_token_id")?; let vocab = self.embed.shape[0]; - window.validate(self.cfg.hidden, vocab, image_token, self.cfg.max_positions)?; + if let Some(features) = device { + DeviceMultimodalInput { + token_ids: window.token_ids, + image_token_indices: window.image_token_indices, + image_embeddings: features, + position_ids: window.position_ids, + } + .validate( + self.cfg.hidden, + vocab, + image_token, + self.cfg.max_positions, + 0, + )?; + features.validate_finite(self.stream)?; + } else { + window.validate(self.cfg.hidden, vocab, image_token, self.cfg.max_positions)?; + } if let Some(pre) = prefix { ensure!( Arc::ptr_eq(&self.prefix_owner, &pre.owner), @@ -1824,7 +1894,9 @@ impl Model { ensure!(pre.initialized, "prefix state has not completed capture"); } let qb = prefix.map_or(0, |pre| pre.len); - let tend = qb + window.token_ids.len(); + let tend = qb + .checked_add(window.token_ids.len()) + .context("window length overflow")?; ensure!( tend <= self.cfg.max_positions, "prompt exceeds the configured maximum length" @@ -1851,24 +1923,21 @@ impl Model { let s = self.scratch.as_ref().unwrap(); self.upload_positions(s, window.position_ids)?; self.embed_tokens_at(s, window.token_ids, qb)?; - self.overwrite_image_rows(s, window, qb)?; + if let Some(features) = device { + self.overwrite_device_rows(s, window.image_token_indices, features, qb)?; + } else { + self.overwrite_image_rows(s, window, qb)?; + } if let Some(cap) = capture.as_deref_mut() { cap.initialized = false; cap.len = tend; } - self.run_window( - s, - qb, - tend, - true, - capture.as_deref_mut(), - prefix, - gemm, - None, - )?; + self.run_window_graph(qb, tend, capture.as_deref_mut(), prefix, gemm)?; + if let Some(cap) = capture { cap.initialized = true; } + let s = self.scratch.as_ref().unwrap(); let hidden = rows .iter() .map(|&r| self.hidden_row(s, r)) @@ -1946,7 +2015,7 @@ impl Model { Ok(logprobs) } - /// Eager multimodal execution using borrowed device-resident image rows. + /// Multimodal execution using borrowed device-resident image rows. /// The producer must finish writes before entry. All queued consumer work is /// completed on every exit path, before the external allocation can be reused. pub fn forward_multimodal_device( @@ -1972,34 +2041,43 @@ impl Model { let s = self.scratch.as_ref().unwrap(); self.upload_positions(s, input.position_ids)?; self.embed_tokens(s, input.token_ids)?; - let indices = input.image_token_indices; - let row_bytes = self.cfg.hidden * BF16; - let mut begin = 0; - while begin < indices.len() { - let mut end = begin + 1; - while end < indices.len() && indices[end] == indices[end - 1] + 1 { - end += 1; - } - // SAFETY: metadata validates source rows, device and allocation - // extent and indices exactly match placeholders inside s.res. - unsafe { - cuda::copy_dd( - s.at(s.res + indices[begin] * row_bytes), - input.image_embeddings.at(begin * row_bytes), - (end - begin) * row_bytes, - self.stream, - )?; - } - begin = end; - } - self.run(s, &[t], true, false)?; - self.last_hidden(s, t) + self.overwrite_device_rows(s, input.image_token_indices, input.image_embeddings, 0)?; + self.run_multimodal(t) })(); let synced = self.synchronize(); synced?; result } + fn overwrite_device_rows( + &self, + s: &Scratch, + indices: &[usize], + features: cuda::DeviceBf16Matrix<'_>, + qb: usize, + ) -> Result<()> { + let row_bytes = self.cfg.hidden * BF16; + let mut begin = 0; + while begin < indices.len() { + let mut end = begin + 1; + while end < indices.len() && indices[end] == indices[end - 1] + 1 { + end += 1; + } + // SAFETY: validated metadata binds the source row extent, device and + // ordered placeholders, with qb+indices inside owned scratch. + unsafe { + cuda::copy_dd( + s.at(s.res + (qb + indices[begin]) * row_bytes), + features.at(begin * row_bytes), + (end - begin) * row_bytes, + self.stream, + )?; + } + begin = end; + } + Ok(()) + } + fn upload_positions(&self, s: &Scratch, positions: [&[i64]; 3]) -> Result<()> { let (cos, sin) = rotary_tables( positions, @@ -2090,9 +2168,11 @@ impl Model { initialized: false, len, cap: len, - attn_kv, - gdn_state, - conv_tail, + buffers: Arc::new(PrefixBuffers { + attn_kv, + gdn_state, + conv_tail, + }), bytes: self.prefix_bytes(len), }) } @@ -2863,7 +2943,7 @@ impl Model { if let Some(suffix) = suffix { (suffix.ops.conv)( p(s.gdn_in + offset * w.gdn_in * BF16), - suffix.prefix.conv_tail[linear_index].at(0), + suffix.prefix.buffers.conv_tail[linear_index].at(0), ld, la.conv.ptr, p(s.lq + offset * kd * BF16), @@ -2921,7 +3001,7 @@ impl Model { hv as i32, cfg.lin_k_heads as i32, (cfg.lin_k_dim as f32).powf(-0.5), - suffix.prefix.gdn_state[linear_index].at(0), + suffix.prefix.buffers.gdn_state[linear_index].at(0), std::ptr::null_mut(), st, ), @@ -3050,8 +3130,8 @@ impl Model { ld, p(s.agate + offset * cfg.heads * cfg.head_dim * BF16), p(s.ao + offset * cfg.heads * cfg.head_dim * BF16), - suffix.prefix.attn_kv[attention_index].0.at(0), - suffix.prefix.attn_kv[attention_index].1.at(0), + suffix.prefix.buffers.attn_kv[attention_index].0.at(0), + suffix.prefix.buffers.attn_kv[attention_index].1.at(0), length, suffix.prefix.len as i32, hq, @@ -3309,15 +3389,16 @@ impl Model { let b = z + vd * BF16; let a = b + hv * BF16; let s_in: *const c_void = - prefix.map_or(std::ptr::null(), |pre| pre.gdn_state[la_i].at(0)); - let s_out: *mut c_void = capture - .as_deref() - .map_or(std::ptr::null_mut(), |cap| cap.gdn_state[la_i].at(0)); + prefix.map_or(std::ptr::null(), |pre| pre.buffers.gdn_state[la_i].at(0)); + let s_out: *mut c_void = + capture.as_deref().map_or(std::ptr::null_mut(), |cap| { + cap.buffers.gdn_state[la_i].at(0) + }); unsafe { if qb > 0 { cuda::copy_dd( p(s.gdn_in + (qb - 3) * ldbb), - prefix.unwrap().conv_tail[la_i].at(0), + prefix.unwrap().buffers.conv_tail[la_i].at(0), 3 * ldbb, st, )?; @@ -3326,7 +3407,7 @@ impl Model { // conv window tail: the last three pre-conv columns, which // include restored prefix columns for a window under three rows. cuda::copy_dd( - cap.conv_tail[la_i].at(0), + cap.buffers.conv_tail[la_i].at(0), p(s.gdn_in + (tend - 3) * ldbb), 3 * ldbb, st, @@ -3435,8 +3516,21 @@ impl Model { let pre = prefix.unwrap(); // Restore the prefix V columns and post-prep K rows in place; // the windowed flash then sees the identical k/v layout. - cuda::copy2d(p(v), ab, pre.attn_kv[fa_i].1.at(0), kvb, kvb, qb, st)?; - cuda::copy_dd(p(s.ak), pre.attn_kv[fa_i].0.at(0), qb * kvb, st)?; + cuda::copy2d( + p(v), + ab, + pre.buffers.attn_kv[fa_i].1.at(0), + kvb, + kvb, + qb, + st, + )?; + cuda::copy_dd( + p(s.ak), + pre.buffers.attn_kv[fa_i].0.at(0), + qb * kvb, + st, + )?; } check( (self @@ -3464,8 +3558,21 @@ impl Model { )?; if let Some(cap) = capture.as_deref() { // K must be post-prep (normed + rotary); V is the raw column. - cuda::copy_dd(cap.attn_kv[fa_i].0.at(0), p(s.ak), tend * kvb, st)?; - cuda::copy2d(cap.attn_kv[fa_i].1.at(0), kvb, p(v), ab, kvb, tend, st)?; + cuda::copy_dd( + cap.buffers.attn_kv[fa_i].0.at(0), + p(s.ak), + tend * kvb, + st, + )?; + cuda::copy2d( + cap.buffers.attn_kv[fa_i].1.at(0), + kvb, + p(v), + ab, + kvb, + tend, + st, + )?; } check( (cuda::api().cs1_attention_gated_prefix)( diff --git a/src/models/qwen3_5/native/src/model/matched.rs b/src/models/qwen3_5/native/src/model/matched.rs index faa3ed9e..15253734 100644 --- a/src/models/qwen3_5/native/src/model/matched.rs +++ b/src/models/qwen3_5/native/src/model/matched.rs @@ -38,6 +38,7 @@ impl Model { let m = lengths[0]; if !self.matched_handles.contains_key(&m) && self.matched_handles.len() == 16 { self.synchronize()?; + self.window_graphs.clear(); // SAFETY: all launches drained before destroying owned workspaces. for (_, h) in self.matched_handles.drain() { unsafe { (cuda::api().cs1_gemm_destroy)(h) }; @@ -87,11 +88,17 @@ impl Model { ) }; if okay != 1 { + if std::env::var("CUA_S1_GRAPH_TRACE").as_deref() == Ok("1") { + eprintln!("CUA_S1_PREFIX event=fallback reference={m}"); + } return Ok(None); } } } } + if std::env::var("CUA_S1_GRAPH_TRACE").as_deref() == Ok("1") { + eprintln!("CUA_S1_PREFIX event=admit reference={m}"); + } self.matched_ref = m; // The worker catches panics and can retain this engine. Restore the // ordinary selection and drain submitted work on every exit path. diff --git a/src/models/qwen3_5/native/src/model/window_graph.rs b/src/models/qwen3_5/native/src/model/window_graph.rs new file mode 100644 index 00000000..f9ad7760 --- /dev/null +++ b/src/models/qwen3_5/native/src/model/window_graph.rs @@ -0,0 +1,89 @@ +//! Window Graphs retain every prefix allocation whose address they record. +use super::*; + +pub(super) struct Entry { + key: [usize; 5], + // Destroy the Graph before releasing recorded allocations. + graph: cuda::Graph, + _states: [Option>; 2], +} + +impl Model { + pub(super) fn run_window_graph( + &mut self, + begin: usize, + end: usize, + mut capture: Option<&mut PrefixState>, + prefix: Option<&PrefixState>, + gemm: *mut c_void, + ) -> Result<()> { + let s = self.scratch.as_ref().unwrap(); + if !self.graph_enabled { + return self.run_window(s, begin, end, true, capture, prefix, gemm, None); + } + let states = [ + prefix.map(|p| Arc::clone(&p.buffers)), + capture.as_deref().map(|p| Arc::clone(&p.buffers)), + ]; + let identities = states + .each_ref() + .map(|p| p.as_ref().map_or(0, |p| Arc::as_ptr(p) as usize)); + let key = [begin, end, identities[0], identities[1], gemm as usize]; + if let Some(entry) = self.window_graphs.iter().find(|entry| entry.key == key) { + entry.graph.launch(self.stream)?; + self.graph_stats.replays += 1; + if std::env::var("CUA_S1_GRAPH_TRACE").as_deref() == Ok("1") { + eprintln!("CUA_S1_GRAPH mode=window event=replay begin={begin} end={end}"); + } + return Ok(()); + } + // Preserve the eager miss result; capture records without executing. + self.run_window( + s, + begin, + end, + true, + capture.as_deref_mut(), + prefix, + gemm, + None, + )?; + cuda::synchronize(self.stream)?; + let recorded = cuda::Graph::capture(self.stream, || { + self.run_window(s, begin, end, true, capture, prefix, gemm, None) + }); + self.cache_window_graph(key, states, recorded) + } + + pub(super) fn cache_window_graph( + &mut self, + key: [usize; 5], + states: [Option>; 2], + recorded: Result, + ) -> Result<()> { + let [begin, end, ..] = key; + match recorded { + Ok(graph) => { + if self.window_graphs.len() == 128 { + cuda::synchronize(self.stream)?; + self.window_graphs.pop_front(); + } + self.window_graphs.push_back(Entry { + key, + graph, + _states: states, + }); + self.graph_stats.captures += 1; + if std::env::var("CUA_S1_GRAPH_TRACE").as_deref() == Ok("1") { + eprintln!("CUA_S1_GRAPH mode=window event=capture begin={begin} end={end}"); + } + } + Err(error) => { + // Existing retirement path drains, disables all language Graph + // modes and retains the already computed eager miss result. + self.cache_graph(&[begin, end], true, Err(error))?; + } + } + Ok(()) + } +} diff --git a/tests/cua_s1/device_controls.rs b/tests/cua_s1/device_controls.rs index 7d5b839f..bcfec86d 100644 --- a/tests/cua_s1/device_controls.rs +++ b/tests/cua_s1/device_controls.rs @@ -66,18 +66,6 @@ fn multimodal_prefix_controls_reject_invalid_and_incompatible_modes_before_loadi "0", "multimodal prefix modes require Graph and device features off", ), - ( - "matched", - "1", - "0", - "multimodal prefix modes require Graph and device features off", - ), - ( - "matched", - "0", - "1", - "multimodal prefix modes require Graph and device features off", - ), ( "fixed", "0", @@ -153,3 +141,27 @@ fn non_unicode_prefix_is_rejected_before_loading() { .contains("CUA_S1_MM_PREFIX must be off, fixed, shared, or matched") ); } + +#[test] +fn matched_prefix_allows_host_device_and_language_graph_combinations() { + for graph in ["0", "1"] { + for device in ["0", "1"] { + let output = std::process::Command::new(env!("CARGO_BIN_EXE_omni-cua-s1-vision")) + .env("CUA_S1_MM_PREFIX", "matched") + .env("CUA_S1_GRAPH", graph) + .env("CUA_S1_DEVICE_FEATURES", device) + .env("CUA_S1_MODEL", "/nonexistent-cua-prefix-model") + .env("CUA_S1_BASE", "/nonexistent-cua-prefix-model") + .env("CUA_S1_VISION_ADAPTER", "/nonexistent-cua-prefix-model") + .env("CUA_S1_CUDA_LIB", "/nonexistent-cua-prefix-model") + .output() + .unwrap(); + let error = String::from_utf8(output.stderr).unwrap(); + assert!(!output.status.success()); + assert!( + error.contains("export the multimodal language checkpoint"), + "graph={graph} device={device}: {error}" + ); + } + } +} diff --git a/tests/cua_s1/multimodal_prefix.rs b/tests/cua_s1/multimodal_prefix.rs index a401037d..fd1022de 100644 --- a/tests/cua_s1/multimodal_prefix.rs +++ b/tests/cua_s1/multimodal_prefix.rs @@ -231,3 +231,84 @@ fn matched_prefix_preserves_default_and_original_m64_after_errors() { } assert!(admitted > 0, "all shared-window probes fell back"); } + +#[test] +#[ignore = "requires pinned CUDA Cua-S1 artifacts; run serially with vision Graph0/1"] +fn resident_matched_windows_equal_ordinary_with_changed_requests_and_error_recovery() { + let mut engine = device_feature_tests::load(); + let mut resident_admissions = 0; + for (width, height, seed, count) in [ + (256, 256, 1, 2), + (512, 256, 2, 4), + (256, 512, 3, 8), + (1024, 1024, 4, 2), + (256, 256, 1, 1), + (256, 256, 7, 4), + ] { + let rgb: Vec<_> = (0..width * height * 3) + .map(|i| ((i * 17 + seed * 31) % 251) as u8) + .collect(); + let questions: Vec<_> = (0..count) + .map(|i| { + let mut q = device_feature_tests::questions(seed + i).remove(i % 2); + q.name = format!("q{i}"); + q + }) + .collect(); + let mut prepared = engine.prepare(width, height, &rgb, &questions).unwrap(); + engine.prefix_mode = prefix::Mode::Off; + engine.device_features = false; + let expected = engine.execute_prepared(&prepared).unwrap(); + let bits = |v: Vec>| { + v.into_iter() + .flatten() + .map(f32::to_bits) + .collect::>() + }; + engine.prefix_mode = prefix::Mode::Matched; + for device in [true, false, true] { + let had_storage = engine.prefix_storage.is_some(); + let before = engine.language.graph_stats(); + engine.device_features = device; + assert_eq!( + bits(engine.execute_prepared(&prepared).unwrap()), + bits(expected.clone()) + ); + if device && !had_storage && engine.prefix_storage.is_some() { + resident_admissions += 1; + if before.requested { + assert!( + before.enabled, + "requested Graph was disabled before resident admission" + ); + assert_eq!(before.fallbacks, 0, "unexpected prior capture failure"); + assert!( + engine.language.graph_stats().captures > before.captures, + "first admitted resident request must capture windows" + ); + let before = engine.language.graph_stats(); + assert_eq!( + bits(engine.execute_prepared(&prepared).unwrap()), + bits(expected.clone()) + ); + assert!( + engine.language.graph_stats().replays > before.replays, + "same admitted resident request must replay windows" + ); + } + } + } + let token = *prepared.prompts[0].token_ids.last().unwrap(); + *prepared.prompts[0].token_ids.last_mut().unwrap() = u32::MAX; + assert!(engine.execute_prepared(&prepared).is_err()); + *prepared.prompts[0].token_ids.last_mut().unwrap() = token; + assert_eq!( + bits(engine.execute_prepared(&prepared).unwrap()), + bits(expected) + ); + } + assert!( + resident_admissions > 0, + "resident production routing never admitted prefix reuse" + ); +} diff --git a/tests/omnijev/window_graph.rs b/tests/omnijev/window_graph.rs new file mode 100644 index 00000000..625feae9 --- /dev/null +++ b/tests/omnijev/window_graph.rs @@ -0,0 +1,184 @@ +//! Exercise OmniJev's production prefix/option routing with language Graphs. +use super::*; + +#[path = "common.rs"] +mod common; + +// Replace the model only after dropping its prefix states and the old model. +// This selects both modes without mutating process-wide environment variables. +fn language_mode(executor: Executor, dir: &Path, library: &Path, graph: bool) -> Executor { + let Executor { + vision, + language, + heads, + calibration, + request, + question, + } = executor; + drop((request, question, language)); + Executor { + vision, + language: Model::load_with_graph(dir, library, graph).unwrap(), + heads, + calibration, + request: None, + question: None, + } +} + +fn request( + processor: &Processor, + width: u32, + height: u32, + image_changed: bool, + question_changed: bool, + growth: bool, +) -> PreparedRequest { + let colour = if image_changed { + [30, 180, 90] + } else { + [200, 40, 40] + }; + let instruction = if growth { + format!( + "Look{} at the image. Which colour dominates?", + " again".repeat(1200) + ) + } else if question_changed { + "Which colour occupies most of this picture?".into() + } else { + "Which colour dominates the image?".into() + }; + let options = if question_changed { + json!([{"key": "green", "text": "green"}, {"key": "red", "text": "red"}, + {"key": "blue", "text": "blue"}]) + } else { + json!([{"key": "red", "text": "red"}, {"key": "green", "text": "green"}]) + }; + let raw = serde_json::to_vec(&json!({ + "model": contract::MODEL_ID, + "state": {"images": [common::png_url(width, height, colour)]}, + "questions": { + "colour": {"type": "choice", "instructions": instruction, "options": options}, + "red": {"type": "noul", "instructions": "The image is mostly red."}, + "warm": {"type": "score", "instructions": "How warm is the colour?", + "levels": ["cold", "neutral", "warm", "hot"]} + } + })) + .unwrap(); + processor.prepare(&raw).unwrap() +} + +fn bits(outputs: &[Vec]) -> Vec { + outputs.iter().flatten().map(|v| v.to_bits()).collect() +} + +fn decisions(mut answers: Vec) -> Vec { + for answer in &mut answers { + let answer = answer.as_object_mut().unwrap(); + answer.remove("latency_s"); + answer.remove("latency_total_s"); + } + answers +} + +#[test] +#[ignore = "requires CUDA, OMNIJEV_EXPORT and OMNIJEV_CUDA_LIB; run serially"] +fn window_graph_matches_eager_across_requests_growth_and_recovery() { + let dir = + std::path::PathBuf::from(std::env::var_os("OMNIJEV_EXPORT").expect("set OMNIJEV_EXPORT")); + let library = std::path::PathBuf::from( + std::env::var_os("OMNIJEV_CUDA_LIB").expect("set OMNIJEV_CUDA_LIB"), + ); + let (executor, processor) = Executor::load(&dir, &library).unwrap(); + let cases = [ + request(&processor, 320, 200, false, false, false), + request(&processor, 320, 200, true, false, false), + request(&processor, 320, 200, true, true, false), + request(&processor, 1280, 720, false, false, true), + ]; + let mut eager = language_mode(executor, &dir, &library, false); + let expected: Vec<_> = cases + .iter() + .map(|prepared| { + let shared = eager.head_outputs(prepared, true).unwrap(); + let plain = eager.head_outputs(prepared, false).unwrap(); + assert_eq!(bits(&shared), bits(&plain), "eager prefix/row parity"); + (bits(&shared), decisions(eager.execute(prepared).unwrap())) + }) + .collect(); + assert_ne!( + expected[0].0, expected[1].0, + "changed request must change head outputs" + ); + assert_eq!(eager.language.graph_stats().captures, 0); + let mut graph = language_mode(eager, &dir, &library, true); + assert_eq!(cases[0].inputs[0].token_ids, cases[1].inputs[0].token_ids); + let mut previous_captures = 0; + for (visit, case) in [0, 1, 2, 0, 3, 0].into_iter().enumerate() { + let prepared = &cases[case]; + let before = graph.language.graph_stats(); + let actual = bits(&graph.head_outputs(prepared, true).unwrap()); + assert_eq!( + actual, expected[case].0, + "Graph/eager head parity, visit {visit}" + ); + let captured = graph.language.graph_stats(); + assert!(captured.enabled, "Graph unexpectedly disabled"); + assert!( + captured.captures > 0, + "OmniJev windows must capture a Graph" + ); + assert_eq!(captured.fallbacks, 0, "unexpected capture failure"); + if case == 1 { + assert_eq!( + captured.captures, previous_captures, + "same-layout new pixels must reuse Graphs" + ); + } + // Growth can retire the request-prefix Graph midway through the first + // pass. Stabilize once before asserting warm-cache reuse. + assert_eq!(bits(&graph.head_outputs(prepared, true).unwrap()), actual); + let stable = graph.language.graph_stats(); + assert_eq!(bits(&graph.head_outputs(prepared, true).unwrap()), actual); + let replayed = graph.language.graph_stats(); + assert!( + replayed.replays > stable.replays, + "warm OmniJev windows must replay" + ); + assert_eq!( + replayed.captures, stable.captures, + "warm windows must reuse captures" + ); + let answers = decisions(graph.execute(prepared).unwrap()); + assert_eq!( + answers, expected[case].1, + "Graph/eager decisions, visit {visit}" + ); + if case == 3 { + assert!( + replayed.invalidations > before.invalidations, + "large request must retire scratch Graphs" + ); + } + previous_captures = graph.language.graph_stats().captures; + eprintln!( + "OMNIJEV_WINDOW_GRAPH {}", + json!({ + "visit": visit, "case": case, "head_bits": actual, "answers": answers, + "captures": replayed.captures, "replays": replayed.replays, + "invalidations": replayed.invalidations, "fallbacks": replayed.fallbacks + }) + ); + } + let mut malformed = request(&processor, 320, 200, false, false, false); + malformed.inputs[0].token_ids[0] = u32::MAX; + assert!(graph.head_outputs(&malformed, true).is_err()); + assert_eq!( + bits(&graph.head_outputs(&cases[0], true).unwrap()), + expected[0].0 + ); + assert_eq!(decisions(graph.execute(&cases[0]).unwrap()), expected[0].1); + assert!(graph.language.graph_stats().enabled); + assert_eq!(graph.language.graph_stats().fallbacks, 0); +} diff --git a/tests/qwen3_5/device_cuda.rs b/tests/qwen3_5/device_cuda.rs index 39553f75..dfefe1e1 100644 --- a/tests/qwen3_5/device_cuda.rs +++ b/tests/qwen3_5/device_cuda.rs @@ -31,3 +31,47 @@ fn device_finite_scan_rejects_all_nonfinite_encodings_and_recovers() { (api().cs1_stream_destroy)(stream); } } + +#[test] +#[ignore = "requires CUDA and CUA_S1_CUDA_LIB; run serially"] +fn device_row_subviews_check_extent_offset_and_only_scan_selected_rows() { + load(&std::path::PathBuf::from( + std::env::var_os("CUA_S1_CUDA_LIB").unwrap(), + )) + .unwrap(); + set_device(0).unwrap(); + let stream = new_stream().unwrap(); + let buffer = DeviceBuffer::new(24).unwrap(); + let values: Vec<_> = (0..12).map(|i| half::bf16::from_f32(i as f32)).collect(); + let bytes: Vec<_> = values.iter().flat_map(|v| v.to_le_bytes()).collect(); + unsafe { + upload(buffer.at(0), &bytes, stream).unwrap(); + } + synchronize(stream).unwrap(); + let full = buffer.bf16_matrix(4, 3).unwrap(); + let middle = full.slice_rows(1..3).unwrap().unwrap(); + assert_eq!(middle.download().unwrap(), values[3..9]); + assert_eq!( + middle + .slice_rows(1..2) + .unwrap() + .unwrap() + .download() + .unwrap(), + values[6..9] + ); + assert!(full.slice_rows(4..4).unwrap().is_none()); + let reversed = std::ops::Range { start: 3, end: 2 }; + for range in [reversed, 0..5, 5..5] { + assert!(full.slice_rows(range).is_err()); + } + unsafe { + upload(buffer.at(0), &half::bf16::NAN.to_le_bytes(), stream).unwrap(); + } + assert!(full.validate_finite(stream).is_err()); + middle.validate_finite(stream).unwrap(); + synchronize(stream).unwrap(); + unsafe { + (api().cs1_stream_destroy)(stream); + } +} diff --git a/tests/qwen3_5/multimodal_graph.rs b/tests/qwen3_5/multimodal_graph.rs index 7b9f7771..68f88053 100644 --- a/tests/qwen3_5/multimodal_graph.rs +++ b/tests/qwen3_5/multimodal_graph.rs @@ -364,3 +364,348 @@ fn graph_recording_failure_keeps_eager_result_and_disables_both_modes() -> Resul } Ok(()) } + +#[test] +#[ignore = "requires the Cua-S1 language checkpoint and real CUDA library; run serially"] +fn device_multimodal_graph_uses_current_rows_and_releases_external_storage() -> Result<()> { + let mut model = load_model()?; + let image_token = model.cfg.image_token_id.context("image_token_id")?; + let text = text_tokens(&model, 4)?; + let mut last = Vec::new(); + for seed in 0..4 { + let ids = [text[seed], image_token, image_token, text[(seed + 1) % 4]]; + let indices = [1, 2]; + let temporal = [0, 1 + seed as i64, 1 + seed as i64, 5]; + let height = [0, 1, 2 + seed as i64, 5]; + let width = [0, 2 + seed as i64, 1, 5]; + let features: Vec = (0..2 * model.cfg.hidden) + .map(|i| bf16::from_f32(((i + seed * 7) % 19) as f32 / 18.0 - 0.5)) + .collect(); + let input = MultimodalInput { + token_ids: &ids, + image_token_indices: &indices, + image_embeddings: &features, + position_ids: [&temporal, &height, &width], + }; + model.graph_enabled = false; + let expected = model.forward_multimodal(&input)?; + if !last.is_empty() { + assert_ne!(expected, last, "changed input must exercise new data"); + } + let bytes: Vec<_> = features.iter().flat_map(|v| v.to_le_bytes()).collect(); + let source = DeviceBuffer::new(bytes.len())?; + unsafe { cuda::upload(source.at(0), &bytes, model.stream)? }; + model.synchronize()?; + let device = DeviceMultimodalInput { + token_ids: &ids, + image_token_indices: &indices, + image_embeddings: source.bf16_matrix(2, model.cfg.hidden)?, + position_ids: input.position_ids, + }; + model.graph_enabled = true; + assert_hidden_eq( + &model.forward_multimodal_device(&device)?, + &expected, + "device capture/replay", + ); + assert_eq!( + model.mm_graphs.len(), + 1, + "device path must cache the shared MM graph" + ); + assert_eq!( + model.graph_stats().captures, + 1, + "same shape must not recapture" + ); + assert_eq!(model.graph_stats().replays, 2 * seed as u64); + // A capture must not retain the external producer address. Releasing it + // and using host input must replay the same graph safely. + drop(source); + assert_hidden_eq( + &model.forward_multimodal(&input)?, + &expected, + "host replay after producer release", + ); + assert_eq!(model.graph_stats().replays, 2 * seed as u64 + 1); + last = expected; + } + model.prepare_scratch(model.scratch.as_ref().unwrap().cap + 1)?; + assert!(model.mm_graphs.is_empty(), "growth must retire captures"); + let ids = [text[0], image_token, image_token, text[1]]; + let positions = [0, 1, 1, 2]; + let source = DeviceBuffer::new(2 * model.cfg.hidden * BF16)?; + let bytes = vec![0; 2 * model.cfg.hidden * BF16]; + unsafe { cuda::upload(source.at(0), &bytes, model.stream)? }; + model.synchronize()?; + let device = DeviceMultimodalInput { + token_ids: &ids, + image_token_indices: &[1, 2], + image_embeddings: source.bf16_matrix(2, model.cfg.hidden)?, + position_ids: [&positions; 3], + }; + let captured = model.forward_multimodal_device(&device)?; + assert_eq!(model.mm_graphs.len(), 1); + assert_eq!(model.graph_stats().captures, 2); + let failed = cuda::Graph::capture(model.stream, || { + anyhow::bail!("injected device record failure") + }); + assert!(failed.is_err()); + model.cache_graph(&[4], true, failed)?; + assert!(!model.graph_enabled); + assert!(model.mm_graphs.is_empty()); + assert_hidden_eq( + &model.forward_multimodal_device(&device)?, + &captured, + "device eager recovery", + ); + assert!(!model.graph_enabled); + Ok(()) +} + +#[test] +#[ignore = "requires pinned Cua-S1 language weights and CUDA; run serially"] +fn device_window_graph_stages_current_rows_and_owns_prefix_storage_until_retirement() -> Result<()> +{ + let mut model = load_model()?; + let image = model.cfg.image_token_id.context("image token")?; + let text = text_tokens(&model, 132)?; + let mut ids = text.clone(); + ids[32] = image; + ids[33] = image; + let indices = [32, 33]; + let mut positions: [Vec; 3] = std::array::from_fn(|_| (0..132).collect()); + let mut state = model.alloc_prefix(128)?; + let mut prior = Vec::new(); + for seed in 0..3 { + ids[0] = text[seed]; + positions[seed][30] += 1; + let features: Vec<_> = (0..2 * model.cfg.hidden) + .map(|i| bf16::from_f32(((i + seed * 7) % 19) as f32 / 18. - 0.5)) + .collect(); + let full = MultimodalInput { + token_ids: &ids, + image_token_indices: &indices, + image_embeddings: &features, + position_ids: std::array::from_fn(|a| positions[a].as_slice()), + }; + model.graph_enabled = false; + let expected = model + .readout_window(&full, None, None, &[131], &[])? + .hidden + .remove(0); + if !prior.is_empty() { + assert_ne!(expected, prior); + } + let buffer = DeviceBuffer::new(features.len() * BF16)?; + let bytes: Vec<_> = features.iter().flat_map(|v| v.to_le_bytes()).collect(); + unsafe { + cuda::upload(buffer.at(0), &bytes, model.stream)?; + } + model.synchronize()?; + let prefix = DeviceMultimodalInput { + token_ids: &ids[..128], + image_token_indices: &indices, + image_embeddings: buffer.bf16_matrix(2, model.cfg.hidden)?, + position_ids: std::array::from_fn(|a| &positions[a][..128]), + }; + model.graph_enabled = true; + model.readout_window_device(&prefix, None, Some(&mut state), &[], &[])?; + // The window Graph cannot retain the producer allocation. + drop(buffer); + let suffix = MultimodalInput { + token_ids: &ids[128..], + image_token_indices: &[], + image_embeddings: &[], + position_ids: std::array::from_fn(|a| &positions[a][128..]), + }; + let actual = model + .readout_window(&suffix, Some(&state), None, &[131], &[])? + .hidden + .remove(0); + assert_hidden_eq(&actual, &expected, "current device prefix + host suffix"); + assert_eq!(model.window_graphs.len(), 2); + prior = expected; + } + assert_eq!(model.graph_stats().captures, 2); + assert_eq!(model.graph_stats().replays, 4); + let retired = Arc::downgrade(&state.buffers); + drop(state); + assert!( + retired.upgrade().is_some(), + "live graphs must own recorded prefix storage" + ); + // New same-sized storage must not alias the retired prefix's cache key. + let mut replacement = model.alloc_prefix(128)?; + let prefix = MultimodalInput { + token_ids: &text[..128], + image_token_indices: &[], + image_embeddings: &[], + position_ids: std::array::from_fn(|a| &positions[a][..128]), + }; + model.readout_window(&prefix, None, Some(&mut replacement), &[], &[])?; + assert_eq!(model.window_graphs.len(), 3); + assert_eq!(model.graph_stats().captures, 3); + model.prepare_scratch(1025)?; + assert!(model.window_graphs.is_empty()); + assert!( + retired.upgrade().is_none(), + "growth must retire Graph-owned storage" + ); + // Exercise the production retirement helper with a real failed recording. + model.readout_window(&prefix, None, Some(&mut replacement), &[], &[])?; + let failed = cuda::Graph::capture(model.stream, || { + anyhow::bail!("injected window recording failure") + }); + assert!(failed.is_err()); + model.cache_window_graph([0, 128, 0, 0, 0], [None, None], failed)?; + assert!(!model.graph_enabled); + assert!(model.window_graphs.is_empty()); + model.readout_window(&prefix, None, Some(&mut replacement), &[], &[])?; + Ok(()) +} + +#[test] +#[ignore = "requires pinned Cua-S1 language weights and CUDA; run serially"] +fn window_graph_retains_and_keys_both_prefix_and_capture_allocations() -> Result<()> { + let mut model = load_model()?; + let ids = text_tokens(&model, 132)?; + let positions: Vec = (0..132).collect(); + let input = |range: std::ops::Range| MultimodalInput { + token_ids: &ids[range.clone()], + image_token_indices: &[], + image_embeddings: &[], + position_ids: [&positions[range.clone()]; 3], + }; + let expected = model + .readout_window(&input(0..132), None, None, &[131], &[])? + .hidden + .remove(0); + let mut source = model.alloc_prefix(64)?; + let mut first = model.alloc_prefix(128)?; + let mut second = model.alloc_prefix(128)?; + model.graph_enabled = true; + model.readout_window(&input(0..64), None, Some(&mut source), &[], &[])?; + model.readout_window(&input(64..128), Some(&source), Some(&mut first), &[], &[])?; + let retired = Arc::downgrade(&first.buffers); + drop(first); + assert!( + retired.upgrade().is_some(), + "Graph must retain destination as well as source" + ); + model.readout_window(&input(64..128), Some(&source), Some(&mut second), &[], &[])?; + assert_eq!( + model.window_graphs.len(), + 3, + "a different destination requires a new Graph" + ); + let actual = model + .readout_window(&input(128..132), Some(&second), None, &[131], &[])? + .hidden + .remove(0); + assert_hidden_eq( + &actual, + &expected, + "continuation from second Graph destination", + ); + model.prepare_scratch(1025)?; + assert!(retired.upgrade().is_none()); + Ok(()) +} + +#[test] +#[ignore = "requires pinned Cua-S1 language weights and CUDA; run serially"] +fn window_graph_fifo_and_matched_workspace_retirement_release_owned_storage() -> Result<()> { + let mut model = load_model()?; + model.graph_enabled = true; + let state = model.alloc_prefix(64)?; + let retired = Arc::downgrade(&state.buffers); + // Real empty CUDA Graphs isolate cache ownership from arithmetic. They are + // never replayed; the subsequent matched section uses real window kernels. + for i in 0..128 { + let graph = cuda::Graph::capture(model.stream, || Ok(()))?; + let held = if i == 0 { + Some(Arc::clone(&state.buffers)) + } else { + None + }; + model.cache_window_graph([i, 64, 0, 0, 0], [None, held], Ok(graph))?; + } + drop(state); + assert_eq!(model.window_graphs.len(), 128); + assert!(retired.upgrade().is_some()); + let graph = cuda::Graph::capture(model.stream, || Ok(()))?; + model.cache_window_graph([128, 64, 0, 0, 0], [None, None], Ok(graph))?; + assert_eq!(model.window_graphs.len(), 128); + assert!( + retired.upgrade().is_none(), + "129th insertion must release the oldest entry" + ); + model.prepare_scratch(1025)?; + assert!(model.window_graphs.is_empty()); + let ids = text_tokens(&model, 64)?; + let positions: Vec = (0..64).collect(); + let input = MultimodalInput { + token_ids: &ids, + image_token_indices: &[], + image_embeddings: &[], + position_ids: [&positions; 3], + }; + let mut state = model.alloc_prefix(64)?; + let retired = Arc::downgrade(&state.buffers); + assert!( + model + .with_matched_prefix(&[64, 64], &[64], |m| m.readout_window( + &input, + None, + Some(&mut state), + &[], + &[] + ))? + .is_some() + ); + assert_eq!(model.window_graphs.len(), 1); + drop(state); + assert!(retired.upgrade().is_some()); + for m in 1..=17 { + let _ = model.with_matched_prefix(&[m, m], &[m], |_| Ok(()))?; + } + assert!( + model.window_graphs.is_empty(), + "matched workspace retirement must clear dependent Graphs" + ); + assert!(retired.upgrade().is_none()); + let captures = model.graph_stats().captures; + let mut state = model.alloc_prefix(64)?; + assert!( + model + .with_matched_prefix(&[64, 64], &[64], |m| m.readout_window( + &input, + None, + Some(&mut state), + &[], + &[] + ))? + .is_some() + ); + assert_eq!( + model.graph_stats().captures, + captures + 1, + "retired shape must recapture" + ); + let replays = model.graph_stats().replays; + assert!( + model + .with_matched_prefix(&[64, 64], &[64], |m| { + m.readout_window(&input, None, Some(&mut state), &[], &[]) + })? + .is_some() + ); + assert_eq!(model.graph_stats().replays, replays + 1); + assert_eq!( + model.graph_stats().captures, + captures + 1, + "same shape must replay after recapture" + ); + Ok(()) +}