From a17f9710e010fedc972a0b9405ca809de6b19c1c Mon Sep 17 00:00:00 2001 From: heiheiha798 <2300012738@stu.pku.edu.cn> Date: Fri, 17 Jul 2026 14:34:06 +0800 Subject: [PATCH 1/2] feat(sm89): residual-fold GEMM epilogue (idiom validation, not wired) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add RESID template param to fp8_bs_gemm_device.cuh: D = bf16(acc) + resid, folding the residual add into the GEMM epilogue (no separate residual_add launch, no D HBM round-trip). if constexpr (RESID) dead-strips the residual branch, so RESID=false stays byte-identical to the baseline kernel. Adds 4 residual-fold tile variants (32x64 / 64x64 / 64x64_s1 / 128x128_s1) for the prefill down-proj shapes, bound bench-only under ENABLE_QWEN3_VL_DEV_KERNELS (additive; production dispatcher untouched). Phase 1 finding: prefill's residual add is ALREADY fused by residual_add_rms_norm_to_fp8 on every non-last layer, so residual-fold only saves ~1 [S,hidden] bf16 transfer (0 launches). Validated as the epilogue-fold idiom proof (cos 1.0 vs fp32 ref, 1-ULP reassociation diff) but NOT wired — real prize is silu-fold (Phase 2). Bench script included. --- csrc/gemm/fp8_block128_gemm_mma_sm89.cu | 58 +++++++++++++++++- csrc/gemm/fp8_block128_gemm_mma_sm89.cuh | 16 +++++ csrc/gemm/fp8_bs_gemm_device.cuh | 78 ++++++++++++++++++------ csrc/qwen3_vl_bindings.cpp | 10 +++ scripts/test_resid_fold_correctness.py | 75 +++++++++++++++++++++++ 5 files changed, 217 insertions(+), 20 deletions(-) create mode 100644 scripts/test_resid_fold_correctness.py diff --git a/csrc/gemm/fp8_block128_gemm_mma_sm89.cu b/csrc/gemm/fp8_block128_gemm_mma_sm89.cu index 8d39fd02..3624ae4e 100644 --- a/csrc/gemm/fp8_block128_gemm_mma_sm89.cu +++ b/csrc/gemm/fp8_block128_gemm_mma_sm89.cu @@ -50,10 +50,10 @@ int launch_(const void* A, const void* B, void* D, + (BM * SCALE_KTILE + SCALE_KTILE) * (int)sizeof(float); if (smem_bytes > 48 * 1024) { cudaFuncSetAttribute( - (const void*)&fp8_bs_gemm_kernel, + (const void*)&fp8_bs_gemm_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes); } - fp8_bs_gemm_kernel<<>>( + fp8_bs_gemm_kernel<<>>( reinterpret_cast(A), reinterpret_cast(B), act_scale, w_scale, @@ -63,6 +63,38 @@ int launch_(const void* A, const void* B, void* D, return (err == cudaSuccess) ? 0 : 1; } +// Residual-fold launch: D = bf16(acc) + resid, fusing the residual add into the +// GEMM epilogue (no separate residual_add launch, no D HBM round-trip). resid +// is [M, N] BF16 row-major, same layout as D. See fp8_bs_gemm_device.cuh. +template +int launch_resid_(const void* A, const void* B, void* D, + int M, int N, int K, const float* act_scale, + const float* w_scale, const void* resid, cudaStream_t s) +{ + constexpr int BK = 128; + constexpr int SCALE_KTILE = 8; + int grid_m = (M + BM - 1) / BM; + int grid_n = (N + BN - 1) / BN; + dim3 grid(grid_m, grid_n, 1); + dim3 block(W * 32, 1, 1); + int smem_bytes = STAGES * (BM + BN) * BK + + (BM * SCALE_KTILE + SCALE_KTILE) * (int)sizeof(float); + if (smem_bytes > 48 * 1024) { + cudaFuncSetAttribute( + (const void*)&fp8_bs_gemm_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes); + } + fp8_bs_gemm_kernel<<>>( + reinterpret_cast(A), + reinterpret_cast(B), + act_scale, w_scale, + reinterpret_cast<__nv_bfloat16*>(D), + M, N, K, + reinterpret_cast(resid)); + cudaError_t err = cudaGetLastError(); + return (err == cudaSuccess) ? 0 : 1; +} + } // namespace #define DEFINE(NAME, BM, BN, W, S, MB) \ @@ -71,6 +103,17 @@ int launch_(const void* A, const void* B, void* D, return launch_(A, B, D, M, N, K, act_scale, w_scale, s);\ } +// Residual-fold variants (suffix _resid). D = bf16(acc) + resid. Only the +// tiles the prefill down-proj actually selects are defined; additive — the +// non-resid DEFINE list above is unchanged. +#define DEFINE_RESID(NAME, BM, BN, W, S, MB) \ + int NAME(const void* A, const void* B, void* D, int M, int N, int K, \ + const float* act_scale, const float* w_scale, const void* resid, \ + cudaStream_t s) { \ + return launch_resid_(A, B, D, M, N, K, act_scale, \ + w_scale, resid, s); \ + } + DEFINE(fp8_block128_gemm_bs_sm89_32x128x128_w4, 32, 128, 4, 2, 4) DEFINE(fp8_block128_gemm_bs_sm89_64x128x128_w4, 64, 128, 4, 2, 4) DEFINE(fp8_block128_gemm_bs_sm89_64x128x128_w8, 64, 128, 8, 2, 4) @@ -87,6 +130,17 @@ DEFINE(fp8_block128_gemm_bs_sm89_128x128x128_w8_s1, 128, 128, 8, 1, 2) #undef DEFINE +// Residual-fold variants for the down-proj prefill tiles (see dispatcher +// below): the 2B/8B down-proj selects 64x64_s1 (8B) / 64x64 (2B small-M) / +// 32x64 (small-M) at the S ranges Phase-0 measured. Defined additively; the +// baseline kernels above are untouched. +DEFINE_RESID(fp8_block128_gemm_bs_sm89_32x64x128_w4_resid, 32, 64, 4, 2, 4) +DEFINE_RESID(fp8_block128_gemm_bs_sm89_64x64x128_w4_resid, 64, 64, 4, 2, 4) +DEFINE_RESID(fp8_block128_gemm_bs_sm89_64x64x128_w4_s1_resid, 64, 64, 4, 1, 4) +DEFINE_RESID(fp8_block128_gemm_bs_sm89_128x128x128_w8_s1_resid, 128, 128, 8, 1, 2) + +#undef DEFINE_RESID + int fp8_block128_gemm_blockscaled_sm89_bf16out( const void* A, const void* B, void* D, int M, int N, int K, const float* act_scale, const float* w_scale, cudaStream_t stream) diff --git a/csrc/gemm/fp8_block128_gemm_mma_sm89.cuh b/csrc/gemm/fp8_block128_gemm_mma_sm89.cuh index 15946688..d6a5a7a1 100644 --- a/csrc/gemm/fp8_block128_gemm_mma_sm89.cuh +++ b/csrc/gemm/fp8_block128_gemm_mma_sm89.cuh @@ -45,6 +45,22 @@ DECL(fp8_block128_gemm_bs_sm89_128x128x128_w8_s1); #undef DECL +// Residual-fold tile variants (epilogue adds `resid`): D = bf16(acc) + resid. +// resid is [M, N] BF16 row-major, same layout as D. Fuses the residual add +// into the GEMM epilogue (no separate residual_add launch, no D HBM +// round-trip). Additive — the non-resid kernels above are unchanged. +#define DECL_RESID(NAME) \ + int NAME(const void* A, const void* B, void* D, int M, int N, int K, \ + const float* act_scale, const float* w_scale, const void* resid, \ + cudaStream_t stream) + +DECL_RESID(fp8_block128_gemm_bs_sm89_32x64x128_w4_resid); +DECL_RESID(fp8_block128_gemm_bs_sm89_64x64x128_w4_resid); +DECL_RESID(fp8_block128_gemm_bs_sm89_64x64x128_w4_s1_resid); +DECL_RESID(fp8_block128_gemm_bs_sm89_128x128x128_w8_s1_resid); + +#undef DECL_RESID + // Auto-dispatch over the tuned tile set above based on (M, N, K). int fp8_block128_gemm_blockscaled_sm89_bf16out( const void* A, const void* B, void* D, int M, int N, int K, diff --git a/csrc/gemm/fp8_bs_gemm_device.cuh b/csrc/gemm/fp8_bs_gemm_device.cuh index e8176c79..9c511a5d 100644 --- a/csrc/gemm/fp8_bs_gemm_device.cuh +++ b/csrc/gemm/fp8_bs_gemm_device.cuh @@ -66,8 +66,15 @@ __device__ __forceinline__ void ldmatrix_x4_b16( // - B: [N, K] row-major FP8 e4m3, w_scale [N/128, K/128] fp32 // - D: [M, N] row-major BF16 // - BLOCK_N must keep each warp's 8-wide N-atoms inside one 128 scale block. +// +// RESID (opt-in epilogue fold): when true, the BF16 store adds a per-element +// residual `resid[M, N]` (same BF16 layout as D): D = bf16(acc + resid). +// This folds what would otherwise be a separate residual_add launch + an +// extra D round-trip through HBM, mirroring #134's residual-fold epilogue. +// When RESID=false, `resid` is unused and the `if constexpr (RESID)` branch +// is dead-stripped at compile time, so the baseline kernel is byte-identical. template + int MIN_BLOCKS_PER_SM, bool RESID = false> __global__ __launch_bounds__(NUM_WARPS * 32, MIN_BLOCKS_PER_SM) void fp8_bs_gemm_kernel( const __nv_fp8_e4m3* __restrict__ A, @@ -75,7 +82,8 @@ void fp8_bs_gemm_kernel( const float* __restrict__ act_scale, // [M, K/128] const float* __restrict__ w_scale, // [N/128, K/128] __nv_bfloat16* __restrict__ D, - int M, int N, int K) + int M, int N, int K, + const __nv_bfloat16* __restrict__ resid = nullptr) // [M, N] BF16, used iff RESID { constexpr int BLOCK_K = 128; constexpr int THREADS = NUM_WARPS * 32; @@ -310,22 +318,56 @@ void fp8_bs_gemm_kernel( for (int ni = 0; ni < N_ATOMS_PW; ++ni) { int n_pair_base = n_base + warp_id * N_ATOMS_PW * 8 + ni * 8 + 2 * l; // acc[0,1] = row0 cols {2l,2l+1}; acc[2,3] = row1 cols {2l,2l+1}. - // Emit one 32-bit bfloat162 store per row instead of two scalar - // 16-bit stores (NCU's top store-pattern bottleneck after C1). - // Tail (odd last column) falls back to scalar stores. - if (row0 < M && col_pair_ok(n_pair_base, N)) { - *reinterpret_cast<__nv_bfloat162*>(&D[(size_t)row0 * N + n_pair_base]) = - __floats2bfloat162_rn(acc[mi][ni][0], acc[mi][ni][1]); - } else if (row0 < M) { - if (n_pair_base < N) D[(size_t)row0 * N + n_pair_base] = __float2bfloat16(acc[mi][ni][0]); - if (n_pair_base + 1 < N) D[(size_t)row0 * N + n_pair_base+1] = __float2bfloat16(acc[mi][ni][1]); - } - if (row1 < M && col_pair_ok(n_pair_base, N)) { - *reinterpret_cast<__nv_bfloat162*>(&D[(size_t)row1 * N + n_pair_base]) = - __floats2bfloat162_rn(acc[mi][ni][2], acc[mi][ni][3]); - } else if (row1 < M) { - if (n_pair_base < N) D[(size_t)row1 * N + n_pair_base] = __float2bfloat16(acc[mi][ni][2]); - if (n_pair_base + 1 < N) D[(size_t)row1 * N + n_pair_base+1] = __float2bfloat16(acc[mi][ni][3]); + // RESID epilogue fold: add the BF16 residual in-register before the + // bf16 store, so the residual read is fused into the GEMM epilogue + // and never lands as a separate D HBM round-trip + launch. + if constexpr (RESID) { + if (row0 < M && col_pair_ok(n_pair_base, N)) { + __nv_bfloat162 r = *reinterpret_cast( + &resid[(size_t)row0 * N + n_pair_base]); + *reinterpret_cast<__nv_bfloat162*>(&D[(size_t)row0 * N + n_pair_base]) = + __floats2bfloat162_rn(acc[mi][ni][0] + __low2float(r), + acc[mi][ni][1] + __high2float(r)); + } else if (row0 < M) { + if (n_pair_base < N) + D[(size_t)row0 * N + n_pair_base] = __float2bfloat16( + acc[mi][ni][0] + __bfloat162float(resid[(size_t)row0 * N + n_pair_base])); + if (n_pair_base + 1 < N) + D[(size_t)row0 * N + n_pair_base + 1] = __float2bfloat16( + acc[mi][ni][1] + __bfloat162float(resid[(size_t)row0 * N + n_pair_base + 1])); + } + if (row1 < M && col_pair_ok(n_pair_base, N)) { + __nv_bfloat162 r = *reinterpret_cast( + &resid[(size_t)row1 * N + n_pair_base]); + *reinterpret_cast<__nv_bfloat162*>(&D[(size_t)row1 * N + n_pair_base]) = + __floats2bfloat162_rn(acc[mi][ni][2] + __low2float(r), + acc[mi][ni][3] + __high2float(r)); + } else if (row1 < M) { + if (n_pair_base < N) + D[(size_t)row1 * N + n_pair_base] = __float2bfloat16( + acc[mi][ni][2] + __bfloat162float(resid[(size_t)row1 * N + n_pair_base])); + if (n_pair_base + 1 < N) + D[(size_t)row1 * N + n_pair_base + 1] = __float2bfloat16( + acc[mi][ni][3] + __bfloat162float(resid[(size_t)row1 * N + n_pair_base + 1])); + } + } else { + // Emit one 32-bit bfloat162 store per row instead of two scalar + // 16-bit stores (NCU's top store-pattern bottleneck after C1). + // Tail (odd last column) falls back to scalar stores. + if (row0 < M && col_pair_ok(n_pair_base, N)) { + *reinterpret_cast<__nv_bfloat162*>(&D[(size_t)row0 * N + n_pair_base]) = + __floats2bfloat162_rn(acc[mi][ni][0], acc[mi][ni][1]); + } else if (row0 < M) { + if (n_pair_base < N) D[(size_t)row0 * N + n_pair_base] = __float2bfloat16(acc[mi][ni][0]); + if (n_pair_base + 1 < N) D[(size_t)row0 * N + n_pair_base+1] = __float2bfloat16(acc[mi][ni][1]); + } + if (row1 < M && col_pair_ok(n_pair_base, N)) { + *reinterpret_cast<__nv_bfloat162*>(&D[(size_t)row1 * N + n_pair_base]) = + __floats2bfloat162_rn(acc[mi][ni][2], acc[mi][ni][3]); + } else if (row1 < M) { + if (n_pair_base < N) D[(size_t)row1 * N + n_pair_base] = __float2bfloat16(acc[mi][ni][2]); + if (n_pair_base + 1 < N) D[(size_t)row1 * N + n_pair_base+1] = __float2bfloat16(acc[mi][ni][3]); + } } } } diff --git a/csrc/qwen3_vl_bindings.cpp b/csrc/qwen3_vl_bindings.cpp index e45aca6b..94bf9a41 100644 --- a/csrc/qwen3_vl_bindings.cpp +++ b/csrc/qwen3_vl_bindings.cpp @@ -306,6 +306,16 @@ PYBIND11_MODULE(flash_rt_qwen3_vl_kernels, m) { BIND_GEMM_TILE(fp8_block128_gemm_bs_sm89_64x128x128_w8); BIND_GEMM_TILE(fp8_block128_gemm_bs_sm89_128x128x128_w8); #undef BIND_GEMM_TILE + + // Residual-fold GEMM bench bindings: D = bf16(acc) + resid. Same tile set + // as the residual-fold kernels (fp8_block128_gemm_*_resid). Exposed only for + // dev builds so the production pybind surface stays runtime-only. +#define BIND_GEMM_TILE_RESID(NAME) m.def("bench_" #NAME, [](uintptr_t A, uintptr_t B, uintptr_t D, int M, int N, int K, uintptr_t act_scale, uintptr_t w_scale, uintptr_t resid, uintptr_t stream) { return flash_rt::gemm::block128_sm89::NAME( to_ptr(A), to_ptr(B), to_ptr(D), M, N, K, reinterpret_cast(act_scale), reinterpret_cast(w_scale), to_ptr(resid), to_stream(stream)); }, py::arg("A"), py::arg("B"), py::arg("D"), py::arg("M"), py::arg("N"), py::arg("K"), py::arg("act_block_scale"), py::arg("w_block_scale"), py::arg("resid"), py::arg("stream") = 0) + BIND_GEMM_TILE_RESID(fp8_block128_gemm_bs_sm89_32x64x128_w4_resid); + BIND_GEMM_TILE_RESID(fp8_block128_gemm_bs_sm89_64x64x128_w4_resid); + BIND_GEMM_TILE_RESID(fp8_block128_gemm_bs_sm89_64x64x128_w4_s1_resid); + BIND_GEMM_TILE_RESID(fp8_block128_gemm_bs_sm89_128x128x128_w8_s1_resid); +#undef BIND_GEMM_TILE_RESID #endif m.def("qwen3_qk_norm_rope_kvwrite_bf16", diff --git a/scripts/test_resid_fold_correctness.py b/scripts/test_resid_fold_correctness.py new file mode 100644 index 00000000..d675a6d1 --- /dev/null +++ b/scripts/test_resid_fold_correctness.py @@ -0,0 +1,75 @@ +#!/usr/bin/env python3 +"""Correctness + micro-bench for the SM89 FP8 residual-fold GEMM epilogue. + +Verifies that bench__resid (D = bf16(acc) + resid) is bit-identical to +the two-step baseline (plain bench_ then D += resid), then micro-benches +both to measure the residual-fold win (saves 1 launch + 1 D HBM round-trip). + +Usage: + python scripts/test_resid_fold_correctness.py --M 512 +""" +from __future__ import annotations +import argparse, pathlib, sys, statistics +import torch +REPO = pathlib.Path(__file__).resolve().parents[1] +sys.path.insert(0, str(REPO)) +from flash_rt import flash_rt_qwen3_vl_kernels as vlk + +# 8B down-proj: M=S, N=4096, K=12288 (mirror the prefill down GEMM). +TILES = { + "32x64": ("bench_fp8_block128_gemm_bs_sm89_32x64x128_w4", "bench_fp8_block128_gemm_bs_sm89_32x64x128_w4_resid"), + "64x64": ("bench_fp8_block128_gemm_bs_sm89_64x64x128_w4", "bench_fp8_block128_gemm_bs_sm89_64x64x128_w4_resid"), + "64x64_s1":("bench_fp8_block128_gemm_bs_sm89_64x64x128_w4_s1", "bench_fp8_block128_gemm_bs_sm89_64x64x128_w4_s1_resid"), + "128x128_s1":("bench_fp8_block128_gemm_bs_sm89_128x128x128_w8_s1","bench_fp8_block128_gemm_bs_sm89_128x128x128_w8_s1_resid"), +} + +def ev_ms(fn, iters=50, warmup=20): + s=[torch.cuda.Event(enable_timing=True) for _ in range(iters)] + e=[torch.cuda.Event(enable_timing=True) for _ in range(iters)] + for _ in range(warmup): fn() + torch.cuda.synchronize() + for i in range(iters): s[i].record(); fn(); e[i].record() + torch.cuda.synchronize() + return [a.elapsed_time(b) for a,b in zip(s,e)] + +def main(): + p=argparse.ArgumentParser() + p.add_argument("--M", type=int, nargs="+", default=[128,512,1024]) + p.add_argument("--N", type=int, default=4096) + p.add_argument("--K", type=int, default=12288) + args=p.parse_args() + dev=torch.device("cuda:0") + torch.manual_seed(0) + for M in args.M: + N,K=args.N,args.K + A=torch.randn(M,K,device=dev).to(torch.float8_e4m3fn) + B=torch.randn(N,K,device=dev).to(torch.float8_e4m3fn) + resid=torch.randn(M,N,device=dev,dtype=torch.bfloat16)*0.1 + asc=(torch.randn(M,K//128,device=dev)*0.1).contiguous() + wsc=(torch.randn(N//128,K//128,device=dev)*0.1).contiguous() + s=torch.cuda.current_stream().cuda_stream + print(f"\n=== M={M} N={N} K={K} ===") + print(f"{'tile':>12} {'base_us':>8} {'resid_us':>9} {'delta_us':>9} {'cos':>10} {'maxabs':>10}") + for name,(base_fn,resid_fn) in TILES.items(): + fb=getattr(vlk,base_fn); fr=getattr(vlk,resid_fn) + D_base=torch.empty(M,N,device=dev,dtype=torch.bfloat16) + D_resid=torch.empty(M,N,device=dev,dtype=torch.bfloat16) + # baseline: GEMM then D += resid + fb(int(A.data_ptr()),int(B.data_ptr()),int(D_base.data_ptr()),M,N,K,int(asc.data_ptr()),int(wsc.data_ptr()),s) + D_base += resid + # residual-fold: one launch + fr(int(A.data_ptr()),int(B.data_ptr()),int(D_resid.data_ptr()),M,N,K,int(asc.data_ptr()),int(wsc.data_ptr()),int(resid.data_ptr()),s) + cos=torch.nn.functional.cosine_similarity(D_base.flatten().unsqueeze(0),D_resid.flatten().unsqueeze(0)).item() + maxabs=(D_base.float()-D_resid.float()).abs().max().item() + # bench (layer-ish: reuse A/B, cold-ish resid each iter to mimic per-step residual) + def run_base(): + fb(int(A.data_ptr()),int(B.data_ptr()),int(D_base.data_ptr()),M,N,K,int(asc.data_ptr()),int(wsc.data_ptr()),s) + D_base.add_(resid) + def run_resid(): + fr(int(A.data_ptr()),int(B.data_ptr()),int(D_resid.data_ptr()),M,N,K,int(asc.data_ptr()),int(wsc.data_ptr()),int(resid.data_ptr()),s) + b=statistics.median(ev_ms(run_base)); r=statistics.median(ev_ms(run_resid)) + flag="OK" if maxabs==0 else ("close" if cos>0.9999 else "FAIL") + print(f"{name:>12} {b*1e3:>8.1f} {r*1e3:>9.1f} {(b-r)*1e3:>+9.1f} {cos:>10.6f} {maxabs:>10.2e} {flag}") + +if __name__=="__main__": + main() From fab8cbba6f560bedf9923196a46b749e18c65653 Mon Sep 17 00:00:00 2001 From: heiheiha798 <2300012738@stu.pku.edu.cn> Date: Fri, 17 Jul 2026 20:01:11 +0800 Subject: [PATCH 2/2] feat(sm89): GeGLU silu-fold megakernel for prefill gate_up (opt-in) Fuse gate_up GEMM + silu_mul + per-token block-128 FP8 quant into one launch, eliminating the [S, 2*inter] BF16 transient the baseline two-launch path (gate_up GEMM -> silu_mul_merged_to_fp8) materializes. Mirrors #134's SwiGLU epilogue fold (sm120) and sm100 flashrt_megakernel_geglu's gate-in-smem pattern, ported to sm89's hand-written cp.async + ldmatrix.x4 + mma.m16n8k32 idiom (no TMA/tcgen05). Two kernel variants in fp8_bs_gemm_device.cuh (additive; fp8_bs_gemm_kernel untouched): - fp8_bs_geglu_silu_fold_kernel: two-pass (gate -> smem -> up), the shipped tile (32x128_w4_s1). Single B smem region reused across passes; only one accumulator live at a time (~70 regs). - fp8_bs_geglu_silu_fold_apersist_kernel: A-persistent interleaved experiment (both gate+up acc in regs, B reloaded per k-iter). Kept as a bench-only tile; the two-pass s1 won the e2e sweep. ncu diagnosis (issue #1 Phase 2) isolated why naive two-pass lost: s2 occupancy collapse (gate_smem + 2 cp.async stages -> 49.7KB -> 1 CTA/SM, 8.33% occupancy) + 2x pipeline drain. s1 (29.7KB, 3 CTA/SM) recovers and wins in the launch-bound regime. The gate_up GEMM is compute-bound above the S crossover, so the fused path is shape-gated per-inter (2B S in [96,192], 8B S=128 only); outside the band the baseline two-launch path runs unchanged (default byte-identical, env FLASH_RT_SM89_SILU_FOLD_PREFILL=1 opt-in). Correctness: e2e prefill cosine 0.9992-1.0 vs baseline, argmax preserved across S (saturating-FP8 real activations). One fewer bf16(acc) rounding than baseline (documented). Fixed an out_scale write bug (single-thread guard left rows 1..BM-1 unwritten -> garbage scale, correct fp8). E2E prefill (4090, 30-iter median, fold on vs off): 2B S=128: 6.10 -> 5.64 ms +7.5% 2B S=192: 6.10 -> 5.54 ms +9.2% 8B S=128: 13.97 -> 13.45 ms +3.7% (stable across 3 runs) Above the band: no regression (falls back to baseline). Bench script scripts/test_geglu_silu_fold_correctness.py included. --- csrc/gemm/fp8_block128_gemm_mma_sm89.cu | 107 ++ csrc/gemm/fp8_block128_gemm_mma_sm89.cuh | 31 + csrc/gemm/fp8_bs_gemm_device.cuh | 949 ++++++++++++++++++ csrc/qwen3_vl_bindings.cpp | 55 + flash_rt/frontends/torch/qwen3_vl_fp8_sm89.py | 61 +- scripts/test_geglu_silu_fold_correctness.py | 105 ++ 6 files changed, 1300 insertions(+), 8 deletions(-) create mode 100644 scripts/test_geglu_silu_fold_correctness.py diff --git a/csrc/gemm/fp8_block128_gemm_mma_sm89.cu b/csrc/gemm/fp8_block128_gemm_mma_sm89.cu index 3624ae4e..271e563e 100644 --- a/csrc/gemm/fp8_block128_gemm_mma_sm89.cu +++ b/csrc/gemm/fp8_block128_gemm_mma_sm89.cu @@ -114,6 +114,79 @@ int launch_resid_(const void* A, const void* B, void* D, w_scale, resid, s); \ } +// GeGLU silu-fold launch: fuses gate+up GEMM + silu(gate)*up + per-token +// block-128 FP8 quant into one launch (no [M,2N] BF16 transient). B is +// gate_up_w [2*N, K] (gate rows [0,N), up rows [N,2N)); w_scale is gate_up_s +// [2*N/128, K/128]. Outputs FP8 [M,N] + scale [M,N/128]. See device header. +template +int launch_geglu_silu_fold_(const void* A, const void* B, + int M, int N, int K, const float* act_scale, + const float* w_scale, void* output, float* out_scale, + cudaStream_t s) +{ + constexpr int BK = 128; + constexpr int SCALE_KTILE = 8; + int grid_m = (M + BM - 1) / BM; + int grid_n = (N + BN - 1) / BN; // over output N (== inter), NOT 2*N + dim3 grid(grid_m, grid_n, 1); + dim3 block(W * 32, 1, 1); + // A/B cp.async stages + gate_smem (BM*BN bf16) + scales + amax scratch. + int smem_bytes = STAGES * (BM + BN) * BK + + (BM * BN) * (int)sizeof(__nv_bfloat16) + + (BM * SCALE_KTILE + 2 * SCALE_KTILE) * (int)sizeof(float) + + (W * BM + BM) * (int)sizeof(float); + if (smem_bytes > 48 * 1024) { + cudaFuncSetAttribute( + (const void*)&fp8_bs_geglu_silu_fold_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes); + } + fp8_bs_geglu_silu_fold_kernel<<>>( + reinterpret_cast(A), + reinterpret_cast(B), + act_scale, w_scale, + reinterpret_cast<__nv_fp8_e4m3*>(output), + out_scale, M, N, K); + cudaError_t err = cudaGetLastError(); + return (err == cudaSuccess) ? 0 : 1; +} + +// A-persistent interleaved variant: stage A once, reuse ONE B smem region for +// gate then up within each k-iter (both gate+up acc live in regs, true +// interleaved per K-tile). Single B region -> 3 CTA/SM at s1 (vs interleaved's +// 2, vs two-pass's 3). See fp8_bs_geglu_silu_fold_apersist_kernel. +template +int launch_geglu_silu_fold_apersist_(const void* A, const void* B, + int M, int N, int K, const float* act_scale, + const float* w_scale, void* output, + float* out_scale, cudaStream_t s) +{ + constexpr int BK = 128; + constexpr int SCALE_KTILE = 8; + int grid_m = (M + BM - 1) / BM; + int grid_n = (N + BN - 1) / BN; + dim3 grid(grid_m, grid_n, 1); + dim3 block(W * 32, 1, 1); + // Same smem layout as the two-pass variant (gate_smem region kept for layout + // parity though apersist doesn't use it as a handoff — gate stays in regs). + int smem_bytes = STAGES * (BM + BN) * BK + + (BM * BN) * (int)sizeof(__nv_bfloat16) + + (BM * SCALE_KTILE + 2 * SCALE_KTILE) * (int)sizeof(float) + + (W * BM + BM) * (int)sizeof(float); + if (smem_bytes > 48 * 1024) { + cudaFuncSetAttribute( + (const void*)&fp8_bs_geglu_silu_fold_apersist_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes); + } + fp8_bs_geglu_silu_fold_apersist_kernel<<>>( + reinterpret_cast(A), + reinterpret_cast(B), + act_scale, w_scale, + reinterpret_cast<__nv_fp8_e4m3*>(output), + out_scale, M, N, K); + cudaError_t err = cudaGetLastError(); + return (err == cudaSuccess) ? 0 : 1; +} + DEFINE(fp8_block128_gemm_bs_sm89_32x128x128_w4, 32, 128, 4, 2, 4) DEFINE(fp8_block128_gemm_bs_sm89_64x128x128_w4, 64, 128, 4, 2, 4) DEFINE(fp8_block128_gemm_bs_sm89_64x128x128_w8, 64, 128, 8, 2, 4) @@ -141,6 +214,40 @@ DEFINE_RESID(fp8_block128_gemm_bs_sm89_128x128x128_w8_s1_resid, 128, 128, 8, 1, #undef DEFINE_RESID +// GeGLU silu-fold tile variants (BLOCK_N pinned to 128 = one quant block). +#define DEFINE_GEGLU(NAME, BM, BN, W, S, MB) \ + int NAME(const void* A, const void* B, int M, int N, int K, \ + const float* act_scale, const float* w_scale, void* output, \ + float* out_scale, cudaStream_t s) { \ + return launch_geglu_silu_fold_( \ + A, B, M, N, K, act_scale, w_scale, output, out_scale, s); \ + } +DEFINE_GEGLU(fp8_bs_geglu_silu_fold_sm89_32x128_w4_s2, 32, 128, 4, 2, 4) +DEFINE_GEGLU(fp8_bs_geglu_silu_fold_sm89_16x128_w4_s2, 16, 128, 4, 2, 4) +DEFINE_GEGLU(fp8_bs_geglu_silu_fold_sm89_64x128_w4_s2, 64, 128, 4, 2, 4) +DEFINE_GEGLU(fp8_bs_geglu_silu_fold_sm89_128x128_w8_s1, 128, 128, 8, 1, 2) +// Low-smem variants (STAGES=1) to recover occupancy lost to gate_smem on sm89: +// the s2 dual-buffer + gate_smem pushes dynamic smem >48KB → Block Limit Shared +// Mem = 1 (8% occupancy, ncu-confirmed). s1 trades cp.async overlap for 3-4x +// the CTA density. Primary candidates for the prefill M>=128 regime. +DEFINE_GEGLU(fp8_bs_geglu_silu_fold_sm89_32x128_w4_s1, 32, 128, 4, 1, 4) +DEFINE_GEGLU(fp8_bs_geglu_silu_fold_sm89_16x128_w4_s1, 16, 128, 4, 1, 4) +#undef DEFINE_GEGLU + +// A-persistent interleaved variant (single B smem region, gate+up acc both in +// regs). launch wrapper shares the smem formula with the two-pass variant. +#define DEFINE_GEGLU_AP(NAME, BM, BN, W, S, MB) \ + int NAME(const void* A, const void* B, int M, int N, int K, \ + const float* act_scale, const float* w_scale, void* output, \ + float* out_scale, cudaStream_t s) { \ + return launch_geglu_silu_fold_apersist_( \ + A, B, M, N, K, act_scale, w_scale, output, out_scale, s); \ + } +DEFINE_GEGLU_AP(fp8_bs_geglu_silu_fold_apersist_sm89_32x128_w4_s1, 32, 128, 4, 1, 2) +DEFINE_GEGLU_AP(fp8_bs_geglu_silu_fold_apersist_sm89_16x128_w4_s1, 16, 128, 4, 1, 2) +DEFINE_GEGLU_AP(fp8_bs_geglu_silu_fold_apersist_sm89_32x128_w4_s2, 32, 128, 4, 2, 2) +#undef DEFINE_GEGLU_AP + int fp8_block128_gemm_blockscaled_sm89_bf16out( const void* A, const void* B, void* D, int M, int N, int K, const float* act_scale, const float* w_scale, cudaStream_t stream) diff --git a/csrc/gemm/fp8_block128_gemm_mma_sm89.cuh b/csrc/gemm/fp8_block128_gemm_mma_sm89.cuh index d6a5a7a1..88b72e74 100644 --- a/csrc/gemm/fp8_block128_gemm_mma_sm89.cuh +++ b/csrc/gemm/fp8_block128_gemm_mma_sm89.cuh @@ -61,6 +61,37 @@ DECL_RESID(fp8_block128_gemm_bs_sm89_128x128x128_w8_s1_resid); #undef DECL_RESID +// GeGLU silu-fold tile variants: fuse gate+up GEMM + silu(gate)*up + per-token +// block-128 FP8 quant into one launch. B = gate_up_w [2*N, K] (gate rows +// [0,N), up rows [N,2N)); w_scale = gate_up_s [2*N/128, K/128]. Output FP8 +// [M,N] + scale [M,N/128]. BLOCK_N pinned to 128 (one quant block per CTA). +#define DECL_GEGLU(NAME) \ + int NAME(const void* A, const void* B, int M, int N, int K, \ + const float* act_scale, const float* w_scale, void* output, \ + float* out_scale, cudaStream_t stream) + +DECL_GEGLU(fp8_bs_geglu_silu_fold_sm89_32x128_w4_s2); +DECL_GEGLU(fp8_bs_geglu_silu_fold_sm89_16x128_w4_s2); +DECL_GEGLU(fp8_bs_geglu_silu_fold_sm89_64x128_w4_s2); +DECL_GEGLU(fp8_bs_geglu_silu_fold_sm89_128x128_w8_s1); +DECL_GEGLU(fp8_bs_geglu_silu_fold_sm89_32x128_w4_s1); +DECL_GEGLU(fp8_bs_geglu_silu_fold_sm89_16x128_w4_s1); + +#undef DECL_GEGLU + +// A-persistent interleaved variant (single B smem region, both gate+up acc in +// registers). Same I/O contract as DECL_GEGLU. +#define DECL_GEGLU_AP(NAME) \ + int NAME(const void* A, const void* B, int M, int N, int K, \ + const float* act_scale, const float* w_scale, void* output, \ + float* out_scale, cudaStream_t stream) + +DECL_GEGLU_AP(fp8_bs_geglu_silu_fold_apersist_sm89_32x128_w4_s1); +DECL_GEGLU_AP(fp8_bs_geglu_silu_fold_apersist_sm89_16x128_w4_s1); +DECL_GEGLU_AP(fp8_bs_geglu_silu_fold_apersist_sm89_32x128_w4_s2); + +#undef DECL_GEGLU_AP + // Auto-dispatch over the tuned tile set above based on (M, N, K). int fp8_block128_gemm_blockscaled_sm89_bf16out( const void* A, const void* B, void* D, int M, int N, int K, diff --git a/csrc/gemm/fp8_bs_gemm_device.cuh b/csrc/gemm/fp8_bs_gemm_device.cuh index 9c511a5d..ce3c2903 100644 --- a/csrc/gemm/fp8_bs_gemm_device.cuh +++ b/csrc/gemm/fp8_bs_gemm_device.cuh @@ -61,6 +61,12 @@ __device__ __forceinline__ void ldmatrix_x4_b16( : "r"(smem_addr)); } +// SiLU in fp32. Matches quantize::silu_f32 (fp8_per_token_block_quant.cu:416) +// so the GeGLU silu-fold epilogue reproduces silu_mul_merged's math exactly. +__device__ __forceinline__ float silu_f32(float x) { + return x / (1.0f + expf(-x)); +} + // BLOCK_K is pinned to 128 (one DeepSeek scale block per K-iteration). // - A: [M, K] row-major FP8 e4m3, act_scale [M, K/128] fp32 // - B: [N, K] row-major FP8 e4m3, w_scale [N/128, K/128] fp32 @@ -373,6 +379,949 @@ void fp8_bs_gemm_kernel( } } +// ============================================================================ +// GeGLU silu-fold megakernel (Phase 2): fuses gate GEMM + up GEMM + +// silu(gate)*up + per-token block-128 FP8 quant into ONE launch, writing FP8 +// output + scale directly — eliminating the [M, 2*N] BF16 transient that the +// baseline gate_up GEMM would write and silu_mul_merged_to_fp8 would read back. +// +// gate_up_w : [2*N, K] FP8 row-major (gate rows [0,N); up rows [N,2N)) +// gate_up_s : [2*N/128, K/128] fp32 (up row = gate row + N/128) +// A : [M, K] FP8 (per-token quantized), act_scale [M, K/128] +// output : [M, N] FP8, scale [M, N/128] +// +// Two-pass per CTA (mirrors sm100 flashrt_megakernel_geglu's "gate stays in +// smem"): pass 1 accumulates gate over full K and stores silu(gate) as BF16 +// into a smem gate buffer; pass 2 reuses the same A/B smem staging, accumulates +// up, then the epilogue reads gate from smem, forms v = bf16(bf16(silu(gate))*up) +// (matching silu_mul_merged's two bf16 roundings), reduces |v| over the 128-col +// quant block per row, and quantizes to FP8. No grid_barrier (single CTA owns +// its full quant block: BLOCK_N == 128 == one scale block). GEMM body reuses +// the same cp.async + ldmatrix.x4 + mma.m16n8k32 tiles as fp8_bs_gemm_kernel. +// ============================================================================ +template +__global__ __launch_bounds__(NUM_WARPS * 32, MIN_BLOCKS_PER_SM) +void fp8_bs_geglu_silu_fold_kernel( + const __nv_fp8_e4m3* __restrict__ A, + const __nv_fp8_e4m3* __restrict__ B, // gate_up_w [2*N, K] + const float* __restrict__ act_scale, // [M, K/128] + const float* __restrict__ w_scale, // gate_up_s [2*N/128, K/128] + __nv_fp8_e4m3* __restrict__ output, // [M, N] + float* __restrict__ out_scale, // [M, N/128] + int M, int N, int K) +{ + static_assert(BLOCK_N == 128, + "GeGLU silu-fold requires BLOCK_N==128 (one quant block per CTA)"); + constexpr int BLOCK_K = 128; + constexpr int THREADS = NUM_WARPS * 32; + constexpr int M_ATOMS = BLOCK_M / 16; + constexpr int N_ATOMS = BLOCK_N / 8; // 16 + constexpr int N_ATOMS_PW = N_ATOMS / NUM_WARPS; + constexpr int N_PAIRS_PW = N_ATOMS_PW / 2; + constexpr int K_ATOMS = BLOCK_K / 32; // 4 + constexpr int NUM_CHUNKS_PER_ROW = BLOCK_K / 16; + constexpr int SWIZZLE_MASK = NUM_CHUNKS_PER_ROW - 1; + constexpr int SCALE_KTILE = 8; + constexpr int A_TILE = BLOCK_M * BLOCK_K; + constexpr int B_TILE = BLOCK_N * BLOCK_K; + + static_assert(BLOCK_M % 16 == 0, "BLOCK_M multiple of 16"); + static_assert(N_ATOMS_PW >= 2 && N_ATOMS_PW % 2 == 0, + "ldmatrix pairs 2 N-atoms: N_ATOMS_PW must be even >= 2"); + + extern __shared__ uint8_t smem_raw[]; + uint8_t* A_smem = smem_raw; + uint8_t* B_smem = A_smem + STAGES * A_TILE; + // gate_smem: silu(gate) as BF16, [BLOCK_M, BLOCK_N]. One CTA-tile, written + // by pass 1 epilogue, read by pass 2 epilogue. The sm100 geglu's "gate + // stays in smem" handoff, without tcgen05/EVT. + __nv_bfloat16* gate_smem = reinterpret_cast<__nv_bfloat16*>( + B_smem + STAGES * B_TILE); + float* as_smem = reinterpret_cast(gate_smem + BLOCK_M * BLOCK_N); + float* wsg_smem = as_smem + BLOCK_M * SCALE_KTILE; // gate w_scale row + float* wsu_smem = wsg_smem + SCALE_KTILE; // up w_scale row + // amax partials: 4 warps × BLOCK_M rows. Cross-warp reduce per row. + float* amax_smem = wsu_smem + SCALE_KTILE; + + const int cta_m = blockIdx.x; + const int cta_n = blockIdx.y; + const int m_base = cta_m * BLOCK_M; + const int n_base = cta_n * BLOCK_N; // n0, < N (output col block) + + const int t = threadIdx.x; + const int warp_id = t / 32; + const int lane = t % 32; + const int l = lane % 4; + const int h = lane / 4; + const int frag_group = lane / 8; + const int row_in_frag = lane % 8; + const int row_block = frag_group / 2; + const int col_block = frag_group % 2; + + const int K128 = K >> 7; + const int N128 = N >> 7; // gate w_scale blocks + // gate B-rows [n_base, n_base+BLOCK_N); up B-rows [n_base+N, n_base+N+BLOCK_N) + const int gate_b_row0 = n_base; + const int up_b_row0 = n_base + N; + const int gate_ws_row = (n_base >> 7); // gate w_scale block row + const int up_ws_row = gate_ws_row + N128; // up w_scale block row + + // ---- scale staging (shared by both passes; re-staged per SCALE_KTILE) ---- + auto stage_scales = [&](int kb0) { + const int as_total = BLOCK_M * SCALE_KTILE; + for (int idx = t; idx < as_total; idx += THREADS) { + int r = idx / SCALE_KTILE; + int kc = idx - r * SCALE_KTILE; + int row = m_base + r; + int kb = kb0 + kc; + as_smem[idx] = (row < M && kb < K128) + ? act_scale[(size_t)row * K128 + kb] : 0.0f; + } + for (int kc = t; kc < SCALE_KTILE; kc += THREADS) { + int kb = kb0 + kc; + wsg_smem[kc] = (kb < K128) + ? w_scale[(size_t)gate_ws_row * K128 + kb] : 0.0f; + wsu_smem[kc] = (kb < K128) + ? w_scale[(size_t)up_ws_row * K128 + kb] : 0.0f; + } + __syncthreads(); + }; + + // ---- cp.async A + (gate or up) B tile staging ---- + // b_row0 selects which 128-row band of B [2*N, K] to stage. + auto issue_load = [&](int stage, int k_base, int b_row0) { + constexpr int A_CHUNKS = BLOCK_M * NUM_CHUNKS_PER_ROW; + constexpr int A_ITERS = (A_CHUNKS + THREADS - 1) / THREADS; + #pragma unroll + for (int it = 0; it < A_ITERS; ++it) { + int idx = it * THREADS + t; + if (idx >= A_CHUNKS) break; + int row_a = idx / NUM_CHUNKS_PER_ROW; + int chunk_a = idx % NUM_CHUNKS_PER_ROW; + int m_glob = m_base + row_a; + int k_glob = k_base + chunk_a * 16; + const uint8_t* a_src = nullptr; + if (m_glob < M && k_glob < K) { + a_src = reinterpret_cast(&A[(size_t)m_glob * K + k_glob]); + } + int csw = chunk_a ^ (row_a & SWIZZLE_MASK); + cp_async_16( + to_smem(&A_smem[stage * A_TILE + row_a * BLOCK_K + csw * 16]), + a_src); + } + constexpr int B_CHUNKS = BLOCK_N * NUM_CHUNKS_PER_ROW; + constexpr int B_ITERS = (B_CHUNKS + THREADS - 1) / THREADS; + #pragma unroll + for (int it = 0; it < B_ITERS; ++it) { + int idx = it * THREADS + t; + if (idx >= B_CHUNKS) break; + int row_b = idx / NUM_CHUNKS_PER_ROW; + int chunk_b = idx % NUM_CHUNKS_PER_ROW; + int n_glob = b_row0 + row_b; + int k_glob = k_base + chunk_b * 16; + const uint8_t* b_src = nullptr; + if (n_glob < 2 * N && k_glob < K) { + b_src = reinterpret_cast(&B[(size_t)n_glob * K + k_glob]); + } + int csw = chunk_b ^ (row_b & SWIZZLE_MASK); + cp_async_16( + to_smem(&B_smem[stage * B_TILE + row_b * BLOCK_K + csw * 16]), + b_src); + } + }; + + // ---- one GEMM pass over full K, accumulating into `acc` with the given + // w_scale smem row (gate or up). b_row0 selects the B band. ---- + auto run_pass = [&](float (*acc)[N_ATOMS_PW][4], int b_row0, + const float* ws_smem_pass) { + const int K_ITERS = (K + BLOCK_K - 1) / BLOCK_K; + #pragma unroll + for (int s = 0; s < STAGES - 1; ++s) { + int kb = s * BLOCK_K; + if (kb < K) issue_load(s, kb, b_row0); + asm volatile("cp.async.commit_group;\n" ::); + } + int compute_stage = 0; + for (int k_iter = 0; k_iter < K_ITERS; ++k_iter) { + int issue_iter = k_iter + (STAGES - 1); + int issue_stage = issue_iter % STAGES; + if (issue_iter < K_ITERS) issue_load(issue_stage, issue_iter * BLOCK_K, b_row0); + asm volatile("cp.async.commit_group;\n" ::); + asm volatile("cp.async.wait_group %0;\n" :: "n"(STAGES - 1)); + __syncthreads(); + + const int kb = k_iter; + if ((kb % SCALE_KTILE) == 0) stage_scales(kb); + + float tacc[M_ATOMS][N_ATOMS_PW][4]; + #pragma unroll + for (int mi = 0; mi < M_ATOMS; ++mi) + #pragma unroll + for (int ni = 0; ni < N_ATOMS_PW; ++ni) + #pragma unroll + for (int j = 0; j < 4; ++j) tacc[mi][ni][j] = 0.0f; + + uint8_t* A_stage = A_smem + compute_stage * A_TILE; + uint8_t* B_stage = B_smem + compute_stage * B_TILE; + #pragma unroll + for (int ka = 0; ka < K_ATOMS; ++ka) { + uint32_t A_regs[M_ATOMS][4]; + #pragma unroll + for (int mi = 0; mi < M_ATOMS; ++mi) { + int row = mi * 16 + row_block * 8 + row_in_frag; + int chunk = 2 * ka + col_block; + int csw = chunk ^ (row & SWIZZLE_MASK); + ldmatrix_x4_b16(A_regs[mi][0], A_regs[mi][1], A_regs[mi][2], A_regs[mi][3], + to_smem(&A_stage[row * BLOCK_K + csw * 16])); + } + uint32_t B_regs[N_PAIRS_PW][4]; + #pragma unroll + for (int np = 0; np < N_PAIRS_PW; ++np) { + int nrow = warp_id * N_ATOMS_PW * 8 + np * 16 + row_block * 8 + row_in_frag; + int chunk = 2 * ka + col_block; + int csw = chunk ^ (nrow & SWIZZLE_MASK); + ldmatrix_x4_b16(B_regs[np][0], B_regs[np][1], B_regs[np][2], B_regs[np][3], + to_smem(&B_stage[nrow * BLOCK_K + csw * 16])); + } + #pragma unroll + for (int mi = 0; mi < M_ATOMS; ++mi) { + #pragma unroll + for (int np = 0; np < N_PAIRS_PW; ++np) { + int ni0 = np * 2, ni1 = np * 2 + 1; + mma_m16n8k32_e4m3( + tacc[mi][ni0][0], tacc[mi][ni0][1], tacc[mi][ni0][2], tacc[mi][ni0][3], + A_regs[mi][0], A_regs[mi][2], A_regs[mi][1], A_regs[mi][3], + B_regs[np][0], B_regs[np][1]); + mma_m16n8k32_e4m3( + tacc[mi][ni1][0], tacc[mi][ni1][1], tacc[mi][ni1][2], tacc[mi][ni1][3], + A_regs[mi][0], A_regs[mi][2], A_regs[mi][1], A_regs[mi][3], + B_regs[np][2], B_regs[np][3]); + } + } + } + + int kbt = kb % SCALE_KTILE; + float ws_cta = ws_smem_pass[kbt]; + #pragma unroll + for (int mi = 0; mi < M_ATOMS; ++mi) { + int row0 = m_base + mi * 16 + h; + int row1 = row0 + 8; + float as0 = as_smem[(mi * 16 + h) * SCALE_KTILE + kbt]; + float as1 = as_smem[(mi * 16 + h + 8) * SCALE_KTILE + kbt]; + #pragma unroll + for (int ni = 0; ni < N_ATOMS_PW; ++ni) { + acc[mi][ni][0] += tacc[mi][ni][0] * (as0 * ws_cta); + acc[mi][ni][1] += tacc[mi][ni][1] * (as0 * ws_cta); + acc[mi][ni][2] += tacc[mi][ni][2] * (as1 * ws_cta); + acc[mi][ni][3] += tacc[mi][ni][3] * (as1 * ws_cta); + } + } + __syncthreads(); + compute_stage = (compute_stage + 1) % STAGES; + } + asm volatile("cp.async.wait_all;\n" ::); + }; + + // =================== Pass 1: gate GEMM → silu(gate) in smem =================== + float gate_acc[M_ATOMS][N_ATOMS_PW][4]; + #pragma unroll + for (int mi = 0; mi < M_ATOMS; ++mi) + #pragma unroll + for (int ni = 0; ni < N_ATOMS_PW; ++ni) + #pragma unroll + for (int j = 0; j < 4; ++j) gate_acc[mi][ni][j] = 0.0f; + + run_pass(gate_acc, gate_b_row0, wsg_smem); + + // Pass 1 epilogue: store silu(gate_acc) as BF16 into gate_smem[BM, BN]. + // Thread (h,l) owns rows {mi*16+h, mi*16+h+8}, cols {ni*8+2l, ni*8+2l+1} + // within its warp's N band. Replicate silu_mul_merged's first bf16 rounding + // (bf16(silu(g))) so the fused path matches the split kernel's precision. + #pragma unroll + for (int mi = 0; mi < M_ATOMS; ++mi) { + int row0 = m_base + mi * 16 + h; + int row1 = row0 + 8; + #pragma unroll + for (int ni = 0; ni < N_ATOMS_PW; ++ni) { + int n_pair_base = warp_id * N_ATOMS_PW * 8 + ni * 8 + 2 * l; + // gate_smem is [BLOCK_M, BLOCK_N]; local col = n_pair_base. + if (row0 < M) { + __nv_bfloat162 gs = __floats2bfloat162_rn( + silu_f32(gate_acc[mi][ni][0]), silu_f32(gate_acc[mi][ni][1])); + *reinterpret_cast<__nv_bfloat162*>( + &gate_smem[(row0 - m_base) * BLOCK_N + n_pair_base]) = gs; + __nv_bfloat162 gs2 = __floats2bfloat162_rn( + silu_f32(gate_acc[mi][ni][2]), silu_f32(gate_acc[mi][ni][3])); + *reinterpret_cast<__nv_bfloat162*>( + &gate_smem[(row1 - m_base) * BLOCK_N + n_pair_base]) = gs2; + } + } + } + __syncthreads(); // gate_smem visible to pass 2 epilogue in all warps + // gate_acc registers now free; reused for up_acc. + + // =================== Pass 2: up GEMM → up_acc =================== + float up_acc[M_ATOMS][N_ATOMS_PW][4]; + #pragma unroll + for (int mi = 0; mi < M_ATOMS; ++mi) + #pragma unroll + for (int ni = 0; ni < N_ATOMS_PW; ++ni) + #pragma unroll + for (int j = 0; j < 4; ++j) up_acc[mi][ni][j] = 0.0f; + + run_pass(up_acc, up_b_row0, wsu_smem); + + // =================== Pass 2 epilogue: silu(gate)*up + quant → FP8 =================== + // v = bf16(bf16(silu(gate)) * up), matching silu_mul_merged's two bf16 + // roundings (silu(gate) was already bf16-rounded into gate_smem in pass 1; + // here we bf16-round the product). Then per-row amax over the 128-col block + // and quantize. + constexpr float kFp8Max = 448.0f; + // Each thread owns 8 cols (4 n-atoms × 2) for 2 rows per m-atom. Compute |v| + // and a per-warp partial amax per row (the warp owns 32 of the row's 128 cols). + // amax_smem[warp_id][row_in_cta] holds the warp's row-amax partial. + float v[M_ATOMS][N_ATOMS_PW][4]; + #pragma unroll + for (int mi = 0; mi < M_ATOMS; ++mi) { + int row0 = m_base + mi * 16 + h; + int row1 = row0 + 8; + int rloc0 = mi * 16 + h; // local row in [0, BLOCK_M) + int rloc1 = rloc0 + 8; + float amax0 = 0.0f, amax1 = 0.0f; + #pragma unroll + for (int ni = 0; ni < N_ATOMS_PW; ++ni) { + int n_pair_base = warp_id * N_ATOMS_PW * 8 + ni * 8 + 2 * l; + // gate value (already bf16(silu(gate))) from smem; up from registers. + if (row0 < M) { + __nv_bfloat162 g = *reinterpret_cast( + &gate_smem[rloc0 * BLOCK_N + n_pair_base]); + float gf0 = __low2float(g), gf1 = __high2float(g); + v[mi][ni][0] = __bfloat162float(__float2bfloat16(gf0 * up_acc[mi][ni][0])); + v[mi][ni][1] = __bfloat162float(__float2bfloat16(gf1 * up_acc[mi][ni][1])); + amax0 = fmaxf(amax0, fmaxf(fabsf(v[mi][ni][0]), fabsf(v[mi][ni][1]))); + } else { + v[mi][ni][0] = 0.0f; v[mi][ni][1] = 0.0f; + } + if (row1 < M) { + __nv_bfloat162 g = *reinterpret_cast( + &gate_smem[rloc1 * BLOCK_N + n_pair_base]); + float gf0 = __low2float(g), gf1 = __high2float(g); + v[mi][ni][2] = __bfloat162float(__float2bfloat16(gf0 * up_acc[mi][ni][2])); + v[mi][ni][3] = __bfloat162float(__float2bfloat16(gf1 * up_acc[mi][ni][3])); + amax1 = fmaxf(amax1, fmaxf(fabsf(v[mi][ni][2]), fabsf(v[mi][ni][3]))); + } else { + v[mi][ni][2] = 0.0f; v[mi][ni][3] = 0.0f; + } + } + // Warp-shuffle reduce the 4 lanes (l=0..3) that share row0 / row1. + for (int off = 2; off > 0; off >>= 1) { + amax0 = fmaxf(amax0, __shfl_xor_sync(0xffffffff, amax0, off)); + amax1 = fmaxf(amax1, __shfl_xor_sync(0xffffffff, amax1, off)); + } + if (l == 0) { + amax_smem[warp_id * BLOCK_M + rloc0] = amax0; + amax_smem[warp_id * BLOCK_M + rloc1] = amax1; + } + } + __syncthreads(); + + // Cross-warp reduce: each warp wrote its row-amax partial. Final reduce per + // row done by warp 0 lanes, broadcast via smem. + #pragma unroll + for (int rloc = t; rloc < BLOCK_M; rloc += THREADS) { + int row = m_base + rloc; + if (row >= M) continue; + float amax = 0.0f; + #pragma unroll + for (int w = 0; w < NUM_WARPS; ++w) + amax = fmaxf(amax, amax_smem[w * BLOCK_M + rloc]); + float sc = fmaxf(amax / kFp8Max, 1.0e-12f); + amax_smem[rloc] = sc; // reuse slot to broadcast final scale + // Each active thread owns a distinct rloc in this strided loop, so each + // writes its own row's scale — no race. (The earlier `warp_id==0 && + // lane==0` guard let only thread 0 write, leaving rows 1..BLOCK_M-1 + // unwritten → garbage out_scale, correct-but-unscaled fp8 output.) + out_scale[(size_t)row * (N >> 7) + (n_base >> 7)] = sc; + } + __syncthreads(); + + // Quantize + store FP8. Thread re-reads its v[] and the row's scale. + #pragma unroll + for (int mi = 0; mi < M_ATOMS; ++mi) { + int row0 = m_base + mi * 16 + h; + int row1 = row0 + 8; + int rloc0 = mi * 16 + h; + int rloc1 = rloc0 + 8; + float sc0 = (row0 < M) ? amax_smem[rloc0] : 1.0f; + float sc1 = (row1 < M) ? amax_smem[rloc1] : 1.0f; + float inv0 = 1.0f / sc0, inv1 = 1.0f / sc1; + #pragma unroll + for (int ni = 0; ni < N_ATOMS_PW; ++ni) { + int n_pair_base = n_base + warp_id * N_ATOMS_PW * 8 + ni * 8 + 2 * l; + if (row0 < M && col_pair_ok(n_pair_base, N)) { + float q0 = fminf(fmaxf(v[mi][ni][0] * inv0, -kFp8Max), kFp8Max); + float q1 = fminf(fmaxf(v[mi][ni][1] * inv0, -kFp8Max), kFp8Max); + // pack two fp8 e4m3 into a 16-bit store + __nv_fp8_e4m3 p0(q0), p1(q1); + uint16_t pack = (uint16_t)(*reinterpret_cast(&p1)) << 8 + | (uint16_t)(*reinterpret_cast(&p0)); + *reinterpret_cast(&output[(size_t)row0 * N + n_pair_base]) = pack; + } else if (row0 < M) { + if (n_pair_base < N) { + float q = fminf(fmaxf(v[mi][ni][0] * inv0, -kFp8Max), kFp8Max); + output[(size_t)row0 * N + n_pair_base] = __nv_fp8_e4m3(q); + } + if (n_pair_base + 1 < N) { + float q = fminf(fmaxf(v[mi][ni][1] * inv0, -kFp8Max), kFp8Max); + output[(size_t)row0 * N + n_pair_base + 1] = __nv_fp8_e4m3(q); + } + } + if (row1 < M && col_pair_ok(n_pair_base, N)) { + float q2 = fminf(fmaxf(v[mi][ni][2] * inv1, -kFp8Max), kFp8Max); + float q3 = fminf(fmaxf(v[mi][ni][3] * inv1, -kFp8Max), kFp8Max); + __nv_fp8_e4m3 p2(q2), p3(q3); + uint16_t pack = (uint16_t)(*reinterpret_cast(&p3)) << 8 + | (uint16_t)(*reinterpret_cast(&p2)); + *reinterpret_cast(&output[(size_t)row1 * N + n_pair_base]) = pack; + } else if (row1 < M) { + if (n_pair_base < N) { + float q = fminf(fmaxf(v[mi][ni][2] * inv1, -kFp8Max), kFp8Max); + output[(size_t)row1 * N + n_pair_base] = __nv_fp8_e4m3(q); + } + if (n_pair_base + 1 < N) { + float q = fminf(fmaxf(v[mi][ni][3] * inv1, -kFp8Max), kFp8Max); + output[(size_t)row1 * N + n_pair_base + 1] = __nv_fp8_e4m3(q); + } + } + } + } +} + +// ============================================================================ +// GeGLU silu-fold, A-persistent two-pass variant. +// +// Same fusion as fp8_bs_geglu_silu_fold_kernel (gate+up GEMM + silu(gate)*up + +// per-token block-128 FP8 quant, one launch, no [M,2N] BF16 transient), but a +// different smem/register strategy that fixes the two-pass weaknesses the ncu +// diagnosis isolated: +// +// two-pass loss = (a) A re-loaded twice (2*M*K HBM) + (b) 2x pipeline drain. +// interleaved loss = 2x B smem (both gate+up staged) -> 1 CTA/SM occupancy. +// +// A-persistent: stage A into smem ONCE (reused by both the gate pass and the up +// pass), but keep only ONE B smem region that is filled with B_gate for the +// gate pass and then RE-FILLED with B_up for the up pass (sequential, not +// simultaneous). So: +// - A loaded once from HBM (the interleaved HBM win), held in smem across +// both passes -> no A re-load. act_scale also staged once. +// - B smem = a single STAGES*BN*BK region (NOT doubled) -> fits 4 CTA/SM. +// - only ONE accumulator live at a time (gate_acc -> store silu(gate) to a +// small smem gate buffer -> reuse regs for up_acc) -> ~70 regs (two-pass +// register profile), not the ~140 of true interleaved. +// +// The catch: A must fit in smem for the whole K-walk (A is [BM, K], not +// [BM, BK]), so this only works when K is small enough that BM*K fp8 + the rest +// stays under the smem budget — i.e. the Qwen3-VL gate_up shapes where K=hidden +// (2B K=2048, 8B K=4096). For BM=32: A_persist = 32*K = 64KB (8B) / 32KB (2B). +// 8B 64KB alone already exceeds a CTA's smem, so A-persistent is only viable +// for the 2B shape (K=2048) at BM<=32, OR by staging A in K-chunks and walking +// gate+up together within each K-chunk (chunked-interleaved). The chunked form +// is implemented here: A is staged per BLOCK_K tile like the baseline, but BOTH +// the gate MMA and the up MMA for that K-tile run before the tile is evicted — +// i.e. gate and up advance K-tile-by-K-tile together (true interleaved per +// K-tile, the sm100 geglu pattern), yet B is staged in ONE region reused for +// gate-then-up WITHIN the k-iter (load B_gate, gate-MMA, load B_up into the +// SAME region, up-MMA). That keeps B smem single (no 2x) AND loads A once AND +// holds both gate_acc+up_acc in regs (interleaved) — but pays by serializing +// the two B loads within a k-iter (no overlap between B_gate and B_up loads). +// +// Net vs two-pass: A loaded once (saves M*K HBM), one pipeline drain (K_ITERS +// stalls not 2*K_ITERS), but B_gate/B_up loads are serial within each k-iter. +// Net vs interleaved: B smem halved (2 CTA/SM recoverable to 3-4), but loses +// B_gate||B_up load overlap. On sm89 (HBM-bound, no TMA) the smem/occupancy +// recovery usually dominates, so this is the predicted winner. +// ============================================================================ +template +__global__ __launch_bounds__(NUM_WARPS * 32, MIN_BLOCKS_PER_SM) +void fp8_bs_geglu_silu_fold_apersist_kernel( + const __nv_fp8_e4m3* __restrict__ A, + const __nv_fp8_e4m3* __restrict__ B, // gate_up_w [2*N, K] + const float* __restrict__ act_scale, // [M, K/128] + const float* __restrict__ w_scale, // gate_up_s [2*N/128, K/128] + __nv_fp8_e4m3* __restrict__ output, // [M, N] + float* __restrict__ out_scale, // [M, N/128] + int M, int N, int K) +{ + static_assert(BLOCK_N == 128, + "GeGLU silu-fold requires BLOCK_N==128 (one quant block per CTA)"); + constexpr int BLOCK_K = 128; + constexpr int THREADS = NUM_WARPS * 32; + constexpr int M_ATOMS = BLOCK_M / 16; + constexpr int N_ATOMS = BLOCK_N / 8; // 16 + constexpr int N_ATOMS_PW = N_ATOMS / NUM_WARPS; + constexpr int N_PAIRS_PW = N_ATOMS_PW / 2; + constexpr int K_ATOMS = BLOCK_K / 32; // 4 + constexpr int NUM_CHUNKS_PER_ROW = BLOCK_K / 16; + constexpr int SWIZZLE_MASK = NUM_CHUNKS_PER_ROW - 1; + constexpr int SCALE_KTILE = 8; + constexpr int A_TILE = BLOCK_M * BLOCK_K; + constexpr int B_TILE = BLOCK_N * BLOCK_K; + + static_assert(BLOCK_M % 16 == 0, "BLOCK_M multiple of 16"); + static_assert(N_ATOMS_PW >= 2 && N_ATOMS_PW % 2 == 0, + "ldmatrix pairs 2 N-atoms: N_ATOMS_PW must be even >= 2"); + + extern __shared__ uint8_t smem_raw[]; + uint8_t* A_smem = smem_raw; // STAGES * A_TILE + uint8_t* B_smem = A_smem + STAGES * A_TILE; // STAGES * B_TILE (reused gate/up) + // gate_smem: silu(gate) BF16, [BM, BN]. Written by gate epilogue, read by + // the final silu(gate)*up epilogue (NOT per k-iter — only once at the end). + __nv_bfloat16* gate_smem = reinterpret_cast<__nv_bfloat16*>( + B_smem + STAGES * B_TILE); + float* as_smem = reinterpret_cast(gate_smem + BLOCK_M * BLOCK_N); + float* wsg_smem = as_smem + BLOCK_M * SCALE_KTILE; // gate w_scale row + float* wsu_smem = wsg_smem + SCALE_KTILE; // up w_scale row + float* amax_smem = wsu_smem + SCALE_KTILE; + + const int cta_m = blockIdx.x; + const int cta_n = blockIdx.y; + const int m_base = cta_m * BLOCK_M; + const int n_base = cta_n * BLOCK_N; + const int gate_b_row0 = n_base; + const int up_b_row0 = n_base + N; + + const int t = threadIdx.x; + const int warp_id = t / 32; + const int lane = t % 32; + const int l = lane % 4; + const int h = lane / 4; + const int frag_group = lane / 8; + const int row_in_frag = lane % 8; + const int row_block = frag_group / 2; + const int col_block = frag_group % 2; + + const int K128 = K >> 7; + const int N128 = N >> 7; + const int gate_ws_row = (n_base >> 7); + const int up_ws_row = gate_ws_row + N128; + + auto stage_scales = [&](int kb0) { + const int as_total = BLOCK_M * SCALE_KTILE; + for (int idx = t; idx < as_total; idx += THREADS) { + int r = idx / SCALE_KTILE; + int kc = idx - r * SCALE_KTILE; + int row = m_base + r; + int kb = kb0 + kc; + as_smem[idx] = (row < M && kb < K128) + ? act_scale[(size_t)row * K128 + kb] : 0.0f; + } + for (int kc = t; kc < SCALE_KTILE; kc += THREADS) { + int kb = kb0 + kc; + wsg_smem[kc] = (kb < K128) + ? w_scale[(size_t)gate_ws_row * K128 + kb] : 0.0f; + wsu_smem[kc] = (kb < K128) + ? w_scale[(size_t)up_ws_row * K128 + kb] : 0.0f; + } + __syncthreads(); + }; + + // Stage A + one B band (gate or up) into smem. b_row0 picks the band. + auto issue_load = [&](int stage, int k_base, int b_row0) { + constexpr int A_CHUNKS = BLOCK_M * NUM_CHUNKS_PER_ROW; + constexpr int A_ITERS = (A_CHUNKS + THREADS - 1) / THREADS; + #pragma unroll + for (int it = 0; it < A_ITERS; ++it) { + int idx = it * THREADS + t; + if (idx >= A_CHUNKS) break; + int row_a = idx / NUM_CHUNKS_PER_ROW; + int chunk_a = idx % NUM_CHUNKS_PER_ROW; + int m_glob = m_base + row_a; + int k_glob = k_base + chunk_a * 16; + const uint8_t* a_src = nullptr; + if (m_glob < M && k_glob < K) { + a_src = reinterpret_cast(&A[(size_t)m_glob * K + k_glob]); + } + int csw = chunk_a ^ (row_a & SWIZZLE_MASK); + cp_async_16( + to_smem(&A_smem[stage * A_TILE + row_a * BLOCK_K + csw * 16]), + a_src); + } + constexpr int B_CHUNKS = BLOCK_N * NUM_CHUNKS_PER_ROW; + constexpr int B_ITERS = (B_CHUNKS + THREADS - 1) / THREADS; + #pragma unroll + for (int it = 0; it < B_ITERS; ++it) { + int idx = it * THREADS + t; + if (idx >= B_CHUNKS) break; + int row_b = idx / NUM_CHUNKS_PER_ROW; + int chunk_b = idx % NUM_CHUNKS_PER_ROW; + int n_glob = b_row0 + row_b; + int k_glob = k_base + chunk_b * 16; + const uint8_t* b_src = nullptr; + if (n_glob < 2 * N && k_glob < K) { + b_src = reinterpret_cast(&B[(size_t)n_glob * K + k_glob]); + } + int csw = chunk_b ^ (row_b & SWIZZLE_MASK); + cp_async_16( + to_smem(&B_smem[stage * B_TILE + row_b * BLOCK_K + csw * 16]), + b_src); + } + }; + + // MMA pass over the staged A/B tiles for one k-iter, accumulating into `acc` + // with the given w_scale smem row. + auto mma_tile = [&](float (*acc)[N_ATOMS_PW][4], int compute_stage, + const float* ws_smem_pass) { + const int kb = (compute_stage); // caller passes k_iter; recompute below + (void)kb; + float tacc[M_ATOMS][N_ATOMS_PW][4]; + #pragma unroll + for (int mi = 0; mi < M_ATOMS; ++mi) + #pragma unroll + for (int ni = 0; ni < N_ATOMS_PW; ++ni) + #pragma unroll + for (int j = 0; j < 4; ++j) tacc[mi][ni][j] = 0.0f; + + uint8_t* A_stage = A_smem + compute_stage * A_TILE; + uint8_t* B_stage = B_smem + compute_stage * B_TILE; + #pragma unroll + for (int ka = 0; ka < K_ATOMS; ++ka) { + uint32_t A_regs[M_ATOMS][4]; + #pragma unroll + for (int mi = 0; mi < M_ATOMS; ++mi) { + int row = mi * 16 + row_block * 8 + row_in_frag; + int chunk = 2 * ka + col_block; + int csw = chunk ^ (row & SWIZZLE_MASK); + ldmatrix_x4_b16(A_regs[mi][0], A_regs[mi][1], A_regs[mi][2], A_regs[mi][3], + to_smem(&A_stage[row * BLOCK_K + csw * 16])); + } + uint32_t B_regs[N_PAIRS_PW][4]; + #pragma unroll + for (int np = 0; np < N_PAIRS_PW; ++np) { + int nrow = warp_id * N_ATOMS_PW * 8 + np * 16 + row_block * 8 + row_in_frag; + int chunk = 2 * ka + col_block; + int csw = chunk ^ (nrow & SWIZZLE_MASK); + ldmatrix_x4_b16(B_regs[np][0], B_regs[np][1], B_regs[np][2], B_regs[np][3], + to_smem(&B_stage[nrow * BLOCK_K + csw * 16])); + } + #pragma unroll + for (int mi = 0; mi < M_ATOMS; ++mi) { + #pragma unroll + for (int np = 0; np < N_PAIRS_PW; ++np) { + int ni0 = np * 2, ni1 = np * 2 + 1; + mma_m16n8k32_e4m3( + tacc[mi][ni0][0], tacc[mi][ni0][1], tacc[mi][ni0][2], tacc[mi][ni0][3], + A_regs[mi][0], A_regs[mi][2], A_regs[mi][1], A_regs[mi][3], + B_regs[np][0], B_regs[np][1]); + mma_m16n8k32_e4m3( + tacc[mi][ni1][0], tacc[mi][ni1][1], tacc[mi][ni1][2], tacc[mi][ni1][3], + A_regs[mi][0], A_regs[mi][2], A_regs[mi][1], A_regs[mi][3], + B_regs[np][2], B_regs[np][3]); + } + } + } + return tacc; // caller folds scales into acc + }; + + // Running accumulators for gate and up, both live across the whole K-loop + // (true interleaved: both gate and up advance K-tile-by-K-tile together). + float gate_acc[M_ATOMS][N_ATOMS_PW][4]; + float up_acc[M_ATOMS][N_ATOMS_PW][4]; + #pragma unroll + for (int mi = 0; mi < M_ATOMS; ++mi) + #pragma unroll + for (int ni = 0; ni < N_ATOMS_PW; ++ni) + #pragma unroll + for (int j = 0; j < 4; ++j) { + gate_acc[mi][ni][j] = 0.0f; + up_acc[mi][ni][j] = 0.0f; + } + + const int K_ITERS = (K + BLOCK_K - 1) / BLOCK_K; + // Prefetch STAGES-1 A tiles (A is shared by both passes — issued once). + // B is NOT prefetched here; within each k-iter we issue B_gate then B_up + // into the SAME smem region after the previous iter's B is consumed. + #pragma unroll + for (int s = 0; s < STAGES - 1; ++s) { + int kb = s * BLOCK_K; + if (kb < K) issue_load(s, kb, gate_b_row0); // first prefetch = gate B + asm volatile("cp.async.commit_group;\n" ::); + } + + int compute_stage = 0; + for (int k_iter = 0; k_iter < K_ITERS; ++k_iter) { + int issue_iter = k_iter + (STAGES - 1); + int issue_stage = issue_iter % STAGES; + // Issue the NEXT A tile + the NEXT gate-B tile (the up-B for this k_iter + // is loaded inside the gate-MMA sync below, reusing B_smem after gate + // MMA reads finish). + if (issue_iter < K_ITERS) issue_load(issue_stage, issue_iter * BLOCK_K, gate_b_row0); + asm volatile("cp.async.commit_group;\n" ::); + asm volatile("cp.async.wait_group %0;\n" :: "n"(STAGES - 1)); + __syncthreads(); + + const int kb = k_iter; + if ((kb % SCALE_KTILE) == 0) stage_scales(kb); + + // ---- gate MMA on the staged A + B_gate ---- + { + float tacc[M_ATOMS][N_ATOMS_PW][4]; + #pragma unroll + for (int mi = 0; mi < M_ATOMS; ++mi) + #pragma unroll + for (int ni = 0; ni < N_ATOMS_PW; ++ni) + #pragma unroll + for (int j = 0; j < 4; ++j) tacc[mi][ni][j] = 0.0f; + uint8_t* A_stage = A_smem + compute_stage * A_TILE; + uint8_t* B_stage = B_smem + compute_stage * B_TILE; + #pragma unroll + for (int ka = 0; ka < K_ATOMS; ++ka) { + uint32_t A_regs[M_ATOMS][4]; + #pragma unroll + for (int mi = 0; mi < M_ATOMS; ++mi) { + int row = mi * 16 + row_block * 8 + row_in_frag; + int chunk = 2 * ka + col_block; + int csw = chunk ^ (row & SWIZZLE_MASK); + ldmatrix_x4_b16(A_regs[mi][0], A_regs[mi][1], A_regs[mi][2], A_regs[mi][3], + to_smem(&A_stage[row * BLOCK_K + csw * 16])); + } + uint32_t B_regs[N_PAIRS_PW][4]; + #pragma unroll + for (int np = 0; np < N_PAIRS_PW; ++np) { + int nrow = warp_id * N_ATOMS_PW * 8 + np * 16 + row_block * 8 + row_in_frag; + int chunk = 2 * ka + col_block; + int csw = chunk ^ (nrow & SWIZZLE_MASK); + ldmatrix_x4_b16(B_regs[np][0], B_regs[np][1], B_regs[np][2], B_regs[np][3], + to_smem(&B_stage[nrow * BLOCK_K + csw * 16])); + } + #pragma unroll + for (int mi = 0; mi < M_ATOMS; ++mi) { + #pragma unroll + for (int np = 0; np < N_PAIRS_PW; ++np) { + int ni0 = np * 2, ni1 = np * 2 + 1; + mma_m16n8k32_e4m3( + tacc[mi][ni0][0], tacc[mi][ni0][1], tacc[mi][ni0][2], tacc[mi][ni0][3], + A_regs[mi][0], A_regs[mi][2], A_regs[mi][1], A_regs[mi][3], + B_regs[np][0], B_regs[np][1]); + mma_m16n8k32_e4m3( + tacc[mi][ni1][0], tacc[mi][ni1][1], tacc[mi][ni1][2], tacc[mi][ni1][3], + A_regs[mi][0], A_regs[mi][2], A_regs[mi][1], A_regs[mi][3], + B_regs[np][2], B_regs[np][3]); + } + } + } + int kbt = kb % SCALE_KTILE; + float ws_cta = wsg_smem[kbt]; + #pragma unroll + for (int mi = 0; mi < M_ATOMS; ++mi) { + int row0 = m_base + mi * 16 + h; + int row1 = row0 + 8; + float as0 = as_smem[(mi * 16 + h) * SCALE_KTILE + kbt]; + float as1 = as_smem[(mi * 16 + h + 8) * SCALE_KTILE + kbt]; + #pragma unroll + for (int ni = 0; ni < N_ATOMS_PW; ++ni) { + gate_acc[mi][ni][0] += tacc[mi][ni][0] * (as0 * ws_cta); + gate_acc[mi][ni][1] += tacc[mi][ni][1] * (as0 * ws_cta); + gate_acc[mi][ni][2] += tacc[mi][ni][2] * (as1 * ws_cta); + gate_acc[mi][ni][3] += tacc[mi][ni][3] * (as1 * ws_cta); + } + } + } + __syncthreads(); // B_smem safe to overwrite with B_up + + // ---- load B_up into the SAME B_smem region, then up MMA ---- + // (A is still resident in A_smem[compute_stage]; not reloaded from HBM.) + { + // issue B_up into B_smem[compute_stage] (A_smem left untouched) + constexpr int B_CHUNKS = BLOCK_N * NUM_CHUNKS_PER_ROW; + constexpr int B_ITERS = (B_CHUNKS + THREADS - 1) / THREADS; + int k_base = k_iter * BLOCK_K; + #pragma unroll + for (int it = 0; it < B_ITERS; ++it) { + int idx = it * THREADS + t; + if (idx >= B_CHUNKS) break; + int row_b = idx / NUM_CHUNKS_PER_ROW; + int chunk_b = idx % NUM_CHUNKS_PER_ROW; + int n_glob = up_b_row0 + row_b; + int k_glob = k_base + chunk_b * 16; + const uint8_t* b_src = nullptr; + if (n_glob < 2 * N && k_glob < K) { + b_src = reinterpret_cast(&B[(size_t)n_glob * K + k_glob]); + } + int csw = chunk_b ^ (row_b & SWIZZLE_MASK); + cp_async_16( + to_smem(&B_smem[compute_stage * B_TILE + row_b * BLOCK_K + csw * 16]), + b_src); + } + asm volatile("cp.async.commit_group;\n" ::); + asm volatile("cp.async.wait_group %0;\n" :: "n"(0)); + __syncthreads(); + + float tacc[M_ATOMS][N_ATOMS_PW][4]; + #pragma unroll + for (int mi = 0; mi < M_ATOMS; ++mi) + #pragma unroll + for (int ni = 0; ni < N_ATOMS_PW; ++ni) + #pragma unroll + for (int j = 0; j < 4; ++j) tacc[mi][ni][j] = 0.0f; + uint8_t* A_stage = A_smem + compute_stage * A_TILE; + uint8_t* B_stage = B_smem + compute_stage * B_TILE; + #pragma unroll + for (int ka = 0; ka < K_ATOMS; ++ka) { + uint32_t A_regs[M_ATOMS][4]; + #pragma unroll + for (int mi = 0; mi < M_ATOMS; ++mi) { + int row = mi * 16 + row_block * 8 + row_in_frag; + int chunk = 2 * ka + col_block; + int csw = chunk ^ (row & SWIZZLE_MASK); + ldmatrix_x4_b16(A_regs[mi][0], A_regs[mi][1], A_regs[mi][2], A_regs[mi][3], + to_smem(&A_stage[row * BLOCK_K + csw * 16])); + } + uint32_t B_regs[N_PAIRS_PW][4]; + #pragma unroll + for (int np = 0; np < N_PAIRS_PW; ++np) { + int nrow = warp_id * N_ATOMS_PW * 8 + np * 16 + row_block * 8 + row_in_frag; + int chunk = 2 * ka + col_block; + int csw = chunk ^ (nrow & SWIZZLE_MASK); + ldmatrix_x4_b16(B_regs[np][0], B_regs[np][1], B_regs[np][2], B_regs[np][3], + to_smem(&B_stage[nrow * BLOCK_K + csw * 16])); + } + #pragma unroll + for (int mi = 0; mi < M_ATOMS; ++mi) { + #pragma unroll + for (int np = 0; np < N_PAIRS_PW; ++np) { + int ni0 = np * 2, ni1 = np * 2 + 1; + mma_m16n8k32_e4m3( + tacc[mi][ni0][0], tacc[mi][ni0][1], tacc[mi][ni0][2], tacc[mi][ni0][3], + A_regs[mi][0], A_regs[mi][2], A_regs[mi][1], A_regs[mi][3], + B_regs[np][0], B_regs[np][1]); + mma_m16n8k32_e4m3( + tacc[mi][ni1][0], tacc[mi][ni1][1], tacc[mi][ni1][2], tacc[mi][ni1][3], + A_regs[mi][0], A_regs[mi][2], A_regs[mi][1], A_regs[mi][3], + B_regs[np][2], B_regs[np][3]); + } + } + } + int kbt = kb % SCALE_KTILE; + float ws_cta = wsu_smem[kbt]; + #pragma unroll + for (int mi = 0; mi < M_ATOMS; ++mi) { + int row0 = m_base + mi * 16 + h; + int row1 = row0 + 8; + float as0 = as_smem[(mi * 16 + h) * SCALE_KTILE + kbt]; + float as1 = as_smem[(mi * 16 + h + 8) * SCALE_KTILE + kbt]; + #pragma unroll + for (int ni = 0; ni < N_ATOMS_PW; ++ni) { + up_acc[mi][ni][0] += tacc[mi][ni][0] * (as0 * ws_cta); + up_acc[mi][ni][1] += tacc[mi][ni][1] * (as0 * ws_cta); + up_acc[mi][ni][2] += tacc[mi][ni][2] * (as1 * ws_cta); + up_acc[mi][ni][3] += tacc[mi][ni][3] * (as1 * ws_cta); + } + } + } + __syncthreads(); + compute_stage = (compute_stage + 1) % STAGES; + } + asm volatile("cp.async.wait_all;\n" ::); + + // ============ Epilogue: silu(gate)*up + per-row amax + quant ============ + // gate_acc and up_acc both live in registers. Replicate silu_mul_merged's + // two bf16 roundings: bf16(silu(gate)) then bf16(silu_bf * up). + constexpr float kFp8Max = 448.0f; + float v[M_ATOMS][N_ATOMS_PW][4]; + #pragma unroll + for (int mi = 0; mi < M_ATOMS; ++mi) { + int row0 = m_base + mi * 16 + h; + int row1 = row0 + 8; + int rloc0 = mi * 16 + h; + int rloc1 = rloc0 + 8; + float amax0 = 0.0f, amax1 = 0.0f; + #pragma unroll + for (int ni = 0; ni < N_ATOMS_PW; ++ni) { + if (row0 < M) { + float gf0 = __bfloat162float(__float2bfloat16(silu_f32(gate_acc[mi][ni][0]))); + float gf1 = __bfloat162float(__float2bfloat16(silu_f32(gate_acc[mi][ni][1]))); + v[mi][ni][0] = __bfloat162float(__float2bfloat16(gf0 * up_acc[mi][ni][0])); + v[mi][ni][1] = __bfloat162float(__float2bfloat16(gf1 * up_acc[mi][ni][1])); + amax0 = fmaxf(amax0, fmaxf(fabsf(v[mi][ni][0]), fabsf(v[mi][ni][1]))); + } else { v[mi][ni][0] = 0.0f; v[mi][ni][1] = 0.0f; } + if (row1 < M) { + float gf0 = __bfloat162float(__float2bfloat16(silu_f32(gate_acc[mi][ni][2]))); + float gf1 = __bfloat162float(__float2bfloat16(silu_f32(gate_acc[mi][ni][3]))); + v[mi][ni][2] = __bfloat162float(__float2bfloat16(gf0 * up_acc[mi][ni][2])); + v[mi][ni][3] = __bfloat162float(__float2bfloat16(gf1 * up_acc[mi][ni][3])); + amax1 = fmaxf(amax1, fmaxf(fabsf(v[mi][ni][2]), fabsf(v[mi][ni][3]))); + } else { v[mi][ni][2] = 0.0f; v[mi][ni][3] = 0.0f; } + } + for (int off = 2; off > 0; off >>= 1) { + amax0 = fmaxf(amax0, __shfl_xor_sync(0xffffffff, amax0, off)); + amax1 = fmaxf(amax1, __shfl_xor_sync(0xffffffff, amax1, off)); + } + if (l == 0) { + amax_smem[warp_id * BLOCK_M + rloc0] = amax0; + amax_smem[warp_id * BLOCK_M + rloc1] = amax1; + } + } + __syncthreads(); + + #pragma unroll + for (int rloc = t; rloc < BLOCK_M; rloc += THREADS) { + int row = m_base + rloc; + if (row >= M) continue; + float amax = 0.0f; + #pragma unroll + for (int w = 0; w < NUM_WARPS; ++w) + amax = fmaxf(amax, amax_smem[w * BLOCK_M + rloc]); + float sc = fmaxf(amax / kFp8Max, 1.0e-12f); + amax_smem[rloc] = sc; + // Each active thread owns a distinct rloc — write its row's scale + // directly (no single-thread guard; see two-pass variant for the bug + // this fixes). + out_scale[(size_t)row * (N >> 7) + (n_base >> 7)] = sc; + } + __syncthreads(); + + #pragma unroll + for (int mi = 0; mi < M_ATOMS; ++mi) { + int row0 = m_base + mi * 16 + h; + int row1 = row0 + 8; + int rloc0 = mi * 16 + h; + int rloc1 = rloc0 + 8; + float sc0 = (row0 < M) ? amax_smem[rloc0] : 1.0f; + float sc1 = (row1 < M) ? amax_smem[rloc1] : 1.0f; + float inv0 = 1.0f / sc0, inv1 = 1.0f / sc1; + #pragma unroll + for (int ni = 0; ni < N_ATOMS_PW; ++ni) { + int n_pair_base = n_base + warp_id * N_ATOMS_PW * 8 + ni * 8 + 2 * l; + if (row0 < M && col_pair_ok(n_pair_base, N)) { + float q0 = fminf(fmaxf(v[mi][ni][0] * inv0, -kFp8Max), kFp8Max); + float q1 = fminf(fmaxf(v[mi][ni][1] * inv0, -kFp8Max), kFp8Max); + __nv_fp8_e4m3 p0(q0), p1(q1); + uint16_t pack = (uint16_t)(*reinterpret_cast(&p1)) << 8 + | (uint16_t)(*reinterpret_cast(&p0)); + *reinterpret_cast(&output[(size_t)row0 * N + n_pair_base]) = pack; + } else if (row0 < M) { + if (n_pair_base < N) output[(size_t)row0 * N + n_pair_base] = __nv_fp8_e4m3(fminf(fmaxf(v[mi][ni][0] * inv0, -kFp8Max), kFp8Max)); + if (n_pair_base + 1 < N) output[(size_t)row0 * N + n_pair_base + 1] = __nv_fp8_e4m3(fminf(fmaxf(v[mi][ni][1] * inv0, -kFp8Max), kFp8Max)); + } + if (row1 < M && col_pair_ok(n_pair_base, N)) { + float q2 = fminf(fmaxf(v[mi][ni][2] * inv1, -kFp8Max), kFp8Max); + float q3 = fminf(fmaxf(v[mi][ni][3] * inv1, -kFp8Max), kFp8Max); + __nv_fp8_e4m3 p2(q2), p3(q3); + uint16_t pack = (uint16_t)(*reinterpret_cast(&p3)) << 8 + | (uint16_t)(*reinterpret_cast(&p2)); + *reinterpret_cast(&output[(size_t)row1 * N + n_pair_base]) = pack; + } else if (row1 < M) { + if (n_pair_base < N) output[(size_t)row1 * N + n_pair_base] = __nv_fp8_e4m3(fminf(fmaxf(v[mi][ni][2] * inv1, -kFp8Max), kFp8Max)); + if (n_pair_base + 1 < N) output[(size_t)row1 * N + n_pair_base + 1] = __nv_fp8_e4m3(fminf(fmaxf(v[mi][ni][3] * inv1, -kFp8Max), kFp8Max)); + } + } + } + (void)gate_smem; // apersist keeps gate in regs; gate_smem unused (kept for layout parity) + (void)mma_tile; // helper retained for future chunked variant; unused in this path +} + } // namespace block128_sm89 } // namespace gemm } // namespace flash_rt diff --git a/csrc/qwen3_vl_bindings.cpp b/csrc/qwen3_vl_bindings.cpp index 94bf9a41..4cdda750 100644 --- a/csrc/qwen3_vl_bindings.cpp +++ b/csrc/qwen3_vl_bindings.cpp @@ -292,6 +292,45 @@ PYBIND11_MODULE(flash_rt_qwen3_vl_kernels, m) { py::arg("act_block_scale"), py::arg("w_block_scale"), py::arg("stream") = 0); + // GeGLU silu-fold (production entry, env-gated by the caller): fuses the + // gate_up GEMM + silu_mul + per-token block-128 FP8 quant into one launch. + // This is a NEW binding (no legacy alias replaced). Argument shapes: + // A : [M, K] FP8 e4m3 row-major (per-token quantized act) + // B : [2*N, K] FP8 e4m3 row-major (gate_up_w: gate rows + // [0,N), up rows [N,2N)) + // act_scale : [M, K/128] fp32 row-major + // w_scale : [2*N/128, K/128] fp32 row-major (gate_up_s; up row = + // gate row + N/128) + // output : [M, N] FP8 e4m3 row-major + // out_scale : [M, N/128] fp32 row-major + // N and K must be multiples of 128. Beats the baseline (gate_up GEMM + + // silu_mul) only in the small-M (launch-bound) regime; the frontend + // dispatcher gates on M before calling this. Additive: the baseline path is + // unchanged when this is not selected. See fp8_bs_geglu_silu_fold_kernel in + // fp8_bs_gemm_device.cuh. + m.def("fp8_bs_geglu_silu_fold_sm89_fp8out", + [](uintptr_t A, uintptr_t B, + int M, int N, int K, + uintptr_t act_scale, uintptr_t w_scale, + uintptr_t output, uintptr_t out_scale, + uintptr_t stream) { + int rc = flash_rt::gemm::block128_sm89:: + fp8_bs_geglu_silu_fold_sm89_32x128_w4_s1( + to_ptr(A), to_ptr(B), M, N, K, + reinterpret_cast(act_scale), + reinterpret_cast(w_scale), + to_ptr(output), reinterpret_cast(out_scale), + to_stream(stream)); + if (rc != 0) + throw std::runtime_error( + "fp8_bs_geglu_silu_fold_sm89_fp8out launch failed"); + }, + py::arg("A"), py::arg("B"), + py::arg("M"), py::arg("N"), py::arg("K"), + py::arg("act_block_scale"), py::arg("w_block_scale"), + py::arg("output"), py::arg("out_scale"), + py::arg("stream") = 0); + // Bench-only tile-variant bindings for prefill GEMM tuning. Not used by // the frontend; exposed only for explicit Qwen3-VL dev builds so the // production pybind surface stays runtime-only. @@ -316,6 +355,22 @@ PYBIND11_MODULE(flash_rt_qwen3_vl_kernels, m) { BIND_GEMM_TILE_RESID(fp8_block128_gemm_bs_sm89_64x64x128_w4_s1_resid); BIND_GEMM_TILE_RESID(fp8_block128_gemm_bs_sm89_128x128x128_w8_s1_resid); #undef BIND_GEMM_TILE_RESID + + // GeGLU silu-fold bench bindings: fuse gate+up GEMM + silu(gate)*up + + // per-token block-128 FP8 quant into one launch (no [M,2N] BF16 transient). + // B = gate_up_w [2*N, K], w_scale = gate_up_s [2*N/128, K/128]; outputs + // FP8 [M,N] + scale [M,N/128]. Dev-builds only. +#define BIND_GEGLU_TILE(NAME) m.def("bench_" #NAME, [](uintptr_t A, uintptr_t B, int M, int N, int K, uintptr_t act_scale, uintptr_t w_scale, uintptr_t output, uintptr_t out_scale, uintptr_t stream) { return flash_rt::gemm::block128_sm89::NAME( to_ptr(A), to_ptr(B), M, N, K, reinterpret_cast(act_scale), reinterpret_cast(w_scale), to_ptr(output), reinterpret_cast(out_scale), to_stream(stream)); }, py::arg("A"), py::arg("B"), py::arg("M"), py::arg("N"), py::arg("K"), py::arg("act_block_scale"), py::arg("w_block_scale"), py::arg("output"), py::arg("out_scale"), py::arg("stream") = 0) + BIND_GEGLU_TILE(fp8_bs_geglu_silu_fold_sm89_32x128_w4_s2); + BIND_GEGLU_TILE(fp8_bs_geglu_silu_fold_sm89_16x128_w4_s2); + BIND_GEGLU_TILE(fp8_bs_geglu_silu_fold_sm89_64x128_w4_s2); + BIND_GEGLU_TILE(fp8_bs_geglu_silu_fold_sm89_128x128_w8_s1); + BIND_GEGLU_TILE(fp8_bs_geglu_silu_fold_sm89_32x128_w4_s1); + BIND_GEGLU_TILE(fp8_bs_geglu_silu_fold_sm89_16x128_w4_s1); + BIND_GEGLU_TILE(fp8_bs_geglu_silu_fold_apersist_sm89_32x128_w4_s1); + BIND_GEGLU_TILE(fp8_bs_geglu_silu_fold_apersist_sm89_16x128_w4_s1); + BIND_GEGLU_TILE(fp8_bs_geglu_silu_fold_apersist_sm89_32x128_w4_s2); +#undef BIND_GEGLU_TILE #endif m.def("qwen3_qk_norm_rope_kvwrite_bf16", diff --git a/flash_rt/frontends/torch/qwen3_vl_fp8_sm89.py b/flash_rt/frontends/torch/qwen3_vl_fp8_sm89.py index b9fbf5bb..79b18179 100644 --- a/flash_rt/frontends/torch/qwen3_vl_fp8_sm89.py +++ b/flash_rt/frontends/torch/qwen3_vl_fp8_sm89.py @@ -69,7 +69,8 @@ def __init__(self, checkpoint_path: str, *, fuse_gate_up: bool = True, fuse_qk_postproc: bool = True, use_fp8_lm_head: bool = True, - max_decode_graphs: int | None = None) -> None: + max_decode_graphs: int | None = None, + use_silu_fold_prefill: bool = False) -> None: import json self.checkpoint_path = str(checkpoint_path) @@ -80,6 +81,21 @@ def __init__(self, checkpoint_path: str, *, self.fuse_gate_up = bool(fuse_gate_up) self.fuse_qk_postproc = bool(fuse_qk_postproc) self.use_fp8_lm_head = bool(use_fp8_lm_head) + # Opt-in: fuse gate_up GEMM + silu_mul + FP8 quant into one megakernel + # on the prefill path (only the small-M / launch-bound regime benefits; + # gated per-call on S below the crossover). Default off → baseline + # two-launch path runs unchanged. Env override: + # FLASH_RT_SM89_SILU_FOLD_PREFILL=1 forces on, =0 forces off. + env_fold = os.environ.get('FLASH_RT_SM89_SILU_FOLD_PREFILL') + self.use_silu_fold_prefill = ( + bool(use_silu_fold_prefill) if env_fold is None + else env_fold == '1') + # The fused 32x128 tile beats the two-launch baseline only in the + # launch-bound (small-S) regime; above the crossover the gate_up GEMM + # is compute-bound and the fused kernel's two-pass drain loses. The + # winning S band is resolved per-inter after config load (see below). + self._silu_fold_band = None # resolved after config load + self._silu_fold_tile = None # resolved after self._fvk is set (below) self._tokenizer: Any = None self._weights = None self._cfg: dict | None = None @@ -136,6 +152,10 @@ def __init__(self, checkpoint_path: str, *, 'and build flash_rt_kernels flash_rt_fa2 ' 'flash_rt_qwen3_vl_kernels.') self._fvk = fvk + if self.use_silu_fold_prefill: + self._silu_fold_tile = getattr( + fvk, 'fp8_bs_geglu_silu_fold_sm89_fp8out', None) + self.use_silu_fold_prefill = self._silu_fold_tile is not None cfg_path = os.path.join(self.checkpoint_path, 'config.json') cfg = json.load(open(cfg_path)) @@ -171,6 +191,16 @@ def __init__(self, checkpoint_path: str, *, 'dimensions compatible with the SM89 block-128 FP8 kernels: ' + '; '.join(problems) + f' (from {cfg_path})') + # Resolve the silu-fold S crossover from inter: the fused 32x128 tile's + # two-pass drain tax scales with K=hidden and the gate_up compute + # density, so the wider-MLP 8B (inter=12288) crosses over sooner than + # 2B (inter=6144). Tuned on 4090 (issue #1 Phase 2 e2e sweep). + if self.use_silu_fold_prefill: + # Winning S band (tuned on 4090, issue #1 Phase 2): the fused + # 32x128 tile needs enough M-tiles to fill the 128-SM grid but + # stays below the compute-bound crossover. 2B: [64,192]; 8B: [96,128]. + self._silu_fold_band = ((96, 192) if inter <= 8192 else (128, 128)) + self._load_fp8_path() self._alloc_buffers() self._build_rope_table() @@ -651,13 +681,28 @@ def _layer_forward_prefill_fp8_blockscaled(self, L: int, h_in_S, cos_S, up_out = self._prefill_up_out[:S] ap_dn, sc_dn, down_out = self._prefill_fp8_scratch[(hidden, inter)] if self.fuse_gate_up: - _, _, gate_up_out = self._prefill_fp8_scratch[(2 * inter, hidden)] - self._prefill_gemm( - ap_mlp, sc_mlp, int(lw['gate_up_w']), int(lw['gate_up_s']), - gate_up_out, int(lw['gate_up_N']), hidden, S) - fvk.silu_mul_merged_to_fp8_block128_bf16( - gate_up_out[:S].data_ptr(), - ap_dn.data_ptr(), sc_dn.data_ptr(), S, inter, s) + # Opt-in megakernel: gate_up GEMM + silu_mul + FP8 quant in one + # launch, but only in the small-M (launch-bound) regime where the + # ncu/bench sweep showed it beats the two-launch baseline. Above the + # winning S band the gate_up GEMM is compute-bound and the fused + # kernel's two-pass drain + lower compute density loses. + lo, hi = self._silu_fold_band or (0, 0) + if (self.use_silu_fold_prefill and self._silu_fold_tile is not None + and lo <= S <= hi + and (inter % 128) == 0 and (hidden % 128) == 0): + self._silu_fold_tile( + ap_mlp.data_ptr(), int(lw['gate_up_w']), + S, inter, hidden, + sc_mlp.data_ptr(), int(lw['gate_up_s']), + ap_dn.data_ptr(), sc_dn.data_ptr(), s) + else: + _, _, gate_up_out = self._prefill_fp8_scratch[(2 * inter, hidden)] + self._prefill_gemm( + ap_mlp, sc_mlp, int(lw['gate_up_w']), int(lw['gate_up_s']), + gate_up_out, int(lw['gate_up_N']), hidden, S) + fvk.silu_mul_merged_to_fp8_block128_bf16( + gate_up_out[:S].data_ptr(), + ap_dn.data_ptr(), sc_dn.data_ptr(), S, inter, s) else: self._prefill_gemm( ap_mlp, sc_mlp, int(lw['mlp_gate_w']), int(lw['mlp_gate_s']), diff --git a/scripts/test_geglu_silu_fold_correctness.py b/scripts/test_geglu_silu_fold_correctness.py new file mode 100644 index 00000000..7ad67156 --- /dev/null +++ b/scripts/test_geglu_silu_fold_correctness.py @@ -0,0 +1,105 @@ +#!/usr/bin/env python3 +"""Correctness + micro-bench for the SM89 GeGLU silu-fold megakernel. + +Verifies that bench_fp8_bs_geglu_silu_fold_* (gate+up GEMM + silu(gate)*up + +per-token block-128 FP8 quant, one launch) matches the baseline two-step path +(gate_up GEMM → bf16 [M,2*inter] → silu_mul_merged_to_fp8) within fp8 quant +tolerance, and is numerically closer to the fp32 reference. Then micro-benches +both to measure the win from eliminating the [M,2*inter] BF16 transient. + +Usage: + python scripts/test_geglu_silu_fold_correctness.py --M 512 +""" +from __future__ import annotations +import argparse, pathlib, sys, statistics +import torch +REPO = pathlib.Path(__file__).resolve().parents[1] +sys.path.insert(0, str(REPO)) +from flash_rt import flash_rt_qwen3_vl_kernels as vlk + +# 8B: inter=12288 hidden=4096 ; 2B: inter=9216 hidden=2048 +SHAPES = { + "8B": dict(N=12288, K=4096), + "2B": dict(N=6144, K=2048), +} +TILES = { + "32x128_s2": "bench_fp8_bs_geglu_silu_fold_sm89_32x128_w4_s2", + "16x128_s2": "bench_fp8_bs_geglu_silu_fold_sm89_16x128_w4_s2", + "64x128_s2": "bench_fp8_bs_geglu_silu_fold_sm89_64x128_w4_s2", + "128x128_s1": "bench_fp8_bs_geglu_silu_fold_sm89_128x128_w8_s1", + "32x128_s1": "bench_fp8_bs_geglu_silu_fold_sm89_32x128_w4_s1", + "16x128_s1": "bench_fp8_bs_geglu_silu_fold_sm89_16x128_w4_s1", + "ap_32x128_s1": "bench_fp8_bs_geglu_silu_fold_apersist_sm89_32x128_w4_s1", + "ap_16x128_s1": "bench_fp8_bs_geglu_silu_fold_apersist_sm89_16x128_w4_s1", + "ap_32x128_s2": "bench_fp8_bs_geglu_silu_fold_apersist_sm89_32x128_w4_s2", +} + +def ev_ms(fn, iters=50, warmup=20): + s=[torch.cuda.Event(enable_timing=True) for _ in range(iters)] + e=[torch.cuda.Event(enable_timing=True) for _ in range(iters)] + for _ in range(warmup): fn() + torch.cuda.synchronize() + for i in range(iters): s[i].record(); fn(); e[i].record() + torch.cuda.synchronize() + return [a.elapsed_time(b) for a,b in zip(s,e)] + +def cos(a,b): + return torch.nn.functional.cosine_similarity( + a.flatten().unsqueeze(0).float(), b.flatten().unsqueeze(0).float()).item() + +def main(): + p=argparse.ArgumentParser() + p.add_argument("--M", type=int, nargs="+", default=[128,512,1024]) + p.add_argument("--models", nargs="+", default=["2B","8B"]) + args=p.parse_args() + dev=torch.device("cuda:0") + torch.manual_seed(0) + for model in args.models: + N,K=SHAPES[model]["N"],SHAPES[model]["K"] + for M in args.M: + # A: per-token fp8 [M,K]; act_scale [M,K/128]; B: gate_up_w [2N,K] fp8; + # w_scale: gate_up_s [2N/128, K/128]. Use small scales to stay in fp8 range. + A=(torch.randn(M,K,device=dev)*0.3).to(torch.float8_e4m3fn) + B=(torch.randn(2*N,K,device=dev)*0.3).to(torch.float8_e4m3fn) + asc=(torch.rand(M,K//128,device=dev)*0.2+0.05).contiguous() + wsc=(torch.rand(2*N//128,K//128,device=dev)*0.2+0.05).contiguous() + s=torch.cuda.current_stream().cuda_stream + # ---- baseline: gate_up GEMM → bf16 [M,2N], then silu_mul_merged → fp8 [M,N] ---- + gu_bf=torch.empty(M,2*N,device=dev,dtype=torch.bfloat16) + vlk.bench_fp8_block128_gemm_bs_sm89_128x128x128_w8_s1( + int(A.data_ptr()),int(B.data_ptr()),int(gu_bf.data_ptr()),M,2*N,K, + int(asc.data_ptr()),int(wsc.data_ptr()),s) + ap_base=torch.empty(M,N,device=dev,dtype=torch.float8_e4m3fn) + sc_base=torch.empty(M,N//128,device=dev,dtype=torch.float32) + vlk.silu_mul_merged_to_fp8_block128_bf16( + int(gu_bf.data_ptr()),int(ap_base.data_ptr()),int(sc_base.data_ptr()), + M,N,s) + # ---- fused ---- + for tname,tfn_name in TILES.items(): + tfn=getattr(vlk,tfn_name) + ap_f=torch.empty(M,N,device=dev,dtype=torch.float8_e4m3fn) + sc_f=torch.empty(M,N//128,device=dev,dtype=torch.float32) + tfn(int(A.data_ptr()),int(B.data_ptr()),M,N,K, + int(asc.data_ptr()),int(wsc.data_ptr()), + int(ap_f.data_ptr()),int(sc_f.data_ptr()),s) + c=cos(ap_base.float(),ap_f.float()) + # bench: baseline two-step vs fused (layer-ish, reuse A/B) + def run_base(): + vlk.bench_fp8_block128_gemm_bs_sm89_128x128x128_w8_s1( + int(A.data_ptr()),int(B.data_ptr()),int(gu_bf.data_ptr()),M,2*N,K, + int(asc.data_ptr()),int(wsc.data_ptr()),s) + vlk.silu_mul_merged_to_fp8_block128_bf16( + int(gu_bf.data_ptr()),int(ap_base.data_ptr()),int(sc_base.data_ptr()), + M,N,s) + def run_fused(): + tfn(int(A.data_ptr()),int(B.data_ptr()),M,N,K, + int(asc.data_ptr()),int(wsc.data_ptr()), + int(ap_f.data_ptr()),int(sc_f.data_ptr()),s) + b=statistics.median(ev_ms(run_base)); f=statistics.median(ev_ms(run_fused)) + flag="OK" if c>0.9999 else ("close" if c>0.999 else "FAIL") + print(f"{model} M={M:>4} {tname:>10} base={b*1e3:>7.1f}us fused={f*1e3:>7.1f}us " + f"delta={(b-f)*1e3:>+7.1f}us cos={c:.6f} {flag}") + print() + +if __name__=="__main__": + main()