cua_s1: add native multimodal language CUDA Graph replay

recipe/cua_s1/native.md

@@ -25,7 +25,7 @@ CUA_S1_MODEL=weights/cua-s1-4b-0.2-text-merged target/release/omni-cua-s1-native
 
 For the local CUDA Graph experiment, also set `CUA_S1_GRAPH=1`. The first use of
 each exact ordered prompt-length vector warms the GEMM plans and captures the forward pass;
-later requests replay it with freshly uploaded token ids. The text worker keeps a 64-entry FIFO cache. Growing the scratch allocation clears the captures before freeing
+later requests replay it with freshly uploaded token ids. Each language mode keeps a separate 64-entry FIFO cache. Growing the scratch allocation clears the captures before freeing
 their buffers. Capture adds first-use latency; leave the variable unset to use
 the eager control. Rebuild both the worker and CUDA library together (ABI 6).
 If capture fails, the worker returns the completed eager result and disables

recipe/cua_s1/native_vision_graph.md

@@ -4,8 +4,8 @@ The `omni-cua-s1-vision` screenshot worker can capture the complete vision
 encoder, including all 24 blocks, unmerged FP32 LoRA branches and the merger.
 Set `CUA_S1_VISION_GRAPH=1` before loading the worker to enable it. Leave it
 unset or set it to `0` for eager execution. This switch is independent of
-`CUA_S1_GRAPH`, which controls language execution; this change does not enable
-multimodal language Graph replay.
+`CUA_S1_GRAPH`, which controls text and multimodal language execution. Enable
+each switch separately; keep both unset for eager execution.
 
 Both modes retain one scratch allocation for the latest exact `[T,H,W]` patch
 grid. Equal patch counts with different heights or widths replace the scratch
@@ -57,9 +57,9 @@ Export the same `CUA_S1_BASE`, `CUA_S1_VISION_ADAPTER`, `CUA_S1_MODEL`, and
 ```sh
 cargo test --release --locked -p omni-cua-s1-native --lib vision_graph_ \
   -- --ignored --test-threads=1 --nocapture
-CUA_S1_VISION_GRAPH=0 cargo run --release --locked -p omni-cua-s1-native \
+CUA_S1_GRAPH=0 CUA_S1_VISION_GRAPH=0 cargo run --release --locked -p omni-cua-s1-native \
   --example vision_graph_bench -- /tmp/vision-eager.json
-CUA_S1_VISION_GRAPH=1 cargo run --release --locked -p omni-cua-s1-native \
+CUA_S1_GRAPH=0 CUA_S1_VISION_GRAPH=1 cargo run --release --locked -p omni-cua-s1-native \
   --example vision_graph_bench -- /tmp/vision-graph.json
 ```
 

src/models/qwen3_5/native/src/model.rs

@@ -598,9 +598,10 @@ pub struct Model {
     gemm: *mut c_void,
     /// Buffers for the largest packed token count so far; grows as needed.
     scratch: Option<Scratch>,
-    /// Opt-in replay with at most 64 captures keyed by ordered sequence lengths.
+    /// Opt-in replay with at most 64 captures per mode, keyed by sequence lengths.
     graph_enabled: bool,
     graphs: VecDeque<(Vec<usize>, cuda::Graph)>,
+    mm_graphs: VecDeque<(Vec<usize>, cuda::Graph)>,
 }
 
 // SAFETY: the raw pointers are device addresses and a cuBLASLt handle owned by the
@@ -610,6 +611,7 @@ unsafe impl Send for Model {}
 impl Drop for Model {
     fn drop(&mut self) {
         self.graphs.clear();
+        self.mm_graphs.clear();
         // SAFETY: created by cs1_gemm_create and not destroyed before.
         unsafe { (cuda::api().cs1_gemm_destroy)(self.gemm) };
     }
@@ -692,6 +694,7 @@ impl Model {
             scratch: None,
             graph_enabled: std::env::var("CUA_S1_GRAPH").as_deref() == Ok("1"),
             graphs: VecDeque::new(),
+            mm_graphs: VecDeque::new(),
         };
         Ok(model)
     }
