diff --git a/docs/architecture.md b/docs/architecture.md index 5cfcaf14..f878c6f6 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -68,7 +68,12 @@ FLA chunk and either FlashInfer dispatch. CuteDSL and CUDA graphs are not selected. Cua-S1 and Open-Jev compute their decision heads on the CPU after downloading -the final hidden state. Decider projects its selected LM-head rows on CUDA. +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, +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 FP32 adapters and framework-reference arithmetic. It computes selected BF16 label-head logits on CUDA, then calibrates typed answers on the CPU. JEMM admits diff --git a/recipe/cua_s1/native_screenshot.md b/recipe/cua_s1/native_screenshot.md index d591a6df..ef32802c 100644 --- a/recipe/cua_s1/native_screenshot.md +++ b/recipe/cua_s1/native_screenshot.md @@ -94,3 +94,148 @@ Recorded RTX 4090 evidence covers native screenshot parity and Graph regressions the resident-feature comparison covers host/device output equality on its fixed real-screenshot corpus. These scopes do not establish general GUI task accuracy or native equivalence to every reference precision and environment. + +## Request-local language prefix reuse + +`CUA_S1_MM_PREFIX=shared` opts into sharing the aligned language prefix of all +questions in one screenshot request. Unset or `off` keeps the existing independent +forwards. `fixed` runs complete prompts independently with the same fixed GEMM +algorithms as `shared`, providing the numerical control. The `matched` mode +below instead checks ordinary full-prompt algorithms. Other values fail before +checkpoint or CUDA loading. + +```sh +export CUA_S1_MM_PREFIX=shared +export CUA_S1_GRAPH=0 CUA_S1_DEVICE_FEATURES=0 +# Start omni-cua-s1-vision using the pinned artifacts above. +``` + +The planner compares prepared token IDs and all three position axes, rounds the +common prefix down to 64 tokens, and leaves a nonempty suffix for each question. +It captures attention K/V, GDN recurrent state and convolution history once, then +continues each suffix independently in request order. The engine retains one +snapshot allocation, growing it when a larger shared prefix needs capacity and +freeing it with the engine. Every shared request captures fresh contents before +reading them; a repeated or changed screenshot is always recomputed. One question, +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, +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 +cross-request content cache is included. Fixed algorithms still use reference +M=64; the experimental M=256 profile is not part of this implementation. + +On one RTX 4090 (SM8.9), a matched eight-payload synthetic HTTP comparison with +vision cache 8 entries/1024 MiB measured shared M64 p50/p95 149.213/223.530 ms +before allocation reuse and 138.404/214.262 ms after it. All response bytes +matched; sampled process maxima were 10282/10288 MiB. Each side excluded an +8-request growth pass, an 8-request all-shape stabilization pass and 5 warmups, +then measured 3 rounds of 8 requests, concurrency 1. The clock includes full +HTTP response parsing and schema validation. The ordinary off control was +118.804/275.748 ms, so this result does not justify enabling shared mode by default. +Source pins, raw samples, exclusions and reproduction inputs are in the +[follow-up qualification](https://github.com/Levius-Fubuki/system1-omni/releases/download/cua-mm-qualification-20261025/cua-mm-followup-qualification-20261025.tar.gz). +These historical measurements apply to the fixed/shared implementation at +`d4908d95c`, before the matched mode below. Other GPUs and general GUI task +success remain unverified. + +CPU control/planner tests run in the ordinary Cua suite. On a reserved GPU with +the four artifact variables above, run the exact model and recovery checks: + +```sh +cargo test --release --locked -p omni-cua-s1-native --lib prefix_tests \ + -- --ignored --test-threads=1 --nocapture +``` + +The screenshot corpus is a numerical and execution regression; it does not +establish general GUI task accuracy or a speedup for every request. + +### Full-prompt-matched algorithms + +`CUA_S1_MM_PREFIX=matched` shares the same request-local aligned prefix using +algorithms selected at the first question's complete prompt length, with the +ordinary BF16 reduction policy. Before any prefix execution, every layer's +projection must have identical algorithm descriptors for every complete question; +the selected algorithms must also support every prefix and suffix window within +the workspace limit. A mismatch or unsupported window falls back to ordinary +independent full passes. A single question or no aligned prefix takes that path +directly. Allocation/setup and execution errors remain errors. + +```sh +export CUA_S1_MM_PREFIX=matched +export CUA_S1_GRAPH=0 CUA_S1_DEVICE_FEATURES=0 +# Rebuild the CUDA library with src/backends/cuda/qwen3_5/build.sh. +``` + +This mode requires the optional additive ABI 9 symbols `cs1_gemm_create_matched` +and `cs1_gemm_compatible`. Both are checked during the single-question startup +warmup, before the worker binds its health endpoint; an incomplete library fails +with a rebuild error. Older libraries still serve the existing modes. +The engine retains at most 16 matched GEMM workspaces, 32 MiB each. On a new +reference length beyond that cap, it drains the language stream and releases the +previous generation before allocating another. This is separate from retained +prefix-state capacity and vision-cache limits. No screenshot contents survive +as reusable results between requests. Errors and panics drain submitted work +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. + +```sh +cargo test --release --locked -p omni-qwen3-5-native --test kernels \ + matched_gemm_rows_equal_complete_ordinary_projection \ + -- --ignored --test-threads=1 --nocapture +# Also run the prefix_tests command above with vision Graph 0 and 1. +``` + +At runtime source `02d75b3b4`, on RTX 4090/SM8.9, driver 595.71.05, +CUDA 13.0/cuBLASLt 13.1.1.3, a comparison against `d4908d95c` used concurrency 1 +and a direct native-worker HTTP +POST through full JSON/schema validation. Each configuration ran in forward and +reverse order, four timed traversals per visit, after excluding growth, +all-shape stabilization and five warmups. The frozen real suite uses 12 screenshots +(desktop/mobile/web, four each), four correlated click-choice permutations each. + +| Workload / mode | timed requests | mean ms | p50 ms | p95 ms | sampled max MiB | +| --- | ---: | ---: | ---: | ---: | ---: | +| synthetic / off | 64 | 133.30 | 119.31 | 277.13 | 10138 | +| synthetic / shared M64 | 64 | 134.73 | 139.35 | 215.54 | 10288 | +| synthetic / matched | 64 | 107.17 | 103.24 | 173.69 | 10448 | +| real / off | 96 | 256.36 | 265.33 | 271.98 | 10210 | +| real / shared M64 | 96 | 205.53 | 209.87 | 215.85 | 10318 | +| real / matched | 96 | 152.83 | 152.94 | 160.08 | 10606 | + +On this real suite, matched reduces mean latency 25.6% versus shared M64 and +40.4% versus off; p95 reductions are 25.8% and 41.1%. The paired-workload-round +bootstrap 95% intervals for mean reduction are 25.1–26.1% and 39.9–40.7%; these +describe repeated timings, not general task performance. Synthetic mean +reductions are 20.5% and 19.6%. One small two-question synthetic case regresses +4.8% versus off. There is no claim of a speedup for every shape. Real-suite sampled +maximum increases 288 MiB versus M64 and 396 MiB versus off. Memory is sampled +at about 100 ms, not an exact allocator peak. Vision-cache controls are 8 entries/ +1024 MiB for synthetic and 16 entries/2048 MiB for real, not new defaults. + +Five profiles each completed 80 real/stress layouts twice and 32 single-question +extractions (960 calls total), plus malformed-input/recovery probes. Matched and +candidate off have exactly the ordinary-control FP32 logit bits, complete +responses, token IDs and all three position axes; candidate shared retains the +original M64 results. All single-question extractions match their parent answers. +The 320 projection-window comparisons and prefix model regressions with vision +Graph 0/1 passed on this GPU. Existing default-versus-reference limits remain: +the adapted 48-screenshot/192-correlated-choice set scores 119 for off/matched +and 120 for M64. This is not official ScreenSpot accuracy or evidence of improved +reference quality. Defaults remain off; other GPUs, libraries and concurrency +need qualification. + +Raw timings, exclusions, logits, model/source/binary hashes, reproduction inputs +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. diff --git a/src/backends/cuda/qwen3_5/README.md b/src/backends/cuda/qwen3_5/README.md index 3becc3b8..090e54fb 100644 --- a/src/backends/cuda/qwen3_5/README.md +++ b/src/backends/cuda/qwen3_5/README.md @@ -176,3 +176,22 @@ does not establish a fresh device pass or whole-model 27B qualification. Shared preprocessing and vision adaptation notices accompany the Qwen crate in [THIRD_PARTY_NOTICES.md](../../../models/qwen3_5/native/THIRD_PARTY_NOTICES.md). + +## Optional full-prompt-matched GEMMs + +`cs1_gemm_create_matched(workspace_bytes, reference_m)` and +`cs1_gemm_compatible(ordinary, matched, M, N, K)` are additive ABI 9 symbols. +Existing consumers do not resolve them. The matched handle retains the ordinary +BF16 reduction mask and chooses one algorithm per weight shape at `reference_m`; +legacy `cs1_gemm_create_fixed` still excludes split-K reductions. Compatibility +returns 1 only for equal complete algorithm descriptors and supported layouts +within the workspace budget. Passing a null ordinary handle checks window +support without requiring another heuristic to choose the same algorithm. +Neither compatibility check executes a GEMM. + +Rust owns admission, request scope and the bounded workspace cache. Descriptor +matching is not a guarantee of row-invariant arithmetic on arbitrary hardware; +run the matched projection regression in `tests/qwen3_5/kernels.rs` and the +[Cua screenshot qualification](../../../../recipe/cua_s1/native_screenshot.md#full-prompt-matched-algorithms) +for each supported GPU/library/model combination. Other numerical profiles, +vision GEMMs and the existing default/fixed policies stay on their own handles. diff --git a/src/backends/cuda/qwen3_5/gemm.cu b/src/backends/cuda/qwen3_5/gemm.cu index 0f55faee..92ebb54b 100644 --- a/src/backends/cuda/qwen3_5/gemm.cu +++ b/src/backends/cuda/qwen3_5/gemm.cu @@ -63,6 +63,7 @@ struct Gemm { void* workspace = nullptr; size_t workspace_bytes = 0; bool reference = false; + bool matched = false; int reference_m = 0; // > 0: one algorithm per weight shape, chosen at this M std::map plans; std::map, cublasLtMatmulAlgo_t> fixed; // N, K, ldy, FP32, bias @@ -126,9 +127,12 @@ int fixed_choice(Gemm& g, int N, int K, int ldy, bool fp32, bool bias, Plan& p) if (it == g.fixed.end()) { Plan r; int rc = describe(g.reference_m, N, K, ldy, r, fp32, bias); - if (rc == 0) rc = first_choice(g, r, CUBLASLT_REDUCTION_SCHEME_NONE); + const uint32_t schemes = g.matched + ? CUBLASLT_REDUCTION_SCHEME_MASK & ~CUBLASLT_REDUCTION_SCHEME_INPLACE + : CUBLASLT_REDUCTION_SCHEME_NONE; + if (rc == 0) rc = first_choice(g, r, schemes); if (rc == 0) { - // Check that the choice really has no split-K reduction. + // Legacy fixed excludes split-K; matched retains ordinary deterministic reductions. int32_t splits = 0; uint32_t scheme = 0; size_t written = 0; @@ -137,7 +141,8 @@ int fixed_choice(Gemm& g, int N, int K, int ldy, bool fp32, bool bias, Plan& p) sizeof(splits), &written) == CUBLAS_STATUS_SUCCESS && cublasLtMatmulAlgoConfigGetAttribute(&r.algo, CUBLASLT_ALGO_CONFIG_REDUCTION_SCHEME, &scheme, sizeof(scheme), &written) == CUBLAS_STATUS_SUCCESS; - if (!read || splits != 1 || scheme != CUBLASLT_REDUCTION_SCHEME_NONE) + if (!read || (!g.matched && (splits != 1 || scheme != CUBLASLT_REDUCTION_SCHEME_NONE)) || + scheme == CUBLASLT_REDUCTION_SCHEME_INPLACE) rc = status(CUBLAS_STATUS_NOT_SUPPORTED); } if (rc == 0) it = g.fixed.emplace(key, r.algo).first; @@ -148,6 +153,7 @@ int fixed_choice(Gemm& g, int N, int K, int ldy, bool fp32, bool bias, Plan& p) cublasLtMatmulHeuristicResult_t check{}; const cublasStatus_t s = cublasLtMatmulAlgoCheck(g.handle, p.op, p.a, p.b, p.c, p.c, &p.algo, &check); if (s != CUBLAS_STATUS_SUCCESS) return status(s); + if (g.matched && check.state != CUBLAS_STATUS_SUCCESS) return status(check.state); if (check.workspaceSize > g.workspace_bytes) return status(CUBLAS_STATUS_NOT_SUPPORTED); return 0; } @@ -432,3 +438,20 @@ extern "C" void* cs1_reference_gemm_create(size_t workspace_bytes) { if(g) g->reference=true; return g; } + +// Additive ABI 9 entry points. Ordinary and fixed-M64 handles keep their policy. +extern "C" void* cs1_gemm_create_matched(size_t bytes, int reference_m) { + auto* g = static_cast(cs1_gemm_create_fixed(bytes, reference_m)); + if (g) g->matched = true; + return g; +} + +// ordinary == nullptr checks that the selected algorithm supports a subwindow. +extern "C" int cs1_gemm_compatible(void* ordinary, void* matched, int M, int N, int K) { + if (!matched || M <= 0 || N <= 0 || K <= 0) return 0; + Plan *a = nullptr, *b = nullptr; + if (plan_for(*static_cast(matched), M, N, K, N, b)) return 0; + if (!ordinary) return 1; + if (plan_for(*static_cast(ordinary), M, N, K, N, a)) return 0; + return std::memcmp(&a->algo, &b->algo, sizeof(a->algo)) == 0; +} diff --git a/src/backends/cuda/qwen3_5/ops.h b/src/backends/cuda/qwen3_5/ops.h index 77b8bd1e..9ad9c739 100644 --- a/src/backends/cuda/qwen3_5/ops.h +++ b/src/backends/cuda/qwen3_5/ops.h @@ -165,6 +165,11 @@ void* cs1_gemm_create(size_t workspace_bytes); // runs on). Null if reference_m <= 0 or setup fails. cs1_gemm with this handle returns a // cuBLAS status for an M the algorithm cannot serve, rather than switching algorithms. void* cs1_gemm_create_fixed(size_t workspace_bytes, int reference_m); +// Optional additive ABI 9: choose the ordinary reduction mask at reference_m. +// Compatibility returns 1 for equal complete descriptors, else 0. A null ordinary +// handle tests support only (used for prefix/suffix windows); no GEMM executes. +void* cs1_gemm_create_matched(size_t workspace_bytes, int reference_m); +int cs1_gemm_compatible(void* ordinary, void* matched, int M, int N, int K); void cs1_gemm_destroy(void* gemm); int cs1_gemm(void* gemm, const void* x, const void* w, void* y, int M, int N, int K, int ldy, void* stream); diff --git a/src/models/cua_s1/README.md b/src/models/cua_s1/README.md index 72945c23..28a39e20 100644 --- a/src/models/cua_s1/README.md +++ b/src/models/cua_s1/README.md @@ -159,3 +159,24 @@ The historical cumulative A800 investigation and frozen-corpus results are recor - More than 26 options per question. - The Metal backend. - The 0.1 adapters, which are a separate artifact. + +## Native screenshot prefix modes + +The screenshot worker optionally shares a request's exact token/3D-position +prefix with `CUA_S1_MM_PREFIX=shared`. The `fixed` control uses independent full +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 +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 +rounding is a separate comparison from shared-versus-fixed continuation parity. +The matched handle cache retains at most 16 workspaces of 32 MiB, in addition +to snapshot storage and vision-cache budgets. Descriptor matching is admission, +not a portable numerical guarantee; GPU/library/model combinations need explicit +row-invariance and model qualification. The original M64 controls retain their +own arithmetic. See the [recipe](../../../recipe/cua_s1/native_screenshot.md#request-local-language-prefix-reuse). diff --git a/src/models/cua_s1/native/src/vision_engine.rs b/src/models/cua_s1/native/src/vision_engine.rs index a529f2ea..941a31cb 100644 --- a/src/models/cua_s1/native/src/vision_engine.rs +++ b/src/models/cua_s1/native/src/vision_engine.rs @@ -1,4 +1,5 @@ //! Native screenshot-to-decision orchestration, with request-local vision reuse. +mod prefix; pub use crate::vision_processing::PreparedRequest; use crate::vision_processing::{VisionProcessor, probabilities}; use crate::{ @@ -24,6 +25,8 @@ pub struct VisionEngine { language: Model, letters: Vec, device_features: bool, + prefix_mode: prefix::Mode, + prefix_storage: Option, } pub struct Readout { @@ -41,6 +44,7 @@ impl VisionEngine { anyhow::bail!("CUA_S1_DEVICE_FEATURES must be unset, 0, or 1") } }; + let prefix_mode = prefix::Mode::from_env(device_features)?; let marker: Value = serde_json::from_slice( &std::fs::read(language.join("cua_s1_language_export.json")) .context("export the multimodal language checkpoint first")?, @@ -82,6 +86,8 @@ impl VisionEngine { language: model, letters, device_features, + prefix_mode, + prefix_storage: None, }) } @@ -209,7 +215,42 @@ impl VisionEngine { prepared.prompts.len() == counts.len() && !counts.is_empty(), "question/prompt count mismatch" ); - if self.device_features { + 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) + } 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)?; + prefix::execute( + &mut self.language, + &self.letters, + prepared, + &features, + self.prefix_mode, + &mut self.prefix_storage, + ) + } else if self.device_features { let features = self.vision.forward_device(&prepared.image)?; // Each consumer synchronizes before returning, including errors; // the view remains borrowed across all questions in this request. @@ -325,3 +366,7 @@ mod device_copy_failure_tests; #[cfg(test)] #[path = "../../../../../tests/cua_s1/device_transient_failure_gpu.rs"] mod device_transient_failure_tests; + +#[cfg(test)] +#[path = "../../../../../tests/cua_s1/multimodal_prefix.rs"] +mod prefix_tests; diff --git a/src/models/cua_s1/native/src/vision_engine/prefix.rs b/src/models/cua_s1/native/src/vision_engine/prefix.rs new file mode 100644 index 00000000..f6f0466e --- /dev/null +++ b/src/models/cua_s1/native/src/vision_engine/prefix.rs @@ -0,0 +1,192 @@ +//! Request-local screenshot prefixes; all branches retain the Cua letter head. +use super::*; +use crate::model::PrefixState; + +#[derive(Clone, Copy, PartialEq, Eq)] +pub(super) enum Mode { + Off, + Fixed, + Shared, + Matched, +} + +impl Mode { + pub(super) fn from_env(device_features: bool) -> Result { + let mode = match std::env::var("CUA_S1_MM_PREFIX").as_deref() { + Err(std::env::VarError::NotPresent) | Ok("off") => Self::Off, + Ok("fixed") => Self::Fixed, + Ok("shared") => Self::Shared, + Ok("matched") => Self::Matched, + _ => anyhow::bail!("CUA_S1_MM_PREFIX must be off, fixed, shared, or matched"), + }; + ensure!( + mode == Self::Off + || (!device_features && std::env::var("CUA_S1_GRAPH").as_deref() != Ok("1")), + "multimodal prefix modes require Graph and device features off" + ); + Ok(mode) + } +} + +/// Compare actual prepared tokens and all rotary axes, leaving a nonempty tail. +pub(super) fn common_prefix(prompts: &[ImagePrompt]) -> usize { + if prompts.len() < 2 { + return 0; + } + let first = &prompts[0]; + let limit = prompts + .iter() + .map(|p| p.token_ids.len()) + .min() + .unwrap() + .saturating_sub(1); + let count = (0..limit) + .take_while(|&i| { + prompts[1..].iter().all(|p| { + p.token_ids[i] == first.token_ids[i] + && (0..3).all(|axis| p.position_ids[axis][i] == first.position_ids[axis][i]) + }) + }) + .count(); + count / 64 * 64 +} + +fn window( + model: &mut Model, + prompt: &ImagePrompt, + features: &[bf16], + range: std::ops::Range, + prefix: Option<&PrefixState>, + capture: Option<&mut PrefixState>, + read_last: bool, +) -> Result> { + let begin = prompt + .image_token_indices + .partition_point(|&i| i < range.start); + let end = prompt + .image_token_indices + .partition_point(|&i| i < range.end); + let indices: Vec<_> = prompt.image_token_indices[begin..end] + .iter() + .map(|i| i - range.start) + .collect(); + let rows = if read_last { + vec![range.end - 1] + } else { + vec![] + }; + 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], + position_ids: std::array::from_fn(|axis| &prompt.position_ids[axis][range.clone()]), + }; + let mut result = model.readout_window(&input, prefix, capture, &rows, &[])?; + Ok(result.hidden.pop().unwrap_or_default()) +} + +pub(super) fn execute( + model: &mut Model, + letters: &[f32], + prepared: &PreparedRequest, + features: &[bf16], + mode: Mode, + storage: &mut Option, +) -> Result>> { + // PreparedRequest is public. Reject malformed layouts before planning or + // slicing; readout_window validates token values and window positions. + let image_token = model + .cfg + .image_token_id + .context("checkpoint has no image_token_id")?; + for prompt in &prepared.prompts { + ensure!( + !prompt.token_ids.is_empty() + && prompt + .position_ids + .iter() + .all(|axis| axis.len() == prompt.token_ids.len()), + "invalid prepared position shape" + ); + ensure!( + prompt.image_token_indices.iter().copied().eq(prompt + .token_ids + .iter() + .enumerate() + .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" + ); + } + let shared = if mode == Mode::Shared { + common_prefix(&prepared.prompts) + } else { + 0 + }; + // Reuse allocation only. Every request captures fresh contents, and a + // request without a common prefix must not read the retained snapshot. + let mut state = if shared > 0 { + if storage.as_ref().is_none_or(|s| s.capacity() < shared) { + *storage = Some(model.alloc_prefix(shared)?); + } + storage.as_mut() + } else { + None + }; + if let Some(state) = &mut state { + window( + model, + &prepared.prompts[0], + features, + 0..shared, + None, + Some(state), + false, + )?; + } + prepared + .prompts + .iter() + .zip(prepared.context.option_counts()) + .map(|(p, count)| { + let hidden = window( + model, + p, + features, + shared..p.token_ids.len(), + state.as_deref(), + None, + true, + )?; + VisionEngine::letter_logits(letters, &hidden, count) + }) + .collect() +} + +/// Validate before prefix discovery; PreparedRequest is a public boundary. +pub(super) fn matched_windows(prompts: &[ImagePrompt]) -> Result<(Vec, Vec)> { + for prompt in prompts { + ensure!( + !prompt.token_ids.is_empty() + && prompt + .position_ids + .iter() + .all(|axis| axis.len() == prompt.token_ids.len()), + "invalid prepared position shape" + ); + } + let lengths: Vec<_> = prompts.iter().map(|p| p.token_ids.len()).collect(); + let shared = common_prefix(prompts); + let mut windows = if shared > 0 { vec![shared] } else { vec![] }; + if shared > 0 { + windows.extend(lengths.iter().map(|n| n - shared)); + } + Ok((lengths, windows)) +} diff --git a/src/models/qwen3_5/native/src/cuda.rs b/src/models/qwen3_5/native/src/cuda.rs index 74cd16e1..612d49da 100644 --- a/src/models/qwen3_5/native/src/cuda.rs +++ b/src/models/qwen3_5/native/src/cuda.rs @@ -420,6 +420,15 @@ pub type GatedAttention = unsafe extern "C" fn( ) -> c_int; impl Api { + pub fn matched_create(&self) -> Result *mut c_void> { + unsafe { Ok(*self._lib.get(b"cs1_gemm_create_matched\0")?) } + } + pub fn matched_compatible( + &self, + ) -> Result i32> { + unsafe { Ok(*self._lib.get(b"cs1_gemm_compatible\0")?) } + } + /// Prompt-bidirectional full attention. Absent on CUDA libraries built before /// this additive ABI 9 symbol; causal workers never resolve it. pub fn attention_gated_bidirectional(&self) -> Result { diff --git a/src/models/qwen3_5/native/src/model.rs b/src/models/qwen3_5/native/src/model.rs index 01cefd5d..6368a98d 100644 --- a/src/models/qwen3_5/native/src/model.rs +++ b/src/models/qwen3_5/native/src/model.rs @@ -12,6 +12,8 @@ //! `Model::forward_shared` runs prompts that share token prefixes, running each prefix //! once and every branch from the state after it (see `SharedPrompts`). +mod matched; + use std::collections::{HashMap, VecDeque}; use std::ffi::c_void; use std::path::Path; @@ -847,6 +849,8 @@ pub struct Model { gemm: *mut c_void, /// One GEMM algorithm per weight shape for `forward_shared`; created on first use. fixed_gemm: *mut c_void, + matched_ref: usize, + matched_handles: HashMap, reference: Option<&'static cuda::Reference>, fp32_gdn: Option, batched_gdn: Option, @@ -891,6 +895,9 @@ impl Drop for Model { self.fp32_gdn = None; // SAFETY: both GEMM handles are owned by this model and destroyed once. unsafe { + for (_, h) in self.matched_handles.drain() { + (cuda::api().cs1_gemm_destroy)(h); + } (cuda::api().cs1_gemm_destroy)(self.gemm); if !self.fixed_gemm.is_null() { (cuda::api().cs1_gemm_destroy)(self.fixed_gemm); @@ -1019,6 +1026,8 @@ impl Model { stream, gemm, fixed_gemm: std::ptr::null_mut(), + matched_ref: 0, + matched_handles: HashMap::new(), reference: None, fp32_gdn: None, batched_gdn: None, @@ -2353,6 +2362,10 @@ impl Model { /// The GEMM handle with one algorithm per weight shape, created on first use. fn fixed_gemm(&mut self) -> Result<*mut c_void> { + if self.matched_ref > 0 { + return Ok(self.matched_handles[&self.matched_ref]); + } + if self.fixed_gemm.is_null() { // SAFETY: plain allocation; checked for null below. self.fixed_gemm = diff --git a/src/models/qwen3_5/native/src/model/matched.rs b/src/models/qwen3_5/native/src/model/matched.rs new file mode 100644 index 00000000..faa3ed9e --- /dev/null +++ b/src/models/qwen3_5/native/src/model/matched.rs @@ -0,0 +1,109 @@ +//! Request-scoped selection of ordinary-compatible full-shape GEMM plans. +use super::*; + +impl Model { + /// Run aligned prefix windows with the first full prompt's GEMM algorithms. + /// Returns `None` without invoking `execute` if independent full prompts use + /// different descriptors or a window is unsupported. Legacy fixed-M64 calls + /// outside this scope keep their own handle, including after errors. + /// Both CUDA symbols are required even for a non-sharing warmup, so an + /// explicitly selected matched worker fails before reporting readiness. + /// + /// Descriptor equality is an admission check, not a portable cuBLAS numerical + /// guarantee: requalify row invariance on each GPU/library/model combination. + pub fn with_matched_prefix( + &mut self, + lengths: &[usize], + windows: &[usize], + execute: impl FnOnce(&mut Self) -> Result, + ) -> Result> { + self.ensure_fixed_profile()?; + ensure!(self.matched_ref == 0, "nested matched execution"); + ensure!( + lengths + .iter() + .chain(windows) + .all(|&m| m > 0 && m <= self.cfg.max_positions && m <= i32::MAX as usize), + "invalid matched prompt/window length" + ); + let create = cuda::api() + .matched_create() + .context("rebuild libqwen3_5_cuda for matched prefix execution")?; + let compatible = cuda::api() + .matched_compatible() + .context("rebuild libqwen3_5_cuda for matched prefix execution")?; + if lengths.len() < 2 || windows.is_empty() { + return Ok(None); + } + let m = lengths[0]; + if !self.matched_handles.contains_key(&m) && self.matched_handles.len() == 16 { + self.synchronize()?; + // SAFETY: all launches drained before destroying owned workspaces. + for (_, h) in self.matched_handles.drain() { + unsafe { (cuda::api().cs1_gemm_destroy)(h) }; + } + } + let handle = if let Some(&h) = self.matched_handles.get(&m) { + h + } else { + // SAFETY: validated positive i32 length; this model owns the handle. + let h = unsafe { create(GEMM_WORKSPACE, m as i32) }; + ensure!(!h.is_null(), "matched cuBLASLt setup failed"); + self.matched_handles.insert(m, h); + h + }; + let mut checked = std::collections::HashSet::new(); + for layer in &self.layers { + let mixer = match &layer.mixer { + Mixer::Linear(l) => [&l.in_proj, &l.out], + Mixer::Full(l) => [&l.qkv, &l.o], + }; + for w in mixer.into_iter().chain([&layer.gate_up, &layer.down]) { + ensure!( + w.shape.iter().all(|&n| n <= i32::MAX as usize), + "matched weight exceeds CUDA layout" + ); + for (rows, ordinary) in lengths + .iter() + .map(|&n| (n, true)) + .chain(windows.iter().map(|&n| (n, false))) + { + if !checked.insert((rows, w.shape[0], w.shape[1], ordinary)) { + continue; + } + // SAFETY: owned handles and checked dimensions. Null ordinary + // asks only whether a prefix/suffix window is supported. + let okay = unsafe { + compatible( + if ordinary { + self.gemm + } else { + std::ptr::null_mut() + }, + handle, + rows as i32, + w.shape[0] as i32, + w.shape[1] as i32, + ) + }; + if okay != 1 { + return Ok(None); + } + } + } + } + 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. + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| execute(self))); + let synced = self.synchronize(); + self.matched_ref = 0; + match result { + Ok(value) => { + synced?; + value.map(Some) + } + Err(panic) => std::panic::resume_unwind(panic), + } + } +} diff --git a/tests/cua_s1/device_controls.rs b/tests/cua_s1/device_controls.rs index a1806e70..7d5b839f 100644 --- a/tests/cua_s1/device_controls.rs +++ b/tests/cua_s1/device_controls.rs @@ -44,3 +44,112 @@ fn non_unicode_feature_control_fails_before_checkpoint_or_cuda_loading() { use std::os::unix::ffi::OsStrExt; rejects_before_loading(std::ffi::OsStr::from_bytes(&[0xff])); } + +#[test] +fn multimodal_prefix_controls_reject_invalid_and_incompatible_modes_before_loading() { + for (prefix, graph, device, expected) in [ + ( + "", + "0", + "0", + "CUA_S1_MM_PREFIX must be off, fixed, shared, or matched", + ), + ( + "true", + "0", + "0", + "CUA_S1_MM_PREFIX must be off, fixed, shared, or matched", + ), + ( + "shared", + "1", + "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", + "1", + "multimodal prefix modes require Graph and device features off", + ), + ] { + let missing = + std::env::temp_dir().join(format!("cua-prefix-missing-{}", std::process::id())); + assert!(!missing.exists()); + let output = std::process::Command::new(env!("CARGO_BIN_EXE_omni-cua-s1-vision")) + .env("CUA_S1_MM_PREFIX", prefix) + .env("CUA_S1_GRAPH", graph) + .env("CUA_S1_DEVICE_FEATURES", device) + .env("CUA_S1_BASE", &missing) + .env("CUA_S1_VISION_ADAPTER", &missing) + .env("CUA_S1_MODEL", &missing) + .env("CUA_S1_CUDA_LIB", &missing) + .output() + .unwrap(); + assert!(!output.status.success()); + let error = String::from_utf8(output.stderr).unwrap(); + assert!(error.contains(expected), "{error}"); + } +} + +#[test] +fn matched_prefix_is_opt_in_and_recognized_before_checkpoint_loading() { + for value in [ + None, + Some("off"), + Some("fixed"), + Some("shared"), + Some("matched"), + ] { + let mut cmd = std::process::Command::new(env!("CARGO_BIN_EXE_omni-cua-s1-vision")); + cmd.env_remove("CUA_S1_MM_PREFIX") + .env("CUA_S1_GRAPH", "0") + .env("CUA_S1_DEVICE_FEATURES", "0") + .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"); + if let Some(value) = value { + cmd.env("CUA_S1_MM_PREFIX", value); + } + let output = cmd.output().unwrap(); + let error = String::from_utf8(output.stderr).unwrap(); + assert!(!output.status.success()); + assert!( + error.contains("export the multimodal language checkpoint"), + "{value:?}: {error}" + ); + } +} + +#[cfg(unix)] +#[test] +fn non_unicode_prefix_is_rejected_before_loading() { + use std::os::unix::ffi::OsStrExt; + let output = std::process::Command::new(env!("CARGO_BIN_EXE_omni-cua-s1-vision")) + .env("CUA_S1_MM_PREFIX", std::ffi::OsStr::from_bytes(&[0xff])) + .env("CUA_S1_DEVICE_FEATURES", "0") + .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(); + assert!( + String::from_utf8(output.stderr) + .unwrap() + .contains("CUA_S1_MM_PREFIX must be off, fixed, shared, or matched") + ); +} diff --git a/tests/cua_s1/multimodal_prefix.rs b/tests/cua_s1/multimodal_prefix.rs new file mode 100644 index 00000000..a401037d --- /dev/null +++ b/tests/cua_s1/multimodal_prefix.rs @@ -0,0 +1,233 @@ +use super::*; + +fn prompt(tokens: usize) -> ImagePrompt { + ImagePrompt { + token_ids: vec![7; tokens], + image_token_indices: (16..tokens.min(144)).collect(), + position_ids: std::array::from_fn(|_| (0..tokens as i64).collect()), + } +} + +#[test] +fn shared_prefix_requires_multiple_prompts_and_aligns_tokens_and_all_positions() { + assert_eq!(prefix::common_prefix(&[]), 0); + assert_eq!(prefix::common_prefix(&[prompt(193)]), 0); + let mut prompts = vec![prompt(193), prompt(193)]; + assert_eq!(prefix::common_prefix(&prompts), 192); + prompts[1].token_ids[130] = 9; + assert_eq!(prefix::common_prefix(&prompts), 128); + for axis in 0..3 { + let mut prompts = vec![prompt(193), prompt(193)]; + prompts[1].position_ids[axis][70] += 1; + assert_eq!(prefix::common_prefix(&prompts), 64); + } + assert_eq!(prefix::common_prefix(&[prompt(64), prompt(64)]), 0); + assert_eq!(prefix::common_prefix(&[prompt(65), prompt(193)]), 64); +} + +#[test] +#[ignore = "requires CUDA and pinned CUA_S1_* weights; run serially"] +fn screenshot_prefix_matches_fixed_and_does_not_retain_previous_request_state() { + let mut engine = device_feature_tests::load(); + for (width, height, seed, count) in [ + (256, 256, 1, 1), + (256, 256, 2, 2), + (512, 256, 3, 4), + (256, 512, 4, 8), + (256, 256, 1, 2), + (256, 256, 8, 1), + ] { + 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 prepared = engine.prepare(width, height, &rgb, &questions).unwrap(); + if count > 1 { + assert!(prefix::common_prefix(&prepared.prompts) >= 64); + } + let features = engine.vision.forward(&prepared.image).unwrap(); + let fixed = prefix::execute( + &mut engine.language, + &engine.letters, + &prepared, + &features, + prefix::Mode::Fixed, + &mut None, + ) + .unwrap(); + let shared = prefix::execute( + &mut engine.language, + &engine.letters, + &prepared, + &features, + prefix::Mode::Shared, + &mut engine.prefix_storage, + ) + .unwrap(); + assert_eq!( + shared, fixed, + "shared letter logits must equal complete fixed passes" + ); + engine.prefix_mode = prefix::Mode::Shared; + let response = engine.predict_prepared(&prepared).unwrap(); + engine.prefix_mode = prefix::Mode::Fixed; + assert_eq!(engine.predict_prepared(&prepared).unwrap(), response); + assert_eq!(response, prepared.context.finish(fixed).unwrap()); + } +} + +#[test] +#[ignore = "requires CUDA and pinned CUA_S1_* weights; run serially"] +fn prefix_rejects_malformed_prepared_inputs_and_recovers() { + let mut engine = device_feature_tests::load(); + engine.prefix_mode = prefix::Mode::Shared; + let rgb = vec![80; 256 * 256 * 3]; + let mut prepared = engine + .prepare(256, 256, &rgb, &device_feature_tests::questions(1)) + .unwrap(); + let last = prepared.prompts[0].position_ids[0].pop().unwrap(); + assert!(engine.predict_prepared(&prepared).is_err()); + prepared.prompts[0].position_ids[0].push(last); + let index = prepared.prompts[0].image_token_indices[0]; + prepared.prompts[0].image_token_indices[0] = usize::MAX; + assert!(engine.predict_prepared(&prepared).is_err()); + prepared.prompts[0].image_token_indices[0] = index; + let other = prepared.prompts[1].image_token_indices[0]; + prepared.prompts[1].image_token_indices[0] = other - 1; + assert!(engine.predict_prepared(&prepared).is_err()); + prepared.prompts[1].image_token_indices[0] = other; + // Reject a suffix token after the common prefix has already been captured, + // then recover without retaining that request's state. + let last = prepared.prompts[0].token_ids.last_mut().unwrap(); + let token = *last; + *last = u32::MAX; + assert!(engine.predict_prepared(&prepared).is_err()); + *prepared.prompts[0].token_ids.last_mut().unwrap() = token; + assert!(engine.predict_prepared(&prepared).is_ok()); +} + +#[test] +fn matched_windows_validate_all_axes_and_leave_nonempty_aligned_suffixes() { + let mut prompts = vec![prompt(193), prompt(219)]; + assert_eq!( + prefix::matched_windows(&prompts).unwrap(), + (vec![193, 219], vec![192, 1, 27]) + ); + prompts[1].token_ids[70] = 9; + assert_eq!( + prefix::matched_windows(&prompts).unwrap().1, + vec![64, 129, 155] + ); + for axis in 0..3 { + let mut invalid = vec![prompt(193), prompt(219)]; + invalid[1].position_ids[axis].pop(); + assert!(prefix::matched_windows(&invalid).is_err()); + } + assert!(prefix::matched_windows(&[prompt(0), prompt(193)]).is_err()); + assert!( + prefix::matched_windows(&[prompt(64), prompt(64)]) + .unwrap() + .1 + .is_empty() + ); + assert!( + prefix::matched_windows(&[prompt(193)]) + .unwrap() + .1 + .is_empty() + ); +} + +#[test] +#[ignore = "requires CUDA and pinned CUA_S1_* weights; run serially with vision Graph0/1"] +fn matched_prefix_preserves_default_and_original_m64_after_errors() { + let mut engine = device_feature_tests::load(); + let mut admitted = 0; + for (width, height, seed, count) in [ + (256, 256, 1, 2), + (512, 256, 2, 4), + (256, 512, 3, 8), + (256, 256, 1, 1), + (256, 256, 4, 2), + ] { + 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; + let ordinary = engine.execute_prepared(&prepared).unwrap(); + engine.prefix_mode = prefix::Mode::Fixed; + let m64 = engine.execute_prepared(&prepared).unwrap(); + let bits = |rows: Vec>| { + rows.into_iter() + .map(|r| r.into_iter().map(f32::to_bits).collect::>()) + .collect::>() + }; + let (lengths, windows) = prefix::matched_windows(&prepared.prompts).unwrap(); + let features = engine.vision.forward(&prepared.image).unwrap(); + if let Some(logits) = engine + .language + .with_matched_prefix(&lengths, &windows, |model| { + prefix::execute( + model, + &engine.letters, + &prepared, + &features, + prefix::Mode::Shared, + &mut engine.prefix_storage, + ) + }) + .unwrap() + { + admitted += 1; + assert_eq!(bits(logits), bits(ordinary.clone())); + } + engine.prefix_mode = prefix::Mode::Matched; + for _ in 0..2 { + assert_eq!( + bits(engine.execute_prepared(&prepared).unwrap()), + bits(ordinary.clone()) + ); + } + 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(ordinary.clone()) + ); + let pos = prepared.prompts[0].position_ids[2].pop().unwrap(); + assert!(engine.execute_prepared(&prepared).is_err()); + prepared.prompts[0].position_ids[2].push(pos); + let index = prepared.prompts[0].image_token_indices[0]; + prepared.prompts[0].image_token_indices[0] = usize::MAX; + assert!(engine.execute_prepared(&prepared).is_err()); + prepared.prompts[0].image_token_indices[0] = index; + assert_eq!( + bits(engine.execute_prepared(&prepared).unwrap()), + bits(ordinary.clone()) + ); + engine.prefix_mode = prefix::Mode::Shared; + assert_eq!(bits(engine.execute_prepared(&prepared).unwrap()), bits(m64)); + engine.prefix_mode = prefix::Mode::Off; + assert_eq!( + bits(engine.execute_prepared(&prepared).unwrap()), + bits(ordinary) + ); + } + assert!(admitted > 0, "all shared-window probes fell back"); +} diff --git a/tests/qwen3_5/graph_loading.rs b/tests/qwen3_5/graph_loading.rs index 1f7dd249..4b245d88 100644 --- a/tests/qwen3_5/graph_loading.rs +++ b/tests/qwen3_5/graph_loading.rs @@ -78,6 +78,10 @@ fn explicit_graph_mode_and_adapter_models_retire_the_stream_once_after_allocatio .trim() .trim_end_matches(';'); let mut fake = include_str!("graph_loading_fake.c").to_owned(); + let missing = std::env::var("OMNI_TEST_MATCHED_MISSING_SYMBOL").ok(); + if let Some(symbol) = &missing { + fake = fake.replace(symbol, &format!("omitted_{symbol}")); + } let implemented = [ "cs1_abi_version", "cs1_error_string", @@ -130,6 +134,30 @@ fn explicit_graph_mode_and_adapter_models_retire_the_stream_once_after_allocatio let metrics = unsafe { libloading::Library::new(&library).unwrap() }; let metric = |name: &[u8]| unsafe { metrics.get:: i32>(name).unwrap()() }; + if let Some(symbol) = missing { + let mut model = Model::load_with_graph(&directory, &library, false).unwrap(); + // Matched warmup must reject either missing capability even though one + // question does not share a prefix. No matched workspace is allocated. + let error = model + .with_matched_prefix(&[64], &[], |_| Ok(())) + .expect_err("matched warmup accepted an incomplete CUDA library"); + let detail = format!("{error:#}"); + assert!(detail.contains("rebuild libqwen3_5_cuda"), "{detail}"); + assert!(detail.contains(&symbol), "{detail}"); + assert_eq!(metric(b"loading_gemms\0"), 1); + // Ordinary and legacy fixed execution still reach their kernel calls; + // those fake kernels fail deliberately, without needing matched symbols. + for result in [model.forward(&[0]), model.forward_fixed(&[0])] { + let error = result.unwrap_err().to_string(); + assert!(error.contains("unexpected host-test CUDA call"), "{error}"); + } + drop(model); + assert_eq!(metric(b"loading_live\0"), 0); + assert_eq!(metric(b"loading_gemms\0"), 0); + assert_eq!(metric(b"loading_early\0"), 0); + std::fs::remove_dir_all(directory).unwrap(); + return; + } for graph in [false, true] { let mut model = Model::load_with_graph(&directory, &library, graph).unwrap(); assert_eq!(model.graph_stats().requested, graph); @@ -138,6 +166,129 @@ fn explicit_graph_mode_and_adapter_models_retire_the_stream_once_after_allocatio // created before this deliberate dispatch error and must retire safely. assert!(model.forward_fixed(&[0]).is_err()); assert_eq!(metric(b"loading_gemms\0"), 2); + // Execute the real model's bounded selection/scope logic through a + // host-only CUDA boundary. These assertions do not test GEMM arithmetic. + let set_compatible = unsafe { + metrics + .get::(b"loading_compatible\0") + .unwrap() + }; + let submit = unsafe { + metrics + .get::(b"loading_submit\0") + .unwrap() + }; + assert_eq!( + model + .with_matched_prefix(&[64, 64], &[64], |_| Ok(17)) + .unwrap(), + Some(17) + ); + assert!( + model + .with_matched_prefix::<()>(&[64, 64], &[64], |_| { + unsafe { submit() }; + anyhow::bail!("suffix failure") + }) + .is_err() + ); + let panic = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + let _ = model.with_matched_prefix::<()>(&[64, 64], &[64], |_| { + unsafe { submit() }; + panic!("suffix panic"); + }); + })); + assert!(panic.is_err()); + assert_eq!(metric(b"loading_pending\0"), 0); + assert_eq!( + model + .with_matched_prefix(&[64, 64], &[64], |_| Ok(23)) + .unwrap(), + Some(23) + ); + assert!( + model + .with_matched_prefix(&[64, 64], &[64], |m| { + m.with_matched_prefix(&[64, 64], &[64], |_| Ok(())) + .map(|_| ()) + }) + .unwrap_err() + .to_string() + .contains("nested matched execution") + ); + // Returning to M64 must select its original handle after a failed scope. + assert!(model.forward_fixed(&[0]).is_err()); + unsafe { set_compatible(0) }; + assert_eq!( + model + .with_matched_prefix(&[64, 64], &[64], |_| -> anyhow::Result<()> { + panic!("incompatible plan executed") + }) + .unwrap(), + None + ); + unsafe { set_compatible(1) }; + let set_window = unsafe { + metrics + .get::(b"loading_window_supported\0") + .unwrap() + }; + unsafe { set_window(0) }; + assert_eq!( + model + .with_matched_prefix(&[64, 64], &[1, 63], |_| -> anyhow::Result<()> { + panic!("unsupported window executed") + }) + .unwrap(), + None + ); + unsafe { set_window(1) }; + let reject = unsafe { + metrics + .get::(b"loading_reject\0") + .unwrap() + }; + // Earlier full shapes and suffixes pass; later failures must still + // prevent any window execution, including in a later weight projection. + for (ordinary, rows, n, k) in [(1, 62, 0, 0), (0, 7, 0, 0), (1, 64, 8, 256)] { + unsafe { reject(ordinary, rows, n, k) }; + assert_eq!( + model + .with_matched_prefix(&[64, 62], &[1, 63, 7], |_| -> anyhow::Result<()> { + panic!("later incompatible shape executed"); + }) + .unwrap(), + None + ); + } + unsafe { reject(-1, 0, 0, 0) }; + assert_eq!( + model.with_matched_prefix(&[64], &[64], |_| Ok(1)).unwrap(), + None + ); + assert!( + model + .with_matched_prefix(&[0, 64], &[64], |_| Ok(1)) + .is_err() + ); + assert!( + model + .with_matched_prefix(&[64, 65], &[64], |_| Ok(1)) + .is_err() + ); + // Seventeen distinct reference lengths must retire the first bounded + // generation before reusing its cuBLAS workspace slots. + for m in 1..=17 { + // Represent queued device work just before workspace retirement. + unsafe { submit() }; + assert_eq!( + model.with_matched_prefix(&[m, m], &[m], |_| Ok(m)).unwrap(), + Some(m) + ); + assert!(metric(b"loading_gemms\0") <= 18); + assert_eq!(metric(b"loading_pending\0"), 0); + assert_eq!(metric(b"loading_unsynced_destroy\0"), 0); + } drop(model); assert_eq!(metric(b"loading_streams\0"), metric(b"loading_destroyed\0")); assert_eq!(metric(b"loading_live\0"), 0); @@ -154,6 +305,13 @@ fn explicit_graph_mode_and_adapter_models_retire_the_stream_once_after_allocatio std::env::var("CUA_S1_GRAPH").as_deref() == Ok("1") ); let allocations = metric(b"loading_live\0"); + assert!( + model + .with_matched_prefix(&[64, 64], &[64], |_| Ok(())) + .unwrap_err() + .to_string() + .contains("only merged weights") + ); let error = model.forward_fixed(&[0]).unwrap_err(); assert!(error.to_string().contains("only merged weights")); let error = model @@ -201,4 +359,22 @@ fn explicit_graph_mode_and_adapter_models_retire_the_stream_once_after_allocatio "failed construction retires its stream" ); std::fs::remove_dir_all(directory).unwrap(); + // CUDA API loading is process-global, so isolate each incomplete library. + for symbol in ["cs1_gemm_create_matched", "cs1_gemm_compatible"] { + let output = Command::new(std::env::current_exe().unwrap()) + .args([ + "--exact", + "explicit_graph_mode_and_adapter_models_retire_the_stream_once_after_allocations", + "--nocapture", + ]) + .env("OMNI_TEST_MATCHED_MISSING_SYMBOL", symbol) + .output() + .unwrap(); + assert!( + output.status.success(), + "{symbol}: {}\n{}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr) + ); + } } diff --git a/tests/qwen3_5/graph_loading_fake.c b/tests/qwen3_5/graph_loading_fake.c index ec22a5ee..cff2ca57 100644 --- a/tests/qwen3_5/graph_loading_fake.c +++ b/tests/qwen3_5/graph_loading_fake.c @@ -2,7 +2,10 @@ #include #include #include -static int live, streams, destroyed, early, gemms; +static int live, streams, destroyed, early, gemms, pending, unsynced_destroy; +void loading_submit(void) { pending=1; } +int loading_pending(void) { return pending; } +int loading_unsynced_destroy(void) { return unsynced_destroy; } int loading_live(void) { return live; } int loading_streams(void) { return streams; } int loading_destroyed(void) { return destroyed; } @@ -15,9 +18,24 @@ int cs1_set_device(int device) { return device == 0 ? 0 : 1; } int cs1_malloc(void** p, size_t n) { *p=malloc(n); if (!*p) return 1; ++live; return 0; } int cs1_free(void* p) { --live; free(p); return 0; } int cs1_stream_create(void** p) { *p=(void*)1; ++streams; return 0; } -int cs1_stream_sync(void* stream) { (void)stream; return 0; } +int cs1_stream_sync(void* stream) { (void)stream; pending=0; return 0; } int cs1_stream_destroy(void* stream) { (void)stream; ++destroyed; if (live || gemms) ++early; return 0; } -int cs1_upload(void* dst, const void* src, size_t n, void* stream) { (void)stream; memcpy(dst,src,n); return 0; } +int cs1_upload(void* dst, const void* src, size_t n, void* stream) { (void)stream; memcpy(dst,src,n); pending=1; return 0; } void* cs1_gemm_create(size_t n) { (void)n; ++gemms; return (void*)2; } void* cs1_gemm_create_fixed(size_t n, int m) { (void)n; if(m<=0) return NULL; ++gemms; return (void*)3; } -void cs1_gemm_destroy(void* p) { (void)p; --gemms; } +void cs1_gemm_destroy(void* p) { (void)p; if (pending) ++unsynced_destroy; --gemms; } + +static int compatible = 1, window_supported = 1; +static int reject_ordinary=-1, reject_m, reject_n, reject_k; +void loading_reject(int ordinary, int m, int n, int k) { + reject_ordinary=ordinary; reject_m=m; reject_n=n; reject_k=k; +} +void loading_window_supported(int value) { window_supported = value; } +void loading_compatible(int value) { compatible = value; } +void* cs1_gemm_create_matched(size_t n, int m) { return cs1_gemm_create_fixed(n, m); } +int cs1_gemm_compatible(void* ordinary, void* matched, int m, int n, int k) { + (void)matched; + if (reject_ordinary == (ordinary != NULL) && m == reject_m && + (!reject_n || n == reject_n) && (!reject_k || k == reject_k)) return 0; + return ordinary ? compatible : window_supported; +} diff --git a/tests/qwen3_5/kernels.rs b/tests/qwen3_5/kernels.rs index 0f626a23..5173a122 100644 --- a/tests/qwen3_5/kernels.rs +++ b/tests/qwen3_5/kernels.rs @@ -1490,6 +1490,97 @@ fn fixed_gemm_rows_do_not_depend_on_m() { #[path = "vision_v2_kernels.rs"] mod vision_v2; +#[test] +#[ignore = "needs sm89 CUDA, rebuilt matched symbols; requalify on other libraries/devices"] +fn matched_gemm_rows_equal_complete_ordinary_projection() { + let st = setup(); + let create = api().matched_create().unwrap(); + let compatible = api().matched_compatible().unwrap(); + // Actual 4B weight shapes and complete-prompt lengths from qualification. + for (n, k) in [ + (12352usize, 2560usize), + (2560, 4096), + (10240, 2560), + (18432, 2560), + (2560, 9216), + ] { + for reference in [214usize, 219, 245, 309, 437, 867, 1179, 2293] { + // SAFETY: plain positive-size cuBLAS context creation, destroyed below. + let ordinary = unsafe { (api().cs1_gemm_create)(32 << 20) }; + let matched = unsafe { create(32 << 20, reference as i32) }; + assert!(!ordinary.is_null() && !matched.is_null()); + let x = random(reference * k, reference as u64, 2.0); + let w = random(n * k, n as u64, 0.05); + let (xd, wd) = (to_device(&x, st), to_device(&w, st)); + let y = DeviceBuffer::new(reference * n * 2).unwrap(); + let project = |handle, at: usize, rows: usize| { + let mut output = vec![0; rows * n * 2]; + // SAFETY: each checked window is inside x; y has reference rows. + unsafe { + check( + (api().cs1_gemm)( + handle, + xd.at(at * k * 2), + wd.at(0), + y.at(0), + rows as i32, + n as i32, + k as i32, + n as i32, + st, + ), + "matched row probe", + ) + .unwrap(); + cuda::download(&mut output, y.at(0), st).unwrap(); + } + output + }; + // SAFETY: valid handles/dimensions; comparison does not run kernels. + assert_eq!( + unsafe { compatible(ordinary, matched, reference as i32, n as i32, k as i32) }, + 1 + ); + let full = project(ordinary, 0, reference); + for (at, rows) in [ + (0, 1), + (0, 63), + (0, 64), + (1, 64), + (17, 64), + (reference - 1, 1), + (reference / 2, reference - reference / 2), + (0, reference), + ] { + // Probe only legal windows of this admitted full matrix. + assert!(at + rows <= reference); + assert_eq!( + unsafe { + compatible( + std::ptr::null_mut(), + matched, + rows as i32, + n as i32, + k as i32, + ) + }, + 1 + ); + assert_eq!( + project(matched, at, rows), + full[at * n * 2..(at + rows) * n * 2], + "N={n} K={k} reference={reference} at={at} rows={rows}" + ); + } + // SAFETY: downloads drained all uses, both contexts owned here. + unsafe { + (api().cs1_gemm_destroy)(matched); + (api().cs1_gemm_destroy)(ordinary); + } + } + } +} + mod batched_gdn; mod packed_gemm;