refactor: dedupe fp8 kernel helpers and trim comments

- merge the quantize launchers into one Tiled template; extract shared cvt_fp8/publish_amax helpers and replace the dtype x format ladder with two-level template dispatch - fold the gemm interior/generic operand loads into one kInterior template and the fast/generic async loads into load_async<kFast>; Policy carries the smem budget - compress kernel comments to the load-bearing invariants, dropping measured-number essays; Policy signature and kernel code unchanged Benchmark: NVIDIA L20, 1.2B model train step fwd+bwd+CE - M=8192: fp8 532.2 -> 530.4 ms (1.26x, noise); tests/extension 65 passed, quantize layouts byte-exact, NT routing diff 0.0
2026-08-28 16:45:21 +08:00
parent 8a353117ea
commit bf239d194c
4 changed files with 376 additions and 651 deletions
@@ -11,30 +11,23 @@
 namespace astrai {
 namespace fp8 {
-// Compile-time FP8 format: E4M3 (forward / high precision, max 448) or
+// Compile-time FP8 format: E4M3 (forward, max 448) or E5M2 (gradients,
-// E5M2 (gradient / large dynamic range, max 57344).
+// max 57344).
 enum class FP8Format : int {
    E4M3 = 0,
    E5M2 = 1,
 };
-// Operand memory layouts as types (CUTLASS-style tags). The tag names the
+// Operand storage tags (CUTLASS-style) relative to the canonical matrices
-// storage order of the raw buffer relative to the operand's canonical GEMM
+// A [M][K] / B [K][N]: A RowMajor = [M][K] (default), A ColMajor = [K][M],
-// matrix — A is [M][K], B is [K][N]:
+// B RowMajor = [K][N], B ColMajor = [N][K] (the nn.Linear weight). Selection
-//   A RowMajor = [M][K] storage (K-contiguous rows; the default)
+// is by type at compile time (see gemm.cuh's stage loads).
 //   A ColMajor = [K][M] storage (M-contiguous; A^T)
 //   B RowMajor = [K][N] storage (N-contiguous; the plain a @ b operand)
 //   B ColMajor = [N][K] storage (K-contiguous; the nn.Linear weight layout)
 // Empty tags: selection happens by type at compile time (see load_operand_tile).
 struct RowMajor {};
 struct ColMajor {};
-// Compile-time tile configuration, mirroring KernelTraits<HEAD_DIM, BC,
+// Compile-time tile configuration, mirroring KernelTraits in the attention
-// WARPS, STAGES> in the attention kernels. `Fmt` selects the FP8 conversion
+// kernels: CTA tile, warp tile (WarpM x WarpN — e.g. 64x32 on the 128x128
-// and the MMA PTX mnemonic; the remaining parameters shape the CTA tile, the
+// CTA, 32x32 on the 64x64 small CTA) and cp.async pipeline depth.
 // warp tile (WarpM x WarpN — e.g. 64x32 on the 128x128 CTA, or 32x32 on the
 // cuBLAS-style 64x64 small CTA that lifts small-shape occupancy) and the
 // cp.async pipeline depth.
 template <FP8Format Fmt, int BlockM, int BlockN, int K, int Stages,
          int WarpM = 64, int WarpN = 32>
 struct Fp8GemmTraits {
@@ -50,10 +43,8 @@ struct Fp8GemmTraits {
        kIsE5M2 ? __NV_E5M2 : __NV_E4M3;
    static constexpr float kFp8Max = kIsE5M2 ? 57344.0f : 448.0f;
-    // Derived launch geometry: WarpM x WarpN warp tiles tile the CTA. The
+    // Derived geometry: warp tiles tile the CTA. The smem budget is
-    // shared-memory budget is layout-aware (crosswise operands add K-major
+    // layout-aware, so it lives in Fp8GemmSmem (gemm.cuh).
    // staging + a canonical buffer), so it lives in Fp8GemmSmem in gemm.cuh
    // together with the resident-CTA hint for __launch_bounds__.
    static constexpr int kWarpsM = BlockM / WarpM;
    static constexpr int kWarpsN = BlockN / WarpN;
    static constexpr int kCtaThreads = kWarpsM * kWarpsN * 32;
@@ -63,73 +54,54 @@ struct Fp8GemmTraits {
                  "warp tile must be a multiple of the m16n8 MMA shape");
 };
-// Quantize-kernel parameter POD: float input (bf16 / fp16 / fp32) -> FP8
+// Quantize-kernel parameter POD: float input -> FP8 with fused amax.
 // with fused amax.
 struct FP8QuantizeParams {
    // Float input and FP8 output buffers; scale is the quantization
    // multiplier (device scalar). amax (may be null) is zero-initialized by
    // the binding and receives the raw-domain absolute maximum.
    const void* __restrict__ input_ptr = nullptr;
    void* __restrict__ output_ptr = nullptr;
-    void* __restrict__ output_transposed_ptr = nullptr;
+    void* __restrict__ output_transposed_ptr = nullptr;  // [cols][rows]
-    // Transposed-output destination ([cols][rows]); the output-layout modes:
+    // Output layout: 0 = row-major only, 1 = transposed only, 2 = both from
-    //   0 = row-major only (output_ptr; the vectorized elementwise kernel)
+    // a single read. Modes 1/2 produce K-contiguous operands so crosswise
-    //   1 = transposed only (output_transposed_ptr; the tiled kernel)
+    // consumers (backward grad_x / grad_w) route through the NT fast path.
    //   2 = both destinations in one read of the input (the tiled kernel)
    // Modes 1/2 exist so crosswise-layout GEMM operands (NN grad_x, TT
    // grad_w) can be produced K-contiguous instead, routing every training
    // GEMM through the dual-congruous NT fast path.
    int out_layout = 0;
-    const float* __restrict__ scale = nullptr;
+    const float* __restrict__ scale = nullptr;  // device multiplier
-    float* __restrict__ amax = nullptr;
+    float* __restrict__ amax = nullptr;         // raw-domain max out
-    // Element count (only the elementwise quantize kernel uses it); the
+    // Element count (elementwise kernel); the tiled kernel views the same
-    // tiled kernel views the same buffer as [rows][cols] row-major.
+    // buffer as [rows][cols] row-major.
    int total = 0;
    int rows = 0;
    int cols = 0;
 };
 // Unified GEMM parameter POD, mirroring AttentionParams: one struct flows
-// through the pre-quantized GEMM kernels. Each kernel touches only the
+// through the kernels; each kernel touches only the fields it needs.
 // fields it needs; buffers are raw pointers packed by the torch binding.
 // Pointer members default to null so optional paths cannot hold garbage.
 struct FP8Params {
-    // Inputs: a/b are FP8 for the pre-quantized path. Scales are
+    // FP8 operands + output; scales are quantization steps (device
-    // quantization steps (device scalars).
+    // scalars). Optional bf16 bias fuses into the epilogue (fp32 add before
-    // Optional bf16 bias broadcast over output rows (fused into the epilogue
+    // the single bf16 rounding); null disables.
    // before the bf16 rounding, so it adds in fp32 — one rounding fewer than
    // the separate out + bias elementwise kernel it replaces). Null disables.
    const void* __restrict__ a_ptr = nullptr;
    const void* __restrict__ b_ptr = nullptr;
    const void* __restrict__ bias_ptr = nullptr;
    void* __restrict__ out_ptr = nullptr;
    const float* __restrict__ scale = nullptr;
-    // Transposed-output mode (set by dispatch_fp8_gemm's swap for NN
+    // NN-swap mode (canonicalize_gemm): the kernel computes the transposed
-    // problems): the kernel computes E[N'][M'] over swapped operands and the
+    // problem and the epilogue scatters D[row][col] to out[col * p.m + row]
-    // epilogue scatters into the caller's [M][N] row-major buffer, so
+    // in the caller's [M][N] buffer. Zero in the plain orientation.
    // D[row][col] lives at out[col * p.m + row] — p.m/p.n are the swapped
    // problem's dims and the D row stride is p.m. Zero in the plain
    // orientation.
    int out_transposed = 0;
-    // Shapes. `int` covers every realistic LLM shape; the kernels promote
+    int m, n, k;  // int covers LLM shapes; kernels promote to int64
    // to int64 for all pointer arithmetic.
    int m, n, k;
-    // Batched (bmm) geometry: grid.z slices step the operand/output pointers
+    // Batched (bmm) geometry: grid.z steps these element strides (0
-    // by these element strides (0 broadcasts the operand across batches).
+    // broadcasts the operand across batches).
    int batch = 1;
    int64_t a_batch_stride = 0;
    int64_t b_batch_stride = 0;
    int64_t out_batch_stride = 0;
-    // Physical leading dimensions (column count, i.e. row stride) of A and
+    // Physical leading dims (row strides) of A and B; the binding packs
-    // B. For a non-transposed operand the stride equals the contract dim;
+    // them so the kernel reads each buffer naturally or transposed per the
-    // for a transposed operand it is the operand's own column count. The
+    // LayoutA/LayoutB tags.
    // binding packs these so the kernel reads both buffers either naturally
    // or transposed depending on the LayoutA/LayoutB tags (see gemm.cuh).
    int a_ld, b_ld;
 };
@@ -1,9 +1,8 @@
 #pragma once
-// FP8 GEMM device code — pure CUDA, no torch. Mirrors the attention kernel
+// FP8 GEMM device code — pure CUDA, no torch. Kernels take the FP8Params
-// layout (attn_*_mma.cuh): kernels take the FP8Params POD, tile shape and
+// POD; tile shape, formats and layout tags ride on one Policy template
-// FP8 format ride on compile-time template parameters, and launchers are
+// parameter (CUTLASS-style), and launchers are plain functions shared by
-// plain functions usable from both the torch binding and pure C tests. The
+// the torch binding and the C tests.
 // quantize kernel lives in quantize.cuh.
 #include <cuda_bf16.h>
 #include <cuda_fp8.h>
@@ -21,7 +20,7 @@ namespace fp8 {
 // m16n8k32 (see astrai::mma_shape<fp8 type>::k in common/mma.cuh)
 constexpr int kMmaK = 32;
-// log2 of a compile-time power of two (for tile_at's swizzle shift).
+// log2 of a compile-time power of two (for the swizzle shifts).
 template <int N, int Acc = 0>
 struct log2_const : log2_const<(N >> 1), Acc + 1> {};
 template <int Acc>
@@ -32,27 +31,14 @@ struct log2_const<1, Acc> {
 // ---------------------------------------------------------------------------
 // Shared device helpers
 // ---------------------------------------------------------------------------
 // The FP8 MMA lives in astrai::mma_sync (common/mma.cuh), instantiated with
 // the kernel's T8 and accumulating in-place. The cp.async primitives live
 // in common/cp_async.cuh.
-// FP8 MMA lives in the shared astrai::mma_sync template (common/mma.cuh);
+// Swizzled address inside a flat [rows * K] staging tile: the 16B chunk
-// instantiate it with the kernel's T8. Accumulates in-place: callers pass
+// index is XORed with the row bits at [3, 3+log2(kChunks)) so a warp's
-// the same accumulator array as both `d` and `c`.
+// ldmatrix fragment load (8 consecutive rows x 16B) hits all 32 banks
-// The cp.async pipeline primitives (predicated 16-byte copy, commit_group,
+// exactly once; chunks stay contiguous, so cp.async staging is unaffected.
 // wait_group + runtime dispatch) live in common/cp_async.cuh.
 // ---------------------------------------------------------------------------
 // Swizzled address inside a flat [rows * K] staging tile: the 16-byte chunk
 // 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 <int K, typename T8>
 __device__ __forceinline__ T8* tile_at(T8* tile, int row, int col) {
    constexpr int kChunks = K / 16;  // 16B chunks per row
@@ -63,11 +49,15 @@ __device__ __forceinline__ T8* tile_at(T8* tile, int row, int col) {
           ((((col >> 4) ^ ((row >> kShift) & (kChunks - 1))) << 4) + (col & 15));
 }
-// Stage-load a CONGRUOUS operand (stored [rows][contract], contract-
+// Stage-load a CONGRUOUS operand (contract-contiguous storage — the only
-// contiguous — the only cp.async-able shape for the canonical tile) into the
+// cp.async-able shape) into the flat [rows * K] swizzled tile. kInterior
-// flat [rows * K] shared tile via tile_at's swizzle. Crosswise operands go
+// drops all predication: valid only for a fully interior CTA (whole rows,
-// through load_crosswise_direct instead.
+// 16B-aligned base|ld, k_base + K <= contract — the fast_cta peel
-template <typename T8, int K, int RowsTile, int kThreads>
+// guarantees these); a thread's chunk run is swizzle-invariant
 // ((n+j)^swz == (n^swz)^j), so the address math folds to one immediate XOR
 // per chunk. Crosswise operands go through load_crosswise_direct instead.
 template <typename T8, int K, int RowsTile, int kThreads,
          bool kInterior = false>
 __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,
@@ -76,15 +66,23 @@ load_operand_tile(T8* tile, const T8* __restrict__ operand, int64_t rows,
    static_assert(RowsTile * kChunks % kThreads == 0,
                  "tile chunks must divide evenly across threads");
    constexpr int kCpt = RowsTile * kChunks / kThreads;  // chunks per thread
    // Linear chunk mapping: thread covers kCpt consecutive 16B chunks of
    // one row (K=64: a contiguous 32B pair; K=32: a single chunk).
    constexpr int kCpr = kChunks / kCpt;  // chunks per row slice
    const int r = tid / kCpr;
    const int c0 = (tid % kCpr) * kCpt * 16;
    if constexpr (kInterior) {
        const char* src = reinterpret_cast<const char*>(
            operand + (block_row + r) * ld + k_base + c0);
        const uintptr_t dst =
            reinterpret_cast<uintptr_t>(tile_at<K>(tile, r, c0));
 #pragma unroll
        for (int j = 0; j < kCpt; ++j)
            astrai::cp_async_16(reinterpret_cast<T8*>(dst ^ (j << 4)),
                                src + j * 16);
    } else {
        const int64_t row = block_row + r;
        const bool row_ok = row < rows;
-    // k_base and every c are multiples of 16, so the per-chunk sources
+        // k_base and every c are multiples of 16, so all chunks share the
-    // share the row base's alignment.
+        // row base's alignment verdict.
        const auto* src = operand + row * ld + k_base;
        const bool chunk_aligned = (reinterpret_cast<uintptr_t>(src) & 15) == 0;
 #pragma unroll
@@ -94,7 +92,7 @@ load_operand_tile(T8* tile, const T8* __restrict__ operand, int64_t rows,
            if (row_ok && chunk_aligned && k_base + c + 15 < contract) {
                astrai::cp_async_16(dst, src + c);
            } else {
-            // Tail chunk (or misaligned base): predicated scalar fill.
+                // Tail chunk / misaligned base / OOB row: scalar fill.
 #pragma unroll
                for (int i = 0; i < 16; ++i)
                    dst[i] =
@@ -102,47 +100,15 @@ load_operand_tile(T8* tile, const T8* __restrict__ operand, int64_t rows,
            }
        }
    }
 // Interior-tile congruous load: zero predication. Valid when
 // block_row + RowsTile <= rows, k_base + K <= contract and
 // (operand base | ld | k_base) is 16B-aligned — the kernel's fast_cta peel
 // guarantees all three. With n = a thread's first chunk a multiple of kCpt,
 // (n+j)^swz == (n^swz)^j, so the swizzled destination of chunk j is the
 // base pointer XOR (j << 4): the whole address math folds into one
 // immediate XOR per chunk (~3 inst/chunk vs ~9 predicated).
 template <typename T8, int K, int RowsTile, int kThreads>
 __device__ __forceinline__ void
 load_operand_tile_interior(T8* tile, const T8* __restrict__ operand,
                           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;
    constexpr int kCpr = kChunks / kCpt;
    const int r = tid / kCpr;
    const int c0 = (tid % kCpr) * kCpt * 16;
    const char* src = reinterpret_cast<const char*>(
        operand + (block_row + r) * ld + k_base + c0);
    const uintptr_t dst = reinterpret_cast<uintptr_t>(tile_at<K>(tile, r, c0));
 #pragma unroll
    for (int j = 0; j < kCpt; ++j)
        astrai::cp_async_16(reinterpret_cast<T8*>(dst ^ (j << 4)),
                            src + j * 16);
 }
-// Loop-carried prefetch state for one congruous operand ring (perf 6.2):
+// Loop-carried prefetch state for one congruous operand ring: per-thread
-// per-thread (r, c0) of the interior copy — the same mapping
+// (r, c0) mapping with the swizzled stage destination and global source
-// load_operand_tile_interior uses — with the swizzled stage destination and
+// pointer carried across k-tiles, so each prefetch chunk is one LDGSTS
-// the global source pointer both carried across k-tiles, so each prefetch
+// issued straight from registers. The guard is a property of the operand's
-// chunk is one LDGSTS at [wr ^ (j << 4)] / [src + j*16] issued straight from
+// layout, so it lives in the type: the false specialization (crosswise
-// registers.
+// operand) is an empty no-op — no dead declarations, no if constexpr at
-//
+// the use sites.
 // Whether an operand has a carry is a property of its layout, so the guard
 // lives in the type: the false specialization (crosswise operand — direct
 // LDG+PRMT staging, no cp.async) is an empty no-op. Crosswise kernel
 // instantiations therefore compile no dead declarations and use sites need
 // no `if constexpr` and no [[maybe_unused]].
 template <bool kAsync, typename T8, int kK, int kRowsTile, int kThreads>
 struct PrefetchCarry;
@@ -172,9 +138,8 @@ struct PrefetchCarry<true, T8, kK, kRowsTile, kThreads> {
              (int64_t)firstTile * kK;
    }
-    // Emit this thread's chunks for the current tile. `pf` false (loop
+    // Emit this thread's chunks for the current tile; pf false (loop tail)
-    // tail) zero-fills: src_size=0 reads nothing, and the destination is
+    // zero-fills into the slot compute(i-1) already released.
    // the slot compute(i-1) already released.
    __device__ __forceinline__ void emit(bool pf) const {
 #pragma unroll
        for (int j = 0; j < kCpt; ++j)
@@ -197,21 +162,16 @@ struct PrefetchCarry<false, T8, kK, kRowsTile, kThreads> {
 };
 // ---------------------------------------------------------------------------
-// Pre-quantized GEMM kernel: FP8 A/B read straight into shared memory, FP32
+// Pre-quantized GEMM kernel: FP8 A/B staged into shared memory, FP32
-// accumulation, BF16 or FP8 output. The input format follows Traits; the
+// accumulation, BF16 output. Operands materialize in the compact canonical
-// tile is compact (row = kK bytes) so MMA fragments read directly — no
+// [rows][kK] tile so MMA fragments read directly — no in-kernel transpose.
 // in-kernel transpose of the operands (the binding handles transposes).
 // ---------------------------------------------------------------------------
 // Direct (synchronous) crosswise load into a canonical rotating stage:
 // LDG.128 x4 (4 consecutive contract bytes x 16 rows) + in-register PRMT
 // transpose + 16 STS.32. Crosswise operands cannot cp.async into the
-// canonical [rows][contract] tile (a 16B global run holds one contract byte
+// canonical tile (a 16B global run holds one contract byte for each of 16
-// for each of 16 rows), so they take this path. A staged variant
+// rows), so they take this path.
 // (cp.async into K-major staging + per-tile smem->smem transpose) measured
 // 15-20% SLOWER than this direct load across every probed shape, including
 // DRAM-streaming B operands — see git history (5745c2f) if it ever needs
 // revisiting for other SKUs.
 template <typename T8, int K, int RowsTile, int kThreads>
 __device__ __forceinline__ void
 load_crosswise_direct(T8* tile, const T8* __restrict__ operand, int64_t rows,
@@ -220,9 +180,8 @@ load_crosswise_direct(T8* tile, const T8* __restrict__ operand, int64_t rows,
    constexpr int kQuads = K / 4;    // 4-byte contract quads per tile
    constexpr int kGroups = RowsTile / 16;
    constexpr int kTChunks = kQuads * kGroups;  // 64B chunks per tile
-    // r0 is always a multiple of 16 (block_row is a multiple of RowsTile and
+    // r0 is a multiple of 16 and p*ld preserves alignment whenever ld has
-    // each group covers 16 rows), and p*ld keeps the base 16B-aligned
+    // it, so every run of a chunk shares one alignment verdict.
    // whenever ld is, so every run of a chunk shares one alignment verdict.
    const bool run_aligned =
        ((reinterpret_cast<uintptr_t>(operand) | ld) & 15) == 0;
    for (int chunk = tid; chunk < kTChunks; chunk += kThreads) {
@@ -247,8 +206,7 @@ load_crosswise_direct(T8* tile, const T8* __restrict__ operand, int64_t rows,
 #pragma unroll
            for (int i = 0; i < 16; ++i) {
                // word i = row r0+i's quad: byte i of each of the four runs
-                // [v0.b(i), v1.b(i), v2.b(i), v3.b(i)]. Byte i of a uint4
+                // [v0.b(i), v1.b(i), v2.b(i), v3.b(i)].
                // lives in its (i>>2)-th 32-bit register.
                const unsigned nib = i & 3;
                const unsigned sel = nib | ((nib + 4) << 4);
                const unsigned w01 =
@@ -286,14 +244,9 @@ load_crosswise_direct(T8* tile, const T8* __restrict__ operand, int64_t rows,
 // Layout-aware shared-memory budget and occupancy hint. Every operand ring
 // holds kStages+1 buffers: the load for tile i+kStages targets slot
-// (i-1)%(kStages+1) — already consumed — so neither the congruous cp.async
+// (i-1)%(kStages+1) — already consumed — so neither load path needs a
-// path nor the direct-crosswise path needs a post-compute barrier (one
+// post-compute barrier (one __syncthreads per k-tile). The 48KB static
-// __syncthreads per k-tile). (A lean kStages-deep ring traded that barrier
+// watermark picks the resident-CTA hint for __launch_bounds__.
 // for a 4th resident CTA and measured slower — 1280³ +5..9% — so the knob
 // was removed; see git history if a small-SKU variant is ever needed.)
 // The 48KB static-smem watermark picks the resident-CTA hint for
 // __launch_bounds__ (sm_89: 100KB smem per SM, so two CTAs fit while each
 // stays within the static budget).
 template <typename Traits, typename LayoutA, typename LayoutB>
 struct Fp8GemmSmem {
    // Crosswise (direct-load) operands: A ColMajor storage, B RowMajor
@@ -307,11 +260,9 @@ struct Fp8GemmSmem {
 };
 // ---------------------------------------------------------------------------
-// Kernel policy: one type per kernel instantiation. The CTA/K-tile/pipeline
+// Kernel policy: one type per kernel instantiation (CUTLASS-style
-// shape rides on Fp8GemmTraits and the operand layouts + scheduling knobs
+// consolidation) — traits + layout tags + scheduling knobs, the single
-// hang beside them — this is the single template parameter fp8_gemm_kernel
+// template parameter the kernel and both collectives take.
 // (and both collectives) take, mirroring CUTLASS's kernel-policy
 // consolidation.
 template <FP8Format Fmt_, int BlockM_, int BlockN_, typename LayoutA_,
          typename LayoutB_, int WarpM_, int WarpN_, int kK_, int Stages_,
          int GroupRaster_, bool StreamOut_ = false, bool FastLoop_ = false>
@@ -323,30 +274,22 @@ struct Fp8GemmPolicy {
    static constexpr int kGroupRaster = GroupRaster_;
    static constexpr bool kStreamOut = StreamOut_;
    static constexpr bool kFastLoop = FastLoop_;
    using Smem = Fp8GemmSmem<Traits, LayoutA_, LayoutB_>;
    // Flattened for __launch_bounds__, which takes no dependent type names.
    static constexpr int kCtaThreads = Traits::kCtaThreads;
-    static constexpr int kMinCtas =
+    static constexpr int kMinCtas = Smem::kMinCtas;
-        Fp8GemmSmem<Traits, LayoutA_, LayoutB_>::kMinCtas;
+    static constexpr int kSmemBytes = Smem::kBytes;
 };
-// LayoutA / LayoutB tag the operands' storage (CUTLASS-style, see common.h):
+// LayoutA / LayoutB tag the operands' storage; the kernel always computes
-// A RowMajor = [M][K] / ColMajor = [K][M]; B RowMajor = [K][N] /
+// out[m][n] = sum_p tileA[m][p] * tileB[n][p] with tiles materialized in
-// ColMajor = [N][K]. The kernel always computes
+// the canonical [M][kK] / [N][kK] layout, so the tags only change how the
-//   out[m][n] = sum_p tileA[m][p] * tileB[n][p]
+// stage-load gathers from global memory. With p.out_transposed set (the
-// with the tiles materialized in the canonical [M][kK] / [N][kK] layout, so the
+// swap dispatch for NN problems) the kernel runs the transposed problem
-// MMA fragments are read identically regardless of layout. The tags only
+// E = B^T * A^T and the epilogue scatters D[m][n] = E[n][m]; bias then
-// change how the stage-load gathers the operand from global memory:
+// indexes D-cols, i.e. the kernel's rows.
 //   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]
 // With p.out_transposed set (the swap dispatch for NN problems, see
 // dispatch_fp8_gemm) the kernel runs the transposed problem E = B^T * A^T
 // over swapped operands and the epilogue scatters D[m][n] = E[n][m] into the
 // caller's [M][N] row-major buffer — bias then indexes D-cols, i.e. the
 // kernel's rows (see the epilogue).
 // BlockM x BlockN CTA as (BlockM/64) x (BlockN/32) warps of 64x32 warp tiles
 // (mt x nt = 4x4 MMA each).
 //
-// The kernel decomposes CUTLASS-style into three collectives below:
+// The kernel decomposes CUTLASS-style into three collectives:
 //   Fp8GemmTileScheduler   — CTA id -> (block_m, block_n) raster order
 //   Fp8CollectiveMainloop  — stage rings, gmem->smem loads, mma.sync loop
 //   Fp8CollectiveEpilogue  — fused bias + bf16 scatter + coalesced copy-out
@@ -354,19 +297,9 @@ struct Fp8GemmPolicy {
 // ---------------------------------------------------------------------------
 // Tile scheduler: the linear CTA id maps to (block_m, block_n) in grouped
-// (L2-friendly, CUTLASS-style) or plain raster order — the grouped order
+// (L2-friendly) raster — consecutive CTAs share one B column stripe — or
-// makes consecutive CTAs cover a group of kRasterGroup M-tiles before
+// plain N-fastest raster (kRasterGroup=0, the measured best for dX's
-// advancing along N, so all CTAs of one group share the same B column
+// crosswise-B layouts where grouping was neutral).
 // stripe and B tiles stay hot in L2 across the wave (the plain N-fastest
 // order makes each wave touch every B tile instead; kRasterGroup=0 selects
 // plain, the measured best for dX's crosswise-B layouts where grouping
 // measured neutral).
 // Persistent schedules (static round-robin and an atomic ticket dispenser,
 // grid capped at the resident CTAs) were both measured and rejected on L20:
 // the stride desynchronizes the in-flight window (-4..-8%), and the ticket
 // variant recovers the L2 locality but lands within noise of plain waves
 // (its loop-head barrier costs what the CTA-restart overlap saves). Keep
 // the classic retiring-wave launch.
 template <int kRasterGroup>
 struct Fp8GemmTileScheduler {
    static __device__ int2 tile(const uint3& block, const dim3& blocks) {
@@ -405,20 +338,12 @@ struct Fp8CollectiveMainloop {
    static constexpr bool kDirectB = Smem::kDirectB;
    static_assert(kStages >= 1 && kStages <= 8,
                  "FP8 GEMM stages must be in [1, 8]");
-    // CTA = (BlockM/WarpM) x (BlockN/WarpN) warps of WarpM x WarpN tiles,
+    // CTA = (BlockM/WarpM) x (BlockN/WarpN) warps, each warp computing
-    // each warp computing (WarpM/16) x (WarpN/8) m16n8k32 MMAs (mt x nt).
+    // kMt x kNt m16n8k32 MMAs. Rings rotate kStages+1 buffers (see
-    // The default 128x128 CTA runs 8 warps of 64x32 (mt x nt = 4x4); the
+    // Fp8GemmSmem) — one __syncthreads per k-tile.
    // small-shape path uses 64x64 CTAs of 32x32 warps (cuBLAS-style) so more
    // CTAs fit per SM (see launch_fp8_gemm).
    static constexpr int kMt = Traits::kWarpM / 16;  // 16-row MMA tiles per warp
    static constexpr int kNt = Traits::kWarpN / 8;   // 8-col MMA tiles per warp
    static constexpr int kSegs = kK / kMmaK;  // mma-sized k segments per tile
    // Both operands rotate kStages+1 buffers: the load for tile i+kStages
    // targets slot (i-1)%(kStages+1), which compute finished reading before
    // this iteration's barrier 1 — the direct-crosswise prefetch (issued
    // right after barrier 1) and the congruous cp.async prefetch alike —
    // so NO post-compute barrier is needed (one __syncthreads per k-tile,
    // the classic multistage rotation).
    static constexpr int kARing = Smem::kRingDepth;
    static constexpr int kBRing = Smem::kRingDepth;
    static constexpr int kAStageBytes = kBlockM * kK;
@@ -435,14 +360,11 @@ struct Fp8CollectiveMainloop {
    const int a_row0;  // + mt * 16 in the loop
    const int b_row0;  // + nt * 8
    const int64_t tile_count;
-    // Interior-CTA peel (kFastLoop instantiations only): when both operands
+    // Interior-CTA peel (kFastLoop instantiations only): whole-CTA,
-    // are congruous, whole-CTA, 16B-aligned and K has no tail, the mainloop
+    // 16B-aligned, K without tail — the mainloop then runs a compile-time
-    // runs a compile-time-specialized copy whose loads carry no predication
+    // specialized copy with no per-chunk predication (measured +4.5..10% on
-    // — the per-chunk guards cost ~6 of ~100 instructions per warp per
+    // the issue-bound small CTA; the 128x128 CTA regressed, so only the
-    // k-tile, and the small-CTA path is issue-bound there (measured
+    // small CTA opts in). The verdict is uniform per CTA.
    // +4.5..10% on 256³..1024³; the 128x128 kernel regressed ~3% with the
    // same change, so only the small CTA opts in). All verdicts are uniform
    // per CTA: one branch picks the loop copy.
    const bool fast_cta;
    __device__ Fp8CollectiveMainloop(char* smem, const T8* a, const T8* b,
@@ -465,49 +387,33 @@ struct Fp8CollectiveMainloop {
                   ((reinterpret_cast<uintptr_t>(b) | (uint64_t)b_ld) & 15) == 0 &&
                   (k % kK) == 0) {}
-    // Stage-slot helpers: the rings rotate one slot per k-tile, so callers
+    // Stage-slot helpers: rings rotate one slot per k-tile, so callers
    // either compute the slot from the tile index (prologue, generic loop)
-    // or carry an advancing pointer (steady-state fast loop below).
+    // or carry an advancing pointer (steady-state fast loop).
    __device__ __forceinline__ T8* a_stage_of(int64_t tile) const {
        return a_base + (size_t)(tile % kARing) * kAStageBytes;
    }
    __device__ __forceinline__ T8* b_stage_of(int64_t tile) const {
        return b_base + (size_t)(tile % kBRing) * kBStageBytes;
    }
-    // Asynchronous loads for tile `tile`: congruous operands cp.async into
+    // Asynchronous congruous loads for one k-tile: cp.async into the
-    // their canonical rings. Called after the post-compute barrier, alongside
+    // canonical rings; kFast selects the predication-free interior copy
-    // the commit.
+    // (fast_cta admits only congruous operands). Called after the
    // post-compute barrier, alongside the commit.
    template <bool kFast = false>
    __device__ __forceinline__ void load_async(T8* a_stage, T8* b_stage,
                                               int64_t k_base) const {
        if constexpr (!kDirectA)
-            load_operand_tile<T8, kK, kBlockM, kCtaThreads>(
+            load_operand_tile<T8, kK, kBlockM, kCtaThreads, kFast>(
                a_stage, a, m, k, a_ld, tid, k_base, block_m * kBlockM);
        if constexpr (!kDirectB)
-            load_operand_tile<T8, kK, kBlockN, kCtaThreads>(
+            load_operand_tile<T8, kK, kBlockN, kCtaThreads, kFast>(
                b_stage, b, n, k, b_ld, tid, k_base, block_n * kBlockN);
    }
-    // Predication-free interior variant of load_async: congruous operands
+    // Synchronous direct-crosswise loads for one k-tile. In the steady
-    // with full CTA rows, aligned (base | ld), k_base + kK <= k. fast_cta
+    // state this runs right after barrier 1, so the LDG latency and the
-    // admits only congruous operands, so no crosswise fallback is needed.
+    // PRMT transpose overlap the MMA phase instead of stalling the
-    __device__ __forceinline__ void load_async_fast(T8* a_stage, T8* b_stage,
+    // inter-barrier window (which dominated the dX/dW stall profile).
                                                    int64_t k_base) const {
        if constexpr (!kDirectA)
            load_operand_tile_interior<T8, kK, kBlockM, kCtaThreads>(
                a_stage, a, a_ld, tid, k_base, block_m * kBlockM);
        if constexpr (!kDirectB)
            load_operand_tile_interior<T8, kK, kBlockN, kCtaThreads>(
                b_stage, b, b_ld, tid, k_base, block_n * kBlockN);
    }
    // Synchronous direct-crosswise loads for tile `tile` into the operand's
    // (kStages+1)-deep canonical ring. In the steady state this runs right
    // after barrier 1, so the LDG latency and the PRMT transpose overlap the
    // MMA phase of the current tile instead of stalling the inter-barrier
    // window (which dominated the dX/dW stall profile: barrier 3.7-4.1 +
    // long-scoreboard 1.6-1.8 stalls per issue on the production shapes).
    // Ring safety: the write targets buffer (i+kStages)%(kStages+1) =
    // (i-1)%(kStages+1), which compute(i-1) finished reading before the
    // previous barrier and compute(i+kStages) does not touch until several
    // barriers later.
    __device__ __forceinline__ void load_direct(T8* a_stage, T8* b_stage,
                                                int64_t k_base) const {
        if constexpr (kDirectA)
@@ -518,21 +424,17 @@ struct Fp8CollectiveMainloop {
                b_stage, b, n, k, b_ld, tid, k_base, block_n * kBlockN);
    }
-    // Prime the pipeline. Each committed group occupies one circular shared
+    // Prime the pipeline: kStages committed groups, one per stage slot.
-    // memory stage. The commit is unconditional: when K is shorter than the
+    // The commit is unconditional — when K is shorter than the pipeline the
-    // pipeline (tile_count < kStages) the skipped stages commit empty groups
+    // skipped stages commit empty groups, so the group sequence stays
-    // so the group sequence stays tile-indexed — the steady-state
+    // tile-indexed and the steady-state wait count never needs a runtime
-    // wait_group<kStages-1> below is then correct for every iteration and
+    // dispatch.
    // no runtime wait-count dispatch is needed (the dispatch ladder cost 16
    // instructions per k-tile: ISETP/SEL chains picking DEPBAR immediates).
    // Direct loads run synchronously here (back to back with their commit);
    // the steady state below overlaps them with the compute phase.
    __device__ __forceinline__ void prologue() const {
 #pragma unroll
        for (int stage = 0; stage < kStages; ++stage) {
            if (stage < tile_count) {
                if (fast_cta)
-                    load_async_fast(a_stage_of(stage), b_stage_of(stage),
+                    load_async<true>(a_stage_of(stage), b_stage_of(stage),
                                     (int64_t)stage * kK);
                else
                    load_async(a_stage_of(stage), b_stage_of(stage),
@@ -544,32 +446,27 @@ struct Fp8CollectiveMainloop {
        }
    }
-    // Steady-state mainloop, compile-time specialized on fast_cta: the fast
+    // Steady-state mainloop, compile-time specialized on kFast: the fast
-    // copy runs predication-free loads; the generic copy keeps full
+    // copy runs predication-free loads with loop-carried read/write
-    // predication. kFastLoop=false instantiates only the generic copy —
+    // pointers; the generic copy keeps full predication. kFastLoop=false
-    // codegen identical to the pre-peel kernel.
+    // instantiates only the generic copy.
    template <bool kFast>
    __device__ __forceinline__ void run_loop(float acc[kNt][kMt][4]) const {
        const int lane = tid & 31;
-        // Fast-path write carries: one per congruous operand (see
+        // Fast-path write carries: one per congruous operand (crosswise
-        // PrefetchCarry; crosswise operands get the empty no-op type).
+        // operands get the empty no-op type), targeting the first
-        // Construction targets the first prefetched tile (kStages).
+        // prefetched tile (kStages).
        PrefetchCarry<!kDirectA, T8, kK, kBlockM, kCtaThreads> carry_a(
            a_base, kARing, kAStageBytes, a, a_ld, block_m * kBlockM, tid,
            kStages);
        PrefetchCarry<!kDirectB, T8, kK, kBlockN, kCtaThreads> carry_b(
            b_base, kBRing, kBStageBytes, b, b_ld, block_n * kBlockN, tid,
            kStages);
        // Interleaved prefetch (cuBLAS/CUTLASS loop shape): the next tile's
        // LDGSTS chunks ride inside the MMA phase so their issue slots fill
        // the tensor-pipe gaps ptxas otherwise pads with NOPs (23 NOPs per
        // 32 QMMA here versus 0 in the cuBLAS loop).
        // Steady-state read carries: the LDSM base of the current k-tile's
        // stage with the lane offset folded in, advanced one stage per
-        // iteration with an equality wrap (the add sequence is exact). This
+        // iteration with an equality wrap — replaces the per-k-tile
-        // replaces the per-k-tile (tile % ring) * stage_bytes recomputation —
+        // (tile % ring) * stage_bytes recomputation (a UIMAD.WIDE
-        // its SASS form was a UIMAD.WIDE magic-division ladder, ~10
+        // magic-division ladder in SASS).
        // uniform-pipe instructions per operand per k-tile (perf 6.2).
        const unsigned a_rd0 = __cvta_generic_to_shared(a_base) + a_lane_off(lane);
        const unsigned b_rd0 =
            __cvta_generic_to_shared(b_base) +
@@ -579,8 +476,8 @@ struct Fp8CollectiveMainloop {
        unsigned a_rd = a_rd0, b_rd = b_rd0;
        for (int64_t tile_index = 0; tile_index < tile_count; ++tile_index) {
        // In the steady state exactly kStages-1 younger groups are in flight
-        // when this fires; the tail's unconditional (possibly empty) commits
+        // when this fires; the tail's unconditional (possibly empty)
-        // keep that invariant true for every iteration.
+        // commits keep that invariant true for every iteration.
        const bool prefetch = tile_index + kStages < tile_count;
        astrai::cp_async_wait_group<kStages - 1>();
        // Barrier 1: every thread's cp.async for this stage is complete
@@ -607,12 +504,9 @@ struct Fp8CollectiveMainloop {
        }
        // kNt ldmatrix.x2 (B) + kMt ldmatrix.x4 (A) feed kMt*kNt*2 mma.sync
-        // per k_seg — 0.5 load instructions per MMA, versus 4.5 scalar LDS
+        // per k_seg — 0.5 load instructions per MMA. B fragments
-        // per MMA in the 128x64-tile version (the kernel was LSU-issue-bound
+        // double-buffer across k_segs; kPairB folds the two adjacent nt
-        // there). B fragments double-buffer across k_segs. kPairB folds the
+        // fragments of one pair into a single x4 (see b4_lane_off).
        // two adjacent nt fragments of one pair into a single x4 (see
        // b4_lane_off): kNt/2 x4 loads, regs {r0,r1}/{r2,r3} feeding
        // the even/odd nt MMAs respectively.
        unsigned b_frag[2][kNt][2];
        unsigned b_frag4[2][kNt / 2][4];
        load_b_frags(b_frag[0][0], b_frag4[0][0], b_seg[0]);
@@ -622,14 +516,9 @@ struct Fp8CollectiveMainloop {
            if (k_seg + 1 < kSegs)
                load_b_frags(b_frag[bnext][0], b_frag4[bnext][0],
                             b_seg[k_seg + 1]);
-        // Software-pipelined A fragments: the ldmatrix.x4 for row mt+1
+        // Software-pipelined A fragments: the ldmatrix.x4 for row mt+1 is
-        // is issued before the MMAs consuming row mt, so the LDS fixed
+        // issued before the MMAs consuming row mt, so the LDS latency hides
-        // latency hides behind tensor-pipe work (cuts the `wait` stall,
+        // behind tensor-pipe work. Costs 4 extra registers.
        // ~2.3 cycles/issue before this). Costs 4 extra registers.
        // (Cross-k_seg prefetch of row 0 was tried and reverted: the
        // register handoff broke ptxas's software pipelining — 171T → 95T
        // at 2048³; the tensor pipe is issue-bound and the seg-start LDS
        // already hides behind the b-fragment issue order.)
        unsigned a_frag[kMt + 1][4];
        astrai::ldmatrix_x4_lane(a_frag[0], a_seg[k_seg]);
 #pragma unroll
@@ -654,8 +543,7 @@ struct Fp8CollectiveMainloop {
        }
        }
        // Generic loop (no interleaved prefetch): the next tile's predicated
-        // loads run after the MMA phase — the fast loop's carries already
+        // loads run after the MMA phase.
        // emitted inside it.
        if constexpr (!kFast) {
            if (prefetch) {
                load_async(a_stage_of(tile_index + kStages),
@@ -664,12 +552,8 @@ struct Fp8CollectiveMainloop {
            }
        }
        // Unconditional commit: empty in the tail, it pads the group
-        // sequence so the fixed wait above stays correct (and the
+        // sequence so the fixed wait above stays correct.
        // predicated-off chunks' zero-fill lands in the slot compute(i-1)
        // released — nothing reads it again before the epilogue drain).
        astrai::cp_async_commit_group();
        // Advance the carries: one stage slot forward, wrapping on the
        // exact ring boundary.
        a_rd += (unsigned)kAStageBytes;
        if (a_rd == a_rd_end) a_rd = a_rd0;
        b_rd += (unsigned)kBStageBytes;
@@ -693,23 +577,14 @@ struct Fp8CollectiveMainloop {
    }
  private:
-    // Per-lane ldmatrix fragment addressing (base-pair scheme, mirrored from
+    // Per-lane ldmatrix fragment addressing (base-pair scheme, mirrored
-    // the cuBLAS SASS: one base register per operand per k_seg, every
+    // from the cuBLAS SASS): one base register per operand per k_seg,
-    // fragment offset an LDSM immediate — zero address arithmetic inside the
+    // every fragment offset an LDSM immediate — zero address arithmetic
-    // MMA phase). The closure works because the XOR swizzle's source bits
+    // inside the MMA phase. The XOR swizzle's source bits come only from
-    // come only from the lane's row-within-matrix (r7): the 8- and 16-row
+    // the lane's row-within-matrix (r7), so the 8/16-row fragment steps
-    // fragment steps (nt*8, mt*16) never reach them, so
+    // never reach them and
-    //   addr(s, mt) = lane_base + mt*(16*kK)  ^ (s<<5)     [A, x4 fragment]
+    //   addr(s, mt) = lane_base + mt*(16*kK)  ^ (s<<5)     [A, x4]
-    //   addr(s, nt) = lane_base + nt*(8*kK)   ^ (s<<5)     [B, x2 fragment]
+    //   addr(s, nt) = lane_base + nt*(8*kK)   ^ (s<<5)     [B, x2 / x4]
    // where the ^ (s<<5) lands inside the 16B-chunk swizzle field (each k_seg
    // advances the chunk index by 2 = 32B) and the step lands outside it.
    //   kChunks=4 (K=64):  swizzle bits = row[2:1] = r7[2:1]
    //   kChunks=8 (K=128): swizzle bits = row[2:0] = r7[2:0]
    //   kChunks=2 (K=32):  swizzle bit  = row[2]   = r7[2]  (single k_seg)
    // Replaces the former a_off[kSegs][kMt]/b_off[kSegs][kNt] runtime tables
    // (16 registers + one IADD per LDSM): at 131 regs the tables spilled and
    // ptxas rematerialized every address each k-tile (~55 of 146 hot-loop
    // instructions were LOP3/IMAD address math; cuBLAS's inner loop has ~0).
    __device__ __forceinline__ unsigned a_lane_off(int lane) const {
        const int r7 = lane & 7;          // row within the 8-row matrix
        const int rh8 = (lane >> 3) & 1;  // +8 rows (A: lanes 8-15, 24-31)
@@ -718,10 +593,8 @@ struct Fp8CollectiveMainloop {
        constexpr int kShift = 3 - log2_const<kChunks>::value;  // tile_at's shift
        const unsigned lswz =
            static_cast<unsigned>((r7 >> kShift) & (kChunks - 1));
-        // Stage-relative, loop-invariant per-lane base (added to each ring
+        // Stage-relative, loop-invariant per-lane base; A's fragment row
-        // slot's converted base once per k-tile). A's fragment row carries
+        // carries the +8-row (rh8) and +1-chunk (rh16) halves.
        // the +8-row half (rh8) and the +1-chunk half (rh16) — matching the
        // m16n8k32 operand layouts above.
        return static_cast<unsigned>((a_row0 + rh8 * 8 + r7) * kK +
                                     ((rh16 ^ lswz) << 4));
    }
@@ -734,16 +607,12 @@ struct Fp8CollectiveMainloop {
            static_cast<unsigned>((r7 >> kShift) & (kChunks - 1));
        return static_cast<unsigned>((b_row0 + r7) * kK + ((rh8 ^ lswz) << 4));
    }
-    // x4-paired B loads (cuBLAS/CUTLASS loop shape): one ldmatrix.x4 feeds
+    // x4-paired B loads: one ldmatrix.x4 feeds the two adjacent nt
-    // the two adjacent nt fragments — 2 x4 instead of 4 x2 per k_seg (12
+    // fragments. Lane contract: lanes 0-7 address rows n0..n7 chunk c,
-    // LDSM per k-tile instead of 16). Lane contract: lanes 0-7 address
+    // lanes 8-15 rows n0..n7 chunk c+1, lanes 16-23 rows n8..n15 chunk c,
-    // rows n0..n7 chunk c, lanes 8-15 rows n0..n7 chunk c+1, lanes 16-23
+    // lanes 24-31 rows n8..n15 chunk c+1. The +8-row step never reaches
-    // rows n8..n15 chunk c, lanes 24-31 rows n8..n15 chunk c+1; regs
+    // the swizzle source bits for kK <= 64; kK=128 swizzles on row[2:0]
-    // {r0,r1} are the even nt's k-halves, {r2,r3} the odd nt's. The +8-row
+    // where +8 flips bits, so that config keeps the x2 loads.
    // step never reaches the swizzle source bits for kK <= 64 (kChunks<=4:
    // bits row[2:1]), so lanes 16-31 reuse the same lswz and each pair
    // address is the even-nt base + p*(16*kK). kK=128 swizzles on row[2:0]
    // where +8 flips bits — that config keeps the x2 loads.
    static constexpr unsigned kMtStep = 16 * kK;  // bytes per m-tile row step
    static constexpr unsigned kNtStep = 8 * kK;   // bytes per n-tile row step
    static constexpr unsigned kSegXor = 32;       // chunk-index +2 per k_seg
@@ -755,10 +624,8 @@ struct Fp8CollectiveMainloop {
    }
    // One k_seg's B-fragment loads, shared by the initial fill and the
-    // double-buffer's next-seg fill: kNt/2 paired ldmatrix.x4 (kPairB) or
+    // double-buffer's next-seg fill. frag2/frag4 are the flat bases of one
-    // kNt ldmatrix.x2 from the seg's base-pair address. frag2/frag4 are the
+    // b_frag / b_frag4 buffer (the unused one is never touched).
    // flat bases of one b_frag / b_frag4 buffer (the unused one of the pair
    // is never touched).
    __device__ __forceinline__ void
    load_b_frags(unsigned* frag2, unsigned* frag4, unsigned seg_base) const {
 #pragma unroll
@@ -811,42 +678,30 @@ struct Fp8CollectiveEpilogue {
          thread_in_group(tid & 3),
          block_m(block_m), block_n(block_n) {}
-    // Swizzled address of one 16B chunk (row r, chunk c) of the staged tile.
+    // Swizzled address of one 16B chunk (row r, chunk c) of the staged
-    // Plain orientation: D-local, kBlockM rows of kBlockN elems.
+    // tile. Plain orientation: kBlockM rows of kBlockN elems; out-
-    // Out-transposed (swap dispatch): the tile stages D-local rows over the
+    // transposed (swap dispatch): rows and row length trade places. Both
-    // swapped problem, so it has kBlockN rows of kBlockM — rows and row
+    // row-chunk counts are powers of two, keeping the XOR swizzle
-    // length trade places. Both row-chunk counts are powers of two, keeping
+    // well-defined.
    // the 16B-chunk XOR swizzle well-defined.
    __device__ __forceinline__ __nv_bfloat16* out_chunk(int r, int c) const {
        return tile_out + (size_t)r * row_elems +
               ((c ^ (r & (row_chunks - 1))) * 8);
    }
    // Address of one element of the staged tile.
    __device__ __forceinline__ __nv_bfloat16* out_elem(int r, int v) const {
        return out_chunk(r, v >> 3) + (v & 7);
    }
-    // Scatter the accumulators into the staging tile. The direct bf16
+    // Scatter the accumulators into the staging tile: the operand rings are
-    // epilogue goes through the operand shared memory: the A/B rings are
+    // dead once the mainloop ends, so their space stages the bf16 output
-    // dead once the mainloop ends, so their space stages the output tile
+    // tile. Threads scatter (STS.32 of bf16x2 pairs), a barrier makes the
-    // (kBlockM x kBlockN bf16, always <= the ring budget). Threads first
+    // tile coherent, then the whole CTA copies it out in fully-coalesced
-    // scatter their accumulators into the tile (STS.32 of bf16x2 pairs), a
+    // 16B chunks. The 16B-chunk XOR swizzle keeps both the scatter and the
-    // barrier makes the tile coherent, then the whole CTA copies it out in
+    // gather conflict-free.
    // fully-coalesced 16B chunks. The direct per-thread stores this replaces
    // hit 8 disjoint 16B segments per warp (rows are n*2 bytes apart), ~50%
    // write efficiency — measurable at 2048+ where the epilogue is ~8% of
    // runtime. The 16B-chunk XOR swizzle (chunk index ^ row) keeps both the
    // scatter and the gather conflict-free: a lane quad's chunk and the 8
    // rows of one gather phase map to distinct 4-bank groups.
    __device__ __forceinline__ void stage(float acc[kNt][kMt][4]) const {
-        // Fused bias (idea B): added to the fp32 accumulator before the
+        // Fused bias: added to the fp32 accumulator before the single bf16
-        // single bf16 rounding — one fewer rounding than the out + bias
+        // rounding. The per-lane loads are L1 broadcasts; rows past the
-        // elementwise pass this replaces, and no extra kernel launch / m*n
+        // edge skip the load (their smem slots never copy out). Under
-        // round-trip. The per-lane loads (2 per nt, kMt-times re-read) are
+        // out_transposed the bias indexes D-cols = the kernel's rows.
        // L1 broadcasts; rows past the N edge skip the load (their smem
        // slots never copy out). Under out_transposed the bias indexes
        // D-cols = the kernel's rows, so one load per r0 broadcasts across
        // the row's cols instead.
        const int local_col0 = warp_n * Traits::kWarpN + thread_in_group * 2;
        const int64_t bias_col0 = block_n * kBlockN;
        const int64_t bias_row0 = block_m * kBlockM;
@@ -865,8 +720,8 @@ struct Fp8CollectiveEpilogue {
                    const int r0 = warp_m * Traits::kWarpM + group + mt * 16;
                    const float* tile_acc = acc[nt][mt];
                    // Two bf16x2 stores per accumulator tile: rows g and
-                    // g+8 of the m16n8 output, columns tig*2 and tig*2+1
+                    // g+8 of the m16n8 output, columns tig*2/tig*2+1 inside
-                    // inside one 16B chunk.
+                    // one 16B chunk.
                    const int off = col & 7;  // element offset in the chunk
                    *reinterpret_cast<__nv_bfloat162*>(
                        out_chunk(r0, col >> 3) + off) =
@@ -881,11 +736,9 @@ struct Fp8CollectiveEpilogue {
        } else {
            // Transposed scatter: accumulator (kernel row r0, col) is
            // D[col0_global + col][row0_global + r0], staged at T[col][r0].
-            // The acc pair spans two staged rows, so these are scalar stores
+            // The acc pair spans two staged rows, so these are scalar
-            // (4 per (nt, mt) vs the packed bf16x2 pair — the swap path is
+            // stores (the swap path is the rare NN layout). OOB elements
-            // the rare NN layout); the row swizzle keeps the quad's stores
+            // store dead lanes of the tile, never copied out.
            // bank-spread. OOB elements store dead lanes of the tile, never
            // copied out.
 #pragma unroll
            for (int nt = 0; nt < kNt; ++nt) {
                const int col = local_col0 + nt * 8;
@@ -912,8 +765,7 @@ struct Fp8CollectiveEpilogue {
    // Coalesced copy-out: thread -> one 16B chunk; consecutive threads walk
    // a row so each global transaction covers a full 128B line. Under the
    // swap the staged rows are D-rows counted from block_n's stripe while
-    // the row length is kernel m', so row/stride flip to the swapped dims
+    // the row length is kernel m', so row/stride flip to the swapped dims.
    // (D[row][col] = out[row * p.m + col]).
    __device__ __forceinline__ void store(__nv_bfloat16* out_bf16) const {
        constexpr int kTotalChunks =
            kBlockM * (kBlockN / 8);  // == kBlockN * (kBlockM/8)
@@ -934,18 +786,15 @@ struct Fp8CollectiveEpilogue {
            if (col + 8 <= row_stride &&
                (reinterpret_cast<uintptr_t>(dst) & 15) == 0) {
                if constexpr (kStreamOut) {
-                    // Evict-first streaming store knob. Measured neutral on
+                    // Evict-first streaming store knob: neutral on L20
-                    // L20 squares and -3..4% on rects (the evict-first
+                    // squares, -3..4% on rects; kept for other SKUs.
                    // policy hurts more than the L2 B-tile protection helps
                    // at these sizes); kept as a template knob for other
                    // SKUs. Default off.
                    __stcs(reinterpret_cast<uint4*>(dst), v);
                } else {
                    *reinterpret_cast<uint4*>(dst) = v;
                }
            } else {
                // Row-edge chunk or an odd-stride row base: spill the
-                // elements that survive the row edge (and stay aligned).
+                // elements that survive the row edge.
                const __nv_bfloat16* elems =
                    reinterpret_cast<const __nv_bfloat16*>(&v);
                for (int e = 0; e < 8 && col + e < row_stride; ++e)
@@ -971,12 +820,8 @@ __global__ void __launch_bounds__(Policy::kCtaThreads, Policy::kMinCtas)
    using Traits = typename Policy::Traits;
    using Mainloop = Fp8CollectiveMainloop<Policy>;
    using Epilogue = Fp8CollectiveEpilogue<Policy>;
-    // Tiles are flat [rows * kK] with a 16B-chunk XOR swizzle (tile_at):
+    // Stages live in dynamic shared memory so deep pipelines (> 48KB
-    // ldmatrix reads whole 16B chunks through the same mapping the staging
+    // static limit) opt in via cudaFuncSetAttribute in the launcher.
    // writes, and the swizzle removes the bank conflict the unswizzled
    // 8-word row stride caused (see tile_at). The stages live in dynamic
    // shared memory so deep pipelines (kStages * (kBlockM + kBlockN) * kK >
    // 48KB static limit) opt in via cudaFuncSetAttribute in the launcher.
    extern __shared__ __align__(16) char fp8_gemm_smem[];
    // Batch slice (grid.z): broadcast operands carry a 0 stride, so the
@@ -999,9 +844,7 @@ __global__ void __launch_bounds__(Policy::kCtaThreads, Policy::kMinCtas)
    float acc[Mainloop::kNt][Mainloop::kMt][4] = {};  // [nt][mt][acc]
    mainloop.prologue();
    mainloop.accumulate(acc);
-    // Drain the pipeline before the epilogue reclaims the operand rings for
+    // Drain the pipeline before the epilogue reclaims the operand rings.
    // output staging: the loop's last commits (possibly only zero-filling
    // predicated-off chunks) are nobody's wait target anymore.
    astrai::cp_async_wait_all();
    Epilogue(fp8_gemm_smem, p, bn.x, bn.y, threadIdx.x).run(acc, out_bf16);
 }
@@ -1011,8 +854,7 @@ __global__ void __launch_bounds__(Policy::kCtaThreads, Policy::kMinCtas)
 // ---------------------------------------------------------------------------
 // SM count of the current device (cached per device; benign init race —
-// every writer stores the same value). Host-side only: feeds the
+// every writer stores the same value).
 // device-adaptive dispatch thresholds.
 inline int device_sm_count() {
    static int cached[64] = {};
    int dev = 0;
@@ -1027,15 +869,12 @@ inline int device_sm_count() {
    return sms;
 }
-
+// Launch one kernel instantiation with its shared-memory budget: budgets
-// Launch one kernel instantiation with its shared-memory budget: stages live
+// beyond the 48KB static limit opt in once per instantiation via
-// in dynamic smem, so budgets beyond the 48KB static limit opt in once per
+// cudaFuncSetAttribute. Templated on the kernel *value* (auto NTTP) so
-// instantiation via cudaFuncSetAttribute (see AGENTS.md "dynamic shared
+// every instantiation owns its own armed flag — same-signature kernels
-// memory"). Templated on the kernel *value* (auto NTTP) so every
+// must not share it. A failed opt-in arms nothing, so the launch below
-// instantiation owns its own armed flag — same-signature kernels must not
+// fails loudly through the caller's error checks.
 // share it (the attribute is per-function). A failed opt-in arms nothing, so
 // the launch below fails loudly through the caller's error checks instead of
 // silently running with an undersized stage buffer.
 template <auto Kernel, typename... Args>
 void launch_with_smem(int smem_bytes, dim3 grid, dim3 block,
                      cudaStream_t stream, Args... args) {
@@ -1051,25 +890,10 @@ void launch_with_smem(int smem_bytes, dim3 grid, dim3 block,
    Kernel<<<grid, block, smem_bytes, stream>>>(args...);
 }
-// Pre-quantized GEMM tile config: 128x128 CTA (8 warps x 64x32 warp tiles).
+// Padding-driven small-CTA rule: m or n <= 64 wastes half a 128-row CTA's
-// kK selects the K tile (32 / 64 / 128; larger kK halves the __syncthreads
+// MMA work, and a non-128-divisible shape drags its edge tiles through the
-// count per K and doubles the MMA work per stage at more smem per stage).
+// predicated generic path — when 64 divides both dims, the 64x64 CTA tiles
-// Stages is the cp.async pipeline depth (smem = Stages * (BM + BN) * kK
+// exactly and wins that band.
 // bytes for congruous layouts; deep pipelines are dynamic-smem backed, 1
 // CTA/SM past 48KB).
 // Crosswise operands always take load_crosswise_direct — the alternative
 // staging+transpose pipeline measured 15-20% slower everywhere probed
 // (contract k 2048..32768, DRAM-streaming B included) and was removed.
 // Padding-driven small-CTA rule. m <= 64 (and n <= 64 symmetric): a 128-row
 // CTA would waste half its MMA work on predicated-off rows. Divisibility:
 // a 128x128 CTA that is NOT exactly tiled (m or n not a multiple of 128)
 // runs its edge tiles on the predicated generic path, and with a single
 // in-flight wave the runtime is the slowest CTA — the edge tiles drag the
 // whole shape down (1088^3: 76T vs 93T with the 64x64 CTA, whose grid tiles
 // exactly and overlaps waves; measured sweep, perf 5.1). When 64 divides
 // both dims, the 64x64 small CTA wins the non-128-divisible band by
 // 23..67%.
 inline bool small_cta_padding(int64_t m, int64_t n) {
    if (m <= 64 || n <= 64) return true;
    const bool big_div = (m % 128 == 0) && (n % 128 == 0);
@@ -1078,23 +902,20 @@ inline bool small_cta_padding(int64_t m, int64_t n) {
 }
 // Launch configuration — a pure function of the problem (unit-testable
-// without a GPU call; the measured crossover rules live in plan_gemm's
+// without a GPU). Raster order is not a plan field: every canonical layout
-// comments). Raster order is not a plan field: every canonical layout
+// runs grouped raster; the plain-raster knob stays available through
-// runs grouped raster (see gemm); the plain-raster knob stays available
+// launch_plan's GroupRaster parameter for experiments.
 // through launch_plan's GroupRaster template parameter for experiments.
 struct Fp8GemmPlan {
    enum class Cta { kSmall64, kNarrow128x64, kBig128 };
    Cta cta;
    bool small_s3;  // kSmall64 only: cp.async pipeline depth (2 vs 3 stages)
 };
-// crosswise_ops: how many operands take the direct crosswise load (A
+// crosswise_ops counts the operands taking the direct crosswise load
-// ColMajor storage / B RowMajor storage, see Fp8GemmSmem). 0 = the
+// (A ColMajor / B RowMajor storage): 0 = dual-congruous NT, 1 = TN and the
-// dual-congruous NT problem, 1 = TN and the NN swap, 2 = TT. The layout
+// NN swap, 2 = TT. The layout shifts the crossovers: the small CTA hides
-// shifts the crossovers: the small CTA hides the crosswise LDG+PRMT
+// the crosswise LDG+PRMT latency far better, while the big CTA's operand
-// latency far better (more resident CTAs, deeper pipeline), while the big
+// reuse buys back load bandwidth the crosswise path does not traffic in.
 // CTA's operand reuse mostly buys back load bandwidth the crosswise path
 // does not traffic in.
 inline Fp8GemmPlan plan_gemm(const FP8Params& p, int crosswise_ops = 0) {
    const int64_t sm = device_sm_count();
    const int64_t tiles_128 =
@@ -1111,81 +932,56 @@ inline Fp8GemmPlan plan_gemm(const FP8Params& p, int crosswise_ops = 0) {
    // Padding rules first: predication waste beats any wave-fill effect.
    if (small_cta_padding(p.m, p.n)) return small(crosswise_ops > 0);
    if (crosswise_ops > 0) {
-        // Crosswise ladder (L20 measured, N=K=4096 band + squares): the
+        // Crosswise ladder (L20 measured): the small s3 CTA holds ~3/4 of
-        // narrow CTA never wins — below the wave band the small CTA beats
+        // the big CTA's per-SM throughput but tiles 4x finer, so it owns
-        // it (M=128: 82.8 vs 72.3T), above it the big CTA does. The small
+        // the whole sub-wave band and past it; the big CTA takes over once
-        // s3 CTA holds ~3/4 of the big CTA's per-SM throughput but tiles 4x
+        // its grid fills ~1.5 waves.
        // finer, so it owns the whole sub-wave band and past it: +15% at
        // M=256 (129.7 vs 113.1), +17% at M=384, +13% at 1024^3 (107.2 vs
        // 94.8). The big CTA takes over once its grid fills ~1.5 waves:
        // +18% at 1536^3 (134.4 vs 114.0), +20% at M=640 (161.2 vs 134.4);
        // the 1.5-wave boundary itself is within noise either way (M=512:
        // 133.8 vs 129.7 small; M=768: 135.8 vs 132.8 small).
        if (tiles_128 >= sm * 3 / 2) return big();
        return small(true);
    }
    if (tiles_128 >= sm) {
        // Wave band: pick by the wave-quantization cost ceil(tiles/sm) *
        // T_tile. The narrow tile carries half the big tile's MMA work at
-        // ~94% of its per-SM efficiency, i.e. T_narrow ~= 0.53 * T_big
+        // ~94% of its per-SM efficiency (T_narrow ~= 0.53 * T_big,
-        // (L20 measured; integer-scaled by 100 below). This formula
+        // integer-scaled by 100 below) — reproduces every measured
-        // reproduces every measured crossover: narrow wins the poor-fill
+        // crossover.
        // big grids (M=384: 134.3 vs 114.4, M=512: 171.7 vs 152.5, M=768:
        // 164.6 vs 153.7), the big CTA wins the well-filled ones (M=640:
        // 187.7 vs 167.3, M=1024: 202.5 vs 178.8, 4096^3: 210.7 vs 192.5).
        const int64_t tiles_narrow =
            (int64_t)p.batch * ((p.m + 127) / 128) * ((p.n + 63) / 64);
        const auto waves = [sm](int64_t tiles) { return (tiles + sm - 1) / sm; };
        if (waves(tiles_narrow) * 53 < waves(tiles_128) * 100) return narrow();
        return big();
    }
-    // Sub-wave band: the 128x64 narrow CTA fills the wave with N-tiles at
+    // Sub-wave band: the narrow CTA fills the wave with N-tiles at full
-    // full warp depth — measured sm_120, it beats the small CTA by +7..77%
+    // warp depth once its grid passes ~3/8 of a wave; below that the plain
-    // across the band once the narrow grid passes ~3/8 of a wave (128x4096
+    // 64x64 CTA's extra parallelism wins, and past ~5/8 of a wave of
-    // 132 vs 123T, 1024^3 174 vs 131T, 4096x384 242 vs 147T, 8192x128 233
+    // 128x128 tiles the big CTA's operand reuse wins instead.
    // vs 131T); below that fill the plain 64x64 CTA's extra parallelism
    // wins (2048x128: small 99 vs 87T). Past ~5/8 of a wave of 128x128
    // tiles the big CTA's operand reuse wins instead (each A element
    // multiplies 128 B columns in-CTA; L20 M=256: 143.7 vs 135.0 narrow,
    // 1024^3: 115.7 vs 113.3; forcing the 64x64 CTA there measured 2048^3
    // 123->171 TF on L20 and 164 vs 308T on sm_120).
    if (tiles_128 >= sm * 5 / 8) return big();
    const int64_t tiles_narrow =
        (int64_t)p.batch * ((p.m + 127) / 128) * ((p.n + 63) / 64);
    if (tiles_narrow >= sm * 3 / 8) return narrow();
-    // Full-ring small CTAs — ONE __syncthreads per k-tile, cuBLAS's barrier
+    // Full-ring small CTAs: the 24KB s2 variant keeps 4 CTAs/SM while the
-    // structure. Two depths by grid shape: the 24KB 3-slot s2 variant keeps
+    // whole grid stays resident; past that the 32KB s3 variant's deeper
-    // 4 CTAs/SM while the whole grid stays resident (<= one 3-CTA wave);
+    // pipeline wins on multi-wave grids.
    // past that the 32KB s3 variant's deeper cp.async pipeline wins on
    // multi-wave grids (measured 1280³: 107T vs 98T; sub-wave grids tie
    // within +-1%).
    const int64_t tiles_64 =
        (int64_t)p.batch * ((p.m + 63) / 64) * ((p.n + 63) / 64);
    return small(tiles_64 > sm * 3);
 }
-// Grid + launch for one concrete Policy — the only place a GEMM kernel
+// Grid + launch for one concrete Policy — the only place a GEMM kernel goes
-// goes to the wire.
+// to the wire.
 template <typename Policy>
 void launch_policy(const FP8Params& p, cudaStream_t stream) {
    using Traits = typename Policy::Traits;
-    constexpr int kBM = Traits::kBlockM;
+    dim3 grid((p.n + Traits::kBlockN - 1) / Traits::kBlockN,
-    constexpr int kBN = Traits::kBlockN;
+              (p.m + Traits::kBlockM - 1) / Traits::kBlockM, p.batch);
    dim3 grid((p.n + kBN - 1) / kBN, (p.m + kBM - 1) / kBM, p.batch);
    launch_with_smem<fp8_gemm_kernel<Policy>>(
-        Fp8GemmSmem<typename Policy::Traits, typename Policy::LayoutTagA,
+        Policy::kSmemBytes, grid, dim3(Traits::kCtaThreads), stream, p);
                    typename Policy::LayoutTagB>::kBytes,
        grid, dim3(Traits::kCtaThreads), stream, p);
 }
 // Plan -> Policy: the production-tuned configs. Big CTA: 128x128 of 8 warps
-// x 64x32, kK=64, 2-stage full ring, interior fast loop only for
+// x 64x32, kK=64, 2-stage full ring, fast loop only for dual-congruous
-// dual-congruous layouts (crosswise instantiations keep the single generic
+// layouts. Narrow: 128x64. Small CTA: 64x64 of 4 warps x 32x32, kK=64,
-// body — no dead second loop in their I-cache). Small CTA: 64x64 of 4 warps
+// kFastLoop always on.
 // x 32x32, kK=64, kFastLoop always on (the predication-free interior load
 // is where the small CTA's issue budget goes). Full rings everywhere — the
 // lean kStages-deep ring traded a second barrier for a 4th resident CTA and
 // measured slower (1280³ +5..9%), so the knob was removed.
 template <FP8Format Fmt, typename LayoutA, typename LayoutB, int GroupRaster>
 void launch_plan(const FP8Params& p, const Fp8GemmPlan& plan,
                 cudaStream_t stream) {
@@ -1221,16 +1017,13 @@ void launch_plan(const FP8Params& p, const Fp8GemmPlan& plan,
    }
 }
-// Pure problem rewrite: the dual-N-contiguous problem — trans_a/trans_b
+// Pure problem rewrite: the dual-N-contiguous problem (trans_a/trans_b both
-// both false — has no dedicated instantiation. It runs as its transpose
+// false) has no dedicated instantiation — it runs as its transpose
-// E[N][M] = B^T @ A^T (CUTLASS-sm90's is_swapAB): new A = B^T is
+// E[N][M] = B^T @ A^T (CUTLASS-sm90's is_swapAB) over swapped operands,
-// M'-contiguous (crosswise load), new B = A^T is K-contiguous (congruous
+// with p.out_transposed making the epilogue scatter into the caller's
 // load), and p.out_transposed makes the epilogue scatter into the caller's
 // [M][N] row-major buffer. The rewritten trans flags become the layout tags
-// the launcher instantiates. One instantiation fewer per (format,
+// the launcher instantiates; the NN path pays a scalar-store scatter, which
-// tile-config); the NN path pays a scalar-store scatter, which its rare
+// its rare usage makes the right trade.
 // usage (no LLM-linear operand pair is dual-N-contiguous) makes the right
 // trade.
 inline void canonicalize_gemm(FP8Params& p, bool& trans_a, bool& trans_b) {
    if (!trans_a && !trans_b) {
        FP8Params s = p;  // E = B^T * A^T: swap roles, M <-> N
@@ -1249,8 +1042,7 @@ inline void canonicalize_gemm(FP8Params& p, bool& trans_a, bool& trans_b) {
 }
 // Entry point: canonicalize the problem, plan the launch, wire the layout
-// tags through. (Every reachable tag combination is grouped-raster: the
+// tags through.
 // plain-raster knob stays available through launch_plan for experiments.)
 template <FP8Format Fmt>
 void gemm(FP8Params p, cudaStream_t stream, bool trans_a, bool trans_b) {
    canonicalize_gemm(p, trans_a, trans_b);
@@ -6,7 +6,6 @@
 #include <cstdint>
 #include <mutex>
 #include <tuple>
 #include <unordered_map>
 #include "../common/device.cuh"
@@ -46,32 +45,13 @@ void check_scale(const torch::Tensor& scale, const torch::Tensor& input) {
                "scale must be a CUDA float32 scalar on the input device");
 }
 void pack_gemm(FP8Params& p, const void* a, const void* b, void* output,
               const torch::Tensor& scale, int64_t m, int64_t n, int64_t k,
               int64_t a_ld, int64_t b_ld) {
    p.a_ptr = a;
    p.b_ptr = b;
    p.out_ptr = output;
    p.scale = scale.data_ptr<float>();
    p.m = static_cast<int>(m);
    p.n = static_cast<int>(n);
    p.k = static_cast<int>(k);
    p.a_ld = static_cast<int>(a_ld);
    p.b_ld = static_cast<int>(b_ld);
 }
 // Layout dispatch (the NN swap in canonicalize_gemm) and launch planning
 // (plan_gemm/launch_plan) live in gemm.cuh behind fp8::gemm — pure CUDA,
 // shared with the C test suite.
 // Inner-layout resolution for one GEMM operand. The user flag names the
-// math (0 = tensor's last two dims are [rows][contract], 1 = transposed);
+// math (0 = last two dims are [rows][contract], 1 = transposed); the
-// the storage may independently be a col-major view (.t() of a contiguous
+// storage may independently be a col-major view (.t() of a contiguous
 // buffer), which folds into the returned dispatch flag at zero copy — the
 // kernel's LayoutA/LayoutB tags cover both storages. m/n/k derive from the
-// user flag only; the fold never swaps them (see the layout table in
+// user flag only. Tensors whose inner dims are neither natural layout fall
-// gemm.cuh). Tensors whose inner dims are neither natural layout fall back
+// back to .contiguous().
 // to .contiguous().
 bool resolve_operand(const torch::Tensor& t_in, bool flag, int64_t& ld,
                     int64_t& batch_stride, torch::Tensor& storage) {
    torch::Tensor t = t_in;
@@ -89,43 +69,36 @@ bool resolve_operand(const torch::Tensor& t_in, bool flag, int64_t& ld,
    return flag ^ col_major;
 }
-// Dtype x format switch shared by both quantize kernels; Tiled selects the
+// Dtype dispatch over the unified quantize launcher.
-// transpose kernel (out_layout 1/2) over the vectorized elementwise one.
+template <bool Tiled, FP8Format Fmt>
-template <bool Tiled, FP8Format Fmt, typename InT>
+void launch_for_dtype(const torch::Tensor& x, const FP8QuantizeParams& p,
-void launch_one(const FP8QuantizeParams& p, cudaStream_t stream) {
+                      cudaStream_t stream) {
-    if constexpr (Tiled)
+    switch (x.scalar_type()) {
-        launch_fp8_quantize_tiled<Fmt, InT>(p, stream);
+    case torch::kHalf:
-    else
+        launch_fp8_quantize<Fmt, __half, Tiled>(p, stream);
-        launch_fp8_quantize<Fmt, InT>(p, stream);
+        break;
    case torch::kFloat32:
        launch_fp8_quantize<Fmt, float, Tiled>(p, stream);
        break;
    default:
        launch_fp8_quantize<Fmt, __nv_bfloat16, Tiled>(p, stream);
    }
 }
 template <bool Tiled>
 void launch_quantize_for(const torch::Tensor& x, const FP8QuantizeParams& p,
                         bool e5m2, cudaStream_t stream) {
    if (x.scalar_type() == torch::kHalf) {
    if (e5m2)
-            launch_one<Tiled, FP8Format::E5M2, __half>(p, stream);
+        launch_for_dtype<Tiled, FP8Format::E5M2>(x, p, stream);
    else
-            launch_one<Tiled, FP8Format::E4M3, __half>(p, stream);
+        launch_for_dtype<Tiled, FP8Format::E4M3>(x, p, stream);
    } else if (x.scalar_type() == torch::kFloat32) {
        if (e5m2)
            launch_one<Tiled, FP8Format::E5M2, float>(p, stream);
        else
            launch_one<Tiled, FP8Format::E4M3, float>(p, stream);
    } else {
        if (e5m2)
            launch_one<Tiled, FP8Format::E5M2, __nv_bfloat16>(p, stream);
        else
            launch_one<Tiled, FP8Format::E4M3, __nv_bfloat16>(p, stream);
    }
 }
 }  // namespace
-// Output-layout dispatch: 0 = [rows][cols] row-major (the historic 2-tuple
+// Output-layout dispatch: 0 = [rows][cols] row-major (2-tuple return),
-// return), 1 = transposed [cols][rows] only (2-tuple), 2 = both orientations
+// 1 = transposed [cols][rows] only (2-tuple), 2 = both orientations from a
-// from a single read of the input (3-tuple). Layouts 1/2 feed the NT GEMM
+// single read (3-tuple). Layouts 1/2 feed the NT GEMM fast path.
 // fast path from crosswise consumers (backward grad_x / grad_w).
 py::object quantize(torch::Tensor x, torch::Tensor scale, int64_t fmt,
                    int64_t layout) {
    TORCH_CHECK(x.is_cuda(), "CUDA tensors required");
@@ -219,8 +192,15 @@ torch::Tensor mm_fp8(torch::Tensor a, torch::Tensor b, torch::Tensor scale,
            ? torch::empty({batch, m, n}, a.options().dtype(torch::kBFloat16))
            : torch::empty({m, n}, a.options().dtype(torch::kBFloat16));
    FP8Params p;
-    pack_gemm(p, a_st.data_ptr(), b_st.data_ptr(), output.data_ptr(), scale,
+    p.a_ptr = a_st.data_ptr();
-              m, n, k, a_ld, b_ld);
+    p.b_ptr = b_st.data_ptr();
    p.out_ptr = output.data_ptr();
    p.scale = scale.data_ptr<float>();
    p.m = static_cast<int>(m);
    p.n = static_cast<int>(n);
    p.k = static_cast<int>(k);
    p.a_ld = static_cast<int>(a_ld);
    p.b_ld = static_cast<int>(b_ld);
    // Fused epilogue bias (bf16, broadcast over rows and batches). An
    // undefined or 0-element tensor keeps the plain scaled output.
    if (bias.defined() && bias.numel() > 0) {
@@ -243,9 +223,8 @@ torch::Tensor mm_fp8(torch::Tensor a, torch::Tensor b, torch::Tensor scale,
    return output;
 }
-// mm_fp8 binding: Python None and an omitted argument both mean "no bias"
+// mm_fp8 binding: Python None and an omitted argument both mean "no bias",
-// (resolved to an undefined tensor here, so every Python layer can pass its
+// so every Python layer can pass its bias argument through untouched.
 // bias argument through untouched instead of normalizing it host-side).
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("quantize", &quantize, py::arg("x"), py::arg("scale"),
          py::arg("fmt"), py::arg("layout") = 0);
@@ -1,10 +1,7 @@
 #pragma once
-// FP8 quantize device code — pure CUDA, no torch. Any float input element
+// FP8 quantize device code — pure CUDA, no torch: kernels take the
-// type (bf16 / fp16 / fp32) converts to E4M3 or E5M2 with a fused amax over
+// FP8QuantizeParams POD, format and input type ride on template parameters,
-// the raw (unscaled) values. Mirrors the GEMM file's split: kernels take the
+// and the launcher is shared by the torch binding and the C tests.
 // FP8QuantizeParams POD, formats and input types ride on template parameters,
 // and the launcher is a plain function usable from both the torch binding and
 // pure C tests.
 #include <cuda_bf16.h>
 #include <cuda_fp16.h>
@@ -18,8 +15,8 @@
 namespace astrai {
 namespace fp8 {
-// Input element type traits: one element -> float, and the vectorized
+// Input element type traits: one element -> float, and the unpack of one
-// unpack of one 16-byte load into kVecElems floats.
+// 16-byte load into kVecElems floats.
 template <typename InT>
 struct quant_in_traits;
@@ -31,12 +28,13 @@ struct quant_in_traits<__nv_bfloat16> {
    }
    static __device__ __forceinline__ void load_vec(const uint4& raw,
                                                    float* f) {
-        const unsigned w[4] = {raw.x, raw.y, raw.z, raw.w};
+        const __nv_bfloat162* b2 =
            reinterpret_cast<const __nv_bfloat162*>(&raw);
 #pragma unroll
        for (int j = 0; j < 4; ++j) {
-            f[2 * j] =
+            const float2 p = __bfloat1622float2(b2[j]);
-                __bfloat162float(__ushort_as_bfloat16(w[j] & 0xffffu));
+            f[2 * j] = p.x;
-            f[2 * j + 1] = __bfloat162float(__ushort_as_bfloat16(w[j] >> 16));
+            f[2 * j + 1] = p.y;
        }
    }
 };
@@ -65,15 +63,22 @@ struct quant_in_traits<float> {
    static __device__ __forceinline__ float to_float(float v) { return v; }
    static __device__ __forceinline__ void load_vec(const uint4& raw,
                                                    float* f) {
-        f[0] = __uint_as_float(raw.x);
+        const unsigned* w = reinterpret_cast<const unsigned*>(&raw);
-        f[1] = __uint_as_float(raw.y);
+#pragma unroll
-        f[2] = __uint_as_float(raw.z);
+        for (int j = 0; j < 4; ++j) f[j] = __uint_as_float(w[j]);
        f[3] = __uint_as_float(raw.w);
    }
 };
-// Convert one float pair to one packed fp8 pair. The stored bytes see
+// One float -> one fp8 byte (round-nearest-even + satfinite).
-// value * mult (round-nearest-even + satfinite).
+template <FP8Format Fmt>
 __device__ __forceinline__ uint8_t cvt_fp8(float v) {
    if constexpr (Fmt == FP8Format::E5M2)
        return __nv_fp8_e5m2(v).__x;
    else
        return __nv_fp8_e4m3(v).__x;
 }
 // One float pair -> one packed fp8x2 word (round-nearest-even + satfinite).
 template <FP8Format Fmt>
 __device__ __forceinline__ unsigned cvt_fp8x2(float a, float b) {
    constexpr __nv_fp8_interpretation_t kFmt =
@@ -82,23 +87,36 @@ __device__ __forceinline__ unsigned cvt_fp8x2(float a, float b) {
        make_float2(a, b), __NV_SATFINITE, kFmt));
 }
-// Quantize kernel: float input -> FP8 (E4M3 or E5M2), fused amax over raw
+// Block-wide amax reduce -> one atomic per block: warp-reduce, park one
-// values.
+// value per warp, thread 0 folds. kWarps must cover the block's warp count.
 template <int kWarps>
 __device__ __forceinline__ void publish_amax(float* amax, float v) {
    v = warp_reduce_max(v);
    __shared__ float slots[kWarps];
    const int tid = threadIdx.y * blockDim.x + threadIdx.x;
    if ((tid & 31) == 0) slots[tid >> 5] = v;
    __syncthreads();
    if (tid == 0) {
 #pragma unroll
        for (int w = 1; w < kWarps; ++w) v = fmaxf(v, slots[w]);
        atomic_max_float(amax, v);
    }
 }
 // Elementwise quantize kernel (out_layout 0): vectorized 16B loads -> fp8
 // stores, fused amax over raw values.
 template <FP8Format Fmt, typename InT>
 __global__ void fp8_quantize_kernel(FP8QuantizeParams p) {
    const float mult = *p.scale;
    const auto* x = static_cast<const InT*>(p.input_ptr);
-    void* x8 = p.output_ptr;
+    uint8_t* x8 = static_cast<uint8_t*>(p.output_ptr);
    float* amax = p.amax;
    float local_amax = 0.0f;
    const int64_t stride = (int64_t)blockDim.x * gridDim.x;
-    // Vectorized body: one 16B load -> kVecElems fp8 bytes per step (8
+    // One 16B load -> kVecElems bytes per step. Torch allocations are >=16B
-    // elements for 16-bit inputs, 4 for fp32). Torch allocations are >=16B
+    // aligned, so element 0 keeps the uint4 access natural; a misaligned
-    // aligned and the binding passes freshly allocated contiguous buffers,
+    // base (odd storage offset view) falls to the scalar tail via
-    // so element 0 keeps the uint4 access natural; a misaligned base
+    // total_vec = 0.
    // (contiguous view with an odd storage offset) falls back to the scalar
    // loop below via total_vec = 0.
    constexpr int kVecElems = quant_in_traits<InT>::kVecElems;
    const bool aligned =
        ((reinterpret_cast<uintptr_t>(x) |
@@ -125,8 +143,7 @@ __global__ void fp8_quantize_kernel(FP8QuantizeParams p) {
            packed[j] = (lo & 0xffffu) | (hi << 16);
        }
        if constexpr (kVecElems == 8)
-            reinterpret_cast<uint2*>(x8)[i] =
+            reinterpret_cast<uint2*>(x8)[i] = make_uint2(packed[0], packed[1]);
                make_uint2(packed[0], packed[1]);
        else
            reinterpret_cast<unsigned*>(x8)[i] = packed[0];
    }
@@ -136,51 +153,20 @@ __global__ void fp8_quantize_kernel(FP8QuantizeParams p) {
         i < p.total; i += stride) {
        const float v = quant_in_traits<InT>::to_float(x[i]);
        local_amax = fmaxf(local_amax, fabsf(v));
-        if constexpr (Fmt == FP8Format::E5M2) {
+        x8[i] = cvt_fp8<Fmt>(v * mult);
            reinterpret_cast<__nv_fp8_e5m2*>(x8)[i] =
                __nv_fp8_e5m2(v * mult);
        } else {
            reinterpret_cast<__nv_fp8_e4m3*>(x8)[i] =
                __nv_fp8_e4m3(v * mult);
        }
    }
    if (amax) {
        local_amax = warp_reduce_max(local_amax);
        __shared__ float slots[32];
        if ((threadIdx.x & 31) == 0) slots[threadIdx.x >> 5] = local_amax;
        __syncthreads();
        if (threadIdx.x == 0) {
            float v = 0.0f;
            for (int w = 0; w < (blockDim.x >> 5); ++w)
                v = fmaxf(v, slots[w]);
            atomic_max_float(amax, v);
        }
    }
    if (p.amax) publish_amax<8>(p.amax, local_amax);
 }
-template <FP8Format Fmt, typename InT>
+// Tiled transpose quantize (out_layout 1/2): reads the [rows][cols] input
-void launch_fp8_quantize(const FP8QuantizeParams& p, cudaStream_t stream) {
+// once and writes the fp8 bytes transposed ([cols][rows], so the contract
-    constexpr int kThreads = 256;
+// dim lands K-contiguous for NT GEMM operands) and, in mode 2, the row-major
-    // One block per 256 vectors; at least one block so the scalar tail of a
+// copy too. A 32x32 tile stages through shared memory: loads and writes
-    // tiny / misaligned tensor is still covered.
+// both stay coalesced, and the byte-wide staging is conflict-free — the +4
-    constexpr int kVecElems = quant_in_traits<InT>::kVecElems;
+// pad makes the store stride 9 words (coprime with the 32 banks) and the
-    int64_t blocks = (p.total / kVecElems + kThreads - 1) / kThreads;
+// read is a 32-byte broadcast segment. (A 64x64 split-half variant measured
-    if (blocks < 1) blocks = 1;
+// +21% L2-resident but -3..5% DRAM-bound; the real step mix ties, so the
-    fp8_quantize_kernel<Fmt, InT><<<blocks, kThreads, 0, stream>>>(p);
+// simpler tile stays.)
 }
 // Tiled transpose quantize (out_layout 1/2): reads the [rows][cols]
 // row-major input once and writes the fp8 bytes transposed ([cols][rows],
 // so the contract dim lands K-contiguous for NT GEMM operands) and, in
 // mode 2, the plain row-major copy too. A 32x32 tile stages through shared
 // memory: input-row-major loads and output writes both stay coalesced, and
 // the byte-wide staging is conflict-free — the +4 pad makes the store
 // stride 9 (words) coprime with the 32 banks and the load is a 32-byte
 // broadcast segment. A 64x64 split-half variant (16 elems/thread, paired
 // 2-byte scatter stores) measured +21% on L2-resident shapes but -3..5%
 // on the DRAM-bound ones that carry the training traffic (occupancy and
 // memory-level parallelism, not instruction count, gate the DRAM regime);
 // weighted by the real step's mix the two tie, so the simpler tile stays.
 template <FP8Format Fmt, typename InT>
 __global__ void fp8_quantize_tiled_kernel(FP8QuantizeParams p) {
    constexpr int kTile = 32;
@@ -201,10 +187,7 @@ __global__ void fp8_quantize_tiled_kernel(FP8QuantizeParams p) {
            const float v =
                quant_in_traits<InT>::to_float(x[(int64_t)(r + j) * p.cols + c]);
            local_amax = fmaxf(local_amax, fabsf(v));
-            if constexpr (Fmt == FP8Format::E5M2)
+            q[j] = cvt_fp8<Fmt>(v * mult);
                q[j] = __nv_fp8_e5m2(v * mult).__x;
            else
                q[j] = __nv_fp8_e4m3(v * mult).__x;
        }
    }
    if (p.out_layout == 2) {
@@ -218,9 +201,9 @@ __global__ void fp8_quantize_tiled_kernel(FP8QuantizeParams p) {
    for (int j = 0; j < 4; ++j) tile[threadIdx.x][threadIdx.y * 4 + j] = q[j];
    __syncthreads();
    // Transposed scatter: output element (c, r) lives at c * rows + r; r
-    // tracks threadIdx.x so each warp writes one contiguous run. The read
+    // tracks threadIdx.x so each warp writes one contiguous run. tile was
-    // swaps the staging indices — tile[col][row] was written, so the value
+    // written as tile[col][row], so input (r0+tx, c0+ty*4+j) reads back
-    // for input (r0+tx, c0+ty*4+j) sits at tile[ty*4+j][tx].
+    // from tile[ty*4+j][tx].
    uint8_t* out_t = static_cast<uint8_t*>(p.output_transposed_ptr);
 #pragma unroll
    for (int j = 0; j < 4; ++j) {
@@ -229,28 +212,27 @@ __global__ void fp8_quantize_tiled_kernel(FP8QuantizeParams p) {
            out_t[(int64_t)oc * p.rows + r0 + threadIdx.x] =
                tile[threadIdx.y * 4 + j][threadIdx.x];
    }
-    if (p.amax) {
+    if (p.amax) publish_amax<8>(p.amax, local_amax);
        local_amax = warp_reduce_max(local_amax);
        __shared__ float slots[8];
        // blockDim.x is 32, so warp id == threadIdx.y; only complete warps
        // exist (blockDim.y == 8).
        if (threadIdx.x == 0) slots[threadIdx.y] = local_amax;
        __syncthreads();
        if (threadIdx.x == 0 && threadIdx.y == 0) {
            float v = 0.0f;
            for (int w = 0; w < (int)blockDim.y; ++w) v = fmaxf(v, slots[w]);
            atomic_max_float(p.amax, v);
        }
    }
 }
-template <FP8Format Fmt, typename InT>
+// Unified quantize launcher: Tiled selects the transpose kernel (out_layout
-void launch_fp8_quantize_tiled(const FP8QuantizeParams& p,
+// 1/2) over the vectorized elementwise one.
-                               cudaStream_t stream) {
+template <FP8Format Fmt, typename InT, bool Tiled = false>
 void launch_fp8_quantize(const FP8QuantizeParams& p, cudaStream_t stream) {
    if constexpr (Tiled) {
        const dim3 grid((p.cols + 31) / 32, (p.rows + 31) / 32);
        if (grid.x == 0 || grid.y == 0) return;
        fp8_quantize_tiled_kernel<Fmt, InT>
            <<<grid, dim3(32, 8), 0, stream>>>(p);
    } else {
        constexpr int kThreads = 256;
        constexpr int kVecElems = quant_in_traits<InT>::kVecElems;
        // One block per 256 vectors; at least one block so a tiny or
        // misaligned tensor's scalar tail is still covered.
        int64_t blocks = (p.total / kVecElems + kThreads - 1) / kThreads;
        if (blocks < 1) blocks = 1;
        fp8_quantize_kernel<Fmt, InT><<<blocks, kThreads, 0, stream>>>(p);
    }
 }
 }  // namespace fp8