From f1458f1a0472db8529902cc535bd73504ca3b286 Mon Sep 17 00:00:00 2001 From: levius <2114377220@qq.com> Date: Sun, 1 Nov 2026 09:49:31 +0800 Subject: [PATCH 1/3] perf(qwen): reuse normalized RGB8 temporal patches --- .../qwen3_5/native/src/image_preprocess.rs | 37 ++++++++++----- tests/qwen3_5/image_preprocess.rs | 45 +++++++++++++++++++ 2 files changed, 72 insertions(+), 10 deletions(-) create mode 100644 tests/qwen3_5/image_preprocess.rs diff --git a/src/models/qwen3_5/native/src/image_preprocess.rs b/src/models/qwen3_5/native/src/image_preprocess.rs index 52dec59c..a10ca2a4 100644 --- a/src/models/qwen3_5/native/src/image_preprocess.rs +++ b/src/models/qwen3_5/native/src/image_preprocess.rs @@ -16,6 +16,17 @@ const PATCH_SIZE: usize = 16; const MERGE_SIZE: usize = 2; const FACTOR: usize = PATCH_SIZE * MERGE_SIZE; const PATCH_VALUES: usize = 3 * 2 * PATCH_SIZE * PATCH_SIZE; +// Preserve the reference's float32 subtraction and division for every RGB8 value. +// Packing only reorders bytes, so each normalized value can be reused exactly. +const RGB8_NORMALIZED: [f32; 256] = { + let mut values = [0.0; 256]; + let mut pixel = 0; + while pixel < values.len() { + values[pixel] = (pixel as f32 - 127.5) / 127.5; + pixel += 1; + } + values +}; /// Source validation and upstream smart-resize limits. Prompt budgeting belongs to callers. /// /// Shared by the cua-s1 and jemm consumers, so the limit surface stays @@ -169,18 +180,20 @@ pub fn preprocess_rgb8_with_limits( for merge_y in 0..MERGE_SIZE { for merge_x in 0..MERGE_SIZE { for channel in 0..3 { - for _temporal in 0..2 { - for patch_y in 0..PATCH_SIZE { - for patch_x in 0..PATCH_SIZE { - let y = block_y * FACTOR + merge_y * PATCH_SIZE + patch_y; - let x = block_x * FACTOR + merge_x * PATCH_SIZE + patch_x; - let pixel = resized[(y * resized_width + x) * 3 + channel]; - // Match the fused float32 torchvision normalization, - // including its operation order (no reciprocal multiply). - pixel_values.push((f32::from(pixel) - 127.5) / 127.5); - } + let mut plane = [0.0; PATCH_SIZE * PATCH_SIZE]; + for patch_y in 0..PATCH_SIZE { + let y = block_y * FACTOR + merge_y * PATCH_SIZE + patch_y; + let x = block_x * FACTOR + merge_x * PATCH_SIZE; + let start = (y * resized_width + x) * 3 + channel; + for patch_x in 0..PATCH_SIZE { + let pixel = resized[start + patch_x * 3]; + plane[patch_y * PATCH_SIZE + patch_x] = + RGB8_NORMALIZED[usize::from(pixel)]; } } + // Still images repeat the same patch in both temporal slots. + pixel_values.extend_from_slice(&plane); + pixel_values.extend_from_slice(&plane); } } } @@ -291,3 +304,7 @@ fn cubic(x: f64) -> f64 { 0.0 } } + +#[cfg(test)] +#[path = "../../../../../tests/qwen3_5/image_preprocess.rs"] +mod tests; diff --git a/tests/qwen3_5/image_preprocess.rs b/tests/qwen3_5/image_preprocess.rs new file mode 100644 index 00000000..5f08df0a --- /dev/null +++ b/tests/qwen3_5/image_preprocess.rs @@ -0,0 +1,45 @@ +use super::*; + +#[test] +fn rgb8_lookup_preserves_float32_subtract_then_divide_for_every_byte() { + for pixel in 0..=255u8 { + // Keep both operations and their order: reciprocal multiplication can + // differ in the low bits consumed by the BF16 vision encoder. + let expected = (f32::from(pixel) - 127.5) / 127.5; + assert_eq!( + RGB8_NORMALIZED[usize::from(pixel)].to_bits(), + expected.to_bits() + ); + } +} + +#[test] +fn spatial_merge_channel_and_duplicate_temporal_planes_match_the_tensor_layout() { + let limits = ImageLimits { + min_pixels: 1024, + ..ImageLimits::cua() + }; + for (width, height) in [(32, 32), (64, 32), (32, 64), (256, 256)] { + let rgb: Vec = (0..width * height * 3) + .map(|i| ((i * 17 + i / 251) % 256) as u8) + .collect(); + let image = preprocess_rgb8_with_limits(width, height, &rgb, &limits).unwrap(); + assert_eq!(image.image_grid_thw, [1, height / 16, width / 16]); + assert_eq!(image.pixel_values.len(), width * height * 6); + // Independently index the output tensor as [block_y, block_x, + // merge_y, merge_x, channel, temporal, patch_y, patch_x]. + for y in 0..height { + for x in 0..width { + let patch = ((y / 32 * (width / 32) + x / 32) * 2 + y % 32 / 16) * 2 + x % 32 / 16; + for channel in 0..3 { + let expected = (f32::from(rgb[(y * width + x) * 3 + channel]) - 127.5) / 127.5; + for temporal in 0..2 { + let index = + ((patch * 3 + channel) * 2 + temporal) * 256 + y % 16 * 16 + x % 16; + assert_eq!(image.pixel_values[index].to_bits(), expected.to_bits()); + } + } + } + } + } +} From 4696f78b451f1b1e20b051628ef17a422ab7448a Mon Sep 17 00:00:00 2001 From: levius <2114377220@qq.com> Date: Sun, 1 Nov 2026 13:09:36 +0800 Subject: [PATCH 2/3] perf(cua): batch screenshot questions with fresh shared prefixes --- recipe/cua_s1/native_screenshot.md | 77 ++- src/backends/cuda/qwen3_5/gdn_prefill.cu | 39 +- src/backends/cuda/qwen3_5/ops.h | 7 + src/models/cua_s1/README.md | 8 + .../cua_s1/native/src/image_preprocess.rs | 17 + src/models/cua_s1/native/src/vision_engine.rs | 53 +- .../src/vision_engine/question_batch.rs | 122 ++++ .../cua_s1/native/src/vision_processing.rs | 35 +- src/models/qwen3_5/native/src/cuda.rs | 39 ++ .../qwen3_5/native/src/image_preprocess.rs | 139 ++++- .../native/src/image_preprocess/x86.rs | 76 +++ src/models/qwen3_5/native/src/model.rs | 184 +++++- .../native/src/model/multimodal_batch.rs | 310 ++++++++++ .../native/src/model/multimodal_shared.rs | 579 ++++++++++++++++++ tests/cua_s1/device_controls.rs | 61 ++ tests/cua_s1/question_batch.rs | 130 ++++ tests/cua_s1/vision_processing.rs | 12 + tests/qwen3_5/batched_gdn.rs | 2 +- tests/qwen3_5/image_preprocess.rs | 91 +++ tests/qwen3_5/kernels.rs | 2 + tests/qwen3_5/multimodal_batch.rs | 196 ++++++ tests/qwen3_5/multimodal_shared.rs | 221 +++++++ tests/qwen3_5/packed_gemm.rs | 11 +- tests/qwen3_5/seeded_gdn.rs | 184 ++++++ 24 files changed, 2551 insertions(+), 44 deletions(-) create mode 100644 src/models/cua_s1/native/src/vision_engine/question_batch.rs create mode 100644 src/models/qwen3_5/native/src/image_preprocess/x86.rs create mode 100644 src/models/qwen3_5/native/src/model/multimodal_batch.rs create mode 100644 src/models/qwen3_5/native/src/model/multimodal_shared.rs create mode 100644 tests/cua_s1/question_batch.rs create mode 100644 tests/cua_s1/vision_processing.rs create mode 100644 tests/qwen3_5/multimodal_batch.rs create mode 100644 tests/qwen3_5/multimodal_shared.rs create mode 100644 tests/qwen3_5/seeded_gdn.rs diff --git a/recipe/cua_s1/native_screenshot.md b/recipe/cua_s1/native_screenshot.md index ef32802c..0347c17e 100644 --- a/recipe/cua_s1/native_screenshot.md +++ b/recipe/cua_s1/native_screenshot.md @@ -79,7 +79,7 @@ multimodal adapter and reports prompt-token usage with zero generated tokens. `CUA_S1_GRAPH=1` enables host-feature language Graph replay; `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 +directly into language scratch. Its separate-question language path remains eager; 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. @@ -95,6 +95,81 @@ 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 question batches + +`CUA_S1_MM_BATCH=1` opts into batching questions from one admitted screenshot +request. Unset or `0` retains separate question execution; other values fail at +startup. Set `CUA_S1_MM_PREFIX=off` when enabling it. The batch adapter supports +host or resident image features and independent language/vision Graph controls: + +```sh +export CUA_S1_MM_PREFIX=off +export CUA_S1_MM_BATCH=1 +export CUA_S1_GRAPH=1 +export CUA_S1_VISION_GRAPH=1 +export CUA_S1_DEVICE_FEATURES=1 +``` + +Preparation validates up to eight questions, 4096 tokens per prompt and 26 options +per question. Compatible questions share an exact 64-token-aligned prefix and +pack independent suffixes while retaining the complete prompts' GEMM algorithms, +three rotary axes and attention/GDN boundaries. Unsupported cohorts execute +serially. Every request recomputes prefix tensors; retained allocations, shape +plans and Graphs do not cache screenshot results. Question/option order, prompt +token accounting, CPU letter readout and response interpretation are unchanged. +GPU work drains before the request releases admission, including failures. + +When every retained original algorithm supports the combined row count, each +layer projects `[prefix, suffix 0, ...]` together. Prefix mixers produce that +layer's fresh FP32 GDN state and attention KV before independent suffix mixers +consume them. Unsupported or oversized combinations keep the separate +prefix/suffix passes. This does not join attention across questions or reuse +prefix results from another request. + +A rebuilt CUDA library can batch suffix GDN continuation from the same readonly +FP32 prefix state. This is an optional additive entry point; libraries without +it keep the existing per-suffix continuation loop. Other model adapters retain +their existing dispatch. Qualify BF16 outputs and FP32 final states as well as +end-to-end responses when enabling the rebuilt library: + +```sh +cargo test --release --locked -p omni-qwen3-5-native --test kernels \ + seeded_gdn_batch_matches_independent_outputs_and_fp32_states_after_replay \ + -- --ignored --test-threads=1 --nocapture +``` + +The shared native adapter bounds each of at most 16 prompts to 8192 tokens, +scratch to 8192 rows and compact suffix segments to 4096 rows. Duplicate prefix +rows do not count as simultaneously allocated scratch. Without a shared prefix, +the Cua planner splits groups at 8192 combined tokens. Qualify exact output bits +on the intended GPU/cuBLAS/model combination before deployment; this mode does +not establish general GUI task accuracy or a speedup for every request shape. + +`CUA_S1_RESIZE_WORKERS=1` independently opts the Cua processor into RGB integer +resize with runtime SSE4.1/AVX2 dispatch on supported x86 CPUs and a scalar +fallback. Unset/`0` keeps the existing resize entry point; `1`–`4` cap active CPU +workers per request, and other values fail at startup. Coefficients, accumulation +order, rounding, output pixels and patch layout are preserved. Small work may +remain serial. This switch does not select the accelerated entry point for JEMM, +OmniJev or Jev-VL. Measure CPU contention before increasing the worker count. + +With the pinned model and library variables from the setup above: + +```sh +cargo test -p omni-qwen3-5-native --lib image_preprocess +CUA_S1_MM_BATCH=0 CUA_S1_MM_PREFIX=off \ + cargo test --release --locked -p omni-cua-s1-native --lib \ + question_batch_keeps_logits_response_order_and_changed_pixels \ + -- --ignored --test-threads=1 --nocapture +cargo test --release --locked -p omni-qwen3-5-native --lib \ + multimodal_shared_suffix_batch_keeps_hidden_bits_after_pixel_changes \ + -- --ignored --test-threads=1 --nocapture +``` + +Repeat the Cua GPU regression with each language/vision Graph setting (`0`/`1`). +The native shared-prefix regression exercises both eager and Graph execution, +including batches whose duplicate raw prefixes exceed 8192 total tokens. + ## Request-local language prefix reuse `CUA_S1_MM_PREFIX=shared` opts into sharing the aligned language prefix of all diff --git a/src/backends/cuda/qwen3_5/gdn_prefill.cu b/src/backends/cuda/qwen3_5/gdn_prefill.cu index 7201307f..b867d311 100644 --- a/src/backends/cuda/qwen3_5/gdn_prefill.cu +++ b/src/backends/cuda/qwen3_5/gdn_prefill.cu @@ -665,6 +665,13 @@ __global__ void __launch_bounds__(ST_THREADS) gdn_batch_state(Batch batch) { const auto sequence = batch.sequences[blockIdx.z]; gdn_chunk_state_body(sequence.ws, (sequence.tokens + C - 1) / C, nullptr, nullptr); } +// Every sequence reads the same immutable FP32 prefix; final states are disjoint. +__global__ void __launch_bounds__(ST_THREADS) gdn_batch_seeded_state( + Batch batch, const float* initial, float* finals, int H) { + const auto sequence = batch.sequences[blockIdx.z]; + float* out = finals ? finals + static_cast(blockIdx.z) * H * K * V : nullptr; + gdn_chunk_state_body(sequence.ws, (sequence.tokens + C - 1) / C, initial, out); +} __global__ void __launch_bounds__(OUT_THREADS) gdn_batch_out(Batch batch, bf16* o, int H) { const auto sequence = batch.sequences[blockIdx.z]; if (blockIdx.x >= (sequence.tokens + C - 1) / C) return; @@ -786,9 +793,9 @@ size_t cs1_gdn_batch_workspace_floats(int tokens, int H, int count) { return (Layout(tokens + (count - 1) * (C - 1), H).total + count * 9 * 256 + 3) / 4; } -int cs1_gdn_prefill_batch(const void* q, const void* k, const void* v, const float* g, +static int gdn_batch_run(const void* q, const void* k, const void* v, const float* g, const void* beta, void* o, float* workspace, size_t workspace_bytes, - const size_t* lengths, size_t count, int H, int HK, float scale, void* stream) { + const size_t* lengths, size_t count, int H, int HK, float scale, const void* initial_state, void* final_states, void* stream) { if (!q || !k || !v || !g || !beta || !o || !workspace || !lengths || !count || count > MAX_SEQUENCES || H <= 0 || HK <= 0 || H % HK != 0) return cudaErrorInvalidValue; Batch batch{}; @@ -813,13 +820,39 @@ int cs1_gdn_prefill_batch(const void* q, const void* k, const void* v, const flo return e; }(); if (configured != cudaSuccess) return configured; + if (initial_state) { + static const cudaError_t seeded_configured = cudaFuncSetAttribute( + gdn_batch_seeded_state, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)SMEM2_BYTES); + if (seeded_configured != cudaSuccess) return seeded_configured; + } const auto st = static_cast(stream); gdn_batch_prep<<>>( static_cast(q), static_cast(k), static_cast(v), g, static_cast(beta), batch, H, HK, scale); - gdn_batch_state<<>>(batch); + if (initial_state) { + gdn_batch_seeded_state<<>>( + batch, static_cast(initial_state), static_cast(final_states), H); + } else { + gdn_batch_state<<>>(batch); + } gdn_batch_out<<>>(batch, static_cast(o), H); return cudaGetLastError(); } +int cs1_gdn_prefill_batch(const void* q, const void* k, const void* v, const float* g, + const void* beta, void* o, float* workspace, size_t workspace_bytes, + const size_t* lengths, size_t count, int H, int HK, float scale, void* stream) { + return gdn_batch_run(q, k, v, g, beta, o, workspace, workspace_bytes, + lengths, count, H, HK, scale, nullptr, nullptr, stream); +} + +int cs1_gdn_prefill_seeded_batch(const void* q, const void* k, const void* v, const float* g, + const void* beta, void* o, float* workspace, size_t workspace_bytes, + const size_t* lengths, size_t count, int H, int HK, float scale, + const void* initial_state, void* final_states, void* stream) { + if (!initial_state) return cudaErrorInvalidValue; + return gdn_batch_run(q, k, v, g, beta, o, workspace, workspace_bytes, + lengths, count, H, HK, scale, initial_state, final_states, stream); +} + } // extern "C" diff --git a/src/backends/cuda/qwen3_5/ops.h b/src/backends/cuda/qwen3_5/ops.h index 9ad9c739..de24b141 100644 --- a/src/backends/cuda/qwen3_5/ops.h +++ b/src/backends/cuda/qwen3_5/ops.h @@ -92,6 +92,13 @@ size_t cs1_gdn_batch_workspace_floats(int tokens, int H, int count); int cs1_gdn_prefill_batch(const void* q, const void* k, const void* v, const float* g, const void* beta, void* o, float* workspace, size_t workspace_bytes, const size_t* lengths, size_t count, int H, int HK, float scale, void* stream); +// Optional suffix batching: initial_state is a readonly FP32 [H,128,128] matrix. +// final_states may be null, or disjoint FP32 [count,H,128,128] output matrices. +// All state/input/output/workspace allocations remain live through stream drain. +int cs1_gdn_prefill_seeded_batch(const void* q, const void* k, const void* v, const float* g, + const void* beta, void* o, float* workspace, size_t workspace_bytes, + const size_t* lengths, size_t count, int H, int HK, float scale, + const void* initial_state, void* final_states, void* stream); size_t cs1_gdn_workspace_floats(int T, int H); // Optional no-FLA text precision mode: FP32 normalized Q/K and recurrent state, // BF16 output. Context and scratch remain model-owned; prefix continuation is unsupported. diff --git a/src/models/cua_s1/README.md b/src/models/cua_s1/README.md index 28a39e20..861d3a86 100644 --- a/src/models/cua_s1/README.md +++ b/src/models/cua_s1/README.md @@ -162,6 +162,14 @@ The historical cumulative A800 investigation and frozen-corpus results are recor ## Native screenshot prefix modes +The independent opt-in `CUA_S1_MM_BATCH=1` batches questions within one screenshot +request, including compatible fresh prefixes and packed suffixes. It requires +`CUA_S1_MM_PREFIX=off` and supports host/resident features and independent +language/vision Graph controls. `CUA_S1_RESIZE_WORKERS=1`–`4` separately selects +the Cua processor's integer RGB resize acceleration; both switches default off. +See the [batch and resize controls](../../../recipe/cua_s1/native_screenshot.md#request-local-question-batches) +for bounds, compatibility, qualification and regression commands. + 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. diff --git a/src/models/cua_s1/native/src/image_preprocess.rs b/src/models/cua_s1/native/src/image_preprocess.rs index 72809b2d..03dbfebb 100644 --- a/src/models/cua_s1/native/src/image_preprocess.rs +++ b/src/models/cua_s1/native/src/image_preprocess.rs @@ -1,2 +1,19 @@ //! Cua-S1 retains its original RGB8 processor limits through the shared implementation. pub use omni_qwen3_5_native::image_preprocess::{ProcessedImage, preprocess_rgb8}; + +/// Cua-only explicit accelerated CPU preparation. Other model adapters retain +/// their existing shared entry points and limits. +pub fn preprocess_rgb8_workers( + width: usize, + height: usize, + rgb: &[u8], + workers: usize, +) -> anyhow::Result { + omni_qwen3_5_native::image_preprocess::preprocess_rgb8_with_workers( + width, + height, + rgb, + &omni_qwen3_5_native::image_preprocess::ImageLimits::cua(), + workers, + ) +} diff --git a/src/models/cua_s1/native/src/vision_engine.rs b/src/models/cua_s1/native/src/vision_engine.rs index 941a31cb..7aa98065 100644 --- a/src/models/cua_s1/native/src/vision_engine.rs +++ b/src/models/cua_s1/native/src/vision_engine.rs @@ -1,5 +1,6 @@ //! Native screenshot-to-decision orchestration, with request-local vision reuse. mod prefix; +mod question_batch; pub use crate::vision_processing::PreparedRequest; use crate::vision_processing::{VisionProcessor, probabilities}; use crate::{ @@ -25,6 +26,7 @@ pub struct VisionEngine { language: Model, letters: Vec, device_features: bool, + batch_questions: bool, prefix_mode: prefix::Mode, prefix_storage: Option, } @@ -45,6 +47,20 @@ impl VisionEngine { } }; let prefix_mode = prefix::Mode::from_env(device_features)?; + let batch_questions = match std::env::var("CUA_S1_MM_BATCH") { + Ok(value) => question_batch::enabled(Some(&value), prefix_mode)?, + Err(std::env::VarError::NotPresent) => false, + Err(std::env::VarError::NotUnicode(_)) => { + anyhow::bail!("CUA_S1_MM_BATCH must be unset, 0, or 1") + } + }; + let resize_workers = match std::env::var("CUA_S1_RESIZE_WORKERS") { + Ok(value) => crate::vision_processing::resize_workers(Some(&value))?, + Err(std::env::VarError::NotPresent) => 0, + Err(std::env::VarError::NotUnicode(_)) => { + anyhow::bail!("CUA_S1_RESIZE_WORKERS must be unset or 0..=4") + } + }; let marker: Value = serde_json::from_slice( &std::fs::read(language.join("cua_s1_language_export.json")) .context("export the multimodal language checkpoint first")?, @@ -69,7 +85,10 @@ impl VisionEngine { Ok(enc.get_ids()[0]) }) .collect::>>()?; - let model = Model::load(language, library)?; + let mut model = Model::load(language, library)?; + if batch_questions { + model.enable_multimodal_batching()?; + } ensure!( tokenizer.token_to_id("<|image_pad|>") == model.cfg.image_token_id && model.cfg.image_token_id.is_some(), @@ -78,14 +97,15 @@ impl VisionEngine { let letters = letter_rows(language, &ids, model.cfg.hidden)?; let vision = VisionModel::load(base, adapter, library)?; Ok(Self { - processor: Arc::new(VisionProcessor::new( - tokenizer, - model.cfg.image_token_id.unwrap(), - )), + processor: Arc::new( + VisionProcessor::new(tokenizer, model.cfg.image_token_id.unwrap()) + .with_resize_workers(resize_workers)?, + ), vision, language: model, letters, device_features, + batch_questions, prefix_mode, prefix_storage: None, }) @@ -215,7 +235,28 @@ impl VisionEngine { prepared.prompts.len() == counts.len() && !counts.is_empty(), "question/prompt count mismatch" ); - if self.prefix_mode == prefix::Mode::Matched { + if self.batch_questions { + let groups = question_batch::plan(prepared)?; + if self.device_features { + let features = self.vision.forward_device(&prepared.image)?; + question_batch::execute( + &mut self.language, + &self.letters, + prepared, + &groups, + question_batch::Features::Device(features), + ) + } else { + let features = self.vision.forward(&prepared.image)?; + question_batch::execute( + &mut self.language, + &self.letters, + prepared, + &groups, + question_batch::Features::Host(&features), + ) + } + } else 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 diff --git a/src/models/cua_s1/native/src/vision_engine/question_batch.rs b/src/models/cua_s1/native/src/vision_engine/question_batch.rs new file mode 100644 index 00000000..d664de60 --- /dev/null +++ b/src/models/cua_s1/native/src/vision_engine/question_batch.rs @@ -0,0 +1,122 @@ +//! Bounded question groups within one admitted screenshot request. +use super::*; +use crate::model::MAX_MULTIMODAL_BATCH_TOKENS; +use std::ops::Range; + +pub(super) fn enabled(value: Option<&str>, mode: prefix::Mode) -> Result { + let enabled = match value { + None | Some("0") => false, + Some("1") => true, + _ => anyhow::bail!("CUA_S1_MM_BATCH must be unset, 0, or 1"), + }; + ensure!( + !enabled || mode == prefix::Mode::Off, + "CUA_S1_MM_BATCH requires CUA_S1_MM_PREFIX=off" + ); + Ok(enabled) +} + +fn groups(lengths: &[usize], shared_prefix: bool) -> Result>> { + ensure!( + (1..=8).contains(&lengths.len()), + "expected 1 to 8 questions" + ); + let mut groups = Vec::new(); + let (mut begin, mut tokens) = (0, 0); + for (index, &length) in lengths.iter().enumerate() { + ensure!( + (1..=crate::multimodal::MAX_TOKENS).contains(&length), + "empty or oversized question prompt" + ); + if !shared_prefix && tokens + length > MAX_MULTIMODAL_BATCH_TOKENS { + groups.push(begin..index); + begin = index; + tokens = 0; + } + tokens += length; + } + groups.push(begin..lengths.len()); + Ok(groups) +} + +pub(super) fn plan(prepared: &PreparedRequest) -> Result>> { + ensure!( + prepared.prompts.iter().all(|p| p + .position_ids + .iter() + .all(|axis| axis.len() == p.token_ids.len())), + "position axis length must match prompt tokens" + ); + ensure!( + prepared + .context + .option_counts() + .all(|n| (1..=26).contains(&n)), + "expected 1 to 26 options" + ); + groups( + &prepared + .prompts + .iter() + .map(|p| p.token_ids.len()) + .collect::>(), + prefix::common_prefix(&prepared.prompts) > 0, + ) +} + +pub(super) enum Features<'a> { + Host(&'a [bf16]), + Device(DeviceBf16Matrix<'a>), +} + +pub(super) fn execute( + model: &mut Model, + letters: &[f32], + prepared: &PreparedRequest, + groups: &[Range], + features: Features<'_>, +) -> Result>> { + let counts: Vec<_> = prepared.context.option_counts().collect(); + let mut logits = Vec::with_capacity(prepared.prompts.len()); + for group in groups { + let prompts = &prepared.prompts[group.clone()]; + let hidden = match features { + Features::Host(features) => { + let inputs: Vec<_> = prompts + .iter() + .map(|p| MultimodalInput { + token_ids: &p.token_ids, + image_token_indices: &p.image_token_indices, + image_embeddings: features, + position_ids: [&p.position_ids[0], &p.position_ids[1], &p.position_ids[2]], + }) + .collect(); + model + .forward_multimodal_batch_with_prefix(&inputs, prefix::common_prefix(prompts))? + } + Features::Device(features) => { + let inputs: Vec<_> = prompts + .iter() + .map(|p| DeviceMultimodalInput { + token_ids: &p.token_ids, + image_token_indices: &p.image_token_indices, + image_embeddings: features, + position_ids: [&p.position_ids[0], &p.position_ids[1], &p.position_ids[2]], + }) + .collect(); + model.forward_multimodal_device_batch_with_prefix( + &inputs, + prefix::common_prefix(prompts), + )? + } + }; + for (row, &count) in hidden.iter().zip(&counts[group.clone()]) { + logits.push(VisionEngine::letter_logits(letters, row, count)?); + } + } + Ok(logits) +} + +#[cfg(test)] +#[path = "../../../../../../tests/cua_s1/question_batch.rs"] +mod tests; diff --git a/src/models/cua_s1/native/src/vision_processing.rs b/src/models/cua_s1/native/src/vision_processing.rs index 9d765367..0020ad54 100644 --- a/src/models/cua_s1/native/src/vision_processing.rs +++ b/src/models/cua_s1/native/src/vision_processing.rs @@ -12,6 +12,7 @@ use tokenizers::Tokenizer; pub struct VisionProcessor { tokenizer: Tokenizer, image_token: u32, + resize_workers: usize, } pub struct PreparedRequest { @@ -30,9 +31,18 @@ impl VisionProcessor { Self { tokenizer, image_token, + resize_workers: 0, } } + /// Explicit CPU resize scheduling for this processor; zero retains the + /// existing shared serial entry point. All child work joins during prepare. + pub fn with_resize_workers(mut self, workers: usize) -> Result { + ensure!(workers <= 4, "resize workers must be 0..=4"); + self.resize_workers = workers; + Ok(self) + } + /// Validate every question and prompt before submitting model work. pub fn prepare( &self, @@ -42,7 +52,16 @@ impl VisionProcessor { questions: &[Question], ) -> Result { crate::multimodal::validate_questions(questions)?; - let image = preprocess_rgb8(width, height, rgb)?; + let image = if self.resize_workers == 0 { + preprocess_rgb8(width, height, rgb)? + } else { + crate::image_preprocess::preprocess_rgb8_workers( + width, + height, + rgb, + self.resize_workers, + )? + }; let prompts = questions .iter() .map(|q| prepare_prompt(&self.tokenizer, q, image.image_grid_thw, self.image_token)) @@ -90,3 +109,17 @@ pub(crate) fn probabilities(logits: &[f32]) -> Result> { let total: f64 = exps.iter().sum(); Ok(exps.iter().map(|x| (x / total) as f32).collect()) } + +/// Serving switch is local to the Cua screenshot adapter; no process-wide pool +/// or behavior change in other shared-preprocessor consumers. +pub(crate) fn resize_workers(value: Option<&str>) -> Result { + match value { + None => Ok(0), + Some(v @ ("0" | "1" | "2" | "3" | "4")) => Ok(usize::from(v.as_bytes()[0] - b'0')), + _ => anyhow::bail!("CUA_S1_RESIZE_WORKERS must be unset or 0..=4"), + } +} + +#[cfg(test)] +#[path = "../../../../../tests/cua_s1/vision_processing.rs"] +mod tests; diff --git a/src/models/qwen3_5/native/src/cuda.rs b/src/models/qwen3_5/native/src/cuda.rs index 612d49da..a74ba136 100644 --- a/src/models/qwen3_5/native/src/cuda.rs +++ b/src/models/qwen3_5/native/src/cuda.rs @@ -149,6 +149,27 @@ pub struct BatchedGdn { ) -> c_int, } +/// Optional independent suffixes seeded from one readonly FP32 prefix state. +/// Final states, when requested, occupy disjoint sequence-major allocations. +pub type SeededGdnBatch = unsafe extern "C" fn( + *const c_void, + *const c_void, + *const c_void, + *const f32, + *const c_void, + *mut c_void, + *mut f32, + usize, + *const usize, + usize, + c_int, + c_int, + f32, + *const c_void, + *mut c_void, + Stream, +) -> c_int; + /// Additive dimension-aware vision entry points. Libraries at the required ABI may omit these optional symbols. pub struct VisionV2 { pub position: unsafe extern "C" fn( @@ -463,6 +484,16 @@ impl Api { } } + pub fn seeded_gdn_batch(&self) -> Option { + // SAFETY: the optional signature matches ops.h; the API retains the library. + unsafe { + self._lib + .get::(b"cs1_gdn_prefill_seeded_batch\0") + .ok() + .map(|symbol| *symbol) + } + } + pub fn shared_suffix(&self) -> Option { // SAFETY: additive signatures match ops.h; the API retains the library. unsafe { @@ -679,6 +710,14 @@ pub struct DeviceBf16Matrix<'a> { } impl DeviceBf16Matrix<'_> { + /// Request-local identity only: producer contents may change between calls. + pub(crate) fn same_storage(&self, other: &Self) -> bool { + std::ptr::eq(self.buffer, other.buffer) + && self.layout.rows == other.layout.rows + && self.layout.columns == other.layout.columns + && self.layout.device == other.layout.device + } + pub fn layout(&self) -> DeviceFeatureLayout { self.layout } diff --git a/src/models/qwen3_5/native/src/image_preprocess.rs b/src/models/qwen3_5/native/src/image_preprocess.rs index a10ca2a4..0bedafbe 100644 --- a/src/models/qwen3_5/native/src/image_preprocess.rs +++ b/src/models/qwen3_5/native/src/image_preprocess.rs @@ -10,6 +10,9 @@ use std::borrow::Cow; +#[cfg(target_arch = "x86_64")] +mod x86; + use anyhow::{Result, ensure}; const PATCH_SIZE: usize = 16; @@ -84,6 +87,30 @@ pub fn preprocess_rgb8_with_limits( height: usize, rgb: &[u8], limits: &ImageLimits, +) -> Result { + preprocess_impl(width, height, rgb, limits, 0) +} + +/// Explicit RGB convolution/row scheduling opt-in. Existing entry points retain +/// their channel-wise serial resize. Uses at most four active CPU workers and +/// preserves every per-channel integer multiply/add, shift and saturation. +pub fn preprocess_rgb8_with_workers( + width: usize, + height: usize, + rgb: &[u8], + limits: &ImageLimits, + workers: usize, +) -> Result { + ensure!((1..=4).contains(&workers), "resize workers must be 1..=4"); + preprocess_impl(width, height, rgb, limits, workers) +} + +fn preprocess_impl( + width: usize, + height: usize, + rgb: &[u8], + limits: &ImageLimits, + workers: usize, ) -> Result { ensure!( limits.min_pixels >= FACTOR * FACTOR @@ -148,28 +175,68 @@ pub fn preprocess_rgb8_with_limits( if resized_width != width { let axis = AxisWeights::new(width, resized_width); let mut horizontal = vec![0; horizontal_len]; - for y in 0..height { - for (x, kernel) in axis.kernels.iter().enumerate() { - for channel in 0..3 { - horizontal[(y * resized_width + x) * 3 + channel] = - axis.apply(kernel, |source_x| rgb[(y * width + source_x) * 3 + channel]); + if workers == 0 { + for y in 0..height { + for (x, kernel) in axis.kernels.iter().enumerate() { + for channel in 0..3 { + horizontal[(y * resized_width + x) * 3 + channel] = axis + .apply(kernel, |source_x| rgb[(y * width + source_x) * 3 + channel]); + } } } + } else { + resize_rows(&mut horizontal, resized_width, workers, |y, row| { + #[cfg(target_arch = "x86_64")] + if std::arch::is_x86_feature_detected!("sse4.1") { + // SAFETY: checked feature, complete source row and axis/output extents. + unsafe { + x86::horizontal(&rgb[y * width * 3..(y + 1) * width * 3], row, &axis); + } + return; + } + for (pixel, kernel) in row.as_chunks_mut::<3>().0.iter_mut().zip(&axis.kernels) { + pixel.copy_from_slice(&axis.apply_rgb(kernel, |source_x| { + let i = (y * width + source_x) * 3; + [rgb[i], rgb[i + 1], rgb[i + 2]] + })); + } + }); } resized = Cow::Owned(horizontal); } if resized_height != height { let axis = AxisWeights::new(height, resized_height); let mut vertical = vec![0; resized_len]; - for (y, kernel) in axis.kernels.iter().enumerate() { - for x in 0..resized_width { - for channel in 0..3 { - vertical[(y * resized_width + x) * 3 + channel] = axis - .apply(kernel, |source_y| { - resized[(source_y * resized_width + x) * 3 + channel] - }); + if workers == 0 { + for (y, kernel) in axis.kernels.iter().enumerate() { + for x in 0..resized_width { + for channel in 0..3 { + vertical[(y * resized_width + x) * 3 + channel] = axis + .apply(kernel, |source_y| { + resized[(source_y * resized_width + x) * 3 + channel] + }); + } } } + } else { + resize_rows(&mut vertical, resized_width, workers, |y, row| { + let kernel = &axis.kernels[y]; + #[cfg(target_arch = "x86_64")] + if std::arch::is_x86_feature_detected!("avx2") { + // SAFETY: checked feature, each source row matches row.len(), + // and the axis builder bounds every source row in this kernel. + unsafe { + x86::vertical(&resized, row, kernel, axis.precision); + } + return; + } + for (x, pixel) in row.as_chunks_mut::<3>().0.iter_mut().enumerate() { + pixel.copy_from_slice(&axis.apply_rgb(kernel, |source_y| { + let i = (source_y * resized_width + x) * 3; + [resized[i], resized[i + 1], resized[i + 2]] + })); + } + }); } resized = Cow::Owned(vertical); } @@ -284,6 +351,18 @@ impl AxisWeights { Self { kernels, precision } } + fn apply_rgb(&self, kernel: &Kernel, pixel: impl Fn(usize) -> [u8; 3]) -> [u8; 3] { + let mut accumulators = [1_i32 << (self.precision - 1); 3]; + for (offset, &weight) in kernel.weights.iter().enumerate() { + let rgb = pixel(kernel.start + offset); + let weight = i32::from(weight); + for (sum, value) in accumulators.iter_mut().zip(rgb) { + *sum += i32::from(value) * weight; + } + } + accumulators.map(|sum| (sum >> self.precision).clamp(0, 255) as u8) + } + fn apply(&self, kernel: &Kernel, pixel: impl Fn(usize) -> u8) -> u8 { let mut accumulator = 1_i32 << (self.precision - 1); for (offset, &weight) in kernel.weights.iter().enumerate() { @@ -293,6 +372,42 @@ impl AxisWeights { } } +/// Disjoint output rows: coefficient and source reads are immutable. The caller +/// processes the first chunk while at most workers-1 scoped threads process the +/// rest; both passes join before their intermediate buffer can be reused. +fn resize_rows( + output: &mut [u8], + width: usize, + workers: usize, + operation: impl Fn(usize, &mut [u8]) + Sync, +) { + let row_bytes = width * 3; + let height = output.len() / row_bytes; + let workers = workers.min(height); + if workers <= 1 || output.len() < 256 * 1024 { + for (y, row) in output.chunks_exact_mut(row_bytes).enumerate() { + operation(y, row); + } + return; + } + let chunk_rows = height.div_ceil(workers); + std::thread::scope(|scope| { + let mut chunks = output.chunks_mut(chunk_rows * row_bytes); + let first = chunks.next().unwrap(); + let op = &operation; + for (index, chunk) in chunks.enumerate() { + scope.spawn(move || { + for (y, row) in chunk.chunks_exact_mut(row_bytes).enumerate() { + op((index + 1) * chunk_rows + y, row); + } + }); + } + for (y, row) in first.chunks_exact_mut(row_bytes).enumerate() { + operation(y, row); + } + }); +} + fn cubic(x: f64) -> f64 { let x = x.abs(); const A: f64 = -0.5; diff --git a/src/models/qwen3_5/native/src/image_preprocess/x86.rs b/src/models/qwen3_5/native/src/image_preprocess/x86.rs new file mode 100644 index 00000000..c525b633 --- /dev/null +++ b/src/models/qwen3_5/native/src/image_preprocess/x86.rs @@ -0,0 +1,76 @@ +//! Explicit resize path only: preserve independent integer channel accumulators. +use super::{AxisWeights, Kernel}; +use std::arch::x86_64::*; + +/// # Safety +/// The caller checks SSE4.1. Source spans and output RGB rows match `axis`. +#[target_feature(enable = "sse4.1")] +pub(super) unsafe fn horizontal(source: &[u8], output: &mut [u8], axis: &AxisWeights) { + // Intrinsics operate on registers under the checked target feature. + // RGB reads/stores use checked slices and never read a fourth byte. + let zero = _mm_setzero_si128(); + let maximum = _mm_set1_epi32(255); + let initial = _mm_set1_epi32(1 << (axis.precision - 1)); + let shift = _mm_cvtsi32_si128(axis.precision as i32); + for (pixel, kernel) in output.as_chunks_mut::<3>().0.iter_mut().zip(&axis.kernels) { + let mut sum = initial; + for (offset, &weight) in kernel.weights.iter().enumerate() { + let i = (kernel.start + offset) * 3; + let rgb = u32::from_le_bytes([source[i], source[i + 1], source[i + 2], 0]); + let values = _mm_cvtepu8_epi32(_mm_cvtsi32_si128(rgb as i32)); + sum = _mm_add_epi32( + sum, + _mm_mullo_epi32(values, _mm_set1_epi32(i32::from(weight))), + ); + } + let shifted = _mm_min_epi32(_mm_sra_epi32(sum, shift), maximum); + let words = _mm_packus_epi32(shifted, zero); + let bytes = (_mm_cvtsi128_si32(_mm_packus_epi16(words, zero)) as u32).to_le_bytes(); + pixel.copy_from_slice(&bytes[..3]); + } +} + +/// # Safety +/// The caller checks AVX2. Each source row has output.len() bytes, and every +/// kernel row lies inside source. Source and output are distinct allocations. +#[target_feature(enable = "avx2")] +pub(super) unsafe fn vertical(source: &[u8], output: &mut [u8], kernel: &Kernel, precision: u32) { + let stride = output.len(); + let mut lane = 0; + // SAFETY: lane+8<=stride bounds every 8-byte load/store within its checked + // row. The axis builder bounds source rows. Each lane follows scalar tap order. + unsafe { + let initial = _mm256_set1_epi32(1 << (precision - 1)); + let shift = _mm256_set1_epi32(precision as i32); + let maximum = _mm256_set1_epi32(255); + let zero = _mm_setzero_si128(); + while lane + 8 <= stride { + let mut sum = initial; + for (offset, &weight) in kernel.weights.iter().enumerate() { + let i = (kernel.start + offset) * stride + lane; + let values = _mm256_cvtepu8_epi32(_mm_loadl_epi64(source.as_ptr().add(i).cast())); + sum = _mm256_add_epi32( + sum, + _mm256_mullo_epi32(values, _mm256_set1_epi32(i32::from(weight))), + ); + } + let shifted = _mm256_min_epi32(_mm256_srav_epi32(sum, shift), maximum); + let words = _mm_packus_epi32( + _mm256_castsi256_si128(shifted), + _mm256_extracti128_si256(shifted, 1), + ); + _mm_storel_epi64( + output.as_mut_ptr().add(lane).cast(), + _mm_packus_epi16(words, zero), + ); + lane += 8; + } + } + for (index, pixel) in output.iter_mut().enumerate().skip(lane) { + let mut sum = 1_i32 << (precision - 1); + for (offset, &weight) in kernel.weights.iter().enumerate() { + sum += i32::from(source[(kernel.start + offset) * stride + index]) * i32::from(weight); + } + *pixel = (sum >> precision).clamp(0, 255) as u8; + } +} diff --git a/src/models/qwen3_5/native/src/model.rs b/src/models/qwen3_5/native/src/model.rs index 05f8c6cb..e5983691 100644 --- a/src/models/qwen3_5/native/src/model.rs +++ b/src/models/qwen3_5/native/src/model.rs @@ -13,6 +13,9 @@ //! once and every branch from the state after it (see `SharedPrompts`). mod matched; +mod multimodal_batch; +mod multimodal_shared; +pub use multimodal_batch::MAX_MULTIMODAL_BATCH_TOKENS; use std::collections::{HashMap, VecDeque}; use std::ffi::c_void; @@ -705,14 +708,20 @@ pub struct PrefixState { bytes: usize, } -/// A packed suffix batch uses one freshly captured prefix and original algorithm. +/// A suffix batch consumes a captured prefix, or produces it in its leading sequence, +/// while retaining the complete prompts' original GEMM algorithms. struct SuffixBatch<'a> { prefix: &'a PrefixState, ops: cuda::SharedSuffix, gemm: cuda::ReuseGemm, original_rows: usize, + seeded_gdn: Option, + /// The leading sequence produces this layer's prefix before suffix mixing. + include_prefix: bool, } +const MAX_SUFFIX_BATCH_ROWS: usize = 4096; + /// Bound only the suffix work; oversized suffixes retain individual windows. fn suffix_segments(indices: &[usize], lengths: &[usize], prefix: usize) -> Vec> { let mut groups = Vec::new(); @@ -720,7 +729,7 @@ fn suffix_segments(indices: &[usize], lengths: &[usize], prefix: usize) -> Vec 4096) { + if !group.is_empty() && (group.len() == 16 || tokens + suffix > MAX_SUFFIX_BATCH_ROWS) { groups.push(std::mem::take(&mut group)); tokens = 0; } @@ -855,6 +864,7 @@ pub struct Model { fp32_gdn: Option, batched_gdn: Option, packed_gemm: Option, + multimodal_packed: Option, /// Buffers for the largest packed token count so far; grows as needed. scratch: Option, /// The request and group prefix states of `forward_shared`; grow as needed. @@ -867,6 +877,8 @@ pub struct Model { graph_stats: GraphStats, graphs: VecDeque<(Vec, cuda::Graph)>, mm_graphs: VecDeque<(Vec, cuda::Graph)>, + mm_batch_windows: VecDeque, + mm_batch_prefix: Option<(usize, PrefixState)>, // The stream retires after every model allocation, including adapter resources. _stream_owner: ModelStream, } @@ -889,6 +901,8 @@ impl Drop for Model { let _ = cuda::synchronize(self.stream); self.graphs.clear(); self.mm_graphs.clear(); + self.mm_batch_windows.clear(); + self.mm_batch_prefix = None; self.scratch = None; self.shared = [None, None]; // Retire the FP32 cuBLAS context while its last stream is still valid. @@ -1032,6 +1046,7 @@ impl Model { fp32_gdn: None, batched_gdn: None, packed_gemm: None, + multimodal_packed: None, scratch: None, shared: [None, None], vocab_buffers: None, @@ -1043,6 +1058,8 @@ impl Model { graphs: VecDeque::new(), _stream_owner: stream_owner, mm_graphs: VecDeque::new(), + mm_batch_windows: VecDeque::new(), + mm_batch_prefix: None, }; Ok(model) } @@ -1050,7 +1067,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.mm_batch_windows.len(), ..self.graph_stats } } @@ -1059,7 +1076,10 @@ impl Model { /// Select before the first forward; other model consumers retain separate GEMMs. pub fn enable_packed_projections(&mut self) -> Result<()> { ensure!( - self.scratch.is_none() && self.reference.is_none() && self.lora.is_none(), + self.scratch.is_none() + && self.reference.is_none() + && self.lora.is_none() + && self.multimodal_packed.is_none(), "select packed projections on a merged model before the first forward" ); self.packed_gemm = Some(cuda::api().packed_gemm()?); @@ -1259,9 +1279,22 @@ impl Model { w: &Tensor, y: usize, lengths: &[usize], + ) -> Result<()> { + self.gemm_sequences_with(s, x, w, y, lengths, self.packed_gemm) + } + + #[allow(clippy::too_many_arguments)] + fn gemm_sequences_with( + &self, + s: &Scratch, + x: usize, + w: &Tensor, + y: usize, + lengths: &[usize], + packed: Option, ) -> Result<()> { let (n, k) = (w.shape[0], w.shape[1]); - if let Some(packed) = self.packed_gemm + if let Some(packed) = packed && self.reference.is_none() && self.lora.is_none() { @@ -1332,11 +1365,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.mm_batch_windows.is_empty() + { self.graph_stats.invalidations += 1; } self.graphs.clear(); self.mm_graphs.clear(); + self.mm_batch_windows.clear(); self.scratch = None; self.scratch = Some(Scratch::new( &self.cfg, @@ -1627,6 +1664,8 @@ impl Model { ops: ops.unwrap(), gemm: reuse, original_rows: original_rows[indices[0]], + seeded_gdn: None, + include_prefix: false, }), )?; let mut end = 0; @@ -1753,6 +1792,7 @@ impl Model { self.graph_enabled = false; self.graphs.clear(); self.mm_graphs.clear(); + self.mm_batch_windows.clear(); } } Ok(()) @@ -2772,6 +2812,20 @@ impl Model { None }; let t: usize = lengths.iter().sum(); + let fresh_prefix = suffix.is_some_and(|batch| batch.include_prefix); + let prefix_rows = suffix + .filter(|batch| batch.include_prefix) + .map_or(0, |batch| batch.prefix.len); + ensure!( + !fresh_prefix + || (custom_positions + && !bidirectional + && prefix_rows >= 64 + && prefix_rows.is_multiple_of(64) + && (2..=17).contains(&lengths.len()) + && lengths[0] == prefix_rows), + "invalid combined prefix/suffix layout" + ); let mut start = 0; let sequences: Vec<(usize, i32)> = lengths .iter() @@ -2814,6 +2868,10 @@ impl Model { "packed suffix projection", ) } + } else if custom_positions && lengths.len() > 1 && self.multimodal_packed.is_some() { + // Preserve each original prompt's algorithm for every projection, + // including input and gate/up, instead of selecting at combined M. + self.gemm_sequences_with(s, x, weight, y, lengths, self.multimodal_packed) } else if let Some(group) = group { self.gemm_group(s, x, weight, y, t, lengths, group) } else { @@ -2858,9 +2916,20 @@ impl Model { let b = z + vd * BF16; let a = b + hv * BF16; unsafe { + if fresh_prefix { + // Capture current-layer pre-convolution history before suffixes read it. + cuda::copy_dd( + suffix.unwrap().prefix.conv_tail[linear_index].at(0), + p(s.gdn_in + (prefix_rows - 3) * w.gdn_in * BF16), + 3 * w.gdn_in * BF16, + st, + )?; + } for &(offset, length) in &sequences { check( - if let Some(suffix) = suffix { + if let Some(suffix) = + suffix.filter(|batch| !batch.include_prefix || offset != 0) + { (suffix.ops.conv)( p(s.gdn_in + offset * w.gdn_in * BF16), suffix.prefix.conv_tail[linear_index].at(0), @@ -2907,17 +2976,41 @@ impl Model { "gdn gates", )?; if let Some(suffix) = suffix { - for &(offset, length) in &sequences { + if fresh_prefix { check( (cuda::api().cs1_gdn_prefill_x)( - p(s.lq + offset * kd * BF16), - p(s.lk + offset * kd * BF16), - p(s.lv + offset * vd * BF16), - p(s.g + offset * hv * F32).cast(), - p(s.beta + offset * hv * BF16), - p(s.lo + offset * vd * BF16), + p(s.lq), + p(s.lk), + p(s.lv), + p(s.g).cast(), + p(s.beta), + p(s.lo), p(s.workspace).cast(), - length, + prefix_rows as i32, + hv as i32, + cfg.lin_k_heads as i32, + (cfg.lin_k_dim as f32).powf(-0.5), + std::ptr::null(), + suffix.prefix.gdn_state[linear_index].at(0), + st, + ), + "current-layer prefix GDN", + )?; + } + let suffix_lengths = &lengths[usize::from(fresh_prefix)..]; + if let Some(seeded) = suffix.seeded_gdn { + check( + seeded( + p(s.lq + prefix_rows * kd * BF16), + p(s.lk + prefix_rows * kd * BF16), + p(s.lv + prefix_rows * vd * BF16), + p(s.g + prefix_rows * hv * F32).cast(), + p(s.beta + prefix_rows * hv * BF16), + p(s.lo + prefix_rows * vd * BF16), + p(s.workspace).cast(), + s.workspace_bytes, + suffix_lengths.as_ptr(), + suffix_lengths.len(), hv as i32, cfg.lin_k_heads as i32, (cfg.lin_k_dim as f32).powf(-0.5), @@ -2925,8 +3018,30 @@ impl Model { std::ptr::null_mut(), st, ), - "suffix GDN prefill", + "seeded suffix GDN batch", )?; + } else { + for &(offset, length) in &sequences[usize::from(fresh_prefix)..] { + check( + (cuda::api().cs1_gdn_prefill_x)( + p(s.lq + offset * kd * BF16), + p(s.lk + offset * kd * BF16), + p(s.lv + offset * vd * BF16), + p(s.g + offset * hv * F32).cast(), + p(s.beta + offset * hv * BF16), + p(s.lo + offset * vd * BF16), + p(s.workspace).cast(), + length, + 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), + std::ptr::null_mut(), + st, + ), + "suffix GDN prefill", + )?; + } } } else if let Some(gdn) = &self.batched_gdn && (2..=16).contains(&lengths.len()) @@ -3025,8 +3140,18 @@ impl Model { ld, fa.q_norm.ptr, fa.k_norm.ptr, - p(cos), - p(sin), + p(cos + + if custom_positions { + offset * cfg.rotary_half * BF16 + } else { + 0 + }), + p(sin + + if custom_positions { + offset * cfg.rotary_half * BF16 + } else { + 0 + }), p(s.aq + offset * cfg.heads * cfg.head_dim * BF16), p(s.agate + offset * cfg.heads * cfg.head_dim * BF16), p(s.ak + offset * cfg.kv_heads * cfg.head_dim * BF16), @@ -3040,9 +3165,30 @@ impl Model { ), "attention prep", )?; + if fresh_prefix && offset == 0 { + let state = suffix.unwrap().prefix; + let bytes = prefix_rows * cfg.kv_heads * cfg.head_dim * BF16; + cuda::copy_dd( + state.attn_kv[attention_index].0.at(0), + p(s.ak), + bytes, + st, + )?; + cuda::copy2d( + state.attn_kv[attention_index].1.at(0), + cfg.kv_heads * cfg.head_dim * BF16, + p(v), + w.attn_in * BF16, + cfg.kv_heads * cfg.head_dim * BF16, + prefix_rows, + st, + )?; + } let scale = (cfg.head_dim as f32).powf(-0.5); check( - if let Some(suffix) = suffix { + if let Some(suffix) = + suffix.filter(|batch| !batch.include_prefix || offset != 0) + { (suffix.ops.attention)( p(s.aq + offset * cfg.heads * cfg.head_dim * BF16), p(s.ak + offset * cfg.kv_heads * cfg.head_dim * BF16), diff --git a/src/models/qwen3_5/native/src/model/multimodal_batch.rs b/src/models/qwen3_5/native/src/model/multimodal_batch.rs new file mode 100644 index 00000000..f3f24d1d --- /dev/null +++ b/src/models/qwen3_5/native/src/model/multimodal_batch.rs @@ -0,0 +1,310 @@ +//! Independent multimodal prompts, sharing weights and launches within one call. +use super::*; + +/// Maximum scratch rows or tokens in one unshared batch, including image placeholders. +pub const MAX_MULTIMODAL_BATCH_TOKENS: usize = 8192; + +fn batch_lengths(inputs: &[MultimodalInput<'_>], prefix: usize) -> Result> { + ensure!( + (1..=16).contains(&inputs.len()), + "multimodal batch requires 1..=16 prompts" + ); + let mut total = 0usize; + let lengths = inputs + .iter() + .map(|input| { + let n = input.token_ids.len(); + ensure!( + (1..=MAX_MULTIMODAL_BATCH_TOKENS).contains(&n), + "empty or oversized multimodal batch prompt" + ); + total = total + .checked_add(n) + .context("multimodal batch token overflow")?; + ensure!( + prefix > 0 || total <= MAX_MULTIMODAL_BATCH_TOKENS, + "multimodal batch exceeds token budget" + ); + Ok(n) + }) + .collect::>>()?; + Ok(lengths) +} + +impl Model { + /// Enable bounded multimodal packing before the first forward. Other callers + /// keep their existing projection/conv dispatch; only independent multimodal + /// batches use the packed GEMM path and the existing batched GDN operator. + pub fn enable_multimodal_batching(&mut self) -> Result<()> { + ensure!( + self.scratch.is_none() && self.packed_gemm.is_none(), + "select multimodal batching before the first forward, without text packing" + ); + self.ensure_fixed_profile()?; + let packed = cuda::api().packed_gemm()?; + let gdn = cuda::api().batched_gdn()?; + self.multimodal_packed = Some(packed); + self.batched_gdn = Some(gdn); + Ok(()) + } + + /// Independent causal multimodal prompts, in input order. Each prompt retains + /// its original GEMM algorithms, T/H/W positions and attention/GDN boundaries. + /// Supports at most 16 prompts and 8192 combined tokens on merged/default + /// numerics. Call `enable_multimodal_batching` before any execution. + pub fn forward_multimodal_batch( + &mut self, + inputs: &[MultimodalInput<'_>], + ) -> Result>> { + self.forward_multimodal_batch_with_prefix(inputs, 0) + } + + /// Share an exact chunk-aligned prefix within compatible cohorts, then pack + /// suffixes retaining the complete prompts' algorithms. Unsupported cohorts + /// retain serial execution. Prefix image values and all three rotary axes + /// must agree; every request computes fresh prefix state. Each of at most + /// 16 prompts is bounded to 8192 tokens. Shared execution allocates at most + /// 8192 scratch rows; raw duplicate prefix rows do not consume that budget. + pub fn forward_multimodal_batch_with_prefix( + &mut self, + inputs: &[MultimodalInput<'_>], + prefix_len: usize, + ) -> Result>> { + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + self.run_multimodal_batch(inputs, None, prefix_len) + })); + self.finish_multimodal(result) + } + + /// Independent prompts using borrowed adapted device rows. The producer + /// finishes writes before entry; all consumer work drains on return or unwind. + pub fn forward_multimodal_device_batch( + &mut self, + inputs: &[DeviceMultimodalInput<'_>], + ) -> Result>> { + self.forward_multimodal_device_batch_with_prefix(inputs, 0) + } + + /// Device counterpart of `forward_multimodal_batch_with_prefix`. Sharing + /// image rows requires the same borrowed allocation for the shared prefix. + /// Captures retain only model-owned scratch/prefix addresses. + pub fn forward_multimodal_device_batch_with_prefix( + &mut self, + inputs: &[DeviceMultimodalInput<'_>], + prefix_len: usize, + ) -> Result>> { + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + let image_token = self + .cfg + .image_token_id + .context("checkpoint has no image_token_id")?; + for input in inputs { + input.validate( + self.cfg.hidden, + self.embed.shape[0], + image_token, + self.cfg.max_positions, + 0, + )?; + } + let host: Vec<_> = inputs + .iter() + .map(|input| MultimodalInput { + token_ids: input.token_ids, + image_token_indices: input.image_token_indices, + image_embeddings: &[], + position_ids: input.position_ids, + }) + .collect(); + let devices: Vec<_> = inputs.iter().map(|input| input.image_embeddings).collect(); + self.run_multimodal_batch(&host, Some(&devices), prefix_len) + })); + self.finish_multimodal(result) + } + + fn finish_multimodal(&self, result: std::thread::Result>) -> Result { + let synced = self.synchronize(); + match result { + Ok(value) => { + synced?; + value + } + Err(panic) => std::panic::resume_unwind(panic), + } + } + + pub(super) fn run_multimodal_batch( + &mut self, + inputs: &[MultimodalInput<'_>], + devices: Option<&[cuda::DeviceBf16Matrix<'_>]>, + prefix_len: usize, + ) -> Result>> { + ensure!( + self.multimodal_packed.is_some(), + "enable multimodal batching before execution" + ); + self.ensure_fixed_profile()?; + ensure!( + self.matched_ref == 0, + "multimodal batching does not use a matched-prefix handle" + ); + let lengths = batch_lengths(inputs, prefix_len)?; + let image_token = self + .cfg + .image_token_id + .context("checkpoint has no image_token_id")?; + if devices.is_none() { + for input in inputs { + input.validate( + self.cfg.hidden, + self.embed.shape[0], + image_token, + self.cfg.max_positions, + )?; + } + } + if prefix_len > 0 { + super::multimodal_shared::common_header(inputs, devices, prefix_len, self.cfg.hidden)?; + } + cuda::set_device(0)?; + if let Some(devices) = devices { + ensure!(devices.len() == inputs.len(), "device batch count mismatch"); + for (index, features) in devices.iter().enumerate() { + // Immutable producer views can repeat across questions. Scan each + // allocation/device/extent once in this call, never across calls. + if !devices[..index] + .iter() + .any(|previous| features.same_storage(previous)) + { + features.validate_finite(self.stream)?; + } + } + } + if prefix_len > 0 { + if let Some(rows) = self.try_multimodal_shared(inputs, devices, &lengths, prefix_len)? { + return Ok(rows); + } + return self.serial_multimodal(inputs, devices); + } + let t = lengths.iter().sum(); + self.prepare_scratch(t)?; + let s = self.scratch.as_ref().unwrap(); + let positions: [Vec; 3] = std::array::from_fn(|axis| { + inputs + .iter() + .flat_map(|input| input.position_ids[axis].iter().copied()) + .collect() + }); + self.upload_positions(s, [&positions[0], &positions[1], &positions[2]])?; + let ids: Vec<_> = inputs + .iter() + .flat_map(|input| input.token_ids.iter().copied()) + .collect(); + self.embed_tokens(s, &ids)?; + let mut offset = 0; + for (index, input) in inputs.iter().enumerate() { + self.stage_multimodal_rows( + s, + input, + devices.map(|d| d[index]), + 0..input.token_ids.len(), + offset, + )?; + offset += input.token_ids.len(); + } + if self.graph_enabled { + if let Some((_, graph)) = self.mm_graphs.iter().find(|(shape, _)| *shape == lengths) { + 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=multimodal event=replay tokens={t} sequences={}", + lengths.len() + ); + } + } else { + self.run(s, &lengths, true, false)?; + cuda::synchronize(self.stream)?; + let captured = + cuda::Graph::capture(self.stream, || self.run(s, &lengths, true, false)); + self.cache_graph(&lengths, true, captured)?; + } + } else { + self.run(s, &lengths, true, false)?; + } + let s = self.scratch.as_ref().unwrap(); + let mut end = 0; + lengths + .iter() + .map(|&n| { + end += n; + self.last_hidden(s, end) + }) + .collect() + } + + pub(super) fn stage_multimodal_rows( + &self, + s: &Scratch, + input: &MultimodalInput<'_>, + device: Option>, + range: std::ops::Range, + offset: usize, + ) -> Result<()> { + let indices = input.image_token_indices; + let first = indices.partition_point(|&i| i < range.start); + let limit = indices.partition_point(|&i| i < range.end); + let row_bytes = self.cfg.hidden * BF16; + if let Some(device) = device { + let mut begin = first; + while begin < limit { + let mut end = begin + 1; + while end < limit && indices[end] == indices[end - 1] + 1 { + end += 1; + } + // SAFETY: complete input metadata and source extents were checked; + // the validated window lies in owned scratch. Borrow lasts to drain. + unsafe { + cuda::copy_dd( + s.at(s.res + (offset + indices[begin] - range.start) * row_bytes), + device.at(begin * row_bytes), + (end - begin) * row_bytes, + self.stream, + )?; + } + begin = end; + } + Ok(()) + } else { + let indices: Vec<_> = indices[first..limit] + .iter() + .map(|&i| i - range.start) + .collect(); + let window = MultimodalInput { + token_ids: &input.token_ids[range.clone()], + image_token_indices: &indices, + image_embeddings: &input.image_embeddings + [first * self.cfg.hidden..limit * self.cfg.hidden], + position_ids: input.position_ids.map(|axis| &axis[range.clone()]), + }; + self.overwrite_image_rows(s, &window, offset) + } + } + + pub(super) fn serial_multimodal( + &mut self, + inputs: &[MultimodalInput<'_>], + devices: Option<&[cuda::DeviceBf16Matrix<'_>]>, + ) -> Result>> { + (0..inputs.len()) + .map(|i| { + self.run_multimodal_batch(&inputs[i..i + 1], devices.map(|d| &d[i..i + 1]), 0) + .map(|mut rows| rows.pop().unwrap()) + }) + .collect() + } +} + +#[cfg(test)] +#[path = "../../../../../../tests/qwen3_5/multimodal_batch.rs"] +mod tests; diff --git a/src/models/qwen3_5/native/src/model/multimodal_shared.rs b/src/models/qwen3_5/native/src/model/multimodal_shared.rs new file mode 100644 index 00000000..f5c04b1d --- /dev/null +++ b/src/models/qwen3_5/native/src/model/multimodal_shared.rs @@ -0,0 +1,579 @@ +//! Fresh request-local multimodal prefixes with independent compact suffixes. +use super::*; + +#[derive(Clone, Debug, PartialEq, Eq)] +pub(super) enum PassKey { + Prefix { + prefix: usize, + reference: usize, + }, + Window { + prefix: usize, + suffix: usize, + reference: usize, + }, + Compact { + prefix: usize, + lengths: Vec, + reference: usize, + }, + Batch { + prefix: usize, + lengths: Vec, + reference: usize, + }, +} + +pub(super) struct Entry { + pub(super) key: PassKey, + graph: cuda::Graph, +} + +fn compact_rows(lengths: &[usize], prefix: usize) -> Option { + if !(2..=16).contains(&lengths.len()) + || prefix < 64 + || !prefix.is_multiple_of(64) + || lengths.iter().any(|&n| n <= prefix) + { + return None; + } + let rows = lengths + .iter() + .try_fold(prefix, |rows, &n| rows.checked_add(n.checked_sub(prefix)?))?; + (rows <= super::multimodal_batch::MAX_MULTIMODAL_BATCH_TOKENS + && rows - prefix <= MAX_SUFFIX_BATCH_ROWS) + .then_some(rows) +} + +pub(super) fn common_header( + inputs: &[MultimodalInput<'_>], + devices: Option<&[cuda::DeviceBf16Matrix<'_>]>, + prefix: usize, + hidden: usize, +) -> Result<()> { + ensure!( + inputs.len() >= 2 && prefix >= 64 && prefix.is_multiple_of(64), + "shared multimodal prefix requires multiple prompts and a positive chunk multiple" + ); + ensure!( + inputs.iter().all(|p| p.token_ids.len() > prefix + && p.position_ids + .iter() + .all(|axis| axis.len() == p.token_ids.len())), + "invalid shared prefix boundary" + ); + let first = &inputs[0]; + let rows = first.image_token_indices.partition_point(|&i| i < prefix); + let elements = rows + .checked_mul(hidden) + .context("shared image extent overflow")?; + for (index, input) in inputs.iter().enumerate().skip(1) { + ensure!( + input.token_ids[..prefix] == first.token_ids[..prefix] + && (0..3) + .all(|axis| input.position_ids[axis][..prefix] + == first.position_ids[axis][..prefix]), + "multimodal prefix tokens or positions differ" + ); + let count = input.image_token_indices.partition_point(|&i| i < prefix); + ensure!( + input.image_token_indices[..count] == first.image_token_indices[..rows], + "multimodal prefix image indices differ" + ); + if let Some(devices) = devices { + ensure!(devices.len() == inputs.len(), "device batch count mismatch"); + ensure!( + rows == 0 || devices[index].at(0) == devices[0].at(0), + "shared image prefix requires the same resident allocation" + ); + } else { + let a = first + .image_embeddings + .get(..elements) + .context("shared image extent")?; + let b = input + .image_embeddings + .get(..elements) + .context("shared image extent")?; + ensure!( + a.iter().zip(b).all(|(a, b)| a.to_bits() == b.to_bits()), + "multimodal prefix image values differ" + ); + } + } + Ok(()) +} + +impl Model { + pub(super) fn try_multimodal_shared( + &mut self, + inputs: &[MultimodalInput<'_>], + devices: Option<&[cuda::DeviceBf16Matrix<'_>]>, + lengths: &[usize], + prefix: usize, + ) -> Result>>> { + let Some(reuse) = cuda::api().reuse_gemm() else { + return Ok(None); + }; + let mut shapes = std::collections::BTreeSet::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 weight in mixer.into_iter().chain([&layer.gate_up, &layer.down]) { + shapes.insert((weight.shape[0], weight.shape[1])); + } + } + let compatible = |originals: &[usize], targets: &[usize]| -> Result { + for &(n, k) in &shapes { + // SAFETY: bounded positive host row arrays and validated weights; + // capability checks do not produce device outputs or prefix state. + let rc = unsafe { + (reuse.check)( + self.gemm, + originals.as_ptr(), + originals.len(), + targets.as_ptr(), + targets.len(), + n as i32, + k as i32, + ) + }; + if rc == -1 { + return Ok(false); + } + check(rc, "multimodal suffix compatibility")?; + } + Ok(true) + }; + let ops = cuda::api().shared_suffix(); + let Some(cohorts) = compatible_prefix_groups(inputs.len(), |indices| { + let originals: Vec<_> = indices.iter().map(|&i| lengths[i]).collect(); + if ops.is_some() + && let Some(rows) = compact_rows(&originals, prefix) + && compatible(&originals, &[rows])? + { + return Ok(true); + } + let mut targets: Vec<_> = originals.iter().map(|&n| n - prefix).collect(); + targets.push(prefix); + compatible(&originals, &targets) + })? + else { + return Ok(None); + }; + let mut plans = Vec::new(); + let mut required = *lengths.iter().max().unwrap(); + for cohort in cohorts { + let mut segments = Vec::new(); + let originals: Vec<_> = cohort.iter().map(|&i| lengths[i]).collect(); + let compact = if ops.is_some() + && let Some(rows) = compact_rows(&originals, prefix) + && compatible(&originals, &[rows])? + { + required = required.max(rows); + true + } else { + false + }; + if cohort.len() > 1 && !compact { + for indices in suffix_segments(&cohort, lengths, prefix) { + let rows: usize = indices.iter().map(|&i| lengths[i] - prefix).sum(); + let packed = indices.len() > 1 + && ops.is_some() + && compatible(&[lengths[indices[0]]], &[rows])?; + if packed { + required = required.max(rows); + } + segments.push((indices, packed)); + } + } + plans.push((cohort, segments, compact)); + } + // Window Graphs retain this model-owned prefix allocation. Retire those + // Graphs before replacing it, including after failures or capacity growth. + if self + .mm_batch_prefix + .as_ref() + .is_none_or(|(cap, _)| *cap < prefix) + { + self.synchronize()?; + if !self.mm_batch_windows.is_empty() { + self.graph_stats.invalidations += 1; + } + self.mm_batch_windows.clear(); + self.mm_batch_prefix = None; + self.mm_batch_prefix = Some((prefix, self.alloc_prefix(prefix)?)); + } + self.prepare_scratch(required)?; + let (capacity, mut state) = self.mm_batch_prefix.take().unwrap(); + state.len = prefix; + state.initialized = false; + // Restore ownership even on unwind: worker panic recovery may retain the + // model, so cached Graphs must never outlive recorded prefix allocations. + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + let mut hidden = vec![Vec::new(); inputs.len()]; + for (cohort, segments, compact) in plans { + if cohort.len() == 1 { + let index = cohort[0]; + hidden[index] = self + .run_multimodal_batch( + &inputs[index..index + 1], + devices.map(|d| &d[index..index + 1]), + 0, + )? + .pop() + .unwrap(); + continue; + } + let first = cohort[0]; + let reference = lengths[first]; + state.initialized = false; + if compact { + let batch = SuffixBatch { + prefix: &state, + ops: ops.unwrap(), + gemm: reuse, + original_rows: reference, + seeded_gdn: cuda::api().seeded_gdn_batch(), + include_prefix: true, + }; + let rows = self.multimodal_compact_pass(inputs, devices, &cohort, &batch)?; + for (index, row) in cohort.into_iter().zip(rows) { + hidden[index] = row; + } + continue; + } + let s = self.scratch.as_ref().unwrap(); + self.upload_positions(s, inputs[first].position_ids.map(|axis| &axis[..prefix]))?; + self.embed_tokens(s, &inputs[first].token_ids[..prefix])?; + self.stage_multimodal_rows( + s, + &inputs[first], + devices.map(|d| d[first]), + 0..prefix, + 0, + )?; + self.multimodal_prefix_pass(&mut state, reference, reuse)?; + state.initialized = true; + for (indices, packed) in segments { + if packed { + let lens: Vec<_> = indices.iter().map(|&i| lengths[i] - prefix).collect(); + let ids: Vec<_> = indices + .iter() + .flat_map(|&i| inputs[i].token_ids[prefix..].iter().copied()) + .collect(); + let positions: [Vec; 3] = std::array::from_fn(|axis| { + indices + .iter() + .flat_map(|&i| { + inputs[i].position_ids[axis][prefix..].iter().copied() + }) + .collect() + }); + let s = self.scratch.as_ref().unwrap(); + self.upload_positions(s, [&positions[0], &positions[1], &positions[2]])?; + self.embed_tokens(s, &ids)?; + let mut offset = 0; + for (&index, &length) in indices.iter().zip(&lens) { + self.stage_multimodal_rows( + s, + &inputs[index], + devices.map(|d| d[index]), + prefix..lengths[index], + offset, + )?; + offset += length; + } + self.multimodal_suffix_pass( + &state, + reference, + &lens, + Some(ops.unwrap()), + reuse, + )?; + let s = self.scratch.as_ref().unwrap(); + let mut end = 0; + for (&index, &length) in indices.iter().zip(&lens) { + end += length; + hidden[index] = self.last_hidden(s, end)?; + } + } else { + for index in indices { + let s = self.scratch.as_ref().unwrap(); + self.upload_positions( + s, + inputs[index].position_ids.map(|axis| &axis[prefix..]), + )?; + self.embed_tokens_at(s, &inputs[index].token_ids[prefix..], prefix)?; + self.stage_multimodal_rows( + s, + &inputs[index], + devices.map(|d| d[index]), + prefix..lengths[index], + prefix, + )?; + self.multimodal_suffix_pass( + &state, + reference, + &[lengths[index] - prefix], + None, + reuse, + )?; + hidden[index] = + self.last_hidden(self.scratch.as_ref().unwrap(), lengths[index])?; + } + } + } + } + Ok(hidden) + })); + let synced = self.synchronize(); + state.initialized = false; + self.mm_batch_prefix = Some((capacity, state)); + if synced.is_err() || !matches!(&result, Ok(Ok(_))) { + self.mm_batch_windows.clear(); + } + match result { + Ok(value) => { + synced?; + value.map(Some) + } + Err(panic) => std::panic::resume_unwind(panic), + } + } + + /// Project each layer's prefix and suffix rows together, then seed that + /// layer's independent mixers from freshly produced prefix state. + fn multimodal_compact_pass( + &mut self, + inputs: &[MultimodalInput<'_>], + devices: Option<&[cuda::DeviceBf16Matrix<'_>]>, + indices: &[usize], + batch: &SuffixBatch<'_>, + ) -> Result>> { + let prefix = batch.prefix.len; + let originals: Vec<_> = indices.iter().map(|&i| inputs[i].token_ids.len()).collect(); + compact_rows(&originals, prefix).context("invalid compact cohort")?; + ensure!(batch.include_prefix, "compact pass must produce its prefix"); + let lengths: Vec<_> = originals.iter().map(|&n| n - prefix).collect(); + let first = indices[0]; + let ids: Vec<_> = inputs[first].token_ids[..prefix] + .iter() + .copied() + .chain( + indices + .iter() + .flat_map(|&i| inputs[i].token_ids[prefix..].iter().copied()), + ) + .collect(); + let positions: [Vec; 3] = std::array::from_fn(|axis| { + inputs[first].position_ids[axis][..prefix] + .iter() + .copied() + .chain( + indices + .iter() + .flat_map(|&i| inputs[i].position_ids[axis][prefix..].iter().copied()), + ) + .collect() + }); + { + let s = self.scratch.as_ref().unwrap(); + self.upload_positions(s, [&positions[0], &positions[1], &positions[2]])?; + self.embed_tokens(s, &ids)?; + self.stage_multimodal_rows(s, &inputs[first], devices.map(|d| d[first]), 0..prefix, 0)?; + let mut offset = prefix; + for (&index, &length) in indices.iter().zip(&lengths) { + self.stage_multimodal_rows( + s, + &inputs[index], + devices.map(|d| d[index]), + prefix..inputs[index].token_ids.len(), + offset, + )?; + offset += length; + } + } + let key = PassKey::Compact { + prefix, + lengths: lengths.clone(), + reference: batch.original_rows, + }; + if !self.replay_multimodal_pass(&key)? { + // Only suffixes enter the 16-sequence GDN descriptor. The leading + // prefix is processed first on the same stream, before they read it. + let segments: Vec<_> = std::iter::once(prefix) + .chain(lengths.iter().copied()) + .collect(); + let s = self.scratch.as_ref().unwrap(); + self.run_with_suffix(s, &segments, true, false, Some(batch))?; + if self.graph_enabled { + self.synchronize()?; + let captured = cuda::Graph::capture(self.stream, || { + self.run_with_suffix(s, &segments, true, false, Some(batch)) + }); + self.remember_multimodal_pass(key, captured)?; + } + } + let s = self.scratch.as_ref().unwrap(); + let mut end = prefix; + lengths + .iter() + .map(|&n| { + end += n; + self.last_hidden(s, end) + }) + .collect() + } + + fn replay_multimodal_pass(&mut self, key: &PassKey) -> Result { + if !self.graph_enabled { + return Ok(false); + } + let Some(entry) = self.mm_batch_windows.iter().find(|entry| &entry.key == key) else { + return Ok(false); + }; + 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=multimodal-batch event=replay key={key:?}"); + } + Ok(true) + } + + fn remember_multimodal_pass( + &mut self, + key: PassKey, + captured: Result, + ) -> Result<()> { + match captured { + Ok(graph) => { + if self.mm_batch_windows.len() == 64 { + self.synchronize()?; + self.mm_batch_windows.pop_front(); + } + self.graph_stats.captures += 1; + if std::env::var("CUA_S1_GRAPH_TRACE").as_deref() == Ok("1") { + eprintln!("CUA_S1_GRAPH mode=multimodal-batch event=capture key={key:?}"); + } + self.mm_batch_windows.push_back(Entry { key, graph }); + } + Err(error) => { + self.cache_graph(&[], true, Err(error))?; + } + } + Ok(()) + } + + fn multimodal_prefix_pass( + &mut self, + state: &mut PrefixState, + reference: usize, + reuse: cuda::ReuseGemm, + ) -> Result<()> { + let key = PassKey::Prefix { + prefix: state.len, + reference, + }; + if self.replay_multimodal_pass(&key)? { + return Ok(()); + } + let s = self.scratch.as_ref().unwrap(); + self.run_window( + s, + 0, + state.len, + true, + Some(&mut *state), + None, + self.gemm, + Some((reuse, reference)), + )?; + if self.graph_enabled { + self.synchronize()?; + let captured = cuda::Graph::capture(self.stream, || { + self.run_window( + s, + 0, + state.len, + true, + Some(&mut *state), + None, + self.gemm, + Some((reuse, reference)), + ) + }); + self.remember_multimodal_pass(key, captured)?; + } + Ok(()) + } + + fn multimodal_suffix_pass( + &mut self, + state: &PrefixState, + reference: usize, + lengths: &[usize], + ops: Option, + reuse: cuda::ReuseGemm, + ) -> Result<()> { + let key = if ops.is_some() { + PassKey::Batch { + prefix: state.len, + lengths: lengths.to_vec(), + reference, + } + } else { + PassKey::Window { + prefix: state.len, + suffix: lengths[0], + reference, + } + }; + if self.replay_multimodal_pass(&key)? { + return Ok(()); + } + let s = self.scratch.as_ref().unwrap(); + let run = || { + if let Some(ops) = ops { + self.run_with_suffix( + s, + lengths, + true, + false, + Some(&SuffixBatch { + prefix: state, + ops, + gemm: reuse, + original_rows: reference, + seeded_gdn: cuda::api().seeded_gdn_batch(), + include_prefix: false, + }), + ) + } else { + self.run_window( + s, + state.len, + state.len + lengths[0], + true, + None, + Some(state), + self.gemm, + Some((reuse, reference)), + ) + } + }; + run()?; + if self.graph_enabled { + self.synchronize()?; + let captured = cuda::Graph::capture(self.stream, run); + self.remember_multimodal_pass(key, captured)?; + } + Ok(()) + } +} + +#[cfg(test)] +#[path = "../../../../../../tests/qwen3_5/multimodal_shared.rs"] +mod tests; diff --git a/tests/cua_s1/device_controls.rs b/tests/cua_s1/device_controls.rs index 7d5b839f..a1deca45 100644 --- a/tests/cua_s1/device_controls.rs +++ b/tests/cua_s1/device_controls.rs @@ -153,3 +153,64 @@ fn non_unicode_prefix_is_rejected_before_loading() { .contains("CUA_S1_MM_PREFIX must be off, fixed, shared, or matched") ); } + +#[test] +fn question_batch_controls_fail_before_checkpoint_or_cuda_loading() { + for (value, prefix, expected) in [ + ("", "off", "CUA_S1_MM_BATCH must be unset, 0, or 1"), + ("true", "off", "CUA_S1_MM_BATCH must be unset, 0, or 1"), + ("2", "off", "CUA_S1_MM_BATCH must be unset, 0, or 1"), + ( + "1", + "matched", + "CUA_S1_MM_BATCH requires CUA_S1_MM_PREFIX=off", + ), + ] { + let missing = + std::env::temp_dir().join(format!("cua-batch-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_BATCH", value) + .env("CUA_S1_MM_PREFIX", prefix) + .env("CUA_S1_GRAPH", "0") + .env("CUA_S1_DEVICE_FEATURES", "0") + .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}"); + assert!(!error.contains("export the multimodal language checkpoint")); + } +} + +#[test] +fn resize_worker_controls_reject_invalid_values_before_loading() { + for value in ["", "5", "true", "+1", "01", " 1"] { + let missing = + std::env::temp_dir().join(format!("cua-resize-missing-{}", std::process::id())); + assert!(!missing.exists()); + let out = std::process::Command::new(env!("CARGO_BIN_EXE_omni-cua-s1-vision")) + .env("CUA_S1_RESIZE_WORKERS", value) + .env("CUA_S1_MM_BATCH", "0") + .env("CUA_S1_MM_PREFIX", "off") + .env("CUA_S1_DEVICE_FEATURES", "0") + .env("CUA_S1_GRAPH", "0") + .env("CUA_S1_BASE", &missing) + .env("CUA_S1_MODEL", &missing) + .env("CUA_S1_VISION_ADAPTER", &missing) + .env("CUA_S1_CUDA_LIB", &missing) + .output() + .unwrap(); + assert!(!out.status.success()); + let error = String::from_utf8(out.stderr).unwrap(); + assert!( + error.contains("CUA_S1_RESIZE_WORKERS must be unset or 0..=4"), + "{error}" + ); + assert!(!error.contains("export the multimodal language checkpoint")); + } +} diff --git a/tests/cua_s1/question_batch.rs b/tests/cua_s1/question_batch.rs new file mode 100644 index 00000000..ff12546c --- /dev/null +++ b/tests/cua_s1/question_batch.rs @@ -0,0 +1,130 @@ +use super::*; + +#[test] +fn question_batch_opt_in_is_strict_and_does_not_mix_prefix_modes() { + assert!(!enabled(None, prefix::Mode::Off).unwrap()); + assert!(!enabled(Some("0"), prefix::Mode::Matched).unwrap()); + assert!(enabled(Some("1"), prefix::Mode::Off).unwrap()); + assert!(enabled(Some("2"), prefix::Mode::Off).is_err()); + assert!(enabled(Some("1"), prefix::Mode::Matched).is_err()); +} + +#[test] +fn question_batch_groups_preserve_order_and_bound_total_tokens() { + assert_eq!(groups(&[4096, 4096, 4096, 1], false).unwrap(), [0..2, 2..4]); + let all = groups(&[100; 8], false).unwrap(); + assert_eq!(all.len(), 1); + assert_eq!(all[0], 0..8); + for bad in [vec![], vec![0], vec![4097], vec![1; 9]] { + assert!(groups(&bad, false).is_err()); + } +} + +#[test] +#[ignore = "requires CUDA and pinned CUA_S1_* weights; run serially with MM_PREFIX=off/MM_BATCH=0"] +fn question_batch_keeps_logits_response_order_and_changed_pixels() { + use crate::vision_engine::device_feature_tests::{load, questions}; + let mut engine = load(); + engine.language.enable_multimodal_batching().unwrap(); + engine.batch_questions = true; + for (width, height, count, seed) in [ + (256, 256, 1, 1), + (256, 256, 4, 1), + (256, 256, 4, 2), + (512, 256, 8, 3), + (1024, 1024, 8, 4), + (256, 256, 4, 1), + ] { + let rgb: Vec<_> = (0..width * height * 3) + .map(|j| ((j * 17 + seed * 31 + j / 251) % 256) as u8) + .collect(); + let qs: Vec<_> = (0..count) + .map(|i| { + let mut q = questions(seed).remove(i % 2); + q.name = format!("q{i}"); + let options = [1, 2, 26][i % 3]; + q.keys = (0..options).map(|j| format!("k{j}")).collect(); + q.labels = (0..options) + .map(|j| format!("Decision number {j}")) + .collect(); + q + }) + .collect(); + let mut prepared = engine.prepare(width, height, &rgb, &qs).unwrap(); + let host = engine.vision.forward(&prepared.image).unwrap(); + let expected: Vec<_> = prepared + .prompts + .iter() + .zip(prepared.context.option_counts()) + .map(|(p, n)| engine.score(p, &host, n).unwrap().logits) + .collect(); + let expected_response = prepared.context.finish(expected.clone()).unwrap(); + for device in [false, true] { + engine.device_features = device; + for _ in 0..2 { + let actual = engine.execute_prepared(&prepared).unwrap(); + for (a, b) in actual.iter().zip(&expected) { + assert_eq!( + a.iter().map(|x| x.to_bits()).collect::>(), + b.iter().map(|x| x.to_bits()).collect::>(), + "{width}x{height} count={count} seed={seed} device={device}" + ); + } + assert_eq!(prepared.context.finish(actual).unwrap(), expected_response); + } + } + // A late malformed question must not poison captures or request recovery. + let last = prepared.prompts.last_mut().unwrap(); + let old = last.position_ids[2].pop().unwrap(); + assert!(engine.execute_prepared(&prepared).is_err()); + prepared.prompts.last_mut().unwrap().position_ids[2].push(old); + assert_eq!( + engine.predict_prepared(&prepared).unwrap(), + expected_response + ); + } +} + +#[test] +fn shared_prefix_questions_keep_one_group_when_compact_work_is_bounded() { + assert_eq!(groups(&[1030; 8], false).unwrap(), vec![0..7, 7..8]); + let shared = groups(&[1030; 8], true).unwrap(); + assert_eq!(shared.len(), 1); + assert_eq!(shared[0], 0..8); +} + +#[test] +fn short_position_axes_return_an_error_before_prefix_discovery() { + use crate::contract::Question; + use tokenizers::{AddedToken, Tokenizer, models::wordlevel::WordLevel}; + let vocab = [("[UNK]".to_owned(), 0), ("<|image_pad|>".to_owned(), 1)] + .into_iter() + .collect(); + let mut tokenizer = Tokenizer::new( + WordLevel::builder() + .vocab(vocab) + .unk_token("[UNK]".into()) + .build() + .unwrap(), + ); + tokenizer.add_special_tokens(&[AddedToken::from("<|image_pad|>", true)]); + let processor = VisionProcessor::new(tokenizer, 1); + let questions: Vec<_> = (0..2) + .map(|i| Question { + name: format!("q{i}"), + goal: String::new(), + keys: vec!["only".into()], + labels: vec!["Only".into()], + }) + .collect(); + let mut prepared = processor + .prepare(256, 256, &vec![0; 256 * 256 * 3], &questions) + .unwrap(); + assert!(plan(&prepared).is_ok()); + for axis in 0..3 { + let saved = std::mem::take(&mut prepared.prompts[1].position_ids[axis]); + assert!(plan(&prepared).is_err()); + prepared.prompts[1].position_ids[axis] = saved; + } + assert!(plan(&prepared).is_ok()); +} diff --git a/tests/cua_s1/vision_processing.rs b/tests/cua_s1/vision_processing.rs new file mode 100644 index 00000000..3a0f377a --- /dev/null +++ b/tests/cua_s1/vision_processing.rs @@ -0,0 +1,12 @@ +use super::*; + +#[test] +fn resize_workers_are_explicit_bounded_and_default_off() { + assert_eq!(resize_workers(None).unwrap(), 0); + for value in 0..=4 { + assert_eq!(resize_workers(Some(&value.to_string())).unwrap(), value); + } + for value in ["", "5", "true", "+1", "01", " 1"] { + assert!(resize_workers(Some(value)).is_err()); + } +} diff --git a/tests/qwen3_5/batched_gdn.rs b/tests/qwen3_5/batched_gdn.rs index 59080d5a..beee7c26 100644 --- a/tests/qwen3_5/batched_gdn.rs +++ b/tests/qwen3_5/batched_gdn.rs @@ -6,7 +6,7 @@ use super::*; fn batched_gdn_matches_serial_with_partial_chunks_and_graph_replay() { let stream = setup(); let batched = api().batched_gdn().unwrap(); - for (heads, key_heads) in [(4usize, 2usize), (48, 16)] { + for (heads, key_heads) in [(4usize, 2usize), (32, 16), (48, 16)] { for lengths in [ vec![1usize, 63, 64, 65, 129], vec![1; 16], diff --git a/tests/qwen3_5/image_preprocess.rs b/tests/qwen3_5/image_preprocess.rs index 5f08df0a..5a5a1cda 100644 --- a/tests/qwen3_5/image_preprocess.rs +++ b/tests/qwen3_5/image_preprocess.rs @@ -43,3 +43,94 @@ fn spatial_merge_channel_and_duplicate_temporal_planes_match_the_tensor_layout() } } } + +#[test] +fn explicit_resize_workers_match_default_pixels_and_metadata() { + for (width, height) in [(17, 19), (33, 49), (97, 65), (512, 500), (1001, 733)] { + let rgb: Vec<_> = (0..width * height * 3) + .map(|j| ((j * 17 + j / 251 + 31) % 256) as u8) + .collect(); + let expected = preprocess_rgb8(width, height, &rgb).unwrap(); + for workers in [1, 2, 4] { + let actual = + preprocess_rgb8_with_workers(width, height, &rgb, &ImageLimits::cua(), workers) + .unwrap(); + assert_eq!( + ( + actual.resized_width, + actual.resized_height, + actual.image_grid_thw + ), + ( + expected.resized_width, + expected.resized_height, + expected.image_grid_thw + ) + ); + assert_eq!(actual.pixel_values.len(), expected.pixel_values.len()); + assert!( + actual + .pixel_values + .iter() + .zip(&expected.pixel_values) + .all(|(a, b)| a.to_bits() == b.to_bits()), + "{width}x{height} workers={workers}" + ); + } + } + for workers in [0, 5, usize::MAX] { + assert!( + preprocess_rgb8_with_workers(1, 1, &[0, 0, 0], &ImageLimits::cua(), workers).is_err() + ); + } +} + +#[cfg(target_arch = "x86_64")] +#[test] +fn simd_resize_rows_match_scalar_at_short_and_unaligned_edges() { + if std::arch::is_x86_feature_detected!("sse4.1") { + for source_width in [1, 2, 7, 17] { + for width in [1, 2, 3, 7, 8, 9, 17, 33] { + let source: Vec<_> = (0..source_width * 3) + .map(|j| ((j * 193 + 11) % 256) as u8) + .collect(); + let axis = AxisWeights::new(source_width, width); + let mut expected = Vec::new(); + for kernel in &axis.kernels { + for channel in 0..3 { + expected.push(axis.apply(kernel, |i| source[i * 3 + channel])); + } + } + let mut actual = vec![0; width * 3]; + // SAFETY: runtime feature check and axis/source/output extents match. + unsafe { + x86::horizontal(&source, &mut actual, &axis); + } + assert_eq!(actual, expected, "horizontal {source_width}->{width}"); + } + } + } + if std::arch::is_x86_feature_detected!("avx2") { + for width in [1, 2, 3, 7, 8, 9, 17, 33] { + for height in [1, 2, 7] { + for target in [1, 3, 9] { + let source: Vec<_> = (0..width * height * 3) + .map(|j| ((j * 193 + 11) % 256) as u8) + .collect(); + let axis = AxisWeights::new(height, target); + for kernel in &axis.kernels { + let expected: Vec<_> = (0..width * 3) + .map(|x| axis.apply(kernel, |y| source[y * width * 3 + x])) + .collect(); + let mut actual = vec![0; width * 3]; + // SAFETY: runtime feature check and every source row has width*3 bytes. + unsafe { + x86::vertical(&source, &mut actual, kernel, axis.precision); + } + assert_eq!(actual, expected, "vertical {width}x{height}->{target}"); + } + } + } + } + } +} diff --git a/tests/qwen3_5/kernels.rs b/tests/qwen3_5/kernels.rs index 5173a122..48078362 100644 --- a/tests/qwen3_5/kernels.rs +++ b/tests/qwen3_5/kernels.rs @@ -1592,3 +1592,5 @@ mod shared_suffix; mod vector_conv; mod wide_gdn; + +mod seeded_gdn; diff --git a/tests/qwen3_5/multimodal_batch.rs b/tests/qwen3_5/multimodal_batch.rs new file mode 100644 index 00000000..d008a5e0 --- /dev/null +++ b/tests/qwen3_5/multimodal_batch.rs @@ -0,0 +1,196 @@ +use super::*; +use half::bf16; + +fn input<'a>( + ids: &'a [u32], + indices: &'a [usize], + features: &'a [bf16], + positions: [&'a [i64]; 3], +) -> MultimodalInput<'a> { + MultimodalInput { + token_ids: ids, + image_token_indices: indices, + image_embeddings: features, + position_ids: positions, + } +} + +#[test] +fn multimodal_batch_layout_bounds_prompt_count_and_total_tokens() { + let ids = [1, 99, 2]; + let indices = [1]; + let features = [bf16::ZERO; 4]; + let positions = [0, 1, 2]; + let one = || input(&ids, &indices, &features, [&positions; 3]); + assert!(batch_lengths(&[], 0).is_err()); + assert_eq!(batch_lengths(&[one(), one()], 0).unwrap(), [3, 3]); + assert!(batch_lengths(&(0..17).map(|_| one()).collect::>(), 0).is_err()); + let long = vec![1; MAX_MULTIMODAL_BATCH_TOKENS + 1]; + let too_big = input(&long, &[], &[], [&[]; 3]); + assert!(batch_lengths(&[too_big], 0).is_err()); +} + +#[test] +#[ignore = "requires CUDA and pinned CUA_S1_MODEL/CUA_S1_CUDA_LIB; run serially"] +fn multimodal_batch_retains_serial_hidden_bits_and_current_inputs() { + use std::path::PathBuf; + let dir = PathBuf::from(std::env::var_os("CUA_S1_MODEL").expect("CUA_S1_MODEL")); + let library = PathBuf::from(std::env::var_os("CUA_S1_CUDA_LIB").expect("CUA_S1_CUDA_LIB")); + for graph in [false, true] { + let mut model = Model::load_with_graph(&dir, &library, graph).unwrap(); + model.enable_multimodal_batching().unwrap(); + let image_token = model.cfg.image_token_id.unwrap(); + for (lengths, seed) in [ + (vec![65, 129, 257], 1), + (vec![65, 129, 257], 2), + (vec![512; 8], 3), + (vec![65, 129, 257], 1), + ] { + let ids: Vec> = lengths + .iter() + .enumerate() + .map(|(i, &n)| { + let mut ids: Vec<_> = (0..n) + .map(|j| (10 + (j * 17 + i * 31 + seed) % 1000) as u32) + .collect(); + ids[8..12].fill(image_token); + ids + }) + .collect(); + let features: Vec> = lengths + .iter() + .enumerate() + .map(|(i, _)| { + (0..4 * model.cfg.hidden) + .map(|j| { + bf16::from_f32(((j * 13 + seed * 17 + i * 7) % 97) as f32 / 97. - 0.5) + }) + .collect() + }) + .collect(); + let positions: Vec<[Vec; 3]> = lengths + .iter() + .map(|&n| { + std::array::from_fn(|axis| (0..n).map(|p| (p + axis * 3) as i64).collect()) + }) + .collect(); + let indices = [8, 9, 10, 11]; + let inputs: Vec<_> = ids + .iter() + .zip(&features) + .zip(&positions) + .map(|((ids, features), p)| input(ids, &indices, features, [&p[0], &p[1], &p[2]])) + .collect(); + let expected: Vec<_> = inputs + .iter() + .map(|p| model.forward_multimodal(p).unwrap()) + .collect(); + for _ in 0..2 { + let actual = model.forward_multimodal_batch(&inputs).unwrap(); + assert_eq!(actual.len(), expected.len()); + for (a, b) in actual.iter().zip(&expected) { + let differences = a + .iter() + .zip(b) + .filter(|(a, b)| a.to_bits() != b.to_bits()) + .count(); + assert_eq!( + differences, 0, + "lengths={lengths:?} seed={seed} graph={graph}" + ); + } + } + } + if graph { + let stats = model.graph_stats(); + assert!(stats.enabled && stats.captures > 0 && stats.replays > 0); + assert_eq!(stats.fallbacks, 0); + } + } +} + +#[test] +fn shared_prefix_bounds_each_prompt_without_counting_duplicate_raw_rows_as_scratch() { + let ids = vec![1; 1030]; + let inputs: Vec<_> = (0..8).map(|_| input(&ids, &[], &[], [&[]; 3])).collect(); + assert!(batch_lengths(&inputs, 0).is_err()); + assert_eq!(batch_lengths(&inputs, 576).unwrap(), vec![1030; 8]); + let oversized = vec![1; MAX_MULTIMODAL_BATCH_TOKENS + 1]; + assert!(batch_lengths(&[input(&oversized, &[], &[], [&[]; 3])], 576).is_err()); +} + +#[test] +#[ignore = "requires CUDA and pinned CUA_S1_MODEL/CUA_S1_CUDA_LIB; run serially"] +fn repeated_device_features_still_validate_a_late_distinct_allocation_and_recover() { + use std::path::PathBuf; + let dir = PathBuf::from(std::env::var_os("CUA_S1_MODEL").unwrap()); + let lib = PathBuf::from(std::env::var_os("CUA_S1_CUDA_LIB").unwrap()); + let mut model = Model::load_with_graph(&dir, &lib, false).unwrap(); + model.enable_multimodal_batching().unwrap(); + let hidden = model.cfg.hidden; + let mut ids = vec![1; 65]; + ids[8] = model.cfg.image_token_id.unwrap(); + let indices = [8]; + let positions: Vec<_> = (0..65).map(i64::from).collect(); + let features = vec![bf16::from_f32(0.25); hidden]; + let bytes: Vec<_> = features.iter().flat_map(|x| x.to_le_bytes()).collect(); + let first = cuda::DeviceBuffer::new(bytes.len()).unwrap(); + let second = cuda::DeviceBuffer::new(bytes.len()).unwrap(); + let mut bad_bytes = bytes.clone(); + bad_bytes[..2].copy_from_slice(&bf16::NAN.to_le_bytes()); + // SAFETY: exact allocation sizes, serialized uploads, no consumer is active. + unsafe { + cuda::upload(first.at(0), &bytes, model.stream).unwrap(); + cuda::upload(second.at(0), &bad_bytes, model.stream).unwrap(); + } + model.synchronize().unwrap(); + let make = |image_embeddings| DeviceMultimodalInput { + token_ids: &ids, + image_token_indices: &indices, + image_embeddings, + position_ids: [&positions; 3], + }; + { + let a = first.bf16_matrix(1, hidden).unwrap(); + let b = second.bf16_matrix(1, hidden).unwrap(); + assert!(a.same_storage(&a)); + assert!(!a.same_storage(&b)); + assert!(!a.same_storage(&first.bf16_matrix(2, hidden / 2).unwrap())); + let error = model + .forward_multimodal_device_batch(&[make(a), make(a), make(b)]) + .unwrap_err(); + assert!( + error.to_string().contains("non-finite image embedding"), + "{error:#}" + ); + } + // SAFETY: the failed consumer drained; this producer allocation is reusable. + unsafe { cuda::upload(second.at(0), &bytes, model.stream).unwrap() }; + model.synchronize().unwrap(); + let a = first.bf16_matrix(1, hidden).unwrap(); + let b = second.bf16_matrix(1, hidden).unwrap(); + let expected = model + .forward_multimodal(&input(&ids, &indices, &features, [&positions; 3])) + .unwrap(); + let actual = model + .forward_multimodal_device_batch(&[make(a), make(a), make(b)]) + .unwrap(); + assert_eq!(actual.len(), 3); + for row in actual { + assert_eq!( + row.iter().map(|x| x.to_bits()).collect::>(), + expected.iter().map(|x| x.to_bits()).collect::>() + ); + } + // SAFETY: the previous consumer drained. Revalidate this same allocation + // after its producer changes it; no finite-check result survives a call. + unsafe { cuda::upload(second.at(0), &bad_bytes, model.stream).unwrap() }; + model.synchronize().unwrap(); + let error = model + .forward_multimodal_device_batch(&[ + make(first.bf16_matrix(1, hidden).unwrap()), + make(second.bf16_matrix(1, hidden).unwrap()), + ]) + .unwrap_err(); + assert!(error.to_string().contains("non-finite image embedding")); +} diff --git a/tests/qwen3_5/multimodal_shared.rs b/tests/qwen3_5/multimodal_shared.rs new file mode 100644 index 00000000..fd030dc3 --- /dev/null +++ b/tests/qwen3_5/multimodal_shared.rs @@ -0,0 +1,221 @@ +use super::*; +use half::bf16; + +#[test] +fn shared_multimodal_header_checks_tokens_positions_and_image_values() { + let ids = vec![1; 129]; + let indices = [8]; + let features = [bf16::ZERO; 4]; + let p: Vec<_> = (0..129).map(i64::from).collect(); + let make = |ids, features, positions| MultimodalInput { + token_ids: ids, + image_token_indices: &indices, + image_embeddings: features, + position_ids: positions, + }; + assert!( + common_header( + &[ + make(&ids, &features, [&p; 3]), + make(&ids, &features, [&p; 3]) + ], + None, + 128, + 4 + ) + .is_ok() + ); + let mut changed = ids.clone(); + changed[12] = 2; + assert!( + common_header( + &[ + make(&ids, &features, [&p; 3]), + make(&changed, &features, [&p; 3]) + ], + None, + 128, + 4 + ) + .is_err() + ); + let different = [bf16::ONE; 4]; + assert!( + common_header( + &[ + make(&ids, &features, [&p; 3]), + make(&ids, &different, [&p; 3]) + ], + None, + 128, + 4 + ) + .is_err() + ); + let signed_zero = [bf16::from_bits(0x8000); 4]; + assert!( + common_header( + &[ + make(&ids, &features, [&p; 3]), + make(&ids, &signed_zero, [&p; 3]) + ], + None, + 128, + 4 + ) + .is_err() + ); + let mut axis = p.clone(); + axis[9] += 1; + assert!( + common_header( + &[ + make(&ids, &features, [&p; 3]), + make(&ids, &features, [&p, &axis, &p]) + ], + None, + 128, + 4 + ) + .is_err() + ); + assert!( + common_header( + &[ + make(&ids, &features, [&p; 3]), + make(&ids, &features, [&p; 3]) + ], + None, + 127, + 4 + ) + .is_err() + ); +} + +#[test] +#[ignore = "requires CUDA and pinned CUA_S1_MODEL/CUA_S1_CUDA_LIB; run serially"] +fn multimodal_shared_suffix_batch_keeps_hidden_bits_after_pixel_changes() { + use std::path::PathBuf; + let dir = PathBuf::from(std::env::var_os("CUA_S1_MODEL").expect("CUA_S1_MODEL")); + let lib = PathBuf::from(std::env::var_os("CUA_S1_CUDA_LIB").expect("CUA_S1_CUDA_LIB")); + for graph in [false, true] { + let mut model = Model::load_with_graph(&dir, &lib, graph).unwrap(); + model.enable_multimodal_batching().unwrap(); + for (prefix, seed) in [ + (64, 1), + (128, 2), + (576, 3), + (576, 4), + (8064, 6), + (512, 5), + (64, 1), + ] { + let lengths = if prefix == 8064 { + // Per-prompt limit, oversized compact layout, then smaller replay. + vec![8192; 16] + } else if prefix == 576 { + vec![1030; 8] + } else if seed == 5 { + vec![513; 16] + } else { + vec![192, 193, 192, 192] + }; + let ids: Vec> = lengths + .into_iter() + .enumerate() + .map(|(i, n)| { + let mut ids: Vec<_> = (0..n) + .map(|j| { + (10 + (j * 17 + if j >= prefix { i * 31 + seed } else { 0 }) % 1000) + as u32 + }) + .collect(); + ids[8..12].fill(model.cfg.image_token_id.unwrap()); + ids + }) + .collect(); + let features: Vec<_> = (0..4 * model.cfg.hidden) + .map(|j| bf16::from_f32(((j * 7 + seed * 13) % 71) as f32 / 71. - 0.5)) + .collect(); + let positions: Vec<[Vec; 3]> = ids + .iter() + .map(|ids| { + std::array::from_fn(|axis| { + (0..ids.len()).map(|p| (p + axis * 3) as i64).collect() + }) + }) + .collect(); + let indices = [8, 9, 10, 11]; + let inputs: Vec<_> = ids + .iter() + .zip(&positions) + .map(|(ids, p)| MultimodalInput { + token_ids: ids, + image_token_indices: &indices, + image_embeddings: &features, + position_ids: [&p[0], &p[1], &p[2]], + }) + .collect(); + let expected: Vec<_> = inputs + .iter() + .map(|p| model.forward_multimodal(p).unwrap()) + .collect(); + for _ in 0..2 { + let actual = model + .forward_multimodal_batch_with_prefix(&inputs, prefix) + .unwrap(); + for (i, (a, b)) in actual.iter().zip(&expected).enumerate() { + assert_eq!( + a.iter() + .zip(b) + .filter(|(a, b)| a.to_bits() != b.to_bits()) + .count(), + 0, + "question={i} prefix={prefix} seed={seed} graph={graph}" + ); + } + } + } + if graph { + let stats = model.graph_stats(); + assert!(stats.enabled); + assert_eq!(stats.fallbacks, 0); + assert!( + model.mm_batch_windows.iter().any(|entry| matches!( + &entry.key, PassKey::Compact { lengths, .. } if lengths.len() == 16 + )), + "qualify the compact path with seventeen internal segments" + ); + assert!( + model + .mm_batch_windows + .iter() + .any(|e| matches!(e.key, PassKey::Batch { .. } | PassKey::Compact { .. })) + ); + } + } +} + +#[test] +fn compact_projection_rows_keep_prefix_once_and_bound_suffixes_and_scratch() { + assert_eq!(compact_rows(&[1030; 8], 576), Some(4208)); + assert_eq!(compact_rows(&[320; 16], 64), Some(4160)); + for (lengths, prefix) in [ + (vec![129], 64), + (vec![129; 17], 64), + (vec![0, 129], 64), + (vec![64; 2], 64), + (vec![129; 2], 63), + (vec![129; 2], 0), + (vec![4096; 8], 64), + (vec![8192; 2], 8128), + (vec![usize::MAX; 2], 64), + ] { + assert_eq!( + compact_rows(&lengths, prefix), + None, + "lengths={lengths:?} prefix={prefix}" + ); + } +} diff --git a/tests/qwen3_5/packed_gemm.rs b/tests/qwen3_5/packed_gemm.rs index b75fe0a2..dabc2de3 100644 --- a/tests/qwen3_5/packed_gemm.rs +++ b/tests/qwen3_5/packed_gemm.rs @@ -9,13 +9,22 @@ fn packed_projections_match_individual_shapes_and_changed_input_replay() { // SAFETY: model-equivalent workspace, released after the final readback. let gemm = unsafe { (api().cs1_gemm_create)(32 << 20) }; assert!(!gemm.is_null()); - for (n, k) in [(5120usize, 6144usize), (5120, 17408)] { + for (n, k) in [ + (5120usize, 6144usize), + (5120, 17408), + (12352, 2560), + (10240, 2560), + (18432, 2560), + (2560, 4096), + (2560, 9216), + ] { let weights = to_device(&random(n * k, 91, 0.25), stream); for lengths in [ vec![78, 78, 82, 80], vec![64; 4], vec![1; 3], vec![129, 63, 130], + vec![65, 129, 257], vec![1024; 2], ] { let rows: usize = lengths.iter().sum(); diff --git a/tests/qwen3_5/seeded_gdn.rs b/tests/qwen3_5/seeded_gdn.rs new file mode 100644 index 00000000..8f9fcea4 --- /dev/null +++ b/tests/qwen3_5/seeded_gdn.rs @@ -0,0 +1,184 @@ +//! Seed each independent suffix from the same readonly FP32 prefix state. +use super::*; + +#[test] +#[ignore = "needs a GPU and the optional seeded GDN batch symbol"] +fn seeded_gdn_batch_matches_independent_outputs_and_fp32_states_after_replay() { + let st = setup(); + let seeded = api().seeded_gdn_batch().expect("seeded GDN batch symbol"); + let batch = api().batched_gdn().unwrap(); + for (h, hk, lengths) in [ + (32usize, 16usize, vec![65usize]), + (32, 16, vec![65, 129, 257]), + ( + 32, + 16, + vec![ + 1, 63, 64, 65, 127, 128, 129, 257, 1, 63, 64, 65, 127, 128, 129, 257, + ], + ), + (48, 16, vec![63, 65]), + ] { + let total: usize = lengths.iter().sum(); + let sizes = [ + total * hk * 128, + total * hk * 128, + total * h * 128, + total * h, + ]; + let inputs = sizes.map(|n| to_device(&random(n, 1801, 0.25), st)); + let gate_values: Vec<_> = random(total * h, 1802, 0.5) + .iter() + .map(|x| (x.to_f32() - 0.5) * 0.01) + .collect(); + let gates = f32_to_device(&gate_values, st); + let state_values: Vec<_> = random(h * 128 * 128, 1803, 0.5) + .iter() + .map(|x| x.to_f32() * 0.01731 + 0.00017) + .collect(); + assert!( + state_values + .iter() + .any(|&x| bf16::from_f32(x).to_f32().to_bits() != x.to_bits()) + ); + let initial = f32_to_device(&state_values, st); + let output_elements = total * h * 128; + let outputs = [ + DeviceBuffer::new(output_elements * 2).unwrap(), + DeviceBuffer::new(output_elements * 2).unwrap(), + ]; + let state_bytes = lengths.len() * h * 128 * 128 * 4; + let states = [ + DeviceBuffer::new(state_bytes).unwrap(), + DeviceBuffer::new(state_bytes).unwrap(), + ]; + // SAFETY: positive bounded dimensions; byte counts mirror the kernel layouts. + let serial_floats = unsafe { + (api().cs1_gdn_workspace_floats)(*lengths.iter().max().unwrap() as i32, h as i32) + }; + let batch_floats = + unsafe { (batch.workspace_floats)(total as i32, h as i32, lengths.len() as i32) }; + let serial_ws = DeviceBuffer::new(serial_floats * 4).unwrap(); + let batch_ws = DeviceBuffer::new(batch_floats * 4).unwrap(); + let run_batch = || unsafe { + // SAFETY: all sequences have disjoint workspaces, outputs and final states; + // the shared FP32 initial allocation is immutable throughout this launch. + check( + seeded( + inputs[0].at(0), + inputs[1].at(0), + inputs[2].at(0), + gates.at(0).cast(), + inputs[3].at(0), + outputs[1].at(0), + batch_ws.at(0).cast(), + batch_floats * 4, + lengths.as_ptr(), + lengths.len(), + h as i32, + hk as i32, + 128f32.powf(-0.5), + initial.at(0), + states[1].at(0), + st, + ), + "seeded GDN batch", + ) + }; + run_batch().unwrap(); + cuda::synchronize(st).unwrap(); + let graph = cuda::Graph::capture(st, run_batch).unwrap(); + for visit in 0..2 { + for (index, size) in sizes.into_iter().enumerate() { + let values: Vec<_> = random(size, 1810 + index as u64 + visit * 13, 0.5) + .into_iter() + .map(|x| { + if index == 3 { + bf16::from_f32(x.to_f32() + 0.5) + } else { + x + } + }) + .collect(); + let bytes: Vec<_> = values.iter().flat_map(|x| x.to_le_bytes()).collect(); + // SAFETY: replace complete live input allocations before consumers. + unsafe { cuda::upload(inputs[index].at(0), &bytes, st).unwrap() }; + } + let values: Vec<_> = random(total * h, 1830 + visit, 0.5) + .iter() + .flat_map(|x| ((x.to_f32() - 0.5) * 0.01).to_le_bytes()) + .collect(); + let fresh_state: Vec<_> = random(h * 128 * 128, 1840 + visit, 0.5) + .iter() + .flat_map(|x| (x.to_f32() * 0.01731 + 0.00017).to_le_bytes()) + .collect(); + unsafe { + // SAFETY: exact FP32 allocations; consumers follow on this same stream. + cuda::upload(gates.at(0), &values, st).unwrap(); + cuda::upload(initial.at(0), &fresh_state, st).unwrap(); + } + let mut offset = 0; + for (index, &length) in lengths.iter().enumerate() { + unsafe { + // SAFETY: sequence-local row offsets and separate final FP32 states. + check( + (api().cs1_gdn_prefill_x)( + inputs[0].at(offset * hk * 128 * 2), + inputs[1].at(offset * hk * 128 * 2), + inputs[2].at(offset * h * 128 * 2), + gates.at(offset * h * 4).cast(), + inputs[3].at(offset * h * 2), + outputs[0].at(offset * h * 128 * 2), + serial_ws.at(0).cast(), + length as i32, + h as i32, + hk as i32, + 128f32.powf(-0.5), + initial.at(0), + states[0].at(index * h * 128 * 128 * 4), + st, + ), + "serial seeded GDN", + ) + .unwrap(); + } + offset += length; + } + if visit == 0 { + run_batch().unwrap() + } else { + graph.launch(st).unwrap() + } + let expected: Vec<_> = from_device(&outputs[0], output_elements, st) + .iter() + .map(|x| x.to_bits()) + .collect(); + let actual: Vec<_> = from_device(&outputs[1], output_elements, st) + .iter() + .map(|x| x.to_bits()) + .collect(); + assert_eq!( + actual, expected, + "BF16 h={h} lengths={lengths:?} visit={visit}" + ); + let download = |buffer: &DeviceBuffer, bytes: usize| { + let mut out = vec![0u8; bytes]; + // SAFETY: buffer owns the requested bytes; download drains the stream. + unsafe { cuda::download(&mut out, buffer.at(0), st).unwrap() }; + out + }; + assert_eq!( + download(&states[1], state_bytes), + download(&states[0], state_bytes), + "FP32 state h={h} lengths={lengths:?} visit={visit}" + ); + assert_eq!( + download(&initial, fresh_state.len()), + fresh_state, + "initial state was mutated" + ); + } + } + // SAFETY: every consumer and download completed; the stream belongs to this test. + unsafe { (api().cs1_stream_destroy)(st) }; +} From c665bea2d9757b2e34f60e4ba9089da8e03e4316 Mon Sep 17 00:00:00 2001 From: levius <2114377220@qq.com> Date: Sun, 1 Nov 2026 21:05:15 +0800 Subject: [PATCH 3/3] fix(cuda): include first-byte address in resident feature identity --- src/models/qwen3_5/native/src/cuda.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/src/models/qwen3_5/native/src/cuda.rs b/src/models/qwen3_5/native/src/cuda.rs index a74ba136..b9dc094a 100644 --- a/src/models/qwen3_5/native/src/cuda.rs +++ b/src/models/qwen3_5/native/src/cuda.rs @@ -713,6 +713,7 @@ impl DeviceBf16Matrix<'_> { /// Request-local identity only: producer contents may change between calls. pub(crate) fn same_storage(&self, other: &Self) -> bool { std::ptr::eq(self.buffer, other.buffer) + && self.at(0) == other.at(0) && self.layout.rows == other.layout.rows && self.layout.columns == other.layout.columns && self.layout.device == other.layout.device