From 4654433116dead8be3cc4b9a399ea47df1d1b3db Mon Sep 17 00:00:00 2001 From: Tom Turney Date: Wed, 22 Jul 2026 22:42:45 -0500 Subject: [PATCH 1/2] sync(rust): land in-flight CUDA engine work from the GB10 working tree Carries the uncommitted engine state the GB10 box has been running: CUDA-graph capture trait plumbing (begin/end capture, graph_launch), fp8 projection microbench and MoE grouped-MMA test updates, and the device trait additions they depend on. Precedes the Laguna port commits that build on these interfaces. --- rust/crates/backends/ffai-cuda/Cargo.toml | 4 + rust/crates/backends/ffai-cuda/src/imp.rs | 36 ++++ .../ffai-cuda/tests/moe_grouped_mma_test.rs | 179 +++++++++++++++++- .../ffai-cuda/tests/proj_fp8_bench.rs | 4 +- rust/crates/ffai-core/src/lib.rs | 28 +++ 5 files changed, 246 insertions(+), 5 deletions(-) diff --git a/rust/crates/backends/ffai-cuda/Cargo.toml b/rust/crates/backends/ffai-cuda/Cargo.toml index 33cc80f0..258aa8bd 100644 --- a/rust/crates/backends/ffai-cuda/Cargo.toml +++ b/rust/crates/backends/ffai-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/ffai-cuda/src/imp.rs b/rust/crates/backends/ffai-cuda/src/imp.rs index 61806fa1..7a1360f0 100644 --- a/rust/crates/backends/ffai-cuda/src/imp.rs +++ b/rust/crates/backends/ffai-cuda/src/imp.rs @@ -412,6 +412,42 @@ impl Device for CudaDevice { let ob = out.as_any().downcast_ref::().ok_or_else(|| Error::Msg("moe_grouped_cutlass: out not CudaBuffer".into()))?; self.dev.moe_grouped_cutlass(ab.ptr, wb.ptr, ob.ptr, group_rows, expert_ids, n, k).map_err(dispatch_err) } + #[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<()> { + let g = |x: &dyn DeviceBuffer, nm: &str| x.as_any().downcast_ref::().map(|b| b.ptr) + .ok_or_else(|| Error::Msg(format!("moe_marlin_gemm: {nm} not CudaBuffer"))); + self.dev.moe_marlin_gemm( + g(a,"a")?, g(b,"b")?, g(c,"c")?, g(c_tmp,"c_tmp")?, + g(b_scales,"b_scales")?, g(global_scale,"global_scale")?, g(topk_weights,"topk_weights")?, + g(sorted_token_ids,"sorted_token_ids")?, g(expert_ids,"expert_ids")?, g(num_tokens_past_padded,"ntpp")?, + g(workspace,"workspace")?, + moe_block, num_experts, top_k, mul_topk_weights, prob_m, prob_n, prob_k, num_groups, group_size, sms, use_fp32_reduce, + ).map_err(dispatch_err) + } + + #[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<()> { + let g = |x: &dyn DeviceBuffer, nm: &str| x.as_any().downcast_ref::().map(|b| b.ptr) + .ok_or_else(|| Error::Msg(format!("marlin_repack: {nm} not CudaBuffer"))); + self.dev.marlin_repack(g(b_q,"b_q")?, g(out,"out")?, size_k, size_n, sms, max_shared_mem).map_err(dispatch_err) + } + #[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<()>{ + let gg=|x:&dyn DeviceBuffer,nm:&str| x.as_any().downcast_ref::().map(|b|b.ptr).ok_or_else(||Error::Msg(format!("marlin_build_routing: {nm} not CudaBuffer"))); + self.dev.marlin_build_routing(gg(off,"off")?,n_exp,blk,mt,gg(stid,"stid")?,gg(eid,"eid")?,gg(ntpp,"ntpp")?).map_err(dispatch_err) + } fn moe_grouped_cutlass_fp4( &self, diff --git a/rust/crates/backends/ffai-cuda/tests/moe_grouped_mma_test.rs b/rust/crates/backends/ffai-cuda/tests/moe_grouped_mma_test.rs index 6454afd6..a4bd5321 100644 --- a/rust/crates/backends/ffai-cuda/tests/moe_grouped_mma_test.rs +++ b/rust/crates/backends/ffai-cuda/tests/moe_grouped_mma_test.rs @@ -6,10 +6,10 @@ use ffai_cuda::CudaDevice; use ffai_core::{DType, Device, Tensor}; -use ffai_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 ffai_ops::{gemm_cublas_f32out, fwht128, gemm_fp8, fp8_quant, moe_q4_grouped_mma_dev, quantize_q4, moe_grouped_gemm_mma, moe_grouped_gemm_cutlass, cast_f32_f16, cast_f16_f32, lt_fp4_quant_g, lt_fp4_quant_g_46, dec_ue4m3, 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]; @@ -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,176 @@ 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)"); +} + +#[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/ffai-cuda/tests/proj_fp8_bench.rs b/rust/crates/backends/ffai-cuda/tests/proj_fp8_bench.rs index 1fc70b01..8a632bbf 100644 --- a/rust/crates/backends/ffai-cuda/tests/proj_fp8_bench.rs +++ b/rust/crates/backends/ffai-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/ffai-core/src/lib.rs b/rust/crates/ffai-core/src/lib.rs index 5defce5c..56721a48 100644 --- a/rust/crates/ffai-core/src/lib.rs +++ b/rust/crates/ffai-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 + From d81f16d8b344ffdfa05999f3db5b9efe970684f8 Mon Sep 17 00:00:00 2001 From: Tom Turney Date: Wed, 22 Jul 2026 22:43:02 -0500 Subject: [PATCH 2/2] feat(laguna): Laguna-S-2.1 decode on the CUDA engine Full single-stream decode for the 117.55B/8.14B-active hybrid MoE: sigmoid top-10 router with score-correction bias and shared expert, per-head softplus attention gate, per-head QK RMSNorm, period-4 full/sliding-window attention with a 512-slot ring KV cache, YaRN partial-rotary rope on full layers (scaling read from GGUF metadata, this export is a 256K/factor-32 checkpoint), Q8 LM head, and a Q4-requantized weight load straight from GGUF. New ops: rope_yarn_partial (+posbuf variant), gate_softplus_mul_perhead, argmax_f32_device, sdpa_decode_nbuf (n_kv from a device buffer), write_u32, moe_gather_q4 and moe_gather_q4_swiglu wrappers. Decode runs eagerly or as a captured CUDA graph (FFAI_LAGUNA_GRAPH=1, warmup-capture-replay via LagunaDecodeCtx with a persistent workspace), with optional fused MoE gate+up (FFAI_LAGUNA_FUSE=1) and micro-fusions (FFAI_LAGUNA_MICRO=1: concatenated QKVG projection, shared expert via the fused gather, o_proj residual accumulate). GB10 single-stream greedy: 21.2 tok/s reference baseline, 36.5 tok/s here (graph+fuse+micro), greedy tokens verified identical to the reference implementation. Env-gated integration smoke: FFAI_LAGUNA_GGUF plus laguna tests in ffai-cuda (unit kernels run without weights). --- .../crates/backends/ffai-cuda/tests/laguna.rs | 279 +++ rust/crates/ffai-models/src/laguna.rs | 1171 +++++++++++++ rust/crates/ffai-models/src/lib.rs | 1 + rust/crates/ffai-modeltests/src/lib.rs | 1027 ++++++++++- rust/crates/ffai-ops/src/lib.rs | 1503 ++++++++++++++++- 5 files changed, 3922 insertions(+), 59 deletions(-) create mode 100644 rust/crates/backends/ffai-cuda/tests/laguna.rs create mode 100644 rust/crates/ffai-models/src/laguna.rs diff --git a/rust/crates/backends/ffai-cuda/tests/laguna.rs b/rust/crates/backends/ffai-cuda/tests/laguna.rs new file mode 100644 index 00000000..89f72663 --- /dev/null +++ b/rust/crates/backends/ffai-cuda/tests/laguna.rs @@ -0,0 +1,279 @@ +// 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 FFAI_LAGUNA_GGUF to the +//! model's GGUF path (the 71G checkpoint), otherwise the test skips. +use ffai_cuda::CudaDevice; +#[test] +fn laguna_host_dequant() { + ffai_modeltests::verify_laguna_host_dequant(); +} + +#[test] +fn laguna_host_layer0() { + ffai_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 ffai_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 = ffai_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 = ffai_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 ffai_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 = ffai_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"); +} + +/// `sdpa_decode_nbuf` (the CUDA-graph-capturable raw-CUDA fallback used by +/// Laguna's decode path under `FFAI_LAGUNA_GRAPH=1`) must match the registry +/// `ops::sdpa_decode` kernel's output — same dense-causal single-pass +/// 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 ffai_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 = ffai_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 = ffai_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"); +} + +/// `ffai_moe_gather_q4_swiglu` (the fused gate+up+SwiGLU MoE gather kernel +/// used by Laguna's decode path under `FFAI_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. +#[test] +fn laguna_moe_swiglu_fused_unit() { + use ffai_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) = ffai_ops::quantize_q4(&gate_w, n_exp * inter, hid); + let (up_qs, up_sc) = ffai_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 = ffai_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 = ffai_ops::moe_gather_q4(d, &gate_qs_t, &gate_sc_t, &x_t, &idx_t, top_k, inter, hid).unwrap(); + let ups = ffai_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, "ffai_moe_gather_q4_swiglu wrong vs unfused compose ({bad} bad elems)"); + eprintln!("moe gather q4 fused swiglu vs unfused compose: OK"); +} + +#[test] +fn laguna_on_cuda() { + let Some(dev) = CudaDevice::create().expect("cuda") else { eprintln!("no CUDA, skip"); return; }; + ffai_modeltests::verify_laguna(dev.as_ref(), "GB10 sm_121"); +} diff --git a/rust/crates/ffai-models/src/laguna.rs b/rust/crates/ffai-models/src/laguna.rs new file mode 100644 index 00000000..4e4288ad --- /dev/null +++ b/rust/crates/ffai-models/src/laguna.rs @@ -0,0 +1,1171 @@ +// 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-only: single-token forward with a persistent KV cache (growing for +//! full-attention layers, a fixed-512 ring buffer for sliding layers). No +//! prefill / batched-token path here. +//! +//! Weights load straight from GGUF (see [`ffai_loader::gguf::Gguf`]): +//! dequantize each tensor to f32 once, then requantize into this engine's Q4 +//! format ([`ffai_ops::quantize_q4`]) and drop the f32 — the GGUF k-quant +//! (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 ffai_core::{DType, Device, Error, Result, Tensor}; +use ffai_loader::gguf::Gguf; +use ffai_ops as ops; + +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 `ffai_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() +} + +/// 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 { + fn upload(dev: &dyn Device, flat: &[f32], m: usize, k: usize) -> Result { + let (qs, scales) = ops::quantize_q4(flat, m, k); + Ok(Q4Dense { + qs: Tensor::new(dev.upload(&u32s_to_bytes(&qs))?, vec![qs.len()], DType::U32), + scales: Tensor::new(dev.upload(&f32s_to_f16_bytes(&scales))?, vec![scales.len()], DType::F16), + m, + k, + }) + } + 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) + } +} + +/// `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 { + /// `flat` is the dequantized fused GGUF tensor, expert-major contiguous + /// (`flat[e*m*k .. (e+1)*m*k]` is expert `e`'s `[m, k]` row-major slab — + /// true for both `ffn_{gate,up}_exps` `[n_embd, n_ff_exp, n_expert]` and + /// `ffn_down_exps` `[n_ff_exp, n_embd, n_expert]` given GGUF's + /// fastest-dim-first tensor layout). + fn upload(dev: &dyn Device, flat: &[f32], n_expert: usize, m: usize, k: usize) -> Result { + let per_expert = m * k; + if flat.len() != per_expert * n_expert { + return Err(Error::Msg(format!( + "Q4ExpertStack: flat len {} != n_expert*m*k ({} * {} * {})", + flat.len(), n_expert, m, k + ))); + } + let qw = per_expert / 32 * 4; // quantize_q4: m*k/32 blocks, 4 u32/block + let sw = per_expert / 32; // 1 f32 scale/block + let mut qs_all = Vec::with_capacity(qw * n_expert); + let mut sc_all = Vec::with_capacity(sw * n_expert); + for e in 0..n_expert { + let slab = &flat[e * per_expert..(e + 1) * per_expert]; + let (qs, sc) = ops::quantize_q4(slab, m, k); + qs_all.extend_from_slice(&qs); + sc_all.extend_from_slice(&sc); + } + Ok(Q4ExpertStack { + qs: Tensor::new(dev.upload(&u32s_to_bytes(&qs_all))?, vec![qs_all.len()], DType::U32), + scales: Tensor::new(dev.upload(&f32s_to_f16_bytes(&sc_all))?, vec![sc_all.len()], DType::F16), + m, + k, + qs_words_per_expert: qw, + scale_elems_per_expert: sw, + }) + } + /// Device-side view of expert `e`'s `(qs, scales)` — a byte-offset slice + /// 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) + } +} + +/// 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, + }, + } + } +} + +/// FFAI_LAGUNA_MICRO=1 enables three decode-path micro-fusions (QKVG concat +/// gemv, shared-expert fused swiglu gather, o_proj residual accum). The env +/// is read once at load time (to decide which weight layout to build below — +/// [`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("FFAI_LAGUNA_MICRO").map(|v| v == "1").unwrap_or(false) +} + +/// QKV(+gate) projection weights. `Fused` concatenates all four projections' +/// rows into one `Q4Dense` at load time (built ONLY under +/// `FFAI_LAGUNA_MICRO=1`, so a layer never carries both layouts resident at +/// once — that would double the ~2.4GB of attention-projection weights). +/// 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 +} + +/// 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] +} + +/// 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 `FFAI_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), +} + +/// A routed MoE FFN block (layers `n_layer_dense_lead..`). +struct MoeFfn { + ffn_norm: Tensor, // [hidden] f32 + router: Tensor, // [n_expert, hidden] f32 — kept dense for routing precision + bias: Tensor, // [n_expert] f32 — e_score_correction_bias (selection only) + gate: Q4ExpertStack, // per-expert [n_ff_exp, hidden] + up: Q4ExpertStack, // per-expert [n_ff_exp, hidden] + down: Q4ExpertStack, // per-expert [hidden, n_ff_exp] + shared_expert: SharedExpert, +} + +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 `ffai_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 }) + } +} + +fn dims_out_in(t: &ffai_loader::gguf::GgufTensor, name: &str) -> Result<(usize, usize)> { + // GGUF dims are fastest-first: a 2-D `[out, in]` weight is stored as + // `[in, out]` (`ggml`'s `create_tensor({k, m}, ...)` convention), which + // 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)) +} + +fn load_dense(dev: &dyn Device, g: &Gguf, name: &str) -> Result { + let t = g.tensor(name).ok_or_else(|| Error::Msg(format!("gguf: tensor '{name}' missing")))?; + let (m, k) = dims_out_in(t, name)?; + let flat = g.dequant_f32(name)?; + Q4Dense::upload(dev, &flat, m, k) +} + +/// Like [`load_dense`] but returns the dequantized f32 slab (not yet +/// requantized/uploaded) plus its `(m, k)` — used by the QKVG-concat fusion, +/// which needs to concatenate rows across four tensors BEFORE `quantize_q4`. +fn load_dense_raw(g: &Gguf, name: &str) -> Result<(Vec, usize, usize)> { + let t = g.tensor(name).ok_or_else(|| Error::Msg(format!("gguf: tensor '{name}' missing")))?; + let (m, k) = dims_out_in(t, name)?; + let flat = g.dequant_f32(name)?; + Ok((flat, m, k)) +} + +fn load_vec_f32(dev: &dyn Device, g: &Gguf, name: &str) -> Result { + let flat = g.dequant_f32(name)?; + let n = flat.len(); + Ok(Tensor::new(dev.upload(&f32s_to_bytes(&flat))?, vec![n], DType::F32)) +} + +fn load_expert_stack(dev: &dyn Device, g: &Gguf, name: &str, n_expert: usize, m: usize, k: usize) -> Result { + let flat = g.dequant_f32(name)?; + Q4ExpertStack::upload(dev, &flat, n_expert, m, k) +} + +/// Load Laguna-S-2.1 straight from a GGUF file: parse `laguna.*` metadata for +/// the config, then stream every tensor through `dequant_f32` → requantize +/// (Q4 for projections/experts, dense f32 for embeddings/norms/router) → +/// upload, dropping the f32 immediately after each tensor. +pub fn load_gguf(dev: &dyn Device, g: &Gguf) -> Result { + let cfg = LagunaConfig::from_gguf(g); + + // token_embd / output are [vocab, hidden] (GGUF fastest-first dims already + // give the flat data in row-major [vocab, hidden] order — see `dims_out_in`). + let load_embed_like = |name: &str| -> Result { + let flat = g.dequant_f32(name)?; + Ok(Tensor::new(dev.upload(&f32s_to_bytes(&flat))?, vec![cfg.vocab, cfg.hidden], DType::F32)) + }; + let tok_embd = load_embed_like("token_embd.weight")?; + + let out_norm = load_vec_f32(dev, g, "output_norm.weight")?; + let head_name = if g.tensor("output.weight").is_some() { "output.weight" } else { "token_embd.weight" }; + let lm_head = { + let (qs, sc, m, k) = g.q8_repack(head_name)?; + Q8Head { + qs: Tensor::new(dev.upload(&u32s_to_bytes(&qs))?, vec![qs.len()], DType::U32), + // gemv_q8 reads f32 scales (unlike the q4 family's f16), matching + // the proven resident-Q8 loader path. + scales: Tensor::new(dev.upload(&f32s_to_bytes(&sc))?, vec![sc.len()], DType::F32), + m, + k, + } + }; + let one = Tensor::new(dev.upload(&1.0f32.to_le_bytes())?, vec![1], DType::F32); + + let micro = laguna_micro_enabled(); + + let mut layers = Vec::with_capacity(cfg.n_layers); + for i in 0..cfg.n_layers { + let p = format!("blk.{i}"); + + // n_head is derived from wq's own GGUF dims (48 full-attention / 72 + // sliding for Laguna-S-2.1), not the LagunaConfig formula — see the + // AttnWeights doc comment. + let (wq_flat, n_q_rows, hidden_k) = load_dense_raw(g, &format!("{p}.attn_q.weight"))?; + let (wk_flat, n_kv_rows, _) = load_dense_raw(g, &format!("{p}.attn_k.weight"))?; + let (wv_flat, n_kv_rows_v, _) = load_dense_raw(g, &format!("{p}.attn_v.weight"))?; + let (wg_flat, n_g_rows, _) = load_dense_raw(g, &format!("{p}.attn_gate.weight"))?; + debug_assert_eq!(n_kv_rows, n_kv_rows_v, "wk/wv row count mismatch"); + let n_head = n_q_rows / cfg.head_dim; + + // FFAI_LAGUNA_MICRO=1: concatenate the four dequantized f32 slabs + // BEFORE quantize_q4 (row order q,k,v,gate; all share k=hidden), so + // decode can do one gemv_q4 instead of four. Only ONE of the two + // layouts is ever built per layer — building both would double the + // resident attention-projection weights (~2.4GB total). + let proj = if micro { + let n_rows = n_q_rows + n_kv_rows + n_kv_rows_v + n_g_rows; + let mut concat = Vec::with_capacity(n_rows * hidden_k); + concat.extend_from_slice(&wq_flat); + concat.extend_from_slice(&wk_flat); + concat.extend_from_slice(&wv_flat); + concat.extend_from_slice(&wg_flat); + let wqkvg = Q4Dense::upload(dev, &concat, n_rows, hidden_k)?; + AttnProj::Fused { wqkvg, n_q_rows, n_kv_rows, n_g_rows } + } else { + AttnProj::Split { + wq: Q4Dense::upload(dev, &wq_flat, n_q_rows, hidden_k)?, + wk: Q4Dense::upload(dev, &wk_flat, n_kv_rows, hidden_k)?, + wv: Q4Dense::upload(dev, &wv_flat, n_kv_rows_v, hidden_k)?, + wgate: Q4Dense::upload(dev, &wg_flat, n_g_rows, hidden_k)?, + } + }; + + let attn = AttnWeights { + n_head, + attn_norm: load_vec_f32(dev, g, &format!("{p}.attn_norm.weight"))?, + proj, + wo: load_dense(dev, g, &format!("{p}.attn_output.weight"))?, + q_norm: load_vec_f32(dev, g, &format!("{p}.attn_q_norm.weight"))?, + k_norm: load_vec_f32(dev, g, &format!("{p}.attn_k_norm.weight"))?, + }; + + let ffn = if cfg.is_dense_layer(i) { + LayerFfn::Dense(DenseFfn { + ffn_norm: load_vec_f32(dev, g, &format!("{p}.ffn_norm.weight"))?, + gate: load_dense(dev, g, &format!("{p}.ffn_gate.weight"))?, + up: load_dense(dev, g, &format!("{p}.ffn_up.weight"))?, + down: load_dense(dev, g, &format!("{p}.ffn_down.weight"))?, + }) + } else { + let router_flat = g.dequant_f32(&format!("{p}.ffn_gate_inp.weight"))?; + let router = Tensor::new( + dev.upload(&f32s_to_bytes(&router_flat))?, + vec![cfg.n_expert, cfg.hidden], + DType::F32, + ); + + // FFAI_LAGUNA_MICRO=1: repack the shared expert's gate+up as + // single-expert (n_expert=1) stacks so decode can drive them + // through the same moe_gather_q4_swiglu the routed experts use + // (top_k=1) instead of two plain gemvs + a host-side swiglu. + // `down` stays a plain Q4Dense either way. Only one layout is + // ever built per layer (same memory-doubling rationale as + // AttnProj above). + let shared_expert = if micro { + let gate = load_expert_stack(dev, g, &format!("{p}.ffn_gate_shexp.weight"), 1, cfg.n_ff_shexp, cfg.hidden)?; + let up = load_expert_stack(dev, g, &format!("{p}.ffn_up_shexp.weight"), 1, cfg.n_ff_shexp, cfg.hidden)?; + let down = load_dense(dev, g, &format!("{p}.ffn_down_shexp.weight"))?; + let zero_idx = Tensor::new(dev.upload(&0u32.to_le_bytes())?, vec![1], DType::U32); + SharedExpert::Fused(SharedExpertFused { gate, up, down, zero_idx }) + } else { + SharedExpert::Split { + gate: load_dense(dev, g, &format!("{p}.ffn_gate_shexp.weight"))?, + up: load_dense(dev, g, &format!("{p}.ffn_up_shexp.weight"))?, + down: load_dense(dev, g, &format!("{p}.ffn_down_shexp.weight"))?, + } + }; + + LayerFfn::Moe(MoeFfn { + ffn_norm: load_vec_f32(dev, g, &format!("{p}.ffn_norm.weight"))?, + router, + bias: load_vec_f32(dev, g, &format!("{p}.exp_probs_b.bias"))?, + gate: load_expert_stack(dev, g, &format!("{p}.ffn_gate_exps.weight"), cfg.n_expert, cfg.n_ff_exp, cfg.hidden)?, + up: load_expert_stack(dev, g, &format!("{p}.ffn_up_exps.weight"), cfg.n_expert, cfg.n_ff_exp, cfg.hidden)?, + down: load_expert_stack(dev, g, &format!("{p}.ffn_down_exps.weight"), cfg.n_expert, cfg.hidden, cfg.n_ff_exp)?, + shared_expert, + }) + }; + + layers.push(LagunaLayer { attn, ffn }); + } + + 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)] +/// FFAI_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("FFAI_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. + // + // FFAI_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); } + // FFAI_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("FFAI_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. + // + // FFAI_LAGUNA_FUSE=1 replaces the two gathers + host swiglu with + // one fused gate+up+SwiGLU gather kernel: same `h2` read once + // instead of twice, one launch instead of two gathers plus the + // elementwise pass. Default (unset) keeps the unfused path. + let act = if std::env::var("FFAI_LAGUNA_FUSE").map(|v| v == "1").unwrap_or(false) { + ops::moe_gather_q4_swiglu( + dev, &f.gate.qs, &f.gate.scales, &f.up.qs, &f.up.scales, &h2, &idx, + cfg.n_expert_used, cfg.n_ff_exp, cfg.hidden, + )? + } else { + let gates = ops::moe_gather_q4(dev, &f.gate.qs, &f.gate.scales, &h2, &idx, cfg.n_expert_used, cfg.n_ff_exp, cfg.hidden)?; + let ups = ops::moe_gather_q4(dev, &f.up.qs, &f.up.scales, &h2, &idx, cfg.n_expert_used, cfg.n_ff_exp, cfg.hidden)?; + ops::swiglu(dev, &gates, &ups)? + }; + 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). + // + // FFAI_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) +} + +/// 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 `FFAI_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`. +/// +/// `FFAI_LAGUNA_GRAPH=1` switches `decode_layer`'s RoPE/SDPA to the +/// buffer-driven (capture-safe) kernels; without it, behavior is byte-for-byte +/// identical to before this refactor (same kernels, same scalar arguments, +/// just reading `token_id`/`pos` off a persistent buffer instead of a fresh +/// one — the values are the same either way). +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("FFAI_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(); + + // FFAI_LAGUNA_DEBUG=1: per-layer residual stats (L2 norm, NaN count) to + // bisect a divergence to the first bad layer. Hot path unaffected when + // the env is unset. Downloads mid-step, so it is INCOMPATIBLE with graph + // capture/replay — `LagunaDecodeCtx::step` refuses to run at all when + // 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("FFAI_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 +/// `FFAI_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 +/// `ffai-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. +/// +/// `FFAI_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("FFAI_LAGUNA_DEBUG").is_ok() { + return Err(Error::Msg( + "LagunaDecodeCtx::step: FFAI_LAGUNA_DEBUG is incompatible with CUDA-graph \ + capture/replay (the debug path downloads mid-step, which the CUDA driver \ + forbids during capture). Unset FFAI_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/ffai-models/src/lib.rs b/rust/crates/ffai-models/src/lib.rs index dec04bb0..ec321b0a 100644 --- a/rust/crates/ffai-models/src/lib.rs +++ b/rust/crates/ffai-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/ffai-modeltests/src/lib.rs b/rust/crates/ffai-modeltests/src/lib.rs index 5db1369e..0a253b20 100644 --- a/rust/crates/ffai-modeltests/src/lib.rs +++ b/rust/crates/ffai-modeltests/src/lib.rs @@ -31,7 +31,14 @@ pub fn verify_gpt2(d: &dyn Device, plat: &str) { 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 }; @@ -94,7 +101,14 @@ pub fn verify_mamba2(d: &dyn Device, plat: &str) { 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) }; @@ -160,6 +174,628 @@ 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 `FFAI_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 FFAI_LAGUNA_GGUF, no device needed. +pub fn verify_laguna_host_dequant() { + let Ok(path) = std::env::var("FFAI_LAGUNA_GGUF") else { + eprintln!("FFAI_LAGUNA_GGUF not set, skipping host dequant check"); + return; + }; + let g = match ffai_loader::gguf::Gguf::open(&path) { + Ok(g) => g, + Err(e) => { eprintln!("gguf open failed: {e:?}"); return; } + }; + // FFAI_LAGUNA_DECODE_IDS="69377,13829": detokenize ids without a model + // load (fast oracle-comparison helper). + if let Ok(ids_s) = std::env::var("FFAI_LAGUNA_DECODE_IDS") { + if let Ok(tok) = ffai_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 +/// (`FFAI_LAGUNA_DEBUG=2` in `ffai_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 +/// [`ffai_ops::quantize_q4`] and dequantized back on host, with the block +/// scale pushed through an f32->f16->f32 roundtrip first, matching how the +/// device stores Q4 scales). The device output should be compared against +/// `roundtrip`, not `exact` — decode always runs on Q4 weights regardless of +/// the GGUF's original per-tensor dtype (see `ffai_models::laguna::load_dense`). +/// +/// Env-gated on FFAI_LAGUNA_GGUF, no device/CUDA needed. +pub fn verify_laguna_host_layer0() { + let Ok(path) = std::env::var("FFAI_LAGUNA_GGUF") else { + eprintln!("FFAI_LAGUNA_GGUF not set, skipping host layer-0 reference"); + return; + }; + let g = match ffai_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 ffai_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 + // (`ffai_models::laguna::f32s_to_f16_bytes`) — the gemv_q4 kernels read + // per-block Q4 scales as f16. + fn f32_to_f16(f: f32) -> u16 { + let x = f.to_bits(); + 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) = ffai_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 ffai_models::gguf_tokenizer::GgufTokenizer; + use ffai_models::laguna; + + let Ok(path) = std::env::var("FFAI_LAGUNA_GGUF") else { + eprintln!("FFAI_LAGUNA_GGUF not set — skipping Laguna-S-2.1 decode smoke test"); + return; + }; + let Ok(g) = ffai_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; + } + }; + + 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("FFAI_LAGUNA_GEN").ok().and_then(|v| v.parse::().ok()).unwrap_or(16) + + 1; + let Ok(kv) = laguna::LagunaKvCache::new(d, &model.cfg, max_seq) else { + eprintln!("Laguna: KV cache alloc failed — skipping"); + return; + }; + + // FFAI_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("FFAI_LAGUNA_SWEEP") { + let n_gen: usize = std::env::var("FFAI_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 ffai_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; + } + + // FFAI_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 `ffai_models::laguna::LagunaDecodeCtx`). Default (unset): plain + // eager `decode_step`, now reading a persistent per-step workspace + // instead of uploading fresh token-id/position buffers every call — + // same kernels, same scalars, byte-identical output to before. + let graph_mode = std::env::var("FFAI_LAGUNA_GRAPH").is_ok(); + let mode_label = if graph_mode { "graph" } else { "eager" }; + 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. + // FFAI_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("FFAI_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; + // FFAI_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("FFAI_LAGUNA_FAST").map(|v| v == "1").unwrap_or(false); + for pos in 0..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) = ffai_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 = ffai_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 @@ -179,7 +815,14 @@ pub fn verify_pythia(d: &dyn Device, plat: &str) { 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; @@ -214,7 +857,14 @@ pub fn verify_olmo2(d: &dyn Device, plat: &str) { 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; @@ -256,7 +906,14 @@ pub fn verify_gptneo(d: &dyn Device, plat: &str) { 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; @@ -293,7 +950,14 @@ pub fn verify_gemma2(d: &dyn Device, plat: &str) { 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) }; @@ -333,7 +997,14 @@ pub fn verify_phi(d: &dyn Device, plat: &str) { 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 { @@ -370,7 +1041,14 @@ pub fn verify_stablelm2(d: &dyn Device, plat: &str) { 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() }; @@ -409,7 +1087,14 @@ pub fn verify_olmoe(d: &dyn Device, plat: &str) { 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) }; @@ -466,7 +1151,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() }; @@ -685,7 +1377,7 @@ pub fn verify_nemotron(d: &dyn Device, plat: &str) { /// Env: NEMOTRON_DECODE (steps, default 32), NEMOTRON_PREFILL (warm the cache to /// this context length before timing, default 0). pub fn bench_nemotron(d: &dyn Device, plat: &str) { - use ffai_ops::{add, cast_f16_f32, cast_f32_f16, conv1d_causal_prefill, conv1d_causal_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 ffai_ops::{add, cast_f16_f32, cast_f32_f16, conv1d_causal_prefill, conv1d_causal_prefill_split, conv1d_causal_step, dequant_q4, dequant_q4_off, gather, gated_group_rmsnorm, gated_group_rmsnorm_batched, gated_group_rmsnorm_batched_f16out, gemm_cublas, gemm_cublas_f32out, gemm_q4_mpp, gemv, gemv_q4, gemv_q4_accum, gemv_q4_relu2, gemv_q8, gemv_q8_relu2, gemv_q8_accum, kv_append, kv_append_many, mamba_split_conv, mamba_split_proj, matmul, moe_bgemm_q4_bm64, moe_expert_ids_from_offsets, moe_fused_ffn, moe_gather_down, moe_gather_up_relu2, moe_grouped_gemm, moe_grouped_gemm_cutlass, moe_q4_grouped_mma, moe_q4_grouped_mma_dev, moe_router_device, moe_scatter_add_det_dev, moe_scatter_add_topk_det_f16, moe_scatter_add_topk_det_f16_x2, moe_scatter_add, moe_scatter_add_det, moe_w4a16, moe_w4a16_marlin, moe_weighted_sum, nvfp4_roundtrip, permute_q4_to_marlin, quantize_q4, quantize_q8, relu2, relu2_scale_f16, rms_norm, rope_llama, rope_llama_many, sdpa_multi, sdpa_multi_tc, sdpa_multi_tc_varlen, silu, slice, sdpa_decode, sdpa_decode_2pass, sdpa_decode_2pass_bc4, sdpa_decode_2pass_tiled, softplus_add, softplus_add_rows, ssm_prefill_scan, ssm_prefill_scan_chunked, ssm_prefill_scan_ssd, ssm_prefill_scan_ssd_portable, ssm_step, strided_col_copy}; use std::collections::HashMap; use std::time::Instant; const PATTERN: &str = "MEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; @@ -716,7 +1408,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) }; @@ -804,6 +1503,8 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { 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). @@ -1120,6 +1821,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) = ffai_ops::quantize_nvfp4(&w, n, k); + let sv: Vec = (0..n * ng).map(|i| ffai_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 = ffai_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 = ffai_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: @@ -1799,6 +2574,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: &ffai_core::Tensor| -> ffai_core::Tensor { + if bf16_stream { ffai_ops::cast_f32_bf16(d, t).unwrap() } else { ffai_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: &ffai_core::Tensor| -> ffai_core::Tensor { + match t.dtype { DType::F16 => ffai_ops::cast_f16_f32(d, t).unwrap(), + DType::BF16 => ffai_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 +2602,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 +2610,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 (vLLM-style) + // 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 @@ -1850,6 +2646,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). @@ -1928,6 +2732,36 @@ 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) — + // FFAI proved 1.81x @ argmax near-tie; A/B vs the per-channel default. + if std::env::var("NEMOTRON_FP8_PT").is_ok() { + // NEMOTRON_FP8_HAD: block-Hadamard rotate (128) along K before the + // per-tensor fp8 quant — spreads activation outliers so per-tensor + // 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 { ffai_ops::fwht128(d, &wf, m, k).unwrap() } else { wf }; + let t = ffai_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, ffai_ops::fwht128_amaxonly(d, xh, rows, k).unwrap()); + pf!(15, 0.0, ffai_ops::fwht128_quant_pre(d, xh, rows, k, &amx, clipv).unwrap()) + } else { + pf!(15, 0.0, ffai_ops::fp8_quant(d, xh, rows*k, clipv).unwrap()) + }; + return pf!(2, 2.0*rows as f64*m as f64*k as f64, + ffai_ops::gemm_fp8(d, &xq, &xsc, &wq, &wsc, rows, m, k, true).unwrap()); + } // PER-CHANNEL weight (rows=m) + PER-TOKEN act (rows) FP8, scales // folded post-GEMM (D[r,c]*=xsc[r]·wsc[c]) — the quality- // preserving granularity (per-tensor was cosine ~0.94). @@ -1947,6 +2781,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); @@ -2019,11 +2861,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' => { @@ -2043,14 +2895,12 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { 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. // 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 +2932,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 { @@ -2099,9 +2953,18 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { ssm_state[l] = Some(so); // 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) let nw = if l + 1 < PATTERN.len() { fwd.get(&format!("language_model.backbone.layers.{}.norm.weight", l + 1)) } else { fwd.get("norm_f") }; @@ -2242,6 +3105,10 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { // 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 @@ -2264,7 +3131,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) = ffai_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) = + ffai_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 { + ffai_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 @@ -2422,7 +3296,7 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { 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() { @@ -2707,7 +3581,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) = ffai_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, + ffai_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, + ffai_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]); @@ -2826,7 +3718,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 +3757,12 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { ffai_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")]; @@ -2894,7 +3796,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) @@ -2922,9 +3829,14 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { // 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(); @@ -3324,6 +4236,10 @@ 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 @@ -3369,6 +4285,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, ffai_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, ffai_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, ffai_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,7 +4306,7 @@ 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, ffai_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 ── @@ -3454,9 +4382,15 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { 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) = ffai_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,6 +4402,7 @@ 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) let nw = if l + 1 < PATTERN.len() { fwd.get(&format!("language_model.backbone.layers.{}.norm.weight", l + 1)) } else { fwd.get("norm_f") }; @@ -3490,16 +4425,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()); @@ -3591,8 +4535,11 @@ pub fn bench_nemotron(d: &dyn Device, plat: &str) { // 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())); diff --git a/rust/crates/ffai-ops/src/lib.rs b/rust/crates/ffai-ops/src/lib.rs index 9aad9d3d..2f614720 100644 --- a/rust/crates/ffai-ops/src/lib.rs +++ b/rust/crates/ffai-ops/src/lib.rs @@ -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 ffai_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 ffai_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", "ffai_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 { "ffai_add_f16" } else { "ffai_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, "ffai_add") } @@ -200,6 +244,14 @@ extern "C" __global__ void ffai_add3( // left-assoc a+b+c flipped the d32768 argmax via last-ulp drift. for(;i +#include +extern "C" __global__ void ffai_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 "ffai_add3_f16", DType::BF16 => "ffai_add3_bf16", _ => "ffai_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) } @@ -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 { ("mt_rms_norm", (n / 4) as u32) } else { ("mt_rms_norm_wide", 256u32) @@ -930,6 +986,152 @@ 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 /// the [s, in_proj_out] projection matrix on device (NEMOTRON_CONV_DEVICE path). @@ -1021,6 +1223,92 @@ 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 ffai_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", "ffai_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. +pub fn moe_gather_q4(dev: &dyn Device, qs: &Tensor, scales: &Tensor, x: &Tensor, idx: &Tensor, top_k: usize, inter: usize, hid: usize) -> Result { + let k = cached_ir("ffai_moe_gather_q4", x.dtype, || { let mut k = metaltile_std::ffai::gemv_q8::ffai_moe_gather_q4::kernel_ir_for(x.dtype); k.mode = metaltile_core::ir::KernelMode::Reduction; k }); + let out = Tensor::empty(dev, vec![top_k * inter], x.dtype)?; + let u = |v: u32| Binding::Scalar(v.to_le_bytes().to_vec()); + let rpt: u32 = std::env::var("MT_MOE_RPT").ok().and_then(|v| v.parse().ok()).filter(|&r| r >= 1).unwrap_or(2); + 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) +} + +/// 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)] +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 { + let k = cached_ir("ffai_moe_gather_q4_swiglu", x.dtype, || { let mut k = metaltile_std::ffai::gemv_q8::ffai_moe_gather_q4_swiglu::kernel_ir_for(x.dtype); k.mode = metaltile_core::ir::KernelMode::Reduction; k }); + let out = Tensor::empty(dev, vec![top_k * inter], x.dtype)?; + let u = |v: u32| Binding::Scalar(v.to_le_bytes().to_vec()); + let rpt: u32 = std::env::var("MT_MOE_RPT").ok().and_then(|v| v.parse().ok()).filter(|&r| r >= 1).unwrap_or(2); + dev.dispatch(&k, &[ + Binding::Buffer(gate_qs.buffer.clone()), + Binding::Buffer(gate_scales.buffer.clone()), + Binding::Buffer(up_qs.buffer.clone()), + Binding::Buffer(up_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) +} + /// 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)] @@ -1336,6 +1624,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 +1691,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("ffai_moe_weighted_sum", acc.dtype, || { let mut k = metaltile_std::ffai::gemv_q8::ffai_moe_weighted_sum::kernel_ir_for(acc.dtype); k.mode = metaltile_core::ir::KernelMode::Grid3D; k }); let u = |v: u32| Binding::Scalar(v.to_le_bytes().to_vec()); @@ -1433,7 +1783,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 +1803,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; @@ -1561,6 +1917,54 @@ pub fn dec_ue4m3(b: u8) -> f32 { /// 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; @@ -1976,6 +2380,7 @@ const MOE_GROUPED_GEMM_MMA_SRC: &str = r#" __device__ __forceinline__ unsigned pk2(__half a, __half b){ return (unsigned)__half_as_ushort(a) | ((unsigned)__half_as_ushort(b) << 16); } +__device__ __forceinline__ unsigned ld32(const __half* p){ return *reinterpret_cast(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 +2425,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 @@ -2140,7 +2554,7 @@ extern "C" __global__ void moe_q4_grouped_mma( __half* __restrict__ Out, int N, int K) { - const int BM=64,BN=128,BK=64,NS=16,NKS=4,GS=32,WPR=BK/8,GPR=BK/GS; + 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) int bx=blockIdx.x, by=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; @@ -2151,7 +2565,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> 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", "ffai_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 (MLX `sdpa_vector_2pass` port, 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 @@ -4579,6 +5140,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 @@ -4662,6 +5276,34 @@ 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)] +/// 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, + 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, @@ -4701,6 +5343,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 +5411,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( @@ -4960,6 +5667,315 @@ 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 +// llama-context.cpp 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 ffai_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", "ffai_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 ffai_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 +/// `FFAI_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", "ffai_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) +} + +/// Write a single `u32` scalar into an existing 1-element device buffer, in +/// place, without reallocating. Laguna's CUDA-graph decode path +/// (`FFAI_LAGUNA_GRAPH=1`) keeps a small persistent "workspace" of such +/// buffers (token id, RoPE position, KV length) whose device POINTERS get +/// captured into the graph and must never move; only their contents change +/// from token to token. Call this OUTSIDE any graph-capture region — it must +/// 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", "ffai_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 ffai_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 ffai_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", "ffai_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) +} + /// Softmax over the last dim, row-wise. Dispatches `mt_softmax`. Row width /// `n` must be a multiple of 1024 (the kernel's 4-elems/thread loop). pub fn softmax(dev: &dyn Device, x: &Tensor) -> Result { @@ -5287,12 +6303,125 @@ pub fn mamba_split_proj( // = 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) ─ // @@ -5339,6 +6468,82 @@ 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. +pub fn rms_norm_f16out(dev: &dyn Device, x: &Tensor, weight: &Tensor, eps: f32) -> Result { + let n = *x.shape.last().ok_or_else(|| Error::Msg("rms_norm_f16out: scalar input".into()))?; + if std::env::var("NEMOTRON_RMS_NORM_F16OUT_OFF").is_ok() + || x.dtype != DType::F32 + || weight.dtype != DType::F32 + || n % 128 != 0 + || n > 4096 + { + let out = rms_norm(dev, x, weight, eps)?; + return cast_f32_f16(dev, &out); + } + let rows = x.elem_count() / n; + let out = Tensor::empty(dev, x.shape.clone(), DType::F16)?; + let eps_buf = scalar_buf(dev, eps.to_bits())?; + let k = cached_ir("mt_rms_norm_f16out", DType::F32, || { + let mut k = metaltile_std::mlx::rms_norm::mt_rms_norm_f16out::kernel_ir_for(DType::F32); + k.mode = metaltile_core::ir::KernelMode::Reduction; + k + }); + let grid = Grid { grid: [rows as u32, 1, 1], block: [(n / 4) as u32, 1, 1] }; + dev.dispatch( + &k, + &[ + Binding::Buffer(x.buffer.clone()), + Binding::Buffer(weight.buffer.clone()), + Binding::Buffer(out.buffer.clone()), + Binding::Buffer(eps_buf), + Binding::Scalar((n as u32).to_le_bytes().to_vec()), + ], + grid, + )?; + Ok(out) +} + +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 { + if std::env::var("NEMOTRON_GATED_NORM_F16OUT_OFF").is_ok() { + let out = gated_group_rmsnorm_batched(dev, y, z, w, eps, s, di, gs)?; + return cast_f32_f16(dev, &out); + } + let ng = di / gs; + let kernel = cached_ir("gated_group_rmsnorm_batched_f16out", DType::F32, || { + use metaltile_core::ir::KernelMode; + let mut k = metaltile_std::ffai::ssm::gated_group_rmsnorm_batched_f16out::kernel_ir_for(); + k.mode = KernelMode::Reduction; + k + }); + let out = Tensor::empty(dev, vec![s * di], DType::F16)?; + let eps_buf = Tensor::new(scalar_buf(dev, eps.to_bits())?, vec![1], DType::F32); + let u = |v: u32| Binding::Scalar(v.to_le_bytes().to_vec()); + dev.dispatch( + &kernel, + &[ + Binding::Buffer(y.buffer.clone()), + Binding::Buffer(z.buffer.clone()), + Binding::Buffer(w.buffer.clone()), + Binding::Buffer(out.buffer.clone()), + Binding::Buffer(eps_buf.buffer.clone()), + u(gs as u32), + u(ng as u32), + ], + Grid { grid: [(s * ng) as u32, 1, 1], block: [(gs / 4) as u32, 1, 1] }, + )?; + Ok(out) +} + pub fn gated_group_rmsnorm_batched( dev: &dyn Device, y: &Tensor, @@ -6926,9 +8131,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) @@ -6993,6 +8198,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) @@ -7021,6 +8281,51 @@ extern "C" __global__ void fp8_quant( } } +// 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 @@ -7690,6 +8995,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], @@ -7747,7 +9100,7 @@ pub fn w8a8_packw( /// Quantize an f16 tensor to e4m3 with a dynamic per-tensor scale. /// Returns `(q, sc)` — e4m3 bytes `[n]` + a 2-f32 scale buffer /// `[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); @@ -7760,7 +9113,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", @@ -7770,6 +9123,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)] @@ -8102,3 +9536,34 @@ 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(){ + let (n,k)=(1856usize,2688usize); let ng=k/16; let kp=k/8; + let w0=rf32("/tmp/mdata/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("/tmp/mdata/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("/tmp/mdata/bscales.bin").unwrap(); + let bm=s0.iter().zip(&bref[..ng*n]).filter(|(a,b)|a==b).count(); + let gref=rf32("/tmp/mdata/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"); + } +}