perf(cua): batch screenshot questions and accelerate opt-in resize

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

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<size_t>(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<cudaStream_t>(stream);
     gdn_batch_prep<<<dim3(chunks, H, count), THREADS, SMEM1_BYTES, st>>>(
         static_cast<const bf16*>(q), static_cast<const bf16*>(k), static_cast<const bf16*>(v), g,
         static_cast<const bf16*>(beta), batch, H, HK, scale);
-    gdn_batch_state<<<dim3(H, V / BVS, count), ST_THREADS, SMEM2_BYTES, st>>>(batch);
+    if (initial_state) {
+        gdn_batch_seeded_state<<<dim3(H, V / BVS, count), ST_THREADS, SMEM2_BYTES, st>>>(
+            batch, static_cast<const float*>(initial_state), static_cast<float*>(final_states), H);
+    } else {
+        gdn_batch_state<<<dim3(H, V / BVS, count), ST_THREADS, SMEM2_BYTES, st>>>(batch);
+    }
     gdn_batch_out<<<dim3(chunks, H, count), OUT_THREADS, SMEM3_BYTES, st>>>(batch, static_cast<bf16*>(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"

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.

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.

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<ProcessedImage> {
+    omni_qwen3_5_native::image_preprocess::preprocess_rgb8_with_workers(
+        width,
+        height,
+        rgb,
+        &omni_qwen3_5_native::image_preprocess::ImageLimits::cua(),
+        workers,
+    )
+}

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<f32>,
     device_features: bool,
+    batch_questions: bool,
     prefix_mode: prefix::Mode,
     prefix_storage: Option<crate::model::PrefixState>,
 }
@@ -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::<Result<Vec<_>>>()?;
-        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