From 589a96c94535e8f0bb3f675571e1e9de2b773101 Mon Sep 17 00:00:00 2001 From: levius <2114377220@qq.com> Date: Thu, 8 Oct 2026 16:18:41 +0800 Subject: [PATCH 1/2] qwen: add shared 27B vision and multimodal CUDA Graph replay --- src/models/qwen3_5/native/src/model.rs | 79 ++++- src/models/qwen3_5/native/src/vision/model.rs | 273 +++++++++++++----- 2 files changed, 263 insertions(+), 89 deletions(-) diff --git a/src/models/qwen3_5/native/src/model.rs b/src/models/qwen3_5/native/src/model.rs index 6d37cd2f..b243f229 100644 --- a/src/models/qwen3_5/native/src/model.rs +++ b/src/models/qwen3_5/native/src/model.rs @@ -563,9 +563,10 @@ pub struct Model { gemm: *mut c_void, /// Buffers for the largest packed token count so far; grows as needed. scratch: Option, - /// 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, cuda::Graph)>, + mm_graphs: VecDeque<(Vec, cuda::Graph)>, } // SAFETY: the raw pointers are device addresses and a cuBLASLt handle owned by the @@ -575,6 +576,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) }; } @@ -656,6 +658,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) } @@ -709,6 +712,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, @@ -757,28 +761,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() @@ -833,8 +830,58 @@ impl Model { } begin = end; } - 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) { + 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<()> { diff --git a/src/models/qwen3_5/native/src/vision/model.rs b/src/models/qwen3_5/native/src/vision/model.rs index c2821b61..c4f71638 100644 --- a/src/models/qwen3_5/native/src/vision/model.rs +++ b/src/models/qwen3_5/native/src/vision/model.rs @@ -7,7 +7,11 @@ use crate::{ use anyhow::{Result, ensure}; use half::bf16; use safetensors::{Dtype, tensor::TensorView}; -use std::{collections::BTreeMap, ffi::c_void, path::Path}; +use std::{ + collections::{BTreeMap, VecDeque}, + ffi::c_void, + path::Path, +}; type Trace<'a> = Option<&'a mut dyn FnMut(&str, &[bf16]) -> Result<()>>; @@ -30,10 +34,13 @@ impl Drop for Gemm { /// One request at a time; returns row-major [image_tokens, config.out_hidden_size]. pub struct VisionModel { config: VisionConfig, + // Captures and their buffers retire before weights, GEMM workspace and stream. + scratch: VecDeque, + graph_enabled: bool, base: BTreeMap, lora: BTreeMap, - stream: OwnedStream, gemm: Gemm, + stream: OwnedStream, } impl VisionModel { pub fn load(dir: impl AsRef, library: &Path) -> Result { @@ -119,6 +126,8 @@ impl VisionModel { ensure!(!gemm.0.is_null(), "cannot create vision cuBLAS handle"); let mut model = Self { config, + scratch: VecDeque::new(), + graph_enabled: std::env::var("CUA_S1_VISION_GRAPH").as_deref() == Ok("1"), base: BTreeMap::new(), lora: BTreeMap::new(), stream, @@ -308,11 +317,15 @@ impl VisionModel { } fn run(&mut self, image: &ProcessedImage, mut callback: Trace<'_>) -> Result> { cuda::set_device(0)?; - let geo = VisionGeometry::new(image.image_grid_thw, &self.config)?; - let h = self.config.hidden_size; - let intermediate = self.config.intermediate_size; - let merged = h * 4; - let output = self.config.out_hidden_size; + let cached = self + .scratch + .iter() + .position(|s| s.grid == image.image_grid_thw); + let geometry = if cached.is_none() { + Some(VisionGeometry::new(image.image_grid_thw, &self.config)?) + } else { + None + }; let n = image.image_grid_thw[1] * image.image_grid_thw[2]; ensure!( image.pixel_values.len() == n * 1536, @@ -332,50 +345,82 @@ impl VisionModel { .iter() .flat_map(|v| bf16::from_f32(*v).to_bits().to_le_bytes()) .collect(); - let pixels = self.upload(&pixels)?; - let indices = self.upload( - &geo.indices - .iter() - .flat_map(|v| v.to_le_bytes()) - .collect::>(), - )?; - let weights = self.upload( - &geo.weights - .iter() - .flat_map(|v| v.to_le_bytes()) - .collect::>(), - )?; - let co = self.upload( - &geo.cos - .iter() - .flat_map(|v| v.to_le_bytes()) - .collect::>(), - )?; - let si = self.upload( - &geo.sin - .iter() - .flat_map(|v| v.to_le_bytes()) - .collect::>(), - )?; - let w = Work::new( - if self.lora.is_empty() { 0 } else { n }, - intermediate.max(merged).max(output), - )?; - let attention_workspace = DeviceBuffer::new(if h == 1024 { - 0 + let index = if let Some(index) = cached { + index } else { - n * self.config.num_heads * 80 * 4 * 2 - })?; - let x = DeviceBuffer::new(n * h * 2)?; - let norm = DeviceBuffer::new(n * h * 2)?; - let qkv = DeviceBuffer::new(n * h * 3 * 2)?; - let q = DeviceBuffer::new(n * h * 2)?; - let k = DeviceBuffer::new(n * h * 2)?; - let attn = DeviceBuffer::new(n * h * 2)?; - let delta = DeviceBuffer::new(n * h * 2)?; - let mlp = DeviceBuffer::new(n * intermediate.max(merged) * 2)?; - self.linear("patch_embed.proj", &pixels, &x, n, h, 1536, &w)?; - self.trace(&mut callback, "patch_embed", &x, n * h)?; + self.synchronize()?; + // FIFO is bounded to four exact grids; graph drops before its buffers. + if self.scratch.len() == 4 { + self.scratch.pop_front(); + } + let scratch = Scratch::new(self, image.image_grid_thw, geometry.unwrap())?; + self.scratch.push_back(scratch); + self.scratch.len() - 1 + }; + let scratch = &self.scratch[index]; + // SAFETY: the current image has the validated exact resident buffer size. + unsafe { + cuda::upload(scratch.pixels.at(0), &pixels, self.stream.0)?; + } + if callback.is_none() + && self.graph_enabled + && let Some(graph) = &scratch.graph + { + graph.launch(self.stream.0)?; + graph_trace("replayed", scratch.grid); + return self.read(&scratch.out, n / 4 * self.config.out_hidden_size); + } + self.enqueue(scratch, &mut callback)?; + let result = self.read(&scratch.out, n / 4 * self.config.out_hidden_size)?; + if let Some(callback) = callback.as_mut() { + callback("merger.output", &result)?; + } else if self.graph_enabled { + // Readback completed warmup. Capture records without executing; keep + // the eager result even when recording/instantiation fails. + let captured = cuda::Graph::capture(self.stream.0, || self.enqueue(scratch, &mut None)); + match captured { + Ok(graph) => { + self.scratch[index].graph = Some(graph); + graph_trace("captured", image.image_grid_thw); + } + Err(error) => { + eprintln!("Vision CUDA Graph capture failed; using eager execution: {error:#}"); + self.graph_enabled = false; + for scratch in &mut self.scratch { + scratch.graph = None; + } + } + } + } + Ok(result) + } + fn enqueue(&self, scratch: &Scratch, callback: &mut Trace<'_>) -> Result<()> { + let n = scratch.grid[1] * scratch.grid[2]; + let h = self.config.hidden_size; + let intermediate = self.config.intermediate_size; + let merged = h * 4; + let output = self.config.out_hidden_size; + let Scratch { + pixels, + indices, + weights, + co, + si, + w, + attention_workspace, + x, + norm, + qkv, + q, + k, + attn, + delta, + mlp, + out, + .. + } = scratch; + self.linear("patch_embed.proj", pixels, x, n, h, 1536, w)?; + self.trace(callback, "patch_embed", x, n * h)?; unsafe { if h == 1024 { check( @@ -404,11 +449,11 @@ impl VisionModel { )?; } } - self.trace(&mut callback, "position", &x, n * h)?; + self.trace(callback, "position", x, n * h)?; for i in 0..self.config.depth { let p = format!("blocks.{i}"); - self.norm(&format!("{p}.norm1"), &x, &norm, n)?; - self.linear(&format!("{p}.attn.qkv"), &norm, &qkv, n, h * 3, h, &w)?; + self.norm(&format!("{p}.norm1"), x, norm, n)?; + self.linear(&format!("{p}.attn.qkv"), norm, qkv, n, h * 3, h, w)?; unsafe { if h == 1024 { check( @@ -466,22 +511,22 @@ impl VisionModel { )?; } } - self.linear(&format!("{p}.attn.proj"), &attn, &delta, n, h, h, &w)?; + self.linear(&format!("{p}.attn.proj"), attn, delta, n, h, h, w)?; unsafe { check( (api().cs1_vision_add)(x.at(0), delta.at(0), n * h, self.stream.0), "vision attention residual", )?; } - self.norm(&format!("{p}.norm2"), &x, &norm, n)?; + self.norm(&format!("{p}.norm2"), x, norm, n)?; self.linear( &format!("{p}.mlp.linear_fc1"), - &norm, - &mlp, + norm, + mlp, n, intermediate, h, - &w, + w, )?; unsafe { check( @@ -491,12 +536,12 @@ impl VisionModel { } self.linear( &format!("{p}.mlp.linear_fc2"), - &mlp, - &delta, + mlp, + delta, n, h, intermediate, - &w, + w, )?; unsafe { check( @@ -504,28 +549,110 @@ impl VisionModel { "vision MLP residual", )?; } - self.trace(&mut callback, &p, &x, n * h)?; + self.trace(callback, &p, x, n * h)?; } - self.norm("merger.norm", &x, &norm, n)?; - self.trace(&mut callback, "merger.norm", &norm, n * h)?; + self.norm("merger.norm", x, norm, n)?; + self.trace(callback, "merger.norm", norm, n * h)?; // Consecutive groups of four patches already have the required 2x2 merge order. - self.linear("merger.linear_fc1", &norm, &mlp, n / 4, merged, merged, &w)?; - self.trace(&mut callback, "merger.linear_fc1", &mlp, n * h)?; + self.linear("merger.linear_fc1", norm, mlp, n / 4, merged, merged, w)?; + self.trace(callback, "merger.linear_fc1", mlp, n * h)?; unsafe { check( (api().cs1_vision_gelu)(mlp.at(0), n * h, 1, self.stream.0), "vision exact GELU", )?; } - let out = DeviceBuffer::new(n / 4 * output * 2)?; - self.linear("merger.linear_fc2", &mlp, &out, n / 4, output, merged, &w)?; - let result = self.read(&out, n / 4 * output)?; - if let Some(callback) = callback.as_mut() { - callback("merger.output", &result)?; - } - Ok(result) + self.linear("merger.linear_fc2", mlp, out, n / 4, output, merged, w)?; + Ok(()) + } +} +impl Drop for VisionModel { + fn drop(&mut self) { + // Even failure paths retain every referenced address through synchronization. + let _ = self.synchronize(); + } +} +fn graph_trace(action: &str, grid: [usize; 3]) { + if std::env::var("CUA_S1_GRAPH_TRACE").as_deref() == Ok("1") { + eprintln!("Vision CUDA Graph {action} grid={grid:?}"); } } +struct Scratch { + graph: Option, + grid: [usize; 3], + pixels: DeviceBuffer, + indices: DeviceBuffer, + weights: DeviceBuffer, + co: DeviceBuffer, + si: DeviceBuffer, + w: Work, + attention_workspace: DeviceBuffer, + x: DeviceBuffer, + norm: DeviceBuffer, + qkv: DeviceBuffer, + q: DeviceBuffer, + k: DeviceBuffer, + attn: DeviceBuffer, + delta: DeviceBuffer, + mlp: DeviceBuffer, + out: DeviceBuffer, +} +impl Scratch { + fn new(model: &VisionModel, grid: [usize; 3], geo: VisionGeometry) -> Result { + let n = grid[1] * grid[2]; + let h = model.config.hidden_size; + let width = model + .config + .intermediate_size + .max(h * 4) + .max(model.config.out_hidden_size); + Ok(Self { + graph: None, + grid, + pixels: DeviceBuffer::new(n * 1536 * 2)?, + indices: model.upload( + &geo.indices + .iter() + .flat_map(|v| v.to_le_bytes()) + .collect::>(), + )?, + weights: model.upload( + &geo.weights + .iter() + .flat_map(|v| v.to_le_bytes()) + .collect::>(), + )?, + co: model.upload( + &geo.cos + .iter() + .flat_map(|v| v.to_le_bytes()) + .collect::>(), + )?, + si: model.upload( + &geo.sin + .iter() + .flat_map(|v| v.to_le_bytes()) + .collect::>(), + )?, + w: Work::new(if model.lora.is_empty() { 0 } else { n }, width)?, + attention_workspace: DeviceBuffer::new(if h == 1024 { + 0 + } else { + n * model.config.num_heads * 80 * 4 * 2 + })?, + x: DeviceBuffer::new(n * h * 2)?, + norm: DeviceBuffer::new(n * h * 2)?, + qkv: DeviceBuffer::new(n * h * 3 * 2)?, + q: DeviceBuffer::new(n * h * 2)?, + k: DeviceBuffer::new(n * h * 2)?, + attn: DeviceBuffer::new(n * h * 2)?, + delta: DeviceBuffer::new(n * h * 2)?, + mlp: DeviceBuffer::new(n * width * 2)?, + out: DeviceBuffer::new(n / 4 * model.config.out_hidden_size * 2)?, + }) + } +} + struct Work { float_input: DeviceBuffer, rank: DeviceBuffer, From 89746b8fc891c4aac219a7421f48ac33e110f9e0 Mon Sep 17 00:00:00 2001 From: levius <2114377220@qq.com> Date: Fri, 9 Oct 2026 08:52:44 +0800 Subject: [PATCH 2/2] docs(jemm): retain graph replay switches and cache limits --- src/models/jemm/README.md | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/src/models/jemm/README.md b/src/models/jemm/README.md index b1083c63..7979f292 100644 --- a/src/models/jemm/README.md +++ b/src/models/jemm/README.md @@ -79,6 +79,22 @@ covers the entire request. Both streams synchronize before resource release; caught execution errors/panics clear atomic readiness and retire loaded state. There is no cross-request batching or prefix cache. +## Optional CUDA Graph replay + +`CUA_S1_GRAPH=1` enables language replay, with separate 64-entry FIFO caches +for text ordered lengths and multimodal `[tokens]`. `CUA_S1_VISION_GRAPH=1` +independently enables vision replay; its four-slot FIFO uses exact `[T,H,W]` +grids. Both switches default off. Current IDs/features/positions/pixels are +uploaded before replay. An eager miss completes before recording, preserving its +answer. Capture failure logs unconditionally, disables that graph path and keeps +the eager result. Eviction synchronizes before graph/buffer destruction. + +Vision scratch persists for up to four grids even with graphs disabled, changing +resident activation memory from per-call allocation. This bounds grid slots, +not weights or GEMM plan memory. Startup's image warmup can now capture its own +shape when enabled; it does not prewarm every user grid. These mechanisms do not +establish lower latency on every workload or fix the multi-image numerical gap. + ## Validation and known multi-image limitation The historical A800 BF16 corpus (13 requests / 16 questions, text and single-image