perf: speed up fp8 gemm tiles and scheduling

- K tile 32->64 (new default): fewer barriers, more MMA per stage; generalize tile_at swizzle and load_operand_tile accordingly
- 64x128 small-M CTA for m<=64 (2x at 64x4096x4096)
- L2 rasterization for crosswise-A layouts (+6..21%)
- micro-bench: NT 4096^3 +35%; linear fwd 1.24-1.76x, bwd 1.71-2.27x vs bf16
- add csrc/tests/fp8_test.cu (single MMA demo + GEMM layouts x K-tiles vs CPU reference)

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 <int N, int Acc = 0>
+struct log2_const : log2_const<(N >> 1), Acc + 1> {};
+template <int Acc>
+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 <FP8Format Fmt>
 struct fp8_input {
@@ -118,21 +126,25 @@ __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
+// index is XORed with a row-dependent slice so a warp's fragment load (8
-// spans only 8 words, so a warp's fragment load (8 consecutive rows x 4B,
+// consecutive rows x 16B) hits all 32 banks exactly once. With kChunks
-// e.g. a_row0+0..7) maps rows r and r+4 onto the same banks — a 2-way
+// power-of-two chunks per row, the XOR source is the top log2(kChunks) bits
-// conflict on every LDS. XORing the chunk index with row bit 2 shifts rows
+// of the row index within each group of 8:
-// 4..7 by one chunk so each warp's 32-word read hits all 32 banks exactly
+//   kChunks=2 -> row bits [3]   (K=32:  rows r and r+4 diverge)
-// once. Chunks stay contiguous, so the cp.async 16B staging path is
+//   kChunks=4 -> row bits [2:1] (K=64:  rows diverge every 2)
-// unaffected. Validated for kK=32 (2 chunks); larger power-of-two chunk
+//   kChunks=8 -> row bits [2:0] (K=128: every row)
-// counts compile but need their own bank analysis.
+// (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 <int K, typename T8>
 __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<kChunks>::value;
     return tile + row * K
-           + ((((col >> 4) ^ ((row >> 2) & (kChunks - 1))) << 4)
+           + ((((col >> 4) ^ ((row >> kShift) & (kChunks - 1))) << 4)
               + (col & 15));
 }
 
@@ -142,76 +154,90 @@ __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
+// the tile's rows. RowsTile is the tile's row capacity (kBlockM / kBlockN)
-// dim; the caller restricts which threads invoke it (all threads for A, the
+// and kThreads the CTA size; the runtime `rows` bound may be smaller (tail
-// first 128 for B).
+// predication). `block_row` is this block's origin in the operand's row dim.
-template <typename T8, int K, typename Layout>
+template <typename T8, int K, typename Layout, int RowsTile, int kThreads>
 __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<Layout, ColMajor>) {
-        // Operand stored [contract][rows]: contiguous along the non-contract dim.
+        // Operand stored [contract][rows]: contiguous along the non-contract
-        const int rg = tid >> 5;  // Rows / 16 row-groups
+        // dim. Each thread scatters one 16-byte run per K/32 pass; when the
-        const int kl = tid & 31;  // K lanes
+        // tile has more 16-row groups than warps (RowsTile > kThreads/2),
-        const int64_t k_idx = k_base + kl;
+        // 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 < 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 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;
-        const bool aligned = (reinterpret_cast<uintptr_t>(src) & 15) == 0;
+                if (k_idx < contract && r0 + 15 < rows &&
-        if (k_idx < contract && r0 + 15 < rows && aligned) {
+                    (reinterpret_cast<uintptr_t>(src) & 15) == 0) {
                     const uint4 v = *reinterpret_cast<const uint4*>(src);
                     const auto* bytes = reinterpret_cast<const T8*>(&v);
-            // Scatter 16 bytes along the tile rows. The swizzle bit flips
+                    // Scatter 16 bytes along the tile rows through tile_at's
-            // every 4 rows ((rg*16 + i) >> 2 & 1 == (i >> 2) & 1), and the
+                    // swizzle. Rows sharing a physical chunk form groups of
-            // physical column of row group g is kl ^ (16 * (g & 1)) — so the
+                    // (8 / kChunks) consecutive rows (see tile_at), so each
-            // whole 16-byte scatter is one base pointer plus two alternating
+                    // group is one tile_at address plus a K-byte row stride.
-            // column offsets, no per-byte XOR in the address math.
+                    constexpr int kGrp = 8 / kChunks;
 #pragma unroll
-            for (int g = 0; g < 4; ++g) {
+                    for (int j = 0; j < 16 / kGrp; ++j) {
-                T8* p = tile + (rg * 16 + 4 * g) * K
+                        T8* p = tile_at<K>(tile,
-                        + (g & 1 ? (kl ^ 16) : kl);
+                                           rg * 16 + j * kGrp, col);
-                p[0] = bytes[4 * g];
+#pragma unroll
-                p[K] = bytes[4 * g + 1];
+                        for (int i = 0; i < kGrp; ++i)
-                p[2 * K] = bytes[4 * g + 2];
+                            p[i * K] = bytes[j * kGrp + i];
-                p[3 * K] = bytes[4 * g + 3];
                     }
                 } else {
                     // Predicated fallback: same layout, byte-granular gather.
-            const int col = kl;
 #pragma unroll
-            for (int g = 0; g < 4; ++g) {
+                    for (int i = 0; i < 16; ++i) {
-                const T8* src_g = operand + k_idx * ld + r0 + 4 * g;
+                        const int64_t r_idx = r0 + i;
-                T8* p = tile + (rg * 16 + 4 * g) * K
+                        *tile_at<K>(tile, rg * 16 + i, col) =
-                        + (g & 1 ? (col ^ 16) : col);
+                            (r_idx < rows && k_idx < contract)
-#pragma unroll
+                                ? operand[k_idx * ld + r_idx]
-                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);
                     }
                 }
             }
+        }
     } else {
         // Operand stored [rows][contract]: contiguous along the contract dim.
-        const int r = tid >> 1;
+        // Linear chunk mapping: thread covers kCpt consecutive 16B chunks of
-        const int c = (tid & 1) * 16;
+        // one row (K=64: a contiguous 32B pair; K=32: a single chunk).
+        const int r = tid / (kChunks / kCpt);
+#pragma unroll
+        for (int j = 0; j < kCpt; ++j) {
+            const int c = ((tid % (kChunks / kCpt)) * kCpt + j) * 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<K>(tile, r, c + i).
-        T8* dst = tile_at<K>(tile, r, c);
             const auto* src = operand + row * ld + k_base + c;
-        const bool full = k_base + c + 15 < contract;
+            T8* dst = tile_at<K>(tile, r, c);
-        if (row < rows && full &&
+            if (row < rows && k_base + c + 15 < contract &&
                 (reinterpret_cast<uintptr_t>(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]
+                    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
+// pointer for ldmatrix. Valid for kK in {32, 64} (the swizzle itself lives
-// XOR itself lives only in tile_at; this wrapper just converts the element
+// only in tile_at; this wrapper just converts the element address).
-// address it returns.
 template <typename T8, int kK>
 __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<kK>(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 <typename Traits, bool OutFp8 = false, typename LayoutA = RowMajor, typename LayoutB = RowMajor>
-__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<Traits::kIsE5M2, __nv_fp8_e5m2, __nv_fp8_e4m3>;
     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<LayoutA, ColMajor>) {
+        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.
+    // combinations. A's tag already names the operand view ([M][K] =
-    // A's tag already names the operand view ([M][K] = [rows][contract]);
+    // [rows][contract]); B's tag is relative to the canonical [K][N], so the
-    // B's tag is relative to the canonical [K][N], so the stage-load sees its
+    // stage-load sees its transpose (transpose_layout_t, see common.h).
-    // transpose (transpose_layout_t, see common.h).
     auto load_tile = [&](int stage, int64_t k_base) {
-        load_operand_tile<T8, kK, LayoutA>(
+        load_operand_tile<T8, kK, LayoutA, kBlockM, kCtaThreads>(
-            a_smem[stage], a, m, k, a_ld, tid, k_base, blockIdx.y * kBlockM);
+            a_smem[stage], a, m, k, a_ld, tid, k_base,
-        load_operand_tile<T8, kK, transpose_layout_t<LayoutB>>(
+            (int64_t)block_m * kBlockM);
-            b_smem[stage], b, n, k, b_ld, tid, k_base, blockIdx.x * kBlockN);
+        load_operand_tile<T8, kK, transpose_layout_t<LayoutB>, 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).
-#pragma unroll
+        constexpr int kSegs = kK / kMmaK;
-        for (int k_seg = 0; k_seg < kK / kMmaK; ++k_seg) {
+        // B fragments double-buffered across k_segs: the next k_seg's B load
-            unsigned b_frag[4][2];
+        // 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 nt = 0; nt < 4; ++nt) {
             const int row = b_row0 + nt * 8 + r7;
-                astrai::ldmatrix_x2_lane(b_frag[nt],
+            astrai::ldmatrix_x2_lane(b_frag[0][nt],
+                    frag_addr<T8, kK>(b_smem[stage], row, rh8));
+        }
+#pragma unroll
+        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<T8, kK>(b_smem[stage], row,
-                                          k_seg * 2 + rh8));
+                                              (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<T8>(acc[nt][mt], a_frag[mt], b_frag[nt],
+                    astrai::mma_sync<T8>(acc[nt][mt], a_frag[mt],
-                                         acc[nt][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<Fmt><<<blocks, kThreads, 0, stream>>>(p);
 }
 
-// Pre-quantized GEMM tile config: 128x128 CTA (8 warps x 64x32 warp tiles),
+// 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
+// kK selects the K tile (32 or 64; 64 halves the __syncthreads count per K
-// ldmatrix fragments lifts the LSU-issue bound of the old 128x64 config.
+// and doubles the MMA work per stage, at 2x the smem per stage — measured
-// Stages remains an explicit template override for tuning. LayoutA/LayoutB
+// 10-35% across shapes, so 64 is the default). Stages=2 with kK=64 keeps the
-// mirror the kernel template (defaults keep the NN layout: out = a @ b).
+// 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 <FP8Format Fmt, bool OutFp8 = false, typename LayoutA = RowMajor,
-          typename LayoutB = RowMajor, int Stages = 3>
+          typename LayoutB = RowMajor, int kK = 64, int Stages = 2>
 void launch_fp8_gemm(const FP8Params& p, cudaStream_t stream) {
-    using Traits = Fp8GemmTraits<Fmt, 128, 128, 32, Stages>;
+    dim3 grid((p.n + 127) / 128, (p.m + 127) / 128);
-    dim3 grid((p.n + Traits::kBlockN - 1) / Traits::kBlockN,
+    if (p.m <= 64) {
-              (p.m + Traits::kBlockM - 1) / Traits::kBlockM);
+        using Traits = Fp8GemmTraits<Fmt, 64, 128, kK, Stages>;
-    fp8_gemm_kernel<Traits, OutFp8, LayoutA, LayoutB><<<grid, kWarps * 32, 0, stream>>>(p);
+        fp8_gemm_kernel<Traits, OutFp8, LayoutA, LayoutB>
+            <<<grid, (64 / 64) * (128 / 32) * 32, 0, stream>>>(p);
+    } else {
+        using Traits = Fp8GemmTraits<Fmt, 128, 128, kK, Stages>;
+        fp8_gemm_kernel<Traits, OutFp8, LayoutA, LayoutB>
+            <<<grid, (128 / 64) * (128 / 32) * 32, 0, stream>>>(p);
+    }
 }
 
 }  // namespace fp8

csrc/tests/fp8_mma_test.cu

@@ -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 <cuda_fp8.h>
-
-#include "../kernels/common/mma.cuh"
-
-#include <algorithm>
-#include <vector>
-
-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<unsigned>(q0.__x) |
-           (static_cast<unsigned>(q1.__x) << 8) |
-           (static_cast<unsigned>(q2.__x) << 16) |
-           (static_cast<unsigned>(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<float>(__nv_fp8_e4m3(value));
-}
-
-int main() {
-    srand(0);
-    std::vector<float> a(M * K), b(N * K), reference(M * N, 0.0f);
-    std::vector<bf16> 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;
-}

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 <cuda_fp8.h>
+#include <algorithm>
+#include <cmath>
+#include <cstdio>
+#include <cstdlib>
+#include <cuda_runtime.h>
+#include <type_traits>
+#include <vector>
+
+#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<unsigned>(q0.__x) |
+           (static_cast<unsigned>(q1.__x) << 8) |
+           (static_cast<unsigned>(q2.__x) << 16) |
+           (static_cast<unsigned>(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<float>(__nv_fp8_e4m3(value));
+}
+
+static bool test_single_mma() {
+    srand(0);
+    std::vector<float> a(kMmaM * kMmaK), b(kMmaN * kMmaK),
+        reference(kMmaM * kMmaN, 0.0f);
+    std::vector<bf16> 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 <typename LA, typename LB, int kK, int Stages>
+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<unsigned char> 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<FP8Format::E4M3, false, LA, LB, kK, Stages>(p, 0);
+    cudaError_t e = cudaDeviceSynchronize();
+    if (e != cudaSuccess) {
+        printf("  CUDA err: %s\n", cudaGetErrorString(e));
+        return false;
+    }
+    std::vector<unsigned short> 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<LA, ColMajor>
+                               ? (float)__nv_fp8_e4m3(ha[kk * m + i])
+                               : (float)__nv_fp8_e4m3(ha[i * k + kk]);
+                float bv;
+                if (std::is_same_v<LB, ColMajor>)
+                    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<RowMajor, ColMajor, 32, 3>(ha, hb_colmajor, c.m,
+                                                        c.n, c.k, c.k, c.k);
+        printf(" NT K64:");
+        all &= run_gemm_case<RowMajor, ColMajor, 64, 2>(ha, hb_colmajor, c.m,
+                                                        c.n, c.k, c.k, c.k);
+        printf(" NN K32:");
+        all &= run_gemm_case<RowMajor, RowMajor, 32, 3>(ha, hb_rowmajor, c.m,
+                                                        c.n, c.k, c.k, c.n);
+        printf(" NN K64:");
+        all &= run_gemm_case<RowMajor, RowMajor, 64, 2>(ha, hb_rowmajor, c.m,
+                                                        c.n, c.k, c.k, c.n);
+        printf(" TN K32:");
+        all &= run_gemm_case<ColMajor, ColMajor, 32, 3>(ha_t, hb_colmajor, c.m,
+                                                        c.n, c.k, c.m, c.k);
+        printf(" TN K64:");
+        all &= run_gemm_case<ColMajor, ColMajor, 64, 2>(ha_t, hb_colmajor, c.m,
+                                                        c.n, c.k, c.m, c.k);
+        printf(" TT K64:");
+        all &= run_gemm_case<ColMajor, RowMajor, 64, 2>(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;
+}