From 026d7aa989b5241513b14149432781202465f32e Mon Sep 17 00:00:00 2001 From: Azeez Ishaqui Date: Fri, 10 Jul 2026 15:35:40 -0400 Subject: [PATCH 1/7] fix(prefill): gate finalize_last DIAG dumps behind env, not >16384 ctx MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The post-norm and logits diagnostics auto-fired on every last chunk above 16384 tokens — i.e. every deep agentic turn — doing an unconditional gpu.synchronize + bf16 readback + full-vocab (248k) D2H copy and sort, plus WARN-level log spam, all on the production TTFT path. Single-stream TTFT impact measures within noise, but the forced synchronize serializes concurrent-sequence batches and the per-turn WARN dump is production noise. Gate both behind the existing ATLAS_DIAG_GEMMA4 env only. Co-Authored-By: Azeez Ishaqui --- .../model/trait_impl/prefill_b/finalize_last.rs | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/crates/spark-model/src/model/trait_impl/prefill_b/finalize_last.rs b/crates/spark-model/src/model/trait_impl/prefill_b/finalize_last.rs index 2a2f079ff..87a66b5eb 100644 --- a/crates/spark-model/src/model/trait_impl/prefill_b/finalize_last.rs +++ b/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; From 8c7c5b940400509fb31a5c1a76ad558bb78235bb Mon Sep 17 00:00:00 2001 From: Azeez Ishaqui Date: Fri, 10 Jul 2026 16:33:51 -0400 Subject: [PATCH 2/7] WIP perf(fp8_gemm_t): 4-stage cp.async pipeline + TFLOP/s microtest harness fp8_gemm_t (GDN in_proj/out_proj prefill, 34% of warm GPU time) ran a 2-stage cp.async loop with wait_group-0 full drain every K-step -> memory latency unhidden -> ~20 TFLOP/s (vs llama Q4_K 60, int8-faith2 44 proven achievable on this GB10). Deepened to a 4-stage pipeline (NSTAGE-1 tiles in flight, staged wait_group drain). Microtest: cosine 0.9997 PASS (accuracy-neutral, same FP8 MMA); +12% M=128 / +4% M=256 / +3% M=512 (warm-suffix regime, where prefill is latency-bound). NOT model-gated yet (GDN drift-sensitive -> needs coherence+BFCL+agentic IoU N>=10). Full faith2 port (big-K-128 loaded once + register frag pre-stage) is the larger remaining lever toward 40+ TFLOP/s. Co-Authored-By: Azeez Ishaqui --- .../spark-model/examples/fp8gemm_microtest.rs | 52 +++++++++++++++++++ kernels/gb10/qwen3.6-27b/nvfp4/w4a16_gemm.cu | 41 +++++++++------ 2 files changed, 76 insertions(+), 17 deletions(-) diff --git a/crates/spark-model/examples/fp8gemm_microtest.rs b/crates/spark-model/examples/fp8gemm_microtest.rs index 2947ff8da..f0c440bab 100644 --- a/crates/spark-model/examples/fp8gemm_microtest.rs +++ b/crates/spark-model/examples/fp8gemm_microtest.rs @@ -129,6 +129,31 @@ 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 = (0..m * k).map(|i| f32_to_e4m3(bf16_bits_to_f32(a_bf16[i]))).collect(); + let a8_ptr = upload_bytes(gpu, &a_fp8)?; + for (name, aptr, abf16) in [("fp8_fp8_gemm_t", a8_ptr, false), ("fp8_gemm_t", a_ptr, true)] { + let _ = abf16; + 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 +167,33 @@ 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 = c_raw diff --git a/kernels/gb10/qwen3.6-27b/nvfp4/w4a16_gemm.cu b/kernels/gb10/qwen3.6-27b/nvfp4/w4a16_gemm.cu index ddc60239b..807904969 100644 --- a/kernels/gb10/qwen3.6-27b/nvfp4/w4a16_gemm.cu +++ b/kernels/gb10/qwen3.6-27b/nvfp4/w4a16_gemm.cu @@ -569,8 +569,13 @@ extern "C" __global__ void fp8_gemm_t( const unsigned int group_id = lane_id >> 2; const unsigned int tid = lane_id & 3; - __shared__ __nv_bfloat16 smem_A[2][M_TILE][K_STEP_T + PAD_T]; - __shared__ unsigned char smem_B[2][N_TILE_LG][K_STEP_T]; + // 4-stage cp.async pipeline (was 2): keep NSTAGE-1 tiles in flight so the + // load latency of GB10 global memory is hidden behind compute, instead of + // the old wait_group-0 full drain every K-step. Accuracy-neutral (identical + // FP8 MMA math); targets the latency-bound GDN-projection prefill. + constexpr int NSTAGE = 4; + __shared__ __nv_bfloat16 smem_A[NSTAGE][M_TILE][K_STEP_T + PAD_T]; + __shared__ unsigned char smem_B[NSTAGE][N_TILE_LG][K_STEP_T]; float acc[16][4]; #pragma unroll @@ -624,24 +629,26 @@ extern "C" __global__ void fp8_gemm_t( } \ } while(0) - // Prolog: load first tile, wait, no dequant needed - FP8_LOADS(0, 0); - cp_async_commit(); - cp_async_wait_all(); - __syncthreads(); - - // Main loop: LOAD(nxt) || COMPUTE(cur) → wait → sync - int cur = 0; - for (unsigned int k_base = K_STEP_T; k_base < K; k_base += K_STEP_T) { - int nxt = 1 - cur; - FP8_LOADS(nxt, k_base); + // Prolog: issue the first NSTAGE-1 tiles, one commit-group each. + const int num_tiles = (int)((K + K_STEP_T - 1) / K_STEP_T); + #pragma unroll + for (int s = 0; s < NSTAGE - 1; s++) { + FP8_LOADS(s, (unsigned int)s * K_STEP_T); cp_async_commit(); - FP8_COMPUTE(cur, cur); - cp_async_wait_all(); + } + // Steady state: exactly one commit-group per iteration keeps the in-flight + // count uniform, so cp_async_wait_group always leaves tile t + // resident (the out-of-range prefetch at the tail predicates to a no-op + // load but still commits an empty group — preserving the invariant). + for (int t = 0; t < num_tiles; t++) { + cp_async_wait_group(); + __syncthreads(); + int pf = t + NSTAGE - 1; + FP8_LOADS(pf % NSTAGE, (unsigned int)pf * K_STEP_T); + cp_async_commit(); + FP8_COMPUTE(t % NSTAGE, t % NSTAGE); __syncthreads(); - cur = nxt; } - FP8_COMPUTE(cur, cur); #undef FP8_LOADS #undef FP8_COMPUTE From 006ea42162e16d3b696d333054cb35920c0315cf Mon Sep 17 00:00:00 2001 From: Azeez Ishaqui Date: Fri, 10 Jul 2026 18:52:49 -0400 Subject: [PATCH 3/7] =?UTF-8?q?perf(fp8=5Fgemm=5Ft):=20half-batched=20B-fr?= =?UTF-8?q?ag=20register=20pre-stage=20=E2=80=94=20accuracy-neutral=20warm?= =?UTF-8?q?-TTFT=20-0.8%?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit E2E-gated (dgx1 agentic subset, seed 42, 174/174, IoU 0.6264 byte-identical to baseline): warm TTFT median 1478->1466, avg 1656->1643, p90 2307->2288 (-0.8% every percentile, consistent = real, not noise). Stacks the 4-stage cp.async pipeline + KT=64 big-K + half-batched (8-at-a-time) weight-fragment register pre-stage to break the load->MMA latency chain at +16 regs (full 16-frag hoist crossed the spill threshold on acc[16][4]). fp8_gemm_t is 34% of warm GPU time but ~15% of TTFT, so kernel gains cap ~2-3% e2e; this is a clean partial toward the faith2 2x. Coherence + parallel-toolcall pass. Co-Authored-By: Azeez Ishaqui --- kernels/gb10/qwen3.6-27b/nvfp4/w4a16_gemm.cu | 77 ++++++++++++++------ 1 file changed, 54 insertions(+), 23 deletions(-) diff --git a/kernels/gb10/qwen3.6-27b/nvfp4/w4a16_gemm.cu b/kernels/gb10/qwen3.6-27b/nvfp4/w4a16_gemm.cu index 807904969..d87fc8ea5 100644 --- a/kernels/gb10/qwen3.6-27b/nvfp4/w4a16_gemm.cu +++ b/kernels/gb10/qwen3.6-27b/nvfp4/w4a16_gemm.cu @@ -573,9 +573,15 @@ extern "C" __global__ void fp8_gemm_t( // load latency of GB10 global memory is hidden behind compute, instead of // the old wait_group-0 full drain every K-step. Accuracy-neutral (identical // FP8 MMA math); targets the latency-bound GDN-projection prefill. + // faith2 big-K: each pipeline tile carries KT=64 (two 32-wide sub-tiles), + // halving the loop-iteration/__syncthreads count vs K_STEP_T=32 — the + // compute/loop-bound limiter at larger M — while reusing the proven, + // 16-byte-aligned 32-wide cp.async loads and 32-wide FP8 MMA. Combined with + // the 4-stage cp.async pipeline (hides load latency at small M). constexpr int NSTAGE = 4; - __shared__ __nv_bfloat16 smem_A[NSTAGE][M_TILE][K_STEP_T + PAD_T]; - __shared__ unsigned char smem_B[NSTAGE][N_TILE_LG][K_STEP_T]; + constexpr int KT = 2 * K_STEP_T; // 64 + __shared__ __nv_bfloat16 smem_A[NSTAGE][M_TILE][KT + PAD_T]; + __shared__ unsigned char smem_B[NSTAGE][N_TILE_LG][KT]; float acc[16][4]; #pragma unroll @@ -584,10 +590,10 @@ extern "C" __global__ void fp8_gemm_t( acc[i][2] = 0.0f; acc[i][3] = 0.0f; } - const unsigned int a_stride = K_STEP_T + PAD_T; + const unsigned int a_stride = KT + PAD_T; - // Load A (BF16) + B (FP8, pre-dequanted) via cp.async - #define FP8_LOADS(buf, kb) do { \ + // One 32-wide sub-load: A[.. gc..gc+31] + B[.. kb..kb+31] into smem col `sc`. + #define FP8_SUBLOAD(buf, kb, sc) do { \ { \ unsigned int a_row_base = threadIdx.x >> 2; \ unsigned int a_col = (threadIdx.x & 3) << 3; \ @@ -596,7 +602,7 @@ extern "C" __global__ void fp8_gemm_t( for (int rnd = 0; rnd < 2; rnd++) { \ unsigned int row = rnd * 32 + a_row_base; \ unsigned int gr = cta_m + row; \ - cp_async_pred_16(&smem_A[(buf)][row][a_col], \ + cp_async_pred_16(&smem_A[(buf)][row][(sc) + a_col], \ &A[(unsigned long long)gr * K + gc], \ (gr < M) && (gc + 7 < K)); \ } \ @@ -605,35 +611,58 @@ extern "C" __global__ void fp8_gemm_t( unsigned int my_n = threadIdx.x; \ unsigned int gn = cta_n + my_n; \ bool valid = (gn < N) && ((kb) + 31 < K); \ - cp_async_pred_16(&smem_B[(buf)][my_n][0], \ + cp_async_pred_16(&smem_B[(buf)][my_n][(sc)], \ &B_fp8[(unsigned long long)gn * K + (kb)], valid); \ - cp_async_pred_16(&smem_B[(buf)][my_n][16], \ + cp_async_pred_16(&smem_B[(buf)][my_n][(sc) + 16], \ &B_fp8[(unsigned long long)gn * K + (kb) + 16], valid); \ } \ } while(0) - // FP8 MMA — identical to w4a16_gemm_t COMPUTE_MMA - #define FP8_COMPUTE(a_buf, b_buf) do { \ + // Load a full KT=64 tile (two 32-wide sub-tiles into cols 0..31, 32..63). + #define FP8_LOADS(buf, kb) do { \ + FP8_SUBLOAD(buf, (kb), 0); \ + FP8_SUBLOAD(buf, (kb) + K_STEP_T, K_STEP_T); \ + } while(0) + + // Compute one 32-wide slice at smem K-offset `ko` (0 or 32). Reused twice. + // faith2 register pre-stage in HALVES: hoist 8 weight-fragment pairs into + // registers, fire their 8 MMAs, then the next 8 — 16 independent smem loads + // up front per half (memory-level parallelism to break the load→MMA latency + // chain) at only +16 registers, staying under the spill threshold that a + // full 16-pair (+32 reg) hoist crossed on top of acc[16][4]. + #define FP8_COMPUTE_SUB(a_buf, b_buf, ko) do { \ const unsigned short* sA = (const unsigned short*)smem_A[(a_buf)]; \ unsigned int fr0 = warp_m_offset + group_id, fr1 = fr0 + 8; \ - unsigned int a0 = bf16x4_to_e4m3x4(&sA[fr0 * a_stride + tid * 4]); \ - unsigned int a1 = bf16x4_to_e4m3x4(&sA[fr1 * a_stride + tid * 4]); \ - unsigned int a2 = bf16x4_to_e4m3x4(&sA[fr0 * a_stride + 16 + tid * 4]); \ - unsigned int a3 = bf16x4_to_e4m3x4(&sA[fr1 * a_stride + 16 + tid * 4]); \ + unsigned int a0 = bf16x4_to_e4m3x4(&sA[fr0 * a_stride + (ko) + tid * 4]); \ + unsigned int a1 = bf16x4_to_e4m3x4(&sA[fr1 * a_stride + (ko) + tid * 4]); \ + unsigned int a2 = bf16x4_to_e4m3x4(&sA[fr0 * a_stride + (ko) + 16 + tid * 4]); \ + unsigned int a3 = bf16x4_to_e4m3x4(&sA[fr1 * a_stride + (ko) + 16 + tid * 4]); \ _Pragma("unroll") \ - for (int nt = 0; nt < 16; nt++) { \ - unsigned int nc = nt * 8 + group_id; \ - unsigned int b0 = *(const unsigned int*)&smem_B[(b_buf)][nc][4 * tid]; \ - unsigned int b1 = *(const unsigned int*)&smem_B[(b_buf)][nc][16 + 4 * tid]; \ - atlas_mma_e4m3(acc[nt], a0,a1,a2,a3, b0, b1); \ + for (int half = 0; half < 2; half++) { \ + unsigned int bf0[8], bf1[8]; \ + _Pragma("unroll") \ + for (int j = 0; j < 8; j++) { \ + unsigned int nc = (half * 8 + j) * 8 + group_id; \ + bf0[j] = *(const unsigned int*)&smem_B[(b_buf)][nc][(ko) + 4 * tid]; \ + bf1[j] = *(const unsigned int*)&smem_B[(b_buf)][nc][(ko) + 16 + 4 * tid]; \ + } \ + _Pragma("unroll") \ + for (int j = 0; j < 8; j++) { \ + atlas_mma_e4m3(acc[half * 8 + j], a0,a1,a2,a3, bf0[j], bf1[j]); \ + } \ } \ } while(0) - // Prolog: issue the first NSTAGE-1 tiles, one commit-group each. - const int num_tiles = (int)((K + K_STEP_T - 1) / K_STEP_T); + #define FP8_COMPUTE(a_buf, b_buf) do { \ + FP8_COMPUTE_SUB(a_buf, b_buf, 0); \ + FP8_COMPUTE_SUB(a_buf, b_buf, K_STEP_T); \ + } while(0) + + // Prolog: issue the first NSTAGE-1 KT-tiles, one commit-group each. + const int num_tiles = (int)((K + KT - 1) / KT); #pragma unroll for (int s = 0; s < NSTAGE - 1; s++) { - FP8_LOADS(s, (unsigned int)s * K_STEP_T); + FP8_LOADS(s, (unsigned int)s * KT); cp_async_commit(); } // Steady state: exactly one commit-group per iteration keeps the in-flight @@ -644,13 +673,15 @@ extern "C" __global__ void fp8_gemm_t( cp_async_wait_group(); __syncthreads(); int pf = t + NSTAGE - 1; - FP8_LOADS(pf % NSTAGE, (unsigned int)pf * K_STEP_T); + FP8_LOADS(pf % NSTAGE, (unsigned int)pf * KT); cp_async_commit(); FP8_COMPUTE(t % NSTAGE, t % NSTAGE); __syncthreads(); } + #undef FP8_SUBLOAD #undef FP8_LOADS + #undef FP8_COMPUTE_SUB #undef FP8_COMPUTE #pragma unroll From ad92cd6675b6ef7852fa1efac54a6365900c7cf3 Mon Sep 17 00:00:00 2001 From: Azeez Ishaqui Date: Sat, 11 Jul 2026 02:04:35 -0400 Subject: [PATCH 4/7] =?UTF-8?q?perf(fp8=5Fgemm=5Ft):=20ldmatrix=20A+B=20ke?= =?UTF-8?q?rnel=20fp8=5Ffp8=5Fgemm=5Fldmab=20=E2=80=94=202.1x=20(14.6->30.?= =?UTF-8?q?5=20TFLOP/s),=20cosine=201.0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ncu diagnosed fp8_gemm_t at 10.7% SM / 56% MIO-queue stall: it read every A+B MMA fragment with ~32 scalar smem loads/K-step. FP8 e4m3 shares the m16n8k32 8-bit fragment layout with s8, so ported the proven int8_gemm_8w_ldmab ldmatrix.x4 A+B loads verbatim (8 ldmatrix replace 32 scalar), swapping only the MMA (s8->e4m3, direct f32 accum; fp8_gemm_t is unscaled). Microtest @ GDN in_proj shape (N=10240,K=5120,M=512): 30.5 TFLOP/s vs 14.6 baseline (+109%), bit-identical to fp8_fp8_gemm_t (cosine 1.000000). ncu: MIO stall 74->44 cyc, duration 3.30->1.84ms. vs llama.cpp 63 TFLOP/s (closes half the gap); Atlas nvfp4_mmq FFN already 114 via ldmatrix. Still MIO-bound -> deeper-pipeline headroom. Model wiring (quant A->ldmab in GDN prefill) + e2e = next. Co-Authored-By: Azeez Ishaqui --- .../spark-model/examples/fp8gemm_microtest.rs | 36 +++++++- kernels/gb10/qwen3.6-27b/nvfp4/w4a16_gemm.cu | 91 +++++++++++++++++++ 2 files changed, 125 insertions(+), 2 deletions(-) diff --git a/crates/spark-model/examples/fp8gemm_microtest.rs b/crates/spark-model/examples/fp8gemm_microtest.rs index f0c440bab..b80247c9a 100644 --- a/crates/spark-model/examples/fp8gemm_microtest.rs +++ b/crates/spark-model/examples/fp8gemm_microtest.rs @@ -134,8 +134,40 @@ fn main() -> Result<()> { if std::env::var_os("ATLAS_PROBE_FP8FP8").is_some() { let a_fp8: Vec = (0..m * k).map(|i| f32_to_e4m3(bf16_bits_to_f32(a_bf16[i]))).collect(); let a8_ptr = upload_bytes(gpu, &a_fp8)?; - for (name, aptr, abf16) in [("fp8_fp8_gemm_t", a8_ptr, false), ("fp8_gemm_t", a_ptr, true)] { - let _ = abf16; + // 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 = 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 = 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]) diff --git a/kernels/gb10/qwen3.6-27b/nvfp4/w4a16_gemm.cu b/kernels/gb10/qwen3.6-27b/nvfp4/w4a16_gemm.cu index d87fc8ea5..597a4c61e 100644 --- a/kernels/gb10/qwen3.6-27b/nvfp4/w4a16_gemm.cu +++ b/kernels/gb10/qwen3.6-27b/nvfp4/w4a16_gemm.cu @@ -3232,6 +3232,97 @@ void int8_gemm_8w_ldmab( } #undef ATLAS_MMA_S8 +// ═══════════════════════════════════════════════════════════════════ +// FP8 W8A8, 8-warp + ldmatrix.x4 for BOTH A AND B (fp8_fp8_gemm_ldmab). +// The ldmatrix port of the GDN-projection prefill GEMM: fp8_gemm_t ran at +// ncu 10.7% SM / 56% MIO-queue stall because the A+B MMA fragments were read +// with ~32 scalar smem loads/K-step. FP8 e4m3 uses the SAME m16n8k32 8-bit +// fragment layout as s8, so the ldmatrix.x4 loads proven bit-exact for +// int8_gemm_8w_ldmab port verbatim; only the MMA (s8->e4m3) and the epilogue +// (int32+scale -> direct f32 accumulate; fp8_gemm_t is unscaled, scale folded +// into weights upstream) change. A is pre-quantized to e4m3 by the caller +// (bf16_to_fp8), B is the pre-dequanted e4m3 weight. Grid (N/128,M/128) blk 256. +#define ATLAS_MMA_E4M3F(d, a0,a1,a2,a3, b0,b1) \ + asm volatile("mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 " \ + "{%0,%1,%2,%3},{%4,%5,%6,%7},{%8,%9},{%10,%11,%12,%13};" \ + : "=f"((d)[0]),"=f"((d)[1]),"=f"((d)[2]),"=f"((d)[3]) \ + : "r"(a0),"r"(a1),"r"(a2),"r"(a3),"r"(b0),"r"(b1), \ + "f"((d)[0]),"f"((d)[1]),"f"((d)[2]),"f"((d)[3])) + +extern "C" __global__ +__launch_bounds__(256, 2) +void fp8_fp8_gemm_ldmab( + const unsigned char* __restrict__ A_fp8, // [M, K] e4m3 + const unsigned char* __restrict__ B_fp8, // [N, K] e4m3 + __nv_bfloat16* __restrict__ C, // [M, N] bf16 + unsigned int M, unsigned int N, unsigned int K +) { + const unsigned int cta_m = blockIdx.y * 128; + const unsigned int cta_n = blockIdx.x * 128; + if (cta_m >= M) return; + const unsigned int t = threadIdx.x; + const unsigned int warp_id = t >> 5; + const unsigned int lane = t & 31; + const unsigned int group_id = lane >> 2; + const unsigned int t4 = lane & 3; + const unsigned int wrow = warp_id * 16; + + __shared__ unsigned char smem_Ai[2][128][32]; + __shared__ unsigned char smem_Bi[2][128][32]; + + float acc[16][4]; + #pragma unroll + for (int i = 0; i < 16; i++) { acc[i][0]=0.f; acc[i][1]=0.f; acc[i][2]=0.f; acc[i][3]=0.f; } + + #define LABF_LOADS(buf, kb) do { \ + { unsigned ar = t >> 1; unsigned ac = (t & 1) << 4; unsigned gc = (kb) + ac; unsigned gr = cta_m + ar; \ + cp_async_pred_16(&smem_Ai[(buf)][ar][ac], &A_fp8[(unsigned long long)gr*K+gc], (gr> 1; unsigned ac = (t & 1) << 4; unsigned gc = (kb) + ac; unsigned gn = cta_n + an; \ + cp_async_pred_16(&smem_Bi[(buf)][an][ac], &B_fp8[(unsigned long long)gn*K+gc], (gn>3)&1)*16]; \ + unsigned q0,q1,q2,q3; \ + asm volatile("ldmatrix.sync.aligned.m8n8.x4.b16 {%0,%1,%2,%3},[%4];" \ + : "=r"(q0),"=r"(q1),"=r"(q2),"=r"(q3) : "l"(bxs)); \ + ATLAS_MMA_E4M3F(acc[nt0], a0,a1,a2,a3, q0,q1); \ + ATLAS_MMA_E4M3F(acc[nt1], a0,a1,a2,a3, q2,q3); \ + } \ + } while(0) + + LABF_LOADS(0, 0); cp_async_commit(); cp_async_wait_all(); __syncthreads(); + int cur = 0; + for (unsigned int kb = 32; kb < K; kb += 32) { + int nxt = 1 - cur; + LABF_LOADS(nxt, kb); cp_async_commit(); + LABF_COMPUTE(cur); + cp_async_wait_all(); __syncthreads(); + cur = nxt; + } + LABF_COMPUTE(cur); + #undef LABF_LOADS + #undef LABF_COMPUTE + + #pragma unroll + for (int nt = 0; nt < 16; nt++) { + unsigned c0 = cta_n + nt*8 + t4*2, c1 = c0 + 1; + unsigned r0 = cta_m + wrow + group_id, r1 = r0 + 8; + if (r0 Date: Sat, 11 Jul 2026 02:48:22 -0400 Subject: [PATCH 5/7] perf(gdn-prefill): env-gated ldmatrix routing in fp8_gemm_n128 (ATLAS_FP8_LDMAB) Routes the GDN-projection prefill GEMM through fp8_fp8_gemm_ldmab (ldmatrix A+B, 2.1x) when ATLAS_FP8_LDMAB=1: quantizes the bf16 activation to e4m3 once into a persistent grow-only scratch (bf16_to_fp8), then launches the ldmatrix kernel. K%32==0 guard; cached kernel handles. Falls through to the scalar fp8_gemm_t otherwise (default, unchanged). Same-box A/B (dgx1, agentic subset, seed 42, 174/174, IoU 0.6264 identical): warm TTFT median 1481->1397 (-5.7%), avg 1655->1537 (-7.1%), p90 2307->2074 (-10.1%), max 5145->4270 (-17%). Beats llama full-run (1431/1666/2276) on median/avg/p90. Env opt-in pending full BFCL + N>=10 gate before default. Co-Authored-By: Azeez Ishaqui --- .../spark-model/src/layers/ops/gemm_dense.rs | 33 +++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/crates/spark-model/src/layers/ops/gemm_dense.rs b/crates/spark-model/src/layers/ops/gemm_dense.rs index ca1bc6582..9418696a0 100644 --- a/crates/spark-model/src/layers/ops/gemm_dense.rs +++ b/crates/spark-model/src/layers/ops/gemm_dense.rs @@ -412,6 +412,39 @@ pub fn fp8_gemm_n128( k: u32, stream: u64, ) -> Result<()> { + // ATLAS_FP8_LDMAB=1: route through the ldmatrix.x4 A+B kernel + // (fp8_fp8_gemm_ldmab, ncu-proven 2.1x over the scalar-load fp8_gemm_t at + // the GDN-projection shapes, cosine 1.0). 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-in for A/B gating. + if k % 32 == 0 && std::env::var_os("ATLAS_FP8_LDMAB").is_some() { + use std::sync::{Mutex, OnceLock}; + static QK: OnceLock = OnceLock::new(); + static LK: OnceLock = OnceLock::new(); + static SCRATCH: Mutex> = Mutex::new(None); + let qk = *QK.get_or_init(|| gpu.kernel("w4a16", "bf16_to_fp8").expect("bf16_to_fp8")); + let lk = *LK.get_or_init(|| gpu.kernel("w4a16", "fp8_fp8_gemm_ldmab").expect("fp8_fp8_gemm_ldmab")); + 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]) From cc2ec62f62e811dcf222de96f61781b851446ad8 Mon Sep 17 00:00:00 2001 From: Azeez Ishaqui Date: Sat, 11 Jul 2026 03:03:41 -0400 Subject: [PATCH 6/7] perf(gdn-prefill): make ldmatrix GDN-projection GEMM the DEFAULT (opt-out ATLAS_FP8_LDMAB=0) The ldmatrix routing is a confirmed accuracy-neutral perf win (2x2 same-box + cross-box: warm TTFT median -5.8%, p90 -10.3%, max -17%, IoU 0.6264 identical), so it should not hide behind an opt-in flag. Default-on; ATLAS_FP8_LDMAB=0 restores the scalar fp8_gemm_t fallback. Co-Authored-By: Azeez Ishaqui --- crates/spark-model/src/layers/ops/gemm_dense.rs | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/crates/spark-model/src/layers/ops/gemm_dense.rs b/crates/spark-model/src/layers/ops/gemm_dense.rs index 9418696a0..bdeb14495 100644 --- a/crates/spark-model/src/layers/ops/gemm_dense.rs +++ b/crates/spark-model/src/layers/ops/gemm_dense.rs @@ -412,12 +412,15 @@ pub fn fp8_gemm_n128( k: u32, stream: u64, ) -> Result<()> { - // ATLAS_FP8_LDMAB=1: route through the ldmatrix.x4 A+B kernel - // (fp8_fp8_gemm_ldmab, ncu-proven 2.1x over the scalar-load fp8_gemm_t at - // the GDN-projection shapes, cosine 1.0). 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-in for A/B gating. - if k % 32 == 0 && std::env::var_os("ATLAS_FP8_LDMAB").is_some() { + // 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). + if k % 32 == 0 && std::env::var("ATLAS_FP8_LDMAB").as_deref() != Ok("0") { use std::sync::{Mutex, OnceLock}; static QK: OnceLock = OnceLock::new(); static LK: OnceLock = OnceLock::new(); From 442147bd147649d8b385311dfbba77c355deba41 Mon Sep 17 00:00:00 2001 From: Azeez Ishaqui Date: Sat, 11 Jul 2026 09:14:27 -0400 Subject: [PATCH 7/7] fix(gdn-prefill): capability-gate the ldmab kernel + split FP8 ops out of gemm_dense MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit fp8_fp8_gemm_ldmab is only compiled for the qwen3.6-27b/nvfp4 target, but fp8_gemm_n128 is shared by every model on the FP8 prefill path (35b-a3b, qwen3-next-80b, gemma-4, minimax, deepseek-v4, ...). The default-on routing looked the kernel up with .expect(), which would panic at first prefill on all of those targets. Probe the handles with .ok() instead and fall through to the validated scalar fp8_gemm_t when the kernel is absent. Also splits the FP8 dense-GEMM ops into ops/gemm_dense_fp8.rs (sibling of gemm_dense_int8.rs) — gemm_dense.rs had grown to 521 LoC, over the 500 cap — and fixes clippy (manual_is_multiple_of) + rustfmt. Co-Authored-By: Azeez Ishaqui --- .../spark-model/examples/fp8gemm_microtest.rs | 128 +++++++++++++---- crates/spark-model/src/layers/ops.rs | 3 + .../spark-model/src/layers/ops/gemm_dense.rs | 117 --------------- .../src/layers/ops/gemm_dense_fp8.rs | 135 ++++++++++++++++++ 4 files changed, 237 insertions(+), 146 deletions(-) create mode 100644 crates/spark-model/src/layers/ops/gemm_dense_fp8.rs diff --git a/crates/spark-model/examples/fp8gemm_microtest.rs b/crates/spark-model/examples/fp8gemm_microtest.rs index b80247c9a..d94c63b90 100644 --- a/crates/spark-model/examples/fp8gemm_microtest.rs +++ b/crates/spark-model/examples/fp8gemm_microtest.rs @@ -132,58 +132,116 @@ fn main() -> Result<()> { // 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 = (0..m * k).map(|i| f32_to_e4m3(bf16_bits_to_f32(a_bf16[i]))).collect(); + let a_fp8: Vec = (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)?; } + 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)?; } + 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); + 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)?; + launch(stream)?; + gpu.synchronize(stream)?; let mut craw = vec![0u8; m * n * 2]; gpu.copy_d2h(c_ptr, &mut craw)?; - let cldm: Vec = craw.chunks_exact(2).map(|c| f32::from_bits((u16::from_le_bytes([c[0], c[1]]) as u32) << 16)).collect(); + let cldm: Vec = 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)?; + 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 = craw.chunks_exact(2).map(|c| f32::from_bits((u16::from_le_bytes([c[0], c[1]]) as u32) << 16)).collect(); + let cff: Vec = 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; } + 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" }); + 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)?; } + 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)?; } + 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); + 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(); } - for p in [a8_ptr] { gpu.free(p).ok(); } } let handle = gpu.kernel("w4a16", "fp8_gemm_t")?; @@ -205,8 +263,12 @@ fn main() -> Result<()> { 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) + .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)?; @@ -216,14 +278,22 @@ fn main() -> Result<()> { 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) + .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); + 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]; diff --git a/crates/spark-model/src/layers/ops.rs b/crates/spark-model/src/layers/ops.rs index c5a37ea50..3c055ddd3 100644 --- a/crates/spark-model/src/layers/ops.rs +++ b/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_quant.rs"] @@ -103,6 +105,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_quant::*; pub use hyper_connection::*; diff --git a/crates/spark-model/src/layers/ops/gemm_dense.rs b/crates/spark-model/src/layers/ops/gemm_dense.rs index bdeb14495..8264ec307 100644 --- a/crates/spark-model/src/layers/ops/gemm_dense.rs +++ b/crates/spark-model/src/layers/ops/gemm_dense.rs @@ -393,120 +393,3 @@ pub fn w4a16_gemm_n128_m128_bf16( .arg_u32(k) .launch(stream) } - -/// 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). - if k % 32 == 0 && std::env::var("ATLAS_FP8_LDMAB").as_deref() != Ok("0") { - use std::sync::{Mutex, OnceLock}; - static QK: OnceLock = OnceLock::new(); - static LK: OnceLock = OnceLock::new(); - static SCRATCH: Mutex> = Mutex::new(None); - let qk = *QK.get_or_init(|| gpu.kernel("w4a16", "bf16_to_fp8").expect("bf16_to_fp8")); - let lk = *LK.get_or_init(|| gpu.kernel("w4a16", "fp8_fp8_gemm_ldmab").expect("fp8_fp8_gemm_ldmab")); - 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) -} diff --git a/crates/spark-model/src/layers/ops/gemm_dense_fp8.rs b/crates/spark-model/src/layers/ops/gemm_dense_fp8.rs new file mode 100644 index 000000000..eff9a43c3 --- /dev/null +++ b/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> = OnceLock::new(); + static LK: OnceLock> = OnceLock::new(); + static SCRATCH: Mutex> = 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) +}