feat(rust): Laguna batched prefill, 156 to 559 pp tok/s, plus weight cache and dtype-cache fix

rust/Cargo.toml

@@ -23,6 +23,7 @@ authors = ["Eric Kryski (@ekryski)", "Tom Turney (@TheTom)"]
 repository = "https://github.com/TheTom/ffai"
 
 [workspace.dependencies]
+rayon = "1"
 # ── engine crates ────────────────────────────────────────────────────
 ffai-core = { path = "crates/ffai-core" }
 ffai-ops = { path = "crates/ffai-ops" }

rust/crates/backends/ffai-cuda/src/imp.rs

@@ -2,8 +2,8 @@
 // SPDX-License-Identifier: Apache-2.0
 
 //! Real CUDA backend (compiled under `--features cuda`). Wraps
-//! `metaltile_runtime::CudaDevice` — the NVRTC → PTX → driver path that
-//! runs the kernel corpus bit-accurately on GB10 — behind the shared
+//! `metaltile_runtime::CudaDevice`, the NVRTC → PTX → driver path that
+//! runs the kernel corpus bit-accurately on GB10, behind the shared
 //! [`ffai_core::Device`] trait, so the engine layer above is identical to
 //! the Metal/Vulkan paths.
 
@@ -27,12 +27,12 @@ fn dispatch_err(e: MetalTileError) -> Error {
 ///
 /// Holds an `Arc<MtCudaDevice>` so the CUDA CONTEXT outlives this module. A
 /// `CudaModule`'s `Drop` calls `cuModuleUnload`, which faults
-/// (STATUS_ACCESS_VIOLATION, 0xc0000005) if the context was already destroyed —
+/// (STATUS_ACCESS_VIOLATION, 0xc0000005) if the context was already destroyed,
 /// which is exactly what happened at process teardown: this device's `dev` field
 /// (its only context-keepalive) dropped before the `modules` cache, destroying
 /// the context, then the cached modules unloaded into the dead context. Pinning
 /// the context here (and dropping `module` before `_dev` via field order below)
-/// guarantees the module always unloads while its context is still live — the
+/// guarantees the module always unloads while its context is still live, the
 /// same invariant `CudaBuffer` already upholds for device allocations.
 struct CachedModule {
     // Field DECLARATION ORDER is DROP ORDER: `module` (cuModuleUnload) must run
@@ -51,7 +51,7 @@ pub struct CudaDevice {
     name: String,
     /// Compile-once cache: kernel name → loaded module.
     modules: RwLock<HashMap<String, Arc<CachedModule>>>,
-    /// Memoized shared-mem bytes per (kernel, block_x) — avoids re-walking the
+    /// Memoized shared-mem bytes per (kernel, block_x), avoids re-walking the
     /// kernel IR on every dispatch (hot in decode: ~hundreds of dispatches/token).
     shared: RwLock<HashMap<(String, u32), u32>>,
 }
@@ -74,8 +74,26 @@ impl CudaDevice {
         }
     }
 
+    /// Cache key for a compiled module. The bare kernel name is NOT unique:
+    /// registry kernels share one name across their dtype variants (the ops
+    /// layer keys its IR cache by (name, dtype), but the same name reaches
+    /// this backend for every variant). Keying compiled modules by name alone
+    /// silently reused the first-compiled dtype's binary for every other
+    /// dtype, which reads with the wrong element stride and faults (or worse,
+    /// corrupts silently). Fold each param's dtype into the key.
+    fn module_key(kernel: &Kernel) -> String {
+        use std::fmt::Write as _;
+        let mut key = kernel.name.clone();
+        key.push('|');
+        for p in &kernel.params {
+            let _ = write!(key, "{:?},", p.dtype);
+        }
+        key
+    }
+
     fn module_for(&self, kernel: &Kernel) -> Result<Arc<CachedModule>> {
-        if let Some(m) = self.modules.read().unwrap().get(&kernel.name) {
+        let mkey = Self::module_key(kernel);
+        if let Some(m) = self.modules.read().unwrap().get(&mkey) {
             return Ok(m.clone());
         }
         let cg = CudaGenerator::new();
@@ -95,7 +113,7 @@ impl CudaDevice {
         self.modules
             .write()
             .unwrap()
-            .insert(kernel.name.clone(), cached.clone());
+            .insert(mkey, cached.clone());
         Ok(cached)
     }
 }
@@ -124,7 +142,7 @@ impl DeviceBuffer for CudaBuffer {
 impl Drop for CudaBuffer {
     fn drop(&mut self) {
         // Return to the device's size-bucketed pool instead of a synchronous
-        // cuMemFree — the hot decode path reuses these every token.
+        // cuMemFree, the hot decode path reuses these every token.
         self.dev.free_raw_pooled(self.ptr, self.len);
     }
 }
@@ -162,7 +180,7 @@ impl Device for CudaDevice {
     }
 
     fn alloc_zeroed(&self, len: usize) -> Result<Arc<dyn DeviceBuffer>> {
-        // Stream-ordered device memset — no host zero buffer, no pageable
+        // Stream-ordered device memset, no host zero buffer, no pageable
         // H2D copy. A [s,hid] f32 MoE accumulator is ~22 MB; uploading zeros
         // for it cost ~4-6 ms of host-blocking copy per E-layer.
         let ptr = self.dev.alloc_raw(len).map_err(dispatch_err)?;
@@ -181,7 +199,9 @@ impl Device for CudaDevice {
     fn dispatch(&self, kernel: &Kernel, bindings: &[Binding], grid: Grid) -> Result<()> {
         let module = self.module_for(kernel)?;
         let func = module.module.function(&kernel.name).map_err(dispatch_err)?;
-        let skey = (kernel.name.clone(), grid.block[0]);
+        // Same dtype-collision hazard as the module cache: shared-memory sizing
+        // depends on element widths, so fold the param dtypes in.
+        let skey = (Self::module_key(kernel), grid.block[0]);
         let shared = if let Some(&s) = self.shared.read().unwrap().get(&skey) {
             s
         } else {
@@ -230,7 +250,7 @@ impl Device for CudaDevice {
             }
         }
 
-        // Async launch (no per-dispatch cuCtxSynchronize) — kernels pipeline on the
+        // Async launch (no per-dispatch cuCtxSynchronize), kernels pipeline on the
         // ordered default stream; `download`/`synchronize` sync when results are read.
         self.dev
             .launch_async(func, grid.grid, grid.block, shared, &mut args)
@@ -311,7 +331,7 @@ impl Device for CudaDevice {
 
         // Use cooperative launch only when NOT inside a CUDA graph capture
         // (cuLaunchCooperativeKernel is not capturable). Fall back to regular
-        // launch during capture — the caller (moe_fused_ffn) must not use
+        // launch during capture, the caller (moe_fused_ffn) must not use
         // grid.sync() in that code path (handled by NEMOTRON_GRAPH exclusion).
         if cooperative && !self.dev.is_capturing() {
             self.dev.launch_async_coop(func, grid, block, shared_bytes, &mut args).map_err(dispatch_err)

rust/crates/backends/ffai-cuda/tests/colcopy.rs

@@ -0,0 +1,99 @@
+#![cfg(feature = "cuda")]
+//! strided_col_copy launch-geometry regression: the op dispatches with a
+//! fixed 64-thread block. Before the fix, the grid was `n / 64` (integer
+//! division, no ceiling), so whenever `s * width` wasn't a multiple of 64
+//! the tail elements were silently dropped (never written into dst, left
+//! as zeroed garbage). Covers a deliberately non-64-aligned shape (the
+//! Laguna-style head-count case, e.g. width=72) plus a 64-aligned control
+//! shape to confirm the fix didn't regress the common case.
+use ffai_core::{DType, Device, Tensor};
+use ffai_cuda::CudaDevice;
+use ffai_ops::strided_col_copy;
+
+fn tb(v: &[f32]) -> Vec<u8> { v.iter().flat_map(|x| x.to_le_bytes()).collect() }
+fn fb(b: &[u8]) -> Vec<f32> { b.chunks_exact(4).map(|c| f32::from_le_bytes(c.try_into().unwrap())).collect() }
+fn tn(d: &dyn Device, v: &[f32]) -> Tensor { Tensor::new(d.upload(&tb(v)).unwrap(), vec![v.len()], DType::F32) }
+
+// Host-computed expectation: dst[ti*width + ci] = src[ti*stride + col_off + ci].
+fn host_expected(src: &[f32], s: usize, stride: usize, col_off: usize, width: usize) -> Vec<f32> {
+    (0..s)
+        .flat_map(|ti| (0..width).map(move |ci| src[ti * stride + col_off + ci]))
+        .collect()
+}
+
+fn run_case(name: &str, s: usize, stride: usize, col_off: usize, width: usize) {
+    let Some(dev) = CudaDevice::create().expect("cuda") else {
+        eprintln!("no CUDA, skipping {name}");
+        return;
+    };
+    let d = dev.as_ref();
+    let src: Vec<f32> = (0..s * stride).map(|i| ((i as f32) * 0.019).sin()).collect();
+    let src_t = tn(d, &src);
+
+    let out = strided_col_copy(d, &src_t, s, stride, col_off, width)
+        .unwrap_or_else(|e| panic!("{name}: strided_col_copy dispatch failed: {e:?}"));
+    d.synchronize().unwrap();
+
+    let mut bytes = vec![0u8; s * width * 4];
+    d.download(out.buffer.as_ref(), &mut bytes).unwrap();
+    let got = fb(&bytes);
+    let want = host_expected(&src, s, stride, col_off, width);
+
+    assert_eq!(got.len(), want.len(), "{name}: output length mismatch");
+    let mut max_diff = 0f32;
+    let mut first_bad: Option<usize> = None;
+    for (i, (g, w)) in got.iter().zip(want.iter()).enumerate() {
+        let diff = (g - w).abs();
+        if diff > max_diff { max_diff = diff; }
+        if diff > 1e-6 && first_bad.is_none() { first_bad = Some(i); }
+    }
+    eprintln!("[{name}] s={s} stride={stride} col_off={col_off} width={width} n={} max|GPU-CPU|={max_diff:e}", s * width);
+    if let Some(i) = first_bad {
+        panic!(
+            "{name}: mismatch at flat index {i} (ti={}, ci={}): got={} want={}",
+            i / width, i % width, got[i], want[i]
+        );
+    }
+}
+
+/// Deliberately non-64-aligned shape: s*width = 3*72 = 216, not a multiple of
+/// 64. This is the Laguna-style case (head counts 48/72) that exposed the
+/// launch-geometry bug: with the old `n / 64` grid, only 192 of 216 elements
+/// were ever dispatched, so the last 24 (rows 2's tail) came back as
+/// zero-initialized garbage instead of the real column values.
+#[test]
+fn strided_col_copy_non_64_aligned_shape() {
+    run_case("non_64_aligned", 3, 100, 4, 72);
+}
+
+/// Control: a 64-aligned shape (s*width divisible by 64) exercising the
+/// previously-working path, to confirm the div_ceil + in-kernel guard fix
+/// didn't regress the common (aligned) case.
+#[test]
+fn strided_col_copy_64_aligned_control() {
+    run_case("64_aligned_control", 4, 200, 8, 64);
+}
+
+/// Micro-repro for the prefill layer-1 fault: f16 table gather at the exact
+/// failing shapes (table [2, 3072] f16 from a cast, indices [20] u32 of 0/1).
+#[test]
+fn gather_f16_prefill_shapes() {
+    use ffai_core::{DType, Tensor};
+    use ffai_cuda::CudaDevice;
+    let Some(dev) = CudaDevice::create().expect("cuda") else { eprintln!("no CUDA, skip"); return; };
+    let d = dev.as_ref();
+    let t = 2usize;
+    let hidden = 3072usize;
+    let x: Vec<f32> = (0..t * hidden).map(|i| (i as f32) * 0.001 - 3.0).collect();
+    let xb: Vec<u8> = x.iter().flat_map(|v| v.to_le_bytes()).collect();
+    let xf32 = Tensor::new(d.upload(&xb).unwrap(), vec![t, hidden], DType::F32);
+    let xf16 = ffai_ops::cast_f32_f16(d, &xf32).unwrap();
+    let idx: Vec<u32> = vec![0, 0, 1, 1, 1, 0, 0, 1, 0, 1, 0, 0, 1, 1, 1, 0, 1, 1, 0, 0];
+    let ib: Vec<u8> = idx.iter().flat_map(|v| v.to_le_bytes()).collect();
+    let ti = Tensor::new(d.upload(&ib).unwrap(), vec![idx.len()], DType::U32);
+    let out = ffai_ops::gather(d, &xf16, &ti).expect("gather f16 dispatch");
+    d.synchronize().expect("post-gather sync");
+    let mut ob = vec![0u8; idx.len() * hidden * 2];
+    d.download(out.buffer.as_ref(), &mut ob).expect("download");
+    eprintln!("gather f16 prefill shapes: OK (first half-word {:02x}{:02x})", ob[0], ob[1]);
+}