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-gates/src/bin/gate_raon1_e2e.rs b/crates/gpu-gates/src/bin/gate_raon1_e2e.rs
index ec8c7e215..3e4ccf7ec 100644
--- a/crates/gpu-gates/src/bin/gate_raon1_e2e.rs
+++ b/crates/gpu-gates/src/bin/gate_raon1_e2e.rs
@@ -25,21 +25,39 @@
//! 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.
+//! - `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.
//! - `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
//! 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.
@@ -93,15 +111,71 @@ 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, 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));
+
+ /// 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 +401,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 +433,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 +623,92 @@ 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()
+ }
+
+ /// 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) 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 +718,74 @@ 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 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());
+ 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_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!(
- "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 +810,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)
@@ -520,6 +849,7 @@ mod gate {
}
pub fn run() -> Result<(), GateError> {
+ within_bands_separate()?;
let sets: Vec = SETS
.iter()
.map(|&n| open_set(n))
@@ -559,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)?;
@@ -623,6 +963,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)?;
diff --git a/crates/gpu/src/arch/raon1/body.rs b/crates/gpu/src/arch/raon1/body.rs
index a6ae3d1b8..9a33134b6 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;
@@ -43,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;
@@ -108,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.
@@ -175,9 +199,50 @@ 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,
+ nope: mla.nope,
+ dense_ff: hp.dense_ff,
+ shared_ff: e.shared_ff(),
+ expert_ff: e.ff,
+ ff_rows: wide_ff_rows(hp.dense_ff, &e),
+ })
+}
+
+/// The width of the widest block the dense forms' buffers hold rows of: the dense
+/// layers' and the shared experts' (`n_shared` experts of `ff` run as one block
+/// of `ff · n_shared` values).
+fn dense_block_ff(dense_ff: usize, e: &Experts) -> 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`].
-struct Scratch {
+pub(super) struct Scratch {
/// The streams the layer reads and the ones it writes, `STREAMS · n` a
/// row.
x: DeviceBuffer,
@@ -228,7 +293,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 [`mla::ATTN_CHUNK`] 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(mla::ATTN_CHUNK));
let (heads, row) = (m.heads, m.row());
Ok((
[
@@ -311,7 +378,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 +464,11 @@ pub struct Body {
pass: Scratch,
ffn_rows: FfnRows,
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,
@@ -491,8 +563,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 +573,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
@@ -510,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,
@@ -620,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)));
@@ -632,6 +733,12 @@ 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,
+ &wd,
+ Scratch::new(stream, &hp, (&kda, &mla), ctx, WIDE_ROWS)?,
+ )?,
+ wide_off,
rows: RowState::new(stream, lanes)?,
kda,
mla,
@@ -652,6 +759,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 +803,8 @@ impl Body {
ffn_rows,
stage,
rows,
+ wide,
+ wide_off,
ctx,
..
} = self;
@@ -722,6 +832,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 +851,52 @@ 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"),
+ ));
+ }
+ 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 {
+ 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 +1005,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 +1039,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 +1066,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)?;
@@ -1009,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 20745b373..ad5bf1ace 100644
--- a/crates/gpu/src/arch/raon1/mla.rs
+++ b/crates/gpu/src/arch/raon1/mla.rs
@@ -28,8 +28,12 @@
//!
//! and last, 11. `attn_output` into `out`.
//!
-//! [`step`] runs one row. [`rows`] runs a pass of 1 to [`PASS_ROWS`] rows in
-//! position order — a verify's rows or a prompt chunk — through the same
+//! [`step`] runs one row. [`rows_wide`] runs a prompt unit of up to
+//! [`WIDE_ROWS`] rows: the projections as GEMMs, the cache append, the gather,
+//! the absorb ([`absorb_heads`], one 32-value GEMM over every (row, head)),
+//! the value projection and the output gate once over the unit, and the
+//! latent attention in chunks of [`ATTN_CHUNK`] query tokens. [`rows`] runs a
+//! pass of 1 to [`PASS_ROWS`] rows in position order — a verify's rows or a prompt chunk — through the same
//! launches, each once over the pass's rows: the projections' m-column
//! launches ([`super::rda::project_rows`]), the cache rows of every row
//! appended before any attends, the absorb and a Q8_0 value projection over
@@ -48,6 +52,10 @@
use super::rda::{RowsKit, kgemv_rows, project_rows};
use super::state::PASS_ROWS;
+use super::wide::{
+ 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;
use crate::flash::{FlashGeom, FlashInputs, FlashSplitArgs, partials_ms_len, partials_v_len};
@@ -56,7 +64,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};
@@ -68,6 +76,15 @@ use model::arch::raon1::names;
/// What the refusals of this file name.
const WHAT: &str = "raon1 mla";
+/// The query tokens a wide unit's latent flash takes a launch: the partials
+/// ([`partials_lens_rows`]) hold this many tokens' rows a segment, and an
+/// even count keeps each launch's groups of sixteen rows in the same token
+/// pairs whatever the chunk (`heads` is 40 rows a token).
+pub(super) const ATTN_CHUNK: usize = 32;
+
+const _: () = assert!(ATTN_CHUNK >= 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,
@@ -217,11 +234,19 @@ 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
-/// 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 +264,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,20 +849,208 @@ 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, 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,
+ 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 (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(
+ 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,
+ )?;
+ 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(ATTN_CHUNK) {
+ let mc = ATTN_CHUNK.min(m - c0);
+ let n_keys = Window::::of(b.n_keys, (c0 + mc - 1) * u, 1)?;
+ 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)?;
+ 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,
+ },
+ )?;
+ drop(att);
+ if let ValueW::Q8_0(vqs, vd) = value {
+ 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 {
+ 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 {
+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};
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;
@@ -900,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
@@ -1106,4 +1354,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/mod.rs b/crates/gpu/src/arch/raon1/mod.rs
index 66f2ed968..528b44dcd 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,9 @@ pub mod mla;
pub mod rda;
mod rows;
mod state;
+#[cfg(test)]
+mod synth;
+mod wide;
pub use body::Body;
pub use ffn::sigmoid_128_6;
diff --git a/crates/gpu/src/arch/raon1/rda.rs b/crates/gpu/src/arch/raon1/rda.rs
index 3ae6c03cc..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,12 @@ 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 {
+ 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.
diff --git a/crates/gpu/src/arch/raon1/rows.rs b/crates/gpu/src/arch/raon1/rows.rs
index 36f061878..080683cbc 100644
--- a/crates/gpu/src/arch/raon1/rows.rs
+++ b/crates/gpu/src/arch/raon1/rows.rs
@@ -4,9 +4,12 @@
//! 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 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
//! layers a draft names, laid in by every walk, one row a position.
//!
@@ -17,6 +20,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 +61,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 +73,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 +83,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 +307,16 @@ 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 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
+ /// 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 +335,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 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)?;
+ 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 +380,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 +429,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 +459,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 +490,7 @@ impl GpuModel {
gpu.stream(),
layers,
n_layer,
- (hidden, GROUP),
+ (hidden, WIDE_ROWS),
free,
)?);
Ok(())
diff --git a/crates/gpu/src/arch/raon1/synth.rs b/crates/gpu/src/arch/raon1/synth.rs
new file mode 100644
index 000000000..1a434960c
--- /dev/null
+++ b/crates/gpu/src/arch/raon1/synth.rs
@@ -0,0 +1,317 @@
+//! 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::{Lcg, f32s, q6k, q8_0, same};
+use crate::model::GpuModel;
+use crate::{Gpu, GpuError};
+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;
+
+/// `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 => 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,
+ _ => 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 {
+ 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`.
+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
new file mode 100644
index 000000000..1cf0c27ae
--- /dev/null
+++ b/crates/gpu/src/arch/raon1/wide.rs
@@ -0,0 +1,1616 @@
+//! 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 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;
+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, 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, Form, 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 };
+ }
+
+ /// `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 and of the window copy.
+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(())
+ }
+
+ /// 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
+/// 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 queries' nope parts, a column a (row, head), into the key
+ /// absorption.
+ nope32: GemmAct32,
+}
+
+/// 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, 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,
+ /// The widest feed-forward block's gate rows a unit holds: the dense
+ /// block's, the shared expert'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)?,
+ nope32: GemmAct32::new(stream, rows * d.heads, d.nope)?,
+ })
+ }
+
+ 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)
+ + rows * d.heads * act32_col_bytes(d.nope)
+ }
+}
+
+/// 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 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, nope: 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)?,
+ nope: DeviceBuffer::zeroed(stream, slots * nope)?,
+ lat: DeviceBuffer::zeroed(stream, slots * LATENT)?,
+ })
+ }
+
+ 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.nope.num_bytes() + self.lat.num_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, d.nope)?,
+ 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, d.nope)
+ + 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()
+ + self.nope32.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 queries' nope parts into the key absorption ([`absorb_heads`]).
+ Nope,
+ /// 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::Nope => (None, Some(&forms.nope32)),
+ 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 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 ([`route_heads`]), 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);
+ heads_filled(&b.heads.route, m * heads)?;
+ gpu.enqueue_quantize_gemm(att, m * heads, &mut b.forms.lat128, 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> {
+ out_runs(ty)?;
+ 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> {
+ 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> {
+ 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)"),
+ ));
+ }
+ 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
+/// 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);
+ rda_runs(s)?;
+ 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)",
+ )),
+ }
+}
+
+/// 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
+/// 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 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,
+ },
+ Some(DownForm::Q32) => SwigluRows {
+ q128: None,
+ q32: Some(&mut forms.shared32),
+ f32: None,
+ },
+ None => return Err(no_down_form(l, g)),
+ };
+ 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 RoutedStacks { ty, gate, up, down } = routed_stacks(w, l, &r.stacks)?;
+ 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)
+}
+
+#[cfg(test)]
+mod tests {
+ 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::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
+ /// [`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/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/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/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),
diff --git a/docs/cards/raonpf-pp.card b/docs/cards/raonpf-pp.card
new file mode 100644
index 000000000..625b53cec
--- /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 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?
+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
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 15bca5a8f..2a1a6d954 100644
--- a/tools/ref/ptx-shapes.tsv
+++ b/tools/ref/ptx-shapes.tsv
@@ -379,6 +379,8 @@ 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 raon_row_window 0 0
generate_ds41 relu_sqr 0 0
generate_ds41 relu_sqr_quant32 0 0
generate_ds41 rms_norm 0 0
diff --git a/tools/unsafe-ratchet.txt b/tools/unsafe-ratchet.txt
index 1e29b4bef..e7b4504fd 100644
--- a/tools/unsafe-ratchet.txt
+++ b/tools/unsafe-ratchet.txt
@@ -415,7 +415,12 @@ 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.
+# 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.