diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 6cb05026..03e4ea87 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -532,6 +532,7 @@ dependencies = [ name = "wh-butter-models" version = "0.1.0" dependencies = [ + "rayon", "serde_json", "wh-butter-core", "wh-butter-loader", diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 6433fad5..d7844862 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -23,6 +23,7 @@ authors = ["Eric Kryski (@ekryski)", "Tom Turney (@TheTom)"] repository = "https://github.com/thewafflehaus/butter" [workspace.dependencies] +rayon = "1" # ── engine crates ──────────────────────────────────────────────────── wh-butter-core = { path = "crates/wh-butter-core" } wh-butter-ops = { path = "crates/wh-butter-ops" } diff --git a/rust/crates/backends/wh-butter-cuda/Cargo.toml b/rust/crates/backends/wh-butter-cuda/Cargo.toml index 3705b6e9..4c24d8c2 100644 --- a/rust/crates/backends/wh-butter-cuda/Cargo.toml +++ b/rust/crates/backends/wh-butter-cuda/Cargo.toml @@ -112,3 +112,7 @@ required-features = ["cuda"] [[test]] name = "all_models" required-features = ["cuda"] + +[[test]] +name = "proj_fp8_bench" +required-features = ["cuda"] diff --git a/rust/crates/backends/wh-butter-cuda/src/imp.rs b/rust/crates/backends/wh-butter-cuda/src/imp.rs index 3d88e221..a871aee2 100644 --- a/rust/crates/backends/wh-butter-cuda/src/imp.rs +++ b/rust/crates/backends/wh-butter-cuda/src/imp.rs @@ -27,12 +27,12 @@ fn dispatch_err(e: IronError) -> 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) @@ -413,6 +433,19 @@ impl Device for CudaDevice { self.dev.moe_grouped_cutlass(ab.ptr, wb.ptr, ob.ptr, group_rows, expert_ids, n, k).map_err(dispatch_err) } + // NOTE(rebrand-merge): PR #70 also implemented moe_marlin_gemm / + // marlin_repack / marlin_build_routing here, calling through to + // `self.dev.{moe_marlin_gemm,marlin_repack,marlin_build_routing}` on the + // raw `wh_iron_runtime::CudaDevice`. Those methods don't exist on + // `thewafflehaus/iron@dev`'s CudaDevice — the marlin/NVFP4 CUDA runtime + // support was apparently only ever on the pre-rebrand kernel-toolchain + // fork this branch was originally pinned to. Left unimplemented here + // (falls through to the wh-butter-core default "unsupported on this + // backend" stub) until the iron runtime crate grows real marlin + // support; re-add the three pass-through impls verbatim from this + // crate's pre-rebrand-merge git history (PR #70 / a27773ec, `src/imp.rs`) + // once it does. + fn moe_grouped_cutlass_fp4( &self, a: &dyn DeviceBuffer, diff --git a/rust/crates/backends/wh-butter-cuda/tests/all_models.rs b/rust/crates/backends/wh-butter-cuda/tests/all_models.rs index bc1c5155..2b23222a 100644 --- a/rust/crates/backends/wh-butter-cuda/tests/all_models.rs +++ b/rust/crates/backends/wh-butter-cuda/tests/all_models.rs @@ -8,11 +8,10 @@ fn all_models_on_cuda() { wh_butter_modeltests::run_all(dev.as_ref(), "GB10 sm_121"); } -/// FAST FUSE_SLICECAST lever A/B on a small same-family model (loads in seconds). -/// See `wh_butter_modeltests::bench_smallmodel_fuse_slicecast`. Default mamba2-130m; -/// FUSE_SC_MODEL=falcon for Falcon-H1-0.5B. Off-vs-on bit-exact + timing. -#[test] -fn smallmodel_fuse_slicecast_ab() { - let Some(dev) = CudaDevice::create().expect("cuda") else { eprintln!("no CUDA — skip"); return; }; - wh_butter_modeltests::bench_smallmodel_fuse_slicecast(dev.as_ref(), "GB10 sm_121"); -} +// NOTE(rebrand-merge): this file used to also have a `smallmodel_fuse_slicecast_ab` +// test calling `wh_butter_modeltests::bench_smallmodel_fuse_slicecast`, but that +// function doesn't exist anywhere in the modeltests crate on the original +// PR #70 (a27773ec) commit either (pre-existing dead reference, not +// something this cherry-pick introduced — the FUSE_SLICECAST bench was +// apparently never actually landed in wh-butter-modeltests). Dropped rather +// than stubbed since there's no reference implementation to port. diff --git a/rust/crates/backends/wh-butter-cuda/tests/colcopy.rs b/rust/crates/backends/wh-butter-cuda/tests/colcopy.rs new file mode 100644 index 00000000..ceeef2f2 --- /dev/null +++ b/rust/crates/backends/wh-butter-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 wh_butter_core::{DType, Device, Tensor}; +use wh_butter_cuda::CudaDevice; +use wh_butter_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 wh_butter_core::{DType, Tensor}; + use wh_butter_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 = wh_butter_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 = wh_butter_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/wh-butter-cuda/tests/f16norm_f32in.rs b/rust/crates/backends/wh-butter-cuda/tests/f16norm_f32in.rs deleted file mode 100644 index 1a92798d..00000000 --- a/rust/crates/backends/wh-butter-cuda/tests/f16norm_f32in.rs +++ /dev/null @@ -1,70 +0,0 @@ -#![cfg(feature = "cuda")] -//! Isolate add_rms_norm_f16norm with REAL in-model f32 inputs (a/b/w all f32). -//! Compares residual + normed against a CPU oracle AND the working add_rms_norm. -use wh_butter_core::{DType, Device, Tensor}; -use wh_butter_cuda::CudaDevice; -use wh_butter_ops::{add_rms_norm, add_rms_norm_f16norm}; - -fn tb(v:&[f32])->Vec{v.iter().flat_map(|x|x.to_le_bytes()).collect()} -fn fb32(b:&[u8])->Vec{b.chunks_exact(4).map(|c|f32::from_le_bytes(c.try_into().unwrap())).collect()} -fn f16_to_f32(bits:u16)->f32{ - let sign=((bits>>15)&1) as u32; let exp=((bits>>10)&0x1f) as u32; let mant=(bits&0x3ff) as u32; - let f=if exp==0 { if mant==0 { sign<<31 } else { // subnormal - let mut e=-1i32; let mut m=mant; while m & 0x400 ==0 { m<<=1; e-=1; } m&=0x3ff; - ((sign<<31)|(((e+127+(-14))as u32)<<23)|(m<<13)) } } - else if exp==0x1f { (sign<<31)|0x7f800000|(mant<<13) } - else { (sign<<31)|((exp+(127-15))<<23)|(mant<<13) }; - f32::from_bits(f) -} -fn fb16(b:&[u8])->Vec{b.chunks_exact(2).map(|c|f16_to_f32(u16::from_le_bytes([c[0],c[1]]))).collect()} -fn tn(d:&dyn Device,v:&[f32],sh:Vec)->Tensor{Tensor::new(d.upload(&tb(v)).unwrap(),sh,DType::F32)} - -#[test] -fn f16norm_f32_inputs_matches_oracle(){ - let Some(dev)=CudaDevice::create().expect("cuda") else { eprintln!("no CUDA — skip"); return; }; - let d = dev.as_ref(); - let (rows, n) = (8usize, 4096usize); // real hidden size - let eps = 1e-5f32; - let a: Vec = (0..rows*n).map(|i| ((i % 13) as f32 * 0.07 - 0.4)).collect(); - let b: Vec = (0..rows*n).map(|i| ((i % 7) as f32 * 0.05 - 0.15)).collect(); - let w: Vec = (0..n).map(|i| 0.5 + (i % 5) as f32 * 0.1).collect(); - // CPU oracle - let mut res_ref = vec![0f32; rows*n]; let mut nrm_ref = vec![0f32; rows*n]; - for r in 0..rows { let mut ssq=0f32; for c in 0..n { let s=a[r*n+c]+b[r*n+c]; res_ref[r*n+c]=s; ssq+=s*s; } - let rms=1f32/(ssq/n as f32+eps).sqrt(); for c in 0..n { nrm_ref[r*n+c]=res_ref[r*n+c]*rms*w[c]; } } - - let at=tn(d,&a,vec![rows,n]); let bt=tn(d,&b,vec![rows,n]); let wt=tn(d,&w,vec![n]); - // f32 control - let (r32, n32) = add_rms_norm(d, &at, &bt, &wt, eps).unwrap(); - // fused f16norm - let (rf, nf) = add_rms_norm_f16norm(d, &at, &bt, &wt, eps).unwrap(); - d.synchronize().unwrap(); - assert_eq!(rf.dtype, DType::F32, "residual must be f32"); - assert_eq!(nf.dtype, DType::F16, "normed must be f16"); - - let mut g=vec![0u8;rows*n*4]; d.download(rf.buffer.as_ref(),&mut g).unwrap(); let res_got=fb32(&g); - let mut h=vec![0u8;rows*n*2]; d.download(nf.buffer.as_ref(),&mut h).unwrap(); let nrm_got=fb16(&h); - let mut c32=vec![0u8;rows*n*4]; d.download(r32.buffer.as_ref(),&mut c32).unwrap(); let res_c=fb32(&c32); - let mut cn=vec![0u8;rows*n*4]; d.download(n32.buffer.as_ref(),&mut cn).unwrap(); let nrm_c=fb32(&cn); - - // EXACT-FUSION CHECK: fused f16 normed must equal cast_f32_f16(f32wrap normed) bit-for-bit. - let nrm_c_f16 = wh_butter_ops::cast_f32_f16(d, &n32).unwrap(); - d.synchronize().unwrap(); - let mut hc=vec![0u8;rows*n*2]; d.download(nrm_c_f16.buffer.as_ref(),&mut hc).unwrap(); - let mut hf=vec![0u8;rows*n*2]; d.download(nf.buffer.as_ref(),&mut hf).unwrap(); - let bitdiff = hc.iter().zip(hf.iter()).filter(|(a,b)| a!=b).count(); - eprintln!("FUSED vs cast_f32_f16(f32wrap): differing f16 BYTES = {bitdiff} / {}", rows*n*2); - let res_nan=res_got.iter().any(|x|!x.is_finite()); - let nrm_nan=nrm_got.iter().any(|x|!x.is_finite()); - let mut res_err=0f32; for i in 0..rows*n { res_err=res_err.max((res_got[i]-res_ref[i]).abs()); } - let mut nrm_err=0f32; for i in 0..rows*n { let d0=(nrm_got[i]-nrm_ref[i]).abs(); let den=nrm_ref[i].abs().max(1e-3); nrm_err=nrm_err.max(d0/den); } - let mut res_vs32=0f32; for i in 0..rows*n { res_vs32=res_vs32.max((res_got[i]-res_c[i]).abs()); } - eprintln!("f16norm f32-in: res_nan={res_nan} nrm_nan={nrm_nan} res|Δ|oracle={res_err:.3e} nrm_rel|Δ|oracle={nrm_err:.3e} res|Δ|vs_f32wrap={res_vs32:.3e}"); - eprintln!(" res_got[0..4]={:?} ref={:?}", &res_got[0..4], &res_ref[0..4]); - eprintln!(" nrm_got[0..4]={:?} ref={:?}", &nrm_got[0..4], &nrm_ref[0..4]); - eprintln!(" f32wrap nrm[0..4]={:?}", &nrm_c[0..4]); - assert!(!res_nan && !nrm_nan, "NaN in fused outputs"); - assert!(res_err < 1e-3, "residual wrong vs oracle: {res_err}"); - assert!(nrm_err < 2e-2, "normed wrong vs oracle (f16 tol): {nrm_err}"); - assert!(res_vs32 < 1e-3, "residual diverges from f32 wrapper: {res_vs32}"); -} diff --git a/rust/crates/backends/wh-butter-cuda/tests/laguna.rs b/rust/crates/backends/wh-butter-cuda/tests/laguna.rs new file mode 100644 index 00000000..4313a322 --- /dev/null +++ b/rust/crates/backends/wh-butter-cuda/tests/laguna.rs @@ -0,0 +1,507 @@ +// Copyright 2026 Eric Kryski (@ekryski) and Tom Turney (@TheTom) +// SPDX-License-Identifier: Apache-2.0 +//! Laguna-S-2.1 decode smoke on CUDA. Env-gated: set BUTTER_LAGUNA_GGUF to the +//! model's GGUF path (the 71G checkpoint), otherwise the test skips. +use wh_butter_cuda::CudaDevice; +#[test] +fn laguna_host_dequant() { + wh_butter_modeltests::verify_laguna_host_dequant(); +} + +#[test] +fn laguna_host_layer0() { + wh_butter_modeltests::verify_laguna_host_layer0(); +} + +/// Isolated rope_yarn_partial check with synthetic data. At position 0 the +/// rotation angle is zero for every dim, so rotated dims must equal +/// input * mscale (mscale = attn_factor exactly; the wrapper applies no +/// internal re-multiply) and passthrough dims must be bit-identical. +#[test] +fn laguna_rope_yarn_unit() { + use wh_butter_core::{DType, Tensor}; + let Some(dev) = CudaDevice::create().expect("cuda") else { eprintln!("no CUDA, skip"); return; }; + let d = dev.as_ref(); + let n_heads = 3usize; + let hd = 8usize; + let n_rot = 4u32; + let n: usize = n_heads * hd; + let x: Vec = (0..n).map(|i| (i as f32) * 0.25 - 2.0).collect(); + let xb: Vec = x.iter().flat_map(|v| v.to_le_bytes()).collect(); + let t = Tensor::new(d.upload(&xb).unwrap(), vec![n_heads, hd], DType::F32); + + let attn_factor = 1.4852030263919618f32; + let freq_scale = 1.0f32 / 128.0; + let out = wh_butter_ops::rope_yarn_partial(d, &t, 0, n_rot, 500000.0, freq_scale, 1.0, attn_factor, 32.0, 1.0, 8192).unwrap(); + let mut ob = vec![0u8; n * 4]; + d.synchronize().unwrap(); + d.download(out.buffer.as_ref(), &mut ob).unwrap(); + let o: Vec = ob.chunks_exact(4).map(|c| f32::from_le_bytes(c.try_into().unwrap())).collect(); + + let mscale = attn_factor; // attn_factor is the final mscale, no internal re-multiply + let mut bad = 0; + for h in 0..n_heads { + for i in 0..hd { + let idx = h * hd + i; + let expect = if (i as u32) < n_rot { x[idx] * mscale } else { x[idx] }; + if (o[idx] - expect).abs() > 1e-4 * expect.abs().max(1.0) { + eprintln!("pos0 MISMATCH h={h} i={i}: got {} expect {expect}", o[idx]); + bad += 1; + } + } + } + assert_eq!(bad, 0, "rope_yarn_partial wrong at position 0 ({bad} bad elems)"); + eprintln!("rope pos-0 identity*mscale: OK (mscale={mscale})"); + + // Exact model geometry: 48 heads x 128 dims, n_rot 64, position 0. The + // same identity*mscale invariant must hold; this catches launch-geometry + // bugs the tiny shape above cannot (block y = 64, grid x = 48). + let n_heads2 = 48usize; + let hd2 = 128usize; + let n2 = n_heads2 * hd2; + let x2: Vec = (0..n2).map(|i| ((i * 2654435761usize) % 1000) as f32 * 0.002 - 1.0).collect(); + let xb2: Vec = x2.iter().flat_map(|v| v.to_le_bytes()).collect(); + let t2 = Tensor::new(d.upload(&xb2).unwrap(), vec![n_heads2, hd2], DType::F32); + let out2 = wh_butter_ops::rope_yarn_partial(d, &t2, 0, 64, 500000.0, freq_scale, 1.0, attn_factor, 32.0, 1.0, 8192).unwrap(); + let mut ob2 = vec![0u8; n2 * 4]; + d.synchronize().unwrap(); + d.download(out2.buffer.as_ref(), &mut ob2).unwrap(); + let o2: Vec = ob2.chunks_exact(4).map(|c| f32::from_le_bytes(c.try_into().unwrap())).collect(); + let mut bad2 = 0; + for h in 0..n_heads2 { + for i in 0..hd2 { + let idx = h * hd2 + i; + let expect = if i < 64 { x2[idx] * mscale } else { x2[idx] }; + if (o2[idx] - expect).abs() > 1e-4 * expect.abs().max(1.0) { + if bad2 < 6 { eprintln!("model-geom pos0 MISMATCH h={h} i={i}: got {} expect {expect}", o2[idx]); } + bad2 += 1; + } + } + } + assert_eq!(bad2, 0, "rope_yarn_partial wrong at model geometry pos 0 ({bad2} bad elems)"); + eprintln!("rope model-geometry pos-0: OK"); +} + +/// Isolated per-head softplus gate check with synthetic data. +#[test] +fn laguna_gate_unit() { + use wh_butter_core::{DType, Tensor}; + let Some(dev) = CudaDevice::create().expect("cuda") else { eprintln!("no CUDA, skip"); return; }; + let d = dev.as_ref(); + let n_heads = 4usize; + let hd = 8usize; + let attn: Vec = (0..n_heads * hd).map(|i| (i as f32) * 0.1 - 1.0).collect(); + let g: Vec = vec![-2.0, 0.0, 1.5, 30.0]; + 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 ta = up(&attn).reshaped(vec![n_heads, hd]); + let tg = up(&g); + let out = wh_butter_ops::gate_softplus_mul_perhead(d, &ta, &tg).unwrap(); + let mut ob = vec![0u8; attn.len() * 4]; + d.synchronize().unwrap(); + d.download(out.buffer.as_ref(), &mut ob).unwrap(); + let o: Vec = ob.chunks_exact(4).map(|c| f32::from_le_bytes(c.try_into().unwrap())).collect(); + let mut bad = 0; + for h in 0..n_heads { + let sp = if g[h] > 20.0 { g[h] } else { (1.0 + g[h].exp()).ln() }; + for i in 0..hd { + let idx = h * hd + i; + let expect = attn[idx] * sp; + if (o[idx] - expect).abs() > 1e-5 * expect.abs().max(1.0) { + eprintln!("gate MISMATCH h={h} i={i}: got {} expect {expect}", o[idx]); + bad += 1; + } + } + } + assert_eq!(bad, 0, "gate_softplus_mul_perhead wrong ({bad} bad elems)"); + 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 wh_butter_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 = wh_butter_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 = wh_butter_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 wh_butter_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 = wh_butter_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 = wh_butter_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 `BUTTER_LAGUNA_GRAPH=1`) must match the registry +/// `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}: +/// Laguna's actual GQA ratios (48 full-attention / 72 sliding-window Q heads +/// over 8 KV heads). +#[test] +fn laguna_sdpa_decode_nbuf_matches_registry() { + use wh_butter_core::{DType, Tensor}; + let Some(dev) = CudaDevice::create().expect("cuda") else { eprintln!("no CUDA, skip"); return; }; + let d = dev.as_ref(); + let hd = 128usize; + let n_kv_heads = 4usize; + let kv_stride = 512u32; + 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) + }; + + for &heads_per_group in &[6usize, 9usize] { + let n_q_heads = n_kv_heads * heads_per_group; + let q = ramp(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![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]); + + for &n_kv in &[1u32, 7, 512] { + let expect = wh_butter_ops::sdpa_decode(d, &tq, &tk, &tv, hd, n_kv, kv_stride, heads_per_group as u32, scale).unwrap(); + let mut eb = vec![0u8; n_q_heads * hd * 4]; + d.synchronize().unwrap(); + d.download(expect.buffer.as_ref(), &mut eb).unwrap(); + let ev: Vec = eb.chunks_exact(4).map(|c| f32::from_le_bytes(c.try_into().unwrap())).collect(); + + let n_kv_buf = Tensor::new(d.upload(&n_kv.to_le_bytes()).unwrap(), vec![1], DType::U32); + let got = wh_butter_ops::sdpa_decode_nbuf(d, &tq, &tk, &tv, hd, &n_kv_buf, kv_stride, heads_per_group as u32, scale).unwrap(); + let mut gb = vec![0u8; n_q_heads * hd * 4]; + d.synchronize().unwrap(); + d.download(got.buffer.as_ref(), &mut gb).unwrap(); + let gv: Vec = gb.chunks_exact(4).map(|c| f32::from_le_bytes(c.try_into().unwrap())).collect(); + + let mut bad = 0; + for i in 0..ev.len() { + if (ev[i] - gv[i]).abs() > 1e-4 * ev[i].abs().max(1.0) { + if bad < 6 { + eprintln!("MISMATCH hpg={heads_per_group} n_kv={n_kv} i={i}: nbuf={} registry={}", gv[i], ev[i]); + } + bad += 1; + } + } + assert_eq!(bad, 0, "sdpa_decode_nbuf mismatch vs registry sdpa_decode (heads_per_group={heads_per_group}, n_kv={n_kv}, {bad} bad elems)"); + } + } + eprintln!("sdpa_decode_nbuf matches registry sdpa_decode: OK"); +} + +/// `iron_moe_gather_q4_swiglu` (the fused gate+up+SwiGLU MoE gather kernel +/// used by Laguna's decode path under `BUTTER_LAGUNA_FUSE=1`) must match the +/// unfused compose: `moe_gather_q4` run twice (gate stack, up stack) then a +/// host-side SwiGLU on the downloaded pair. Synthetic 4-expert stack, top_k +/// 3, inter 8, hid 32 (small but exercises the same Q4 block-32 / f16 +/// per-block-scale layout the real Laguna MoE stacks use). Modeled on +/// `laguna_gate_unit` above: synthetic buffers, upload, compare with a +/// relative tolerance. +/// +/// TODO(rebrand-merge): `wh_butter_ops::moe_gather_q4` / `moe_gather_q4_swiglu` +/// dispatch `iron_moe_gather_q4` / `iron_moe_gather_q4_swiglu` via +/// `wh_iron_std`'s kernel IR, but the pinned `thewafflehaus/iron@dev` kernel +/// library only ships `iron_moe_gather_q4_relu2/_down/_down_accum` — the +/// plain (non-relu2) and `_swiglu` variants this test needs are not present +/// upstream yet. Ignored until those kernels land in the iron repo; do not +/// silently unignore without confirming they exist. +#[ignore = "blocked: iron_moe_gather_q4 / iron_moe_gather_q4_swiglu kernels not yet in thewafflehaus/iron@dev"] +#[test] +fn laguna_moe_swiglu_fused_unit() { + use wh_butter_core::{DType, Tensor}; + let Some(dev) = CudaDevice::create().expect("cuda") else { eprintln!("no CUDA, skip"); return; }; + let d = dev.as_ref(); + + let n_exp = 4usize; + let inter = 8usize; + let hid = 32usize; + let top_k = 3usize; + + // f32 -> f16 bit conversion, exactly as the Laguna GGUF loader converts + // `quantize_q4`'s host f32 per-block scales before upload (the gather + // kernels read scales as f16; uploading raw f32 bits would silently + // reinterpret them as garbage f16s). + fn half_bits(f: f32) -> u16 { + let x = f.to_bits(); + let sign = ((x >> 16) & 0x8000) as u16; + let e = ((x >> 23) & 0xff) as i32 - 112; + if e <= 0 { return sign; } + if e >= 0x1f { return sign | 0x7c00; } + let m = (x >> 13) & 0x3ff; + let round = (x >> 12) & 1; + let v = ((e as u32) << 10) | m; + sign | ((v + round) as u16) + } + let f16_bytes = |v: &[f32]| -> Vec { v.iter().flat_map(|&f| half_bits(f).to_le_bytes()).collect() }; + let f32_bytes = |v: &[f32]| -> Vec { v.iter().flat_map(|x| x.to_le_bytes()).collect() }; + let u32_bytes = |v: &[u32]| -> Vec { v.iter().flat_map(|x| x.to_le_bytes()).collect() }; + + let mut s = 0x9e3779b9u32; + let mut rng = || { s ^= s << 13; s ^= s >> 17; s ^= s << 5; (s as f32 / u32::MAX as f32) - 0.5 }; + + let gate_w: Vec = (0..n_exp * inter * hid).map(|_| rng()).collect(); + let up_w: Vec = (0..n_exp * inter * hid).map(|_| rng()).collect(); + let x: Vec = (0..hid).map(|_| rng()).collect(); + + let (gate_qs, gate_sc) = wh_butter_ops::quantize_q4(&gate_w, n_exp * inter, hid); + let (up_qs, up_sc) = wh_butter_ops::quantize_q4(&up_w, n_exp * inter, hid); + + let gate_qs_t = Tensor::new(d.upload(&u32_bytes(&gate_qs)).unwrap(), vec![gate_qs.len()], DType::U32); + let gate_sc_t = Tensor::new(d.upload(&f16_bytes(&gate_sc)).unwrap(), vec![gate_sc.len()], DType::F16); + let up_qs_t = Tensor::new(d.upload(&u32_bytes(&up_qs)).unwrap(), vec![up_qs.len()], DType::U32); + let up_sc_t = Tensor::new(d.upload(&f16_bytes(&up_sc)).unwrap(), vec![up_sc.len()], DType::F16); + let x_t = Tensor::new(d.upload(&f32_bytes(&x)).unwrap(), vec![hid], DType::F32); + + let idx: Vec = vec![2, 0, 3]; + assert_eq!(idx.len(), top_k); + let idx_t = Tensor::new(d.upload(&u32_bytes(&idx)).unwrap(), vec![top_k], DType::U32); + + // Fused: one dispatch, gate+up+silu(gate)*up in the kernel. + let fused = wh_butter_ops::moe_gather_q4_swiglu(d, &gate_qs_t, &gate_sc_t, &up_qs_t, &up_sc_t, &x_t, &idx_t, top_k, inter, hid).unwrap(); + let mut fb = vec![0u8; top_k * inter * 4]; + d.synchronize().unwrap(); + d.download(fused.buffer.as_ref(), &mut fb).unwrap(); + let fv: Vec = fb.chunks_exact(4).map(|c| f32::from_le_bytes(c.try_into().unwrap())).collect(); + + // Unfused compose: two plain gathers + host-side SwiGLU on the download. + let gates = wh_butter_ops::moe_gather_q4(d, &gate_qs_t, &gate_sc_t, &x_t, &idx_t, top_k, inter, hid).unwrap(); + let ups = wh_butter_ops::moe_gather_q4(d, &up_qs_t, &up_sc_t, &x_t, &idx_t, top_k, inter, hid).unwrap(); + let mut gb = vec![0u8; top_k * inter * 4]; + let mut ub = vec![0u8; top_k * inter * 4]; + d.synchronize().unwrap(); + d.download(gates.buffer.as_ref(), &mut gb).unwrap(); + d.download(ups.buffer.as_ref(), &mut ub).unwrap(); + let gv: Vec = gb.chunks_exact(4).map(|c| f32::from_le_bytes(c.try_into().unwrap())).collect(); + let uv: Vec = ub.chunks_exact(4).map(|c| f32::from_le_bytes(c.try_into().unwrap())).collect(); + + let mut bad = 0; + for i in 0..top_k * inter { + let silu_g = gv[i] / (1.0 + (-gv[i]).exp()); + let expect = silu_g * uv[i]; + if (fv[i] - expect).abs() > 1e-3 * expect.abs().max(1.0) { + eprintln!("swiglu-fused MISMATCH i={i}: fused={} expect={expect} (gate={} up={})", fv[i], gv[i], uv[i]); + bad += 1; + } + } + assert_eq!(bad, 0, "iron_moe_gather_q4_swiglu wrong vs unfused compose ({bad} bad elems)"); + 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 wh_butter_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 = wh_butter_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 = wh_butter_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; }; + wh_butter_modeltests::verify_laguna(dev.as_ref(), "GB10 sm_121"); +} diff --git a/rust/crates/backends/wh-butter-cuda/tests/moe_grouped_mma_test.rs b/rust/crates/backends/wh-butter-cuda/tests/moe_grouped_mma_test.rs index 38340185..e4e6e418 100644 --- a/rust/crates/backends/wh-butter-cuda/tests/moe_grouped_mma_test.rs +++ b/rust/crates/backends/wh-butter-cuda/tests/moe_grouped_mma_test.rs @@ -6,10 +6,10 @@ use wh_butter_cuda::CudaDevice; use wh_butter_core::{DType, Device, Tensor}; -use wh_butter_ops::{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 wh_butter_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}; + lt_fp4_quant_slab, lt_fp4_quant_grouped, moe_grouped_gemm_cutlass_fp4, rms_norm, gemm_cublas, cast_f32_bf16, cast_bf16_f32}; fn rng(n: usize, seed: u64) -> Vec { let mut v = vec![0f32; n]; @@ -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)]); } @@ -284,7 +284,7 @@ fn moe_grouped_mma_bench() { let Some(d) = CudaDevice::create().expect("cuda") else { eprintln!("no CUDA"); return; }; let dev = d.as_ref(); // Full S=2048 MoE up-proj: 128 experts × 96 rows, N=inter=1856, K=hid=2688. - let (n_exp, n, k, rows) = (128usize, 1856usize, 2688usize, 96usize); + let n_exp: usize = std::env::var("NEXP").ok().and_then(|v|v.parse().ok()).unwrap_or(128); let (n, k, rows) = (1856usize, 2688usize, 96usize); let mut g_starts = vec![0usize]; let mut eids = vec![]; for e in 0..n_exp { eids.push(e); g_starts.push(g_starts.last().unwrap() + rows); } let mt = *g_starts.last().unwrap(); @@ -344,3 +344,277 @@ fn lt_fp4_46_unit() { let rms=|v:&[f32]|(v.iter().map(|x|(*x as f64)*(*x as f64)).sum::()/v.len() as f64).sqrt(); eprintln!(" rms: input={:.4} base={:.4} _46={:.4}", rms(&xf), rms(&rec0), rms(&rec6)); } + +#[test] +fn moe_grouped_cutlass_bench() { + let Some(d) = CudaDevice::create().expect("cuda") else { eprintln!("no CUDA"); return; }; + let dev = d.as_ref(); + let n_exp: usize = std::env::var("NEXP").ok().and_then(|v|v.parse().ok()).unwrap_or(128); let (n, k, rows) = (1856usize, 2688usize, 96usize); + let mut g_starts = vec![0usize]; let mut eids = vec![]; + for e in 0..n_exp { eids.push(e); g_starts.push(g_starts.last().unwrap() + rows); } + let mt = *g_starts.last().unwrap(); + 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(&rng(mt*k, 1), vec![mt, k])).unwrap().reshaped(vec![mt, k]); + let w_dev = cast_f32_f16(dev, &up(&rng(n_exp*n*k, 2), vec![n_exp, n, k])).unwrap().reshaped(vec![n_exp, n, k]); + let (wp4, wsf4, gw4) = lt_fp4_quant_slab(dev, &w_dev, n_exp, n, k).unwrap(); + let (ap4, asf4, aoff4, ags4) = lt_fp4_quant_grouped(dev, &a_dev, mt, k, &g_starts).unwrap(); + let _ = moe_grouped_gemm_cutlass_fp4(dev, &ap4, &asf4, &aoff4, &ags4, &wp4, &wsf4, &gw4, &g_starts, &eids, n, k).unwrap(); + dev.synchronize().ok(); + let iters = 50; + let t = std::time::Instant::now(); + for _ in 0..iters { let _ = moe_grouped_gemm_cutlass_fp4(dev, &ap4, &asf4, &aoff4, &ags4, &wp4, &wsf4, &gw4, &g_starts, &eids, n, k).unwrap(); } + dev.synchronize().ok(); + let ms = t.elapsed().as_secs_f64() * 1000.0 / iters as f64; + let tflops = 2.0 * mt as f64 * n as f64 * k as f64 / (ms * 1e-3) / 1e12; + eprintln!("moe_grouped_CUTLASS_bench: mt={mt} N={n} K={k} -> {ms:.3} ms/call, {tflops:.1} TFLOP/s (dense-fp4 ref ~240-306)"); +} + +#[test] +fn f16_gemm_ceiling_bench() { + let Some(d) = CudaDevice::create().expect("cuda") else { eprintln!("no CUDA"); return; }; + let dev = d.as_ref(); + let up = |v: &[f32], sh: Vec| -> Tensor { + Tensor::new(dev.upload(&tb(v)).unwrap(), vec![v.len()], DType::F32).reshaped(sh) }; + for (m,n,k) in [(2048usize,2048usize,2048usize),(4096,4096,4096),(2048,10304,2688),(8192,8192,8192)] { + let x = cast_f32_f16(dev, &up(&rng(m*k, 1), vec![m,k])).unwrap().reshaped(vec![m,k]); + let w = cast_f32_f16(dev, &up(&rng(n*k, 2), vec![n,k])).unwrap().reshaped(vec![n,k]); + let _ = gemm_cublas_f32out(dev, &x, &w, m, n, k).unwrap(); + dev.synchronize().ok(); + let iters=30; let t=std::time::Instant::now(); + for _ in 0..iters { let _=gemm_cublas_f32out(dev,&x,&w,m,n,k).unwrap(); } + dev.synchronize().ok(); + let ms=t.elapsed().as_secs_f64()*1000.0/iters as f64; + let tf=2.0*m as f64*n as f64*k as f64/(ms*1e-3)/1e12; + eprintln!("f16_ceiling m={m} n={n} k={k} -> {ms:.3} ms, {tf:.1} TFLOP/s"); + } +} + +#[test] +fn fwht_rotate_invariance() { + let Some(d) = CudaDevice::create().expect("cuda") else { eprintln!("no CUDA"); return; }; + let dev = d.as_ref(); + let (m,n,k) = (64usize, 64usize, 256usize); // k multiple of 128 + let up = |v: &[f32], sh: Vec| -> Tensor { + Tensor::new(dev.upload(&tb(v)).unwrap(), vec![v.len()], DType::F32).reshaped(sh) }; + let x = cast_f32_f16(dev, &up(&rng(m*k,1), vec![m,k])).unwrap().reshaped(vec![m,k]); + let w = cast_f32_f16(dev, &up(&rng(n*k,2), vec![n,k])).unwrap().reshaped(vec![n,k]); + let y1 = dl(dev, &gemm_cublas_f32out(dev, &x, &w, m, n, k).unwrap()); + let xh = fwht128(dev, &x, m, k).unwrap(); + let wh = fwht128(dev, &w, n, k).unwrap(); + let y2 = dl(dev, &gemm_cublas_f32out(dev, &xh, &wh, m, n, k).unwrap()); + let c = cos(&y1, &y2); + eprintln!("fwht_rotate_invariance: cos(y1,y2)={c:.6} (want >0.999)"); + assert!(c > 0.999, "rotate-invariance broken cos={c}"); +} + +#[test] +fn moe_q4_grouped_bench() { + let Some(d) = CudaDevice::create().expect("cuda") else { eprintln!("no CUDA"); return; }; + let dev = d.as_ref(); + let n_exp: usize = std::env::var("NEXP").ok().and_then(|v|v.parse().ok()).unwrap_or(128); let (n, k, rows) = (1856usize, 2688usize, 96usize); + let mt = n_exp * rows; + let up = |v: &[f32], sh: Vec| -> Tensor { + Tensor::new(dev.upload(&tb(v)).unwrap(), vec![v.len()], DType::F32).reshaped(sh) }; + let a = cast_f32_f16(dev, &up(&rng(mt*k, 1), vec![mt,k])).unwrap().reshaped(vec![mt,k]); + let (qs_v, sc_v) = quantize_q4(&rng(n_exp*n*k, 2), 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(); + let offs: Vec = (0..=n_exp).flat_map(|e| ((e*rows) as u32).to_le_bytes()).collect(); + let off = Tensor::new(dev.upload(&offs).unwrap(), vec![n_exp+1], DType::U32); + let _ = moe_q4_grouped_mma_dev(dev, &a, &qs, &sc, &off, n_exp, mt, n, k).unwrap(); + dev.synchronize().ok(); + let iters = 50; + let t = std::time::Instant::now(); + for _ in 0..iters { let _ = moe_q4_grouped_mma_dev(dev, &a, &qs, &sc, &off, n_exp, mt, n, k).unwrap(); } + dev.synchronize().ok(); + let ms = t.elapsed().as_secs_f64() * 1000.0 / iters as f64; + let tf = 2.0 * mt as f64 * n as f64 * k as f64 / (ms * 1e-3) / 1e12; + 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)"); +} + +// BUTTER_LAGUNA_SCHED=1 A/B: descending-group-order tile emission + banded +// N-block/tile grid transpose (see moe_q4_grouped_mma / _dev in wh-butter-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 BUTTER_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: BUTTER_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("BUTTER_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("BUTTER_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: BUTTER_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("BUTTER_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("BUTTER_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: BUTTER_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 BUTTER_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; }; + let dev = d.as_ref(); + let n_exp: usize = std::env::var("NEXP").ok().and_then(|v|v.parse().ok()).unwrap_or(128); let (n, k, rows) = (1856usize, 2688usize, 96usize); + let mut g_starts = vec![0usize]; let mut eids = vec![]; + for e in 0..n_exp { eids.push(e); g_starts.push(g_starts.last().unwrap() + rows); } + let mt = *g_starts.last().unwrap(); + let up = |v: &[f32], sh: Vec| -> Tensor { + Tensor::new(dev.upload(&tb(v)).unwrap(), vec![v.len()], DType::F32).reshaped(sh) }; + let a = cast_f32_f16(dev, &up(&rng(mt*k, 1), vec![mt,k])).unwrap().reshaped(vec![mt,k]); + let w = cast_f32_f16(dev, &up(&rng(n_exp*n*k, 2), vec![n_exp*n, k])).unwrap().reshaped(vec![n_exp*n, k]); + let _ = moe_grouped_gemm_cutlass(dev, &a, &w, &g_starts, &eids, n, k).unwrap(); + dev.synchronize().ok(); + let iters = 50; + let t = std::time::Instant::now(); + for _ in 0..iters { let _ = moe_grouped_gemm_cutlass(dev, &a, &w, &g_starts, &eids, n, k).unwrap(); } + dev.synchronize().ok(); + let ms = t.elapsed().as_secs_f64() * 1000.0 / iters as f64; + let tf = 2.0 * mt as f64 * n as f64 * k as f64 / (ms * 1e-3) / 1e12; + eprintln!("moe_cutlass_f16_bench: mt={mt} N={n} K={k} -> {ms:.3} ms, {tf:.1} TFLOP/s (q4-grouped ref 13.7, cutlass-fp4 60)"); +} + +// Isolates the bf16-stream NaN: run rms_norm + gemm_tc in f32 vs bf16, report cosine + NaN. +#[test] +fn bf16_kernel_isolate() { + let Some(d) = CudaDevice::create().expect("cuda") else { eprintln!("no CUDA"); return; }; + let dev = d.as_ref(); + let (rows, hid) = (2048usize, 2688usize); + let up = |v: &[f32], sh: Vec| -> Tensor { + Tensor::new(dev.upload(&tb(v)).unwrap(), vec![v.len()], DType::F32).reshaped(sh) }; + let x = up(&rng(rows*hid, 7), vec![rows, hid]); + let w = up(&rng(hid, 9), vec![hid]); + let cos = |a: &[f32], b: &[f32]| -> f32 { + let (mut dot, mut na, mut nb) = (0f64,0f64,0f64); + for i in 0..a.len() { dot+=a[i] as f64*b[i] as f64; na+=(a[i] as f64).powi(2); nb+=(b[i] as f64).powi(2); } + (dot/(na.sqrt()*nb.sqrt())) as f32 }; + let has_nan = |v: &[f32]| v.iter().any(|x| x.is_nan()); + + // --- rms_norm f32 vs bf16(wide) --- + let n_f32 = dl(dev, &rms_norm(dev, &x, &w, 1e-5).unwrap()); + let xb = cast_f32_bf16(dev, &x).unwrap(); + let wb = cast_f32_bf16(dev, &w).unwrap(); + let n_bf = rms_norm(dev, &xb, &wb, 1e-5).unwrap(); + let n_bf_f = dl(dev, &cast_bf16_f32(dev, &n_bf).unwrap()); + eprintln!("RMS_NORM bf16-wide: NaN={} cosine_vs_f32={:.6}", has_nan(&n_bf_f), cos(&n_f32, &n_bf_f)); + + // --- gemm_tc f32-ref vs bf16 --- out[rows,hid] = x[rows,hid] . W[hid,hid]^T + let wsq = up(&rng(hid*hid, 11), vec![hid, hid]); + let xh = cast_f32_f16(dev, &x).unwrap(); + let wh = cast_f32_f16(dev, &wsq).unwrap(); + let g_f16 = dl(dev, &cast_f16_f32(dev, &gemm_cublas(dev, &xh, &wh, rows, hid, hid).unwrap()).unwrap()); + let xbb = cast_f32_bf16(dev, &x).unwrap(); + let wbb = cast_f32_bf16(dev, &wsq).unwrap(); + let g_bf = dl(dev, &cast_bf16_f32(dev, &gemm_cublas(dev, &xbb, &wbb, rows, hid, hid).unwrap()).unwrap()); + eprintln!("GEMM_TC bf16: NaN={} cosine_vs_f16={:.6}", has_nan(&g_bf), cos(&g_f16, &g_bf)); +} + +#[test] +fn fp8_rate_bench() { + let Some(d)=CudaDevice::create().expect("cuda") else { eprintln!("no CUDA"); return; }; + let dev=d.as_ref(); + let up=|v:&[f32],sh:Vec|->Tensor{ Tensor::new(dev.upload(&tb(v)).unwrap(),vec![v.len()],DType::F32).reshaped(sh) }; + for (m,n,k) in [(2048usize,2688usize,2688usize),(2048,8192,2688),(4096,4096,4096)] { + let x=up(&rng(m*k,1),vec![m,k]); let w=up(&rng(n*k,2),vec![n,k]); + // fp8 path + let (xq,xsc)=fp8_quant(dev,&cast_f32_f16(dev,&x).unwrap(),m*k,1.0).unwrap(); + let (wq,wsc)=fp8_quant(dev,&cast_f32_f16(dev,&w).unwrap(),n*k,1.0).unwrap(); + let _=gemm_fp8(dev,&xq,&xsc,&wq,&wsc,m,n,k,true).unwrap(); dev.synchronize().ok(); + let it=50; let t=std::time::Instant::now(); + for _ in 0..it { let _=gemm_fp8(dev,&xq,&xsc,&wq,&wsc,m,n,k,true).unwrap(); } + dev.synchronize().ok(); let ms8=t.elapsed().as_secs_f64()*1000.0/it as f64; + let tf8=2.0*m as f64*n as f64*k as f64/(ms8*1e-3)/1e12; + // f16 path + let xh=cast_f32_f16(dev,&x).unwrap(); let wh=cast_f32_f16(dev,&w).unwrap(); + let _=gemm_cublas(dev,&xh,&wh,m,n,k).unwrap(); dev.synchronize().ok(); + let t=std::time::Instant::now(); + for _ in 0..it { let _=gemm_cublas(dev,&xh,&wh,m,n,k).unwrap(); } + dev.synchronize().ok(); let ms16=t.elapsed().as_secs_f64()*1000.0/it as f64; + let tf16=2.0*m as f64*n as f64*k as f64/(ms16*1e-3)/1e12; + eprintln!("fp8_rate_bench m={m} n={n} k={k}: fp8 {ms8:.3}ms {tf8:.1} TFLOP/s | f16 {ms16:.3}ms {tf16:.1} TFLOP/s | fp8/f16 speedup {:.2}x", tf8/tf16); + } +} diff --git a/rust/crates/backends/wh-butter-cuda/tests/proj_fp8_bench.rs b/rust/crates/backends/wh-butter-cuda/tests/proj_fp8_bench.rs index 606945f0..4f97aab9 100644 --- a/rust/crates/backends/wh-butter-cuda/tests/proj_fp8_bench.rs +++ b/rust/crates/backends/wh-butter-cuda/tests/proj_fp8_bench.rs @@ -55,8 +55,8 @@ fn bench_shape(dev: &dyn Device, label: &str, m: usize, n: usize, k: usize, n_wa // ---- W8A8 per-TENSOR e4m3 (gemm_fp8) — GEMM-only timing (quant excluded) ---- // Pre-quantize once (weights fixed; per-tensor act scale recomputed per call // in the real engine, but here we time the GEMM core in isolation). - let (wq_pt, wsc_pt) = fp8_quant(dev, &wh, n * k).unwrap(); - let (xq_pt, xsc_pt) = fp8_quant(dev, &xh, m * k).unwrap(); + let (wq_pt, wsc_pt) = fp8_quant(dev, &wh, n * k, 1.0).unwrap(); + let (xq_pt, xsc_pt) = fp8_quant(dev, &xh, m * k, 1.0).unwrap(); let pt_ref = gemm_fp8(dev, &xq_pt, &xsc_pt, &wq_pt, &wsc_pt, m, n, k, true).unwrap(); for _ in 0..n_warm { let _ = gemm_fp8(dev, &xq_pt, &xsc_pt, &wq_pt, &wsc_pt, m, n, k, true).unwrap(); } dev.synchronize().ok(); diff --git a/rust/crates/wh-butter-core/src/lib.rs b/rust/crates/wh-butter-core/src/lib.rs index b906af0c..a70042f1 100644 --- a/rust/crates/wh-butter-core/src/lib.rs +++ b/rust/crates/wh-butter-core/src/lib.rs @@ -269,6 +269,34 @@ pub trait Device: Send + Sync { ) -> Result<()> { Err(Error::Msg("moe_grouped_cutlass unsupported on this backend".into())) } + /// Validated NVFP4 grouped-MoE GEMM (marlin). Device buffers; routing arrays + /// (`sorted_token_ids`/`expert_ids`/`num_tokens_past_padded`) are DEVICE i32. + #[allow(clippy::too_many_arguments)] + fn moe_marlin_gemm( + &self, + _a: &dyn DeviceBuffer, _b: &dyn DeviceBuffer, _c: &dyn DeviceBuffer, _c_tmp: &dyn DeviceBuffer, + _b_scales: &dyn DeviceBuffer, _global_scale: &dyn DeviceBuffer, _topk_weights: &dyn DeviceBuffer, + _sorted_token_ids: &dyn DeviceBuffer, _expert_ids: &dyn DeviceBuffer, _num_tokens_past_padded: &dyn DeviceBuffer, + _workspace: &dyn DeviceBuffer, + _moe_block: usize, _num_experts: usize, _top_k: usize, _mul_topk_weights: bool, + _prob_m: usize, _prob_n: usize, _prob_k: usize, _num_groups: usize, _group_size: usize, + _sms: i32, _use_fp32_reduce: bool, + ) -> Result<()> { + Err(Error::Msg("moe_marlin_gemm unsupported on this backend".into())) + } + + /// Repack one expert's GPTQ-packed `[K/8,N]` u32 weights into marlin tile layout. + #[allow(clippy::too_many_arguments)] + fn marlin_repack( + &self, _b_q: &dyn DeviceBuffer, _out: &dyn DeviceBuffer, + _size_k: usize, _size_n: usize, _sms: i32, _max_shared_mem: i32, + ) -> Result<()> { + Err(Error::Msg("marlin_repack unsupported on this backend".into())) + } + #[allow(clippy::too_many_arguments)] + fn marlin_build_routing(&self, _off:&dyn DeviceBuffer,_n_exp:usize,_blk:usize,_mt:usize,_stid:&dyn DeviceBuffer,_eid:&dyn DeviceBuffer,_ntpp:&dyn DeviceBuffer)->Result<()>{ + Err(Error::Msg("marlin_build_routing unsupported on this backend".into())) + } /// CUTLASS grouped block-scaled NVFP4 MoE GEMM: `out[t,n] = Σ_k a[t,k]· /// w[eid(t)][n,k]` over sorted-token groups, both operands packed e2m1 + diff --git a/rust/crates/wh-butter-loader/src/gguf.rs b/rust/crates/wh-butter-loader/src/gguf.rs index a9d6d7cf..b41b6ffd 100644 --- a/rust/crates/wh-butter-loader/src/gguf.rs +++ b/rust/crates/wh-butter-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 `wh_butter_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/wh-butter-models/Cargo.toml b/rust/crates/wh-butter-models/Cargo.toml index 05a1ccc9..b414b84e 100644 --- a/rust/crates/wh-butter-models/Cargo.toml +++ b/rust/crates/wh-butter-models/Cargo.toml @@ -11,3 +11,4 @@ wh-butter-core.workspace = true wh-butter-ops.workspace = true wh-butter-loader.workspace = true serde_json.workspace = true +rayon.workspace = true diff --git a/rust/crates/wh-butter-models/src/laguna.rs b/rust/crates/wh-butter-models/src/laguna.rs new file mode 100644 index 00000000..539526f1 --- /dev/null +++ b/rust/crates/wh-butter-models/src/laguna.rs @@ -0,0 +1,2979 @@ +// Copyright 2026 Eric Kryski (@ekryski) and Tom Turney (@TheTom) +// SPDX-License-Identifier: Apache-2.0 + +//! Laguna-S-2.1 (GGUF arch "laguna"): a 117.55B / 8.14B-active sigmoid-routed +//! MoE transformer with a per-head softplus attention-output gate, per-head +//! QK-RMSNorm, and hybrid full/sliding-window attention (period-4, dense- +//! first). Layer 0 is a dense SwiGLU MLP (`mlp_only_layers = [0]`); layers +//! 1..47 are MoE (256 experts, top-10, sigmoid router with a score- +//! correction bias, one always-on shared expert). Full-attention layers use +//! 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: 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 [`wh_butter_loader::gguf::Gguf`]): +//! dequantize each tensor to f32 once, then requantize into this engine's Q4 +//! format ([`wh_butter_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). + +use wh_butter_core::{DType, Device, Error, Result, Tensor}; +use wh_butter_loader::gguf::Gguf; +use wh_butter_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() +} +fn f32s_to_bytes(v: &[f32]) -> Vec { + v.iter().flat_map(|x| x.to_le_bytes()).collect() +} +/// The `iron_gemv_q4_coalesced*` kernels read per-block scales as f16 (the +/// `d_f32: Tensor` param, name notwithstanding); uploading raw f32 bits +/// there silently reinterprets them as garbage f16s. Convert on the host at +/// load, exactly like the validated `qload(.., f16=true)` path in modeltests. +fn f32s_to_f16_bytes(v: &[f32]) -> Vec { + fn half_bits(f: f32) -> u16 { + let x = f.to_bits(); + let sign = ((x >> 16) & 0x8000) as u16; + let e = ((x >> 23) & 0xff) as i32 - 112; + if e <= 0 { return sign; } + if e >= 0x1f { return sign | 0x7c00; } + let m = (x >> 13) & 0x3ff; + let round = (x >> 12) & 1; + let v = ((e as u32) << 10) | m; + sign | ((v + round) as u16) + } + 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 BUTTER_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"; // "Butter 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 +/// `BUTTER_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 { + /// `BUTTER_LAGUNA_CACHE` overrides the cache directory; otherwise it lives + /// alongside the GGUF at `.butter-cache/`. + pub fn new(g: &Gguf) -> Self { + let path = g.path(); + let dir = match std::env::var("BUTTER_LAGUNA_CACHE") { + Ok(d) if !d.is_empty() => std::path::PathBuf::from(d), + _ => std::path::PathBuf::from(format!("{path}.butter-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 = ["BUTTER_LAGUNA_MICRO", "BUTTER_LAGUNA_MOEFUSE", "BUTTER_LAGUNA_MARLIN", "BUTTER_LAGUNA_W4A8", "BUTTER_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). +struct Q4Dense { + qs: Tensor, + scales: Tensor, + m: usize, + k: usize, +} +impl Q4Dense { + // 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) + } + fn gemv_accum(&self, dev: &dyn Device, x: &Tensor, acc: &Tensor, scale: &Tensor) -> Result<()> { + ops::gemv_q4_accum(dev, &self.qs, &self.scales, x, acc, scale, self.m, self.k, self.m) + } +} + +/// 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 `BUTTER_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) + } +} + +/// BUTTER_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("BUTTER_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, +/// only the (tiny) expert-id lookup needs to be on the host. +struct Q4ExpertStack { + qs: Tensor, + scales: Tensor, + m: usize, + k: usize, + qs_words_per_expert: usize, + scale_elems_per_expert: usize, +} +impl Q4ExpertStack { + // 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; + let sw = self.scale_elems_per_expert; + let qs = Tensor { + buffer: self.qs.buffer.clone(), + offset: self.qs.offset + e * qw * 4, + shape: vec![qw], + dtype: DType::U32, + }; + let sc = Tensor { + buffer: self.scales.buffer.clone(), + offset: self.scales.offset + e * sw * 2, // f16 scales: 2 bytes each + shape: vec![sw], + dtype: DType::F16, + }; + (qs, sc) + } + fn gemv(&self, dev: &dyn Device, e: usize, x: &Tensor) -> Result { + let (qs, sc) = self.expert_view(e); + ops::gemv_q4(dev, &qs, &sc, x, self.m, self.k, self.m) + } + fn gemv_accum(&self, dev: &dyn Device, e: usize, x: &Tensor, acc: &Tensor, scale: &Tensor) -> Result<()> { + let (qs, sc) = self.expert_view(e); + ops::gemv_q4_accum(dev, &qs, &sc, x, acc, scale, self.m, self.k, self.m) + } +} + +/// BUTTER_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 +/// (`wh-butter-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, +} + +/// BUTTER_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 { + pub n_rot: u32, + pub theta: f32, + pub freq_scale: f32, + pub ext_factor: f32, + pub attn_factor: f32, + pub beta_fast: f32, + pub beta_slow: f32, + pub n_orig_ctx: u32, +} + +/// 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 +/// wiring a new checkpoint). +#[derive(Debug, Clone)] +pub struct LagunaConfig { + pub hidden: usize, + pub n_layers: usize, + pub head_dim: usize, + pub n_kv_heads: usize, + /// Q head count on a full-attention layer (`i % swa_period == 0`). + pub n_head_full: usize, + /// Q head count on a sliding-window layer. + pub n_head_sliding: usize, + /// Hybrid attention period: layer `i` is full-attention iff `i % swa_period == 0`. + pub swa_period: usize, + pub sliding_window: usize, + pub eps: f32, + pub vocab: usize, + /// Leading dense (non-MoE) block count (Laguna-S-2.1: 1, i.e. layer 0 only). + pub n_layer_dense_lead: usize, + pub n_ff_dense: usize, // layer-0 dense MLP intermediate size + pub n_expert: usize, + pub n_expert_used: usize, // top_k + pub n_ff_exp: usize, + pub n_ff_shexp: usize, + pub routed_scaling: f32, + pub rope_full: RopeParams, + pub rope_sliding: RopeParams, +} + +impl LagunaConfig { + pub fn is_full_layer(&self, i: usize) -> bool { + i % self.swa_period == 0 + } + pub fn is_dense_layer(&self, i: usize) -> bool { + i < self.n_layer_dense_lead + } + pub fn n_head(&self, i: usize) -> usize { + if self.is_full_layer(i) { self.n_head_full } else { self.n_head_sliding } + } + pub fn heads_per_group(&self, i: usize) -> u32 { + (self.n_head(i) / self.n_kv_heads) as u32 + } + pub fn attn_scale(&self) -> f32 { + 1.0 / (self.head_dim as f32).sqrt() + } + pub fn rope_for(&self, i: usize) -> RopeParams { + if self.is_full_layer(i) { self.rope_full } else { self.rope_sliding } + } + + /// The literal Laguna-S-2.1 spec values (48 layers), used as fallbacks + /// when a GGUF key is missing. + fn defaults() -> Self { + LagunaConfig { + hidden: 3072, + n_layers: 48, + head_dim: 128, + n_kv_heads: 8, + n_head_full: 48, + n_head_sliding: 72, + swa_period: 4, + sliding_window: 512, + eps: 1e-6, + vocab: 100352, + n_layer_dense_lead: 1, + n_ff_dense: 12288, + n_expert: 256, + n_expert_used: 10, + n_ff_exp: 1024, + n_ff_shexp: 1024, + routed_scaling: 2.5, + rope_full: RopeParams { + n_rot: 64, + theta: 500000.0, + freq_scale: 1.0 / 128.0, + ext_factor: 1.0, + attn_factor: 1.4852030263919618, + beta_fast: 32.0, + beta_slow: 1.0, + n_orig_ctx: 8192, + }, + rope_sliding: RopeParams { + n_rot: 128, + theta: 10000.0, + freq_scale: 1.0, + ext_factor: 0.0, + attn_factor: 1.0, + beta_fast: 32.0, + beta_slow: 1.0, + n_orig_ctx: 8192, + }, + } + } + + /// Read `laguna.*` GGUF metadata, falling back to [`Self::defaults`] for + /// 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(); + let u = |k: &str, def: usize| g.meta_u32(k).map(|v| v as usize).unwrap_or(def); + let f = |k: &str, def: f32| g.meta_f32(k).unwrap_or(def); + + let n_layers = u("laguna.block_count", d.n_layers); + let sliding_window = u("laguna.attention.sliding_window", d.sliding_window); + let n_rot_full = u("laguna.rope.dimension_count", d.rope_full.n_rot as usize) as u32; + let n_rot_sliding = u("laguna.rope.dimension_count_swa", d.rope_sliding.n_rot as usize) as u32; + + LagunaConfig { + hidden: u("laguna.embedding_length", d.hidden), + n_layers, + head_dim: u("laguna.attention.key_length", d.head_dim), + n_kv_heads: u("laguna.attention.head_count_kv", d.n_kv_heads), + // The per-layer-type Q head split isn't a single scalar GGUF key + // in this loader (see the struct doc); keep the spec constants. + n_head_full: d.n_head_full, + n_head_sliding: d.n_head_sliding, + swa_period: d.swa_period, + sliding_window, + eps: f("laguna.attention.layer_norm_rms_epsilon", d.eps), + vocab: u("laguna.vocab_size", d.vocab), + n_layer_dense_lead: u("laguna.leading_dense_block_count", d.n_layer_dense_lead), + n_ff_dense: u("laguna.feed_forward_length", d.n_ff_dense), + n_expert: u("laguna.expert_count", d.n_expert), + n_expert_used: u("laguna.expert_used_count", d.n_expert_used), + n_ff_exp: u("laguna.expert_feed_forward_length", d.n_ff_exp), + n_ff_shexp: u("laguna.expert_shared_feed_forward_length", d.n_ff_shexp), + routed_scaling: f("laguna.expert_weights_scale", d.routed_scaling), + rope_full: { + // The GGUF is authoritative over the HF config: this export is + // 256K (factor 32, yarn_attn_factor 1 + 0.1 ln 32 = 1.34657), + // NOT the HF card's 1M / factor 128. Read the stored scaling + // keys; the spec constants are only a last-resort fallback. + let factor = f("laguna.rope.scaling.factor", 1.0 / d.rope_full.freq_scale); + let attn_factor = g.meta_f32("laguna.rope.scaling.yarn_attn_factor") + .unwrap_or_else(|| 1.0 + 0.1 * factor.ln()); + RopeParams { + n_rot: n_rot_full, + theta: f("laguna.rope.freq_base", d.rope_full.theta), + freq_scale: 1.0 / factor, + ext_factor: d.rope_full.ext_factor, + attn_factor, + beta_fast: f("laguna.rope.scaling.yarn_beta_fast", d.rope_full.beta_fast), + beta_slow: f("laguna.rope.scaling.yarn_beta_slow", d.rope_full.beta_slow), + n_orig_ctx: u("laguna.rope.scaling.original_context_length", d.rope_full.n_orig_ctx as usize) as u32, + } + }, + rope_sliding: RopeParams { + n_rot: n_rot_sliding, + theta: f("laguna.rope.freq_base_swa", d.rope_sliding.theta), + freq_scale: d.rope_sliding.freq_scale, + ext_factor: d.rope_sliding.ext_factor, + attn_factor: d.rope_sliding.attn_factor, + beta_fast: d.rope_sliding.beta_fast, + beta_slow: d.rope_sliding.beta_slow, + n_orig_ctx: d.rope_sliding.n_orig_ctx, + }, + } + } +} + +/// BUTTER_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, +/// [`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 +/// [`GraphBufs`]), so load-time and decode-time never disagree. +fn laguna_micro_enabled() -> bool { + std::env::var("BUTTER_LAGUNA_MICRO").map(|v| v == "1").unwrap_or(false) +} + +/// BUTTER_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("BUTTER_LAGUNA_MOEFUSE").map(|v| v == "1").unwrap_or(false) +} + +/// BUTTER_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 (`wh-butter-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("BUTTER_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 +/// `BUTTER_LAGUNA_MICRO=1`, so a layer never carries both layouts resident at +/// 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 +/// (`quantize_q4` blocks are per-row, block size 32 along `k`). +enum AttnProj { + Split { wq: Q4Dense, wk: Q4Dense, wv: Q4Dense, wgate: Q4Dense }, + Fused { wqkvg: Q4Dense, n_q_rows: usize, n_kv_rows: usize, n_g_rows: usize }, +} + +/// Per-layer attention weights (Q4, except the two `[head_dim]` per-head +/// norms which stay f32). +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 + /// wins over the spec constants if the two ever disagree. + n_head: usize, + attn_norm: Tensor, // [hidden] f32 + proj: AttnProj, + wo: Q4Dense, // [hidden, n_head*head_dim] + q_norm: Tensor, // [head_dim] f32 + k_norm: Tensor, // [head_dim] f32 + /// BUTTER_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. +struct DenseFfn { + ffn_norm: Tensor, // [hidden] f32 + gate: Q4Dense, // [n_ff_dense, hidden] + up: Q4Dense, // [n_ff_dense, hidden] + down: Q4Dense, // [hidden, n_ff_dense] + /// BUTTER_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 +/// as single-expert (`n_expert=1`) [`Q4ExpertStack`]s at load time (built +/// ONLY under `BUTTER_LAGUNA_MICRO=1`, same one-layout-resident rule as +/// [`AttnProj::Fused`]) plus a persistent one-element `idx=[0]` device +/// 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. +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 + /// "expert index" `moe_gather_q4_swiglu` needs for a 1-expert stack. + zero_idx: Tensor, +} +enum SharedExpert { + Split { gate: Q4Dense, up: Q4Dense, down: Q4Dense }, + Fused(SharedExpertFused), +} + +/// Routed-expert gate+up weight layout. `Split` is the plain per-projection +/// layout ([`Q4ExpertStack`] each). `Fused` is built ONLY under +/// `BUTTER_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 `BUTTER_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-BUTTER_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] +} + +/// BUTTER_LAGUNA_W4A8=1 mirror of [`RoutedGateUp`], packed as +/// [`W4a8ExpertStack`]s instead of [`Q4ExpertStack`]s. Which variant is +/// built follows the SAME `BUTTER_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] +} + +/// BUTTER_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 `BUTTER_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 `BUTTER_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) + gateup: RoutedGateUp, + down: Q4ExpertStack, // per-expert [hidden, n_ff_exp] + shared_expert: SharedExpert, + /// BUTTER_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, + /// BUTTER_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 { + Dense(DenseFfn), + Moe(MoeFfn), +} + +struct LagunaLayer { + attn: AttnWeights, + ffn: LayerFfn, +} + +/// LM head resident as Q8 (the f32 head was a 1.2GB read per token, ~10% of +/// the whole decode bandwidth budget; Q8 quarters it at negligible logit +/// error). Same layout `wh_butter_ops::gemv_q8` consumes: 4 codes/u32, f32 scale +/// per 32-block. +struct Q8Head { + qs: Tensor, + scales: Tensor, + m: usize, + k: usize, +} + +/// Full Laguna-S-2.1 weights, resident on `dev`. +pub struct LagunaModel { + pub cfg: LagunaConfig, + tok_embd: Tensor, // [vocab, hidden] f32 + layers: Vec, + out_norm: Tensor, // [hidden] f32 + lm_head: Q8Head, // [vocab, hidden] q8 + /// Cached device-resident `1.0f32` scalar, reused as the shared expert's + /// `gemv_q4_accum` weight (unconditional, unweighted accumulate). + one: Tensor, +} + +/// 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 +/// over the cached K/V set is order-invariant once every slot is filled. +struct KvSlot { + k: Tensor, // [n_kv_heads, cap, head_dim] + v: Tensor, + cap: usize, +} +pub struct LagunaKvCache { + slots: Vec, +} +impl LagunaKvCache { + pub fn new(dev: &dyn Device, cfg: &LagunaConfig, max_seq: usize) -> Result { + let mut slots = Vec::with_capacity(cfg.n_layers); + for i in 0..cfg.n_layers { + let cap = if cfg.is_full_layer(i) { max_seq } else { cfg.sliding_window }; + let n = cfg.n_kv_heads * cap * cfg.head_dim; + let k = Tensor::new(dev.alloc_zeroed(n * 4)?, vec![cfg.n_kv_heads, cap, cfg.head_dim], DType::F32); + let v = Tensor::new(dev.alloc_zeroed(n * 4)?, vec![cfg.n_kv_heads, cap, cfg.head_dim], DType::F32); + slots.push(KvSlot { k, v, cap }); + } + Ok(LagunaKvCache { slots }) + } +} + +/// 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: &wh_butter_loader::gguf::GgufTensor, name: &str) -> Result<(usize, usize)> { + // GGUF dims are fastest-first: a 2-D `[out, in]` weight is stored as + // `[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))); + } + Ok((t.dims[1] as usize, t.dims[0] as usize)) +} + +/// 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")))?; + dims_out_in(t, name) +} + +// ── 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 } + +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, +} + +struct DenseFfnBytes { + ffn_norm: Vec, + gate: Q4DenseBytes, + up: Q4DenseBytes, + down: Q4DenseBytes, + gateup_marlin: Option, + down_marlin: Option, +} + +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), +} + +struct MoeFfnBytes { + ffn_norm: Vec, + router: Vec, + bias: Vec, + /// `None` under BUTTER_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, +} + +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 + }; + Ok((dense, marlin)) +} + +/// 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()?))) +} + +/// 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 BUTTER_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; + + // BUTTER_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(&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))?, + } + }; + + // BUTTER_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 { + 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) { + // BUTTER_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"); + + // BUTTER_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))?, + } + }; + + // BUTTER_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 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"); + + // BUTTER_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 { + // BUTTER_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 }) +} + +/// BUTTER_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 BUTTER_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(); + // BUTTER_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("BUTTER_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 BUTTER_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 BUTTER_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: 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()?, + }) + } + 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: 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 +/// `BUTTER_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); + + 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. BUTTER_LAGUNA_LOAD_WINDOW overrides. + let window: usize = std::env::var("BUTTER_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 }) +} + +/// One decode layer: RMSNorm → QKV proj (+ per-head gate proj on the same +/// normed input) → per-head QK-RMSNorm → RoPE → KV-cache append → SDPA → +/// per-head softplus gate → o_proj → residual → RMSNorm → FFN (dense or +/// MoE) → residual. `pos` is the token's absolute sequence position. +#[allow(clippy::too_many_arguments)] +/// BUTTER_LAGUNA_DEBUG=2: dump a tensor's stats (first-bad-op bisection inside +/// a layer). No-op cost is one env read per call site when disabled upstream. +fn dbg_stats(dev: &dyn Device, tag: &str, t: &Tensor) { + let n = t.elem_count(); + let mut b = vec![0u8; n * 4]; + if dev.synchronize().is_err() || dev.download(t.buffer.as_ref(), &mut b).is_err() { + eprintln!(" {tag}: "); + return; + } + let f: Vec = b.chunks_exact(4).map(|c| f32::from_le_bytes(c.try_into().unwrap())).collect(); + let n_nan = f.iter().filter(|v| v.is_nan()).count(); + let l2 = f.iter().map(|v| (*v as f64) * (*v as f64)).sum::().sqrt(); + let amax = f.iter().filter(|v| v.is_finite()).fold(0f32, |a, &v| a.max(v.abs())); + eprintln!(" {tag}: n={n} |x|2={l2:.3e} amax={amax:.3e} nan={n_nan} head={:?}", &f[..4.min(n)]); +} + +/// Per-step device buffers that make `decode_layer` CUDA-graph-capturable: +/// a captured graph replays the exact same recorded kernel launches every +/// token, so any value that changes per token (RoPE position, KV length) +/// must be read from a buffer whose CONTENTS the caller refreshes between +/// 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 +/// 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. +struct GraphBufs<'a> { + n_kv_full: &'a Tensor, + n_kv_sliding: &'a Tensor, +} + +#[allow(clippy::too_many_arguments)] +fn decode_layer( + dev: &dyn Device, + cfg: &LagunaConfig, + layer_idx: usize, + w: &LagunaLayer, + kv: &KvSlot, + x: &Tensor, + pos: u32, + full_posbuf: &Tensor, + sliding_posbuf: &Tensor, + one: &Tensor, + graph_bufs: Option<&GraphBufs>, + micro: bool, +) -> Result { + let deep_dbg = layer_idx == 0 + && std::env::var("BUTTER_LAGUNA_DEBUG").map(|v| v == "2").unwrap_or(false); + let hd = cfg.head_dim; + 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; + + // ── attention ──────────────────────────────────────────────────── + if deep_dbg { dbg_stats(dev, "x_in", x); } + 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. + // + // BUTTER_LAGUNA_MICRO=1 (fusion 1): one gemv_q4 on the concatenated + // [wq;wk;wv;wgate] weight instead of four separate gemvs, then four + // device-side slices (cheap 36KB/4KB/4KB/0.2KB copies) split the result + // back into q/k/v/gate_raw. `ops::slice` is used rather than zero-copy + // Tensor offset views because the gemv/gather kernels bind buffer BASE + // pointers and ignore Tensor.offset. + let (q, k, v, gate_raw) = match &w.attn.proj { + AttnProj::Fused { wqkvg, n_q_rows, n_kv_rows, n_g_rows } => { + let qkvg = wqkvg.gemv(dev, &h)?; + let q = ops::slice(dev, &qkvg, 0, *n_q_rows)?; + let k = ops::slice(dev, &qkvg, *n_q_rows, *n_kv_rows)?; + let v = ops::slice(dev, &qkvg, n_q_rows + n_kv_rows, *n_kv_rows)?; + let gate_raw = ops::slice(dev, &qkvg, n_q_rows + 2 * n_kv_rows, *n_g_rows)?; + (q, k, v, gate_raw) + } + AttnProj::Split { wq, wk, wv, wgate } => { + let q = wq.gemv(dev, &h)?; + let k = wk.gemv(dev, &h)?; + let v = wv.gemv(dev, &h)?; + let gate_raw = wgate.gemv(dev, &h)?; // [n_head] + (q, k, v, gate_raw) + } + }; + if deep_dbg { + dbg_stats(dev, "q_proj", &q); + dbg_stats(dev, "k_proj", &k); + dbg_stats(dev, "v_proj", &v); + dbg_stats(dev, "g_proj", &gate_raw); + } + + let q = q.reshaped(vec![n_head, hd]); + let k = k.reshaped(vec![n_kv_heads, hd]); + let v = v.reshaped(vec![n_kv_heads, hd]); + + // Per-head QK RMSNorm (head_dim-level), before RoPE. + let q = ops::rms_norm(dev, &q, &w.attn.q_norm, cfg.eps)?; + let k = ops::rms_norm(dev, &k, &w.attn.k_norm, cfg.eps)?; + if deep_dbg { + dbg_stats(dev, "q_norm", &q); + dbg_stats(dev, "k_norm", &k); + } + + // 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, + // 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 }; + let q = ops::rope_yarn_partial_posbuf( + dev, &q, posbuf, rope.n_rot, rope.theta, rope.freq_scale, rope.ext_factor, + rope.attn_factor, rope.beta_fast, rope.beta_slow, rope.n_orig_ctx, + )?; + let k = ops::rope_yarn_partial_posbuf( + dev, &k, posbuf, rope.n_rot, rope.theta, rope.freq_scale, rope.ext_factor, + rope.attn_factor, rope.beta_fast, rope.beta_slow, rope.n_orig_ctx, + )?; + (q, k) + } else { + let q = ops::rope_yarn_partial( + dev, &q, pos, rope.n_rot, rope.theta, rope.freq_scale, rope.ext_factor, + rope.attn_factor, rope.beta_fast, rope.beta_slow, rope.n_orig_ctx, + )?; + let k = ops::rope_yarn_partial( + dev, &k, pos, rope.n_rot, rope.theta, rope.freq_scale, rope.ext_factor, + rope.attn_factor, rope.beta_fast, rope.beta_slow, rope.n_orig_ctx, + )?; + (q, k) + }; + if deep_dbg { + dbg_stats(dev, "q_rope", &q); + dbg_stats(dev, "k_rope", &k); + } + + // 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. + 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, + // 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 }; + ops::sdpa_decode_nbuf( + dev, &q, &kv.k, &kv.v, hd, n_kv_buf, kv.cap as u32, heads_per_group, cfg.attn_scale(), + )? + } else { + let n_kv_used = if is_full { pos + 1 } else { (pos + 1).min(cfg.sliding_window as u32) }; + ops::sdpa_decode( + dev, &q, &kv.k, &kv.v, hd, n_kv_used, kv.cap as u32, heads_per_group, cfg.attn_scale(), + )? + }; // [n_head, hd] + if deep_dbg { dbg_stats(dev, "sdpa", &attn); } + + // Per-head softplus gate, before o_proj. + let gated = ops::gate_softplus_mul_perhead(dev, &attn, &gate_raw)?; + let gated_flat = gated.reshaped(vec![n_head * hd]); + if deep_dbg { dbg_stats(dev, "gated", &gated); } + // BUTTER_LAGUNA_MICRO=1 (fusion 3): copy x into a fresh accumulator (there + // 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` + // 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 { + let acc = ops::slice(dev, x, 0, cfg.hidden)?; + w.attn.wo.gemv_accum(dev, &gated_flat, &acc, one)?; + if deep_dbg { dbg_stats(dev, "x_attn_res", &acc); } + acc + } else { + let o = w.attn.wo.gemv(dev, &gated_flat)?; + let x1 = ops::add(dev, x, &o)?; + if deep_dbg { + dbg_stats(dev, "o_proj", &o); + dbg_stats(dev, "x_attn_res", &x1); + } + x1 + }; + + // ── FFN ────────────────────────────────────────────────────────── + let x2 = match &w.ffn { + LayerFfn::Dense(f) => { + let h2 = ops::rms_norm(dev, &x1, &f.ffn_norm, cfg.eps)?; + let gate = f.gate.gemv(dev, &h2)?; + let up = f.up.gemv(dev, &h2)?; + let act = ops::swiglu(dev, &gate, &up)?; + let down = f.down.gemv(dev, &act)?; + if deep_dbg { + dbg_stats(dev, "ffn_norm", &h2); + dbg_stats(dev, "ffn_act", &act); + dbg_stats(dev, "ffn_down", &down); + } + ops::add(dev, &x1, &down)? + } + LayerFfn::Moe(f) => { + let moe_dbg = (1..=2).contains(&layer_idx) + && std::env::var("BUTTER_LAGUNA_DEBUG").map(|v| v == "2").unwrap_or(false); + let h2 = ops::rms_norm(dev, &x1, &f.ffn_norm, cfg.eps)?; + let logits = ops::gemv(dev, &f.router, &h2)?; // [n_expert] + if moe_dbg { + dbg_stats(dev, &format!("moe{layer_idx}_x_in"), x); + dbg_stats(dev, &format!("moe{layer_idx}_x_attn_res"), &x1); + dbg_stats(dev, &format!("moe{layer_idx}_ffn_norm"), &h2); + dbg_stats(dev, &format!("moe{layer_idx}_router_logits"), &logits); + } + let (idx, wts) = ops::moe_router_device( + dev, &logits, &f.bias, cfg.n_expert, cfg.n_expert_used, cfg.routed_scaling, + )?; + if moe_dbg { + let mut wb = vec![0u8; cfg.n_expert_used * 4]; + let _ = dev.synchronize(); + let _ = dev.download(wts.buffer.as_ref(), &mut wb); + let w: Vec = wb.chunks_exact(4).map(|c| f32::from_le_bytes(c.try_into().unwrap())).collect(); + let mut ib = vec![0u8; cfg.n_expert_used * 4]; + let _ = dev.download(idx.buffer.as_ref(), &mut ib); + let iv: Vec = ib.chunks_exact(4).map(|c| u32::from_le_bytes(c.try_into().unwrap())).collect(); + eprintln!(" moe{layer_idx}_topk idx={iv:?} wts={w:?}"); + } + + // Fully device-side expert path: gather-GEMV the top_k gate and up + // stacks (the gather kernels index the expert slab by `idx` inside + // the kernel; a byte-offset Tensor view does NOT work here because + // dispatch binds buffer BASE pointers and drops Tensor.offset), + // SwiGLU elementwise on the [top_k*inter] pair, gather the down + // stack, then router-weighted sum into the accumulator. + // + // BUTTER_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. 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 (BUTTER_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-BUTTER_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("BUTTER_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); + ops::moe_weighted_sum(dev, &downs, &wts, &acc, cfg.hidden, cfg.n_expert_used)?; + + 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). + // + // BUTTER_LAGUNA_MICRO=1 (fusion 2): the gate+up projections were + // repacked at load time as single-expert (n_expert=1) stacks, so + // one moe_gather_q4_swiglu call (top_k=1, the persistent + // zero_idx buffer selecting "expert 0") replaces the separate + // gate gemv + up gemv + host-orchestrated swiglu. `down` is + // unchanged either way (plain gemv_accum into the routed acc). + match &f.shared_expert { + SharedExpert::Fused(se) => { + let sact = ops::moe_gather_q4_swiglu( + dev, &se.gate.qs, &se.gate.scales, &se.up.qs, &se.up.scales, + &h2, &se.zero_idx, 1, cfg.n_ff_shexp, cfg.hidden, + )?; + se.down.gemv_accum(dev, &sact, &acc, one)?; + } + SharedExpert::Split { gate, up, down } => { + let sgate = gate.gemv(dev, &h2)?; + let sup = up.gemv(dev, &h2)?; + let sact = ops::swiglu(dev, &sgate, &sup)?; + down.gemv_accum(dev, &sact, &acc, one)?; + } + } + + if moe_dbg { dbg_stats(dev, &format!("moe{layer_idx}_acc_final"), &acc); } + ops::add(dev, &x1, &acc)? + } + }; + + 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)] + +/// BUTTER_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 + // BUTTER_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; + // BUTTER_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("BUTTER_LAGUNA_PP_PROF").is_ok(); + let psync = |tag: &str| -> Result<()> { + if prof || std::env::var("BUTTER_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. BUTTER_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()) { + // BUTTER_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 \ + BUTTER_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()) + { + // BUTTER_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 wh-iron-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("BUTTER_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 wh-butter-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 + + // BUTTER_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 wh-butter-modeltests NEMOTRON_W4A8_MOE call site + // exactly, no device-side descriptor-building variant of this + // GEMM exists in wh-butter-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 { + // BUTTER_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 (BUTTER_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()) { + // BUTTER_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); + + // BUTTER_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("BUTTER_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("BUTTER_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 +/// 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`/ +/// `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` +/// are only consulted by `decode_layer` when `BUTTER_LAGUNA_GRAPH=1`. +pub struct DecodeWorkspace { + token_id: Tensor, + full_pos: Tensor, + sliding_pos: Tensor, + n_kv_full: Tensor, + n_kv_sliding: Tensor, +} + +impl DecodeWorkspace { + pub fn new(dev: &dyn Device) -> Result { + let zero = || -> Result { + Ok(Tensor::new(dev.upload(&0u32.to_le_bytes())?, vec![1], DType::U32)) + }; + Ok(DecodeWorkspace { + token_id: zero()?, + full_pos: zero()?, + sliding_pos: zero()?, + n_kv_full: zero()?, + n_kv_sliding: zero()?, + }) + } + + /// Refresh every workspace buffer for the token at `pos`. Five tiny + /// (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<()> { + // TODO(perf): 5 tiny kernel launches per step (one per scalar). Still + // far cheaper than 5 cuMemAlloc+H2D uploads, but a single fused + // "write 5 u32s" raw kernel (one launch, 5 scalar args) would shave a + // bit more host-side launch overhead off the outside-the-graph cost + // once this becomes the next bottleneck. + ops::write_u32(dev, &self.token_id, token_id)?; + ops::write_u32(dev, &self.full_pos, pos)?; + let sliding_pos = pos % cfg.sliding_window as u32; + ops::write_u32(dev, &self.sliding_pos, sliding_pos)?; + ops::write_u32(dev, &self.n_kv_full, pos + 1)?; + let n_kv_sliding = (pos + 1).min(cfg.sliding_window as u32); + ops::write_u32(dev, &self.n_kv_sliding, n_kv_sliding)?; + Ok(()) + } +} + +/// Single-token decode: embed → every layer (self-attention + cache append + +/// FFN) → final RMSNorm → lm_head → next-token logits `[vocab]`. `pos` is +/// this token's absolute sequence position (0 for the first token); the +/// caller advances it and re-calls per generated token, reusing `kv` across +/// calls so the cache persists. +/// +/// 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 +/// `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 +/// replays the exact recorded kernel launches, so every per-token input has +/// to sit behind a buffer POINTER that stays fixed across replays. +/// +/// 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 +/// the already-recorded launches via `Device::graph_launch`. +/// +/// `BUTTER_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). +pub fn decode_step( + dev: &dyn Device, + model: &LagunaModel, + kv: &LagunaKvCache, + ws: &DecodeWorkspace, + pos: u32, +) -> Result { + let cfg = &model.cfg; + let mut x = ops::gather(dev, &model.tok_embd, &ws.token_id)?.reshaped(vec![cfg.hidden]); + + let graph_mode = std::env::var("BUTTER_LAGUNA_GRAPH").is_ok(); + let graph_bufs = if graph_mode { + Some(GraphBufs { n_kv_full: &ws.n_kv_full, n_kv_sliding: &ws.n_kv_sliding }) + } else { + None + }; + // 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, + // 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(); + + // BUTTER_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 + // 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. + let debug = std::env::var("BUTTER_LAGUNA_DEBUG").map(|v| v == "1").unwrap_or(false); + for (i, layer) in model.layers.iter().enumerate() { + x = decode_layer(dev, cfg, i, layer, &kv.slots[i], &x, pos, &ws.full_pos, &ws.sliding_pos, &model.one, graph_bufs.as_ref(), micro)?; + if debug { + let mut xb = vec![0u8; cfg.hidden * 4]; + dev.synchronize()?; + dev.download(x.buffer.as_ref(), &mut xb)?; + let xf: Vec = xb.chunks_exact(4).map(|c| f32::from_le_bytes(c.try_into().unwrap())).collect(); + let n_nan = xf.iter().filter(|v| v.is_nan()).count(); + let l2 = xf.iter().map(|v| (*v as f64) * (*v as f64)).sum::().sqrt(); + let amax = xf.iter().fold(0f32, |a, &v| if v.abs() > a { v.abs() } else { a }); + eprintln!(" layer {i:2} ({}{}): |x|2={l2:.3e} amax={amax:.3e} nan={n_nan} head={:?}", + if cfg.is_full_layer(i) { "full" } else { "swa " }, + if cfg.is_dense_layer(i) { ", dense" } else { "" }, + &xf[..4]); + } + } + + let xn = ops::rms_norm(dev, &x, &model.out_norm, cfg.eps)?; + 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) +} + +/// CUDA-graph-capturing decode driver: collapses Laguna's ~1000+ per-token +/// kernel launches (48 layers x ~22 ops) into ONE `cuGraphLaunch` replay per +/// token, removing host-launch-gap overhead. Enabled by running with +/// `BUTTER_LAGUNA_GRAPH=1` (also read by `decode_step`/`decode_layer` to pick +/// the buffer-driven RoPE/SDPA kernels — see their docs). +/// +/// Protocol (mirrors the NemotronH `NEMOTRON_GRAPH` precedent in +/// `wh-butter-modeltests`): the first two `step()` calls run [`decode_step`] +/// 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 +/// 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 +/// is what every subsequent `graph_launch` writes into, and what `step()` +/// returns (a cheap `Arc` clone) every time thereafter. +/// +/// `BUTTER_LAGUNA_DEBUG` downloads mid-step and is incompatible with capture; +/// `step()` returns an error immediately if it's set, rather than silently +/// producing a broken capture. +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 + /// `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). + step_count: usize, +} + +impl LagunaDecodeCtx { + /// Eager warmup steps before the capture step (brief: "warm up N=2 full + /// steps normally"). Capture happens on call number `WARMUP_STEPS + 1`. + const WARMUP_STEPS: usize = 2; + + /// `model`/`kv`/`max_seq` are accepted (matching the `LagunaKvCache::new` + /// shape) for parity with the decode driver's usual construction + /// arguments and to leave room for future workspace sizing; today's + /// workspace is five fixed 1-element buffers, independent of `max_seq`. + pub fn new(dev: &dyn Device, model: &LagunaModel, kv: &LagunaKvCache, max_seq: usize) -> Result { + let _ = (model, kv, max_seq); + Ok(LagunaDecodeCtx { ws: DecodeWorkspace::new(dev)?, graph: None, logits: None, step_count: 0 }) + } + + /// Decode one token. Identical call shape to [`decode_step`] plus `self`; + /// internally dispatches to eager warmup, one-time capture, or graph + /// replay depending on how many steps have run. See the struct doc for + /// the full protocol. + pub fn step( + &mut self, + dev: &dyn Device, + model: &LagunaModel, + kv: &LagunaKvCache, + token_id: u32, + pos: u32, + ) -> Result { + if std::env::var("BUTTER_LAGUNA_DEBUG").is_ok() { + return Err(Error::Msg( + "LagunaDecodeCtx::step: BUTTER_LAGUNA_DEBUG is incompatible with CUDA-graph \ + capture/replay (the debug path downloads mid-step, which the CUDA driver \ + forbids during capture). Unset BUTTER_LAGUNA_DEBUG or use plain decode_step." + .into(), + )); + } + + // 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)?; + + if let Some(exec) = self.graph { + dev.graph_launch(exec)?; + return Ok(self + .logits + .clone() + .expect("LagunaDecodeCtx: graph captured but no logits tensor saved")); + } + + 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). + 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 + // `end_capture` actually computes it (against the workspace values + // written above, so it's the correct result for THIS token). + dev.begin_capture()?; + let logits = decode_step(dev, model, kv, &self.ws, pos)?; + let exec = dev.end_capture()?; + self.graph = Some(exec); + self.logits = Some(logits.clone()); + dev.graph_launch(exec)?; + Ok(logits) + } +} diff --git a/rust/crates/wh-butter-models/src/lib.rs b/rust/crates/wh-butter-models/src/lib.rs index 9a7640e4..a8e6ea3b 100644 --- a/rust/crates/wh-butter-models/src/lib.rs +++ b/rust/crates/wh-butter-models/src/lib.rs @@ -22,5 +22,6 @@ pub trait Model: Send + Sync { pub mod dsv4; pub mod gguf_tokenizer; +pub mod laguna; pub mod llama; pub mod moe; diff --git a/rust/crates/wh-butter-modeltests/src/lib.rs b/rust/crates/wh-butter-modeltests/src/lib.rs index 460808e5..ccca5304 100644 --- a/rust/crates/wh-butter-modeltests/src/lib.rs +++ b/rust/crates/wh-butter-modeltests/src/lib.rs @@ -27,11 +27,18 @@ 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(); - let g = |name: &str| -> Vec { st.tensor_f32(name).unwrap().0 }; + let g = |name: &str| -> Vec { + if std::env::var("NEMOTRON_MARLIN_CKPT").is_ok() { + if let Ok(b) = std::fs::read(format!("/home/pidtom/ckpt-f32/{name}.bin")) { + return b.chunks_exact(4).map(|c| f32::from_le_bytes([c[0],c[1],c[2],c[3]])).collect(); + } + } + st.tensor_f32(name).unwrap().0 + }; let up = |v: &[f32], sh: Vec| -> Tensor { Tensor::new(d.upload(&tb(v)).unwrap(), sh, DType::F32) }; let dl = |t: &Tensor, n: usize| -> Vec { let mut b = vec![0u8; n * 4]; d.download(t.buffer.as_ref(), &mut b).unwrap(); fb(&b) }; let conv_t = |w: &[f32], nin: usize, nout: usize| -> Vec { let mut o = vec![0.0f32; nin * nout]; for i in 0..nin { for j in 0..nout { o[j * nin + i] = w[i * nout + j]; } } o }; @@ -89,12 +96,19 @@ pub fn verify_mamba2(d: &dyn Device, plat: &str) { use wh_butter_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; let n_layers = 24usize; - let g = |name: &str| -> Vec { st.tensor_f32(name).unwrap().0 }; + let g = |name: &str| -> Vec { + if std::env::var("NEMOTRON_MARLIN_CKPT").is_ok() { + if let Ok(b) = std::fs::read(format!("/home/pidtom/ckpt-f32/{name}.bin")) { + return b.chunks_exact(4).map(|c| f32::from_le_bytes([c[0],c[1],c[2],c[3]])).collect(); + } + } + st.tensor_f32(name).unwrap().0 + }; let up = |v: &[f32]| -> Tensor { Tensor::new(d.upload(&tb(v)).unwrap(), vec![v.len()], DType::F32) }; let upm = |v: &[f32], sh: Vec| -> Tensor { Tensor::new(d.upload(&tb(v)).unwrap(), sh, DType::F32) }; let dl = |t: &Tensor, n: usize| -> Vec { let mut b = vec![0u8; n * 4]; d.download(t.buffer.as_ref(), &mut b).unwrap(); fb(&b) }; @@ -147,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); @@ -160,9 +174,735 @@ pub fn run_all(d: &dyn Device, plat: &str) { verify_olmoe(d, plat); verify_mamba2(d, plat); verify_falcon_h1(d, plat); + verify_laguna(d, plat); +} + +/// Laguna-S-2.1 decode smoke test: env-gated on `BUTTER_LAGUNA_GGUF` (path to +/// 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 +/// 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 +/// layer-0 weights. A healthy checkpoint has projection weights of O(0.01..1); +/// amax in the hundreds+ means the k-quant dequant path is mis-unpacking +/// (scale/min layout). Env-gated on BUTTER_LAGUNA_GGUF, no device needed. +pub fn verify_laguna_host_dequant() { + let Ok(path) = std::env::var("BUTTER_LAGUNA_GGUF") else { + eprintln!("BUTTER_LAGUNA_GGUF not set, skipping host dequant check"); + return; + }; + let g = match wh_butter_loader::gguf::Gguf::open(&path) { + Ok(g) => g, + Err(e) => { eprintln!("gguf open failed: {e:?}"); return; } + }; + // BUTTER_LAGUNA_DECODE_IDS="69377,13829": detokenize ids without a model + // load (fast oracle-comparison helper). + if let Ok(ids_s) = std::env::var("BUTTER_LAGUNA_DECODE_IDS") { + if let Ok(tok) = wh_butter_models::gguf_tokenizer::GgufTokenizer::from_gguf(&g) { + let ids: Vec = ids_s.split(',').filter_map(|s| s.trim().parse().ok()).collect(); + for &id in &ids { + eprintln!("token {id} -> {:?}", tok.decode(&[id])); + } + eprintln!("joined: {:?}", tok.decode(&ids)); + } + return; + } + for name in [ + "token_embd.weight", + "blk.0.attn_norm.weight", + "blk.0.attn_q.weight", + "blk.0.attn_k.weight", + "blk.0.attn_gate.weight", + "blk.0.ffn_gate.weight", + "blk.0.attn_q_norm.weight", + ] { + let Some(t) = g.tensor(name) else { eprintln!("{name}: MISSING"); continue; }; + let ty = t.ggml_type; + match g.dequant_f32(name) { + Ok(f) => { + let n = f.len(); + let amax = f.iter().fold(0f32, |a, &v| a.max(v.abs())); + let rms = (f.iter().map(|v| (*v as f64) * (*v as f64)).sum::() / n as f64).sqrt(); + let n_nan = f.iter().filter(|v| !v.is_finite()).count(); + eprintln!("{name}: dtype={ty:?} n={n} amax={amax:.4e} rms={rms:.4e} nonfinite={n_nan} head={:?}", &f[..4.min(n)]); + } + Err(e) => eprintln!("{name}: dequant failed: {e:?}"), + } + } +} + +/// Laguna-S-2.1 EXACT host f32 reference for layer 0, single decode step, +/// token id 2 (BOS) at position 0. Companion to the device instrumented run +/// (`BUTTER_LAGUNA_DEBUG=2` in `wh_butter_models::laguna::decode_layer`): every +/// sub-step below prints the same head/|x|2/amax stats so the two logs can be +/// diffed by eye to find the first place the device path goes wrong. +/// +/// Every projection is computed TWO ways: `exact` (plain f32 matvec on the +/// GGUF-dequantized weight) and `roundtrip` (that same weight pushed through +/// [`wh_butter_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 +/// the GGUF's original per-tensor dtype (see `wh_butter_models::laguna::load_dense`). +/// +/// Env-gated on BUTTER_LAGUNA_GGUF, no device/CUDA needed. +pub fn verify_laguna_host_layer0() { + let Ok(path) = std::env::var("BUTTER_LAGUNA_GGUF") else { + eprintln!("BUTTER_LAGUNA_GGUF not set, skipping host layer-0 reference"); + return; + }; + let g = match wh_butter_loader::gguf::Gguf::open(&path) { + Ok(g) => g, + Err(e) => { eprintln!("gguf open failed: {e:?}"); return; } + }; + + // Laguna-S-2.1 layer-0 spec constants (dense SwiGLU MLP, full-attention + // layer, see wh_butter_models::laguna::LagunaConfig::defaults). + let hidden = 3072usize; + let head_dim = 128usize; + let n_head = 48usize; + let n_kv = 8usize; + let heads_per_group = n_head / n_kv; // 6 + let n_ff = 12288usize; + let eps = 1e-6f32; + let token = 2usize; + + // Position-0 YaRN detail: the full-attention layer's rope has ext_factor + // != 0, so the kernel re-multiplies attn_factor by (1 + 0.1*ln(1/freq_scale)) + // (freq_scale = 1/128, so 1/freq_scale = 128). At position 0 every + // rotation angle is zero (theta_extrap = position * ... = 0), so this + // collapses to a scalar multiply by mscale over the first n_rot dims of + // every head; dims [n_rot, head_dim) pass through untouched. + let attn_factor = 1.4852030263919618f32; + // mscale = attn_factor exactly: the reference stack pre-divides on the + // host to cancel its kernel's internal re-multiply, so the net magnitude + // scale equals the stored attention_factor. Our kernel applies it directly. + let mscale = attn_factor; + let n_rot = 64usize; // full-layer rope.dimension_count + + let dq = |name: &str| -> Vec { + g.dequant_f32(name).unwrap_or_else(|e| panic!("dequant {name} failed: {e:?}")) + }; + + // f32 <-> f16 (round-to-nearest), matching the device's scale-upload path + // (`wh_butter_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(); + let sign = ((x >> 16) & 0x8000) as u16; + let e = ((x >> 23) & 0xff) as i32 - 112; + if e <= 0 { return sign; } + if e >= 0x1f { return sign | 0x7c00; } + let m = (x >> 13) & 0x3ff; + let round = (x >> 12) & 1; + let v = ((e as u32) << 10) | m; + sign | ((v + round) as u16) + } + fn f16_to_f32(bits: u16) -> f32 { + let sign = ((bits >> 15) & 1) as u32; + let exp = ((bits >> 10) & 0x1f) as u32; + let mant = (bits & 0x3ff) as u32; + let out = if exp == 0 { + if mant == 0 { sign << 31 } else { + let mut e = -1i32; + let mut m = mant; + while m & 0x400 == 0 { m <<= 1; e -= 1; } + (sign << 31) | (((e + 127 - 15) as u32) << 23) | ((m & 0x3ff) << 13) + } + } else if exp == 0x1f { + (sign << 31) | (0xff << 23) | (mant << 13) + } else { + (sign << 31) | ((exp + 127 - 15) << 23) | (mant << 13) + }; + f32::from_bits(out) + } + + // Q4 round-trip: quantize_q4 (block 32, symmetric int4 in [-7,7], scale = + // 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) = wh_butter_ops::quantize_q4(w, m, k); + let bpr = k / 32; + let mut out = vec![0f32; m * k]; + for r in 0..m { + for b in 0..bpr { + let d = f16_to_f32(f32_to_f16(scales[r * bpr + b])); + for word in 0..4 { + let packed = qs[r * bpr * 4 + b * 4 + word]; + for i in 0..8 { + let nib = ((packed >> (i * 4)) & 0xf) as i32; + let q = if nib >= 8 { nib - 16 } else { nib }; + out[r * k + b * 32 + word * 8 + i] = q as f32 * d; + } + } + } + } + out + }; + + // Row-major [m,k] @ [k] matvec, f64 accumulation (reference precision). + let matvec = |w: &[f32], m: usize, k: usize, x: &[f32]| -> Vec { + let mut out = vec![0f32; m]; + for r in 0..m { + let row = &w[r * k..(r + 1) * k]; + let mut acc = 0f64; + for i in 0..k { acc += row[i] as f64 * x[i] as f64; } + out[r] = acc as f32; + } + out + }; + + let rms_norm = |x: &[f32], w: &[f32], eps: f32| -> Vec { + let n = x.len(); + let ms = x.iter().map(|&v| (v as f64) * (v as f64)).sum::() / n as f64; + let inv = 1.0 / (ms + eps as f64).sqrt(); + (0..n).map(|i| (x[i] as f64 * inv) as f32 * w[i]).collect() + }; + + // Per-head RMSNorm: same formula, applied independently over each + // head_dim-length row (q_norm/k_norm weight is shared across heads). + let rms_norm_perhead = |x: &[f32], w: &[f32], n_heads: usize, hd: usize, eps: f32| -> Vec { + let mut out = vec![0f32; n_heads * hd]; + for hh in 0..n_heads { + let row = rms_norm(&x[hh * hd..(hh + 1) * hd], w, eps); + out[hh * hd..(hh + 1) * hd].copy_from_slice(&row); + } + out + }; + + let softplus = |x: f32| -> f32 { if x > 20.0 { x } else { (1.0 + x.exp()).ln() } }; + let silu = |x: f32| -> f32 { x / (1.0 + (-x).exp()) }; + + let stats = |label: &str, v: &[f32]| { + let l2 = v.iter().map(|&x| (x as f64) * (x as f64)).sum::().sqrt(); + let amax = v.iter().fold(0f32, |a, &x| a.max(x.abs())); + eprintln!(" {label}: head={:?} |x|2={l2:.3e} amax={amax:.3e}", &v[..4.min(v.len())]); + }; + + // Flags a sub-step whose first 4 values differ from the device dbg line + // (pasted from the instrumented CUDA run) by more than 5% relative. + let check = |label: &str, v: &[f32], device_head: &[f32]| { + let mut worst = 0f32; + for i in 0..device_head.len().min(v.len()) { + let dv = device_head[i]; + let rel = if dv.abs() > 1e-12 { (v[i] - dv).abs() / dv.abs() } else { (v[i] - dv).abs() }; + worst = worst.max(rel); + } + if worst > 0.05 { + eprintln!(" {label}: MISMATCH? max rel diff over head vs device = {worst:.3}"); + } + }; + + eprintln!("── Laguna-S-2.1 host f32 reference, layer 0, token {token}, pos 0 ──"); + + // (a) embed: dequant_f32 gives row-major [vocab, hidden] (GGUF fastest- + // dim-first == hidden fastest, matching load_gguf's load_embed_like). + let embd = dq("token_embd.weight"); + let x: Vec = embd[token * hidden..(token + 1) * hidden].to_vec(); + stats("x_in", &x); + check("x_in", &x, &[0.28873444, -0.12119293, 0.08377075, -0.12119293]); + + // (b) attn_norm + let attn_norm_w = dq("blk.0.attn_norm.weight"); + let h = rms_norm(&x, &attn_norm_w, eps); + stats("attn_norm", &h); + + // weights + let wq = dq("blk.0.attn_q.weight"); + let wk = dq("blk.0.attn_k.weight"); + let wv = dq("blk.0.attn_v.weight"); + let wo = dq("blk.0.attn_output.weight"); + let wg = dq("blk.0.attn_gate.weight"); + let q_norm_w = dq("blk.0.attn_q_norm.weight"); + let k_norm_w = dq("blk.0.attn_k_norm.weight"); + let ffn_norm_w = dq("blk.0.ffn_norm.weight"); + let w_gate_ffn = dq("blk.0.ffn_gate.weight"); + let w_up_ffn = dq("blk.0.ffn_up.weight"); + let w_down_ffn = dq("blk.0.ffn_down.weight"); + + let qdim = n_head * head_dim; // 6144 + let kvdim = n_kv * head_dim; // 1024 + + let wq_rt = q4_roundtrip(&wq, qdim, hidden); + let wk_rt = q4_roundtrip(&wk, kvdim, hidden); + let wv_rt = q4_roundtrip(&wv, kvdim, hidden); + let wg_rt = q4_roundtrip(&wg, n_head, hidden); + + // (d) q/k/v/g projections + let q_exact = matvec(&wq, qdim, hidden, &h); + let q_rt = matvec(&wq_rt, qdim, hidden, &h); + stats("q_proj exact", &q_exact); + stats("q_proj roundtrip", &q_rt); + check("q_proj", &q_rt, &[-0.0017298572, 0.041349545, -0.016936049, -0.0050829584]); + + let k_exact = matvec(&wk, kvdim, hidden, &h); + let k_rt = matvec(&wk_rt, kvdim, hidden, &h); + stats("k_proj exact", &k_exact); + stats("k_proj roundtrip", &k_rt); + + let v_exact = matvec(&wv, kvdim, hidden, &h); + let v_rt = matvec(&wv_rt, kvdim, hidden, &h); + stats("v_proj exact", &v_exact); + stats("v_proj roundtrip", &v_rt); + + let g_exact = matvec(&wg, n_head, hidden, &h); + let g_rt = matvec(&wg_rt, n_head, hidden, &h); + stats("g_proj exact", &g_exact); + stats("g_proj roundtrip", &g_rt); + + // (e) per-head QK RMSNorm + let q_norm_exact = rms_norm_perhead(&q_exact, &q_norm_w, n_head, head_dim, eps); + let q_norm_rt = rms_norm_perhead(&q_rt, &q_norm_w, n_head, head_dim, eps); + let k_norm_exact = rms_norm_perhead(&k_exact, &k_norm_w, n_kv, head_dim, eps); + let k_norm_rt = rms_norm_perhead(&k_rt, &k_norm_w, n_kv, head_dim, eps); + stats("q_norm exact", &q_norm_exact); + stats("q_norm roundtrip", &q_norm_rt); + check("q_norm", &q_norm_rt, &[-0.04173411, 1.265691, -0.5209586, -0.06744661]); + stats("k_norm exact", &k_norm_exact); + stats("k_norm roundtrip", &k_norm_rt); + + // (f) RoPE at position 0 (see mscale comment above). + let apply_rope0 = |v: &[f32], n_heads: usize| -> Vec { + let mut out = v.to_vec(); + for hh in 0..n_heads { + for i in 0..n_rot { out[hh * head_dim + i] *= mscale; } + } + out + }; + let q_rope_exact = apply_rope0(&q_norm_exact, n_head); + let q_rope_rt = apply_rope0(&q_norm_rt, n_head); + let k_rope_exact = apply_rope0(&k_norm_exact, n_kv); + let k_rope_rt = apply_rope0(&k_norm_rt, n_kv); + stats("q_rope exact", &q_rope_exact); + stats("q_rope roundtrip", &q_rope_rt); + stats("k_rope exact", &k_rope_exact); + stats("k_rope roundtrip", &k_rope_rt); + + // (g) attention, n_kv_used = 1: softmax over a single cached position is + // always 1, so every q-head's output is just its group's cached v. + let sdpa = |v_h: &[f32]| -> Vec { + let mut out = vec![0f32; n_head * head_dim]; + for qh in 0..n_head { + let kvh = qh / heads_per_group; + out[qh * head_dim..(qh + 1) * head_dim] + .copy_from_slice(&v_h[kvh * head_dim..(kvh + 1) * head_dim]); + } + out + }; + let sdpa_exact = sdpa(&v_exact); + let sdpa_rt = sdpa(&v_rt); + stats("sdpa exact", &sdpa_exact); + stats("sdpa roundtrip", &sdpa_rt); + check("sdpa", &sdpa_rt, &[-0.0042726835, -0.119657904, 0.03584555, -0.02707379]); + + // (h) per-head softplus gate, applied before o_proj. + let gate_mul = |attn: &[f32], gate: &[f32]| -> Vec { + let mut out = vec![0f32; attn.len()]; + for hh in 0..n_head { + let sp = softplus(gate[hh]); + for i in 0..head_dim { out[hh * head_dim + i] = attn[hh * head_dim + i] * sp; } + } + out + }; + let gated_exact = gate_mul(&sdpa_exact, &g_exact); + let gated_rt = gate_mul(&sdpa_rt, &g_rt); + stats("gated exact", &gated_exact); + stats("gated roundtrip", &gated_rt); + check("gated", &gated_rt, &[-0.0016357758, -0.04581044, 0.013723293, -0.010365067]); + + // (i) o_proj + let wo_rt = q4_roundtrip(&wo, hidden, qdim); + let o_exact = matvec(&wo, hidden, qdim, &gated_exact); + let o_rt = matvec(&wo_rt, hidden, qdim, &gated_rt); + stats("o_proj exact", &o_exact); + stats("o_proj roundtrip", &o_rt); + check("o_proj", &o_rt, &[-0.037135, 0.013525429, -0.012531518, -0.044919115]); + + // (j) residual + let x1_exact: Vec = (0..hidden).map(|i| x[i] + o_exact[i]).collect(); + let x1_rt: Vec = (0..hidden).map(|i| x[i] + o_rt[i]).collect(); + stats("x_attn_res exact", &x1_exact); + stats("x_attn_res roundtrip", &x1_rt); + check("x_attn_res", &x1_rt, &[0.25159943, -0.107667506, 0.07123923, -0.16611205]); + + // (k) dense SwiGLU FFN + let h2_exact = rms_norm(&x1_exact, &ffn_norm_w, eps); + let h2_rt = rms_norm(&x1_rt, &ffn_norm_w, eps); + stats("ffn_norm exact", &h2_exact); + stats("ffn_norm roundtrip", &h2_rt); + check("ffn_norm", &h2_rt, &[0.0009642678, 0.00045515582, 0.00072807475, 0.0006906503]); + + let w_gate_rt = q4_roundtrip(&w_gate_ffn, n_ff, hidden); + let w_up_rt = q4_roundtrip(&w_up_ffn, n_ff, hidden); + let w_down_rt = q4_roundtrip(&w_down_ffn, hidden, n_ff); + + let gate_exact = matvec(&w_gate_ffn, n_ff, hidden, &h2_exact); + let gate_rt = matvec(&w_gate_rt, n_ff, hidden, &h2_rt); + let up_exact = matvec(&w_up_ffn, n_ff, hidden, &h2_exact); + let up_rt = matvec(&w_up_rt, n_ff, hidden, &h2_rt); + + let act_exact: Vec = (0..n_ff).map(|i| silu(gate_exact[i]) * up_exact[i]).collect(); + let act_rt: Vec = (0..n_ff).map(|i| silu(gate_rt[i]) * up_rt[i]).collect(); + stats("ffn_act exact", &act_exact); + stats("ffn_act roundtrip", &act_rt); + check("ffn_act", &act_rt, &[5.9853085e-7, -0.00017139639, 5.882993e-9, -4.645081e-5]); + + let down_exact = matvec(&w_down_ffn, hidden, n_ff, &act_exact); + let down_rt = matvec(&w_down_rt, hidden, n_ff, &act_rt); + stats("ffn_down exact", &down_exact); + stats("ffn_down roundtrip", &down_rt); + check("ffn_down", &down_rt, &[-0.0008143601, 0.00016136278, -0.00094114855, -0.00012908175]); + + let x2_exact: Vec = (0..hidden).map(|i| x1_exact[i] + down_exact[i]).collect(); + let x2_rt: Vec = (0..hidden).map(|i| x1_rt[i] + down_rt[i]).collect(); + stats("x2 (layer0 out) exact", &x2_exact); + stats("x2 (layer0 out) roundtrip", &x2_rt); + + eprintln!("── Laguna-S-2.1 host layer-0 reference complete ──"); +} + +pub fn verify_laguna(d: &dyn Device, plat: &str) { + use wh_butter_models::gguf_tokenizer::GgufTokenizer; + use wh_butter_models::laguna; + + let Ok(path) = std::env::var("BUTTER_LAGUNA_GGUF") else { + eprintln!("BUTTER_LAGUNA_GGUF not set, skipping Laguna-S-2.1 decode smoke test"); + return; + }; + let Ok(g) = wh_butter_loader::gguf::Gguf::open(&path) else { + 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"); + return; + } + }; + + // BUTTER_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). BUTTER_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("BUTTER_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(); + // BUTTER_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("BUTTER_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 BUTTER_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). + let mut tokens = vec![bos]; + if let Ok(tok) = GgufTokenizer::from_gguf(&g) { + tokens.extend(tok.encode("Hello")); + } else { + tokens.extend_from_slice(&[1, 2]); + } + + let max_seq = tokens.len() + + std::env::var("BUTTER_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"); + return; + }; + + // BUTTER_LAGUNA_SWEEP="moe_rpt,gemv_rpt;moe_rpt,gemv_rpt;...": bench each + // warp-count combo in ONE process (the 71G load dominates a bench cycle, + // so re-capturing a fresh graph per combo amortizes it). Each combo gets + // a fresh KV cache + decode ctx; graph capture bakes the env values read + // at capture time, so set_var before ctx creation takes effect per combo. + if let Ok(sweep) = std::env::var("BUTTER_LAGUNA_SWEEP") { + let n_gen: usize = std::env::var("BUTTER_LAGUNA_GEN").ok().and_then(|v| v.parse().ok()).unwrap_or(64); + for combo in sweep.split(';').filter(|s| !s.is_empty()) { + let parts: Vec<&str> = combo.split(',').collect(); + if parts.len() != 2 { eprintln!("sweep: bad combo {combo:?}"); continue; } + // set_var is unsafe in this edition; single-threaded bench context. + unsafe { + std::env::set_var("MT_MOE_RPT", parts[0].trim()); + std::env::set_var("MT_GEMV_RPT", parts[1].trim()); + } + let Ok(kv2) = laguna::LagunaKvCache::new(d, &model.cfg, tokens.len() + n_gen + 1) else { + eprintln!("sweep: kv alloc failed"); return; + }; + let Ok(mut ctx2) = laguna::LagunaDecodeCtx::new(d, &model, &kv2, tokens.len() + n_gen + 1) else { + eprintln!("sweep: ctx failed"); return; + }; + let mut cur2: Option = None; + let mut t0 = std::time::Instant::now(); + let total2 = tokens.len() + n_gen; + let mut first_gen: Vec = Vec::new(); + let mut died = false; + for pos in 0..total2 { + let tok = if pos < tokens.len() { tokens[pos] } else { cur2.unwrap() }; + match ctx2.step(d, &model, &kv2, tok, pos as u32) { + Ok(logits) => { + let am = match wh_butter_ops::argmax_f32_device(d, &logits) { + Ok(a) => a, + Err(e) => { eprintln!("sweep argmax failed: {e:?}"); died = true; break; } + }; + let mut ib = [0u8; 4]; + let _ = d.synchronize(); + if d.download(am.buffer.as_ref(), &mut ib).is_err() { died = true; break; } + let a = u32::from_le_bytes(ib); + if pos >= tokens.len() && first_gen.len() < 6 { first_gen.push(a); } + cur2 = Some(a); + if pos + 1 == tokens.len() { t0 = std::time::Instant::now(); } + } + Err(e) => { eprintln!("sweep combo {combo}: step failed: {e:?}"); died = true; break; } + } + } + if !died { + let tps = n_gen as f64 / t0.elapsed().as_secs_f64(); + eprintln!("SWEEP moe_rpt={} gemv_rpt={}: {tps:.2} tok/s first_gen={first_gen:?}", parts[0].trim(), parts[1].trim()); + } + } + eprintln!("✅ Laguna sweep complete ({plat})."); + return; + } + + // BUTTER_LAGUNA_GRAPH=1: capture one full decode step as a CUDA graph and + // replay it per token instead of relaunching ~1000+ kernels every step + // (see `wh_butter_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, + // same kernels, same scalars, byte-identical output to before. + let graph_mode = std::env::var("BUTTER_LAGUNA_GRAPH").is_ok(); + let mode_label = if graph_mode { "graph" } else { "eager" }; + let ws = if graph_mode { + None + } else { + match laguna::DecodeWorkspace::new(d) { + Ok(w) => Some(w), + Err(e) => { + eprintln!("Laguna: decode workspace alloc failed: {e:?}, skipping"); + return; + } + } + }; + let mut ctx = if graph_mode { + match laguna::LagunaDecodeCtx::new(d, &model, &kv, max_seq) { + Ok(c) => Some(c), + Err(e) => { + eprintln!("Laguna: LagunaDecodeCtx::new failed: {e:?}, skipping"); + return; + } + } + } else { + None + }; + + // Greedy generation: feed the prompt, then feed each argmax back in. + // BUTTER_LAGUNA_GEN overrides the generated-token count (default 16). + // Prints the detokenized continuation for direct comparison against the + // reference engine's temp-0 output on the same prompt, plus decode t/s + // over the generation phase (each step is one full 48-layer forward). + let n_gen: usize = std::env::var("BUTTER_LAGUNA_GEN").ok().and_then(|v| v.parse().ok()).unwrap_or(16); + let tok_dec = GgufTokenizer::from_gguf(&g).ok(); + let mut generated: Vec = Vec::new(); + let mut cur: Option = None; + let mut gen_t0 = std::time::Instant::now(); + let total = tokens.len() + n_gen; + // BUTTER_LAGUNA_FAST=1: device-side argmax, download 4 bytes/step instead + // 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("BUTTER_LAGUNA_FAST").map(|v| v == "1").unwrap_or(false); + + // BUTTER_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()`). + // BUTTER_LAGUNA_PREFILL_CHUNK overrides the chunk size (default 1024). + let prefill_mode = std::env::var("BUTTER_LAGUNA_PREFILL").map(|v| v == "1").unwrap_or(false); + let prefill_chunk: usize = std::env::var("BUTTER_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 = wh_butter_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: + // the argmax/download below still happen outside either path, + // exactly as before. + let step_result = if let Some(ctx) = ctx.as_mut() { + ctx.step(d, &model, &kv, tok, pos as u32) + } else { + let w = ws.as_ref().expect("eager path always has a workspace"); + if let Err(e) = w.update(d, &model.cfg, tok, pos as u32) { + eprintln!("Laguna step {pos}: workspace update failed: {e:?}"); + return; + } + laguna::decode_step(d, &model, &kv, w, pos as u32) + }; + match step_result { + Ok(logits) => { + if fast { + let Ok(am) = wh_butter_ops::argmax_f32_device(d, &logits) else { + eprintln!("Laguna step {pos}: device argmax failed"); + return; + }; + let mut ib = [0u8; 4]; + let _ = d.synchronize(); + if d.download(am.buffer.as_ref(), &mut ib).is_err() { + eprintln!("Laguna step {pos}: argmax download failed"); + return; + } + let argmax = u32::from_le_bytes(ib); + if pos >= tokens.len() { + generated.push(argmax); + } + cur = Some(argmax); + if pos + 1 == tokens.len() { + gen_t0 = std::time::Instant::now(); + } + if pos + 1 == total { + let dt = gen_t0.elapsed().as_secs_f64(); + let tps = n_gen as f64 / dt; + let text = tok_dec.as_ref().map(|t| t.decode(&generated)).unwrap_or_default(); + eprintln!("Laguna greedy continuation ({n_gen} tokens, {tps:.2} tok/s device-argmax, {mode_label}): {text:?}"); + } + continue; + } + 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 step {pos} (token {tok}): logits download failed"); + return; + } + let lf: Vec = lb.chunks_exact(4).map(|c| f32::from_le_bytes(c.try_into().unwrap())).collect(); + let argmax = wh_butter_runtime::argmax(&lf); + // Logit sanity: NaN/inf counts, range, and top-5 ids. An argmax + // parked at the last vocab id is the classic all-NaN signature. + let n_nan = lf.iter().filter(|v| v.is_nan()).count(); + let n_inf = lf.iter().filter(|v| v.is_infinite()).count(); + let (mut mn, mut mx) = (f32::INFINITY, f32::NEG_INFINITY); + for &v in &lf { + if v.is_finite() { + if v < mn { mn = v; } + if v > mx { mx = v; } + } + } + let mut order: Vec = (0..lf.len()).collect(); + order.sort_by(|&a, &b| lf[b].total_cmp(&lf[a])); + let top5: Vec<(usize, f32)> = order.iter().take(5).map(|&i| (i, lf[i])).collect(); + eprintln!("Laguna-S-2.1 decode step {pos} (token {tok}) on {plat}: argmax = {argmax}"); + eprintln!(" logits: nan={n_nan} inf={n_inf} finite range [{mn:.4}, {mx:.4}] top5={top5:?}"); + if pos >= tokens.len() { + generated.push(argmax as u32); + } + cur = Some(argmax as u32); + if pos + 1 == tokens.len() { + gen_t0 = std::time::Instant::now(); // time the pure-generation phase + } + if pos + 1 == total { + let dt = gen_t0.elapsed().as_secs_f64(); + let tps = n_gen as f64 / dt; + let text = tok_dec.as_ref().map(|t| t.decode(&generated)).unwrap_or_default(); + eprintln!("Laguna greedy continuation ({n_gen} tokens, {tps:.2} tok/s incl logit download, {mode_label}): {text:?}"); + } + } + Err(e) => { + eprintln!("Laguna step {pos} (token {tok}) failed: {e:?}"); + return; + } + } + } + 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); @@ -176,10 +916,17 @@ 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 { st.tensor_f32(name).unwrap().0 }; + let g = |name: &str| -> Vec { + if std::env::var("NEMOTRON_MARLIN_CKPT").is_ok() { + if let Ok(b) = std::fs::read(format!("/home/pidtom/ckpt-f32/{name}.bin")) { + return b.chunks_exact(4).map(|c| f32::from_le_bytes([c[0],c[1],c[2],c[3]])).collect(); + } + } + st.tensor_f32(name).unwrap().0 + }; let up = |v: &[f32], sh: Vec| -> Tensor { Tensor::new(d.upload(&tb(v)).unwrap(), sh, DType::F32) }; let dl = |t: &Tensor, n: usize| -> Vec { let mut b = vec![0u8; n * 4]; d.download(t.buffer.as_ref(), &mut b).unwrap(); fb(&b) }; let token = 5usize; @@ -211,10 +958,17 @@ pub fn verify_pythia(d: &dyn Device, plat: &str) { pub fn verify_olmo2(d: &dyn Device, plat: &str) { use wh_butter_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 { st.tensor_f32(name).unwrap().0 }; + let g = |name: &str| -> Vec { + if std::env::var("NEMOTRON_MARLIN_CKPT").is_ok() { + if let Ok(b) = std::fs::read(format!("/home/pidtom/ckpt-f32/{name}.bin")) { + return b.chunks_exact(4).map(|c| f32::from_le_bytes([c[0],c[1],c[2],c[3]])).collect(); + } + } + st.tensor_f32(name).unwrap().0 + }; let up = |v: &[f32], sh: Vec| -> Tensor { Tensor::new(d.upload(&tb(v)).unwrap(), sh, DType::F32) }; let dl = |t: &Tensor, n: usize| -> Vec { let mut b = vec![0u8; n*4]; d.download(t.buffer.as_ref(), &mut b).unwrap(); fb(&b) }; let token = 9707usize; @@ -253,10 +1007,17 @@ 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 { st.tensor_f32(name).unwrap().0 }; + let g = |name: &str| -> Vec { + if std::env::var("NEMOTRON_MARLIN_CKPT").is_ok() { + if let Ok(b) = std::fs::read(format!("/home/pidtom/ckpt-f32/{name}.bin")) { + return b.chunks_exact(4).map(|c| f32::from_le_bytes([c[0],c[1],c[2],c[3]])).collect(); + } + } + st.tensor_f32(name).unwrap().0 + }; let up = |v: &[f32], sh: Vec| -> Tensor { Tensor::new(d.upload(&tb(v)).unwrap(), sh, DType::F32) }; let dl = |t: &Tensor, n: usize| -> Vec { let mut b = vec![0u8; n*4]; d.download(t.buffer.as_ref(), &mut b).unwrap(); fb(&b) }; let token = 5usize; @@ -290,10 +1051,17 @@ pub fn verify_gptneo(d: &dyn Device, plat: &str) { pub fn verify_gemma2(d: &dyn Device, plat: &str) { use wh_butter_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 { st.tensor_f32(name).unwrap().0 }; + let g = |name: &str| -> Vec { + if std::env::var("NEMOTRON_MARLIN_CKPT").is_ok() { + if let Ok(b) = std::fs::read(format!("/home/pidtom/ckpt-f32/{name}.bin")) { + return b.chunks_exact(4).map(|c| f32::from_le_bytes([c[0],c[1],c[2],c[3]])).collect(); + } + } + st.tensor_f32(name).unwrap().0 + }; let g1 = |name: &str| -> Vec { st.tensor_f32(name).unwrap().0.iter().map(|w| w+1.0).collect() }; let up = |v: &[f32], sh: Vec| -> Tensor { Tensor::new(d.upload(&tb(v)).unwrap(), sh, DType::F32) }; let dl = |t: &Tensor, n: usize| -> Vec { let mut b = vec![0u8; n*4]; d.download(t.buffer.as_ref(), &mut b).unwrap(); fb(&b) }; @@ -330,10 +1098,17 @@ 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 { st.tensor_f32(name).unwrap().0 }; + let g = |name: &str| -> Vec { + if std::env::var("NEMOTRON_MARLIN_CKPT").is_ok() { + if let Ok(b) = std::fs::read(format!("/home/pidtom/ckpt-f32/{name}.bin")) { + return b.chunks_exact(4).map(|c| f32::from_le_bytes([c[0],c[1],c[2],c[3]])).collect(); + } + } + st.tensor_f32(name).unwrap().0 + }; let up = |v: &[f32], sh: Vec| -> Tensor { Tensor::new(d.upload(&tb(v)).unwrap(), sh, DType::F32) }; let dl = |t: &Tensor, n: usize| -> Vec { let mut b = vec![0u8; n*4]; d.download(t.buffer.as_ref(), &mut b).unwrap(); fb(&b) }; let proj = |w: &str, b: &str, x: &Tensor, m: usize, inn: usize| -> Tensor { @@ -367,10 +1142,17 @@ pub fn verify_phi(d: &dyn Device, plat: &str) { pub fn verify_stablelm2(d: &dyn Device, plat: &str) { use wh_butter_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 { st.tensor_f32(name).unwrap().0 }; + let g = |name: &str| -> Vec { + if std::env::var("NEMOTRON_MARLIN_CKPT").is_ok() { + if let Ok(b) = std::fs::read(format!("/home/pidtom/ckpt-f32/{name}.bin")) { + return b.chunks_exact(4).map(|c| f32::from_le_bytes([c[0],c[1],c[2],c[3]])).collect(); + } + } + st.tensor_f32(name).unwrap().0 + }; let up = |v: &[f32], sh: Vec| -> Tensor { Tensor::new(d.upload(&tb(v)).unwrap(), sh, DType::F32) }; let dl = |t: &Tensor, n: usize| -> Vec { let mut b = vec![0u8; n*4]; d.download(t.buffer.as_ref(), &mut b).unwrap(); fb(&b) }; let bproj = |w: &str, b: &str, x: &Tensor, m: usize| -> Tensor { add(d, &gemv(d, &up(&g(w), vec![m, hid]), x).unwrap(), &up(&g(b), vec![m])).unwrap() }; @@ -406,10 +1188,17 @@ pub fn verify_stablelm2(d: &dyn Device, plat: &str) { pub fn verify_olmoe(d: &dyn Device, plat: &str) { use wh_butter_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 { st.tensor_f32(name).unwrap().0 }; + let g = |name: &str| -> Vec { + if std::env::var("NEMOTRON_MARLIN_CKPT").is_ok() { + if let Ok(b) = std::fs::read(format!("/home/pidtom/ckpt-f32/{name}.bin")) { + return b.chunks_exact(4).map(|c| f32::from_le_bytes([c[0],c[1],c[2],c[3]])).collect(); + } + } + st.tensor_f32(name).unwrap().0 + }; let up = |v: &[f32]| -> Tensor { Tensor::new(d.upload(&tb(v)).unwrap(), vec![v.len()], DType::F32) }; let upm = |v: &[f32], sh: Vec| -> Tensor { Tensor::new(d.upload(&tb(v)).unwrap(), sh, DType::F32) }; let dl = |t: &Tensor, n: usize| -> Vec { let mut b = vec![0u8; n*4]; d.download(t.buffer.as_ref(), &mut b).unwrap(); fb(&b) }; @@ -454,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 wh_butter_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; @@ -466,7 +1255,14 @@ pub fn verify_falcon_h1(d: &dyn Device, plat: &str) { let (ssm_in, ssm_out, attn_out) = (1.25f32, 0.23570226039551587f32, 0.9375f32); let (gate_mult, down_mult, embed_mult) = (0.8838834764831844f32, 0.5859375f32, 5.656854249492381f32); let ssm_m = [0.3535533905932738f32, 0.25, 0.3535533905932738, 0.5, 0.3535533905932738]; - let g = |name: &str| -> Vec { st.tensor_f32(name).unwrap().0 }; + let g = |name: &str| -> Vec { + if std::env::var("NEMOTRON_MARLIN_CKPT").is_ok() { + if let Ok(b) = std::fs::read(format!("/home/pidtom/ckpt-f32/{name}.bin")) { + return b.chunks_exact(4).map(|c| f32::from_le_bytes([c[0],c[1],c[2],c[3]])).collect(); + } + } + st.tensor_f32(name).unwrap().0 + }; let up = |v: &[f32], sh: Vec| -> Tensor { Tensor::new(d.upload(&tb(v)).unwrap(), sh, DType::F32) }; let dl = |t: &Tensor, n: usize| -> Vec { let mut b = vec![0u8; n*4]; d.download(t.buffer.as_ref(), &mut b).unwrap(); fb(&b) }; let softplus = |x: f32| if x > 20.0 { x } else { (1.0 + x.exp()).ln() }; @@ -534,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); @@ -678,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 wh_butter_ops::{add, cast_f16_f32, cast_f32_f16, conv1d_causal_prefill, conv1d_causal_step, dequant_q4, dequant_q4_off, gather, gated_group_rmsnorm, gated_group_rmsnorm_batched, 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, 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 wh_butter_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 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; @@ -716,7 +1512,14 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { sign | ((v + round) as u16) } let tb_f16 = |v: &[f32]| -> Vec { v.iter().flat_map(|&f| f32_to_f16(f).to_le_bytes()).collect() }; - let g = |name: &str| -> Vec { st.tensor_f32(name).unwrap().0 }; + let g = |name: &str| -> Vec { + if std::env::var("NEMOTRON_MARLIN_CKPT").is_ok() { + if let Ok(b) = std::fs::read(format!("/home/pidtom/ckpt-f32/{name}.bin")) { + return b.chunks_exact(4).map(|c| f32::from_le_bytes([c[0],c[1],c[2],c[3]])).collect(); + } + } + st.tensor_f32(name).unwrap().0 + }; let up = |v: &[f32]| -> Tensor { Tensor::new(d.upload(&tb(v)).unwrap(), vec![v.len()], DType::F32) }; let upm = |v: &[f32], sh: Vec| -> Tensor { Tensor::new(d.upload(&tb(v)).unwrap(), sh, DType::F32) }; let dl = |t: &Tensor, n: usize| -> Vec { let owned; let src = if t.dtype == DType::F16 { owned = cast_f16_f32(d, t).unwrap(); &owned } else { t }; let mut b = vec![0u8; n * 4]; d.download(src.buffer.as_ref(), &mut b).unwrap(); fb(&b) }; @@ -777,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); @@ -798,12 +1601,14 @@ 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(); let mut w8a8w: HashMap = HashMap::new(); let use_w8a8_env = std::env::var("NEMOTRON_W8A8_MOE").is_ok(); + let use_marlin_moe = std::env::var("NEMOTRON_MARLIN_MOE").is_ok(); + let mut marlinw: std::collections::HashMap = std::collections::HashMap::new(); // Persistent CUTLASS grouped-FP4 handles per layer: (up_handle, dn_handle, // alpha_up, alpha_dn). Prepared lazily on first use; alpha buffers are // fixed addresses the prepared blob points at (values rewritten per call). @@ -818,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")) @@ -863,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. @@ -960,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); @@ -1010,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 { @@ -1084,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"); @@ -1120,6 +1925,80 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { buildw8("up_proj", inter, hid); buildw8("down_proj", hid, inter); } + // NVFP4 marlin MoE weights (NEMOTRON_MARLIN_MOE). RTN via + // quantize_nvfp4, OR (NEMOTRON_MARLIN_CKPT) calibrated codes + // from NVIDIA's NVFP4 checkpoint blobs. repack + S0E5M3 + global. + if use_marlin_moe { + let sms_i = 48i32; let max_sh = 232448i32; // GB10 sm_121a + let ckpt = std::env::var("NEMOTRON_MARLIN_CKPT").is_ok(); + let mut buildm = |which: &str, n: usize, k: usize| { + let ng = k / 16; let kp = k / 8; + let key = format!("{p}.moe_{which}_marlin.{}", if ckpt { "ckpt" } else { "v1" }); + let bm_bytes = n_exp * (k / 16) * (n * 2) * 4; + let bs_bytes = n_exp * ng * n; + let gl_bytes = n_exp * 4; + let cpath = cache_path(&key); + if let Ok(b) = std::fs::read(&cpath) { + if b.len() == bm_bytes + bs_bytes + gl_bytes { + let bt = Tensor::new(d.upload(&b[..bm_bytes]).unwrap(), vec![bm_bytes / 4], DType::U32); + let bst = Tensor::new(d.upload(&b[bm_bytes..bm_bytes + bs_bytes]).unwrap(), vec![bs_bytes / 4], DType::U32); + let glt = Tensor::new(d.upload(&b[bm_bytes + bs_bytes..]).unwrap(), vec![n_exp], DType::F32); + marlinw.insert(key, (bt, bst, glt)); + return; + } + } + // source calibrated codes/scales from the checkpoint blob, if enabled + let (ck_codes, ck_sv, ck_gl): (Vec, Vec, Vec) = if ckpt { + let lnum: usize = p.rsplit('.').next().unwrap().parse().unwrap(); + let blob = std::fs::read(format!("/home/pidtom/marlin-ckpt-blobs/L{lnum}_{which}.bin")) + .unwrap_or_else(|_| panic!("missing ckpt blob L{lnum}_{which}")); + let nc = n_exp * n * kp; let nsv = n_exp * n * ng; + let codes: Vec = blob[..nc * 4].chunks_exact(4).map(|c| u32::from_le_bytes([c[0], c[1], c[2], c[3]])).collect(); + let sv: Vec = blob[nc * 4..(nc + nsv) * 4].chunks_exact(4).map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])).collect(); + let gl: Vec = blob[(nc + nsv) * 4..].chunks_exact(4).map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])).collect(); + (codes, sv, gl) + } else { (Vec::new(), Vec::new(), Vec::new()) }; + let mut bm_host: Vec = Vec::with_capacity(n_exp * (k / 16) * (n * 2)); + let mut bs_host: Vec = vec![0u8; n_exp * ng * n]; + let mut gl: Vec = vec![0f32; n_exp]; + for e in 0..n_exp { + let (codes_e, sv_e, global_e): (Vec, Vec, f32) = if ckpt { + (ck_codes[e * n * kp..(e + 1) * n * kp].to_vec(), + ck_sv[e * n * ng..(e + 1) * n * ng].to_vec(), + ck_gl[e]) + } else { + let w = g(&format!("{p}.mixer.experts.{e}.{which}.weight")); + let (codes, uesc, global) = wh_butter_ops::quantize_nvfp4(&w, n, k); + let sv: Vec = (0..n * ng).map(|i| wh_butter_ops::dec_ue4m3(uesc[i])).collect(); + (codes, sv, global) + }; + let mut qt = vec![0u32; kp * n]; + for r in 0..n { for c in 0..kp { qt[c * n + r] = codes_e[r * kp + c]; } } + let qtt = Tensor::new(d.upload(&tbu(&qt)).unwrap(), vec![kp * n], DType::U32); + let bm = wh_butter_ops::marlin_repack(d, &qtt, k, n, sms_i, max_sh).unwrap(); + let mut bmh = vec![0u8; (k / 16) * (n * 2) * 4]; + d.download(bm.buffer.as_ref(), &mut bmh).unwrap(); + bm_host.extend(bmh.chunks_exact(4).map(|c| u32::from_le_bytes([c[0], c[1], c[2], c[3]]))); + let mut svm = vec![0f32; ng * n]; + for r in 0..n { for gi in 0..ng { svm[gi * n + r] = sv_e[r * ng + gi]; } } + let s0 = wh_butter_ops::marlin_s0e5m3_scales(&svm, ng, n); + bs_host[e * ng * n..(e + 1) * ng * n].copy_from_slice(&s0); + gl[e] = global_e * 128.0; + } + let mut blob2: Vec = Vec::with_capacity(bm_bytes + bs_bytes + gl_bytes); + blob2.extend(tbu(&bm_host)); + blob2.extend(&bs_host); + blob2.extend(gl.iter().flat_map(|x| x.to_le_bytes())); + let _ = std::fs::write(&cpath, &blob2); + let bt = Tensor::new(d.upload(&tbu(&bm_host)).unwrap(), vec![bm_host.len()], DType::U32); + let bst = Tensor::new(d.upload(&bs_host).unwrap(), vec![bs_host.len() / 4], DType::U32); + let glb: Vec = gl.iter().flat_map(|x| x.to_le_bytes()).collect(); + let glt = Tensor::new(d.upload(&glb).unwrap(), vec![n_exp], DType::F32); + marlinw.insert(key, (bt, bst, glt)); + }; + buildm("up_proj", inter, hid); + buildm("down_proj", hid, inter); + } // NEMOTRON_FP4_SHARED_FOLD=1: shared expert folded into the routed // grouped GEMM as 2 extra always-selected groups. Up: shared up // [2*inter, hid] appended as two [inter, hid] expert slabs. Down: @@ -1251,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() @@ -1269,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 ── @@ -1295,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 @@ -1312,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 @@ -1320,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(); @@ -1387,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}"); @@ -1630,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, @@ -1653,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 @@ -1672,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(); @@ -1680,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()); @@ -1689,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 @@ -1713,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()); @@ -1724,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()); @@ -1747,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 { @@ -1764,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() @@ -1799,6 +2678,23 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { // stack-local host Vec, so a CUDA-graph capture recorded a memcpy node // from a pointer that's dangling by replay time. Cached, the pre-warm // forwards fill it and the captured forward does NO embed H2D at all. + let bf16_stream = is_cuda && std::env::var("NEMOTRON_BF16_STREAM").is_ok(); + let f16_stream = is_cuda && std::env::var("NEMOTRON_F16_STREAM").is_ok(); + let stream16 = bf16_stream || f16_stream; + // NEMOTRON_TOPK_SCATTER=1: fixed-top_k deterministic scatter reuses + // router rowpos/weights and skips the per-token CSR rebuild. + let topk_scatter = is_cuda && std::env::var("NEMOTRON_TOPK_SCATTER").is_ok(); + let stream_dt = if bf16_stream { DType::BF16 } else if f16_stream { DType::F16 } else { DType::F32 }; + let cast_stream = |d: &dyn Device, t: &wh_butter_core::Tensor| -> wh_butter_core::Tensor { + if bf16_stream { wh_butter_ops::cast_f32_bf16(d, t).unwrap() } else { wh_butter_ops::cast_f32_f16(d, t).unwrap() } + }; + // 16-bit stream -> f32 (for ops without a 16-bit kernel: conv/SSD/gated-norm/attn). + let to_f32 = |d: &dyn Device, t: &wh_butter_core::Tensor| -> wh_butter_core::Tensor { + match t.dtype { DType::F16 => wh_butter_ops::cast_f16_f32(d, t).unwrap(), + DType::BF16 => wh_butter_ops::cast_bf16_f32(d, t).unwrap(), + _ => t.clone() } + }; + let act_dt = stream_dt; let _ = (act_dt, stream16); let mut xt = { use std::hash::{Hash, Hasher}; let mut h = std::collections::hash_map::DefaultHasher::new(); @@ -1810,6 +2706,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { upm(&xh, vec![s, hid]) }) }; + if stream16 { xt = cast_stream(d, &xt); } // Tensor-core projection GEMM (Path A): dequant Q4 weight → f16 once, // cast x → f16, cuBLAS GemmEx (real tensor cores, f32 accumulate), // cast result → f32. Replaces the software-emulated coop_tile MMA @@ -1817,6 +2714,9 @@ 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 (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 // top prefill cost at small S. On Metal we dequant→f16 and run the per-expert @@ -1827,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 @@ -1837,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()); @@ -1850,6 +2750,14 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { let deq16 = |key: &str, qs: &Tensor, sc: &Tensor, m: usize, k: usize, blk_off: usize| -> Tensor { deq16c(key, qs, sc, m, k, blk_off, true) }; + // BF16 weight dequant (cached) for the BF16-coherent FFN under stream16. + let deq16bf = |key: &str, qs: &Tensor, sc: &Tensor, m: usize, k: usize| -> Tensor { + let ck = format!("{key}.bf16w"); + if let Some(w) = w16.borrow().get(&ck) { return w.clone(); } + let w = pf!(3, (m * k) as f64, dequant_q4_off(d, qs, sc, m, k, DType::BF16, 0).unwrap()); + w16.borrow_mut().insert(ck, w.clone()); + w + }; // FP4 dense projections (NEMOTRON_FP4_DENSE=1): quantize the f16 // activation per call, weight pack cached cross-forward; cublasLt // block-scaled NVFP4 GEMM with f32 output (fused output cast). @@ -1863,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, @@ -1914,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(); @@ -1928,8 +2836,38 @@ 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), + // BUTTER 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 + // 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; + let wkey = if had { format!("{name}.fp8pth") } else { format!("{name}.fp8pt") }; + let (wq, wsc) = { + if let Some(t) = w4lt.borrow().get(&wkey) { t.clone() } else { + let wf = dequant_q4_off(d, qs, sc, m, k, DType::F16, 0).unwrap(); + let wf = if had { wh_butter_ops::fwht128(d, &wf, m, k).unwrap() } else { wf }; + let t = wh_butter_ops::fp8_quant(d, &wf, m*k, 1.0).unwrap(); + w4lt.borrow_mut().insert(wkey.clone(), t.clone()); + t + } + }; + let clipv = std::env::var("NEMOTRON_FP8_CLIP").ok().and_then(|v| v.parse().ok()).unwrap_or(1.0f32); + let (xq, xsc) = if had { + // Fused: amax-only rotate (no y write) + rotate-and-quant reading + // xh directly. Kills the [rows,k] f16 rotated-tensor round-trip. + let amx = pf!(15, 0.0, wh_butter_ops::fwht128_amaxonly(d, xh, rows, k).unwrap()); + pf!(15, 0.0, wh_butter_ops::fwht128_quant_pre(d, xh, rows, k, &amx, clipv).unwrap()) + } else { + pf!(15, 0.0, wh_butter_ops::fp8_quant(d, xh, rows*k, clipv).unwrap()) + }; + return pf!(2, 2.0*rows as f64*m as f64*k as f64, + wh_butter_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 { @@ -1947,6 +2885,14 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { let (qs, sc, m, k) = &qw[name]; let rows = x.elem_count() / *k; let flops = 2.0 * rows as f64 * *m as f64 * *k as f64; + // BF16-STREAM: dequant Q4->bf16 (cached) + native bf16 GEMM, no f16 cast. + if stream16 && rows > 1 { + let wkey = format!("{name}.s16w"); + let wb = if let Some(w) = w16.borrow().get(&wkey) { w.clone() } + else { let w = pf!(3, (*m * *k) as f64, dequant_q4_off(d, qs, sc, *m, *k, stream_dt, 0).unwrap()); w16.borrow_mut().insert(wkey.clone(), w.clone()); w }; + let xb = if x.dtype == stream_dt { x.clone() } else { pf!(15, 0.0, cast_stream(d, x)) }; + return pf!(2, flops, gemm_cublas(d, &xb, &wb, rows, *m, *k).unwrap()); + } if fp8_dense && rows > 1 && fp8_name_ok(name) { let xh = pf!(15, 0.0, cast_f32_f16(d, x).unwrap()); return qmm_fp8d(&xh, name, rows, *m, *k, qs, sc); @@ -1960,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] @@ -1969,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]; @@ -1994,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 @@ -2006,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(); @@ -2019,11 +2965,21 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { // xn carried from the previous layer's fused add+rms_norm (None on layer 0 // or after a non-fusable tail). let mut xn_carry: Option = None; + let stream_dbg = std::env::var("NEMOTRON_STREAM_DBG").is_ok(); for (l, mix) in PATTERN.chars().enumerate() { let p = format!("language_model.backbone.layers.{l}"); + if stream_dbg { + let xf32 = if xt.dtype == DType::F32 { xt.clone() } else { to_f32(d, &xt) }; + let mut hb = vec![0u8; xf32.elem_count()*4]; d.download(xf32.buffer.as_ref(), &mut hb).unwrap(); + let vals: Vec = hb.chunks(4).map(|c| f32::from_le_bytes([c[0],c[1],c[2],c[3]])).collect(); + let nnan = vals.iter().filter(|x| x.is_nan()).count(); + let amax = vals.iter().filter(|x| x.is_finite()).fold(0f32, |a,&b| a.max(b.abs())); + eprintln!("DBG layer {l} ({mix}) xt.dtype={:?} nan={nnan} amax={amax:.3}", xt.dtype); + } let xn = match xn_carry.take() { Some(t) => t, - None => pf!(1, 0.0, rms_norm(d, &xt, &fwd[&format!("{p}.norm.weight")], eps).unwrap()), // [s, hid] + None => { if stream16 { let xtf = to_f32(d, &xt); pf!(1, 0.0, rms_norm(d, &xtf, &fwd[&format!("{p}.norm.weight")], eps).unwrap()) } + else { pf!(1, 0.0, rms_norm(d, &xt, &fwd[&format!("{p}.norm.weight")], eps).unwrap()) } }, }; match mix { 'M' => { @@ -2041,16 +2997,14 @@ 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()); - // On-device causal depthwise conv1d + silu over all S tokens. - // Result: yc_silu [s * conv_dim] flat. - let yc_silu = pf!(4, 0.0, conv1d_causal_prefill(d, &xbc_t, - &fwd[&format!("{p}.mixer.conv1d.weight")], - &fwd[&format!("{p}.mixer.conv1d.bias")], - s, conv_dim, kc).unwrap()); + // conv1d+silu+x/B/C split FUSED into one raw-float4 kernel + // (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; let ring_off = (s.saturating_sub(kc - 1)) * conv_dim; @@ -2082,9 +3036,13 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { // cost more than the split's streaming write+read. Machinery kept in // ssm_prefill_scan_ssd_strided; compact split path stays. { - // PART B: fused split — 1 dispatch replaces 3 strided_col_copy calls. - let (x_dev, b_dev, c_dev) = pf!(15, 0.0, - mamba_split_conv(d, &yc_silu, s, conv_dim, di, ng * ds).unwrap()); + // Fused conv1d+silu+split (raw float4): 1 launch, no temp conv buffer. + // Escape NEMOTRON_CONV_PREFILL_SPLIT_OFF=1 -> conv+split fallback. + let (x_dev, b_dev, c_dev) = pf!(4, 0.0, + conv1d_causal_prefill_split(d, &xbc_t, + &fwd[&format!("{p}.mixer.conv1d.weight")], + &fwd[&format!("{p}.mixer.conv1d.bias")], + s, conv_dim, di, ng * ds, kc).unwrap()); if use_ssd { ssd_split_out = Some(pf!(5, ssm_flops, ssm_prefill_scan_ssd(d, &x_dev, &fwd[&format!("{p}.mixer.A_log")], &b_dev, &c_dev, &fwd[&format!("{p}.mixer.D")], &dt_dev, ssm_state[l].as_ref().unwrap(), s as u32, m_dh as u32, ds as u32, m_nh as u32, ng as u32, ssd_l, seg_reset_dev.as_ref()).unwrap())); } else if use_ssd_port { @@ -2097,13 +3055,22 @@ 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). - let ynt = pf!(6, 0.0, gated_group_rmsnorm_batched(d, &y_dev, &z_t, - &up(&fw[&format!("{p}.mixer.norm.weight")]), eps, s, di, gs).unwrap()); - let out = qmm(&ynt.reshaped(vec![s, di]), &format!("{p}.mixer.out_proj.weight")); + // cast-fusion: write gated-norm directly as f16 + consume via qmm_h + // (skips a cast_f32_f16 per Mamba layer). Bit-identical f16. Escape: + // NEMOTRON_GATED_NORM_F16OUT_OFF=1 → original f32 gated-norm + qmm. + let out = if is_cuda && std::env::var("NEMOTRON_GATED_NORM_F16OUT_OFF").is_err() { + let ynt = pf!(6, 0.0, gated_group_rmsnorm_batched_f16out(d, &y_dev, &z_t, + &up(&fw[&format!("{p}.mixer.norm.weight")]), eps, s, di, gs).unwrap()); + qmm_h(&ynt.reshaped(vec![s, di]), &format!("{p}.mixer.out_proj.weight")) + } else { + let ynt = pf!(6, 0.0, gated_group_rmsnorm_batched(d, &y_dev, &z_t, + &up(&fw[&format!("{p}.mixer.norm.weight")]), eps, s, di, gs).unwrap()); + 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) = wh_butter_ops::add_rms_norm(d, &xt, &out, w, eps).unwrap(); xt = r; xn_carry = Some(nx); } @@ -2215,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) = wh_butter_ops::add_rms_norm(d, &xt, &out, w, eps).unwrap(); xt = r; xn_carry = Some(nx); } @@ -2230,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. @@ -2239,9 +3206,13 @@ 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) + // 16-bit stream: MoE router/gather/fp4-quant/shared-expert all assume + // f32 xn; cast once at arm entry so the whole MoE path stays f32. + let xn = if stream16 { to_f32(d, &xn) } else { xn }; let mut xn_f16_router: Option = None; // f16 xn from the router, reused by the gather // Fully-on-device MoE (descriptors+gather+scatter all device) → SKIP // the host triples reconstruction entirely. mt = s*top_k is @@ -2252,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[..])); @@ -2264,7 +3235,14 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { let rl_dev = pf!(10, 2.0 * s as f64 * n_exp as f64 * hid as f64, gemm_cublas_f32out(d, &xn_f16_e, gate_w16, s, n_exp, hid).unwrap()); xn_f16_router = Some(xn_f16_e); - let (st_dev, sw_dev, off_dev) = wh_butter_ops::moe_route_sort_device(d, &rl_dev, &bias_dev, n_exp, top_k, scale_f).unwrap(); + let (st_dev, sw_dev, off_dev) = if topk_scatter { + let (st_dev, sw_dev, off_dev, rowpos_dev, wt_tok_dev) = + wh_butter_ops::moe_route_sort_device_rowpos(d, &rl_dev, &bias_dev, n_exp, top_k, scale_f).unwrap(); + moe_tok = Some((rowpos_dev, wt_tok_dev)); + (st_dev, sw_dev, off_dev) + } else { + wh_butter_ops::moe_route_sort_device(d, &rl_dev, &bias_dev, n_exp, top_k, scale_f).unwrap() + }; let m = s * top_k; let tr: Vec<(u32, usize, f32)> = if devdesc_skip { Vec::new() // host-free: descriptors/gather/scatter consume device tensors @@ -2303,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; @@ -2312,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. @@ -2328,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") @@ -2344,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). @@ -2364,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 @@ -2390,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. @@ -2418,11 +3396,11 @@ 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() - && fp4w.contains_key(&format!("{p}.moe_up_proj_fp4")) && !use_w4a16; + && fp4w.contains_key(&format!("{p}.moe_up_proj_fp4")) && !use_w4a16 && !use_fp4_cutlass_env; let (xs, xs_dev_f16_opt) = if fp4_skip_gather { (Vec::new(), None) } else if use_w4a16_marlin && moe_dev.is_some() { @@ -2456,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(), @@ -2493,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); @@ -2652,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) }); @@ -2678,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() @@ -2694,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")); @@ -2707,7 +3685,25 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { // same det-scatter. Quality-holding (256 act levels). ── let w8a8_moe = use_w8a8_env && w8a8w.contains_key(&format!("{p}.moe_up_proj_w8a8")); let w4a8_moe = use_w4a8_env && w4a8w.contains_key(&format!("{p}.moe_up_proj_w4a8")); - if w8a8_moe { + let marlin_moe = use_marlin_moe && marlinw.contains_key(&format!("{p}.moe_up_proj_marlin")); + if marlin_moe { + let (off_dev, st_dev, sw_dev) = moe_dev.as_ref().unwrap(); + let xn2d = xn_f16_router.as_ref().unwrap().reshaped(vec![s, hid]); + let xs_g = pf!(12, 0.0, gather(d, &xn2d, st_dev).unwrap()); + let blk = 16usize; + let (stid_t, eids_t, ntpp_t) = wh_butter_ops::marlin_build_routing(d, off_dev, n_exp, blk, mt).unwrap(); + let topk1 = Tensor::new(d.upload(&vec![1.0f32; mt.max(1)].iter().flat_map(|x| x.to_le_bytes()).collect::>()).unwrap(), vec![mt.max(1)], DType::F32); + let (bu, bsu, gu) = &marlinw[&format!("{p}.moe_up_proj_marlin")]; + let (bd, bsd, gd) = &marlinw[&format!("{p}.moe_down_proj_marlin")]; + let up_out = pf!(11, 2.0 * mt as f64 * inter as f64 * hid as f64, + wh_butter_ops::moe_marlin_gemm(d, &xs_g, bu, bsu, gu, &topk1, &stid_t, &eids_t, &ntpp_t, mt, inter, hid, n_exp, 1, 16, blk, 48).unwrap()); + let a2 = pf!(3, 0.0, relu2_scale_f16(d, &up_out, inv).unwrap()); + let dn_out = pf!(11, 2.0 * mt as f64 * hid as f64 * inter as f64, + wh_butter_ops::moe_marlin_gemm(d, &a2, bd, bsd, gd, &topk1, &stid_t, &eids_t, &ntpp_t, mt, hid, inter, n_exp, 1, 16, blk, 48).unwrap()); + let acc_dev = Tensor::new(d.alloc_zeroed(s * hid * 4).unwrap(), vec![s, hid], DType::F32); + moe_scatter_add_det_dev(d, &dn_out, st_dev, sw_dev, &acc_dev, s, mt, hid, 256.0f32, true).unwrap(); + if ondevice_moe { acc_dev_keep = Some(acc_dev); } else { acc_h = dl(&acc_dev, s * hid); } + } else if w8a8_moe { let off_dev = off_dev_ref.unwrap(); let (_, st_dev, sw_dev) = moe_dev.as_ref().unwrap(); let xn2d = xn_f16_router.as_ref().unwrap().reshaped(vec![s, hid]); @@ -2804,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() { @@ -2826,7 +3822,12 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { drop(hb); 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(); - moe_scatter_add_det_dev(d, &dn_out, st_dev, sw_dev, &acc_dev, s, mt, hid, 256.0f32, true).unwrap(); + if topk_scatter && moe_tok.is_some() { + let (rowpos_dev, wt_tok_dev) = moe_tok.as_ref().unwrap(); + if std::env::var("NEMOTRON_TOPK_SCATTER_X2_OFF").is_err() { moe_scatter_add_topk_det_f16_x2(d, &dn_out, rowpos_dev, wt_tok_dev, &acc_dev, s, top_k, hid, 256.0f32).unwrap(); } else { moe_scatter_add_topk_det_f16(d, &dn_out, rowpos_dev, wt_tok_dev, &acc_dev, s, top_k, hid, 256.0f32).unwrap(); } + } else { + moe_scatter_add_det_dev(d, &dn_out, st_dev, sw_dev, &acc_dev, s, mt, hid, 256.0f32, true).unwrap(); + } if ondevice_moe { acc_dev_keep = Some(acc_dev); } else { acc_h = dl(&acc_dev, s * hid); } } else if fp4_cutlass { // ── CUTLASS grouped NVFP4 GEMM (NEMOTRON_CUTLASS_MOE=1 @@ -2860,7 +3861,12 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { wh_butter_ops::moe_grouped_gemm_cutlass_fp4(d, &bp, &bsf, &boff, &bgs, dwp, dwsf, dgw, &g_starts_c, &eids_c, hid, inter).unwrap()); 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(); - moe_scatter_add_det_dev(d, &dn_out, st_dev, sw_dev, &acc_dev, s, mt, hid, 256.0f32, true).unwrap(); + if topk_scatter && moe_tok.is_some() { + let (rowpos_dev, wt_tok_dev) = moe_tok.as_ref().unwrap(); + if std::env::var("NEMOTRON_TOPK_SCATTER_X2_OFF").is_err() { moe_scatter_add_topk_det_f16_x2(d, &dn_out, rowpos_dev, wt_tok_dev, &acc_dev, s, top_k, hid, 256.0f32).unwrap(); } else { moe_scatter_add_topk_det_f16(d, &dn_out, rowpos_dev, wt_tok_dev, &acc_dev, s, top_k, hid, 256.0f32).unwrap(); } + } else { + moe_scatter_add_det_dev(d, &dn_out, st_dev, sw_dev, &acc_dev, s, mt, hid, 256.0f32, true).unwrap(); + } if ondevice_moe { acc_dev_keep = Some(acc_dev); } else { acc_h = dl(&acc_dev, s * hid); } } else if fold_keys { let (uwp, uwsc, ugw) = &fp4w[&format!("{p}.moe_up_proj_fp4_fold")]; @@ -2881,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 }; @@ -2894,7 +3900,12 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { // 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(); - moe_scatter_add_det_dev(d, &dn_out, st_dev, sw_dev, &acc_dev, s, mt, hid, 256.0f32, true).unwrap(); + if topk_scatter && moe_tok.is_some() { + let (rowpos_dev, wt_tok_dev) = moe_tok.as_ref().unwrap(); + if std::env::var("NEMOTRON_TOPK_SCATTER_X2_OFF").is_err() { moe_scatter_add_topk_det_f16_x2(d, &dn_out, rowpos_dev, wt_tok_dev, &acc_dev, s, top_k, hid, 256.0f32).unwrap(); } else { moe_scatter_add_topk_det_f16(d, &dn_out, rowpos_dev, wt_tok_dev, &acc_dev, s, top_k, hid, 256.0f32).unwrap(); } + } else { + moe_scatter_add_det_dev(d, &dn_out, st_dev, sw_dev, &acc_dev, s, mt, hid, 256.0f32, true).unwrap(); + } if ondevice_moe { acc_dev_keep = Some(acc_dev); } else { acc_h = dl(&acc_dev, s * hid); } } else { // UP: [mt,inter] = xs[mt,hid] · Wup_q4[eid; inter,hid] (n=inter,k=hid) @@ -2915,16 +3926,21 @@ 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 // device token-idx + weight tensors → atomic scatter, NO host // tidx_h/wts_h (host-sync-free, toward graph capture). let (_, st_dev, sw_dev) = moe_dev.as_ref().unwrap(); - // DETERMINISTIC device scatter (device CSR build) → host-sync- - // free AND argmax-stable (the atomic scatter flipped 1104/1776). - moe_scatter_add_det_dev(d, &dn_out, st_dev, sw_dev, &acc_dev, s, mt, hid, 256.0f32, true).unwrap(); + // DETERMINISTIC scatter: top-k rowpos path skips the token CSR + // rebuild; fallback preserves the existing device CSR path. + if topk_scatter && moe_tok.is_some() { + let (rowpos_dev, wt_tok_dev) = moe_tok.as_ref().unwrap(); + if std::env::var("NEMOTRON_TOPK_SCATTER_X2_OFF").is_err() { moe_scatter_add_topk_det_f16_x2(d, &dn_out, rowpos_dev, wt_tok_dev, &acc_dev, s, top_k, hid, 256.0f32).unwrap(); } else { moe_scatter_add_topk_det_f16(d, &dn_out, rowpos_dev, wt_tok_dev, &acc_dev, s, top_k, hid, 256.0f32).unwrap(); } + } else { + moe_scatter_add_det_dev(d, &dn_out, st_dev, sw_dev, &acc_dev, s, mt, hid, 256.0f32, true).unwrap(); + } } else { let wts_h: Vec = triples.iter().map(|(_,_,w)| *w).collect(); let tidx_h: Vec = triples.iter().map(|(_,t,_)| *t as u32).collect(); @@ -2983,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 @@ -3018,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. @@ -3030,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 @@ -3054,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. @@ -3142,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() { @@ -3153,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 { @@ -3234,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 (iron_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()); @@ -3295,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 @@ -3306,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 @@ -3324,10 +4340,14 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { Some(t) => t.clone(), None => pf!(15, 0.0, cast_f32_f16(d, &xn).unwrap()), }; + if std::env::var("NEMOTRON_STREAM_DBG").is_ok() && l <= 2 { + let amx = |t: &Tensor| { let tf = if t.dtype==DType::F32 { t.clone() } else { to_f32(d,t) }; let mut hb=vec![0u8;tf.elem_count()*4]; d.download(tf.buffer.as_ref(),&mut hb).unwrap(); hb.chunks(4).map(|c| f32::from_le_bytes([c[0],c[1],c[2],c[3]]).abs()).filter(|x| x.is_finite()).fold(0f32,f32::max) }; + eprintln!(" DBG-SHARED l={l} xn.dt={:?} amax={:.3} xnh.dt={:?} amax={:.3} shared_inter={shared_inter}", xn.dtype, amx(&xn), xnh.dtype, amx(&xnh)); + } 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")]; @@ -3369,6 +4389,18 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { let sd_lt = pf!(12, 2.0*s as f64*hid as f64*shared_inter as f64, wh_butter_ops::gemm_fp4(d, &aq, &asf, &dwq, &dwsf, s, *sm_dn, *sk_dn, true, Some(&ds_d)).unwrap()); cast_f32_f16(d, &sd_lt).unwrap() + } else if stream16 { + // BF16-coherent shared FFN: up bf16, relu2 FP32-compute/BF16-store + // (no f16 clamp/saturation), down FP32-accum/BF16-out. Matches the + // bf16 residual stream so ReLU2 amplification stays representable. + let wsu = deq16bf(&format!("{p}.shup"), suqs, susc, *sm_up, *sk_up); + let xnb = pf!(15, 0.0, wh_butter_ops::cast_f32_bf16(d, &xn).unwrap()); + let sa = pf!(12, 2.0*s as f64*shared_inter as f64*hid as f64, + gemm_cublas(d, &xnb, &wsu, s, *sm_up, *sk_up).unwrap()); // bf16 + let sa2 = pf!(3, 0.0, wh_butter_ops::relu2_scale_bf16(d, &sa, 1.0f32).unwrap()); // bf16 + let wsd = deq16bf(&format!("{p}.shdn"), sdqs, sdsc, *sm_dn, *sk_dn); + pf!(12, 2.0*s as f64*hid as f64*shared_inter as f64, + gemm_cublas(d, &sa2, &wsd, s, *sm_dn, *sk_dn).unwrap()) // [s,hid] bf16 } else { let wsu = deq16(&format!("{p}.shup"), suqs, susc, *sm_up, *sk_up, 0); let sa = pf!(12, 2.0*s as f64*shared_inter as f64*hid as f64, @@ -3378,14 +4410,14 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { pf!(12, 2.0*s as f64*hid as f64*shared_inter as f64, gemm_cublas(d, &sa2, &wsd, s, *sm_dn, *sk_dn).unwrap()) // [s,hid] f16 }; - sd_dev_opt = Some(pf!(15, 0.0, cast_f16_f32(d, &sd).unwrap())); // device f32, full scale + sd_dev_opt = Some(if sd.dtype == DType::BF16 { pf!(15, 0.0, wh_butter_ops::cast_bf16_f32(d, &sd).unwrap()) } else { pf!(15, 0.0, cast_f16_f32(d, &sd).unwrap()) }); Vec::new() } else if metal_mma_shared { // ── Portable Metal shared-expert MMA: gemm_q4_mpp up/down, host relu2 ── // 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, @@ -3400,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()); @@ -3446,17 +4478,23 @@ 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 + let acc_dev = if stream16 { cast_stream(d, &acc_dev) } else { acc_dev }; xt = add(d, &xt, &acc_dev).unwrap(); } else { let sd_dev = sd_dev_opt.take().unwrap_or_else(|| upm(&sd_h, vec![s, hid])); + if stream_dbg { + let amx = |t: &Tensor| { let tf = if t.dtype==DType::F32 { t.clone() } else { to_f32(d,t) }; let mut hb=vec![0u8;tf.elem_count()*4]; d.download(tf.buffer.as_ref(),&mut hb).unwrap(); hb.chunks(4).map(|c| f32::from_le_bytes([c[0],c[1],c[2],c[3]]).abs()).filter(|x| x.is_finite()).fold(0f32,f32::max) }; + eprintln!(" DBG-MERGE l={l} acc_dev.dt={:?} amax={:.3} sd_dev.dt={:?} amax={:.3}", acc_dev.dtype, amx(&acc_dev), sd_dev.dtype, amx(&sd_dev)); + } + let (acc_dev, sd_dev) = if stream16 { (cast_stream(d, &acc_dev), cast_stream(d, &sd_dev)) } else { (acc_dev, sd_dev) }; match nw { Some(w) if is_cuda && std::env::var("NEMOTRON_FUSED_ADDNORM").is_ok() => { let merged = add(d, &acc_dev, &sd_dev).unwrap(); let (r, nx) = wh_butter_ops::add_rms_norm(d, &xt, &merged, w, eps).unwrap(); xt = r; xn_carry = Some(nx); } // one launch, no intermediate: residual + routed acc + shared out @@ -3468,8 +4506,9 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { } else { for i in 0..s*hid { acc_h[i] += sd_h[i]; } 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) = wh_butter_ops::add_rms_norm(d, &xt, &acc_up, w, eps).unwrap(); xt = r; xn_carry = Some(nx); } @@ -3479,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")), @@ -3490,16 +4529,25 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { qmm(&xn, &format!("{p}.mixer.k_proj.weight")), qmm(&xn, &format!("{p}.mixer.v_proj.weight"))) }; + // 16-bit stream: attention runs f32 (kvcache + rope + sdpa f32). + let (q_all, k_all, v_all) = if stream16 { + (to_f32(d, &q_all), to_f32(d, &k_all), to_f32(d, &v_all)) + } else { (q_all, k_all, v_all) }; // [s, qdim]=[s,nq*hd]; k/v [s, kvdim]=[s,nkv*hd] // PART C: eliminate dl(q)/dl(k)/dl(v) + per-token rope dispatches. // Build positions [s] once, use batched rope_llama_many + kv_append_many. if kvcache[l].is_none() { kvcache[l] = Some((Tensor::new(d.alloc_zeroed(nkv*cap*hd*4).unwrap(), vec![nkv*cap*hd], DType::F32), Tensor::new(d.alloc_zeroed(nkv*cap*hd*4).unwrap(), vec![nkv*cap*hd], DType::F32))); } - let positions_h: Vec = (0..s).map(|ti| (base + ti) as u32).collect(); - let positions_dev = Tensor::new( - d.upload(unsafe { std::slice::from_raw_parts(positions_h.as_ptr() as *const u8, positions_h.len() * 4) }).unwrap(), - vec![s], DType::U32); + // Cache the [s] position ids by (base,s): deterministic, so upload + // ONCE and reuse across all attention layers + forwards (removes the + // per-attention-layer H2D copy). + let positions_dev = cdev(format!("__positions.{base}.{s}"), &|| { + let positions_h: Vec = (0..s).map(|ti| (base + ti) as u32).collect(); + Tensor::new( + d.upload(unsafe { std::slice::from_raw_parts(positions_h.as_ptr() as *const u8, positions_h.len() * 4) }).unwrap(), + vec![s], DType::U32) + }); // Batched rope: [s, n_heads, hd] each → rotated in one dispatch. let qr = pf!(8, 0.0, rope_llama_many(d, &q_all.reshaped(vec![s, nq, hd]), &positions_dev, nq, hd, rope_theta, 1.0, 1.0, 1.0, 8192.0).unwrap()); let kr = pf!(8, 0.0, rope_llama_many(d, &k_all.reshaped(vec![s, nkv, hd]), &positions_dev, nkv, hd, rope_theta, 1.0, 1.0, 1.0, 8192.0).unwrap()); @@ -3538,14 +4586,14 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { pf!(7, attn_flops, wh_butter_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, wh_butter_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()) @@ -3566,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) = wh_butter_ops::add_rms_norm(d, &xt, &o, w, eps).unwrap(); xt = r; xn_carry = Some(nx); } @@ -3585,14 +4633,17 @@ 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). let xf = match xn_carry.take() { Some(t) => t, - None => rms_norm(d, &xt, &fwd["norm_f"], eps).unwrap(), // [s, hid] + None => { if stream16 { let xtf = to_f32(d, &xt); rms_norm(d, &xtf, &fwd["norm_f"], eps).unwrap() } + else { rms_norm(d, &xt, &fwd["norm_f"], eps).unwrap() } }, // [s, hid] }; + // tail (slice + logits) runs f32: cast the 16-bit stream back once. + let xf = if stream16 { to_f32(d, &xf) } else { xf }; let last_dev = slice(d, &xf, (s-1)*hid, hid).unwrap(); // [hid] device let logits_dev = pf!(14, 2.0 * vocab as f64 * hid as f64, qmv(&last_dev, "language_model.lm_head.weight")); LAST_BATCHED_LOGITS_DEV.with(|c| *c.borrow_mut() = Some(logits_dev.clone())); @@ -3607,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 @@ -3703,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 }; @@ -3755,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() { @@ -3810,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 @@ -3852,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() { @@ -3931,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"); @@ -3953,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"); @@ -4125,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(); @@ -4144,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 @@ -4157,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/wh-butter-ops/src/lib.rs b/rust/crates/wh-butter-ops/src/lib.rs index 4811e717..24887eca 100644 --- a/rust/crates/wh-butter-ops/src/lib.rs +++ b/rust/crates/wh-butter-ops/src/lib.rs @@ -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); @@ -180,7 +180,51 @@ fn elementwise( } /// Elementwise sum `a + b` (e.g. residual connections). +const ADD16_SRC: &str = r#" +#include +#include +extern "C" __global__ void butter_add_f16(const __half* __restrict__ a, const __half* __restrict__ b, __half* __restrict__ out, long n){ + long i=(long)blockIdx.x*blockDim.x+threadIdx.x; long stride=(long)gridDim.x*blockDim.x; + for(;i +extern "C" __global__ void butter_mul_scalar_f16(const __half* __restrict__ a, __half* __restrict__ out, float sc, long n){ + long i=(long)blockIdx.x*blockDim.x+threadIdx.x; long stride=(long)gridDim.x*blockDim.x; + for(;i Result { + let out = Tensor::empty(dev, a.shape.clone(), DType::F16)?; + let n = a.elem_count(); + dev.dispatch_raw_cuda(MULSCALAR16_SRC, "mul_scalar_f16.cu", "butter_mul_scalar_f16", + &[(a.buffer.as_ref(), a.offset), (out.buffer.as_ref(), 0)], + &[scale.to_le_bytes().to_vec(), (n as i64).to_le_bytes().to_vec()], + [(n as u32).div_ceil(256).max(1), 1, 1], [256, 1, 1], 0, false)?; + Ok(out) +} + pub fn add(dev: &dyn Device, a: &Tensor, b: &Tensor) -> Result { + // 16-bit DSL elementwise vec4 codegen IMAs; route f16/bf16 through a raw kernel. + if a.dtype == DType::F16 || a.dtype == DType::BF16 { + let out = Tensor::empty(dev, a.shape.clone(), a.dtype)?; + let n = a.elem_count(); + let fname = if a.dtype == DType::F16 { "butter_add_f16" } else { "butter_add_bf16" }; + dev.dispatch_raw_cuda(ADD16_SRC, "add16.cu", fname, + &[(a.buffer.as_ref(), a.offset), (b.buffer.as_ref(), b.offset), (out.buffer.as_ref(), 0)], + &[(n as i64).to_le_bytes().to_vec()], + [(n as u32).div_ceil(256).max(1), 1, 1], [256, 1, 1], 0, false)?; + return Ok(out); + } elementwise(dev, a, b, BinOpKind::Add, "iron_add") } @@ -196,24 +240,34 @@ extern "C" __global__ void butter_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 +#include +extern "C" __global__ void butter_add3_f16(const __half* __restrict__ a, const __half* __restrict__ b, const __half* __restrict__ c, __half* __restrict__ out, long n){ + long i=(long)blockIdx.x*blockDim.x+threadIdx.x; long stride=(long)gridDim.x*blockDim.x; + for(;i Result { let out = Tensor::empty(dev, a.shape.clone(), a.dtype)?; let n = a.elem_count(); let l = |v: i64| v.to_le_bytes().to_vec(); - dev.dispatch_raw_cuda(ADD3_SRC, "add3.cu", "butter_add3", + let fname = match a.dtype { DType::F16 => "butter_add3_f16", DType::BF16 => "butter_add3_bf16", _ => "butter_add3" }; + let per = if a.dtype == DType::F32 { 256 * 4 } else { 256 }; + dev.dispatch_raw_cuda(ADD3_SRC, "add3.cu", fname, &[(a.buffer.as_ref(), a.offset), (b.buffer.as_ref(), b.offset), (c.buffer.as_ref(), c.offset), (out.buffer.as_ref(), 0)], &[l(n as i64)], - [(n as u32).div_ceil(256 * 4).max(1), 1, 1], [256, 1, 1], 0, false)?; + [(n as u32).div_ceil(per).max(1), 1, 1], [256, 1, 1], 0, false)?; Ok(out) } @@ -243,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; @@ -403,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 `iron_dsv4_mhc_collapse`. pub fn dsv4_mhc_collapse( @@ -481,7 +535,9 @@ pub fn rms_norm(dev: &dyn Device, x: &Tensor, weight: &Tensor, eps: f32) -> Resu // Fast path: 4 elems/thread, needs n a multiple of 128 and ≤ 4096. // Otherwise the strided wide variant handles any row width. - let (kname, block) = if n % 128 == 0 && n <= 4096 { + // 16-bit (f16/bf16) DSL vec4 fast-path codegen is broken (IMAs); route them + // through the scalar wide variant. f32 keeps the fast path. + let (kname, block) = if x.dtype == DType::F32 && n % 128 == 0 && n <= 4096 { ("iron_rms_norm", (n / 4) as u32) } else { ("iron_rms_norm_wide", 256u32) @@ -534,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` @@ -640,7 +696,7 @@ 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 { @@ -793,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("iron_softplus_add", DType::F32, || { let mut k = wh_iron_std::kernels::ssm::scan::iron_softplus_add::kernel_ir(); k.mode = wh_iron_core::ir::KernelMode::Grid3D; k }); let n = a.elem_count(); @@ -930,8 +986,154 @@ pub fn conv1d_causal_prefill( Ok(y) } +const CONV1D_PREFILL_V4_SRC: &str = r#" +#include + +__device__ __forceinline__ float silu_f32(float x) { + return x / (1.0f + expf(-x)); +} + +extern "C" __global__ void conv1d_causal_prefill_v4( + const float* __restrict__ xbc, + const float* __restrict__ w, + const float* __restrict__ bias, + float* __restrict__ y, + int conv_dim4, int kc, int total4) +{ + int idx = (int)(blockIdx.x * blockDim.x + threadIdx.x); + if (idx >= total4) return; + int ti = idx / conv_dim4; + int ci4 = idx - ti * conv_dim4; + const float4* __restrict__ x4 = reinterpret_cast(xbc); + const float4* __restrict__ w4 = reinterpret_cast(w); + const float4* __restrict__ b4 = reinterpret_cast(bias); + float4 acc = b4[ci4]; + #pragma unroll + for (int k = 0; k < 8; ++k) { + if (k >= kc) break; + int src_t = ti - (kc - 1 - k); + if (src_t >= 0) { + float4 xv = x4[(size_t)src_t * conv_dim4 + ci4]; + float4 wv = w4[(size_t)k * conv_dim4 + ci4]; + acc.x += xv.x * wv.x; + acc.y += xv.y * wv.y; + acc.z += xv.z * wv.z; + acc.w += xv.w * wv.w; + } + } + acc.x = silu_f32(acc.x); + acc.y = silu_f32(acc.y); + acc.z = silu_f32(acc.z); + acc.w = silu_f32(acc.w); + reinterpret_cast(y)[idx] = acc; +} + +extern "C" __global__ void conv1d_causal_prefill_split_v4( + const float* __restrict__ xbc, + const float* __restrict__ w, + const float* __restrict__ bias, + float* __restrict__ x_out, + float* __restrict__ b_out, + float* __restrict__ c_out, + int conv_dim4, int di4, int ng_ds4, int kc, int total4) +{ + int idx = (int)(blockIdx.x * blockDim.x + threadIdx.x); + if (idx >= total4) return; + int ti = idx / conv_dim4; + int ci4 = idx - ti * conv_dim4; + const float4* __restrict__ x4 = reinterpret_cast(xbc); + const float4* __restrict__ w4 = reinterpret_cast(w); + const float4* __restrict__ b4 = reinterpret_cast(bias); + float4 acc = b4[ci4]; + #pragma unroll + for (int k = 0; k < 8; ++k) { + if (k >= kc) break; + int src_t = ti - (kc - 1 - k); + if (src_t >= 0) { + float4 xv = x4[(size_t)src_t * conv_dim4 + ci4]; + float4 wv = w4[(size_t)k * conv_dim4 + ci4]; + acc.x += xv.x * wv.x; + acc.y += xv.y * wv.y; + acc.z += xv.z * wv.z; + acc.w += xv.w * wv.w; + } + } + acc.x = silu_f32(acc.x); + acc.y = silu_f32(acc.y); + acc.z = silu_f32(acc.z); + acc.w = silu_f32(acc.w); + if (ci4 < di4) { + reinterpret_cast(x_out)[(size_t)ti * di4 + ci4] = acc; + } else if (ci4 < di4 + ng_ds4) { + int j = ci4 - di4; + reinterpret_cast(b_out)[(size_t)ti * ng_ds4 + j] = acc; + } else { + int j = ci4 - di4 - ng_ds4; + reinterpret_cast(c_out)[(size_t)ti * ng_ds4 + j] = acc; + } +} +"#; + +/// Fused raw-float4 causal conv1d+SiLU that writes the post-conv x/B/C split +/// surfaces directly, deleting the temp [s,conv_dim] conv output + the split +/// launch. Equivalent to conv1d_causal_prefill + mamba_split_conv. Escape with +/// NEMOTRON_CONV_PREFILL_SPLIT_OFF=1 (falls back to that pair). +pub fn conv1d_causal_prefill_split( + dev: &dyn Device, + xbc_in: &Tensor, + w: &Tensor, + bias: &Tensor, + s: usize, + conv_dim: usize, + di: usize, + ng_ds: usize, + kc: usize, +) -> Result<(Tensor, Tensor, Tensor)> { + let raw_ok = std::env::var("NEMOTRON_CONV_PREFILL_RAW_OFF").is_err() + && std::env::var("NEMOTRON_CONV_PREFILL_SPLIT_OFF").is_err() + && xbc_in.dtype == DType::F32 && w.dtype == DType::F32 && bias.dtype == DType::F32 + && xbc_in.offset % 16 == 0 && w.offset % 16 == 0 && bias.offset % 16 == 0 + && di + 2 * ng_ds == conv_dim + && di % 4 == 0 && ng_ds % 4 == 0 && conv_dim % 4 == 0; + if raw_ok { + let x_out = Tensor::empty(dev, vec![s * di], DType::F32)?; + let b_out = Tensor::empty(dev, vec![s * ng_ds], DType::F32)?; + let c_out = Tensor::empty(dev, vec![s * ng_ds], DType::F32)?; + let i = |v: i32| v.to_le_bytes().to_vec(); + let total4 = (s * conv_dim / 4) as u32; + if dev.dispatch_raw_cuda( + CONV1D_PREFILL_V4_SRC, + "conv1d_prefill_v4.cu", + "conv1d_causal_prefill_split_v4", + &[ + (xbc_in.buffer.as_ref(), xbc_in.offset), + (w.buffer.as_ref(), w.offset), + (bias.buffer.as_ref(), bias.offset), + (x_out.buffer.as_ref(), 0), + (b_out.buffer.as_ref(), 0), + (c_out.buffer.as_ref(), 0), + ], + &[ + i((conv_dim / 4) as i32), + i((di / 4) as i32), + i((ng_ds / 4) as i32), + i(kc as i32), + i(total4 as i32), + ], + [total4.div_ceil(256), 1, 1], + [256, 1, 1], + 0, + false, + ).is_ok() { + return Ok((x_out, b_out, c_out)); + } + } + let y = conv1d_causal_prefill(dev, xbc_in, w, bias, s, conv_dim, kc)?; + mamba_split_conv(dev, &y, s, conv_dim, di, ng_ds) +} + /// 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, @@ -943,10 +1145,10 @@ pub fn strided_col_copy( ) -> Result { let kernel = lookup("iron_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] @@ -960,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, &[ @@ -969,16 +1172,97 @@ pub fn strided_col_copy( u(col_off as u32), u(width as u32), ], - // 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] } }, + // The kernel has no internal idx guard, so the grid must cover exactly + // s*width threads. Laguna's gate-column extraction is not always + // 64-aligned (head counts 48/72): use the largest power-of-two block + // (<= 64) that divides the total — 64 for the common aligned shapes. + { + let b = 1u32 << total.trailing_zeros().min(6); + Grid { grid: [total / b, 1, 1], block: [b, 1, 1] } + }, + )?; + Ok(dst) +} + +const EXTRACT_COLS_SRC: &str = r#" +extern "C" __global__ void butter_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 butter_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", "butter_extract_cols"), + DType::F16 => (DType::F16.size_bytes(), EXTRACT_COLS_F16_SRC, "extract_cols_f16.cu", "butter_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, @@ -1021,6 +1305,82 @@ pub fn moe_gather_up_relu2(dev: &dyn Device, qs: &Tensor, scales: &Tensor, x: &T dev.dispatch(&k, &[Binding::Buffer(qs.buffer.clone()), Binding::Buffer(scales.buffer.clone()), Binding::Buffer(x.buffer.clone()), Binding::Buffer(idx.buffer.clone()), Binding::Buffer(out.buffer.clone()), u(hid as u32), u(inter as u32), u(rpt)], Grid { grid: [(inter as u32).div_ceil(rpt), top_k as u32, 1], block: [32 * rpt, 1, 1] })?; Ok(out) } +// Single-block strided argmax over an f32 vector (the 100K-vocab logit row). +// Ties resolve to the LOWEST index, matching the host-side scan. One block of +// 1024 threads striding n, then a shared-memory tree reduce: microseconds at +// vocab scale, and it shrinks the per-step D2H from the full logit row to 4 +// bytes for greedy decode. +const ARGMAX_F32_SRC: &str = r#" +extern "C" __global__ void butter_argmax_f32( + const float* __restrict__ x, unsigned* __restrict__ out, int n) +{ + __shared__ float sv[1024]; + __shared__ unsigned si[1024]; + int tid = threadIdx.x; + float best = -3.4028235e38f; + unsigned bi = 0u; + for (int i = tid; i < n; i += blockDim.x) { + float v = x[i]; + if (v > best || (v == best && (unsigned)i < bi)) { best = v; bi = (unsigned)i; } + } + sv[tid] = best; si[tid] = bi; + __syncthreads(); + for (int s = blockDim.x / 2; s > 0; s >>= 1) { + if (tid < s) { + if (sv[tid + s] > sv[tid] || (sv[tid + s] == sv[tid] && si[tid + s] < si[tid])) { + sv[tid] = sv[tid + s]; si[tid] = si[tid + s]; + } + } + __syncthreads(); + } + if (tid == 0) out[0] = si[0]; +} +"#; + +/// Device argmax of an f32 vector: returns a 1-element u32 tensor holding the +/// index of the maximum (ties: lowest index). Greedy decode downloads these 4 +/// bytes instead of the whole logit row. +pub fn argmax_f32_device(dev: &dyn Device, x: &Tensor) -> Result { + let n = x.elem_count(); + let out = Tensor::empty(dev, vec![1], DType::U32)?; + let i = |v: i32| v.to_le_bytes().to_vec(); + dev.dispatch_raw_cuda( + ARGMAX_F32_SRC, "argmax_f32.cu", "butter_argmax_f32", + &[(x.buffer.as_ref(), x.offset), (out.buffer.as_ref(), 0)], + &[i(n as i32)], + [1, 1, 1], [1024, 1, 1], 0, false, + )?; + Ok(out) +} + +/// Batched MoE expert gather GEMV, plain (no fused activation): one +/// `[top_k*inter]` output from the contiguous `[n_exp*inter, hid]` Q4 stack. +/// The SwiGLU MoE path calls this for the gate and up stacks, then applies +/// the activation elementwise before the down gather. F16 per-block scales. +/// +/// TODO(rebrand-merge): blocked — dispatches `iron_moe_gather_q4` via +/// `wh_iron_std`'s kernel IR, but the pinned `thewafflehaus/iron@dev` kernel +/// library only ships `iron_moe_gather_q4_relu2/_down/_down_accum` — the +/// plain (non-relu2) variant this needs is not present upstream yet. Kept as +/// an always-error stub until that kernel lands in the iron repo. +pub fn moe_gather_q4(_dev: &dyn Device, _qs: &Tensor, _scales: &Tensor, _x: &Tensor, _idx: &Tensor, _top_k: usize, _inter: usize, _hid: usize) -> Result { + Err(Error::Msg("moe_gather_q4: blocked — iron_moe_gather_q4 kernel not yet in thewafflehaus/iron@dev (only _relu2/_down/_down_accum variants exist)".into())) +} + +/// Batched MoE GATE+UP gather with FUSED inline SwiGLU: reads the SAME +/// activation `x` once and computes both the gate and up dot products for +/// each (slot, inter-row), writing `silu(gate)·up` straight to a single +/// `[top_k*inter]` output. Replaces two `moe_gather_q4` calls (gate stack, +/// up stack) plus a host-side elementwise SwiGLU pass with one dispatch. +/// F16 per-block scales on both `gate_scales` and `up_scales`. +#[allow(clippy::too_many_arguments)] +/// TODO(rebrand-merge): blocked — same iron-kernel gap as `moe_gather_q4` +/// above (`iron_moe_gather_q4_swiglu` not in `thewafflehaus/iron@dev` yet). +#[allow(clippy::too_many_arguments)] +pub fn moe_gather_q4_swiglu(_dev: &dyn Device, _gate_qs: &Tensor, _gate_scales: &Tensor, _up_qs: &Tensor, _up_scales: &Tensor, _x: &Tensor, _idx: &Tensor, _top_k: usize, _inter: usize, _hid: usize) -> Result { + Err(Error::Msg("moe_gather_q4_swiglu: blocked — iron_moe_gather_q4_swiglu kernel not yet in thewafflehaus/iron@dev".into())) +} + /// Batched MoE down + router-weighted accumulate into `acc[hid]`. `qs` is the /// contiguous `[n_exp*hid, inter]` Q4 weight; `x` is the `[top_k*inter]` up output. #[allow(clippy::too_many_arguments)] @@ -1112,7 +1472,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; @@ -1232,7 +1592,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 wh_iron_core::ir::KernelMode; @@ -1254,12 +1614,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, @@ -1336,6 +1696,32 @@ extern "C" __global__ void moe_scatter_sort( base += __popc(mask); } } + +extern "C" __global__ void moe_scatter_sort_rowpos( + const unsigned* __restrict__ idx, const float* __restrict__ wt, + const unsigned* __restrict__ offsets, unsigned* __restrict__ sorted_tok, + float* __restrict__ sorted_wt, unsigned* __restrict__ rowpos, + int mt, int top_k, int n_exp) +{ + // Same stable expert-major ordering as moe_scatter_sort, plus a token-major + // slot -> sorted-row map for deterministic fixed-top_k output scatter. + int e = (blockIdx.x*blockDim.x + threadIdx.x) >> 5; if (e >= n_exp) return; + int lane = threadIdx.x & 31; + unsigned base = offsets[e]; + for (int w = 0; w < mt; w += 32) { + int i = w + lane; + bool match = (i < mt) && (idx[i] == (unsigned)e); + unsigned mask = __ballot_sync(0xffffffffu, match); + if (match) { + unsigned rank = __popc(mask & ((1u << lane) - 1u)); + unsigned pos = base + rank; + sorted_tok[pos] = (unsigned)(i / top_k); + sorted_wt[pos] = wt[i]; + rowpos[i] = pos; + } + base += __popc(mask); + } +} "#; #[allow(clippy::type_complexity)] @@ -1377,6 +1763,42 @@ pub fn moe_route_sort_device( Ok((sorted_tok, sorted_wt, offsets)) } +#[allow(clippy::type_complexity)] +pub fn moe_route_sort_device_rowpos( + dev: &dyn Device, gate_logits: &Tensor, bias: &Tensor, n_exp: usize, top_k: usize, scale: f32, +) -> Result<(Tensor, Tensor, Tensor, Tensor, Tensor)> { + let s = gate_logits.elem_count() / n_exp; + let mt = s * top_k; + let idx = Tensor::empty(dev, vec![mt], DType::U32)?; + let wt = Tensor::empty(dev, vec![mt], DType::F32)?; + let counts = Tensor::new(dev.alloc_zeroed(n_exp * 4).map_err(|e| Error::Msg(format!("{e:?}")))?, vec![n_exp], DType::U32); + let offsets = Tensor::empty(dev, vec![n_exp + 1], DType::U32)?; + let cursor = Tensor::empty(dev, vec![n_exp], DType::U32)?; + let sorted_tok = Tensor::empty(dev, vec![mt], DType::U32)?; + let sorted_wt = Tensor::empty(dev, vec![mt], DType::F32)?; + let rowpos = Tensor::empty(dev, vec![mt], DType::U32)?; + let i = |x: i32| x.to_le_bytes().to_vec(); + let f = |x: f32| x.to_le_bytes().to_vec(); + let blk = if n_exp >= 256 { 256 } else { n_exp.max(32) } as u32; + dev.dispatch_raw_cuda(MOE_BATCHED_ROUTER_SRC, "moe_route_sort.cu", "moe_batched_router", + &[(gate_logits.buffer.as_ref(), gate_logits.offset), (bias.buffer.as_ref(), bias.offset), + (idx.buffer.as_ref(), 0), (wt.buffer.as_ref(), 0)], + &[i(s as i32), i(n_exp as i32), i(top_k as i32), f(scale)], + [s as u32, 1, 1], [blk, 1, 1], 0, false)?; + dev.dispatch_raw_cuda(MOE_HIST_SRC, "moe_hist.cu", "moe_hist", + &[(idx.buffer.as_ref(), 0), (counts.buffer.as_ref(), 0)], + &[i(mt as i32)], [(mt as u32).div_ceil(256), 1, 1], [256, 1, 1], 0, false)?; + dev.dispatch_raw_cuda(MOE_PREFIX_SRC, "moe_prefix.cu", "moe_prefix", + &[(counts.buffer.as_ref(), 0), (offsets.buffer.as_ref(), 0), (cursor.buffer.as_ref(), 0)], + &[i(n_exp as i32)], [1, 1, 1], [32, 1, 1], 0, false)?; + let _ = &cursor; + dev.dispatch_raw_cuda(MOE_SCATTER_SORT_SRC, "moe_scatter_sort.cu", "moe_scatter_sort_rowpos", + &[(idx.buffer.as_ref(), 0), (wt.buffer.as_ref(), 0), (offsets.buffer.as_ref(), 0), + (sorted_tok.buffer.as_ref(), 0), (sorted_wt.buffer.as_ref(), 0), (rowpos.buffer.as_ref(), 0)], + &[i(mt as i32), i(top_k as i32), i(n_exp as i32)], [((n_exp * 32) as u32).div_ceil(128), 1, 1], [128, 1, 1], 0, false)?; + Ok((sorted_tok, sorted_wt, offsets, rowpos, wt)) +} + pub fn moe_weighted_sum(dev: &dyn Device, downs: &Tensor, wts: &Tensor, acc: &Tensor, hid: usize, top_k: usize) -> Result<()> { let k = cached_ir("iron_moe_weighted_sum", acc.dtype, || { let mut k = wh_iron_std::kernels::moe::moe_gather_q4::iron_moe_weighted_sum::kernel_ir_for(acc.dtype); k.mode = wh_iron_core::ir::KernelMode::Grid3D; k }); let u = |v: u32| Binding::Scalar(v.to_le_bytes().to_vec()); @@ -1384,7 +1806,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). @@ -1422,7 +1844,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); BUTTER_GEMV_SHARED_RPT overrides for experimentation. @@ -1433,7 +1855,10 @@ fn gemv_q4_dispatch(dev: &dyn Device, kernel: &str, qs: &Tensor, scales: &Tensor dev.dispatch(&k, &b, grid)?; Ok(out) } -/// Q4 matvec `out = Wq4·x`. (block 32, int4 [-7,7], f32 scale/block.) +/// Q4 matvec `out = Wq4·x`. Block 32, int4 [-7,7], ONE F16 SCALE PER BLOCK: +/// the coalesced/relu2/accum kernels read `scales` as f16, so convert the f32 +/// scales from [`quantize_q4`] with an f32-to-f16 pass before upload (raw f32 +/// bytes get silently reinterpreted as garbage f16 and inflate every output). pub fn gemv_q4(dev: &dyn Device, qs: &Tensor, scales: &Tensor, x: &Tensor, m_out: usize, k_in: usize, rpg: usize) -> Result { gemv_q4_dispatch(dev, "plain", qs, scales, x, None, None, m_out, k_in, rpg) } @@ -1450,6 +1875,9 @@ pub fn gemv_q4_accum(dev: &dyn Device, qs: &Tensor, scales: &Tensor, x: &Tensor, /// Quantize a row-major `[m,k]` F32 weight to Q4 blocks (block 32, symmetric int4 /// in [-7,7], f32 scale `amax/7`). Returns `(qs_u32, scales)`: qs `[m·k/32·4]` /// (8 nibbles/u32, 4 u32/block), scales `[m·k/32]`. `k` must be a multiple of 32. +/// NOTE: the returned scales are host f32; the `gemv_q4`/`gemv_q4_relu2`/ +/// `gemv_q4_accum` device kernels expect them uploaded as F16 (see the +/// `gemv_q4` doc). Only `MT_GEMV_VEC`'s vectorized kernel takes f32 scales. pub fn quantize_q4(w: &[f32], m: usize, k: usize) -> (Vec, Vec) { assert_eq!(k % 32, 0, "quantize_q4: k must be a multiple of 32"); let bpr = k / 32; @@ -1523,7 +1951,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; } @@ -1555,12 +1983,60 @@ 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. /// /// Example: `let (codes, scales, g) = quantize_nvfp4(&w, n_rows, k_in);` +#[inline] +fn marlin_f32_to_f16(f: f32) -> u16 { + let x = f.to_bits(); + let sign = ((x >> 16) & 0x8000) as u16; + let e = ((x >> 23) & 0xff) as i32 - 112; // 127 - 15 + if e <= 0 { return sign; } + if e >= 0x1f { return sign | 0x7c00; } + let m = (x >> 13) & 0x3ff; + let round = (x >> 12) & 1; + let v = ((e as u32) << 10) | m; + sign | ((v + round) as u16) +} + +/// Convert per-16 NVFP4 block-scale VALUES `[ngroup, n]` (row-major; value = +/// dec_ue4m3(byte)*global) into the marlin S0E5M3 fp8 byte layout the GEMM reads. +/// Mirrors marlin_permute_scales(group Vec { + let total = ngroup * n; + assert_eq!(vals.len(), total); + assert_eq!(total % 64, 0, "marlin_s0e5m3_scales: ngroup*n must be a multiple of 64"); + // scale_perm (64): [i + 8*j for i in 0..8 for j in 0..8] + let mut sp = [0usize; 64]; + let mut idx = 0; + for i in 0..8 { for j in 0..8 { sp[idx] = i + 8 * j; idx += 1; } } + // A: 64-chunk column permute (marlin_permute_scales) + let mut s1 = vec![0f32; total]; + for r in 0..(total / 64) { + for j in 0..64 { s1[r * 64 + j] = vals[r * 64 + sp[j]]; } + } + // B: 4-chunk interleave [0,2,1,3] + let il = [0usize, 2, 1, 3]; + let mut s2 = vec![0f32; total]; + for r in 0..(total / 4) { + for j in 0..4 { s2[r * 4 + j] = s1[r * 4 + il[j]]; } + } + // C: x2^7, zero-if<2, f16 bits <<1, high byte = S0E5M3 + let mut out = vec![0u8; total]; + for p in 0..total { + let mut v = s2[p] * 128.0; + if v < 2.0 { v = 0.0; } + let shifted = marlin_f32_to_f16(v) << 1; + out[p] = (shifted >> 8) as u8; + } + out +} + pub fn quantize_nvfp4(w: &[f32], m: usize, k: usize) -> (Vec, Vec, f32) { assert_eq!(k % 16, 0, "quantize_nvfp4: k must be a multiple of 16"); let nblk = k / 16; @@ -1611,7 +2087,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`). @@ -1621,7 +2097,7 @@ 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 /// `iron_gemm` (weight read once per tile, reused across the row block — the /// many-token path that a `gemv`-per-row would re-stream). Requires @@ -1662,7 +2138,7 @@ pub fn matmul(dev: &dyn Device, weight: &Tensor, input: &Tensor) -> Result(p); } #define CPA(SD, SP, V) asm volatile("cp.async.cg.shared.global [%0],[%1],16,%2;\n" \ :: "r"(SD), "l"(SP), "r"((V)?16:0)) // BM=64, BN=128: each block computes 128 N-cols so the loaded A tile is reused @@ -2020,14 +2497,14 @@ extern "C" __global__ void moe_grouped_gemm_mma( __syncthreads(); int buf = c&1; for (int ks=0; ks>(nib*4))&0xf) ^ 0x8) | 0x6400u); return __hmul(__hsub(__ushort_as_half(t), __float2half(1032.f)), sc); } +// Vectorized dequant: two adjacent nibbles (low,high of the byte at bit-offset sh) +// -> packed f16x2, SIMD affine. Bit-exact to pk2(dqn(n0),dqn(n1)). sh=(nib_even)*4. +__device__ __forceinline__ unsigned dq2(unsigned w, int sh, __half sc){ + unsigned two = (w >> sh) & 0xFFu; + unsigned packed = ((((two & 0xf0u) << 12) | (two & 0xfu)) ^ 0x00080008u) | 0x64006400u; + __half2 h = *reinterpret_cast<__half2*>(&packed); + h = __hmul2(__hsub2(h, __float2half2_rn(1032.0f)), __half2half2(sc)); + return *reinterpret_cast(&h); +} #define CPA(SD,SP) asm volatile("cp.async.cg.shared.global [%0],[%1],16,16;\n" :: "r"(SD), "l"(SP)) // BN=128 (NS=16): each block does 128 N-cols → the loaded A tile is reused over // 2x the weights, halving redundant A reads (the bottleneck). Weight/scale loads @@ -2138,10 +2624,18 @@ extern "C" __global__ void moe_q4_grouped_mma( const __half* __restrict__ Sc, const int* __restrict__ tok0_arr, const int* __restrict__ eid_arr, const int* __restrict__ gend_arr, __half* __restrict__ Out, - int N, int K) + int N, int K, int band_n) { - const int BM=64,BN=128,BK=64,NS=16,NKS=4,GS=32,WPR=BK/8,GPR=BK/GS; - int bx=blockIdx.x, by=blockIdx.y; + const int BM=64,BN=128,BK=32,NS=16,NKS=2,GS=32,WPR=BK/8,GPR=BK/GS; // BK32 occ-tune (smem 25->12.5KB) + // 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 (BUTTER_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]; @@ -2151,7 +2645,7 @@ extern "C" __global__ void moe_q4_grouped_mma( int nchunk=K/BK, WPB=K/32; long ebase=(long)eid*N; #define LOADC(C,BUF) do{ int kc=(C)*BK; \ - for(int e=tid;e bool { + std::env::var("BUTTER_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`. +/// +/// `BUTTER_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, @@ -2217,9 +2727,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. BUTTER_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; } @@ -2231,13 +2753,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) } @@ -2246,14 +2772,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` (BUTTER_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++; } @@ -2266,6 +2811,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. +/// +/// `BUTTER_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, @@ -2282,24 +2832,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) } @@ -2565,11 +3124,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#" @@ -2660,7 +3219,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] @@ -2768,7 +3327,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 @@ -3151,7 +3710,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 @@ -3373,7 +3932,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() { @@ -3432,7 +3991,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( @@ -3552,7 +4111,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; @@ -3640,7 +4199,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), @@ -3694,9 +4253,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, @@ -3714,6 +4274,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; @@ -3766,21 +4336,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", @@ -3789,19 +4396,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), @@ -3823,7 +4432,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 @@ -4013,7 +4622,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, @@ -4040,7 +4649,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]; @@ -4078,7 +4687,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, @@ -4112,7 +4721,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. @@ -4239,13 +4848,160 @@ pub fn sdpa_decode( 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 +// ── CUDA-graph-capturable dense decode SDPA (raw CUDA, head_dim 128) ──── +// +// The registry `iron_sdpa_decode` kernel above takes `n_kv` as a launch-time +// 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 iron registry kernel (invasive: `n_kv` there is a `#[constexpr]` +// used to size the unroll-friendly loop bound), this is a hand-written raw +// CUDA fallback that reads `n_kv` from a 1-element `u32` device buffer at +// kernel runtime instead. Same single-pass online-softmax algorithm as the +// registry kernel's dense path (`iron_sdpa_decode` with `sink_end=0, +// window_start=0, has_sink=0`): block per Q-head, 1024 threads (32 warps x +// 32 lanes), each lane owns 4 of the 128 head dims, `__shfl_xor_sync` +// butterfly reduction (broadcasts to every lane, matching MSL's +// all-to-all `simd_sum`/`simd_max`), f32 accumulation throughout. Verified +// against `sdpa_decode`'s output in `crates/backends/wh-butter-cuda/tests/laguna.rs`. +const SDPA_DECODE_NBUF_SRC: &str = r#" +extern "C" __global__ void butter_sdpa_decode_nbuf( + const float* __restrict__ q, const float* __restrict__ k, const float* __restrict__ v, + float* __restrict__ out, const unsigned int* __restrict__ n_kv_buf, + int head_dim, int kv_stride, int heads_per_group, float scale) +{ + int q_head = blockIdx.x; + int kv_head = q_head / heads_per_group; + int tid = threadIdx.x; + int lane = tid & 31; + int sg = tid >> 5; + int ns = blockDim.x >> 5; + unsigned int n_kv = n_kv_buf[0]; + + __shared__ float tg_max[32]; + __shared__ float tg_sum[32]; + __shared__ float tg_out0[1056]; + __shared__ float tg_out1[1056]; + __shared__ float tg_out2[1056]; + __shared__ float tg_out3[1056]; + + int q_off = q_head * head_dim; + long kv_head_base = (long)kv_head * kv_stride * head_dim; + int d0 = lane * 4; + + float qs[4]; + #pragma unroll + for (int i = 0; i < 4; i++) qs[i] = q[q_off + d0 + i] * scale; + + float run_max = -3.4028235e38f; + float run_sum = 0.0f; + float os[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + + for (unsigned int t = sg; t < n_kv; t += ns) { + long base = kv_head_base + (long)t * head_dim; + long kv0 = base + d0; + float partial = 0.0f; + #pragma unroll + for (int i = 0; i < 4; i++) partial += qs[i] * k[kv0 + i]; + for (int mask = 16; mask > 0; mask >>= 1) partial += __shfl_xor_sync(0xffffffffu, partial, mask); + float score = partial; + float new_max = fmaxf(score, run_max); + float factor = expf(run_max - new_max); + float weight = expf(score - new_max); + run_sum = run_sum * factor + weight; + run_max = new_max; + #pragma unroll + for (int i = 0; i < 4; i++) os[i] = os[i] * factor + weight * v[kv0 + i]; + } + + if (lane == 0) { tg_max[sg] = run_max; tg_sum[sg] = run_sum; } + __syncthreads(); + if (sg == 0) { + float g_max_raw = (lane < ns) ? tg_max[lane] : -3.4028235e38f; + float g_max = g_max_raw; + for (int mask = 16; mask > 0; mask >>= 1) g_max = fmaxf(g_max, __shfl_xor_sync(0xffffffffu, g_max, mask)); + float g_sum_in = (lane < ns) ? tg_sum[lane] * expf(g_max_raw - g_max) : 0.0f; + float g_sum = g_sum_in; + for (int mask = 16; mask > 0; mask >>= 1) g_sum += __shfl_xor_sync(0xffffffffu, g_sum, mask); + if (lane == 0) { tg_max[0] = g_max; tg_sum[0] = g_sum; } + } + __syncthreads(); + float g_max = tg_max[0]; + float g_sum = tg_sum[0]; + float rescale = (g_sum > 0.0f) ? expf(run_max - g_max) / g_sum : 0.0f; + int stride = ns + 1; + int idx = lane * stride + sg; + tg_out0[idx] = os[0] * rescale; + tg_out1[idx] = os[1] * rescale; + tg_out2[idx] = os[2] * rescale; + tg_out3[idx] = os[3] * rescale; + __syncthreads(); + if (sg == 0) { + float so0 = 0.0f, so1 = 0.0f, so2 = 0.0f, so3 = 0.0f; + for (int g = 0; g < ns; g++) { + int ri = lane * stride + g; + so0 += tg_out0[ri]; + so1 += tg_out1[ri]; + so2 += tg_out2[ri]; + so3 += tg_out3[ri]; + } + int out_off = q_off + d0; + out[out_off] = so0; + out[out_off + 1] = so1; + out[out_off + 2] = so2; + out[out_off + 3] = so3; + } +} +"#; + +/// 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 +/// 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. +#[allow(clippy::too_many_arguments)] +pub fn sdpa_decode_nbuf( + dev: &dyn Device, + q: &Tensor, + k: &Tensor, + v: &Tensor, + head_dim: usize, + n_kv_buf: &Tensor, + kv_stride: u32, + heads_per_group: u32, + scale: f32, +) -> Result { + if head_dim != 128 { + return Err(Error::Msg(format!("sdpa_decode_nbuf: only head_dim 128 is implemented, got {head_dim}"))); + } + let n_q_heads = q.elem_count() / head_dim; + let out = Tensor::empty(dev, vec![n_q_heads, head_dim], DType::F32)?; + let i = |x: i32| x.to_le_bytes().to_vec(); + let f = |x: f32| x.to_le_bytes().to_vec(); + dev.dispatch_raw_cuda( + SDPA_DECODE_NBUF_SRC, "sdpa_decode_nbuf.cu", "butter_sdpa_decode_nbuf", + &[ + (q.buffer.as_ref(), q.offset), + (k.buffer.as_ref(), k.offset), + (v.buffer.as_ref(), v.offset), + (out.buffer.as_ref(), 0), + (n_kv_buf.buffer.as_ref(), n_kv_buf.offset), + ], + &[i(head_dim as i32), i(kv_stride as i32), i(heads_per_group as i32), f(scale)], + [n_q_heads as u32, 1, 1], [1024, 1, 1], 0, false, + )?; + Ok(out) +} + +/// 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_butter_sdpa_decode_2pass_combined`, f32/f16/bf16). #[allow(clippy::too_many_arguments)] @@ -4412,7 +5168,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) @@ -4423,7 +5179,7 @@ pub fn silu(dev: &dyn Device, x: &Tensor) -> Result { elementwise_kernel(dev, "iron_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 `iron_relu2`. pub fn relu2(dev: &dyn Device, x: &Tensor) -> Result { @@ -4520,8 +5276,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 @@ -4579,6 +5335,59 @@ extern "C" __global__ void moe_scatter_add_det_f16( acc[t * hid + h] = sum; } } + +// f16-input fixed-top_k variant: rowpos maps token-major router slots +// (t*top_k+j) to expert-major sorted rows. This skips the token CSR rebuild. +extern "C" __global__ void moe_scatter_add_topk_det_f16( + const __half* __restrict__ dn, + const unsigned* __restrict__ rowpos, + const float* __restrict__ wt_token_major, + float* acc, + int s, int top_k, int hid, float unscale) +{ + int t = (int)blockIdx.x; + int h = (int)(blockIdx.y * BLOCK_H + threadIdx.x); + if (t < s && h < hid) { + float sum = 0.f; + int base = t * top_k; + for (int j = 0; j < top_k; ++j) { + unsigned r = rowpos[base + j]; + sum += __half2float(dn[(int)r * hid + h]) * (wt_token_major[base + j] * unscale); + } + acc[t * hid + h] = sum; + } +} + +extern "C" __global__ void moe_scatter_add_topk_det_f16_x2( + const __half* __restrict__ dn, + const unsigned* __restrict__ rowpos, + const float* __restrict__ wts_tok, + float* acc, + int s, int top_k, int hid, float unscale) +{ + int t = (int)blockIdx.x; + int h = ((int)blockIdx.y * BLOCK_H + threadIdx.x) * 2; + if (t < s && h < hid) { + float sum0 = 0.f; + float sum1 = 0.f; + int base = t * top_k; + #pragma unroll + for (int j = 0; j < 8; ++j) { + if (j < top_k) { + int r = (int)rowpos[base + j]; + float w = wts_tok[base + j] * unscale; + long idx = (long)r * hid + h; + __half2 hv = *reinterpret_cast(dn + idx); + float2 fv = __half22float2(hv); + sum0 += fv.x * w; + sum1 += fv.y * w; + } + } + long out = (long)t * hid + h; + acc[out] = sum0; + acc[out + 1] = sum1; + } +} "#; /// Deterministic MoE scatter-add (drop-in for [`moe_scatter_add`] on the prefill @@ -4586,7 +5395,7 @@ extern "C" __global__ void moe_scatter_add_det_f16( /// 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, @@ -4662,11 +5471,39 @@ extern "C" __global__ void moe_token_csr( /// runs the deterministic moe_scatter_add_det kernel. Same result as the host /// `moe_scatter_add_det` but host-sync-free → graph-capturable. #[allow(clippy::too_many_arguments)] -pub fn moe_scatter_add_det_dev( +/// x2 variant of moe_scatter_add_topk_det_f16: 2 hidden-dims/thread via __half2 +/// vectorized loads. Exact; falls back to the x1 kernel on odd hid. +pub fn moe_scatter_add_topk_det_f16_x2( dev: &dyn Device, dn: &Tensor, - sorted_tok: &Tensor, - wts: &Tensor, + rowpos: &Tensor, + wts_tok: &Tensor, + acc: &Tensor, + s: usize, + top_k: usize, + hid: usize, + unscale: f32, +) -> Result<()> { + if hid % 2 != 0 { + return moe_scatter_add_topk_det_f16(dev, dn, rowpos, wts_tok, acc, s, top_k, hid, unscale); + } + let i = |x: i32| x.to_le_bytes().to_vec(); + let pairs = (hid as u32).div_ceil(2); + dev.dispatch_raw_cuda( + MOE_SCATTER_ADD_DET_SRC, + "moe_scatter_add_det.cu", + "moe_scatter_add_topk_det_f16_x2", + &[(dn.buffer.as_ref(), 0), (rowpos.buffer.as_ref(), rowpos.offset), + (wts_tok.buffer.as_ref(), wts_tok.offset), (acc.buffer.as_ref(), 0)], + &[i(s as i32), i(top_k as i32), i(hid as i32), unscale.to_le_bytes().to_vec()], + [s as u32, pairs.div_ceil(128), 1], [128, 1, 1], 0, false) +} + +pub fn moe_scatter_add_det_dev( + dev: &dyn Device, + dn: &Tensor, + sorted_tok: &Tensor, + wts: &Tensor, acc: &Tensor, s: usize, mt: usize, @@ -4701,6 +5538,39 @@ pub fn moe_scatter_add_det_dev( Ok(()) } +#[allow(clippy::too_many_arguments)] +pub fn moe_scatter_add_topk_det_f16( + dev: &dyn Device, + dn: &Tensor, + rowpos: &Tensor, + wt_token_major: &Tensor, + acc: &Tensor, + s: usize, + top_k: usize, + hid: usize, + unscale: f32, +) -> Result<()> { + let i = |x: i32| x.to_le_bytes().to_vec(); + let f = |x: f32| x.to_le_bytes().to_vec(); + dev.dispatch_raw_cuda( + MOE_SCATTER_ADD_DET_SRC, + "moe_scatter_add_det.cu", + "moe_scatter_add_topk_det_f16", + &[ + (dn.buffer.as_ref(), 0), + (rowpos.buffer.as_ref(), rowpos.offset), + (wt_token_major.buffer.as_ref(), wt_token_major.offset), + (acc.buffer.as_ref(), 0), + ], + &[i(s as i32), i(top_k as i32), i(hid as i32), f(unscale)], + [s as u32, (hid as u32).div_ceil(128), 1], + [128, 1, 1], + 0, + false, + )?; + Ok(()) +} + /// Expand per-expert CSR offsets `[n_exp+1]` into a per-row expert-id vector /// `[mt]` on device: `idx[r] = e` for every `r` in `[off[e], off[e+1])`. /// Graph-safe replacement for the host `triples.map(|(e,_,_)| e)` build -- lets @@ -4736,6 +5606,38 @@ pub fn moe_expert_ids_from_offsets(dev: &dyn Device, off: &Tensor, n_exp: usize, /// scale, and cast-back-to-f16 that the MoE prefill loop previously did by /// downloading to host, computing on CPU, and re-uploading. One kernel replaces /// the full host-round-trip between the up and down GEMMs. +const RELU2_SCALE_BF16_SRC: &str = r#" +#include +extern "C" __global__ void relu2_scale_bf16( + const __nv_bfloat16* __restrict__ inp, + __nv_bfloat16* __restrict__ out, + int n, float scale) +{ + int i = (int)(blockIdx.x * blockDim.x + threadIdx.x); + if (i < n) { + float v = __bfloat162float(inp[i]); + if (v < 0.f) v = 0.f; + // bf16 range (~3e38) holds relu^2 for this model; no f16-style clamp needed. + out[i] = __float2bfloat16(v * v * scale); + } +} +"#; + +/// relu(x)^2 * scale for a BF16 tensor: reads BF16, squares in FP32, writes BF16. +/// Unlike the f16 variant there is no square-safe clamp (bf16 exponent range holds +/// the squared activations), so the FFN intermediate is not range-distorted. +pub fn relu2_scale_bf16(dev: &dyn Device, x: &Tensor, scale: f32) -> Result { + let n = x.elem_count(); + let out = Tensor::empty(dev, x.shape.clone(), DType::BF16)?; + let block = 256u32; + let grid = (n as u32).div_ceil(block); + dev.dispatch_raw_cuda(RELU2_SCALE_BF16_SRC, "relu2_scale_bf16.cu", "relu2_scale_bf16", + &[(x.buffer.as_ref(), x.offset), (out.buffer.as_ref(), 0)], + &[(n as i32).to_le_bytes().to_vec(), scale.to_le_bytes().to_vec()], + [grid, 1, 1], [block, 1, 1], 0, false)?; + Ok(out) +} + const RELU2_SCALE_F16_SRC: &str = r#" #include extern "C" __global__ void relu2_scale_f16( @@ -4810,7 +5712,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 `iron_fma_inplace`. @@ -4961,6 +5863,526 @@ pub fn dsv4_partial_rope( Ok(Tensor::new(qk.buffer.clone(), qk.shape.clone(), qk.dtype)) } +// ── Laguna YaRN partial-rotary RoPE (raw CUDA) ─────────────────────────── +// +// Laguna's full-attention layers rotate only the first `n_rot` (partial- +// rotary) dims of each head with NEOX pairing (`j`, `j+n_rot/2`) and a YaRN +// frequency/magnitude correction; the tail dims pass through unrotated. The +// sliding-window layers reuse the same kernel with `n_rot == head_dim`, +// `freq_scale = 1.0`, `ext_factor = 0.0`, `attn_factor = 1.0` (plain RoPE, +// no passthrough, no magnitude scaling). +// +// Ports the corr_dims + ramp + mscale math from the reference C++ engine's +// YaRN rope kernel (`ggml_rope_yarn_corr_dims` + `rope_yarn` in that +// engine's CUDA rope source) bit-for-bit: +// corr_dim(n_rotations) = n_rot * ln(n_orig_ctx / (n_rotations * 2π)) / (2 * ln(theta_base)) +// corr_low = max(0, floor(corr_dim(beta_fast))) +// corr_high = min(n_rot - 1, ceil(corr_dim(beta_slow))) +// ramp(j) = 1 - clamp((j - corr_low) / max(0.001, corr_high - corr_low), 0, 1) +// theta = lerp(freq_scale * theta_extrap, theta_extrap, ramp(j) * ext_factor) +// mscale = attn_factor * (ext_factor != 0 ? 1 + 0.1 * ln(1 / freq_scale) : 1) +// +// 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)` +// correction internally exactly as the reference engine does (the engine's +// 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 +// the correction a second time in caller code. +const ROPE_YARN_PARTIAL_SRC: &str = r#" +extern "C" __global__ void butter_rope_yarn_partial( + const float* __restrict__ qk, float* __restrict__ out, + int head_dim, int n_rot, int position, + float theta_base, float freq_scale, float ext_factor, float attn_factor, + float corr_low, float corr_high) +{ + int h = blockIdx.x; + int tid = threadIdx.y; + int half_rot = n_rot / 2; + const float* row_in = qk + (long)h * head_dim; + float* row_out = out + (long)h * head_dim; + + 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; + // attn_factor IS the final magnitude scale (mscale). The reference + // stack cancels its kernel's internal 1 + 0.1*ln(1/freq_scale) + // re-multiply by pre-dividing on the host, so the net scale equals the + // model's stored attention_factor; we just apply it directly. + 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 { + // Passthrough tail: copy the two contiguous elements at [n_rot + 2p, + // n_rot + 2p + 1) unrotated (matches the reference NEOX kernel's + // untouched-tail convention). + 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]; + } + } +} +"#; + +/// YaRN-scaled partial-rotary RoPE (NEOX pairing) for a `[n_heads, head_dim]` +/// 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 +/// unused (the ramp contributes zero once `ext_factor` is zero). f32 only. +#[allow(clippy::too_many_arguments)] +pub fn rope_yarn_partial( + dev: &dyn Device, + qk: &Tensor, + position: u32, + 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: scalar input".into()))?; + let n_heads = qk.elem_count() / head_dim; + if n_rot == 0 || n_rot as usize > head_dim || n_rot % 2 != 0 { + return Err(Error::Msg(format!( + "rope_yarn_partial: n_rot {n_rot} invalid for head_dim {head_dim}" + ))); + } + // ggml_rope_yarn_corr_dims: the [low, high] band (in rotated-pair-index + // units) where the rotary frequency ramps from interpolated (YaRN-scaled) + // to extrapolated (raw). Degenerate but harmless when ext_factor == 0. + 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_SRC, "rope_yarn_partial.cu", "butter_rope_yarn_partial", + &[(qk.buffer.as_ref(), qk.offset), (out.buffer.as_ref(), 0)], + &[ + i(head_dim as i32), i(n_rot as i32), i(position as i32), + f(theta_base), f(freq_scale), f(ext_factor), f(attn_factor), + f(corr_low), f(corr_high), + ], + [n_heads as u32, 1, 1], [1, half, 1], 0, false, + )?; + Ok(out) +} + +const ROPE_YARN_PARTIAL_POSBUF_SRC: &str = r#" +extern "C" __global__ void butter_rope_yarn_partial_posbuf( + const float* __restrict__ qk, float* __restrict__ out, + const unsigned int* __restrict__ pos_buf, + int head_dim, int n_rot, + float theta_base, float freq_scale, float ext_factor, float attn_factor, + float corr_low, float corr_high) +{ + int position = (int)pos_buf[0]; + int h = blockIdx.x; + int tid = threadIdx.y; + int half_rot = n_rot / 2; + const float* row_in = qk + (long)h * head_dim; + float* row_out = out + (long)h * head_dim; + + 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]; + } + } +} +"#; + +/// CUDA-graph-capturable variant of [`rope_yarn_partial`]: identical math, +/// but `position` is read from a 1-element `u32` device buffer at kernel +/// runtime instead of being baked in as a launch-time scalar. A captured +/// graph replays the exact same recorded launch every token, so any value +/// that changes per token (the sequence position) has to live behind a +/// stable pointer whose CONTENTS the caller updates before each replay, +/// rather than as an argument frozen at capture time. `corr_low`/`corr_high` +/// depend only on the (per-layer-type, capture-time-constant) RoPE constants, +/// so those stay ordinary scalars. Used by Laguna's decode path under +/// `BUTTER_LAGUNA_GRAPH=1`; the plain scalar-position `rope_yarn_partial` above +/// remains the default (eager, non-graph) path. +#[allow(clippy::too_many_arguments)] +pub fn rope_yarn_partial_posbuf( + dev: &dyn Device, + qk: &Tensor, + pos_buf: &Tensor, + 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_posbuf: scalar input".into()))?; + let n_heads = qk.elem_count() / head_dim; + if n_rot == 0 || n_rot as usize > head_dim || n_rot % 2 != 0 { + return Err(Error::Msg(format!( + "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 + // 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() + / (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_POSBUF_SRC, "rope_yarn_partial_posbuf.cu", "butter_rope_yarn_partial_posbuf", + &[ + (qk.buffer.as_ref(), qk.offset), + (out.buffer.as_ref(), 0), + (pos_buf.buffer.as_ref(), pos_buf.offset), + ], + &[ + i(head_dim as i32), i(n_rot as i32), + f(theta_base), f(freq_scale), f(ext_factor), f(attn_factor), + f(corr_low), f(corr_high), + ], + [n_heads as u32, 1, 1], [1, half, 1], 0, false, + )?; + Ok(out) +} + +const ROPE_YARN_PARTIAL_MANY_SRC: &str = r#" +extern "C" __global__ void butter_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", "butter_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 +/// (`BUTTER_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 +/// 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<()> { + dev.dispatch_raw_cuda( + WRITE_U32_SRC, "write_u32.cu", "butter_write_u32", + &[(buf.buffer.as_ref(), buf.offset)], + &[val.to_le_bytes().to_vec()], + [1, 1, 1], [1, 1, 1], 0, false, + ) +} + +const WRITE_U32_SRC: &str = r#" +extern "C" __global__ void butter_write_u32(unsigned int* __restrict__ buf, unsigned int val) +{ + buf[0] = val; +} +"#; + +// ── Laguna per-head softplus attention gate (raw CUDA) ─────────────────── +// +// `g` holds one raw (pre-activation) gate scalar per head, projected from +// 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`. +const GATE_SOFTPLUS_MUL_PERHEAD_SRC: &str = r#" +extern "C" __global__ void butter_gate_softplus_mul_perhead( + const float* __restrict__ attn, const float* __restrict__ g, + float* __restrict__ out, int head_dim, int n_heads) +{ + int h = blockIdx.x; + int i = threadIdx.x; + if (h >= n_heads || i >= head_dim) return; + float gv = g[h]; + // softplus(x) = log1p(exp(x)); linear for large x where exp(x) overflows. + float sp = gv > 20.0f ? gv : log1pf(expf(gv)); + out[h * head_dim + i] = attn[h * head_dim + i] * sp; +} +"#; + +/// Per-head softplus attention gate: `out[h, i] = attn[h, i] * softplus(g[h])` +/// for `attn` `[n_heads, head_dim]` and `g` `[n_heads]` (one raw gate scalar +/// per head, broadcast over `head_dim`). Laguna's XS.2/S.2 attention-output +/// gate, applied before `o_proj`. f32 only. +pub fn gate_softplus_mul_perhead(dev: &dyn Device, attn: &Tensor, g: &Tensor) -> Result { + let head_dim = *attn.shape.last().ok_or_else(|| Error::Msg("gate_softplus_mul_perhead: scalar attn".into()))?; + let n_heads = attn.elem_count() / head_dim; + if g.elem_count() != n_heads { + return Err(Error::Msg(format!( + "gate_softplus_mul_perhead: g has {} elems, expected {n_heads} (one per head)", + g.elem_count() + ))); + } + 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_SRC, "gate_softplus_mul_perhead.cu", "butter_gate_softplus_mul_perhead", + &[(attn.buffer.as_ref(), attn.offset), (g.buffer.as_ref(), g.offset), (out.buffer.as_ref(), 0)], + &[i(head_dim as i32), i(n_heads as i32)], + [n_heads as u32, 1, 1], [head_dim as u32, 1, 1], 0, false, + )?; + Ok(out) +} + +const GATE_SOFTPLUS_MUL_PERHEAD_MANY_SRC: &str = r#" +extern "C" __global__ void butter_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", "butter_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", "butter_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 butter_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 `iron_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 { @@ -5290,18 +6712,131 @@ 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. { let nt = (s * in_proj_out) as u32; - let b = if nt % 64 == 0 { 64u32 } else if nt % 32 == 0 { 32u32 } else { 1u32 }; + let b = if nt % 512 == 0 { 512u32 } else if nt % 256 == 0 { 256u32 } else if nt % 64 == 0 { 64u32 } else if nt % 32 == 0 { 32u32 } else { 1u32 }; Grid { grid: [nt / b, 1, 1], block: [b, 1, 1] } }, )?; Ok((z_out, xbc_out, dt_out)) } +const MAMBA_SPLIT_V4_SRC: &str = r#" +#include + +extern "C" __global__ void mamba_split_proj_v4( + const float* __restrict__ proj, + float* __restrict__ z_out, + float* __restrict__ xbc_out, + float* __restrict__ dt_out, + int in_proj_out4, int di4, int conv_dim4, int m_nh4, int total4) +{ + int idx = (int)(blockIdx.x * blockDim.x + threadIdx.x); + if (idx >= total4) return; + int row4 = di4 + conv_dim4 + m_nh4; + int ti = idx / row4; + int ci4 = idx - ti * row4; + const float4* __restrict__ src4 = reinterpret_cast(proj); + float4 v = src4[(size_t)ti * in_proj_out4 + ci4]; + if (ci4 < di4) { + reinterpret_cast(z_out)[(size_t)ti * di4 + ci4] = v; + } else if (ci4 < di4 + conv_dim4) { + int j = ci4 - di4; + reinterpret_cast(xbc_out)[(size_t)ti * conv_dim4 + j] = v; + } else { + int j = ci4 - di4 - conv_dim4; + reinterpret_cast(dt_out)[(size_t)ti * m_nh4 + j] = v; + } +} + +__device__ __forceinline__ float mamba_softplus(float x) { + return x > 20.0f ? x : logf(1.0f + expf(x)); +} + +extern "C" __global__ void mamba_split_proj_softplus_v4( + const float* __restrict__ proj, + const float* __restrict__ bias, + float* __restrict__ z_out, + float* __restrict__ xbc_out, + float* __restrict__ dt_out, + int in_proj_out4, int di4, int conv_dim4, int m_nh4, int total4) +{ + int idx = (int)(blockIdx.x * blockDim.x + threadIdx.x); + if (idx >= total4) return; + int row4 = di4 + conv_dim4 + m_nh4; + int ti = idx / row4; + int ci4 = idx - ti * row4; + const float4* __restrict__ src4 = reinterpret_cast(proj); + float4 v = src4[(size_t)ti * in_proj_out4 + ci4]; + if (ci4 < di4) { + reinterpret_cast(z_out)[(size_t)ti * di4 + ci4] = v; + } else if (ci4 < di4 + conv_dim4) { + int j = ci4 - di4; + reinterpret_cast(xbc_out)[(size_t)ti * conv_dim4 + j] = v; + } else { + int j = ci4 - di4 - conv_dim4; + int b = j << 2; + v.x = mamba_softplus(v.x + bias[b + 0]); + v.y = mamba_softplus(v.y + bias[b + 1]); + v.z = mamba_softplus(v.z + bias[b + 2]); + v.w = mamba_softplus(v.w + bias[b + 3]); + reinterpret_cast(dt_out)[(size_t)ti * m_nh4 + j] = v; + } +} +"#; + +/// Fused raw-float4 in_proj split (z/xbc/dt) WITH dt-bias+softplus folded into +/// the dt band. Replaces mamba_split_proj + softplus_add_rows. Escape with +/// NEMOTRON_MAMBA_SPLIT_RAW_OFF=1 or NEMOTRON_MAMBA_SPLIT_SOFTPLUS_OFF=1. +pub fn mamba_split_proj_softplus( + dev: &dyn Device, + proj: &Tensor, + bias: &Tensor, + s: usize, + in_proj_out: usize, + di: usize, + conv_dim: usize, + m_nh: usize, +) -> Result<(Tensor, Tensor, Tensor)> { + let raw_ok = std::env::var("NEMOTRON_MAMBA_SPLIT_RAW_OFF").is_err() + && std::env::var("NEMOTRON_MAMBA_SPLIT_SOFTPLUS_OFF").is_err() + && proj.dtype == DType::F32 && bias.dtype == DType::F32 + && proj.offset % 16 == 0 && bias.offset % 4 == 0 + && di % 4 == 0 && conv_dim % 4 == 0 && m_nh % 4 == 0 && in_proj_out % 4 == 0; + if raw_ok { + let z_out = Tensor::empty(dev, vec![s * di], DType::F32)?; + let xbc_out = Tensor::empty(dev, vec![s * conv_dim], DType::F32)?; + let dt_out = Tensor::empty(dev, vec![s * m_nh], DType::F32)?; + let i = |v: i32| v.to_le_bytes().to_vec(); + let total4 = (s * (di + conv_dim + m_nh) / 4) as u32; + if dev.dispatch_raw_cuda( + MAMBA_SPLIT_V4_SRC, + "mamba_split_v4.cu", + "mamba_split_proj_softplus_v4", + &[ + (proj.buffer.as_ref(), proj.offset), + (bias.buffer.as_ref(), bias.offset), + (z_out.buffer.as_ref(), 0), + (xbc_out.buffer.as_ref(), 0), + (dt_out.buffer.as_ref(), 0), + ], + &[i((in_proj_out / 4) as i32), i((di / 4) as i32), i((conv_dim / 4) as i32), i((m_nh / 4) as i32), i(total4 as i32)], + [total4.div_ceil(256), 1, 1], + [256, 1, 1], + 0, + false, + ).is_ok() { + return Ok((z_out, xbc_out, dt_out)); + } + } + let (z, xbc, dt_raw) = mamba_split_proj(dev, proj, s, in_proj_out, di, conv_dim, m_nh)?; + let dt = softplus_add_rows(dev, &dt_raw, bias, s, m_nh)?; + Ok((z, xbc, dt)) +} + // ── Fused Mamba conv-output split (Part-B: 3 strided_col_copy → 1 dispatch) ─ // @@ -5348,6 +6883,35 @@ pub fn mamba_split_conv( // Processes S tokens in one dispatch. Each thread-group handles one (token, // norm-group) pair. Replaces the host dl+compute+upload loop in the // CONV_DEVICE prefill path (TODO comment at line ~1543 of modeltests/lib.rs). +/// RMSNorm with f16 output only. Equivalent to `cast_f32_f16(rms_norm(...))`, +/// but avoids the intermediate f32 norm tensor and the separate cast launch. +/// TODO(rebrand-merge): the fast-path fused kernel (`iron_rms_norm_f16out`) +/// isn't in `thewafflehaus/iron@dev` yet, so this always takes the +/// `rms_norm` + `cast_f32_f16` fallback below (same numerics, one extra +/// dispatch) instead of the single-pass fused kernel. +pub fn rms_norm_f16out(dev: &dyn Device, x: &Tensor, weight: &Tensor, eps: f32) -> Result { + let out = rms_norm(dev, x, weight, eps)?; + cast_f32_f16(dev, &out) +} + +/// TODO(rebrand-merge): the fused kernel (`iron_gated_group_rmsnorm_batched_f16out`) +/// isn't in `thewafflehaus/iron@dev` yet (only the plain, non-f16out +/// `iron_gated_group_rmsnorm_batched` is) — always takes the +/// `gated_group_rmsnorm_batched` + `cast_f32_f16` fallback below. +pub fn gated_group_rmsnorm_batched_f16out( + dev: &dyn Device, + y: &Tensor, + z: &Tensor, + w: &Tensor, + eps: f32, + s: usize, + di: usize, + gs: usize, +) -> Result { + let out = gated_group_rmsnorm_batched(dev, y, z, w, eps, s, di, gs)?; + cast_f32_f16(dev, &out) +} + pub fn gated_group_rmsnorm_batched( dev: &dyn Device, y: &Tensor, @@ -5449,14 +7013,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]; @@ -5617,8 +7181,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) @@ -5770,9 +7334,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]; @@ -6187,7 +7751,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 @@ -6274,7 +7838,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`). @@ -6328,7 +7892,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, @@ -6364,10 +7928,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; @@ -6399,7 +7963,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 } @@ -6514,7 +8078,7 @@ 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 IRON_CUDA_ARCH= @@ -6524,7 +8088,7 @@ pub fn moe_grouped_gemm( 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; @@ -6758,7 +8322,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 butter_cp16(void* sm, const void* gm){ @@ -6778,7 +8342,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] @@ -6864,7 +8428,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, @@ -6935,9 +8499,9 @@ extern "C" __global__ void moe_extend_groups( // FP8 e4m3 per-tensor quant. Pass 1 = nvfp4_amax_slab (reused, mode 0). // fp8_scale: amax bits -> scale = amax/448 and inv (one thread). extern "C" __global__ void fp8_scale( - const unsigned* __restrict__ amax_bits, float* __restrict__ sc) + const unsigned* __restrict__ amax_bits, float* __restrict__ sc, float clip) { - float am=__uint_as_float(*amax_bits); + float am=__uint_as_float(*amax_bits)*((clip>0.f)?clip:1.f); float s=(am>0.f)?(am/448.f):1.f; sc[0]=s; // dequant scale (consumed by the GEMM) sc[1]=1.f/s; // quant scale (consumed by fp8_quant) @@ -7002,6 +8566,61 @@ extern "C" __global__ void fp8_rowcol_scale_f32( for(;e= rows * (long)nb) return; + int t = threadIdx.x; + __shared__ float s[128]; + long base = bid * 128; + s[t] = __half2float(x[base + t]); + __syncthreads(); + for (int len = 1; len < 128; len <<= 1) { + float a = s[t]; float b = s[t ^ len]; + __syncthreads(); + s[t] = (t & len) ? (b - a) : (a + b); + __syncthreads(); + } + float r = s[t] * 0.08838834764831845f; + y[base + t] = __float2half(r); + // block-reduce |r| then one atomicMax to the global amax + s[t] = fabsf(r); + __syncthreads(); + for (int o = 64; o > 0; o >>= 1) { if (t < o) s[t] = fmaxf(s[t], s[t + o]); __syncthreads(); } + if (t == 0) atomicMax(amax, __float_as_uint(s[0])); +} + +// Block Walsh-Hadamard transform (size 128) along the K dim, normalized by +// 1/sqrt(128). Each CUDA block handles one (row, kblock) of 128 contiguous +// elems. Orthonormal: rotating both activation and weight along K leaves the +// GEMM x.Wt unchanged, but spreads per-channel activation outliers so a +// per-tensor fp8 quant holds quality (QuaRot/Hadamard recipe). K % 128 == 0. +extern "C" __global__ void fwht128( + const __half* __restrict__ x, __half* __restrict__ y, long rows, int K) +{ + int nb = K / 128; + long bid = (long)blockIdx.x; + if (bid >= rows * (long)nb) return; + int t = threadIdx.x; // 0..127 + __shared__ float s[128]; + long base = bid * 128; + s[t] = __half2float(x[base + t]); + __syncthreads(); + for (int len = 1; len < 128; len <<= 1) { + float a = s[t]; float b = s[t ^ len]; + __syncthreads(); + s[t] = (t & len) ? (b - a) : (a + b); + __syncthreads(); + } + y[base + t] = __float2half(s[t] * 0.08838834764831845f); // 1/sqrt(128) +} + extern "C" __global__ void fp8_quant( const __half* __restrict__ x, unsigned char* __restrict__ q, const float* __restrict__ sc, long n) @@ -7030,7 +8649,52 @@ extern "C" __global__ void fp8_quant( } } -// f32-input variant of lt_fp4_quant — reads the residual-stream f32 tensor +// Pass-1 of the fused fp8 act path: block-Hadamard(128) along K + global amax, +// WITHOUT materializing the rotated tensor (no y write). Halves the HBM of the +// old fwht128_amax (which wrote a full [rows,K] f16 that fp8_quant then re-read). +extern "C" __global__ void fwht128_amaxonly( + const __half* __restrict__ x, unsigned* __restrict__ amax, long rows, int K) +{ + int nb = K / 128; long bid = (long)blockIdx.x; + if (bid >= rows * (long)nb) return; + int t = threadIdx.x; __shared__ float s[128]; long base = bid * 128; + s[t] = __half2float(x[base + t]); __syncthreads(); + for (int len = 1; len < 128; len <<= 1) { + float a = s[t]; float b = s[t ^ len]; __syncthreads(); + s[t] = (t & len) ? (b - a) : (a + b); __syncthreads(); + } + float r = s[t] * 0.08838834764831845f; + s[t] = fabsf(r); __syncthreads(); + for (int o = 64; o > 0; o >>= 1) { if (t < o) s[t] = fmaxf(s[t], s[t + o]); __syncthreads(); } + if (t == 0) atomicMax(amax, __float_as_uint(s[0])); +} + +// Pass-2 of the fused fp8 act path: re-rotate x (Hadamard 128) AND quant to e4m3 +// in one kernel reading x directly (no rotated-tensor round-trip). Uses the +// precomputed scale sc[1]. Same e4m3 encoding as fp8_quant. +extern "C" __global__ void fwht128_quant( + const __half* __restrict__ x, unsigned char* __restrict__ q, + const float* __restrict__ sc, long rows, int K) +{ + int nb = K / 128; long bid = (long)blockIdx.x; + if (bid >= rows * (long)nb) return; + int t = threadIdx.x; __shared__ float s[128]; long base = bid * 128; + s[t] = __half2float(x[base + t]); __syncthreads(); + for (int len = 1; len < 128; len <<= 1) { + float a = s[t]; float b = s[t ^ len]; __syncthreads(); + s[t] = (t & len) ? (b - a) : (a + b); __syncthreads(); + } + float v = s[t] * 0.08838834764831845f * sc[1]; + v = fminf(fmaxf(v, -448.f), 448.f); + unsigned sgn=(v<0.f)?0x80u:0u; float a=fabsf(v); unsigned c; + if(a<0.0009765625f){ int m=(int)rintf(a*512.f); if(m>7)m=7; c=(unsigned)m; } + else { int e; float fr=frexpf(a,&e); int E=e+6; int m=(int)rintf((2.f*fr-1.f)*8.f); + if(m>7){m=0;E+=1;} if(E<1){ int ms=(int)rintf(a*512.f); if(ms>7)ms=7; c=(unsigned)ms; } + else { if(E>15)E=15; if(E==15&&m>6)m=6; c=(unsigned)((E<<3)|m); } } + q[base + t]=(unsigned char)(sgn|c); +} + +// 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). @@ -7108,7 +8772,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)> { @@ -7456,7 +9120,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 @@ -7489,7 +9153,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 { @@ -7506,7 +9170,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( @@ -7566,12 +9230,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. @@ -7699,6 +9363,54 @@ pub fn w4a8_packw( /// `sfb_exp_bytes` = per-expert SFB stride in bytes (from `w8a8_packw`). #[allow(clippy::too_many_arguments)] +/// Validated NVFP4 (f16-act) grouped-MoE GEMM (marlin). `a` f16 `[m,K]`, `b` fp4 +/// marlin-tile u32 `[n_exp,K/16,N*2]`, `b_scales` S0E5M3 fp8 (u32-viewed), +/// `global_scale` f32 `[n_exp]`, `topk_weights` f32 `[m*top_k]`; routing +/// (`sorted_token_ids`/`expert_ids`/`num_tokens_past_padded`) device i32. +/// Returns f16 `[m*top_k, N]`. mul_topk_weights + fp32 reduce on. +#[allow(clippy::too_many_arguments)] +pub fn moe_marlin_gemm( + dev: &dyn Device, + a: &Tensor, b: &Tensor, b_scales: &Tensor, global_scale: &Tensor, topk_weights: &Tensor, + sorted_token_ids: &Tensor, expert_ids: &Tensor, num_tokens_past_padded: &Tensor, + prob_m: usize, prob_n: usize, prob_k: usize, num_experts: usize, top_k: usize, + group_size: usize, moe_block: usize, sms: i32, +) -> Result { + let mt = (prob_m * top_k).max(1); + let out = Tensor::empty(dev, vec![mt, prob_n.max(1)], DType::F16)?; + let num_groups = prob_k / group_size; + let ctmp = ((sms as usize) * 4 * moe_block * 256).max(1); + let c_tmp = Tensor::new(dev.alloc_zeroed(ctmp * 4)?, vec![ctmp], DType::F32); + let wsn = ((sms as usize) * 4).max(1); + let ws = Tensor::new(dev.alloc_zeroed(wsn * 4)?, vec![wsn], DType::U32); + dev.moe_marlin_gemm( + a.buffer.as_ref(), b.buffer.as_ref(), out.buffer.as_ref(), c_tmp.buffer.as_ref(), + b_scales.buffer.as_ref(), global_scale.buffer.as_ref(), topk_weights.buffer.as_ref(), + sorted_token_ids.buffer.as_ref(), expert_ids.buffer.as_ref(), num_tokens_past_padded.buffer.as_ref(), + ws.buffer.as_ref(), + moe_block, num_experts, top_k, true, prob_m, prob_n, prob_k, num_groups, group_size, sms, true, + )?; + Ok(out) +} + +/// Repack one expert's GPTQ-packed `[K/8,N]` u32 weights -> marlin tile `[K/16,N*2]` u32 (device). +pub fn marlin_repack(dev: &dyn Device, b_q: &Tensor, size_k: usize, size_n: usize, sms: i32, max_shared_mem: i32) -> Result { + let out_elems = ((size_k / 16) * (size_n * 2)).max(1); + let out = Tensor::new(dev.alloc(out_elems * 4)?, vec![out_elems], DType::U32); + dev.marlin_repack(b_q.buffer.as_ref(), out.buffer.as_ref(), size_k, size_n, sms, max_shared_mem)?; + Ok(out) +} + +/// On-device marlin routing build from per-expert offsets `off` [n_exp+1] (device i32). +pub fn marlin_build_routing(dev:&dyn Device, off:&Tensor, n_exp:usize, blk:usize, mt:usize)->Result<(Tensor,Tensor,Tensor)>{ + let max_blocks = mt/blk + n_exp + 1; + let stid=Tensor::new(dev.alloc((max_blocks*blk).max(1)*4)?, vec![(max_blocks*blk).max(1)], DType::U32); + let eid=Tensor::new(dev.alloc(max_blocks.max(1)*4)?, vec![max_blocks.max(1)], DType::U32); + let ntpp=Tensor::new(dev.alloc_zeroed(4)?, vec![1], DType::U32); + dev.marlin_build_routing(off.buffer.as_ref(), n_exp, blk, mt, stid.buffer.as_ref(), eid.buffer.as_ref(), ntpp.buffer.as_ref())?; + Ok((stid,eid,ntpp)) +} + pub fn moe_grouped_gemm_w8a8( dev: &dyn Device, a_pack: &Tensor, a_sf: &Tensor, sfa_off: &[i64], @@ -7754,9 +9466,9 @@ 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) -> Result<(Tensor, Tensor)> { +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(); let l = |v: i64| v.to_le_bytes().to_vec(); let amax = Tensor::new(dev.alloc_zeroed(4)?, vec![1], DType::U32); @@ -7769,7 +9481,7 @@ pub fn fp8_quant(dev: &dyn Device, x: &Tensor, n: usize) -> Result<(Tensor, Tens dev.dispatch_raw_cuda( MOE_FP4_SRC, "moe_fp4.cu", "fp8_scale", &[(amax.buffer.as_ref(), 0), (sc.buffer.as_ref(), 0)], - &[], [1, 1, 1], [1, 1, 1], 0, false)?; + &[clip.to_le_bytes().to_vec()], [1, 1, 1], [1, 1, 1], 0, false)?; let q = Tensor::new(dev.alloc(n)?, vec![n], DType::U32); dev.dispatch_raw_cuda( MOE_FP4_SRC, "moe_fp4.cu", "fp8_quant", @@ -7779,6 +9491,87 @@ pub fn fp8_quant(dev: &dyn Device, x: &Tensor, n: usize) -> Result<(Tensor, Tens Ok((q, sc)) } +/// Fused block-Hadamard(128) along K + global amax. Returns `(rotated [rows,k] +/// f16, amax [1] u32 bits)` for a subsequent fp8 quant that skips its own amax. +pub fn fwht128_amax(dev: &dyn Device, x: &Tensor, rows: usize, k: usize) -> Result<(Tensor, Tensor)> { + if k % 128 != 0 { return Err(Error::Msg(format!("fwht128_amax: k {k} %128 != 0"))); } + let y = Tensor::new(dev.alloc(rows * k * 2)?, vec![rows * k], DType::F16); + let amax = Tensor::new(dev.alloc_zeroed(4)?, vec![1], DType::U32); + let nb = (rows * (k / 128)) as u32; + let i = |v: i32| v.to_le_bytes().to_vec(); + let l = |v: i64| v.to_le_bytes().to_vec(); + dev.dispatch_raw_cuda(MOE_FP4_SRC, "moe_fp4.cu", "fwht128_amax", + &[(x.buffer.as_ref(), x.offset), (y.buffer.as_ref(), 0), (amax.buffer.as_ref(), 0)], + &[l(rows as i64), i(k as i32)], + [nb, 1, 1], [128, 1, 1], 0, false)?; + Ok((y.reshaped(vec![rows, k]), amax)) +} + +/// Per-tensor fp8 quant using a PRECOMPUTED amax (from `fwht128_amax`), with +/// optional clip. Returns `(q bytes, sc [2] f32)`. Skips the amax pass. +pub fn fp8_quant_pre(dev: &dyn Device, x: &Tensor, n: usize, amax: &Tensor, clip: f32) -> Result<(Tensor, Tensor)> { + let i = |v: i32| v.to_le_bytes().to_vec(); + let l = |v: i64| v.to_le_bytes().to_vec(); + let sc = Tensor::new(dev.alloc(8)?, vec![2], DType::F32); + dev.dispatch_raw_cuda(MOE_FP4_SRC, "moe_fp4.cu", "fp8_scale", + &[(amax.buffer.as_ref(), 0), (sc.buffer.as_ref(), 0)], + &[clip.to_le_bytes().to_vec()], [1, 1, 1], [1, 1, 1], 0, false)?; + let q = Tensor::new(dev.alloc(n)?, vec![n], DType::U32); + dev.dispatch_raw_cuda(MOE_FP4_SRC, "moe_fp4.cu", "fp8_quant", + &[(x.buffer.as_ref(), x.offset), (q.buffer.as_ref(), 0), (sc.buffer.as_ref(), 0)], + &[l(n as i64)], + [(n as u32).div_ceil(256 * 4).max(1), 1, 1], [256, 1, 1], 0, false)?; + Ok((q, sc)) +} + +/// Fused-path pass 1: block-Hadamard(128) along K + global amax, NO rotated +/// output (saves the [rows,k] f16 write+read vs `fwht128_amax`). Returns amax bits. +pub fn fwht128_amaxonly(dev: &dyn Device, x: &Tensor, rows: usize, k: usize) -> Result { + if k % 128 != 0 { return Err(Error::Msg(format!("fwht128_amaxonly: k {k} %128 != 0"))); } + let amax = Tensor::new(dev.alloc_zeroed(4)?, vec![1], DType::U32); + let nb = (rows * (k / 128)) as u32; + dev.dispatch_raw_cuda(MOE_FP4_SRC, "moe_fp4.cu", "fwht128_amaxonly", + &[(x.buffer.as_ref(), x.offset), (amax.buffer.as_ref(), 0)], + &[(rows as i64).to_le_bytes().to_vec(), (k as i32).to_le_bytes().to_vec()], + [nb, 1, 1], [128, 1, 1], 0, false)?; + Ok(amax) +} + +/// Fused-path pass 2: re-rotate `x` (Hadamard 128 along K) AND quant to e4m3 in +/// one kernel reading `x` directly (no rotated-tensor round-trip), using a +/// precomputed amax. Returns `(q bytes, sc [2] f32)`. +pub fn fwht128_quant_pre(dev: &dyn Device, x: &Tensor, rows: usize, k: usize, amax: &Tensor, clip: f32) -> Result<(Tensor, Tensor)> { + if k % 128 != 0 { return Err(Error::Msg(format!("fwht128_quant_pre: k {k} %128 != 0"))); } + let sc = Tensor::new(dev.alloc(8)?, vec![2], DType::F32); + dev.dispatch_raw_cuda(MOE_FP4_SRC, "moe_fp4.cu", "fp8_scale", + &[(amax.buffer.as_ref(), 0), (sc.buffer.as_ref(), 0)], + &[clip.to_le_bytes().to_vec()], [1, 1, 1], [1, 1, 1], 0, false)?; + let n = rows * k; + let q = Tensor::new(dev.alloc(n)?, vec![n], DType::U32); + let nb = (rows * (k / 128)) as u32; + dev.dispatch_raw_cuda(MOE_FP4_SRC, "moe_fp4.cu", "fwht128_quant", + &[(x.buffer.as_ref(), x.offset), (q.buffer.as_ref(), 0), (sc.buffer.as_ref(), 0)], + &[(rows as i64).to_le_bytes().to_vec(), (k as i32).to_le_bytes().to_vec()], + [nb, 1, 1], [128, 1, 1], 0, false)?; + Ok((q, sc)) +} + +/// Block Walsh-Hadamard transform (128) along K of an f16 `[rows, k]` tensor. +/// k must be a multiple of 128. Orthonormal; used to spread activation outliers +/// before a per-tensor fp8 quant (and applied identically to the weight offline). +pub fn fwht128(dev: &dyn Device, x: &Tensor, rows: usize, k: usize) -> Result { + if k % 128 != 0 { return Err(Error::Msg(format!("fwht128: k {k} %128 != 0"))); } + let y = Tensor::new(dev.alloc(rows * k * 2)?, vec![rows * k], DType::F16); + let nb = (rows * (k / 128)) as u32; + let i = |v: i32| v.to_le_bytes().to_vec(); + let l = |v: i64| v.to_le_bytes().to_vec(); + dev.dispatch_raw_cuda(MOE_FP4_SRC, "moe_fp4.cu", "fwht128", + &[(x.buffer.as_ref(), x.offset), (y.buffer.as_ref(), 0)], + &[l(rows as i64), i(k as i32)], + [nb, 1, 1], [128, 1, 1], 0, false)?; + Ok(y.reshaped(vec![rows, k])) +} + /// FP8 e4m3 dense GEMM with per-tensor scales (vendor-recipe precision for /// the Mamba in/out projections + shared expert on this model family). #[allow(clippy::too_many_arguments)] @@ -7797,7 +9590,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(); @@ -7811,7 +9604,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)] @@ -7822,7 +9615,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( @@ -7924,7 +9717,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"))); } @@ -7955,7 +9748,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)> { @@ -7993,7 +9786,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)] @@ -8054,7 +9847,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; @@ -8075,7 +9868,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, @@ -8111,3 +9904,39 @@ pub fn add_rms_norm( )?; Ok((residual, normed)) } + +#[cfg(test)] +mod marlin_prep_validate { + use super::*; + use std::fs; + fn rf32(p:&str)->Vec{let b=fs::read(p).unwrap(); b.chunks(4).map(|c|f32::from_le_bytes([c[0],c[1],c[2],c[3]])).collect()} + #[test] + fn marlin_weight_prep_matches_python(){ + // Reference fixtures come from env, never a hardcoded path; unset -> quiet skip. + let Ok(dir)=std::env::var("BUTTER_MARLIN_FIXTURE_DIR") else { + println!("[skipped] BUTTER_MARLIN_FIXTURE_DIR unset"); + return; + }; + let (n,k)=(1856usize,2688usize); let ng=k/16; let kp=k/8; + let w0=rf32(&format!("{dir}/W0.bin")); assert_eq!(w0.len(), n*k); + let (codes,uesc,global)=quantize_nvfp4(&w0,n,k); + // codes[N,kp] -> transpose [kp,N] + let mut qw=vec![0u32; kp*n]; + for r in 0..n { for c in 0..kp { qw[c*n+r]=codes[r*kp+c]; } } + let qref=fs::read(format!("{dir}/qweight.bin")).unwrap(); + let qref0:Vec=qref[..kp*n*4].chunks(4).map(|c|u32::from_le_bytes([c[0],c[1],c[2],c[3]])).collect(); + let qm=qw.iter().zip(&qref0).filter(|(a,b)|a==b).count(); + // scales: dec_ue4m3[N,ng] -> values, transpose [ng,N], s0e5m3 + let mut sv=vec![0f32; ng*n]; + for r in 0..n { for g in 0..ng { sv[g*n+r]=dec_ue4m3(uesc[r*ng+g]); } } + let s0=marlin_s0e5m3_scales(&sv,ng,n); + let bref=fs::read(format!("{dir}/bscales.bin")).unwrap(); + let bm=s0.iter().zip(&bref[..ng*n]).filter(|(a,b)|a==b).count(); + let gref=rf32(&format!("{dir}/global.bin")); + let grust=global*128.0; + eprintln!("QW match {}/{} BSCALE match {}/{} global rust {} py {} (ratio {})", qm, kp*n, bm, ng*n, grust, gref[0], grust/gref[0]); + assert!(qm==kp*n, "qweight mismatch"); + assert!(bm>=(ng*n*999/1000), "bscale mismatch"); + assert!((grust/gref[0]-1.0).abs()<0.02, "global mismatch"); + } +}