perf(gdn-prefill): ldmatrix A+B GDN-projection GEMM — 2.1x kernel, warm-TTFT median -5.8% / p90 -10.3%

crates/spark-model/examples/fp8gemm_microtest.rs

@@ -129,6 +129,121 @@ fn main() -> Result<()> {
     let b_ptr = upload_bytes(gpu, &b_fp8)?;
     let c_ptr = gpu.alloc(m * n * 2)?;
 
+    // A/B probe: time fp8_fp8_gemm_t (FP8 activation, no in-loop convert) vs
+    // fp8_gemm_t (BF16 activation). Accuracy-neutral routing candidate.
+    if std::env::var_os("ATLAS_PROBE_FP8FP8").is_some() {
+        let a_fp8: Vec<u8> = (0..m * k)
+            .map(|i| f32_to_e4m3(bf16_bits_to_f32(a_bf16[i])))
+            .collect();
+        let a8_ptr = upload_bytes(gpu, &a_fp8)?;
+        // ldmab: FP8 A, FP8 B, grid (N/128, M/128), block 256 — vs scalar-load fp8_gemm_t.
+        {
+            let h = gpu.kernel("w4a16", "fp8_fp8_gemm_ldmab")?;
+            let launch = |s| {
+                KernelLaunch::new(gpu, h)
+                    .grid([div_ceil(n as u32, 128), div_ceil(m as u32, 128), 1])
+                    .block([256, 1, 1])
+                    .arg_ptr(a8_ptr)
+                    .arg_ptr(b_ptr)
+                    .arg_ptr(c_ptr)
+                    .arg_u32(m as u32)
+                    .arg_u32(n as u32)
+                    .arg_u32(k as u32)
+                    .launch(s)
+            };
+            for _ in 0..8 {
+                launch(stream)?;
+            }
+            gpu.synchronize(stream)?;
+            let iters = 60u32;
+            let t0 = std::time::Instant::now();
+            for _ in 0..iters {
+                launch(stream)?;
+            }
+            gpu.synchronize(stream)?;
+            let secs = t0.elapsed().as_secs_f64() / iters as f64;
+            let flop = 2.0 * m as f64 * n as f64 * k as f64;
+            println!(
+                "PROBE fp8_fp8_gemm_ldmab: M={m} N={n} K={k} {:.3} ms {:.1} TFLOP/s",
+                secs * 1e3,
+                flop / secs / 1e12
+            );
+            // CORRECTNESS: ldmab vs fp8_fp8_gemm_t (both fp8xfp8, same math) — cosine ~1.0 iff layout correct.
+            launch(stream)?;
+            gpu.synchronize(stream)?;
+            let mut craw = vec![0u8; m * n * 2];
+            gpu.copy_d2h(c_ptr, &mut craw)?;
+            let cldm: Vec<f32> = craw
+                .chunks_exact(2)
+                .map(|c| f32::from_bits((u16::from_le_bytes([c[0], c[1]]) as u32) << 16))
+                .collect();
+            let hff = gpu.kernel("w4a16", "fp8_fp8_gemm_t")?;
+            KernelLaunch::new(gpu, hff)
+                .grid([div_ceil(n as u32, 128), div_ceil(m as u32, 64), 1])
+                .block([128, 1, 1])
+                .arg_ptr(a8_ptr)
+                .arg_ptr(b_ptr)
+                .arg_ptr(c_ptr)
+                .arg_u32(m as u32)
+                .arg_u32(n as u32)
+                .arg_u32(k as u32)
+                .launch(stream)?;
+            gpu.synchronize(stream)?;
+            gpu.copy_d2h(c_ptr, &mut craw)?;
+            let cff: Vec<f32> = craw
+                .chunks_exact(2)
+                .map(|c| f32::from_bits((u16::from_le_bytes([c[0], c[1]]) as u32) << 16))
+                .collect();
+            let (mut d, mut na, mut nb) = (0f64, 0f64, 0f64);
+            for i in 0..cldm.len() {
+                let (x, y) = (cldm[i] as f64, cff[i] as f64);
+                d += x * y;
+                na += x * x;
+                nb += y * y;
+            }
+            let cos = d / (na.sqrt() * nb.sqrt() + 1e-30);
+            println!(
+                "PROBE ldmab_vs_fp8fp8 cosine={cos:.6}  ({}=PASS)",
+                if cos >= 0.99 { "OK" } else { "FAIL" }
+            );
+        }
+        for (name, aptr) in [("fp8_fp8_gemm_t", a8_ptr), ("fp8_gemm_t", a_ptr)] {
+            let h = gpu.kernel("w4a16", name)?;
+            let launch = |s| {
+                KernelLaunch::new(gpu, h)
+                    .grid([div_ceil(n as u32, 128), div_ceil(m as u32, 64), 1])
+                    .block([128, 1, 1])
+                    .arg_ptr(aptr)
+                    .arg_ptr(b_ptr)
+                    .arg_ptr(c_ptr)
+                    .arg_u32(m as u32)
+                    .arg_u32(n as u32)
+                    .arg_u32(k as u32)
+                    .launch(s)
+            };
+            for _ in 0..8 {
+                launch(stream)?;
+            }
+            gpu.synchronize(stream)?;
+            let iters = 60u32;
+            let t0 = std::time::Instant::now();
+            for _ in 0..iters {
+                launch(stream)?;
+            }
+            gpu.synchronize(stream)?;
+            let secs = t0.elapsed().as_secs_f64() / iters as f64;
+            let flop = 2.0 * m as f64 * n as f64 * k as f64;
+            println!(
+                "PROBE {name}: M={m} N={n} K={k} {:.3} ms {:.1} TFLOP/s",
+                secs * 1e3,
+                flop / secs / 1e12
+            );
+        }
+        for p in [a8_ptr] {
+            gpu.free(p).ok();
+        }
+    }
+
     let handle = gpu.kernel("w4a16", "fp8_gemm_t")?;
     KernelLaunch::new(gpu, handle)
         .grid([div_ceil(n as u32, 128), div_ceil(m as u32, 64), 1])
@@ -142,6 +257,45 @@ fn main() -> Result<()> {
         .launch(stream)?;
     gpu.synchronize(stream)?;
 
+    // Timed loop → TFLOP/s (accuracy-neutral throughput probe). Warmup 5, time 50.
+    {
+        for _ in 0..5 {
+            KernelLaunch::new(gpu, handle)
+                .grid([div_ceil(n as u32, 128), div_ceil(m as u32, 64), 1])
+                .block([128, 1, 1])
+                .arg_ptr(a_ptr)
+                .arg_ptr(b_ptr)
+                .arg_ptr(c_ptr)
+                .arg_u32(m as u32)
+                .arg_u32(n as u32)
+                .arg_u32(k as u32)
+                .launch(stream)?;
+        }
+        gpu.synchronize(stream)?;
+        let iters = 50u32;
+        let t0 = std::time::Instant::now();
+        for _ in 0..iters {
+            KernelLaunch::new(gpu, handle)
+                .grid([div_ceil(n as u32, 128), div_ceil(m as u32, 64), 1])
+                .block([128, 1, 1])
+                .arg_ptr(a_ptr)
+                .arg_ptr(b_ptr)
+                .arg_ptr(c_ptr)
+                .arg_u32(m as u32)
+                .arg_u32(n as u32)
+                .arg_u32(k as u32)
+                .launch(stream)?;
+        }
+        gpu.synchronize(stream)?;
+        let secs = t0.elapsed().as_secs_f64() / iters as f64;
+        let flop = 2.0 * m as f64 * n as f64 * k as f64;
+        println!(
+            "TIMING: M={m} N={n} K={k} {:.3} ms/iter {:.1} TFLOP/s",
+            secs * 1e3,
+            flop / secs / 1e12
+        );
+    }
+
     let mut c_raw = vec![0u8; m * n * 2];
     gpu.copy_d2h(c_ptr, &mut c_raw)?;
     let c_gpu: Vec<u16> = c_raw

crates/spark-model/src/layers/ops.rs

@@ -31,6 +31,8 @@ mod fp8_moe_batch_b;
 pub mod gdn_flashinfer;
 #[path = "ops/gemm_dense.rs"]
 mod gemm_dense;
+#[path = "ops/gemm_dense_fp8.rs"]
+mod gemm_dense_fp8;
 #[path = "ops/gemm_dense_int8.rs"]
 mod gemm_dense_int8;
 #[path = "ops/gemm_fp4.rs"]
@@ -111,6 +113,7 @@ pub use fp8_moe::*;
 pub use fp8_moe_batch_a::*;
 pub use fp8_moe_batch_b::*;
 pub use gemm_dense::*;
+pub use gemm_dense_fp8::*;
 pub use gemm_dense_int8::*;
 pub use gemm_fp4::*;
 pub use gemm_fp8_prefill::*;

crates/spark-model/src/layers/ops/gemm_dense_fp8.rs

@@ -0,0 +1,135 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+
+//! FP8 (E4M3) dense GEMM ops — split out of `gemm_dense.rs` to keep each file
+//! under the 500-LoC cap. Sibling of `gemm_dense_int8.rs`; re-exported through
+//! `ops::*`, so call sites keep using `ops::fp8_gemm_n128(..)` unchanged.
+
+use anyhow::Result;
+use spark_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
+use spark_runtime::kernel_args::{KernelLaunch, div_ceil};
+
+/// Pre-dequanted FP8 GEMM (prefill): C = A @ B_fp8.
+///
+/// A: [M, K] BF16, B_fp8: [N, K] FP8 E4M3 (pre-dequanted from NVFP4), C: [M, N] BF16.
+/// Eliminates runtime NVFP4→FP8 dequant — only LOAD + FP8 MMA per K step.
+///
+/// Grid: (ceil(N/128), ceil(M/64), 1)  Block: (128, 1, 1)
+#[allow(clippy::too_many_arguments)]
+pub fn fp8_gemm_n128(
+    gpu: &dyn GpuBackend,
+    kernel: KernelHandle,
+    input: DevicePtr,
+    b_fp8: DevicePtr,
+    output: DevicePtr,
+    m: u32,
+    n: u32,
+    k: u32,
+    stream: u64,
+) -> Result<()> {
+    // DEFAULT-ON: route the GDN-projection prefill GEMM through the ldmatrix.x4
+    // A+B kernel (fp8_fp8_gemm_ldmab). ncu-proven 2.1x over the scalar-load
+    // fp8_gemm_t (10.7%->higher SM, MIO stall 74->44 cyc), cosine 1.000000 vs
+    // fp8_fp8_gemm_t, and a confirmed same-box+cross-box e2e warm-TTFT win
+    // (median -5.8%, p90 -10.3%, IoU 0.6264 unchanged). Quantizes the bf16
+    // activation to e4m3 once into a persistent scratch, then launches the
+    // ldmatrix GEMM. K must be a multiple of 32 (the ldmab K-tile). Opt-OUT with
+    // ATLAS_FP8_LDMAB=0 (falls through to the scalar path below).
+    //
+    // CAPABILITY-GATED: `fp8_fp8_gemm_ldmab` is currently only compiled for the
+    // qwen3.6-27b/nvfp4 target, but this fn is shared by every model that takes the
+    // FP8 prefill path (35b-a3b, qwen3-next-80b, gemma-4, minimax, deepseek-v4, ...).
+    // So the handles are probed with `.ok()`, NOT `.expect()`: a target without the
+    // kernel transparently falls through to the validated scalar `fp8_gemm_t` below
+    // instead of panicking at first prefill.
+    if k.is_multiple_of(32) && std::env::var("ATLAS_FP8_LDMAB").as_deref() != Ok("0") {
+        use std::sync::{Mutex, OnceLock};
+        static QK: OnceLock<Option<KernelHandle>> = OnceLock::new();
+        static LK: OnceLock<Option<KernelHandle>> = OnceLock::new();
+        static SCRATCH: Mutex<Option<(DevicePtr, usize)>> = Mutex::new(None);
+        let qk = *QK.get_or_init(|| gpu.kernel("w4a16", "bf16_to_fp8").ok());
+        let lk = *LK.get_or_init(|| gpu.kernel("w4a16", "fp8_fp8_gemm_ldmab").ok());
+        if let (Some(qk), Some(lk)) = (qk, lk) {
+            let need = (m as usize) * (k as usize); // e4m3 bytes
+            let a8 = {
+                let mut g = SCRATCH.lock().unwrap();
+                if g.map(|(_, sz)| sz < need).unwrap_or(true) {
+                    let p = gpu.alloc(need)?; // grow-only; old ptr leaked (rare, per-run)
+                    *g = Some((p, need));
+                }
+                g.unwrap().0
+            };
+            bf16_to_fp8(gpu, qk, input, a8, m * k, stream)?;
+            return KernelLaunch::new(gpu, lk)
+                .grid([div_ceil(n, 128), div_ceil(m, 128), 1])
+                .block([256, 1, 1])
+                .arg_ptr(a8)
+                .arg_ptr(b_fp8)
+                .arg_ptr(output)
+                .arg_u32(m)
+                .arg_u32(n)
+                .arg_u32(k)
+                .launch(stream);
+        }
+    }
+    KernelLaunch::new(gpu, kernel)
+        .grid([div_ceil(n, 128), div_ceil(m, 64), 1])
+        .block([128, 1, 1])
+        .arg_ptr(input)
+        .arg_ptr(b_fp8)
+        .arg_ptr(output)
+        .arg_u32(m)
+        .arg_u32(n)
+        .arg_u32(k)
+        .launch(stream)
+}
+
+/// Pre-dequant NVFP4 → FP8 E4M3.  One-time conversion at model load.
+///
+/// Reads B_packed[N, K/2] + B_scale[N, K/GROUP_SIZE] + scale2 → B_fp8[N, K].
+///
+/// Grid: (ceil(N*K/2 / 256), 1, 1)  Block: (256, 1, 1)
+#[allow(clippy::too_many_arguments)]
+pub fn predequant_nvfp4_to_fp8(
+    gpu: &dyn GpuBackend,
+    kernel: KernelHandle,
+    b_packed: DevicePtr,
+    b_scale: DevicePtr,
+    scale2: f32,
+    b_fp8: DevicePtr,
+    n: u32,
+    k: u32,
+    stream: u64,
+) -> Result<()> {
+    let total = n * k / 2;
+    KernelLaunch::new(gpu, kernel)
+        .grid([div_ceil(total, 256), 1, 1])
+        .block([256, 1, 1])
+        .arg_ptr(b_packed)
+        .arg_ptr(b_scale)
+        .arg_f32(scale2)
+        .arg_ptr(b_fp8)
+        .arg_u32(n)
+        .arg_u32(k)
+        .launch(stream)
+}
+
+/// Convert BF16 activations to FP8 E4M3 for FP8×FP8 GEMM.
+///
+/// Grid: (ceil(total_elements/2 / 256), 1, 1)  Block: (256, 1, 1)
+pub fn bf16_to_fp8(
+    gpu: &dyn GpuBackend,
+    kernel: KernelHandle,
+    src: DevicePtr,
+    dst: DevicePtr,
+    total_elements: u32,
+    stream: u64,
+) -> Result<()> {
+    let threads_needed = total_elements / 2;
+    KernelLaunch::new(gpu, kernel)
+        .grid([div_ceil(threads_needed, 256), 1, 1])
+        .block([256, 1, 1])
+        .arg_ptr(src)
+        .arg_ptr(dst)
+        .arg_u32(total_elements)
+        .launch(stream)
+}

crates/spark-model/src/model/trait_impl/prefill_b/finalize_last.rs

@@ -142,10 +142,10 @@ impl TransformerModel {
             stream,
         )?;
 
-        // Diagnostic: post-norm hidden state
-        if (chunk_start + chunk_len) > 16384
-            || std::env::var("ATLAS_DIAG_GEMMA4").is_ok_and(|v| v == "1" || v == "true")
-        {
+        // Diagnostic: post-norm hidden state (env-gated; the former >16384-ctx
+        // auto-trigger fired a synchronize + readback on EVERY deep agentic
+        // turn's last chunk — pure TTFT waste on the warm path).
+        if std::env::var("ATLAS_DIAG_GEMMA4").is_ok_and(|v| v == "1" || v == "true") {
             self.gpu.synchronize(stream)?;
             let (vals, norm) = self.readback_bf16(normed, h.min(16))?;
             tracing::warn!(
@@ -241,10 +241,10 @@ impl TransformerModel {
             tracing::info!("ATLAS_NEMO_DUMP: top-10 logits = {top:?}");
         }
 
-        // Diagnostic: logits stats
-        if (chunk_start + chunk_len) > 16384
-            || std::env::var("ATLAS_DIAG_GEMMA4").is_ok_and(|v| v == "1" || v == "true")
-        {
+        // Diagnostic: logits stats (env-gated; see above — the former
+        // >16384-ctx auto-trigger did a synchronize + full-vocab (248k) D2H +
+        // sort on every deep turn's last chunk, all off the TTFT critical path).
+        if std::env::var("ATLAS_DIAG_GEMMA4").is_ok_and(|v| v == "1" || v == "true") {
             self.gpu.synchronize(stream)?;
             let logits_ptr = self.buffers.logits();
             let n_logits = self.config.vocab_size;