diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 0cb2aec6..fe496e5b 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -23,6 +23,7 @@ authors = ["Eric Kryski (@ekryski)", "Tom Turney (@TheTom)"] repository = "https://github.com/TheTom/ffai" [workspace.dependencies] +rayon = "1" # ── engine crates ──────────────────────────────────────────────────── ffai-core = { path = "crates/ffai-core" } ffai-ops = { path = "crates/ffai-ops" } diff --git a/rust/crates/backends/ffai-cuda/src/imp.rs b/rust/crates/backends/ffai-cuda/src/imp.rs index 7a1360f0..97743283 100644 --- a/rust/crates/backends/ffai-cuda/src/imp.rs +++ b/rust/crates/backends/ffai-cuda/src/imp.rs @@ -2,8 +2,8 @@ // SPDX-License-Identifier: Apache-2.0 //! Real CUDA backend (compiled under `--features cuda`). Wraps -//! `metaltile_runtime::CudaDevice` — the NVRTC → PTX → driver path that -//! runs the kernel corpus bit-accurately on GB10 — behind the shared +//! `metaltile_runtime::CudaDevice`, the NVRTC → PTX → driver path that +//! runs the kernel corpus bit-accurately on GB10, behind the shared //! [`ffai_core::Device`] trait, so the engine layer above is identical to //! the Metal/Vulkan paths. @@ -27,12 +27,12 @@ fn dispatch_err(e: MetalTileError) -> Error { /// /// Holds an `Arc` so the CUDA CONTEXT outlives this module. A /// `CudaModule`'s `Drop` calls `cuModuleUnload`, which faults -/// (STATUS_ACCESS_VIOLATION, 0xc0000005) if the context was already destroyed — +/// (STATUS_ACCESS_VIOLATION, 0xc0000005) if the context was already destroyed, /// which is exactly what happened at process teardown: this device's `dev` field /// (its only context-keepalive) dropped before the `modules` cache, destroying /// the context, then the cached modules unloaded into the dead context. Pinning /// the context here (and dropping `module` before `_dev` via field order below) -/// guarantees the module always unloads while its context is still live — the +/// guarantees the module always unloads while its context is still live, the /// same invariant `CudaBuffer` already upholds for device allocations. struct CachedModule { // Field DECLARATION ORDER is DROP ORDER: `module` (cuModuleUnload) must run @@ -51,7 +51,7 @@ pub struct CudaDevice { name: String, /// Compile-once cache: kernel name → loaded module. modules: RwLock>>, - /// Memoized shared-mem bytes per (kernel, block_x) — avoids re-walking the + /// Memoized shared-mem bytes per (kernel, block_x), avoids re-walking the /// kernel IR on every dispatch (hot in decode: ~hundreds of dispatches/token). shared: RwLock>, } @@ -74,8 +74,26 @@ impl CudaDevice { } } + /// Cache key for a compiled module. The bare kernel name is NOT unique: + /// registry kernels share one name across their dtype variants (the ops + /// layer keys its IR cache by (name, dtype), but the same name reaches + /// this backend for every variant). Keying compiled modules by name alone + /// silently reused the first-compiled dtype's binary for every other + /// dtype, which reads with the wrong element stride and faults (or worse, + /// corrupts silently). Fold each param's dtype into the key. + fn module_key(kernel: &Kernel) -> String { + use std::fmt::Write as _; + let mut key = kernel.name.clone(); + key.push('|'); + for p in &kernel.params { + let _ = write!(key, "{:?},", p.dtype); + } + key + } + fn module_for(&self, kernel: &Kernel) -> Result> { - if let Some(m) = self.modules.read().unwrap().get(&kernel.name) { + let mkey = Self::module_key(kernel); + if let Some(m) = self.modules.read().unwrap().get(&mkey) { return Ok(m.clone()); } let cg = CudaGenerator::new(); @@ -95,7 +113,7 @@ impl CudaDevice { self.modules .write() .unwrap() - .insert(kernel.name.clone(), cached.clone()); + .insert(mkey, cached.clone()); Ok(cached) } } @@ -124,7 +142,7 @@ impl DeviceBuffer for CudaBuffer { impl Drop for CudaBuffer { fn drop(&mut self) { // Return to the device's size-bucketed pool instead of a synchronous - // cuMemFree — the hot decode path reuses these every token. + // cuMemFree, the hot decode path reuses these every token. self.dev.free_raw_pooled(self.ptr, self.len); } } @@ -162,7 +180,7 @@ impl Device for CudaDevice { } fn alloc_zeroed(&self, len: usize) -> Result> { - // Stream-ordered device memset — no host zero buffer, no pageable + // Stream-ordered device memset, no host zero buffer, no pageable // H2D copy. A [s,hid] f32 MoE accumulator is ~22 MB; uploading zeros // for it cost ~4-6 ms of host-blocking copy per E-layer. let ptr = self.dev.alloc_raw(len).map_err(dispatch_err)?; @@ -181,7 +199,9 @@ impl Device for CudaDevice { fn dispatch(&self, kernel: &Kernel, bindings: &[Binding], grid: Grid) -> Result<()> { let module = self.module_for(kernel)?; let func = module.module.function(&kernel.name).map_err(dispatch_err)?; - let skey = (kernel.name.clone(), grid.block[0]); + // Same dtype-collision hazard as the module cache: shared-memory sizing + // depends on element widths, so fold the param dtypes in. + let skey = (Self::module_key(kernel), grid.block[0]); let shared = if let Some(&s) = self.shared.read().unwrap().get(&skey) { s } else { @@ -230,7 +250,7 @@ impl Device for CudaDevice { } } - // Async launch (no per-dispatch cuCtxSynchronize) — kernels pipeline on the + // Async launch (no per-dispatch cuCtxSynchronize), kernels pipeline on the // ordered default stream; `download`/`synchronize` sync when results are read. self.dev .launch_async(func, grid.grid, grid.block, shared, &mut args) @@ -311,7 +331,7 @@ impl Device for CudaDevice { // Use cooperative launch only when NOT inside a CUDA graph capture // (cuLaunchCooperativeKernel is not capturable). Fall back to regular - // launch during capture — the caller (moe_fused_ffn) must not use + // launch during capture, the caller (moe_fused_ffn) must not use // grid.sync() in that code path (handled by NEMOTRON_GRAPH exclusion). if cooperative && !self.dev.is_capturing() { self.dev.launch_async_coop(func, grid, block, shared_bytes, &mut args).map_err(dispatch_err) diff --git a/rust/crates/backends/ffai-cuda/tests/colcopy.rs b/rust/crates/backends/ffai-cuda/tests/colcopy.rs new file mode 100644 index 00000000..4cee679c --- /dev/null +++ b/rust/crates/backends/ffai-cuda/tests/colcopy.rs @@ -0,0 +1,99 @@ +#![cfg(feature = "cuda")] +//! strided_col_copy launch-geometry regression: the op dispatches with a +//! fixed 64-thread block. Before the fix, the grid was `n / 64` (integer +//! division, no ceiling), so whenever `s * width` wasn't a multiple of 64 +//! the tail elements were silently dropped (never written into dst, left +//! as zeroed garbage). Covers a deliberately non-64-aligned shape (the +//! Laguna-style head-count case, e.g. width=72) plus a 64-aligned control +//! shape to confirm the fix didn't regress the common case. +use ffai_core::{DType, Device, Tensor}; +use ffai_cuda::CudaDevice; +use ffai_ops::strided_col_copy; + +fn tb(v: &[f32]) -> Vec { v.iter().flat_map(|x| x.to_le_bytes()).collect() } +fn fb(b: &[u8]) -> Vec { b.chunks_exact(4).map(|c| f32::from_le_bytes(c.try_into().unwrap())).collect() } +fn tn(d: &dyn Device, v: &[f32]) -> Tensor { Tensor::new(d.upload(&tb(v)).unwrap(), vec![v.len()], DType::F32) } + +// Host-computed expectation: dst[ti*width + ci] = src[ti*stride + col_off + ci]. +fn host_expected(src: &[f32], s: usize, stride: usize, col_off: usize, width: usize) -> Vec { + (0..s) + .flat_map(|ti| (0..width).map(move |ci| src[ti * stride + col_off + ci])) + .collect() +} + +fn run_case(name: &str, s: usize, stride: usize, col_off: usize, width: usize) { + let Some(dev) = CudaDevice::create().expect("cuda") else { + eprintln!("no CUDA, skipping {name}"); + return; + }; + let d = dev.as_ref(); + let src: Vec = (0..s * stride).map(|i| ((i as f32) * 0.019).sin()).collect(); + let src_t = tn(d, &src); + + let out = strided_col_copy(d, &src_t, s, stride, col_off, width) + .unwrap_or_else(|e| panic!("{name}: strided_col_copy dispatch failed: {e:?}")); + d.synchronize().unwrap(); + + let mut bytes = vec![0u8; s * width * 4]; + d.download(out.buffer.as_ref(), &mut bytes).unwrap(); + let got = fb(&bytes); + let want = host_expected(&src, s, stride, col_off, width); + + assert_eq!(got.len(), want.len(), "{name}: output length mismatch"); + let mut max_diff = 0f32; + let mut first_bad: Option = None; + for (i, (g, w)) in got.iter().zip(want.iter()).enumerate() { + let diff = (g - w).abs(); + if diff > max_diff { max_diff = diff; } + if diff > 1e-6 && first_bad.is_none() { first_bad = Some(i); } + } + eprintln!("[{name}] s={s} stride={stride} col_off={col_off} width={width} n={} max|GPU-CPU|={max_diff:e}", s * width); + if let Some(i) = first_bad { + panic!( + "{name}: mismatch at flat index {i} (ti={}, ci={}): got={} want={}", + i / width, i % width, got[i], want[i] + ); + } +} + +/// Deliberately non-64-aligned shape: s*width = 3*72 = 216, not a multiple of +/// 64. This is the Laguna-style case (head counts 48/72) that exposed the +/// launch-geometry bug: with the old `n / 64` grid, only 192 of 216 elements +/// were ever dispatched, so the last 24 (rows 2's tail) came back as +/// zero-initialized garbage instead of the real column values. +#[test] +fn strided_col_copy_non_64_aligned_shape() { + run_case("non_64_aligned", 3, 100, 4, 72); +} + +/// Control: a 64-aligned shape (s*width divisible by 64) exercising the +/// previously-working path, to confirm the div_ceil + in-kernel guard fix +/// didn't regress the common (aligned) case. +#[test] +fn strided_col_copy_64_aligned_control() { + run_case("64_aligned_control", 4, 200, 8, 64); +} + +/// Micro-repro for the prefill layer-1 fault: f16 table gather at the exact +/// failing shapes (table [2, 3072] f16 from a cast, indices [20] u32 of 0/1). +#[test] +fn gather_f16_prefill_shapes() { + use ffai_core::{DType, Tensor}; + use ffai_cuda::CudaDevice; + let Some(dev) = CudaDevice::create().expect("cuda") else { eprintln!("no CUDA, skip"); return; }; + let d = dev.as_ref(); + let t = 2usize; + let hidden = 3072usize; + let x: Vec = (0..t * hidden).map(|i| (i as f32) * 0.001 - 3.0).collect(); + let xb: Vec = x.iter().flat_map(|v| v.to_le_bytes()).collect(); + let xf32 = Tensor::new(d.upload(&xb).unwrap(), vec![t, hidden], DType::F32); + let xf16 = ffai_ops::cast_f32_f16(d, &xf32).unwrap(); + let idx: Vec = vec![0, 0, 1, 1, 1, 0, 0, 1, 0, 1, 0, 0, 1, 1, 1, 0, 1, 1, 0, 0]; + let ib: Vec = idx.iter().flat_map(|v| v.to_le_bytes()).collect(); + let ti = Tensor::new(d.upload(&ib).unwrap(), vec![idx.len()], DType::U32); + let out = ffai_ops::gather(d, &xf16, &ti).expect("gather f16 dispatch"); + d.synchronize().expect("post-gather sync"); + let mut ob = vec![0u8; idx.len() * hidden * 2]; + d.download(out.buffer.as_ref(), &mut ob).expect("download"); + eprintln!("gather f16 prefill shapes: OK (first half-word {:02x}{:02x})", ob[0], ob[1]); +} diff --git a/rust/crates/backends/ffai-cuda/tests/laguna.rs b/rust/crates/backends/ffai-cuda/tests/laguna.rs index 89f72663..343dcf84 100644 --- a/rust/crates/backends/ffai-cuda/tests/laguna.rs +++ b/rust/crates/backends/ffai-cuda/tests/laguna.rs @@ -119,12 +119,141 @@ fn laguna_gate_unit() { eprintln!("gate softplus per-head: OK"); } +/// Batched-prefill `rope_yarn_partial_many` must match a T-loop of the +/// scalar-position `rope_yarn_partial` (the proven decode kernel), one call +/// per token row at that row's absolute position. Exercises both a YaRN +/// full-attention-layer shape (n_rot < head_dim, ext_factor != 0) and a +/// plain sliding-layer shape (n_rot == head_dim, ext_factor == 0), matching +/// Laguna's two RoPE configs. +#[test] +fn laguna_rope_yarn_many_matches_looped() { + use ffai_core::{DType, Tensor}; + let Some(dev) = CudaDevice::create().expect("cuda") else { eprintln!("no CUDA, skip"); return; }; + let d = dev.as_ref(); + let up = |v: &[f32]| -> Tensor { + let b: Vec = v.iter().flat_map(|x| x.to_le_bytes()).collect(); + Tensor::new(d.upload(&b).unwrap(), vec![v.len()], DType::F32) + }; + let dl = |t: &Tensor, n: usize| -> Vec { + let mut b = vec![0u8; n * 4]; + d.synchronize().unwrap(); + d.download(t.buffer.as_ref(), &mut b).unwrap(); + b.chunks_exact(4).map(|c| f32::from_le_bytes(c.try_into().unwrap())).collect() + }; + + // (n_rot, theta, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow, n_orig_ctx, label) + let configs: [(u32, f32, f32, f32, f32, f32, f32, u32, &str); 2] = [ + (64, 500000.0, 1.0 / 128.0, 1.0, 1.4852030263919618, 32.0, 1.0, 8192, "full/yarn"), + (128, 10000.0, 1.0, 0.0, 1.0, 32.0, 1.0, 8192, "sliding/plain"), + ]; + + for (n_rot, theta, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow, n_orig_ctx, label) in configs { + let t = 5usize; + let n_heads = 6usize; + let hd = 128usize; + let n = t * n_heads * hd; + let x: Vec = (0..n).map(|i| ((i * 2654435761usize) % 2000) as f32 * 0.001 - 1.0).collect(); + let positions: Vec = (0..t).map(|i| (7 + i * 13) as u32).collect(); // non-contiguous, non-zero-based + + // Batched: one dispatch over [T, n_heads, hd]. + let xt = up(&x).reshaped(vec![t, n_heads, hd]); + let posb = Tensor::new( + d.upload(&positions.iter().flat_map(|p| p.to_le_bytes()).collect::>()).unwrap(), + vec![t], DType::U32, + ); + let got = ffai_ops::rope_yarn_partial_many( + d, &xt, &posb, n_heads, n_rot, theta, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow, n_orig_ctx, + ).unwrap(); + let gv = dl(&got, n); + + // Looped reference: one rope_yarn_partial call per token row at its + // own absolute position, over that row's [n_heads, hd] slab. + let mut ev = vec![0f32; n]; + for (ti, &pos) in positions.iter().enumerate() { + let row = &x[ti * n_heads * hd..(ti + 1) * n_heads * hd]; + let rt = up(row).reshaped(vec![n_heads, hd]); + let ro = ffai_ops::rope_yarn_partial( + d, &rt, pos, n_rot, theta, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow, n_orig_ctx, + ).unwrap(); + let rov = dl(&ro, n_heads * hd); + ev[ti * n_heads * hd..(ti + 1) * n_heads * hd].copy_from_slice(&rov); + } + + let mut bad = 0; + for i in 0..n { + if (gv[i] - ev[i]).abs() > 1e-4 * ev[i].abs().max(1.0) { + if bad < 6 { + eprintln!("[{label}] MISMATCH i={i}: many={} looped={}", gv[i], ev[i]); + } + bad += 1; + } + } + assert_eq!(bad, 0, "rope_yarn_partial_many mismatch vs looped rope_yarn_partial ({label}, {bad} bad elems)"); + eprintln!("rope_yarn_partial_many vs looped [{label}]: OK"); + } +} + +/// Batched-prefill `gate_softplus_mul_perhead_many` must match a T-loop of +/// the scalar-row `gate_softplus_mul_perhead` (the proven decode kernel), +/// one call per token row. +#[test] +fn laguna_gate_many_matches_looped() { + use ffai_core::{DType, Tensor}; + let Some(dev) = CudaDevice::create().expect("cuda") else { eprintln!("no CUDA, skip"); return; }; + let d = dev.as_ref(); + let up = |v: &[f32]| -> Tensor { + let b: Vec = v.iter().flat_map(|x| x.to_le_bytes()).collect(); + Tensor::new(d.upload(&b).unwrap(), vec![v.len()], DType::F32) + }; + let dl = |t: &Tensor, n: usize| -> Vec { + let mut b = vec![0u8; n * 4]; + d.synchronize().unwrap(); + d.download(t.buffer.as_ref(), &mut b).unwrap(); + b.chunks_exact(4).map(|c| f32::from_le_bytes(c.try_into().unwrap())).collect() + }; + + let t = 4usize; + let n_heads = 5usize; // deliberately not 64-aligned (Laguna's real head counts: 48/72) + let hd = 8usize; + let n = t * n_heads * hd; + let attn: Vec = (0..n).map(|i| (i as f32) * 0.05 - 2.0).collect(); + let g: Vec = (0..t * n_heads).map(|i| (i as f32) * 3.7 - 6.0).collect(); // spans <0, ~0, >20 branches + + let at = up(&attn).reshaped(vec![t, n_heads, hd]); + let gt = up(&g); + let got = ffai_ops::gate_softplus_mul_perhead_many(d, &at, >, n_heads).unwrap(); + let gv = dl(&got, n); + + let mut ev = vec![0f32; n]; + for ti in 0..t { + let arow = &attn[ti * n_heads * hd..(ti + 1) * n_heads * hd]; + let grow = &g[ti * n_heads..(ti + 1) * n_heads]; + let art = up(arow).reshaped(vec![n_heads, hd]); + let grt = up(grow); + let ro = ffai_ops::gate_softplus_mul_perhead(d, &art, &grt).unwrap(); + let rov = dl(&ro, n_heads * hd); + ev[ti * n_heads * hd..(ti + 1) * n_heads * hd].copy_from_slice(&rov); + } + + let mut bad = 0; + for i in 0..n { + if (gv[i] - ev[i]).abs() > 1e-5 * ev[i].abs().max(1.0) { + if bad < 6 { + eprintln!("MISMATCH i={i}: many={} looped={}", gv[i], ev[i]); + } + bad += 1; + } + } + assert_eq!(bad, 0, "gate_softplus_mul_perhead_many mismatch vs looped gate_softplus_mul_perhead ({bad} bad elems)"); + eprintln!("gate_softplus_mul_perhead_many vs looped: OK"); +} + /// `sdpa_decode_nbuf` (the CUDA-graph-capturable raw-CUDA fallback used by /// Laguna's decode path under `FFAI_LAGUNA_GRAPH=1`) must match the registry -/// `ops::sdpa_decode` kernel's output — same dense-causal single-pass +/// `ops::sdpa_decode` kernel's output, same dense-causal single-pass /// online-softmax math, just reading `n_kv` from a device buffer at kernel /// runtime instead of a launch-time baked-in scalar. Checked at `n_kv` in -/// {1, 7, 512} (kv_stride fixed at 512) x `heads_per_group` in {6, 9} — +/// {1, 7, 512} (kv_stride fixed at 512) x `heads_per_group` in {6, 9}: /// Laguna's actual GQA ratios (48 full-attention / 72 sliding-window Q heads /// over 8 KV heads). #[test] @@ -272,6 +401,96 @@ fn laguna_moe_swiglu_fused_unit() { eprintln!("moe gather q4 fused swiglu vs unfused compose: OK"); } +/// `sdpa_multi_tc_varlen`'s new sliding-window block skip (`win != 0`) must +/// match the `win = 0` full-range/mask-only path exactly, the skip only +/// shrinks which query rows each KV block's GEMMs cover, the per-element +/// `seg_lo` mask inside `sdpa_tc_softmax_varlen` is still the correctness +/// source of truth either way, so both should reach the same result. +/// +/// `T = 16` query rows, sliding window 4, but `base_kv = 4080` so this +/// chunk continues a much longer prefix (`kv_stride = n_kv = 4096`). That +/// makes `n_kv` exceed the kernel's fixed 2048-wide KV block (`bk = +/// 2048.min(n_kv)`), so the KV range splits into two blocks: block 0 +/// `[0, 2048)` and block 1 `[2048, 4096)`. With window 4 every query's +/// `seg_lo` sits at `4077..4092`, entirely inside block 1, so under `win=4` +/// block 0 should be skipped ENTIRELY (q_cnt=0) while block 1 is computed +/// for every row, genuinely exercising the row-range skip, not just +/// threading the new parameter through as a no-op. Modeled on +/// `laguna_sdpa_decode_nbuf_matches_registry`'s synthetic-ramp-data + +/// download + relative-tolerance-compare harness. +#[test] +fn laguna_sdpa_varlen_window_skip_matches_full_range() { + use ffai_core::{DType, Tensor}; + let Some(dev) = CudaDevice::create().expect("cuda") else { eprintln!("no CUDA, skip"); return; }; + let d = dev.as_ref(); + + let hd = 8usize; + let n_kv_heads = 1usize; + let heads_per_group = 2usize; + let n_q_heads = n_kv_heads * heads_per_group; + let t = 16usize; // n_query (this chunk's rows) + let base_kv = 4080u32; // this chunk continues a long prefix + let kv_stride = 4096u32; // n_kv = base_kv + t = 4096 > bk(2048) => 2 blocks + let window = 4u32; + let scale = 1.0 / (hd as f32).sqrt(); + + let ramp = |n: usize, step: f32, start: f32| -> Vec { + (0..n).map(|i| ((start + i as f32 * step) % 2.0) - 1.0).collect() + }; + let up = |v: &[f32]| -> Tensor { + let b: Vec = v.iter().flat_map(|x| x.to_le_bytes()).collect(); + Tensor::new(d.upload(&b).unwrap(), vec![v.len()], DType::F32) + }; + + let q = ramp(t * n_q_heads * hd, 0.013, -0.4); + let k = ramp(n_kv_heads * kv_stride as usize * hd, 0.011, -0.5); + let v = ramp(n_kv_heads * kv_stride as usize * hd, 0.007, -0.3); + let tq = up(&q).reshaped(vec![t, n_q_heads, hd]); + let tk = up(&k).reshaped(vec![n_kv_heads, kv_stride as usize, hd]); + let tv = up(&v).reshaped(vec![n_kv_heads, kv_stride as usize, hd]); + + // Exactly Laguna's sliding-layer prefill convention: + // seg_lo[i] = max(0, (base_kv+i) - (window-1)). + let seg_lo_h: Vec = (0..t) + .map(|i| ((base_kv as i64 + i as i64) - (window as i64 - 1)).max(0) as u32) + .collect(); + let seg_lo = Tensor::new( + d.upload(&seg_lo_h.iter().flat_map(|x| x.to_le_bytes()).collect::>()).unwrap(), + vec![t], DType::U32, + ); + + let out_full = ffai_ops::sdpa_multi_tc_varlen( + d, &tq, &tk, &tv, &seg_lo, 0, hd, n_q_heads as u32, + base_kv, t as u32, kv_stride, heads_per_group as u32, true, scale, 0, + ).unwrap(); + let out_win = ffai_ops::sdpa_multi_tc_varlen( + d, &tq, &tk, &tv, &seg_lo, 0, hd, n_q_heads as u32, + base_kv, t as u32, kv_stride, heads_per_group as u32, true, scale, window, + ).unwrap(); + + let n = t * n_q_heads * hd; + let dl = |ten: &Tensor| -> Vec { + let mut b = vec![0u8; n * 4]; + d.synchronize().unwrap(); + d.download(ten.buffer.as_ref(), &mut b).unwrap(); + b.chunks_exact(4).map(|c| f32::from_le_bytes(c.try_into().unwrap())).collect() + }; + let fv = dl(&out_full); + let wv = dl(&out_win); + + let mut bad = 0; + for i in 0..n { + if (fv[i] - wv[i]).abs() > 1e-4 * fv[i].abs().max(1.0) { + if bad < 6 { + eprintln!("MISMATCH i={i}: win=0 {} vs win={window} {}", fv[i], wv[i]); + } + bad += 1; + } + } + assert_eq!(bad, 0, "sdpa_multi_tc_varlen window skip (win={window}) mismatch vs win=0 full-range ({bad} bad elems)"); + eprintln!("sdpa_multi_tc_varlen window skip vs full-range: OK"); +} + #[test] fn laguna_on_cuda() { let Some(dev) = CudaDevice::create().expect("cuda") else { eprintln!("no CUDA, skip"); return; }; diff --git a/rust/crates/backends/ffai-cuda/tests/moe_grouped_mma_test.rs b/rust/crates/backends/ffai-cuda/tests/moe_grouped_mma_test.rs index a4bd5321..7180f1a3 100644 --- a/rust/crates/backends/ffai-cuda/tests/moe_grouped_mma_test.rs +++ b/rust/crates/backends/ffai-cuda/tests/moe_grouped_mma_test.rs @@ -6,7 +6,7 @@ use ffai_cuda::CudaDevice; use ffai_core::{DType, Device, Tensor}; -use ffai_ops::{gemm_cublas_f32out, fwht128, gemm_fp8, fp8_quant, moe_q4_grouped_mma_dev, quantize_q4, moe_grouped_gemm_mma, moe_grouped_gemm_cutlass, cast_f32_f16, cast_f16_f32, lt_fp4_quant_g, lt_fp4_quant_g_46, dec_ue4m3, +use ffai_ops::{gemm_cublas_f32out, fwht128, gemm_fp8, fp8_quant, moe_q4_grouped_mma, moe_q4_grouped_mma_dev, quantize_q4, moe_grouped_gemm_mma, moe_grouped_gemm_cutlass, cast_f32_f16, cast_f16_f32, lt_fp4_quant_g, lt_fp4_quant_g_46, dec_ue4m3, moe_grouped_gemm_w4a8, w4a8_actquant_grouped, w4a8_packw, w4a8_sfa_offsets, moe_grouped_gemm_w8a8, w8a8_actquant_grouped, w8a8_packw, lt_fp4_quant_slab, lt_fp4_quant_grouped, moe_grouped_gemm_cutlass_fp4, rms_norm, gemm_cublas, cast_f32_bf16, cast_bf16_f32}; @@ -79,7 +79,7 @@ fn run_cutlass(d: &dyn Device, n_exp: usize, k: usize, n: usize, groups: &[(usiz let w_dev = cast_f32_f16(d, &up(&w_f, vec![n_exp, n, k])).unwrap().reshaped(vec![n_exp, n, k]); let out = match moe_grouped_gemm_cutlass(d, &a_dev, &w_dev, &g_starts, &expert_ids, n, k) { Ok(o) => o, - Err(e) => { eprintln!("cutlass unavailable ({e:?}) — skip"); return; } + Err(e) => { eprintln!("cutlass unavailable ({e:?}), skip"); return; } }; let got = dl(d, &cast_f16_f32(d, &out).unwrap()); let mut exp = vec![0f32; mt*n]; @@ -112,11 +112,11 @@ fn run_w4a8(d: &dyn Device, n_exp: usize, k: usize, n: usize, groups: &[(usize, // W4A8: fp8 e4m3 acts (group-aware per-32 SFA) x mxfp4 weights (per-32 SFB). let (sfa_off, sfa_bytes) = w4a8_sfa_offsets(&g_starts, k); let (a_q, a_sf) = match w4a8_actquant_grouped(d, &a_dev, &g_starts.iter().zip(g_starts.iter().skip(1)).map(|(s,e)|(e-s) as i32).collect::>(), &sfa_off, sfa_bytes, mt, k) { - Ok(o) => o, Err(e) => { eprintln!("w4a8 actquant unavailable ({e:?}) — skip"); return; } }; + Ok(o) => o, Err(e) => { eprintln!("w4a8 actquant unavailable ({e:?}): skip"); return; } }; let (w_pack, w_sf, sfb_eb) = match w4a8_packw(d, &w_dev, n_exp, n, k) { - Ok(o) => o, Err(e) => { eprintln!("w4a8 packw unavailable ({e:?}) — skip"); return; } }; + Ok(o) => o, Err(e) => { eprintln!("w4a8 packw unavailable ({e:?}): skip"); return; } }; let out = match moe_grouped_gemm_w4a8(d, &a_q, &a_sf, &sfa_off, &w_pack, &w_sf, sfb_eb, &g_starts, &expert_ids, n, k) { - Ok(o) => o, Err(e) => { eprintln!("w4a8 gemm unavailable ({e:?}) — skip"); return; } }; + Ok(o) => o, Err(e) => { eprintln!("w4a8 gemm unavailable ({e:?}): skip"); return; } }; let got = dl(d, &cast_f16_f32(d, &out).unwrap()); // host f32 reference + cosine @@ -156,11 +156,11 @@ fn run_w8a8(d: &dyn Device, n_exp: usize, k: usize, n: usize, groups: &[(usize, // W4A8: fp8 e4m3 acts (group-aware per-32 SFA) x mxfp4 weights (per-32 SFB). let (sfa_off, sfa_bytes) = w4a8_sfa_offsets(&g_starts, k); let (a_q, a_sf) = match w8a8_actquant_grouped(d, &a_dev, &g_starts.iter().zip(g_starts.iter().skip(1)).map(|(s,e)|(e-s) as i32).collect::>(), &sfa_off, sfa_bytes, mt, k) { - Ok(o) => o, Err(e) => { eprintln!("w8a8 actquant unavailable ({e:?}) — skip"); return; } }; + Ok(o) => o, Err(e) => { eprintln!("w8a8 actquant unavailable ({e:?}): skip"); return; } }; let (w_pack, w_sf, sfb_eb) = match w8a8_packw(d, &w_dev, n_exp, n, k) { - Ok(o) => o, Err(e) => { eprintln!("w8a8 packw unavailable ({e:?}) — skip"); return; } }; + Ok(o) => o, Err(e) => { eprintln!("w8a8 packw unavailable ({e:?}): skip"); return; } }; let out = match moe_grouped_gemm_w8a8(d, &a_q, &a_sf, &sfa_off, &w_pack, &w_sf, sfb_eb, &g_starts, &expert_ids, n, k) { - Ok(o) => o, Err(e) => { eprintln!("w8a8 gemm unavailable ({e:?}) — skip"); return; } }; + Ok(o) => o, Err(e) => { eprintln!("w8a8 gemm unavailable ({e:?}): skip"); return; } }; let got = dl(d, &cast_f16_f32(d, &out).unwrap()); // host f32 reference + cosine @@ -261,7 +261,7 @@ fn moe_grouped_w4a8_down() { #[test] fn moe_grouped_w4a8_down_pad() { - // DOWN-proj K padded to 1920 (=15*128) — verifies K%128 is the constraint. + // DOWN-proj K padded to 1920 (=15*128), verifies K%128 is the constraint. let Some(d) = CudaDevice::create().expect("cuda") else { eprintln!("no CUDA"); return; }; run_w4a8(d.as_ref(), 8, 1920, 2688, &[(0,96),(3,40),(5,128),(7,17)]); } @@ -434,6 +434,107 @@ fn moe_q4_grouped_bench() { eprintln!("moe_q4_grouped_bench: mt={mt} N={n} K={k} -> {ms:.3} ms, {tf:.1} TFLOP/s (f16 ceiling ~85, cutlass-fp4 ref 60)"); } +// FFAI_LAGUNA_SCHED=1 A/B: descending-group-order tile emission + banded +// N-block/tile grid transpose (see moe_q4_grouped_mma / _dev in ffai-ops). +// Deliberately skewed groups (one big 200-row group, several tiny 3-row +// groups, an empty group) so the reorder + banding actually rearranges tiles +// relative to the default expert-order/tile-major launch. +// +// The real correctness bar here is FFAI_LAGUNA_SCHED=1 vs the default path, +// NOT vs the host f32 reference: both device paths run the identical Q4 +// dequant + mma math per tile (same tile inputs, same K-reduction order +// inside a tile), the only thing the flag changes is which CTA runs which +// (tile, N-block) pair and in what order. Since tile-to-output-element +// mapping is unchanged (disjoint writes per tile) and per-tile accumulation +// order is unchanged (task constraint: "float accumulation order within a +// tile is unchanged"), the two orderings should be near bit-exact: 1e-3 +// relative is generous headroom, not a quantization-noise budget. The +// host-f32 reference is used only as a loose sanity check (cosine) that the +// Q4 pipeline itself is producing a real signal, not to gate the A/B. +#[test] +fn moe_q4_grouped_laguna_sched_ab() { + let Some(d) = CudaDevice::create().expect("cuda") else { eprintln!("no CUDA"); return; }; + let dev = d.as_ref(); + let (n_exp, k, n) = (6usize, 128usize, 128usize); + // (expert_id, rows): order is deliberately skewed and includes a + // genuinely empty group to exercise the padding-tile path. + let groups: [(usize, usize); 6] = [(0, 200), (1, 3), (2, 0), (3, 3), (4, 3), (5, 37)]; + let mut g_starts = vec![0usize]; + let mut expert_ids = vec![]; + for &(e, r) in &groups { expert_ids.push(e); g_starts.push(g_starts.last().unwrap() + r); } + let mt = *g_starts.last().unwrap(); + + let a_f = rng(mt * k, 41); + let w_f = rng(n_exp * n * k, 43); + let up = |v: &[f32], sh: Vec| -> Tensor { + Tensor::new(dev.upload(&tb(v)).unwrap(), vec![v.len()], DType::F32).reshaped(sh) }; + let a_dev = cast_f32_f16(dev, &up(&a_f, vec![mt, k])).unwrap().reshaped(vec![mt, k]); + let (qs_v, sc_v) = quantize_q4(&w_f, n_exp * n, k); + let qs_b: Vec = qs_v.iter().flat_map(|x| x.to_le_bytes()).collect(); + let qs = Tensor::new(dev.upload(&qs_b).unwrap(), vec![qs_v.len()], DType::U32); + let sc = cast_f32_f16(dev, &up(&sc_v, vec![sc_v.len()])).unwrap(); + + // host f32 reference (loose sanity check only, see comment above) + let mut exp = vec![0f32; mt * n]; + for g in 0..groups.len() { + let (eid, _) = groups[g]; + for t in g_starts[g]..g_starts[g + 1] { + for nn in 0..n { + let mut acc = 0f32; + for kk in 0..k { acc += a_f[t * k + kk] * w_f[(eid * n + nn) * k + kk]; } + exp[t * n + nn] = acc; + } + } + } + let sanity_cos = cos(&exp, &dl(dev, &cast_f16_f32(dev, &moe_q4_grouped_mma( + dev, &a_dev, &qs, &sc, &g_starts, &expert_ids, n, k).unwrap()).unwrap())); + eprintln!("moe_q4_grouped_laguna_sched_ab: sanity cosine (Q4 default vs f32 host ref) = {sanity_cos:.6}"); + assert!(sanity_cos > 0.99, "moe_q4_grouped_mma default path looks broken, cosine {sanity_cos:.6} vs host ref"); + + // the actual A/B: FFAI_LAGUNA_SCHED=1 vs default, max relative diff. + let rel_diff = |a: &[f32], b: &[f32]| -> f32 { + let mut maxr = 0f32; + for (x, y) in a.iter().zip(b.iter()) { + let rel = (x - y).abs() / x.abs().max(1e-2); + if rel > maxr { maxr = rel; } + } + maxr + }; + + // --- host-descriptor path (moe_q4_grouped_mma) --- + let out_default = dl(dev, &cast_f16_f32(dev, &moe_q4_grouped_mma( + dev, &a_dev, &qs, &sc, &g_starts, &expert_ids, n, k).unwrap()).unwrap()); + unsafe { std::env::set_var("FFAI_LAGUNA_SCHED", "1"); } + let out_sched = dl(dev, &cast_f16_f32(dev, &moe_q4_grouped_mma( + dev, &a_dev, &qs, &sc, &g_starts, &expert_ids, n, k).unwrap()).unwrap()); + unsafe { std::env::remove_var("FFAI_LAGUNA_SCHED"); } + let r_host = rel_diff(&out_default, &out_sched); + eprintln!("moe_q4_grouped_laguna_sched_ab (host path): sched-vs-default max_rel={r_host:.2e}"); + assert!(r_host < 1e-3, "host path: FFAI_LAGUNA_SCHED=1 vs default max_rel {r_host:.2e} > 1e-3"); + + // --- fully-on-device descriptor path (moe_q4_grouped_mma_dev) --- + let offs: Vec = g_starts.iter().flat_map(|&s| (s as u32).to_le_bytes()).collect(); + let off = Tensor::new(dev.upload(&offs).unwrap(), vec![g_starts.len()], DType::U32); + let out_dev_default = dl(dev, &cast_f16_f32(dev, &moe_q4_grouped_mma_dev( + dev, &a_dev, &qs, &sc, &off, n_exp, mt, n, k).unwrap()).unwrap()); + unsafe { std::env::set_var("FFAI_LAGUNA_SCHED", "1"); } + let out_dev_sched = dl(dev, &cast_f16_f32(dev, &moe_q4_grouped_mma_dev( + dev, &a_dev, &qs, &sc, &off, n_exp, mt, n, k).unwrap()).unwrap()); + unsafe { std::env::remove_var("FFAI_LAGUNA_SCHED"); } + let r_dev = rel_diff(&out_dev_default, &out_dev_sched); + eprintln!("moe_q4_grouped_laguna_sched_ab (dev path): sched-vs-default max_rel={r_dev:.2e}"); + assert!(r_dev < 1e-3, "dev path: FFAI_LAGUNA_SCHED=1 vs default max_rel {r_dev:.2e} > 1e-3"); + + // the dev-path descriptor builder (on-device, offsets-only) must also + // agree with the host-descriptor builder for both orderings. + let r_cross_default = rel_diff(&out_default, &out_dev_default); + let r_cross_sched = rel_diff(&out_sched, &out_dev_sched); + eprintln!("moe_q4_grouped_laguna_sched_ab: host-vs-dev max_rel default={r_cross_default:.2e} sched={r_cross_sched:.2e}"); + assert!(r_cross_default < 1e-3 && r_cross_sched < 1e-3, "host vs dev descriptor builders disagree"); + + eprintln!("moe_q4_grouped_laguna_sched_ab: PASS (default vs FFAI_LAGUNA_SCHED=1 agree within 1e-3 on skewed groups)"); +} + #[test] fn moe_cutlass_f16_bench() { let Some(d) = CudaDevice::create().expect("cuda") else { eprintln!("no CUDA"); return; }; diff --git a/rust/crates/ffai-loader/src/gguf.rs b/rust/crates/ffai-loader/src/gguf.rs index c6cb3073..d017f948 100644 --- a/rust/crates/ffai-loader/src/gguf.rs +++ b/rust/crates/ffai-loader/src/gguf.rs @@ -14,7 +14,7 @@ use std::collections::BTreeMap; use crate::iq2xxs_tables::{IQ2XXS_GRID, KSIGNS}; -/// GGML tensor type (subset). +/// GGUF tensor dtype code (subset). #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum GgmlType { F32, @@ -50,7 +50,7 @@ pub struct GgufTensor { pub ggml_type: GgmlType, pub offset: u64, // within the data section /// Which file part this tensor's data lives in (0 for single-file GGUF; - /// the split index for llama.cpp multi-part `*-NNNNN-of-MMMMM.gguf`). + /// the split index for the reference C++ implementation's multi-part `*-NNNNN-of-MMMMM.gguf` naming). pub part: usize, } @@ -62,6 +62,11 @@ pub struct Gguf { parts: Vec, /// Per-part data-section start offset (after that part's header). data_starts: Vec, + /// The path passed to `open` (part 0 for a split GGUF), kept so callers + /// that need to identify the source file on disk (e.g. an on-disk cache + /// keyed off the file's size/mtime) don't have to thread the original + /// path alongside every `&Gguf` themselves. + source_path: String, tensors: BTreeMap, pub metadata_u32: BTreeMap, pub metadata_str: BTreeMap, @@ -162,7 +167,7 @@ fn f16_to_f32(bits: u16) -> f32 { } /// Unpack the j-th 6-bit (scale, min) pair from a Q4_K/Q5_K block's packed -/// 12-byte `scales` array. This is the canonical GGML `get_scale_min_k4`: +/// 12-byte `scales` array. This mirrors the reference C++ implementation's canonical `get_scale_min_k4`: /// the 8 sub-block scales/mins are stored 6 bits each, split across the first /// 8 bytes (low 6 bits) and the last 4 bytes (high 2 bits + extra 4-bit field). #[inline] @@ -287,7 +292,7 @@ fn parse_header(bytes: &[u8], path: &str) -> Result { }) } -/// If `path` matches the llama.cpp split pattern `*-NNNNN-of-MMMMM.gguf`, +/// If `path` matches the reference C++ implementation's split pattern `*-NNNNN-of-MMMMM.gguf`, /// return the list of all part paths in order (1..=MMMMM). Otherwise return /// `None` (single-file GGUF). The given `path` may be any one of the parts. fn split_parts(path: &str) -> Option> { @@ -354,6 +359,7 @@ impl Gguf { Ok(Gguf { parts, data_starts, + source_path: path.to_string(), tensors, metadata_u32: m.metadata_u32, metadata_str: m.metadata_str, @@ -364,6 +370,12 @@ impl Gguf { }) } + /// The path this `Gguf` was opened from (part 0's path for a split GGUF, + /// see [`Self::open`]). + pub fn path(&self) -> &str { + &self.source_path + } + /// Read a `u32`/int scalar from metadata under `key`. pub fn meta_u32(&self, key: &str) -> Option { self.metadata_u32.get(key).copied() @@ -601,7 +613,7 @@ impl Gguf { /// codes packed 4-per-u32 (8 u32/block) at `qs[r*(k/32)*8 + b*8 + i/4]`. /// /// For a Q8_0 tensor this is a *lossless* repack of the on-disk blocks (the - /// f16 block scale is widened to f32, the 32 int8 are re-packed) — no second + /// f16 block scale is widened to f32, the 32 int8 are re-packed): no second /// quantization. F16/F32 tensors are quantized per-32-block via amax/127, the /// identical scheme to `ffai_ops::quantize_q8`. `k` (the in-dim, fastest GGUF /// dim) must be a multiple of 32. @@ -623,7 +635,7 @@ impl Gguf { GgmlType::Q8_0 => { // On-disk block of 32: f16 scale (2 bytes) + 32 int8. The block // order is row-major over [out, in], so block (r,b) is at linear - // block index r*bpr + b — exactly the gemv_q8 ordering. Lossless. + // block index r*bpr + b, exactly the gemv_q8 ordering. Lossless. let nblocks = m * bpr; let b = self.raw(t, nblocks * 34); for blk in 0..nblocks { diff --git a/rust/crates/ffai-models/Cargo.toml b/rust/crates/ffai-models/Cargo.toml index ef80d790..baca3f13 100644 --- a/rust/crates/ffai-models/Cargo.toml +++ b/rust/crates/ffai-models/Cargo.toml @@ -11,3 +11,4 @@ ffai-core.workspace = true ffai-ops.workspace = true ffai-loader.workspace = true serde_json.workspace = true +rayon.workspace = true diff --git a/rust/crates/ffai-models/src/laguna.rs b/rust/crates/ffai-models/src/laguna.rs index 4e4288ad..73455578 100644 --- a/rust/crates/ffai-models/src/laguna.rs +++ b/rust/crates/ffai-models/src/laguna.rs @@ -10,13 +10,14 @@ //! YaRN RoPE over the first half of each head's dims; sliding-window layers //! (window 512) use plain RoPE over the full head dim. //! -//! Decode-only: single-token forward with a persistent KV cache (growing for -//! full-attention layers, a fixed-512 ring buffer for sliding layers). No -//! prefill / batched-token path here. +//! Decode: single-token forward with a persistent KV cache (growing for +//! full-attention layers, a fixed-512 ring buffer for sliding layers). +//! Prefill: chunked batched forward (see [`prefill`]) with a linear sliding +//! window scratch compacted into the decode ring afterward. //! //! Weights load straight from GGUF (see [`ffai_loader::gguf::Gguf`]): //! dequantize each tensor to f32 once, then requantize into this engine's Q4 -//! format ([`ffai_ops::quantize_q4`]) and drop the f32 — the GGUF k-quant +//! format ([`ffai_ops::quantize_q4`]) and drop the f32, the GGUF k-quant //! (Q4_K/Q6_K/...) never needs to match our block layout. Embeddings, the LM //! head, and all norms stay dense f32 (small relative to the expert stacks //! and precision-sensitive for routing/final logits). @@ -24,6 +25,8 @@ use ffai_core::{DType, Device, Error, Result, Tensor}; use ffai_loader::gguf::Gguf; use ffai_ops as ops; +use rayon::prelude::*; +use std::sync::atomic::{AtomicU64, Ordering}; fn u32s_to_bytes(v: &[u32]) -> Vec { v.iter().flat_map(|x| x.to_le_bytes()).collect() @@ -50,6 +53,227 @@ fn f32s_to_f16_bytes(v: &[f32]) -> Vec { v.iter().flat_map(|&f| half_bits(f).to_le_bytes()).collect() } +// ───────────────────────────────────────────────────────────────────────── +// On-disk cache of converted engine-format weight blobs. +// +// `load_gguf` is the ~7.5min-cold-load path this exists to shortcut: per +// tensor it dequantizes GGUF (Q4_K/Q6_K/F16) to f32 on one CPU thread, +// requantizes to this engine's Q4 format (`ops::quantize_q4`), optionally +// permutes/interleaves/packs for FFAI_LAGUNA_MARLIN / _MOEFUSE / _W4A8, then +// uploads. Every one of those UPLOAD-READY byte blobs is cached here, one +// file per blob, so a warm reload skips straight to `read file -> upload`. +// +// LOUD REMINDER: bump `FORMAT_VERSION` whenever ANY of that byte-shaping +// logic changes (quantize_q4's block layout, permute_q4_to_marlin, the +// moefuse interleave order, w4a8_packw's packed layout, the f16 scale +// conversion, ...). Bumping it changes every blob's cache key, so stale +// pre-bump blobs are simply never read back (treated as a miss and rebuilt) +// instead of silently reinterpreted. +const FORMAT_VERSION: u32 = 1; + +const CACHE_MAGIC: [u8; 4] = *b"FLC1"; // "FFAI Laguna Cache v1" + +// `dtype_tag` is purely descriptive (a debugging aid if someone inspects a +// cache dir by hand), the cache never interprets a blob's bytes itself, it +// only ever hands them back verbatim to `dev.upload`. +const DTYPE_U32: u8 = 0; +const DTYPE_F32: u8 = 1; +const DTYPE_F16: u8 = 2; +/// Opaque packed bytes (Marlin-tile-permuted qs, W4A8 mxfp4 packed weights / +/// scale factors, ...): anything that isn't a flat array of one of the +/// above elementary types. +const DTYPE_RAW: u8 = 3; + +fn fnv1a64(bytes: &[u8]) -> u64 { + let mut h: u64 = 0xcbf29ce484222325; + for &b in bytes { + h ^= b as u64; + h = h.wrapping_mul(0x100000001b3); + } + h +} + +/// Header: magic (4B) + FORMAT_VERSION (u32 LE) + dtype tag (1B) + payload +/// len (u64 LE), then the payload. Written atomically (tmp file + rename) so +/// a crash mid-write never leaves a corrupt file at the final path, readers +/// only ever see a complete blob or no file at all. +fn write_cache_blob(path: &std::path::Path, dtype_tag: u8, payload: &[u8]) -> std::io::Result<()> { + use std::io::Write; + let tmp = path.with_extension(format!("tmp-{}-{:?}", std::process::id(), std::thread::current().id())); + { + let mut f = std::fs::File::create(&tmp)?; + f.write_all(&CACHE_MAGIC)?; + f.write_all(&FORMAT_VERSION.to_le_bytes())?; + f.write_all(&[dtype_tag])?; + f.write_all(&(payload.len() as u64).to_le_bytes())?; + f.write_all(payload)?; + } + std::fs::rename(&tmp, path) +} + +/// Reads back a blob written by [`write_cache_blob`]. Returns `None` for any +/// validation failure (missing file, bad magic, mismatched FORMAT_VERSION, +/// truncated payload), every failure mode is treated as a plain cache miss, +/// never an error, so a corrupt or foreign-version cache dir just gets +/// silently rebuilt. +fn read_cache_blob(path: &std::path::Path) -> Option> { + let bytes = std::fs::read(path).ok()?; + if bytes.len() < 17 || bytes[0..4] != CACHE_MAGIC { + return None; + } + let fv = u32::from_le_bytes(bytes[4..8].try_into().unwrap()); + if fv != FORMAT_VERSION { + return None; + } + let len = u64::from_le_bytes(bytes[9..17].try_into().unwrap()) as usize; + let payload = &bytes[17..]; + if payload.len() != len { + return None; + } + Some(payload.to_vec()) +} + +/// On-disk cache of converted engine-format weight blobs, keyed by a +/// caller-chosen artifact key (conventionally the GGUF tensor name, possibly +/// suffixed for composite/variant artifacts) folded together with +/// [`FORMAT_VERSION`], the source GGUF's identity (size + mtime), and every +/// `FFAI_LAGUNA_*` load-shaping env flag, so a changed checkpoint, a changed +/// requant flag combo, or a bumped `FORMAT_VERSION` each land in their own +/// namespace within the same cache directory rather than colliding. +/// +/// `Send + Sync` (every field is): built once per [`load_gguf`] call and +/// shared (via `&CacheCtx`) across the rayon worker threads that build the +/// per-layer byte blobs in parallel on a cache miss. +pub struct CacheCtx { + /// `None` when the cache directory couldn't be created, every lookup + /// then unconditionally misses (rebuild every artifact, every load), so + /// a read-only or unwritable filesystem degrades to "cache disabled" + /// instead of failing the load. + dir: Option, + version_tag: u64, + hits: AtomicU64, + misses: AtomicU64, +} + +impl CacheCtx { + /// `FFAI_LAGUNA_CACHE` overrides the cache directory; otherwise it lives + /// alongside the GGUF at `.ffai-cache/`. + pub fn new(g: &Gguf) -> Self { + let path = g.path(); + let dir = match std::env::var("FFAI_LAGUNA_CACHE") { + Ok(d) if !d.is_empty() => std::path::PathBuf::from(d), + _ => std::path::PathBuf::from(format!("{path}.ffai-cache")), + }; + let dir = match std::fs::create_dir_all(&dir) { + Ok(()) => Some(dir), + Err(e) => { + eprintln!( + "Laguna cache: could not create cache dir {} ({e}), caching disabled, every load will fully rebuild", + dir.display() + ); + None + } + }; + // Source-GGUF identity: (size, mtime) rather than a content hash, + // hashing an 80GB+ checkpoint on every load would itself dominate + // load time. A changed file (new export, re-quant, ...) almost + // always changes at least one of these; a false-negative (same + // size+mtime, different content) is the same class of risk `make` + // and every mtime-based build cache accepts. + let (size, mtime) = std::fs::metadata(path) + .map(|m| { + let mtime = m + .modified() + .ok() + .and_then(|t| t.duration_since(std::time::UNIX_EPOCH).ok()) + .map(|d| d.as_secs()) + .unwrap_or(0); + (m.len(), mtime) + }) + .unwrap_or((0, 0)); + let flags = ["FFAI_LAGUNA_MICRO", "FFAI_LAGUNA_MOEFUSE", "FFAI_LAGUNA_MARLIN", "FFAI_LAGUNA_W4A8", "FFAI_LAGUNA_W4A8_ONLY"] + .iter() + .map(|k| format!("{k}={}", std::env::var(k).unwrap_or_default())) + .collect::>() + .join("|"); + let version_tag = fnv1a64(format!("v{FORMAT_VERSION}|{size}|{mtime}|{flags}").as_bytes()); + CacheCtx { dir, version_tag, hits: AtomicU64::new(0), misses: AtomicU64::new(0) } + } + + fn blob_path(&self, key: &str) -> Option { + let dir = self.dir.as_ref()?; + let mut sanitized = String::with_capacity(key.len()); + for c in key.chars() { + if c.is_ascii_alphanumeric() || matches!(c, '.' | '_' | '-') { + sanitized.push(c); + } else { + sanitized.push('_'); + } + } + let h = fnv1a64(format!("{key}\u{0}{}", self.version_tag).as_bytes()); + Some(dir.join(format!("{sanitized}-{h:016x}.bin"))) + } + + /// One cached byte blob. `build` runs (and its result gets written to + /// disk) only on a miss. + pub fn get_or_build(&self, key: &str, dtype_tag: u8, build: impl FnOnce() -> Result>) -> Result> { + if let Some(path) = self.blob_path(key) { + if let Some(bytes) = read_cache_blob(&path) { + self.hits.fetch_add(1, Ordering::Relaxed); + return Ok(bytes); + } + } + let bytes = build()?; + self.misses.fetch_add(1, Ordering::Relaxed); + if let Some(path) = self.blob_path(key) { + if let Err(e) = write_cache_blob(&path, dtype_tag, &bytes) { + eprintln!("Laguna cache: failed writing '{key}': {e}"); + } + } + Ok(bytes) + } + + /// Same contract as [`Self::get_or_build`] but for an artifact whose + /// bytes are naturally co-produced by ONE expensive build step (e.g. a + /// `Q4Dense`'s `qs` + `scales`, both products of a single + /// `quantize_q4` call, or a `Q4Dense`+`MarlinDense` pair sharing one + /// `quantize_q4` pass). Checks every key first; `build` runs at most + /// once, only if at least one key is missing: so a full hit never + /// redoes the shared work, and a full miss never redoes it either. + pub fn get_or_build_multi(&self, keys: &[(&str, u8)], build: impl FnOnce() -> Result>>) -> Result>> { + let paths: Vec> = keys.iter().map(|(k, _)| self.blob_path(k)).collect(); + let cached: Option>> = paths.iter().map(|p| p.as_ref().and_then(|p| read_cache_blob(p))).collect(); + if let Some(all) = cached { + self.hits.fetch_add(keys.len() as u64, Ordering::Relaxed); + return Ok(all); + } + let blobs = build()?; + if blobs.len() != keys.len() { + return Err(Error::Msg(format!("Laguna cache: build produced {} blobs, expected {} for keys {:?}", blobs.len(), keys.len(), keys.iter().map(|(k, _)| *k).collect::>()))); + } + self.misses.fetch_add(keys.len() as u64, Ordering::Relaxed); + for ((key, dtype), (path, bytes)) in keys.iter().zip(paths.iter().zip(blobs.iter())) { + if let Some(path) = path { + if let Err(e) = write_cache_blob(path, *dtype, bytes) { + eprintln!("Laguna cache: failed writing '{key}': {e}"); + } + } + } + Ok(blobs) + } + + pub fn stats(&self) -> (u64, u64) { + (self.hits.load(Ordering::Relaxed), self.misses.load(Ordering::Relaxed)) + } + + pub fn dir_display(&self) -> String { + match &self.dir { + Some(d) => d.display().to_string(), + None => "".to_string(), + } + } +} + /// One `[m, k]` row-major f32 weight requantized to this engine's Q4 format /// and uploaded. `m`/`k` match the `gemv_q4` convention (`k` must be a /// multiple of 32). @@ -60,15 +284,12 @@ struct Q4Dense { k: usize, } impl Q4Dense { - fn upload(dev: &dyn Device, flat: &[f32], m: usize, k: usize) -> Result { - let (qs, scales) = ops::quantize_q4(flat, m, k); - Ok(Q4Dense { - qs: Tensor::new(dev.upload(&u32s_to_bytes(&qs))?, vec![qs.len()], DType::U32), - scales: Tensor::new(dev.upload(&f32s_to_f16_bytes(&scales))?, vec![scales.len()], DType::F16), - m, - k, - }) - } + // Construction (dequant -> quantize_q4 -> upload, or the cache-hit + // shortcut straight from disk) now lives in the free functions around + // `q4dense_bytes_from_flat_cached` / `upload_q4dense_bytes` below, so + // every construction site (plain, Marlin-shared, cached, parallel-phase) + // goes through the same on-disk-cache-aware path instead of each having + // its own uncached `dev.upload` call. fn gemv(&self, dev: &dyn Device, x: &Tensor) -> Result { ops::gemv_q4(dev, &self.qs, &self.scales, x, self.m, self.k, self.m) } @@ -77,9 +298,79 @@ impl Q4Dense { } } +/// One `[n_out, k_in]` row-major f32 weight, requantized to Q4 and permuted +/// into the tensor-core Marlin tile-major layout ([`ops::permute_q4_to_marlin`]) +/// for [`ops::moe_w4a16_marlin`]. Built ONLY under `FFAI_LAGUNA_MARLIN=1`, as +/// an ADDITIONAL resident copy alongside the plain [`Q4Dense`] decode uses, +/// this struct is prefill-only (see `prefill_layer`'s marlin branches). +/// +/// `n_out` must be a multiple of 64 (`permute_q4_to_marlin`'s tile height); +/// `k_in` a multiple of 32 (Q4 block size). A dense `[n_out, k_in]` matrix is +/// fed to the (expert-batched) Marlin GEMM as a single "expert" (`n_exp=1`): +/// every row of `x` carries expert-id 0 via a persistent all-zero `indices` +/// buffer (see `prefill`'s `marlin_zero_idx`). +struct MarlinDense { + qs: Tensor, // Marlin tile-major Q4: [(n_out/64)*(k_in/32)*256] u32 + scales: Tensor, // f16, STANDARD (non-permuted) per-block layout: [n_out*(k_in/32)]. + // `moe_w4a16_marlin`'s kernel doc is explicit that scales are read in the + // SAME layout `quantize_q4` produces; only `qs` gets tile-permuted. + n_out: usize, + k_in: usize, +} +impl MarlinDense { + // Construction lives in `marlindense_bytes_from_flat_cached` / + // `upload_marlindense_bytes` below (same rationale as `Q4Dense`'s impl + // block doc above). + /// `x` is `[t, k_in]` f16; `zero_idx` is a persistent all-zero `u32` + /// buffer with at least `t` elements (the 1-expert "every row is expert + /// 0" trick). Returns `[t, n_out]` f16, cast back with `ops::cast_f16_f32` + /// wherever the caller needs f32 (see `prefill_layer`). + fn call(&self, dev: &dyn Device, x_f16: &Tensor, zero_idx: &Tensor, t: usize) -> Result { + ops::moe_w4a16_marlin(dev, x_f16, &self.qs, &self.scales, zero_idx, t, self.n_out, self.k_in) + } +} + +/// FFAI_LAGUNA_MARLIN=1: route dense prefill projections (QKV concat, o_proj, +/// layer-0 dense FFN, shared-expert FFN) through the tensor-core +/// [`MarlinDense`] GEMM instead of `gemm_q4_mpp`. Read once at load (to build +/// the extra Marlin-layout weight copies below) and once per `prefill_layer` +/// call (mirrors `laguna_micro_enabled`'s load/call split). Decode is +/// entirely unaffected, it never reads this env. +fn laguna_marlin_enabled() -> bool { + std::env::var("FFAI_LAGUNA_MARLIN").map(|v| v == "1").unwrap_or(false) +} + +/// Marlin-packed `[wq;wk;wv]` concat (gate excluded, its row count, 48 or 72, +/// is never 64-aligned, so it stays on `gemm_q4_mpp` via `AttnWeights::g_proj`). +/// `n_q_rows + 2*n_kv_rows` is 64-aligned for both Laguna-S-2.1 layer types +/// (full: 6144+2*1024=8192, sliding: 9216+2*1024=11264), verified at build +/// time rather than assumed, see `load_gguf`. +struct QkvMarlin { + mat: MarlinDense, // [n_q_rows + 2*n_kv_rows, hidden] + n_q_rows: usize, + n_kv_rows: usize, +} + +/// Marlin-packed `[gate;up]` concat (both halves share `n_ff` rows and +/// `hidden` columns), used for the dense layer-0 FFN and the always-on +/// shared expert, both of whose `2*n_ff` is 64-aligned for Laguna-S-2.1 +/// (24576 and 2048 respectively). +struct GateUpMarlin { + mat: MarlinDense, // [2*n_ff, hidden] + n_ff: usize, +} + +/// The always-on shared expert's Marlin copies (gate/up concat + down), +/// built regardless of whether decode's [`SharedExpert`] resolved to +/// `Split` or `Fused`, this struct reads the shexp GGUF tensors directly. +struct SharedMarlin { + gateup: GateUpMarlin, + down: MarlinDense, // [hidden, n_ff] +} + /// `n_expert` stacked `[m, k]` Q4 weights, requantized expert-by-expert and /// concatenated into one device buffer per projection. A per-expert view is -/// a byte-offset slice — no host round-trip to select an expert's weights, +/// a byte-offset slice, no host round-trip to select an expert's weights, /// only the (tiny) expert-id lookup needs to be on the host. struct Q4ExpertStack { qs: Tensor, @@ -90,39 +381,11 @@ struct Q4ExpertStack { scale_elems_per_expert: usize, } impl Q4ExpertStack { - /// `flat` is the dequantized fused GGUF tensor, expert-major contiguous - /// (`flat[e*m*k .. (e+1)*m*k]` is expert `e`'s `[m, k]` row-major slab — - /// true for both `ffn_{gate,up}_exps` `[n_embd, n_ff_exp, n_expert]` and - /// `ffn_down_exps` `[n_ff_exp, n_embd, n_expert]` given GGUF's - /// fastest-dim-first tensor layout). - fn upload(dev: &dyn Device, flat: &[f32], n_expert: usize, m: usize, k: usize) -> Result { - let per_expert = m * k; - if flat.len() != per_expert * n_expert { - return Err(Error::Msg(format!( - "Q4ExpertStack: flat len {} != n_expert*m*k ({} * {} * {})", - flat.len(), n_expert, m, k - ))); - } - let qw = per_expert / 32 * 4; // quantize_q4: m*k/32 blocks, 4 u32/block - let sw = per_expert / 32; // 1 f32 scale/block - let mut qs_all = Vec::with_capacity(qw * n_expert); - let mut sc_all = Vec::with_capacity(sw * n_expert); - for e in 0..n_expert { - let slab = &flat[e * per_expert..(e + 1) * per_expert]; - let (qs, sc) = ops::quantize_q4(slab, m, k); - qs_all.extend_from_slice(&qs); - sc_all.extend_from_slice(&sc); - } - Ok(Q4ExpertStack { - qs: Tensor::new(dev.upload(&u32s_to_bytes(&qs_all))?, vec![qs_all.len()], DType::U32), - scales: Tensor::new(dev.upload(&f32s_to_f16_bytes(&sc_all))?, vec![sc_all.len()], DType::F16), - m, - k, - qs_words_per_expert: qw, - scale_elems_per_expert: sw, - }) - } - /// Device-side view of expert `e`'s `(qs, scales)` — a byte-offset slice + // Construction lives in `q4_expert_stack_bytes_cached` / + // `upload_expert_stack_bytes` below (same rationale as `Q4Dense`'s impl + // block doc above); the per-expert `quantize_q4` loop that used to live + // here is now inside that cached builder's closure. + /// Device-side view of expert `e`'s `(qs, scales)`, a byte-offset slice /// into the shared stack, no copy. fn expert_view(&self, e: usize) -> (Tensor, Tensor) { let qw = self.qs_words_per_expert; @@ -151,6 +414,92 @@ impl Q4ExpertStack { } } +/// FFAI_LAGUNA_W4A8=1: one `[n_expert, n, k]` dequantized f32 expert-major +/// slab packed into the mxfp4 (e2m1 nibbles + per-32-block ue8m0 scale +/// factor) layout `ops::moe_grouped_gemm_w4a8` consumes as its weight +/// operand. Mirrors the NEMOTRON_W4A8_MOE loader's `buildw` closure +/// (`ffai-modeltests/src/lib.rs`, around `NEMOTRON_W4A8_MOE`) exactly: f32 → +/// f16 upload, zero-pad `k` up to the mxf8f6f4 SF atom width (128) via +/// `ops::pad_rows_f16` (a no-op clone when `k` is already 128-aligned, true +/// for both Laguna-S-2.1 K dims, `hidden=3072` and `n_ff_exp=1024`, so this +/// is a documented no-op on the shipped checkpoint, not dead code for a +/// hypothetical one), then `ops::w4a8_packw`. Neither this struct nor the +/// nemotron recipe it mirrors applies any outlier clipping or per-channel +/// rescale before packing, the packer's own per-32 ue8m0 block scale is the +/// only quantization-error mitigation on the weight side (mxfp4 is a +/// coarser format than the engine's own Q4: e2m1 has only 16 representable +/// magnitudes per block vs Q4's calibrated int4 codebook, so this is +/// expected to be lossier, the box's greedy-continuation gate decides +/// acceptance, per this file's module doc). +struct W4a8ExpertStack { + packed: Tensor, // e2m1 nibble-packed, `ops::w4a8_packw`'s `packed` output + sf: Tensor, // per-32-block ue8m0 scale factors, atom-padded SFB layout + /// Per-expert SFB stride in bytes (`ops::w4a8_packw`'s `sfb_exp_elems` + /// return, fed straight through as `moe_grouped_gemm_w4a8`'s + /// `sfb_exp_bytes`, see that fn's doc, the name is elem-count but the + /// GEMM wants it as a byte stride, matching the nemotron call site). + sfb_exp_bytes: i64, + /// Original (unpadded) input width, `hidden` for gate/up (incl. the + /// fused `2*n_ff_exp`-row stack), `n_ff_exp` for down. Not read at + /// prefill time (only `kpad` drives the padding/quant calls below), + /// kept for debugging/asserts, mirroring `Q4Dense`/`Q4ExpertStack`'s own + /// `k` fields. + #[allow(dead_code)] + k: usize, + /// `k` rounded up to the mxf8f6f4 SF atom width (128), the width + /// actually packed and what every `w4a8_sfa_offsets`/ + /// `w4a8_actquant_grouped` call at prefill time must pad activations to + /// before quantizing against this stack. + kpad: usize, +} + +/// FFAI_LAGUNA_W4A8=1: pack a dequantized `[n_expert, n, k]` f32 expert-major +/// slab (same `flat` a caller would otherwise only feed to `Q4ExpertStack`'s +/// builder, see the call sites in `build_moe_w4a8_sequential`, which +/// dequantize ONCE and hand the same slab to both) into a [`W4a8ExpertStack`], +/// threading the packed output through `cache`. See that struct's doc for the +/// exact recipe (f32→f16, pad-to-128, `ops::w4a8_packw`) and why it applies +/// no extra outlier handling. +/// +/// UNLIKE every other cached artifact in this file, the pack itself +/// (`ops::pad_rows_f16` + `ops::w4a8_packw`) runs ON THE DEVICE, not on the +/// host, it needs `dev`, so it can never be moved into the rayon-parallel +/// per-layer CPU pass (`build_layer_bytes`); it always runs sequentially, on +/// whichever thread calls it (see `build_moe_w4a8_sequential`'s doc). On a +/// cache MISS this does exactly what the pre-cache `pack_w4a8_expert_stack` +/// did (upload raw f16, pad, pack) plus a download-back-to-host of the +/// result to write the cache blob. On a HIT, all of that: including the +/// f16 upload and both GPU kernels, is skipped: the packed bytes are +/// re-uploaded straight from disk. +fn pack_w4a8_expert_stack_cached(dev: &dyn Device, cache: &CacheCtx, key: &str, flat: &[f32], n_expert: usize, n: usize, k: usize) -> Result { + let kpad = k.div_ceil(128) * 128; + let packed_key = format!("{key}.w4a8.packed"); + let sf_key = format!("{key}.w4a8.sf"); + let meta_key = format!("{key}.w4a8.meta"); + let blobs = cache.get_or_build_multi( + &[(packed_key.as_str(), DTYPE_RAW), (sf_key.as_str(), DTYPE_RAW), (meta_key.as_str(), DTYPE_RAW)], + || -> Result>> { + let wt = Tensor::new(dev.upload(&f32s_to_f16_bytes(flat))?, vec![n_expert * n, k], DType::F16); + let wt_pad = ops::pad_rows_f16(dev, &wt, n_expert * n, k, kpad)?; + let (packed, sf, sfb_exp_bytes) = ops::w4a8_packw(dev, &wt_pad, n_expert, n, kpad)?; + dev.synchronize()?; + let mut packed_bytes = vec![0u8; packed.elem_count() * 4]; + dev.download(packed.buffer.as_ref(), &mut packed_bytes)?; + let mut sf_bytes = vec![0u8; sf.elem_count() * 4]; + dev.download(sf.buffer.as_ref(), &mut sf_bytes)?; + Ok(vec![packed_bytes, sf_bytes, sfb_exp_bytes.to_le_bytes().to_vec()]) + }, + )?; + let mut it = blobs.into_iter(); + let packed_bytes = it.next().unwrap(); + let sf_bytes = it.next().unwrap(); + let meta_bytes = it.next().unwrap(); + let packed = Tensor::new(dev.upload(&packed_bytes)?, vec![packed_bytes.len() / 4], DType::U32); + let sf = Tensor::new(dev.upload(&sf_bytes)?, vec![sf_bytes.len() / 4], DType::U32); + let sfb_exp_bytes = i64::from_le_bytes(meta_bytes[..8].try_into().map_err(|_| Error::Msg("Laguna cache: truncated w4a8 meta blob".into()))?); + Ok(W4a8ExpertStack { packed, sf, sfb_exp_bytes, k, kpad }) +} + /// YaRN RoPE parameters for one layer-type (full-attention or sliding). #[derive(Debug, Clone, Copy)] pub struct RopeParams { @@ -167,7 +516,7 @@ pub struct RopeParams { /// Architecture config for Laguna-S-2.1, read from GGUF `laguna.*` metadata /// (arch spec fields fall back to the model's known-fixed values when a key /// is absent, e.g. weightless test fixtures or a schema this loader hasn't -/// seen yet — always cross-check against a real GGUF's `laguna.*` keys when +/// seen yet, always cross-check against a real GGUF's `laguna.*` keys when /// wiring a new checkpoint). #[derive(Debug, Clone)] pub struct LagunaConfig { @@ -261,7 +610,7 @@ impl LagunaConfig { } /// Read `laguna.*` GGUF metadata, falling back to [`Self::defaults`] for - /// any key this loader doesn't find (see the module doc — verify key + /// any key this loader doesn't find (see the module doc, verify key /// names against a real GGUF before trusting a mismatch silently). pub fn from_gguf(g: &Gguf) -> Self { let d = Self::defaults(); @@ -328,19 +677,52 @@ impl LagunaConfig { /// FFAI_LAGUNA_MICRO=1 enables three decode-path micro-fusions (QKVG concat /// gemv, shared-expert fused swiglu gather, o_proj residual accum). The env -/// is read once at load time (to decide which weight layout to build below — +/// is read once at load time (to decide which weight layout to build below, /// [`AttnProj`] / [`SharedExpert`]) and once per [`decode_step`] (to decide /// the runtime code path). Both reads see the same value: the env is fixed -/// for the process lifetime (required for CUDA-graph capture anyway — see +/// for the process lifetime (required for CUDA-graph capture anyway, see /// [`GraphBufs`]), so load-time and decode-time never disagree. fn laguna_micro_enabled() -> bool { std::env::var("FFAI_LAGUNA_MICRO").map(|v| v == "1").unwrap_or(false) } +/// FFAI_LAGUNA_MOEFUSE=1 fuses the routed MoE gate and up expert stacks +/// (`ffn_gate_exps`, `ffn_up_exps`) into ONE per-layer Q4 stack at load time, +/// so `prefill_layer`'s two `ops::moe_q4_grouped_mma_dev` calls over the gate +/// stack and the up stack (same sorted activation `a`, same `offsets`) +/// collapse into a single call with `n_out = 2*n_ff_exp`. See +/// [`RoutedGateUp::Fused`] for the exact per-expert row layout and the +/// decode-side tradeoff this implies. Read once at load time (to decide which +/// weight layout to build, same one-layout-resident rule as +/// [`laguna_micro_enabled`]); `decode_layer`/`prefill_layer` never re-read the +/// env, they just match on the [`RoutedGateUp`] variant `load_gguf` already +/// picked. +fn laguna_moefuse_enabled() -> bool { + std::env::var("FFAI_LAGUNA_MOEFUSE").map(|v| v == "1").unwrap_or(false) +} + +/// FFAI_LAGUNA_W4A8=1: route prefill's routed-MoE grouped GEMMs (gate/up, +/// down, the `moe_gateup_gemm`/`moe_down_gemm` stages that dominate the pp +/// profile, ~73% combined) through the W4A8 grouped-MoE GEMM (fp8 e4m3 +/// activations x mxfp4 weights, cutlass mxf8f6f4) instead of the W4A16 +/// `ops::moe_q4_grouped_mma_dev` path, mirroring the NEMOTRON_W4A8_MOE +/// recipe (`ffai-modeltests`'s `NEMOTRON_W4A8_MOE` build closure, and its +/// `w4a8_moe` prefill branch) exactly. Read once at load time (to build the +/// extra [`W4a8ExpertStack`]-packed weight copies below, ADDITIONAL to the +/// resident Q4 stacks decode still needs, see [`MoeFfnW4a8`]'s doc for the +/// resident-memory cost) and once per `prefill_layer` call is NOT needed: +/// the load-time-built `Option` on `MoeFfn` already encodes +/// whether this layer has a W4A8 path, so `prefill_layer` just matches on +/// it (same convention as `shared_marlin`/`qkv_marlin`). Decode is entirely +/// unaffected, it never reads this env or the `w4a8` field. +fn laguna_w4a8_enabled() -> bool { + std::env::var("FFAI_LAGUNA_W4A8").map(|v| v == "1").unwrap_or(false) +} + /// QKV(+gate) projection weights. `Fused` concatenates all four projections' /// rows into one `Q4Dense` at load time (built ONLY under /// `FFAI_LAGUNA_MICRO=1`, so a layer never carries both layouts resident at -/// once — that would double the ~2.4GB of attention-projection weights). +/// once, that would double the ~2.4GB of attention-projection weights). /// Row order in `wqkvg` is `[wq rows; wk rows; wv rows; wgate rows]`, all /// sharing `k = hidden`; concatenation happens on the dequantized f32 slabs /// BEFORE `quantize_q4`, so each row's Q4 block boundaries are unaffected @@ -355,7 +737,7 @@ enum AttnProj { struct AttnWeights { /// This layer's Q head count, derived from the loaded Q projection's /// actual row count (`n_q_rows / head_dim`) rather than `LagunaConfig`'s - /// hardcoded 48/72 split — so a real GGUF's per-layer head count always + /// hardcoded 48/72 split, so a real GGUF's per-layer head count always /// wins over the spec constants if the two ever disagree. n_head: usize, attn_norm: Tensor, // [hidden] f32 @@ -363,6 +745,15 @@ struct AttnWeights { wo: Q4Dense, // [hidden, n_head*head_dim] q_norm: Tensor, // [head_dim] f32 k_norm: Tensor, // [head_dim] f32 + /// FFAI_LAGUNA_MARLIN=1 prefill-only extras (see [`QkvMarlin`]/ + /// [`MarlinDense`]'s docs); `None` when the env is unset OR (defensively) + /// when a dimension isn't 64-row-aligned. `g_proj` is a plain standalone + /// `Q4Dense` (NOT Marlin, gate's 48/72 rows are never 64-aligned) built + /// only under the same env, so the marlin prefill branch has gate access + /// independent of whether `proj` above resolved to `Split` or `Fused`. + qkv_marlin: Option, + g_proj: Option, + wo_marlin: Option, } /// Layer-0's dense SwiGLU MLP. @@ -371,6 +762,10 @@ struct DenseFfn { gate: Q4Dense, // [n_ff_dense, hidden] up: Q4Dense, // [n_ff_dense, hidden] down: Q4Dense, // [hidden, n_ff_dense] + /// FFAI_LAGUNA_MARLIN=1 prefill-only extras, see [`GateUpMarlin`]/ + /// [`MarlinDense`]'s docs. `None` when the env is unset. + gateup_marlin: Option, + down_marlin: Option, } /// The always-on shared expert's gate+up projections. `Fused` repacks both @@ -380,12 +775,12 @@ struct DenseFfn { /// buffer, so decode can drive both through the same /// `ops::moe_gather_q4_swiglu` the routed experts use (`top_k=1`) instead of /// two plain gemvs + a host-orchestrated swiglu. `down` stays a plain -/// `Q4Dense` either way — only gate+up fuse into the gather. +/// `Q4Dense` either way, only gate+up fuse into the gather. struct SharedExpertFused { gate: Q4ExpertStack, // n_expert=1, [n_ff_shexp, hidden] up: Q4ExpertStack, // n_expert=1, [n_ff_shexp, hidden] down: Q4Dense, // [hidden, n_ff_shexp] - /// Persistent device buffer holding the constant `u32` `0` — the single + /// Persistent device buffer holding the constant `u32` `0`, the single /// "expert index" `moe_gather_q4_swiglu` needs for a 1-expert stack. zero_idx: Tensor, } @@ -394,15 +789,83 @@ enum SharedExpert { Fused(SharedExpertFused), } +/// Routed-expert gate+up weight layout. `Split` is the plain per-projection +/// layout ([`Q4ExpertStack`] each). `Fused` is built ONLY under +/// `FFAI_LAGUNA_MOEFUSE=1` and REPLACES `Split` (see +/// [`laguna_moefuse_enabled`]'s doc, same one-layout-resident rule as +/// [`AttnProj`]/[`SharedExpert`]): expert `e`'s slab is `[2*n_ff_exp, hidden]` +/// with gate rows at local offset `[0, n_ff_exp)` and up rows at local offset +/// `[n_ff_exp, 2*n_ff_exp)`, i.e. per-expert interleaved (gate_0, up_0, +/// gate_1, up_1, ...), NOT all gate experts followed by all up experts. This +/// is the SAME single stack that `prefill_layer` and `decode_layer` both +/// read, so fusing gate+up does not add a second resident copy of the expert +/// weights. +/// +/// Decode tradeoff: the unfused layout lets decode drive gate+up through one +/// `ops::moe_gather_q4_swiglu` call under `FFAI_LAGUNA_FUSE=1` (fused inline +/// SwiGLU, reads the activation once). The fused stack's per-expert row +/// layout does not fit that kernel's `row = e*inter + local` indexing for two +/// SEPARATE logical projections sharing one physical stack, so decode falls +/// back to one `ops::moe_gather_q4` (inter=2*n_ff_exp) plus a host-orchestrated +/// column split and `ops::swiglu` instead (see `decode_layer`'s MoE branch). +/// That is roughly the pre-FFAI_LAGUNA_FUSE decode path (about 1 tok/s slower +/// than the fused gather), traded for prefill's halved grouped-GEMM call +/// count. Default OFF. +enum RoutedGateUp { + Split { gate: Q4ExpertStack, up: Q4ExpertStack }, + Fused(Q4ExpertStack), // per-expert [2*n_ff_exp, hidden] +} + +/// FFAI_LAGUNA_W4A8=1 mirror of [`RoutedGateUp`], packed as +/// [`W4a8ExpertStack`]s instead of [`Q4ExpertStack`]s. Which variant is +/// built follows the SAME `FFAI_LAGUNA_MOEFUSE` choice `RoutedGateUp` made +/// for this layer (`load_gguf` builds both from the identical dequantized +/// f32 slab, so `Fused` here always pairs with `RoutedGateUp::Fused` and +/// vice versa), `prefill_layer` never needs to check the two enums agree, +/// they're constructed together. +enum RoutedGateUpW4a8 { + Split { gate: W4a8ExpertStack, up: W4a8ExpertStack }, + Fused(W4a8ExpertStack), // per-expert [2*n_ff_exp, hidden] +} + +/// FFAI_LAGUNA_W4A8=1 prefill-only routed-expert extras: mxfp4-packed +/// [`RoutedGateUpW4a8`] + down [`W4a8ExpertStack`], built ADDITIONALLY to +/// (never replacing) `MoeFfn`'s existing `gateup`/`down` [`Q4ExpertStack`]s +///, decode always drives the Q4 stacks, this is prefill-only, mirroring +/// every other `FFAI_LAGUNA_*` weight-layout env in this file. Unlike those, +/// this ADDS roughly 30GB of resident memory across all MoE layers (the +/// routed expert stacks are the bulk of the model, Q4 stays resident for +/// decode, so this genuinely doubles the routed-expert footprint rather +/// than trading one layout for another); acceptable for a prefill-throughput +/// win on hardware with the headroom, but unlike `FFAI_LAGUNA_MARLIN` +/// (which is roughly memory-neutral, permuting the SAME quantized weight) +/// this is NOT a "just flip it on" env on a memory-constrained box. +struct MoeFfnW4a8 { + gateup: RoutedGateUpW4a8, + down: W4a8ExpertStack, // per-expert [hidden, n_ff_exp] +} + /// A routed MoE FFN block (layers `n_layer_dense_lead..`). struct MoeFfn { ffn_norm: Tensor, // [hidden] f32 - router: Tensor, // [n_expert, hidden] f32 — kept dense for routing precision - bias: Tensor, // [n_expert] f32 — e_score_correction_bias (selection only) - gate: Q4ExpertStack, // per-expert [n_ff_exp, hidden] - up: Q4ExpertStack, // per-expert [n_ff_exp, hidden] + router: Tensor, // [n_expert, hidden] f32: kept dense for routing precision + bias: Tensor, // [n_expert] f32: e_score_correction_bias (selection only) + gateup: RoutedGateUp, down: Q4ExpertStack, // per-expert [hidden, n_ff_exp] shared_expert: SharedExpert, + /// FFAI_LAGUNA_MARLIN=1 prefill-only shared-expert extras, see + /// [`SharedMarlin`]'s doc. `None` when the env is unset. The routed + /// per-expert `gate`/`up`/`down` stacks above stay on the existing + /// grouped-MMA prefill path (`ops::moe_q4_grouped_mma_dev`), that's + /// already a tensor-core grouped GEMM, not the GEMV-class kernel this + /// Marlin work targets. + shared_marlin: Option, + /// FFAI_LAGUNA_W4A8=1 prefill-only routed-expert extras, see + /// [`MoeFfnW4a8`]'s doc. `None` when the env is unset; when `Some`, + /// `prefill_layer` routes the routed gate/up/down grouped GEMMs through + /// `ops::moe_grouped_gemm_w4a8` instead of `ops::moe_q4_grouped_mma_dev` + /// (the shared expert and everything else in the layer is untouched). + w4a8: Option, } enum LayerFfn { @@ -441,7 +904,7 @@ pub struct LagunaModel { /// Persistent decode-time KV cache: one entry per layer, growing (`cap = /// max_seq`) for full-attention layers, a fixed 512-slot ring for sliding /// layers. RoPE is applied (with the token's absolute position) before the -/// K row is cached, so ring order doesn't matter for attention — softmax +/// K row is cached, so ring order doesn't matter for attention: softmax /// over the cached K/V set is order-invariant once every slot is filled. struct KvSlot { k: Tensor, // [n_kv_heads, cap, head_dim] @@ -465,9 +928,59 @@ impl LagunaKvCache { } } +/// Per-sliding-layer LINEAR K/V scratch used ONLY during [`prefill`]. The +/// existing sliding-layer decode cache (a `KvSlot` inside [`LagunaKvCache`]) +/// is a fixed `sliding_window`-slot RING, safe for single-token decode +/// (write one slot, evict the oldest), but NOT safe to batch-write during +/// prefill: a `kv_append_many` writing a whole T-token chunk into the ring +/// would evict entries that EARLIER queries in that SAME chunk still need to +/// attend to (their window hasn't slid that far yet). Prefill instead writes +/// every sliding-layer token into a LINEAR buffer sized to the whole prompt +/// (`[n_kv_heads, prompt_cap, head_dim]`, absolute positions, no +/// wraparound), attends with `seg_lo[i] = max(0, pos_i - (sliding_window - +/// 1))` via `ops::sdpa_multi_tc_varlen`, and once the whole prompt has been +/// processed, [`prefill`] compacts the last `sliding_window` positions into +/// the decode ring (`ops::kv_compact_ring`) so `decode_step` / +/// `LagunaDecodeCtx` can continue from `pos = tokens.len()` exactly as if +/// every prompt token had gone through the per-token decode path. +/// +/// Full-attention layers need no scratch, they reuse `LagunaKvCache`'s +/// already-growing per-layer buffer directly (a batched write there is safe: +/// nothing is ever evicted). +/// +/// Memory: `n_kv_heads * prompt_cap * head_dim * 4 bytes * 2 (K+V)` per +/// sliding layer. At `prompt_cap = 32768` (this v1's documented supported +/// prompt-length ceiling) that is `8 * 32768 * 128 * 4 * 2 ≈ 256MiB` per +/// sliding layer; Laguna-S-2.1 has 36 sliding layers (48 total, minus 12 +/// full-attention layers every 4th), so a full-size scratch is `≈9.7GB`. +/// Allocate ONE scratch sized to the largest prompt the caller expects to +/// prefill, not per call. +pub struct LagunaPrefillScratch { + /// Indexed by layer; `None` for full-attention layers (`i % swa_period == 0`). + slots: Vec>, + prompt_cap: usize, +} + +impl LagunaPrefillScratch { + pub fn new(dev: &dyn Device, cfg: &LagunaConfig, prompt_cap: usize) -> Result { + let mut slots = Vec::with_capacity(cfg.n_layers); + for i in 0..cfg.n_layers { + if cfg.is_full_layer(i) { + slots.push(None); + continue; + } + let n = cfg.n_kv_heads * prompt_cap * cfg.head_dim; + let k = Tensor::new(dev.alloc_zeroed(n * 4)?, vec![cfg.n_kv_heads, prompt_cap, cfg.head_dim], DType::F32); + let v = Tensor::new(dev.alloc_zeroed(n * 4)?, vec![cfg.n_kv_heads, prompt_cap, cfg.head_dim], DType::F32); + slots.push(Some(KvSlot { k, v, cap: prompt_cap })); + } + Ok(LagunaPrefillScratch { slots, prompt_cap }) + } +} + fn dims_out_in(t: &ffai_loader::gguf::GgufTensor, name: &str) -> Result<(usize, usize)> { // GGUF dims are fastest-first: a 2-D `[out, in]` weight is stored as - // `[in, out]` (`ggml`'s `create_tensor({k, m}, ...)` convention), which + // `[in, out]` (the reference C++ implementation's `create_tensor({k, m}, ...)` convention), which // is exactly the `(m, k)` `gemv_q4` wants with zero transposition. if t.dims.len() != 2 { return Err(Error::Msg(format!("{name}: expected 2-D tensor, got dims {:?}", t.dims))); @@ -475,162 +988,767 @@ fn dims_out_in(t: &ffai_loader::gguf::GgufTensor, name: &str) -> Result<(usize, Ok((t.dims[1] as usize, t.dims[0] as usize)) } -fn load_dense(dev: &dyn Device, g: &Gguf, name: &str) -> Result { +/// Read `name`'s `(m, k)` from its GGUF header dims alone, cheap (metadata +/// only, no mmap data touched), so this is always safe to call eagerly, even +/// when the tensor's actual bytes end up not being dequantized at all (a +/// full cache hit for that tensor's requantized artifact). +fn tensor_dims(g: &Gguf, name: &str) -> Result<(usize, usize)> { let t = g.tensor(name).ok_or_else(|| Error::Msg(format!("gguf: tensor '{name}' missing")))?; - let (m, k) = dims_out_in(t, name)?; - let flat = g.dequant_f32(name)?; - Q4Dense::upload(dev, &flat, m, k) + dims_out_in(t, name) } -/// Like [`load_dense`] but returns the dequantized f32 slab (not yet -/// requantized/uploaded) plus its `(m, k)` — used by the QKVG-concat fusion, -/// which needs to concatenate rows across four tensors BEFORE `quantize_q4`. -fn load_dense_raw(g: &Gguf, name: &str) -> Result<(Vec, usize, usize)> { - let t = g.tensor(name).ok_or_else(|| Error::Msg(format!("gguf: tensor '{name}' missing")))?; - let (m, k) = dims_out_in(t, name)?; - let flat = g.dequant_f32(name)?; - Ok((flat, m, k)) +// ── Bytes-mirror types ────────────────────────────────────────────────── +// +// Host-side (no `Tensor`, no `dev`) mirrors of the device-resident weight +// structs above, holding UPLOAD-READY byte blobs instead of `Tensor`s. This +// is the payload of `build_layer_bytes` (the rayon-parallel, cache-checked, +// pure-CPU half of loading one layer), `upload_layer` below turns one of +// these into the real `LagunaLayer` by doing nothing but `dev.upload` + +// `Tensor::new` calls, sequentially. + +struct Q4DenseBytes { qs: Vec, scales: Vec, m: usize, k: usize } +struct MarlinDenseBytes { qs: Vec, scales: Vec, n_out: usize, k_in: usize } +struct Q4ExpertStackBytes { qs: Vec, scales: Vec, m: usize, k: usize } + +enum AttnProjBytes { + Fused { wqkvg: Q4DenseBytes, n_q_rows: usize, n_kv_rows: usize, n_g_rows: usize }, + Split { wq: Q4DenseBytes, wk: Q4DenseBytes, wv: Q4DenseBytes, wgate: Q4DenseBytes }, } +struct QkvMarlinBytes { mat: MarlinDenseBytes, n_q_rows: usize, n_kv_rows: usize } +struct GateUpMarlinBytes { mat: MarlinDenseBytes, n_ff: usize } +struct SharedMarlinBytes { gateup: GateUpMarlinBytes, down: MarlinDenseBytes } -fn load_vec_f32(dev: &dyn Device, g: &Gguf, name: &str) -> Result { - let flat = g.dequant_f32(name)?; - let n = flat.len(); - Ok(Tensor::new(dev.upload(&f32s_to_bytes(&flat))?, vec![n], DType::F32)) +struct AttnWeightsBytes { + n_head: usize, + attn_norm: Vec, + proj: AttnProjBytes, + wo: Q4DenseBytes, + q_norm: Vec, + k_norm: Vec, + qkv_marlin: Option, + g_proj: Option, + wo_marlin: Option, } -fn load_expert_stack(dev: &dyn Device, g: &Gguf, name: &str, n_expert: usize, m: usize, k: usize) -> Result { - let flat = g.dequant_f32(name)?; - Q4ExpertStack::upload(dev, &flat, n_expert, m, k) +struct DenseFfnBytes { + ffn_norm: Vec, + gate: Q4DenseBytes, + up: Q4DenseBytes, + down: Q4DenseBytes, + gateup_marlin: Option, + down_marlin: Option, } -/// Load Laguna-S-2.1 straight from a GGUF file: parse `laguna.*` metadata for -/// the config, then stream every tensor through `dequant_f32` → requantize -/// (Q4 for projections/experts, dense f32 for embeddings/norms/router) → -/// upload, dropping the f32 immediately after each tensor. -pub fn load_gguf(dev: &dyn Device, g: &Gguf) -> Result { - let cfg = LagunaConfig::from_gguf(g); +struct SharedExpertFusedBytes { gate: Q4ExpertStackBytes, up: Q4ExpertStackBytes, down: Q4DenseBytes } +enum SharedExpertBytes { + Split { gate: Q4DenseBytes, up: Q4DenseBytes, down: Q4DenseBytes }, + Fused(SharedExpertFusedBytes), +} +enum RoutedGateUpBytes { + Split { gate: Q4ExpertStackBytes, up: Q4ExpertStackBytes }, + Fused(Q4ExpertStackBytes), +} - // token_embd / output are [vocab, hidden] (GGUF fastest-first dims already - // give the flat data in row-major [vocab, hidden] order — see `dims_out_in`). - let load_embed_like = |name: &str| -> Result { - let flat = g.dequant_f32(name)?; - Ok(Tensor::new(dev.upload(&f32s_to_bytes(&flat))?, vec![cfg.vocab, cfg.hidden], DType::F32)) - }; - let tok_embd = load_embed_like("token_embd.weight")?; +struct MoeFfnBytes { + ffn_norm: Vec, + router: Vec, + bias: Vec, + /// `None` under FFAI_LAGUNA_W4A8=1: the routed gate/up/down stacks are + /// built afterward, sequentially, in `upload_layer` / + /// `build_moe_w4a8_sequential` instead of here, W4A8 packing + /// (`ops::pad_rows_f16` / `ops::w4a8_packw`) is a GPU op and can't run on + /// a rayon worker thread, and it needs the SAME dequantized slab the Q4 + /// stack does, so keeping the two together (both sequential) avoids + /// dequantizing the routed tensor twice. + gateup: Option, + down: Option, + shared_expert: SharedExpertBytes, + shared_marlin: Option, +} - let out_norm = load_vec_f32(dev, g, "output_norm.weight")?; - let head_name = if g.tensor("output.weight").is_some() { "output.weight" } else { "token_embd.weight" }; - let lm_head = { - let (qs, sc, m, k) = g.q8_repack(head_name)?; - Q8Head { - qs: Tensor::new(dev.upload(&u32s_to_bytes(&qs))?, vec![qs.len()], DType::U32), - // gemv_q8 reads f32 scales (unlike the q4 family's f16), matching - // the proven resident-Q8 loader path. - scales: Tensor::new(dev.upload(&f32s_to_bytes(&sc))?, vec![sc.len()], DType::F32), - m, - k, +enum LayerFfnBytes { Dense(DenseFfnBytes), Moe(MoeFfnBytes) } +struct LayerBytes { attn: AttnWeightsBytes, ffn: LayerFfnBytes } + +// ── Cache-aware builders (host bytes in, cache-checked, no `dev`) ────────── + +/// The core cached-requant primitive: `build_flat` (a DEFERRED dequant, it +/// only runs on a miss, so a cache hit never touches the GGUF mmap) produces +/// the `[m, k]` f32 slab, which gets `quantize_q4`'d and cached as its +/// `qs`/`scales` byte pair under `key`. +fn q4dense_bytes_from_flat_cached(cache: &CacheCtx, key: &str, m: usize, k: usize, build_flat: impl FnOnce() -> Result>) -> Result { + let qs_key = format!("{key}.q4.qs"); + let sc_key = format!("{key}.q4.scales"); + let blobs = cache.get_or_build_multi(&[(qs_key.as_str(), DTYPE_U32), (sc_key.as_str(), DTYPE_F16)], || { + let flat = build_flat()?; + let (qs, scales) = ops::quantize_q4(&flat, m, k); + Ok(vec![u32s_to_bytes(&qs), f32s_to_f16_bytes(&scales)]) + })?; + let mut it = blobs.into_iter(); + Ok(Q4DenseBytes { qs: it.next().unwrap(), scales: it.next().unwrap(), m, k }) +} + +/// Marlin-only counterpart of [`q4dense_bytes_from_flat_cached`]: quantizes +/// AND Marlin-permutes inside the (deferred, miss-only) build closure. +fn marlindense_bytes_from_flat_cached(cache: &CacheCtx, key: &str, n_out: usize, k_in: usize, build_flat: impl FnOnce() -> Result>) -> Result { + let qs_key = format!("{key}.marlin.qs"); + let sc_key = format!("{key}.marlin.scales"); + let blobs = cache.get_or_build_multi(&[(qs_key.as_str(), DTYPE_U32), (sc_key.as_str(), DTYPE_F16)], || { + let flat = build_flat()?; + let (qs_std, scales) = ops::quantize_q4(&flat, n_out, k_in); + let qs_mar = ops::permute_q4_to_marlin(&qs_std, 1, n_out, k_in); + Ok(vec![u32s_to_bytes(&qs_mar), f32s_to_f16_bytes(&scales)]) + })?; + let mut it = blobs.into_iter(); + Ok(MarlinDenseBytes { qs: it.next().unwrap(), scales: it.next().unwrap(), n_out, k_in }) +} + +/// Cached counterpart of the old `load_dense_and_marlin`: one deferred +/// `quantize_q4` pass feeds BOTH the plain `Q4DenseBytes` (`.q4.*` keys, +/// always built) and, when `want_marlin && m % 64 == 0`, an additional +/// `MarlinDenseBytes` (`.marlin.*` keys) permuted from the SAME `qs_std`. A +/// full cache hit (all keys present) skips `build_flat`, and therefore the +/// GGUF dequant, entirely. +fn q4dense_and_marlin_bytes_cached(cache: &CacheCtx, key: &str, m: usize, k: usize, want_marlin: bool, build_flat: impl FnOnce() -> Result>) -> Result<(Q4DenseBytes, Option)> { + let do_marlin = want_marlin && m % 64 == 0; + let qs_key = format!("{key}.q4.qs"); + let sc_key = format!("{key}.q4.scales"); + let mqs_key = format!("{key}.marlin.qs"); + let msc_key = format!("{key}.marlin.scales"); + let mut keys: Vec<(&str, u8)> = vec![(qs_key.as_str(), DTYPE_U32), (sc_key.as_str(), DTYPE_F16)]; + if do_marlin { + keys.push((mqs_key.as_str(), DTYPE_U32)); + keys.push((msc_key.as_str(), DTYPE_F16)); + } + let blobs = cache.get_or_build_multi(&keys, || { + let flat = build_flat()?; + let (qs_std, scales) = ops::quantize_q4(&flat, m, k); + let mut out = vec![u32s_to_bytes(&qs_std), f32s_to_f16_bytes(&scales)]; + if do_marlin { + let qs_mar = ops::permute_q4_to_marlin(&qs_std, 1, m, k); + out.push(u32s_to_bytes(&qs_mar)); + out.push(f32s_to_f16_bytes(&scales)); } + Ok(out) + })?; + let mut it = blobs.into_iter(); + let dense = Q4DenseBytes { qs: it.next().unwrap(), scales: it.next().unwrap(), m, k }; + let marlin = if do_marlin { + Some(MarlinDenseBytes { qs: it.next().unwrap(), scales: it.next().unwrap(), n_out: m, k_in: k }) + } else { + None }; - let one = Tensor::new(dev.upload(&1.0f32.to_le_bytes())?, vec![1], DType::F32); + Ok((dense, marlin)) +} - let micro = laguna_micro_enabled(); +/// Cached dense f32 vector (norms, router, bias, embeddings, ...), the +/// upload bytes are literally `f32s_to_bytes(flat)`, so caching mostly buys +/// skipping `build_flat`'s dequant on a hit (which matters a lot for the +/// two big ones, `token_embd.weight` / `output.weight`, and is free for the +/// tiny per-layer norms). +fn f32vec_bytes_cached(cache: &CacheCtx, key: &str, build_flat: impl FnOnce() -> Result>) -> Result> { + cache.get_or_build(&format!("{key}.f32vec"), DTYPE_F32, || Ok(f32s_to_bytes(&build_flat()?))) +} - let mut layers = Vec::with_capacity(cfg.n_layers); - for i in 0..cfg.n_layers { - let p = format!("blk.{i}"); - - // n_head is derived from wq's own GGUF dims (48 full-attention / 72 - // sliding for Laguna-S-2.1), not the LagunaConfig formula — see the - // AttnWeights doc comment. - let (wq_flat, n_q_rows, hidden_k) = load_dense_raw(g, &format!("{p}.attn_q.weight"))?; - let (wk_flat, n_kv_rows, _) = load_dense_raw(g, &format!("{p}.attn_k.weight"))?; - let (wv_flat, n_kv_rows_v, _) = load_dense_raw(g, &format!("{p}.attn_v.weight"))?; - let (wg_flat, n_g_rows, _) = load_dense_raw(g, &format!("{p}.attn_gate.weight"))?; - debug_assert_eq!(n_kv_rows, n_kv_rows_v, "wk/wv row count mismatch"); - let n_head = n_q_rows / cfg.head_dim; - - // FFAI_LAGUNA_MICRO=1: concatenate the four dequantized f32 slabs - // BEFORE quantize_q4 (row order q,k,v,gate; all share k=hidden), so - // decode can do one gemv_q4 instead of four. Only ONE of the two - // layouts is ever built per layer — building both would double the - // resident attention-projection weights (~2.4GB total). - let proj = if micro { - let n_rows = n_q_rows + n_kv_rows + n_kv_rows_v + n_g_rows; +/// Cached [`Q4ExpertStack`] bytes: `build_flat` (deferred) produces the +/// `[n_expert, m, k]` expert-major f32 slab; the per-expert `quantize_q4` +/// loop (the heavy part) only runs on a miss. +fn q4_expert_stack_bytes_cached(cache: &CacheCtx, key: &str, n_expert: usize, m: usize, k: usize, build_flat: impl FnOnce() -> Result>) -> Result { + let qs_key = format!("{key}.q4x.qs"); + let sc_key = format!("{key}.q4x.scales"); + let blobs = cache.get_or_build_multi(&[(qs_key.as_str(), DTYPE_U32), (sc_key.as_str(), DTYPE_F16)], || { + let flat = build_flat()?; + let per_expert = m * k; + if flat.len() != per_expert * n_expert { + return Err(Error::Msg(format!( + "Q4ExpertStack '{key}': flat len {} != n_expert*m*k ({} * {} * {})", + flat.len(), n_expert, m, k + ))); + } + let qw = per_expert / 32 * 4; // quantize_q4: m*k/32 blocks, 4 u32/block + let sw = per_expert / 32; // 1 f32 scale/block + let mut qs_all = Vec::with_capacity(qw * n_expert); + let mut sc_all = Vec::with_capacity(sw * n_expert); + for e in 0..n_expert { + let slab = &flat[e * per_expert..(e + 1) * per_expert]; + let (qs, sc) = ops::quantize_q4(slab, m, k); + qs_all.extend_from_slice(&qs); + sc_all.extend_from_slice(&sc); + } + Ok(vec![u32s_to_bytes(&qs_all), f32s_to_f16_bytes(&sc_all)]) + })?; + let mut it = blobs.into_iter(); + Ok(Q4ExpertStackBytes { qs: it.next().unwrap(), scales: it.next().unwrap(), m, k }) +} + +/// Cached [`Q8Head`] bytes via [`Gguf::q8_repack`] (deferred, only runs on a +/// miss). `(m, k)` come from the tensor's own GGUF dims (matches +/// `q8_repack`'s internal derivation exactly, both read `dims[1]`/`dims[0]` +/// the same way `dims_out_in` does), so they're known without calling +/// `q8_repack` and therefore available even on a full cache hit. +fn q8head_bytes_cached(cache: &CacheCtx, g: &Gguf, name: &str, key: &str) -> Result<(Vec, Vec, usize, usize)> { + let (m, k) = tensor_dims(g, name)?; + let qs_key = format!("{key}.q8.qs"); + let sc_key = format!("{key}.q8.scales"); + let blobs = cache.get_or_build_multi(&[(qs_key.as_str(), DTYPE_U32), (sc_key.as_str(), DTYPE_F32)], || { + let (qs, sc, _m, _k) = g.q8_repack(name)?; + Ok(vec![u32s_to_bytes(&qs), f32s_to_bytes(&sc)]) + })?; + let mut it = blobs.into_iter(); + Ok((it.next().unwrap(), it.next().unwrap(), m, k)) +} + +// ── Upload-phase helpers (sequential, `dev`-only, no CPU-heavy work) ─────── + +fn upload_q4dense_bytes(dev: &dyn Device, b: Q4DenseBytes) -> Result { + Ok(Q4Dense { + qs: Tensor::new(dev.upload(&b.qs)?, vec![b.qs.len() / 4], DType::U32), + scales: Tensor::new(dev.upload(&b.scales)?, vec![b.scales.len() / 2], DType::F16), + m: b.m, + k: b.k, + }) +} +fn upload_marlindense_bytes(dev: &dyn Device, b: MarlinDenseBytes) -> Result { + Ok(MarlinDense { + qs: Tensor::new(dev.upload(&b.qs)?, vec![b.qs.len() / 4], DType::U32), + scales: Tensor::new(dev.upload(&b.scales)?, vec![b.scales.len() / 2], DType::F16), + n_out: b.n_out, + k_in: b.k_in, + }) +} +fn upload_expert_stack_bytes(dev: &dyn Device, b: Q4ExpertStackBytes) -> Result { + let per_expert = b.m * b.k; + Ok(Q4ExpertStack { + qs: Tensor::new(dev.upload(&b.qs)?, vec![b.qs.len() / 4], DType::U32), + scales: Tensor::new(dev.upload(&b.scales)?, vec![b.scales.len() / 2], DType::F16), + m: b.m, + k: b.k, + qs_words_per_expert: per_expert / 32 * 4, + scale_elems_per_expert: per_expert / 32, + }) +} +fn upload_f32_bytes(dev: &dyn Device, bytes: Vec) -> Result { + let n = bytes.len() / 4; + Ok(Tensor::new(dev.upload(&bytes)?, vec![n], DType::F32)) +} + +/// Pure-CPU half of loading one layer: dequantize (deferred inside every +/// cache closure, a cache HIT never touches the GGUF mmap) + requantize / +/// permute / interleave every weight this layer needs, each artifact +/// individually cache-checked. Touches neither `dev` nor any device buffer, +/// so this is exactly the per-layer unit `load_gguf` runs across a rayon +/// thread pool on a cache miss (`into_par_iter` over layer indices, see +/// that fn's doc). Under FFAI_LAGUNA_W4A8=1 the routed gate/up/down stacks +/// are left `None` here and built afterward, sequentially, by +/// `build_moe_w4a8_sequential` (see [`MoeFfnBytes::gateup`]'s doc for why). +#[allow(clippy::too_many_arguments)] +fn build_layer_bytes(g: &Gguf, cache: &CacheCtx, cfg: &LagunaConfig, i: usize, micro: bool, marlin: bool, moefuse: bool, w4a8: bool) -> Result { + let p = format!("blk.{i}"); + let q_name = format!("{p}.attn_q.weight"); + let k_name = format!("{p}.attn_k.weight"); + let v_name = format!("{p}.attn_v.weight"); + let gate_name = format!("{p}.attn_gate.weight"); + + // n_head is derived from wq's own GGUF dims (48 full-attention / 72 + // sliding for Laguna-S-2.1), not the LagunaConfig formula: see the + // AttnWeights doc comment. Dims-only reads (no dequant) so this is cheap + // regardless of cache state. + let (n_q_rows, hidden_k) = tensor_dims(g, &q_name)?; + let (n_kv_rows, _) = tensor_dims(g, &k_name)?; + let (n_kv_rows_v, _) = tensor_dims(g, &v_name)?; + let (n_g_rows, _) = tensor_dims(g, &gate_name)?; + debug_assert_eq!(n_kv_rows, n_kv_rows_v, "wk/wv row count mismatch"); + let n_head = n_q_rows / cfg.head_dim; + + // FFAI_LAGUNA_MICRO=1: concatenate the four dequantized f32 slabs BEFORE + // quantize_q4 (row order q,k,v,gate; all share k=hidden), so decode can + // do one gemv_q4 instead of four. Only ONE of the two layouts is ever + // built per layer, building both would double the resident + // attention-projection weights (~2.4GB total). + let proj = if micro { + let n_rows = n_q_rows + n_kv_rows + n_kv_rows_v + n_g_rows; + let wqkvg = q4dense_bytes_from_flat_cached(cache, &format!("{p}.attn_qkvg_fused"), n_rows, hidden_k, || { let mut concat = Vec::with_capacity(n_rows * hidden_k); - concat.extend_from_slice(&wq_flat); - concat.extend_from_slice(&wk_flat); - concat.extend_from_slice(&wv_flat); - concat.extend_from_slice(&wg_flat); - let wqkvg = Q4Dense::upload(dev, &concat, n_rows, hidden_k)?; - AttnProj::Fused { wqkvg, n_q_rows, n_kv_rows, n_g_rows } + concat.extend_from_slice(&g.dequant_f32(&q_name)?); + concat.extend_from_slice(&g.dequant_f32(&k_name)?); + concat.extend_from_slice(&g.dequant_f32(&v_name)?); + concat.extend_from_slice(&g.dequant_f32(&gate_name)?); + Ok(concat) + })?; + AttnProjBytes::Fused { wqkvg, n_q_rows, n_kv_rows, n_g_rows } + } else { + AttnProjBytes::Split { + wq: q4dense_bytes_from_flat_cached(cache, &q_name, n_q_rows, hidden_k, || g.dequant_f32(&q_name))?, + wk: q4dense_bytes_from_flat_cached(cache, &k_name, n_kv_rows, hidden_k, || g.dequant_f32(&k_name))?, + wv: q4dense_bytes_from_flat_cached(cache, &v_name, n_kv_rows_v, hidden_k, || g.dequant_f32(&v_name))?, + wgate: q4dense_bytes_from_flat_cached(cache, &gate_name, n_g_rows, hidden_k, || g.dequant_f32(&gate_name))?, + } + }; + + // FFAI_LAGUNA_MARLIN=1: pack [wq;wk;wv] (gate excluded, see QkvMarlin's + // doc) into a Marlin-tile-major MarlinDense for the tensor-core + // batched-prefill GEMM, built ALONGSIDE (never instead of) the AttnProj + // decode copy above. `g_proj` is a standalone plain Q4Dense so the + // marlin prefill branch has gate access independent of whether `proj` + // above resolved to Split or Fused. + let qkv_marlin = if marlin { + let n_rows_qkv = n_q_rows + n_kv_rows + n_kv_rows_v; + if n_rows_qkv % 64 == 0 { + let mat = marlindense_bytes_from_flat_cached(cache, &format!("{p}.attn_qkv"), n_rows_qkv, hidden_k, || { + let mut concat = Vec::with_capacity(n_rows_qkv * hidden_k); + concat.extend_from_slice(&g.dequant_f32(&q_name)?); + concat.extend_from_slice(&g.dequant_f32(&k_name)?); + concat.extend_from_slice(&g.dequant_f32(&v_name)?); + Ok(concat) + })?; + Some(QkvMarlinBytes { mat, n_q_rows, n_kv_rows }) } else { - AttnProj::Split { - wq: Q4Dense::upload(dev, &wq_flat, n_q_rows, hidden_k)?, - wk: Q4Dense::upload(dev, &wk_flat, n_kv_rows, hidden_k)?, - wv: Q4Dense::upload(dev, &wv_flat, n_kv_rows_v, hidden_k)?, - wgate: Q4Dense::upload(dev, &wg_flat, n_g_rows, hidden_k)?, + None + } + } else { + None + }; + let g_proj = if marlin { + Some(q4dense_bytes_from_flat_cached(cache, &format!("{gate_name}.standalone"), n_g_rows, hidden_k, || g.dequant_f32(&gate_name))?) + } else { + None + }; + + let wo_name = format!("{p}.attn_output.weight"); + let (wo_m, wo_k) = tensor_dims(g, &wo_name)?; + let (wo, wo_marlin) = q4dense_and_marlin_bytes_cached(cache, &wo_name, wo_m, wo_k, marlin, || g.dequant_f32(&wo_name))?; + + let attn_norm_name = format!("{p}.attn_norm.weight"); + let q_norm_name = format!("{p}.attn_q_norm.weight"); + let k_norm_name = format!("{p}.attn_k_norm.weight"); + let attn = AttnWeightsBytes { + n_head, + attn_norm: f32vec_bytes_cached(cache, &attn_norm_name, || g.dequant_f32(&attn_norm_name))?, + proj, + wo, + q_norm: f32vec_bytes_cached(cache, &q_norm_name, || g.dequant_f32(&q_norm_name))?, + k_norm: f32vec_bytes_cached(cache, &k_norm_name, || g.dequant_f32(&k_norm_name))?, + qkv_marlin, + g_proj, + wo_marlin, + }; + + let ffn = if cfg.is_dense_layer(i) { + // FFAI_LAGUNA_MARLIN=1: [gate;up] concat (2*n_ff_dense=24576, + // 64-aligned) plus down, both Marlin-packed alongside the plain + // per-projection Q4Dense copies decode/gemm_q4_mpp use. + let gate_name2 = format!("{p}.ffn_gate.weight"); + let up_name2 = format!("{p}.ffn_up.weight"); + let down_name2 = format!("{p}.ffn_down.weight"); + let (ff_m, ff_k) = tensor_dims(g, &gate_name2)?; + let gate = q4dense_bytes_from_flat_cached(cache, &gate_name2, ff_m, ff_k, || g.dequant_f32(&gate_name2))?; + let up = q4dense_bytes_from_flat_cached(cache, &up_name2, ff_m, ff_k, || g.dequant_f32(&up_name2))?; + let (down_m, down_k) = tensor_dims(g, &down_name2)?; + let (down, down_marlin) = q4dense_and_marlin_bytes_cached(cache, &down_name2, down_m, down_k, marlin, || g.dequant_f32(&down_name2))?; + let gateup_marlin = if marlin && (2 * ff_m) % 64 == 0 { + let mat = marlindense_bytes_from_flat_cached(cache, &format!("{p}.ffn_gateup"), 2 * ff_m, ff_k, || { + let mut concat = Vec::with_capacity(2 * ff_m * ff_k); + concat.extend_from_slice(&g.dequant_f32(&gate_name2)?); + concat.extend_from_slice(&g.dequant_f32(&up_name2)?); + Ok(concat) + })?; + Some(GateUpMarlinBytes { mat, n_ff: ff_m }) + } else { + None + }; + let ffn_norm_name = format!("{p}.ffn_norm.weight"); + LayerFfnBytes::Dense(DenseFfnBytes { + ffn_norm: f32vec_bytes_cached(cache, &ffn_norm_name, || g.dequant_f32(&ffn_norm_name))?, + gate, + up, + down, + gateup_marlin, + down_marlin, + }) + } else { + let router_name = format!("{p}.ffn_gate_inp.weight"); + let router = f32vec_bytes_cached(cache, &router_name, || g.dequant_f32(&router_name))?; + + let gate_shexp = format!("{p}.ffn_gate_shexp.weight"); + let up_shexp = format!("{p}.ffn_up_shexp.weight"); + let down_shexp = format!("{p}.ffn_down_shexp.weight"); + + // FFAI_LAGUNA_MICRO=1: repack the shared expert's gate+up as + // single-expert (n_expert=1) stacks so decode can drive them through + // the same moe_gather_q4_swiglu the routed experts use (top_k=1) + // instead of two plain gemvs + a host-side swiglu. `down` stays a + // plain Q4Dense either way. Only one layout is ever built per layer + // (same memory-doubling rationale as AttnProj above). + let shared_expert = if micro { + let gate = q4_expert_stack_bytes_cached(cache, &format!("{gate_shexp}.se1"), 1, cfg.n_ff_shexp, cfg.hidden, || g.dequant_f32(&gate_shexp))?; + let up = q4_expert_stack_bytes_cached(cache, &format!("{up_shexp}.se1"), 1, cfg.n_ff_shexp, cfg.hidden, || g.dequant_f32(&up_shexp))?; + let (dm, dk) = tensor_dims(g, &down_shexp)?; + let down = q4dense_bytes_from_flat_cached(cache, &down_shexp, dm, dk, || g.dequant_f32(&down_shexp))?; + SharedExpertBytes::Fused(SharedExpertFusedBytes { gate, up, down }) + } else { + let (gm, gk) = tensor_dims(g, &gate_shexp)?; + let (um, uk) = tensor_dims(g, &up_shexp)?; + let (dm, dk) = tensor_dims(g, &down_shexp)?; + SharedExpertBytes::Split { + gate: q4dense_bytes_from_flat_cached(cache, &gate_shexp, gm, gk, || g.dequant_f32(&gate_shexp))?, + up: q4dense_bytes_from_flat_cached(cache, &up_shexp, um, uk, || g.dequant_f32(&up_shexp))?, + down: q4dense_bytes_from_flat_cached(cache, &down_shexp, dm, dk, || g.dequant_f32(&down_shexp))?, } }; - let attn = AttnWeights { - n_head, - attn_norm: load_vec_f32(dev, g, &format!("{p}.attn_norm.weight"))?, - proj, - wo: load_dense(dev, g, &format!("{p}.attn_output.weight"))?, - q_norm: load_vec_f32(dev, g, &format!("{p}.attn_q_norm.weight"))?, - k_norm: load_vec_f32(dev, g, &format!("{p}.attn_k_norm.weight"))?, + // FFAI_LAGUNA_MARLIN=1: the shared expert's [gate;up] concat + // (2*n_ff_shexp=2048, 64-aligned) plus down, Marlin-packed straight + // from the shexp GGUF tensors, independent of whether + // `shared_expert` above resolved to Split or Fused. The routed + // per-expert stacks (gate/up/down just below) are NOT Marlin-packed: + // prefill already drives them through the grouped-MMA GEMM + // (`ops::moe_q4_grouped_mma_dev`), a different tensor-core kernel + // this Marlin work doesn't target. + let shared_marlin = if marlin { + let (sm, sk) = tensor_dims(g, &gate_shexp)?; + let (sdm, sdk) = tensor_dims(g, &down_shexp)?; + if (2 * sm) % 64 == 0 && sdm % 64 == 0 { + let mat = marlindense_bytes_from_flat_cached(cache, &format!("{p}.shexp_gateup"), 2 * sm, sk, || { + let mut concat = Vec::with_capacity(2 * sm * sk); + concat.extend_from_slice(&g.dequant_f32(&gate_shexp)?); + concat.extend_from_slice(&g.dequant_f32(&up_shexp)?); + Ok(concat) + })?; + let down = marlindense_bytes_from_flat_cached(cache, &format!("{down_shexp}.standalone"), sdm, sdk, || g.dequant_f32(&down_shexp))?; + Some(SharedMarlinBytes { gateup: GateUpMarlinBytes { mat, n_ff: sm }, down }) + } else { + None + } + } else { + None }; - let ffn = if cfg.is_dense_layer(i) { + let gate_exps_name = format!("{p}.ffn_gate_exps.weight"); + let up_exps_name = format!("{p}.ffn_up_exps.weight"); + let down_exps_name = format!("{p}.ffn_down_exps.weight"); + + // FFAI_LAGUNA_W4A8=1: the routed gate/up/down stacks are built + // sequentially afterward (`build_moe_w4a8_sequential`), see + // `MoeFfnBytes::gateup`'s doc. + let (gateup, down) = if w4a8 { + (None, None) + } else if moefuse { + // FFAI_LAGUNA_MOEFUSE=1: build ONLY the fused [gate;up] stack + // (see RoutedGateUp::Fused's doc for the row layout), never both + // layouts at once, the routed expert stacks are the bulk of + // the model's resident memory and a second copy isn't + // affordable. + let per_expert = cfg.n_ff_exp * cfg.hidden; + let n_expert = cfg.n_expert; + let (n_ff_exp, hidden) = (cfg.n_ff_exp, cfg.hidden); + let fused = q4_expert_stack_bytes_cached(cache, &format!("{p}.ffn_gateup_exps_fused"), n_expert, 2 * n_ff_exp, hidden, || { + let gate_flat = g.dequant_f32(&gate_exps_name)?; + let up_flat = g.dequant_f32(&up_exps_name)?; + if gate_flat.len() != per_expert * n_expert || up_flat.len() != per_expert * n_expert { + return Err(Error::Msg(format!( + "ffn_gate/up_exps len mismatch gate={} up={} expected {} each (n_expert={n_expert} n_ff_exp={n_ff_exp} hidden={hidden})", + gate_flat.len(), up_flat.len(), per_expert * n_expert, + ))); + } + let mut fused = Vec::with_capacity(2 * per_expert * n_expert); + for e in 0..n_expert { + fused.extend_from_slice(&gate_flat[e * per_expert..(e + 1) * per_expert]); + fused.extend_from_slice(&up_flat[e * per_expert..(e + 1) * per_expert]); + } + Ok(fused) + })?; + let down = q4_expert_stack_bytes_cached(cache, &down_exps_name, cfg.n_expert, cfg.hidden, cfg.n_ff_exp, || g.dequant_f32(&down_exps_name))?; + (Some(RoutedGateUpBytes::Fused(fused)), Some(down)) + } else { + let gate = q4_expert_stack_bytes_cached(cache, &gate_exps_name, cfg.n_expert, cfg.n_ff_exp, cfg.hidden, || g.dequant_f32(&gate_exps_name))?; + let up = q4_expert_stack_bytes_cached(cache, &up_exps_name, cfg.n_expert, cfg.n_ff_exp, cfg.hidden, || g.dequant_f32(&up_exps_name))?; + let down = q4_expert_stack_bytes_cached(cache, &down_exps_name, cfg.n_expert, cfg.hidden, cfg.n_ff_exp, || g.dequant_f32(&down_exps_name))?; + (Some(RoutedGateUpBytes::Split { gate, up }), Some(down)) + }; + + let ffn_norm_name = format!("{p}.ffn_norm.weight"); + let bias_name = format!("{p}.exp_probs_b.bias"); + LayerFfnBytes::Moe(MoeFfnBytes { + ffn_norm: f32vec_bytes_cached(cache, &ffn_norm_name, || g.dequant_f32(&ffn_norm_name))?, + router, + bias: f32vec_bytes_cached(cache, &bias_name, || g.dequant_f32(&bias_name))?, + gateup, + down, + shared_expert, + shared_marlin, + }) + }; + + Ok(LayerBytes { attn, ffn }) +} + +/// FFAI_LAGUNA_W4A8=1 sequential (non-rayon-parallel) build of one MoE +/// layer's routed gate/up/down. Mirrors the original inline `load_gguf` +/// logic exactly: dequantize each routed GGUF tensor ONCE and feed BOTH the +/// resident Q4 stack decode needs and the ADDITIONAL W4A8-packed stack +/// prefill needs (see [`MoeFfnW4a8`]'s doc), the Q4 side is routed through +/// [`CacheCtx`] (pure CPU; kept here rather than in the parallel +/// `build_layer_bytes` pass purely so it shares that one dequant with the +/// W4A8 pack below instead of paying for it twice) and the W4A8 pack itself +/// through [`pack_w4a8_expert_stack_cached`] (a GPU op, must run on this +/// thread, see that fn's doc). Not parallelized across layers: acceptable +/// since FFAI_LAGUNA_W4A8 is an opt-in, default-off, prefill-throughput +/// experiment, not the path the "fast load" goal targets. +fn build_moe_w4a8_sequential(dev: &dyn Device, g: &Gguf, cache: &CacheCtx, cfg: &LagunaConfig, i: usize) -> Result<(RoutedGateUp, Q4ExpertStack, Option)> { + let p = format!("blk.{i}"); + let gate_exps_name = format!("{p}.ffn_gate_exps.weight"); + let up_exps_name = format!("{p}.ffn_up_exps.weight"); + let down_exps_name = format!("{p}.ffn_down_exps.weight"); + let moefuse = laguna_moefuse_enabled(); + // FFAI_LAGUNA_W4A8_ONLY=1: PREFILL BENCH MODE. Dual expert residency (Q4 + // for decode + mxfp4 for W4A8 prefill) is ~112GB and cannot fit; this + // mode uploads a 1-expert DUMMY Q4 stack (decode would be wrong, never + // run it in this mode) so the W4A8 prefill speed can be measured. The + // real fix is a single shared expert format both paths read. + let w4a8_only = std::env::var("FFAI_LAGUNA_W4A8_ONLY").map(|v| v == "1").unwrap_or(false); + + let gate_flat = g.dequant_f32(&gate_exps_name)?; + let up_flat = g.dequant_f32(&up_exps_name)?; + + let (gateup, gateup_w4a8) = if moefuse { + let per_expert = cfg.n_ff_exp * cfg.hidden; + if gate_flat.len() != per_expert * cfg.n_expert || up_flat.len() != per_expert * cfg.n_expert { + return Err(Error::Msg(format!( + "ffn_gate/up_exps len mismatch gate={} up={} expected {} each (n_expert={} n_ff_exp={} hidden={})", + gate_flat.len(), up_flat.len(), per_expert * cfg.n_expert, cfg.n_expert, cfg.n_ff_exp, cfg.hidden + ))); + } + let mut fused = Vec::with_capacity(2 * per_expert * cfg.n_expert); + for e in 0..cfg.n_expert { + fused.extend_from_slice(&gate_flat[e * per_expert..(e + 1) * per_expert]); + fused.extend_from_slice(&up_flat[e * per_expert..(e + 1) * per_expert]); + } + let w4a8s = pack_w4a8_expert_stack_cached(dev, cache, &format!("{p}.ffn_gateup_exps_fused"), &fused, cfg.n_expert, 2 * cfg.n_ff_exp, cfg.hidden)?; + let q4 = if w4a8_only { + let b = q4_expert_stack_bytes_cached(cache, &format!("{p}.ffn_gateup_exps_fused.w4a8only1"), 1, 2 * cfg.n_ff_exp, cfg.hidden, || Ok(fused[..2 * per_expert].to_vec()))?; + upload_expert_stack_bytes(dev, b)? + } else { + let b = q4_expert_stack_bytes_cached(cache, &format!("{p}.ffn_gateup_exps_fused"), cfg.n_expert, 2 * cfg.n_ff_exp, cfg.hidden, || Ok(fused))?; + upload_expert_stack_bytes(dev, b)? + }; + (RoutedGateUp::Fused(q4), Some(RoutedGateUpW4a8::Fused(w4a8s))) + } else { + let gate_b = q4_expert_stack_bytes_cached(cache, &gate_exps_name, cfg.n_expert, cfg.n_ff_exp, cfg.hidden, || Ok(gate_flat.clone()))?; + let up_b = q4_expert_stack_bytes_cached(cache, &up_exps_name, cfg.n_expert, cfg.n_ff_exp, cfg.hidden, || Ok(up_flat.clone()))?; + let gate_q4 = upload_expert_stack_bytes(dev, gate_b)?; + let up_q4 = upload_expert_stack_bytes(dev, up_b)?; + let gate_w4a8 = pack_w4a8_expert_stack_cached(dev, cache, &gate_exps_name, &gate_flat, cfg.n_expert, cfg.n_ff_exp, cfg.hidden)?; + let up_w4a8 = pack_w4a8_expert_stack_cached(dev, cache, &up_exps_name, &up_flat, cfg.n_expert, cfg.n_ff_exp, cfg.hidden)?; + (RoutedGateUp::Split { gate: gate_q4, up: up_q4 }, Some(RoutedGateUpW4a8::Split { gate: gate_w4a8, up: up_w4a8 })) + }; + + let down_flat = g.dequant_f32(&down_exps_name)?; + let down_w4a8 = pack_w4a8_expert_stack_cached(dev, cache, &down_exps_name, &down_flat, cfg.n_expert, cfg.hidden, cfg.n_ff_exp)?; + let down = if w4a8_only { + let b = q4_expert_stack_bytes_cached(cache, &format!("{down_exps_name}.w4a8only1"), 1, cfg.hidden, cfg.n_ff_exp, || Ok(down_flat[..cfg.hidden * cfg.n_ff_exp].to_vec()))?; + upload_expert_stack_bytes(dev, b)? + } else { + let b = q4_expert_stack_bytes_cached(cache, &down_exps_name, cfg.n_expert, cfg.hidden, cfg.n_ff_exp, || Ok(down_flat))?; + upload_expert_stack_bytes(dev, b)? + }; + + // Both halves are always built together in this fn (it only ever runs + // under FFAI_LAGUNA_W4A8=1), so `gateup_w4a8` is always `Some` here: + // `.expect` just makes that invariant explicit instead of a silent + // `.unwrap()`. + let w4a8_moe = Some(MoeFfnW4a8 { gateup: gateup_w4a8.expect("w4a8 branch always builds a gateup pack"), down: down_w4a8 }); + Ok((gateup, down, w4a8_moe)) +} + +/// Sequential half of loading one layer: upload every byte blob +/// `build_layer_bytes` produced (or, under FFAI_LAGUNA_W4A8=1, that this fn +/// itself builds via `build_moe_w4a8_sequential`, see [`MoeFfnBytes::gateup`]'s +/// doc) and assemble the `Tensor`-based [`LagunaLayer`]. Never parallelized +/// across layers, `dev.upload`/device dispatch calls only ever happen on +/// the caller's thread (the loop in `load_gguf`), since backends aren't +/// guaranteed safe to drive concurrently from multiple threads. +fn upload_layer(dev: &dyn Device, g: &Gguf, cache: &CacheCtx, cfg: &LagunaConfig, i: usize, w4a8: bool, lb: LayerBytes) -> Result { + let attn = { + let a = lb.attn; + let proj = match a.proj { + AttnProjBytes::Fused { wqkvg, n_q_rows, n_kv_rows, n_g_rows } => { + AttnProj::Fused { wqkvg: upload_q4dense_bytes(dev, wqkvg)?, n_q_rows, n_kv_rows, n_g_rows } + } + AttnProjBytes::Split { wq, wk, wv, wgate } => AttnProj::Split { + wq: upload_q4dense_bytes(dev, wq)?, + wk: upload_q4dense_bytes(dev, wk)?, + wv: upload_q4dense_bytes(dev, wv)?, + wgate: upload_q4dense_bytes(dev, wgate)?, + }, + }; + let wo = upload_q4dense_bytes(dev, a.wo)?; + let qkv_marlin = a.qkv_marlin.map(|qm| -> Result { + Ok(QkvMarlin { mat: upload_marlindense_bytes(dev, qm.mat)?, n_q_rows: qm.n_q_rows, n_kv_rows: qm.n_kv_rows }) + }).transpose()?; + let g_proj = a.g_proj.map(|b| upload_q4dense_bytes(dev, b)).transpose()?; + let wo_marlin = a.wo_marlin.map(|b| upload_marlindense_bytes(dev, b)).transpose()?; + AttnWeights { + n_head: a.n_head, + attn_norm: upload_f32_bytes(dev, a.attn_norm)?, + proj, + wo, + q_norm: upload_f32_bytes(dev, a.q_norm)?, + k_norm: upload_f32_bytes(dev, a.k_norm)?, + qkv_marlin, + g_proj, + wo_marlin, + } + }; + + let ffn = match lb.ffn { + LayerFfnBytes::Dense(d) => { + let gateup_marlin = d.gateup_marlin.map(|gm| -> Result { + Ok(GateUpMarlin { mat: upload_marlindense_bytes(dev, gm.mat)?, n_ff: gm.n_ff }) + }).transpose()?; LayerFfn::Dense(DenseFfn { - ffn_norm: load_vec_f32(dev, g, &format!("{p}.ffn_norm.weight"))?, - gate: load_dense(dev, g, &format!("{p}.ffn_gate.weight"))?, - up: load_dense(dev, g, &format!("{p}.ffn_up.weight"))?, - down: load_dense(dev, g, &format!("{p}.ffn_down.weight"))?, + ffn_norm: upload_f32_bytes(dev, d.ffn_norm)?, + gate: upload_q4dense_bytes(dev, d.gate)?, + up: upload_q4dense_bytes(dev, d.up)?, + down: upload_q4dense_bytes(dev, d.down)?, + gateup_marlin, + down_marlin: d.down_marlin.map(|b| upload_marlindense_bytes(dev, b)).transpose()?, }) - } else { - let router_flat = g.dequant_f32(&format!("{p}.ffn_gate_inp.weight"))?; - let router = Tensor::new( - dev.upload(&f32s_to_bytes(&router_flat))?, - vec![cfg.n_expert, cfg.hidden], - DType::F32, - ); - - // FFAI_LAGUNA_MICRO=1: repack the shared expert's gate+up as - // single-expert (n_expert=1) stacks so decode can drive them - // through the same moe_gather_q4_swiglu the routed experts use - // (top_k=1) instead of two plain gemvs + a host-side swiglu. - // `down` stays a plain Q4Dense either way. Only one layout is - // ever built per layer (same memory-doubling rationale as - // AttnProj above). - let shared_expert = if micro { - let gate = load_expert_stack(dev, g, &format!("{p}.ffn_gate_shexp.weight"), 1, cfg.n_ff_shexp, cfg.hidden)?; - let up = load_expert_stack(dev, g, &format!("{p}.ffn_up_shexp.weight"), 1, cfg.n_ff_shexp, cfg.hidden)?; - let down = load_dense(dev, g, &format!("{p}.ffn_down_shexp.weight"))?; - let zero_idx = Tensor::new(dev.upload(&0u32.to_le_bytes())?, vec![1], DType::U32); - SharedExpert::Fused(SharedExpertFused { gate, up, down, zero_idx }) - } else { - SharedExpert::Split { - gate: load_dense(dev, g, &format!("{p}.ffn_gate_shexp.weight"))?, - up: load_dense(dev, g, &format!("{p}.ffn_up_shexp.weight"))?, - down: load_dense(dev, g, &format!("{p}.ffn_down_shexp.weight"))?, + } + LayerFfnBytes::Moe(m) => { + let shared_expert = match m.shared_expert { + SharedExpertBytes::Split { gate, up, down } => SharedExpert::Split { + gate: upload_q4dense_bytes(dev, gate)?, + up: upload_q4dense_bytes(dev, up)?, + down: upload_q4dense_bytes(dev, down)?, + }, + SharedExpertBytes::Fused(f) => { + let zero_idx = Tensor::new(dev.upload(&0u32.to_le_bytes())?, vec![1], DType::U32); + SharedExpert::Fused(SharedExpertFused { + gate: upload_expert_stack_bytes(dev, f.gate)?, + up: upload_expert_stack_bytes(dev, f.up)?, + down: upload_q4dense_bytes(dev, f.down)?, + zero_idx, + }) } }; + let shared_marlin = m.shared_marlin.map(|sm| -> Result { + Ok(SharedMarlin { + gateup: GateUpMarlin { mat: upload_marlindense_bytes(dev, sm.gateup.mat)?, n_ff: sm.gateup.n_ff }, + down: upload_marlindense_bytes(dev, sm.down)?, + }) + }).transpose()?; + + let (gateup, down, w4a8_moe) = if w4a8 { + build_moe_w4a8_sequential(dev, g, cache, cfg, i)? + } else { + let gateup = match m.gateup.expect("non-W4A8 build_layer_bytes always sets gateup") { + RoutedGateUpBytes::Fused(q) => RoutedGateUp::Fused(upload_expert_stack_bytes(dev, q)?), + RoutedGateUpBytes::Split { gate, up } => { + RoutedGateUp::Split { gate: upload_expert_stack_bytes(dev, gate)?, up: upload_expert_stack_bytes(dev, up)? } + } + }; + let down = upload_expert_stack_bytes(dev, m.down.expect("non-W4A8 build_layer_bytes always sets down"))?; + (gateup, down, None) + }; LayerFfn::Moe(MoeFfn { - ffn_norm: load_vec_f32(dev, g, &format!("{p}.ffn_norm.weight"))?, - router, - bias: load_vec_f32(dev, g, &format!("{p}.exp_probs_b.bias"))?, - gate: load_expert_stack(dev, g, &format!("{p}.ffn_gate_exps.weight"), cfg.n_expert, cfg.n_ff_exp, cfg.hidden)?, - up: load_expert_stack(dev, g, &format!("{p}.ffn_up_exps.weight"), cfg.n_expert, cfg.n_ff_exp, cfg.hidden)?, - down: load_expert_stack(dev, g, &format!("{p}.ffn_down_exps.weight"), cfg.n_expert, cfg.hidden, cfg.n_ff_exp)?, + ffn_norm: upload_f32_bytes(dev, m.ffn_norm)?, + router: Tensor::new(dev.upload(&m.router)?, vec![cfg.n_expert, cfg.hidden], DType::F32), + bias: upload_f32_bytes(dev, m.bias)?, + gateup, + down, shared_expert, + shared_marlin, + w4a8: w4a8_moe, }) - }; + } + }; + + Ok(LagunaLayer { attn, ffn }) +} + +/// Load Laguna-S-2.1 straight from a GGUF file: parse `laguna.*` metadata for +/// the config, then stream every tensor through `dequant_f32` → requantize +/// (Q4 for projections/experts, dense f32 for embeddings/norms/router) → +/// upload, dropping the f32 immediately after each tensor. +/// +/// Every uploaded byte blob goes through an on-disk [`CacheCtx`] first (see +/// that struct's doc): a warm reload of an unchanged GGUF under the same +/// `FFAI_LAGUNA_*` flags skips straight to `read file -> upload` for every +/// weight, instead of redoing the dequant+requant work. On a cache MISS, the +/// pure-CPU half of that work (`build_layer_bytes`) runs across a rayon +/// thread pool, one task per layer: layers are fully independent (no shared +/// mutable state besides the cache's own atomic hit/miss counters and the +/// read-only `g`/`cfg`), so this is an embarrassingly parallel fan-out. +/// Every `dev.upload` call happens afterward, sequentially, in layer order +/// (`upload_layer`), backends aren't guaranteed safe to drive from multiple +/// threads at once, so nothing in the parallel phase ever touches `dev`. +pub fn load_gguf(dev: &dyn Device, g: &Gguf) -> Result { + let load_t0 = std::time::Instant::now(); + let cfg = LagunaConfig::from_gguf(g); + let cache = CacheCtx::new(g); + + // token_embd / output are [vocab, hidden] (GGUF fastest-first dims + // already give the flat data in row-major [vocab, hidden] order, see + // `dims_out_in`). + let tok_embd = { + let bytes = f32vec_bytes_cached(&cache, "token_embd.weight", || g.dequant_f32("token_embd.weight"))?; + Tensor::new(dev.upload(&bytes)?, vec![cfg.vocab, cfg.hidden], DType::F32) + }; + + let out_norm = { + let bytes = f32vec_bytes_cached(&cache, "output_norm.weight", || g.dequant_f32("output_norm.weight"))?; + upload_f32_bytes(dev, bytes)? + }; + let head_name = if g.tensor("output.weight").is_some() { "output.weight" } else { "token_embd.weight" }; + let lm_head = { + let (qs_bytes, sc_bytes, m, k) = q8head_bytes_cached(&cache, g, head_name, head_name)?; + Q8Head { + qs: Tensor::new(dev.upload(&qs_bytes)?, vec![qs_bytes.len() / 4], DType::U32), + // gemv_q8 reads f32 scales (unlike the q4 family's f16), matching + // the proven resident-Q8 loader path. + scales: Tensor::new(dev.upload(&sc_bytes)?, vec![sc_bytes.len() / 4], DType::F32), + m, + k, + } + }; + // Tiny process-lifetime constant, not worth routing through the cache. + let one = Tensor::new(dev.upload(&1.0f32.to_le_bytes())?, vec![1], DType::F32); - layers.push(LagunaLayer { attn, ffn }); + let micro = laguna_micro_enabled(); + let marlin = laguna_marlin_enabled(); + let moefuse = laguna_moefuse_enabled(); + let w4a8 = laguna_w4a8_enabled(); + + // Pure-CPU half, rayon-parallel WITHIN a bounded window of layers. + // Building all 48 layers before uploading holds tens of GB of blob bytes + // plus each worker's multi-GB f32 dequant transients at once, which + // OOM-kills the process on the shared 128GB. Window of 4: at most 4 + // layers' transients + blobs in flight, uploaded and dropped before the + // next window starts. FFAI_LAGUNA_LOAD_WINDOW overrides. + let window: usize = std::env::var("FFAI_LAGUNA_LOAD_WINDOW") + .ok().and_then(|v| v.parse().ok()).filter(|&w| w >= 1).unwrap_or(4); + let mut layers = Vec::with_capacity(cfg.n_layers); + let mut start_l = 0usize; + while start_l < cfg.n_layers { + let end_l = (start_l + window).min(cfg.n_layers); + let batch: Vec = (start_l..end_l) + .into_par_iter() + .map(|i| build_layer_bytes(g, &cache, &cfg, i, micro, marlin, moefuse, w4a8)) + .collect::>>()?; + // Sequential half: every dev.upload / device dispatch call, in order. + for (off, lb) in batch.into_iter().enumerate() { + layers.push(upload_layer(dev, g, &cache, &cfg, start_l + off, w4a8, lb)?); + } + start_l = end_l; } + let (hits, misses) = cache.stats(); + eprintln!( + "Laguna: load done in {:.2}s ({hits} cache hits, {misses} misses, cache dir {})", + load_t0.elapsed().as_secs_f64(), + cache.dir_display(), + ); + Ok(LagunaModel { cfg, tok_embd, layers, out_norm, lm_head, one }) } @@ -662,7 +1780,7 @@ fn dbg_stats(dev: &dyn Device, tag: &str, t: &Tensor) { /// replays, rather than passed as a scalar baked in at capture time. The /// buffer POINTERS are fixed for the lifetime of a `LagunaDecodeCtx`; only /// the u32 values inside them change. `n_kv_full`/`n_kv_sliding` hold -/// `pos+1` / `min(pos+1, sliding_window)` respectively — which one a given +/// `pos+1` / `min(pos+1, sliding_window)` respectively, which one a given /// layer reads is chosen at call time by `cfg.is_full_layer(layer_idx)`, /// which never changes across replays, so the choice of POINTER is itself /// capture-stable even though the VALUE behind it isn't. @@ -700,7 +1818,7 @@ fn decode_layer( let h = ops::rms_norm(dev, x, &w.attn.attn_norm, cfg.eps)?; if deep_dbg { dbg_stats(dev, "attn_norm", &h); } // g_proj consumes the SAME pre-attention normed input as q/k/v (not the - // attention output) — matches the reference engine's graph exactly. + // attention output), matches the reference engine's graph exactly. // // FFAI_LAGUNA_MICRO=1 (fusion 1): one gemv_q4 on the concatenated // [wq;wk;wv;wgate] weight instead of four separate gemvs, then four @@ -746,7 +1864,7 @@ fn decode_layer( // YaRN (full layers) / plain (sliding layers) partial-rotary RoPE. // Graph mode reads `position` from the (capture-stable-pointer) posbuf - // instead of taking it as a scalar baked in at capture time — same math, + // instead of taking it as a scalar baked in at capture time, same math, // see `ops::rope_yarn_partial_posbuf`'s doc. let (q, k) = if graph_bufs.is_some() { let posbuf = if is_full { full_posbuf } else { sliding_posbuf }; @@ -778,13 +1896,13 @@ fn decode_layer( // KV cache append: full layers grow (posbuf = pos), sliding layers write // a ring slot (posbuf = pos % sliding_window). RoPE already used the // ABSOLUTE position above, so ring order doesn't affect the attention - // math — only which slot gets overwritten. + // math, only which slot gets overwritten. let posbuf = if is_full { full_posbuf } else { sliding_posbuf }; ops::kv_append(dev, &k, &kv.k, posbuf, hd, kv.cap, n_kv_heads * hd)?; ops::kv_append(dev, &v, &kv.v, posbuf, hd, kv.cap, n_kv_heads * hd)?; // Graph mode reads the filled KV length from a (capture-stable-pointer) - // device buffer instead of a baked-in scalar — same dense-causal math, + // device buffer instead of a baked-in scalar, same dense-causal math, // see `ops::sdpa_decode_nbuf`'s doc. let attn = if let Some(gb) = graph_bufs { let n_kv_buf = if is_full { gb.n_kv_full } else { gb.n_kv_sliding }; @@ -807,7 +1925,7 @@ fn decode_layer( // is no tensor-clone op; `ops::slice(dev, x, 0, hidden)` is a device copy) // then `wo.gemv_accum` computes x1 = x + wo·gated in one kernel, instead // of materializing `o = wo.gemv(...)` and adding it separately. Same op - // count (copy+accum vs gemv+add) — this fusion trades materializing `o` + // count (copy+accum vs gemv+add), this fusion trades materializing `o` // for materializing the copy, so it wins only if the accum kernel avoids // enough bandwidth to offset that; benchmarking decides, not this diff. let x1 = if micro { @@ -875,16 +1993,40 @@ fn decode_layer( // FFAI_LAGUNA_FUSE=1 replaces the two gathers + host swiglu with // one fused gate+up+SwiGLU gather kernel: same `h2` read once // instead of twice, one launch instead of two gathers plus the - // elementwise pass. Default (unset) keeps the unfused path. - let act = if std::env::var("FFAI_LAGUNA_FUSE").map(|v| v == "1").unwrap_or(false) { - ops::moe_gather_q4_swiglu( - dev, &f.gate.qs, &f.gate.scales, &f.up.qs, &f.up.scales, &h2, &idx, - cfg.n_expert_used, cfg.n_ff_exp, cfg.hidden, - )? - } else { - let gates = ops::moe_gather_q4(dev, &f.gate.qs, &f.gate.scales, &h2, &idx, cfg.n_expert_used, cfg.n_ff_exp, cfg.hidden)?; - let ups = ops::moe_gather_q4(dev, &f.up.qs, &f.up.scales, &h2, &idx, cfg.n_expert_used, cfg.n_ff_exp, cfg.hidden)?; - ops::swiglu(dev, &gates, &ups)? + // elementwise pass. Default (unset) keeps the unfused path. Only + // available for RoutedGateUp::Split: `moe_gather_q4_swiglu` reads + // gate and up as two SEPARATE `[n_exp*inter, hid]` stacks with + // `row = e*inter + local` indexing, which the fused stack's + // per-expert-interleaved rows do not satisfy. + // + // RoutedGateUp::Fused (FFAI_LAGUNA_MOEFUSE=1) instead reads the + // ONE fused stack with a single moe_gather_q4 call at + // inter=2*n_ff_exp (gathers both gate and up rows for the + // selected experts in one launch), then splits the result's + // columns into gate/up with extract_cols and runs the same + // host-orchestrated swiglu as the pre-FFAI_LAGUNA_FUSE path + // above, see RoutedGateUp::Fused's doc for the tradeoff this + // implies (loses the fused-gather micro-win, about 1 tok/s). + let act = match &f.gateup { + RoutedGateUp::Fused(gu) => { + let inter = 2 * cfg.n_ff_exp; + let both = ops::moe_gather_q4(dev, &gu.qs, &gu.scales, &h2, &idx, cfg.n_expert_used, inter, cfg.hidden)?; + let gates = ops::extract_cols(dev, &both, cfg.n_expert_used, inter, 0, cfg.n_ff_exp)?; + let ups = ops::extract_cols(dev, &both, cfg.n_expert_used, inter, cfg.n_ff_exp, cfg.n_ff_exp)?; + ops::swiglu(dev, &gates, &ups)? + } + RoutedGateUp::Split { gate, up } => { + if std::env::var("FFAI_LAGUNA_FUSE").map(|v| v == "1").unwrap_or(false) { + ops::moe_gather_q4_swiglu( + dev, &gate.qs, &gate.scales, &up.qs, &up.scales, &h2, &idx, + cfg.n_expert_used, cfg.n_ff_exp, cfg.hidden, + )? + } else { + let gates = ops::moe_gather_q4(dev, &gate.qs, &gate.scales, &h2, &idx, cfg.n_expert_used, cfg.n_ff_exp, cfg.hidden)?; + let ups = ops::moe_gather_q4(dev, &up.qs, &up.scales, &h2, &idx, cfg.n_expert_used, cfg.n_ff_exp, cfg.hidden)?; + ops::swiglu(dev, &gates, &ups)? + } + } }; let downs = ops::moe_gather_down(dev, &f.down.qs, &f.down.scales, &act, &idx, cfg.n_expert_used, cfg.n_ff_exp, cfg.hidden)?; let acc = Tensor::new(dev.alloc_zeroed(cfg.hidden * 4)?, vec![cfg.hidden], DType::F32); @@ -893,7 +2035,7 @@ fn decode_layer( if moe_dbg { dbg_stats(dev, &format!("moe{layer_idx}_acc_routed"), &acc); } // Always-on shared expert: same accumulator, unweighted (scale = 1, - // the model-wide cached `one` scalar — no gating/router weight). + // the model-wide cached `one` scalar, no gating/router weight). // // FFAI_LAGUNA_MICRO=1 (fusion 2): the gate+up projections were // repacked at load time as single-expert (n_expert=1) stacks, so @@ -925,15 +2067,681 @@ fn decode_layer( Ok(x2) } +/// One BATCHED prefill layer: `T` prompt-token rows processed in one forward +/// instead of `T` sequential [`decode_layer`] calls. Same op sequence as +/// `decode_layer` (RMSNorm → QKV(+gate) proj → per-head QK-RMSNorm → RoPE → +/// KV-cache write → SDPA → per-head softplus gate → o_proj → residual → +/// RMSNorm → FFN (dense or MoE) → residual), just every op runs on `[T, ...]` +/// batched tensors through the batched-prefill ops (`gemm_q4_mpp`, +/// `rope_yarn_partial_many`, `kv_append_many`, `sdpa_multi_tc_varlen`, +/// `gate_softplus_mul_perhead_many`, the on-device MoE +/// route+sort+grouped-GEMM+scatter pipeline) instead of the per-token gemv +/// family. `x` is `[T, hidden]`; `start` is the ABSOLUTE sequence position of +/// row 0 in this chunk (`positions[i] == start + i`, uploaded once by the +/// caller and shared across every layer of this chunk). +/// +/// Full-attention layers write straight into `kv`'s persistent growing cache +/// (batched writes there are safe, nothing is evicted) and attend the full +/// causal prefix (`seg_lo` all-zero). Sliding-window layers write into +/// `scratch`'s LINEAR per-layer buffer instead (see [`LagunaPrefillScratch`]'s +/// doc for why the ring can't be batch-written) and attend a per-row window +/// `seg_lo[i] = max(0, (start+i) - (sliding_window-1))`. +#[allow(clippy::too_many_arguments)] + +/// FFAI_LAGUNA_PP_PROF=1: coarse stage profile of the prefill path. psync +/// points synchronize, so the elapsed time of each sync attributes the queued +/// work since the previous point to the CURRENT tag. Coarse but names the wall. +static PROF: std::sync::Mutex>> = + std::sync::Mutex::new(None); +fn prof_add(tag: &str, d: std::time::Duration) { + let mut g = PROF.lock().unwrap(); + let m = g.get_or_insert_with(Default::default); + *m.entry(tag.to_string()).or_default() += d; +} +/// Print and reset the accumulated stage profile. +pub fn prof_report() { + if let Some(m) = PROF.lock().unwrap().take() { + let total: f64 = m.values().map(|d| d.as_secs_f64()).sum(); + eprintln!(" prefill stage profile (sync-attributed, total {total:.3}s):"); + let mut rows: Vec<_> = m.into_iter().collect(); + rows.sort_by(|a, b| b.1.cmp(&a.1)); + for (tag, d) in rows { + eprintln!(" {tag:<20} {:>8.3}s {:>5.1}%", d.as_secs_f64(), 100.0 * d.as_secs_f64() / total.max(1e-9)); + } + } +} + +fn prefill_layer( + dev: &dyn Device, + cfg: &LagunaConfig, + layer_idx: usize, + w: &LagunaLayer, + kv: &LagunaKvCache, + scratch: &LagunaPrefillScratch, + x: &Tensor, // [T, hidden] f32 + t: usize, + start: usize, + positions: &Tensor, // [T] u32, positions[i] = start + i + // FFAI_LAGUNA_MARLIN=1: a persistent all-zero `u32` buffer with at least + // `t` elements, shared by every Marlin call in this layer (and every + // other layer/chunk of the same `prefill()` call, see that fn's doc). + // `None` when the env is unset; `Some` implies "check each per-weight + // `Option` before using it" since an individual weight may + // still be `None` (dimension not 64-aligned). + marlin_zero_idx: Option<&Tensor>, +) -> Result { + let hd = cfg.head_dim; + let hidden = cfg.hidden; + let n_head = w.attn.n_head; + let n_kv_heads = cfg.n_kv_heads; + let is_full = cfg.is_full_layer(layer_idx); + let rope = cfg.rope_for(layer_idx); + let heads_per_group = (n_head / n_kv_heads) as u32; + // FFAI_LAGUNA_PREFILL_SYNC=1: synchronize after each stage so an async + // device fault surfaces AT the guilty op instead of a later module load. + let prof = std::env::var("FFAI_LAGUNA_PP_PROF").is_ok(); + let psync = |tag: &str| -> Result<()> { + if prof || std::env::var("FFAI_LAGUNA_PREFILL_SYNC").is_ok() { + let t0 = std::time::Instant::now(); + dev.synchronize().map_err(|e| { + Error::Msg(format!("prefill fault layer {layer_idx} stage {tag}: {e:?}")) + })?; + if prof { + prof_add(tag, t0.elapsed()); + } + } + Ok(()) + }; + + // ── attention ──────────────────────────────────────────────────── + let h = ops::rms_norm(dev, x, &w.attn.attn_norm, cfg.eps) + .map_err(|e| Error::Msg(format!("prefill layer {layer_idx} attn_norm: {e:?}")))?; // [T, hidden] + psync("attn_norm")?; + + // g_proj consumes the SAME pre-attention normed input as q/k/v, exactly + // like decode_layer. FFAI_LAGUNA_MICRO=1 (Fused): one gemm_q4_mpp on the + // concatenated [wq;wk;wv;wgate] weight, then COLUMN-range extraction + // (`ops::extract_cols`, NOT `ops::slice`, which only slices a flat 1-D + // buffer; the fused output here is `[T, n_rows_total]` row-major, so each + // sub-projection is a strided column range, not a contiguous byte range). + let (q, k, v, gate_raw) = if let (Some(zero_idx), Some(qkvm)) = (marlin_zero_idx, w.attn.qkv_marlin.as_ref()) { + // FFAI_LAGUNA_MARLIN=1: tensor-core Marlin GEMM on the [wq;wk;wv] + // concat instead of gemm_q4_mpp's scattered-nibble-read path. Marlin + // consumes/emits f16. `extract_cols` is now f16-native, so extract + // straight out of `qkv_f16` instead of casting the whole [T, + // n_rows_qkv] result to f32 first, the extraction itself then moves + // half the bytes (f16 in/out vs f32 in/out). Every one of q/k/v + // still ends up cast to f32 individually right after, though: + // q/k RMSNorm reads `w.attn.q_norm`/`k_norm`, which are loaded + // f32-only (no f16 weight copy exists, see `AttnWeights`), + // `rope_yarn_partial_many`'s raw-CUDA kernel is hardcoded `float*` + // (f32 only, see its doc), and the persistent KV cache + // (`LagunaKvCache`/`LagunaPrefillScratch`) is allocated f32-only, + // so kv_append_many's destination forces v to f32 too. Net effect: + // same total cast bytes as before (one f32-only chain either way), + // but the extract itself is now cheaper, and it's one whole-buffer + // cast fewer to schedule. + // g_proj stays on gemm_q4_mpp, its 48/72 rows are never 64-aligned + //, reading the standalone `g_proj` Q4Dense built for exactly this + // (independent of whether `proj` above resolved to Split or Fused). + let h_f16 = ops::cast_f32_f16(dev, &h) + .map_err(|e| Error::Msg(format!("prefill layer {layer_idx} marlin qkv cast_f16: {e:?}")))?; + let qkv_f16 = qkvm.mat.call(dev, &h_f16, zero_idx, t) + .map_err(|e| Error::Msg(format!("prefill layer {layer_idx} marlin qkv gemm: {e:?}")))?; + psync("qkv_proj_marlin")?; + let n_rows_qkv = qkvm.n_q_rows + 2 * qkvm.n_kv_rows; + let q_f16 = ops::extract_cols(dev, &qkv_f16, t, n_rows_qkv, 0, qkvm.n_q_rows)?; + let k_f16 = ops::extract_cols(dev, &qkv_f16, t, n_rows_qkv, qkvm.n_q_rows, qkvm.n_kv_rows)?; + let v_f16 = ops::extract_cols(dev, &qkv_f16, t, n_rows_qkv, qkvm.n_q_rows + qkvm.n_kv_rows, qkvm.n_kv_rows)?; + let q = ops::cast_f16_f32(dev, &q_f16) + .map_err(|e| Error::Msg(format!("prefill layer {layer_idx} marlin q cast_f32: {e:?}")))?; + let k = ops::cast_f16_f32(dev, &k_f16) + .map_err(|e| Error::Msg(format!("prefill layer {layer_idx} marlin k cast_f32: {e:?}")))?; + let v = ops::cast_f16_f32(dev, &v_f16) + .map_err(|e| Error::Msg(format!("prefill layer {layer_idx} marlin v cast_f32: {e:?}")))?; + let g_proj = w.attn.g_proj.as_ref().expect( + "AttnWeights::qkv_marlin and ::g_proj are both built together, gated on the same \ + FFAI_LAGUNA_MARLIN load-time read, see load_gguf", + ); + let gate_raw = ops::gemm_q4_mpp(dev, &h, &g_proj.qs, &g_proj.scales, t, g_proj.m, g_proj.k) + .map_err(|e| Error::Msg(format!("prefill layer {layer_idx} g_proj gemm: {e:?}")))?; + (q, k, v, gate_raw) + } else { + match &w.attn.proj { + AttnProj::Fused { wqkvg, n_q_rows, n_kv_rows, n_g_rows } => { + let n_rows_total = n_q_rows + 2 * n_kv_rows + n_g_rows; + let qkvg = ops::gemm_q4_mpp(dev, &h, &wqkvg.qs, &wqkvg.scales, t, n_rows_total, hidden)?; // [T, n_rows_total] + let q = ops::extract_cols(dev, &qkvg, t, n_rows_total, 0, *n_q_rows)?; + let k = ops::extract_cols(dev, &qkvg, t, n_rows_total, *n_q_rows, *n_kv_rows)?; + let v = ops::extract_cols(dev, &qkvg, t, n_rows_total, n_q_rows + n_kv_rows, *n_kv_rows)?; + let gate_raw = ops::extract_cols(dev, &qkvg, t, n_rows_total, n_q_rows + 2 * n_kv_rows, *n_g_rows)?; + (q, k, v, gate_raw) + } + AttnProj::Split { wq, wk, wv, wgate } => { + let q = ops::gemm_q4_mpp(dev, &h, &wq.qs, &wq.scales, t, wq.m, wq.k) + .map_err(|e| Error::Msg(format!("prefill layer {layer_idx} q_proj gemm: {e:?}")))?; + psync("q_proj")?; + let k = ops::gemm_q4_mpp(dev, &h, &wk.qs, &wk.scales, t, wk.m, wk.k) + .map_err(|e| Error::Msg(format!("prefill layer {layer_idx} k_proj gemm: {e:?}")))?; + let v = ops::gemm_q4_mpp(dev, &h, &wv.qs, &wv.scales, t, wv.m, wv.k) + .map_err(|e| Error::Msg(format!("prefill layer {layer_idx} v_proj gemm: {e:?}")))?; + let gate_raw = ops::gemm_q4_mpp(dev, &h, &wgate.qs, &wgate.scales, t, wgate.m, wgate.k) + .map_err(|e| Error::Msg(format!("prefill layer {layer_idx} g_proj gemm: {e:?}")))?; // [T, n_head] + (q, k, v, gate_raw) + } + } + }; + + // Per-head QK RMSNorm (rms_norm normalizes the last dim, so a + // [T*heads, head_dim] view is correct per (token,head)), then batched + // RoPE, YaRN on full layers, plain on sliding (`cfg.rope_for` already + // picked the right params for each, same as decode_layer). + psync("qkvg_proj")?; + let q2 = ops::rms_norm(dev, &q.reshaped(vec![t * n_head, hd]), &w.attn.q_norm, cfg.eps) + .map_err(|e| Error::Msg(format!("prefill layer {layer_idx} q_norm: {e:?}")))?; + let k2 = ops::rms_norm(dev, &k.reshaped(vec![t * n_kv_heads, hd]), &w.attn.k_norm, cfg.eps) + .map_err(|e| Error::Msg(format!("prefill layer {layer_idx} k_norm: {e:?}")))?; + psync("qk_norm")?; + let q3 = ops::rope_yarn_partial_many( + dev, &q2.reshaped(vec![t, n_head, hd]), positions, n_head, + rope.n_rot, rope.theta, rope.freq_scale, rope.ext_factor, rope.attn_factor, + rope.beta_fast, rope.beta_slow, rope.n_orig_ctx, + ).map_err(|e| Error::Msg(format!("prefill layer {layer_idx} q_rope: {e:?}")))?; + psync("q_rope")?; + let k3 = ops::rope_yarn_partial_many( + dev, &k2.reshaped(vec![t, n_kv_heads, hd]), positions, n_kv_heads, + rope.n_rot, rope.theta, rope.freq_scale, rope.ext_factor, rope.attn_factor, + rope.beta_fast, rope.beta_slow, rope.n_orig_ctx, + ).map_err(|e| Error::Msg(format!("prefill layer {layer_idx} k_rope: {e:?}")))?; + psync("k_rope")?; + let v3 = v.reshaped(vec![t, n_kv_heads, hd]); + + // KV cache write + attention. Full layers grow the persistent cache + // directly (safe: nothing is ever evicted). Sliding layers write the + // LINEAR prefill scratch instead, see `LagunaPrefillScratch`'s doc for + // why the decode ring cannot be batch-written mid-prompt. + let attn = if is_full { + let slot = &kv.slots[layer_idx]; + ops::kv_append_many(dev, &k3, positions, &slot.k, n_kv_heads, hd, slot.cap) + .map_err(|e| Error::Msg(format!("prefill layer {layer_idx} kv_append k: {e:?}")))?; + ops::kv_append_many(dev, &v3, positions, &slot.v, n_kv_heads, hd, slot.cap) + .map_err(|e| Error::Msg(format!("prefill layer {layer_idx} kv_append v: {e:?}")))?; + psync("kv_write")?; + // Every query attends the full causal prefix: seg_lo = 0 for all rows. + let seg_lo = Tensor::new(dev.alloc_zeroed(t * 4)?, vec![t], DType::U32); + // Full-attention layer: no sliding window, so the block skip stays + // off (win=0), every KV block is reachable from every query row. + ops::sdpa_multi_tc_varlen( + dev, &q3, &slot.k, &slot.v, &seg_lo, 0, hd, n_head as u32, + start as u32, t as u32, slot.cap as u32, heads_per_group, true, cfg.attn_scale(), 0, + ).map_err(|e| Error::Msg(format!("prefill layer {layer_idx} sdpa_varlen_1: {e:?}")))? + } else { + let slot = scratch.slots[layer_idx].as_ref().ok_or_else(|| { + Error::Msg(format!("prefill_layer: layer {layer_idx} is sliding-window but scratch has no slot for it")) + })?; + ops::kv_append_many(dev, &k3, positions, &slot.k, n_kv_heads, hd, slot.cap) + .map_err(|e| Error::Msg(format!("prefill layer {layer_idx} kv_append k: {e:?}")))?; + ops::kv_append_many(dev, &v3, positions, &slot.v, n_kv_heads, hd, slot.cap) + .map_err(|e| Error::Msg(format!("prefill layer {layer_idx} kv_append v: {e:?}")))?; + psync("kv_write")?; + // Per-row sliding window: query at absolute position p attends + // [max(0, p - (window-1)), p]. All values are >= 0, so a u32 tensor + // (the convention every other seg_lo caller in this codebase uses) + // carries them losslessly even though the raw-CUDA kernel reads them + // as `int`. + let window = cfg.sliding_window as i64; + let seg_lo_h: Vec = (0..t) + .map(|i| ((start + i) as i64 - (window - 1)).max(0) as u32) + .collect(); + let seg_lo = Tensor::new(dev.upload(&u32s_to_bytes(&seg_lo_h))?, vec![t], DType::U32); + // Sliding-window layer: `seg_lo` above was built with exactly the + // `seg_lo[i] = max(0, (start+i) - (window-1))` convention the + // window-mode block skip expects, so pass the same window width as + // `win`, it prunes the KV blocks each chunk of query rows can't + // reach, on top of the (already-applied) per-element mask. + ops::sdpa_multi_tc_varlen( + dev, &q3, &slot.k, &slot.v, &seg_lo, 0, hd, n_head as u32, + start as u32, t as u32, slot.cap as u32, heads_per_group, true, cfg.attn_scale(), + cfg.sliding_window as u32, + ).map_err(|e| Error::Msg(format!("prefill layer {layer_idx} sdpa_varlen_0: {e:?}")))? + }; // [T, n_head, hd] + + // Per-head softplus gate, o_proj, residual. + psync("sdpa")?; + let gated = ops::gate_softplus_mul_perhead_many(dev, &attn, &gate_raw, n_head) + .map_err(|e| Error::Msg(format!("prefill layer {layer_idx} gate_many: {e:?}")))?; // [T, n_head, hd] + psync("gate")?; + let gated_flat = gated.reshaped(vec![t, n_head * hd]); + let o = if let (Some(zero_idx), Some(wo_m)) = (marlin_zero_idx, w.attn.wo_marlin.as_ref()) { + let gated_f16 = ops::cast_f32_f16(dev, &gated_flat) + .map_err(|e| Error::Msg(format!("prefill layer {layer_idx} marlin o_proj cast_f16: {e:?}")))?; + let o_f16 = wo_m.call(dev, &gated_f16, zero_idx, t) + .map_err(|e| Error::Msg(format!("prefill layer {layer_idx} marlin o_proj gemm: {e:?}")))?; + ops::cast_f16_f32(dev, &o_f16) + .map_err(|e| Error::Msg(format!("prefill layer {layer_idx} marlin o_proj cast_f32: {e:?}")))? + } else { + ops::gemm_q4_mpp(dev, &gated_flat, &w.attn.wo.qs, &w.attn.wo.scales, t, w.attn.wo.m, w.attn.wo.k) + .map_err(|e| Error::Msg(format!("prefill layer {layer_idx} o_gemm: {e:?}")))? // [T, hidden] + }; + let x1 = ops::add(dev, &x.reshaped(vec![t * hidden]), &o.reshaped(vec![t * hidden]))?.reshaped(vec![t, hidden]); + psync("o_proj_res")?; + + // ── FFN ────────────────────────────────────────────────────────── + let x2 = match &w.ffn { + LayerFfn::Dense(f) => { + let h2 = ops::rms_norm(dev, &x1, &f.ffn_norm, cfg.eps)?; // [T, hidden] + let down = if let (Some(zero_idx), Some(gum), Some(down_m)) = + (marlin_zero_idx, f.gateup_marlin.as_ref(), f.down_marlin.as_ref()) + { + // FFAI_LAGUNA_MARLIN=1: [gate;up] concat as one Marlin GEMM, + // split + SwiGLU straight in f16 (extract_cols is f16-native + // and mt_swiglu is registered for f16 in metaltile-std), then + // a second Marlin GEMM for down, `down_m.call` wants f16 + // input directly. Unlike the attention QKV split above, this + // chain has no f32-only consumer in the middle (no RMSNorm + // weight, no rope, no f32 KV cache write), so the old + // cast_f16_f32(gateup)/cast_f32_f16(act) pair is pure + // overhead and is removed entirely; only `h2`'s initial + // f32->f16 cast (Marlin's input requirement) and the final + // f16->f32 cast (the f32 residual stream) remain. + let h2_f16 = ops::cast_f32_f16(dev, &h2)?; + let gateup_f16 = gum.mat.call(dev, &h2_f16, zero_idx, t)?; // [T, 2*n_ff] f16 + let gate = ops::extract_cols(dev, &gateup_f16, t, 2 * gum.n_ff, 0, gum.n_ff)?; + let up = ops::extract_cols(dev, &gateup_f16, t, 2 * gum.n_ff, gum.n_ff, gum.n_ff)?; + let act_f16 = ops::swiglu(dev, &gate, &up)?; // [T, n_ff] f16 + let down_f16 = down_m.call(dev, &act_f16, zero_idx, t)?; // [T, hidden] f16 + ops::cast_f16_f32(dev, &down_f16)? + } else { + let gate = ops::gemm_q4_mpp(dev, &h2, &f.gate.qs, &f.gate.scales, t, f.gate.m, f.gate.k)?; + let up = ops::gemm_q4_mpp(dev, &h2, &f.up.qs, &f.up.scales, t, f.up.m, f.up.k)?; + let act = ops::swiglu(dev, &gate, &up)?; + ops::gemm_q4_mpp(dev, &act, &f.down.qs, &f.down.scales, t, f.down.m, f.down.k)? + }; + ops::add(dev, &x1.reshaped(vec![t * hidden]), &down.reshaped(vec![t * hidden]))?.reshaped(vec![t, hidden]) + } + LayerFfn::Moe(f) => { + let h2 = ops::rms_norm(dev, &x1, &f.ffn_norm, cfg.eps)?; // [T, hidden] + + // Batched sigmoid router + on-device counting-sort: the SAME + // sigmoid+bias+top_k math as decode's `moe_router_device`, over + // all T rows in one pass. `logits` is the small dense f32 router + // GEMM (router stays f32, precision-sensitive); route+sort is + // fully on device (no host sync, mirrors the NemotronH + // NEMOTRON_DEVSORT/devdesc prefill pattern). + let logits = ops::matmul(dev, &f.router, &h2) + .map_err(|e| Error::Msg(format!("prefill layer {layer_idx} router_matmul: {e:?}")))?; // [T, n_expert] + psync("moe_router_logits")?; + let (sorted_tok, sorted_wt, offsets) = ops::moe_route_sort_device( + dev, &logits, &f.bias, cfg.n_expert, cfg.n_expert_used, cfg.routed_scaling, + )?; + let mt = t * cfg.n_expert_used; + if std::env::var("FFAI_LAGUNA_PREFILL_SYNC").is_ok() { + let _ = dev.synchronize(); + let mut stb = vec![0u8; mt * 4]; + let _ = dev.download(sorted_tok.buffer.as_ref(), &mut stb); + let st: Vec = stb.chunks_exact(4).map(|c| u32::from_le_bytes(c.try_into().unwrap())).collect(); + let mut ofb = vec![0u8; (cfg.n_expert + 1) * 4]; + let _ = dev.download(offsets.buffer.as_ref(), &mut ofb); + let of: Vec = ofb.chunks_exact(4).map(|c| u32::from_le_bytes(c.try_into().unwrap())).collect(); + eprintln!(" prefill moe layer {layer_idx}: sorted_tok={st:?} offsets_last={} (expect {mt})", of[cfg.n_expert]); + } + + // Gather the sorted-order activation (f16, the grouped-MMA + // input dtype) once, then TWO grouped Q4 GEMMs (gate stack, up + // stack) over the sorted rows, elementwise SwiGLU, one grouped + // Q4 GEMM down. `moe_q4_grouped_mma_dev` builds its tile + // descriptors ON DEVICE from `offsets`, no host g_starts/ + // expert_ids vectors, no host sync. + // + // NOTE ON DTYPE (RoutedGateUp::Fused branch below): the down + // grouped-GEMM's raw CUDA kernel reads its activation input as + // `const __half*` unconditionally (see `MOE_Q4_GROUPED_MMA_SRC` + // in ffai-ops), f16 in, no f32 path exists. Combined with + // `moe_q4_grouped_mma_dev` always emitting f16 and `extract_cols` + // now being f16-native, the gate/up split + SwiGLU chain has NO + // f32-only consumer anywhere in it, so it runs start to finish in + // f16 with zero casts (the old cast_f16_f32/cast_f32_f16 pair + // existed only because `extract_cols`/`ops::swiglu` used to + // require f32 input; both now take f16 directly). + psync("moe_route_sort")?; + let h2_f16 = ops::cast_f32_f16(dev, &h2) + .map_err(|e| Error::Msg(format!("prefill layer {layer_idx} cast_f16: {e:?}")))?; // [T, hidden] + let a = ops::gather(dev, &h2_f16, &sorted_tok) + .map_err(|e| Error::Msg(format!("prefill layer {layer_idx} moe_gather_rows: {e:?}")))?; // [mt, hidden] f16, sorted-token order + + // FFAI_LAGUNA_W4A8=1: route the routed gate/up/down grouped + // GEMMs through the W4A8 grouped-MoE GEMM (fp8 e4m3 activations + // x mxfp4 weights, cutlass mxf8f6f4) instead of + // `moe_q4_grouped_mma_dev`'s W4A16 path, see `MoeFfnW4a8`'s doc + // for the weight side. `f.w4a8` was decided once at load time + // (`laguna_w4a8_enabled`'s doc), so this just matches the + // `Option` rather than re-reading the env. + // + // Unlike `moe_q4_grouped_mma_dev` (device-built tile + // descriptors, zero host sync: see the NOTE ON DTYPE below), + // `ops::w4a8_sfa_offsets`/`w4a8_actquant_grouped`/ + // `moe_grouped_gemm_w4a8` all take HOST `g_starts`/`expert_ids` + // (mirroring their ffai-modeltests NEMOTRON_W4A8_MOE call site + // exactly, no device-side descriptor-building variant of this + // GEMM exists in ffai-ops today). So this branch pays ONE small + // device->host download per layer (`offsets` is `n_expert+1` + // u32s, 1028 bytes for Laguna-S-2.1's 256 experts) to turn the + // on-device routing offsets into host group-starts, the one + // host sync this batched-prefill path otherwise avoids + // entirely (decode-era work eliminated exactly this kind of + // per-layer host round-trip for the W4A16 grouped path; W4A8 + // reintroduces a tiny one because `ops::w4a8_*` has no + // on-device-descriptor sibling yet). `Device::download` blocks + // until prior queued work completes, so no separate + // `dev.synchronize()` call is needed first (same convention as + // the NEMOTRON_W4A8_MOE call site). + let down_out = if let Some(w4a8) = f.w4a8.as_ref() { + let mut ofb = vec![0u8; (cfg.n_expert + 1) * 4]; + dev.download(offsets.buffer.as_ref(), &mut ofb) + .map_err(|e| Error::Msg(format!("prefill layer {layer_idx} w4a8 offsets download: {e:?}")))?; + let offs: Vec = ofb.chunks_exact(4) + .map(|c| u32::from_le_bytes(c.try_into().unwrap()) as usize).collect(); + let mut g_starts: Vec = vec![0]; + let mut expert_ids: Vec = Vec::new(); + for e in 0..cfg.n_expert { + if offs[e + 1] > offs[e] { + expert_ids.push(e); + g_starts.push(offs[e + 1]); + } + } + let group_rows: Vec = g_starts.windows(2).map(|w| (w[1] - w[0]) as i32).collect(); + psync("w4a8_offsets_dl")?; + + // gate/up both read the same pre-FFN-norm activation, so + // they share `k = hidden` and therefore the same `kpad`, + // grab it off whichever stack this layer built + // (`RoutedGateUpW4a8::Fused`'s single stack, or `::Split`'s + // `gate` stack; both give the identical padded width). + let hid_pad = match &w4a8.gateup { + RoutedGateUpW4a8::Fused(s) => s.kpad, + RoutedGateUpW4a8::Split { gate, .. } => gate.kpad, + }; + let a_pad = ops::pad_rows_f16(dev, &a, mt, cfg.hidden, hid_pad) + .map_err(|e| Error::Msg(format!("prefill layer {layer_idx} w4a8 gateup pad: {e:?}")))?; + let (uoff, ubytes) = ops::w4a8_sfa_offsets(&g_starts, hid_pad); + let (a_q, a_sf) = ops::w4a8_actquant_grouped(dev, &a_pad, &group_rows, &uoff, ubytes, mt, hid_pad) + .map_err(|e| Error::Msg(format!("prefill layer {layer_idx} w4a8 gateup actquant: {e:?}")))?; + psync("w4a8_gateup_actquant")?; + + // Same fused-vs-split shape as the W4A16 branch below: ONE + // grouped GEMM to `2*n_ff_exp` for `Fused`, TWO separate + // grouped GEMMs (sharing the same quantized activation) for + // `Split`, then the identical SwiGLU either way. + let act = match &w4a8.gateup { + RoutedGateUpW4a8::Fused(gu) => { + let n_ff = cfg.n_ff_exp; + let gateup_out = ops::moe_grouped_gemm_w4a8( + dev, &a_q, &a_sf, &uoff, &gu.packed, &gu.sf, gu.sfb_exp_bytes, + &g_starts, &expert_ids, 2 * n_ff, hid_pad, + ).map_err(|e| Error::Msg(format!("prefill layer {layer_idx} w4a8 gateup gemm: {e:?}")))?; // [mt, 2*n_ff_exp] f16 + psync("w4a8_gateup_gemm")?; + let gate_out = ops::extract_cols(dev, &gateup_out, mt, 2 * n_ff, 0, n_ff)?; + let up_out = ops::extract_cols(dev, &gateup_out, mt, 2 * n_ff, n_ff, n_ff)?; + ops::swiglu(dev, &gate_out, &up_out) + .map_err(|e| Error::Msg(format!("prefill layer {layer_idx} w4a8 gateup swiglu: {e:?}")))? // [mt, n_ff_exp] f16 + } + RoutedGateUpW4a8::Split { gate, up } => { + let gate_out = ops::moe_grouped_gemm_w4a8( + dev, &a_q, &a_sf, &uoff, &gate.packed, &gate.sf, gate.sfb_exp_bytes, + &g_starts, &expert_ids, cfg.n_ff_exp, hid_pad, + ).map_err(|e| Error::Msg(format!("prefill layer {layer_idx} w4a8 gate gemm: {e:?}")))?; // [mt, n_ff_exp] f16 + let up_out = ops::moe_grouped_gemm_w4a8( + dev, &a_q, &a_sf, &uoff, &up.packed, &up.sf, up.sfb_exp_bytes, + &g_starts, &expert_ids, cfg.n_ff_exp, hid_pad, + ).map_err(|e| Error::Msg(format!("prefill layer {layer_idx} w4a8 up gemm: {e:?}")))?; // [mt, n_ff_exp] f16 + psync("w4a8_gateup_gemm")?; + ops::swiglu(dev, &gate_out, &up_out) + .map_err(|e| Error::Msg(format!("prefill layer {layer_idx} w4a8 gateup swiglu: {e:?}")))? // [mt, n_ff_exp] f16 + } + }; + + let inter_pad = w4a8.down.kpad; + let act_pad = ops::pad_rows_f16(dev, &act, mt, cfg.n_ff_exp, inter_pad) + .map_err(|e| Error::Msg(format!("prefill layer {layer_idx} w4a8 down pad: {e:?}")))?; + let (doff, dbytes) = ops::w4a8_sfa_offsets(&g_starts, inter_pad); + let (d_q, d_sf) = ops::w4a8_actquant_grouped(dev, &act_pad, &group_rows, &doff, dbytes, mt, inter_pad) + .map_err(|e| Error::Msg(format!("prefill layer {layer_idx} w4a8 down actquant: {e:?}")))?; + psync("w4a8_down_actquant")?; + let down_out = ops::moe_grouped_gemm_w4a8( + dev, &d_q, &d_sf, &doff, &w4a8.down.packed, &w4a8.down.sf, w4a8.down.sfb_exp_bytes, + &g_starts, &expert_ids, cfg.hidden, inter_pad, + ).map_err(|e| Error::Msg(format!("prefill layer {layer_idx} w4a8 down gemm: {e:?}")))?; // [mt, hidden] f16 + down_out + } else { + // FFAI_LAGUNA_MOEFUSE=1 (RoutedGateUp::Fused): ONE grouped Q4 GEMM + // with n_out = 2*n_ff_exp over the fused [gate;up] stack instead + // of two separate grouped GEMMs, halving the grouped-GEMM call + // count (each call is itself two kernel launches, an on-device + // tile-descriptor build plus the GEMM, so this also halves the + // launch count). See the NOTE ON DTYPE below: the whole + // split+SwiGLU chain runs in f16 with no cast, mirroring the + // shared-expert Marlin [gate;up] split above. + let act = match &f.gateup { + RoutedGateUp::Fused(gu) => { + let n_ff = cfg.n_ff_exp; + let gateup_out = ops::moe_q4_grouped_mma_dev( + dev, &a, &gu.qs, &gu.scales, &offsets, cfg.n_expert, mt, 2 * n_ff, cfg.hidden, + ) + .map_err(|e| Error::Msg(format!("prefill layer {layer_idx} moe_gateup_gemm_fused: {e:?}")))?; // [mt, 2*n_ff_exp] f16 + psync("moe_gateup_gemm")?; + // f16 straight through, see the NOTE ON DTYPE below. + let gate_out = ops::extract_cols(dev, &gateup_out, mt, 2 * n_ff, 0, n_ff)?; + let up_out = ops::extract_cols(dev, &gateup_out, mt, 2 * n_ff, n_ff, n_ff)?; + ops::swiglu(dev, &gate_out, &up_out) + .map_err(|e| Error::Msg(format!("prefill layer {layer_idx} moe_act_swiglu: {e:?}")))? // [mt, n_ff_exp] f16 + } + RoutedGateUp::Split { gate, up } => { + let gate_out = ops::moe_q4_grouped_mma_dev( + dev, &a, &gate.qs, &gate.scales, &offsets, cfg.n_expert, mt, cfg.n_ff_exp, cfg.hidden, + )?; // [mt, n_ff_exp] f16 + let up_out = ops::moe_q4_grouped_mma_dev( + dev, &a, &up.qs, &up.scales, &offsets, cfg.n_expert, mt, cfg.n_ff_exp, cfg.hidden, + )?; // [mt, n_ff_exp] f16 + psync("moe_gateup_gemm")?; + ops::swiglu(dev, &gate_out, &up_out)? // [mt, n_ff_exp] f16 + } + }; + let down_out = ops::moe_q4_grouped_mma_dev( + dev, &act, &f.down.qs, &f.down.scales, &offsets, cfg.n_expert, mt, cfg.hidden, cfg.n_ff_exp, + )?; // [mt, hidden] f16 + down_out + }; + + // Weighted scatter-add back to token order: acc[t,h] = sum over + // this token's top_k sorted rows of down_out[r,h] * sorted_wt[r] + // (router weight already includes routed_scaling, no extra + // `unscale` factor needed here, unlike Nemotron's ReLU²/256 trick). + let acc = Tensor::new(dev.alloc_zeroed(t * cfg.hidden * 4)?, vec![t, cfg.hidden], DType::F32); + psync("moe_down_gemm")?; + ops::moe_scatter_add_det_dev( + dev, &down_out, &sorted_tok, &sorted_wt, &acc, t, mt, cfg.hidden, 1.0, true, + )?; + + // Always-on shared expert: same batched gate→up→SwiGLU→down as a + // routed expert, just over ALL T rows unweighted (decode's + // unweighted `gemv_accum(scale=1)` equivalent). Handles both + // weight layouts (FFAI_LAGUNA_MICRO=1's single-expert + // Q4ExpertStack fusion, or the plain Q4Dense split) uniformly by + // pulling a (qs, scales, m, k) view for gate/up either way, + // `expert_view(0)`'s byte offset is always zero for a 1-expert + // stack, so it's safe to feed into `gemm_q4_mpp` even though that + // op's `Binding::Buffer` ignores a non-zero `Tensor.offset`. + psync("moe_scatter")?; + let sdown_out = if let (Some(zero_idx), Some(sm)) = (marlin_zero_idx, f.shared_marlin.as_ref()) { + // FFAI_LAGUNA_MARLIN=1: same [gate;up]-concat-then-split + // Marlin pattern as the dense-layer FFN above. `h2_f16` was + // already computed for the routed-expert gather just above, + // reuse it instead of re-casting. Split + SwiGLU stay in f16 + // end to end (extract_cols f16-native, mt_swiglu f16-native, + // `sm.down.call` wants f16 in), the old cast_f16_f32(gateup) + // / cast_f32_f16(sact) pair is removed; only the final + // f16->f32 cast remains, needed because `sdown_out` is added + // into the f32 MoE accumulator (`acc2`) below. + let gateup_f16 = sm.gateup.mat.call(dev, &h2_f16, zero_idx, t)?; // [T, 2*n_ff_shexp] f16 + let sgate = ops::extract_cols(dev, &gateup_f16, t, 2 * sm.gateup.n_ff, 0, sm.gateup.n_ff)?; + let sup = ops::extract_cols(dev, &gateup_f16, t, 2 * sm.gateup.n_ff, sm.gateup.n_ff, sm.gateup.n_ff)?; + let sact_f16 = ops::swiglu(dev, &sgate, &sup)?; // [T, n_ff_shexp] f16 + let sdown_f16 = sm.down.call(dev, &sact_f16, zero_idx, t)?; // [T, hidden] f16 + ops::cast_f16_f32(dev, &sdown_f16)? // [T, hidden] + } else { + let (sg_qs, sg_sc, sg_m, sg_k, su_qs, su_sc, su_m, su_k, sdown) = match &f.shared_expert { + SharedExpert::Split { gate, up, down } => ( + gate.qs.clone(), gate.scales.clone(), gate.m, gate.k, + up.qs.clone(), up.scales.clone(), up.m, up.k, + down, + ), + SharedExpert::Fused(se) => { + let (gqs, gsc) = se.gate.expert_view(0); + let (uqs, usc) = se.up.expert_view(0); + (gqs, gsc, se.gate.m, se.gate.k, uqs, usc, se.up.m, se.up.k, &se.down) + } + }; + let sgate = ops::gemm_q4_mpp(dev, &h2, &sg_qs, &sg_sc, t, sg_m, sg_k)?; + let sup = ops::gemm_q4_mpp(dev, &h2, &su_qs, &su_sc, t, su_m, su_k)?; + let sact = ops::swiglu(dev, &sgate, &sup)?; + ops::gemm_q4_mpp(dev, &sact, &sdown.qs, &sdown.scales, t, sdown.m, sdown.k)? // [T, hidden] + }; + let acc2 = ops::add(dev, &acc.reshaped(vec![t * cfg.hidden]), &sdown_out.reshaped(vec![t * cfg.hidden]))?; + + ops::add(dev, &x1.reshaped(vec![t * hidden]), &acc2)?.reshaped(vec![t, hidden]) + } + }; + + Ok(x2) +} + +/// Batched (multi-token) **prefill**: process the whole `tokens` prompt +/// `chunk` tokens at a time (default recommendation: 1024), writing every +/// token's K/V into `kv`/`scratch`, and return the LAST token's next-token +/// logits, the prefill counterpart of [`decode_step`]'s per-token forward. +/// Each chunk's `T` rows run through every layer as batched GEMMs +/// ([`ops::gemm_q4_mpp`], the grouped-MMA MoE pipeline, `sdpa_multi_tc_varlen`) +/// instead of `T` sequential [`decode_step`] calls. +/// +/// `scratch` must have been sized (via [`LagunaPrefillScratch::new`]) for AT +/// LEAST `tokens.len()`. `kv` must have been sized (via +/// [`LagunaKvCache::new`]) for at least `tokens.len()` plus however many +/// decode tokens the caller intends to generate afterward, same contract as +/// feeding the same prompt through `decode_step` one token at a time. +/// +/// After the LAST chunk, this compacts each sliding-window layer's linear +/// prefill scratch (the last `min(sliding_window, tokens.len())` positions) +/// into `kv`'s decode ring, so `decode_step` / `LagunaDecodeCtx` can continue +/// from `pos = tokens.len()` and produce the IDENTICAL greedy continuation a +/// pure per-token `decode_step` loop over the same prompt would. +/// +/// `chunk` only bounds how many prompt tokens are embedded/projected/attended +/// in one forward pass (memory for the batched intermediates scales with +/// it); correctness is independent of the choice. +pub fn prefill( + dev: &dyn Device, + model: &LagunaModel, + kv: &LagunaKvCache, + scratch: &LagunaPrefillScratch, + tokens: &[u32], + chunk: usize, +) -> Result { + let cfg = &model.cfg; + let total = tokens.len(); + if total == 0 { + return Err(Error::Msg("prefill: empty prompt".into())); + } + if total > scratch.prompt_cap { + return Err(Error::Msg(format!( + "prefill: prompt len {total} exceeds scratch prompt_cap {}", scratch.prompt_cap + ))); + } + let chunk = chunk.max(1); + + // FFAI_LAGUNA_MARLIN=1: ONE persistent all-zero `u32` indices buffer, + // sized to `chunk` (every chunk's row count `t = chunk.min(total-start)` + // is `<= chunk`), reused by every `MarlinDense::call` in every layer of + // every chunk this `prefill()` invocation processes. The Marlin GEMM's + // grouped-MoE ABI wants a per-row expert id; a dense `[n_out, k_in]` + // matrix is fed as a single "expert 0", so every row's id is 0: see + // `MarlinDense`'s doc. Built even if a given layer's per-weight + // `Option` ends up `None` (dimension not 64-aligned); the + // cost is one small zeroed alloc, and `prefill_layer` only reads through + // this when a `Some(MarlinDense)` is also present. + let marlin_zero_idx = if laguna_marlin_enabled() { + Some(Tensor::new(dev.alloc_zeroed(chunk * 4)?, vec![chunk], DType::U32)) + } else { + None + }; + + let mut x_last: Option = None; + let mut t_last = 0usize; + let mut start = 0usize; + while start < total { + let t = chunk.min(total - start); + let chunk_tokens = &tokens[start..start + t]; + + let ids = Tensor::new(dev.upload(&u32s_to_bytes(chunk_tokens))?, vec![t], DType::U32); + let mut x = ops::gather(dev, &model.tok_embd, &ids) + .map_err(|e| Error::Msg(format!("prefill embed gather: {e:?}")))? + .reshaped(vec![t, cfg.hidden]); + + let pos_vec: Vec = (0..t).map(|i| (start + i) as u32).collect(); + let positions = Tensor::new(dev.upload(&u32s_to_bytes(&pos_vec))?, vec![t], DType::U32); + + if std::env::var("FFAI_LAGUNA_PREFILL_SYNC").is_ok() { + dev.synchronize().map_err(|e| Error::Msg(format!("prefill fault pre-layers (embed/positions): {e:?}")))?; + } + for (i, layer) in model.layers.iter().enumerate() { + x = prefill_layer(dev, cfg, i, layer, kv, scratch, &x, t, start, &positions, marlin_zero_idx.as_ref())? + .into(); + if std::env::var("FFAI_LAGUNA_PREFILL_SYNC").is_ok() { + dev.synchronize().map_err(|e| Error::Msg(format!("prefill fault after layer {i}: {e:?}")))?; + } + } + + x_last = Some(x); + t_last = t; + start += t; + } + let x = x_last.expect("prefill: total > 0 was checked above, so at least one chunk always runs"); + + // Final norm + q8 lm_head on the LAST token of the LAST chunk only (the + // next-token logits), cheaper than decode's reference-engine-style "norm all rows + // then slice" since RMSNorm is per-row-independent anyway. + let last_row = ops::slice(dev, &x.reshaped(vec![t_last * cfg.hidden]), (t_last - 1) * cfg.hidden, cfg.hidden)?; + let xn = ops::rms_norm(dev, &last_row, &model.out_norm, cfg.eps)?; + let logits = ops::gemv_q8( + dev, &model.lm_head.qs, &model.lm_head.scales, &xn, model.lm_head.m, model.lm_head.k, model.lm_head.m, + )?; + + // Compact each sliding-window layer's linear prefill scratch into the + // persistent decode ring, one-time cost, so decode can continue from + // `pos = total` exactly as it does today. + let copy_len = total.min(cfg.sliding_window); + let copy_start = total - copy_len; + for i in 0..cfg.n_layers { + if cfg.is_full_layer(i) { + continue; + } + let sslot = scratch.slots[i].as_ref().expect("sliding layer always has a scratch slot"); + let dslot = &kv.slots[i]; + ops::kv_compact_ring(dev, &sslot.k, &dslot.k, cfg.n_kv_heads, sslot.cap, cfg.head_dim, dslot.cap, copy_start, copy_len)?; + ops::kv_compact_ring(dev, &sslot.v, &dslot.v, cfg.n_kv_heads, sslot.cap, cfg.head_dim, dslot.cap, copy_start, copy_len)?; + } + + Ok(logits) +} + /// Persistent per-step device scalars decode reads instead of uploading a /// fresh 1-element buffer every call. Allocated ONCE (in [`DecodeWorkspace::new`]) -/// and refreshed in place every step via [`DecodeWorkspace::update`] — the +/// and refreshed in place every step via [`DecodeWorkspace::update`], the /// refresh always happens BEFORE the (eager or captured-graph) decode work /// that reads them, and never inside a graph-capture region, so it's safe /// regardless of whether the caller is warming up, capturing, or replaying. /// /// `full_pos`/`sliding_pos` are decode_layer's existing posbuf inputs (KV -/// cache write index — full layers grow, sliding layers ring); `n_kv_full`/ +/// cache write index, full layers grow, sliding layers ring); `n_kv_full`/ /// `n_kv_sliding` are the NEW graph-mode buffers `sdpa_decode_nbuf` reads /// (`pos+1` / `min(pos+1, sliding_window)`). Eager (non-graph) decode only /// ever reads `token_id`/`full_pos`/`sliding_pos`; `n_kv_full`/`n_kv_sliding` @@ -961,7 +2769,7 @@ impl DecodeWorkspace { } /// Refresh every workspace buffer for the token at `pos`. Five tiny - /// (single-`u32`) raw-CUDA writes, in place — no reallocation, so the + /// (single-`u32`) raw-CUDA writes, in place: no reallocation, so the /// buffers' device pointers (what a captured graph binds) never move. /// MUST be called outside any graph-capture region. pub fn update(&self, dev: &dyn Device, cfg: &LagunaConfig, token_id: u32, pos: u32) -> Result<()> { @@ -989,7 +2797,7 @@ impl DecodeWorkspace { /// /// Reads `token_id`/RoPE-position/KV-length from `ws` rather than uploading /// fresh 1-element buffers every call (the former per-step host upload this -/// function used to do) — the caller must have already called +/// function used to do), the caller must have already called /// `ws.update(dev, &model.cfg, token_id, pos)` for this step. This is both a /// standalone eager-mode win (no per-step `cuMemAlloc`+H2D for those /// scalars) and what makes the whole step CUDA-graph-capturable: a graph @@ -998,14 +2806,14 @@ impl DecodeWorkspace { /// /// This is the path used both by plain eager decode (no env set) and by /// [`LagunaDecodeCtx`]'s warmup steps and one-time capture step; graph -/// REPLAY (after capture) does not call this function at all — it re-runs +/// REPLAY (after capture) does not call this function at all, it re-runs /// the already-recorded launches via `Device::graph_launch`. /// /// `FFAI_LAGUNA_GRAPH=1` switches `decode_layer`'s RoPE/SDPA to the /// buffer-driven (capture-safe) kernels; without it, behavior is byte-for-byte /// identical to before this refactor (same kernels, same scalar arguments, /// just reading `token_id`/`pos` off a persistent buffer instead of a fresh -/// one — the values are the same either way). +/// one, the values are the same either way). pub fn decode_step( dev: &dyn Device, model: &LagunaModel, @@ -1025,7 +2833,7 @@ pub fn decode_step( // Read once per decode_step, passed down as a bool: all three micro- // fusions (QKVG concat gemv, shared-expert fused swiglu gather, o_proj // residual accum) are gated on this single env var. Fixed for the - // process lifetime (same requirement CUDA-graph capture already has — + // process lifetime (same requirement CUDA-graph capture already has, // see GraphBufs's doc), so this per-step read always agrees with the // load-time read that picked AttnProj/SharedExpert's Fused vs Split. let micro = laguna_micro_enabled(); @@ -1033,7 +2841,7 @@ pub fn decode_step( // FFAI_LAGUNA_DEBUG=1: per-layer residual stats (L2 norm, NaN count) to // bisect a divergence to the first bad layer. Hot path unaffected when // the env is unset. Downloads mid-step, so it is INCOMPATIBLE with graph - // capture/replay — `LagunaDecodeCtx::step` refuses to run at all when + // capture/replay, `LagunaDecodeCtx::step` refuses to run at all when // this is set (see its doc); a bare `decode_step` call with both env // vars set would otherwise try to `dev.download` from inside a capture // region, which is forbidden by the CUDA driver. @@ -1063,24 +2871,24 @@ pub fn decode_step( /// kernel launches (48 layers x ~22 ops) into ONE `cuGraphLaunch` replay per /// token, removing host-launch-gap overhead. Enabled by running with /// `FFAI_LAGUNA_GRAPH=1` (also read by `decode_step`/`decode_layer` to pick -/// the buffer-driven RoPE/SDPA kernels — see their docs). +/// the buffer-driven RoPE/SDPA kernels, see their docs). /// /// Protocol (mirrors the NemotronH `NEMOTRON_GRAPH` precedent in /// `ffai-modeltests`): the first two `step()` calls run [`decode_step`] -/// eagerly ("warm up N=2 full steps normally" — the pool's bucket free-list +/// eagerly ("warm up N=2 full steps normally", the pool's bucket free-list /// reaches a deterministic steady state after the first call, since every /// step walks the identical fixed 48-layer op sequence regardless of /// `token_id`/`pos`, so the second call already allocates nothing new). The /// THIRD call captures: `Device::begin_capture` + one more `decode_step` + /// `Device::end_capture`. Capture only RECORDS the stream's kernel launches -/// — it does not execute them — so that third call's `decode_step` result is +/// (it does not execute them), so that third call's `decode_step` result is /// not yet real data; `step()` immediately follows the capture with one /// `Device::graph_launch` to actually compute that token's logits. Every /// call after that just refreshes the workspace and replays the graph. /// /// The logits `Tensor` returned from the capture step is kept alive for the /// lifetime of this `LagunaDecodeCtx` (never dropped) so its device buffer's -/// pointer can never be recycled by the allocator's pool — the same buffer +/// pointer can never be recycled by the allocator's pool, the same buffer /// is what every subsequent `graph_launch` writes into, and what `step()` /// returns (a cheap `Arc` clone) every time thereafter. /// @@ -1090,13 +2898,13 @@ pub fn decode_step( pub struct LagunaDecodeCtx { ws: DecodeWorkspace, /// Executable-graph handle from `Device::end_capture` (an opaque `u64`; - /// see the `Device` trait doc — CUDA's `CUgraphExec` cast through + /// see the `Device` trait doc, CUDA's `CUgraphExec` cast through /// `usize`). `None` until the capture step has run. graph: Option, /// The fixed output buffer graph replays write into (see the struct doc). logits: Option, /// Steps run so far via this ctx (warmup + capture; stops incrementing - /// once capture has happened — replay doesn't touch it). + /// once capture has happened, replay doesn't touch it). step_count: usize, } @@ -1135,7 +2943,7 @@ impl LagunaDecodeCtx { )); } - // Refresh the workspace BEFORE any capture/replay — always outside + // Refresh the workspace BEFORE any capture/replay, always outside // the captured region, whether we're warming up, about to capture, // or replaying an already-captured graph. self.ws.update(dev, &model.cfg, token_id, pos)?; @@ -1151,13 +2959,13 @@ impl LagunaDecodeCtx { self.step_count += 1; if self.step_count <= Self::WARMUP_STEPS { // Eager warmup: reaches pool steady state (deterministic alloc - // order — every step walks the same fixed layer sequence). + // order, every step walks the same fixed layer sequence). return decode_step(dev, model, kv, &self.ws, pos); } // step_count == WARMUP_STEPS + 1: capture this step's op sequence. // Capture only RECORDS launches; it does not execute them, so - // `logits` is not yet valid data — one `graph_launch` right after + // `logits` is not yet valid data, one `graph_launch` right after // `end_capture` actually computes it (against the workspace values // written above, so it's the correct result for THIS token). dev.begin_capture()?; diff --git a/rust/crates/ffai-modeltests/src/lib.rs b/rust/crates/ffai-modeltests/src/lib.rs index 0a253b20..d44efb25 100644 --- a/rust/crates/ffai-modeltests/src/lib.rs +++ b/rust/crates/ffai-modeltests/src/lib.rs @@ -4,7 +4,7 @@ //! //! Each `verify_*(dev: &dyn Device)` holds a model's forward and its HF oracle //! ONCE. The per-backend test files (`ffai-metal/tests/*`, `ffai-cuda/tests/*`) -//! are thin wrappers that build their device and call these — so a model's +//! are thin wrappers that build their device and call these, so a model's //! logic lives in exactly one place, not a Metal test + a sed'd CUDA twin. use ffai_core::{Backend, DType, Device, Tensor}; use ffai_loader::SafeTensors; @@ -27,7 +27,7 @@ fn model_dir(env: &str, hub: &str) -> Option { pub fn verify_gpt2(d: &dyn Device, plat: &str) { let dir = model_dir("GPT2_DIR", "models--gpt2").unwrap_or_default(); let path = format!("{dir}/model.safetensors"); - let Ok(st) = SafeTensors::open(&path) else { eprintln!("no model at {path} — skipping"); return; }; + let Ok(st) = SafeTensors::open(&path) else { eprintln!("no model at {path}, skipping"); return; }; let (hid, nh, hd, n_layers, vocab, eps) = (768usize, 12usize, 64usize, 12usize, 50257usize, 1e-5f32); let scale = 1.0 / (hd as f32).sqrt(); @@ -96,7 +96,7 @@ pub fn verify_mamba2(d: &dyn Device, plat: &str) { use ffai_ops::{conv1d_causal_step, rms_norm, silu, ssm_step}; let dir = model_dir("MAMBA2_DIR", "models--AntonV--mamba2-130m-hf").unwrap_or_default(); let path = format!("{dir}/model.safetensors"); - let Ok(st) = SafeTensors::open(&path) else { eprintln!("no model at {path} — skipping"); return; }; + let Ok(st) = SafeTensors::open(&path) else { eprintln!("no model at {path}, skipping"); return; }; let (hid, di, nh, dh, ds, ng, kc, vocab, eps) = (768usize, 1536usize, 24usize, 64usize, 128usize, 1usize, 4usize, 50288usize, 1e-5f32); let conv_dim = di + 2 * ng * ds; @@ -161,7 +161,7 @@ pub fn verify_mamba2(d: &dyn Device, plat: &str) { } /// Run the whole shared model suite against any backend. A new backend (ROCm, -/// Vulkan, …) implements `Device`, then calls this from one test file — and +/// Vulkan, …) implements `Device`, then calls this from one test file: and /// inherits every model with zero model code. pub fn run_all(d: &dyn Device, plat: &str) { verify_gpt2(d, plat); @@ -178,11 +178,11 @@ pub fn run_all(d: &dyn Device, plat: &str) { } /// Laguna-S-2.1 decode smoke test: env-gated on `FFAI_LAGUNA_GGUF` (path to -/// the model's GGUF file). Loads straight from GGUF (no HF safetensors dir — +/// the model's GGUF file). Loads straight from GGUF (no HF safetensors dir, /// this model doesn't have a small reference checkpoint to compare against), /// runs a handful of decode steps from BOS + a short token sequence with a /// persistent KV cache, and prints the logits argmax per step. No HF oracle -/// assertion (unlike the other `verify_*` — there's no cheap reference to +/// assertion (unlike the other `verify_*`, there's no cheap reference to /// compare against for a 117B model), just a load-bearing "it runs and /// produces finite, plausible-looking output" smoke check. /// Host-only Laguna GGUF dequant sanity: dump amax/rms per tensor for a few @@ -245,7 +245,7 @@ pub fn verify_laguna_host_dequant() { /// [`ffai_ops::quantize_q4`] and dequantized back on host, with the block /// scale pushed through an f32->f16->f32 roundtrip first, matching how the /// device stores Q4 scales). The device output should be compared against -/// `roundtrip`, not `exact` — decode always runs on Q4 weights regardless of +/// `roundtrip`, not `exact`: decode always runs on Q4 weights regardless of /// the GGUF's original per-tensor dtype (see `ffai_models::laguna::load_dense`). /// /// Env-gated on FFAI_LAGUNA_GGUF, no device/CUDA needed. @@ -260,7 +260,7 @@ pub fn verify_laguna_host_layer0() { }; // Laguna-S-2.1 layer-0 spec constants (dense SwiGLU MLP, full-attention - // layer — see ffai_models::laguna::LagunaConfig::defaults). + // layer, see ffai_models::laguna::LagunaConfig::defaults). let hidden = 3072usize; let head_dim = 128usize; let n_head = 48usize; @@ -288,7 +288,7 @@ pub fn verify_laguna_host_layer0() { }; // f32 <-> f16 (round-to-nearest), matching the device's scale-upload path - // (`ffai_models::laguna::f32s_to_f16_bytes`) — the gemv_q4 kernels read + // (`ffai_models::laguna::f32s_to_f16_bytes`), the gemv_q4 kernels read // per-block Q4 scales as f16. fn f32_to_f16(f: f32) -> u16 { let x = f.to_bits(); @@ -321,7 +321,7 @@ pub fn verify_laguna_host_layer0() { } // Q4 round-trip: quantize_q4 (block 32, symmetric int4 in [-7,7], scale = - // amax/7), then push the scale through f16 before dequantizing back — + // amax/7), then push the scale through f16 before dequantizing back: // mirrors exactly what the CUDA gemv_q4 kernel reads for its scales. let q4_roundtrip = |w: &[f32], m: usize, k: usize| -> Vec { let (qs, scales) = ffai_ops::quantize_q4(w, m, k); @@ -569,27 +569,88 @@ pub fn verify_laguna(d: &dyn Device, plat: &str) { use ffai_models::laguna; let Ok(path) = std::env::var("FFAI_LAGUNA_GGUF") else { - eprintln!("FFAI_LAGUNA_GGUF not set — skipping Laguna-S-2.1 decode smoke test"); + eprintln!("FFAI_LAGUNA_GGUF not set, skipping Laguna-S-2.1 decode smoke test"); return; }; let Ok(g) = ffai_loader::gguf::Gguf::open(&path) else { - eprintln!("Laguna: failed to open GGUF at {path} — skipping"); + eprintln!("Laguna: failed to open GGUF at {path}, skipping"); return; }; let model = match laguna::load_gguf(d, &g) { Ok(m) => m, Err(e) => { - eprintln!("Laguna: load_gguf failed: {e:?} — skipping"); + eprintln!("Laguna: load_gguf failed: {e:?}, skipping"); return; } }; + // FFAI_LAGUNA_PP=: standalone batched-prefill throughput bench over a + // synthetic n-token prompt (token id sequence i % 50000, any valid ids + // exercise the same code path; this is a speed bench, not a semantic + // eval). FFAI_LAGUNA_PREFILL_CHUNK overrides the chunk size (default + // 1024). Prints "pp tok/s" and returns, no generation, no decode. + if let Ok(pp_s) = std::env::var("FFAI_LAGUNA_PP") { + // Comma-separated sizes bench the whole ladder on ONE model load + // (the 71G requant load dominates a cycle). + let sizes: Vec = pp_s.split(',').filter_map(|v| v.trim().parse().ok()).collect(); + // FFAI_LAGUNA_PREFILL_CHUNK also accepts a comma list: the harness + // benches every (size, chunk) combo on the single model load. + let chunks: Vec = std::env::var("FFAI_LAGUNA_PREFILL_CHUNK") + .ok().map(|v| v.split(',').filter_map(|c| c.trim().parse().ok()).collect()) + .filter(|v: &Vec| !v.is_empty()) + .unwrap_or_else(|| vec![1024]); + for &(mut chunk) in chunks.iter() { + chunk = chunk.max(1); + for &n in sizes.iter().filter(|&&n| n > 0) { + let synth_tokens: Vec = (0..n).map(|i| (i % 50000) as u32).collect(); + let Ok(kv_pp) = laguna::LagunaKvCache::new(d, &model.cfg, n + 1) else { + eprintln!("Laguna PP bench: KV cache alloc failed, skipping"); + continue; + }; + let Ok(scratch_pp) = laguna::LagunaPrefillScratch::new(d, &model.cfg, n) else { + eprintln!("Laguna PP bench: prefill scratch alloc failed, skipping"); + continue; + }; + // One warmup pass (first-touch device allocs / kernel resolution), + // then the timed pass, matches the bench-honest convention used + // elsewhere in this harness (e.g. the FFAI_LAGUNA_SWEEP loop above). + if laguna::prefill(d, &model, &kv_pp, &scratch_pp, &synth_tokens, chunk).is_err() { + eprintln!("Laguna PP bench: warmup prefill failed, skipping timed run"); + continue; + } + let _ = d.synchronize(); + let kv_pp2 = match laguna::LagunaKvCache::new(d, &model.cfg, n + 1) { + Ok(k) => k, + Err(e) => { eprintln!("Laguna PP bench: KV cache realloc failed: {e:?}"); continue; } + }; + let scratch_pp2 = match laguna::LagunaPrefillScratch::new(d, &model.cfg, n) { + Ok(s) => s, + Err(e) => { eprintln!("Laguna PP bench: scratch realloc failed: {e:?}"); continue; } + }; + let t0 = std::time::Instant::now(); + match laguna::prefill(d, &model, &kv_pp2, &scratch_pp2, &synth_tokens, chunk) { + Ok(logits) => { + let _ = d.synchronize(); + let dt = t0.elapsed().as_secs_f64(); + let tps = n as f64 / dt; + let mut lb = vec![0u8; model.cfg.vocab * 4]; + let _ = d.download(logits.buffer.as_ref(), &mut lb); // liveness check on the returned tensor + eprintln!("Laguna prefill bench (n={n}, chunk={chunk}, {plat}): {tps:.2} pp tok/s"); + laguna::prof_report(); + } + Err(e) => eprintln!("Laguna PP bench: timed prefill failed: {e:?}"), + } + } + } + return; + } + let bos = g.meta_u32("tokenizer.ggml.bos_token_id").unwrap_or(0); // A short prompt: BOS + whatever the tokenizer produces for a trivial // string, falling back to a couple of low-valued ids if the tokenizer // metadata is unavailable (this is a decode-plumbing smoke test, not a - // semantic eval — any valid token ids exercise the same code path). + // semantic eval, any valid token ids exercise the same code path). let mut tokens = vec![bos]; if let Ok(tok) = GgufTokenizer::from_gguf(&g) { tokens.extend(tok.encode("Hello")); @@ -601,7 +662,7 @@ pub fn verify_laguna(d: &dyn Device, plat: &str) { + std::env::var("FFAI_LAGUNA_GEN").ok().and_then(|v| v.parse::().ok()).unwrap_or(16) + 1; let Ok(kv) = laguna::LagunaKvCache::new(d, &model.cfg, max_seq) else { - eprintln!("Laguna: KV cache alloc failed — skipping"); + eprintln!("Laguna: KV cache alloc failed, skipping"); return; }; @@ -663,7 +724,7 @@ pub fn verify_laguna(d: &dyn Device, plat: &str) { // replay it per token instead of relaunching ~1000+ kernels every step // (see `ffai_models::laguna::LagunaDecodeCtx`). Default (unset): plain // eager `decode_step`, now reading a persistent per-step workspace - // instead of uploading fresh token-id/position buffers every call — + // instead of uploading fresh token-id/position buffers every call, // same kernels, same scalars, byte-identical output to before. let graph_mode = std::env::var("FFAI_LAGUNA_GRAPH").is_ok(); let mode_label = if graph_mode { "graph" } else { "eager" }; @@ -673,7 +734,7 @@ pub fn verify_laguna(d: &dyn Device, plat: &str) { match laguna::DecodeWorkspace::new(d) { Ok(w) => Some(w), Err(e) => { - eprintln!("Laguna: decode workspace alloc failed: {e:?} — skipping"); + eprintln!("Laguna: decode workspace alloc failed: {e:?}, skipping"); return; } } @@ -682,7 +743,7 @@ pub fn verify_laguna(d: &dyn Device, plat: &str) { match laguna::LagunaDecodeCtx::new(d, &model, &kv, max_seq) { Ok(c) => Some(c), Err(e) => { - eprintln!("Laguna: LagunaDecodeCtx::new failed: {e:?} — skipping"); + eprintln!("Laguna: LagunaDecodeCtx::new failed: {e:?}, skipping"); return; } } @@ -705,11 +766,54 @@ pub fn verify_laguna(d: &dyn Device, plat: &str) { // of the full 400KB logit row (the bench-honest greedy path). Default // keeps the full download for logit stats/top-5 debugging. let fast = std::env::var("FFAI_LAGUNA_FAST").map(|v| v == "1").unwrap_or(false); - for pos in 0..total { + + // FFAI_LAGUNA_PREFILL=1: run the whole prompt through the batched + // `laguna::prefill` path in ONE (chunked) forward instead of feeding it + // token-by-token through decode_step/ctx.step below, then resume the + // IDENTICAL per-token decode loop from `pos = tokens.len()`. This is the + // correctness gate for the batched-prefill work: the greedy continuation + // must come out byte-identical to the decode-only path (same prompt, + // same seed token handed to decode at `pos = tokens.len()`). + // FFAI_LAGUNA_PREFILL_CHUNK overrides the chunk size (default 1024). + let prefill_mode = std::env::var("FFAI_LAGUNA_PREFILL").map(|v| v == "1").unwrap_or(false); + let prefill_chunk: usize = std::env::var("FFAI_LAGUNA_PREFILL_CHUNK") + .ok().and_then(|v| v.parse().ok()).unwrap_or(2048); + let start_pos = if prefill_mode { + let Ok(scratch) = laguna::LagunaPrefillScratch::new(d, &model.cfg, tokens.len()) else { + eprintln!("Laguna: prefill scratch alloc failed, skipping"); + return; + }; + let logits = match laguna::prefill(d, &model, &kv, &scratch, &tokens, prefill_chunk) { + Ok(l) => l, + Err(e) => { + eprintln!("Laguna: prefill failed: {e:?}"); + return; + } + }; + let mut lb = vec![0u8; model.cfg.vocab * 4]; + let _ = d.synchronize(); + if d.download(logits.buffer.as_ref(), &mut lb).is_err() { + eprintln!("Laguna: prefill logits download failed"); + return; + } + let lf: Vec = lb.chunks_exact(4).map(|c| f32::from_le_bytes(c.try_into().unwrap())).collect(); + let argmax = ffai_runtime::argmax(&lf); + eprintln!( + "Laguna prefill (prompt_len={}, chunk={prefill_chunk}, {plat}): last-token argmax = {argmax}", + tokens.len() + ); + cur = Some(argmax as u32); + gen_t0 = std::time::Instant::now(); // generation phase starts here, same convention as decode-only mode + tokens.len() + } else { + 0 + }; + + for pos in start_pos..total { let tok = if pos < tokens.len() { tokens[pos] } else { cur.unwrap() }; // Graph mode: LagunaDecodeCtx internally warms up (eager), captures // once, then replays the captured graph. Eager mode: refresh the - // persistent workspace by hand, then call decode_step directly — + // persistent workspace by hand, then call decode_step directly: // the argmax/download below still happen outside either path, // exactly as before. let step_result = if let Some(ctx) = ctx.as_mut() { @@ -798,7 +902,7 @@ pub fn verify_laguna(d: &dyn Device, plat: &str) { eprintln!("✅ Laguna-S-2.1 decode smoke test completed ({plat})."); } -// exact-erf GELU (Abramowitz-Stegun) — shared by GPT-NeoX / Whisper-style nets +// exact-erf GELU (Abramowitz-Stegun), shared by GPT-NeoX / Whisper-style nets fn erf(x: f32) -> f32 { let s = x.signum(); let x = x.abs(); let t = 1.0 / (1.0 + 0.3275911 * x); @@ -812,7 +916,7 @@ fn gelu_erf(x: f32) -> f32 { 0.5 * x * (1.0 + erf(x * std::f32::consts::FRAC_1_S pub fn verify_pythia(d: &dyn Device, plat: &str) { let dir = model_dir("PYTHIA_DIR", "models--EleutherAI--pythia-160m").unwrap_or_default(); let path = format!("{dir}/model.safetensors"); - let Ok(st) = SafeTensors::open(&path) else { eprintln!("no model at {path} — skipping"); return; }; + let Ok(st) = SafeTensors::open(&path) else { eprintln!("no model at {path}, skipping"); return; }; let (hid, nh, hd, inter, n_layers, vocab, eps) = (768usize, 12usize, 64usize, 3072usize, 12usize, 50304usize, 1e-5f32); let scale = 1.0 / (hd as f32).sqrt(); let g = |name: &str| -> Vec { @@ -854,7 +958,7 @@ pub fn verify_pythia(d: &dyn Device, plat: &str) { pub fn verify_olmo2(d: &dyn Device, plat: &str) { use ffai_ops::{rms_norm, swiglu}; let dir = model_dir("OLMO2_DIR", "models--allenai--OLMo-2-0425-1B").unwrap_or_default(); - let Ok(st) = SafeTensors::open_dir(&dir) else { eprintln!("no model at {dir} — skipping"); return; }; + let Ok(st) = SafeTensors::open_dir(&dir) else { eprintln!("no model at {dir}, skipping"); return; }; let (hid, nq, nkv, hd, inter, n_layers, vocab, eps) = (2048usize, 16usize, 16usize, 128usize, 8192usize, 16usize, 100352usize, 1e-6f32); let (qdim, kvdim) = (nq*hd, nkv*hd); let scale = 1.0/(hd as f32).sqrt(); let g = |name: &str| -> Vec { @@ -903,7 +1007,7 @@ pub fn verify_olmo2(d: &dyn Device, plat: &str) { pub fn verify_gptneo(d: &dyn Device, plat: &str) { let dir = model_dir("GPTNEO_DIR", "models--EleutherAI--gpt-neo-125m").unwrap_or_default(); let path = format!("{dir}/model.safetensors"); - let Ok(st) = SafeTensors::open(&path) else { eprintln!("no model at {path} — skipping"); return; }; + let Ok(st) = SafeTensors::open(&path) else { eprintln!("no model at {path}, skipping"); return; }; let (hid, nh, hd, inter, n_layers, vocab, eps) = (768usize, 12usize, 64usize, 3072usize, 12usize, 50257usize, 1e-5f32); let scale = 1.0f32; let g = |name: &str| -> Vec { @@ -947,7 +1051,7 @@ pub fn verify_gptneo(d: &dyn Device, plat: &str) { pub fn verify_gemma2(d: &dyn Device, plat: &str) { use ffai_ops::{mul, rms_norm}; let dir = model_dir("GEMMA2_DIR", "models--unsloth--gemma-2-2b-it").unwrap_or_default(); - let Ok(st) = SafeTensors::open_dir(&dir) else { eprintln!("no model at {dir} — skipping"); return; }; + let Ok(st) = SafeTensors::open_dir(&dir) else { eprintln!("no model at {dir}, skipping"); return; }; let (hid, nq, nkv, hd, inter, n_layers, vocab, eps) = (2304usize, 8usize, 4usize, 256usize, 9216usize, 26usize, 256000usize, 1e-6f32); let (qdim, kvdim) = (nq*hd, nkv*hd); let scale = 1.0/(256.0f32).sqrt(); let embed_scale = (hid as f32).sqrt(); let g = |name: &str| -> Vec { @@ -994,7 +1098,7 @@ pub fn verify_gemma2(d: &dyn Device, plat: &str) { /// partial rotary, gelu_new. vs HF top-3 [11,13,546]. pub fn verify_phi(d: &dyn Device, plat: &str) { let dir = model_dir("PHI_DIR", "models--microsoft--phi-1_5").unwrap_or_default(); - let Ok(st) = SafeTensors::open_dir(&dir) else { eprintln!("no model at {dir} — skipping"); return; }; + let Ok(st) = SafeTensors::open_dir(&dir) else { eprintln!("no model at {dir}, skipping"); return; }; let (hid, nh, hd, inter, n_layers, vocab, eps) = (2048usize, 32usize, 64usize, 8192usize, 24usize, 51200usize, 1e-5f32); let scale = 1.0/(hd as f32).sqrt(); let g = |name: &str| -> Vec { @@ -1038,7 +1142,7 @@ pub fn verify_phi(d: &dyn Device, plat: &str) { pub fn verify_stablelm2(d: &dyn Device, plat: &str) { use ffai_ops::swiglu; let dir = model_dir("STABLELM2_DIR", "models--stabilityai--stablelm-2-1_6b").unwrap_or_default(); - let Ok(st) = SafeTensors::open_dir(&dir) else { eprintln!("no model at {dir} — skipping"); return; }; + let Ok(st) = SafeTensors::open_dir(&dir) else { eprintln!("no model at {dir}, skipping"); return; }; let (hid, nh, hd, inter, n_layers, vocab, eps) = (2048usize, 32usize, 64usize, 5632usize, 24usize, 100352usize, 1e-5f32); let scale = 1.0/(hd as f32).sqrt(); let g = |name: &str| -> Vec { @@ -1084,7 +1188,7 @@ pub fn verify_stablelm2(d: &dyn Device, plat: &str) { pub fn verify_olmoe(d: &dyn Device, plat: &str) { use ffai_ops::{rms_norm, swiglu}; let dir = model_dir("OLMOE_DIR", "models--allenai--OLMoE-1B-7B-0924").unwrap_or_default(); - let Ok(st) = SafeTensors::open_dir(&dir) else { eprintln!("no model at {dir} — skipping"); return; }; + let Ok(st) = SafeTensors::open_dir(&dir) else { eprintln!("no model at {dir}, skipping"); return; }; let (hid, nh, hd, inter, n_exp, top_k, n_layers, vocab, eps) = (2048usize, 16usize, 128usize, 1024usize, 64usize, 8usize, 16usize, 50304usize, 1e-5f32); let scale = 1.0/(hd as f32).sqrt(); let g = |name: &str| -> Vec { @@ -1139,11 +1243,11 @@ pub fn verify_olmoe(d: &dyn Device, plat: &str) { } /// Falcon-H1-0.5B: hybrid Mamba2 mixer ∥ GQA attention per layer + µP. vs HF -/// top-3 [593,531,587]. (token 5 is a reserved zero-embedding slot — use 9707.) +/// top-3 [593,531,587]. (token 5 is a reserved zero-embedding slot, use 9707.) pub fn verify_falcon_h1(d: &dyn Device, plat: &str) { use ffai_ops::{conv1d_causal_step, rms_norm, silu, ssm_step}; let dir = model_dir("FALCON_H1_DIR", "models--tiiuae--Falcon-H1-0.5B-Base").unwrap_or_default(); - let Ok(st) = SafeTensors::open_dir(&dir) else { eprintln!("no model at {dir} — skipping"); return; }; + let Ok(st) = SafeTensors::open_dir(&dir) else { eprintln!("no model at {dir}, skipping"); return; }; let (hid, nq, nkv, ahd, inter, n_layers, vocab, eps) = (1024usize, 8usize, 2usize, 64usize, 2048usize, 36usize, 32784usize, 1e-5f32); let (d_ssm, m_nh, m_dh, d_state, n_groups, d_conv) = (1536usize, 24usize, 64usize, 128usize, 1usize, 4usize); let conv_dim = d_ssm + 2*n_groups*d_state; let proj_dim = 2*d_ssm + 2*n_groups*d_state + m_nh; @@ -1226,7 +1330,7 @@ pub fn verify_nemotron(d: &dyn Device, plat: &str) { const PATTERN: &str = "MEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; let dir = std::env::var("NEMOTRON_DIR") .unwrap_or_else(|_| "/home/pidtom/models/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16".into()); - let Ok(st) = SafeTensors::open_dir(&dir) else { eprintln!("no model at {dir} — skipping"); return; }; + let Ok(st) = SafeTensors::open_dir(&dir) else { eprintln!("no model at {dir}, skipping"); return; }; let (hid, vocab, eps) = (2688usize, 131072usize, 1e-5f32); // Mamba2: d_inner 4096 (64h×64), n_groups 8, d_state 128, conv_kernel 4. let (di, m_nh, m_dh, ds, ng, kc) = (4096usize, 64usize, 64usize, 128usize, 8usize, 4usize); @@ -1370,20 +1474,20 @@ pub fn verify_nemotron(d: &dyn Device, plat: &str) { } /// NemotronH-Nano-30B-A3B **resident Q8 decode benchmark**: quantize every big -/// matrix to Q8_0 and upload it ONCE (resident), then run a real decode loop — +/// matrix to Q8_0 and upload it ONCE (resident), then run a real decode loop: /// per-token forward reuses the resident weights via [`gemv_q8`], carries an /// attention KV cache (6 layers) + Mamba2 conv/SSM state (23 layers), and times /// steady-state tok/s. This is the path off the 0.003 tok/s naive-reload floor. /// Env: NEMOTRON_DECODE (steps, default 32), NEMOTRON_PREFILL (warm the cache to /// this context length before timing, default 0). pub fn bench_nemotron(d: &dyn Device, plat: &str) { - use ffai_ops::{add, cast_f16_f32, cast_f32_f16, conv1d_causal_prefill, conv1d_causal_prefill_split, conv1d_causal_step, dequant_q4, dequant_q4_off, gather, gated_group_rmsnorm, gated_group_rmsnorm_batched, gated_group_rmsnorm_batched_f16out, gemm_cublas, gemm_cublas_f32out, gemm_q4_mpp, gemv, gemv_q4, gemv_q4_accum, gemv_q4_relu2, gemv_q8, gemv_q8_relu2, gemv_q8_accum, kv_append, kv_append_many, mamba_split_conv, mamba_split_proj, matmul, moe_bgemm_q4_bm64, moe_expert_ids_from_offsets, moe_fused_ffn, moe_gather_down, moe_gather_up_relu2, moe_grouped_gemm, moe_grouped_gemm_cutlass, moe_q4_grouped_mma, moe_q4_grouped_mma_dev, moe_router_device, moe_scatter_add_det_dev, moe_scatter_add_topk_det_f16, moe_scatter_add_topk_det_f16_x2, moe_scatter_add, moe_scatter_add_det, moe_w4a16, moe_w4a16_marlin, moe_weighted_sum, nvfp4_roundtrip, permute_q4_to_marlin, quantize_q4, quantize_q8, relu2, relu2_scale_f16, rms_norm, rope_llama, rope_llama_many, sdpa_multi, sdpa_multi_tc, sdpa_multi_tc_varlen, silu, slice, sdpa_decode, sdpa_decode_2pass, sdpa_decode_2pass_bc4, sdpa_decode_2pass_tiled, softplus_add, softplus_add_rows, ssm_prefill_scan, ssm_prefill_scan_chunked, ssm_prefill_scan_ssd, ssm_prefill_scan_ssd_portable, ssm_step, strided_col_copy}; + use ffai_ops::{add, cast_f16_f32, cast_f32_f16, conv1d_causal_prefill_split, conv1d_causal_step, dequant_q4, dequant_q4_off, gather, gated_group_rmsnorm, gated_group_rmsnorm_batched, gated_group_rmsnorm_batched_f16out, gemm_cublas, gemm_cublas_f32out, gemm_q4_mpp, gemv, gemv_q4, gemv_q4_accum, gemv_q4_relu2, gemv_q8, gemv_q8_relu2, gemv_q8_accum, kv_append, kv_append_many, mamba_split_proj, matmul, moe_bgemm_q4_bm64, moe_expert_ids_from_offsets, moe_fused_ffn, moe_gather_down, moe_gather_up_relu2, moe_grouped_gemm, moe_grouped_gemm_cutlass, moe_q4_grouped_mma, moe_q4_grouped_mma_dev, moe_router_device, moe_scatter_add_det_dev, moe_scatter_add_topk_det_f16, moe_scatter_add_topk_det_f16_x2, moe_scatter_add, moe_scatter_add_det, moe_w4a16, moe_w4a16_marlin, moe_weighted_sum, nvfp4_roundtrip, permute_q4_to_marlin, quantize_q4, quantize_q8, relu2_scale_f16, rms_norm, rope_llama, rope_llama_many, sdpa_multi, sdpa_multi_tc, sdpa_multi_tc_varlen, silu, slice, sdpa_decode, sdpa_decode_2pass, sdpa_decode_2pass_bc4, sdpa_decode_2pass_tiled, softplus_add, softplus_add_rows, ssm_prefill_scan, ssm_prefill_scan_chunked, ssm_prefill_scan_ssd, ssm_prefill_scan_ssd_portable, ssm_step}; use std::collections::HashMap; use std::time::Instant; const PATTERN: &str = "MEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; let dir = std::env::var("NEMOTRON_DIR") .unwrap_or_else(|_| "/home/pidtom/models/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16".into()); - let Ok(st) = SafeTensors::open_dir(&dir) else { eprintln!("no model at {dir} — skipping"); return; }; + let Ok(st) = SafeTensors::open_dir(&dir) else { eprintln!("no model at {dir}, skipping"); return; }; let (hid, vocab, eps) = (2688usize, 131072usize, 1e-5f32); let (di, m_nh, m_dh, ds, ng, kc) = (4096usize, 64usize, 64usize, 128usize, 8usize, 4usize); let conv_dim = di + 2 * ng * ds; @@ -1476,7 +1580,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { // ── NVFP4 recipe flag ────────────────────────────────────────────────────── // NEMOTRON_NVFP4_RECIPE=1: switch Mamba in_proj/out_proj, shared-expert up/down, // and attention o_proj from Q4 to Q8. Routed MoE experts, q/k/v proj, and - // lm_head stay Q4. Result is ~4.98 bits/weight, ~20.9 GB resident — matching + // lm_head stay Q4. Result is ~4.98 bits/weight, ~20.9 GB resident: matching // NVIDIA's published NVFP4 mixed-precision recipe for apples-to-apples comparison. let use_nvfp4_recipe = std::env::var("NEMOTRON_NVFP4_RECIPE").map(|v| v != "0" && v != "false").unwrap_or(false); @@ -1497,7 +1601,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { let mut fp4w_cutlass: HashMap = HashMap::new(); // W4A8 MoE weights (NEMOTRON_W4A8_MOE): mxfp4 weights (per-32 ue8m0) packed // ONCE at load → name → (packed e2m1, sf ue8m0, sfb_exp_elems). Fed to the - // mxf8f6f4 grouped GEMM with fp8 e4m3 acts — the quality-holding fast path + // mxf8f6f4 grouped GEMM with fp8 e4m3 acts, the quality-holding fast path // (256 act levels vs W4A4's 16; weights stay 4-bit so HBM/speed ~= W4A4). let mut w4a8w: HashMap = HashMap::new(); let use_w4a8_env = std::env::var("NEMOTRON_W4A8_MOE").is_ok(); @@ -1519,7 +1623,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { // ── Q8 cache helpers (parallel to the Q4 cache, uses .q8b extension) ────── // Cache format: [8B m][8B k][N*4 bytes qs u32 LE][M*4 bytes scales f32 LE] - // (Q8 scales are always f32; no f16 flag byte needed — layout is fixed.) + // (Q8 scales are always f32; no f16 flag byte needed, layout is fixed.) let cache_path_q8 = |name: &str| -> std::path::PathBuf { let safe: String = name.chars().map(|c| if c.is_alphanumeric() || c == '-' { c } else { '_' }).collect(); std::path::PathBuf::from(&cache_dir).join(format!("{safe}.q8b")) @@ -1564,7 +1668,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { qw8.insert(name.to_string(), (qt, sct, m, k)); }; - // f16: true for weights read by the PLAIN gemv kernel (qmv) — its scale param + // f16: true for weights read by the PLAIN gemv kernel (qmv), its scale param // is f16. false for shared-expert weights (relu2/accum kernels, f32 scale). let qload = |qw: &mut HashMap, name: &str, m: usize, k: usize, f16: bool| { // Try cache hit first: skip BF16→F32 + quantize if available. @@ -1661,7 +1765,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { Tensor::new(d.upload(&dsb).unwrap(), vec![dsb.len() / 2], DType::F16), n_exp * hid, inter)); } } else { - // Dimension mismatch — fall through to rebuild. + // Dimension mismatch, fall through to rebuild. let (mut uqs, mut usc, mut dqs, mut dsc) = (Vec::new(), Vec::new(), Vec::new(), Vec::new()); for e in 0..n_exp { let (q, s) = quantize_q4(&g(&format!("{p}.mixer.experts.{e}.up_proj.weight")), inter, hid); @@ -1711,7 +1815,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { let (wp_len, wsc_len) = (n_exp * nt2 * kt2 * 32 * 2, n_exp * nt2 * kt2 * 8); let cpath = cache_path(&format!("{key}.v1")); // disk-cache hit: upload the packed blobs directly (skips the - // f32 collect + f16 convert + device pack — the slow part). + // f32 collect + f16 convert + device pack, the slow part). if let Ok(b) = std::fs::read(&cpath) { let expect = (wp_len + wsc_len + n_exp) * 4; if b.len() == expect { @@ -1785,7 +1889,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { // W4A8: pack the SAME dequanted f16 experts as mxfp4 (per-32 // ue8m0) once at load. sfb_exp_elems returned by w4a8_packw // is the per-expert SF stride the GEMM consumes (no disk cache - // — packing is fast and the stride must come from the packer). + //, packing is fast and the stride must come from the packer). if use_w4a8_env { let mut buildw = |which: &str, n: usize, k: usize| { let key = format!("{p}.moe_{which}_w4a8"); @@ -2026,8 +2130,8 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { // ── Fused MoE FFN (NEMOTRON_MOE_FUSED=1) ───────────────────────────────── // Pre-allocate the scratch buffer once. All 23 MoE layers share it (they're - // serial — layer N+1 starts only after N finishes). - // NOTE: NEMOTRON_GRAPH + NEMOTRON_MOE_FUSED are mutually exclusive — the + // serial, layer N+1 starts only after N finishes). + // NOTE: NEMOTRON_GRAPH + NEMOTRON_MOE_FUSED are mutually exclusive, the // fused kernel uses cuLaunchCooperativeKernel which can't be captured into a // CUDA graph. If both are set, fusion is silently disabled. let use_moe_fused = std::env::var("NEMOTRON_MOE_FUSED").is_ok() @@ -2044,7 +2148,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { eprintln!(" On GB10 (48 SMs), max coop blocks = 288 at 256 threads, need 336 for hid=2688."); eprintln!(" This will FAIL with 'too many blocks'. See NEMOTRON_MOE_FUSED analysis."); } else if std::env::var("NEMOTRON_MOE_FUSED").is_ok() && std::env::var("NEMOTRON_GRAPH").is_ok() { - eprintln!("NEMOTRON_MOE_FUSED=1: disabled (NEMOTRON_GRAPH is set — cooperative launch not capturable)"); + eprintln!("NEMOTRON_MOE_FUSED=1: disabled (NEMOTRON_GRAPH is set, cooperative launch not capturable)"); } // ── DECODE: per-token forward reusing resident weights, KV + Mamba state ── @@ -2070,12 +2174,12 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { // needed for the graph throughput measurement). let skip_dl = std::cell::Cell::new(false); // F16 KV cache: the clock-locked "+11-27%" was a thermal/order artifact (the - // internal A/B shows it neutral-to-negative — casts cost ≥ halved-read saves). + // internal A/B shows it neutral-to-negative, casts cost ≥ halved-read saves). // DEFAULT OFF (f32, no casts); opt in with NEMOTRON_F16KV. let f16kv = std::env::var("NEMOTRON_F16KV").is_ok(); // KV-cache capacity (positions). In the batched-prefill path the queries // run at positions [fakectx, fakectx+prefill), so the cache must hold the - // full fakectx prefix PLUS the whole prefill block PLUS the decode tail — + // full fakectx prefix PLUS the whole prefill block PLUS the decode tail, // i.e. fakectx + prefill, NOT max(fakectx, prefill). Using max() under-sized // the cache whenever both were nonzero, so sdpa_multi/kv_append_many walked // past the buffer (base_kv + n_query > kv_stride) → illegal memory access at @@ -2087,7 +2191,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { let cap = fakectx + prefill * pack_n_cap + n_decode + 8; // per-layer state, indexed by absolute layer id. KV cache is now ON-DEVICE // ([nkv,cap,hd] per attn layer), so the growing context never round-trips - // through the host — the 32K decode fix. + // through the host, the 32K decode fix. let mut conv_state: Vec> = vec![Vec::new(); 52]; let conv_dev: std::cell::RefCell>> = std::cell::RefCell::new((0..52).map(|_| None).collect()); // device conv state (NEMOTRON_DEVMAMBA) let mut ssm_state: Vec> = (0..52).map(|_| None).collect(); // recurrent SSM state ON-DEVICE @@ -2095,7 +2199,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { let u32buf = |v: u32| Tensor::new(d.upload(&v.to_le_bytes()).unwrap(), vec![1], DType::U32); let ones_gs = up(&vec![1.0f32; gs]); // grouped-norm: normalize each 512-group weightless, then scale by the real weight // Optional per-op profiling (NEMOTRON_PROFILE=1): synchronize around EACH individual - // op call to attribute GPU time precisely. Adds sync overhead on every op — use only + // op call to attribute GPU time precisely. Adds sync overhead on every op, use only // for profiling runs, not for throughput measurement. let prof = std::env::var("NEMOTRON_PROFILE").is_ok(); @@ -2162,7 +2266,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { }}; } - // residual stream stays ON-DEVICE the whole forward — no per-layer up(x)/dl(out). + // residual stream stays ON-DEVICE the whole forward, no per-layer up(x)/dl(out). let mut xt = up(&embed[token * hid..(token + 1) * hid]); for (l, mix) in PATTERN.chars().enumerate() { let p = format!("language_model.backbone.layers.{l}"); @@ -2405,7 +2509,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { // Final norm + lm_head let xf = pt!(25, hid * 8, rms_norm(d, &xt, &fwd["norm_f"], eps).unwrap()); // Gate the final argmax download: forbidden during CUDA-graph capture (host-sync). - // During capture and graph replay we return a dummy 0 — the bench only measures + // During capture and graph replay we return a dummy 0, the bench only measures // throughput, not token correctness. if skip_dl.get() { pt!(26, vocab * hid / 2 + vocab * 2 + hid * 4, @@ -2428,7 +2532,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { let mut pos = 0usize; // For the 32K measurement, start at `fakectx` (the device KV cache is alloc'd // to `cap` and zero-filled on first use, so the sdpa genuinely reads `fakectx` - // cached positions — the real 32K work — without a 38-min token prefill). + // cached positions, the real 32K work, without a 38-min token prefill). if fakectx > 0 { pos = fakectx; } // warm the cache to `prefill` positions, then time `n_decode` steps. // PHASE-0 BASELINE: time the sequential-prefill loop so we can report the @@ -2447,7 +2551,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { // (the same token the sequential path would produce). Token sequence is the // greedy chain seeded by `tok` (matching the sequential warm-up). thread_local! { static LAST_BATCHED_LOGITS: std::cell::RefCell> = const { std::cell::RefCell::new(Vec::new()) }; } - // Device logits tensor (no dl) + capture flag — for CUDA-graphing the forward. + // Device logits tensor (no dl) + capture flag, for CUDA-graphing the forward. thread_local! { static LAST_BATCHED_LOGITS_DEV: std::cell::RefCell> = const { std::cell::RefCell::new(None) }; } thread_local! { static CAPTURING: std::cell::Cell = const { std::cell::Cell::new(false) }; } let prefill_batched = std::env::var("NEMOTRON_PREFILL_BATCHED").is_ok(); @@ -2455,7 +2559,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { let s = prefill; // number of prompt tokens // f16 weight cache (Path A): dequant each Q4 weight to f16 ONCE (lazily on // first use), keyed by name (+ expert id for MoE). The forward then feeds - // the cached f16 weight straight to cuBLAS — eliminates the ~26%-of-prefill + // the cached f16 weight straight to cuBLAS, eliminates the ~26%-of-prefill // redundant per-forward dequant. ~2× the GEMM-weight VRAM; fine on GB10. // Disable with NEMOTRON_NO_W16CACHE=1 (A/B). let w16: std::cell::RefCell> = std::cell::RefCell::new(HashMap::new()); @@ -2464,7 +2568,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { // resident (not just dense+shared). ~44.8 GB extra f16 on GB10 (62 GB total // resident). Kills the routed-expert per-forward dequant (the remaining // ~42% of prefill) at the cost of streaming a fat f16 expert pool from the - // unified-memory pool. Honest A/B lever — measure dequant%→0 vs the GEMM + // unified-memory pool. Honest A/B lever, measure dequant%→0 vs the GEMM // becoming bandwidth-bound on the larger f16 working set. let resident_w16 = std::env::var("NEMOTRON_RESIDENT_W16").map(|v| v != "0" && v != "false").unwrap_or(false); // Greedy token chain (same as sequential): we need the S token ids. The @@ -2488,7 +2592,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { // Device-resident constants cache: per-layer constant vectors (dt_bias, // router correction bias) + tiny index tensors, uploaded ONCE and reused // across layers AND forwards. These were re-uploaded per layer per - // forward — ~140 small pageable H2D copies/forward, each paying ~0.3ms + // forward, ~140 small pageable H2D copies/forward, each paying ~0.3ms // of driver API time on GB10. let const_dev: std::cell::RefCell> = std::cell::RefCell::new(HashMap::new()); @@ -2499,7 +2603,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { let w4lt: std::cell::RefCell> = std::cell::RefCell::new(HashMap::new()); // Single-slot act-quant dedup: q/k/v (and any same-input projections) - // quantize the SAME activation tensor — cache the last (input, dims) -> + // quantize the SAME activation tensor, cache the last (input, dims) -> // (pack, sf). Holding the input Tensor keeps its buffer alive, so the // pointer-identity check can't alias a pool-recycled buffer. let w4g: std::cell::RefCell> = std::cell::RefCell::new(HashMap::new()); @@ -2522,7 +2626,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { // matmul (f32 MMA) for the per-expert / shared-expert GEMMs. The scan // (ssm_prefill_scan), attention (sdpa_multi), conv (host-bridge) and router // (host top-k) defaults are already portable at d0/S≤4096. CUDA path is - // left byte-for-byte unchanged — gated purely on `is_cuda`. + // left byte-for-byte unchanged, gated purely on `is_cuda`. let is_cuda = d.backend() == Backend::Cuda; // Fetch-or-upload a device-resident constant (cached across forwards). let cdev = |key: String, make: &dyn Fn() -> Tensor| -> Tensor { @@ -2539,7 +2643,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { // proj/MoE/router/norm are token-parallel → unchanged. RoPE is relative // so global positions give correct intra-segment rotation. NOTE: the // causal conv1d is NOT yet segment-aware (kc-1≈3 tokens/segment leak - // across the boundary) — throughput path until the varlen conv lands. + // across the boundary), throughput path until the varlen conv lands. let n_seg = std::env::var("NEMOTRON_PACKED").ok() .and_then(|v| v.parse::().ok()).unwrap_or(1).max(1); let ssd_l_pk = std::env::var("NEMOTRON_SSD_L").ok() @@ -2610,8 +2714,8 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { // NEMOTRON_PREFILL_F32GEMM=1 = f32 scalar matmul (A/B). let use_f32_gemm = std::env::var("NEMOTRON_PREFILL_F32GEMM").is_ok(); let use_mpp_gemm = std::env::var("NEMOTRON_PREFILL_MPPGEMM").is_ok(); - // NEMOTRON_BF16_STREAM: run residual/hidden stream in bf16 (vLLM-style) - // instead of f32 — halves elementwise/residual HBM traffic + native model + // NEMOTRON_BF16_STREAM: run residual/hidden stream in bf16 (a serving-framework-style pattern) + // instead of f32, halves elementwise/residual HBM traffic + native model // precision. GEMMs/norms/adds bf16; conv/SSD/MoE-quant cast at f32 boundary. // ── Metal per-expert MoE GEMM compute dtype (Apple/Metal only) ── // The per-forward Q4→dtype dequant of the routed experts × 23 MoE layers is a @@ -2623,7 +2727,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { // f16-compute is the Metal DEFAULT; set NEMOTRON_METAL_F32_EXPERTS=1 to fall // back to the f32 path for A/B. No resident f16 weight cache: caching the f16 // experts resident was measured a NET LOSS on Metal (bandwidth/residency bound - // — re-dequanting the compact Q4 each forward streams better than a fat f16 + //, re-dequanting the compact Q4 each forward streams better than a fat f16 // working set). CUDA path is unaffected (gated on `is_cuda` below). let metal_f16_experts = !is_cuda && !std::env::var("NEMOTRON_METAL_F32_EXPERTS").is_ok(); // Cached Q4→f16 dequant. On miss, dequant the [m,k] slab (optionally at @@ -2633,7 +2737,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { // weights (small, always used: ~few GB) by default. The 128 routed // experts ×23 layers (~44.8 GB f16) re-dequant per forward UNLESS // NEMOTRON_RESIDENT_W16=1 (`resident_w16`), which also caches them - // resident — fits GB10's 121 GB unified pool (~62 GB total resident). + // resident, fits GB10's 121 GB unified pool (~62 GB total resident). let deq16c = |key: &str, qs: &Tensor, sc: &Tensor, m: usize, k: usize, blk_off: usize, cache: bool| -> Tensor { if no_w16 || !cache { return pf!(3, (m * k) as f64, dequant_q4_off(d, qs, sc, m, k, DType::F16, blk_off).unwrap()); @@ -2667,7 +2771,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { Some("qkv") => name.contains("q_proj") || name.contains("k_proj") || name.contains("v_proj") || name.contains("o_proj"), // audit: in_proj carries the worst FP4 quantization error - // (normed-residual outliers) — exclude it, FP4 the rest. + // (normed-residual outliers), exclude it, FP4 the rest. Some("noin") => !name.contains("in_proj"), Some("out") => name.contains("out_proj"), Some(_) => true, @@ -2718,7 +2822,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { }; // FP8 dense (NEMOTRON_FP8_DENSE): the vendor recipe for this model // family runs Mamba in/out_proj (+ o_proj, shared) at FP8 e4m3 with - // per-tensor scales — far gentler than FP4. Scopes: "mamba" + // per-tensor scales, far gentler than FP4. Scopes: "mamba" // (in/out_proj), "recipe" (+ o_proj), "all". let fp8_dense_scope = if is_cuda { std::env::var("NEMOTRON_FP8_DENSE").ok() } else { None }; let fp8_dense = fp8_dense_scope.is_some(); @@ -2732,11 +2836,11 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { }; let qmm_fp8d = |xh: &Tensor, name: &str, rows: usize, m: usize, k: usize, qs: &Tensor, sc: &Tensor| -> Tensor { - // NEMOTRON_FP8_PT: per-TENSOR fp8 (fast gemm_fp8, scales folded in alpha) — + // NEMOTRON_FP8_PT: per-TENSOR fp8 (fast gemm_fp8, scales folded in alpha), // FFAI proved 1.81x @ argmax near-tie; A/B vs the per-channel default. if std::env::var("NEMOTRON_FP8_PT").is_ok() { // NEMOTRON_FP8_HAD: block-Hadamard rotate (128) along K before the - // per-tensor fp8 quant — spreads activation outliers so per-tensor + // per-tensor fp8 quant, spreads activation outliers so per-tensor // holds quality. Rotate weight offline (cached) + act at runtime; // orthonormal so the GEMM result is unchanged. let had = std::env::var("NEMOTRON_FP8_HAD").is_ok() && k % 128 == 0; @@ -2763,7 +2867,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { ffai_ops::gemm_fp8(d, &xq, &xsc, &wq, &wsc, rows, m, k, true).unwrap()); } // PER-CHANNEL weight (rows=m) + PER-TOKEN act (rows) FP8, scales - // folded post-GEMM (D[r,c]*=xsc[r]·wsc[c]) — the quality- + // folded post-GEMM (D[r,c]*=xsc[r]·wsc[c]), the quality- // preserving granularity (per-tensor was cosine ~0.94). let (wq, wsc) = { if let Some(t) = w4lt.borrow().get(&format!("{name}.fp8c")) { t.clone() } else { @@ -2802,7 +2906,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { pf!(2, flops, matmul(d, &wf, x).unwrap()) // [rows, m] } else if use_mpp_gemm || !is_cuda { // Portable Metal path (also NEMOTRON_PREFILL_MPPGEMM): Q4-native - // cooperative-tensor MMA — no cuBLAS, runs on Apple GPU hardware MMA. + // cooperative-tensor MMA, no cuBLAS, runs on Apple GPU hardware MMA. let xh = pf!(15, 0.0, cast_f32_f16(d, x).unwrap()); let oh = pf!(2, flops, gemm_q4_mpp(d, &xh, qs, sc, rows, *m, *k).unwrap()); pf!(15, 0.0, cast_f16_f32(d, &oh).unwrap()) // [rows, m] @@ -2811,14 +2915,14 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { let wf = deq16(name, qs, sc, *m, *k, 0); let xh = pf!(15, 0.0, cast_f32_f16(d, x).unwrap()); // Fused output cast: cuBLAS writes f32 directly (accumulate is - // already f32) — drops the separate cast_f16_f32, MFU preserved. + // already f32), drops the separate cast_f16_f32, MFU preserved. pf!(2, flops, gemm_cublas_f32out(d, &xh, &wf, rows, *m, *k).unwrap()) // [rows, m] f32 } }; // Like `qmm` but takes a PRE-CAST f16 activation `xh` (skips the input // cast_f32_f16). Lets multiple projections over the SAME normed input // (q/k/v from `xn`, shared-expert up from `xn`) share one f16 cast - // instead of re-casting per projection — kills redundant slice/cast. + // instead of re-casting per projection, kills redundant slice/cast. // cuBLAS tensor-core path only (the fast prefill path). let qmm_h = |xh: &Tensor, name: &str| -> Tensor { let (qs, sc, m, k) = &qw[name]; @@ -2836,7 +2940,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { return pf!(15, 0.0, cast_f16_f32(d, &oh).unwrap()); // [rows, m] } let wf = deq16(name, qs, sc, *m, *k, 0); - // Fused output cast: cuBLAS writes f32 directly — no cast_f16_f32. + // Fused output cast: cuBLAS writes f32 directly, no cast_f16_f32. pf!(2, flops, gemm_cublas_f32out(d, xh, &wf, rows, *m, *k).unwrap()) // [rows, m] f32 }; // Fuse-slice-cast opt-in (NEMOTRON_FUSE_QKV=1): cast xn→f16 once per @@ -2848,7 +2952,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { // DN scratch: [n_exp * hid, inter] f16. // These are bounded: 128 * 1856 * 4096 * 2 ≈ 1.83 GB each, ≈ 3.67 GB total. // Pre-alloc grouped-gemm scratch for GROUPED_GEMM and W4A16 (which falls back - // to grouped_gemm at large S — needs scratch to avoid per-layer cuMemAlloc). + // to grouped_gemm at large S, needs scratch to avoid per-layer cuMemAlloc). let needs_grouped_scratch = (std::env::var("NEMOTRON_GROUPED_GEMM").is_ok() && !std::env::var("NEMOTRON_BGEMM").is_ok()) || std::env::var("NEMOTRON_W4A16").is_ok() || std::env::var("NEMOTRON_W4A16_MARLIN").is_ok(); @@ -2893,13 +2997,13 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { if use_conv_device { let proj_t = pf!(2, 2.0 * s as f64 * in_proj_out as f64 * hid as f64, qmm(&xn, &format!("{p}.mixer.in_proj.weight"))); // [s, in_proj_out] on device - // PART B: fused split — 1 dispatch replaces 3 strided_col_copy calls. + // PART B: fused split, 1 dispatch replaces 3 strided_col_copy calls. // proj layout: [z(di) | xbc(conv_dim) | dt_raw(m_nh)] per row. let proj_t = if stream16 { pf!(15, 0.0, to_f32(d, &proj_t)) } else { proj_t }; let (z_t, xbc_t, dt_raw_t) = pf!(15, 0.0, mamba_split_proj(d, &proj_t, s, in_proj_out, di, conv_dim, m_nh).unwrap()); // conv1d+silu+x/B/C split FUSED into one raw-float4 kernel - // (conv1d_causal_prefill_split) below — deletes the temp + // (conv1d_causal_prefill_split) below, deletes the temp // [s,conv_dim] conv output + the separate split launch. // Persist conv ring state for decode continuity: last (kc-1) rows of xbc. let ring_len = (kc - 1) * conv_dim; @@ -2951,7 +3055,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { }; let (so, y_dev) = ssd_split_out.take().unwrap(); ssm_state[l] = Some(so); - // PART B: on-device batched gated group RMSNorm — eliminates the + // PART B: on-device batched gated group RMSNorm, eliminates the // dl(y)+dl(z)+host loop+up(yn) round-trip (2 × s×di = ~42MB per layer). // cast-fusion: write gated-norm directly as f16 + consume via qmm_h // (skips a cast_f32_f16 per Mamba layer). Bit-identical f16. Escape: @@ -2966,7 +3070,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { qmm(&ynt.reshaped(vec![s, di]), &format!("{p}.mixer.out_proj.weight")) }; { - // fused residual add + NEXT layer's rms_norm (one pass, llama.cpp-style) + // fused residual add + NEXT layer's rms_norm (one pass, in the reference C++ implementation's style) let nw = if l + 1 < PATTERN.len() { fwd.get(&format!("language_model.backbone.layers.{}.norm.weight", l + 1)) } else { fwd.get("norm_f") }; match nw { Some(w) if is_cuda && std::env::var("NEMOTRON_FUSED_ADDNORM").is_ok() => { let (r, nx) = ffai_ops::add_rms_norm(d, &xt, &out, w, eps).unwrap(); xt = r; xn_carry = Some(nx); } @@ -3078,7 +3182,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { let ynt = upm(&yn, vec![s, di]); let out = qmm(&ynt, &format!("{p}.mixer.out_proj.weight")); // [s, hid] { - // fused residual add + NEXT layer's rms_norm (one pass, llama.cpp-style) + // fused residual add + NEXT layer's rms_norm (one pass, in the reference C++ implementation's style) let nw = if l + 1 < PATTERN.len() { fwd.get(&format!("language_model.backbone.layers.{}.norm.weight", l + 1)) } else { fwd.get("norm_f") }; match nw { Some(w) if is_cuda && std::env::var("NEMOTRON_FUSED_ADDNORM").is_ok() => { let (r, nx) = ffai_ops::add_rms_norm(d, &xt, &out, w, eps).unwrap(); xt = r; xn_carry = Some(nx); } @@ -3093,7 +3197,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { // Batched gate gemv: [S, n_exp]. gate.weight is f32 dense. let gate_w = &fwd[&format!("{p}.mixer.gate.weight")]; // NEMOTRON_DEVSORT=1 (CUDA): on-device router + counting-sort - // (moe_route_sort_device) — keeps the gate logits ON DEVICE (no + // (moe_route_sort_device), keeps the gate logits ON DEVICE (no // ~1MB rl_all dl), runs sigmoid+bias+topk + sort on the GPU, then // reconstructs the SAME `triples` via 3 small dls (offsets, sorted // tok, sorted wt). Removes the per-E-layer host sigmoid/topk/sort. @@ -3102,7 +3206,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { let _t_router = std::time::Instant::now(); // off_dev/st_dev/sw_dev hoisted out (device routing tensors) so the // Q4 path can run fully on-device (device descriptors + gather + - // scatter) with NO host sync — the prerequisite for CUDA-graphing + // scatter) with NO host sync, the prerequisite for CUDA-graphing // the forward (kill the host-launch GPU idle, the real 3-5× lever). let mut moe_dev: Option<(Tensor, Tensor, Tensor)> = None; // (offsets, sorted_tok, sorted_wt) let mut moe_tok: Option<(Tensor, Tensor)> = None; // (rowpos, token-major router weights) @@ -3119,7 +3223,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { // ~2× prefill (S=2048 925→1820, S=1024 968→2009, S=3072 461→1865 tok/s; // argmax-matches the host path at S=1024/2048/3072, gold 1104@S2048). // NEMOTRON_Q4_DEVDESC_OFF=1 reverts to the host-triples path (A/B). - // (Decode is unaffected — this is inside the batched-prefill closure.) + // (Decode is unaffected, this is inside the batched-prefill closure.) let devdesc_skip = is_cuda && std::env::var("NEMOTRON_Q4_DEVDESC_OFF").is_err(); let (triples, mt): (Vec<(u32, usize, f32)>, usize) = if devsort { let bias_dev = cdev(format!("{p}.mixer.gate.e_score_correction_bias"), &|| up(&bias[..])); @@ -3177,7 +3281,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { // NEMOTRON_W4A16=1: W4A16 WMMA grouped GEMM (inline Q4 dequant, // scattered nibble reads). Supersedes NEMOTRON_BGEMM when both set. // NEMOTRON_W4A16_MARLIN=1: same as W4A16 but with Marlin coalesced - // tile-major layout — requires weights pre-permuted at load time. + // tile-major layout, requires weights pre-permuted at load time. // Supersedes both NEMOTRON_W4A16 and NEMOTRON_BGEMM. let use_w4a16_marlin = std::env::var("NEMOTRON_W4A16_MARLIN").is_ok(); let use_w4a16 = std::env::var("NEMOTRON_W4A16").is_ok() || use_w4a16_marlin; @@ -3186,10 +3290,10 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { // expert paths all finish with moe_scatter_add (+relu2_scale_f16), // which dispatch_raw_cuda → hard-error on Metal. Force the portable // per-expert path (host scatter, host f32 relu2) whose only CUDA dep - // is the GEMM — and that we swap below for dequant_q4_off + matmul. + // is the GEMM, and that we swap below for dequant_q4_off + matmul. let (use_bgemm, use_w4a16, use_w4a16_marlin) = if is_cuda { (use_bgemm, use_w4a16, use_w4a16_marlin) } else { (false, false, false) }; - // NEMOTRON_GROUPED_GEMM=1: async two-pass all-expert GEMM — all UP + // NEMOTRON_GROUPED_GEMM=1: async two-pass all-expert GEMM, all UP // GEMMs enqueued (no inter-expert sync), on-device relu2_scale_f16, // all DN GEMMs, then on-device scatter-add. Uses the device-gather // path for xs. Eliminates per-expert host syncs + cast overhead. @@ -3202,14 +3306,14 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { && !use_fp4_moe && !use_bgemm && !use_w4a16 && !use_grouped_gemm; // NEMOTRON_Q4_GROUPED=1 (CUDA): Q4-NATIVE single-launch grouped MoE - // GEMM (moe_q4_grouped_mma) — reads the model's signed Q4 weights + // GEMM (moe_q4_grouped_mma), reads the model's signed Q4 weights // DIRECTLY (no f16 slab, no dequant) per up/down over the sorted // tokens; on-device relu²+det-scatter. Collapses ~234 per-expert // cuBLAS launches → 2 AND reads 4.5-bit weights (vs 16-bit) → the // Marlin lever for the BW-bound MoE wall. Uses the device gather. // NEMOTRON_Q4_GROUPED_LASTEXACT=K: run the EXACT cuBLAS per-expert // path for the LAST K MoE layers (output-proximal) and Q4-grouped for - // the earlier ones — a precision/speed hybrid to test whether late- + // the earlier ones, a precision/speed hybrid to test whether late- // layer exactness restores the argmax (1104) while keeping most of // the +35%. moe_idx = this E-layer's 0-based index among all E-layers. let q4_lastexact: usize = std::env::var("NEMOTRON_Q4_GROUPED_LASTEXACT") @@ -3218,7 +3322,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { let moe_idx = PATTERN.chars().take(l).filter(|&c| c == 'E').count(); // NEMOTRON_Q4_GROUPED_FIRSTEXACT=K: exact cuBLAS for the FIRST K MoE // layers (the sensitive early ones), Q4 for the rest. (lastexact failed - // at all K — early error propagates irreducibly; test exact-early.) + // at all K, early error propagates irreducibly; test exact-early.) // Default FIRSTEXACT=4: the first 4 MoE layers run EXACT cuBLAS, the // rest Q4. Validated argmax==1104@S2048 (exact-early restores the // razor 1104/1776 near-tie that pure-Q4 flips; gap 0.125, robust). @@ -3238,12 +3342,12 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { && !use_bgemm && !use_w4a16 && !use_grouped_gemm && !use_cutlass_moe; // NEMOTRON_MOE_GROUPED=1 (Metal only): replace the per-expert // dequant_q4_off+matmul loop with the token-sorted GROUPED Q4 - // GEMM (moe_bgemm_q4_bm64) — ONE MMA dispatch over all sorted + // GEMM (moe_bgemm_q4_bm64), ONE MMA dispatch over all sorted // rows, Q4 dequant fused in the kernel's block prologue (no // per-expert f16 weight materialization). relu²+scale on host // f32 (overflow-safe, matches the validated per-expert numerics), // then a second grouped GEMM for down, then host scatter. The - // Metal analog of the CUDA GROUPED_GEMM/BGEMM device path — but + // Metal analog of the CUDA GROUPED_GEMM/BGEMM device path, but // those finish with the CUDA-only moe_scatter_add, so this path // keeps the portable host relu²/scatter and only swaps the GEMM. // bm64 requires n_out % 64 == 0: inter=1856 (29·64) and @@ -3264,7 +3368,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { // dl(&xn)+host-scatter+per-expert re-upload, and keep the routed // accumulator on device through the shared-expert + residual merge. // Kills the ~22MB D2H × 23 host gather + host scatter + the - // dl(acc)/upm(acc) round-trip — the dominant remaining host wall in + // dl(acc)/upm(acc) round-trip, the dominant remaining host wall in // the conv-device prefill path. argmax-exact (1104/1120/1763), // bit-deterministic, +~70% e2e. Escape: NEMOTRON_ONDEVICE_MOE_OFF=1 // reverts to the host-gather path. @@ -3292,7 +3396,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { None }; // FP4-direct: actq_pack gathers on read from the ungathered - // [s,hid] f16 — skip building the gathered xs entirely. + // [s,hid] f16, skip building the gathered xs entirely. let fp4_skip_gather = is_cuda && moe_dev.is_some() && xn_f16_router.is_some() && std::env::var("NEMOTRON_FP4_DIRECT_OFF").is_err() && std::env::var("NEMOTRON_FP4_SIM").is_err() @@ -3330,7 +3434,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { }; let tok_dev = match &moe_dev { Some((_,st,_)) => st, None => tok_owned.as_ref().unwrap() }; // f16-FIRST: cast [s,hid] once (or reuse the router's cast) - // then gather in f16 — 3.2x less traffic than the old + // then gather in f16, 3.2x less traffic than the old // gather-f32 [mt,hid] + cast-down chain (mt = 6s). let xn_f16g = match &xn_f16_router { Some(t) => t.clone(), @@ -3367,16 +3471,16 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { // residual add stay on-device (no dl(acc)+upm(acc) round-trip). let mut acc_dev_keep: Option = None; // set when the shared expert was folded into the routed grouped - // GEMM (acc already holds routed + shared) — skip the sd block. + // GEMM (acc already holds routed + shared), skip the sd block. let mut shared_folded = false; if use_w4a16 { // ── W4A16 WMMA path (NEMOTRON_W4A16=1 or NEMOTRON_W4A16_MARLIN=1) ── // Small S (mt ≤ W4A16_THRESH, standard path only): inline Q4 dequant - // + WMMA — wins because inline dequant saves 50% weight BW. + // + WMMA, wins because inline dequant saves 50% weight BW. // Large S (mt > W4A16_THRESH, standard path only): fall back to - // grouped_gemm (dequant-once + cuBLAS) — cuBLAS tiles are much larger. + // grouped_gemm (dequant-once + cuBLAS), cuBLAS tiles are much larger. // Marlin path (use_w4a16_marlin=true): always use moe_w4a16_marlin for - // all S (no grouped_gemm fallback — grouped_gemm expects standard + // all S (no grouped_gemm fallback, grouped_gemm expects standard // layout, but Marlin-layout weights would produce wrong results). let w4a16_thresh: usize = std::env::var("NEMOTRON_W4A16_THRESH") .ok().and_then(|v| v.parse().ok()).unwrap_or(12288); @@ -3526,7 +3630,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { // moe_q4_grouped_mma reads the signed Q4 weights directly (no // f16 slab, no dequant). ONE launch per up/down. Device gather // gives xs_f16; on-device relu²(1/256)+det-scatter(×256). - // None in FP4-direct mode (gather skipped upstream) — dummy. + // None in FP4-direct mode (gather skipped upstream), dummy. let xs_f16 = xs_dev_f16_opt.clone().unwrap_or_else(|| { Tensor::new(d.upload(&[0u8; 2]).unwrap(), vec![1], DType::F16) }); @@ -3552,7 +3656,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { let rt = nvfp4_roundtrip(&hf, rows, kk); cast_f32_f16(d, &upm(&rt, vec![rows, kk])).unwrap() }; - // FP4 direct mode: skip the gathered xs entirely — actq_pack + // FP4 direct mode: skip the gathered xs entirely, actq_pack // gathers on read via st_dev from the ungathered [s,hid] f16 // (and the amax pass scans 6x fewer elements). let fp4_direct = !fp4_sim && moe_dev.is_some() && xn_f16_router.is_some() @@ -3568,7 +3672,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { // Shared expert FOLDED into the grouped GEMM as two always- // selected groups (up: two N=inter slabs; down: K-split halves // whose scatter-adds sum to the full down GEMM). Same act quant - // (identity gather rows of xn), same deterministic scatter — + // (identity gather rows of xn), same deterministic scatter: // kills the separate shared up/down f16 GEMMs + relu2 + merge. let fold_keys = fp4_moe && fp4_direct && fp4w.contains_key(&format!("{p}.moe_up_proj_fp4_fold")); @@ -3696,7 +3800,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { } let hb = cutlass_handles.borrow(); let (hu, hd, alpha_u, alpha_d) = hb.get(p.as_str()).unwrap(); - // gather-on-read up-quant from the UNGATHERED xn — kills the + // gather-on-read up-quant from the UNGATHERED xn, kills the // explicit xs gather pass; amax scans s rows instead of mt. let (_, st_dev0, _) = moe_dev.as_ref().unwrap(); let (ap, asf, ags) = if let Some(xnr) = xn_f16_router.as_ref() { @@ -3783,7 +3887,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { let (uwp, uwsc, ugw) = &fp4w[&format!("{p}.moe_up_proj_fp4")]; let (dwp, dwsc, dgw) = &fp4w[&format!("{p}.moe_down_proj_fp4")]; // up-GEMM epilogue also computes the DOWN-GEMM's activation - // amax (of relu²-transformed outputs) — kills the full 45MB + // amax (of relu²-transformed outputs), kills the full 45MB // re-read the down amax pass used to do. let amax_dn = Tensor::new(d.alloc_zeroed(4).unwrap(), vec![1], DType::U32); let up_gidx = if fp4_direct { moe_dev.as_ref().map(|(_, st, _)| st) } else { None }; @@ -3792,7 +3896,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { // relu²·(1/256) fused into the down act-quant (no standalone pass) let dn_out = pf!(11, 2.0*mt as f64*hid as f64*inter as f64, ffai_ops::moe_fp4_grouped_mma_dev(d, &up_out, dwp, dwsc, dgw, off_dev_ref.unwrap(), n_exp, mt, hid, inter, 1, inv, Some(&amax_dn), None, None).unwrap()); - // acc_h is still all-zero on this path — seed the accumulator + // acc_h is still all-zero on this path, seed the accumulator // with a device-side memset instead of uploading 22MB of zeros. let acc_dev = Tensor::new(d.alloc_zeroed(s * hid * 4).unwrap(), vec![s, hid], DType::F32); let (_, st_dev, sw_dev) = moe_dev.as_ref().unwrap(); @@ -3822,7 +3926,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { pf!(11, 2.0*mt as f64*hid as f64*inter as f64, moe_q4_grouped_mma(d, &a2, dqs, dsc, &g_starts, &expert_ids, hid, inter).unwrap()) }; - // acc_h is still all-zero here — device-side memset, no 22MB upload. + // acc_h is still all-zero here, device-side memset, no 22MB upload. let acc_dev = Tensor::new(d.alloc_zeroed(s * hid * 4).unwrap(), vec![s, hid], DType::F32); if devdesc { // Fully on-device scatter: sorted_tok/sorted_wt (st/sw) ARE the @@ -3895,7 +3999,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { // ── Metal GROUPED Q4 GEMM path (NEMOTRON_MOE_GROUPED=1) ─────── // Token-sort already done (triples sorted by expert; xs_f16 is // the [mt, hid] f16 gather in expert order). Run ONE grouped - // Q4 MMA GEMM per pass via moe_bgemm_q4_bm64 — the kernel reads + // Q4 MMA GEMM per pass via moe_bgemm_q4_bm64, the kernel reads // each row's expert id from `idx_dev` and dequants that expert's // Q4 nibbles inline (no per-expert f16 weight slab, no per-expert // dispatch). relu²+scale stays host-f32 (squares can exceed f16 @@ -3930,7 +4034,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { } } else { // ── Per-expert cuBLAS loop ──────────────────────────────────── - // NEMOTRON_TWO_PASS=1: two-pass variant — all UP GEMMs first + // NEMOTRON_TWO_PASS=1: two-pass variant, all UP GEMMs first // (no intermediate CPU sync), then batch relu2_scale, then all // DOWN GEMMs, then batch downloads + scatter. This keeps the // cuBLAS stream uninterrupted for all ~128 experts per layer. @@ -3942,7 +4046,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { // NEMOTRON_FEWER_SYNCS=1: keep the per-expert cuBLAS UP/DOWN // GEMMs (near-best), but fuse relu² ON DEVICE (relu2_scale_f16) // and accumulate each expert group into a DEVICE acc via - // moe_scatter_add — so NOTHING is downloaded per expert. One + // moe_scatter_add, so NOTHING is downloaded per expert. One // dl(acc_dev) at the end of the layer replaces the ~2×128 // per-expert dl()/cuStreamSynchronize pairs. Implies dev relu². // DEFAULT-ON for CUDA: FEWER_SYNCS is now deterministic (the htod @@ -3966,7 +4070,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { // envelope: argmax is bit-stable across most runs but can flip on a // near-tie, because cuBLAS tensor-core GEMM (gemm_tc_off, shared by // GROUPED_GEMM) is itself run-to-run nondeterministic (split-K - // accumulation). This is NOT a race introduced here — proven by the + // accumulation). This is NOT a race introduced here, proven by the // host-relu²/host-scatter variant below jittering identically, and by // GROUPED_GEMM landing on the same alt-tokens. The per-expert host // dl()-per-expert path (default, no flag) stays bit-deterministic. @@ -4054,7 +4158,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { let acc_dev = upm(&acc_h, vec![s, hid]); // acc_h all-zero here // Deterministic scatter by default (atomicAdd is run-to-run // nondeterministic). NEMOTRON_ATOMIC_SCATTER=1 → old kernel. - // Default det path reads dn_all (f16) DIRECTLY — skips the + // Default det path reads dn_all (f16) DIRECTLY, skips the // per-layer cast_f16_f32 + [mt,hid] f32 materialization (~66MB // write/layer). Atomic fallback still needs f32, cast lazily. if std::env::var("NEMOTRON_ATOMIC_SCATTER").is_ok() { @@ -4065,7 +4169,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { moe_scatter_add_det(d, &dn_all, &tidx_h, &wts_dev, &acc_dev, s, mt, hid, 256.0f32, true).unwrap(); } if ondevice_moe { - // Keep the routed accumulator on device — the shared + // Keep the routed accumulator on device, the shared // expert + residual merge below consume it on-device. acc_dev_keep = Some(acc_dev); } else { @@ -4146,7 +4250,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { // the shared up/down qs/sc), cast activation → f16, run the // matmul in f16 (ffai_gemm[F16]). relu2 in host f32 (squares // can exceed f16 max → overflow). EXACT MATCH vs the f32 path - // (argmax unchanged). No resident cache — re-dequant each + // (argmax unchanged). No resident cache, re-dequant each // forward (the f16 cache is a bandwidth loss on Metal). let xg = pf!(15, 0.0, cast_f32_f16(d, &upm(&xs[g0*hid..g1*hid], vec![rows, hid])).unwrap()); // f16 let wup = pf!(3, (inter*hid) as f64, dequant_q4_off(d, uqs, usc, inter, hid, DType::F16, e*inter*up_bpr).unwrap()); @@ -4207,7 +4311,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { // GEMMs through the Q4-native cooperative-tensor MMA (gemm_q4_mpp, // same kernel as the dense projections / per-expert experts) instead // of dequant_q4→f32 + scalar matmul. relu2 stays in HOST f32 (squares - // can exceed f16 max → overflow) — identical to the routed per-expert + // can exceed f16 max → overflow), identical to the routed per-expert // f16 path (argmax-exact). The GEMM itself moving to Q4-native f16 // MMA is the win (~2 → ~14 TFLOP/s). Set NEMOTRON_METAL_F32_SHARED=1 // (or NEMOTRON_MOE_SHARED_MMA=0) to fall back to the f32 dequant+matmul @@ -4218,17 +4322,17 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { && shared_mma_flag.as_deref() != Some("0") && !std::env::var("NEMOTRON_METAL_F32_SHARED").is_ok(); // devdesc: fully ON-DEVICE shared expert (no dl) for graph capture. - // device cuBLAS GEMMs + device relu²(full scale — shared activations + // device cuBLAS GEMMs + device relu²(full scale, shared activations // are small, sa²≪f16max so no /256 needed) + device merge below. let mut sd_dev_opt: Option = None; let sd_h: Vec = if shared_folded { Vec::new() // shared already accumulated in the grouped scatter } else if devdesc_skip && is_cuda { - // devdesc: fully ON-DEVICE shared expert (no dl) — device cuBLAS - // (cublasLt, fast tensor-core) + device relu²(full scale — shared + // devdesc: fully ON-DEVICE shared expert (no dl), device cuBLAS + // (cublasLt, fast tensor-core) + device relu²(full scale, shared // activations small, sa²≪f16max) + device merge below. NOTE: this // is NOT CUDA-graph-capturable (the f16-out cublasLt wmma kernel - // faults on replay), but graphs are a dead end here anyway — the + // faults on replay), but graphs are a dead end here anyway: the // host-free forward alone pegs the GPU at ~98% busy, so the win is // killing the host round-trip idle, NOT graph replay. // reuse the router's f16 cast of xn (same tensor) when present @@ -4243,7 +4347,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { let fp4_shared = fp4w.contains_key(&format!("{p}.sh_up_proj_fp4")); let sd = if fp4_shared { // NVFP4 W4A4 shared expert: same grouped kernel, n_exp=1 - // (offsets [0, s] — one run covering all rows). + // (offsets [0, s], one run covering all rows). let off1 = cdev(format!("__off1.{s}"), &|| Tensor::new(d.upload(&tbu(&[0u32, s as u32])).unwrap(), vec![2], DType::U32)); let (uwp, uwsc, ugw) = &fp4w[&format!("{p}.sh_up_proj_fp4")]; let (dwp, dwsc, dgw) = &fp4w[&format!("{p}.sh_down_proj_fp4")]; @@ -4313,7 +4417,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { // Cast xn → f16 once; up GEMM [s,hid]·Wupᵀ → [s,shared_inter] f16; // relu2+(/256) in host f32 (overflow-safe); re-cast → f16; down GEMM // [s,shared_inter]·Wdnᵀ → [s,hid] f16; (×256) on host. EXACT-MATCH - // numerics vs the f32 path (the f16 up-proj output fits f16 — the + // numerics vs the f32 path (the f16 up-proj output fits f16, the // routed experts already prove this at the same activation scale). let xnh = pf!(15, 0.0, cast_f32_f16(d, &xn).unwrap()); // [s, hid] f16 let sa = pf!(12, 2.0*s as f64*shared_inter as f64*hid as f64, @@ -4328,7 +4432,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { } else if !is_cuda { // ── Portable Metal shared-expert: dequant→f32 + matmul, host relu2 ── // relu2 squares activations (can exceed f16 max) so we keep the - // intermediate in f32 throughout — no overflow, fully portable. + // intermediate in f32 throughout, no overflow, fully portable. // A/B fallback (NEMOTRON_METAL_F32_SHARED=1). let xnf = xn.clone(); // [s, hid] f32 let wsu = pf!(3, (*sm_up * *sk_up) as f64, dequant_q4(d, suqs, susc, *sm_up, *sk_up, DType::F32).unwrap()); @@ -4374,11 +4478,11 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { }; if let Some(acc_dev) = acc_dev_keep.take() { // ONDEVICE_MOE: routed acc is already on device. Merge the - // shared-expert output + residual on device — no dl/upm. Under + // shared-expert output + residual on device, no dl/upm. Under // devdesc sd is ALREADY device (sd_dev_opt, host-sync-free); // else upload the host f32 shared output. { - // fused residual add + NEXT layer's rms_norm (one pass, llama.cpp-style) + // fused residual add + NEXT layer's rms_norm (one pass, in the reference C++ implementation's style) let nw = if l + 1 < PATTERN.len() { fwd.get(&format!("language_model.backbone.layers.{}.norm.weight", l + 1)) } else { fwd.get("norm_f") }; if shared_folded { // shared already accumulated into acc by the grouped scatter @@ -4404,7 +4508,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { let acc_up = upm(&acc_h, vec![s, hid]); let acc_up = if stream16 { cast_stream(d, &acc_up) } else { acc_up }; { - // fused residual add + NEXT layer's rms_norm (one pass, llama.cpp-style) + // fused residual add + NEXT layer's rms_norm (one pass, in the reference C++ implementation's style) let nw = if l + 1 < PATTERN.len() { fwd.get(&format!("language_model.backbone.layers.{}.norm.weight", l + 1)) } else { fwd.get("norm_f") }; match nw { Some(w) if is_cuda && std::env::var("NEMOTRON_FUSED_ADDNORM").is_ok() => { let (r, nx) = ffai_ops::add_rms_norm(d, &xt, &acc_up, w, eps).unwrap(); xt = r; xn_carry = Some(nx); } @@ -4414,7 +4518,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { } } '*' => { - // Cast xn→f16 ONCE, reuse for q/k/v (NEMOTRON_FUSE_QKV) — else per-proj cast. + // Cast xn→f16 ONCE, reuse for q/k/v (NEMOTRON_FUSE_QKV): else per-proj cast. let (q_all, k_all, v_all) = if fuse_qkv { let xnh = pf!(15, 0.0, cast_f32_f16(d, &xn).unwrap()); (qmm_h(&xnh, &format!("{p}.mixer.q_proj.weight")), @@ -4482,14 +4586,14 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { pf!(7, attn_flops, ffai_ops::sdpa_flash_wmma(d, &qr.reshaped(vec![s, nq, hd]), &kcache.reshaped(vec![nkv, cap, hd]), &vcache.reshaped(vec![nkv, cap, hd]), hd, nq as u32, base as u32, s as u32, cap as u32, (nq/nkv) as u32, true, ascale).unwrap()) } _ => { // Packed: block-diagonal varlen flash-attn (each query stays in its segment). - pf!(7, attn_flops, sdpa_multi_tc_varlen(d, &qr.reshaped(vec![s, nq, hd]), &kcache.reshaped(vec![nkv, cap, hd]), &vcache.reshaped(vec![nkv, cap, hd]), seg_lo, pk_seg_len, hd, nq as u32, base as u32, s as u32, cap as u32, (nq/nkv) as u32, true, ascale).unwrap()) + pf!(7, attn_flops, sdpa_multi_tc_varlen(d, &qr.reshaped(vec![s, nq, hd]), &kcache.reshaped(vec![nkv, cap, hd]), &vcache.reshaped(vec![nkv, cap, hd]), seg_lo, pk_seg_len, hd, nq as u32, base as u32, s as u32, cap as u32, (nq/nkv) as u32, true, ascale, 0).unwrap()) } } } else if hd == 128 && std::env::var("NEMOTRON_FLASH_MMA_OFF").is_err() && std::env::var("NEMOTRON_FLASH_WMMA").is_err() && std::env::var("NEMOTRON_FLASH_FUSED").is_err() { // DEFAULT (CUDA, non-packed, hd=128): FlashAttention-2 with - // register-resident O (mma.sync) — beats cuBLAS-TC attention + // register-resident O (mma.sync), beats cuBLAS-TC attention // 1.5-3.3× (no HBM score materialization, causal-skipped, high // occupancy). Argmax-exact. NEMOTRON_FLASH_MMA_OFF=1 → cuBLAS-TC. pf!(7, attn_flops, ffai_ops::sdpa_flash_mma(d, &qr.reshaped(vec![s, nq, hd]), &kcache.reshaped(vec![nkv, cap, hd]), &vcache.reshaped(vec![nkv, cap, hd]), hd, nq as u32, base as u32, s as u32, cap as u32, (nq/nkv) as u32, true, ascale).unwrap()) @@ -4510,7 +4614,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { // attn [s, nq, hd] = [s, qdim]; o_proj batched. let o = qmm(&attn.reshaped(vec![s, qdim]), &format!("{p}.mixer.o_proj.weight")); { - // fused residual add + NEXT layer's rms_norm (one pass, llama.cpp-style) + // fused residual add + NEXT layer's rms_norm (one pass, in the reference C++ implementation's style) let nw = if l + 1 < PATTERN.len() { fwd.get(&format!("language_model.backbone.layers.{}.norm.weight", l + 1)) } else { fwd.get("norm_f") }; match nw { Some(w) if is_cuda && std::env::var("NEMOTRON_FUSED_ADDNORM").is_ok() => { let (r, nx) = ffai_ops::add_rms_norm(d, &xt, &o, w, eps).unwrap(); xt = r; xn_carry = Some(nx); } @@ -4529,7 +4633,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { BATCHED_LAYER_TRACE.with(|c| { let mut v = c.borrow_mut(); if v.len() <= l { v.resize(l+1, Vec::new()); } v[l] = h; }); } } - // final norm + lm_head on the LAST token only — DEVICE-CLEAN (no dl→up + // final norm + lm_head on the LAST token only, DEVICE-CLEAN (no dl→up // roundtrip): slice the last token on device, lm_head GEMM on device → // device logits. Store the device logits for the graph path; dl for argmax // ONLY when not capturing (a dl aborts cuda-graph stream capture). @@ -4554,7 +4658,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { // Fixed prompt id list (deterministic ramp). The correctness gate runs BOTH // the sequential and batched paths over THIS SAME list and compares the - // last-token argmax — the KV cache + conv/SSM final states must agree. + // last-token argmax, the KV cache + conv/SSM final states must agree. let ids: Vec = { let base_ids: Vec = (0..s).map(|i| (tok + i) % vocab).collect(); // NEMOTRON_PACKED=N: replicate the prompt into N packed sequences (s·N @@ -4650,7 +4754,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { // A benign precision flip = the two argmaxes are a reshuffle WITHIN the shared // top-5 (same candidate set), not a wrong prediction. Accumulated f16/Q4 drift // over long S widens the top-1/top-2 gap (~5% at S=2048) while the top-5 stays - // identical (cosine >0.998) — still benign. So key the verdict on top-5 agreement, + // identical (cosine >0.998), still benign. So key the verdict on top-5 agreement, // not just a fixed gap%. (Fixed-2% flagged the S=2048 near-tie as a false "bug".) let top5q = |v: &[f32]| { let mut idx: Vec = (0..v.len()).collect(); idx.sort_by(|&a,&b| v[b].partial_cmp(&v[a]).unwrap()); idx.truncate(5); idx }; @@ -4702,7 +4806,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { eprintln!(" seq-ref top5 = {:?}", t5s.iter().map(|&i| (i, seq_logits[i])).collect::>()); eprintln!(" batched top5 = {:?}", t5b.iter().map(|&i| (i, blog[i])).collect::>()); } else { - eprintln!(" LOGIT-LEVEL: length mismatch seq={} bat={} — skipped", seq_logits.len(), blog.len()); + eprintln!(" LOGIT-LEVEL: length mismatch seq={} bat={}, skipped", seq_logits.len(), blog.len()); } // ── PER-LAYER DIVERGENCE TRACE (NEMOTRON_DUMP_LAYERS=1; REVERT) ─────── if std::env::var("NEMOTRON_DUMP_LAYERS").is_ok() { @@ -4757,8 +4861,8 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { for kv in kvcache.iter_mut() { *kv = None; } ROUTER_HOST_T.with(|c| c.set(0.0)); // NEMOTRON_GRAPH=1 (CUDA, needs NEMOTRON_Q4_DEVDESC for a dl-free forward): - // capture the whole forward into a CUDA graph, then time the REPLAY — kills - // the host-launch-serialized GPU idle (the ~2× prefill lever vs vLLM). + // capture the whole forward into a CUDA graph, then time the REPLAY: kills + // the host-launch-serialized GPU idle (the ~2× prefill lever vs. a reference serving framework). let graph_mode = std::env::var("NEMOTRON_GRAPH").is_ok() && d.backend() == Backend::Cuda; let (next, pf_s) = if graph_mode { // Pre-warm the caching pool to its high-water mark so NO cuMemAlloc fires @@ -4799,7 +4903,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { } if let Ok(p) = std::env::var("NEMOTRON_BF16_REF") { let nf = bl.iter().filter(|x| !x.is_finite()).count(); - if nf > 0 { eprintln!(" [BF16-REF] {nf} non-finite logits — skip"); } + if nf > 0 { eprintln!(" [BF16-REF] {nf} non-finite logits, skip"); } else if let Ok(rb) = std::fs::read(&p) { let rl: Vec = rb.chunks(4).map(|c| f32::from_le_bytes([c[0],c[1],c[2],c[3]])).collect(); if rl.len()==bl.len() && !bl.is_empty() { @@ -4878,7 +4982,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { md.push_str(&format!("- Clean batched throughput: **{:.1} tok/s** ({:.2} ms/tok)\n", s as f64 / pf_s, pf_s * 1000.0 / s as f64)); md.push_str(&format!("- Profiled pass wall (sync-bracketed, inflated): {prof_wall:.3}s; summed op time: {sum_t:.3}s\n")); if is_cuda { - md.push_str(&format!("- vLLM-on-GB10 reference: pp2048@d0=6395, @d8192=4993, @d32768=2734 tok/s\n")); + md.push_str(&format!("- Reference serving-framework numbers on GB10: pp2048@d0=6395, @d8192=4993, @d32768=2734 tok/s\n")); md.push_str(&format!("- Tensor-core peak assumed: {PEAK_TFLOPS:.0} TFLOP/s (bf16 dense)\n\n")); } else { md.push_str("\n"); @@ -4900,7 +5004,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { md.push_str("- **`dequant_q4` ~26% (5500 calls):** Q4 weights are re-dequanted to f16 EVERY forward (per layer + per routed expert). This is pure redundant work → cache the f16-dequanted weights resident once at load (trades VRAM for ~26% of prefill time).\n"); md.push_str("- `sdpa_prefill`: software-MMA coop_tile by default; set NEMOTRON_PREFILL_TCATTN=1 for the cuBLAS tensor-core FlashAttention path (`sdpa_multi_tc`). At deep context (zero-KV synthetic) the TC path cuts this stage ~9.5× @d8192 (1754→185ms, 54.7%→11.3%), ~14× @d16384 (4616→328ms), ~13.7× @d32768 (8958→652ms), lifting end-to-end prefill from 104→165 tok/s @d32768.\n"); md.push_str("- `moe_experts` ~12%: the GEMM is fast (28 TFLOP/s) but the per-expert cuBLAS loop + host relu2/scatter round-trips remain serial → fuse on-device or use cublasGemmStridedBatched.\n"); - md.push_str("- Throughput peaks at S=512 then declines (S=2048→178, S=8192→117): the per-expert host round-trips + dequant grow with S/expert-count. Removing dequant-per-forward + the host MoE bridges should restore monotonic scaling toward the vLLM band.\n"); + md.push_str("- Throughput peaks at S=512 then declines (S=2048→178, S=8192→117): the per-expert host round-trips + dequant grow with S/expert-count. Removing dequant-per-forward + the host MoE bridges should restore monotonic scaling toward the reference framework's throughput band.\n"); } else { md.push_str("\n## Notes (Apple/Metal path)\n"); md.push_str("- GEMMs run on Apple GPU hardware MMA: dense projections via `gemm_q4_mpp` (Q4-native cooperative-tensor MMA), per-expert/shared MoE via dequant_q4_off + `matmul`. No cuBLAS/tensor-core escape hatches (those are CUDA-only).\n"); @@ -4931,7 +5035,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { tok = first; pos += 1; if prefill == 0 && fakectx == 0 { eprintln!("Nemotron decode: token1234 → next argmax {first} (F32 ref argmax=1234; Q8 may perturb the near-tie)"); } // Fast internal A/B: one model load, toggle an env flag (e.g. MT_GEMV_2ROW) - // between alternating decode batches IN-PROCESS — same thermal/clock state, + // between alternating decode batches IN-PROCESS, same thermal/clock state, // order-alternated to cancel drift. Resets pos each batch so the KV cap holds. if let Ok(ab) = std::env::var("NEMOTRON_AB") { // "on" sets the flag to NEMOTRON_AB_VAL (default "1"); "off" unsets it. @@ -5072,16 +5176,16 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { // f. Print eager_tps_graph_run, graph_tps, ratio. if use_graph { if no_devrouter { - eprintln!("NEMOTRON_GRAPH: WARNING — NEMOTRON_DEVROUTER not set. Host router mid-token sync will break capture. Re-run with NEMOTRON_DEVROUTER=1."); + eprintln!("NEMOTRON_GRAPH: WARNING, NEMOTRON_DEVROUTER not set. Host router mid-token sync will break capture. Re-run with NEMOTRON_DEVROUTER=1."); } let gm = n_decode; // use same step count for apples-to-apples // Use `fakectx` as the fixed position for all graph-section measurements. // This keeps `pos` within bounds (KV cap = fakectx + n_decode + 8) AND gives // the correct 32K SDPA length. We don't advance pos in the graph section. let gpos = fakectx.max(prefill); // fixed position for graph-section steps - // (a) 5 warm steps at gpos (untimed) — populates pool, reaches steady clocks. + // (a) 5 warm steps at gpos (untimed), populates pool, reaches steady clocks. for _ in 0..5 { step(tok, gpos, &mut conv_state, &mut ssm_state, &mut kvcache); } - // (b) Eager baseline at fixed gpos (skip_dl=false — argmax fires for comparison). + // (b) Eager baseline at fixed gpos (skip_dl=false, argmax fires for comparison). let t_eager = Instant::now(); for _ in 0..gm { step(tok, gpos, &mut conv_state, &mut ssm_state, &mut kvcache); } d.synchronize().ok(); @@ -5091,10 +5195,10 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { step(tok, gpos, &mut conv_state, &mut ssm_state, &mut kvcache); // (d) Capture: skip_dl = true so no host-sync in step. skip_dl.set(true); - d.begin_capture().expect("NEMOTRON_GRAPH: begin_capture failed — check DEVROUTER and that no cuMemAlloc slipped into the step"); + d.begin_capture().expect("NEMOTRON_GRAPH: begin_capture failed, check DEVROUTER and that no cuMemAlloc slipped into the step"); step(tok, gpos, &mut conv_state, &mut ssm_state, &mut kvcache); - let graph_handle = d.end_capture().expect("NEMOTRON_GRAPH: end_capture failed — check for stray sync/alloc in the captured region"); - // (e) Synchronize then time M graph launches — two modes: + let graph_handle = d.end_capture().expect("NEMOTRON_GRAPH: end_capture failed, check for stray sync/alloc in the captured region"); + // (e) Synchronize then time M graph launches, two modes: // serial: graph_launch (sync-per-token) measures single-token latency // batched: graph_launch_batch (N enqueues, one sync) eliminates per-token // host-GPU handoff, giving maximum pipeline throughput @@ -5104,7 +5208,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { for _ in 0..gm { d.graph_launch(graph_handle).expect("graph_launch failed"); } let graph_serial_dt = t_graph_serial.elapsed().as_secs_f64(); let graph_serial_tps = gm as f64 / graph_serial_dt; - // Batched: N cuGraphLaunch enqueues, one cuStreamSynchronize — removes + // Batched: N cuGraphLaunch enqueues, one cuStreamSynchronize: removes // per-token host-GPU handoff to reveal the true GPU throughput ceiling. d.synchronize().ok(); let t_graph_batch = Instant::now(); diff --git a/rust/crates/ffai-ops/src/lib.rs b/rust/crates/ffai-ops/src/lib.rs index 2f614720..ddcd3ca5 100644 --- a/rust/crates/ffai-ops/src/lib.rs +++ b/rust/crates/ffai-ops/src/lib.rs @@ -3,7 +3,7 @@ //! # ffai-ops //! -//! The **seam** between model code and kernels — the Rust analog of +//! The **seam** between model code and kernels, the Rust analog of //! FFAI-Swift's `Ops/`. Each function takes Tensors, builds (or looks up) //! the corresponding metaltile [`Kernel`](ffai_core::Kernel), and dispatches //! it through the [`Device`](ffai_core::Device) trait. Model code calls @@ -50,7 +50,7 @@ mod ssd_scan_portable; pub use ssd_scan_portable::ssm_prefill_scan_ssd_portable; /// Cache of resolved kernel IR keyed by (name, dtype). Resolving walks the -/// whole test registry building setups, which is expensive — a forward pass +/// whole test registry building setups, which is expensive: a forward pass /// dispatches the same handful of kernels hundreds of times, so cache them. /// FxHashMap (single-cycle hash) + parking_lot Mutex (unpoisoned, inlinable). fn kernel_cache() -> &'static Mutex> { @@ -59,7 +59,7 @@ fn kernel_cache() -> &'static Mutex> { } /// Cache an inline-built kernel IR by (name, dtype). `kernel_ir_for` rebuilds the -/// full IR (~30+ ops) on every call — on the decode hot path that's ~250 +/// full IR (~30+ ops) on every call, on the decode hot path that's ~250 /// rebuilds/token. Build once, clone thereafter. fn cached_ir(name: &str, dtype: DType, build: impl FnOnce() -> Kernel) -> Kernel { let key = (name.to_string(), dtype); @@ -240,7 +240,7 @@ extern "C" __global__ void ffai_add3( { long i=(long)blockIdx.x*blockDim.x+threadIdx.x; long stride=(long)gridDim.x*blockDim.x; - // association matches the unfused path exactly: xt + (acc + sd) — a + // association matches the unfused path exactly: xt + (acc + sd), a // left-assoc a+b+c flipped the d32768 argmax via last-ulp drift. for(;i Result { @@ -297,7 +297,7 @@ pub fn swiglu_limit(dev: &dyn Device, gate: &Tensor, up: &Tensor, limit: f32) -> } /// DSv4 MoE router scoring: `score_unbiased[e] = sqrt(softplus(logit[e]))`, -/// `score_biased[e] = score_unbiased[e] + bias[e]`. Computed **host-side** — +/// `score_biased[e] = score_unbiased[e] + bias[e]`. Computed **host-side**, /// the router operates on the tiny `[n_experts]` logit vector that the MoE /// already downloads for top-k selection, so a GPU kernel buys nothing and /// avoids the multi-output dispatch path. `biased` selects the top-k experts; @@ -457,7 +457,7 @@ pub fn dsv4_mhc_sinkhorn_split( } -/// mHC collapse: `out[d] = Σ_c pre[c] · state[c, d]` — mix the 4-channel +/// mHC collapse: `out[d] = Σ_c pre[c] · state[c, d]`, mix the 4-channel /// residual state `[n_hc, hidden]` down to `[hidden]` using the per-channel /// `pre` weights. Single token. Dispatches `ffai_dsv4_mhc_collapse`. pub fn dsv4_mhc_collapse( @@ -520,13 +520,13 @@ pub fn dsv4_mhc_expand( Ok(state) } -// ── Heavier ops — dispatch the registered metaltile kernels ───────────── +// ── Heavier ops, dispatch the registered metaltile kernels ───────────── /// Row-wise RMS norm: `out[r] = x[r] * rsqrt(mean(x[r]²) + eps) * weight`. -/// Dispatches the registered `mt_rms_norm` reduction kernel — the same one +/// Dispatches the registered `mt_rms_norm` reduction kernel, the same one /// the Swift side runs. The last dim is the row width `n`; the kernel owns 4 /// elements per thread, so `n` must be a multiple of 128 and ≤ 4096 (the -/// `mt_rms_norm_wide` variant lifts this — wired later). +/// `mt_rms_norm_wide` variant lifts this, wired later). pub fn rms_norm(dev: &dyn Device, x: &Tensor, weight: &Tensor, eps: f32) -> Result { let n = *x.shape.last().ok_or_else(|| Error::Msg("rms_norm: scalar input".into()))?; let rows = x.elem_count() / n; @@ -590,7 +590,7 @@ pub fn gemv(dev: &dyn Device, mat: &Tensor, vec: &Tensor) -> Result { Ok(out) } -/// Q8_0 grouped matvec — weights stay QUANTIZED resident (int8, block 32, one +/// Q8_0 grouped matvec, weights stay QUANTIZED resident (int8, block 32, one /// f32 scale/block): `out[r] = Σ_k dequant(qs[r,k], scales[r, k/32]) · /// x[(r/rows_per_group)·k_in + k]`. `qs` is u32-packed (4 int8/u32, 8 u32/block), /// `scales` is `[m_out · k_in/32]` f32. Dense gemv ⇒ `rows_per_group = m_out` @@ -633,11 +633,11 @@ pub fn gemv_q8( Ok(out) } -/// Batched Q8 GEMM — the **prefill** projection path. `out[r,o] = Σ_k W[o,k]·x[r,k]` +/// Batched Q8 GEMM, the **prefill** projection path. `out[r,o] = Σ_k W[o,k]·x[r,k]` /// for `n_rows` tokens at once, dispatching metaltile's `ffai_gemm_q8_mpp` /// (SimdGroup-scope CoopTile 64×64×32 tile). On the Vulkan/RDNA4 backend this /// kernel picks up the gated `VK_KHR_cooperative_matrix` fragment ops when -/// `MT_VK_COOPMAT=1` (≈6× the scalar GEMM) — the realized prefill win. With the +/// `MT_VK_COOPMAT=1` (≈6× the scalar GEMM), the realized prefill win. With the /// gate unset it runs the bit-exact scalar `SoftwareLocalC` tiles. Either way it /// replaces the `n_rows`× sequential `gemv_q8` matvecs of the per-token path. /// @@ -696,19 +696,19 @@ pub fn gemm_q8_mpp( /// `out[r,o] = x[r,o] + bias[o]`. The decode path adds a `[n]` bias to a single /// `[n]` projection with plain [`add`]; the batched prefill path needs the same /// bias on every token row, but [`add`] is shape-strict (no broadcast). This is -/// a fully on-device elementwise kernel (`out[i] = x[i] + bias[i % n]`) — no +/// a fully on-device elementwise kernel (`out[i] = x[i] + bias[i % n]`), no /// host round-trip, so it works on device-local (non-mappable) weight buffers. /// Returns `[n_rows, n]` f32. pub fn add_bias_rows(dev: &dyn Device, x: &Tensor, bias: &Tensor, n_rows: usize, n: usize) -> Result { // Fully on-device per-feature bias broadcast: `i = ProgramId(0)` (global // flat id), row-local column `col = i % n`, `out[i] = x[i] + bias[col]`. - // `n` is a CONSTEXPR push-constant (bound as a `Binding::Scalar`) — NOT a - // baked `Op::Const` literal — so the one cached `ffai_add_bias_rows` + // `n` is a CONSTEXPR push-constant (bound as a `Binding::Scalar`), NOT a + // baked `Op::Const` literal, so the one cached `ffai_add_bias_rows` // pipeline serves every projection width (q: nq*hd, k/v: nkv*hd) instead // of silently reusing the first call's `n`. The DSL `%` lowers to // `BinOpKind::Mod` → GLSL `uint(i) % uint(n)` (integer modulo), validated // bit-exact on Vulkan by `metaltile-std/tests/vulkan_add_bias_rows.rs`. - // No host round-trip — runs straight on the device-local weight buffer. + // No host round-trip, runs straight on the device-local weight buffer. use metaltile_core::ir::ConstExprDecl; use metaltile_core::ConstExpr; let total = n_rows * n; @@ -780,7 +780,7 @@ pub fn add_bias_rows(dev: &dyn Device, x: &Tensor, bias: &Tensor, n_rows: usize, Ok(out.reshaped(vec![n_rows, n])) } -/// Q8 gemv with fused ReLU²: `out[r] = max(0, (Wq·x)[r])²` — a MoE expert's +/// Q8 gemv with fused ReLU²: `out[r] = max(0, (Wq·x)[r])²`, a MoE expert's /// `up` projection + activation in one dispatch. Dispatches `ffai_gemv_q8_coalesced_relu2`. pub fn gemv_q8_relu2(dev: &dyn Device, qs: &Tensor, scales: &Tensor, x: &Tensor, m_out: usize, k_in: usize, rows_per_group: usize) -> Result { let k = cached_ir("ffai_gemv_q8_coalesced_relu2", x.dtype, || { let mut k = metaltile_std::ffai::gemv_q8::ffai_gemv_q8_coalesced_relu2::kernel_ir_for(x.dtype); k.mode = metaltile_core::ir::KernelMode::Reduction; k }); @@ -793,7 +793,7 @@ pub fn gemv_q8_relu2(dev: &dyn Device, qs: &Tensor, scales: &Tensor, x: &Tensor, /// Q8 gemv that scales + accumulates in place: `acc[r] += scale · (Wq · x)[r]`. /// Fuses an MoE expert's `down` projection with its router-weighted sum into the -/// layer accumulator — no per-expert scalar-broadcast upload or separate add. +/// layer accumulator, no per-expert scalar-broadcast upload or separate add. /// `scale_buf` is a 1-element device buffer. Dispatches `ffai_gemv_q8_coalesced_accum`. #[allow(clippy::too_many_arguments)] pub fn gemv_q8_accum( @@ -833,7 +833,7 @@ pub fn gemv_q8_accum( /// Append one decode step's K (or V) into an in-place device KV cache: /// `dst[h, pos, :] = src[h, :]`, where `dst` is `[nkv, cap, hd]`, `src` is /// `[nkv*hd]`, and `posbuf` is a 1-element u32 device buffer. Keeps the growing -/// context entirely on-device — no host reorg/reupload per step (the 32K-context +/// context entirely on-device, no host reorg/reupload per step (the 32K-context /// fix). Dispatches `ffai_kv_append` (runtime `pos` ⇒ compiled once, not per step). pub fn kv_append(dev: &dyn Device, src: &Tensor, dst: &Tensor, posbuf: &Tensor, hd: usize, cap: usize, n: usize) -> Result<()> { let k = cached_ir("ffai_kv_append", src.dtype, || { let mut k = metaltile_std::ffai::gemv_q8::ffai_kv_append::kernel_ir_for(src.dtype); k.mode = metaltile_core::ir::KernelMode::Grid3D; k }); @@ -849,7 +849,7 @@ pub fn slice(dev: &dyn Device, src: &Tensor, off: usize, len: usize) -> Result Result { Ok(out) } -/// Device Mamba dt: `softplus(dt_raw + dt_bias)` — no host round-trip. +/// Device Mamba dt: `softplus(dt_raw + dt_bias)`, no host round-trip. pub fn softplus_add(dev: &dyn Device, a: &Tensor, b: &Tensor) -> Result { let k = cached_ir("ffai_softplus_add", DType::F32, || { let mut k = metaltile_std::ffai::gemv_q8::ffai_softplus_add::kernel_ir(); k.mode = metaltile_core::ir::KernelMode::Grid3D; k }); let n = a.elem_count(); @@ -1133,7 +1133,7 @@ pub fn conv1d_causal_prefill_split( } /// Extract `[s, width]` columns from row-major `src [s, stride]` at column offset `col_off`. -/// Returns `dst [s * width]` — a contiguous slab. Used to carve z, xbc, dt_raw out of +/// Returns `dst [s * width]`, a contiguous slab. Used to carve z, xbc, dt_raw out of /// the [s, in_proj_out] projection matrix on device (NEMOTRON_CONV_DEVICE path). pub fn strided_col_copy( dev: &dyn Device, @@ -1145,10 +1145,10 @@ pub fn strided_col_copy( ) -> Result { let kernel = lookup("strided_col_copy", DType::F32)?; // Defensive bounds check. The kernel reads `src[ti*stride + col_off + ci]` - // for ti in [0, s), ci in [0, width) — a sub-slab addressed purely via the + // for ti in [0, s), ci in [0, width): a sub-slab addressed purely via the // host stride/col_off against the RAW bound buffer (tensor.offset ignored by // design). A caller binding a row-stride/offset/width that overruns the - // bound buffer triggers an out-of-bounds READ — silent garbage on Metal, a + // bound buffer triggers an out-of-bounds READ, silent garbage on Metal, a // deterministic MMU-fault on CUDA. Fail loudly, naming the shapes. if s != 0 && width != 0 { let need = (s - 1) * stride + col_off + width; // max read is src[need - 1] @@ -1162,6 +1162,7 @@ pub fn strided_col_copy( } let dst = Tensor::empty(dev, vec![s * width], DType::F32)?; let u = |v: u32| Binding::Scalar(v.to_le_bytes().to_vec()); + let total = (s * width) as u32; dev.dispatch( &kernel, &[ @@ -1170,17 +1171,96 @@ pub fn strided_col_copy( u(stride as u32), u(col_off as u32), u(width as u32), + u(total), ], - // PERF: 64 threads/block (was 1). `width` is one of {di=4096, ng*ds=1024, - // m_nh=64} — all multiples of 64, so s*width is divisible by 64 → exact grid. - { let n = (s * width) as u32; let b = 64u32; Grid { grid: [n / b, 1, 1], block: [b, 1, 1] } }, + // PERF: 64 threads/block. Grid is div_ceil'd up to a whole block and + // the kernel guards `idx < total` internally, so this is correct for + // any s*width, not just multiples of 64 (e.g. Laguna's non-64-aligned + // gate-column extraction, head counts 48/72). + { let b = 64u32; Grid { grid: [total.div_ceil(b), 1, 1], block: [b, 1, 1] } }, + )?; + Ok(dst) +} + +const EXTRACT_COLS_SRC: &str = r#" +extern "C" __global__ void ffai_extract_cols( + const float* __restrict__ src, float* __restrict__ dst, + int rows, int stride, int col_off, int width) +{ + long i = (long)blockIdx.x * blockDim.x + threadIdx.x; + long total = (long)rows * width; + if (i >= total) return; + int c = (int)(i % width); + long r = i / width; + dst[i] = src[r * stride + col_off + c]; +} +"#; + +// f16 mirror of EXTRACT_COLS_SRC: same indexing, 16-bit loads/stores. Added +// so prefill's Marlin/grouped-MoE-GEMM column splits (both f16-native +// producers) can extract without a whole-buffer cast_f16_f32 first, see +// `extract_cols`'s doc. +const EXTRACT_COLS_F16_SRC: &str = r#" +#include +extern "C" __global__ void ffai_extract_cols_f16( + const __half* __restrict__ src, __half* __restrict__ dst, + int rows, int stride, int col_off, int width) +{ + long i = (long)blockIdx.x * blockDim.x + threadIdx.x; + long total = (long)rows * width; + if (i >= total) return; + int c = (int)(i % width); + long r = i / width; + dst[i] = src[r * stride + col_off + c]; +} +"#; + +/// Extract a `[rows, width]` column range starting at `col_off` from every +/// row of a `[rows, stride]` row-major buffer (`stride >= col_off + width`). +/// Used to split Laguna's fused `[T, n_q_rows+2*n_kv_rows+n_g_rows]` +/// QKVG-concat batched-prefill projection back into its Q/K/V/gate column +/// ranges (and, in f16, to split a Marlin/grouped-MoE-GEMM's fused +/// `[gate;up]` output without a cast first). Originally added as a +/// workaround for a `strided_col_copy` launch-geometry bug (fixed grid +/// dropped the tail whenever `rows*width` wasn't a multiple of 64, true for +/// Laguna's gate column range, whose width is the per-layer head count of 48 +/// or 72, neither 64-aligned, times an arbitrary chunk length). +/// `strided_col_copy` now div_ceils its grid and guards in-kernel, so this +/// raw CUDA path is kept as-is only to avoid touching a call site mid-flight +/// elsewhere; this kernel's grid is sized to cover every element regardless +/// of `rows*width`'s divisibility. f32 and f16 (dispatched on `src.dtype`); +/// any other dtype errors. +pub fn extract_cols(dev: &dyn Device, src: &Tensor, rows: usize, stride: usize, col_off: usize, width: usize) -> Result { + let (elem_bytes, cuda_src, prog_name, fn_name) = match src.dtype { + DType::F32 => (DType::F32.size_bytes(), EXTRACT_COLS_SRC, "extract_cols.cu", "ffai_extract_cols"), + DType::F16 => (DType::F16.size_bytes(), EXTRACT_COLS_F16_SRC, "extract_cols_f16.cu", "ffai_extract_cols_f16"), + other => return Err(Error::Msg(format!("extract_cols: unsupported dtype {other:?} (f32/f16 only)"))), + }; + if rows != 0 && width != 0 { + let need = (rows - 1) * stride + col_off + width; // max read is src[need - 1] + let have = src.buffer.len() / elem_bytes; + if need > have { + return Err(Error::Msg(format!( + "extract_cols OOB: rows={rows} stride={stride} col_off={col_off} width={width} \ + needs {need} {:?} elems but bound buffer holds {have}", src.dtype + ))); + } + } + let dst = Tensor::empty(dev, vec![rows * width], src.dtype)?; + let i = |x: i32| x.to_le_bytes().to_vec(); + let total = (rows * width) as u32; + dev.dispatch_raw_cuda( + cuda_src, prog_name, fn_name, + &[(src.buffer.as_ref(), src.offset), (dst.buffer.as_ref(), 0)], + &[i(rows as i32), i(stride as i32), i(col_off as i32), i(width as i32)], + [total.div_ceil(256).max(1), 1, 1], [256, 1, 1], 0, false, )?; Ok(dst) } /// Batched softplus + tiled bias: `dst[ti*n + hi] = softplus(src[ti*n + hi] + bias[hi])`. /// Converts the raw `[s, m_nh]` dt_raw tensor + dt_bias to `dt_all [s, m_nh]` on device. -/// (NEMOTRON_CONV_DEVICE path — replaces the CPU softplus loop.) +/// (NEMOTRON_CONV_DEVICE path, replaces the CPU softplus loop.) pub fn softplus_add_rows( dev: &dyn Device, src: &Tensor, @@ -1400,7 +1480,7 @@ extern "C" __global__ void moe_fused_ffn( { cg::grid_group gg = cg::this_grid(); - // Global warp id — one warp per acc[h] for phase 2. + // Global warp id, one warp per acc[h] for phase 2. int tid = (int)(blockIdx.x * blockDim.x + threadIdx.x); int lane = threadIdx.x & 31; int warp_id = tid >> 5; @@ -1520,7 +1600,7 @@ pub fn moe_fused_ffn( /// Router-weighted sum of per-expert down outputs into `acc`: `acc[h] += Σ wts·downs[·,h]`. /// Fully ON-DEVICE MoE router (NemotronH sigmoid + e_score_correction_bias, top-k by /// biased score, weights from unbiased renormalized to sum-1, ×routed_scaling_factor). -/// Returns (idx[top_k] U32, wts[top_k] F32) WITHOUT a host round-trip — replaces the +/// Returns (idx[top_k] U32, wts[top_k] F32) WITHOUT a host round-trip, replaces the /// per-layer dl(gate)+host-topk+up(idx) sync that drains the async pipeline 23×/token. pub fn moe_router_device(dev: &dyn Device, gate_logits: &Tensor, bias: &Tensor, n_exp: usize, top_k: usize, scale: f32) -> Result<(Tensor, Tensor)> { use metaltile_core::ir::KernelMode; @@ -1542,12 +1622,12 @@ pub fn moe_router_device(dev: &dyn Device, gate_logits: &Tensor, bias: &Tensor, /// On-device MoE route + counting-sort. Replaces the per-E-layer HOST `triples` /// round-trip (router logits → host → sigmoid+topk+sort → upload) that stalls the /// GPU every MoE layer and blocks CUDA-graph capture of prefill. Given gate logits -/// `[s, n_exp]` (f32) + bias `[n_exp]`, produces — ALL on device: -/// sorted_tok `[mt]` u32 — token index per sorted slot (mt = s*top_k) -/// sorted_wt `[mt]` f32 — router weight per slot (sigmoid, normalized, ·scale) -/// offsets `[n_exp+1]` u32 — group boundaries (expert e's rows = [off[e],off[e+1])) +/// `[s, n_exp]` (f32) + bias `[n_exp]`, produces (ALL on device): +/// sorted_tok `[mt]` u32, token index per sorted slot (mt = s*top_k) +/// sorted_wt `[mt]` f32, router weight per slot (sigmoid, normalized, ·scale) +/// offsets `[n_exp+1]` u32, group boundaries (expert e's rows = [off[e],off[e+1])) /// Tokens grouped contiguously by expert. Order within an expert is atomic-cursor -/// (not stable) — same near-tie nondeterminism as the existing GROUPED_GEMM path. +/// (not stable), same near-tie nondeterminism as the existing GROUPED_GEMM path. const MOE_BATCHED_ROUTER_SRC: &str = r#" extern "C" __global__ void moe_batched_router( const float* __restrict__ logits, const float* __restrict__ bias, @@ -1734,7 +1814,7 @@ pub fn moe_weighted_sum(dev: &dyn Device, downs: &Tensor, wts: &Tensor, acc: &Te Ok(()) } -// ── Q4 (4-bit) family — half the weight DRAM of Q8 (the decode lever). ── +// ── Q4 (4-bit) family, half the weight DRAM of Q8 (the decode lever). ── fn gemv_q4_dispatch(dev: &dyn Device, kernel: &str, qs: &Tensor, scales: &Tensor, x: &Tensor, acc: Option<&Tensor>, scale_buf: Option<&Tensor>, m_out: usize, k_in: usize, rows_per_group: usize) -> Result { // 2-row tiling: 2 output rows per warp, shared activation read, 2 weight // streams in flight (more memory-level-parallelism on the latency-bound Q4 read). @@ -1772,7 +1852,7 @@ fn gemv_q4_dispatch(dev: &dyn Device, kernel: &str, qs: &Tensor, scales: &Tensor } else { // relu2 / accum: multi-warp capable (rows_per_tg warps/TG). The shared- // expert matrices (3712×2688, 2688×3712) are small; rpt>1 gives no - // measured gain (these kernels are NOT latency-bound — the ~50% BW the + // measured gain (these kernels are NOT latency-bound, the ~50% BW the // sync-bracketed profiler showed is a measurement artifact, the kernel // time is small). Default rpt=1 (bit-identical to the original single- // warp form); MT_GEMV_SHARED_RPT overrides for experimentation. @@ -1879,7 +1959,7 @@ fn enc_e2m1(v: f32) -> u8 { } /// Encode a non-negative scale `s` to a UE4M3 byte (4 exp bits bias-7, 3 mantissa, -/// no sign — NVFP4's micro-scale type). Round-to-nearest over representable values. +/// no sign, NVFP4's micro-scale type). Round-to-nearest over representable values. #[inline] fn enc_ue4m3(s: f32) -> u8 { if s <= 0.0 { return 0; } @@ -1911,7 +1991,7 @@ pub fn dec_ue4m3(b: u8) -> f32 { /// /// Returns `(codes, scales, global)`: /// - `codes`: `[m · k/8]` u32, 8 e2m1 nibbles per u32 (nibble j of word w = element -/// `w*8 + j` of the row), row-major — same packing convention as [`quantize_q4`]. +/// `w*8 + j` of the row), row-major: same packing convention as [`quantize_q4`]. /// - `scales`: `[m · k/16]` UE4M3 bytes (one per 16-element K-block, row-major). /// - `global`: per-tensor f32; the true block scale is `dec_ue4m3(scale) * global`. /// Chosen so each block's `amax/6` lands in UE4M3's representable range. @@ -2015,7 +2095,7 @@ pub fn dequantize_nvfp4(codes: &[u32], scales: &[u8], global: f32, m: usize, k: out } -/// Round-trip a row-major `[m, k]` f32 buffer through NVFP4 quantize→dequantize — +/// Round-trip a row-major `[m, k]` f32 buffer through NVFP4 quantize→dequantize, /// i.e. inject exactly the E2M1+UE4M3 block-16 quantization noise the native FP4 /// GEMM would see, without any kernel. Used to de-risk argmax before building the /// real FP4 path (`NEMOTRON_FP4_SIM`). @@ -2025,9 +2105,9 @@ pub fn nvfp4_roundtrip(w: &[f32], m: usize, k: usize) -> Vec { } /// Prefill linear: `out[r, :] = weight · input[r, :]` for a block of rows in -/// one dispatch — `weight` is `[out_dim, in_dim]`, `input` is `[n_rows, in_dim]`, +/// one dispatch, `weight` is `[out_dim, in_dim]`, `input` is `[n_rows, in_dim]`, /// result `[n_rows, out_dim]`. Dispatches the 32×32-tiled Reduction kernel -/// `ffai_gemm` (weight read once per tile, reused across the row block — the +/// `ffai_gemm` (weight read once per tile, reused across the row block: the /// many-token path that a `gemv`-per-row would re-stream). Requires /// `in_dim % 16 == 0` (the K-tile contract); row/out-dim edges are handled /// in-kernel. This is the matmul the prefill + VLM/ViT towers run on. @@ -2066,7 +2146,7 @@ pub fn matmul(dev: &dyn Device, weight: &Tensor, input: &Tensor) -> Result12.5KB) - int bx=blockIdx.x, by=blockIdx.y; + // band_n=0 (default): grid=[tile,nblk] with tile fast-varying -> a full + // tile sweep (crossing experts) happens before the N-block advances. + // band_n=1 (FFAI_LAGUNA_SCHED=1): grid=[nblk,tile] with nblk fast-varying -> + // a BAND of consecutive CTAs sweeps every N-block of the SAME tile (same + // expert weight rows) before moving to the next tile, so the immediately + // prior tile of that same expert (if any, from a >64-row group) just + // touched the identical weight footprint -> L2 reuse across the band. + int bx = band_n ? blockIdx.y : blockIdx.x; + int by = band_n ? blockIdx.x : blockIdx.y; int tok0=tok0_arr[bx], gend=gend_arr[bx], eid=eid_arr[bx], n0=by*BN; int tid=threadIdx.x,warp=tid>>5,lane=tid&31,gid=lane>>2,tid4=lane&3,wrow=warp*16; __shared__ __half As[2][BM][BK]; @@ -2609,11 +2697,27 @@ extern "C" __global__ void moe_q4_grouped_mma( } "#; +// FFAI_LAGUNA_SCHED=1: A/B switch for the two Laguna prefill CTA-scheduling +// levers on the Q4 grouped GEMM (see moe_q4_grouped_mma / _dev below). Default +// off = byte-identical launch geometry to before this change. +#[inline] +fn laguna_sched_on() -> bool { + std::env::var("FFAI_LAGUNA_SCHED").as_deref() != Ok("0") +} + /// Q4-native grouped MoE GEMM: `out[t,n] = Σ_k A[t,k]·dequant(Qs[eid,n,k])` over /// sorted token groups. `a` f16 `[mt,K]`; `qs` u32 + `sc` f16 = the model's /// contiguous `[n_exp*N, K]` Q4 pool (moe_up_all/down_all); `g_starts[g]..[g+1]` /// = group g's token range; `expert_ids[g]` selects the expert. Single launch, /// no dequant/slab. Returns `out[mt,N]` f16. `K%64==0`, `N%64==0`. +/// +/// `FFAI_LAGUNA_SCHED=1` turns on two scheduling levers (both host-side data +/// here already, so no extra device sync): (1) tiles are emitted in +/// DESCENDING group-row order instead of expert order, so the biggest groups' +/// CTAs launch first and small trailing groups' stragglers overlap with them +/// instead of forming a ragged tail; (2) the launch grid is transposed so a +/// band of consecutive CTAs sweeps all N-blocks of one tile (weight-tile L2 +/// reuse) before moving to the next tile, see `band_n` in the kernel source. #[allow(clippy::too_many_arguments)] pub fn moe_q4_grouped_mma( dev: &dyn Device, @@ -2631,9 +2735,21 @@ pub fn moe_q4_grouped_mma( if n % 64 != 0 { return Err(Error::Msg(format!("moe_q4_grouped_mma: N {n} %64!=0"))); } let out = Tensor::empty(dev, vec![mt.max(1), n], DType::F16)?; if n_active == 0 || mt == 0 { return Ok(out); } + let band_n = laguna_sched_on(); + // Default: expert order 0..n_active. FFAI_LAGUNA_SCHED=1: descending group + // row-count, so the largest groups' tiles are emitted (and thus launched) + // first. Purely a host-side Vec reorder, g_starts/expert_ids are already + // host slices, so this adds no new device sync. + let order: Vec = if band_n { + let mut idx: Vec = (0..n_active).collect(); + idx.sort_by_key(|&g| std::cmp::Reverse(g_starts[g + 1] - g_starts[g])); + idx + } else { + (0..n_active).collect() + }; // Per-m-tile descriptors (BM=64 token rows each, never crossing a group). let (mut tok0, mut eid, mut gend) = (Vec::new(), Vec::new(), Vec::new()); - for g in 0..n_active { + for &g in &order { let (s, e) = (g_starts[g], g_starts[g + 1]); let mut t = s; while t < e { tok0.push(t as i32); eid.push(expert_ids[g] as i32); gend.push(e as i32); t += 64; } @@ -2645,13 +2761,17 @@ pub fn moe_q4_grouped_mma( }; let (t0d, eidd, ged) = (up_i32(&tok0)?, up_i32(&eid)?, up_i32(&gend)?); let i = |x: i32| x.to_le_bytes().to_vec(); - let grid = [n_mtiles as u32, (n as u32).div_ceil(128), 1]; + let grid = if band_n { + [(n as u32).div_ceil(128), n_mtiles as u32, 1] + } else { + [n_mtiles as u32, (n as u32).div_ceil(128), 1] + }; dev.dispatch_raw_cuda( MOE_Q4_GROUPED_MMA_SRC, "moe_q4_grouped_mma.cu", "moe_q4_grouped_mma", &[(a.buffer.as_ref(), a.offset), (qs.buffer.as_ref(), qs.offset), (sc.buffer.as_ref(), sc.offset), (t0d.buffer.as_ref(), 0), (eidd.buffer.as_ref(), 0), (ged.buffer.as_ref(), 0), (out.buffer.as_ref(), 0)], - &[i(n as i32), i(k as i32)], + &[i(n as i32), i(k as i32), i(band_n as i32)], grid, [128, 1, 1], 0, false)?; Ok(out) } @@ -2660,14 +2780,33 @@ pub fn moe_q4_grouped_mma( // `offsets[n_exp+1]` (no host sync) → graph-safe. One thread walks the experts, // emitting one descriptor per BM=64-row tile; padding tiles get gend=0 (the GEMM // early-exits on them). maxt = mt/BM + n_exp is a per-S constant → fixed launch. +// +// `desc!=0` (FFAI_LAGUNA_SCHED=1): before walking, the single thread does an +// O(n_exp^2) descending selection-sort of the experts by row-count (`offsets` +// deltas) into `order`, using dynamic shared mem (2*n_exp ints) as scratch, and +// walks experts in that order instead of 0..n_exp. n_exp is at most a few +// hundred (one routing step of one layer/chunk) so this is negligible next to +// the GEMM launch it precedes; no host readback of `offsets` is added. const MOE_BUILD_TILES_SRC: &str = r#" extern "C" __global__ void moe_build_tiles( const unsigned* __restrict__ offsets, int* __restrict__ tok0, - int* __restrict__ eid, int* __restrict__ gend, int n_exp, int BM, int maxt) + int* __restrict__ eid, int* __restrict__ gend, int n_exp, int BM, int maxt, int desc) { if (blockIdx.x || threadIdx.x) return; + extern __shared__ int smem[]; + int* cnt = smem; // [n_exp], only used when desc!=0 + int* order = smem + n_exp; // [n_exp] + if (desc) { + for (int e = 0; e < n_exp; e++) { cnt[e] = (int)(offsets[e+1] - offsets[e]); order[e] = e; } + for (int i2 = 0; i2 < n_exp; i2++) { + int best = i2; + for (int j = i2 + 1; j < n_exp; j++) if (cnt[order[j]] > cnt[order[best]]) best = j; + if (best != i2) { int tmp = order[i2]; order[i2] = order[best]; order[best] = tmp; } + } + } int t = 0; - for (int e = 0; e < n_exp; e++) { + for (int ii = 0; ii < n_exp; ii++) { + int e = desc ? order[ii] : ii; unsigned s = offsets[e], en = offsets[e+1]; for (unsigned x = s; x < en; x += (unsigned)BM) { if (t < maxt) { tok0[t]=(int)x; eid[t]=e; gend[t]=(int)en; t++; } @@ -2680,6 +2819,11 @@ extern "C" __global__ void moe_build_tiles( /// Fully-on-device Q4 grouped MoE GEMM: builds tile descriptors on-device from /// `offsets` (device `[n_exp+1]` u32) with a FIXED launch (grid=[maxt,N/64], /// maxt = mt/64 + n_exp) → NO host sync, graph-capturable. Returns `out[mt,N]` f16. +/// +/// `FFAI_LAGUNA_SCHED=1` turns on the same two scheduling levers as +/// `moe_q4_grouped_mma`: descending group-row tile order (computed ON DEVICE +/// in `moe_build_tiles`, since this path only has a device `offsets` buffer: +/// no new device->host sync is added) and the transposed/banded launch grid. #[allow(clippy::too_many_arguments)] pub fn moe_q4_grouped_mma_dev( dev: &dyn Device, @@ -2696,24 +2840,33 @@ pub fn moe_q4_grouped_mma_dev( if n % 64 != 0 { return Err(Error::Msg(format!("moe_q4_grouped_mma_dev: N {n} %64!=0"))); } let out = Tensor::empty(dev, vec![mt.max(1), n], DType::F16)?; if mt == 0 { return Ok(out); } + let band_n = laguna_sched_on(); let maxt = mt.div_ceil(64) + n_exp; let t0d = Tensor::empty(dev, vec![maxt], DType::U32)?; let eidd = Tensor::empty(dev, vec![maxt], DType::U32)?; let ged = Tensor::empty(dev, vec![maxt], DType::U32)?; let i = |x: i32| x.to_le_bytes().to_vec(); + // 2 int arrays of n_exp each, only needed (and only touched by the + // kernel) when desc!=0. Kept at 0 in the default/off path so that path's + // launch geometry is byte-identical to before this change. + let build_shared = if band_n { (2 * n_exp * std::mem::size_of::()) as u32 } else { 0 }; dev.dispatch_raw_cuda( MOE_BUILD_TILES_SRC, "moe_build_tiles.cu", "moe_build_tiles", &[(offsets.buffer.as_ref(), offsets.offset), (t0d.buffer.as_ref(), 0), (eidd.buffer.as_ref(), 0), (ged.buffer.as_ref(), 0)], - &[i(n_exp as i32), i(64), i(maxt as i32)], - [1, 1, 1], [1, 1, 1], 0, false)?; - let grid = [maxt as u32, (n as u32).div_ceil(128), 1]; + &[i(n_exp as i32), i(64), i(maxt as i32), i(band_n as i32)], + [1, 1, 1], [1, 1, 1], build_shared, false)?; + let grid = if band_n { + [(n as u32).div_ceil(128), maxt as u32, 1] + } else { + [maxt as u32, (n as u32).div_ceil(128), 1] + }; dev.dispatch_raw_cuda( MOE_Q4_GROUPED_MMA_SRC, "moe_q4_grouped_mma.cu", "moe_q4_grouped_mma", &[(a.buffer.as_ref(), a.offset), (qs.buffer.as_ref(), qs.offset), (sc.buffer.as_ref(), sc.offset), (t0d.buffer.as_ref(), 0), (eidd.buffer.as_ref(), 0), (ged.buffer.as_ref(), 0), (out.buffer.as_ref(), 0)], - &[i(n as i32), i(k as i32)], + &[i(n as i32), i(k as i32), i(band_n as i32)], grid, [128, 1, 1], 0, false)?; Ok(out) } @@ -2979,11 +3132,11 @@ extern "C" __global__ void sdpa_tc_finalize( /// FUSED single-kernel causal FlashAttention (no HBM score/prob materialization, /// no qprep/kprep/vprep passes). One warp per (head, query row): the warp holds /// `Q[r]` in registers, streams K/V rows from HBM, and runs the online softmax + -/// `O` accumulation entirely in registers — scores never touch global memory. +/// `O` accumulation entirely in registers, scores never touch global memory. /// Causal is exploited by looping `p` only to `base_kv + r` (the masked upper /// triangle is never computed, ~2× fewer FLOPs than the materialized path). /// hd must be a multiple of 32 (warp lanes own `hd/32` dims each). Q/K/V F32 or -/// F16; output matches `q.dtype`. This is v1 (scalar FMA, no tensor cores) — it +/// F16; output matches `q.dtype`. This is v1 (scalar FMA, no tensor cores), it /// wins by removing the 6-dispatch HBM round-trip + the causal waste; the wmma /// upgrade is the follow-up. NEMOTRON_FLASH_FUSED=1 in the model. const SDPA_FLASH_FUSED_SRC: &str = r#" @@ -3074,7 +3227,7 @@ const SDPA_FLASH_WMMA_SRC: &str = r#" #include using namespace nvcuda; extern "C" __global__ void sdpa_flash_wmma( - const __half* __restrict__ q_in, // [Sq, nq, 128] (already scaled? no — scaled here) + const __half* __restrict__ q_in, // [Sq, nq, 128] (already scaled? no, scaled here) const __half* __restrict__ k_in, // [nkv, kv_stride, 128] const __half* __restrict__ v_in, // [nkv, kv_stride, 128] __half* __restrict__ out, // [Sq, nq, 128] @@ -3182,7 +3335,7 @@ extern "C" __global__ void sdpa_flash_wmma( "#; /// FlashAttention-2 with register-resident O (mma.sync m16n8k16). v3 of the fused -/// path: O lives in registers (acc fragments), not shared — freeing ~16KB so the +/// path: O lives in registers (acc fragments), not shared: freeing ~16KB so the /// 1-warp/16-row block hits high occupancy (v2's 30KB shared capped it at ~11%). /// QKᵀ + P·V via `mma.sync` (manual fragment packing, documented layout); the /// causal mask + online softmax round-trip the small S/P tiles through shared @@ -3565,7 +3718,7 @@ extern "C" __global__ void sdpa_flash_mma_w4s( } // FP8 (e4m3) attention v2: QK and PV via m16n8k32 e4m3 mma (2x f16 rate). -// K/V are PRE-QUANTIZED to e4m3 buffers once per call (kv_to_e4m3 below) — +// K/V are PRE-QUANTIZED to e4m3 buffers once per call (kv_to_e4m3 below), // v1 quantized in-smem on every KV re-read (frexp per element x 32 // row-blocks) which ate the mma savings; v2 staging is byte copies and the // KV re-read traffic halves. BC=32 keys/tile so the PV k32 mma is fully @@ -3787,7 +3940,7 @@ pub fn sdpa_flash_mma( let kv_chunk_target = 8192usize; let n_part = n_kv.div_ceil(kv_chunk_target).min(8); // FP8 attention (NEMOTRON_FLASH_FP8=1): e4m3 QK and PV mma at 2x the f16 - // tensor-core rate — attention at depth is mma-compute-bound. Static + // tensor-core rate, attention at depth is mma-compute-bound. Static // conservative scales (16x: |x|<28 covers the post-softmax-scale Q and the // K/V ranges with e4m3's 448 ceiling). Quant-regime lever. if std::env::var("NEMOTRON_FLASH_FP8").is_ok() { @@ -3846,7 +3999,7 @@ pub fn sdpa_flash_mma( [n_el.div_ceil(256), 1, 1], [256, 1, 1], 0, false)?; return Ok(out); } - // 4-warp blocks: 64 query rows/block sharing the K/V smem stage — + // 4-warp blocks: 64 query rows/block sharing the K/V smem stage, // 2.7x the 1-warp kernel at S=2048, bit-identical output. let grid = [nq as u32, (sq as u32).div_ceil(64), 1]; dev.dispatch_raw_cuda( @@ -3966,7 +4119,7 @@ pub fn sdpa_flash_fused( Ok(out) } -/// Tensor-core prefill SDPA — drop-in faster replacement for [`sdpa_multi`]. +/// Tensor-core prefill SDPA, drop-in faster replacement for [`sdpa_multi`]. /// Same signature/semantics (causal/full, GQA, head_dim=128, Q pre-scaled by /// `scale`), but runs QKᵀ and P·V on the cuBLAS tensor cores with a /// FlashAttention online-softmax loop tiled over KV. Q/K/V may be F32 or F16; @@ -4054,7 +4207,7 @@ pub fn sdpa_multi_tc( // ── softmax + causal mask for this block ──────────────────────────── // scores/p_blk are PACKED at pitch `blk` (head stride sq*blk, row pitch - // blk) — only the first nq*sq*blk f16 of the bk-sized allocation are used. + // blk), only the first nq*sq*blk f16 of the bk-sized allocation are used. // This keeps the GEMM's lda == k == blk valid for the final partial block. dev.dispatch_raw_cuda(SDPA_TC_SOFTMAX_SRC, "sdpa_tc_softmax.cu", "sdpa_tc_softmax", &[(scores.buffer.as_ref(), 0), (p_blk.buffer.as_ref(), 0), @@ -4108,9 +4261,10 @@ pub fn sdpa_multi_tc( /// share one set of GEMMs while attention stays block-diagonal (correct). /// /// NOTE: this computes the full `[Sq, n_kv]` QKᵀ and masks off-segment to -/// −inf — bit-correct, but still O((ΣL)²) QKᵀ. The segment-skip optimisation -/// (only run KV blocks intersecting a query's segment → O(ΣLᵢ²)) is the -/// follow-up that turns correctness into the throughput win. +/// −inf, bit-correct, but still O((ΣL)²) QKᵀ for the packed-segment case. +/// The `seg_len == KV-block-size` branch below skips that; `win` (below) +/// is the analogous skip for a per-row SLIDING WINDOW instead of a packed +/// segment boundary. pub fn sdpa_multi_tc_varlen( dev: &dyn Device, q: &Tensor, @@ -4128,6 +4282,16 @@ pub fn sdpa_multi_tc_varlen( heads_per_group: u32, causal: bool, scale: f32, + // Sliding-window block skip. 0 = off (no effect on any existing caller). + // When `win > 0`, `seg_lo` is expected to encode a sliding window of + // width `win` (`seg_lo[r] = max(0, (base_kv+r) - (win-1))`, exactly what + // Laguna's sliding-layer prefill already uploads) and the per-KV-block + // query-row range fed to the QKᵀ/PV GEMMs is restricted to the rows + // whose window can possibly reach that block, see the derivation at + // the `q_lo`/`q_cnt` computation below. Independent of `seg_len`/the + // packed-segment skip: a caller should pass exactly one of the two + // (Laguna always passes `seg_len = 0` and this `win`). + win: u32, ) -> Result { let hd = head_dim; let nq = n_q_heads as usize; @@ -4180,21 +4344,58 @@ pub fn sdpa_multi_tc_varlen( let mut kb0 = 0usize; while kb0 < n_kv { let blk = bk.min(n_kv - kb0); - // Segment-skip: when each KV-block is exactly one packed segment - // (seg_len == bk), only that segment's query rows [kb0, kb0+blk) attend - // to this block — every other row masks to p=0. Restrict the expensive - // QKᵀ/PV tensor-core GEMMs to that range (O((NL)²)→O(NL²)); the cheap - // full-range softmax/merge correctly no-op the off-segment rows. - let (q_lo, q_cnt) = if seg_len != 0 && seg_len as usize == bk { (kb0, blk) } else { (0, sq) }; - for g in 0..nkv { - let q_off = g * hpg * sq * hd * 2 + q_lo * hd * 2; - let k_off = (g * n_kv + kb0) * hd * 2; - let s_off = g * hpg * sq * blk * 2 + q_lo * blk * 2; - dev.gemm_strided_batched_off( - qh.buffer.as_ref(), q_off, (sq*hd) as i64 * el, - kh.buffer.as_ref(), k_off, 0, - scores.buffer.as_ref(), s_off, (sq*blk) as i64 * el, - q_cnt, blk, hd, hpg, DType::F16)?; + // Query-range skip: restrict the expensive QKᵀ/PV tensor-core GEMMs + // to the rows that can possibly attend anywhere in KV block [kb0, + // kb0+blk), the cheap full-range softmax/merge below still apply + // the real per-element mask, so a range that's a SUPERSET of the + // true set is always safe (the extra rows just mask to p=0); a + // range narrower than the true set would silently drop attention + // mass. Only ever shrink, never tighten past that guarantee. + // + // Packed-segment mode (seg_len == bk): each KV-block is exactly one + // packed segment, so only that segment's own rows [kb0, kb0+blk) + // attend to it, row index and absolute KV position coincide in + // this mode (packed sequences are laid out from base_kv=0). + // + // Sliding-window mode (win != 0): causal means query absolute + // position i_abs can reach KV position kb0 (the block's first + // index) only if i_abs >= kb0. The window means i_abs can reach KV + // position kb0+blk-1 (the block's last index) only if that index is + // still inside the window, i.e. seg_lo[i] = i_abs - win + 1 <= + // kb0+blk-1, i.e. i_abs < kb0 + blk + win. So the absolute-position + // range that can touch this block is [kb0, kb0+blk+win). Row `r` of + // this call sits at absolute position `base + r` (base_kv is this + // chunk's starting position), so convert that absolute range to a + // row-index range and clamp into [0, sq). + let (q_lo, q_cnt) = if win != 0 { + let lo_abs = kb0 as i64; + let hi_abs_excl = kb0 as i64 + blk as i64 + win as i64; + let lo_r = (lo_abs - base as i64).max(0); + let hi_r = (hi_abs_excl - base as i64).min(sq as i64); + if hi_r > lo_r { (lo_r as usize, (hi_r - lo_r) as usize) } else { (0, 0) } + } else if seg_len != 0 && seg_len as usize == bk { + (kb0, blk) + } else { + (0, sq) + }; + // q_cnt == 0: no row in this chunk can reach this block at all (the + // derivation above is a superset of the true reachable set, so an + // empty computed range proves the true set is also empty). Skip the + // QKᵀ/PV GEMMs and vprep entirely; softmax/merge stay unconditional + // below (softmax re-derives bm/bl fresh every iteration straight + // from the mask, and merge's `l_b <= 0` early-return means it never + // reads the skipped o_blk/vt buffers, so this is correctness-safe). + if q_cnt > 0 { + for g in 0..nkv { + let q_off = g * hpg * sq * hd * 2 + q_lo * hd * 2; + let k_off = (g * n_kv + kb0) * hd * 2; + let s_off = g * hpg * sq * blk * 2 + q_lo * blk * 2; + dev.gemm_strided_batched_off( + qh.buffer.as_ref(), q_off, (sq*hd) as i64 * el, + kh.buffer.as_ref(), k_off, 0, + scores.buffer.as_ref(), s_off, (sq*blk) as i64 * el, + q_cnt, blk, hd, hpg, DType::F16)?; + } } // varlen softmax: causal upper (base+r) AND segment lower (seg_lo[r]). dev.dispatch_raw_cuda(SDPA_TC_SOFTMAX_VARLEN_SRC, "sdpa_tc_softmax_varlen.cu", "sdpa_tc_softmax_varlen", @@ -4203,19 +4404,21 @@ pub fn sdpa_multi_tc_varlen( (seg_lo.buffer.as_ref(), seg_lo.offset)], &[i(nq as i32), i(sq as i32), i(blk as i32), i(kb0 as i32), i(base as i32)], [(nq*sq) as u32, 1, 1], [256,1,1], 0, false)?; - dev.dispatch_raw_cuda(SDPA_TC_VPREP_SRC, "sdpa_tc_vprep.cu", "sdpa_tc_vprep", - &[(v.buffer.as_ref(), v.offset), (vt.buffer.as_ref(), 0)], - &[i(nkv as i32), i(kv_stride as i32), i(hd as i32), i(kb0 as i32), i(blk as i32), i(in_f32)], - blk256(nkv*blk*hd), [256,1,1], 0, false)?; - for g in 0..nkv { - let p_off = g * hpg * sq * blk * 2 + q_lo * blk * 2; - let v_off = g * hd * blk * 2; - let o_off = g * hpg * sq * hd * 2 + q_lo * hd * 2; - dev.gemm_strided_batched_off( - p_blk.buffer.as_ref(), p_off, (sq*blk) as i64 * el, - vt.buffer.as_ref(), v_off, 0, - o_blk.buffer.as_ref(), o_off, (sq*hd) as i64 * el, - q_cnt, hd, blk, hpg, DType::F16)?; + if q_cnt > 0 { + dev.dispatch_raw_cuda(SDPA_TC_VPREP_SRC, "sdpa_tc_vprep.cu", "sdpa_tc_vprep", + &[(v.buffer.as_ref(), v.offset), (vt.buffer.as_ref(), 0)], + &[i(nkv as i32), i(kv_stride as i32), i(hd as i32), i(kb0 as i32), i(blk as i32), i(in_f32)], + blk256(nkv*blk*hd), [256,1,1], 0, false)?; + for g in 0..nkv { + let p_off = g * hpg * sq * blk * 2 + q_lo * blk * 2; + let v_off = g * hd * blk * 2; + let o_off = g * hpg * sq * hd * 2 + q_lo * hd * 2; + dev.gemm_strided_batched_off( + p_blk.buffer.as_ref(), p_off, (sq*blk) as i64 * el, + vt.buffer.as_ref(), v_off, 0, + o_blk.buffer.as_ref(), o_off, (sq*hd) as i64 * el, + q_cnt, hd, blk, hpg, DType::F16)?; + } } dev.dispatch_raw_cuda(SDPA_TC_MERGE_SRC, "sdpa_tc_merge.cu", "sdpa_tc_merge", &[(o_blk.buffer.as_ref(), 0), (bm.buffer.as_ref(), 0), (bl.buffer.as_ref(), 0), @@ -4237,7 +4440,7 @@ pub fn sdpa_multi_tc_varlen( /// NemotronH-Nano-30B Mamba2 cell (Dh=64, Ds=128, H=64, G=8, n_per_t=4). /// /// The existing `ssm_step_record_d64_128_64_8` runs a serial loop over T -/// positions per (h, dh) cell — the serial chain of length T is the bottleneck. +/// positions per (h, dh) cell, the serial chain of length T is the bottleneck. /// This kernel replaces it with a **two-pass parallel segment scan**: /// /// **Pass 1 (parallel)**: `N_SEG` warp-segments each process `T/N_SEG` positions @@ -4427,7 +4630,7 @@ void ssm_chunked_prefill( /// Grid3D: `[1, Dh=64, H=64]`, block `[WARP × N_SEG]` = `[32 × 32 = 1024]`. /// Shared mem per block: ~17KB (partial states + alpha + s_in scratch). /// -/// Returns `(state_out, y)` — same signature as `ssm_prefill_scan`. +/// Returns `(state_out, y)`, same signature as `ssm_prefill_scan`. #[allow(clippy::too_many_arguments)] pub fn ssm_prefill_scan_chunked( dev: &dyn Device, @@ -4454,7 +4657,7 @@ pub fn ssm_prefill_scan_chunked( let y = Tensor::empty(dev, vec![t * nh * dhu], x.dtype)?; let state_out = Tensor::empty(dev, state_in.shape.clone(), x.dtype)?; // N_SEG must divide t_total (or kernel pads). 64 divides 512,2048,8192. - // For t_total not divisible by 64, pass as-is — the kernel uses a guard. + // For t_total not divisible by 64, pass as-is: the kernel uses a guard. let t_bytes = t_total.to_le_bytes().to_vec(); // Grid: [1, Dh, H]. Block: [WARP * N_SEG] = [32 * 64] = [2048]. let grid = [1u32, dh, n_heads]; @@ -4492,7 +4695,7 @@ pub fn ssm_prefill_scan_chunked( Ok((state_out, y)) } -/// Mamba2 SSD **batched-prefill scan** — runs the sequential SSD forward over +/// Mamba2 SSD **batched-prefill scan**, runs the sequential SSD forward over /// all `t_total` prompt tokens in ONE dispatch, emitting every per-token /// `y[t]` plus the final recurrent `state_out` (for decode continuity). /// Dispatches `ssm_step_record_d64_128_64_8` (Nemotron cell: Dh=64, Ds=128, @@ -4526,7 +4729,7 @@ pub fn ssm_prefill_scan( let (nh, dhu, dsu) = (n_heads as usize, dh as usize, ds as usize); let y = Tensor::empty(dev, vec![t * nh * dhu], x.dtype)?; let state_out = Tensor::empty(dev, state_in.shape.clone(), x.dtype)?; - // Tape outputs — written by the kernel, unused by prefill. + // Tape outputs, written by the kernel, unused by prefill. let da_log = Tensor::empty(dev, vec![t * nh * dsu], x.dtype)?; let dbx_log = Tensor::empty(dev, vec![t * nh * dhu * dsu], x.dtype)?; // mask buffer: has_mask=0 → kernel ignores contents, but the binding must exist. @@ -4656,7 +4859,7 @@ pub fn sdpa_decode( // ── CUDA-graph-capturable dense decode SDPA (raw CUDA, head_dim 128) ──── // // The registry `ffai_sdpa_decode` kernel above takes `n_kv` as a launch-time -// constexpr — fine for eager decode (a fresh launch every token) but wrong +// constexpr, fine for eager decode (a fresh launch every token) but wrong // for a captured graph, which replays the exact same recorded launch (same // baked `n_kv`) every token. Rather than plumb a buffer-read `n_kv` through // the metaltile registry kernel (invasive: `n_kv` there is a `#[constexpr]` @@ -4762,7 +4965,7 @@ extern "C" __global__ void ffai_sdpa_decode_nbuf( /// CUDA-graph-capturable fallback of [`sdpa_decode`] for `head_dim == 128` /// (Laguna's only head_dim): reads the filled KV length `n_kv` from a /// 1-element `u32` device buffer instead of a baked-in launch scalar. Dense -/// causal only (no sink / sliding-window mask) — Laguna's sliding layers +/// causal only (no sink / sliding-window mask), Laguna's sliding layers /// bound the KV walk via the buffer's VALUE (the ring's current filled /// length), so a separate window mask isn't needed. f32 only. See the /// kernel-source doc above for the algorithm. @@ -4800,13 +5003,13 @@ pub fn sdpa_decode_nbuf( Ok(out) } -/// GQA-aware split-K flash-decode SDPA (MLX `sdpa_vector_2pass` port, head_dim -/// 128). Pass 1: grid `(n_kv_heads, blocks)`, block `(32, gqa_factor)` — one +/// GQA-aware split-K flash-decode SDPA (ported from a reference ML framework's `sdpa_vector_2pass`, head_dim +/// 128). Pass 1: grid `(n_kv_heads, blocks)`, block `(32, gqa_factor)`: one /// simdgroup per Q-head of the group, so the `gqa_factor` heads SHARE each K/V -/// load (the single-pass kernel re-reads the shared KV head `gqa_factor`× — at +/// load (the single-pass kernel re-reads the shared KV head `gqa_factor`×, at /// 32K that's the dominant cost). Each block-row strides a `1/blocks` slice of /// the KV positions, emitting per-block (max, sum, partial_o). Pass 2: grid -/// `(n_q_heads)`, block 1024 — online-softmax merge of the `blocks` partials. +/// `(n_q_heads)`, block 1024: online-softmax merge of the `blocks` partials. /// `blocks` MUST be a multiple of 32 (pass-2 reducer constraint). Kernel math /// is registry-validated (`test_ffai_sdpa_decode_2pass_combined`, f32/f16/bf16). #[allow(clippy::too_many_arguments)] @@ -4973,7 +5176,7 @@ fn elementwise_kernel( bindings.push(Binding::Buffer(out.buffer.clone())); let n = out.elem_count() as u32; let grid = Grid::d1(n.div_ceil(256), 256); - // output binding goes in its signature position — callers pass inputs, + // output binding goes in its signature position, callers pass inputs, // we append the single output last (matches a,…,out param order). dev.dispatch(&k, &bindings, grid)?; Ok(out) @@ -4984,7 +5187,7 @@ pub fn silu(dev: &dyn Device, x: &Tensor) -> Result { elementwise_kernel(dev, "mt_silu", x.dtype, x.shape.clone(), vec![Binding::Buffer(x.buffer.clone())]) } -/// ReLU² activation `out = max(x,0)²`, elementwise — NemotronH MoE expert act. +/// ReLU² activation `out = max(x,0)²`, elementwise: NemotronH MoE expert act. /// Built inline (Relu → square) so the expert `up → relu² → down` chain stays /// ON-DEVICE (no host round-trip between the two Q8 GEMVs). Dispatches `mt_relu2`. pub fn relu2(dev: &dyn Device, x: &Tensor) -> Result { @@ -5081,8 +5284,8 @@ pub fn moe_scatter_add( /// The atomicAdd variant (`moe_scatter_add_f32`) is run-to-run NONdeterministic: /// floating-point `atomicAdd` completes in hardware-arbitrated (nondeterministic) /// order, so when several routed experts for the same token accumulate into the -/// same `acc[token,h]`, the summation order — and thus the low bits of the result -/// — varies between identical runs. That jitter propagates through 52 layers and +/// same `acc[token,h]`, the summation order: and thus the low bits of the result +///, varies between identical runs. That jitter propagates through 52 layers and /// can flip the final argmax at deep context. /// /// This variant is bit-exact reproducible: one block per `(token, h-block)`, and @@ -5200,7 +5403,7 @@ extern "C" __global__ void moe_scatter_add_topk_det_f16_x2( /// host before uploading). Builds a CSR index grouping rows by token in ascending /// row order, uploads it, and runs the atomic-free deterministic kernel. /// -/// `acc` is WRITTEN (not accumulated) — every output token is fully summed here, +/// `acc` is WRITTEN (not accumulated), every output token is fully summed here, /// so it need not be pre-zeroed, but tokens with no routed rows are set to 0. pub fn moe_scatter_add_det( dev: &dyn Device, @@ -5488,7 +5691,7 @@ pub fn relu2_scale_f16(dev: &dyn Device, x: &Tensor, scale: f32) -> Result Kernel { let mut k = Kernel::new(name); @@ -5517,7 +5720,7 @@ fn unary_act_kernel(name: &str, dtype: DType, kind: ActKind) -> Kernel { } /// Fused multiply-add into an accumulator, elementwise IN-PLACE: -/// `acc[i] += x[i] · s[i]`. Lets the MoE expert sum stay ON-DEVICE — each +/// `acc[i] += x[i] · s[i]`. Lets the MoE expert sum stay ON-DEVICE, each /// expert's `down` output is folded into `acc` on the GPU (one final download /// per layer instead of one per expert). `s` is the per-expert weight broadcast /// to `[len]`. Dispatches `mt_fma_inplace`. @@ -5559,7 +5762,7 @@ pub fn gelu(dev: &dyn Device, x: &Tensor) -> Result { /// LayerNorm `out = (x - mean) / sqrt(var + eps) * w + b`, normalized over the /// last dim. Mean-subtracting + bias (unlike RMSNorm). Dispatches the -/// Reduction-mode `mt_layer_norm` (one threadgroup per row, block = n/4 — needs +/// Reduction-mode `mt_layer_norm` (one threadgroup per row, block = n/4: needs /// the row width `n` divisible by `4·lsize`; ViT/SigLIP widths are). Used by /// every transformer with LayerNorm (vision towers, BERT-style, GPT-2, …). pub fn layer_norm(dev: &dyn Device, x: &Tensor, weight: &Tensor, bias: &Tensor, eps: f32) -> Result { @@ -5688,9 +5891,9 @@ pub fn dsv4_partial_rope( // // IMPORTANT: `attn_factor` here is the model's stored `attention_factor` // (e.g. 1.4852030263919618 for Laguna-S-2.1's factor-128 YaRN), passed -// straight through — this kernel re-applies the `1 + 0.1*ln(1/freq_scale)` +// straight through, this kernel re-applies the `1 + 0.1*ln(1/freq_scale)` // correction internally exactly as the reference engine does (the engine's -// llama-context.cpp pre-cancels its OWN internally-derived default mscale +// the reference engine's context handling pre-cancels its OWN internally-derived default mscale // against this same term before folding in the config's `attn_factor`; the // net effect is that `attn_factor` reaches the rope kernel unmodified and // gets this multiplier applied once, INSIDE the kernel). Do not pre-apply @@ -5749,7 +5952,7 @@ extern "C" __global__ void ffai_rope_yarn_partial( /// Q or K tensor at sequence `position`. `n_rot` is the number of rotated /// dims (`<= head_dim`; the rest pass through unrotated). Pass `n_rot == /// head_dim`, `freq_scale = 1.0`, `ext_factor = 0.0`, `attn_factor = 1.0` for -/// plain (non-YaRN) RoPE — `beta_fast`/`beta_slow`/`n_orig_ctx` are then +/// plain (non-YaRN) RoPE, `beta_fast`/`beta_slow`/`n_orig_ctx` are then /// unused (the ramp contributes zero once `ext_factor` is zero). f32 only. #[allow(clippy::too_many_arguments)] pub fn rope_yarn_partial( @@ -5876,7 +6079,7 @@ pub fn rope_yarn_partial_posbuf( "rope_yarn_partial_posbuf: n_rot {n_rot} invalid for head_dim {head_dim}" ))); } - // Same YaRN corr-dims derivation as `rope_yarn_partial` — a function of + // Same YaRN corr-dims derivation as `rope_yarn_partial`, a function of // the RoPE constants only, never the runtime position. let corr_dim = |n_rotations: f32| -> f32 { n_rot as f32 * (n_orig_ctx as f32 / (n_rotations * 2.0 * std::f32::consts::PI)).ln() @@ -5906,12 +6109,127 @@ pub fn rope_yarn_partial_posbuf( Ok(out) } +const ROPE_YARN_PARTIAL_MANY_SRC: &str = r#" +extern "C" __global__ void ffai_rope_yarn_partial_many( + const float* __restrict__ qk, float* __restrict__ out, + const unsigned int* __restrict__ positions, + int head_dim, int n_rot, int n_heads, + float theta_base, float freq_scale, float ext_factor, float attn_factor, + float corr_low, float corr_high) +{ + int t = blockIdx.x; + int h = blockIdx.y; + int tid = threadIdx.z; + int half_rot = n_rot / 2; + int position = (int)positions[t]; + long row_off = ((long)t * n_heads + h) * head_dim; + const float* row_in = qk + row_off; + float* row_out = out + row_off; + + if (tid < half_rot) { + int j = tid; + float theta_scale = powf(theta_base, -2.0f / (float)n_rot); + float theta_extrap = (float)position * powf(theta_scale, (float)j); + float theta_interp = freq_scale * theta_extrap; + float theta = theta_interp; + float mscale = attn_factor; + if (ext_factor != 0.0f) { + float y = ((float)j - corr_low) / fmaxf(0.001f, corr_high - corr_low); + float ramp = 1.0f - fminf(1.0f, fmaxf(0.0f, y)); + float ramp_mix = ramp * ext_factor; + theta = theta_interp * (1.0f - ramp_mix) + theta_extrap * ramp_mix; + } + float cos_t = cosf(theta) * mscale; + float sin_t = sinf(theta) * mscale; + float x0 = row_in[j]; + float x1 = row_in[j + half_rot]; + row_out[j] = x0 * cos_t - x1 * sin_t; + row_out[j + half_rot] = x0 * sin_t + x1 * cos_t; + } else { + int p = tid - half_rot; + int idx = n_rot + 2 * p; + if (idx < head_dim) { + row_out[idx] = row_in[idx]; + row_out[idx + 1] = row_in[idx + 1]; + } + } +} +"#; + +/// Batched [`rope_yarn_partial`]: applies YaRN-scaled partial-rotary RoPE to +/// ALL `T` rows of a `[T, n_heads, head_dim]` Q or K tensor in ONE dispatch, +/// each row `t` rotated at `positions[t]` (one absolute sequence position per +/// token row) instead of a single scalar `position` shared by every row. +/// Replaces the T-loop of per-token `rope_yarn_partial` calls in the Laguna +/// prefill attention path. Same math/params as `rope_yarn_partial` (pass +/// `n_rot == head_dim`, `freq_scale = 1.0`, `ext_factor = 0.0`, +/// `attn_factor = 1.0` for plain, non-YaRN RoPE: Laguna's sliding-window +/// layers use exactly this). Grid `[T, n_heads, 1]`, block `[1, 1, head_dim/2]` +/// (covers both the rotated-pair range and the passthrough-tail range, same +/// as the scalar-position kernel). f32 only. +#[allow(clippy::too_many_arguments)] +pub fn rope_yarn_partial_many( + dev: &dyn Device, + qk: &Tensor, + positions: &Tensor, + n_heads: usize, + n_rot: u32, + theta_base: f32, + freq_scale: f32, + ext_factor: f32, + attn_factor: f32, + beta_fast: f32, + beta_slow: f32, + n_orig_ctx: u32, +) -> Result { + let head_dim = *qk.shape.last().ok_or_else(|| Error::Msg("rope_yarn_partial_many: scalar input".into()))?; + let t = qk.elem_count() / (n_heads * head_dim); + if positions.elem_count() != t { + return Err(Error::Msg(format!( + "rope_yarn_partial_many: positions len {} != T {t}", positions.elem_count() + ))); + } + if n_rot == 0 || n_rot as usize > head_dim || n_rot % 2 != 0 { + return Err(Error::Msg(format!( + "rope_yarn_partial_many: n_rot {n_rot} invalid for head_dim {head_dim}" + ))); + } + // Same YaRN corr-dims derivation as `rope_yarn_partial`, depends only on + // the (per-layer-type, chunk-invariant) RoPE constants, never on position. + let corr_dim = |n_rotations: f32| -> f32 { + n_rot as f32 * (n_orig_ctx as f32 / (n_rotations * 2.0 * std::f32::consts::PI)).ln() + / (2.0 * theta_base.ln()) + }; + let corr_low = corr_dim(beta_fast).floor().max(0.0); + let corr_high = corr_dim(beta_slow).ceil().min(n_rot as f32 - 1.0); + + let out = Tensor::empty(dev, qk.shape.clone(), DType::F32)?; + let i = |x: i32| x.to_le_bytes().to_vec(); + let f = |x: f32| x.to_le_bytes().to_vec(); + let half = (head_dim / 2) as u32; + dev.dispatch_raw_cuda( + ROPE_YARN_PARTIAL_MANY_SRC, "rope_yarn_partial_many.cu", "ffai_rope_yarn_partial_many", + &[ + (qk.buffer.as_ref(), qk.offset), + (out.buffer.as_ref(), 0), + (positions.buffer.as_ref(), positions.offset), + ], + &[ + i(head_dim as i32), i(n_rot as i32), i(n_heads as i32), + f(theta_base), f(freq_scale), f(ext_factor), f(attn_factor), + f(corr_low), f(corr_high), + ], + [t as u32, n_heads as u32, 1], [1, 1, half], 0, false, + )?; + Ok(out) +} + /// Write a single `u32` scalar into an existing 1-element device buffer, in /// place, without reallocating. Laguna's CUDA-graph decode path /// (`FFAI_LAGUNA_GRAPH=1`) keeps a small persistent "workspace" of such /// buffers (token id, RoPE position, KV length) whose device POINTERS get /// captured into the graph and must never move; only their contents change -/// from token to token. Call this OUTSIDE any graph-capture region — it must +/// from token to token. Call this OUTSIDE any graph-capture region, it must /// run eagerly every step, right before the (captured or replayed) decode /// work that reads the buffer. pub fn write_u32(dev: &dyn Device, buf: &Tensor, val: u32) -> Result<()> { @@ -5936,7 +6254,7 @@ extern "C" __global__ void ffai_write_u32(unsigned int* __restrict__ buf, unsign // the same normed hidden state that feeds Q/K/V (`g_proj(attn_norm(x))`, // NOT the attention output). This multiplies `softplus(g[h])` into every // element of head `h`'s attention output, broadcasting the scalar across -// `head_dim` — applied BEFORE `o_proj`. +// `head_dim`, applied BEFORE `o_proj`. const GATE_SOFTPLUS_MUL_PERHEAD_SRC: &str = r#" extern "C" __global__ void ffai_gate_softplus_mul_perhead( const float* __restrict__ attn, const float* __restrict__ g, @@ -5976,6 +6294,102 @@ pub fn gate_softplus_mul_perhead(dev: &dyn Device, attn: &Tensor, g: &Tensor) -> Ok(out) } +const GATE_SOFTPLUS_MUL_PERHEAD_MANY_SRC: &str = r#" +extern "C" __global__ void ffai_gate_softplus_mul_perhead_many( + const float* __restrict__ attn, const float* __restrict__ g, + float* __restrict__ out, int head_dim) +{ + long row = (long)blockIdx.x; // flat (t * n_heads + h) + int i = threadIdx.x; + if (i >= head_dim) return; + float gv = g[row]; + float sp = gv > 20.0f ? gv : log1pf(expf(gv)); + out[row * head_dim + i] = attn[row * head_dim + i] * sp; +} +"#; + +/// Batched [`gate_softplus_mul_perhead`]: `out[t, h, i] = attn[t, h, i] * +/// softplus(g[t, h])` for `attn` `[T, n_heads, head_dim]` and `g` `[T, +/// n_heads]` (one raw gate scalar per token-head, broadcast over +/// `head_dim`). Replaces the T-loop of per-token `gate_softplus_mul_perhead` +/// calls in the Laguna prefill attention path, applied before `o_proj`, +/// same as the decode path. f32 only. +pub fn gate_softplus_mul_perhead_many(dev: &dyn Device, attn: &Tensor, g: &Tensor, n_heads: usize) -> Result { + let head_dim = *attn.shape.last().ok_or_else(|| Error::Msg("gate_softplus_mul_perhead_many: scalar attn".into()))?; + let rows = attn.elem_count() / head_dim; + if g.elem_count() != rows { + return Err(Error::Msg(format!( + "gate_softplus_mul_perhead_many: g has {} elems, expected {rows} (T*n_heads)", + g.elem_count() + ))); + } + let _ = n_heads; // rows already folds T*n_heads; kept for call-site clarity/shape docs + let out = Tensor::empty(dev, attn.shape.clone(), DType::F32)?; + let i = |x: i32| x.to_le_bytes().to_vec(); + dev.dispatch_raw_cuda( + GATE_SOFTPLUS_MUL_PERHEAD_MANY_SRC, "gate_softplus_mul_perhead_many.cu", "ffai_gate_softplus_mul_perhead_many", + &[(attn.buffer.as_ref(), attn.offset), (g.buffer.as_ref(), g.offset), (out.buffer.as_ref(), 0)], + &[i(head_dim as i32)], + [rows as u32, 1, 1], [head_dim as u32, 1, 1], 0, false, + )?; + Ok(out) +} + +/// One-time compaction of a linear prefill K/V scratch buffer into a +/// fixed-size decode ring: copies the LAST `copy_len` absolute positions +/// `[copy_start, copy_start+copy_len)` from `src` `[n_kv_heads, src_cap, +/// head_dim]` into `dst` `[n_kv_heads, ring_cap, head_dim]` at ring slot +/// `abs_pos % ring_cap`. Used once per sliding-window layer after a Laguna +/// prefill finishes the whole prompt, to hand off from the prefill-only +/// linear scratch (`LagunaPrefillScratch`, sized to the prompt) to the +/// persistent decode ring (`LagunaKvCache`'s fixed `sliding_window`-slot +/// buffer) so decode can continue exactly where prefill left off. f32 only. +#[allow(clippy::too_many_arguments)] +pub fn kv_compact_ring( + dev: &dyn Device, + src: &Tensor, + dst: &Tensor, + n_kv_heads: usize, + src_cap: usize, + head_dim: usize, + ring_cap: usize, + copy_start: usize, + copy_len: usize, +) -> Result<()> { + if copy_len == 0 { + return Ok(()); + } + let i = |x: i32| x.to_le_bytes().to_vec(); + let total = (n_kv_heads * copy_len * head_dim) as u32; + dev.dispatch_raw_cuda( + KV_COMPACT_RING_SRC, "kv_compact_ring.cu", "ffai_kv_compact_ring", + &[(src.buffer.as_ref(), src.offset), (dst.buffer.as_ref(), dst.offset)], + &[ + i(n_kv_heads as i32), i(src_cap as i32), i(head_dim as i32), + i(ring_cap as i32), i(copy_start as i32), i(copy_len as i32), + ], + [total.div_ceil(256).max(1), 1, 1], [256, 1, 1], 0, false, + ) +} + +const KV_COMPACT_RING_SRC: &str = r#" +extern "C" __global__ void ffai_kv_compact_ring( + const float* __restrict__ src, float* __restrict__ dst, + int n_kv_heads, int src_cap, int hd, int ring_cap, int copy_start, int copy_len) +{ + long i = (long)blockIdx.x * blockDim.x + threadIdx.x; + long total = (long)n_kv_heads * copy_len * hd; + if (i >= total) return; + int d = (int)(i % hd); + long t = i / hd; + int c = (int)(t % copy_len); // 0..copy_len + int h = (int)(t / copy_len); // kv head + int src_pos = copy_start + c; // absolute position within the scratch + int dst_slot = src_pos % ring_cap; + dst[((long)h * ring_cap + dst_slot) * hd + d] = src[((long)h * src_cap + src_pos) * hd + d]; +} +"#; + /// Softmax over the last dim, row-wise. Dispatches `mt_softmax`. Row width /// `n` must be a multiple of 1024 (the kernel's 4-elems/thread loop). pub fn softmax(dev: &dyn Device, x: &Tensor) -> Result { @@ -6297,7 +6711,7 @@ pub fn mamba_split_proj( u(m_nh as u32), ], // PERF: pack threads/block (was block:[1,1,1] = 1 thread/block, 1/32 of a - // warp — this split was ~its full cost in launch overhead, not memory). The + // warp, this split was ~its full cost in launch overhead, not memory). The // kernel uses a global linear element id (gid_x), so block size is transparent // as long as the grid covers exactly s*in_proj_out threads. in_proj_out (10304 // = 64*161) is a multiple of 64 → s*in_proj_out always divisible by 64. @@ -6645,14 +7059,14 @@ extern "C" __global__ void moe_w4a16( const int m_tile_base = blockIdx.y * 64; const int warp_id = threadIdx.x >> 5; // 2×2 warp layout: warp_id 0→(m=0,n=0), 1→(m=0,n=32), 2→(m=32,n=0), 3→(m=32,n=32) - (void)(threadIdx.x & 31); // lane — unused in this kernel (warp_id sufficient) + (void)(threadIdx.x & 31); // lane, unused in this kernel (warp_id sufficient) const int warp_m_base = (warp_id >> 1) * 32; // 0 or 32 const int warp_n_base = (warp_id & 1) * 32; // 0 or 32 // ── Shared memory layout ─────────────────────────────────────────────── - // smem_X [64×32] f16 = 4096 bytes — row-major [m_local, k_local] - // smem_WT [32×64] f16 = 4096 bytes — row-major [k_local, n_local] = W transposed - // smem_C [64×64] f32 = 16384 bytes — output accumulation tile + // smem_X [64×32] f16 = 4096 bytes, row-major [m_local, k_local] + // smem_WT [32×64] f16 = 4096 bytes, row-major [k_local, n_local] = W transposed + // smem_C [64×64] f32 = 16384 bytes, output accumulation tile __shared__ __half smem_X [64 * 32]; // [m_local][k_local] __shared__ __half smem_WT[32 * 64]; // [k_local][n_local] ← W^T __shared__ float smem_C [64 * 64]; @@ -6813,8 +7227,8 @@ extern "C" __global__ void moe_w4a16( // row_group = j / 32 (0..7, each group = 8 consecutive rows) // nibble i → output row (row_group*8 + i), k = k_tile_base + k_local // -// This layout makes thread `t` load tile_base+t and tile_base+t+128 — two -// consecutive addresses — covering the full 64×32 tile with 4 cache lines +// This layout makes thread `t` load tile_base+t and tile_base+t+128, two +// consecutive addresses, covering the full 64×32 tile with 4 cache lines // per warp per load, vs ~64 scattered cache lines in the standard layout. // // `n_exp` = total experts in the weight pool (all experts, packed contiguously) @@ -6966,9 +7380,9 @@ extern "C" __global__ void moe_w4a16_marlin( const int t = threadIdx.x; // 0..127 // ── Shared memory layout ─────────────────────────────────────────────── - // smem_X [64×32] f16 = 4096 bytes — row-major [m_local, k_local] - // smem_WT [32×64] f16 = 4096 bytes — row-major [k_local, n_local] (W^T) - // smem_C [64×64] f32 = 16384 bytes — output accumulation tile + // smem_X [64×32] f16 = 4096 bytes, row-major [m_local, k_local] + // smem_WT [32×64] f16 = 4096 bytes, row-major [k_local, n_local] (W^T) + // smem_C [64×64] f32 = 16384 bytes, output accumulation tile __shared__ __half smem_X [64 * 32]; __shared__ __half smem_WT[32 * 64]; __shared__ float smem_C [64 * 64]; @@ -7383,7 +7797,7 @@ extern "C" __global__ void moe_w4a16_marlin128( } "#; -/// W4A16 MoE grouped GEMM — Marlin coalesced layout. +/// W4A16 MoE grouped GEMM, Marlin coalesced layout. /// /// Drop-in replacement for `moe_w4a16` but requires `qs` to be in Marlin /// tile-major layout (from `permute_q4_to_marlin`). Achieves coalesced @@ -7470,7 +7884,7 @@ pub fn moe_w4a16_marlin( /// W4A16 MoE grouped GEMM (Marlin-style, WMMA tensor cores, inline Q4 dequant). /// /// Drop-in replacement for `moe_bgemm_q4_bm64`. Uses `mma.sync m16n16k16` with -/// weights staged as W^T in smem — no f16 weight scratch in global memory, no +/// weights staged as W^T in smem, no f16 weight scratch in global memory, no /// 2× weight bandwidth. `n_out % 64 == 0` required. /// /// Returns `[m_total, n_out]` f16 (same contract as `moe_bgemm_q4_bm64`). @@ -7524,7 +7938,7 @@ pub fn moe_grouped_gemm( sc_up: &Tensor, qs_dn: &Tensor, sc_dn: &Tensor, - xs_f16: &Tensor, // [mt, hid] f16 — rows sorted by expert id + xs_f16: &Tensor, // [mt, hid] f16: rows sorted by expert id g_starts: &[usize], // length = n_active + 1 expert_ids: &[usize], // expert id for each group hid: usize, @@ -7560,10 +7974,10 @@ pub fn moe_grouped_gemm( // Defensive bounds check on caller-supplied scratch pools. The chunked // batched-dequant writes f16 expert weights to slot `c0..c1` of the scratch // at byte offset `c0*{up,dn}_slot*2`, and the per-expert GEMM reads/writes - // `slot*{up,dn}_slot*2` — so the scratch MUST hold `n_active` full slots + // `slot*{up,dn}_slot*2`, so the scratch MUST hold `n_active` full slots // (`n_active*{up,dn}_slot` f16 elems). A caller binding an under-sized pool // (e.g. sized for fewer active experts than this token batch routed to) - // triggers an out-of-bounds WRITE — silent corruption on Metal, a + // triggers an out-of-bounds WRITE, silent corruption on Metal, a // deterministic MMU-fault on CUDA. Fail loudly here, naming the shapes. if let Some(scratch) = up_scratch { let need = n_active * up_slot; @@ -7595,7 +8009,7 @@ pub fn moe_grouped_gemm( // arrays serializes the stream. Per-expert cuBLAS (27 TFLOP/s) remains best. // ── UP scratch: full `n_active` slots (distinct regions per chunk so chunk - // N+1's dequant overlaps chunk N's GEMMs — same-buffer reuse would serialize). ── + // N+1's dequant overlaps chunk N's GEMMs, same-buffer reuse would serialize). ── let up_w_owned: Option; let up_w = if let Some(scratch) = up_scratch { Tensor { buffer: scratch.buffer.clone(), offset: 0, shape: vec![n_active * inter, hid], dtype: DType::F16 } @@ -7710,17 +8124,17 @@ pub fn moe_grouped_gemm( } // ════════════════════════════════════════════════════════════════════════════ -// NVFP4 (mxf4nvf4) grouped MoE GEMM — native Blackwell block-scaled FP4 mma. +// NVFP4 (mxf4nvf4) grouped MoE GEMM, native Blackwell block-scaled FP4 mma. // Fragment-packed layout: weights repacked once at load (nvfp4_pack_w), acts // quantized+packed per layer (moe_fp4_actq_pack), one fixed-launch grouped mma // over BM=16 tile descriptors (graph-safe). Requires METALTILE_CUDA_ARCH= -// compute_121a (or _120a) — base-arch PTX rejects the mma. +// compute_121a (or _120a), base-arch PTX rejects the mma. // ════════════════════════════════════════════════════════════════════════════ const MOE_FP4_SRC: &str = r#" #include -// UE4M3 codebook (monotonic): e=0 subnormal m/8*2^-7? NO — matches quantize_nvfp4: +// UE4M3 codebook (monotonic): e=0 subnormal m/8*2^-7? NO, matches quantize_nvfp4: // e==0 → (m/8)*0.0078125 (=2^-7), else (1+m/8)*2^(e-7). __device__ __forceinline__ float dec_ue4m3(unsigned b){ int e=(b>>3)&0xF, m=b&7; @@ -7954,7 +8368,7 @@ extern "C" __global__ void moe_fp4_actq_pack( Asc[fbase*32+lane]=sw; } -// VARIANT: each warp computes 4 consecutive n8-tiles — A frags loaded once, +// VARIANT: each warp computes 4 consecutive n8-tiles, A frags loaded once, // reused 4x; W streams per-subtile. Cuts A traffic 4x, raises mma density. // 16-byte async global->shared copy (per-warp staging; no block syncs). __device__ __forceinline__ void mt_cp16(void* sm, const void* gm){ @@ -7974,7 +8388,7 @@ extern "C" __global__ void moe_fp4_grouped_mma_x4( // direct global->register W loads left the mma latency-bound (~42 TF); // double-buffering the 4 W tiles per warp while the mma consumes the // other stage measures 47.9 TF standalone on GB10 (+14%). Each warp owns - // its smem slot — per-warp wait_group + syncwarp only, NO __syncthreads + // its smem slot, per-warp wait_group + syncwarp only, NO __syncthreads // (the cross-warp shared-W variant with block syncs measured SLOWER). __shared__ unsigned smW[8][2][256]; // [warp][stage][4 tiles x 64 u32] __shared__ unsigned smS[8][2][32]; // [warp][stage][4 tiles x 8 u32] @@ -8060,7 +8474,7 @@ extern "C" __global__ void moe_fp4_grouped_mma_x4( } } -// Grouped mma: one warp = (tile, nt). FIXED launch grid — padding tiles +// Grouped mma: one warp = (tile, nt). FIXED launch grid, padding tiles // (gend==0) early-out → graph-safe. out f16 [mt, N] in sorted-row space. extern "C" __global__ void moe_fp4_grouped_mma( const unsigned* __restrict__ Ap, const unsigned* __restrict__ Asc, @@ -8326,7 +8740,7 @@ extern "C" __global__ void fwht128_quant( q[base + t]=(unsigned char)(sgn|c); } -// f32-input variant of lt_fp4_quant — reads the residual-stream f32 tensor +// f32-input variant of lt_fp4_quant, reads the residual-stream f32 tensor // directly, killing the cast_f32_f16 + f16 intermediate round-trip that // preceded every dense-projection quant (f32 -> f16 -> e2m1 became // f32 -> e2m1; the e2m1 grid subsumes the f16 rounding). @@ -8404,7 +8818,7 @@ extern "C" __global__ void nvfp4_gscale( } // fp4_sfa_offsets: device build of the per-group SFA byte offsets from the -// router offsets — sfa_off[g] = (sum_{j Result< /// e2m1 `[mt*k/2]` bytes (global row-major) + a per-group swizzled ue4m3 /// scale pool where group g's blob (group-LOCAL rows, padded to 128-row /// stripes, zero-seeded) starts at byte `sfa_off[g]`. A per-call activation -/// GLOBAL (amax/(6·448), dynamic, device-computed — no host sync) centers the +/// GLOBAL (amax/(6·448), dynamic, device-computed: no host sync) centers the /// block scales in the ue4m3 normal range; the GEMM folds it back in via the /// per-group alpha (see [`moe_grouped_gemm_cutlass_fp4`]). `g_starts` has /// `n_groups+1` entries (`g_starts[n_groups] == mt`). Returns -/// `(pack, sf, sfa_off, gs)` — `gs` = device `[global, inv_global]` f32 pair. +/// `(pack, sf, sfa_off, gs)`, `gs` = device `[global, inv_global]` f32 pair. pub fn lt_fp4_quant_grouped( dev: &dyn Device, x: &Tensor, mt: usize, k: usize, g_starts: &[usize], ) -> Result<(Tensor, Tensor, Vec, Tensor)> { @@ -8752,7 +9166,7 @@ pub fn lt_fp4_quant_grouped( /// block-scaled NVFP4 in the CUTLASS/cuBLASLt canonical layout: per-expert /// packed e2m1 (`n*k/2` bytes each, contiguous) + per-expert swizzled ue4m3 /// scales (`ceil(n/128)*512*ceil(k/64)` bytes each, contiguous, zero-seeded) -/// + per-expert GLOBALS `gw [n_exp]` f32 (amax/(6·448) — keeps the block +/// + per-expert GLOBALS `gw [n_exp]` f32 (amax/(6·448), keeps the block /// scales in the ue4m3 normal range; the GEMM folds `gw` back in via the /// per-group alpha). One quant dispatch per expert so every expert's SF blob /// is laid out from local row 0 (an `[n_exp*n, k]` whole-slab quant would @@ -8785,7 +9199,7 @@ pub fn lt_fp4_quant_slab( let pack = Tensor::new(dev.alloc(n_exp * pack_exp)?, vec![n_exp * pack_exp], DType::U32); let sf = Tensor::new(dev.alloc_zeroed(n_exp * sf_exp)?, vec![n_exp * sf_exp], DType::U32); let nb = (n * kb) as u32; - // B8a: routed-expert WEIGHT 4/6 (Four-Over-Six) under NEMOTRON_FP4_46 — free + // B8a: routed-expert WEIGHT 4/6 (Four-Over-Six) under NEMOTRON_FP4_46, free // (setup-time, covers ~90% of params). Same per-expert global gw[e] (gs[0]). let kname = if std::env::var("NEMOTRON_FP4_46").is_ok() { "lt_fp4_quant_gs_46" } else { "lt_fp4_quant_gs" }; for e in 0..n_exp { @@ -8802,7 +9216,7 @@ pub fn lt_fp4_quant_slab( /// Device-descriptor variant of [`lt_fp4_quant_grouped`]: per-group SFA /// offsets are computed ON DEVICE from the router offsets (`off_dev`, u32 -/// `[n_groups+1]`) — no host offsets, no uploads, graph-safe. The SFA pool is +/// `[n_groups+1]`), no host offsets, no uploads, graph-safe. The SFA pool is /// sized worst-case (`mt/128 + n_groups` stripes). Returns `(pack, sf, gs)`; /// the GEMM run derives the same per-group offsets from `off_dev`. pub fn lt_fp4_quant_grouped_dev( @@ -8862,12 +9276,12 @@ pub fn fp4_group_alpha_id(dev: &dyn Device, gs: &Tensor, gw: &Tensor, alpha: &Te [(n as u32).div_ceil(64), 1, 1], [64, 1, 1], 0, false) } -/// CUTLASS grouped block-scaled NVFP4 MoE GEMM — `out[t,n](f16) = Σ_k +/// CUTLASS grouped block-scaled NVFP4 MoE GEMM, `out[t,n](f16) = Σ_k /// A[t,k]·W[eid][n,k]` over sorted-token groups with BOTH operands in packed /// e2m1 + swizzled ue4m3 block-16 scales. `a_pack`/`a_sf`/`sfa_off`/`a_gs` /// from [`lt_fp4_quant_grouped`]; `w_pack`/`w_sf`/`w_gw` from /// [`lt_fp4_quant_slab`]. The per-tensor globals are folded back in through -/// the GEMM's per-group alpha (`gA·gw[eid]`, built device-side — no host +/// the GEMM's per-group alpha (`gA·gw[eid]`, built device-side: no host /// sync). `g_starts[g]..g_starts[g+1]` = group g's token rows; /// `expert_ids[g]` selects the weight slab. ONE GEMM launch for all groups. /// Returns `out [mt, n]` f16. CUDA + CUTLASS_DIR sm_120a/121a build only. @@ -9098,7 +9512,7 @@ pub fn w8a8_packw( /// Quantize an f16 tensor to e4m3 with a dynamic per-tensor scale. -/// Returns `(q, sc)` — e4m3 bytes `[n]` + a 2-f32 scale buffer +/// Returns `(q, sc)`, e4m3 bytes `[n]` + a 2-f32 scale buffer /// `[dequant_scale, quant_scale]` (the GEMM consumes word 0). pub fn fp8_quant(dev: &dyn Device, x: &Tensor, n: usize, clip: f32) -> Result<(Tensor, Tensor)> { let i = |v: i32| v.to_le_bytes().to_vec(); @@ -9222,7 +9636,7 @@ pub fn gemm_fp8( } /// Per-row e4m3 quant of an f16/f32 `[r,k]` tensor → `(q [r*k] bytes, scale [r] f32)`. -/// Per-channel (weights, rows=n) / per-token (acts, rows=m) — the quality- +/// Per-channel (weights, rows=n) / per-token (acts, rows=m), the quality- /// preserving FP8 granularity (per-tensor crushed cosine to ~0.94). pub fn fp8_quant_perrow(dev: &dyn Device, x: &Tensor, r: usize, k: usize) -> Result<(Tensor, Tensor)> { let i = |v: i32| v.to_le_bytes().to_vec(); @@ -9236,7 +9650,7 @@ pub fn fp8_quant_perrow(dev: &dyn Device, x: &Tensor, r: usize, k: usize) -> Res } /// FP8 e4m3 GEMM with PER-CHANNEL weight + PER-TOKEN act scales, folded as a -/// post-pass (`D[r,c] *= xsc[r]·wsc[c]`) — sidesteps cuBLASLt's unsupported +/// post-pass (`D[r,c] *= xsc[r]·wsc[c]`), sidesteps cuBLASLt's unsupported /// OUTER_VEC FP8 path. `x_q`/`x_sc` per-token (rows=m), `w_q`/`w_sc` /// per-channel (rows=n). Quality-preserving (cosine >0.98 typical). out[m,n] f32. #[allow(clippy::too_many_arguments)] @@ -9247,7 +9661,7 @@ pub fn gemm_fp8_perchan( m: usize, n: usize, k: usize, ) -> Result { let out = Tensor::empty(dev, vec![m, n], DType::F32)?; - // identity per-tensor scale (1.0) — the real per-channel/per-token scales + // identity per-tensor scale (1.0), the real per-channel/per-token scales // are folded in the post-pass below. D_raw = Wq·Xq in f32. let one = scalar_buf(dev, (1.0f32).to_bits())?; dev.gemm_fp8( @@ -9349,7 +9763,7 @@ pub fn lt_fp4_quant_g(dev: &dyn Device, x: &Tensor, r: usize, k: usize) -> Resul /// No-global FP4 quant: per-16 ue4m3 block scales absorb the FULL magnitude /// (global gs=[1,1]), so the dequant stays ~true-scale and the f16-output GEMM -/// (out_f32=false) does NOT overflow — fast f16 path AND correct (vs the +/// (out_f32=false) does NOT overflow, fast f16 path AND correct (vs the /// two-level global variant whose globally-normalized raw blows past f16 max). pub fn lt_fp4_quant_g_ng(dev: &dyn Device, x: &Tensor, r: usize, k: usize) -> Result<(Tensor, Tensor, Tensor)> { if k % 32 != 0 { return Err(Error::Msg(format!("lt_fp4_quant_g_ng: k {k} %32!=0"))); } @@ -9380,7 +9794,7 @@ pub fn fp4_dscale(dev: &dyn Device, ga: &Tensor, gw: &Tensor) -> Result /// Quantize+fragment-pack MoE expert weights to NVFP4 on device. /// `w_f16`: f16 `[n_exp*n, k]` row-major (concatenated experts). Returns -/// `(wp, wsc, gw)` — fragment codes, scale words, per-expert globals (device). +/// `(wp, wsc, gw)`, fragment codes, scale words, per-expert globals (device). pub fn moe_fp4_pack_weights( dev: &dyn Device, w_f16: &Tensor, n_exp: usize, n: usize, k: usize, ) -> Result<(Tensor, Tensor, Tensor)> { @@ -9418,7 +9832,7 @@ pub fn moe_fp4_pack_weights( /// Fully-on-device NVFP4 grouped MoE GEMM (W4A4): quantizes the f16 sorted /// activations `a [mt,k]` to NVFP4 (dynamic per-call global), packs fragments, /// and runs the Blackwell block-scaled mma over BM=16 tile descriptors built -/// from `offsets` (device `[n_exp+1]`). Fixed launch — graph-safe, no host +/// from `offsets` (device `[n_exp+1]`). Fixed launch, graph-safe, no host /// syncs. Returns `out [mt, n]` f16. #[allow(clippy::too_many_arguments)] #[allow(clippy::too_many_arguments)] @@ -9479,7 +9893,7 @@ pub fn moe_fp4_grouped_mma_dev( i(if gather_idx.is_some() { 1 } else { 0 })], [(warps * 32).div_ceil(256), 1, 1], [256, 1, 1], 0, false)?; // grouped mma: grid [maxt, NT/8] × 8 warps/block - // x4 variant: 4 n8-tiles per warp (A frags register-reused) — +49% on + // x4 variant: 4 n8-tiles per warp (A frags register-reused), +49% on // GB10 vs 1 nt/warp (L2-poorer than desktop Blackwell, where x1 wins). let warps_per_block = 8u32; let ng = ntl.div_ceil(4) as u32; @@ -9500,7 +9914,7 @@ pub fn moe_fp4_grouped_mma_dev( } /// Fused residual-add + RMS norm: `residual_out = a + b`, `normed_out = -/// rms_norm(a + b, w)` — one pass instead of an `add` kernel plus a +/// rms_norm(a + b, w)`, one pass instead of an `add` kernel plus a /// `rms_norm` kernel re-reading the sum (the per-layer residual hot path). pub fn add_rms_norm( dev: &dyn Device, a: &Tensor, b: &Tensor, weight: &Tensor, eps: f32,