gpu, model, gpu-gates, tools: the common pieces a Nemotron-H card body attaches to

crates/gpu-gates/src/bin/gate_qwen3moe_experts.rs

@@ -53,6 +53,9 @@
 //!
 //! And each kernel as a captured graph: one node, the replay equal to the
 //! eager launch.
+//!
+//! Q5_0 embedding rows (`embed_rows_q5_0`) are bit-identical to `dequant_row`
+//! on the same bytes.
 
 #[cfg(feature = "gpu")]
 fn main() -> std::process::ExitCode {
@@ -136,6 +139,7 @@ mod gate {
         let mut ok = true;
         ok &= embed(&gpu, &gguf, &sets)?;
         ok &= embed_kquant(&gpu)?;
+        ok &= embed_q5_0(&gpu)?;
         ok &= gate_up_body(&gpu, &gguf, &ex)?;
         ok &= gate_up_tap(&gpu, &gguf, &ex, &sets)?;
         ok &= combine(&gpu, &ex, &sets, stream)?;
@@ -146,6 +150,63 @@ mod gate {
         Ok(())
     }
 
+    /// Q5_0 rows through `embed_rows_q5_0` against `dequant_row` on the same
+    /// bytes, bit for bit, positions included, one captured node, and an id
+    /// past the table (NaN row, `FaultSite::TokenId`, the other rows intact).
+    fn embed_q5_0(gpu: &Gpu) -> Result<bool, GateError> {
+        const ROWS: usize = 4;
+        // 34 blocks: past the first 32-block window of the packed code plane.
+        const K: usize = 1088;
+        // Reuse synthetic_rows with the Q5_0 scale at byte 0; host emulation
+        // confirms this seed rejects both addressing mutants.
+        let bytes = synthetic_rows((GgmlType::Q5_0, &[0]), (ROWS, K / 32), 0x2545_f491)?;
+        let packed = bloomery_gpu::q5::pack_q5_0(&bytes, K, ROWS)?;
+        let cols = packed.len() / ROWS;
+        let table = DeviceTensor::upload(gpu.stream(), &packed, ROWS, cols)?;
+        let ids = [0u32, (ROWS - 1) as u32, 1];
+        let launch = |ids: &[u32], graph: bool| {
+            gather(
+                gpu,
+                (&table, Some(GgmlType::Q5_0)),
+                ids,
+                POS0,
+                FIRST,
+                K,
+                graph,
+            )
+        };
+        let (got, _) = launch(&ids, false)?;
+        let mut want = vec![0.0f32; ids.len() * K];
+        let rb = (K / 32) * 22;
+        for (t, &id) in ids.iter().enumerate() {
+            let row = &bytes[id as usize * rb..(id as usize + 1) * rb];
+            dequant_row(GgmlType::Q5_0, row, &mut want[t * K..(t + 1) * K])?;
+        }
+        let host_ok = bits_equal(&got.y, &want);
+        let pos_ok = positions_ok(&got, ids.len(), POS0, FIRST);
+        let (replayed, nodes) = launch(&ids, true)?;
+        let graph_ok = bits_equal(&replayed.y, &got.y)
+            && replayed.pos == got.pos
+            && replayed.n_keys == got.n_keys
+            && nodes == 1;
+        let clean = gpu.take_fault()?;
+        let bad = [ids[0], ROWS as u32, ids[2]];
+        let (ybad, _) = launch(&bad, false)?;
+        let fault = gpu.take_fault()?;
+        let want_fault = Fault::at(LAYER_NONE, FaultSite::TokenId);
+        let others =
+            bits_equal(&ybad.y[..K], &got.y[..K]) && bits_equal(&ybad.y[2 * K..], &got.y[2 * K..]);
+        let nan = ybad.y[K..2 * K].iter().filter(|v| v.is_nan()).count();
+        let bad_pos = positions_ok(&ybad, bad.len(), POS0, FIRST);
+        let oor_ok = clean.is_none() && fault == Some(want_fault) && others && nan == K && bad_pos;
+        println!(
+            "embed Q5_0 K={K} rows={ROWS}: vs dequant_row bit_identical={host_ok} positions={pos_ok} \
+             graph_nodes={nodes} graph_same={graph_ok} oor_nan={nan}/{K} {}",
+            verdict(host_ok && pos_ok && graph_ok && oor_ok)
+        );
+        Ok(host_ok && pos_ok && graph_ok && oor_ok)
+    }
+
     // ------------------------------------------------------------ 1. embed
 
     /// One gather's outputs, read back: the rows, each row's position and
@@ -159,7 +220,8 @@ mod gate {
     /// `ids` gathered from `table` (`k` values a row), the rows' positions
     /// from `pos0` and offset `first`: eagerly, or captured as a graph and
     /// replayed once, which also returns its node count. `ty` `None` takes
-    /// the Q4_K entry directly, `Some(t)` the K-quant entry of type `t`.
+    /// the Q4_K entry directly, `Some(Q5_0)` the Q5_0 entry, `Some(t)` the
+    /// K-quant entry of type `t`.
     #[allow(clippy::too_many_arguments, reason = "one launch's inputs, each named")]
     fn gather(
         gpu: &Gpu,
@@ -188,6 +250,7 @@ mod gate {
             };
             match ty {
                 None => gpu.elem().enqueue_embed_rows_q4k(s, args),
+                Some(GgmlType::Q5_0) => gpu.elem().enqueue_embed_rows_q5_0(s, args, k),
                 Some(ty) => gpu.elem().enqueue_embed_rows_kquant(s, ty, args),
             }
         };

crates/gpu-gates/src/bin/gate_qwen3moe_flash.rs

@@ -145,7 +145,7 @@ fn main() -> std::process::ExitCode {
 mod gate {
     use bloomery_gpu::fault::{Fault, FaultSink, FaultSite, LAYER_NONE};
     use bloomery_gpu::flash_gqa::{
-        FlashGqaKernels, GROUP, GqaArgs, GqaQ8Args, HEAD, KEY_TILE, SEG_KEYS, SEGMENTS,
+        FlashGqaKernels, GROUP, GROUP_5, GqaArgs, GqaQ8Args, HEAD, KEY_TILE, SEG_KEYS, SEGMENTS,
         partials_ms_len, partials_v_len, seg_span,
     };
     use bloomery_gpu::flash_gqa_prefill::{
@@ -241,7 +241,7 @@ mod gate {
         let mut pv = DeviceBuffer::<f32>::zeroed(stream, partials_v_len(m, n_head))?;
         let mut pms = DeviceBuffer::<f32>::zeroed(stream, partials_ms_len(m, n_head))?;
         let mut y = DeviceBuffer::<f32>::zeroed(stream, m * n_head * HEAD)?;
-        k.enqueue_pass(
+        k.enqueue_pass::<GROUP>(
             stream,
             GqaArgs {
                 q,
@@ -500,7 +500,7 @@ mod gate {
                             },
                         )
                     } else {
-                        k.enqueue_pass(
+                        k.enqueue_pass::<GROUP>(
                             s,
                             GqaArgs {
                                 q: &inp.q,
@@ -620,7 +620,7 @@ mod gate {
             let mut pv = DeviceBuffer::<f32>::zeroed(stream, partials_v_len(1, n_head))?;
             let mut pms = DeviceBuffer::<f32>::zeroed(stream, partials_ms_len(1, n_head))?;
             let mut y = DeviceBuffer::<f32>::zeroed(stream, n_head * HEAD)?;
-            k.enqueue_pass(
+            k.enqueue_pass::<GROUP>(
                 stream,
                 GqaArgs {
                     q: &qd,
@@ -772,7 +772,7 @@ mod gate {
             let mut pv = DeviceBuffer::<f32>::zeroed(stream, partials_v_len(ROWS, n_head))?;
             let mut pms = DeviceBuffer::<f32>::zeroed(stream, partials_ms_len(ROWS, n_head))?;
             let mut y = DeviceBuffer::<f32>::zeroed(stream, ROWS * n_head * HEAD)?;
-            k.enqueue_pass(
+            k.enqueue_pass::<GROUP>(
                 stream,
                 GqaArgs {
                     q: &rows_q,
@@ -1901,6 +1901,80 @@ mod gate {
         plain: f64,
     }
 
+    /// Head 128, [`GROUP_5`] query heads a key head, the scalar pass,
+    /// against [`exact`]'s head-128 band. Two key heads: a KV index of
+    /// `h / 8` still reads the first plane for the second group.
+    fn group5(k: &FlashGqaKernels, gpu: &Gpu, scale: f32) -> Result<bool, GateError> {
+        const N_KV: usize = 2;
+        const KEYS: usize = 96;
+        const CTX: usize = 128;
+        let n_head = N_KV * GROUP_5;
+        let stream = gpu.stream();
+        let q: Vec<f32> = activations(HEAD, n_head, 5);
+        let kb: Vec<u16> = activations(HEAD, N_KV * CTX, 6)
+            .into_iter()
+            .map(f32_to_f16_bits)
+            .collect();
+        let vb: Vec<u16> = activations(HEAD, N_KV * CTX, 7)
+            .into_iter()
+            .map(f32_to_f16_bits)
+            .collect();
+        let kf: Vec<f32> = kb.iter().copied().map(half_to_f32).collect();
+        let vf: Vec<f32> = vb.iter().copied().map(half_to_f32).collect();
+        let q_d = DeviceBuffer::from_host(stream, &q)?;
+        let kc = DeviceBuffer::from_host(stream, &kb)?;
+        let vc = DeviceBuffer::from_host(stream, &vb)?;
+        let n_keys = DeviceBuffer::from_host(stream, &[KEYS as u32])?;
+        let mut part_v = DeviceBuffer::<f32>::zeroed(stream, partials_v_len(1, n_head))?;
+        let mut part_ms = DeviceBuffer::<f32>::zeroed(stream, partials_ms_len(1, n_head))?;
+        let mut y = DeviceBuffer::<f32>::zeroed(stream, n_head * HEAD)?;
+        k.enqueue_pass::<GROUP_5>(
+            stream,
+            GqaArgs {
+                q: &q_d,
+                kc: &kc,
+                vc: &vc,
+                n_keys: &n_keys,
+                scale,
+                n_kv: N_KV,
+                ctx: CTX,
+                m: 1,
+                part_v: &mut part_v,
+                part_ms: &mut part_ms,
+                fault: gpu.unlabelled_sink(),
+                y: &mut y,
+            },
+            false,
+        )?;
+        stream.synchronize()?;
+        let got = y.to_host_vec(stream)?;
+        let mut worst = 0.0f64;
+        let mut ok = true;
+        for h in 0..n_head {
+            let plane = (h / GROUP_5) * CTX * HEAD;
+            let ex = exact(
+                &q[h * HEAD..(h + 1) * HEAD],
+                &kf[plane..plane + KEYS * HEAD],
+                &vf[plane..plane + KEYS * HEAD],
+                KEYS,
+                scale,
+            );
+            for d in 0..HEAD {
+                let err = (f64::from(got[h * HEAD + d]) - ex.o[d]).abs();
+                let bound = ex.bound_ours[d];
+                if bound > 0.0 {
+                    worst = worst.max(err / bound);
+                }
+                ok &= err <= bound;
+            }
+        }
+        println!(
+            "flash group5 head={HEAD} group={GROUP_5} kv={N_KV} keys={KEYS}: measured/bound {worst:.3e} {}",
+            verdict(ok)
+        );
+        Ok(ok)
+    }
+
     pub fn run() -> Result<(), GateError> {
         // A lever set to a value it does not take, or a retired name that is
         // set, is refused by name before anything loads.
@@ -2142,6 +2216,7 @@ mod gate {
         );
         ok &= seeded_ok;
         ok &= q8_window_refusal(&kp, &gpu, scale)?;
+        ok &= group5(&k, &gpu, scale)?;
         println!("gate_qwen3moe_flash: {}", verdict(ok));
         if !ok {
             return Err(checks_failed());

crates/gpu/src/arch/qwen3moe/body35.rs

@@ -229,7 +229,7 @@ fn checkpoint_lens(kinds: &[Kind35]) -> Vec<usize> {
     kinds
         .iter()
         .filter_map(|k| match k {
-            Kind35::Delta(shape) => Some(RecStore::lens(*shape, PLAIN_LANES)),
+            Kind35::Delta(shape) => Some(shape.lens(PLAIN_LANES)),
             Kind35::Gqa(_) => None,
         })
         .flatten()
@@ -1618,9 +1618,11 @@ impl Body35 {
             stores.push(match *kind {
                 Kind35::Gqa(_) => LayerStore::Kv(KvPlanes::new(stream, &d, o.kv)?),
                 Kind35::Delta(shape) if lanes > PLAIN_LANES => {
-                    LayerStore::Rec(RecStore::stamped(stream, shape, lanes)?)
+                    LayerStore::Rec(RecStore::stamped(stream, shape.lens(lanes), lanes)?)
+                }
+                Kind35::Delta(shape) => {
+                    LayerStore::Rec(RecStore::new(stream, shape.lens(lanes), lanes)?)
                 }
-                Kind35::Delta(shape) => LayerStore::Rec(RecStore::new(stream, shape, lanes)?),
             });
         }
         let forms = Forms::of(&plans, &d);
@@ -3153,7 +3155,7 @@ impl Slots for Body35 {
                         Ok(LayerStore::Kv(KvPlanes::new(stream, &d, kv)?))
                     }
                     (MixerPlan::Delta(dp), LayerStore::Rec(_)) => Ok(LayerStore::Rec(
-                        RecStore::new(stream, dp.shape, PLAIN_LANES)?,
+                        RecStore::new(stream, dp.shape.lens(PLAIN_LANES), PLAIN_LANES)?,
                     )),
                     _ => Err(GpuError::state(
                         WHAT_N,
@@ -3335,7 +3337,7 @@ mod tests {
     /// layer's state (2,097,152 B) and conv ring (360,448 B): 73,728,000 B,
     /// spaced at the 4,096-token default ubatch, so a slot of 2,048 positions
     /// pins one point and one of 10,240 three. The lengths are the stores'
-    /// own ([`RecStore::lens`]), in [`Copied35`]'s order.
+    /// own ([`LinearShape::lens`]), in [`Copied35`]'s order.
     #[test]
     fn a_checkpoint_copies_each_delta_layers_state_and_ring() {
         let delta = Kind35::Delta(LinearShape {

crates/gpu/src/arch/qwen3moe/dispatch.rs

@@ -473,7 +473,7 @@ pub(super) fn mixer(
 /// and each row's position and live key count into `s.pos` and `s.n_keys` —
 /// the rows every later launch of the unit reads. The table's resident type
 /// picks the lookup: Q3_K, Q4_K, Q5_K or Q6_K rows, or Q8_0 planes (the load
-/// admits only these).
+/// admits these). Q5_0 packed rows are served here for a body that admits them.
 pub(super) fn embed_rows(
     gpu: &Gpu,
     w: &Weights,
@@ -486,7 +486,9 @@ pub(super) fn embed_rows(
 
 /// [`embed_rows`] into `(y, pos, n_keys)`, a unit's embedding rows and
 /// their positions and live key counts: the arena's own, or a window of
-/// them where a pass of several slots embeds each slot's rows apart.
+/// them where a pass of several slots embeds each slot's rows apart. Q5_0
+/// packed rows take [`crate::elem::ElemKernels::enqueue_embed_rows_q5_0`]
+/// for a body that admits them.
 pub(super) fn embed_into(
     gpu: &Gpu,
     w: &Weights,
@@ -516,9 +518,30 @@ pub(super) fn embed_into(
             },
         );
     }
+    if let Some(DevWeight::Q5_0 { w, k }) = w.get(&name) {
+        return gpu.elem().enqueue_embed_rows_q5_0(
+            gpu.stream(),
+            EmbedRowsArgs {
+                w,
+                ids: io.ids,
+                pos0: io.pos0,
+                first: io.first,
+                y,
+                pos,
+                n_keys,
+            },
+            *k,
+        );
+    }
     let ty = match w.get(&name) {
         Some(DevWeight::KQuant { ty, .. }) => *ty,
-        _ => return Err(GpuError::tensor(WHAT, &name, "K-quant rows or Q8_0 planes")),
+        _ => {
+            return Err(GpuError::tensor(
+                WHAT,
+                &name,
+                "K-quant rows or Q8_0 planes; Q5_0 packed rows when the body admits them",
+            ));
+        }
     };
     gpu.elem().enqueue_embed_rows_kquant(
         gpu.stream(),