From a4d55170d03d00e65e2946c2741cf6ea2ca71edc Mon Sep 17 00:00:00 2001 From: midagedev Date: Thu, 22 Oct 2026 01:14:58 +0900 Subject: [PATCH 01/10] kernel table: a Q5_0 stack's packed words route through the 32-value GEMM (gemm_q5_0p), so a Q5_0 down has a prompt route Co-Authored-By: Claude Opus 5.5 --- crates/gguf/src/card.rs | 8 ++++++-- crates/gguf/src/qtype.rs | 2 +- crates/gpu/src/card_leg.rs | 22 ++++++++++++++++------ crates/gpu/src/gemm/grouped.rs | 6 +++++- crates/placement/src/kernels.rs | 4 ++++ 5 files changed, 32 insertions(+), 10 deletions(-) diff --git a/crates/gguf/src/card.rs b/crates/gguf/src/card.rs index 81560f119..27bee9de9 100644 --- a/crates/gguf/src/card.rs +++ b/crates/gguf/src/card.rs @@ -30,8 +30,9 @@ //! //! Weights. The routed `_sel`, gate·up and file-layout gemm entries read a //! stack as the file's blocks in the [`CardFormat::KQuant`] word stream; -//! `q5_0_gemv_sel` and the Q5_0 dense site ([`SiteEntry::Q5_0`]) read the -//! packed [`CardFormat::Q5_0`] layout, the Q8_0 dense, head and embedding +//! `q5_0_gemv_sel`, the Q5_0 gemm entry and the Q5_0 dense site +//! ([`SiteEntry::Q5_0`]) read the packed [`CardFormat::Q5_0`] layout, the +//! Q8_0 dense, head and embedding //! entries read [`CardFormat::Q8_0Planes`], and the NVFP4 entries read //! [`CardFormat::Nvfp4Planes`] (a routed stack's with one tensor scale an //! expert). A routed entry's layout is its `layout`, which @@ -196,6 +197,9 @@ pub enum GemmEntry { /// the q8f32 planes) for a dense projection (`SiteEntry::Q8_0`); both in /// `crates/gpu/src/gemm/kernels32.rs`. Q8_0, + /// `Gemm32Weight::Q5_0Packed` (`gemm_q5_0p`, the `pack_q5_0` layout a + /// stack's `q5_0_gemv_sel` reads) in `crates/gpu/src/gemm/kernels32.rs`. + Q5_0, /// `Gemm32Weight::Q5_1File` (`gemm_q5_1`) in /// `crates/gpu/src/gemm/kernels32.rs`. Q5_1, diff --git a/crates/gguf/src/qtype.rs b/crates/gguf/src/qtype.rs index 359c3e585..1ec0a7d2d 100644 --- a/crates/gguf/src/qtype.rs +++ b/crates/gguf/src/qtype.rs @@ -78,7 +78,7 @@ pub const QROWS: &[QRow] = &[ // q5_0: ggml.c:756, ggml-common.h:196-216; q5_1: ggml.c:777; q8_0: ggml.c:819, ggml-common.h:233-238 // (one f16 d and 32 int8 codes) QRow { ty: T::Q5_0, name: "q5_0", blck: 32, bytes: 22, dequant: true, host_act: H::Q8_2X4, - card: card!(Q5_0 Q5_0 - (Q5_0, Q8Blocks32) - Q5_0 - -) }, + card: card!(Q5_0 Q5_0 - (Q5_0, Q8Blocks32) Q5_0 Q5_0 - -) }, QRow { ty: T::Q5_1, name: "q5_1", blck: 32, bytes: 24, dequant: true, host_act: H::Q8_2X4, card: card!(Q5_1 Q5_1 - (Q5_1, Q8Blocks32) Q5_1 - - -) }, QRow { ty: T::Q8_0, name: "q8_0", blck: 32, bytes: 34, dequant: true, host_act: H::Q8_2X4, diff --git a/crates/gpu/src/card_leg.rs b/crates/gpu/src/card_leg.rs index 149524fbc..71ed6081c 100644 --- a/crates/gpu/src/card_leg.rs +++ b/crates/gpu/src/card_leg.rs @@ -254,11 +254,13 @@ impl<'a> StackW<'a> { /// The stack, of file type `ty`, as the route GEMM the kernel table names /// for that type reads it ([`kernels::gemm`]): a K-quant's or an /// i-quant's words through the grouped GEMM, the file's Q8_0, Q5_1 or - /// IQ4_NL blocks and NVFP4's planes through the 32-value GEMM. `None` for - /// a type the table names no GEMM for, or a form that GEMM does not read. + /// IQ4_NL blocks, a Q5_0 stack's packed words and NVFP4's planes through + /// the 32-value GEMM. `None` for a type the table names no GEMM for, or a + /// form that GEMM does not read. pub(crate) fn route(self, ty: GgmlType) -> Option> { Some(match (kernels::gemm(ty)?, self) { (GemmEntry::Q8_0, StackW::Words(w)) => RouteW::Blocks32(Gemm32Weight::Q8_0File(w)), + (GemmEntry::Q5_0, StackW::Q5_0(w)) => RouteW::Blocks32(Gemm32Weight::Q5_0Packed(w)), (GemmEntry::Q5_1, StackW::Words(w)) => RouteW::Blocks32(Gemm32Weight::Q5_1File(w)), (GemmEntry::Iq4nl, StackW::Words(w)) => RouteW::Blocks32(Gemm32Weight::Iq4NlFile(w)), (GemmEntry::Nvfp4, StackW::Nvfp4(s)) => RouteW::Blocks32(Gemm32Weight::Nvfp4Planes(s)), @@ -2485,9 +2487,9 @@ pub(crate) mod tests { /// A stack routes by the GEMM entry the kernel table names for its type /// (`kernels::gemm`): a K-quant's or an i-quant's words through the /// grouped GEMM in the type that GEMM decodes (`GemmWeight::from_ggml`), - /// the file's Q8_0, Q5_1 and IQ4_NL blocks and NVFP4's planes through the - /// 32-value GEMM in their own weights; a stack in the other form, and a - /// type the table names no GEMM for, has no route. + /// the file's Q8_0, Q5_1 and IQ4_NL blocks, a Q5_0 stack's packed words and + /// NVFP4's planes through the 32-value GEMM in their own weights; a stack + /// in the other form, and a type the table names no GEMM for, has no route. #[test] #[ignore = "needs a CUDA device; `just gate-gpu-lib` runs it on the box"] fn hw_a_stack_routes_by_the_tables_gemm_entry() { @@ -2534,6 +2536,14 @@ pub(crate) mod tests { blocks32(words, GgmlType::Q5_1), Some(Gemm32Weight::Q5_1File(_)) )); + // PIN(2026-10-22): a Q5_0 stack's packed words join the 32-value + // routes with the table's GEMM row for the type (`gemm_q5_0p`); the + // file's blocks as words have no route. + assert!(matches!( + blocks32(StackW::Q5_0(&qs), GgmlType::Q5_0), + Some(Gemm32Weight::Q5_0Packed(_)) + )); + assert!(words.route(GgmlType::Q5_0).is_none(), "Q5_0 as words"); assert!(matches!( blocks32(words, GgmlType::IQ4_NL), Some(Gemm32Weight::Iq4NlFile(_)) @@ -2546,7 +2556,7 @@ pub(crate) mod tests { assert!(planes.route(ty).is_none(), "{ty} in planes"); } assert!(words.route(GgmlType::NVFP4).is_none(), "NVFP4 as words"); - for ty in [GgmlType::MXFP4, GgmlType::Q5_0, GgmlType::F32] { + for ty in [GgmlType::MXFP4, GgmlType::F32] { assert!(kernels::gemm(ty).is_none(), "{ty} has a GEMM row"); assert!(words.route(ty).is_none(), "{ty}"); } diff --git a/crates/gpu/src/gemm/grouped.rs b/crates/gpu/src/gemm/grouped.rs index 5fcdfc9d0..390aef80a 100644 --- a/crates/gpu/src/gemm/grouped.rs +++ b/crates/gpu/src/gemm/grouped.rs @@ -203,7 +203,11 @@ impl GemmWeight { GemmEntry::Iq2xxs => Some(GemmWeight::Iq2Xxs), GemmEntry::Iq3xxs => Some(GemmWeight::Iq3Xxs), GemmEntry::Iq4xs => Some(GemmWeight::Iq4Xs), - GemmEntry::Q8_0 | GemmEntry::Q5_1 | GemmEntry::Iq4nl | GemmEntry::Nvfp4 => None, + GemmEntry::Q8_0 + | GemmEntry::Q5_0 + | GemmEntry::Q5_1 + | GemmEntry::Iq4nl + | GemmEntry::Nvfp4 => None, } } diff --git a/crates/placement/src/kernels.rs b/crates/placement/src/kernels.rs index 14c3e6432..49c98646c 100644 --- a/crates/placement/src/kernels.rs +++ b/crates/placement/src/kernels.rs @@ -504,11 +504,15 @@ mod tests { // planes); every other pair is unchanged. // PIN(2026-10-21): and IQ2_XXS (`gemm_iq2xxs`, which `GemmWeight` already // decoded and the latent MoE body launches); every other pair is unchanged. + // PIN(2026-10-22): and Q5_0 (`gemm_q5_0p` over a stack's `pack_q5_0` words, + // which Raon-1's Q5_0 down stacks' prompt route reads); every other pair + // is unchanged. #[test] fn gemm_admits_exactly_the_pinned_pairs() { assert_eq!( rows(gemm), [ + (GgmlType::Q5_0, GemmEntry::Q5_0), (GgmlType::Q5_1, GemmEntry::Q5_1), (GgmlType::Q8_0, GemmEntry::Q8_0), (GgmlType::Q3_K, GemmEntry::Q3k), From 87d538f2e8b78c5c7600b1170a40ef08b4a0bae6 Mon Sep 17 00:00:00 2001 From: midagedev Date: Thu, 22 Oct 2026 01:14:58 +0900 Subject: [PATCH 02/10] raon1: the prompt runs in wide units of up to 256 rows: the grouped GEMM for every projection, the common KDA mixer's chunked path for the RDA layers, the route tables for the experts (#616) Co-Authored-By: Claude Opus 5.5 --- crates/gpu/src/arch/raon1/body.rs | 121 +++- crates/gpu/src/arch/raon1/mla.rs | 203 +++++- crates/gpu/src/arch/raon1/mod.rs | 5 +- crates/gpu/src/arch/raon1/rows.rs | 123 ++-- crates/gpu/src/arch/raon1/wide.rs | 992 ++++++++++++++++++++++++++++++ tools/ref/ptx-shapes.tsv | 1 + 6 files changed, 1389 insertions(+), 56 deletions(-) create mode 100644 crates/gpu/src/arch/raon1/wide.rs diff --git a/crates/gpu/src/arch/raon1/body.rs b/crates/gpu/src/arch/raon1/body.rs index a6ae3d1b8..d6c26e0aa 100644 --- a/crates/gpu/src/arch/raon1/body.rs +++ b/crates/gpu/src/arch/raon1/body.rs @@ -25,6 +25,7 @@ use super::mla::{self, MlaBufs, MlaDims, MlaKernels, MlaSites}; use super::rda::{self, RdaSites, RowKernels, RowsKit}; use super::rows::{RowState, RowViews}; use super::state::{self, PASS_ROWS, Pass}; +use super::wide::{self, WIDE_ROWS, Wide, WideDims}; use crate::arch::hybrid::{Acts, DecodeInput, GqaCx, Site}; use crate::elem::EmbedRowsArgs; use crate::flash_gqa::FlashGqaKernels; @@ -175,9 +176,36 @@ fn dims(hp: &Hparams) -> Result<(KdaDims, MlaDims), GpuError> { Ok((kda, mla)) } +/// The widths a wide unit is cut for, or a named refusal of a shape the +/// wide arm's forms do not cover: the two mixers' output projections read +/// one width, the routed block's slots a card sum of the router's pick count. +fn wide_dims(hp: &Hparams, kda: &KdaDims, mla: &MlaDims) -> Result { + let attn = kda.shape.n_v * HEAD; + if attn != mla.heads * mla.v_head { + return Err(GpuError::shape( + WHAT, + format!( + "the RDA layers' output projection reads {attn} values, the latent layers' {}: the wide arm holds one form of its input", + mla.heads * mla.v_head + ), + )); + } + let e = hp.experts; + Ok(WideDims { + embd: hp.n_embd, + attn, + q_lora: hp.mla.q_lora, + heads: mla.heads, + dense_ff: hp.dense_ff, + shared_ff: e.shared_ff(), + expert_ff: e.ff, + ff_rows: hp.dense_ff.max(ffn::sigmoid_128_6::N_USED * e.ff), + }) +} + /// A walk's buffers, `rows` rows of each width, token-major: the step's /// one row, or a pass's [`PASS_ROWS`]. -struct Scratch { +pub(super) struct Scratch { /// The streams the layer reads and the ones it writes, `STREAMS · n` a /// row. x: DeviceBuffer, @@ -228,7 +256,9 @@ impl Scratch { })?; let lin = d.shape; let (channels, v, n_v) = (lin.channels(), lin.n_v * HEAD, lin.n_v); - let (part_v, part_ms) = mla::partials_lens_rows(m, ctx, rows); + // The latent attention takes at most a pass's rows a launch: a wider + // unit runs it in chunks that share one set of partials. + let (part_v, part_ms) = mla::partials_lens_rows(m, ctx, rows.min(PASS_ROWS)); let (heads, row) = (m.heads, m.row()); Ok(( [ @@ -311,7 +341,7 @@ impl Scratch { }) } - fn bytes(&self) -> usize { + pub(super) fn bytes(&self) -> usize { [ &self.x, &self.x2, &self.hc, &self.part, &self.u, &self.un, &self.br, ] @@ -397,6 +427,8 @@ pub struct Body { pass: Scratch, ffn_rows: FfnRows, stage: DeviceBuffer, + /// A wide unit's rows, forms and kernels ([`super::wide`]). + pub(super) wide: Wide, /// The passes' staging, the lane word, the taps and the position the /// stores stand at. pub(super) rows: RowState, @@ -491,8 +523,9 @@ fn exact(w: &Weights, name: String, len: usize) -> Result { } impl Body { - /// Card bytes of the rows pass's walk buffers ([`PASS_ROWS`] rows, the - /// latent layers' attention partials among them), which every load of + /// Card bytes of the walk buffers of the rows pass ([`PASS_ROWS`] rows, the + /// latent layers' attention partials among them) and of a wide unit + /// ([`WIDE_ROWS`] rows, its forms and feed-forward rows), which every load of /// the file `hp` describes makes beside a plan's terms: what its plan /// keeps out of the expert budget (`PlanInputs::with_prompt_reserve`). /// Counted at a cache taller than one attention segment, the most any @@ -500,7 +533,9 @@ impl Body { pub fn pass_reserve(hp: &Hparams) -> Result { let (kda, mla) = dims(hp)?; let lens = Scratch::lens(hp, (&kda, &mla), usize::MAX, PASS_ROWS)?; - Ok(Scratch::bytes_of(&lens) as u64) + let wide_lens = Scratch::lens(hp, (&kda, &mla), usize::MAX, WIDE_ROWS)?; + let wide = Scratch::bytes_of(&wide_lens) + Wide::bytes_of(&wide_dims(hp, &kda, &mla)?); + Ok((Scratch::bytes_of(&lens) + wide) as u64) } /// The last hyper-connection site's route as the step left it: the last @@ -632,6 +667,11 @@ impl Body { pass: Scratch::new(stream, &hp, (&kda, &mla), ctx, PASS_ROWS)?, ffn_rows: FfnRows::new(stream, n, ff, PASS_ROWS)?, stage: DeviceBuffer::zeroed(stream, stage)?, + wide: Wide::new( + gpu, + &wide_dims(&hp, &kda, &mla)?, + Scratch::new(stream, &hp, (&kda, &mla), ctx, WIDE_ROWS)?, + )?, rows: RowState::new(stream, lanes)?, kda, mla, @@ -652,6 +692,7 @@ impl Body { + self.pass.bytes() + self.ffn_rows.bytes() + self.stage.num_bytes() + + self.wide.bytes() + self.rows.bytes() + self.acts.bytes() + self.k.kgemv.device_bytes() @@ -695,6 +736,7 @@ impl Body { ffn_rows, stage, rows, + wide, ctx, .. } = self; @@ -722,6 +764,7 @@ impl Body { } = rows; let pos_w: Window<'_, u32>; let keys_w: Window<'_, u32>; + let mut wide_kit = None; let (sc, mut ffn_b, pos, keys, (at, m, pass)) = match run { Run::Step => { gpu.elem().enqueue_embed_rows_kquant( @@ -740,6 +783,46 @@ impl Body { let (pos, keys) = (&step_in.pos, &step_in.n_keys); (s, FfnBufs::Step(ffn_s), pos, keys, (0, 1, Pass::InPlace)) } + Run::Rows { at, m, pass } if m > PASS_ROWS => { + if pass != Pass::InPlace { + return Err(GpuError::shape( + STEP, + format!("a {pass:?} pass of {m} rows: only a prompt's units run wide"), + )); + } + wide.route_dense(gpu, m)?; + let (kit, ws, wrin) = wide.kit(&k.g32, &k.ffn.router); + let RowViews { + ids, + pos, + keys, + epos, + ekeys, + } = wrin.views(at, m)?; + gpu.elem().enqueue_embed_rows_kquant( + stream, + *ty, + EmbedRowsArgs { + w: table, + ids: &ids, + pos0: &pos, + first: 0, + y: &mut ws.x, + pos: epos, + n_keys: ekeys, + }, + )?; + pos_w = pos; + keys_w = keys; + wide_kit = Some(kit); + ( + ws, + FfnBufs::Rows(ffn_rows), + &*pos_w, + &*keys_w, + (at, m, pass), + ) + } Run::Rows { at, m, pass } => { let RowViews { ids, @@ -848,6 +931,8 @@ impl Body { }; if m == 1 { rda::step(&cx, &k.linear, acts, sites, (l, &*kd), st, bufs)?; + } else if let Some(kit) = wide_kit.as_mut() { + wide::rda_rows(&cx, kit, &k.linear, sites, (l, &*kd, m), st, bufs)?; } else { let kit = RowsKit { linear: &k.linear, @@ -880,6 +965,8 @@ impl Body { }; if m == 1 { mla::step(&cx, &k.mla, acts, (&**sites, &*md), l, bufs)?; + } else if let Some(kit) = wide_kit.as_mut() { + mla::rows_wide(&cx, &k.mla, kit, (&**sites, &*md), (l, m), bufs)?; } else { let kit = RowsKit { linear: &k.linear, @@ -905,15 +992,19 @@ impl Body { .enqueue_rms_norm(stream, u, gain, hp.rms_eps, n, m, un)?; let rule = (*act, hp.experts.weights_scale); let at_l = (l, &layer.ffn); - match &mut ffn_b { - FfnBufs::Step(f) => ffn::ffn(&cx, &k.ffn, acts, rule, at_l, un, f, br)?, - FfnBufs::Rows(f) => { - let kit = RowsKit { - linear: &k.linear, - rows: &k.rows, - stage: &mut *stage, - }; - ffn::ffn_rows(&cx, &k.ffn, acts, rule, at_l, m, un, f, br, kit)? + if let Some(kit) = wide_kit.as_mut() { + wide::ffn(&cx, kit, rule, at_l, m, un, br)?; + } else { + match &mut ffn_b { + FfnBufs::Step(f) => ffn::ffn(&cx, &k.ffn, acts, rule, at_l, un, f, br)?, + FfnBufs::Rows(f) => { + let kit = RowsKit { + linear: &k.linear, + rows: &k.rows, + stage: &mut *stage, + }; + ffn::ffn_rows(&cx, &k.ffn, acts, rule, at_l, m, un, f, br, kit)? + } } } mhc.enqueue_post(stream, (x2, br, hc), (n, m), x)?; diff --git a/crates/gpu/src/arch/raon1/mla.rs b/crates/gpu/src/arch/raon1/mla.rs index 20745b373..14934610b 100644 --- a/crates/gpu/src/arch/raon1/mla.rs +++ b/crates/gpu/src/arch/raon1/mla.rs @@ -48,6 +48,9 @@ use super::rda::{RowsKit, kgemv_rows, project_rows}; use super::state::PASS_ROWS; +use super::wide::{ + Input, WIDE_ROWS, WideKit, input, quantize_hid, quantize_out, quantize_qb, value_heads, +}; use crate::arch::hybrid::{Acts, GqaCx, Site}; use crate::fault::FaultSink; use crate::flash::{FlashGeom, FlashInputs, FlashSplitArgs, partials_ms_len, partials_v_len}; @@ -56,7 +59,7 @@ use crate::linear::out_gate::{OutGateArgs, OutGateKernels}; use crate::model::{Q8_0GemvHeadsArgs, StepKernels}; use crate::q6k_heads::{Q6kHeadsArgs, Q6kHeadsKernels}; use crate::q8f32::Q8_0GemvHeadsMcolArgs; -use crate::site::{self, Form, SiteTy}; +use crate::site::{self, Form, SiteTy, WideActs}; use crate::tensor::{DeviceTensor, Window, WindowMut}; use crate::weights::{DevWeight, Weights}; use crate::{Gpu, GpuError}; @@ -220,8 +223,8 @@ impl MlaSites { } /// The kernels every MLA layer launches beyond the `Gpu`'s own, and the -/// extra part's gather tables for up to [`PASS_ROWS`] rows -/// (`[(t·heads + h)·head_k + nope + j] → [(t·heads + h)·row + j]`; a pass of +/// extra part's gather tables for up to [`WIDE_ROWS`] rows +/// (`[(t·heads + h)·head_k + nope + j] → [(t·heads + h)·row + j]`; a call of /// `m` rows reads the first `m·heads·extra` pairs). pub struct MlaKernels { step: StepKernels, @@ -239,9 +242,9 @@ impl MlaKernels { let at = |v: usize| { u32::try_from(v).map_err(|_| GpuError::shape(WHAT, "a gather index past u32")) }; - let pairs = PASS_ROWS * d.heads * d.extra; + let pairs = WIDE_ROWS * d.heads * d.extra; let (mut src, mut dst) = (Vec::with_capacity(pairs), Vec::with_capacity(pairs)); - for head in 0..PASS_ROWS * d.heads { + for head in 0..WIDE_ROWS * d.heads { for j in 0..d.extra { src.push(at(head * d.head_k() + d.nope + j)?); dst.push(at(head * d.row() + j)?); @@ -824,6 +827,196 @@ pub(crate) fn rows( project_rows(cx, &mut kit, acts, (l, m), out, b.gated, b.out) } +/// Enqueue MLA layer `l` over a wide unit of `m` rows (above [`PASS_ROWS`]), +/// in position order: the projections as GEMMs over the unit's rows +/// ([`super::wide`]); the cache append, the gather of the extra part and the +/// output gate once over the unit; the absorb and the latent flash in chunks +/// of [`PASS_ROWS`] query tokens on windows of the unit's buffers (the flash +/// takes no more), each chunk over the cache rows the unit's append wrote, +/// bounded by its own live counts; and a Q6_K per-head value projection as one +/// grouped GEMM over every (row, head) of the unit ([`value_heads`]), a Q8_0 +/// one per chunk. `b`'s buffers hold at least `m` rows of their widths, the +/// partials a pass's rows. +pub(super) fn rows_wide( + cx: &GqaCx<'_>, + k: &MlaKernels, + kit: &mut WideKit<'_>, + (sites, d): (&MlaSites, &MlaDims), + (l, m): (usize, usize), + b: MlaBufs<'_>, +) -> Result<(), GpuError> { + let (gpu, w) = (cx.gpu, cx.w); + let stream = gpu.stream(); + let fault = gpu.layer_sink(l)?; + let [norm, q_a_norm, kv_a_norm] = &sites.vecs; + gpu.elem().enqueue_rms_norm( + stream, + b.x, + w.f32_buf(WHAT, norm)?, + d.rms_eps, + d.embd, + m, + b.xn, + )?; + let (wk, wb) = (&kit.k, &mut kit.b); + let tys = [sites.q_a.ty(), sites.kv_a.ty(), sites.gate.ty()]; + quantize_hid( + gpu, + wk.g32, + wb.forms, + (&*b.xn, m), + (fault, WideActs::wants_of(&tys)), + )?; + { + let x_in = input(wb.forms, Input::Hid, &*b.xn); + wk.project(cx, (&sites.q_a, d.q_lora), &x_in, m, b.qa)?; + wk.project(cx, (&sites.kv_a, d.row()), &x_in, m, b.kv)?; + wk.project(cx, (&sites.gate, d.heads * d.v_head), &x_in, m, b.gate)?; + } + gpu.elem().enqueue_rms_norm( + stream, + b.qa, + w.f32_buf(WHAT, q_a_norm)?, + d.latent_eps, + d.q_lora, + m, + b.qr, + )?; + quantize_qb(gpu, wk.g32, wb.forms, (&*b.qr, m), (fault, sites.q_b.ty()))?; + { + let x_in = input(wb.forms, Input::Qb, &*b.qr); + wk.project(cx, (&sites.q_b, d.heads * d.head_k()), &x_in, m, b.q)?; + } + let (hk, row, hl, hv) = ( + d.heads * d.head_k(), + d.heads * d.row(), + d.heads * LATENT, + d.heads * d.v_head, + ); + let (f, u) = (size_of::(), size_of::()); + let kv_norm = w.f32_buf(WHAT, kv_a_norm)?; + k.latent.enqueue_latent_extra_append( + stream, + LatentExtraAppendArgs { + rows: Rows { + x: &*b.kv, + stride: d.row(), + m, + }, + off: 0, + e_off: LATENT, + extra: d.extra, + gain: kv_norm, + pos: &*b.pos, + eps: d.latent_eps, + fault, + cache: &mut *b.cache, + }, + )?; + k.step.enqueue_gather( + stream, + &*b.q, + &k.src, + &k.dst, + m * d.heads * d.extra, + &mut *b.qabs, + )?; + let (qs, dd) = q8_planes(w, &sites.k_b)?; + let value = value_w(w, &sites.v_b)?; + let q8 = gpu.q8f32(); + for c0 in (0..m).step_by(PASS_ROWS) { + let mc = PASS_ROWS.min(m - c0); + let q = Window::::of(&*b.q, c0 * hk * f, mc * hk)?; + let n_keys = Window::::of(b.n_keys, (c0 + mc - 1) * u, 1)?; + let mut qabs = WindowMut::::of_mut(&mut *b.qabs, c0 * row * f, mc * row)?; + let mut att = WindowMut::::of_mut(&mut *b.att, c0 * hl * f, mc * hl)?; + q8.enqueue_q8_0_gemv_heads_mcol( + stream, + Q8_0GemvHeadsMcolArgs { + qs, + d: dd, + x: &q, + rows_per_head: LATENT, + x_head_stride: d.head_k(), + y_head_stride: d.row(), + y_off: d.extra, + m: mc, + x_col_stride: hk, + y_col_stride: row, + y: &mut qabs, + }, + )?; + gpu.flash().enqueue_flash_latent_split( + stream, + FlashSplitArgs { + inputs: FlashInputs { + q: &qabs, + kv: &*b.cache, + n_keys_buf: &n_keys, + kq_scale: 1.0 / (d.head_k() as f32).sqrt(), + geom: FlashGeom { + tokens: mc, + heads: d.heads, + rope_dims: d.extra, + latent_dims: LATENT, + }, + }, + part_v: &mut *b.part_v, + part_ms: &mut *b.part_ms, + y: &mut att, + }, + )?; + if let ValueW::Q8_0(vqs, vd) = value { + let mut av = WindowMut::::of_mut(&mut *b.av, c0 * hv * f, mc * hv)?; + q8.enqueue_q8_0_gemv_heads_mcol( + stream, + Q8_0GemvHeadsMcolArgs { + qs: vqs, + d: vd, + x: &att, + rows_per_head: d.v_head, + x_head_stride: LATENT, + y_head_stride: d.v_head, + y_off: 0, + m: mc, + x_col_stride: hl, + y_col_stride: hv, + y: &mut av, + }, + )?; + } + } + if let ValueW::Q6K(wv) = value { + value_heads( + gpu, + (wk, wb), + wv, + (&*b.att, m), + (d.v_head, fault, &mut *b.av), + )?; + } + k.out_gate.enqueue_out_gate( + stream, + OutGateArgs { + attn: &*b.av, + gate: &*b.gate, + n_head: d.heads, + m, + fault, + y: &mut *b.gated, + }, + )?; + quantize_out( + gpu, + wk.g32, + wb.forms, + (&*b.gated, m), + (fault, sites.out.ty()), + )?; + let x_in = input(wb.forms, Input::Out, &*b.gated); + wk.project(cx, (&sites.out, d.embd), &x_in, m, b.out) +} + #[cfg(test)] mod tests { use super::{MlaBufs, MlaDims, MlaKernels, MlaSites, ValueHeads, partials_lens_rows, rows}; diff --git a/crates/gpu/src/arch/raon1/mod.rs b/crates/gpu/src/arch/raon1/mod.rs index 66f2ed968..62ebf6ea6 100644 --- a/crates/gpu/src/arch/raon1/mod.rs +++ b/crates/gpu/src/arch/raon1/mod.rs @@ -3,7 +3,9 @@ //! ([`ffn`]), the RDA layer ([`rda`]) and the gated latent layer ([`mla`]); //! the passes of several rows — a verify, a prompt call's chunks — and the //! taps a block draft reads ([`rows`]), over the RDA stores' lanes -//! ([`state`]). The model description is `model::arch::raon1`. +//! ([`state`]); and a prompt's units of more rows than a pass, which run the +//! projections and the experts as GEMMs ([`wide`]). The model description is +//! `model::arch::raon1`. mod body; mod ffn; @@ -12,6 +14,7 @@ pub mod mla; pub mod rda; mod rows; mod state; +mod wide; pub use body::Body; pub use ffn::sigmoid_128_6; diff --git a/crates/gpu/src/arch/raon1/rows.rs b/crates/gpu/src/arch/raon1/rows.rs index 36f061878..dcb77d0d2 100644 --- a/crates/gpu/src/arch/raon1/rows.rs +++ b/crates/gpu/src/arch/raon1/rows.rs @@ -4,9 +4,11 @@ //! positions, each row's RDA state in a lane of its own, every row into //! one head of the pass's rows; [`Rollback`] commits the rows kept //! ([`super::state`]); -//! - a prompt call ([`GpuModel::prompt_rows`]): its ids staged on the card -//! [`GROUP`] at a time, each chunk of up to [`PASS_ROWS`] of them one walk -//! in place, the head on the call's last row alone; +//! - a prompt call ([`GpuModel::prompt_rows`]): wide units of up to +//! [`WIDE_ROWS`] rows ([`super::wide`]) while at least [`WIDE_FROM`] rows +//! remain, then the rest's ids staged on the card [`GROUP`] at a time, each +//! chunk of up to [`PASS_ROWS`] of them one walk in place, the head on the +//! call's last row alone; //! - the taps ([`GpuModel::set_draft_taps`]): the streams entering the //! layers a draft names, laid in by every walk, one row a position. //! @@ -17,6 +19,7 @@ use super::body::{Body, Run, Tail}; use super::mhc::STREAMS; use super::state::{Lanes, PASS_ROWS, Pass}; +use super::wide::{WIDE_FROM, WIDE_ROWS}; use crate::draft_taps::{DraftTaps, TapRowsSink}; use crate::head::Head; use crate::hybrid::Chain; @@ -57,7 +60,7 @@ pub(super) struct RowViews<'a> { } impl RowsIn { - fn new(stream: &CudaStream, cap: usize) -> Result { + pub(super) fn new(stream: &CudaStream, cap: usize) -> Result { let len = (STREAMS + 2) * cap; let u = |n: usize| DeviceBuffer::::zeroed(stream, n); Ok(RowsIn { @@ -69,7 +72,7 @@ impl RowsIn { }) } - fn bytes(&self) -> usize { + pub(super) fn bytes(&self) -> usize { [&self.buf, &self.epos, &self.ekeys] .into_iter() .map(DeviceBuffer::num_bytes) @@ -79,7 +82,12 @@ impl RowsIn { /// `ids` at positions `pos, pos + 1, …` into the staging's first rows: /// one copy, which waits for the walks before it. Refused by name for /// no ids, more than the staging holds, or a live count past u32. - fn stage(&mut self, stream: &CudaStream, ids: &[u32], pos: u32) -> Result<(), GpuError> { + pub(super) fn stage( + &mut self, + stream: &CudaStream, + ids: &[u32], + pos: u32, + ) -> Result<(), GpuError> { let c = self.cap; if ids.is_empty() || ids.len() > c { return Err(GpuError::shape( @@ -298,11 +306,14 @@ impl Rollback for Body { } impl Body { - /// A prompt call's walks from `pos` ([`GpuModel::prompt_rows`]): the ids - /// staged a group at a time, each chunk of up to [`PASS_ROWS`] a walk in - /// place, the head on the call's last row; with `tapped`, the taps of the - /// call's last `window` rows handed to its sink after each group that - /// holds any. Returns `true`: the call's last chunk enqueues the head. + /// A prompt call's walks from `pos` ([`GpuModel::prompt_rows`]): wide + /// units of up to [`WIDE_ROWS`] rows ([`super::wide`]) as long as at + /// least [`WIDE_FROM`] rows remain, then the rest with its ids staged a + /// group at a time, each chunk of up to [`PASS_ROWS`] a walk in place + /// (every row the step's bits). The head is on the call's last row; with + /// `tapped`, the taps of the call's last `window` rows are handed to its + /// sink after each unit or group that holds any. Returns `true`: the + /// call's last walk enqueues the head. fn prompt_walk( &mut self, gpu: &Gpu, @@ -321,8 +332,31 @@ impl Body { let from = tapped .as_ref() .map_or(n, |(window, _)| n - (*window).min(n)); - for (g, group) in ids.chunks(GROUP).enumerate() { - let g0 = g * GROUP; + let mut done = 0; + while n - done >= WIDE_FROM { + let m = WIDE_ROWS.min(n - done); + let at = pos + launch_u32(WHAT, "row", done)?; + self.wide.rin.stage(stream, &ids[done..done + m], at)?; + if let Some(t) = self.rows.taps.as_mut() { + t.plan(m); + } + let tail = if done + m == n { + Tail::Last(&mut *head) + } else { + Tail::Open + }; + let run = Run::Rows { + at: 0, + m, + pass: Pass::InPlace, + }; + self.walk(gpu, w, run, tail)?; + self.rows.held = pos + launch_u32(WHAT, "row", done + m)?; + self.hand_over(gpu, tapped.as_mut(), (pos, from), (done, m))?; + done += m; + } + for (g, group) in ids[done..].chunks(GROUP).enumerate() { + let g0 = done + g * GROUP; let gpos = pos + launch_u32(WHAT, "row", g0)?; self.rows.rin.stage(stream, group, gpos)?; if let Some(t) = self.rows.taps.as_mut() { @@ -343,29 +377,48 @@ impl Body { self.walk(gpu, w, run, tail)?; self.rows.held = gpos + launch_u32(WHAT, "row", at + m)?; } - if let Some((_, sink)) = tapped.as_mut() { - let (lo, hi) = (from.max(g0), g0 + group.len()); - if lo < hi { - let taps = self - .rows - .taps - .as_ref() - .ok_or(GpuError::state(WHAT, "taps set (set_draft_taps)"))?; - let rows = taps.read(stream, lo - g0, hi - lo)?; - sink(pos + launch_u32(WHAT, "row", lo)?, &rows)?; - } - } + self.hand_over(gpu, tapped.as_mut(), (pos, from), (g0, group.len()))?; } Ok(true) } + + /// The tapped rows of the call's rows `first .. first + rows` that lie + /// from `from` on, handed to the sink as `(position, rows)`; the call + /// began at position `pos`. Nothing for an untapped call or a unit wholly + /// before `from`. + fn hand_over( + &self, + gpu: &Gpu, + tapped: Option<&mut (usize, &mut TapRowsSink<'_>)>, + (pos, from): (u32, usize), + (first, rows): (usize, usize), + ) -> Result<(), GpuError> { + let Some((_, sink)) = tapped else { + return Ok(()); + }; + let (lo, hi) = (from.max(first), first + rows); + if lo >= hi { + return Ok(()); + } + let taps = self + .rows + .taps + .as_ref() + .ok_or(GpuError::state(WHAT, "taps set (set_draft_taps)"))?; + let read = taps.read(gpu.stream(), lo - first, hi - lo)?; + sink(pos + launch_u32(WHAT, "row", lo)?, &read) + } } impl GpuModel { - /// Run `ids` from the model's position as passes of rows and return the - /// argmax after the last: each chunk of up to [`PASS_ROWS`] ids one walk - /// whose rows are the steps' bits, the head on the call's last row alone - /// (its fault word read there). Refused by name with no ids, and while a - /// verify waits. + /// Run `ids` from the model's position and return the argmax after the + /// last: units of up to [`WIDE_ROWS`] ids as GEMMs (the wide arm: a row + /// is the step's row to the error of the GEMM's sums, not its bits) while + /// at least [`WIDE_FROM`] remain, the rest as passes of up to + /// [`PASS_ROWS`] ids whose rows are the steps' bits; a call of up to + /// [`PASS_ROWS`] ids is those passes alone. The head is on the call's + /// last row alone (its fault word read there). Refused by name with no + /// ids, and while a verify waits. pub fn prompt_rows(&mut self, ids: &[u32]) -> Result { self.prompt_call(ids, None) } @@ -373,8 +426,8 @@ impl GpuModel { /// [`GpuModel::prompt_rows`] with the tapped rows of the call's last /// `window` positions (all of them when the call is shorter) handed to /// `sink` as `(first position, rows)`, one [`GpuModel::draft_tap_width`] - /// row a position, after each group of [`GROUP`] rows that holds any of - /// them — in order, no position twice or missed. The walks are the plain + /// row a position, after each unit or group that holds any of them — in + /// order, no position twice or missed. The walks are the plain /// call's, so its stores and argmax are bit for bit. Refused by name with /// no taps set; a sink that fails ends the call with its error. pub fn prompt_rows_tapped( @@ -403,8 +456,8 @@ impl GpuModel { /// Tap the rows entering `layers` (strictly rising, each below the /// layer count; empty stops tapping) from the next walk on: every walk — - /// a step, a verify's pass, each prompt chunk — lays them into a buffer - /// of [`GROUP`] rows a layer ([`DraftTaps`]), allocated here against the + /// a step, a verify's pass, each prompt chunk or wide unit — lays them + /// into a buffer of [`WIDE_ROWS`] rows a layer ([`DraftTaps`]), allocated here against the /// card's free bytes with [`FIT_RESERVE`](crate::model::fit::FIT_RESERVE) /// left. The captured chains are dropped first (a capture records the /// copies into the buffer of its time), and the next graph step or pass @@ -434,7 +487,7 @@ impl GpuModel { gpu.stream(), layers, n_layer, - (hidden, GROUP), + (hidden, WIDE_ROWS), free, )?); Ok(()) diff --git a/crates/gpu/src/arch/raon1/wide.rs b/crates/gpu/src/arch/raon1/wide.rs new file mode 100644 index 000000000..59b83ef1a --- /dev/null +++ b/crates/gpu/src/arch/raon1/wide.rs @@ -0,0 +1,992 @@ +//! The wide arm of Raon-1's walk: a unit of more than [`PASS_ROWS`] rows, up +//! to [`WIDE_ROWS`], the prompt's bulk. One layer walk as the passes run, each +//! op choosing its launch family from the rows (`qwen3moe::wide`'s shape): the +//! projections through the grouped int8 tensor-core GEMM over a one-expert +//! table ([`site::gemm`]) — a K-quant over q8_1 blocks of 128, a Q8_0 over q8 +//! blocks of 32, an F32 site on the F32 tile — the routed experts through one +//! route table over the unit's slots, the router, the conv, the delta step, the +//! norms and the card sum as the passes launch them (they take any row count). +//! +//! A unit's rows are not each the one-row step's bits: the GEMM sums 128-value +//! integer blocks into one f32 accumulator where the gemv sums lane partials +//! through a warp tree. What stays the step's: the router (the expert ids), the +//! card sum's order and every element-wise op. A unit of up to [`PASS_ROWS`] +//! rows is the pass of [`super::rows`], bit for bit the steps. +//! +//! The pieces by owner: +//! - the RDA layer: [`rda_rows`], the common mixer +//! ([`crate::linear::mixer_kda`]) over [`RdaWide`]'s projections; +//! - the latent layer: [`super::mla::rows_wide`], the projections wide, the +//! cache append, the gather, the per-head value projection ([`value_heads`]) +//! and the output gate over the unit, the absorb and the latent flash in +//! chunks of [`PASS_ROWS`] query tokens (the flash takes no more); +//! - the feed-forward block: [`ffn`], below. + +use super::body::Scratch; +use super::ffn::{Ffn, Glu, Routed, sigmoid_128_6}; +use super::mhc::{MAX_TOKENS, STREAMS}; +use super::rda::RdaSites; +use super::rows::RowsIn; +use super::state::PASS_ROWS; +use crate::arch::hybrid::{GqaCx, Site}; +use crate::card_leg::{RouteW, StackW}; +use crate::card_slots::{CardAcc, CardSlotKernels}; +use crate::fault::FaultSink; +use crate::gemm::{ + Gemm32Args, Gemm32Kernels, GemmAct, GemmAct32, GemmArgs, GemmInput, GemmKernels, GemmRoute, +}; +use crate::kquant::Act; +use crate::latent::LATENT; +use crate::linear::mixer_kda::{self, KdaBufs, KdaDims, KdaIn, KdaProj, KdaState}; +use crate::linear::{HEAD, LinearKernels, LinearShape}; +use crate::model::fit::{act_col_bytes, act32_col_bytes, route_bytes}; +use crate::site::{self, Gated, GatedIo, SiteTy, SiteWeight, SwigluRows, WideActs, WideIn}; +use crate::tensor::DeviceTensor; +use crate::weights::{DevWeight, Weights}; +use crate::{Gpu, GpuError, launch_u32}; +use cuda_core::{CudaContext, CudaStream, DeviceBuffer, LaunchConfig1D}; +use cuda_device::{DisjointSlice, kernel, launch_bounds, launch_contract, thread}; +use cuda_host::cuda_module; +use gguf::GgmlType; +use std::sync::Arc; + +/// What the refusals of this file name. +const WHAT: &str = "raon1 wide"; + +/// The most rows a wide unit holds: its buffers' rows. At most the mHC route +/// launch's tokens ([`MAX_TOKENS`]). +pub(super) const WIDE_ROWS: usize = 256; + +/// The fewest rows of a unit that runs wide; fewer run as the pass does. +pub(super) const WIDE_FROM: usize = PASS_ROWS + 1; + +const _: () = assert!(WIDE_ROWS <= MAX_TOKENS && WIDE_ROWS >= WIDE_FROM); + +#[cuda_module] +mod wide_kernels { + use super::*; + + /// The joined q·k·v rows of the conv: row `t` of `y` is `q[t]`, `k[t]`, + /// `v[t]` end to end (`q` and `k` `wq` values a row, `v` `wv`) — a copy, + /// so exact. One thread per value. + #[kernel] + #[launch_bounds(256)] + #[launch_contract( + domain = 1, + block = (256, 1, 1), + requires = ( + q.len() >= wq * m, + k.len() >= wq * m, + v.len() >= wv * m, + y.len() >= (2 * wq + wv) * m + ) + )] + pub fn raon_join_qkv( + q: &[f32], + k: &[f32], + v: &[f32], + wq: u32, + wv: u32, + m: u32, + mut y: DisjointSlice, + ) { + let i = thread::index_1d().get(); + let (wq, wv, m) = (wq as usize, wv as usize, m as usize); + let row = 2 * wq + wv; + if i >= row * m { + return; + } + let t = i / row; + let c = i - t * row; + // SAFETY: i < row·m, so t < m and c < row; the source index of each + // arm is below its buffer's m rows by the launch contract; thread i is + // y[i]'s only writer. + let x = unsafe { + if c < wq { + *q.get_unchecked(t * wq + c) + } else if c < 2 * wq { + *k.get_unchecked(t * wq + (c - wq)) + } else { + *v.get_unchecked(t * wv + (c - 2 * wq)) + } + }; + // SAFETY: i < row·m <= y.len() by the launch contract. + unsafe { *y.get_unchecked_mut(i) = x }; + } +} + +/// The module of the join. +pub(super) struct JoinKernels { + module: wide_kernels::LoadedModule, +} + +impl JoinKernels { + /// Load this file's device bundle into `ctx`. Load-time only. + pub(super) fn load(ctx: &Arc) -> Result { + // SAFETY: this package owns the embedded device bundle produced for + // the module above; every launcher checks its launch contract. + let module = unsafe { crate::shared_module!(wide_kernels, ctx)? }; + Ok(JoinKernels { module }) + } + + /// Enqueue the join of `m` rows of `q` and `k` (`wq` values a row) and + /// `v` (`wv`) into `y`'s `m` rows of `2 · wq + wv`. + #[allow( + clippy::too_many_arguments, + reason = "the three parts, their widths, the rows and the joined rows (rust-quality R8)" + )] + pub(super) fn enqueue_join_qkv( + &self, + stream: &CudaStream, + (q, k, v): (&DeviceBuffer, &DeviceBuffer, &DeviceBuffer), + (wq, wv): (usize, usize), + m: usize, + y: &mut DeviceBuffer, + ) -> Result<(), GpuError> { + let what = "raon1 wide join"; + let row = 2 * wq + wv; + if m == 0 || q.len() < wq * m || k.len() < wq * m || v.len() < wv * m || y.len() < row * m { + return Err(GpuError::shape( + what, + format!( + "{m} rows of {wq}+{wq}+{wv}: q {}, k {}, v {}, y {}", + q.len(), + k.len(), + v.len(), + y.len() + ), + )); + } + let grid = launch_u32(what, "grid", (row * m).div_ceil(256))?; + let (wq, wv, m) = ( + launch_u32(what, "wq", wq)?, + launch_u32(what, "wv", wv)?, + launch_u32(what, "m", m)?, + ); + let prep = self + .module + .prepare_raon_join_qkv(LaunchConfig1D::new(grid, 256, 0))?; + self.module + .raon_join_qkv(stream, &prep, q, k, v, wq, wv, m, y)?; + Ok(()) + } +} + +/// The activation forms a wide unit's sites read, one buffer per input the +/// walk quantizes: each holds the unit's rows (or slots) of its width. +pub(super) struct Forms { + /// The normed rows: the model width, in q8_1 blocks of 128 (the K-quant + /// sites) and of 32 (a Q8_0 site). + hid128: GemmAct, + hid32: GemmAct32, + /// The mixers' gated rows into the output projection, `heads · v_head`. + out32: GemmAct32, + /// The latent layer's normed query latent into `attn_q_b`. + qb32: GemmAct32, + /// The dense block's SiTU rows into its down, the dense width. + dense128: GemmAct, + /// The shared expert's SiTU rows into its down, the shared width. + shared32: GemmAct32, + /// The routed experts' SiTU rows, a column a slot, into the stacks' down. + slots32: GemmAct32, + /// The latent attention's output, a column a (row, head), into the value + /// projection. + lat128: GemmAct, +} + +/// The widths a wide unit is cut for. +#[derive(Clone, Copy)] +pub(super) struct WideDims { + pub embd: usize, + /// `heads · v_head`: the output projections' input. + pub attn: usize, + pub q_lora: usize, + /// The latent layers' heads. + pub heads: usize, + pub dense_ff: usize, + pub shared_ff: usize, + pub expert_ff: usize, + /// The widest feed-forward block's gate rows a unit holds: the dense + /// block's, or the routed slots'. + pub ff_rows: usize, +} + +impl Forms { + fn new(stream: &CudaStream, d: &WideDims) -> Result { + let rows = WIDE_ROWS; + let slots = rows * sigmoid_128_6::N_USED; + Ok(Forms { + hid128: GemmAct::new(stream, rows, d.embd)?, + hid32: GemmAct32::new(stream, rows, d.embd)?, + out32: GemmAct32::new(stream, rows, d.attn)?, + qb32: GemmAct32::new(stream, rows, d.q_lora)?, + dense128: GemmAct::new(stream, rows, d.dense_ff)?, + shared32: GemmAct32::new(stream, rows, d.shared_ff)?, + slots32: GemmAct32::new(stream, slots, d.expert_ff)?, + lat128: GemmAct::new(stream, rows * d.heads, LATENT)?, + }) + } + + fn bytes_of(d: &WideDims) -> usize { + let rows = WIDE_ROWS; + let slots = rows * sigmoid_128_6::N_USED; + rows * (act_col_bytes(d.embd) + act_col_bytes(d.dense_ff)) + + rows * (act32_col_bytes(d.embd) + act32_col_bytes(d.attn)) + + rows * (act32_col_bytes(d.q_lora) + act32_col_bytes(d.shared_ff)) + + slots * act32_col_bytes(d.expert_ff) + + rows * d.heads * act_col_bytes(LATENT) + } +} + +/// The feed-forward block's buffers over a unit's rows, token-major. +pub(super) struct WideFfn { + /// The gate and up rows of the dense block, the shared expert (a row + /// each) and the routed slots (a slot each): `ff_rows` values a row. + g: DeviceBuffer, + u: DeviceBuffer, + shared: DeviceBuffer, + routed: DeviceBuffer, + /// The slots' down rows, `embd` a slot. + down: DeviceBuffer, + probs: DeviceBuffer, + ids: DeviceBuffer, + weights: DeviceBuffer, + /// The route table over the unit's slots and the experts. + moe: GemmRoute, +} + +impl WideFfn { + fn new(stream: &CudaStream, d: &WideDims) -> Result { + let z = |n: usize| DeviceBuffer::::zeroed(stream, n.max(1)); + let (rows, used) = (WIDE_ROWS, sigmoid_128_6::N_USED); + Ok(WideFfn { + g: z(rows * d.ff_rows)?, + u: z(rows * d.ff_rows)?, + shared: z(rows * d.embd)?, + routed: z(rows * d.embd)?, + down: z(rows * used * d.embd)?, + probs: z(rows * sigmoid_128_6::N_EXPERT)?, + ids: DeviceBuffer::zeroed(stream, rows * used)?, + weights: z(rows * used)?, + moe: GemmRoute::new(stream, rows * used, sigmoid_128_6::N_EXPERT)?, + }) + } + + fn bytes_of(d: &WideDims) -> usize { + let (rows, used) = (WIDE_ROWS, sigmoid_128_6::N_USED); + let f32s = 2 * rows * d.ff_rows + + 2 * rows * d.embd + + rows * used * d.embd + + rows * sigmoid_128_6::N_EXPERT + + rows * used; + 4 * (f32s + rows * used) + route_bytes(rows * used, sigmoid_128_6::N_EXPERT) + } +} + +/// The latent layers' per-head value projection's table: a slot a (row, head), +/// slot `s` routed to expert `s % heads`, the head's rows of the stack. +pub(super) struct HeadSlots { + heads: usize, + ids: DeviceBuffer, + route: GemmRoute, +} + +impl HeadSlots { + fn new(stream: &CudaStream, heads: usize) -> Result { + let slots = WIDE_ROWS * heads; + let ids = (0..slots) + .map(|s| launch_u32(WHAT, "head", s % heads)) + .collect::, _>>()?; + Ok(HeadSlots { + heads, + ids: DeviceBuffer::from_host(stream, &ids)?, + route: GemmRoute::new(stream, slots, heads)?, + }) + } + + fn bytes_of(heads: usize) -> usize { + 4 * WIDE_ROWS * heads + route_bytes(WIDE_ROWS * heads, heads) + } + + fn bytes(&self) -> usize { + self.ids.num_bytes() + self.route.bytes() + } +} + +/// A wide unit's own buffers and kernels beyond the passes': the staging, the +/// walk's rows, the forms and the feed-forward block's buffers, the join's +/// parts and the GEMM family. +pub(super) struct Wide { + pub(super) rin: RowsIn, + pub(super) s: Scratch, + ffn: WideFfn, + forms: Forms, + /// `q`, `k` and `v`'s rows before the join, `attn` values a row each. + parts: [DeviceBuffer; 3], + /// The one-expert table every projection of a unit reads. + dense: GemmRoute, + heads: HeadSlots, + gemm: GemmKernels, + join: JoinKernels, + slots: CardSlotKernels, +} + +/// The kernels and the table a wide launch reads. +#[derive(Clone, Copy)] +pub(super) struct Kern<'a> { + pub(super) gemm: &'a GemmKernels, + pub(super) g32: &'a Gemm32Kernels, + pub(super) join: &'a JoinKernels, + pub(super) slots: &'a CardSlotKernels, + pub(super) router: &'a sigmoid_128_6::RouterKernels, + pub(super) dense: &'a GemmRoute, +} + +/// The buffers a wide launch writes beyond a layer's own. +pub(super) struct Bufs<'a> { + pub(super) forms: &'a mut Forms, + pub(super) parts: &'a mut [DeviceBuffer; 3], + pub(super) ffn: &'a mut WideFfn, + pub(super) heads: &'a mut HeadSlots, +} + +/// What a wide launch borrows beyond a layer's buffers. +pub(super) struct WideKit<'a> { + pub(super) k: Kern<'a>, + pub(super) b: Bufs<'a>, +} + +impl Wide { + /// The unit's buffers for `d`, and the kernels of the GEMM family, over + /// a walk's `scratch` of [`WIDE_ROWS`] rows. Load-time only. + pub(super) fn new(gpu: &Gpu, d: &WideDims, scratch: Scratch) -> Result { + let (stream, ctx) = (gpu.stream(), gpu.context()); + let part = || DeviceBuffer::::zeroed(stream, WIDE_ROWS * d.attn); + Ok(Wide { + rin: RowsIn::new(stream, WIDE_ROWS)?, + s: scratch, + ffn: WideFfn::new(stream, d)?, + forms: Forms::new(stream, d)?, + parts: [part()?, part()?, part()?], + dense: GemmRoute::new(stream, WIDE_ROWS, 1)?, + heads: HeadSlots::new(stream, d.heads)?, + gemm: GemmKernels::load(ctx)?, + join: JoinKernels::load(ctx)?, + slots: CardSlotKernels::load(ctx)?, + }) + } + + /// Card bytes of what [`Wide::new`] allocates beyond the walk's rows. + pub(super) fn bytes_of(d: &WideDims) -> usize { + let stage = (STREAMS + 2) * WIDE_ROWS * size_of::() + 2 * STREAMS * WIDE_ROWS * 4; + Forms::bytes_of(d) + + WideFfn::bytes_of(d) + + 3 * 4 * WIDE_ROWS * d.attn + + route_bytes(WIDE_ROWS, 1) + + HeadSlots::bytes_of(d.heads) + + stage + } + + /// Card bytes held, the walk's rows included. + pub(super) fn bytes(&self) -> usize { + self.s.bytes() + + self.rin.bytes() + + self.forms.bytes() + + self.ffn.bytes() + + self + .parts + .iter() + .map(DeviceBuffer::num_bytes) + .sum::() + + self.dense.bytes() + + self.heads.bytes() + } + + /// The unit's one-expert table for `m` rows, once behind the embedding. + pub(super) fn route_dense(&mut self, gpu: &Gpu, m: usize) -> Result<(), GpuError> { + self.gemm + .enqueue_route_dense(gpu.stream(), m, &mut self.dense, gpu.unlabelled_sink()) + } + + /// The kit over this unit's pieces beside the passes' `g32` and router, + /// the walk's rows and the staging. + pub(super) fn kit<'a>( + &'a mut self, + g32: &'a Gemm32Kernels, + router: &'a sigmoid_128_6::RouterKernels, + ) -> (WideKit<'a>, &'a mut Scratch, &'a mut RowsIn) { + let Wide { + rin, + s, + ffn, + forms, + parts, + dense, + heads, + gemm, + join, + slots, + } = self; + let k = Kern { + gemm, + g32, + join, + slots, + router, + dense, + }; + ( + WideKit { + k, + b: Bufs { + forms, + parts, + ffn, + heads, + }, + }, + s, + rin, + ) + } +} + +impl Forms { + fn bytes(&self) -> usize { + self.hid128.bytes() + + self.hid32.bytes() + + self.out32.bytes() + + self.qb32.bytes() + + self.dense128.bytes() + + self.shared32.bytes() + + self.slots32.bytes() + + self.lat128.bytes() + } +} + +impl WideFfn { + fn bytes(&self) -> usize { + [ + &self.g, + &self.u, + &self.shared, + &self.routed, + &self.down, + &self.probs, + &self.weights, + ] + .into_iter() + .map(DeviceBuffer::num_bytes) + .sum::() + + self.ids.num_bytes() + + self.moe.bytes() + } +} + +/// Which of the unit's forms a projection reads. +#[derive(Clone, Copy)] +pub(super) enum Input { + /// The normed rows' forms ([`quantize_hid`]). + Hid, + /// The gated rows into an output projection ([`quantize_out`]). + Out, + /// The query latent into `attn_q_b` ([`quantize_qb`]). + Qb, + /// The f32 rows alone (an F32 site). + Plain, +} + +/// The input `x` of a projection as `which` reads it. +pub(super) fn input<'f>(forms: &'f Forms, which: Input, x: &'f DeviceBuffer) -> WideIn<'f> { + let (q128, q32) = match which { + Input::Hid => (Some(&forms.hid128), Some(&forms.hid32)), + Input::Out => (None, Some(&forms.out32)), + Input::Qb => (None, Some(&forms.qb32)), + Input::Plain => (None, None), + }; + WideIn { + q128, + q32, + f32: Some(x), + } +} + +impl Kern<'_> { + fn kernels(&self) -> site::WideKernels<'_> { + site::WideKernels { + gemm: self.gemm, + g32: Some(self.g32), + } + } + + /// Site `t` (`rows` outputs) over the unit's `m` rows of `x` in the forms + /// `x_in` carries, into `y`. + pub(super) fn project( + &self, + cx: &GqaCx<'_>, + (t, rows): (&Site, usize), + x_in: &WideIn<'_>, + m: usize, + y: &mut DeviceBuffer, + ) -> Result<(), GpuError> { + site::gemm( + cx.gpu, + &self.kernels(), + (t.ty(), cx.w, SiteWeight::Resident(t.name())), + rows, + x_in, + (self.dense, GemmInput::PerSlot), + m, + y, + ) + } +} + +/// The normed rows `x` (`m` of the model width) in the forms asked for: +/// `(q128, q32)`, as [`WideActs::wants_of`] names them for a set of sites. +pub(super) fn quantize_hid( + gpu: &Gpu, + g32: &Gemm32Kernels, + forms: &mut Forms, + (x, m): (&DeviceBuffer, usize), + (sink, (q128, q32)): (FaultSink, (bool, bool)), +) -> Result<(), GpuError> { + if q128 { + gpu.enqueue_quantize_gemm(x, m, &mut forms.hid128, sink)?; + } + if q32 { + g32.enqueue_quantize_gemm32(gpu.stream(), x, m, &mut forms.hid32, sink)?; + } + Ok(()) +} + +/// The latent layers' per-head value projection over a unit's `m` rows: `att` +/// (`m · heads` rows of [`LATENT`] values, the attention's output) times +/// `wv`, a Q6_K stack of `heads` experts of `rows` rows (`attn_v_b`), into `av` +/// (`m · heads` rows of `rows` values): the rows in q8_1 blocks, a slot a +/// (row, head) routed to its head, one grouped GEMM. +pub(super) fn value_heads( + gpu: &Gpu, + (k, b): (&Kern<'_>, &mut Bufs<'_>), + wv: &DeviceTensor, + (att, m): (&DeviceBuffer, usize), + (rows, sink, av): (usize, FaultSink, &mut DeviceBuffer), +) -> Result<(), GpuError> { + let Some(RouteW::Grouped(ty, w)) = StackW::Words(wv).route(GgmlType::Q6_K) else { + return Err(GpuError::shape( + WHAT, + "the value projection's Q6_K words have no grouped GEMM route", + )); + }; + let (stream, heads) = (gpu.stream(), b.heads.heads); + gpu.enqueue_quantize_gemm(att, m * heads, &mut b.forms.lat128, sink)?; + k.gemm + .enqueue_route(stream, &b.heads.ids, m * heads, &mut b.heads.route, sink)?; + k.gemm.enqueue_gemm( + stream, + GemmArgs { + ty, + w, + rows_per_expert: rows, + act: &b.forms.lat128, + route: &b.heads.route, + input: GemmInput::PerSlot, + y: av, + }, + ) +} + +/// The mixers' gated rows `x` (`m` of `heads · v_head`) in the q8 blocks of 32 +/// the output projection of type `ty` reads. +pub(super) fn quantize_out( + gpu: &Gpu, + g32: &Gemm32Kernels, + forms: &mut Forms, + (x, m): (&DeviceBuffer, usize), + (sink, ty): (FaultSink, SiteTy), +) -> Result<(), GpuError> { + only_blocks32(ty, "an output projection")?; + g32.enqueue_quantize_gemm32(gpu.stream(), x, m, &mut forms.out32, sink) +} + +/// The latent layer's normed query latent `x` in the q8 blocks of 32 +/// `attn_q_b` of type `ty` reads. +pub(super) fn quantize_qb( + gpu: &Gpu, + g32: &Gemm32Kernels, + forms: &mut Forms, + (x, m): (&DeviceBuffer, usize), + (sink, ty): (FaultSink, SiteTy), +) -> Result<(), GpuError> { + only_blocks32(ty, "a query up projection")?; + g32.enqueue_quantize_gemm32(gpu.stream(), x, m, &mut forms.qb32, sink) +} + +/// A named refusal of a projection `what` of type `ty` that reads another +/// form than q8 blocks of 32: the wide arm's forms for those inputs are +/// Q8_0's. +fn only_blocks32(ty: SiteTy, what: &str) -> Result<(), GpuError> { + let (q128, q32) = WideActs::wants(ty); + if q128 || !q32 { + return Err(GpuError::shape( + WHAT, + format!("{what} of {ty}: the wide arm's form for its input is q8 blocks of 32 (Q8_0)"), + )); + } + Ok(()) +} + +/// The RDA layer's wide projections ([`KdaProj`]): q, k, v and the gate over +/// the normed rows' q8_1 form, β's, the write gate's and the forget +/// projection's halves on the F32 tile, the output projection over the gated +/// rows' q8 blocks of 32. +struct RdaWide<'a, 'k> { + cx: &'a GqaCx<'a>, + sites: &'a RdaSites, + shape: LinearShape, + kit: &'a mut WideKit<'k>, +} + +impl KdaProj for RdaWide<'_, '_> { + /// q, k and v into three row sets joined into the conv's channels, `z` + /// from the gate projection, β's and the write gate's raw rows, `fa`, then + /// `f = f_b · fa`; the K-quant sites share one quantization of `xn`. + fn input( + &mut self, + gpu: &Gpu, + l: usize, + m: usize, + xn: &DeviceBuffer, + y: KdaIn<'_>, + ) -> Result<(), GpuError> { + let (cx, s, shape) = (self.cx, self.sites, self.shape); + let sink = gpu.layer_sink(l)?; + let (qk, v) = (shape.n_k * HEAD, shape.n_v * HEAD); + let all = s.sites(); + let tys: Vec = all[..7].iter().map(|t| t.ty()).collect(); + let KdaIn { + qkv, + fa, + beta_raw, + write_raw, + f, + z, + .. + } = y; + let write_raw = write_raw.ok_or(GpuError::state( + WHAT, + "the write projection's rows (the rule has a write gate)", + ))?; + let WideKit { k, b } = &mut *self.kit; + let Bufs { forms, parts, .. } = b; + quantize_hid(gpu, k.g32, forms, (xn, m), (sink, WideActs::wants_of(&tys)))?; + let x_in = input(forms, Input::Hid, xn); + let [pq, pk, pv] = &mut **parts; + for ((t, rows), y) in [(&s.qkv[0], qk), (&s.qkv[1], qk), (&s.qkv[2], v)] + .into_iter() + .zip([&mut *pq, &mut *pk, &mut *pv]) + { + k.project(cx, (t, rows), &x_in, m, y)?; + } + k.project(cx, (&s.g, v), &x_in, m, z)?; + k.project(cx, (&s.beta, shape.n_v), &x_in, m, beta_raw)?; + k.project(cx, (&s.write, shape.n_v), &x_in, m, write_raw)?; + k.project(cx, (&s.f_a, s.f_b.k()), &x_in, m, fa)?; + let fa_in = input(forms, Input::Plain, &*fa); + k.project(cx, (&s.f_b, v), &fa_in, m, f)?; + k.join + .enqueue_join_qkv(gpu.stream(), (pq, pk, pv), (qk, v), m, qkv) + } + + /// `out = o_proj · gated` over the unit's rows. + fn output( + &mut self, + gpu: &Gpu, + l: usize, + m: usize, + gated: &DeviceBuffer, + out: &mut DeviceBuffer, + ) -> Result<(), GpuError> { + let (cx, t) = (self.cx, &self.sites.out); + let sink = gpu.layer_sink(l)?; + let WideKit { k, b } = &mut *self.kit; + quantize_out(gpu, k.g32, b.forms, (gated, m), (sink, t.ty()))?; + let x_in = input(b.forms, Input::Out, gated); + k.project(cx, (t, cx.hidden), &x_in, m, out) + } +} + +/// Enqueue RDA layer `l` over a wide unit of `m` rows (module doc): the +/// common mixer over [`RdaWide`]'s projections, in position order, the store +/// `st` as the unit's lanes name it. +#[allow( + clippy::too_many_arguments, + reason = "the launch context, the unit's kit, the layer and its store and buffers (rust-quality R8)" +)] +pub(super) fn rda_rows( + cx: &GqaCx<'_>, + kit: &mut WideKit<'_>, + linear: &LinearKernels, + sites: &RdaSites, + (l, d, m): (usize, &KdaDims, usize), + st: KdaState<'_>, + bufs: KdaBufs<'_>, +) -> Result<(), GpuError> { + let gpu = cx.gpu; + let vecs = sites.vecs(cx.w)?; + let mut proj = RdaWide { + cx, + sites, + shape: d.shape, + kit, + }; + mixer_kda::enqueue( + gpu, + gpu.stream(), + linear, + &mut proj, + (l, m, d), + &vecs, + st, + bufs, + ) +} + +/// The stack `name` as the route GEMM reads it: its file type and form. +fn stack<'w>(w: &'w Weights, name: &str) -> Result<(GgmlType, StackW<'w>), GpuError> { + match w.get(name) { + Some(DevWeight::KQuant { ty, w, .. }) => Ok((*ty, StackW::Words(w))), + Some(DevWeight::Q5_0 { w, .. }) => Ok((GgmlType::Q5_0, StackW::Q5_0(w))), + _ => Err(GpuError::tensor( + WHAT, + name, + "a resident routed stack (the file's blocks as words, or Q5_0 packed)", + )), + } +} + +/// A dense block (the leading layer's, or a shared expert) over the unit's +/// `m` rows of the normed rows whose forms the caller made: the gate and up +/// GEMMs, the rule's SiTU quantized into the form the down reads, the down +/// GEMM into `y`. `g_buf` and `u_buf` hold the gate and up rows. +#[allow( + clippy::too_many_arguments, + reason = "the block's context, rule, rows, forms and buffers (rust-quality R8)" +)] +fn glu( + cx: &GqaCx<'_>, + k: &Kern<'_>, + rule: Act, + (l, m): (usize, usize), + (g, forms): (&Glu, &mut Forms), + (g_buf, u_buf): (&mut DeviceBuffer, &mut DeviceBuffer), + y: &mut DeviceBuffer, +) -> Result<(), GpuError> { + let blk = Gated { + w: cx.w, + route: k.dense, + gate: (g.gate.ty(), SiteWeight::Resident(g.gate.name())), + up: (g.up.ty(), SiteWeight::Resident(g.up.name())), + down: (g.down.ty(), SiteWeight::Resident(g.down.name())), + ff: g.ff, + hidden: cx.hidden, + rule, + input: GemmInput::PerSlot, + }; + let x = WideIn { + q128: Some(&forms.hid128), + q32: Some(&forms.hid32), + f32: None, + }; + let rows = match WideActs::wants(g.down.ty()) { + (true, _) if forms.dense128.k() == g.ff => SwigluRows { + q128: Some(&mut forms.dense128), + q32: None, + f32: None, + }, + (false, true) if forms.shared32.k() == g.ff => SwigluRows { + q128: None, + q32: Some(&mut forms.shared32), + f32: None, + }, + _ => { + return Err(GpuError::shape( + WHAT, + format!( + "layer {l}: a block of width {} with a down of {}: the wide arm holds q8_1 \ + rows of the dense width and q8 blocks of 32 of the shared width", + g.ff, + g.down.ty() + ), + )); + } + }; + let io = GatedIo { + x: &x, + gate: g_buf, + up: u_buf, + rows, + y, + }; + site::gated( + cx.gpu, + &k.kernels(), + &blk, + (m, m), + io, + cx.gpu.layer_sink(l)?, + ) +} + +/// Layer `l`'s block over the unit's `m` rows of `xn` into `y`: the dense +/// block, or the routed sum plus the shared expert ([`super::ffn::ffn`]'s +/// launches over a route table): the router, the gate and up GEMMs over the +/// slots, the SiTU into q8 blocks of 32 a slot, the down GEMM of the stack's +/// type, the card sum in slot order, the shared expert, the add. +#[allow( + clippy::too_many_arguments, + reason = "ffn's inputs and the unit's kit (rust-quality R8)" +)] +pub(super) fn ffn( + cx: &GqaCx<'_>, + kit: &mut WideKit<'_>, + (rule, scale): (Act, f32), + (l, f): (usize, &Ffn), + m: usize, + xn: &DeviceBuffer, + y: &mut DeviceBuffer, +) -> Result<(), GpuError> { + let gpu = cx.gpu; + let (stream, sink) = (gpu.stream(), gpu.layer_sink(l)?); + let WideKit { k, b } = kit; + let Bufs { forms, ffn: wf, .. } = b; + match f { + Ffn::Dense(gl) => { + let want = WideActs::wants_of(&[gl.gate.ty(), gl.up.ty()]); + quantize_hid(gpu, k.g32, forms, (xn, m), (sink, want))?; + glu(cx, k, rule, (l, m), (gl, forms), (&mut wf.g, &mut wf.u), y) + } + Ffn::Routed(r) => { + // The routed stacks' gate and up read the q8_1 rows; the shared + // expert's sites read theirs. + let (_, q32) = WideActs::wants_of(&[r.shared.gate.ty(), r.shared.up.ty()]); + quantize_hid(gpu, k.g32, forms, (xn, m), (sink, (true, q32)))?; + routed(cx, k, rule, (l, m), (r, forms, wf), xn, scale)?; + glu( + cx, + k, + rule, + (l, m), + (&r.shared, forms), + (&mut wf.g, &mut wf.u), + &mut wf.shared, + )?; + gpu.elem() + .enqueue_add(stream, &wf.routed, &wf.shared, m * cx.hidden, y) + } + } +} + +/// The routed experts of layer `l` over the unit's `m` rows of `xn` into +/// `wf.routed`: the router, the route table over its ids, the gate and up +/// GEMMs over the q8_1 rows, the rule into q8 blocks of 32 a slot, the down +/// GEMM, the card sum. +#[allow( + clippy::too_many_arguments, + reason = "the routed block's context, rule, rows, buffers and the router's scale (rust-quality R8)" +)] +fn routed( + cx: &GqaCx<'_>, + k: &Kern<'_>, + rule: Act, + (l, m): (usize, usize), + (r, forms, wf): (&Routed, &mut Forms, &mut WideFfn), + xn: &DeviceBuffer, + scale: f32, +) -> Result<(), GpuError> { + let (gpu, w) = (cx.gpu, cx.w); + let (stream, sink) = (gpu.stream(), gpu.layer_sink(l)?); + let used = sigmoid_128_6::N_USED; + let (gty, gate) = stack(w, &r.stacks[0])?; + let (uty, up) = stack(w, &r.stacks[1])?; + let (dty, down) = stack(w, &r.stacks[2])?; + let (Some(RouteW::Grouped(ty, gate)), Some(RouteW::Grouped(_, up))) = + (gate.route(gty), up.route(uty)) + else { + return Err(GpuError::shape( + WHAT, + format!( + "layer {l}: gate and up stacks of {gty} and {uty}: the wide arm runs the grouped K-quant GEMM over them" + ), + )); + }; + let Some(RouteW::Blocks32(down)) = down.route(dty) else { + return Err(GpuError::shape( + WHAT, + format!("layer {l}: a down stack of {dty}, which no 32-value route GEMM reads"), + )); + }; + k.router.enqueue_router_rows( + stream, + gpu, + k.g32, + w.f32_tensor(WHAT, &r.router)?, + xn, + w.f32_buf(WHAT, &r.bias)?, + scale, + m, + &mut wf.probs, + &mut wf.ids, + &mut wf.weights, + sink, + )?; + k.gemm + .enqueue_route(stream, &wf.ids, m * used, &mut wf.moe, sink)?; + let ff = forms.slots32.k(); + for (w, y) in [(gate, &mut wf.g), (up, &mut wf.u)] { + k.gemm.enqueue_gemm( + stream, + GemmArgs { + ty, + w, + rows_per_expert: ff, + act: &forms.hid128, + route: &wf.moe, + input: GemmInput::Shared { top_k: used }, + y, + }, + )?; + } + k.g32.enqueue_swiglu_act_quant32( + stream, + &wf.g, + &wf.u, + rule, + m * used, + &mut forms.slots32, + sink, + )?; + k.g32.enqueue_gemm32( + stream, + Gemm32Args { + w: down, + rows_per_expert: cx.hidden, + act: &forms.slots32, + route: &wf.moe, + input: GemmInput::PerSlot, + y: &mut wf.down, + }, + )?; + let acc = CardAcc { + down: &wf.down, + w: &wf.weights, + sel: &wf.ids, + n: cx.hidden, + m, + n_card: sigmoid_128_6::N_EXPERT, + n_used: used, + w_pitch: used, + fault: sink, + }; + k.slots.enqueue_card_acc(stream, &acc, &mut wf.routed) +} diff --git a/tools/ref/ptx-shapes.tsv b/tools/ref/ptx-shapes.tsv index 15bca5a8f..23957e87c 100644 --- a/tools/ref/ptx-shapes.tsv +++ b/tools/ref/ptx-shapes.tsv @@ -379,6 +379,7 @@ generate_ds41 raon_hc_mix 0 0 generate_ds41 raon_hc_mean 0 0 generate_ds41 raon_hc_post 0 0 generate_ds41 raon_hc_route 0 112 +generate_ds41 raon_join_qkv 0 0 generate_ds41 relu_sqr 0 0 generate_ds41 relu_sqr_quant32 0 0 generate_ds41 rms_norm 0 0 From 1ba05b037b184b3334fb974b798dd0d0f5d58b81 Mon Sep 17 00:00:00 2001 From: midagedev Date: Thu, 22 Oct 2026 02:16:42 +0900 Subject: [PATCH 03/10] gate_raon1_e2e: the wide prompt against the vendor's logits, two partitions of one prompt bit-identical, and the long set's taps against the steps Co-Authored-By: Claude Opus 5.5 --- crates/gpu-gates/src/bin/gate_raon1_e2e.rs | 357 ++++++++++++++++++--- 1 file changed, 321 insertions(+), 36 deletions(-) diff --git a/crates/gpu-gates/src/bin/gate_raon1_e2e.rs b/crates/gpu-gates/src/bin/gate_raon1_e2e.rs index ec8c7e215..8ebfdf802 100644 --- a/crates/gpu-gates/src/bin/gate_raon1_e2e.rs +++ b/crates/gpu-gates/src/bin/gate_raon1_e2e.rs @@ -25,16 +25,27 @@ //! Then, for each set, over the same file in the verifying load a seat with a draft opens //! (`GpuModel::open_verifying`), as a session in graph mode unless named: //! -//! - `prompt_rows_is_steps`: the prompt as the rows pass (`GpuModel::prompt_rows`) and one id a -//! step: the same argmax, the same logits row by bits, the same steps after it. +//! - `prompt_rows_is_steps`: the first 8 ids of the prompt as the rows pass +//! (`GpuModel::prompt_rows`, a call of up to a pass's rows) and one id a step: the same argmax, +//! the same logits row by bits, the same steps after it. +//! - `prompt_wide_is_vendors`, on the sets of at least [`WIDE_P`] ids (the Rust set and the 1,536-id +//! wikitext one): the whole prompt as the rows pass, whose units of more than 8 rows run as GEMMs +//! (the wide arm), reaches the vendor's argmax, and its last-position logits row is within +//! [`WIDE_REL_BAND`] of the vendor's `result_output`, and of the steps' row at the full length and +//! at each of [`DRIFT_AT`]; a set with a greedy table also holds the named-tie continuation. The +//! steps' own distance from the vendor's is printed beside it. +//! - `prompt_wide_is_partitions`, on the sets past [`SPLIT_AT`] ids: the same prompt as one call and +//! as two cut at [`SPLIT_AT`] gives the last logits row by bits. //! - `verify_is_steps`, in eager and in graph mode: after the prompt call, a verify of the plain //! run's first 8 ids (and of its first 2) gives the plain run's next ids, each row's logits the //! plain step's by bits. //! - `commit_keeps_steps`, in eager and in graph mode: a verify of 8 rows committed at 1, 3 and 8 //! rows, then steps: the plain run's ids. -//! - `taps_are_steps`: the rows entering the draft's tapped layers, as a step taps them, are by -//! bits the rows the rows pass hands its sink over the prompt (every position once, in order) -//! and the rows a verify taps. +//! - `taps_are_steps`: the rows entering the draft's tapped layers, as a step taps them, are the +//! rows the rows pass hands its sink over the prompt (every position once, in order) and the +//! rows a verify taps: by bits for a call of a pass's rows, for the wide units' within +//! [`TAPS_MEDIAN_BAND`] at the median and [`TAPS_P90_BAND`] at the 90th percentile of a layer's rows. The 1,536-id set runs it with a window of [`TAPS_WINDOW`] positions, whose +//! first row lies inside a unit, and prints the spread of its rows a layer and a block of positions. //! - `speculative_is_plain`: Raon-1-mini's DFlash2 draft beside it (`raon1_dflash`: the serve //! seat's pick and open), every pass a proposal of 7 ids verified by one pass of 8 rows: the //! first [`DRAFTED`] greedy ids are the plain session's (its prompt call, then a step an id), id @@ -93,15 +104,69 @@ mod gate { use model::arch::Arch; use model::arch::dflash2::Hparams as DraftHparams; use refset::arch::prompt_tables::{self, Greedy, Top}; - use refset::arch::raon1::{A, B, C, LCPP}; + use refset::arch::raon1::{A, B, C, D, LCPP}; use runtime::{Advance, RowLogits, Speculative, Target, Verify, Want}; /// The sets the step path is held to, in the order they run. const SETS: &[&str] = &[A, B, C]; /// Rows of each latent layer's cache: the longest prompt and its - /// continuation fit. - const CTX: usize = 256; + /// continuation fit (the 1,536-id set included). + const CTX: usize = 2048; + + /// The rows of a pass: a call of up to this many ids runs as the steps do. + const PASS: usize = 8; + const _: () = assert!(PASS == bloomery_gpu::linear::PASS_ROWS); + + /// The median, over the rows of a wide arm's tapped window (or of a verify after it), of a + /// layer's relative distance from the row a step taps. PIN(2026-10-22): 7.0e-2, 2.2 times the + /// worst measured (3.2e-2, layer 22 of the 1,536-id set; 6.5e-3 at layer 2): a wide unit's + /// activation blocks (128 values) and its GEMM sums differ from the step's, and the rounding + /// chains through the layers; a stage wrong for the rows of a unit moves the median to O(1). + const TAPS_MEDIAN_BAND: f64 = 7.0e-2; + + /// The 90th percentile of the same distances. PIN(2026-10-22): 1.5e-1, 2.1 times the worst + /// measured (7.0e-2, the last 256 positions of the 1,536-id set). The largest rows (4.3e-1 a + /// layer) are a heavy tail, likely tokens whose router took another expert on a + /// rounding-sized score gap (not measured): printed, not pinned; a stage wrong for a quarter + /// of the rows moves the 90th percentile to O(1). + const TAPS_P90_BAND: f64 = 1.5e-1; + + /// The tap window of the long set's `taps_are_steps`: its last rows only, from inside a unit + /// (1,536 ids, so the window opens at position 536 of units of 256). + const TAPS_WINDOW: usize = 1000; + + /// The positions a row of `taps_are_steps`'s profile covers. + const TAPS_BLOCK: usize = 256; + + /// The fewest ids of a set `prompt_wide_is_vendors` runs. + const WIDE_P: usize = 64; + + /// The wide arm's last-position logits row's relative distance from the vendor's, and from the + /// steps' row at the same length. PIN(2026-10-22): 8.0e-2, 1.5 times the worst measured (5.3e-2 + /// from the vendor on the 1,536-id set, where the steps' own distance reads 5.0e-2; 3.9e-2 on + /// the 72-id set beside the steps' 4.2e-2; 3.6e-2 from the steps' row at every prefix): a wide + /// unit quantizes its activations in blocks of 128 values where the step's K-quant dots take + /// 256 (and its per-head value projection takes f32 rows), so the arm is a rounding class away + /// from the steps that the vendor's own quantization sits as far from; a stage that reads + /// another row, position or expert is O(1) away. + const WIDE_REL_BAND: f64 = 8.0e-2; + + /// The prefix lengths of a wide set at which the wide row is also held [`WIDE_REL_BAND`] from + /// the steps' row: the drift with the positions is printed, and a jump at a unit's end names + /// a stage that carries state across units. + const DRIFT_AT: &[usize] = &[64, 256, 512, 1024]; + + /// Where `prompt_wide_is_partitions` cuts a prompt: a multiple of a pass's rows, so the + /// attention chunks of both runs hold the same rows, and not of a unit's. + const SPLIT_AT: usize = 600; + const _: () = assert!(SPLIT_AT.is_multiple_of(PASS)); + + /// The fewest ids of a set `prompt_wide_is_partitions` runs: both calls hold a wide unit. + const SPLIT_FROM: usize = SPLIT_AT + PASS + 1; + + /// The sets the wide arm is held to the vendor on: past [`WIDE_P`] ids. + const WIDE_SETS: &[&str] = &[C, D]; /// The file's layers: the last one's route is printed. const LAYERS: usize = 25; @@ -327,16 +392,17 @@ mod gate { Ok((ids, passes, kept)) } - /// Clause `prompt_rows_is_steps` on set `s`: the prompt one id a step and as the rows pass - /// (`GpuModel::prompt_rows`), each from a reset model: the same argmax, the same logits row by - /// bits, and the same [`AFTER`] steps after it. + /// Clause `prompt_rows_is_steps` on set `s`: the prompt's first [`PASS`] ids one id a step and + /// as the rows pass (`GpuModel::prompt_rows`), each from a reset model: the same argmax, the + /// same logits row by bits, and the same [`AFTER`] steps after it. fn prompt_rows_is_steps(m: &mut GpuModel, s: &Set) -> Result { + let head = &s.prompt[..PASS.min(s.prompt.len())]; let mut run = |rows: bool| -> Result<(Vec, Vec), GateError> { m.reset()?; let first = if rows { - m.prompt_rows(&s.prompt)? + m.prompt_rows(head)? } else { - m.step(&s.prompt)? + m.step(head)? }; let row = m.logits()?; let mut ids = vec![first]; @@ -358,6 +424,135 @@ mod gate { Ok(pass) } + /// Clause `prompt_wide_is_vendors` on set `s` (at least [`WIDE_P`] ids): the prompt as the + /// rows pass from a reset model, its units of more than a pass's rows run wide. The argmax is + /// the vendor's top-1, the last logits row is within [`WIDE_REL_BAND`] of the vendor's + /// `result_output` and of the steps' row at the prompt's length and at each of [`DRIFT_AT`] + /// below it, and a set with a greedy table continues as the vendor does, up to a named tie. + /// The steps' own distance from the vendor's is printed beside it. + fn prompt_wide_is_vendors(m: &mut GpuModel, s: &Set) -> Result { + if s.prompt.len() < WIDE_P { + return Err(format!( + "{}: a prompt of {} ids; the wide clause runs from {WIDE_P}", + s.name, + s.prompt.len() + ) + .into()); + } + let theirs = s + .top + .first() + .ok_or_else(|| format!("{}: an empty top-k table", s.name))?; + m.reset()?; + let first = m.prompt_rows(&s.prompt)?; + let row = m.logits()?; + let d = s + .top + .iter() + .map(|t| (row[t.id as usize] - t.logit).abs()) + .fold(0.0f32, f32::max); + let ours5 = top_ids(&row, TOP_K); + let overlap = s + .top + .iter() + .take(TOP_K) + .filter(|t| ours5.contains(&(t.id as usize))) + .count(); + let logits_rel = rel(&row, &s.logits); + let top1 = first == theirs.id; + let logits_ok = logits_rel <= WIDE_REL_BAND; + let mut pass = top1 && logits_ok; + if !s.greedy.ids.is_empty() && s.greedy.margins.len() == s.greedy.ids.len() { + let (_, diverge) = continue_greedy(m, s, first, row.clone())?; + let (ok, at) = match diverge { + None => (true, "none".to_owned()), + Some((i, our_gap)) => { + let margin = s.greedy.margins[i]; + let tie = margin <= 2.0 * d; + ( + tie, + format!( + "step {i} vendor_margin={margin} our_gap={our_gap} 2d={} ({})", + 2.0 * d, + if tie { "a named tie" } else { "past 2d" } + ), + ) + } + }; + println!( + "prompt_wide_greedy set={} first_divergence={at} {}", + s.name, + verdict(ok) + ); + pass &= ok; + } + m.reset()?; + m.step(&s.prompt)?; + let steps_row = m.logits()?; + let (steps_rel, steps_vendor_rel) = (rel(&row, &steps_row), rel(&steps_row, &s.logits)); + let mut drift = vec![(s.prompt.len(), steps_rel)]; + for &n in DRIFT_AT.iter().filter(|&&n| n < s.prompt.len()) { + m.reset()?; + m.prompt_rows(&s.prompt[..n])?; + let wide = m.logits()?; + m.reset()?; + m.step(&s.prompt[..n])?; + drift.push((n, rel(&wide, &m.logits()?))); + } + drift.sort_by_key(|&(n, _)| n); + let drift_ok = drift.iter().all(|&(_, r)| r <= WIDE_REL_BAND); + pass &= drift_ok; + let drift = drift + .iter() + .map(|(n, r)| format!("{n}:{r:.3e}")) + .collect::>() + .join(","); + println!( + "prompt_wide_is_vendors set={} prompt_ids={} ours={first} vendor={} \ + top{TOP_K}_overlap={overlap}/{TOP_K} d={d} logits_rel={logits_rel:.4e} \ + band={WIDE_REL_BAND:.1e} steps_vendor_rel={steps_vendor_rel:.4e} \ + steps_rel_by_ids=[{drift}] {}", + s.name, + s.prompt.len(), + theirs.id, + verdict(pass) + ); + Ok(pass) + } + + /// Clause `prompt_wide_is_partitions` on set `s`: the prompt as one call and as two cut at + /// [`SPLIT_AT`] reach the same last logits row, by bits. A wide row depends on its own id, + /// position and the carried state alone, so a unit or a call that mislays a position, the + /// conv ring, the delta state or the cache rows moves it. + fn prompt_wide_is_partitions(m: &mut GpuModel, s: &Set) -> Result { + if s.prompt.len() < SPLIT_FROM { + return Err(format!( + "{}: a prompt of {} ids; the split clause runs from {}", + s.name, + s.prompt.len(), + SPLIT_FROM + ) + .into()); + } + m.reset()?; + let whole = m.prompt_rows(&s.prompt)?; + let whole_row = m.logits()?; + m.reset()?; + m.prompt_rows(&s.prompt[..SPLIT_AT])?; + let cut = m.prompt_rows(&s.prompt[SPLIT_AT..])?; + let same_row = bits_equal(&whole_row, &m.logits()?); + let pass = whole == cut && same_row; + println!( + "prompt_wide_is_partitions set={} prompt_ids={} split_at={SPLIT_AT} ids_equal={} \ + logits_same_bits={same_row} {}", + s.name, + s.prompt.len(), + whole == cut, + verdict(pass) + ); + Ok(pass) + } + /// A verify of `M` rows of the plain run `plain` from the prompt's end, through the session /// from a reset: whether the prompt call's id is `plain[0]` and the verify's ids are /// `plain[1..=M]`, and whether each row's logits are, by bits, `rows` one id later (`rows[r]` @@ -419,16 +614,48 @@ mod gate { Ok(ok) } + /// A window's rows against the steps': `(position, distance over every value, distance over + /// each tapped layer)`. + type Rels = Vec<(usize, f64, Vec)>; + + /// The median, 90th percentile and largest of `v`. + fn spread(mut v: Vec) -> (f64, f64, f64) { + v.sort_by(f64::total_cmp); + let at = |q: f64| v[((v.len() - 1) as f64 * q).round() as usize]; + (at(0.5), at(0.9), at(1.0)) + } + + /// The rows `ours` holds, from position `from`, against the steps' rows `theirs` (from + /// position 0): per row its relative distance over all `w` values, and over each of the + /// `layers` segments of `w / layers` values. + fn row_rels(ours: &[f32], theirs: &[f32], w: usize, (from, layers): (usize, usize)) -> Rels { + let seg = w / layers; + let rows = ours.chunks_exact(w).zip(theirs[from * w..].chunks_exact(w)); + rows.enumerate() + .map(|(i, (a, b))| { + let per = (0..layers) + .map(|l| rel(&a[l * seg..(l + 1) * seg], &b[l * seg..(l + 1) * seg])) + .collect(); + (from + i, rel(a, b), per) + }) + .collect() + } + /// Clause `taps_are_steps` on set `s`, taps on `layers`: the rows a step taps /// (`GpuModel::draft_taps` of 1 after each id of the prompt, then of the plain run's first - /// [`VERIFY_ROWS`] ids) are, by bits, the rows the rows pass hands its sink over the prompt - /// (`GpuModel::prompt_rows_tapped`, every position once and in order), then the rows a verify - /// of those ids taps (`draft_taps` of [`VERIFY_ROWS`]). The taps are off again after it. + /// [`VERIFY_ROWS`] ids) against the rows the rows pass hands its sink + /// (`GpuModel::prompt_rows_tapped` of the last `window` positions, every one once and in + /// order), then the rows a verify of those ids taps (`draft_taps` of [`VERIFY_ROWS`]). The + /// first [`PASS`] positions of the prompt, a call of a pass's rows, are by bits; the window's + /// rows, whose wide units are not the steps' bits, and the verify after it hold + /// [`TAPS_MEDIAN_BAND`] and [`TAPS_P90_BAND`] a layer; the spread of the window's rows a layer + /// and a [`TAPS_BLOCK`] of positions is printed. + /// The taps are off again after it. fn taps_are_steps( m: &mut GpuModel, s: &Set, plain: &[u32], - layers: &[usize], + (layers, window): (&[usize], usize), ) -> Result { m.set_draft_taps(layers)?; let w = m.draft_tap_width(); @@ -438,31 +665,76 @@ mod gate { m.step(&[id])?; stepped.extend(m.draft_taps(1)?); } - m.reset()?; + let tapped_call = |m: &mut GpuModel, + ids: &[u32], + window: usize| + -> Result<(Vec, bool), GateError> { + m.reset()?; + let (p, from) = (ids.len(), ids.len() - window.min(ids.len())); + let mut sunk = vec![f32::NAN; (p - from) * w]; + let (mut next, mut in_order) = (from, true); + m.prompt_rows_tapped(ids, window, &mut |first: u32, rows: &[f32]| { + let (first, n) = (first as usize, rows.len() / w); + in_order &= first == next && first + n <= p && rows.len() == n * w; + if in_order { + sunk[(first - from) * w..(first + n - from) * w].copy_from_slice(rows); + } + next = first + n; + Ok::<_, GpuError>(()) + })?; + in_order &= next == p; + Ok((sunk, in_order)) + }; + let head = &s.prompt[..PASS.min(s.prompt.len())]; + let (sunk_head, head_in_order) = tapped_call(m, head, head.len())?; + let head_bits = head_in_order && bits_equal(&sunk_head, &stepped[..head.len() * w]); let p = s.prompt.len(); - let mut sunk = vec![f32::NAN; p * w]; - let (mut next, mut in_order) = (0usize, true); - m.prompt_rows_tapped(&s.prompt, p, &mut |first: u32, rows: &[f32]| { - let (first, n) = (first as usize, rows.len() / w); - in_order &= first == next && first + n <= p && rows.len() == n * w; - if in_order { - sunk[first * w..(first + n) * w].copy_from_slice(rows); - } - next = first + n; - Ok::<_, GpuError>(()) - })?; - in_order &= next == p; + let from = p - window.min(p); + let (sunk, in_order) = tapped_call(m, &s.prompt, window)?; m.step_rows::(<[u32; VERIFY_ROWS]>::try_from(&plain[..VERIFY_ROWS])?)?; let verified = m.draft_taps(VERIFY_ROWS)?; m.reset()?; m.set_draft_taps(&[])?; - let prompt = in_order && bits_equal(&sunk, &stepped[..p * w]); - let verify = bits_equal(&verified, &stepped[p * w..]); - let pass = prompt && verify; + let rels = row_rels(&sunk, &stepped, w, (from, layers.len())); + let within = |v: Vec| { + let (med, p90, _) = spread(v); + med <= TAPS_MEDIAN_BAND && p90 <= TAPS_P90_BAND + }; + let prompt_ok = (0..layers.len()).all(|l| within(rels.iter().map(|r| r.2[l]).collect())); + let by_layer = (0..layers.len()) + .map(|l| { + let (med, p90, max) = spread(rels.iter().map(|r| r.2[l]).collect()); + format!("L{}:{med:.1e}/{p90:.1e}/{max:.1e}", layers[l]) + }) + .collect::>() + .join(","); + let mut by_block: Vec<(usize, Vec)> = Vec::new(); + for &(pos, r, _) in &rels { + match by_block.last_mut() { + Some((f, v)) if *f == pos - pos % TAPS_BLOCK => v.push(r), + _ => by_block.push((pos - pos % TAPS_BLOCK, vec![r])), + } + } + let by_block = by_block + .into_iter() + .map(|(f, v)| { + let (med, p90, max) = spread(v); + format!("{f}:{med:.1e}/{p90:.1e}/{max:.1e}") + }) + .collect::>() + .join(","); + let verify = row_rels(&verified, &stepped, w, (p, layers.len())); + let verify_ok = (0..layers.len()).all(|l| within(verify.iter().map(|r| r.2[l]).collect())); + let (vmed, vp90, vmax) = spread(verify.iter().map(|r| r.1).collect()); + let pass = head_bits && in_order && prompt_ok && verify_ok; println!( - "taps_are_steps set={} layers={layers:?} width={w} prompt_rows={p} in_order={in_order} \ - prompt_same_bits={prompt} verify_same_bits={verify} {}", + "taps_are_steps set={} layers={layers:?} width={w} prompt_rows={p} window={} \ + in_order={in_order} head_same_bits={head_bits} \ + bands=med<={TAPS_MEDIAN_BAND:.1e},p90<={TAPS_P90_BAND:.1e} \ + by_layer(med/p90/max)=[{by_layer}] by_{TAPS_BLOCK}_positions=[{by_block}] \ + verify_rows(med/p90/max)={vmed:.1e}/{vp90:.1e}/{vmax:.1e} {}", s.name, + p - from, verdict(pass) ); Ok(pass) @@ -487,7 +759,13 @@ mod gate { ok &= verify_is_steps(&mut seat, s, word)?; } let (plain, _) = greedy(&mut seat, &s.prompt, VERIFY_ROWS, false)?; - ok &= taps_are_steps(seat.model_mut(), s, &plain, &layers)?; + ok &= taps_are_steps(seat.model_mut(), s, &plain, (&layers, s.prompt.len()))?; + Target::reset(&mut seat)?; + } + for name in WIDE_SETS.iter().filter(|&&n| !SETS.contains(&n)) { + let s = open_set(name)?; + let (plain, _) = greedy(&mut seat, &s.prompt, VERIFY_ROWS, false)?; + ok &= taps_are_steps(seat.model_mut(), &s, &plain, (&layers, TAPS_WINDOW))?; Target::reset(&mut seat)?; } Ok(speculative_is_plain(seat, &path, sets)? && ok) @@ -623,6 +901,13 @@ mod gate { ); ok &= top1 && pass && logits_ok; } + for name in WIDE_SETS { + let s = open_set(name)?; + ok &= prompt_wide_is_vendors(&mut m, &s)?; + if s.prompt.len() >= SPLIT_FROM { + ok &= prompt_wide_is_partitions(&mut m, &s)?; + } + } // The plain load's card bytes go back before the verifying load takes its own. drop(m); ok &= verifying(&sets)?; From a9c75e69793effcfc80c6ebbc08a10412d514100 Mon Sep 17 00:00:00 2001 From: midagedev Date: Thu, 22 Oct 2026 02:16:42 +0900 Subject: [PATCH 04/10] unsafe-ratchet: gpu 2101 -> 2104 blocks for raon_join_qkv's new kernel entry Co-Authored-By: Claude Opus 5.5 --- tools/unsafe-ratchet.txt | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tools/unsafe-ratchet.txt b/tools/unsafe-ratchet.txt index 1e29b4bef..d1a1bdfe4 100644 --- a/tools/unsafe-ratchet.txt +++ b/tools/unsafe-ratchet.txt @@ -415,7 +415,10 @@ gguf 7 2 0 # Eleven audited caller blocks discharge that obligation; no device code or launch changes. # PIN(2026-10-22): gpu 2112 -> 2115 blocks, a new kernel entry, not old code relaxed: q6k_gate_up_sel (q6k_sel.rs), the # Q6_K gate·up the common card leg launches for Body35's placed layers. Each block has its SAFETY line. -gpu 2115 247 10 +# PIN(2026-10-22): gpu 2115 -> 2118 blocks, a new kernel entry, not old code relaxed: raon_join_qkv (arch/raon1/wide.rs) +# joins the wide unit's q, k and v rows into one token-major buffer, its source read and its one-writer store (2), and +# the shared-slot load every module family makes (1). Each block has its SAFETY line. +gpu 2118 247 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 826218b1cca420598552d64509b8b40eb12ba46d Mon Sep 17 00:00:00 2001 From: midagedev Date: Thu, 22 Oct 2026 02:19:35 +0900 Subject: [PATCH 05/10] cards: raonpf-pp times Raon-1-mini's prompt on the wide units (pp512, pp4096, RTX 3090) Co-Authored-By: Claude Opus 5.5 --- docs/cards/raonpf-pp.card | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) create mode 100644 docs/cards/raonpf-pp.card diff --git a/docs/cards/raonpf-pp.card b/docs/cards/raonpf-pp.card new file mode 100644 index 000000000..249a2a4cb --- /dev/null +++ b/docs/cards/raonpf-pp.card @@ -0,0 +1,23 @@ +# Raon-1-mini-preview's prompt on the wide units (branch raonpf): the rows pass of up to 256 rows, the grouped GEMM +# for every projection, the common KDA mixer's chunked path for the RDA layers and the route tables for the experts. +# Derivation [derived], RTX 3090, P = 512 in two 256-row units: int8 GEMM 1.56 TFLOP at the ~40 TOPS the +# Qwen3-30B-A3B pp512 row implies, 36..40 ms; the weights' 17.4 GB at ~900 GB/s, 19 ms, under the GEMM's shadow; +# delta 2 ms, latent 7..10 ms, mHC and element-wise 5 ms: 50..57 ms a unit pair, so pp512 about 4,300 tok/s +# (band 3,000..5,800). The 8-row chunks it replaces ran about 1,150 launches a chunk, 64 chunks at P = 512, and read +# 250..300 tok/s. pp4096 about 2,000 tok/s (band 1,000..3,000): the latent attention grows with depth. +# Timeline per unit: card SMs carry the wall; host issue 0.25 ms a layer is under it; PCIe, DRAM and NVMe 0. +kind: baseline +question: does Raon-1-mini-preview's prompt on the wide units reach the derived rate on one RTX 3090? +term: pp tok/s @ P = 512 (prose corpus ids), Raon-1-mini-preview Q4_K_M, bloomery-serve at the serving defaults, + RTX 3090; P = 4096 in the same lease as the second row +unit: tok/s +predict: 3000..5800 +condition: one lease with the box quiet (no gate on either card); depth-serve.sh, one fresh server an arm and a + round, the base arm main's bloomery-serve and the changed arm raonpf's, D 512 and 4096 at C 4160 with N 64; + rows tagged [cpu-busy], [other-busy] or [cold] are rerun +reason: the 0.3.7 release notes name Raon's prompt rate as its known limit (about 250..300 tok/s against the model's + own llama.cpp fork's pp512 2,675 on the same card); this sitting is the change's one measurement before it lands +decide-in: land the wide units; the README's Raon row and the known limit take the measured pp512 and pp4096 +decide-below: under 3,000: a unit term is larger than derived; read one nsys of a P = 512 prompt (nsys-gpu.sh) and + name the stage before landing +decide-above: beyond 5,800: the GEMM term is smaller than derived; land, and re-derive the int8 rate from the row From 8f2f1155284565f6f47094afd55e4a5cdb22a3bd Mon Sep 17 00:00:00 2001 From: midagedev Date: Thu, 22 Oct 2026 04:27:12 +0900 Subject: [PATCH 06/10] raon1: the wide unit's latent absorb is one GEMM over the attn_k_b Q8_0 planes, and its flash attend takes 32 tokens a launch; check_flash takes any m >= 1, the partials sized by the caller The absorb quantizes its activations to int8 for the 32-value GEMM, which moves the e2e gate's D set from 5.32e-2 to 7.61e-2 against the vendor (band 8.0e-2). The flash change moves no bit: 32 tokens a launch equal every cut by bits. raon_row_window is a new kernel entry (add class). The card bytes added are 185,753,600 (partials, unit buffers, form); the e2e gate's resident bytes rose by exactly that. docs/cards/raonpf2-pp.card holds the prediction for the prompt sitting. Co-Authored-By: Claude Opus 5.5 --- crates/gpu-gates/src/bin/gate_raon1_e2e.rs | 6 +- crates/gpu/src/arch/raon1/body.rs | 7 +- crates/gpu/src/arch/raon1/mla.rs | 211 +++++++++---- crates/gpu/src/arch/raon1/wide.rs | 341 ++++++++++++++++++++- crates/gpu/src/flash.rs | 13 +- docs/cards/raonpf2-pp.card | 26 ++ tools/ref/ptx-shapes.tsv | 1 + 7 files changed, 523 insertions(+), 82 deletions(-) create mode 100644 docs/cards/raonpf2-pp.card diff --git a/crates/gpu-gates/src/bin/gate_raon1_e2e.rs b/crates/gpu-gates/src/bin/gate_raon1_e2e.rs index 8ebfdf802..29cd96cd5 100644 --- a/crates/gpu-gates/src/bin/gate_raon1_e2e.rs +++ b/crates/gpu-gates/src/bin/gate_raon1_e2e.rs @@ -157,8 +157,10 @@ mod gate { /// a stage that carries state across units. const DRIFT_AT: &[usize] = &[64, 256, 512, 1024]; - /// Where `prompt_wide_is_partitions` cuts a prompt: a multiple of a pass's rows, so the - /// attention chunks of both runs hold the same rows, and not of a unit's. + /// Where `prompt_wide_is_partitions` cuts a prompt: a multiple of a pass's rows, and of + /// neither a unit's rows nor a latent attention chunk's, so the second call's launches start + /// at other positions than the first call's: a row's bits do not depend on where its launch + /// begins. const SPLIT_AT: usize = 600; const _: () = assert!(SPLIT_AT.is_multiple_of(PASS)); diff --git a/crates/gpu/src/arch/raon1/body.rs b/crates/gpu/src/arch/raon1/body.rs index d6c26e0aa..04445e61f 100644 --- a/crates/gpu/src/arch/raon1/body.rs +++ b/crates/gpu/src/arch/raon1/body.rs @@ -196,6 +196,7 @@ fn wide_dims(hp: &Hparams, kda: &KdaDims, mla: &MlaDims) -> Result= PASS_ROWS && ATTN_CHUNK.is_multiple_of(2)); +const _: () = assert!(WIDE_ROWS.is_multiple_of(ATTN_CHUNK)); + /// Types a dense projection of the layer may have. const DENSE: &[SiteTy] = &[ SiteTy::Q3K, @@ -829,14 +843,15 @@ pub(crate) fn rows( /// Enqueue MLA layer `l` over a wide unit of `m` rows (above [`PASS_ROWS`]), /// in position order: the projections as GEMMs over the unit's rows -/// ([`super::wide`]); the cache append, the gather of the extra part and the -/// output gate once over the unit; the absorb and the latent flash in chunks -/// of [`PASS_ROWS`] query tokens on windows of the unit's buffers (the flash -/// takes no more), each chunk over the cache rows the unit's append wrote, -/// bounded by its own live counts; and a Q6_K per-head value projection as one -/// grouped GEMM over every (row, head) of the unit ([`value_heads`]), a Q8_0 -/// one per chunk. `b`'s buffers hold at least `m` rows of their widths, the -/// partials a pass's rows. +/// ([`super::wide`]); the cache append, the gather of the extra part, the +/// key absorption ([`absorb_heads`]: one GEMM over every (row, head) slot of +/// the unit), a Q6_K per-head value projection ([`value_heads`]: one grouped +/// GEMM over the same slots) and the output gate once over the unit; and the +/// latent flash in chunks of [`ATTN_CHUNK`] query tokens on windows of the +/// unit's buffers, each chunk over the cache rows the unit's append wrote, +/// bounded by its own live counts (a Q8_0 value projection runs per chunk, in +/// windows of [`PASS_ROWS`] rows). `b`'s buffers hold at least `m` rows of +/// their widths, the partials [`ATTN_CHUNK`] rows. pub(super) fn rows_wide( cx: &GqaCx<'_>, k: &MlaKernels, @@ -887,12 +902,7 @@ pub(super) fn rows_wide( let x_in = input(wb.forms, Input::Qb, &*b.qr); wk.project(cx, (&sites.q_b, d.heads * d.head_k()), &x_in, m, b.q)?; } - let (hk, row, hl, hv) = ( - d.heads * d.head_k(), - d.heads * d.row(), - d.heads * LATENT, - d.heads * d.v_head, - ); + let (row, hl, hv) = (d.heads * d.row(), d.heads * LATENT, d.heads * d.v_head); let (f, u) = (size_of::(), size_of::()); let kv_norm = w.f32_buf(WHAT, kv_a_norm)?; k.latent.enqueue_latent_extra_append( @@ -921,31 +931,26 @@ pub(super) fn rows_wide( m * d.heads * d.extra, &mut *b.qabs, )?; - let (qs, dd) = q8_planes(w, &sites.k_b)?; + route_heads(gpu, (wk, wb), (m, fault))?; + let shape = AbsorbShape { + head_k: d.head_k(), + nope: d.nope, + row: d.row(), + }; + absorb_heads( + gpu, + (wk, wb), + (w, &sites.k_b), + (&*b.q, m), + (shape, fault, &mut *b.qabs), + )?; let value = value_w(w, &sites.v_b)?; let q8 = gpu.q8f32(); - for c0 in (0..m).step_by(PASS_ROWS) { - let mc = PASS_ROWS.min(m - c0); - let q = Window::::of(&*b.q, c0 * hk * f, mc * hk)?; + for c0 in (0..m).step_by(ATTN_CHUNK) { + let mc = ATTN_CHUNK.min(m - c0); let n_keys = Window::::of(b.n_keys, (c0 + mc - 1) * u, 1)?; - let mut qabs = WindowMut::::of_mut(&mut *b.qabs, c0 * row * f, mc * row)?; + let qabs = Window::::of(&*b.qabs, c0 * row * f, mc * row)?; let mut att = WindowMut::::of_mut(&mut *b.att, c0 * hl * f, mc * hl)?; - q8.enqueue_q8_0_gemv_heads_mcol( - stream, - Q8_0GemvHeadsMcolArgs { - qs, - d: dd, - x: &q, - rows_per_head: LATENT, - x_head_stride: d.head_k(), - y_head_stride: d.row(), - y_off: d.extra, - m: mc, - x_col_stride: hk, - y_col_stride: row, - y: &mut qabs, - }, - )?; gpu.flash().enqueue_flash_latent_split( stream, FlashSplitArgs { @@ -966,24 +971,30 @@ pub(super) fn rows_wide( y: &mut att, }, )?; + drop(att); if let ValueW::Q8_0(vqs, vd) = value { - let mut av = WindowMut::::of_mut(&mut *b.av, c0 * hv * f, mc * hv)?; - q8.enqueue_q8_0_gemv_heads_mcol( - stream, - Q8_0GemvHeadsMcolArgs { - qs: vqs, - d: vd, - x: &att, - rows_per_head: d.v_head, - x_head_stride: LATENT, - y_head_stride: d.v_head, - y_off: 0, - m: mc, - x_col_stride: hl, - y_col_stride: hv, - y: &mut av, - }, - )?; + for c1 in (0..mc).step_by(PASS_ROWS) { + let mm = PASS_ROWS.min(mc - c1); + let at = c0 + c1; + let x = Window::::of(&*b.att, at * hl * f, mm * hl)?; + let mut av = WindowMut::::of_mut(&mut *b.av, at * hv * f, mm * hv)?; + q8.enqueue_q8_0_gemv_heads_mcol( + stream, + Q8_0GemvHeadsMcolArgs { + qs: vqs, + d: vd, + x: &x, + rows_per_head: d.v_head, + x_head_stride: LATENT, + y_head_stride: d.v_head, + y_off: 0, + m: mm, + x_col_stride: hl, + y_col_stride: hv, + y: &mut av, + }, + )?; + } } } if let ValueW::Q6K(wv) = value { @@ -1025,12 +1036,13 @@ mod tests { use crate::arch::raon1::rda::tests::{Lcg, f32s, q6k, q8_0, resident, same}; use crate::arch::raon1::rda::{RowKernels, RowsKit, add_row_forms}; use crate::arch::raon1::state::PASS_ROWS; + use crate::flash::{FlashGeom, FlashInputs, FlashSplitArgs}; use crate::flash_gqa::FlashGqaKernels; use crate::gemm::Gemm32Kernels; use crate::latent::LATENT; use crate::linear::LinearKernels; use crate::site::{KGemvKernels, SiteTy}; - use crate::tensor::DeviceTensor; + use crate::tensor::{DeviceTensor, Window, WindowMut}; use crate::weights::Weights; use cuda_core::DeviceBuffer; use gguf::quant::f32_to_f16_bits; @@ -1299,4 +1311,91 @@ mod tests { ); } } + + /// The latent flash over a unit's query tokens is the same bits however + /// the tokens are cut into launches: a row's cut of the keys is its + /// group's, and under the single-span regime (a live count within + /// `SEGMENTS` segments of the floor) every group cuts alike. The split + /// attention (a cache past one segment), chunks of the groups' token + /// period and a cut that falls inside a group. + #[test] + #[ignore = "needs a CUDA device; `just gate-gpu-lib` runs it on the box"] + fn hw_flash_tokens_per_launch_are_bits() { + const TOKENS: usize = 32; + const HEIGHT: usize = 1024; + // The last token's live key count. + const LIVE: usize = 600; + let gpu = Gpu::with_device(0).expect("a Gpu on device 0"); + let s = gpu.stream(); + let mut r = Lcg(0x0066_6c61); + let mut cache_rows: Vec = r + .fill(LIVE * D.row(), -1.0, 1.0) + .into_iter() + .map(f32_to_f16_bits) + .collect(); + cache_rows.resize(HEIGHT * D.row(), 0); + let cache = DeviceTensor::upload(s, &cache_rows, HEIGHT, D.row()).expect("the cache"); + let (row, hl) = (D.heads * D.row(), D.heads * LATENT); + let q = r.fill(TOKENS * row, -1.0, 1.0); + let q = DeviceBuffer::from_host(s, &q).expect("the queries up"); + // The tokens cut into launches of `chunks` tokens each, in order. + let run = |chunks: &[usize]| -> Vec { + assert_eq!( + chunks.iter().sum::(), + TOKENS, + "the cut covers the tokens" + ); + let widest = chunks.iter().copied().max().expect("a chunk"); + let (pv, pm) = partials_lens_rows(&D, HEIGHT, widest); + let mut part_v = DeviceBuffer::::zeroed(s, pv).expect("the partials"); + let mut part_ms = DeviceBuffer::::zeroed(s, pm).expect("the partials"); + let mut y = DeviceBuffer::::zeroed(s, TOKENS * hl).expect("the output"); + let mut c0 = 0; + for &chunk in chunks { + let last = u32::try_from(LIVE - (TOKENS - c0 - chunk)).expect("a u32 count"); + let n_keys = DeviceBuffer::from_host(s, &[last]).expect("the live count up"); + let qw = Window::::of(&q, c0 * row * 4, chunk * row).expect("the queries"); + let mut yw = WindowMut::::of_mut(&mut y, c0 * hl * 4, chunk * hl) + .expect("the output window"); + gpu.flash() + .enqueue_flash_latent_split( + s, + FlashSplitArgs { + inputs: FlashInputs { + q: &qw, + kv: &cache, + n_keys_buf: &n_keys, + kq_scale: 1.0 / (D.head_k() as f32).sqrt(), + geom: FlashGeom { + tokens: chunk, + heads: D.heads, + rope_dims: D.extra, + latent_dims: LATENT, + }, + }, + part_v: &mut part_v, + part_ms: &mut part_ms, + y: &mut yw, + }, + ) + .expect("the latent flash"); + c0 += chunk; + } + y.to_host_vec(s).expect("the output back") + }; + let whole = run(&[TOKENS]); + for cut in [ + &[16, 16][..], + &[8, 8, 8, 8], + &[4; 8], + &[6, 26], + &[1, 31], + &[9, 9, 14], + ] { + assert!( + same(&whole, &run(cut), f32::to_bits), + "{TOKENS} tokens in one launch are the same bits cut as {cut:?}" + ); + } + } } diff --git a/crates/gpu/src/arch/raon1/wide.rs b/crates/gpu/src/arch/raon1/wide.rs index 59b83ef1a..cf52cf187 100644 --- a/crates/gpu/src/arch/raon1/wide.rs +++ b/crates/gpu/src/arch/raon1/wide.rs @@ -17,9 +17,10 @@ //! - the RDA layer: [`rda_rows`], the common mixer //! ([`crate::linear::mixer_kda`]) over [`RdaWide`]'s projections; //! - the latent layer: [`super::mla::rows_wide`], the projections wide, the -//! cache append, the gather, the per-head value projection ([`value_heads`]) -//! and the output gate over the unit, the absorb and the latent flash in -//! chunks of [`PASS_ROWS`] query tokens (the flash takes no more); +//! cache append, the gather, the per-head key absorption ([`absorb_heads`]) +//! and value projection ([`value_heads`]) and the output gate over the +//! unit, the latent flash in chunks of [`super::mla::ATTN_CHUNK`] query +//! tokens (the partials' size bounds them); //! - the feed-forward block: [`ffn`], below. use super::body::Scratch; @@ -113,9 +114,52 @@ mod wide_kernels { // SAFETY: i < row·m <= y.len() by the launch contract. unsafe { *y.get_unchecked_mut(i) = x }; } + + /// `rows` windows of `width` values copied between two row-strided + /// buffers: `y[r·y_stride + y_off + c] = x[r·x_stride + x_off + c]` for + /// `r < rows`, `c < width` — a copy, so exact. One thread per value; the + /// host holds each window inside its stride, so no two threads store one + /// slot. + #[kernel] + #[launch_bounds(256)] + #[launch_contract( + domain = 1, + block = (256, 1, 1), + requires = (x.len() >= rows * x_stride, y.len() >= rows * y_stride) + )] + #[allow( + clippy::too_many_arguments, + reason = "kernel entry: the device ABI takes the flat arguments (rust-quality R8)" + )] + pub fn raon_row_window( + x: &[f32], + x_stride: u32, + x_off: u32, + width: u32, + rows: u32, + y_stride: u32, + y_off: u32, + mut y: DisjointSlice, + ) { + let i = thread::index_1d().get(); + let (width, rows) = (width as usize, rows as usize); + if i >= width * rows { + return; + } + let r = i / width; + let c = i - r * width; + // SAFETY: r < rows and c < width, with x_off + width <= x_stride and + // y_off + width <= y_stride (host-checked), put both indices below + // rows·stride, inside x and y by the launch contract; thread i is the + // only writer of its (r, c) slot of y. + unsafe { + *y.get_unchecked_mut(r * y_stride as usize + y_off as usize + c) = + *x.get_unchecked(r * x_stride as usize + x_off as usize + c); + } + } } -/// The module of the join. +/// The module of the join and of the window copy. pub(super) struct JoinKernels { module: wide_kernels::LoadedModule, } @@ -170,6 +214,53 @@ impl JoinKernels { .raon_join_qkv(stream, &prep, q, k, v, wq, wv, m, y)?; Ok(()) } + + /// Enqueue the copy of `rows` windows of `width` values: row `r` of `x` + /// starts at `r · x_stride + x_off`, row `r` of `y` at `r · y_stride + + /// y_off`. A window past its stride, or a buffer short of its rows, is + /// refused by name. + pub(super) fn enqueue_window( + &self, + stream: &CudaStream, + (x, x_stride, x_off): (&DeviceBuffer, usize, usize), + (width, rows): (usize, usize), + (y, y_stride, y_off): (&mut DeviceBuffer, usize, usize), + ) -> Result<(), GpuError> { + let what = "raon1 wide window"; + if width == 0 + || rows == 0 + || x_off + width > x_stride + || y_off + width > y_stride + || x.len() < rows * x_stride + || y.len() < rows * y_stride + { + return Err(GpuError::shape( + what, + format!( + "{rows} windows of {width}: x stride {x_stride} from {x_off} over {} values, \ + y stride {y_stride} from {y_off} over {}", + x.len(), + y.len() + ), + )); + } + let grid = launch_u32(what, "grid", (width * rows).div_ceil(256))?; + let (x_stride, x_off, width, rows, y_stride, y_off) = ( + launch_u32(what, "x_stride", x_stride)?, + launch_u32(what, "x_off", x_off)?, + launch_u32(what, "width", width)?, + launch_u32(what, "rows", rows)?, + launch_u32(what, "y_stride", y_stride)?, + launch_u32(what, "y_off", y_off)?, + ); + let prep = self + .module + .prepare_raon_row_window(LaunchConfig1D::new(grid, 256, 0))?; + self.module.raon_row_window( + stream, &prep, x, x_stride, x_off, width, rows, y_stride, y_off, y, + )?; + Ok(()) + } } /// The activation forms a wide unit's sites read, one buffer per input the @@ -192,6 +283,9 @@ pub(super) struct Forms { /// The latent attention's output, a column a (row, head), into the value /// projection. lat128: GemmAct, + /// The queries' nope parts, a column a (row, head), into the key + /// absorption. + nope32: GemmAct32, } /// The widths a wide unit is cut for. @@ -201,8 +295,10 @@ pub(super) struct WideDims { /// `heads · v_head`: the output projections' input. pub attn: usize, pub q_lora: usize, - /// The latent layers' heads. + /// The latent layers' heads, and the nope values of a query head the key + /// absorption reads. pub heads: usize, + pub nope: usize, pub dense_ff: usize, pub shared_ff: usize, pub expert_ff: usize, @@ -224,6 +320,7 @@ impl Forms { shared32: GemmAct32::new(stream, rows, d.shared_ff)?, slots32: GemmAct32::new(stream, slots, d.expert_ff)?, lat128: GemmAct::new(stream, rows * d.heads, LATENT)?, + nope32: GemmAct32::new(stream, rows * d.heads, d.nope)?, }) } @@ -235,6 +332,7 @@ impl Forms { + rows * (act32_col_bytes(d.q_lora) + act32_col_bytes(d.shared_ff)) + slots * act32_col_bytes(d.expert_ff) + rows * d.heads * act_col_bytes(LATENT) + + rows * d.heads * act32_col_bytes(d.nope) } } @@ -283,16 +381,22 @@ impl WideFfn { } } -/// The latent layers' per-head value projection's table: a slot a (row, head), -/// slot `s` routed to expert `s % heads`, the head's rows of the stack. +/// The latent layers' per-head projections' table: a slot a (row, head), slot +/// `s` routed to expert `s % heads`, the head's rows of the stack, shared by +/// the key absorption and the value projection; and the rows between the +/// absorption's steps. pub(super) struct HeadSlots { heads: usize, ids: DeviceBuffer, route: GemmRoute, + /// The absorption's input, the queries' nope values a slot, and its output, + /// [`LATENT`] values a slot. + nope: DeviceBuffer, + lat: DeviceBuffer, } impl HeadSlots { - fn new(stream: &CudaStream, heads: usize) -> Result { + fn new(stream: &CudaStream, heads: usize, nope: usize) -> Result { let slots = WIDE_ROWS * heads; let ids = (0..slots) .map(|s| launch_u32(WHAT, "head", s % heads)) @@ -301,15 +405,18 @@ impl HeadSlots { heads, ids: DeviceBuffer::from_host(stream, &ids)?, route: GemmRoute::new(stream, slots, heads)?, + nope: DeviceBuffer::zeroed(stream, slots * nope)?, + lat: DeviceBuffer::zeroed(stream, slots * LATENT)?, }) } - fn bytes_of(heads: usize) -> usize { - 4 * WIDE_ROWS * heads + route_bytes(WIDE_ROWS * heads, heads) + fn bytes_of(heads: usize, nope: usize) -> usize { + let slots = WIDE_ROWS * heads; + 4 * slots + route_bytes(slots, heads) + 4 * slots * (nope + LATENT) } fn bytes(&self) -> usize { - self.ids.num_bytes() + self.route.bytes() + self.ids.num_bytes() + self.route.bytes() + self.nope.num_bytes() + self.lat.num_bytes() } } @@ -369,7 +476,7 @@ impl Wide { forms: Forms::new(stream, d)?, parts: [part()?, part()?, part()?], dense: GemmRoute::new(stream, WIDE_ROWS, 1)?, - heads: HeadSlots::new(stream, d.heads)?, + heads: HeadSlots::new(stream, d.heads, d.nope)?, gemm: GemmKernels::load(ctx)?, join: JoinKernels::load(ctx)?, slots: CardSlotKernels::load(ctx)?, @@ -383,7 +490,7 @@ impl Wide { + WideFfn::bytes_of(d) + 3 * 4 * WIDE_ROWS * d.attn + route_bytes(WIDE_ROWS, 1) - + HeadSlots::bytes_of(d.heads) + + HeadSlots::bytes_of(d.heads, d.nope) + stage } @@ -461,6 +568,7 @@ impl Forms { + self.shared32.bytes() + self.slots32.bytes() + self.lat128.bytes() + + self.nope32.bytes() } } @@ -492,6 +600,8 @@ pub(super) enum Input { Out, /// The query latent into `attn_q_b` ([`quantize_qb`]). Qb, + /// The queries' nope parts into the key absorption ([`absorb_heads`]). + Nope, /// The f32 rows alone (an F32 site). Plain, } @@ -502,6 +612,7 @@ pub(super) fn input<'f>(forms: &'f Forms, which: Input, x: &'f DeviceBuffer Input::Hid => (Some(&forms.hid128), Some(&forms.hid32)), Input::Out => (None, Some(&forms.out32)), Input::Qb => (None, Some(&forms.qb32)), + Input::Nope => (None, Some(&forms.nope32)), Input::Plain => (None, None), }; WideIn { @@ -560,11 +671,92 @@ pub(super) fn quantize_hid( Ok(()) } +/// The head table ([`HeadSlots`]) filled for a unit's `m` rows: `m · heads` +/// slots, which [`absorb_heads`] and [`value_heads`] read. +pub(super) fn route_heads( + gpu: &Gpu, + (k, b): (&Kern<'_>, &mut Bufs<'_>), + (m, sink): (usize, FaultSink), +) -> Result<(), GpuError> { + let h = &mut *b.heads; + k.gemm + .enqueue_route(gpu.stream(), &h.ids, m * h.heads, &mut h.route, sink) +} + +/// A named refusal of a head table not filled for exactly `slots` slots. +fn heads_filled(route: &GemmRoute, slots: usize) -> Result<(), GpuError> { + if route.filled() == Some(slots) { + return Ok(()); + } + Err(GpuError::shape( + WHAT, + format!( + "the head table filled for the unit's {slots} slots (route_heads); it is filled for \ + {:?}", + route.filled() + ), + )) +} + +/// The widths of a latent layer's queries the key absorption reads: a head's +/// query row is `head_k` values, its nope part the first `nope`, and the +/// absorbed row `row` values with the latent last. +#[derive(Clone, Copy)] +pub(super) struct AbsorbShape { + pub head_k: usize, + pub nope: usize, + pub row: usize, +} + +/// The latent layers' per-head key absorption over a unit's `m` rows: the nope +/// part of each (row, head) slot of `q` in q8 blocks of 32, a slot routed to +/// its head ([`route_heads`]), times `k_b`, a Q8_0 stack of `heads` experts of +/// [`LATENT`] rows (`attn_k_b`), one 32-value GEMM, the result into each slot's +/// absorbed row `qabs` after its first `row − LATENT` values, which the caller +/// has laid there. +pub(super) fn absorb_heads( + gpu: &Gpu, + (k, b): (&Kern<'_>, &mut Bufs<'_>), + (w, k_b): (&Weights, &str), + (q, m): (&DeviceBuffer, usize), + (shape, sink, qabs): (AbsorbShape, FaultSink, &mut DeviceBuffer), +) -> Result<(), GpuError> { + let AbsorbShape { head_k, nope, row } = shape; + let (stream, h) = (gpu.stream(), &mut *b.heads); + let slots = m * h.heads; + heads_filled(&h.route, slots)?; + k.join.enqueue_window( + stream, + (q, head_k, 0), + (nope, slots), + (&mut h.nope, nope, 0), + )?; + k.g32 + .enqueue_quantize_gemm32(stream, &h.nope, slots, &mut b.forms.nope32, sink)?; + let x_in = input(b.forms, Input::Nope, &h.nope); + site::gemm( + gpu, + &k.kernels(), + (SiteTy::Q8_0, w, SiteWeight::Resident(k_b)), + LATENT, + &x_in, + (&h.route, GemmInput::PerSlot), + slots, + &mut h.lat, + )?; + k.join.enqueue_window( + stream, + (&h.lat, LATENT, 0), + (LATENT, slots), + (qabs, row, row - LATENT), + ) +} + /// The latent layers' per-head value projection over a unit's `m` rows: `att` /// (`m · heads` rows of [`LATENT`] values, the attention's output) times /// `wv`, a Q6_K stack of `heads` experts of `rows` rows (`attn_v_b`), into `av` /// (`m · heads` rows of `rows` values): the rows in q8_1 blocks, a slot a -/// (row, head) routed to its head, one grouped GEMM. +/// (row, head) routed to its head ([`route_heads`]), one grouped GEMM. pub(super) fn value_heads( gpu: &Gpu, (k, b): (&Kern<'_>, &mut Bufs<'_>), @@ -579,9 +771,8 @@ pub(super) fn value_heads( )); }; let (stream, heads) = (gpu.stream(), b.heads.heads); + heads_filled(&b.heads.route, m * heads)?; gpu.enqueue_quantize_gemm(att, m * heads, &mut b.forms.lat128, sink)?; - k.gemm - .enqueue_route(stream, &b.heads.ids, m * heads, &mut b.heads.route, sink)?; k.gemm.enqueue_gemm( stream, GemmArgs { @@ -990,3 +1181,121 @@ fn routed( }; k.slots.enqueue_card_acc(stream, &acc, &mut wf.routed) } + +#[cfg(test)] +mod tests { + use super::{Gemm32Args, Gemm32Kernels, GemmAct32, GemmInput, GemmKernels, GemmRoute, LATENT}; + use crate::Gpu; + use crate::arch::raon1::rda::tests::{Lcg, q8_0, resident}; + use crate::gemm::Gemm32Weight; + use crate::q8f32::Q8_0GemvHeadsMcolArgs; + use crate::tensor::{Window, WindowMut}; + use cuda_core::DeviceBuffer; + use model::arch::raon1::names; + + /// The query's nope values the absorption reads a head: the file's `attn_k_b` is + /// 128 values (4 blocks of 32) a row. + const NOPE: usize = 128; + + /// The key absorption's 32-value GEMM ([`super::absorb_heads`]'s core) over the file's + /// geometry is the f32 gemv's rows ([`crate::q8f32`]'s `q8_0_gemv_heads_mcol`, the pass's + /// absorb) to the activation's quantization: the rows' relative rms distance stays under + /// [`BAND`], where a step, a scale or a head's rows read wrong is O(1). Both the unit's + /// shape (40 heads, 256 rows: the slab entry's launch) and a small one. + #[test] + #[ignore = "needs a CUDA device; `just gate-gpu-lib` runs it on the box"] + fn hw_absorb_gemm_is_the_gemv_to_the_activation_quantization() { + // 8-bit codes of a block of 32 with the block's own scale: ~0.4 % of the rms of a result + // of 128 products; 1.0e-2 is 2.5 times that. + const BAND: f64 = 1.0e-2; + let gpu = Gpu::with_device(0).expect("a Gpu on device 0"); + let (s, ctx) = (gpu.stream(), gpu.context()); + let gemm = GemmKernels::load(ctx).expect("the route kernels"); + let g32 = Gemm32Kernels::load(ctx).expect("the 32-value GEMM"); + for (heads, rows) in [(40, 256), (6, 40)] { + let mut r = Lcg(0x0061_6273 + heads as u64); + let name = names::attn_k_b(3); + let w = resident( + &gpu, + &[q8_0(&mut r, name.clone(), vec![NOPE, LATENT, heads])], + ); + let (qs, d) = w.q8_planes("test", &name).expect("the planes"); + let slots = rows * heads; + let x = r.fill(slots * NOPE, -1.0, 1.0); + let x = DeviceBuffer::from_host(s, &x).expect("the rows up"); + // The gemv, eight tokens a launch. + let mut want = DeviceBuffer::::zeroed(s, slots * LATENT).expect("the output"); + for t0 in (0..rows).step_by(8) { + let mc = 8.min(rows - t0); + let xw = Window::::of(&x, t0 * heads * NOPE * 4, mc * heads * NOPE) + .expect("the rows"); + let mut yw = WindowMut::::of_mut( + &mut want, + t0 * heads * LATENT * 4, + mc * heads * LATENT, + ) + .expect("the output window"); + gpu.q8f32() + .enqueue_q8_0_gemv_heads_mcol( + s, + Q8_0GemvHeadsMcolArgs { + qs, + d, + x: &xw, + rows_per_head: LATENT, + x_head_stride: NOPE, + y_head_stride: LATENT, + y_off: 0, + m: mc, + x_col_stride: heads * NOPE, + y_col_stride: heads * LATENT, + y: &mut yw, + }, + ) + .expect("the gemv"); + } + // The GEMM: a slot a (row, head), slot s on head s % heads. + let mut act = GemmAct32::new(s, slots, NOPE).expect("the form"); + g32.enqueue_quantize_gemm32(s, &x, slots, &mut act, gpu.unlabelled_sink()) + .expect("the quantization"); + let ids: Vec = (0..slots).map(|i| (i % heads) as u32).collect(); + let ids = DeviceBuffer::from_host(s, &ids).expect("the ids up"); + let mut route = GemmRoute::new(s, slots, heads).expect("the table"); + gemm.enqueue_route(s, &ids, slots, &mut route, gpu.unlabelled_sink()) + .expect("the route"); + let mut got = DeviceBuffer::::zeroed(s, slots * LATENT).expect("the output"); + g32.enqueue_gemm32( + s, + Gemm32Args { + w: Gemm32Weight::Q8_0Plane { qs, d }, + rows_per_expert: LATENT, + act: &act, + route: &route, + input: GemmInput::PerSlot, + y: &mut got, + }, + ) + .expect("the GEMM"); + let (got, want) = ( + got.to_host_vec(s).expect("the GEMM back"), + want.to_host_vec(s).expect("the gemv back"), + ); + assert!( + got.iter().chain(&want).all(|v| v.is_finite()), + "{heads} heads, {rows} rows: finite outputs" + ); + let (num, den) = got + .iter() + .zip(&want) + .fold((0.0f64, 0.0f64), |(n, d), (&g, &w)| { + (n + f64::from(g - w).powi(2), d + f64::from(w).powi(2)) + }); + let rel = (num / den).sqrt(); + println!("absorb gemm vs gemv: {heads} heads, {rows} rows: rel rms {rel:.3e}"); + assert!( + rel <= BAND, + "{heads} heads, {rows} rows: the GEMM is {rel:.3e} (rel rms) from the gemv; band {BAND:.1e}" + ); + } + } +} diff --git a/crates/gpu/src/flash.rs b/crates/gpu/src/flash.rs index c8ef1d9dc..a6416a56a 100644 --- a/crates/gpu/src/flash.rs +++ b/crates/gpu/src/flash.rs @@ -833,7 +833,7 @@ mod flash_kernels { } /// Latent flash attention over the absorbed KV cache, `m` query tokens - /// (1..=8) with every head in one launch. `q` holds `m * n_heads` rows + /// (at least 1) with every head in one launch. `q` holds `m * n_heads` rows /// of `width = rope_dims + latent` f32, row `t * n_heads + h` (the /// `kqv_compressed` column order); `kv` is the `[dst_rows x width]` u16 /// cache; the live key count is `n_keys_buf[0]`. Query `t` attends to @@ -1778,8 +1778,11 @@ pub struct FlashKernels { } /// The query geometry every flash launch over the cache shares: `tokens` -/// query tokens (1..=8) of `heads` heads each, so `tokens * heads` query -/// rows of `rope_dims + latent_dims` f32, row `t * heads + h`. +/// query tokens (at least 1) of `heads` heads each, so `tokens * heads` query +/// rows of `rope_dims + latent_dims` f32, row `t * heads + h`. The kernels +/// bound nothing on `tokens`: the split launch's partials do +/// ([`partials_v_len`], [`partials_ms_len`] at `tokens * heads` rows), and a +/// caller sizes them for the most tokens it launches at once. #[derive(Clone, Copy, Debug)] pub struct FlashGeom { pub tokens: usize, @@ -2507,8 +2510,8 @@ fn check_flash( rope_dims, latent_dims: latent, } = geom; - if !(1..=8).contains(&m) { - return Err(GpuError::shape(what, format!("1 <= m <= 8, got {m}"))); + if m == 0 { + return Err(GpuError::shape(what, "m >= 1 query token")); } if n_heads == 0 { return Err(GpuError::shape(what, "n_heads >= 1")); diff --git a/docs/cards/raonpf2-pp.card b/docs/cards/raonpf2-pp.card new file mode 100644 index 000000000..1e8d2fbf9 --- /dev/null +++ b/docs/cards/raonpf2-pp.card @@ -0,0 +1,26 @@ +# Raon-1-mini-preview's prompt: the key absorption as one 32-value GEMM over every (row, head) of a wide unit, +# and the latent flash in launches of 32 query tokens (branch raonpf2 on PR #627's head raonpf). +# Derivation [derived], RTX 3090, P = 512 in two 256-row units: a unit is a = 105.5 ms plus 19.5 us a key of +# depth (from the two measured rows 2,317.3 and 1,759.6 tok/s); the latent loops are 672 of ~1,780 launches a +# unit, 532 of them leave: 532 x f (16..20 us, the chunk path's floor) = 8.5..10.6 ms, less 2.2 ms of work the +# bigger launches carry: 6.3..8.4 ms a unit, 4.5..6.2 % of the wall. docs: raonpf-design.md section 9. +kind: ab +question: does the wide absorb with 32-token flash launches read the derived gain on the prompt's rate, and so + the floor of a small launch (f = 16..20 us)? +term: pp tok/s @ P = 512 (prose corpus ids), Raon-1-mini-preview Q4_K_M, bloomery-serve at the serving + defaults, RTX 3090; P = 4096 in the same lease as the second row +unit: % +predict: +4.0..+10.0 +rounds: 3 +sd: 0.6 % AGENTS.md Know the ruler, same-binary scatter; the ratio at 3 rounds reads +-1.4 % +condition: one lease with the box quiet (no gate on either card); depth-serve.sh, one fresh server an arm and + a round, the base arm PR #627's bloomery-serve (raonpf) and the changed arm raonpf2's, D 512 and 4096 at C + 4160 with N 64; rows tagged [cpu-busy], [other-busy] or [cold] are rerun +decide-in: land: the README's Raon row takes the measured pp512 and pp4096; f is read from the effect (effect + / 532 launches); the next lever is the unit size (a 512-row unit reads each expert once for 24 columns, + section 9.5), read first by one nsys of a P = 512 prompt +decide-below: under +4 %: the launch floor is under 16 us or the 32-token flash launch carries more work than + derived; read one nsys of a P = 512 prompt (nsys-gpu.sh): flash_latent_mma, flash_merge and the absorb GEMM + against the derived 5..8 ms a unit, before landing +decide-above: beyond +10 %: the floor is above 20 us; the 3 remaining small-launch groups (mHC 225, RDA ~14 a + layer, routed FFN ~10 a layer) are the next bulk lever, each at f; land diff --git a/tools/ref/ptx-shapes.tsv b/tools/ref/ptx-shapes.tsv index 23957e87c..2a1a6d954 100644 --- a/tools/ref/ptx-shapes.tsv +++ b/tools/ref/ptx-shapes.tsv @@ -380,6 +380,7 @@ generate_ds41 raon_hc_mean 0 0 generate_ds41 raon_hc_post 0 0 generate_ds41 raon_hc_route 0 112 generate_ds41 raon_join_qkv 0 0 +generate_ds41 raon_row_window 0 0 generate_ds41 relu_sqr 0 0 generate_ds41 relu_sqr_quant32 0 0 generate_ds41 rms_norm 0 0 From 580f9b4a6b48811117854824516f7432ce390688 Mon Sep 17 00:00:00 2001 From: midagedev Date: Thu, 22 Oct 2026 04:27:12 +0900 Subject: [PATCH 07/10] unsafe-ratchet: raon_row_window's one block (gpu 2104 -> 2105) Co-Authored-By: Claude Opus 5.5 --- tools/unsafe-ratchet.txt | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tools/unsafe-ratchet.txt b/tools/unsafe-ratchet.txt index d1a1bdfe4..e7b4504fd 100644 --- a/tools/unsafe-ratchet.txt +++ b/tools/unsafe-ratchet.txt @@ -418,7 +418,9 @@ gguf 7 2 0 # PIN(2026-10-22): gpu 2115 -> 2118 blocks, a new kernel entry, not old code relaxed: raon_join_qkv (arch/raon1/wide.rs) # joins the wide unit's q, k and v rows into one token-major buffer, its source read and its one-writer store (2), and # the shared-slot load every module family makes (1). Each block has its SAFETY line. -gpu 2118 247 10 +# PIN(2026-10-22): gpu 2118 -> 2119 blocks, a new kernel entry, not old code relaxed: raon_row_window +# (arch/raon1/wide.rs) copies a column window of each of a wide unit's rows; its one-writer store has its SAFETY line. +gpu 2119 247 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 43bebb25b2295cbc1647794074f6d8809f46a769 Mon Sep 17 00:00:00 2001 From: midagedev Date: Thu, 22 Oct 2026 04:27:37 +0900 Subject: [PATCH 08/10] cards: raonpf-pp's derivation counts 50..57 ms a 256-row unit, 100..114 ms at P = 512 Co-Authored-By: Claude Opus 5.5 --- docs/cards/raonpf-pp.card | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/cards/raonpf-pp.card b/docs/cards/raonpf-pp.card index 249a2a4cb..625b53cec 100644 --- a/docs/cards/raonpf-pp.card +++ b/docs/cards/raonpf-pp.card @@ -2,9 +2,9 @@ # for every projection, the common KDA mixer's chunked path for the RDA layers and the route tables for the experts. # Derivation [derived], RTX 3090, P = 512 in two 256-row units: int8 GEMM 1.56 TFLOP at the ~40 TOPS the # Qwen3-30B-A3B pp512 row implies, 36..40 ms; the weights' 17.4 GB at ~900 GB/s, 19 ms, under the GEMM's shadow; -# delta 2 ms, latent 7..10 ms, mHC and element-wise 5 ms: 50..57 ms a unit pair, so pp512 about 4,300 tok/s -# (band 3,000..5,800). The 8-row chunks it replaces ran about 1,150 launches a chunk, 64 chunks at P = 512, and read -# 250..300 tok/s. pp4096 about 2,000 tok/s (band 1,000..3,000): the latent attention grows with depth. +# delta 2 ms, latent 7..10 ms, mHC and element-wise 5 ms: 50..57 ms a 256-row unit, 100..114 ms at P = 512, so pp512 +# about 4,300 tok/s (band 3,000..5,800). The 8-row chunks it replaces ran about 1,150 launches a chunk, 64 chunks at P = +# 512, and read 250..300 tok/s. pp4096 about 2,000 tok/s (band 1,000..3,000): the latent attention grows with depth. # Timeline per unit: card SMs carry the wall; host issue 0.25 ms a layer is under it; PCIe, DRAM and NVMe 0. kind: baseline question: does Raon-1-mini-preview's prompt on the wide units reach the derived rate on one RTX 3090? From 9e8ca16368774a7b20cf9e74023a38827d8381f0 Mon Sep 17 00:00:00 2001 From: midagedev Date: Thu, 22 Oct 2026 05:30:28 +0900 Subject: [PATCH 09/10] raon1: the wide arm is decided at load from the refusals it raises, so a file it cannot run keeps the rows pass at every prompt length; non-finite taps fail the band check; the wide gate rows hold the shared block - Body::wide_ok: the load runs the wide arm's own refusals (wide::rda_runs, MlaSites::wide_runs, wide::ffn_runs) over every layer. A file with, for instance, a Q6_K attn_output runs a 9-id prompt on the rows pass. Routed gate and up stacks of two types are refused by name. - gate_raon1_e2e: a tap distance that is not finite fails within_bands before the percentiles (taps_band_separation), and wide_runs_the_file holds the real file on the wide arm. - wide_ff_rows includes the shared experts' width, as the rows pass does. - synth.rs: a lib test builds a three-layer file at load, a Q6_K output projection against a Q8_0 control. FAIL-first: the loop without wide_ok turns the synthetic test red; the old ff_rows formula turns a_wide_units_gate_rows_hold_the_shared_block red. Co-Authored-By: Claude Opus 5.5 --- crates/gpu-gates/src/bin/gate_raon1_e2e.rs | 72 +++- crates/gpu/src/arch/raon1/body.rs | 202 ++++++++++- crates/gpu/src/arch/raon1/mla.rs | 49 ++- crates/gpu/src/arch/raon1/mod.rs | 2 + crates/gpu/src/arch/raon1/rda.rs | 5 + crates/gpu/src/arch/raon1/rows.rs | 9 +- crates/gpu/src/arch/raon1/synth.rs | 354 ++++++++++++++++++ crates/gpu/src/arch/raon1/wide.rs | 399 ++++++++++++++++++--- 8 files changed, 1035 insertions(+), 57 deletions(-) create mode 100644 crates/gpu/src/arch/raon1/synth.rs diff --git a/crates/gpu-gates/src/bin/gate_raon1_e2e.rs b/crates/gpu-gates/src/bin/gate_raon1_e2e.rs index 29cd96cd5..3e4ccf7ec 100644 --- a/crates/gpu-gates/src/bin/gate_raon1_e2e.rs +++ b/crates/gpu-gates/src/bin/gate_raon1_e2e.rs @@ -36,6 +36,9 @@ //! steps' own distance from the vendor's is printed beside it. //! - `prompt_wide_is_partitions`, on the sets past [`SPLIT_AT`] ids: the same prompt as one call and //! as two cut at [`SPLIT_AT`] gives the last logits row by bits. +//! - `wide_runs_the_file`, on the plain load: the load found every layer of the file runnable by +//! the wide arm (`Body::wide_ok`), so the wide clauses run the wide arm and not the passes a +//! file it cannot run takes. //! - `verify_is_steps`, in eager and in graph mode: after the prompt call, a verify of the plain //! run's first 8 ids (and of its first 2) gives the plain run's next ids, each row's logits the //! plain step's by bits. @@ -51,6 +54,10 @@ //! first [`DRAFTED`] greedy ids are the plain session's (its prompt call, then a step an id), id //! for id. The passes, and the proposed ids the target kept, are printed beside it. //! +//! Before the card loads, the band check of `taps_are_steps` is held to what it refuses: eight +//! rows at a distance inside the bands pass it, and the same rows with one tap of the last row NaN +//! do not. +//! //! A step that raises the card's fault word ends the gate by name. A set //! absent from the box, dumped from another file or build, or unfinished ends //! the gate by name before the card loads. @@ -643,6 +650,50 @@ mod gate { .collect() } + /// Whether one layer's rows' distances `v` hold the bands: every one finite, then the median + /// at most [`TAPS_MEDIAN_BAND`] and the 90th percentile at most [`TAPS_P90_BAND`]. A NaN tap + /// is an infinite distance ([`rel`]), and the percentiles of a few rows skip the largest + /// (eight rows: the median reads the fifth, the 90th percentile the seventh), so finiteness + /// is its own condition. + fn within_bands(v: Vec) -> bool { + if v.iter().any(|r| !r.is_finite()) { + return false; + } + let (med, p90, _) = spread(v); + med <= TAPS_MEDIAN_BAND && p90 <= TAPS_P90_BAND + } + + /// The band check held to what it refuses, on the host before the card loads: [`VERIFY_ROWS`] + /// rows of two layers' taps at a distance of 1 % of their rows' size pass it in both layers, + /// and the same rows with one tap of the last row NaN do not pass it in that tap's layer. + fn within_bands_separate() -> Result<(), GateError> { + const LAYERS_TAPPED: usize = 2; + const SEG: usize = 4; + let w = LAYERS_TAPPED * SEG; + let theirs: Vec = (0..VERIFY_ROWS * w).map(|i| 1.0 + (i % 7) as f32).collect(); + let ours: Vec = theirs.iter().map(|v| v * 1.01).collect(); + let held = |ours: &[f32], layer: usize| { + let rels = row_rels(ours, &theirs, w, (0, LAYERS_TAPPED)); + within_bands(rels.iter().map(|r| r.2[layer]).collect()) + }; + let clean = (0..LAYERS_TAPPED).all(|l| held(&ours, l)); + let mut nan = ours.clone(); + nan[(VERIFY_ROWS - 1) * w + SEG] = f32::NAN; + let refused = !held(&nan, 1) && held(&nan, 0); + println!( + "taps_band_separation rows={VERIFY_ROWS} clean_held={clean} nan_last_row_refused={refused} {}", + verdict(clean && refused) + ); + if clean && refused { + Ok(()) + } else { + Err(format!( + "the taps' band check holds clean rows {clean} and refuses a NaN tap {refused}" + ) + .into()) + } + } + /// Clause `taps_are_steps` on set `s`, taps on `layers`: the rows a step taps /// (`GpuModel::draft_taps` of 1 after each id of the prompt, then of the plain run's first /// [`VERIFY_ROWS`] ids) against the rows the rows pass hands its sink @@ -698,11 +749,8 @@ mod gate { m.reset()?; m.set_draft_taps(&[])?; let rels = row_rels(&sunk, &stepped, w, (from, layers.len())); - let within = |v: Vec| { - let (med, p90, _) = spread(v); - med <= TAPS_MEDIAN_BAND && p90 <= TAPS_P90_BAND - }; - let prompt_ok = (0..layers.len()).all(|l| within(rels.iter().map(|r| r.2[l]).collect())); + let prompt_ok = + (0..layers.len()).all(|l| within_bands(rels.iter().map(|r| r.2[l]).collect())); let by_layer = (0..layers.len()) .map(|l| { let (med, p90, max) = spread(rels.iter().map(|r| r.2[l]).collect()); @@ -726,7 +774,8 @@ mod gate { .collect::>() .join(","); let verify = row_rels(&verified, &stepped, w, (p, layers.len())); - let verify_ok = (0..layers.len()).all(|l| within(verify.iter().map(|r| r.2[l]).collect())); + let verify_ok = + (0..layers.len()).all(|l| within_bands(verify.iter().map(|r| r.2[l]).collect())); let (vmed, vp90, vmax) = spread(verify.iter().map(|r| r.1).collect()); let pass = head_bits && in_order && prompt_ok && verify_ok; println!( @@ -800,6 +849,7 @@ mod gate { } pub fn run() -> Result<(), GateError> { + within_bands_separate()?; let sets: Vec = SETS .iter() .map(|&n| open_set(n)) @@ -839,6 +889,16 @@ mod gate { .into()); } let mut ok = true; + { + let body = m.body("gate_raon1_e2e")?; + let wide_ok = body.wide_ok(); + println!( + "wide_runs_the_file wide_ok={wide_ok} refusal={:?} {}", + body.wide_off(), + verdict(wide_ok) + ); + ok &= wide_ok; + } for s in &sets { m.reset()?; let first = m.step(&s.prompt)?; diff --git a/crates/gpu/src/arch/raon1/body.rs b/crates/gpu/src/arch/raon1/body.rs index 04445e61f..9a33134b6 100644 --- a/crates/gpu/src/arch/raon1/body.rs +++ b/crates/gpu/src/arch/raon1/body.rs @@ -44,6 +44,7 @@ use cuda_core::{CudaStream, DeviceBuffer}; use gguf::Split; use model::arch::Arch; use model::arch::models::KdaGate; +use model::arch::noaux::Experts; use model::arch::raon1::hparams::{Hparams, Kind}; use model::arch::raon1::names::{self, Sub}; use model::placement::CardFormat; @@ -109,6 +110,28 @@ impl Layer { Mixer::Mla { cache, .. } => cache.buf().num_bytes(), } } + + /// `Ok` when the wide arm runs layer `l`: neither its mixer's projections + /// nor its feed-forward block of resident weights `w` meet a refusal the wide + /// arm raises ([`Body::wide_ok`]). + fn wide_runs(&self, l: usize, w: &Weights, d: &WideDims) -> Result<(), GpuError> { + let mixer = match &self.mixer { + Mixer::Rda { sites, .. } => wide::rda_runs(sites), + Mixer::Mla { sites, .. } => sites.wide_runs(), + }; + mixer.map_err(|refusal| GpuError::shape(WHAT, format!("layer {l}'s mixer: {refusal}")))?; + wide::ffn_runs(l, &self.ffn, w, d) + } +} + +/// The first refusal the wide arm raises over `layers` in file order, or `None` +/// when it runs every one ([`Body::wide_ok`]). +fn wide_refusal(layers: &[Layer], w: &Weights, d: &WideDims) -> Option { + layers + .iter() + .enumerate() + .find_map(|(l, layer)| layer.wide_runs(l, w, d).err()) + .map(|refusal| refusal.to_string()) } /// The kernels the step launches beyond the `Gpu`'s own. @@ -200,10 +223,23 @@ fn wide_dims(hp: &Hparams, kda: &KdaDims, mla: &MlaDims) -> Result usize { + dense_ff.max(e.shared_ff()) +} + +/// The gate and up rows' width of a wide unit's feed-forward buffers: the widest +/// dense-form block, or the routed slots' (`N_USED` experts of `ff`). +fn wide_ff_rows(dense_ff: usize, e: &Experts) -> usize { + dense_block_ff(dense_ff, e).max(ffn::sigmoid_128_6::N_USED * e.ff) +} + /// A walk's buffers, `rows` rows of each width, token-major: the step's /// one row, or a pass's [`PASS_ROWS`]. pub(super) struct Scratch { @@ -430,6 +466,9 @@ pub struct Body { stage: DeviceBuffer, /// A wide unit's rows, forms and kernels ([`super::wide`]). pub(super) wide: Wide, + /// What keeps the wide arm off this file's prompts, found at load, or `None` + /// when it runs every layer ([`Body::wide_ok`]). + wide_off: Option, /// The passes' staging, the lane word, the taps and the position the /// stores stand at. pub(super) rows: RowState, @@ -546,6 +585,30 @@ impl Body { Ok(self.s.hc.to_host_vec(stream)?) } + /// Whether a prompt of more than [`PASS_ROWS`] ids runs as wide units + /// ([`super::wide`]), decided once at load over every layer. The wide arm + /// refuses a weight the loader admits — an output projection or `attn_q_b` + /// that is not Q8_0, a forget projection that is not F32, a dense block or + /// shared expert whose down has no form at its width (a K-quant down over + /// the dense width, a Q8_0 down over the shared width), routed stacks that + /// are not one grouped K-quant type with a down the 32-value GEMM reads + /// ([`wide::rda_runs`], `MlaSites::wide_runs`, [`wide::ffn_runs`]). One + /// such layer keeps the whole file on the rows pass for prompts of every + /// length, where a prompt past [`PASS_ROWS`] ids would otherwise stop at + /// that layer by name. + #[must_use] + pub fn wide_ok(&self) -> bool { + self.wide_off.is_none() + } + + /// The refusal the load found that keeps the wide arm off this file's + /// prompts, the first layer's in file order; `None` when + /// [`Body::wide_ok`]. + #[must_use] + pub fn wide_off(&self) -> Option<&str> { + self.wide_off.as_deref() + } + fn load( gpu: &Gpu, file: &Split, @@ -656,7 +719,9 @@ impl Body { &mut acts, layers.iter().flat_map(Layer::mixer_sites), )?; - let ff = (hp.dense_ff.max(e.shared_ff()), e.ff); + let ff = (dense_block_ff(hp.dense_ff, &e), e.ff); + let wd = wide_dims(&hp, &kda, &mla)?; + let wide_off = wide_refusal(&layers, w, &wd); let stage = rda::stage_len(&kda) .max(mla::stage_len(&mla)) .max(ffn::stage_len(ff.0.max(ff.1))); @@ -670,9 +735,10 @@ impl Body { stage: DeviceBuffer::zeroed(stream, stage)?, wide: Wide::new( gpu, - &wide_dims(&hp, &kda, &mla)?, + &wd, Scratch::new(stream, &hp, (&kda, &mla), ctx, WIDE_ROWS)?, )?, + wide_off, rows: RowState::new(stream, lanes)?, kda, mla, @@ -738,6 +804,7 @@ impl Body { stage, rows, wide, + wide_off, ctx, .. } = self; @@ -791,6 +858,12 @@ impl Body { format!("a {pass:?} pass of {m} rows: only a prompt's units run wide"), )); } + if let Some(why) = wide_off { + return Err(GpuError::shape( + STEP, + format!("a wide unit of {m} rows with the wide arm off: {why}"), + )); + } wide.route_dense(gpu, m)?; let (kit, ws, wrin) = wide.kit(&k.g32, &k.ffn.router); let RowViews { @@ -1101,3 +1174,126 @@ impl ChainBody for Body { None } } + +#[cfg(test)] +mod tests { + use super::ffn::sigmoid_128_6::{N_EXPERT, N_USED}; + use super::mla::tests::typed_sites; + use super::{Ffn, Glu, Layer, Mixer, STREAMS, Site, SiteTy, dims, wide_dims, wide_refusal}; + use crate::weights::Weights; + use crate::{DeviceTensor, Gpu}; + use model::arch::kda::Keys; + use model::arch::models::KdaGate; + use model::arch::noaux::Experts; + use model::arch::raon1::hparams::{Hc, Hparams, Kind, Mla, Situ}; + + /// A header with Raon's widths and a dense width, a routed expert's width and a + /// count of shared experts of the caller's. + fn header(dense_ff: usize, ff: usize, n_shared: usize) -> Hparams { + Hparams { + n_layer: 3, + n_embd: 2560, + n_head: 40, + kda: Keys { + conv: 4, + head_dim: 128, + }, + decay_rank: 128, + gate: KdaGate::Sigmoid { lower_bound: -5.0 }, + mla: Mla { + q_lora: 576, + kv_lora: 512, + head_k: 192, + head_v: 128, + pe: 64, + }, + rms_eps: 1e-5, + latent_eps: 1e-6, + hc: Hc { + streams: STREAMS, + eps: 1e-5, + magnitude: 2.0, + }, + situ: Situ { + gate_beta: 4.0, + value_beta: 4.0, + }, + n_vocab: 1024, + n_ctx_train: 4096, + dense_lead: 1, + dense_ff, + experts: Experts { + n_expert: N_EXPERT, + n_used: N_USED, + ff, + n_shared, + weights_scale: 1.0, + weights_norm: true, + }, + kinds: vec![Kind::Rda, Kind::Rda, Kind::Mla], + defaults: Vec::new(), + } + } + + /// A wide unit's gate and up rows hold the widest block that writes them: the + /// dense layer's, a routed layer's shared experts' (`n_shared` of `ff` run as + /// one block) and the routed slots' (`N_USED` of `ff`). Seven shared experts of + /// 256 beside a dense width of 1024 are 1792 values, past the dense width and + /// the slots' 1536. The buffers' size and the plan's reserve + /// ([`Body::pass_reserve`]) both read these widths. + #[test] + fn a_wide_units_gate_rows_hold_the_shared_block() { + for (dense_ff, ff, n_shared) in [(1024, 256, 7), (12032, 1088, 1), (512, 128, 1)] { + let hp = header(dense_ff, ff, n_shared); + let (kda, mla) = dims(&hp).expect("the header's widths"); + let d = wide_dims(&hp, &kda, &mla).expect("the wide arm's widths"); + for (what, width) in [ + ("dense", dense_ff), + ("shared", ff * n_shared), + ("routed slots", N_USED * ff), + ] { + assert!( + d.ff_rows >= width, + "dense {dense_ff}, {n_shared} shared of {ff}: gate rows of {} values hold \ + the {what} block of {width}", + d.ff_rows + ); + } + } + } + + /// One layer the wide arm refuses keeps the whole file off it: a latent layer + /// whose output projection is Q6_K among Q8_0 ones is the first refusal, named + /// by its layer, and a file of Raon's types has none. + #[test] + #[ignore = "needs a CUDA device; `just gate-gpu-lib` runs it on the box"] + fn hw_one_layer_with_a_q6k_output_keeps_the_file_off_the_wide_arm() { + let gpu = Gpu::with_device(0).expect("a Gpu on device 0"); + let hp = header(1024, 256, 1); + let (kda, mla) = dims(&hp).expect("the header's widths"); + let d = wide_dims(&hp, &kda, &mla).expect("the wide arm's widths"); + let site = |name: &str, ty: SiteTy, k: usize| Site::new(name.to_owned(), ty, k); + let layer = |out: SiteTy| Layer { + mixer: Mixer::Mla { + sites: Box::new(typed_sites(SiteTy::Q8_0, out)), + cache: DeviceTensor::zeroed(gpu.stream(), 1, mla.row()).expect("a cache row"), + }, + ffn: Ffn::Dense(Glu { + gate: site("gate", SiteTy::Q4K, hp.n_embd), + up: site("up", SiteTy::Q4K, hp.n_embd), + down: site("down", SiteTy::Q6K, hp.dense_ff), + ff: hp.dense_ff, + }), + hc: std::array::from_fn(|_| std::array::from_fn(|_| String::new())), + }; + let w = Weights::default(); + let raon = [layer(SiteTy::Q8_0), layer(SiteTy::Q8_0)]; + assert_eq!(wide_refusal(&raon, &w, &d), None); + let mixed = [layer(SiteTy::Q8_0), layer(SiteTy::Q6K), layer(SiteTy::F32)]; + let refusal = wide_refusal(&mixed, &w, &d).expect("a refusal"); + assert!( + refusal.contains("layer 1's mixer") && refusal.contains("an output projection of Q6_K"), + "{refusal}" + ); + } +} diff --git a/crates/gpu/src/arch/raon1/mla.rs b/crates/gpu/src/arch/raon1/mla.rs index 658529039..ad5bf1ace 100644 --- a/crates/gpu/src/arch/raon1/mla.rs +++ b/crates/gpu/src/arch/raon1/mla.rs @@ -53,8 +53,8 @@ use super::rda::{RowsKit, kgemv_rows, project_rows}; use super::state::PASS_ROWS; use super::wide::{ - AbsorbShape, Input, WIDE_ROWS, WideKit, absorb_heads, input, quantize_hid, quantize_out, - quantize_qb, route_heads, value_heads, + AbsorbShape, Input, WIDE_ROWS, WideKit, absorb_heads, input, out_runs, qb_runs, quantize_hid, + quantize_out, quantize_qb, route_heads, value_heads, }; use crate::arch::hybrid::{Acts, GqaCx, Site}; use crate::fault::FaultSink; @@ -234,6 +234,14 @@ impl MlaSites { pub(crate) fn sites(&self) -> [&Site; 5] { [&self.q_a, &self.kv_a, &self.gate, &self.q_b, &self.out] } + + /// `Ok` when the wide arm runs the layer's projections ([`rows_wide`]), else + /// the refusal it raises: `attn_q_b` and the output projection read q8 + /// blocks of 32. + pub(super) fn wide_runs(&self) -> Result<(), GpuError> { + qb_runs(self.q_b.ty())?; + out_runs(self.out.ty()) + } } /// The kernels every MLA layer launches beyond the `Gpu`'s own, and the @@ -1029,7 +1037,7 @@ pub(super) fn rows_wide( } #[cfg(test)] -mod tests { +pub(super) mod tests { use super::{MlaBufs, MlaDims, MlaKernels, MlaSites, ValueHeads, partials_lens_rows, rows}; use crate::Gpu; use crate::arch::hybrid::{Acts, GqaCx, Site}; @@ -1105,6 +1113,41 @@ mod tests { (w, sites) } + /// A layer's sites with `q_b` and `out` of the given types, the others Raon's. + pub(crate) fn typed_sites(q_b: SiteTy, out: SiteTy) -> MlaSites { + let site = |name: String, ty: SiteTy, k: usize| Site::new(name, ty, k); + let (e, hv) = (D.embd, D.heads * D.v_head); + MlaSites { + q_a: site(names::attn_q_a(L), SiteTy::Q6K, e), + kv_a: site(names::attn_kv_a_mqa(L), SiteTy::Q6K, e), + gate: site(names::attn_gate(L), SiteTy::Q8_0, e), + q_b: site(names::attn_q_b(L), q_b, D.q_lora), + out: site(names::attn_output(L), out, hv), + k_b: names::attn_k_b(L), + v_b: ValueHeads::Q6K(names::attn_v_b(L)), + vecs: [ + names::attn_norm(L), + names::attn_q_a_norm(L), + names::attn_kv_a_norm(L), + ], + } + } + + /// The wide arm runs a latent layer whose `attn_q_b` and output projection + /// are Q8_0 (Raon's), and refuses by name one whose either is another type: + /// a Q6_K output projection reads q8_1 blocks of 128, and the wide arm + /// quantizes the gated rows in blocks of 32 alone. + #[test] + fn a_latent_layer_runs_wide_only_with_q8_0_query_up_and_output() { + assert!(typed_sites(SiteTy::Q8_0, SiteTy::Q8_0).wide_runs().is_ok()); + for ty in [SiteTy::Q6K, SiteTy::Q4K, SiteTy::F32] { + let e = typed_sites(SiteTy::Q8_0, ty).wide_runs().unwrap_err(); + assert!(e.to_string().contains("an output projection"), "{e}"); + let e = typed_sites(ty, SiteTy::Q8_0).wide_runs().unwrap_err(); + assert!(e.to_string().contains("a query up projection"), "{e}"); + } + } + /// The buffers of a pass of up to [`PASS_ROWS`] rows over a cache of /// `ctx` rows, each its width times [`PASS_ROWS`]: the input, the /// positions and live counts, [`MlaBufs`]' f32 buffers in its order from diff --git a/crates/gpu/src/arch/raon1/mod.rs b/crates/gpu/src/arch/raon1/mod.rs index 62ebf6ea6..528b44dcd 100644 --- a/crates/gpu/src/arch/raon1/mod.rs +++ b/crates/gpu/src/arch/raon1/mod.rs @@ -14,6 +14,8 @@ pub mod mla; pub mod rda; mod rows; mod state; +#[cfg(test)] +mod synth; mod wide; pub use body::Body; diff --git a/crates/gpu/src/arch/raon1/rda.rs b/crates/gpu/src/arch/raon1/rda.rs index 3ae6c03cc..a7ae233da 100644 --- a/crates/gpu/src/arch/raon1/rda.rs +++ b/crates/gpu/src/arch/raon1/rda.rs @@ -617,6 +617,11 @@ pub(super) mod tests { /// contiguous axis first) and bytes. pub(crate) struct Part(String, GgmlType, Vec, Vec); + /// A file tensor of `ty` over `dims` whose bytes are `bytes` as they stand. + pub(crate) fn raw(name: String, ty: GgmlType, dims: Vec, bytes: Vec) -> Part { + Part(name, ty, dims, bytes) + } + /// Q6_K rows over `dims`: random codes, sub-block scales in −4..=4 and a /// small `d` (high byte 0x14), so every weight is finite and small. pub(crate) fn q6k(r: &mut Lcg, name: String, dims: Vec) -> Part { diff --git a/crates/gpu/src/arch/raon1/rows.rs b/crates/gpu/src/arch/raon1/rows.rs index dcb77d0d2..080683cbc 100644 --- a/crates/gpu/src/arch/raon1/rows.rs +++ b/crates/gpu/src/arch/raon1/rows.rs @@ -6,7 +6,8 @@ //! ([`super::state`]); //! - a prompt call ([`GpuModel::prompt_rows`]): wide units of up to //! [`WIDE_ROWS`] rows ([`super::wide`]) while at least [`WIDE_FROM`] rows -//! remain, then the rest's ids staged on the card [`GROUP`] at a time, each +//! remain and the file runs wide ([`Body::wide_ok`]), then the rest's ids +//! staged on the card [`GROUP`] at a time, each //! chunk of up to [`PASS_ROWS`] of them one walk in place, the head on the //! call's last row alone; //! - the taps ([`GpuModel::set_draft_taps`]): the streams entering the @@ -308,7 +309,9 @@ impl Rollback for Body { impl Body { /// A prompt call's walks from `pos` ([`GpuModel::prompt_rows`]): wide /// units of up to [`WIDE_ROWS`] rows ([`super::wide`]) as long as at - /// least [`WIDE_FROM`] rows remain, then the rest with its ids staged a + /// least [`WIDE_FROM`] rows remain and the file runs wide + /// ([`Body::wide_ok`]; a file that does not takes the rest's path for + /// every row), then the rest with its ids staged a /// group at a time, each chunk of up to [`PASS_ROWS`] a walk in place /// (every row the step's bits). The head is on the call's last row; with /// `tapped`, the taps of the call's last `window` rows are handed to its @@ -333,7 +336,7 @@ impl Body { .as_ref() .map_or(n, |(window, _)| n - (*window).min(n)); let mut done = 0; - while n - done >= WIDE_FROM { + while self.wide_ok() && n - done >= WIDE_FROM { let m = WIDE_ROWS.min(n - done); let at = pos + launch_u32(WHAT, "row", done)?; self.wide.rin.stage(stream, &ids[done..done + m], at)?; diff --git a/crates/gpu/src/arch/raon1/synth.rs b/crates/gpu/src/arch/raon1/synth.rs new file mode 100644 index 000000000..cbc53082c --- /dev/null +++ b/crates/gpu/src/arch/raon1/synth.rs @@ -0,0 +1,354 @@ +//! A synthetic Raon-1 file at small widths, written whole and loaded as the engine loads a model +//! file: the case the real file does not hold, a weight the loader admits that the wide arm +//! refuses (a latent layer's Q6_K output projection), and Raon's own types beside it. The +//! widths are the file's where a kernel fixes them (the heads, the latent, the experts) and +//! small where none does; the weights are random with small scales, so every value is finite. + +use super::body::Body; +use super::rda::tests::same; +use crate::model::GpuModel; +use crate::{Gpu, GpuError}; +use gguf::quant::Q8Block; +use gguf::write::{Layout, TensorDecl, Writer}; +use gguf::{GgmlType, Split, Value}; +use std::path::PathBuf; + +/// Two Q6_K super-blocks a row: the K-quant gemv's even count. +const E: usize = 512; +const HEADS: usize = 40; +const HEAD: usize = 128; +const INNER: usize = HEADS * HEAD; +const CONV: usize = 4; +const RANK: usize = 128; +const Q_LORA: usize = 256; +const NOPE: usize = 128; +const PE: usize = 64; +const KV_LORA: usize = 512; +const FF: usize = 512; +const EFF: usize = 256; +const N_EXP: usize = 128; +const N_USED: usize = 6; +const STREAMS: usize = 4; +const VOCAB: usize = 512; +/// A delta layer with the dense block, a delta layer and a latent layer with the routed one. +const LAYERS: usize = 3; +const CTX: usize = 128; + +/// A 64-bit LCG (Knuth's MMIX constants). +struct Lcg(u64); + +impl Lcg { + fn next(&mut self) -> u64 { + self.0 = self + .0 + .wrapping_mul(6_364_136_223_846_793_005) + .wrapping_add(1_442_695_040_888_963_407); + self.0 >> 40 + } + + fn uniform(&mut self, (lo, hi): (f32, f32)) -> f32 { + lo + (hi - lo) * (self.next() as f32 / (1u64 << 24) as f32) + } +} + +/// `n` values of type `ty` as the file holds them: F32 values in `range`, or blocks of random +/// codes under a small scale, so every weight is finite and small. +fn values(r: &mut Lcg, ty: GgmlType, n: usize, range: (f32, f32)) -> Vec { + match ty { + GgmlType::F32 => (0..n) + .flat_map(|_| r.uniform(range).to_le_bytes()) + .collect(), + GgmlType::Q8_0 => { + let mut bytes = Vec::with_capacity(n / 32 * 34); + for _ in 0..n / 32 { + let x: [f32; 32] = std::array::from_fn(|_| r.uniform((-0.05, 0.05))); + let q = Q8Block::quantize_ref(&x).expect("a finite block"); + bytes.extend(q.d.to_le_bytes()); + bytes.extend(q.q.map(|c| c.to_le_bytes()[0])); + } + bytes + } + GgmlType::Q6_K => (0..n / 256 * 210) + .map(|i| match i % 210 { + 192..=207 => ((r.next() % 9) as i8 - 4).to_le_bytes()[0], + 209 => 0x14, + _ => r.next() as u8, + }) + .collect(), + GgmlType::Q4_K => (0..n / 256 * 144) + .map(|i| match i % 144 { + 1 | 3 => 0x14, + _ => r.next() as u8, + }) + .collect(), + GgmlType::Q5_0 => (0..n / 32 * 22) + .map(|i| if i % 22 == 1 { 0x14 } else { r.next() as u8 }) + .collect(), + other => panic!("the synthetic file holds no {other:?} tensor"), + } +} + +/// The file's tensors: name, ggml dims (the contiguous axis first), type and the range of an F32 +/// one. `out` is the type of the latent layer's output projection. +fn tensors(out: GgmlType) -> Vec<(String, Vec, GgmlType, (f32, f32))> { + use GgmlType::{F32, Q4_K, Q5_0, Q6_K, Q8_0}; + let mix = 2 * STREAMS + (STREAMS - 1) * (STREAMS - 1); + let small = (-0.05, 0.05); + let mut t: Vec<(String, Vec, GgmlType, (f32, f32))> = vec![ + ("token_embd.weight".into(), vec![E, VOCAB], Q4_K, small), + ("output.weight".into(), vec![E, VOCAB], Q6_K, small), + ("output_norm.weight".into(), vec![E], F32, (0.5, 1.5)), + ]; + for l in 0..LAYERS { + let mut layer: Vec<(&str, Vec, GgmlType, (f32, f32))> = + vec![("attn_norm.weight", vec![E], F32, (0.5, 1.5))]; + if l == LAYERS - 1 { + layer.extend([ + ("attn_q_a.weight", vec![E, Q_LORA], Q6_K, small), + ("attn_q_a_norm.weight", vec![Q_LORA], F32, (0.5, 1.5)), + ( + "attn_q_b.weight", + vec![Q_LORA, HEADS * (NOPE + PE)], + Q8_0, + small, + ), + ("attn_kv_a_mqa.weight", vec![E, KV_LORA + PE], Q6_K, small), + ("attn_kv_a_norm.weight", vec![KV_LORA], F32, (0.5, 1.5)), + ("attn_k_b.weight", vec![NOPE, KV_LORA, HEADS], Q8_0, small), + ("attn_v_b.weight", vec![KV_LORA, HEAD, HEADS], Q6_K, small), + ("attn_gate.weight", vec![E, HEADS * HEAD], Q8_0, small), + ("attn_output.weight", vec![HEADS * HEAD, E], out, small), + ]); + } else { + layer.extend([ + ("attn_q.weight", vec![E, INNER], Q6_K, small), + ("attn_k.weight", vec![E, INNER], Q6_K, small), + ("attn_v.weight", vec![E, INNER], Q6_K, small), + ( + "ssm_conv1d_q.weight", + vec![CONV, 1, INNER], + F32, + (-0.5, 0.5), + ), + ( + "ssm_conv1d_k.weight", + vec![CONV, 1, INNER], + F32, + (-0.5, 0.5), + ), + ( + "ssm_conv1d_v.weight", + vec![CONV, 1, INNER], + F32, + (-0.5, 0.5), + ), + ("ssm_f_a.weight", vec![E, RANK], F32, small), + ("ssm_f_b.weight", vec![RANK, INNER], F32, (-0.1, 0.1)), + ("ssm_beta.weight", vec![E, HEADS], F32, small), + ("ssm_rda.weight", vec![E, HEADS], F32, small), + ("ssm_a", vec![HEADS], F32, (-2.0, -0.5)), + ("ssm_dt.bias", vec![INNER], F32, (-0.5, 0.5)), + ("ssm_g.weight", vec![E, INNER], Q6_K, small), + ("ssm_norm.weight", vec![HEAD], F32, (0.5, 1.5)), + ("attn_output.weight", vec![INNER, E], Q8_0, small), + ]); + } + layer.push(("ffn_norm.weight", vec![E], F32, (0.5, 1.5))); + if l == 0 { + layer.extend([ + ("ffn_gate.weight", vec![E, FF], Q6_K, small), + ("ffn_up.weight", vec![E, FF], Q6_K, small), + ("ffn_down.weight", vec![FF, E], Q6_K, small), + ]); + } else { + layer.extend([ + ("ffn_gate_inp.weight", vec![E, N_EXP], F32, (-0.1, 0.1)), + ("exp_probs_b.bias", vec![N_EXP], F32, (-0.1, 0.1)), + ("ffn_gate_exps.weight", vec![E, EFF, N_EXP], Q4_K, small), + ("ffn_up_exps.weight", vec![E, EFF, N_EXP], Q4_K, small), + ("ffn_down_exps.weight", vec![EFF, E, N_EXP], Q5_0, small), + ("ffn_gate_shexp.weight", vec![E, EFF], Q6_K, small), + ("ffn_up_shexp.weight", vec![E, EFF], Q6_K, small), + ("ffn_down_shexp.weight", vec![EFF, E], Q8_0, small), + ]); + } + t.extend( + layer + .into_iter() + .map(|(s, d, ty, r)| (format!("blk.{l}.{s}"), d, ty, r)), + ); + for sub in ["attn", "ffn"] { + t.extend([ + ( + format!("blk.{l}.hc_{sub}_fn.weight"), + vec![STREAMS * E, mix], + F32, + (-0.05, 0.05), + ), + ( + format!("blk.{l}.hc_{sub}_base.weight"), + vec![mix], + F32, + (-0.1, 0.1), + ), + ( + format!("blk.{l}.hc_{sub}_scale.weight"), + vec![4], + F32, + (0.3, 0.7), + ), + ]); + } + } + t +} + +/// The file's keys: the widths above, delta layers at 0 and 1 and a latent one at 2, the first +/// layer dense. +fn keys() -> Vec<(String, Value)> { + let u = |v: usize| Value::U32(u32::try_from(v).expect("a width")); + let kv: Vec = (0..LAYERS) + .map(|l| Value::I32(i32::from(l == LAYERS - 1))) + .collect(); + let k = |s: &str, v: Value| (format!("raon1.{s}"), v); + vec![ + ( + gguf::GENERAL_ARCHITECTURE.to_owned(), + Value::String("raon1".to_owned()), + ), + k("block_count", u(LAYERS)), + k("context_length", u(4096)), + k("embedding_length", u(E)), + k("vocab_size", u(VOCAB)), + k("feed_forward_length", u(FF)), + k("leading_dense_block_count", u(1)), + k("attention.head_count", u(HEADS)), + k("attention.head_count_kv", Value::Array(kv)), + k("attention.key_length", u(KV_LORA + PE)), + k("attention.value_length", u(KV_LORA)), + k("attention.key_length_mla", u(NOPE + PE)), + k("attention.value_length_mla", u(HEAD)), + k("attention.q_lora_rank", u(Q_LORA)), + k("attention.kv_lora_rank", u(KV_LORA)), + k("attention.layer_norm_rms_epsilon", Value::F32(1e-5)), + k("attention.latent_norm_rms_epsilon", Value::F32(1e-6)), + k("rope.dimension_count", u(PE)), + k("ssm.conv_kernel", u(CONV)), + k("ssm.time_step_rank", u(RANK)), + k("kda.head_dim", u(HEAD)), + k("kda.gate_lower_bound", Value::F32(-5.0)), + k("expert_count", u(N_EXP)), + k("expert_used_count", u(N_USED)), + k("expert_feed_forward_length", u(EFF)), + k("expert_shared_count", u(1)), + k("expert_gating_func", u(2)), + k("expert_weights_scale", Value::F32(1.0)), + k("expert_weights_norm", Value::Bool(true)), + k("hyper_connection.count", u(STREAMS)), + k("hyper_connection.epsilon", Value::F32(1e-5)), + k("hyper_connection.magnitude", Value::F32(2.0)), + k("activation.situ_beta", Value::F32(4.0)), + k("activation.situ_linear_beta", Value::F32(25.0)), + ] +} + +/// The bytes `n` values of type `ty` take in the file. +fn nbytes(ty: GgmlType, n: usize) -> usize { + match ty { + GgmlType::F32 => 4 * n, + GgmlType::Q8_0 => n / 32 * 34, + GgmlType::Q6_K => n / 256 * 210, + GgmlType::Q4_K => n / 256 * 144, + GgmlType::Q5_0 => n / 32 * 22, + other => panic!("the synthetic file holds no {other:?} tensor"), + } +} + +/// The file of [`tensors`]`(out)` written to a temporary path named by `tag`. +fn write(tag: &str, out: GgmlType) -> PathBuf { + let mut r = Lcg(0x0073_796e); + let table = tensors(out); + let decls: Vec = table + .iter() + .map(|(name, dims, ty, _)| TensorDecl { + name: name.clone(), + dims: dims.iter().map(|&d| d as u64).collect(), + type_id: ty.as_u32(), + nbytes: nbytes(*ty, dims.iter().product()) as u64, + }) + .collect(); + let layout = Layout::new(&keys(), decls).expect("the file's layout"); + let path = std::env::temp_dir().join(format!("raon1-synth-{tag}-{}.gguf", std::process::id())); + let file = std::io::BufWriter::new(std::fs::File::create(&path).expect("a temporary file")); + let mut w = Writer::new(file, layout).expect("the file's writer"); + for (name, dims, ty, range) in &table { + let n: usize = dims.iter().product(); + let bytes = values(&mut r, *ty, n, *range); + assert_eq!(bytes.len(), nbytes(*ty, n), "{name}"); + w.tensor(name, &bytes).expect("a tensor's bytes"); + } + w.finish().expect("the file's end"); + path +} + +/// The model of the file `write(tag, out)` made, loaded whole on a card; the file is removed once +/// it is open. +fn open(tag: &str, out: GgmlType) -> Result, GpuError> { + let path = write(tag, out); + let split = Split::open(&path).unwrap_or_else(|e| panic!("open {}: {e}", path.display())); + let gpu = Gpu::with_device(0).expect("a Gpu on device 0"); + let m = GpuModel::::open(gpu, split, CTX); + let _ = std::fs::remove_file(&path); + m +} + +/// A prompt of `n` ids of the file's vocabulary. +fn prompt(n: usize) -> Vec { + (0..n).map(|i| ((7 * i + 3) % VOCAB) as u32).collect() +} + +/// A file whose latent layer's output projection is Q6_K, a type the loader admits and the wide +/// arm refuses, loads with the wide arm off, names the layer, and runs a prompt of nine ids and +/// one past a staging group as the rows pass: the argmax and the last logits row of the steps, by +/// bits. The same file with Q8_0 there (Raon's) loads with the wide arm on and runs the nine ids +/// as a wide unit, to finite logits. +#[test] +#[ignore = "needs a CUDA device; `just gate-gpu-lib` runs it on the box"] +fn hw_a_file_the_wide_arm_refuses_runs_every_prompt_as_passes() { + let mut m = open("q6k", GgmlType::Q6_K).expect("a file the loader admits"); + let body = m.body("synth").expect("the body"); + assert!( + !body.wide_ok(), + "a Q6_K output projection keeps the file off the wide arm" + ); + let why = body.wide_off().expect("the refusal").to_owned(); + assert!( + why.contains("layer 2's mixer") && why.contains("an output projection of Q6_K"), + "{why}" + ); + for n in [9, 70] { + let ids = prompt(n); + m.reset().expect("a reset"); + let steps = m.step(&ids).expect("the steps"); + let steps_row = m.logits().expect("the steps' row"); + m.reset().expect("a reset"); + let rows = m.prompt_rows(&ids).expect("a prompt of the rows pass"); + let rows_row = m.logits().expect("the rows' row"); + assert_eq!(steps, rows, "{n} ids: the argmax"); + assert!( + steps_row.iter().all(|v| v.is_finite()) && same(&steps_row, &rows_row, f32::to_bits), + "{n} ids: the last logits row by bits" + ); + } + drop(m); + let mut m = open("q8", GgmlType::Q8_0).expect("a file of Raon's types"); + let body = m.body("synth").expect("the body"); + assert!(body.wide_ok(), "{:?}", body.wide_off()); + m.reset().expect("a reset"); + m.prompt_rows(&prompt(9)).expect("nine ids as a wide unit"); + let row = m.logits().expect("the wide row"); + assert!( + row.iter().all(|v| v.is_finite()), + "the wide unit's logits are finite" + ); +} diff --git a/crates/gpu/src/arch/raon1/wide.rs b/crates/gpu/src/arch/raon1/wide.rs index cf52cf187..1cf0c27ae 100644 --- a/crates/gpu/src/arch/raon1/wide.rs +++ b/crates/gpu/src/arch/raon1/wide.rs @@ -34,14 +34,15 @@ use crate::card_leg::{RouteW, StackW}; use crate::card_slots::{CardAcc, CardSlotKernels}; use crate::fault::FaultSink; use crate::gemm::{ - Gemm32Args, Gemm32Kernels, GemmAct, GemmAct32, GemmArgs, GemmInput, GemmKernels, GemmRoute, + Gemm32Args, Gemm32Kernels, Gemm32Weight, GemmAct, GemmAct32, GemmArgs, GemmInput, GemmKernels, + GemmRoute, GemmWeight, }; use crate::kquant::Act; use crate::latent::LATENT; use crate::linear::mixer_kda::{self, KdaBufs, KdaDims, KdaIn, KdaProj, KdaState}; use crate::linear::{HEAD, LinearKernels, LinearShape}; use crate::model::fit::{act_col_bytes, act32_col_bytes, route_bytes}; -use crate::site::{self, Gated, GatedIo, SiteTy, SiteWeight, SwigluRows, WideActs, WideIn}; +use crate::site::{self, Form, Gated, GatedIo, SiteTy, SiteWeight, SwigluRows, WideActs, WideIn}; use crate::tensor::DeviceTensor; use crate::weights::{DevWeight, Weights}; use crate::{Gpu, GpuError, launch_u32}; @@ -303,7 +304,7 @@ pub(super) struct WideDims { pub shared_ff: usize, pub expert_ff: usize, /// The widest feed-forward block's gate rows a unit holds: the dense - /// block's, or the routed slots'. + /// block's, the shared expert's, or the routed slots'. pub ff_rows: usize, } @@ -796,7 +797,7 @@ pub(super) fn quantize_out( (x, m): (&DeviceBuffer, usize), (sink, ty): (FaultSink, SiteTy), ) -> Result<(), GpuError> { - only_blocks32(ty, "an output projection")?; + out_runs(ty)?; g32.enqueue_quantize_gemm32(gpu.stream(), x, m, &mut forms.out32, sink) } @@ -809,16 +810,22 @@ pub(super) fn quantize_qb( (x, m): (&DeviceBuffer, usize), (sink, ty): (FaultSink, SiteTy), ) -> Result<(), GpuError> { - only_blocks32(ty, "a query up projection")?; + qb_runs(ty)?; g32.enqueue_quantize_gemm32(gpu.stream(), x, m, &mut forms.qb32, sink) } +/// Whether a site of type `ty` reads q8 blocks of 32 alone, as the unit's gated +/// rows and query latent are quantized for an output or query up projection. +fn blocks32(ty: SiteTy) -> bool { + let (q128, q32) = WideActs::wants(ty); + !q128 && q32 +} + /// A named refusal of a projection `what` of type `ty` that reads another /// form than q8 blocks of 32: the wide arm's forms for those inputs are /// Q8_0's. fn only_blocks32(ty: SiteTy, what: &str) -> Result<(), GpuError> { - let (q128, q32) = WideActs::wants(ty); - if q128 || !q32 { + if !blocks32(ty) { return Err(GpuError::shape( WHAT, format!("{what} of {ty}: the wide arm's form for its input is q8 blocks of 32 (Q8_0)"), @@ -827,6 +834,34 @@ fn only_blocks32(ty: SiteTy, what: &str) -> Result<(), GpuError> { Ok(()) } +/// `Ok` when the wide arm runs a mixer's output projection of type `ty`. +pub(super) fn out_runs(ty: SiteTy) -> Result<(), GpuError> { + only_blocks32(ty, "an output projection") +} + +/// `Ok` when the wide arm runs the latent layers' `attn_q_b` of type `ty`. +pub(super) fn qb_runs(ty: SiteTy) -> Result<(), GpuError> { + only_blocks32(ty, "a query up projection") +} + +/// `Ok` when the wide arm runs the RDA layer's projections `s`, else the refusal +/// it raises: the output projection reads q8 blocks of 32, and the forget +/// projection's second half the f32 rows of its first half's output. +pub(super) fn rda_runs(s: &RdaSites) -> Result<(), GpuError> { + out_runs(s.out.ty())?; + let ty = s.f_b.ty(); + if WideActs::wants(ty) != (false, false) { + return Err(GpuError::shape( + WHAT, + format!( + "the forget projection's second half of {ty}: the wide arm's form for its input \ + is the f32 rows (F32)" + ), + )); + } + Ok(()) +} + /// The RDA layer's wide projections ([`KdaProj`]): q, k, v and the gate over /// the normed rows' q8_1 form, β's, the write gate's and the forget /// projection's halves on the F32 tile, the output projection over the gated @@ -853,6 +888,7 @@ impl KdaProj for RdaWide<'_, '_> { let (cx, s, shape) = (self.cx, self.sites, self.shape); let sink = gpu.layer_sink(l)?; let (qk, v) = (shape.n_k * HEAD, shape.n_v * HEAD); + rda_runs(s)?; let all = s.sites(); let tys: Vec = all[..7].iter().map(|t| t.ty()).collect(); let KdaIn { @@ -956,6 +992,111 @@ fn stack<'w>(w: &'w Weights, name: &str) -> Result<(GgmlType, StackW<'w>), GpuEr } } +/// The form a dense block's or a shared expert's SwiGLU rows take into its +/// down. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(super) enum DownForm { + /// q8_1 blocks of 128 over the dense width: a K-quant down. + Q128, + /// q8 blocks of 32 over the shared width: a Q8_0 down. + Q32, +} + +/// The form of the SwiGLU rows of a block of width `ff` into its down of type +/// `down`, with the unit's forms cut for `(dense_ff, shared_ff)` values; `None` +/// where the wide arm holds none: a down whose form no wide SwiGLU writes (a +/// Q5_0 down reads such blocks), or a width other than its form's. +pub(super) fn down_form( + down: SiteTy, + ff: usize, + (dense_ff, shared_ff): (usize, usize), +) -> Option { + match down.reads() { + Form::Q8x128 if ff == dense_ff => Some(DownForm::Q128), + Form::Q8x32 if ff == shared_ff => Some(DownForm::Q32), + _ => None, + } +} + +/// A layer's routed stacks as the wide arm's route GEMMs read them. +struct RoutedStacks<'w> { + /// The grouped GEMM's type of the gate and the up stacks, one for both. + ty: GemmWeight, + gate: &'w DeviceTensor, + up: &'w DeviceTensor, + down: Gemm32Weight<'w>, +} + +/// Layer `l`'s stacks `names` (gate, up, down) as the wide arm reads them: the +/// gate and up of one type the grouped K-quant GEMM decodes, the down of a type +/// the 32-value route GEMM reads; any other is refused by name. +fn routed_stacks<'w>( + w: &'w Weights, + l: usize, + names: &[String; 3], +) -> Result, GpuError> { + let (gty, gate) = stack(w, &names[0])?; + let (uty, up) = stack(w, &names[1])?; + let (dty, down) = stack(w, &names[2])?; + if gty != uty { + return Err(GpuError::shape( + WHAT, + format!("layer {l}: gate and up stacks of {gty} and {uty}: one GEMM type reads both"), + )); + } + let (Some(RouteW::Grouped(ty, gate)), Some(RouteW::Grouped(_, up))) = + (gate.route(gty), up.route(uty)) + else { + return Err(GpuError::shape( + WHAT, + format!( + "layer {l}: gate and up stacks of {gty} and {uty}: the wide arm runs the grouped K-quant GEMM over them" + ), + )); + }; + let Some(RouteW::Blocks32(down)) = down.route(dty) else { + return Err(GpuError::shape( + WHAT, + format!("layer {l}: a down stack of {dty}, which no 32-value route GEMM reads"), + )); + }; + Ok(RoutedStacks { ty, gate, up, down }) +} + +/// The named refusal of block `g` of layer `l`: its down has no form at the +/// block's width ([`down_form`]). +fn no_down_form(l: usize, g: &Glu) -> GpuError { + GpuError::shape( + WHAT, + format!( + "layer {l}: a block of width {} with a down of {}: the wide arm holds q8_1 rows of \ + the dense width and q8 blocks of 32 of the shared width", + g.ff, + g.down.ty() + ), + ) +} + +/// `Ok` when the wide arm runs layer `l`'s feed-forward block `f` of resident +/// weights `w` at the widths `d`, else the refusal [`ffn`] raises for it: a +/// dense block or shared expert needs a form for its down ([`down_form`]), a +/// routed layer's stacks ones [`routed_stacks`] reads. +pub(super) fn ffn_runs(l: usize, f: &Ffn, w: &Weights, d: &WideDims) -> Result<(), GpuError> { + let widths = (d.dense_ff, d.shared_ff); + let glu = |g: &Glu| { + down_form(g.down.ty(), g.ff, widths) + .map(drop) + .ok_or_else(|| no_down_form(l, g)) + }; + match f { + Ffn::Dense(g) => glu(g), + Ffn::Routed(r) => { + glu(&r.shared)?; + routed_stacks(w, l, &r.stacks).map(drop) + } + } +} + /// A dense block (the leading layer's, or a shared expert) over the unit's /// `m` rows of the normed rows whose forms the caller made: the gate and up /// GEMMs, the rule's SiTU quantized into the form the down reads, the down @@ -989,28 +1130,19 @@ fn glu( q32: Some(&forms.hid32), f32: None, }; - let rows = match WideActs::wants(g.down.ty()) { - (true, _) if forms.dense128.k() == g.ff => SwigluRows { + let widths = (forms.dense128.k(), forms.shared32.k()); + let rows = match down_form(g.down.ty(), g.ff, widths) { + Some(DownForm::Q128) => SwigluRows { q128: Some(&mut forms.dense128), q32: None, f32: None, }, - (false, true) if forms.shared32.k() == g.ff => SwigluRows { + Some(DownForm::Q32) => SwigluRows { q128: None, q32: Some(&mut forms.shared32), f32: None, }, - _ => { - return Err(GpuError::shape( - WHAT, - format!( - "layer {l}: a block of width {} with a down of {}: the wide arm holds q8_1 \ - rows of the dense width and q8 blocks of 32 of the shared width", - g.ff, - g.down.ty() - ), - )); - } + None => return Err(no_down_form(l, g)), }; let io = GatedIo { x: &x, @@ -1098,25 +1230,7 @@ fn routed( let (gpu, w) = (cx.gpu, cx.w); let (stream, sink) = (gpu.stream(), gpu.layer_sink(l)?); let used = sigmoid_128_6::N_USED; - let (gty, gate) = stack(w, &r.stacks[0])?; - let (uty, up) = stack(w, &r.stacks[1])?; - let (dty, down) = stack(w, &r.stacks[2])?; - let (Some(RouteW::Grouped(ty, gate)), Some(RouteW::Grouped(_, up))) = - (gate.route(gty), up.route(uty)) - else { - return Err(GpuError::shape( - WHAT, - format!( - "layer {l}: gate and up stacks of {gty} and {uty}: the wide arm runs the grouped K-quant GEMM over them" - ), - )); - }; - let Some(RouteW::Blocks32(down)) = down.route(dty) else { - return Err(GpuError::shape( - WHAT, - format!("layer {l}: a down stack of {dty}, which no 32-value route GEMM reads"), - )); - }; + let RoutedStacks { ty, gate, up, down } = routed_stacks(w, l, &r.stacks)?; k.router.enqueue_router_rows( stream, gpu, @@ -1184,19 +1298,220 @@ fn routed( #[cfg(test)] mod tests { - use super::{Gemm32Args, Gemm32Kernels, GemmAct32, GemmInput, GemmKernels, GemmRoute, LATENT}; + use super::{ + DownForm, Ffn, Gemm32Args, Gemm32Kernels, GemmAct32, GemmInput, GemmKernels, GemmRoute, + Glu, LATENT, RdaSites, Routed, WideDims, down_form, ffn_runs, out_runs, qb_runs, rda_runs, + routed_stacks, + }; use crate::Gpu; - use crate::arch::raon1::rda::tests::{Lcg, q8_0, resident}; + use crate::arch::hybrid::Site; + use crate::arch::raon1::rda::tests::{Lcg, q8_0, raw, resident}; use crate::gemm::Gemm32Weight; use crate::q8f32::Q8_0GemvHeadsMcolArgs; + use crate::site::SiteTy; use crate::tensor::{Window, WindowMut}; use cuda_core::DeviceBuffer; + use gguf::GgmlType; use model::arch::raon1::names; /// The query's nope values the absorption reads a head: the file's `attn_k_b` is /// 128 values (4 blocks of 32) a row. const NOPE: usize = 128; + /// A dense block or a shared expert has a form for its down at the widths the + /// unit's forms are cut for: a K-quant down reads q8_1 blocks over the dense + /// width, a Q8_0 down q8 blocks of 32 over the shared width; a Q5_0 down reads + /// blocks no wide SwiGLU writes, and any down at the other form's width has none. + #[test] + fn a_down_has_a_form_only_at_the_width_its_form_is_cut_for() { + let (dense, shared) = (1024, 1792); + for ty in [SiteTy::Q4K, SiteTy::Q5K, SiteTy::Q6K] { + assert_eq!( + down_form(ty, dense, (dense, shared)), + Some(DownForm::Q128), + "{ty} at the dense width" + ); + assert_eq!( + down_form(ty, shared, (dense, shared)), + None, + "{ty} at the shared width" + ); + } + assert_eq!( + down_form(SiteTy::Q8_0, shared, (dense, shared)), + Some(DownForm::Q32) + ); + assert_eq!(down_form(SiteTy::Q8_0, dense, (dense, shared)), None); + for ff in [dense, shared] { + assert_eq!(down_form(SiteTy::Q5_0, ff, (dense, shared)), None, "Q5_0"); + } + let one = (1088, 1088); + assert_eq!(down_form(SiteTy::Q6K, 1088, one), Some(DownForm::Q128)); + assert_eq!(down_form(SiteTy::Q8_0, 1088, one), Some(DownForm::Q32)); + } + + /// An output projection and `attn_q_b` run wide over q8 blocks of 32 alone: + /// Q8_0 and Q5_0 do; every K-quant and every f32-row type is refused by name. + #[test] + fn an_output_projection_runs_wide_only_over_blocks_of_32() { + for &ty in SiteTy::ALL { + let runs = matches!(ty, SiteTy::Q8_0 | SiteTy::Q5_0); + for (what, r) in [ + ("an output projection", out_runs(ty)), + ("a query up projection", qb_runs(ty)), + ] { + assert_eq!(r.is_ok(), runs, "{what} of {ty}"); + if let Err(e) = r { + assert!(e.to_string().contains(what), "{e}"); + } + } + } + } + + /// An RDA layer's sites with `out` and `f_b` of the given types, the others + /// Raon's. + fn rda_sites(out: SiteTy, f_b: SiteTy) -> RdaSites { + let site = |name: &str, ty: SiteTy, k: usize| Site::new(name.to_owned(), ty, k); + RdaSites { + qkv: [ + site("q", SiteTy::Q6K, 512), + site("k", SiteTy::Q6K, 512), + site("v", SiteTy::Q6K, 512), + ], + g: site("g", SiteTy::Q6K, 512), + beta: site("beta", SiteTy::F32, 512), + write: site("write", SiteTy::F32, 512), + f_a: site("f_a", SiteTy::F32, 512), + f_b: site("f_b", f_b, 128), + out: site("out", out, 256), + vecs: std::array::from_fn(|_| String::new()), + } + } + + /// The wide arm runs an RDA layer whose output projection is Q8_0 and whose + /// forget projection's second half is F32 (Raon's), and refuses by name one + /// whose output projection is another type or whose second half reads a + /// quantized form. + #[test] + fn an_rda_layer_runs_wide_with_a_q8_0_output_and_an_f32_forget_projection() { + assert!(rda_runs(&rda_sites(SiteTy::Q8_0, SiteTy::F32)).is_ok()); + for ty in [SiteTy::Q6K, SiteTy::Q4K, SiteTy::F32] { + let e = rda_runs(&rda_sites(ty, SiteTy::F32)).unwrap_err(); + assert!(e.to_string().contains("an output projection"), "{e}"); + } + for ty in [SiteTy::Q8_0, SiteTy::Q6K] { + let e = rda_runs(&rda_sites(SiteTy::Q8_0, ty)).unwrap_err(); + assert!(e.to_string().contains("forget projection"), "{e}"); + } + } + + /// The wide arm runs a dense block by its down (a K-quant down over the dense + /// width, a Q8_0 down over the shared width), and a routed layer by its shared + /// expert's down and by its stacks: a shared expert with a down of no form, or + /// stacks the weights do not hold, are refused by name. + #[test] + fn a_feed_forward_block_runs_wide_by_its_down_and_its_stacks() { + let site = |name: &str, ty: SiteTy, k: usize| Site::new(name.to_owned(), ty, k); + let glu = |down: SiteTy, ff: usize| Glu { + gate: site("gate", SiteTy::Q4K, 512), + up: site("up", SiteTy::Q4K, 512), + down: site("down", down, ff), + ff, + }; + let (dense_ff, shared_ff) = (1024, 1792); + let d = WideDims { + embd: 512, + attn: 512, + q_lora: 256, + heads: 4, + nope: 128, + dense_ff, + shared_ff, + expert_ff: 256, + ff_rows: shared_ff, + }; + let w = crate::weights::Weights::default(); + let routed = |shared: Glu| { + Ffn::Routed(Routed { + router: "router".to_owned(), + bias: "bias".to_owned(), + stacks: ["gate", "up", "down"].map(str::to_owned), + shared, + }) + }; + assert!(ffn_runs(0, &Ffn::Dense(glu(SiteTy::Q6K, dense_ff)), &w, &d).is_ok()); + let e = ffn_runs(0, &Ffn::Dense(glu(SiteTy::Q8_0, dense_ff)), &w, &d).unwrap_err(); + assert!(e.to_string().contains("a block of width 1024"), "{e}"); + let e = ffn_runs(1, &routed(glu(SiteTy::Q5_0, shared_ff)), &w, &d).unwrap_err(); + assert!(e.to_string().contains("a block of width 1792"), "{e}"); + let e = ffn_runs(1, &routed(glu(SiteTy::Q8_0, shared_ff)), &w, &d).unwrap_err(); + assert!(e.to_string().contains("a resident routed stack"), "{e}"); + } + + /// A layer's routed stacks run wide as the gate and up of one grouped + /// K-quant type over a down the 32-value GEMM reads (a packed Q5_0 here); a + /// mixed gate and up, a K-quant down and a missing stack are refused by name. + #[test] + #[ignore = "needs a CUDA device; `just gate-gpu-lib` runs it on the box"] + fn hw_routed_stacks_run_wide_as_one_grouped_type_over_a_32_value_down() { + const ROWS: usize = 256; + const K: usize = 256; + const EXPERTS: usize = 2; + let gpu = Gpu::with_device(0).expect("a Gpu on device 0"); + let names = ["gate", "up", "down"].map(|n| format!("blk.1.ffn_{n}_exps.weight")); + // Every block's scale bytes small (the high byte of an f16 at 0x14), the codes zero. + let stack = |name: &String, ty: GgmlType, block: (usize, usize)| { + let blocks = ROWS * K * EXPERTS / block.0; + let bytes = (0..blocks * block.1) + .map(|i| { + if i % block.1 % 2 == 1 && i % block.1 < 4 { + 0x14 + } else { + 0 + } + }) + .collect(); + raw(name.clone(), ty, vec![K, ROWS, EXPERTS], bytes) + }; + let q4k = (256, 144); + let q5k = (256, 176); + let q5_0 = (32, 22); + let resident_of = |gate: (GgmlType, (usize, usize)), + up: (GgmlType, (usize, usize)), + down: Option<(GgmlType, (usize, usize))>| { + let mut parts = vec![ + stack(&names[0], gate.0, gate.1), + stack(&names[1], up.0, up.1), + ]; + parts.extend(down.map(|d| stack(&names[2], d.0, d.1))); + resident(&gpu, &parts) + }; + let runs = |w: &crate::weights::Weights| routed_stacks(w, 1, &names).map(drop); + let ok = resident_of( + (GgmlType::Q4_K, q4k), + (GgmlType::Q4_K, q4k), + Some((GgmlType::Q5_0, q5_0)), + ); + runs(&ok).expect("Q4_K gate and up over a Q5_0 down"); + let mixed = resident_of( + (GgmlType::Q4_K, q4k), + (GgmlType::Q5_K, q5k), + Some((GgmlType::Q5_0, q5_0)), + ); + let e = runs(&mixed).unwrap_err(); + assert!(e.to_string().contains("one GEMM type reads both"), "{e}"); + let kquant_down = resident_of( + (GgmlType::Q4_K, q4k), + (GgmlType::Q4_K, q4k), + Some((GgmlType::Q4_K, q4k)), + ); + let e = runs(&kquant_down).unwrap_err(); + assert!(e.to_string().contains("no 32-value route GEMM"), "{e}"); + let no_down = resident_of((GgmlType::Q4_K, q4k), (GgmlType::Q4_K, q4k), None); + let e = runs(&no_down).unwrap_err(); + assert!(e.to_string().contains(&names[2]), "{e}"); + } + /// The key absorption's 32-value GEMM ([`super::absorb_heads`]'s core) over the file's /// geometry is the f32 gemv's rows ([`crate::q8f32`]'s `q8_0_gemv_heads_mcol`, the pass's /// absorb) to the activation's quantization: the rows' relative rms distance stays under From 7d4a697824931e100ed7ebaf591701aec0308b9f Mon Sep 17 00:00:00 2001 From: midagedev Date: Thu, 22 Oct 2026 06:48:35 +0900 Subject: [PATCH 10/10] raon1 synth: reuse rda's test generators and gguf's type table The synthetic file's bytes are unchanged: F32, Q8_0 and Q6_K come from rda::tests' generators, which hold the same expressions and call order. Co-Authored-By: Claude Opus 5.5 --- crates/gpu/src/arch/raon1/rda.rs | 4 +-- crates/gpu/src/arch/raon1/synth.rs | 57 ++++++------------------------ 2 files changed, 12 insertions(+), 49 deletions(-) diff --git a/crates/gpu/src/arch/raon1/rda.rs b/crates/gpu/src/arch/raon1/rda.rs index a7ae233da..185728c89 100644 --- a/crates/gpu/src/arch/raon1/rda.rs +++ b/crates/gpu/src/arch/raon1/rda.rs @@ -596,7 +596,7 @@ pub(super) mod tests { pub(crate) struct Lcg(pub(crate) u64); impl Lcg { - fn next(&mut self) -> u64 { + pub(crate) fn next(&mut self) -> u64 { self.0 = self .0 .wrapping_mul(6_364_136_223_846_793_005) @@ -615,7 +615,7 @@ pub(super) mod tests { /// A file tensor of the test's layers: its name, type, ggml dims (the /// contiguous axis first) and bytes. - pub(crate) struct Part(String, GgmlType, Vec, Vec); + pub(crate) struct Part(String, GgmlType, Vec, pub(crate) Vec); /// A file tensor of `ty` over `dims` whose bytes are `bytes` as they stand. pub(crate) fn raw(name: String, ty: GgmlType, dims: Vec, bytes: Vec) -> Part { diff --git a/crates/gpu/src/arch/raon1/synth.rs b/crates/gpu/src/arch/raon1/synth.rs index cbc53082c..1a434960c 100644 --- a/crates/gpu/src/arch/raon1/synth.rs +++ b/crates/gpu/src/arch/raon1/synth.rs @@ -5,10 +5,9 @@ //! small where none does; the weights are random with small scales, so every value is finite. use super::body::Body; -use super::rda::tests::same; +use super::rda::tests::{Lcg, f32s, q6k, q8_0, same}; use crate::model::GpuModel; use crate::{Gpu, GpuError}; -use gguf::quant::Q8Block; use gguf::write::{Layout, TensorDecl, Writer}; use gguf::{GgmlType, Split, Value}; use std::path::PathBuf; @@ -34,47 +33,15 @@ const VOCAB: usize = 512; const LAYERS: usize = 3; const CTX: usize = 128; -/// A 64-bit LCG (Knuth's MMIX constants). -struct Lcg(u64); - -impl Lcg { - fn next(&mut self) -> u64 { - self.0 = self - .0 - .wrapping_mul(6_364_136_223_846_793_005) - .wrapping_add(1_442_695_040_888_963_407); - self.0 >> 40 - } - - fn uniform(&mut self, (lo, hi): (f32, f32)) -> f32 { - lo + (hi - lo) * (self.next() as f32 / (1u64 << 24) as f32) - } -} - /// `n` values of type `ty` as the file holds them: F32 values in `range`, or blocks of random /// codes under a small scale, so every weight is finite and small. fn values(r: &mut Lcg, ty: GgmlType, n: usize, range: (f32, f32)) -> Vec { + let name = String::new(); + let dims = vec![n]; match ty { - GgmlType::F32 => (0..n) - .flat_map(|_| r.uniform(range).to_le_bytes()) - .collect(), - GgmlType::Q8_0 => { - let mut bytes = Vec::with_capacity(n / 32 * 34); - for _ in 0..n / 32 { - let x: [f32; 32] = std::array::from_fn(|_| r.uniform((-0.05, 0.05))); - let q = Q8Block::quantize_ref(&x).expect("a finite block"); - bytes.extend(q.d.to_le_bytes()); - bytes.extend(q.q.map(|c| c.to_le_bytes()[0])); - } - bytes - } - GgmlType::Q6_K => (0..n / 256 * 210) - .map(|i| match i % 210 { - 192..=207 => ((r.next() % 9) as i8 - 4).to_le_bytes()[0], - 209 => 0x14, - _ => r.next() as u8, - }) - .collect(), + GgmlType::F32 => f32s(r, name, dims, range).3, + GgmlType::Q8_0 => q8_0(r, name, dims).3, + GgmlType::Q6_K => q6k(r, name, dims).3, GgmlType::Q4_K => (0..n / 256 * 144) .map(|i| match i % 144 { 1 | 3 => 0x14, @@ -254,14 +221,10 @@ fn keys() -> Vec<(String, Value)> { /// The bytes `n` values of type `ty` take in the file. fn nbytes(ty: GgmlType, n: usize) -> usize { - match ty { - GgmlType::F32 => 4 * n, - GgmlType::Q8_0 => n / 32 * 34, - GgmlType::Q6_K => n / 256 * 210, - GgmlType::Q4_K => n / 256 * 144, - GgmlType::Q5_0 => n / 32 * 22, - other => panic!("the synthetic file holds no {other:?} tensor"), - } + let row = ty.row().expect("a known fixture tensor type"); + let n = u64::try_from(n).expect("fixture element count fits u64"); + let bytes = gguf::tensor_bytes(row.blck, row.bytes, &[n]).expect("fixture byte count fits u64"); + usize::try_from(bytes).expect("fixture byte count fits usize") } /// The file of [`tensors`]`(out)` written to a temporary path named by `tag`.