From 0073585bdcaf79dcd14a2a0804c0d6f7a44dd143 Mon Sep 17 00:00:00 2001 From: midagedev Date: Thu, 15 Oct 2026 08:15:15 +0900 Subject: [PATCH 01/13] gpu, models, model, gpu-gates: the common pieces a Nemotron-H card body attaches to - The recurrent store (RecStore, LayerStore) leaves arch/qwen3moe for gpu::recur and takes its length pair; LinearShape keeps the method that gives its pair, so the GDN bodies size their stores as before. - A Q5_0 embedding row (embed_rows_q5_0) over pack_q5_0's layout; embed_into admits Q5_0. - relu_sqr beside swiglu. - A flash GQA decode entry at head 128, group 5, with its GqaInst row. - Coverage names the rows that are now available. Co-Authored-By: Claude Opus 5.5 --- .../src/bin/gate_qwen3moe_experts.rs | 159 +++++++++- .../gpu-gates/src/bin/gate_qwen3moe_flash.rs | 76 ++++- crates/gpu/src/arch/qwen3moe/body35.rs | 13 +- crates/gpu/src/arch/qwen3moe/dispatch.rs | 30 +- crates/gpu/src/arch/qwen3moe/scratch.rs | 160 +--------- crates/gpu/src/arch/qwen3moe/scratch38.rs | 2 +- crates/gpu/src/elem.rs | 243 +++++++++++++++- crates/gpu/src/flash_gqa.rs | 275 ++++++++++++++++++ crates/gpu/src/lib.rs | 1 + crates/gpu/src/linear/mod.rs | 8 + crates/gpu/src/recur.rs | 131 +++++++++ crates/model/src/arch/coverage.rs | 2 +- crates/models/src/shape.rs | 24 +- 13 files changed, 964 insertions(+), 160 deletions(-) create mode 100644 crates/gpu/src/recur.rs diff --git a/crates/gpu-gates/src/bin/gate_qwen3moe_experts.rs b/crates/gpu-gates/src/bin/gate_qwen3moe_experts.rs index 9c2a0809d..1d1af96cf 100644 --- a/crates/gpu-gates/src/bin/gate_qwen3moe_experts.rs +++ b/crates/gpu-gates/src/bin/gate_qwen3moe_experts.rs @@ -53,6 +53,10 @@ //! //! 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. `relu_sqr` is bit-identical to `max(x, 0)^2`, and a +//! non-finite input is copied through. #[cfg(feature = "gpu")] fn main() -> std::process::ExitCode { @@ -72,7 +76,7 @@ mod gate { topk_ids_logical_within, verdict, }; use cuda_core::{CudaStream, DeviceBuffer}; - use gguf::quant::{GgmlType, dequant_row}; + use gguf::quant::{GgmlType, dequant_row, f32_to_f16_bits}; use model::arch::qwen3moe::names; /// f32 roundings each side's dot is allowed per output, on `Σ|w·x̂|` @@ -136,6 +140,8 @@ mod gate { let mut ok = true; ok &= embed(&gpu, &gguf, &sets)?; ok &= embed_kquant(&gpu)?; + ok &= embed_q5_0(&gpu)?; + ok &= relu_sqr(&gpu)?; ok &= gate_up_body(&gpu, &gguf, &ex)?; ok &= gate_up_tap(&gpu, &gguf, &ex, &sets)?; ok &= combine(&gpu, &ex, &sets, stream)?; @@ -146,6 +152,157 @@ mod gate { Ok(()) } + /// `rows` of Q5_0 (`k` values, two blocks) whose low and high nibbles + /// differ, so a read of the low half's word in place of the high half's + /// disagrees with `dequant_row`. + fn q5_0_bytes(rows: usize, k: usize) -> Vec { + let blocks = k / 32; + let mut bytes = vec![0u8; rows * blocks * 22]; + for r in 0..rows { + for b in 0..blocks { + let blk = &mut bytes[(r * blocks + b) * 22..][..22]; + let d = 0.5f32 + 0.25 * (b + r) as f32; + blk[..2].copy_from_slice(&f32_to_f16_bits(d).to_le_bytes()); + let qh: u32 = 0xA5A5_5A5A ^ (r as u32).wrapping_mul(0x0101_0101); + blk[2..6].copy_from_slice(&qh.to_le_bytes()); + for j in 0..16 { + blk[6 + j] = j as u8 | ((15 - j) as u8) << 4; + } + } + } + bytes + } + + /// 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; + const K: usize = 64; + let bytes = q5_0_bytes(ROWS, K); + 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| -> Result<(Embedded, usize), GateError> { + let stream = gpu.stream(); + let idb = DeviceBuffer::from_host(stream, ids)?; + let p0 = DeviceBuffer::from_host(stream, &[POS0])?; + let mut y = DeviceBuffer::from_host(stream, &vec![SENT; ids.len() * K])?; + let mut pos = DeviceBuffer::from_host(stream, &vec![SENT_POS; ids.len()])?; + let mut n_keys = DeviceBuffer::from_host(stream, &vec![SENT_POS; ids.len()])?; + let mut enqueue = |s: &CudaStream| { + gpu.elem().enqueue_embed_rows_q5_0( + s, + EmbedRowsArgs { + w: &table, + ids: &idb, + pos0: &p0, + first: FIRST, + y: &mut y, + pos: &mut pos, + n_keys: &mut n_keys, + }, + K, + ) + }; + let nodes = if graph { + let g = gpu.capture(&mut enqueue)?; + g.launch(stream)?; + g.node_count() + } else { + enqueue(stream)?; + 0 + }; + stream.synchronize()?; + Ok(( + Embedded { + y: y.to_host_vec(stream)?, + pos: pos.to_host_vec(stream)?, + n_keys: n_keys.to_host_vec(stream)?, + }, + nodes, + )) + }; + 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) + } + + /// Host transcription of `relu_sqr`: a non-finite input stays that value, + /// and a finite one is `max(x, 0)^2`. + fn host_relu_sqr(v: f32) -> f32 { + if !v.is_finite() { + v + } else { + let r = if v > 0.0 { v } else { 0.0 }; + r * r + } + } + + /// `relu_sqr` against [`host_relu_sqr`], bit for bit, including a + /// non-finite input, and one captured node. + fn relu_sqr(gpu: &Gpu) -> Result { + let x = [ + 0.0f32, + 1.5, + -2.0, + f32::NAN, + f32::INFINITY, + f32::NEG_INFINITY, + 3.0, + ]; + let want: Vec = x.iter().copied().map(host_relu_sqr).collect(); + let stream = gpu.stream(); + let xb = DeviceBuffer::from_host(stream, &x)?; + let mut y = DeviceBuffer::::zeroed(stream, x.len())?; + gpu.elem().enqueue_relu_sqr(stream, &xb, x.len(), &mut y)?; + stream.synchronize()?; + let got = y.to_host_vec(stream)?; + let host_ok = bits_equal(&got, &want); + // The eager read ends before this closure borrows `y` again. + let mut enqueue = |s: &CudaStream| gpu.elem().enqueue_relu_sqr(s, &xb, x.len(), &mut y); + let g = gpu.capture(&mut enqueue)?; + g.launch(stream)?; + stream.synchronize()?; + let replay = y.to_host_vec(stream)?; + let nodes = g.node_count(); + let graph_ok = bits_equal(&replay, &want) && nodes == 1; + println!( + "relu_sqr n={}: vs host bit_identical={host_ok} graph_nodes={nodes} graph_same={graph_ok} {}", + x.len(), + verdict(host_ok && graph_ok) + ); + Ok(host_ok && graph_ok) + } + // ------------------------------------------------------------ 1. embed /// One gather's outputs, read back: the rows, each row's position and diff --git a/crates/gpu-gates/src/bin/gate_qwen3moe_flash.rs b/crates/gpu-gates/src/bin/gate_qwen3moe_flash.rs index cfad3b7b5..a79463fdc 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::{ @@ -1901,6 +1901,79 @@ 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_g5( + 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, + }, + )?; + 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 +2215,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..36f559a94 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,12 @@ 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) => { + let [state_f32, ring_f32] = shape.lens(lanes); + LayerStore::Rec(RecStore::new(stream, [state_f32, ring_f32], lanes)?) } - Kind35::Delta(shape) => LayerStore::Rec(RecStore::new(stream, shape, lanes)?), }); } let forms = Forms::of(&plans, &d); @@ -3153,7 +3156,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 +3338,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..0f3c8c43a 100644 --- a/crates/gpu/src/arch/qwen3moe/dispatch.rs +++ b/crates/gpu/src/arch/qwen3moe/dispatch.rs @@ -472,8 +472,8 @@ pub(super) fn mixer( /// A unit's front, one launch: the embedding rows of `io`'s ids into `s.x` /// 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). +/// picks the lookup: Q3_K, Q4_K, Q5_K or Q6_K rows, Q5_0 packed rows, or +/// Q8_0 planes (the load admits only these). pub(super) fn embed_rows( gpu: &Gpu, w: &Weights, @@ -486,7 +486,8 @@ 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`]. pub(super) fn embed_into( gpu: &Gpu, w: &Weights, @@ -516,9 +517,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, Q5_0 packed rows or Q8_0 planes", + )); + } }; 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..f7288abff 100644 --- a/crates/gpu/src/arch/qwen3moe/scratch.rs +++ b/crates/gpu/src/arch/qwen3moe/scratch.rs @@ -29,6 +29,8 @@ use crate::flash_gqa::{ 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, @@ -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>>; @@ -1647,33 +1518,32 @@ impl LaneWord { } /// A layer's state store: K/V planes by position, or a recurrent store. -pub(super) enum LayerStore { - Kv(KvPlanes), - Rec(RecStore), +/// The variants are [`crate::recur::LayerStore`]'s; the methods below are +/// this family's. +pub(super) type LayerStore = crate::recur::LayerStore; + +/// 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 { +impl crate::recur::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..1ff6ffe3e 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,100 @@ 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; + 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`. @@ -1003,6 +1098,36 @@ mod elem_kernels { e.map(silu_mul); } + /// `y[i] = max(x[i], 0)^2` over `n` values. A non-finite input is copied + /// as it is: `max` of a non-finite is not this rule, and the square of + /// `−inf` is not 0. + #[kernel] + #[launch_bounds(256)] + #[launch_contract( + domain = 1, + block = (256, 1, 1), + requires = (x.len() >= n, y.len() >= n) + )] + pub fn relu_sqr(x: &[f32], n: u32, mut y: DisjointSlice) { + let i = thread::index_1d().get(); + if i >= n as usize { + return; + } + // SAFETY: i < n <= x.len() and y.len() by the launch contract; thread + // i owns entry i. + let v = unsafe { *x.get_unchecked(i) }; + let out = if !v.is_finite() { + v + } else { + let r = if v > 0.0 { v } else { 0.0 }; + let yv = r * r; + yv + }; + unsafe { + *y.get_unchecked_mut(i) = out; + } + } + /// `y[i] = a[i] + b[i]`, elementwise over `n` values — exact; the body /// is one add, so nothing is extracted into a core. #[kernel] @@ -1641,6 +1766,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 @@ -1955,6 +2173,29 @@ impl ElemKernels { Ok(()) } + /// Enqueue `y = max(x, 0)^2` over `n` values (flat). A non-finite input + /// stays that value. Asynchronous, allocation-free, capturable. + pub fn enqueue_relu_sqr( + &self, + stream: &CudaStream, + x: &DeviceBuffer, + n: usize, + y: &mut DeviceBuffer, + ) -> Result<(), GpuError> { + if n == 0 || x.len() < n || y.len() < n { + return Err(GpuError::shape( + "enqueue_relu_sqr", + format!("n={n}, x.len() {}, y.len() {}", x.len(), y.len()), + )); + } + let n = launch_u32("enqueue_relu_sqr", "n", n)?; + let prep = self + .module + .prepare_relu_sqr(LaunchConfig1D::new(n.div_ceil(256), 256, 0))?; + self.module.relu_sqr(stream, &prep, x, n, y)?; + Ok(()) + } + /// Enqueue `y = a + b` over `n` values (flat). Exact. Asynchronous, /// allocation-free, capturable. pub fn enqueue_add( diff --git a/crates/gpu/src/flash_gqa.rs b/crates/gpu/src/flash_gqa.rs index f6390aed3..71f6febb6 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_g5`] 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 @@ -2925,6 +2933,206 @@ mod flash_gqa_kernels { } } + /// [`gqa_flash_seg`] for [`GROUP_5`] query heads a key head. The block + /// stays 256 threads so the key and value tile stage as that kernel's; + /// warps `0..5` own the query heads and the pad warps only help stage. + /// Five does not pad onto the tensor-core tile (see [`GROUP_5`]), so + /// this entry is scalar. The partial index is `(t·n_head + h)·segs + seg` + /// with `h = kh·GROUP_5 + w`, and [`gqa_flash_merge`] folds it. + #[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, + mut part_v: DisjointSlice, + mut 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; + + 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 * GROUP_5; + let h0 = kh * GROUP_5; + // The key head this block reads. For a live warp it is the query + // head's group; the pad warps share it so the staged tile is one. + let src = h0 / GROUP_5; + let h = h0 + 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 w < GROUP_5 && lane == 0 { + // SAFETY: w < GROUP_5 gives h < (kh+1)·GROUP_5 <= 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); + + // SAFETY: each `static mut` above is this block's own shared + // allocation; the raw form reaches it without a reference. Every + // index below is inside its array, and every write is ordered before + // its reads by a block barrier (or, for WS, written and read by one + // warp between two barriers). + let (qs, ks, vs, ws) = unsafe { + ( + 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 group's 5·128 query values, four per thread. Threads past the + // group's 160 stage nothing; the pad warps have no query row. + let qb = (t * n_head + h0) * HEAD; + if tid < GROUP_5 * HEAD / 4 { + let mut i = 0usize; + while i < 4 { + // SAFETY: tid < 160, so 4·tid + i < GROUP_5·HEAD, and + // qb + that < (t·n_head + (kh+1)·GROUP_5)·HEAD <= + // m·n_kv·5·128 <= q.len(). + unsafe { *qs.add(4 * tid + i) = *q.get_unchecked(qb + 4 * tid + i) }; + i += 1; + } + } + + let plane = src * 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, as `gqa_flash_seg` does: 32 + // keys × 32 u64 of K and of V. 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 `src`'s plane; a row is 128 u16 = 32 u64 and the + // planes start 8-byte aligned (a device allocation). + 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 < GROUP_5 { + // 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; + // w < GROUP_5, so the query index is inside 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 is GROUP_5·KEY_TILE and w < GROUP_5; 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 < GROUP_5 { + // SAFETY: idx is inside the partials, as on the neutral path; + // lane `lane` owns dims 4·lane .. +3, and lane 0 writes (m, s). + 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]; + if lane == 0 { + *part_ms.get_unchecked_mut(2 * idx) = mx; + *part_ms.get_unchecked_mut(2 * idx + 1) = s_sum; + } + } + } + } + /// 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 @@ -4416,6 +4624,73 @@ impl FlashGqaKernels { Ok(()) } + /// [`FlashGqaKernels::enqueue_pass`] at [`GROUP_5`] query heads a key + /// head: `gqa_flash_seg_g5` (`m · n_kv · [`SEGMENTS`]` blocks of 256) and + /// [`gqa_flash_merge`] over `n_kv · GROUP_5` heads. Scalar only. Two + /// launches. Asynchronous, allocation-free, capturable. + pub fn enqueue_pass_g5(&self, stream: &CudaStream, args: GqaArgs<'_>) -> Result<(), GpuError> { + let what = "flash_gqa::enqueue_g5"; + 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 { + return Err(GpuError::shape( + what, + format!("need n_kv, ctx and m >= 1, got n_kv={n_kv} ctx={ctx} m={m}"), + )); + } + let n_head = n_kv * GROUP_5; + let segs = SEGMENTS; + 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), + ]; + if let Some((name, got, need)) = lens.iter().find(|(_, got, need)| got < need) { + return Err(GpuError::shape( + what, + 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 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 seg_keys = launch_u32(what, "seg_keys", SEG_KEYS)?; + let m = launch_u32(what, "m", m)?; + let cfg = LaunchConfig1D::new(grid, THREADS_U32, 0); + let prep = self.module.prepare_gqa_flash_seg_g5(cfg)?; + self.module.gqa_flash_seg_g5( + stream, &prep, q, kc, vc, n_keys, scale, n_kv, ctx, segs, seg_keys, m, part_v, part_ms, + )?; + let prep = self.module.prepare_gqa_flash_merge(LaunchConfig1D::new( + merge_grid, + MERGE_THREADS, + 0, + ))?; + self.module.gqa_flash_merge( + stream, &prep, part_v, part_ms, n_keys, heads, ctx, segs, seg_keys, m, fault, y, + )?; + Ok(()) + } + /// [`FlashGqaKernels::enqueue_pass`] over heads of [`HEAD_256`] values: /// `args.q` and `args.y` hold `m · n_kv · GROUP` rows of 256, the planes /// `n_kv · ctx` rows of 256 f16, `args.part_v` [`partials_v_len_256`]. 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..781102401 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 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..eeabdec4b --- /dev/null +++ b/crates/gpu/src/recur.rs @@ -0,0 +1,131 @@ +//! 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) +} + +/// A layer's state store: `K` (the architecture's K/V planes) or a +/// [`RecStore`]. +pub(crate) enum LayerStore { + Kv(K), + Rec(RecStore), +} diff --git a/crates/model/src/arch/coverage.rs b/crates/model/src/arch/coverage.rs index 1b1619ff9..9ba5cb2bb 100644 --- a/crates/model/src/arch/coverage.rs +++ b/crates/model/src/arch/coverage.rs @@ -191,7 +191,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 { 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, From 80c58096a6262b0c726fa002bf049881df4b616d Mon Sep 17 00:00:00 2001 From: midagedev Date: Thu, 15 Oct 2026 08:44:50 +0900 Subject: [PATCH 02/13] tools: pin the Q5_0 embed row, relu_sqr and the head-128 group-5 flash entry at spill 0 Co-Authored-By: Claude Opus 5.5 --- tools/ref/ptx-shapes.tsv | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tools/ref/ptx-shapes.tsv b/tools/ref/ptx-shapes.tsv index d4afef493..12783be4e 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 @@ -307,6 +309,7 @@ generate_ds41 qwen3moe_router_norm_m_prefetch 0 0 generate_ds41 qwen3moe_router_norm_prefetch 0 0 generate_ds41 qwen3moe_router_route 0 0 generate_ds41 qwen3moe_token_major 0 0 +generate_ds41 relu_sqr 0 0 generate_ds41 rms_norm 0 0 generate_ds41 rope 0 0 generate_ds41 router_topk 0 0 @@ -349,6 +352,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 +395,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 @@ -547,6 +552,7 @@ gate_e2e qwen3moe_router_norm_m_prefetch 0 0 gate_e2e qwen3moe_router_norm_prefetch 0 0 gate_e2e qwen3moe_router_route 0 0 gate_e2e qwen3moe_token_major 0 0 +gate_e2e relu_sqr 0 0 gate_e2e rms_norm 0 0 gate_e2e rope 0 0 gate_e2e router_topk 0 0 From d2df7462eb5a172390ce0caf041dc691d0c3dc6a Mon Sep 17 00:00:00 2001 From: midagedev Date: Thu, 15 Oct 2026 09:21:10 +0900 Subject: [PATCH 03/13] gpu, gpu-gates: the Q5_0 embed case spans a second code window and reuses gather; LayerStore stays qwen3moe's Co-Authored-By: Claude Opus 5.5 --- .../src/bin/gate_qwen3moe_experts.rs | 72 +++++++------------ crates/gpu/src/arch/qwen3moe/scratch.rs | 9 +-- crates/gpu/src/recur.rs | 7 -- 3 files changed, 31 insertions(+), 57 deletions(-) diff --git a/crates/gpu-gates/src/bin/gate_qwen3moe_experts.rs b/crates/gpu-gates/src/bin/gate_qwen3moe_experts.rs index 1d1af96cf..c630a19d2 100644 --- a/crates/gpu-gates/src/bin/gate_qwen3moe_experts.rs +++ b/crates/gpu-gates/src/bin/gate_qwen3moe_experts.rs @@ -152,9 +152,11 @@ mod gate { Ok(()) } - /// `rows` of Q5_0 (`k` values, two blocks) whose low and high nibbles - /// differ, so a read of the low half's word in place of the high half's - /// disagrees with `dequant_row`. + /// `rows` of Q5_0 (`k` values, `k / 32` blocks) whose low and high + /// nibbles differ and whose codes differ from block to block, so a read + /// of the low half's word in place of the high half's, or of another + /// block's (another 32-block window's) words, disagrees with + /// `dequant_row`. fn q5_0_bytes(rows: usize, k: usize) -> Vec { let blocks = k / 32; let mut bytes = vec![0u8; rows * blocks * 22]; @@ -163,10 +165,14 @@ mod gate { let blk = &mut bytes[(r * blocks + b) * 22..][..22]; let d = 0.5f32 + 0.25 * (b + r) as f32; blk[..2].copy_from_slice(&f32_to_f16_bits(d).to_le_bytes()); - let qh: u32 = 0xA5A5_5A5A ^ (r as u32).wrapping_mul(0x0101_0101); + let qh: u32 = 0xA5A5_5A5A + ^ (r as u32).wrapping_mul(0x0101_0101) + ^ (b as u32).wrapping_mul(0x9E37_79B9); blk[2..6].copy_from_slice(&qh.to_le_bytes()); for j in 0..16 { - blk[6 + j] = j as u8 | ((15 - j) as u8) << 4; + let lo = (j + b) % 16; + let hi = (15 - j + 3 * b) % 16; + blk[6 + j] = lo as u8 | (hi as u8) << 4; } } } @@ -178,51 +184,23 @@ mod gate { /// past the table (NaN row, `FaultSite::TokenId`, the other rows intact). fn embed_q5_0(gpu: &Gpu) -> Result { const ROWS: usize = 4; - const K: usize = 64; + // 34 blocks: past the first 32-block window of the packed code plane. + const K: usize = 1088; let bytes = q5_0_bytes(ROWS, K); 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| -> Result<(Embedded, usize), GateError> { - let stream = gpu.stream(); - let idb = DeviceBuffer::from_host(stream, ids)?; - let p0 = DeviceBuffer::from_host(stream, &[POS0])?; - let mut y = DeviceBuffer::from_host(stream, &vec![SENT; ids.len() * K])?; - let mut pos = DeviceBuffer::from_host(stream, &vec![SENT_POS; ids.len()])?; - let mut n_keys = DeviceBuffer::from_host(stream, &vec![SENT_POS; ids.len()])?; - let mut enqueue = |s: &CudaStream| { - gpu.elem().enqueue_embed_rows_q5_0( - s, - EmbedRowsArgs { - w: &table, - ids: &idb, - pos0: &p0, - first: FIRST, - y: &mut y, - pos: &mut pos, - n_keys: &mut n_keys, - }, - K, - ) - }; - let nodes = if graph { - let g = gpu.capture(&mut enqueue)?; - g.launch(stream)?; - g.node_count() - } else { - enqueue(stream)?; - 0 - }; - stream.synchronize()?; - Ok(( - Embedded { - y: y.to_host_vec(stream)?, - pos: pos.to_host_vec(stream)?, - n_keys: n_keys.to_host_vec(stream)?, - }, - nodes, - )) + 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]; @@ -316,7 +294,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, @@ -345,6 +324,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/src/arch/qwen3moe/scratch.rs b/crates/gpu/src/arch/qwen3moe/scratch.rs index f7288abff..f2bbbd3be 100644 --- a/crates/gpu/src/arch/qwen3moe/scratch.rs +++ b/crates/gpu/src/arch/qwen3moe/scratch.rs @@ -1518,9 +1518,10 @@ impl LaneWord { } /// A layer's state store: K/V planes by position, or a recurrent store. -/// The variants are [`crate::recur::LayerStore`]'s; the methods below are -/// this family's. -pub(super) type LayerStore = crate::recur::LayerStore; +pub(super) enum LayerStore { + Kv(KvPlanes), + Rec(RecStore), +} /// A layer's store as one layer's enqueue borrows it. pub(super) enum StoreMut<'a> { @@ -1528,7 +1529,7 @@ pub(super) enum StoreMut<'a> { Rec(&'a mut RecStore), } -impl crate::recur::LayerStore { +impl LayerStore { pub(super) fn as_mut(&mut self) -> StoreMut<'_> { match self { Self::Kv(k) => StoreMut::Kv(k), diff --git a/crates/gpu/src/recur.rs b/crates/gpu/src/recur.rs index eeabdec4b..ff4f34441 100644 --- a/crates/gpu/src/recur.rs +++ b/crates/gpu/src/recur.rs @@ -122,10 +122,3 @@ pub(crate) fn stamps_at(lanes: usize, lane: u32, at: u32) -> Result, Gp })? = at; Ok(st) } - -/// A layer's state store: `K` (the architecture's K/V planes) or a -/// [`RecStore`]. -pub(crate) enum LayerStore { - Kv(K), - Rec(RecStore), -} From bd926f650f23c1bf25b1cd079776cab98a69fe72 Mon Sep 17 00:00:00 2001 From: midagedev Date: Thu, 15 Oct 2026 10:21:21 +0900 Subject: [PATCH 04/13] gpu, model, tools: the new kernels' SAFETY lines and ratchet pin, and coverage keeps the group-5 row off the Qwen bodies Co-Authored-By: Claude Opus 5.5 --- crates/gpu/src/arch/qwen3moe/body35.rs | 3 +-- crates/gpu/src/elem.rs | 8 ++++++-- crates/model/src/arch/coverage.rs | 15 ++++++++++++++- tools/unsafe-ratchet.txt | 6 +++++- 4 files changed, 26 insertions(+), 6 deletions(-) diff --git a/crates/gpu/src/arch/qwen3moe/body35.rs b/crates/gpu/src/arch/qwen3moe/body35.rs index 36f559a94..b80ffebde 100644 --- a/crates/gpu/src/arch/qwen3moe/body35.rs +++ b/crates/gpu/src/arch/qwen3moe/body35.rs @@ -1621,8 +1621,7 @@ impl Body35 { LayerStore::Rec(RecStore::stamped(stream, shape.lens(lanes), lanes)?) } Kind35::Delta(shape) => { - let [state_f32, ring_f32] = shape.lens(lanes); - LayerStore::Rec(RecStore::new(stream, [state_f32, ring_f32], lanes)?) + LayerStore::Rec(RecStore::new(stream, shape.lens(lanes), lanes)?) } }); } diff --git a/crates/gpu/src/elem.rs b/crates/gpu/src/elem.rs index 1ff6ffe3e..f49cd0bb7 100644 --- a/crates/gpu/src/elem.rs +++ b/crates/gpu/src/elem.rs @@ -777,6 +777,9 @@ mod elem_kernels { // (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 { @@ -1120,9 +1123,10 @@ mod elem_kernels { v } else { let r = if v > 0.0 { v } else { 0.0 }; - let yv = r * r; - yv + r * r }; + // SAFETY: i < n <= y.len() by the launch contract; thread i owns + // entry i. unsafe { *y.get_unchecked_mut(i) = out; } diff --git a/crates/model/src/arch/coverage.rs b/crates/model/src/arch/coverage.rs index 9ba5cb2bb..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) ) }, }, @@ -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/tools/unsafe-ratchet.txt b/tools/unsafe-ratchet.txt index bfffd554c..fafca82d1 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 -> 1788 blocks, 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), relu_sqr (its input and output) and gqa_flash_seg_g5 +# (the head-128 group-5 flash entry over the shared flash core) read and write through raw pointers; each block with +# its SAFETY line. +gpu 1788 195 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. From 1f011169bccf02ec49b13024b131b1f7b66b8788 Mon Sep 17 00:00:00 2001 From: midagedev Date: Thu, 15 Oct 2026 11:47:47 +0900 Subject: [PATCH 05/13] gpu, gpu-gates: one head-128 launch, and the Q5_0 rows reuse the generator enqueue_pass takes the group width. The Q5_0 embed fixture is synthetic_rows at the scale's byte. LinearShape::lens stays in the crate. --- .../src/bin/gate_qwen3moe_experts.rs | 33 +--- .../gpu-gates/src/bin/gate_qwen3moe_flash.rs | 11 +- crates/gpu/src/arch/qwen3moe/scratch.rs | 6 +- crates/gpu/src/flash_gqa.rs | 152 +++++------------- crates/gpu/src/linear/mod.rs | 2 +- 5 files changed, 51 insertions(+), 153 deletions(-) diff --git a/crates/gpu-gates/src/bin/gate_qwen3moe_experts.rs b/crates/gpu-gates/src/bin/gate_qwen3moe_experts.rs index c630a19d2..a2e2c9bbb 100644 --- a/crates/gpu-gates/src/bin/gate_qwen3moe_experts.rs +++ b/crates/gpu-gates/src/bin/gate_qwen3moe_experts.rs @@ -76,7 +76,7 @@ mod gate { topk_ids_logical_within, verdict, }; use cuda_core::{CudaStream, DeviceBuffer}; - use gguf::quant::{GgmlType, dequant_row, f32_to_f16_bits}; + use gguf::quant::{GgmlType, dequant_row}; use model::arch::qwen3moe::names; /// f32 roundings each side's dot is allowed per output, on `Σ|w·x̂|` @@ -152,33 +152,6 @@ mod gate { Ok(()) } - /// `rows` of Q5_0 (`k` values, `k / 32` blocks) whose low and high - /// nibbles differ and whose codes differ from block to block, so a read - /// of the low half's word in place of the high half's, or of another - /// block's (another 32-block window's) words, disagrees with - /// `dequant_row`. - fn q5_0_bytes(rows: usize, k: usize) -> Vec { - let blocks = k / 32; - let mut bytes = vec![0u8; rows * blocks * 22]; - for r in 0..rows { - for b in 0..blocks { - let blk = &mut bytes[(r * blocks + b) * 22..][..22]; - let d = 0.5f32 + 0.25 * (b + r) as f32; - blk[..2].copy_from_slice(&f32_to_f16_bits(d).to_le_bytes()); - let qh: u32 = 0xA5A5_5A5A - ^ (r as u32).wrapping_mul(0x0101_0101) - ^ (b as u32).wrapping_mul(0x9E37_79B9); - blk[2..6].copy_from_slice(&qh.to_le_bytes()); - for j in 0..16 { - let lo = (j + b) % 16; - let hi = (15 - j + 3 * b) % 16; - blk[6 + j] = lo as u8 | (hi as u8) << 4; - } - } - } - bytes - } - /// 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). @@ -186,7 +159,9 @@ mod gate { const ROWS: usize = 4; // 34 blocks: past the first 32-block window of the packed code plane. const K: usize = 1088; - let bytes = q5_0_bytes(ROWS, K); + // 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)?; diff --git a/crates/gpu-gates/src/bin/gate_qwen3moe_flash.rs b/crates/gpu-gates/src/bin/gate_qwen3moe_flash.rs index a79463fdc..b3d5a9ce3 100644 --- a/crates/gpu-gates/src/bin/gate_qwen3moe_flash.rs +++ b/crates/gpu-gates/src/bin/gate_qwen3moe_flash.rs @@ -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, @@ -1928,7 +1928,7 @@ mod gate { 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_g5( + k.enqueue_pass::( stream, GqaArgs { q: &q_d, @@ -1944,6 +1944,7 @@ mod gate { fault: gpu.unlabelled_sink(), y: &mut y, }, + false, )?; stream.synchronize()?; let got = y.to_host_vec(stream)?; diff --git a/crates/gpu/src/arch/qwen3moe/scratch.rs b/crates/gpu/src/arch/qwen3moe/scratch.rs index f2bbbd3be..bcab80ed1 100644 --- a/crates/gpu/src/arch/qwen3moe/scratch.rs +++ b/crates/gpu/src/arch/qwen3moe/scratch.rs @@ -23,8 +23,8 @@ 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; @@ -605,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, diff --git a/crates/gpu/src/flash_gqa.rs b/crates/gpu/src/flash_gqa.rs index 71f6febb6..1d6f48a44 100644 --- a/crates/gpu/src/flash_gqa.rs +++ b/crates/gpu/src/flash_gqa.rs @@ -227,7 +227,7 @@ 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_g5`] takes. +/// 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 @@ -4539,51 +4539,33 @@ 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 head-128 attention; refused counts fault and yield NaNs. + 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( @@ -4591,26 +4573,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( @@ -4619,74 +4607,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, - )?; - Ok(()) - } - - /// [`FlashGqaKernels::enqueue_pass`] at [`GROUP_5`] query heads a key - /// head: `gqa_flash_seg_g5` (`m · n_kv · [`SEGMENTS`]` blocks of 256) and - /// [`gqa_flash_merge`] over `n_kv · GROUP_5` heads. Scalar only. Two - /// launches. Asynchronous, allocation-free, capturable. - pub fn enqueue_pass_g5(&self, stream: &CudaStream, args: GqaArgs<'_>) -> Result<(), GpuError> { - let what = "flash_gqa::enqueue_g5"; - 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 { - return Err(GpuError::shape( - what, - format!("need n_kv, ctx and m >= 1, got n_kv={n_kv} ctx={ctx} m={m}"), - )); - } - let n_head = n_kv * GROUP_5; - let segs = SEGMENTS; - 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), - ]; - if let Some((name, got, need)) = lens.iter().find(|(_, got, need)| got < need) { - return Err(GpuError::shape( - what, - 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 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 seg_keys = launch_u32(what, "seg_keys", SEG_KEYS)?; - let m = launch_u32(what, "m", m)?; - let cfg = LaunchConfig1D::new(grid, THREADS_U32, 0); - let prep = self.module.prepare_gqa_flash_seg_g5(cfg)?; - self.module.gqa_flash_seg_g5( - stream, &prep, q, kc, vc, n_keys, scale, n_kv, ctx, segs, seg_keys, m, part_v, part_ms, - )?; - let prep = self.module.prepare_gqa_flash_merge(LaunchConfig1D::new( - merge_grid, - MERGE_THREADS, - 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/linear/mod.rs b/crates/gpu/src/linear/mod.rs index 781102401..f828b12b5 100644 --- a/crates/gpu/src/linear/mod.rs +++ b/crates/gpu/src/linear/mod.rs @@ -184,7 +184,7 @@ impl LinearShape { /// [`crate::recur::RecStore::new`] allocates. The ring is not multiplied /// by `lanes`. #[must_use] - pub fn lens(&self, lanes: usize) -> [usize; 2] { + pub(crate) fn lens(&self, lanes: usize) -> [usize; 2] { [lanes * self.state_len(), self.ring_len()] } From aeb2c6a1c1f2e104063711d84dfc4273c323aa56 Mon Sep 17 00:00:00 2001 From: midagedev Date: Thu, 15 Oct 2026 13:01:14 +0900 Subject: [PATCH 06/13] gpu: one scalar segment core for both head-128 group widths gqa_flash_seg and gqa_flash_seg_g5 are thin entries over seg_scalar_g. The group of eight still uses every warp of the 256-thread block; a group of five leaves the pad warps to stage the key and value tile. --- crates/gpu/src/flash_gqa.rs | 506 +++++++++++++++--------------------- 1 file changed, 215 insertions(+), 291 deletions(-) diff --git a/crates/gpu/src/flash_gqa.rs b/crates/gpu/src/flash_gqa.rs index 1d6f48a44..cbe39f60d 100644 --- a/crates/gpu/src/flash_gqa.rs +++ b/crates/gpu/src/flash_gqa.rs @@ -2736,6 +2736,182 @@ 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`] +/// and [`flash_gqa_kernels::gqa_flash_seg_g5`]. `G` is the query heads of one +/// key head. The block stays 256 threads: a group of eight uses every warp, +/// and a group of five leaves warps `5..8` to stage the key and value tile. +/// +/// 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, +) { + const { assert!(G == GROUP || G == GROUP_5) }; + // Eight fills the block, so that instance's pad checks are `true`. + const FULL: bool = 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 && (FULL || 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 full group stages four values on every + // thread; a short group stages only the threads that own its rows. + let qb = (t * n_head + kh * G) * HEAD; + if FULL || 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 FULL || 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 FULL || w < G { + // SAFETY: idx is inside the partials, as on the neutral path; lane + // `lane` owns dims 4·lane .. +3, and lane 0 writes (m, s). + 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]; + if lane == 0 { + *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::*; @@ -2748,7 +2924,8 @@ 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 is + /// `seg_scalar_g` at [`GROUP`]. #[allow( clippy::too_many_arguments, reason = "kernel entry: the device ABI takes the arguments flat (rust-quality R8)" @@ -2787,158 +2964,36 @@ mod flash_gqa_kernels { static mut VS: SharedArray = SharedArray::UNINIT; static mut WS: SharedArray = SharedArray::UNINIT; - 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 * GROUP; - let h = kh * GROUP + 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 { - // SAFETY: idx < m·n_kv·GROUP·segs, both slots 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); - // SAFETY: each `static mut` above is this block's own shared - // allocation; the raw form reaches it without a reference. Every - // index below is inside its array, and every write is ordered before - // its reads by a block barrier (or, for WS, written and read by one - // warp between two barriers). - let (qs, ks, vs, ws) = unsafe { - ( + // allocation, sized for GROUP 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), ) }; - - // Row t's group query rows: thread `tid` stages values 4·tid .. +3 of - // the group's 1024. - let qb = (t * n_head + kh * GROUP) * HEAD; - let mut i = 0usize; - while i < 4 { - // SAFETY: qb + 4·tid + i < (t·n_head + (kh + 1)·GROUP)·HEAD <= - // m·n_kv·1024 <= q.len(); the shared index 4·tid + i < GROUP·HEAD. - 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: 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(); - - // 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; the - // query index w·HEAD + 2(wd + p) + 1 < GROUP·HEAD. - 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 and VS are this block's tiles, VS staged before the - // barrier above; w < GROUP and lane < 32. - unsafe { fold_tile(sc, live, &mut mx, &mut s_sum, &mut acc, ws, vs, w, lane) }; - t0 += KEY_TILE; - } - - // SAFETY: idx < m·n_kv·GROUP·segs; the four dims 4·lane .. +3 of the - // partial row are this lane's alone, and lane 0 writes (m, s). - 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]; - if lane == 0 { - *part_ms.get_unchecked_mut(2 * idx) = mx; - *part_ms.get_unchecked_mut(2 * idx + 1) = s_sum; - } - } } - /// [`gqa_flash_seg`] for [`GROUP_5`] query heads a key head. The block - /// stays 256 threads so the key and value tile stage as that kernel's; - /// warps `0..5` own the query heads and the pad warps only help stage. - /// Five does not pad onto the tensor-core tile (see [`GROUP_5`]), so - /// this entry is scalar. The partial index is `(t·n_head + h)·segs + seg` - /// with `h = kh·GROUP_5 + w`, and [`gqa_flash_merge`] folds it. + /// `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)" @@ -2977,160 +3032,29 @@ mod flash_gqa_kernels { static mut VS: SharedArray = SharedArray::UNINIT; static mut WS: SharedArray = SharedArray::UNINIT; - 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 * GROUP_5; - let h0 = kh * GROUP_5; - // The key head this block reads. For a live warp it is the query - // head's group; the pad warps share it so the staged tile is one. - let src = h0 / GROUP_5; - let h = h0 + 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 w < GROUP_5 && lane == 0 { - // SAFETY: w < GROUP_5 gives h < (kh+1)·GROUP_5 <= 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); - // SAFETY: each `static mut` above is this block's own shared - // allocation; the raw form reaches it without a reference. Every - // index below is inside its array, and every write is ordered before - // its reads by a block barrier (or, for WS, written and read by one - // warp between two barriers). - let (qs, ks, vs, ws) = unsafe { - ( + // 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 group's 5·128 query values, four per thread. Threads past the - // group's 160 stage nothing; the pad warps have no query row. - let qb = (t * n_head + h0) * HEAD; - if tid < GROUP_5 * HEAD / 4 { - let mut i = 0usize; - while i < 4 { - // SAFETY: tid < 160, so 4·tid + i < GROUP_5·HEAD, and - // qb + that < (t·n_head + (kh+1)·GROUP_5)·HEAD <= - // m·n_kv·5·128 <= q.len(). - unsafe { *qs.add(4 * tid + i) = *q.get_unchecked(qb + 4 * tid + i) }; - i += 1; - } - } - - let plane = src * 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, as `gqa_flash_seg` does: 32 - // keys × 32 u64 of K and of V. 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 `src`'s plane; a row is 128 u16 = 32 u64 and the - // planes start 8-byte aligned (a device allocation). - 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 < GROUP_5 { - // 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; - // w < GROUP_5, so the query index is inside 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 is GROUP_5·KEY_TILE and w < GROUP_5; 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 < GROUP_5 { - // SAFETY: idx is inside the partials, as on the neutral path; - // lane `lane` owns dims 4·lane .. +3, and lane 0 writes (m, s). - 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]; - if lane == 0 { - *part_ms.get_unchecked_mut(2 * idx) = mx; - *part_ms.get_unchecked_mut(2 * idx + 1) = s_sum; - } - } - } } /// The segment pass with `S = Q·Kᵀ` on the tensor cores: the grid, From 6bc6f647047e2d2eb3d837e2aaa66a12b023fc65 Mon Sep 17 00:00:00 2001 From: midagedev Date: Thu, 15 Oct 2026 13:09:59 +0900 Subject: [PATCH 07/13] gpu: fold the group-8 pad check as an inline const A nested const item cannot name seg_scalar_g's generic parameter, so the group-8 predicate is an inline const. The thin entries no longer mark the partial slices mut; the core still writes them. --- crates/gpu/src/flash_gqa.rs | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/crates/gpu/src/flash_gqa.rs b/crates/gpu/src/flash_gqa.rs index cbe39f60d..20e8408b9 100644 --- a/crates/gpu/src/flash_gqa.rs +++ b/crates/gpu/src/flash_gqa.rs @@ -2771,7 +2771,8 @@ unsafe fn seg_scalar_g( ) { const { assert!(G == GROUP || G == GROUP_5) }; // Eight fills the block, so that instance's pad checks are `true`. - const FULL: bool = G == GROUP; + // An inline const: a nested const item cannot name the generic parameter. + let full = const { G == GROUP }; let b = thread::blockIdx_x() as usize; let nkv = n_kv as usize; @@ -2796,7 +2797,7 @@ unsafe fn seg_scalar_g( let span = seg_span(limit, n_seg, seg_keys as usize); let lo = seg * span; if lo >= limit { - if lane == 0 && (FULL || w < G) { + if lane == 0 && (full || 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 { @@ -2811,7 +2812,7 @@ unsafe fn seg_scalar_g( // Row t's group query rows. A full group stages four values on every // thread; a short group stages only the threads that own its rows. let qb = (t * n_head + kh * G) * HEAD; - if FULL || tid < G * HEAD / 4 { + if full || tid < G * HEAD / 4 { let mut i = 0usize; while i < 4 { // SAFETY: the thread owns four values inside the group's @@ -2861,7 +2862,7 @@ unsafe fn seg_scalar_g( } thread::sync_threads(); - if FULL || w < G { + if full || w < G { // The score of key t0 + lane for head h. let live = t0 + lane < hi; let mut a = [0.0f32; ILP]; @@ -2895,7 +2896,7 @@ unsafe fn seg_scalar_g( t0 += KEY_TILE; } - if FULL || w < G { + if full || w < G { // SAFETY: idx is inside the partials, as on the neutral path; lane // `lane` owns dims 4·lane .. +3, and lane 0 writes (m, s). unsafe { @@ -2956,8 +2957,8 @@ mod flash_gqa_kernels { segs: u32, seg_keys: u32, m: u32, - mut part_v: DisjointSlice, - mut part_ms: DisjointSlice, + part_v: DisjointSlice, + part_ms: DisjointSlice, ) { static mut QS: SharedArray = SharedArray::UNINIT; static mut KS: SharedArray = SharedArray::UNINIT; @@ -3024,8 +3025,8 @@ mod flash_gqa_kernels { segs: u32, seg_keys: u32, m: u32, - mut part_v: DisjointSlice, - mut part_ms: DisjointSlice, + part_v: DisjointSlice, + part_ms: DisjointSlice, ) { static mut QS: SharedArray = SharedArray::UNINIT; static mut KS: SharedArray = SharedArray::UNINIT; From 6e1accd0425ce2f7a49c83e77b76270f039a0593 Mon Sep 17 00:00:00 2001 From: midagedev Date: Thu, 15 Oct 2026 13:58:19 +0900 Subject: [PATCH 08/13] gpu: bound the inlined segment core to its 256 threads The shared body is inlined into the entry, so it does not see the launch bound. The thread index is limited to the block, and a full group's (m, s) store keeps the lane test on its own. --- crates/gpu/src/flash_gqa.rs | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/crates/gpu/src/flash_gqa.rs b/crates/gpu/src/flash_gqa.rs index 20e8408b9..0c7fe8fa3 100644 --- a/crates/gpu/src/flash_gqa.rs +++ b/crates/gpu/src/flash_gqa.rs @@ -2786,6 +2786,9 @@ unsafe fn seg_scalar_g( let seg = sk / nkv; let kh = sk - seg * nkv; let tid = thread::threadIdx_x() as usize; + // The launch is 256 threads. The inlined body does not see the entry's + // launch bound, so the address math is told the same limit here. + unsafe { core::hint::assert_unchecked(tid < THREADS) }; let w = tid / 32; let lane = warp::lane_id() as usize; let n_head = nkv * G; @@ -2898,17 +2901,21 @@ unsafe fn seg_scalar_g( if full || w < G { // SAFETY: idx is inside the partials, as on the neutral path; lane - // `lane` owns dims 4·lane .. +3, and lane 0 writes (m, s). + // `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]; - if lane == 0 { - *part_ms.get_unchecked_mut(2 * idx) = mx; - *part_ms.get_unchecked_mut(2 * idx + 1) = s_sum; - } + } + } + // Lane 0 writes (m, s). A full group keeps this the kernel's own test. + if lane == 0 && (full || 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; } } } From 0eb457a024e22647a64325917cd1287fba5ac3de Mon Sep 17 00:00:00 2001 From: midagedev Date: Thu, 15 Oct 2026 14:11:51 +0900 Subject: [PATCH 09/13] gpu: keep the group-8 segment body in its entry A group of eight fills the 256-thread block, so the launch bound has to see the thread index in the entry. The short group still calls the shared core. --- crates/gpu/src/flash_gqa.rs | 173 ++++++++++++++++++++++++++++++------ 1 file changed, 146 insertions(+), 27 deletions(-) diff --git a/crates/gpu/src/flash_gqa.rs b/crates/gpu/src/flash_gqa.rs index 0c7fe8fa3..93d9a6e44 100644 --- a/crates/gpu/src/flash_gqa.rs +++ b/crates/gpu/src/flash_gqa.rs @@ -2736,10 +2736,10 @@ 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`] -/// and [`flash_gqa_kernels::gqa_flash_seg_g5`]. `G` is the query heads of one -/// key head. The block stays 256 threads: a group of eight uses every warp, -/// and a group of five leaves warps `5..8` to stage the key and value tile. +/// 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 @@ -2786,9 +2786,6 @@ unsafe fn seg_scalar_g( let seg = sk / nkv; let kh = sk - seg * nkv; let tid = thread::threadIdx_x() as usize; - // The launch is 256 threads. The inlined body does not see the entry's - // launch bound, so the address math is told the same limit here. - unsafe { core::hint::assert_unchecked(tid < THREADS) }; let w = tid / 32; let lane = warp::lane_id() as usize; let n_head = nkv * G; @@ -2932,8 +2929,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. The body is - /// `seg_scalar_g` at [`GROUP`]. + /// 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)" @@ -2964,37 +2962,158 @@ mod flash_gqa_kernels { segs: u32, seg_keys: u32, m: u32, - part_v: DisjointSlice, - part_ms: DisjointSlice, + mut part_v: DisjointSlice, + mut 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; + 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 * GROUP; + let h = kh * GROUP + 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 { + // SAFETY: idx < m·n_kv·GROUP·segs, both slots 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); + // SAFETY: each `static mut` above is this block's own shared - // allocation, sized for GROUP 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, + // allocation; the raw form reaches it without a reference. Every + // index below is inside its array, and every write is ordered before + // its reads by a block barrier (or, for WS, written and read by one + // warp between two barriers). + let (qs, ks, vs, ws) = unsafe { + ( 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), ) }; + + // Row t's group query rows: thread `tid` stages values 4·tid .. +3 of + // the group's 1024. + let qb = (t * n_head + kh * GROUP) * HEAD; + let mut i = 0usize; + while i < 4 { + // SAFETY: qb + 4·tid + i < (t·n_head + (kh + 1)·GROUP)·HEAD <= + // m·n_kv·1024 <= q.len(); the shared index 4·tid + i < GROUP·HEAD. + 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: 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(); + + // 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; the + // query index w·HEAD + 2(wd + p) + 1 < GROUP·HEAD. + 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 and VS are this block's tiles, VS staged before the + // barrier above; w < GROUP and lane < 32. + unsafe { fold_tile(sc, live, &mut mx, &mut s_sum, &mut acc, ws, vs, w, lane) }; + t0 += KEY_TILE; + } + + // SAFETY: idx < m·n_kv·GROUP·segs; the four dims 4·lane .. +3 of the + // partial row are this lane's alone, and lane 0 writes (m, s). + 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]; + if lane == 0 { + *part_ms.get_unchecked_mut(2 * idx) = mx; + *part_ms.get_unchecked_mut(2 * idx + 1) = s_sum; + } + } } /// `seg_scalar_g` at [`GROUP_5`]: the same 256-thread segment block as From 19224bce8c918a0ecbdc5688e9c4a667257acb4a Mon Sep 17 00:00:00 2001 From: midagedev Date: Thu, 15 Oct 2026 14:26:01 +0900 Subject: [PATCH 10/13] gpu: drop the unused group-8 path from the short segment core seg_scalar_g is only the five-head entry. The embedding comment names the types the load admits, and Q5_0 rows as the branch for a body that admits them. --- crates/gpu/src/arch/qwen3moe/dispatch.rs | 9 +++++---- crates/gpu/src/flash_gqa.rs | 22 ++++++++++------------ 2 files changed, 15 insertions(+), 16 deletions(-) diff --git a/crates/gpu/src/arch/qwen3moe/dispatch.rs b/crates/gpu/src/arch/qwen3moe/dispatch.rs index 0f3c8c43a..e9902b090 100644 --- a/crates/gpu/src/arch/qwen3moe/dispatch.rs +++ b/crates/gpu/src/arch/qwen3moe/dispatch.rs @@ -472,8 +472,8 @@ pub(super) fn mixer( /// A unit's front, one launch: the embedding rows of `io`'s ids into `s.x` /// 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, Q5_0 packed rows, or -/// Q8_0 planes (the load admits only these). +/// picks the lookup: Q3_K, Q4_K, Q5_K or Q6_K rows, or Q8_0 planes (the load +/// admits these). Q5_0 packed rows are served here for a body that admits them. pub(super) fn embed_rows( gpu: &Gpu, w: &Weights, @@ -487,7 +487,8 @@ 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. Q5_0 -/// packed rows take [`crate::elem::ElemKernels::enqueue_embed_rows_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, @@ -538,7 +539,7 @@ pub(super) fn embed_into( return Err(GpuError::tensor( WHAT, &name, - "K-quant rows, Q5_0 packed rows or Q8_0 planes", + "K-quant rows or Q8_0 planes; Q5_0 packed rows when the body admits them", )); } }; diff --git a/crates/gpu/src/flash_gqa.rs b/crates/gpu/src/flash_gqa.rs index 93d9a6e44..0af2827e0 100644 --- a/crates/gpu/src/flash_gqa.rs +++ b/crates/gpu/src/flash_gqa.rs @@ -2769,10 +2769,8 @@ unsafe fn seg_scalar_g( vs: *mut u64, ws: *mut f32, ) { - const { assert!(G == GROUP || G == GROUP_5) }; - // Eight fills the block, so that instance's pad checks are `true`. - // An inline const: a nested const item cannot name the generic parameter. - let full = const { G == GROUP }; + // 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; @@ -2797,7 +2795,7 @@ unsafe fn seg_scalar_g( let span = seg_span(limit, n_seg, seg_keys as usize); let lo = seg * span; if lo >= limit { - if lane == 0 && (full || w < G) { + 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 { @@ -2809,10 +2807,10 @@ unsafe fn seg_scalar_g( } let hi = (lo + span).min(limit); - // Row t's group query rows. A full group stages four values on every - // thread; a short group stages only the threads that own its rows. + // 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 full || tid < G * HEAD / 4 { + if tid < G * HEAD / 4 { let mut i = 0usize; while i < 4 { // SAFETY: the thread owns four values inside the group's @@ -2862,7 +2860,7 @@ unsafe fn seg_scalar_g( } thread::sync_threads(); - if full || w < G { + if w < G { // The score of key t0 + lane for head h. let live = t0 + lane < hi; let mut a = [0.0f32; ILP]; @@ -2896,7 +2894,7 @@ unsafe fn seg_scalar_g( t0 += KEY_TILE; } - if full || w < G { + if w < G { // SAFETY: idx is inside the partials, as on the neutral path; lane // `lane` owns dims 4·lane .. +3. unsafe { @@ -2907,8 +2905,8 @@ unsafe fn seg_scalar_g( *part_v.get_unchecked_mut(o + 3) = acc[3]; } } - // Lane 0 writes (m, s). A full group keeps this the kernel's own test. - if lane == 0 && (full || w < G) { + // 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; From b879ca9d86828da9fa40c77ea4a5161039627ff4 Mon Sep 17 00:00:00 2001 From: midagedev Date: Thu, 15 Oct 2026 15:16:35 +0900 Subject: [PATCH 11/13] gpu: the head-128 enqueue doc names its pack GqaArgs sizes the query and output rows by PACK, and enqueue_pass keeps the two-launch contract: a bad count faults, and the merge grid is PACK. --- crates/gpu/src/flash_gqa.rs | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/crates/gpu/src/flash_gqa.rs b/crates/gpu/src/flash_gqa.rs index 0af2827e0..ee6cdd92e 100644 --- a/crates/gpu/src/flash_gqa.rs +++ b/crates/gpu/src/flash_gqa.rs @@ -4433,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` @@ -4588,7 +4588,14 @@ impl FlashGqaKernels { Ok(FlashGqaKernels { module }) } - /// Enqueue head-128 attention; refused counts fault and yield NaNs. + /// 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, From 9470897ca2f9b987c674ef560aacfc86492b9401 Mon Sep 17 00:00:00 2001 From: midagedev Date: Thu, 15 Oct 2026 15:33:41 +0900 Subject: [PATCH 12/13] gpu: drop relu_sqr until a body calls it The squared ReLU kernel, its gate and its spill rows were an add with no caller. They return in the commit that first calls them. --- .../src/bin/gate_qwen3moe_experts.rs | 51 +----------------- crates/gpu/src/elem.rs | 54 ------------------- tools/ref/ptx-shapes.tsv | 2 - tools/unsafe-ratchet.txt | 6 +-- 4 files changed, 4 insertions(+), 109 deletions(-) diff --git a/crates/gpu-gates/src/bin/gate_qwen3moe_experts.rs b/crates/gpu-gates/src/bin/gate_qwen3moe_experts.rs index a2e2c9bbb..689ce2ff5 100644 --- a/crates/gpu-gates/src/bin/gate_qwen3moe_experts.rs +++ b/crates/gpu-gates/src/bin/gate_qwen3moe_experts.rs @@ -55,8 +55,7 @@ //! eager launch. //! //! Q5_0 embedding rows (`embed_rows_q5_0`) are bit-identical to `dequant_row` -//! on the same bytes. `relu_sqr` is bit-identical to `max(x, 0)^2`, and a -//! non-finite input is copied through. +//! on the same bytes. #[cfg(feature = "gpu")] fn main() -> std::process::ExitCode { @@ -141,7 +140,6 @@ mod gate { ok &= embed(&gpu, &gguf, &sets)?; ok &= embed_kquant(&gpu)?; ok &= embed_q5_0(&gpu)?; - ok &= relu_sqr(&gpu)?; ok &= gate_up_body(&gpu, &gguf, &ex)?; ok &= gate_up_tap(&gpu, &gguf, &ex, &sets)?; ok &= combine(&gpu, &ex, &sets, stream)?; @@ -209,53 +207,6 @@ mod gate { Ok(host_ok && pos_ok && graph_ok && oor_ok) } - /// Host transcription of `relu_sqr`: a non-finite input stays that value, - /// and a finite one is `max(x, 0)^2`. - fn host_relu_sqr(v: f32) -> f32 { - if !v.is_finite() { - v - } else { - let r = if v > 0.0 { v } else { 0.0 }; - r * r - } - } - - /// `relu_sqr` against [`host_relu_sqr`], bit for bit, including a - /// non-finite input, and one captured node. - fn relu_sqr(gpu: &Gpu) -> Result { - let x = [ - 0.0f32, - 1.5, - -2.0, - f32::NAN, - f32::INFINITY, - f32::NEG_INFINITY, - 3.0, - ]; - let want: Vec = x.iter().copied().map(host_relu_sqr).collect(); - let stream = gpu.stream(); - let xb = DeviceBuffer::from_host(stream, &x)?; - let mut y = DeviceBuffer::::zeroed(stream, x.len())?; - gpu.elem().enqueue_relu_sqr(stream, &xb, x.len(), &mut y)?; - stream.synchronize()?; - let got = y.to_host_vec(stream)?; - let host_ok = bits_equal(&got, &want); - // The eager read ends before this closure borrows `y` again. - let mut enqueue = |s: &CudaStream| gpu.elem().enqueue_relu_sqr(s, &xb, x.len(), &mut y); - let g = gpu.capture(&mut enqueue)?; - g.launch(stream)?; - stream.synchronize()?; - let replay = y.to_host_vec(stream)?; - let nodes = g.node_count(); - let graph_ok = bits_equal(&replay, &want) && nodes == 1; - println!( - "relu_sqr n={}: vs host bit_identical={host_ok} graph_nodes={nodes} graph_same={graph_ok} {}", - x.len(), - verdict(host_ok && graph_ok) - ); - Ok(host_ok && graph_ok) - } - // ------------------------------------------------------------ 1. embed /// One gather's outputs, read back: the rows, each row's position and diff --git a/crates/gpu/src/elem.rs b/crates/gpu/src/elem.rs index f49cd0bb7..b5e7fcddb 100644 --- a/crates/gpu/src/elem.rs +++ b/crates/gpu/src/elem.rs @@ -1101,37 +1101,6 @@ mod elem_kernels { e.map(silu_mul); } - /// `y[i] = max(x[i], 0)^2` over `n` values. A non-finite input is copied - /// as it is: `max` of a non-finite is not this rule, and the square of - /// `−inf` is not 0. - #[kernel] - #[launch_bounds(256)] - #[launch_contract( - domain = 1, - block = (256, 1, 1), - requires = (x.len() >= n, y.len() >= n) - )] - pub fn relu_sqr(x: &[f32], n: u32, mut y: DisjointSlice) { - let i = thread::index_1d().get(); - if i >= n as usize { - return; - } - // SAFETY: i < n <= x.len() and y.len() by the launch contract; thread - // i owns entry i. - let v = unsafe { *x.get_unchecked(i) }; - let out = if !v.is_finite() { - v - } else { - let r = if v > 0.0 { v } else { 0.0 }; - r * r - }; - // SAFETY: i < n <= y.len() by the launch contract; thread i owns - // entry i. - unsafe { - *y.get_unchecked_mut(i) = out; - } - } - /// `y[i] = a[i] + b[i]`, elementwise over `n` values — exact; the body /// is one add, so nothing is extracted into a core. #[kernel] @@ -2177,29 +2146,6 @@ impl ElemKernels { Ok(()) } - /// Enqueue `y = max(x, 0)^2` over `n` values (flat). A non-finite input - /// stays that value. Asynchronous, allocation-free, capturable. - pub fn enqueue_relu_sqr( - &self, - stream: &CudaStream, - x: &DeviceBuffer, - n: usize, - y: &mut DeviceBuffer, - ) -> Result<(), GpuError> { - if n == 0 || x.len() < n || y.len() < n { - return Err(GpuError::shape( - "enqueue_relu_sqr", - format!("n={n}, x.len() {}, y.len() {}", x.len(), y.len()), - )); - } - let n = launch_u32("enqueue_relu_sqr", "n", n)?; - let prep = self - .module - .prepare_relu_sqr(LaunchConfig1D::new(n.div_ceil(256), 256, 0))?; - self.module.relu_sqr(stream, &prep, x, n, y)?; - Ok(()) - } - /// Enqueue `y = a + b` over `n` values (flat). Exact. Asynchronous, /// allocation-free, capturable. pub fn enqueue_add( diff --git a/tools/ref/ptx-shapes.tsv b/tools/ref/ptx-shapes.tsv index 12783be4e..e386a1e16 100644 --- a/tools/ref/ptx-shapes.tsv +++ b/tools/ref/ptx-shapes.tsv @@ -309,7 +309,6 @@ generate_ds41 qwen3moe_router_norm_m_prefetch 0 0 generate_ds41 qwen3moe_router_norm_prefetch 0 0 generate_ds41 qwen3moe_router_route 0 0 generate_ds41 qwen3moe_token_major 0 0 -generate_ds41 relu_sqr 0 0 generate_ds41 rms_norm 0 0 generate_ds41 rope 0 0 generate_ds41 router_topk 0 0 @@ -552,7 +551,6 @@ gate_e2e qwen3moe_router_norm_m_prefetch 0 0 gate_e2e qwen3moe_router_norm_prefetch 0 0 gate_e2e qwen3moe_router_route 0 0 gate_e2e qwen3moe_token_major 0 0 -gate_e2e relu_sqr 0 0 gate_e2e rms_norm 0 0 gate_e2e rope 0 0 gate_e2e router_topk 0 0 diff --git a/tools/unsafe-ratchet.txt b/tools/unsafe-ratchet.txt index fafca82d1..b44c29e24 100644 --- a/tools/unsafe-ratchet.txt +++ b/tools/unsafe-ratchet.txt @@ -225,11 +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. -# PIN(2026-10-15): gpu 1771 -> 1788 blocks, 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), relu_sqr (its input and output) and gqa_flash_seg_g5 +# PIN(2026-10-15): gpu 1771 -> 1786 blocks, 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 over the shared flash core) read and write through raw pointers; each block with # its SAFETY line. -gpu 1788 195 10 +gpu 1786 195 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. From 21e0ea141a7a2449588d8134620887461ba42dc2 Mon Sep 17 00:00:00 2001 From: midagedev Date: Thu, 15 Oct 2026 15:57:53 +0900 Subject: [PATCH 13/13] tools: the unsafe pin counts the shared segment core seg_scalar_g is an unsafe fn, and the group-8 body kept in its entry leaves gpu at 1787 blocks and 196 fns on main's 1771. The relu drop's pin line was short of that count. --- tools/unsafe-ratchet.txt | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/tools/unsafe-ratchet.txt b/tools/unsafe-ratchet.txt index b44c29e24..621a6660e 100644 --- a/tools/unsafe-ratchet.txt +++ b/tools/unsafe-ratchet.txt @@ -225,11 +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. -# PIN(2026-10-15): gpu 1771 -> 1786 blocks, 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 over the shared flash core) read and write through raw pointers; each block with -# its SAFETY line. -gpu 1786 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.