diff --git a/docs/architecture.md b/docs/architecture.md index 2846528b..c4aa209e 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -12,7 +12,8 @@ that share the [Qwen3.5/3.8 executor](../src/models/qwen3_5/native/), which acce single prompts and bounded packed prefill. Cua-S1 and Open-Jev-9B use single-prompt calls; Open-Jev-27B-v1.1 packs candidates within one request for input and gate/up GEMMs while preserving per-sequence mixers and output/down GEMM -shapes. +shapes. The executor can also run prompts that share token prefixes, each prefix +once (below); no worker uses that yet. [Laya's native worker](../src/models/laya/README.md) uses a separate Hopper CUDA backend for one complete padded request. All three coordinate independent processors and executors through @@ -32,6 +33,24 @@ Graph captures Encoder/Decision; gather, scorer/action head and synchronized readback remain outside capture. Native Metal remains planned; Python workers retain their documented reference/serving roles. +The Qwen executor's `forward_shared` takes a request prefix, group prefixes and +branches, ends each prefix at its last multiple of 64 tokens, and returns one final +hidden state per branch in order +([#85](https://github.com/ThinkFlowLab/system1-omni/issues/85)). It runs on the +prefix continuation that JEV-VL's cache uses: each prefix is captured once (per +linear-attention layer the float32 recurrent state and the last three conv inputs, +per full-attention layer the prefix keys and values), and each branch then runs as +its own pass from that state. The two prefix states persist across calls and grow only for a +longer prefix; their contents are valid only within one call, which runs eagerly and +synchronizes its stream before returning, also on error. Its GEMMs keep one +algorithm per weight shape, so on GPUs where that keeps rows independent of M +(checked by `fixed_gemm_rows_do_not_depend_on_m`, so far on sm_89) each result equals +`forward_fixed` on the full prompt bit for bit; it differs from `forward` by +rounding. Both take merged weights with the default numerics only, and refuse an +unmerged adapter, reference numerics or FP32 GDN before any device work. Each +prefix state takes about 50 MiB on Qwen3.5-4B and 9B, plus 32 KiB per position of +its prefix rounded up to 1,024 positions, and the GEMM handle a 32 MiB workspace. + ## Native worker boundaries The native workers separate `processing.rs` from `executor.rs`; their worker diff --git a/src/backends/cuda/qwen3_5/README.md b/src/backends/cuda/qwen3_5/README.md index 56df461f..3a5b918d 100644 --- a/src/backends/cuda/qwen3_5/README.md +++ b/src/backends/cuda/qwen3_5/README.md @@ -44,7 +44,10 @@ entry points. Open-Jev uses this path within requests; Cua-S1 keeps single-promp Prefix continuation retains full-attention KV, three convolution input rows, and FP32 Gated DeltaNet state at 64-token chunk boundaries. It is used by the -experimental JEV-VL worker; existing full-prefill callers keep their own paths. +experimental JEV-VL worker and by the shared executor's `forward_shared`, which +runs prompts that share prefixes once per prefix with the fixed-algorithm GEMM +handle ([#85](https://github.com/ThinkFlowLab/system1-omni/issues/85)); existing +full-prefill callers keep their own paths. ABI 6 adds mandatory copy and prefix entry points after ABI 5 native vision, and ABI 9 adds the fixed-algorithm GEMM handle. Rebuild the library and **all** Qwen worker binaries together; older libraries diff --git a/src/models/qwen3_5/native/Cargo.toml b/src/models/qwen3_5/native/Cargo.toml index 4b3c2ed0..54ec298f 100644 --- a/src/models/qwen3_5/native/Cargo.toml +++ b/src/models/qwen3_5/native/Cargo.toml @@ -25,3 +25,7 @@ path = "../../../../tests/qwen3_5/kernels.rs" [[test]] name = "device_features" path = "../../../../tests/qwen3_5/device_features.rs" + +[[test]] +name = "prefix" +path = "../../../../tests/qwen3_5/prefix.rs" diff --git a/src/models/qwen3_5/native/src/model.rs b/src/models/qwen3_5/native/src/model.rs index d092943d..5aa60e8c 100644 --- a/src/models/qwen3_5/native/src/model.rs +++ b/src/models/qwen3_5/native/src/model.rs @@ -8,6 +8,9 @@ //! their kernels). Text prompts use one position per //! token. The explicit multimodal boundary inserts adapted image rows and supplies //! the interleaved temporal/height/width rotary positions to the same layer loop. +//! +//! `Model::forward_shared` runs prompts that share token prefixes, running each prefix +//! once and every branch from the state after it (see `SharedPrompts`). use std::collections::{HashMap, VecDeque}; use std::ffi::c_void; @@ -25,6 +28,9 @@ const ALIGN: usize = 256; const BF16: usize = 2; const F32: usize = 4; const GEMM_WORKSPACE: usize = 32 << 20; +/// Rows at which `forward_shared` fixes each weight shape's GEMM algorithm (see +/// cs1_gemm_create_fixed); its branches run tens of tokens. +const FIXED_GEMM_M: i32 = 64; #[derive(Debug, Clone)] pub struct Config { @@ -593,6 +599,8 @@ pub struct PrefixState { /// Tokens this prefix covers; always a multiple of 64 (the GDN chunk length), /// which also aligns the flash-attention key tiles of the one-shot pass. len: usize, + /// Tokens the buffers hold; `forward_shared` reuses a state for shorter prefixes. + cap: usize, /// Post-prep K [len, Hk*Dh] and raw V [len, Hk*Dh] rows per full-attention layer. attn_kv: Vec<(DeviceBuffer, DeviceBuffer)>, /// Float32 [H, K, V] recurrent states, one per Gated DeltaNet layer. @@ -615,6 +623,62 @@ impl PrefixState { } } +/// Gated DeltaNet chunk length. `forward_shared` ends prefixes at its multiples, where +/// a continuation repeats the one-shot arithmetic (see cs1_gdn_prefill_x). +const GDN_CHUNK: usize = 64; + +/// Prompts that share token prefixes, for `Model::forward_shared`. Branch `b` of group `g` +/// is the prompt `prefix`, then `groups[g].prefix`, then `groups[g].branches[b]`. How the +/// prompts are grouped does not change any result where the fixed GEMM algorithms keep +/// rows independent of M (see `forward_shared`). +pub struct SharedPrompts<'a> { + pub prefix: &'a [u32], + pub groups: Vec>, +} + +pub struct PromptGroup<'a> { + pub prefix: &'a [u32], + pub branches: Vec<&'a [u32]>, +} + +/// The token runs `forward_shared` executes: the request prefix, then per group its own +/// prefix and branches. Each prefix ends at the last multiple of `GDN_CHUNK` within it; +/// its remaining tokens run with every group or branch after it. +struct Runs { + prefix: Vec, + groups: Vec, +} + +struct GroupRuns { + prefix: Vec, + branches: Vec>, +} + +fn runs(prompts: &SharedPrompts) -> Runs { + let end = |n: usize| n - n % GDN_CHUNK; + let p = end(prompts.prefix.len()); + let groups = prompts + .groups + .iter() + .map(|g| { + let shared = [prompts.prefix, g.prefix].concat(); + let q = end(shared.len()); + GroupRuns { + prefix: shared[p..q].to_vec(), + branches: g + .branches + .iter() + .map(|b| [&shared[q..], b].concat()) + .collect(), + } + }) + .collect(); + Runs { + prefix: prompts.prefix[..p].to_vec(), + groups, + } +} + pub struct Model { pub cfg: Config, prefix_owner: Arc<()>, @@ -627,8 +691,12 @@ pub struct Model { gemm: *mut c_void, reference: Option<&'static cuda::Reference>, fp32_gdn: Option, + /// One GEMM algorithm per weight shape for `forward_shared`; created on first use. + fixed_gemm: *mut c_void, /// 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. + shared: [Option; 2], /// Opt-in replay with at most 64 captures per mode, keyed by sequence lengths. graph_enabled: bool, graphs: VecDeque<(Vec, cuda::Graph)>, @@ -637,7 +705,7 @@ pub struct Model { _lora_stream: Option, } -// SAFETY: the raw pointers are device addresses and a cuBLASLt handle owned by the +// SAFETY: the raw pointers are device addresses and cuBLASLt handles owned by the // model; the engine runs one forward pass at a time behind a mutex. unsafe impl Send for Model {} @@ -662,8 +730,13 @@ impl Drop for Model { } // Retire the FP32 cuBLAS context while its last stream is still valid. self.fp32_gdn = None; - // SAFETY: created by cs1_gemm_create and not destroyed before. - unsafe { (cuda::api().cs1_gemm_destroy)(self.gemm) }; + // SAFETY: created by cs1_gemm_create(_fixed) and not destroyed before. + unsafe { + (cuda::api().cs1_gemm_destroy)(self.gemm); + if !self.fixed_gemm.is_null() { + (cuda::api().cs1_gemm_destroy)(self.fixed_gemm); + } + } } } @@ -771,7 +844,9 @@ impl Model { gemm, reference: None, fp32_gdn: None, + fixed_gemm: std::ptr::null_mut(), scratch: None, + shared: [None, None], graph_enabled: std::env::var("CUA_S1_GRAPH").as_deref() == Ok("1"), graphs: VecDeque::new(), _lora_stream: stream_owner, @@ -836,12 +911,25 @@ impl Model { } fn gemm(&self, s: &Scratch, x: usize, w: &Tensor, y: usize, m: usize) -> Result<()> { + self.gemm_with(self.gemm, s, x, w, y, m) + } + + /// `gemm` with the cuBLASLt handle `handle`. + fn gemm_with( + &self, + handle: *mut c_void, + s: &Scratch, + x: usize, + w: &Tensor, + y: usize, + m: usize, + ) -> Result<()> { let (n, k) = (w.shape[0] as i32, w.shape[1] as i32); // SAFETY: x and y are scratch buffers sized for m rows of w's shape. check( unsafe { (cuda::api().cs1_gemm)( - self.gemm, + handle, s.at(x), w.ptr, s.at(y), @@ -866,9 +954,25 @@ impl Model { m: usize, lengths: &[usize], group: Group, + ) -> Result<()> { + self.gemm_group_with(self.gemm, s, x, w, y, m, lengths, group) + } + + /// `gemm_group` with the cuBLASLt handle `handle`. + #[allow(clippy::too_many_arguments)] // Base projection coordinates and sequence shape. + fn gemm_group_with( + &self, + handle: *mut c_void, + s: &Scratch, + x: usize, + w: &Tensor, + y: usize, + m: usize, + lengths: &[usize], + group: Group, ) -> Result<()> { if self.lora.is_none() { - return self.gemm(s, x, w, y, m); + return self.gemm_with(handle, s, x, w, y, m); } let components = lora::component_rows(&self.cfg, group)?; ensure!( @@ -887,7 +991,7 @@ impl Model { w.shape[1], &components, work, - self.gemm, + handle, self.stream, ) } @@ -1256,6 +1360,7 @@ impl Model { owner: Arc::clone(&self.prefix_owner), initialized: false, len, + cap: len, attn_kv, gdn_state, conv_tail, @@ -1297,7 +1402,7 @@ 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_window(s, 0, t, true, Some(state), None)?; + self.run_window(s, 0, t, true, Some(state), None, self.gemm)?; self.synchronize()?; state.initialized = true; Ok(()) @@ -1339,10 +1444,186 @@ impl Model { self.upload_positions(s, input.position_ids)?; self.embed_tokens_at(s, input.token_ids, prefix.len)?; self.overwrite_image_rows(s, input, prefix.len)?; - self.run_window(s, prefix.len, tend, true, None, Some(prefix))?; + self.run_window(s, prefix.len, tend, true, None, Some(prefix), self.gemm)?; self.last_hidden(s, tend) } + /// `forward` with the GEMM algorithms of `forward_shared`, one per weight shape, run + /// eagerly: the result that `forward_shared` reproduces for this prompt. + pub fn forward_fixed(&mut self, ids: &[u32]) -> Result> { + self.ensure_fixed_profile()?; + let result = self.run_fixed(ids); + if result.is_err() { + // Work already queued may still use the buffers. + let _ = cuda::synchronize(self.stream); + } + result + } + + fn run_fixed(&mut self, ids: &[u32]) -> Result> { + let t = ids.len(); + ensure!( + t > 0 && t <= self.cfg.max_positions, + "empty or oversized prompt" + ); + ensure!( + ids.iter().all(|&i| (i as usize) < self.embed.shape[0]), + "token id outside the vocabulary" + ); + self.prepare_scratch(t)?; + let gemm = self.fixed_gemm()?; + let s = self.scratch.as_ref().unwrap(); + self.embed_tokens(s, ids)?; + self.run_window(s, 0, t, false, None, None, gemm)?; + self.last_hidden(s, t) + } + + /// The final-norm hidden state at the last position of every branch, grouped and + /// ordered as in `prompts`. Each prefix runs once, up to its last multiple of 64 + /// tokens, capturing its state as `forward_multimodal_capture` does, and each branch + /// continues from the state after its prefixes, with its positions and causal mask + /// continuing theirs. A prefix's remaining tokens run with every group or branch + /// after it. + /// + /// On a GPU where the fixed GEMM algorithms keep rows independent of M (see + /// cs1_gemm_create_fixed), each result equals `forward_fixed` on the full prompt bit + /// for bit, however the prompts are grouped. It differs from `forward`, which picks + /// GEMM algorithms per row count, by rounding. + /// + /// Runs eagerly. The prefix states are overwritten by the next call, and the stream + /// is synchronized before this returns, also on error. + pub fn forward_shared(&mut self, prompts: &SharedPrompts) -> Result>>> { + self.ensure_fixed_profile()?; + let result = self.run_shared(prompts); + if result.is_err() { + // Work already queued may still use the buffers. + let _ = cuda::synchronize(self.stream); + } + result + } + + fn run_shared(&mut self, prompts: &SharedPrompts) -> Result>>> { + let vocab = self.embed.shape[0]; + ensure!(!prompts.groups.is_empty(), "no prompts"); + let mut longest = 0; + for g in &prompts.groups { + ensure!(!g.branches.is_empty(), "a prompt group without branches"); + for b in &g.branches { + ensure!(!b.is_empty(), "an empty branch"); + longest = longest.max(prompts.prefix.len() + g.prefix.len() + b.len()); + } + } + ensure!(longest <= self.cfg.max_positions, "oversized prompt"); + let all = prompts + .groups + .iter() + .flat_map(|g| g.prefix.iter().chain(g.branches.iter().copied().flatten())); + ensure!( + prompts + .prefix + .iter() + .chain(all) + .all(|&i| (i as usize) < vocab), + "token id outside the vocabulary" + ); + let runs = runs(prompts); + let p = runs.prefix.len(); + // the longest group prefix that runs on its own, if any + let q = runs + .groups + .iter() + .filter(|g| !g.prefix.is_empty()) + .map(|g| p + g.prefix.len()) + .max() + .unwrap_or(0); + self.prepare_scratch(longest)?; + let gemm = self.fixed_gemm()?; + let [mut request, mut group] = std::mem::take(&mut self.shared); + let result = (|| { + let request = self.shared_state(&mut request, p)?; + let group = self.shared_state(&mut group, q)?; + let s = self.scratch.as_ref().unwrap(); + let request = match request { + Some(state) => { + self.embed_tokens(s, &runs.prefix)?; + self.run_window(s, 0, p, false, Some(&mut *state), None, gemm)?; + Some(&*state) + } + None => None, + }; + let mut group = group; + let mut hidden = Vec::with_capacity(runs.groups.len()); + for run in &runs.groups { + let start = p + run.prefix.len(); + let from = if run.prefix.is_empty() { + request + } else { + let state = group.as_deref_mut().unwrap(); + state.len = start; + self.embed_tokens_at(s, &run.prefix, p)?; + self.run_window(s, p, start, false, Some(&mut *state), request, gemm)?; + Some(&*state) + }; + let rows = run + .branches + .iter() + .map(|b| { + let end = start + b.len(); + self.embed_tokens_at(s, b, start)?; + self.run_window(s, start, end, false, None, from, gemm)?; + self.last_hidden(s, end) + }) + .collect::>>()?; + hidden.push(rows); + } + Ok(hidden) + })(); + self.shared = [request, group]; + result + } + + /// `slot`'s state resized to hold `len` tokens, none for an empty prefix. + fn shared_state<'s>( + &self, + slot: &'s mut Option, + len: usize, + ) -> Result> { + if len == 0 { + return Ok(None); + } + if slot.as_ref().is_none_or(|state| state.cap < len) { + *slot = None; + let cap = len + .next_multiple_of(1024) + .min(self.cfg.max_positions - self.cfg.max_positions % GDN_CHUNK); + *slot = Some(self.alloc_prefix(cap.max(len))?); + } + let state = slot.as_mut().unwrap(); + state.len = len; + Ok(Some(state)) + } + + /// `forward_fixed` and `forward_shared` cover merged weights with the default + /// numerics. Refuse the other profiles before any allocation or CUDA call. + fn ensure_fixed_profile(&self) -> Result<()> { + ensure!( + self.lora.is_none() && self.reference.is_none() && self.fp32_gdn.is_none(), + "fixed-algorithm GEMMs and shared prefixes support only merged weights with the default numerics" + ); + Ok(()) + } + + /// The GEMM handle with one algorithm per weight shape, created on first use. + fn fixed_gemm(&mut self) -> Result<*mut c_void> { + if self.fixed_gemm.is_null() { + // SAFETY: plain allocation; checked for null below. + self.fixed_gemm = + unsafe { (cuda::api().cs1_gemm_create_fixed)(GEMM_WORKSPACE, FIXED_GEMM_M) }; + ensure!(!self.fixed_gemm.is_null(), "cuBLASLt setup failed"); + } + Ok(self.fixed_gemm) + } + fn embed_tokens(&self, s: &Scratch, ids: &[u32]) -> Result<()> { self.embed_tokens_at(s, ids, 0) } @@ -1677,15 +1958,19 @@ impl Model { /// The R2d prefix-cache pass: run rows `[qb, tend)` of one prompt with /// explicit positions, touching per-layer buffers only for those rows. /// - /// `capture` (qb must be 0): collect post-prep K and pre-GEMM V columns of each + /// `capture` (`tend` rows): collect post-prep K and pre-GEMM V columns of each /// full-attention layer plus the float32 GDN state and the three pre-conv - /// projection columns of each Gated DeltaNet layer. `prefix` (qb = prefix.len): + /// projection columns of each Gated DeltaNet layer, for rows `[0, tend)`; with a + /// `prefix` too, the state after it continued. `prefix` (qb = prefix.len): /// seed the window from a captured entry — the conv tail columns are copied /// back into place before the convolution, and the prefix K/V columns before /// the attention. Everything else is `run()` with shifted row bases, so a /// captured+continued pass repeats the one-shot arithmetic token for token /// (qb is a multiple of the 64-row GDN chunk; the only remaining divergence - /// is cuBLASLt's M-shaped plan choice, an f32-accumulation-order ulp). + /// is cuBLASLt's M-shaped plan choice, an f32-accumulation-order ulp, which the + /// fixed-algorithm `gemm` handle removes). Custom positions hold the window's + /// rows from index 0; the text tables are indexed by position. + #[allow(clippy::too_many_arguments)] // One window: rows, positions, states and GEMM handle. fn run_window( &self, s: &Scratch, @@ -1694,13 +1979,16 @@ impl Model { custom_positions: bool, capture: Option<&mut PrefixState>, prefix: Option<&PrefixState>, + gemm: *mut c_void, ) -> Result<()> { ensure!( self.fp32_gdn.is_none(), "FP32 GDN does not support prefix continuation" ); ensure!( - qb <= tend && (qb == 0) == prefix.is_none() && (qb == 0 || capture.is_none()), + qb < tend + && prefix.map_or(qb == 0, |pre| pre.len == qb) + && capture.as_deref().is_none_or(|cap| cap.len == tend), "prefix window shape" ); let cfg = &self.cfg; @@ -1714,7 +2002,8 @@ impl Model { let (cos, sin) = if custom_positions { (s.custom_cos, s.custom_sin) } else { - (s.cos, s.sin) + let rope = qb * cfg.rotary_half * BF16; + (s.cos + rope, s.sin + rope) }; let hb = cfg.hidden * BF16; let ob = cfg.heads * cfg.head_dim * BF16; // post-prep q/gate and ao row bytes @@ -1729,7 +2018,8 @@ impl Model { let mut la_i = 0usize; // SAFETY (every kernel call below): pointers are weights in the arena or // scratch buffers laid out for tend tokens at the widths used here; prefix - // buffers hold exactly the shapes allocated by alloc_prefix(qb). + // buffers hold at least `len` rows (a capture's `tend`) of the shapes + // allocated by alloc_prefix. unsafe { check( (self @@ -1749,7 +2039,8 @@ impl Model { for (i, layer) in self.layers.iter().enumerate() { match &layer.mixer { Mixer::Linear(la) => { - self.gemm_group( + self.gemm_group_with( + gemm, s, s.x + qb * hb, &la.in_proj, @@ -1770,28 +2061,26 @@ impl Model { let z = s.gdn_in + w.conv * BF16; let b = z + vd * BF16; let a = b + hv * BF16; - let (s_in, s_out): (*const c_void, *mut c_void) = - match (&prefix, capture.as_deref()) { - (Some(pre), _) => { - (pre.gdn_state[la_i].at(0).cast_const(), std::ptr::null_mut()) - } - (None, Some(cap)) => (std::ptr::null(), cap.gdn_state[la_i].at(0)), - (None, None) => (std::ptr::null(), std::ptr::null_mut()), - }; + let s_in: *const c_void = + prefix.map_or(std::ptr::null(), |pre| pre.gdn_state[la_i].at(0)); + let s_out: *mut c_void = capture + .as_deref() + .map_or(std::ptr::null_mut(), |cap| cap.gdn_state[la_i].at(0)); unsafe { - if let Some(cap) = capture.as_deref() { - // conv window tail: the last three pre-conv columns. + if qb > 0 { cuda::copy_dd( - cap.conv_tail[la_i].at(0), - p(s.gdn_in + (rows - 3) * ldbb), + p(s.gdn_in + (qb - 3) * ldbb), + prefix.unwrap().conv_tail[la_i].at(0), 3 * ldbb, st, )?; } - if qb > 0 { + if let Some(cap) = capture.as_deref() { + // conv window tail: the last three pre-conv columns, which + // include restored prefix columns for a window under three rows. cuda::copy_dd( - p(s.gdn_in + (qb - 3) * ldbb), - prefix.unwrap().conv_tail[la_i].at(0), + cap.conv_tail[la_i].at(0), + p(s.gdn_in + (tend - 3) * ldbb), 3 * ldbb, st, )?; @@ -1865,7 +2154,7 @@ impl Model { "gated norm", )?; } - self.gemm(s, s.ln + qb * vb, &la.out, s.delta + qb * hb, rows)?; + self.gemm_with(gemm, s, s.ln + qb * vb, &la.out, s.delta + qb * hb, rows)?; self.apply_lora( i, Group::LinearOutput, @@ -1877,7 +2166,8 @@ impl Model { la_i += 1; } Mixer::Full(fa) => { - self.gemm_group( + self.gemm_group_with( + gemm, s, s.x + qb * hb, &fa.qkv, @@ -1953,7 +2243,7 @@ impl Model { "gated attention", )?; } - self.gemm(s, s.ao + qb * ob, &fa.o, s.delta + qb * hb, rows)?; + self.gemm_with(gemm, s, s.ao + qb * ob, &fa.o, s.delta + qb * hb, rows)?; self.apply_lora( i, Group::AttentionOutput, @@ -1982,7 +2272,8 @@ impl Model { "post-attention norm", )?; } - self.gemm_group( + self.gemm_group_with( + gemm, s, s.x + qb * hb, &layer.gate_up, @@ -2012,7 +2303,8 @@ impl Model { "silu mul", )?; } - self.gemm( + self.gemm_with( + gemm, s, s.act + qb * cfg.intermediate * BF16, &layer.down, @@ -2052,6 +2344,10 @@ impl Model { Ok(()) } } + #[cfg(test)] #[path = "../../../../../tests/qwen3_5/multimodal_graph.rs"] mod graph_tests; +#[cfg(test)] +#[path = "../../../../../tests/qwen3_5/model.rs"] +mod tests; diff --git a/src/runtime/README.md b/src/runtime/README.md index bd571b03..fcb2ec0f 100644 --- a/src/runtime/README.md +++ b/src/runtime/README.md @@ -34,8 +34,10 @@ existing failure behavior. This implements serial admission per executor within a worker process. Admitted concurrency is one. Queue-length limits, token budgets, compatibility grouping, shared processing orchestration and GPU batching remain planned. The Qwen -executor still accepts one prompt per forward; this scheduler does not pack -inputs or coordinate separately running workers. +executor runs single prompts, bounded packed prefill (Open-Jev-27B-v1.1's +candidates within a request) or prompts that share prefixes in one call, which no +worker uses yet; this scheduler does not pack inputs or coordinate separately +running workers. CPU tests cover FIFO serialization, queued cancellation, retention of admission and resources after dispatched cancellation, errors/panics, and independent diff --git a/tests/qwen3_5/model.rs b/tests/qwen3_5/model.rs new file mode 100644 index 00000000..6e66fd1c --- /dev/null +++ b/tests/qwen3_5/model.rs @@ -0,0 +1,268 @@ +use super::*; + +fn seq(n: usize, base: u32) -> Vec { + (0..n as u32).map(|i| base + i).collect() +} + +#[test] +fn runs_keep_every_prompt_and_end_prefixes_on_chunks() { + for p in [0usize, 1, 63, 64, 65, 130] { + let prefix = seq(p, 0); + let groups = [ + (seq(0, 1000), vec![seq(1, 2000), seq(70, 3000)]), + (seq(1, 4000), vec![seq(2, 5000)]), + (seq(64, 6000), vec![seq(5, 7000), seq(1, 8000)]), + ]; + let prompts = SharedPrompts { + prefix: &prefix, + groups: groups + .iter() + .map(|(prefix, branches)| PromptGroup { + prefix, + branches: branches.iter().map(Vec::as_slice).collect(), + }) + .collect(), + }; + let r = runs(&prompts); + let chunk = |n: usize| n - n % GDN_CHUNK; + assert_eq!(r.prefix.len(), chunk(p), "p = {p}"); + for ((group, branches), run) in groups.iter().zip(&r.groups) { + assert_eq!(r.prefix.len() + run.prefix.len(), chunk(p + group.len())); + assert_eq!(branches.len(), run.branches.len()); + for (branch, run_branch) in branches.iter().zip(&run.branches) { + assert!(!run_branch.is_empty()); + assert_eq!( + [&r.prefix[..], &run.prefix, run_branch].concat(), + [&prefix[..], group, branch].concat(), + "p = {p}" + ); + } + } + } +} + +/// A two-layer checkpoint with hidden size 128, the smallest whose stacked weights stay +/// contiguous in the arena, and an FP32 adapter on layer 0's down projection. +#[cfg(unix)] +fn write_tiny_checkpoint(dir: &Path) { + use safetensors::{ + Dtype, + tensor::{TensorView, serialize}, + }; + let config = serde_json::json!({ + "text_config": { + "hidden_size": 128, "intermediate_size": 8, "num_hidden_layers": 2, + "layer_types": ["linear_attention", "full_attention"], + "num_attention_heads": 1, "num_key_value_heads": 1, "head_dim": 256, + "linear_num_key_heads": 1, "linear_num_value_heads": 1, + "linear_key_head_dim": 128, "linear_value_head_dim": 128, + "linear_conv_kernel_dim": 4, "hidden_act": "silu", "rms_norm_eps": 1e-6, + "max_position_embeddings": 4096, + "rope_parameters": { + "rope_type": "default", "rope_theta": 1e7, "partial_rotary_factor": 0.25, + "mrope_section": [11, 11, 10], "mrope_interleaved": true + } + } + }); + std::fs::write(dir.join("config.json"), config.to_string()).unwrap(); + let write = |file: &str, dtype: Dtype, tensors: &[(String, Vec)]| { + let size = if dtype == Dtype::F32 { 4 } else { 2 }; + let data: Vec> = tensors + .iter() + .map(|(_, shape)| vec![0; shape.iter().product::() * size]) + .collect(); + let views: Vec<_> = tensors + .iter() + .zip(&data) + .map(|((name, shape), bytes)| { + ( + name.clone(), + TensorView::new(dtype, shape.clone(), bytes).unwrap(), + ) + }) + .collect(); + std::fs::write(dir.join(file), serialize(views, None).unwrap()).unwrap(); + }; + let layer = |index: usize, mixer: &[(&str, Vec)]| { + [ + ("input_layernorm.weight", vec![128]), + ("post_attention_layernorm.weight", vec![128]), + ("mlp.gate_proj.weight", vec![8, 128]), + ("mlp.up_proj.weight", vec![8, 128]), + ("mlp.down_proj.weight", vec![128, 8]), + ] + .into_iter() + .chain(mixer.iter().cloned()) + .map(|(name, shape)| (format!("layers.{index}.{name}"), shape)) + .collect::>() + }; + let weights: Vec<_> = [ + ("embed_tokens.weight".to_owned(), vec![16, 128]), + ("norm.weight".to_owned(), vec![128]), + ] + .into_iter() + .chain(layer( + 0, + &[ + ("linear_attn.in_proj_qkv.weight", vec![384, 128]), + ("linear_attn.in_proj_z.weight", vec![128, 128]), + ("linear_attn.in_proj_b.weight", vec![1, 128]), + ("linear_attn.in_proj_a.weight", vec![1, 128]), + ("linear_attn.conv1d.weight", vec![384, 1, 4]), + ("linear_attn.A_log", vec![1]), + ("linear_attn.dt_bias", vec![1]), + ("linear_attn.norm.weight", vec![128]), + ("linear_attn.out_proj.weight", vec![128, 128]), + ], + )) + .chain(layer( + 1, + &[ + ("self_attn.q_proj.weight", vec![512, 128]), + ("self_attn.k_proj.weight", vec![256, 128]), + ("self_attn.v_proj.weight", vec![256, 128]), + ("self_attn.o_proj.weight", vec![128, 256]), + ("self_attn.q_norm.weight", vec![256]), + ("self_attn.k_norm.weight", vec![256]), + ], + )) + .map(|(name, shape)| (format!("model.language_model.{name}"), shape)) + .collect(); + write("model.safetensors", Dtype::BF16, &weights); + let target = "base_model.model.model.language_model.layers.0.mlp.down_proj"; + write( + "adapter.safetensors", + Dtype::F32, + &[ + (format!("{target}.lora_A.weight"), vec![1, 8]), + (format!("{target}.lora_B.weight"), vec![128, 1]), + ], + ); +} + +#[cfg(unix)] +#[test] +fn cpu_fixed_and_shared_refuse_other_profiles_before_device_work() { + // The library loads once per process, so this runs in a child process. + if std::env::var_os("QWEN_TEST_PROFILE_CHILD").is_none() { + let name = format!( + "{}::cpu_fixed_and_shared_refuse_other_profiles_before_device_work", + module_path!().split_once("::").unwrap().1 + ); + let output = std::process::Command::new(std::env::current_exe().unwrap()) + .args(["--exact", &name, "--nocapture"]) + .env("QWEN_TEST_PROFILE_CHILD", "1") + .output() + .unwrap(); + assert!( + output.status.success(), + "{}\n{}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr) + ); + return; + } + let dir = std::env::temp_dir().join(format!("qwen-profile-{}", std::process::id())); + std::fs::create_dir_all(&dir).unwrap(); + write_tiny_checkpoint(&dir); + let library = dir.join(if cfg!(target_os = "macos") { + "fake.dylib" + } else { + "fake.so" + }); + let cuda_source = include_str!("../../src/models/qwen3_5/native/src/cuda.rs"); + let reference_source = include_str!("../../src/models/qwen3_5/native/src/cuda/reference.rs"); + let abi = cuda_source + .lines() + .find_map(|line| line.strip_prefix("const ABI_VERSION: u32 =")) + .unwrap() + .trim() + .trim_end_matches(';'); + // Every other required and optional symbol counts the call and fails. + let mut code = include_str!("profile_fake.c").to_owned(); + let required = cuda_source + .split("api! {") + .nth(1) + .unwrap() + .split("\n}\n") + .next() + .unwrap() + .lines() + .filter_map(|line| line.trim().strip_prefix("cs1_")); + let optional = [cuda_source, reference_source] + .into_iter() + .flat_map(|source| source.split("b\"cs1_").skip(1)); + let mut names: Vec<_> = required + .chain(optional) + .map(|rest| format!("cs1_{}", rest.split(['(', '\\']).next().unwrap())) + .filter(|name| !code.contains(&format!(" {name}("))) + .collect(); + names.sort(); + names.dedup(); + for name in names { + code.push_str(&format!("int {name}(void) {{ ++calls; return 77; }}\n")); + } + let source = dir.join("fake.c"); + std::fs::write(&source, code).unwrap(); + let output = std::process::Command::new("cc") + .args(["-shared", "-fPIC"]) + .arg(format!("-DPROFILE_FAKE_ABI={abi}")) + .arg(&source) + .arg("-o") + .arg(&library) + .output() + .unwrap(); + assert!( + output.status.success(), + "{}", + String::from_utf8_lossy(&output.stderr) + ); + let mut model = + Model::load_with_lora(&dir, &library, &dir.join("adapter.safetensors"), 2.).unwrap(); + // SAFETY: the test library defines this counter and stays loaded in the process. + let metrics = unsafe { libloading::Library::new(&library).unwrap() }; + let calls = unsafe { + *metrics + .get:: i32>(b"profile_fake_calls\0") + .unwrap() + }; + let prompts = SharedPrompts { + prefix: &[1, 2], + groups: vec![PromptGroup { + prefix: &[3], + branches: vec![&[4], &[5, 6]], + }], + }; + let refuse = |model: &mut Model, profile: &str| { + let before = unsafe { calls() }; + for error in [ + model.forward_fixed(&[1, 2, 3]).unwrap_err(), + model.forward_shared(&prompts).unwrap_err(), + ] { + assert!( + error.to_string().contains("only merged weights"), + "{profile}: {error:#}" + ); + } + assert_eq!(unsafe { calls() }, before, "{profile}: CUDA calls"); + assert!( + model.scratch.is_none() + && model.fixed_gemm.is_null() + && model.shared.iter().all(Option::is_none), + "{profile}: buffers" + ); + }; + refuse(&mut model, "LoRA"); + // Each of the other profiles alone is refused too; without any, the check passes. + let adapter = model.lora.take(); + model.ensure_fixed_profile().unwrap(); + model.fp32_gdn = Some(cuda::api().fp32_gdn().unwrap()); + refuse(&mut model, "FP32 GDN"); + model.fp32_gdn = None; + model.reference = Some(cuda::api().language_reference().unwrap()); + refuse(&mut model, "reference numerics"); + model.reference = None; + model.lora = adapter; + drop(model); + let _ = std::fs::remove_dir_all(&dir); +} diff --git a/tests/qwen3_5/prefix.rs b/tests/qwen3_5/prefix.rs new file mode 100644 index 00000000..1ca07061 --- /dev/null +++ b/tests/qwen3_5/prefix.rs @@ -0,0 +1,205 @@ +//! GPU checks of `Model::forward_shared` against `forward_fixed` on each full prompt. They +//! need a GPU, CUA_S1_CUDA_LIB pointing at libqwen3_5_cuda.so and QWEN3_5_MODEL pointing at +//! a merged Qwen3.5/3.8 checkpoint, such as a Cua-S1 or Open-Jev export: +//! +//! CUA_S1_CUDA_LIB=$PWD/target/release/libqwen3_5_cuda.so QWEN3_5_MODEL= \ +//! cargo test --release -p omni-qwen3-5-native --test prefix -- --ignored --test-threads=1 +//! +//! With CUA_S1_GRAPH=1, repeated plain `forward` calls of one length replay a CUDA Graph, +//! also between shared runs. + +use std::path::PathBuf; + +use omni_qwen3_5_native::model::{Model, PromptGroup, SharedPrompts}; + +fn model() -> Model { + let var = |name: &str| { + std::env::var_os(name) + .map(PathBuf::from) + .unwrap_or_else(|| panic!("{name} must be set")) + }; + Model::load(&var("QWEN3_5_MODEL"), &var("CUA_S1_CUDA_LIB")).unwrap() +} + +/// Token ids in [1000, 100000) from a fixed seed, inside every Qwen3.5 vocabulary. +fn tokens(n: usize, seed: u64) -> Vec { + let mut x = seed.wrapping_mul(0x9e37_79b9_7f4a_7c15) | 1; + (0..n) + .map(|_| { + x ^= x << 13; + x ^= x >> 7; + x ^= x << 17; + 1000 + (x >> 33) as u32 % 99_000 + }) + .collect() +} + +struct Case { + prefix: Vec, + groups: Vec<(Vec, Vec>)>, +} + +impl Case { + /// A request prefix of `p` tokens and groups of (prefix length, branch lengths). + fn new(p: usize, groups: &[(usize, &[usize])], seed: u64) -> Self { + let mut seed = seed * 1000; + let mut next = |n| { + seed += 1; + tokens(n, seed) + }; + Self { + prefix: next(p), + groups: groups + .iter() + .map(|&(q, branches)| (next(q), branches.iter().map(|&b| next(b)).collect())) + .collect(), + } + } + + fn prompts(&self) -> SharedPrompts<'_> { + SharedPrompts { + prefix: &self.prefix, + groups: self + .groups + .iter() + .map(|(prefix, branches)| PromptGroup { + prefix, + branches: branches.iter().map(Vec::as_slice).collect(), + }) + .collect(), + } + } + + fn full(&self, g: usize, b: usize) -> Vec { + let (prefix, branches) = &self.groups[g]; + [&self.prefix[..], prefix, &branches[b]].concat() + } +} + +fn rel_l2(a: &[f32], b: &[f32]) -> f64 { + let diff: f64 = a + .iter() + .zip(b) + .map(|(x, y)| (*x as f64 - *y as f64).powi(2)) + .sum(); + let norm: f64 = b.iter().map(|y| (*y as f64).powi(2)).sum(); + (diff / norm).sqrt() +} + +fn bits(rows: &[Vec>]) -> Vec { + rows.iter() + .flatten() + .flatten() + .map(|v| v.to_bits()) + .collect() +} + +#[test] +#[ignore = "needs a GPU, CUA_S1_CUDA_LIB and QWEN3_5_MODEL"] +fn shared_prefixes_match_full_prompts() { + let mut model = model(); + // Request prefixes around and inside 64-token chunks; groups without a prefix of their + // own, with one token, 70 and 128; branches of 1 to 200 tokens. + let groups: &[(usize, &[usize])] = + &[(0, &[1, 64]), (1, &[2, 65]), (70, &[1, 200]), (128, &[3])]; + for (i, p) in [0usize, 1, 63, 64, 65, 300, 1000].into_iter().enumerate() { + let case = Case::new(p, groups, i as u64 + 1); + let shared = model.forward_shared(&case.prompts()).unwrap(); + let mut to_forward = 0f64; + for (g, rows) in shared.iter().enumerate() { + for (b, row) in rows.iter().enumerate() { + let full = case.full(g, b); + let want = model.forward_fixed(&full).unwrap(); + assert!(row.iter().all(|v| v.is_finite())); + assert!( + row.iter() + .map(|v| v.to_bits()) + .eq(want.iter().map(|v| v.to_bits())), + "prefix {p}, group {g}, branch {b}" + ); + // the GEMM algorithms differ from `forward`'s, not the math + let plain = model.forward(&full).unwrap(); + to_forward = to_forward.max(rel_l2(row, &plain)); + } + } + eprintln!( + "prefix {p}: bit-identical to forward_fixed; largest relative L2 difference from forward {to_forward:.2e}" + ); + // A sanity bound for the change of GEMM algorithms, not a tolerance for reuse. + assert!(to_forward <= 3e-2, "prefix {p}: {to_forward}"); + } +} + +#[test] +#[ignore = "needs a GPU, CUA_S1_CUDA_LIB and QWEN3_5_MODEL"] +fn shared_prefix_state_stays_inside_one_call() { + let mut model = model(); + let a = Case::new(200, &[(10, &[3, 40, 1]), (0, &[64]), (65, &[2, 7])], 11); + // longer than any prompt before it, so the buffers grow in between + let b = Case::new(1100, &[(30, &[9, 100])], 12); + let plain = model.forward(&a.full(0, 1)).unwrap(); + // with CUA_S1_GRAPH=1, the first call captured a graph and this one replays it + assert_eq!(model.forward(&a.full(0, 1)).unwrap(), plain); + let first = model.forward_shared(&a.prompts()).unwrap(); + // the same request again, and after a longer one + let again = model.forward_shared(&a.prompts()).unwrap(); + assert_eq!(bits(&first), bits(&again)); + // the plain forward between shared runs, replayed with CUA_S1_GRAPH=1 + assert_eq!(model.forward(&a.full(0, 1)).unwrap(), plain); + model.forward_shared(&b.prompts()).unwrap(); + let after = model.forward_shared(&a.prompts()).unwrap(); + assert_eq!(bits(&first), bits(&after)); + // groups and branches in reverse order: every branch keeps its result + let reversed = Case { + prefix: a.prefix.clone(), + groups: a + .groups + .iter() + .rev() + .map(|(prefix, branches)| (prefix.clone(), branches.iter().rev().cloned().collect())) + .collect(), + }; + let back = model.forward_shared(&reversed.prompts()).unwrap(); + let mut back: Vec>> = back.into_iter().rev().collect(); + back.iter_mut().for_each(|branches| branches.reverse()); + assert_eq!(bits(&first), bits(&back)); + // the plain forward is unaffected (eager here: the longer request dropped the graphs) + assert_eq!(model.forward(&a.full(0, 1)).unwrap(), plain); + // Invalid requests fail without breaking the model. All of these are refused before + // any GPU work is queued. + let empty: &[u32] = &[]; + let outside: &[u32] = &[5, u32::MAX]; + let group = |prefix, branches| PromptGroup { prefix, branches }; + let bad = [ + SharedPrompts { + prefix: &a.prefix, + groups: vec![], + }, + SharedPrompts { + prefix: &a.prefix, + groups: vec![group(empty, vec![])], + }, + SharedPrompts { + prefix: &a.prefix, + groups: vec![group(empty, vec![empty])], + }, + SharedPrompts { + prefix: &a.prefix, + groups: vec![group(empty, vec![outside])], + }, + SharedPrompts { + prefix: &a.prefix, + groups: vec![group(outside, vec![&a.prefix])], + }, + SharedPrompts { + prefix: outside, + groups: vec![group(empty, vec![&a.prefix])], + }, + ]; + for prompts in &bad { + assert!(model.forward_shared(prompts).is_err()); + } + assert!(model.forward_fixed(outside).is_err()); + let last = model.forward_shared(&a.prompts()).unwrap(); + assert_eq!(bits(&first), bits(&last)); +} diff --git a/tests/qwen3_5/profile_fake.c b/tests/qwen3_5/profile_fake.c new file mode 100644 index 00000000..add89a17 --- /dev/null +++ b/tests/qwen3_5/profile_fake.c @@ -0,0 +1,20 @@ +// CPU stand-in for the CUDA library that counts every call, for the profile test. +#include +#include +#include +static int calls; +int profile_fake_calls(void) { return calls; } +uint32_t cs1_abi_version(void) { ++calls; return PROFILE_FAKE_ABI; } +const char* cs1_error_string(int code) { (void)code; ++calls; return "unexpected fake CUDA call"; } +int cs1_get_device(int* device) { ++calls; if (!device) return 1; *device=0; return 0; } +int cs1_set_device(int device) { (void)device; ++calls; return 0; } +int cs1_malloc(void** p, size_t n) { ++calls; *p=malloc(n ? n : 1); return *p ? 0 : 1; } +int cs1_free(void* p) { ++calls; free(p); return 0; } +int cs1_stream_create(void** stream) { ++calls; *stream=(void*)1; return 0; } +int cs1_stream_sync(void* stream) { (void)stream; ++calls; return 0; } +int cs1_stream_destroy(void* stream) { (void)stream; ++calls; return 0; } +int cs1_upload(void* dst, const void* src, size_t n, void* stream) { (void)stream; ++calls; memcpy(dst,src,n); return 0; } +void* cs1_gemm_create(size_t n) { (void)n; ++calls; return &calls; } +void cs1_gemm_destroy(void* gemm) { (void)gemm; ++calls; } +void* cs1_gdn_fp32_create(size_t n) { (void)n; ++calls; return &calls; } +void cs1_gdn_fp32_destroy(void* handle) { (void)handle; ++calls; }