diff --git a/csrc/kernels/fp8/gemm.cuh b/csrc/kernels/fp8/gemm.cuh index ede2943..e905a32 100644 --- a/csrc/kernels/fp8/gemm.cuh +++ b/csrc/kernels/fp8/gemm.cuh @@ -21,6 +21,14 @@ namespace fp8 { constexpr int kMmaK = 32; constexpr int kWarps = 8; // 128x128 CTA = 8 warps +// log2 of a compile-time power of two (for tile_at's swizzle shift). +template +struct log2_const : log2_const<(N >> 1), Acc + 1> {}; +template +struct log2_const<1, Acc> { + static constexpr int value = Acc; +}; + // Map the FP8Format enum to the CUDA fp8 element type consumed by mma_sync. template struct fp8_input { @@ -118,22 +126,26 @@ __global__ void fp8_quantize_kernel(FP8Params p) { } // Swizzled address inside a flat [rows * K] staging tile: the 16-byte chunk -// index is XORed with row bits starting at bit 2. Unswizzled, a kK=32 row -// spans only 8 words, so a warp's fragment load (8 consecutive rows x 4B, -// e.g. a_row0+0..7) maps rows r and r+4 onto the same banks — a 2-way -// conflict on every LDS. XORing the chunk index with row bit 2 shifts rows -// 4..7 by one chunk so each warp's 32-word read hits all 32 banks exactly -// once. Chunks stay contiguous, so the cp.async 16B staging path is -// unaffected. Validated for kK=32 (2 chunks); larger power-of-two chunk -// counts compile but need their own bank analysis. +// index is XORed with a row-dependent slice so a warp's fragment load (8 +// consecutive rows x 16B) hits all 32 banks exactly once. With kChunks +// power-of-two chunks per row, the XOR source is the top log2(kChunks) bits +// of the row index within each group of 8: +// kChunks=2 -> row bits [3] (K=32: rows r and r+4 diverge) +// kChunks=4 -> row bits [2:1] (K=64: rows diverge every 2) +// kChunks=8 -> row bits [2:0] (K=128: every row) +// (row word-stride is K/4 words = 4*kChunks, so unswizzled rows r and +// r + 8/kChunks collide mod 32 banks; the XOR spreads the 8 rows of one +// ldmatrix matrix across the 8 distinct 4-bank groups.) Chunks stay +// contiguous, so the cp.async 16B staging path is unaffected. template __device__ __forceinline__ T8* tile_at(T8* tile, int row, int col) { constexpr int kChunks = K / 16; // 16B chunks per row static_assert(kChunks >= 1 && (kChunks & (kChunks - 1)) == 0, "swizzle needs a power-of-two 16B-chunk count"); + constexpr int kShift = 3 - log2_const::value; return tile + row * K - + ((((col >> 4) ^ ((row >> 2) & (kChunks - 1))) << 4) - + (col & 15)); + + ((((col >> 4) ^ ((row >> kShift) & (kChunks - 1))) << 4) + + (col & 15)); } // Stage-load one GEMM operand into the canonical flat [rows * K] shared tile @@ -142,74 +154,88 @@ __device__ __forceinline__ T8* tile_at(T8* tile, int row, int col) { // layout: RowMajor (stored [rows][contract]) copies 16-byte K-contiguous runs // with cp.async, while ColMajor (stored [contract][rows]) reads 16-byte runs // along the operand's contiguous non-contract dim and scatters them across -// the tile's rows. `block_row` is this block's origin in the operand's row -// dim; the caller restricts which threads invoke it (all threads for A, the -// first 128 for B). -template +// the tile's rows. RowsTile is the tile's row capacity (kBlockM / kBlockN) +// and kThreads the CTA size; the runtime `rows` bound may be smaller (tail +// predication). `block_row` is this block's origin in the operand's row dim. +template __device__ __forceinline__ void load_operand_tile( T8* tile, const T8* __restrict__ operand, int64_t rows, int64_t contract, int64_t ld, int tid, int64_t k_base, int64_t block_row) { + constexpr int kChunks = K / 16; + static_assert(RowsTile * kChunks % kThreads == 0, + "tile chunks must divide evenly across threads"); + constexpr int kCpt = RowsTile * kChunks / kThreads; // chunks per thread if constexpr (std::is_same_v) { - // Operand stored [contract][rows]: contiguous along the non-contract dim. - const int rg = tid >> 5; // Rows / 16 row-groups - const int kl = tid & 31; // K lanes - const int64_t k_idx = k_base + kl; - const int64_t r0 = block_row + rg * 16; - const auto* src = operand + k_idx * ld + r0; - const bool aligned = (reinterpret_cast(src) & 15) == 0; - if (k_idx < contract && r0 + 15 < rows && aligned) { - const uint4 v = *reinterpret_cast(src); - const auto* bytes = reinterpret_cast(&v); - // Scatter 16 bytes along the tile rows. The swizzle bit flips - // every 4 rows ((rg*16 + i) >> 2 & 1 == (i >> 2) & 1), and the - // physical column of row group g is kl ^ (16 * (g & 1)) — so the - // whole 16-byte scatter is one base pointer plus two alternating - // column offsets, no per-byte XOR in the address math. + // Operand stored [contract][rows]: contiguous along the non-contract + // dim. Each thread scatters one 16-byte run per K/32 pass; when the + // tile has more 16-row groups than warps (RowsTile > kThreads/2), + // each thread covers several groups. + constexpr int kWarpsTile = kThreads / 32; + constexpr int kGroups = RowsTile / 16; + static_assert(kGroups % kWarpsTile == 0, + "row groups must divide evenly across warps"); #pragma unroll - for (int g = 0; g < 4; ++g) { - T8* p = tile + (rg * 16 + 4 * g) * K - + (g & 1 ? (kl ^ 16) : kl); - p[0] = bytes[4 * g]; - p[K] = bytes[4 * g + 1]; - p[2 * K] = bytes[4 * g + 2]; - p[3 * K] = bytes[4 * g + 3]; - } - } else { - // Predicated fallback: same layout, byte-granular gather. - const int col = kl; + for (int g = 0; g < kGroups / kWarpsTile; ++g) { + const int rg = (tid >> 5) + g * kWarpsTile; + const int kl = tid & 31; // byte column within a 32B pass + const int64_t r0 = block_row + rg * 16; #pragma unroll - for (int g = 0; g < 4; ++g) { - const T8* src_g = operand + k_idx * ld + r0 + 4 * g; - T8* p = tile + (rg * 16 + 4 * g) * K - + (g & 1 ? (col ^ 16) : col); + for (int pass = 0; pass < K / 32; ++pass) { + const int col = kl + pass * 32; + const int64_t k_idx = k_base + col; + const auto* src = operand + k_idx * ld + r0; + if (k_idx < contract && r0 + 15 < rows && + (reinterpret_cast(src) & 15) == 0) { + const uint4 v = *reinterpret_cast(src); + const auto* bytes = reinterpret_cast(&v); + // Scatter 16 bytes along the tile rows through tile_at's + // swizzle. Rows sharing a physical chunk form groups of + // (8 / kChunks) consecutive rows (see tile_at), so each + // group is one tile_at address plus a K-byte row stride. + constexpr int kGrp = 8 / kChunks; #pragma unroll - for (int i = 0; i < 4; ++i) { - const int64_t r_idx = r0 + 4 * g + i; - p[i * K] = (r_idx < rows && k_idx < contract) - ? src_g[i] - : T8(0.0f); + for (int j = 0; j < 16 / kGrp; ++j) { + T8* p = tile_at(tile, + rg * 16 + j * kGrp, col); +#pragma unroll + for (int i = 0; i < kGrp; ++i) + p[i * K] = bytes[j * kGrp + i]; + } + } else { + // Predicated fallback: same layout, byte-granular gather. +#pragma unroll + for (int i = 0; i < 16; ++i) { + const int64_t r_idx = r0 + i; + *tile_at(tile, rg * 16 + i, col) = + (r_idx < rows && k_idx < contract) + ? operand[k_idx * ld + r_idx] + : T8(0.0f); + } } } } } else { // Operand stored [rows][contract]: contiguous along the contract dim. - const int r = tid >> 1; - const int c = (tid & 1) * 16; - const int64_t row = block_row + r; - // c is a multiple of 16, so the whole 16-byte run shares one chunk - // and dst[i] addressing below matches tile_at(tile, r, c + i). - T8* dst = tile_at(tile, r, c); - const auto* src = operand + row * ld + k_base + c; - const bool full = k_base + c + 15 < contract; - if (row < rows && full && - (reinterpret_cast(src) & 15) == 0) { - astrai::cp_async_16(dst, src, true); - } else { + // Linear chunk mapping: thread covers kCpt consecutive 16B chunks of + // one row (K=64: a contiguous 32B pair; K=32: a single chunk). + const int r = tid / (kChunks / kCpt); #pragma unroll - for (int i = 0; i < 16; ++i) - dst[i] = row < rows && k_base + c + i < contract ? src[i] - : T8(0.0f); + for (int j = 0; j < kCpt; ++j) { + const int c = ((tid % (kChunks / kCpt)) * kCpt + j) * 16; + const int64_t row = block_row + r; + const auto* src = operand + row * ld + k_base + c; + T8* dst = tile_at(tile, r, c); + if (row < rows && k_base + c + 15 < contract && + (reinterpret_cast(src) & 15) == 0) { + astrai::cp_async_16(dst, src, true); + } else { +#pragma unroll + for (int i = 0; i < 16; ++i) + dst[i] = row < rows && k_base + c + i < contract + ? src[i] + : T8(0.0f); + } } } } @@ -222,13 +248,13 @@ __device__ __forceinline__ void load_operand_tile( // --------------------------------------------------------------------------- // Swizzled 16B-chunk address (tile_at's layout) as a raw shared-memory -// pointer for ldmatrix. Requires kK == 32 (2 chunks/row swizzle). The chunk -// XOR itself lives only in tile_at; this wrapper just converts the element -// address it returns. +// pointer for ldmatrix. Valid for kK in {32, 64} (the swizzle itself lives +// only in tile_at; this wrapper just converts the element address). template __device__ __forceinline__ unsigned frag_addr(const T8* tile, int row, int chunk) { - static_assert(kK == 32, "fragment swizzle offsets assume kK == 32"); + static_assert(kK == 32 || kK == 64, + "fragment swizzle offsets assume kK in {32, 64}"); return __cvta_generic_to_shared(tile_at(tile, row, chunk << 4)); } @@ -241,13 +267,20 @@ __device__ __forceinline__ unsigned frag_addr(const T8* tile, int row, // change how the stage-load gathers the operand from global memory: // A ColMajor: tileA[m][p] = a[p*a_ld + m]; A RowMajor: a[m*a_ld + p] // B RowMajor: tileB[n][p] = b[p*b_ld + n]; B ColMajor: b[n*b_ld + p] +// BlockM x BlockN CTA as (BlockM/64) x (BlockN/32) warps of 64x32 warp tiles +// (mt x nt = 4x4 MMA each). The 64x128 variant runs 4 warps / 128 threads and +// exists for small-M calls: m <= 64 wastes half of every 128-row CTA, so the +// launcher dispatches to it there (see launch_fp8_gemm). template -__global__ void __launch_bounds__(kWarps * 32, 2) fp8_gemm_kernel(FP8Params p) { +__global__ void __launch_bounds__( + (Traits::kBlockM / 64) * (Traits::kBlockN / 32) * 32, 2) +fp8_gemm_kernel(FP8Params p) { using T8 = std::conditional_t; constexpr int kBlockM = Traits::kBlockM; constexpr int kBlockN = Traits::kBlockN; constexpr int kK = Traits::kK; constexpr int kStages = Traits::kStages; + constexpr int kCtaThreads = (kBlockM / 64) * (kBlockN / 32) * 32; static_assert(kStages >= 1 && kStages <= 8, "FP8 GEMM stages must be in the range [1, 8]"); // Tiles are flat [rows * kK] with a 16B-chunk XOR swizzle (tile_at): @@ -269,13 +302,36 @@ __global__ void __launch_bounds__(kWarps * 32, 2) fp8_gemm_kernel(FP8Params p) { const int lane = tid & 31; const int group = lane >> 2; const int thread_in_group = lane & 3; + // L2-friendly rasterization (CUTLASS-style grouped launch order): remap + // the linear block id so consecutive CTAs cover a group of kGroupM M-tiles + // before advancing along N. All CTAs of one group share the same B column + // stripe, so B tiles stay hot in L2 across the wave (the default + // N-fastest order makes each wave touch every B tile instead). + // Measured win for the A-crosswise layouts (10-21% at K>=2048) and loss + // for A-congruous (-17..20%, A's cp.async stream prefers the N-fastest + // order) — so the branch follows LayoutA. + constexpr int kGroupM = 8; + int block_m, block_n; + if constexpr (std::is_same_v) { + const int blocks_m = gridDim.y; + const int bid = blockIdx.y * gridDim.x + blockIdx.x; + const int group_first_m = (bid / (kGroupM * gridDim.x)) * kGroupM; + const int group_rows = + min(blocks_m - group_first_m, kGroupM); // M-tail group is short + block_m = group_first_m + bid % group_rows; + block_n = (bid % (kGroupM * gridDim.x)) / group_rows; + } else { + block_m = blockIdx.y; + block_n = blockIdx.x; + } // 128x128 CTA = 8 warps as 2x4 warp tiles of 64x32 (mt x nt = 4x4 MMA). constexpr int warps_n = kBlockN / 32; const int warp_m = warp / warps_n; const int warp_n = warp % warps_n; - const int64_t row_base = blockIdx.y * kBlockM + warp_m * 64 + group; + const int64_t row_base = + (int64_t)block_m * kBlockM + warp_m * 64 + group; const int64_t output_col = - blockIdx.x * kBlockN + warp_n * 32 + thread_in_group * 2; + (int64_t)block_n * kBlockN + warp_n * 32 + thread_in_group * 2; const int a_row0 = warp_m * 64; // + mt * 16 in the loop const int b_row0 = warp_n * 32; // + nt * 8 const float sa = *p.scale_a; @@ -285,15 +341,17 @@ __global__ void __launch_bounds__(kWarps * 32, 2) fp8_gemm_kernel(FP8Params p) { // Both operands are staged into the canonical [M][kK] / [N][kK] shared // tiles regardless of their global layout (see load_operand_tile), so the // MMA fragment reads below stay unchanged across the four layout - // combinations. Each 128x32 tile is 256 16B chunks: one per thread. - // A's tag already names the operand view ([M][K] = [rows][contract]); - // B's tag is relative to the canonical [K][N], so the stage-load sees its - // transpose (transpose_layout_t, see common.h). + // combinations. A's tag already names the operand view ([M][K] = + // [rows][contract]); B's tag is relative to the canonical [K][N], so the + // stage-load sees its transpose (transpose_layout_t, see common.h). auto load_tile = [&](int stage, int64_t k_base) { - load_operand_tile( - a_smem[stage], a, m, k, a_ld, tid, k_base, blockIdx.y * kBlockM); - load_operand_tile>( - b_smem[stage], b, n, k, b_ld, tid, k_base, blockIdx.x * kBlockN); + load_operand_tile( + a_smem[stage], a, m, k, a_ld, tid, k_base, + (int64_t)block_m * kBlockM); + load_operand_tile, kBlockN, + kCtaThreads>( + b_smem[stage], b, n, k, b_ld, tid, k_base, + (int64_t)block_n * kBlockN); }; const int64_t tile_count = (k + kK - 1) / kK; @@ -339,15 +397,29 @@ __global__ void __launch_bounds__(kWarps * 32, 2) fp8_gemm_kernel(FP8Params p) { // 4 ldmatrix.x2 (B) + 4 ldmatrix.x4 (A) feed 16 mma.sync per k_seg — // 0.5 load instructions per MMA, versus 4.5 scalar LDS per MMA in // the 128x64-tile version (the kernel was LSU-issue-bound there). + constexpr int kSegs = kK / kMmaK; + // B fragments double-buffered across k_segs: the next k_seg's B load + // is issued before the current k_seg's MMA sequence, so its LDS + // latency hides behind the A pipeline + tensor-pipe work (same trick + // as the A mt+1 prefetch below; costs kSegs x 8 registers). + unsigned b_frag[2][4][2]; #pragma unroll - for (int k_seg = 0; k_seg < kK / kMmaK; ++k_seg) { - unsigned b_frag[4][2]; + for (int nt = 0; nt < 4; ++nt) { + const int row = b_row0 + nt * 8 + r7; + astrai::ldmatrix_x2_lane(b_frag[0][nt], + frag_addr(b_smem[stage], row, rh8)); + } #pragma unroll - for (int nt = 0; nt < 4; ++nt) { - const int row = b_row0 + nt * 8 + r7; - astrai::ldmatrix_x2_lane(b_frag[nt], - frag_addr(b_smem[stage], row, - k_seg * 2 + rh8)); + for (int k_seg = 0; k_seg < kSegs; ++k_seg) { + const int bcur = k_seg & 1, bnext = bcur ^ 1; + if (k_seg + 1 < kSegs) { +#pragma unroll + for (int nt = 0; nt < 4; ++nt) { + const int row = b_row0 + nt * 8 + r7; + astrai::ldmatrix_x2_lane(b_frag[bnext][nt], + frag_addr(b_smem[stage], row, + (k_seg + 1) * 2 + rh8)); + } } // Software-pipelined A fragments: the ldmatrix.x4 for row mt+1 // is issued before the MMAs consuming row mt, so the LDS fixed @@ -367,8 +439,8 @@ __global__ void __launch_bounds__(kWarps * 32, 2) fp8_gemm_kernel(FP8Params p) { k_seg * 2 + rh16)); #pragma unroll for (int nt = 0; nt < 4; ++nt) - astrai::mma_sync(acc[nt][mt], a_frag[mt], b_frag[nt], - acc[nt][mt]); + astrai::mma_sync(acc[nt][mt], a_frag[mt], + b_frag[bcur][nt], acc[nt][mt]); } } // Barrier 2: every thread finished reading this stage's tiles before @@ -440,18 +512,27 @@ void launch_fp8_quantize(const FP8Params& p, cudaStream_t stream) { fp8_quantize_kernel<<>>(p); } -// Pre-quantized GEMM tile config: 128x128 CTA (8 warps x 64x32 warp tiles), -// K=32, 3-stage pipeline (24KB smem -> 2 CTAs/SM). The wide warp tile plus -// ldmatrix fragments lifts the LSU-issue bound of the old 128x64 config. -// Stages remains an explicit template override for tuning. LayoutA/LayoutB -// mirror the kernel template (defaults keep the NN layout: out = a @ b). +// Pre-quantized GEMM tile config: 128x128 CTA (8 warps x 64x32 warp tiles). +// kK selects the K tile (32 or 64; 64 halves the __syncthreads count per K +// and doubles the MMA work per stage, at 2x the smem per stage — measured +// 10-35% across shapes, so 64 is the default). Stages=2 with kK=64 keeps the +// pipeline at 32KB smem; deeper pipelines only win on K >= 4096 squares and +// lose elsewhere. LayoutA/LayoutB mirror the kernel template (defaults keep +// the NN layout: out = a @ b). m <= 64 dispatches to the 64x128 CTA — a +// 128-row CTA would waste half its MMA work on predicated-off rows. template + typename LayoutB = RowMajor, int kK = 64, int Stages = 2> void launch_fp8_gemm(const FP8Params& p, cudaStream_t stream) { - using Traits = Fp8GemmTraits; - dim3 grid((p.n + Traits::kBlockN - 1) / Traits::kBlockN, - (p.m + Traits::kBlockM - 1) / Traits::kBlockM); - fp8_gemm_kernel<<>>(p); + dim3 grid((p.n + 127) / 128, (p.m + 127) / 128); + if (p.m <= 64) { + using Traits = Fp8GemmTraits; + fp8_gemm_kernel + <<>>(p); + } else { + using Traits = Fp8GemmTraits; + fp8_gemm_kernel + <<>>(p); + } } } // namespace fp8 diff --git a/csrc/tests/fp8_mma_test.cu b/csrc/tests/fp8_mma_test.cu deleted file mode 100644 index 334c9c2..0000000 --- a/csrc/tests/fp8_mma_test.cu +++ /dev/null @@ -1,146 +0,0 @@ -/* -Single-kernel BF16 -> FP8 MMA -> BF16 demo for Ada (sm_89). - -nvcc -I csrc -arch=sm_89 -std=c++17 -O3 --use_fast_math \ - --ptxas-options=-O3,-v csrc/tests/fp8_mma_test.cu -o fp8_mma_test \ - && ./fp8_mma_test -*/ - -#include "test_utils.cuh" - -#include - -#include "../kernels/common/mma.cuh" - -#include -#include - -constexpr int M = 16; -constexpr int N = 8; -constexpr int K = 32; - -__device__ __forceinline__ unsigned pack_fp8x4(float x0, float x1, float x2, - float x3) { - __nv_fp8_e4m3 q0(x0); - __nv_fp8_e4m3 q1(x1); - __nv_fp8_e4m3 q2(x2); - __nv_fp8_e4m3 q3(x3); - return static_cast(q0.__x) | - (static_cast(q1.__x) << 8) | - (static_cast(q2.__x) << 16) | - (static_cast(q3.__x) << 24); -} - -__device__ __forceinline__ unsigned load_quantize_fp8x4( - const bf16* src, float scale_inv) { - return pack_fp8x4(__bfloat162float(src[0]) * scale_inv, - __bfloat162float(src[1]) * scale_inv, - __bfloat162float(src[2]) * scale_inv, - __bfloat162float(src[3]) * scale_inv); -} - -__global__ void fused_bf16_fp8_mma_kernel( - const bf16* __restrict__ a, const bf16* __restrict__ b, - bf16* __restrict__ out, float scale_a, float scale_b) { - const int lane = threadIdx.x; - const int group = lane >> 2; - const int thread_in_group = lane & 3; - const int k0 = thread_in_group * 4; - - // PTX m16n8k32 A fragment: two rows, two 16-column K partitions. - unsigned a_frag[4]; - a_frag[0] = load_quantize_fp8x4(&a[group * K + k0], 1.0f / scale_a); - a_frag[1] = load_quantize_fp8x4(&a[(group + 8) * K + k0], 1.0f / scale_a); - a_frag[2] = load_quantize_fp8x4(&a[group * K + k0 + 16], 1.0f / scale_a); - a_frag[3] = load_quantize_fp8x4(&a[(group + 8) * K + k0 + 16], - 1.0f / scale_a); - - // B is supplied as row-major [N,K], equivalent to the col-major [K,N] - // operand required by the MMA instruction. - unsigned b_frag[2]; - b_frag[0] = load_quantize_fp8x4(&b[group * K + k0], 1.0f / scale_b); - b_frag[1] = load_quantize_fp8x4(&b[group * K + k0 + 16], 1.0f / scale_b); - - float acc[4] = {0.0f, 0.0f, 0.0f, 0.0f}; - astrai::mma_sync<__nv_fp8_e4m3>(acc, a_frag, b_frag, acc); - - const int col = thread_in_group * 2; - const float output_scale = scale_a * scale_b; - *reinterpret_cast<__nv_bfloat162*>(&out[group * N + col]) = - __floats2bfloat162_rn(acc[0] * output_scale, - acc[1] * output_scale); - *reinterpret_cast<__nv_bfloat162*>(&out[(group + 8) * N + col]) = - __floats2bfloat162_rn(acc[2] * output_scale, - acc[3] * output_scale); -} - -static float quantize_e4m3(float value) { - return static_cast(__nv_fp8_e4m3(value)); -} - -int main() { - srand(0); - std::vector a(M * K), b(N * K), reference(M * N, 0.0f); - std::vector a_bf16(M * K), b_bf16(N * K), output(M * N); - for (float& value : a) value = randf() * 4.0f; - for (float& value : b) value = randf() * 4.0f; - for (int i = 0; i < M * K; ++i) { - a_bf16[i] = f2bf(a[i]); - a[i] = bf2f(a_bf16[i]); - } - for (int i = 0; i < N * K; ++i) { - b_bf16[i] = f2bf(b[i]); - b[i] = bf2f(b_bf16[i]); - } - - const float amax = *std::max_element( - a.begin(), a.end(), [](float x, float y) { return fabsf(x) < fabsf(y); }); - const float bmax = *std::max_element( - b.begin(), b.end(), [](float x, float y) { return fabsf(x) < fabsf(y); }); - const float scale_a = fabsf(amax) / 448.0f; - const float scale_b = fabsf(bmax) / 448.0f; - - for (int row = 0; row < M; ++row) { - for (int col = 0; col < N; ++col) { - float sum = 0.0f; - for (int k = 0; k < K; ++k) { - float qa = quantize_e4m3(a[row * K + k] / scale_a); - float qb = quantize_e4m3(b[col * K + k] / scale_b); - sum = fmaf(qa, qb, sum); - } - reference[row * N + col] = sum * scale_a * scale_b; - } - } - - bf16 *d_a, *d_b, *d_out; - CUDA_CHECK(cudaMalloc(&d_a, a_bf16.size() * sizeof(bf16))); - CUDA_CHECK(cudaMalloc(&d_b, b_bf16.size() * sizeof(bf16))); - CUDA_CHECK(cudaMalloc(&d_out, output.size() * sizeof(bf16))); - CUDA_CHECK(cudaMemcpy(d_a, a_bf16.data(), a_bf16.size() * sizeof(bf16), - cudaMemcpyHostToDevice)); - CUDA_CHECK(cudaMemcpy(d_b, b_bf16.data(), b_bf16.size() * sizeof(bf16), - cudaMemcpyHostToDevice)); - - fused_bf16_fp8_mma_kernel<<<1, 32>>>(d_a, d_b, d_out, scale_a, scale_b); - CUDA_CHECK(cudaDeviceSynchronize()); - CUDA_CHECK(cudaMemcpy(output.data(), d_out, output.size() * sizeof(bf16), - cudaMemcpyDeviceToHost)); - - float max_abs_error = 0.0f; - float max_rel_error = 0.0f; - for (int i = 0; i < M * N; ++i) { - float error = fabsf(bf2f(output[i]) - reference[i]); - max_abs_error = fmaxf(max_abs_error, error); - max_rel_error = fmaxf(max_rel_error, - error / fmaxf(fabsf(reference[i]), 1e-4f)); - } - const bool pass = max_abs_error < 0.05f; - print_test_header(); - print_test_row("M=16 N=8 K=32 fused BF16->E4M3 MMA", max_abs_error, - max_rel_error, pass); - - cudaFree(d_a); - cudaFree(d_b); - cudaFree(d_out); - return pass ? 0 : 1; -} diff --git a/csrc/tests/fp8_test.cu b/csrc/tests/fp8_test.cu new file mode 100644 index 0000000..bc5c608 --- /dev/null +++ b/csrc/tests/fp8_test.cu @@ -0,0 +1,312 @@ +/* +FP8 family tests: single-warp MMA demo + full GEMM correctness. + +Part 1 exercises one bf16 -> fp8 -> mma.sync m16n8k32 instruction pair +(sanity for astrai::mma_sync + the fragment layout contract). +Part 2 checks launch_fp8_gemm across all four operand layouts, both K +tiles, and ragged shapes against an fp32 CPU reference. + +nvcc -I csrc -arch=sm_89 -std=c++17 -O3 csrc/tests/fp8_test.cu -o /tmp/fp8_test \ + && /tmp/fp8_test +*/ + +#include "test_utils.cuh" + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "../kernels/common/mma.cuh" +#include "../kernels/fp8/gemm.cuh" + +using namespace astrai::fp8; + +// --------------------------------------------------------------------------- +// Part 1: single-kernel BF16 -> FP8 MMA -> BF16 demo (m16n8k32) +// --------------------------------------------------------------------------- + +namespace { + +constexpr int kMmaM = 16; +constexpr int kMmaN = 8; +constexpr int kMmaK = 32; + +__device__ __forceinline__ unsigned pack_fp8x4(float x0, float x1, float x2, + float x3) { + __nv_fp8_e4m3 q0(x0); + __nv_fp8_e4m3 q1(x1); + __nv_fp8_e4m3 q2(x2); + __nv_fp8_e4m3 q3(x3); + return static_cast(q0.__x) | + (static_cast(q1.__x) << 8) | + (static_cast(q2.__x) << 16) | + (static_cast(q3.__x) << 24); +} + +__device__ __forceinline__ unsigned load_quantize_fp8x4( + const bf16* src, float scale_inv) { + return pack_fp8x4(__bfloat162float(src[0]) * scale_inv, + __bfloat162float(src[1]) * scale_inv, + __bfloat162float(src[2]) * scale_inv, + __bfloat162float(src[3]) * scale_inv); +} + +__global__ void fused_bf16_fp8_mma_kernel( + const bf16* __restrict__ a, const bf16* __restrict__ b, + bf16* __restrict__ out, float scale_a, float scale_b) { + const int lane = threadIdx.x; + const int group = lane >> 2; + const int thread_in_group = lane & 3; + const int k0 = thread_in_group * 4; + + // PTX m16n8k32 A fragment: two rows, two 16-column K partitions. + unsigned a_frag[4]; + a_frag[0] = load_quantize_fp8x4(&a[group * kMmaK + k0], 1.0f / scale_a); + a_frag[1] = + load_quantize_fp8x4(&a[(group + 8) * kMmaK + k0], 1.0f / scale_a); + a_frag[2] = + load_quantize_fp8x4(&a[group * kMmaK + k0 + 16], 1.0f / scale_a); + a_frag[3] = load_quantize_fp8x4(&a[(group + 8) * kMmaK + k0 + 16], + 1.0f / scale_a); + + // B is supplied as row-major [N,K], equivalent to the col-major [K,N] + // operand required by the MMA instruction. + unsigned b_frag[2]; + b_frag[0] = load_quantize_fp8x4(&b[group * kMmaK + k0], 1.0f / scale_b); + b_frag[1] = + load_quantize_fp8x4(&b[group * kMmaK + k0 + 16], 1.0f / scale_b); + + float acc[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + astrai::mma_sync<__nv_fp8_e4m3>(acc, a_frag, b_frag, acc); + + const int col = thread_in_group * 2; + const float output_scale = scale_a * scale_b; + *reinterpret_cast<__nv_bfloat162*>(&out[group * kMmaN + col]) = + __floats2bfloat162_rn(acc[0] * output_scale, acc[1] * output_scale); + *reinterpret_cast<__nv_bfloat162*>(&out[(group + 8) * kMmaN + col]) = + __floats2bfloat162_rn(acc[2] * output_scale, acc[3] * output_scale); +} + +static float quantize_e4m3(float value) { + return static_cast(__nv_fp8_e4m3(value)); +} + +static bool test_single_mma() { + srand(0); + std::vector a(kMmaM * kMmaK), b(kMmaN * kMmaK), + reference(kMmaM * kMmaN, 0.0f); + std::vector a_bf16(kMmaM * kMmaK), b_bf16(kMmaN * kMmaK), + output(kMmaM * kMmaN); + for (float& value : a) value = randf() * 4.0f; + for (float& value : b) value = randf() * 4.0f; + for (int i = 0; i < kMmaM * kMmaK; ++i) { + a_bf16[i] = f2bf(a[i]); + a[i] = bf2f(a_bf16[i]); + } + for (int i = 0; i < kMmaN * kMmaK; ++i) { + b_bf16[i] = f2bf(b[i]); + b[i] = bf2f(b_bf16[i]); + } + + const float amax = *std::max_element( + a.begin(), a.end(), + [](float x, float y) { return fabsf(x) < fabsf(y); }); + const float bmax = *std::max_element( + b.begin(), b.end(), + [](float x, float y) { return fabsf(x) < fabsf(y); }); + const float scale_a = fabsf(amax) / 448.0f; + const float scale_b = fabsf(bmax) / 448.0f; + + for (int row = 0; row < kMmaM; ++row) { + for (int col = 0; col < kMmaN; ++col) { + float sum = 0.0f; + for (int k = 0; k < kMmaK; ++k) { + float qa = quantize_e4m3(a[row * kMmaK + k] / scale_a); + float qb = quantize_e4m3(b[col * kMmaK + k] / scale_b); + sum = fmaf(qa, qb, sum); + } + reference[row * kMmaN + col] = sum * scale_a * scale_b; + } + } + + bf16 *d_a, *d_b, *d_out; + CUDA_CHECK(cudaMalloc(&d_a, a_bf16.size() * sizeof(bf16))); + CUDA_CHECK(cudaMalloc(&d_b, b_bf16.size() * sizeof(bf16))); + CUDA_CHECK(cudaMalloc(&d_out, output.size() * sizeof(bf16))); + CUDA_CHECK(cudaMemcpy(d_a, a_bf16.data(), a_bf16.size() * sizeof(bf16), + cudaMemcpyHostToDevice)); + CUDA_CHECK(cudaMemcpy(d_b, b_bf16.data(), b_bf16.size() * sizeof(bf16), + cudaMemcpyHostToDevice)); + + fused_bf16_fp8_mma_kernel<<<1, 32>>>(d_a, d_b, d_out, scale_a, scale_b); + CUDA_CHECK(cudaDeviceSynchronize()); + CUDA_CHECK(cudaMemcpy(output.data(), d_out, output.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + + float max_abs_error = 0.0f; + float max_rel_error = 0.0f; + for (int i = 0; i < kMmaM * kMmaN; ++i) { + float error = fabsf(bf2f(output[i]) - reference[i]); + max_abs_error = fmaxf(max_abs_error, error); + max_rel_error = fmaxf( + max_rel_error, error / fmaxf(fabsf(reference[i]), 1e-4f)); + } + const bool pass = max_abs_error < 0.05f; + print_test_row("M=16 N=8 K=32 fused BF16->E4M3 MMA", max_abs_error, + max_rel_error, pass); + + cudaFree(d_a); + cudaFree(d_b); + cudaFree(d_out); + return pass; +} + +// --------------------------------------------------------------------------- +// Part 2: GEMM correctness — layouts x K-tiles vs fp32 CPU reference +// --------------------------------------------------------------------------- + +template +static bool run_gemm_case(const float* ha, const float* hb, int m, int n, + int k, int a_ld, int b_ld) { + __nv_fp8_e4m3 *da, *db; + __nv_bfloat16* dout; + float *dsa, *dsb; + cudaMalloc(&da, (size_t)m * k); + cudaMalloc(&db, (size_t)n * k); + cudaMalloc(&dout, (size_t)m * n * 2); + cudaMalloc(&dsa, 4); + cudaMalloc(&dsb, 4); + float one = 1.0f; + cudaMemcpy(dsa, &one, 4, cudaMemcpyHostToDevice); + cudaMemcpy(dsb, &one, 4, cudaMemcpyHostToDevice); + // quantize inputs to e4m3 on host and upload byte-by-byte + std::vector qa(m * k), qb(n * k); + for (int i = 0; i < m * k; ++i) { + __nv_fp8_e4m3 q(ha[i]); + qa[i] = *(unsigned char*)&q; + } + for (int i = 0; i < n * k; ++i) { + __nv_fp8_e4m3 q(hb[i]); + qb[i] = *(unsigned char*)&q; + } + cudaMemcpy(da, qa.data(), qa.size(), cudaMemcpyHostToDevice); + cudaMemcpy(db, qb.data(), qb.size(), cudaMemcpyHostToDevice); + + FP8Params p = {}; + p.a_ptr = da; + p.b_ptr = db; + p.out_ptr = dout; + p.scale_a = dsa; + p.scale_b = dsb; + p.m = m; + p.n = n; + p.k = k; + p.a_ld = a_ld; + p.b_ld = b_ld; + launch_fp8_gemm(p, 0); + cudaError_t e = cudaDeviceSynchronize(); + if (e != cudaSuccess) { + printf(" CUDA err: %s\n", cudaGetErrorString(e)); + return false; + } + std::vector hb16(m * n); + cudaMemcpy(hb16.data(), dout, (size_t)m * n * 2, cudaMemcpyDeviceToHost); + const float tol = 0.06f; + double max_rel = 0; + bool ok = true; + for (int i = 0; i < m && ok; ++i) { + for (int j = 0; j < n && ok; ++j) { + float ref = 0; + for (int kk = 0; kk < k; ++kk) { + // A reference reads the actual uploaded buffer: LA ColMajor + // means the buffer is [K][M] (ha_t), else [M][K]. + float av = std::is_same_v + ? (float)__nv_fp8_e4m3(ha[kk * m + i]) + : (float)__nv_fp8_e4m3(ha[i * k + kk]); + float bv; + if (std::is_same_v) + bv = (float)__nv_fp8_e4m3(hb[j * k + kk]); + else + bv = (float)__nv_fp8_e4m3(hb[kk * n + j]); + ref += av * bv; + } + float got = + __bfloat162float(__ushort_as_bfloat16(hb16[i * n + j])); + float err = fabsf(got - ref); + float rel = err / fmaxf(fabsf(ref), 0.5f); + if (rel > max_rel) max_rel = rel; + if (err > tol * fmaxf(fabsf(ref), 1.0f)) ok = false; + } + } + printf(" max_rel=%.4f %s\n", max_rel, ok ? "PASS" : "FAIL"); + cudaFree(da); + cudaFree(db); + cudaFree(dout); + cudaFree(dsa); + cudaFree(dsb); + return ok; +} + +static bool test_gemm() { + struct { + int m, n, k; + } cfgs[] = { + {128, 128, 128}, {256, 128, 256}, {128, 256, 64}, + {100, 130, 96}, {64, 64, 160}, {300, 200, 320}, + }; + bool all = true; + for (auto& c : cfgs) { + float* ha = new float[c.m * c.k]; + float* hb_rowmajor = new float[c.k * c.n]; // [K][N] for B RowMajor + float* hb_colmajor = new float[c.n * c.k]; // [N][K] for B ColMajor + for (int i = 0; i < c.m * c.k; ++i) ha[i] = randf(); + for (int i = 0; i < c.k * c.n; ++i) hb_rowmajor[i] = randf(); + for (int i = 0; i < c.k * c.n; ++i) + hb_colmajor[i / c.k * c.k + i % c.k] = hb_rowmajor[i]; + float* ha_t = new float[c.k * c.m]; // [K][M] for A ColMajor + for (int i = 0; i < c.m; ++i) + for (int p = 0; p < c.k; ++p) ha_t[p * c.m + i] = ha[i * c.k + p]; + printf("%dx%dx%d:\n", c.m, c.n, c.k); + printf(" NT K32:"); + all &= run_gemm_case(ha, hb_colmajor, c.m, + c.n, c.k, c.k, c.k); + printf(" NT K64:"); + all &= run_gemm_case(ha, hb_colmajor, c.m, + c.n, c.k, c.k, c.k); + printf(" NN K32:"); + all &= run_gemm_case(ha, hb_rowmajor, c.m, + c.n, c.k, c.k, c.n); + printf(" NN K64:"); + all &= run_gemm_case(ha, hb_rowmajor, c.m, + c.n, c.k, c.k, c.n); + printf(" TN K32:"); + all &= run_gemm_case(ha_t, hb_colmajor, c.m, + c.n, c.k, c.m, c.k); + printf(" TN K64:"); + all &= run_gemm_case(ha_t, hb_colmajor, c.m, + c.n, c.k, c.m, c.k); + printf(" TT K64:"); + all &= run_gemm_case(ha_t, hb_rowmajor, c.m, + c.n, c.k, c.m, c.n); + delete[] ha; + delete[] hb_rowmajor; + delete[] hb_colmajor; + delete[] ha_t; + } + return all; +} + +} // namespace + +int main() { + print_test_header(); + bool ok = test_single_mma(); + ok &= test_gemm(); + printf(ok ? "All PASS\n" : "FAILURES\n"); + return ok ? 0 : 1; +}