diff --git a/crates/gpu-gates/src/bin/gate_qwen3moe_experts.rs b/crates/gpu-gates/src/bin/gate_qwen3moe_experts.rs index 9c2a0809d..689ce2ff5 100644 --- a/crates/gpu-gates/src/bin/gate_qwen3moe_experts.rs +++ b/crates/gpu-gates/src/bin/gate_qwen3moe_experts.rs @@ -53,6 +53,9 @@ //! //! And each kernel as a captured graph: one node, the replay equal to the //! eager launch. +//! +//! Q5_0 embedding rows (`embed_rows_q5_0`) are bit-identical to `dequant_row` +//! on the same bytes. #[cfg(feature = "gpu")] fn main() -> std::process::ExitCode { @@ -136,6 +139,7 @@ mod gate { let mut ok = true; ok &= embed(&gpu, &gguf, &sets)?; ok &= embed_kquant(&gpu)?; + ok &= embed_q5_0(&gpu)?; ok &= gate_up_body(&gpu, &gguf, &ex)?; ok &= gate_up_tap(&gpu, &gguf, &ex, &sets)?; ok &= combine(&gpu, &ex, &sets, stream)?; @@ -146,6 +150,63 @@ mod gate { Ok(()) } + /// Q5_0 rows through `embed_rows_q5_0` against `dequant_row` on the same + /// bytes, bit for bit, positions included, one captured node, and an id + /// past the table (NaN row, `FaultSite::TokenId`, the other rows intact). + fn embed_q5_0(gpu: &Gpu) -> Result { + const ROWS: usize = 4; + // 34 blocks: past the first 32-block window of the packed code plane. + const K: usize = 1088; + // Reuse synthetic_rows with the Q5_0 scale at byte 0; host emulation + // confirms this seed rejects both addressing mutants. + let bytes = synthetic_rows((GgmlType::Q5_0, &[0]), (ROWS, K / 32), 0x2545_f491)?; + let packed = bloomery_gpu::q5::pack_q5_0(&bytes, K, ROWS)?; + let cols = packed.len() / ROWS; + let table = DeviceTensor::upload(gpu.stream(), &packed, ROWS, cols)?; + let ids = [0u32, (ROWS - 1) as u32, 1]; + let launch = |ids: &[u32], graph: bool| { + gather( + gpu, + (&table, Some(GgmlType::Q5_0)), + ids, + POS0, + FIRST, + K, + graph, + ) + }; + let (got, _) = launch(&ids, false)?; + let mut want = vec![0.0f32; ids.len() * K]; + let rb = (K / 32) * 22; + for (t, &id) in ids.iter().enumerate() { + let row = &bytes[id as usize * rb..(id as usize + 1) * rb]; + dequant_row(GgmlType::Q5_0, row, &mut want[t * K..(t + 1) * K])?; + } + let host_ok = bits_equal(&got.y, &want); + let pos_ok = positions_ok(&got, ids.len(), POS0, FIRST); + let (replayed, nodes) = launch(&ids, true)?; + let graph_ok = bits_equal(&replayed.y, &got.y) + && replayed.pos == got.pos + && replayed.n_keys == got.n_keys + && nodes == 1; + let clean = gpu.take_fault()?; + let bad = [ids[0], ROWS as u32, ids[2]]; + let (ybad, _) = launch(&bad, false)?; + let fault = gpu.take_fault()?; + let want_fault = Fault::at(LAYER_NONE, FaultSite::TokenId); + let others = + bits_equal(&ybad.y[..K], &got.y[..K]) && bits_equal(&ybad.y[2 * K..], &got.y[2 * K..]); + let nan = ybad.y[K..2 * K].iter().filter(|v| v.is_nan()).count(); + let bad_pos = positions_ok(&ybad, bad.len(), POS0, FIRST); + let oor_ok = clean.is_none() && fault == Some(want_fault) && others && nan == K && bad_pos; + println!( + "embed Q5_0 K={K} rows={ROWS}: vs dequant_row bit_identical={host_ok} positions={pos_ok} \ + graph_nodes={nodes} graph_same={graph_ok} oor_nan={nan}/{K} {}", + verdict(host_ok && pos_ok && graph_ok && oor_ok) + ); + Ok(host_ok && pos_ok && graph_ok && oor_ok) + } + // ------------------------------------------------------------ 1. embed /// One gather's outputs, read back: the rows, each row's position and @@ -159,7 +220,8 @@ mod gate { /// `ids` gathered from `table` (`k` values a row), the rows' positions /// from `pos0` and offset `first`: eagerly, or captured as a graph and /// replayed once, which also returns its node count. `ty` `None` takes - /// the Q4_K entry directly, `Some(t)` the K-quant entry of type `t`. + /// the Q4_K entry directly, `Some(Q5_0)` the Q5_0 entry, `Some(t)` the + /// K-quant entry of type `t`. #[allow(clippy::too_many_arguments, reason = "one launch's inputs, each named")] fn gather( gpu: &Gpu, @@ -188,6 +250,7 @@ mod gate { }; match ty { None => gpu.elem().enqueue_embed_rows_q4k(s, args), + Some(GgmlType::Q5_0) => gpu.elem().enqueue_embed_rows_q5_0(s, args, k), Some(ty) => gpu.elem().enqueue_embed_rows_kquant(s, ty, args), } }; diff --git a/crates/gpu-gates/src/bin/gate_qwen3moe_flash.rs b/crates/gpu-gates/src/bin/gate_qwen3moe_flash.rs index cfad3b7b5..b3d5a9ce3 100644 --- a/crates/gpu-gates/src/bin/gate_qwen3moe_flash.rs +++ b/crates/gpu-gates/src/bin/gate_qwen3moe_flash.rs @@ -145,7 +145,7 @@ fn main() -> std::process::ExitCode { mod gate { use bloomery_gpu::fault::{Fault, FaultSink, FaultSite, LAYER_NONE}; use bloomery_gpu::flash_gqa::{ - FlashGqaKernels, GROUP, GqaArgs, GqaQ8Args, HEAD, KEY_TILE, SEG_KEYS, SEGMENTS, + FlashGqaKernels, GROUP, GROUP_5, GqaArgs, GqaQ8Args, HEAD, KEY_TILE, SEG_KEYS, SEGMENTS, partials_ms_len, partials_v_len, seg_span, }; use bloomery_gpu::flash_gqa_prefill::{ @@ -241,7 +241,7 @@ mod gate { let mut pv = DeviceBuffer::::zeroed(stream, partials_v_len(m, n_head))?; let mut pms = DeviceBuffer::::zeroed(stream, partials_ms_len(m, n_head))?; let mut y = DeviceBuffer::::zeroed(stream, m * n_head * HEAD)?; - k.enqueue_pass( + k.enqueue_pass::( stream, GqaArgs { q, @@ -500,7 +500,7 @@ mod gate { }, ) } else { - k.enqueue_pass( + k.enqueue_pass::( s, GqaArgs { q: &inp.q, @@ -620,7 +620,7 @@ mod gate { let mut pv = DeviceBuffer::::zeroed(stream, partials_v_len(1, n_head))?; let mut pms = DeviceBuffer::::zeroed(stream, partials_ms_len(1, n_head))?; let mut y = DeviceBuffer::::zeroed(stream, n_head * HEAD)?; - k.enqueue_pass( + k.enqueue_pass::( stream, GqaArgs { q: &qd, @@ -772,7 +772,7 @@ mod gate { let mut pv = DeviceBuffer::::zeroed(stream, partials_v_len(ROWS, n_head))?; let mut pms = DeviceBuffer::::zeroed(stream, partials_ms_len(ROWS, n_head))?; let mut y = DeviceBuffer::::zeroed(stream, ROWS * n_head * HEAD)?; - k.enqueue_pass( + k.enqueue_pass::( stream, GqaArgs { q: &rows_q, @@ -1901,6 +1901,80 @@ mod gate { plain: f64, } + /// Head 128, [`GROUP_5`] query heads a key head, the scalar pass, + /// against [`exact`]'s head-128 band. Two key heads: a KV index of + /// `h / 8` still reads the first plane for the second group. + fn group5(k: &FlashGqaKernels, gpu: &Gpu, scale: f32) -> Result { + const N_KV: usize = 2; + const KEYS: usize = 96; + const CTX: usize = 128; + let n_head = N_KV * GROUP_5; + let stream = gpu.stream(); + let q: Vec = activations(HEAD, n_head, 5); + let kb: Vec = activations(HEAD, N_KV * CTX, 6) + .into_iter() + .map(f32_to_f16_bits) + .collect(); + let vb: Vec = activations(HEAD, N_KV * CTX, 7) + .into_iter() + .map(f32_to_f16_bits) + .collect(); + let kf: Vec = kb.iter().copied().map(half_to_f32).collect(); + let vf: Vec = vb.iter().copied().map(half_to_f32).collect(); + let q_d = DeviceBuffer::from_host(stream, &q)?; + let kc = DeviceBuffer::from_host(stream, &kb)?; + let vc = DeviceBuffer::from_host(stream, &vb)?; + let n_keys = DeviceBuffer::from_host(stream, &[KEYS as u32])?; + let mut part_v = DeviceBuffer::::zeroed(stream, partials_v_len(1, n_head))?; + let mut part_ms = DeviceBuffer::::zeroed(stream, partials_ms_len(1, n_head))?; + let mut y = DeviceBuffer::::zeroed(stream, n_head * HEAD)?; + k.enqueue_pass::( + stream, + GqaArgs { + q: &q_d, + kc: &kc, + vc: &vc, + n_keys: &n_keys, + scale, + n_kv: N_KV, + ctx: CTX, + m: 1, + part_v: &mut part_v, + part_ms: &mut part_ms, + fault: gpu.unlabelled_sink(), + y: &mut y, + }, + false, + )?; + stream.synchronize()?; + let got = y.to_host_vec(stream)?; + let mut worst = 0.0f64; + let mut ok = true; + for h in 0..n_head { + let plane = (h / GROUP_5) * CTX * HEAD; + let ex = exact( + &q[h * HEAD..(h + 1) * HEAD], + &kf[plane..plane + KEYS * HEAD], + &vf[plane..plane + KEYS * HEAD], + KEYS, + scale, + ); + for d in 0..HEAD { + let err = (f64::from(got[h * HEAD + d]) - ex.o[d]).abs(); + let bound = ex.bound_ours[d]; + if bound > 0.0 { + worst = worst.max(err / bound); + } + ok &= err <= bound; + } + } + println!( + "flash group5 head={HEAD} group={GROUP_5} kv={N_KV} keys={KEYS}: measured/bound {worst:.3e} {}", + verdict(ok) + ); + Ok(ok) + } + pub fn run() -> Result<(), GateError> { // A lever set to a value it does not take, or a retired name that is // set, is refused by name before anything loads. @@ -2142,6 +2216,7 @@ mod gate { ); ok &= seeded_ok; ok &= q8_window_refusal(&kp, &gpu, scale)?; + ok &= group5(&k, &gpu, scale)?; println!("gate_qwen3moe_flash: {}", verdict(ok)); if !ok { return Err(checks_failed()); diff --git a/crates/gpu/src/arch/qwen3moe/body35.rs b/crates/gpu/src/arch/qwen3moe/body35.rs index 50a2c2425..b80ffebde 100644 --- a/crates/gpu/src/arch/qwen3moe/body35.rs +++ b/crates/gpu/src/arch/qwen3moe/body35.rs @@ -229,7 +229,7 @@ fn checkpoint_lens(kinds: &[Kind35]) -> Vec { kinds .iter() .filter_map(|k| match k { - Kind35::Delta(shape) => Some(RecStore::lens(*shape, PLAIN_LANES)), + Kind35::Delta(shape) => Some(shape.lens(PLAIN_LANES)), Kind35::Gqa(_) => None, }) .flatten() @@ -1618,9 +1618,11 @@ impl Body35 { stores.push(match *kind { Kind35::Gqa(_) => LayerStore::Kv(KvPlanes::new(stream, &d, o.kv)?), Kind35::Delta(shape) if lanes > PLAIN_LANES => { - LayerStore::Rec(RecStore::stamped(stream, shape, lanes)?) + LayerStore::Rec(RecStore::stamped(stream, shape.lens(lanes), lanes)?) + } + Kind35::Delta(shape) => { + LayerStore::Rec(RecStore::new(stream, shape.lens(lanes), lanes)?) } - Kind35::Delta(shape) => LayerStore::Rec(RecStore::new(stream, shape, lanes)?), }); } let forms = Forms::of(&plans, &d); @@ -3153,7 +3155,7 @@ impl Slots for Body35 { Ok(LayerStore::Kv(KvPlanes::new(stream, &d, kv)?)) } (MixerPlan::Delta(dp), LayerStore::Rec(_)) => Ok(LayerStore::Rec( - RecStore::new(stream, dp.shape, PLAIN_LANES)?, + RecStore::new(stream, dp.shape.lens(PLAIN_LANES), PLAIN_LANES)?, )), _ => Err(GpuError::state( WHAT_N, @@ -3335,7 +3337,7 @@ mod tests { /// layer's state (2,097,152 B) and conv ring (360,448 B): 73,728,000 B, /// spaced at the 4,096-token default ubatch, so a slot of 2,048 positions /// pins one point and one of 10,240 three. The lengths are the stores' - /// own ([`RecStore::lens`]), in [`Copied35`]'s order. + /// own ([`LinearShape::lens`]), in [`Copied35`]'s order. #[test] fn a_checkpoint_copies_each_delta_layers_state_and_ring() { let delta = Kind35::Delta(LinearShape { diff --git a/crates/gpu/src/arch/qwen3moe/dispatch.rs b/crates/gpu/src/arch/qwen3moe/dispatch.rs index ca187d45e..e9902b090 100644 --- a/crates/gpu/src/arch/qwen3moe/dispatch.rs +++ b/crates/gpu/src/arch/qwen3moe/dispatch.rs @@ -473,7 +473,7 @@ pub(super) fn mixer( /// and each row's position and live key count into `s.pos` and `s.n_keys` — /// the rows every later launch of the unit reads. The table's resident type /// picks the lookup: Q3_K, Q4_K, Q5_K or Q6_K rows, or Q8_0 planes (the load -/// admits only these). +/// admits these). Q5_0 packed rows are served here for a body that admits them. pub(super) fn embed_rows( gpu: &Gpu, w: &Weights, @@ -486,7 +486,9 @@ pub(super) fn embed_rows( /// [`embed_rows`] into `(y, pos, n_keys)`, a unit's embedding rows and /// their positions and live key counts: the arena's own, or a window of -/// them where a pass of several slots embeds each slot's rows apart. +/// them where a pass of several slots embeds each slot's rows apart. Q5_0 +/// packed rows take [`crate::elem::ElemKernels::enqueue_embed_rows_q5_0`] +/// for a body that admits them. pub(super) fn embed_into( gpu: &Gpu, w: &Weights, @@ -516,9 +518,30 @@ pub(super) fn embed_into( }, ); } + if let Some(DevWeight::Q5_0 { w, k }) = w.get(&name) { + return gpu.elem().enqueue_embed_rows_q5_0( + gpu.stream(), + EmbedRowsArgs { + w, + ids: io.ids, + pos0: io.pos0, + first: io.first, + y, + pos, + n_keys, + }, + *k, + ); + } let ty = match w.get(&name) { Some(DevWeight::KQuant { ty, .. }) => *ty, - _ => return Err(GpuError::tensor(WHAT, &name, "K-quant rows or Q8_0 planes")), + _ => { + return Err(GpuError::tensor( + WHAT, + &name, + "K-quant rows or Q8_0 planes; Q5_0 packed rows when the body admits them", + )); + } }; gpu.elem().enqueue_embed_rows_kquant( gpu.stream(), diff --git a/crates/gpu/src/arch/qwen3moe/scratch.rs b/crates/gpu/src/arch/qwen3moe/scratch.rs index c8426c62c..bcab80ed1 100644 --- a/crates/gpu/src/arch/qwen3moe/scratch.rs +++ b/crates/gpu/src/arch/qwen3moe/scratch.rs @@ -23,12 +23,14 @@ use super::wide::{GEMV_COLS, Wide}; use crate::GpuError; use crate::fault::FaultSink; use crate::flash_gqa::{ - FlashGqaKernels, GqaArgs, GqaQ8Args, GqaRowsArgs, HEAD, HEAD_256, RowTable, partials_ms_len, - partials_v_len, partials_v_len_256, + FlashGqaKernels, GROUP, GqaArgs, GqaQ8Args, GqaRowsArgs, HEAD, HEAD_256, RowTable, + partials_ms_len, partials_v_len, partials_v_len_256, }; use crate::flash_gqa_prefill::{FlashGqaPrefill, GqaPrefillArgs, GqaPrefillQ8Args}; use crate::gemm::GEMM_MAX_SLOTS; use crate::linear::{self, LinearShape}; + +pub(super) use crate::recur::{RecStore, fresh_stamps, stamps_at}; use crate::rope_neox::{ NeoxArgs, NeoxQ8Args, NeoxRowsArgs, PartialNeoxArgs, PartialNeoxQ8Args, RopeNeoxKernels, q8_plane_lens, @@ -603,7 +605,7 @@ impl KvPlanes { y, } = s; match self { - KvPlanes::F16 { k: kc, v: vc } => flash.enqueue_pass( + KvPlanes::F16 { k: kc, v: vc } => flash.enqueue_pass::( stream, GqaArgs { q, @@ -1407,137 +1409,6 @@ impl GdnArena { } } -/// One delta layer's recurrent store: the state `[lanes][n_v][128][128]` -/// (value-major, `linear`'s layout), read and written in place through the -/// lane word, the conv ring `[RING_ROWS][C]`, indexed by position, and — on -/// a store that keeps a lane per verify row — the lanes' stamps (`linear::delta`'s -/// module doc): the position each lane's state stands at. -pub(super) struct RecStore { - pub(super) state: DeviceBuffer, - pub(super) ring: DeviceBuffer, - pub(super) lanes: usize, - /// `lanes` u32 stamps ([`fresh_stamps`]); `None` on a one-lane store, - /// whose launches are `gdn_delta`'s. - pub(super) stamp: Option>, -} - -impl RecStore { - /// A zeroed store of `lanes` lanes for `shape`, with no stamps. Load-time - /// only. - pub(super) fn new( - stream: &CudaStream, - shape: LinearShape, - lanes: usize, - ) -> Result { - let [state, ring] = RecStore::lens(shape, lanes); - Ok(RecStore { - state: DeviceBuffer::zeroed(stream, state)?, - ring: DeviceBuffer::zeroed(stream, ring)?, - lanes, - stamp: None, - }) - } - - /// A zeroed store of `lanes` lanes for `shape` with its stamps fresh - /// ([`fresh_stamps`]): the store a load that verifies drafted rows keeps. - /// Load-time only. - pub(super) fn stamped( - stream: &CudaStream, - shape: LinearShape, - lanes: usize, - ) -> Result { - let mut r = RecStore::new(stream, shape, lanes)?; - r.stamp = Some(DeviceBuffer::from_host(stream, &fresh_stamps(lanes))?); - Ok(r) - } - - /// The stamps back to a fresh store's ([`fresh_stamps`]); a store with - /// none is refused by name. Never inside a capture. - pub(super) fn restamp(&mut self, stream: &CudaStream) -> Result<(), GpuError> { - let fresh = fresh_stamps(self.lanes); - self.stamp_with(stream, &fresh) - } - - /// The stamps with lane `lane` at position `at` and every other lane - /// never written ([`stamps_at`]): the store after a cut restored a - /// checkpoint into the committed lane. A store with no stamps is - /// refused by name. Never inside a capture. - pub(super) fn restamp_at( - &mut self, - stream: &CudaStream, - lane: u32, - at: u32, - ) -> Result<(), GpuError> { - let st = stamps_at(self.lanes, lane, at)?; - self.stamp_with(stream, &st) - } - - fn stamp_with(&mut self, stream: &CudaStream, st: &[u32]) -> Result<(), GpuError> { - match self.stamp.as_mut() { - Some(stamp) => Ok(stamp.copy_from_host(stream, st)?), - None => Err(GpuError::state( - "qwen3moe::RecStore::restamp", - "a store of stamped lanes", - )), - } - } - - /// The f32s of the state and of the ring of a store of `lanes` lanes for - /// `shape`, the lengths [`RecStore::new`] allocates and a checkpoint of - /// the whole store copies. - pub(super) fn lens(shape: LinearShape, lanes: usize) -> [usize; 2] { - [lanes * shape.state_len(), shape.ring_len()] - } - - /// Every lane of the state and the whole ring back to zero, in stream - /// order through `zeros` (host zeros at least as long as the longer of - /// the two). Never inside a capture. - pub(super) fn clear(&mut self, stream: &CudaStream, zeros: &[f32]) -> Result<(), GpuError> { - for buf in [&mut self.state, &mut self.ring] { - let n = buf.len(); - let src = zeros.get(..n).ok_or_else(|| { - GpuError::shape( - "qwen3moe::RecStore::clear", - format!("{} host zeros for a {n}-value buffer", zeros.len()), - ) - })?; - buf.copy_from_host(stream, src)?; - } - Ok(()) - } - - pub(super) fn bytes(&self) -> usize { - self.state.num_bytes() - + self.ring.num_bytes() - + self.stamp.as_ref().map_or(0, DeviceBuffer::num_bytes) - } -} - -/// A fresh store's stamps over `lanes` lanes: lane 0 holds the zero state of -/// position 0, every other lane was never written ([`linear::delta::NEVER`], a position no -/// call starts at, so reading such a lane raises `FaultSite::DeltaStamp`). -pub(super) fn fresh_stamps(lanes: usize) -> Vec { - let mut st = vec![linear::delta::NEVER; lanes]; - if let Some(first) = st.first_mut() { - *first = 0; - } - st -} - -/// The stamps of a store over `lanes` lanes whose committed lane `lane` -/// stands at position `at` and whose every other lane was never written. A -/// lane at or past `lanes` is refused by name. -pub(super) fn stamps_at(lanes: usize, lane: u32, at: u32) -> Result, GpuError> { - let mut st = vec![linear::delta::NEVER; lanes]; - *st.get_mut(lane as usize).ok_or_else(|| { - GpuError::shape( - "qwen3moe::stamps_at", - format!("lane {lane} of {lanes} lanes"), - ) - })? = at; - Ok(st) -} - /// Non-owning views of a [`LaneWord`]'s device word ([`LaneWord::with_views`]). pub(super) type LaneViews = Vec>>; @@ -1652,28 +1523,28 @@ pub(super) enum LayerStore { Rec(RecStore), } +/// A layer's store as one layer's enqueue borrows it. +pub(super) enum StoreMut<'a> { + Kv(&'a mut KvPlanes), + Rec(&'a mut RecStore), +} + impl LayerStore { pub(super) fn as_mut(&mut self) -> StoreMut<'_> { match self { - LayerStore::Kv(k) => StoreMut::Kv(k), - LayerStore::Rec(r) => StoreMut::Rec(r), + Self::Kv(k) => StoreMut::Kv(k), + Self::Rec(r) => StoreMut::Rec(r), } } pub(super) fn bytes(&self) -> usize { match self { - LayerStore::Kv(k) => k.bytes(), - LayerStore::Rec(r) => r.bytes(), + Self::Kv(k) => k.bytes(), + Self::Rec(r) => r.bytes(), } } } -/// A layer's store as one layer's enqueue borrows it. -pub(super) enum StoreMut<'a> { - Kv(&'a mut KvPlanes), - Rec(&'a mut RecStore), -} - /// The committed lane of a recurrent store of one lane: the lane word every /// record carries while a load holds one lane. pub(super) const LANE: u32 = 0; diff --git a/crates/gpu/src/arch/qwen3moe/scratch38.rs b/crates/gpu/src/arch/qwen3moe/scratch38.rs index fb6170935..b721a6abe 100644 --- a/crates/gpu/src/arch/qwen3moe/scratch38.rs +++ b/crates/gpu/src/arch/qwen3moe/scratch38.rs @@ -80,7 +80,7 @@ impl Store38 { /// 0 and every other never written. pub(super) fn rec(stream: &CudaStream) -> Result { Ok(Store38::Rec { - rec: RecStore::new(stream, GDN, LANES)?, + rec: RecStore::new(stream, GDN.lens(LANES), LANES)?, stamp: DeviceBuffer::from_host(stream, &fresh_stamps(LANES))?, }) } diff --git a/crates/gpu/src/elem.rs b/crates/gpu/src/elem.rs index 889b88d8c..b5e7fcddb 100644 --- a/crates/gpu/src/elem.rs +++ b/crates/gpu/src/elem.rs @@ -12,7 +12,8 @@ //! v`, the layout of the reference dump's `q_rope`/`k_rope` views); //! `weighted_sum` on `[rows, n_exp, m]` down-projections with `[n_exp, m]` //! weights; the embedding table is Q3_K rows of 2048 values (880 bytes = 220 -//! u32 words each) or Q3_K, Q4_K, Q5_K or Q6_K rows of whole super-blocks. +//! u32 words each), Q3_K, Q4_K, Q5_K or Q6_K rows of whole super-blocks, or +//! Q5_0 rows in the `pack_q5_0` layout. //! Extents are launch arguments, never buffer lengths: scratch buffers may //! be larger than the shape in flight. @@ -696,6 +697,103 @@ mod elem_kernels { } } + /// Dequantize `ids.len()` rows of a Q5_0 table in the `pack_q5_0` layout + /// (`q_stride + k/32` u32 words a row; byte `j % 4` of word + /// `256·(b>>5) + (b&31) + 32·(j/4)` of block `b = v/32` is the low half's + /// code when `v%32 < 16`, and the high half lives `4` word-groups later) + /// into `y`, with each row's position and live key count as + /// `embed_rows_q4k`. The code byte is the 5-bit code minus 16, and the + /// block's scale is the f32 bits after the code plane, so a value is that + /// byte as `i8` times the scale. Q5_0's rows are 32-value blocks split + /// into a code plane and a scale plane, not 256-value super-blocks, so + /// this does not share [`embed_rows_sb`]. An id past `n_rows` raises + /// [`FaultSite::TokenId`] and writes a NaN row, as `embed_rows` does. + #[allow( + clippy::too_many_arguments, + reason = "kernel entry: the device ABI takes the arguments flat (rust-quality R8)" + )] + #[kernel] + #[launch_bounds(256)] + #[launch_contract( + domain = 1, + block = (256, 1, 1), + requires = ( + w.len() >= (q_stride + k_blocks) * n_rows, + y.len() >= k * ids.len(), + pos0.len() >= 1, + pos.len() >= ids.len(), + n_keys.len() >= ids.len() + ) + )] + pub fn embed_rows_q5_0( + w: &[u32], + ids: &[u32], + pos0: &[u32], + first: u32, + n_rows: u32, + k: u32, + q_stride: u32, + k_blocks: u32, + fault: FaultSink, + mut y: DisjointSlice, + mut pos: DisjointSlice, + mut n_keys: DisjointSlice, + ) { + let i = thread::index_1d().get(); + if i < ids.len() { + // SAFETY: pos0 holds a word by the launch contract. + let p = unsafe { *pos0.get_unchecked(0) } as usize + first as usize + i; + // SAFETY: i < ids.len() <= pos.len(), n_keys.len() by the launch + // contract; thread i owns entry i of both. + unsafe { + *pos.get_unchecked_mut(i) = position_word(p); + *n_keys.get_unchecked_mut(i) = position_word(p + 1); + } + } + let width = k as usize; + if i >= ids.len() * width { + return; + } + let t = i / width; + let v = i % width; + // SAFETY: t < ids.len() by the guard. + let id = unsafe { *ids.get_unchecked(t) }; + let out = if id < n_rows { + let b = v / 32; + let within = v % 32; + let j = within % 16; + // The high half of the block (values 16..32) is the second group + // of four words; the low half is the first. + let word_i = j / 4 + 4 * usize::from(within >= 16); + let byte = j % 4; + let row_words = q_stride as usize + k_blocks as usize; + let row = id as usize * row_words; + let wbase = 256 * (b >> 5) + (b & 31); + // SAFETY: id < n_rows, b < k/32 = k_blocks (v < k and the launcher + // passes k = 32·k_blocks), and word_i < 8, so + // wbase + 32·word_i < q_stride (a window is 256 words and the + // addressed word sits in the last 256 of that window). The scale + // word is the next k_blocks words. The row ends at + // (id+1)·(q_stride+k_blocks) <= w.len() by the launch contract. + let word = unsafe { *w.get_unchecked(row + wbase + 32 * word_i) }; + let code = ((word >> (8 * byte)) & 0xff) as u8 as i8; + // SAFETY: b < k_blocks, so the scale word row + q_stride + b sits + // inside the row, which ends at (id+1)·(q_stride+k_blocks) <= + // w.len() by the launch contract. + let scale = f32::from_bits(unsafe { *w.get_unchecked(row + q_stride as usize + b) }); + (code as f32) * scale + } else { + if v == 0 { + fault.raise(FaultSite::TokenId); + } + f32::NAN + }; + // SAFETY: i < ids.len() * k <= y.len() by the launch contract. + unsafe { + *y.get_unchecked_mut(i) = out; + } + } + /// Dequantize `ids.len()` rows of the Q5_K embedding table `w` (`44 · /// n_sb` u32 words per row) into `y` with their positions and live key /// counts ([`embed_rows_sb`] over [`Q5kRows`]), as `embed_rows_q4k`. @@ -1641,6 +1739,99 @@ impl ElemKernels { Ok(()) } + /// Enqueue the Q5_0 embedding lookup of one unit of rows: `args.ids` each + /// select one row of `args.w` in the `pack_q5_0` layout (`k` values, `k` + /// a positive multiple of 32, `q_stride + k/32` words a row with + /// `q_stride = 256·ceil((k/32)/32)`), dequantized into `args.y` + /// token-major, bit-identical to `gguf::quant::dequant_row` on the same + /// row bytes, with each row's position and live key count as + /// [`Self::enqueue_embed_rows_q4k`]. An id past the table raises + /// [`FaultSite::TokenId`]. A table whose word count is not that layout's + /// is refused by name. Asynchronous, allocation-free, capturable. + pub fn enqueue_embed_rows_q5_0( + &self, + stream: &CudaStream, + args: EmbedRowsArgs<'_>, + k: usize, + ) -> Result<(), GpuError> { + let what = "enqueue_embed_rows_q5_0"; + let EmbedRowsArgs { + w, + ids, + pos0, + first, + y, + pos, + n_keys, + } = args; + if k == 0 || !k.is_multiple_of(32) { + return Err(GpuError::shape( + what, + format!("k must be a positive multiple of 32, got {k}"), + )); + } + let k_blocks = k / 32; + let q_stride = 256 * k_blocks.div_ceil(32); + if w.rows() == 0 || w.cols() != q_stride + k_blocks { + return Err(GpuError::shape( + what, + format!( + "a Q5_0 table is {q_stride} + {k_blocks} words a row, got {}x{}", + w.rows(), + w.cols() + ), + )); + } + if ids.is_empty() { + return Err(GpuError::shape(what, "empty ids")); + } + if y.len() < k * ids.len() { + return Err(GpuError::shape( + what, + format!("y.len() {} < {k}*{}", y.len(), ids.len()), + )); + } + if pos0.is_empty() || pos.len() < ids.len() || n_keys.len() < ids.len() { + return Err(GpuError::shape( + what, + format!( + "pos0 holds {} words (want 1), pos {} and n_keys {} (want {} each)", + pos0.len(), + pos.len(), + n_keys.len(), + ids.len() + ), + )); + } + let grid = launch_u32(what, "grid", (ids.len() * k).div_ceil(256))?; + let first = launch_u32(what, "first", first)?; + let n_rows = launch_u32(what, "w.rows()", w.rows())?; + let k = launch_u32(what, "k", k)?; + let q_stride = launch_u32(what, "q_stride", q_stride)?; + let k_blocks = launch_u32(what, "k_blocks", k_blocks)?; + let prep = self + .module + .prepare_embed_rows_q5_0(LaunchConfig1D::new(grid, 256, 0))?; + let fault = crate::sink_over(&self.fault, LAYER_NONE); + self.module.embed_rows_q5_0( + stream, + &prep, + w.buf(), + ids, + pos0, + first, + n_rows, + k, + q_stride, + k_blocks, + fault, + y, + pos, + n_keys, + )?; + Ok(()) + } + /// Enqueue the embedding lookup of one unit of rows from a K-quant table /// of type `ty`, as [`Self::enqueue_embed_rows_q4k`] does for Q4_K: the /// rows dequantized token-major into `args.y`, bit-identical to diff --git a/crates/gpu/src/flash_gqa.rs b/crates/gpu/src/flash_gqa.rs index f6390aed3..ee6cdd92e 100644 --- a/crates/gpu/src/flash_gqa.rs +++ b/crates/gpu/src/flash_gqa.rs @@ -227,6 +227,14 @@ const THREADS_P12_U32: u32 = THREADS_P12 as u32; const _: () = assert!(PACK_12 == 12 && THREADS_P12_U32 == 384); const _: () = assert!(PACK_12 + PAD_12 == MMA_ROWS && PAD_12 <= PACK_12 && PACK_12 > MMA_ROWS / 2); +/// Query heads one block of [`FlashGqaKernels::enqueue_pass`] takes. +/// The tensor-core tile is [`MMA_ROWS`] rows and pads a pack only when the +/// pad is at most the pack and the pack is past half the tile, which five +/// is not, and the packed scalar wants a pack that divides the key tile's +/// staged words, which five does not. The pass is scalar. +pub const GROUP_5: usize = 5; +const _: () = assert!(GROUP_5 == 5); + /// Segments every decode launch cuts a row's keys into, whatever the cache /// height: per (row, key head) unit about one block for each SM of either /// workstation card (82 and 84 SMs), and whole merge batches. The grid is @@ -2728,6 +2736,185 @@ unsafe fn merge_sink_body( unsafe { *y.get_unchecked_mut(b * HEAD + d) = out }; } +/// The head-128 scalar segment pass of [`flash_gqa_kernels::gqa_flash_seg_g5`]. +/// `G` is the query heads of one key head. The block stays 256 threads: a +/// group of five leaves warps `5..8` to stage the key and value tile. The +/// group-of-eight entry keeps its own body, so this core is the short group. +/// +/// SAFETY: the entry's launch contract at `G` query heads per key head, a +/// block of 256 threads, and the four tiles this block's shared memory as the +/// entry declared them: the group's query rows (`G · HEAD` f32), the key tile +/// (`KEY_TILE · K_STRIDE` u32), the value tile (`KEY_TILE · ROW_QWORDS` u64) +/// and the weight tile (`G · KEY_TILE` f32). +#[inline(always)] +#[allow( + clippy::too_many_arguments, + reason = "a kernel entry's flat arguments, handed on (rust-quality R8)" +)] +unsafe fn seg_scalar_g( + q: &[f32], + kc: &[u16], + vc: &[u16], + n_keys_buf: &[u32], + scale: f32, + n_kv: u32, + ctx: u32, + segs: u32, + seg_keys: u32, + m: u32, + mut part_v: DisjointSlice, + mut part_ms: DisjointSlice, + qs: *mut f32, + ks: *mut u32, + vs: *mut u64, + ws: *mut f32, +) { + // The group-of-eight entry keeps its own body. This core is a shorter group. + const { assert!(G < GROUP) }; + + let b = thread::blockIdx_x() as usize; + let nkv = n_kv as usize; + let n_seg = segs as usize; + let rows = m as usize; + if b >= rows * nkv * n_seg { + return; // block-uniform + } + let t = b % rows; + let sk = b / rows; + let seg = sk / nkv; + let kh = sk - seg * nkv; + let tid = thread::threadIdx_x() as usize; + let w = tid / 32; + let lane = warp::lane_id() as usize; + let n_head = nkv * G; + let h = kh * G + w; + let idx = (t * n_head + h) * n_seg + seg; + let ctx = ctx as usize; + // SAFETY: t < m <= n_keys_buf.len() by the launch contract. + let limit = live_keys(unsafe { *n_keys_buf.get_unchecked(t) }, ctx); + let span = seg_span(limit, n_seg, seg_keys as usize); + let lo = seg * span; + if lo >= limit { + if lane == 0 && w < G { + // SAFETY: a live warp's h < (kh + 1)·G <= n_head, so idx < + // m·n_head·segs and both slots are inside part_ms. + unsafe { + *part_ms.get_unchecked_mut(2 * idx) = f32::NEG_INFINITY; + *part_ms.get_unchecked_mut(2 * idx + 1) = 0.0; + } + } + return; // block-uniform: lo and limit are the block's + } + let hi = (lo + span).min(limit); + + // Row t's group query rows. A short group stages only the threads that + // own its rows. + let qb = (t * n_head + kh * G) * HEAD; + if tid < G * HEAD / 4 { + let mut i = 0usize; + while i < 4 { + // SAFETY: the thread owns four values inside the group's + // G·HEAD, and qb + that < (t·n_head + (kh + 1)·G)·HEAD <= + // m·n_kv·G·HEAD <= q.len(). + unsafe { *qs.add(4 * tid + i) = *q.get_unchecked(qb + 4 * tid + i) }; + i += 1; + } + } + + let plane = kh * ctx * HEAD; + let k64 = kc.as_ptr() as *const u64; + let v64 = vc.as_ptr() as *const u64; + let mut mx = f32::NEG_INFINITY; + let mut s_sum = 0.0f32; + let mut acc = [0.0f32; 4]; + let mut t0 = lo; + while t0 < hi { + thread::sync_threads(); + // Stage the tile with every thread: 32 keys × 32 u64 of K and of V, + // four of each per thread. A key at or past `hi` stages zeros. + let mut i = 0usize; + while i < 4 { + let e = tid + THREADS * i; + let key = e / ROW_QWORDS; + let word = e - key * ROW_QWORDS; + let (kw, vw) = if t0 + key < hi { + // SAFETY: t0 + key < hi <= ctx, so the row is inside key + // head kh's plane; a row is 128 u16 = 32 u64 and the planes + // start 8-byte aligned (a device allocation), so the u64 + // read is aligned and inside both planes. + unsafe { + let r = (plane + (t0 + key) * HEAD) / 4 + word; + (*k64.add(r), *v64.add(r)) + } + } else { + (0u64, 0u64) + }; + // SAFETY: key < 32 and word < 32: 2·word + 1 < K_STRIDE and + // key·32 + word < KEY_TILE·ROW_QWORDS. + unsafe { + *ks.add(key * K_STRIDE + 2 * word) = kw as u32; + *ks.add(key * K_STRIDE + 2 * word + 1) = (kw >> 32) as u32; + *vs.add(key * ROW_QWORDS + word) = vw; + } + i += 1; + } + thread::sync_threads(); + + if w < G { + // The score of key t0 + lane for head h. + let live = t0 + lane < hi; + let mut a = [0.0f32; ILP]; + let mut wd = 0usize; + while wd < ROW_WORDS { + let mut p = 0usize; + while p < ILP { + // SAFETY: lane < 32, wd + p < ROW_WORDS: inside KS; a + // live warp's query index is inside the group's QS. + let (kw, q0, q1) = unsafe { + ( + *ks.add(lane * K_STRIDE + wd + p), + *qs.add(w * HEAD + 2 * (wd + p)), + *qs.add(w * HEAD + 2 * (wd + p) + 1), + ) + }; + let k0 = half_bits_to_f32(kw as u16); + let k1 = half_bits_to_f32((kw >> 16) as u16); + a[p] = f32::mul_add(q0, k0, a[p]); + a[p] = f32::mul_add(q1, k1, a[p]); + p += 1; + } + wd += ILP; + } + let dot = (a[0] + a[1]) + (a[2] + a[3]); + let sc = if live { dot * scale } else { f32::NEG_INFINITY }; + // SAFETY: WS holds G·KEY_TILE and a live warp has w < G; VS was + // staged before the barrier above. The warp calls it together. + unsafe { fold_tile(sc, live, &mut mx, &mut s_sum, &mut acc, ws, vs, w, lane) }; + } + t0 += KEY_TILE; + } + + if w < G { + // SAFETY: idx is inside the partials, as on the neutral path; lane + // `lane` owns dims 4·lane .. +3. + unsafe { + let o = idx * HEAD + 4 * lane; + *part_v.get_unchecked_mut(o) = acc[0]; + *part_v.get_unchecked_mut(o + 1) = acc[1]; + *part_v.get_unchecked_mut(o + 2) = acc[2]; + *part_v.get_unchecked_mut(o + 3) = acc[3]; + } + } + // Lane 0 of a live warp writes (m, s). + if lane == 0 && w < G { + // SAFETY: idx is inside part_ms, as on the neutral path. + unsafe { + *part_ms.get_unchecked_mut(2 * idx) = mx; + *part_ms.get_unchecked_mut(2 * idx + 1) = s_sum; + } + } +} + #[cuda_module] mod flash_gqa_kernels { use super::*; @@ -2740,7 +2927,9 @@ mod flash_gqa_kernels { /// kh·GROUP + w`, whose partials land at index `(t·n_head + h)·segs + /// seg`. `n_keys` is [`live_keys`] of `n_keys_buf[t]`: a refused row's /// segments are all neutral. The rows of one segment are neighbouring - /// blocks, so they read its key rows while they are in L2. + /// blocks, so they read its key rows while they are in L2. The body stays + /// in this entry: a group of eight fills the block, and the launch bound + /// applies to the thread index here. #[allow( clippy::too_many_arguments, reason = "kernel entry: the device ABI takes the arguments flat (rust-quality R8)" @@ -2925,6 +3114,74 @@ mod flash_gqa_kernels { } } + /// `seg_scalar_g` at [`GROUP_5`]: the same 256-thread segment block as + /// [`gqa_flash_seg`], with warps `0..5` owning the query heads and the + /// pad warps only staging the key and value tile. Five does not pad onto + /// the tensor-core tile (see [`GROUP_5`]), so this entry is scalar. + /// [`gqa_flash_merge`] folds its partials. + #[allow( + clippy::too_many_arguments, + reason = "kernel entry: the device ABI takes the arguments flat (rust-quality R8)" + )] + #[kernel] + #[launch_bounds(256)] + #[launch_contract( + domain = 1, + block = (256, 1, 1), + requires = ( + n_keys_buf.len() >= m, + segs >= 1, + q.len() >= m * n_kv * 5 * 128, + kc.len() >= n_kv * ctx * 128, + vc.len() >= n_kv * ctx * 128, + part_v.len() >= m * n_kv * 5 * segs * 128, + part_ms.len() >= m * n_kv * 5 * segs * 2 + ) + )] + pub fn gqa_flash_seg_g5( + q: &[f32], + kc: &[u16], + vc: &[u16], + n_keys_buf: &[u32], + scale: f32, + n_kv: u32, + ctx: u32, + segs: u32, + seg_keys: u32, + m: u32, + part_v: DisjointSlice, + part_ms: DisjointSlice, + ) { + static mut QS: SharedArray = SharedArray::UNINIT; + static mut KS: SharedArray = SharedArray::UNINIT; + static mut VS: SharedArray = SharedArray::UNINIT; + static mut WS: SharedArray = SharedArray::UNINIT; + + // SAFETY: each `static mut` above is this block's own shared + // allocation, sized for GROUP_5 query heads; the raw form reaches it + // without a reference. The body is `seg_scalar_g` at that group. + unsafe { + seg_scalar_g::( + q, + kc, + vc, + n_keys_buf, + scale, + n_kv, + ctx, + segs, + seg_keys, + m, + part_v, + part_ms, + SharedArray::as_raw_mut_ptr(&raw mut QS), + SharedArray::as_raw_mut_ptr(&raw mut KS), + SharedArray::as_raw_mut_ptr(&raw mut VS), + SharedArray::as_raw_mut_ptr(&raw mut WS), + ) + }; + } + /// The segment pass with `S = Q·Kᵀ` on the tensor cores: the grid, /// the partials and every step after the scores are [`gqa_flash_seg`]'s /// (`fold_tile`), so [`gqa_flash_merge`] serves both. The group's eight @@ -4176,7 +4433,7 @@ impl PackedSeg { } } -/// [`FlashGqaKernels::enqueue_pass`]'s arguments: `m` rows of `n_kv · GROUP` +/// [`FlashGqaKernels::enqueue_pass`]'s arguments: `m` rows of `n_kv · PACK` /// query heads (roped, unscaled, token-major), the layer's two planes of /// `n_kv · ctx` rows, each row's live key count on the device (`m` of them), /// the partials scratch ([`partials_v_len`], [`partials_ms_len`] of `m` @@ -4331,51 +4588,40 @@ impl FlashGqaKernels { Ok(FlashGqaKernels { module }) } - /// Enqueue `m` rows' attention: the segment pass (`m · n_kv · - /// [`SEGMENTS`]` blocks; the tensor-core one when `mma`) and the - /// merge (`m · n_kv · GROUP` blocks). Each block cuts its row's range - /// from the row's own count, so every count the kernels accept is - /// walked in full; a count - /// of zero or past `ctx` raises [`FaultSite::KeyCount`] on `args.fault` - /// and its row is NaN. Two launches. Asynchronous, allocation-free, - /// capturable. - pub fn enqueue_pass( + /// Enqueue `m` rows' attention for `PACK` query heads a key head + /// (`GROUP`, or `GROUP_5` with `mma` false): the segment pass + /// (`m · n_kv · [`SEGMENTS`]` blocks; the tensor-core one when `mma`) + /// and the merge (`m · n_kv · PACK` blocks). Each block cuts its row's + /// range from the row's own count, so every count the kernels accept is + /// walked in full; a count of zero or past `ctx` raises + /// [`FaultSite::KeyCount`] on `a.fault` and its row is NaN. Two launches. + /// Asynchronous, allocation-free, capturable. + pub fn enqueue_pass( &self, stream: &CudaStream, - args: GqaArgs<'_>, + a: GqaArgs<'_>, mma: bool, ) -> Result<(), GpuError> { + const { assert!(PACK == GROUP || PACK == GROUP_5) }; let what = "flash_gqa::enqueue"; - let GqaArgs { - q, - kc, - vc, - n_keys, - scale, - n_kv, - ctx, - m, - part_v, - part_ms, - fault, - y, - } = args; - if n_kv == 0 || ctx == 0 || m == 0 { + if PACK == GROUP_5 && mma { + return Err(GpuError::shape(what, "group 5 requires the scalar pass")); + } + if a.n_kv == 0 || a.ctx == 0 || a.m == 0 { return Err(GpuError::shape( what, - format!("need n_kv, ctx and m >= 1, got n_kv={n_kv} ctx={ctx} m={m}"), + format!("n_kv={}, ctx={}, m={} must be positive", a.n_kv, a.ctx, a.m,), )); } - let n_head = n_kv * GROUP; - let segs = SEGMENTS; + let n_head = a.n_kv * PACK; let lens = [ - ("q", q.len(), m * n_head * HEAD), - ("kc", kc.len(), n_kv * ctx * HEAD), - ("vc", vc.len(), n_kv * ctx * HEAD), - ("n_keys", n_keys.len(), m), - ("part_v", part_v.len(), partials_v_len(m, n_head)), - ("part_ms", part_ms.len(), partials_ms_len(m, n_head)), - ("y", y.len(), m * n_head * HEAD), + ("q", a.q.len(), a.m * n_head * HEAD), + ("kc", a.kc.len(), a.n_kv * a.ctx * HEAD), + ("vc", a.vc.len(), a.n_kv * a.ctx * HEAD), + ("n_keys", a.n_keys.len(), a.m), + ("part_v", a.part_v.len(), partials_v_len(a.m, n_head)), + ("part_ms", a.part_ms.len(), partials_ms_len(a.m, n_head)), + ("y", a.y.len(), a.m * n_head * HEAD), ]; if let Some((name, got, need)) = lens.iter().find(|(_, got, need)| got < need) { return Err(GpuError::shape( @@ -4383,26 +4629,32 @@ impl FlashGqaKernels { format!("{name}.len() {got} < {need}"), )); } - let grid = launch_u32(what, "grid", m * n_kv * segs)?; - let merge_grid = launch_u32(what, "merge grid", m * n_head)?; + let grid = launch_u32(what, "grid", a.m * a.n_kv * SEGMENTS)?; + let merge_grid = launch_u32(what, "merge grid", a.m * n_head)?; let heads = launch_u32(what, "n_head", n_head)?; - let n_kv = launch_u32(what, "n_kv", n_kv)?; - let ctx = launch_u32(what, "ctx", ctx)?; - let segs = launch_u32(what, "segs", segs)?; + let n_kv = launch_u32(what, "n_kv", a.n_kv)?; + let ctx = launch_u32(what, "ctx", a.ctx)?; + let segs = launch_u32(what, "segs", SEGMENTS)?; let seg_keys = launch_u32(what, "seg_keys", SEG_KEYS)?; - let m = launch_u32(what, "m", m)?; + let m = launch_u32(what, "m", a.m)?; let cfg = LaunchConfig1D::new(grid, THREADS_U32, 0); - if mma { + if PACK == GROUP_5 { + let prep = self.module.prepare_gqa_flash_seg_g5(cfg)?; + self.module.gqa_flash_seg_g5( + stream, &prep, a.q, a.kc, a.vc, a.n_keys, a.scale, n_kv, ctx, segs, seg_keys, m, + a.part_v, a.part_ms, + )?; + } else if mma { let prep = self.module.prepare_gqa_flash_seg_mma(cfg)?; self.module.gqa_flash_seg_mma( - stream, &prep, q, kc, vc, n_keys, scale, n_kv, ctx, segs, seg_keys, m, part_v, - part_ms, + stream, &prep, a.q, a.kc, a.vc, a.n_keys, a.scale, n_kv, ctx, segs, seg_keys, m, + a.part_v, a.part_ms, )?; } else { let prep = self.module.prepare_gqa_flash_seg(cfg)?; self.module.gqa_flash_seg( - stream, &prep, q, kc, vc, n_keys, scale, n_kv, ctx, segs, seg_keys, m, part_v, - part_ms, + stream, &prep, a.q, a.kc, a.vc, a.n_keys, a.scale, n_kv, ctx, segs, seg_keys, m, + a.part_v, a.part_ms, )?; } let prep = self.module.prepare_gqa_flash_merge(LaunchConfig1D::new( @@ -4411,7 +4663,8 @@ impl FlashGqaKernels { 0, ))?; self.module.gqa_flash_merge( - stream, &prep, part_v, part_ms, n_keys, heads, ctx, segs, seg_keys, m, fault, y, + stream, &prep, a.part_v, a.part_ms, a.n_keys, heads, ctx, segs, seg_keys, m, a.fault, + a.y, )?; Ok(()) } diff --git a/crates/gpu/src/lib.rs b/crates/gpu/src/lib.rs index e0fad872c..4ee46ab8f 100644 --- a/crates/gpu/src/lib.rs +++ b/crates/gpu/src/lib.rs @@ -75,6 +75,7 @@ pub mod q6k_sel; pub mod q8_0_sel32; pub mod q8f32; pub mod qsa; +pub(crate) mod recur; pub mod rope_neox; pub mod rope_table; pub mod route_body; diff --git a/crates/gpu/src/linear/mod.rs b/crates/gpu/src/linear/mod.rs index 63c969b4a..f828b12b5 100644 --- a/crates/gpu/src/linear/mod.rs +++ b/crates/gpu/src/linear/mod.rs @@ -180,6 +180,14 @@ impl LinearShape { RING_ROWS * self.channels() } + /// f32s of the state of `lanes` lanes and of the conv ring: the pair + /// [`crate::recur::RecStore::new`] allocates. The ring is not multiplied + /// by `lanes`. + #[must_use] + pub(crate) fn lens(&self, lanes: usize) -> [usize; 2] { + [lanes * self.state_len(), self.ring_len()] + } + /// Refuse by name a shape the kernels have no geometry for. pub(crate) fn check(&self, what: &'static str) -> Result<(), GpuError> { if self.n_k == 0 || self.n_v == 0 || !self.n_v.is_multiple_of(self.n_k) { diff --git a/crates/gpu/src/recur.rs b/crates/gpu/src/recur.rs new file mode 100644 index 000000000..ff4f34441 --- /dev/null +++ b/crates/gpu/src/recur.rs @@ -0,0 +1,124 @@ +//! Recurrent state a layer keeps on the card: the state lanes and the conv +//! ring. A delta layer passes [`crate::linear::LinearShape::lens`]; a later +//! scan passes its own lengths. The length pair is the only architecture +//! fact this module takes. + +use crate::GpuError; +use crate::linear::delta::NEVER; +use cuda_core::{CudaStream, DeviceBuffer}; + +/// One layer's recurrent store: `state_f32` f32s of state (every lane) and +/// `ring_f32` f32s of conv ring, and — when a load keeps a lane per verify +/// row — the lanes' stamps ([`fresh_stamps`]). +pub(crate) struct RecStore { + pub(crate) state: DeviceBuffer, + pub(crate) ring: DeviceBuffer, + pub(crate) lanes: usize, + /// `lanes` u32 stamps ([`fresh_stamps`]); `None` on a store whose + /// launches do not read stamps. + pub(crate) stamp: Option>, +} + +impl RecStore { + /// A zeroed store of `lanes` lanes over `lens` (`[state_f32, ring_f32]`), + /// with no stamps. Load-time only. + pub(crate) fn new( + stream: &CudaStream, + lens: [usize; 2], + lanes: usize, + ) -> Result { + let [state, ring] = lens; + Ok(RecStore { + state: DeviceBuffer::zeroed(stream, state)?, + ring: DeviceBuffer::zeroed(stream, ring)?, + lanes, + stamp: None, + }) + } + + /// [`RecStore::new`] with its stamps fresh ([`fresh_stamps`]). Load-time + /// only. + pub(crate) fn stamped( + stream: &CudaStream, + lens: [usize; 2], + lanes: usize, + ) -> Result { + let mut r = RecStore::new(stream, lens, lanes)?; + r.stamp = Some(DeviceBuffer::from_host(stream, &fresh_stamps(lanes))?); + Ok(r) + } + + /// The stamps back to a fresh store's ([`fresh_stamps`]); a store with + /// none is refused by name. Never inside a capture. + pub(crate) fn restamp(&mut self, stream: &CudaStream) -> Result<(), GpuError> { + let fresh = fresh_stamps(self.lanes); + self.stamp_with(stream, &fresh) + } + + /// The stamps with lane `lane` at position `at` and every other lane + /// never written ([`stamps_at`]). A store with no stamps is refused by + /// name. Never inside a capture. + pub(crate) fn restamp_at( + &mut self, + stream: &CudaStream, + lane: u32, + at: u32, + ) -> Result<(), GpuError> { + let st = stamps_at(self.lanes, lane, at)?; + self.stamp_with(stream, &st) + } + + fn stamp_with(&mut self, stream: &CudaStream, st: &[u32]) -> Result<(), GpuError> { + match self.stamp.as_mut() { + Some(stamp) => Ok(stamp.copy_from_host(stream, st)?), + None => Err(GpuError::state( + "recur::RecStore::restamp", + "a store of stamped lanes", + )), + } + } + + /// Every lane of the state and the whole ring back to zero, in stream + /// order through `zeros` (host zeros at least as long as the longer of + /// the two). Never inside a capture. + pub(crate) fn clear(&mut self, stream: &CudaStream, zeros: &[f32]) -> Result<(), GpuError> { + for buf in [&mut self.state, &mut self.ring] { + let n = buf.len(); + let src = zeros.get(..n).ok_or_else(|| { + GpuError::shape( + "recur::RecStore::clear", + format!("{} host zeros for a {n}-value buffer", zeros.len()), + ) + })?; + buf.copy_from_host(stream, src)?; + } + Ok(()) + } + + pub(crate) fn bytes(&self) -> usize { + self.state.num_bytes() + + self.ring.num_bytes() + + self.stamp.as_ref().map_or(0, DeviceBuffer::num_bytes) + } +} + +/// A fresh store's stamps over `lanes` lanes: lane 0 holds position 0, every +/// other lane was never written ([`NEVER`]). +pub(crate) fn fresh_stamps(lanes: usize) -> Vec { + let mut st = vec![NEVER; lanes]; + if let Some(first) = st.first_mut() { + *first = 0; + } + st +} + +/// The stamps of a store over `lanes` lanes whose lane `lane` stands at +/// position `at` and whose every other lane was never written. A lane at or +/// past `lanes` is refused by name. +pub(crate) fn stamps_at(lanes: usize, lane: u32, at: u32) -> Result, GpuError> { + let mut st = vec![NEVER; lanes]; + *st.get_mut(lane as usize).ok_or_else(|| { + GpuError::shape("recur::stamps_at", format!("lane {lane} of {lanes} lanes")) + })? = at; + Ok(st) +} diff --git a/crates/model/src/arch/coverage.rs b/crates/model/src/arch/coverage.rs index 1b1619ff9..8bc1b26ce 100644 --- a/crates/model/src/arch/coverage.rs +++ b/crates/model/src/arch/coverage.rs @@ -108,7 +108,9 @@ pub const AVAILABLE: &[Available] = &[ window: None, sinks: false, scaled: false, - } if value == head && gqa_row(*head, *group).is_some() + // The pack-5 row (head 128, group 5) is launched by no Qwen + // body; the program that launches it names its own entry. + } if value == head && gqa_row(*head, *group).is_some_and(|r| r.pack != 5) ) }, }, @@ -191,7 +193,7 @@ pub const AVAILABLE: &[Available] = &[ }, Available { programs: &[Program::Qwen35Body], - at: "gpu/src/arch/qwen3moe/scratch.rs RecStore (Body35's delta stores)", + at: "gpu/src/recur.rs RecStore (Body35's delta stores)", runs: |n| matches!(n, Need::RecurrentState), }, Available { @@ -1136,6 +1138,17 @@ mod tests { unreachable!() }; for (what, need) in [ + ( + "group 5, a row no Qwen body launches", + Need::Gqa { + head, + value: head, + group: 5, + window: None, + sinks: false, + scaled: false, + }, + ), ( "a narrower value head", Need::Gqa { diff --git a/crates/models/src/shape.rs b/crates/models/src/shape.rs index c052bf5c9..305f5a01a 100644 --- a/crates/models/src/shape.rs +++ b/crates/models/src/shape.rs @@ -534,6 +534,16 @@ pub const GQA: &[GqaInst] = &[ value_scale: false, at: "gpu/src/flash_gqa.rs HEAD, GROUP", }, + GqaInst { + head: 128, + value: 128, + pack: 5, + group: GroupRule::One, + window: false, + sinks: false, + value_scale: false, + at: "gpu/src/flash_gqa.rs GROUP_5 (gqa_flash_seg_g5, scalar)", + }, GqaInst { head: 256, value: 256, @@ -913,7 +923,7 @@ mod tests { e.why, Refusal::GqaGroup { group: 12, - served: vec!["exactly 8".to_string()] + served: vec!["exactly 8".to_string(), "exactly 5".to_string()] } ); let e = select_gqa(attn(16, 2, 160)).expect_err("head 160"); @@ -930,6 +940,18 @@ mod tests { assert!(e.to_string().contains("group 3"), "{e}"); } + /// Head 128 at five query heads a key head (10/2) selects the scalar + /// group-5 row: one block a group, no window, sinks or value scale. + #[test] + fn a_head_128_group_of_5_selects_the_scalar_row() { + let s = attn(10, 2, 128); + let row = select_gqa(s).unwrap_or_else(|e| panic!("{e}")); + assert_eq!(row.pack, 5); + assert_eq!(row.group, GroupRule::One); + assert!(!row.window && !row.sinks && !row.value_scale); + assert_eq!(row.packs(5), 1); + } + fn mimo(n_head: u32, n_kv: u32, window: Option, sinks: bool) -> AttnShape { AttnShape { n_head, diff --git a/tools/ref/ptx-shapes.tsv b/tools/ref/ptx-shapes.tsv index d4afef493..e386a1e16 100644 --- a/tools/ref/ptx-shapes.tsv +++ b/tools/ref/ptx-shapes.tsv @@ -102,6 +102,7 @@ generate_ds41 ds41_tap_means 0 0 generate_ds41 embed_rows 0 0 generate_ds41 embed_rows_q3k 0 0 generate_ds41 embed_rows_q4k 0 0 +generate_ds41 embed_rows_q5_0 0 0 generate_ds41 embed_rows_q5k 0 0 generate_ds41 embed_rows_q6k 0 0 generate_ds41 embed_rows_q8_0 0 0 @@ -147,6 +148,7 @@ generate_ds41 gqa_flash_seg 0 0 generate_ds41 gqa_flash_seg_256 0 0 generate_ds41 gqa_flash_seg_256_p2 0 0 generate_ds41 gqa_flash_seg_256_p4 0 0 +generate_ds41 gqa_flash_seg_g5 0 0 generate_ds41 gqa_flash_seg_k192 0 0 generate_ds41 gqa_flash_seg_k192_ring 0 0 generate_ds41 gqa_flash_seg_mma 0 0 @@ -349,6 +351,7 @@ gate_e2e ds41_markov 0 0 gate_e2e embed_rows 0 0 gate_e2e embed_rows_q3k 0 0 gate_e2e embed_rows_q4k 0 0 +gate_e2e embed_rows_q5_0 0 0 gate_e2e embed_rows_q5k 0 0 gate_e2e embed_rows_q6k 0 0 gate_e2e embed_rows_q8_0 0 0 @@ -391,6 +394,7 @@ gate_e2e gqa_flash_seg 0 0 gate_e2e gqa_flash_seg_256 0 0 gate_e2e gqa_flash_seg_256_p2 0 0 gate_e2e gqa_flash_seg_256_p4 0 0 +gate_e2e gqa_flash_seg_g5 0 0 gate_e2e gqa_flash_seg_k192 0 0 gate_e2e gqa_flash_seg_k192_ring 0 0 gate_e2e gqa_flash_seg_mma 0 0 diff --git a/tools/unsafe-ratchet.txt b/tools/unsafe-ratchet.txt index bfffd554c..621a6660e 100644 --- a/tools/unsafe-ratchet.txt +++ b/tools/unsafe-ratchet.txt @@ -225,7 +225,11 @@ gguf 9 2 0 # layers' conv and the Mamba-2 step now share (one unsafe fn, its one block); each block with its SAFETY line. # PIN(2026-10-14): gpu 1769 -> 1771 blocks, a new kernel entry, not old code relaxed: gqa_prefill_flash_256_p2_lower's # call into prefill_256_p and the walk's read of a row's first key (k_first_buf), each with its SAFETY line. -gpu 1771 195 10 +# PIN(2026-10-15): gpu 1771 -> 1787 blocks, 195 -> 196 fns, new code, not old code relaxed: embed_rows_q5_0 (the Q5_0 +# embedding gather: its position words, the code and scale words, the output row) and gqa_flash_seg_g5 (the head-128 +# group-5 flash entry) read and write through raw pointers, and seg_scalar_g (the scalar core that entry calls) is one +# unsafe fn. The group-8 entry keeps its own body. Each block with its SAFETY line. +gpu 1787 196 10 # PIN(2026-10-01): gpu-deepseek41 459 -> 464 blocks, 51 -> 52 fns, a new kernel entry, not old code relaxed: # ds41_ffn_card_acc_8_tier (the tier form of the card sum at 8 slots, chain/ffn/batch.rs) adds its two blocks and # card_acc_tier_at (the one unsafe fn) with its three, each with its SAFETY line.