@@ -745,6 +748,7 @@ impl Model {
         cuda::set_device(0)?;
         if self.scratch.as_ref().is_none_or(|s| t > s.cap) {
             self.graphs.clear();
+            self.mm_graphs.clear();
             self.scratch = None;
             self.scratch = Some(Scratch::new(
                 &self.cfg,
@@ -793,28 +797,21 @@ impl Model {
         if self.graph_enabled {
             if let Some((_, graph)) = self.graphs.iter().find(|(shape, _)| *shape == lengths) {
                 graph.launch(self.stream)?;
+                if std::env::var("CUA_S1_GRAPH_TRACE").as_deref() == Ok("1") {
+                    eprintln!("CUA_S1_GRAPH mode=text event=replay tokens={t}");
+                }
             } else {
                 // Warm plans and keep the eager result: run() advances the residual
                 // in place, so replaying on this cache miss would advance it twice.
                 self.run(s, &lengths, false)?;
                 cuda::synchronize(self.stream)?;
-                match cuda::Graph::capture(self.stream, || self.run(s, &lengths, false)) {
-                    Ok(graph) => {
-                        if self.graphs.len() == 64 {
-                            self.graphs.pop_front();
-                        }
-                        self.graphs.push_back((lengths.clone(), graph));
-                    }
-                    Err(error) => {
-                        eprintln!("CUDA Graph capture failed; using eager execution: {error:#}");
-                        self.graph_enabled = false;
-                        self.graphs.clear();
-                    }
-                }
+                let captured = cuda::Graph::capture(self.stream, || self.run(s, &lengths, false));
+                self.cache_graph(&lengths, false, captured);
             }
         } else {
             self.run(s, &lengths, false)?;
         }
+        let s = self.scratch.as_ref().unwrap();
         let mut end = 0;
         lengths
             .iter()
@@ -845,8 +842,58 @@ impl Model {
         self.upload_positions(s, input.position_ids)?;
         self.embed_tokens(s, input.token_ids)?;
         self.overwrite_image_rows(s, input, 0)?;
-        self.run(s, &[t], true)?;
-        self.last_hidden(s, t)
+        if self.graph_enabled {
+            if let Some((_, graph)) = self
+                .mm_graphs
+                .iter()
+                .find(|(shape, _)| shape.as_slice() == [t])
+            {
+                graph.launch(self.stream)?;
+                if std::env::var("CUA_S1_GRAPH_TRACE").as_deref() == Ok("1") {
+                    eprintln!("CUA_S1_GRAPH mode=multimodal event=replay tokens={t}");
+                }
+            } else {
+                // Uploads and embedding are outside capture. Keep the warmed eager
+                // result; launching now would advance the residual a second time.
+                self.run(s, &[t], true)?;
+                cuda::synchronize(self.stream)?;
+                let captured = cuda::Graph::capture(self.stream, || self.run(s, &[t], true));
+                self.cache_graph(&[t], true, captured);
+            }
+        } else {
+            self.run(s, &[t], true)?;
+        }
+        self.last_hidden(self.scratch.as_ref().unwrap(), t)
+    }
+
+    /// Both modes retain the eager miss result; capture records without executing.
+    fn cache_graph(&mut self, lengths: &[usize], multimodal: bool, captured: Result<cuda::Graph>) {
+        let mode = if multimodal { "multimodal" } else { "text" };
+        match captured {
+            Ok(graph) => {
+                let cache = if multimodal {
+                    &mut self.mm_graphs
+                } else {
+                    &mut self.graphs
+                };
+                if cache.len() == 64 {
+                    cache.pop_front();
+                }
+                cache.push_back((lengths.to_vec(), graph));
+                if std::env::var("CUA_S1_GRAPH_TRACE").as_deref() == Ok("1") {
+                    let t: usize = lengths.iter().sum();
+                    eprintln!("CUA_S1_GRAPH mode={mode} event=capture tokens={t}");
+                }
+            }
+            Err(error) => {
+                eprintln!(
+                    "CUDA Graph capture failed in {mode} mode; using eager execution: {error:#}"
+                );
+                self.graph_enabled = false;
+                self.graphs.clear();
+                self.mm_graphs.clear();
+            }
+        }
     }
 
     fn upload_positions(&self, s: &Scratch, positions: [&[i64]; 3]) -> Result<()> {