ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor: dedupe fp8 kernel helpers and trim comments

Browse Source 
- 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
This commit is contained in:
ViperEkura
2026-08-28 16:45:21 +08:00
parent 8a353117ea
commit bf239d194c
4 changed files with 376 additions and 651 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/kernels/fp8/common.h
+33 -61
View File
@@ -11,30 +11,23 @@
namespace astrai {
namespace fp8 {
// Compile-time FP8 format: E4M3 (forward / high precision, max 448) or
// E5M2 (gradient / large dynamic range, max 57344).
// Compile-time FP8 format: E4M3 (forward, max 448) or E5M2 (gradients,
// max 57344).
enum class FP8Format : int {
E4M3 = 0,
E5M2 = 1,
};
// Operand memory layouts as types (CUTLASS-style tags). The tag names the
// storage order of the raw buffer relative to the operand's canonical GEMM
// matrix — A is [M][K], B is [K][N]:
// A RowMajor = [M][K] storage (K-contiguous rows; the default)
// 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).
// Operand storage tags (CUTLASS-style) relative to the canonical matrices
// A [M][K] / B [K][N]: A RowMajor = [M][K] (default), A ColMajor = [K][M],
// B RowMajor = [K][N], B ColMajor = [N][K] (the nn.Linear weight). Selection
// is by type at compile time (see gemm.cuh's stage loads).
struct RowMajor {};
struct ColMajor {};
// Compile-time tile configuration, mirroring KernelTraits<HEAD_DIM, BC,
// WARPS, STAGES> in the attention kernels. `Fmt` selects the FP8 conversion
// and the MMA PTX mnemonic; the remaining parameters shape the CTA tile, the
// 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.
// Compile-time tile configuration, mirroring KernelTraits in the attention
// kernels: CTA tile, warp tile (WarpM x WarpN — e.g. 64x32 on the 128x128
// CTA, 32x32 on the 64x64 small CTA) and 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
// shared-memory budget is layout-aware (crosswise operands add K-major
// staging + a canonical buffer), so it lives in Fp8GemmSmem in gemm.cuh
// together with the resident-CTA hint for __launch_bounds__.
// Derived geometry: warp tiles tile the CTA. The smem budget is
// layout-aware, so it lives in Fp8GemmSmem (gemm.cuh).
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
// with fused amax.
// Quantize-kernel parameter POD: float input -> FP8 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;
// Transposed-output destination ([cols][rows]); the output-layout modes:
// 0 = row-major only (output_ptr; the vectorized elementwise kernel)
// 1 = transposed only (output_transposed_ptr; the tiled kernel)
// 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.
void* __restrict__ output_transposed_ptr = nullptr; // [cols][rows]
// Output layout: 0 = row-major only, 1 = transposed only, 2 = both from
// a single read. Modes 1/2 produce K-contiguous operands so crosswise
// consumers (backward grad_x / grad_w) route through the NT fast path.
int out_layout = 0;
const float* __restrict__ scale = nullptr;
float* __restrict__ amax = nullptr;
const float* __restrict__ scale = nullptr; // device multiplier
float* __restrict__ amax = nullptr; // raw-domain max out
// Element count (only the elementwise quantize kernel uses it); the
// tiled kernel views the same buffer as [rows][cols] row-major.
// Element count (elementwise kernel); the tiled kernel views the same
// 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
// fields it needs; buffers are raw pointers packed by the torch binding.
// Pointer members default to null so optional paths cannot hold garbage.
// through the kernels; each kernel touches only the fields it needs.
struct FP8Params {
// Inputs: a/b are FP8 for the pre-quantized path. Scales are
// quantization steps (device scalars).
// Optional bf16 bias broadcast over output rows (fused into the epilogue
// before the bf16 rounding, so it adds in fp32 — one rounding fewer than
// the separate out + bias elementwise kernel it replaces). Null disables.
// FP8 operands + output; scales are quantization steps (device
// scalars). Optional bf16 bias fuses into the epilogue (fp32 add before
// the single bf16 rounding); 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
// problems): the kernel computes E[N'][M'] over swapped operands and the
// epilogue scatters into the caller's [M][N] row-major buffer, so
// 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.
// NN-swap mode (canonicalize_gemm): the kernel computes the transposed
// problem and the epilogue scatters D[row][col] to out[col * p.m + row]
// in the caller's [M][N] buffer. Zero in the plain orientation.
int out_transposed = 0;
// Shapes. `int` covers every realistic LLM shape; the kernels promote
// to int64 for all pointer arithmetic.
int m, n, k;
int m, n, k; // int covers LLM shapes; kernels promote to int64
// Batched (bmm) geometry: grid.z slices step the operand/output pointers
// by these element strides (0 broadcasts the operand across batches).
// Batched (bmm) geometry: grid.z steps these element strides (0
// 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
// B. For a non-transposed operand the stride equals the contract dim;
// for a transposed operand it is the operand's own column count. The
// binding packs these so the kernel reads both buffers either naturally
// or transposed depending on the LayoutA/LayoutB tags (see gemm.cuh).
// Physical leading dims (row strides) of A and B; the binding packs
// them so the kernel reads each buffer naturally or transposed per the
// LayoutA/LayoutB tags.
int a_ld, b_ld;
};
csrc/kernels/fp8/gemm.cuh
+212 -420
View File
@@ -1,9 +1,8 @@
#pragma once
// FP8 GEMM device code — pure CUDA, no torch. Mirrors the attention kernel
// layout (attn_*_mma.cuh): kernels take the FP8Params POD, tile shape and
// FP8 format ride on compile-time template parameters, and launchers are
// plain functions usable from both the torch binding and pure C tests. The
// quantize kernel lives in quantize.cuh.
// FP8 GEMM device code — pure CUDA, no torch. Kernels take the FP8Params
// POD; tile shape, formats and layout tags ride on one Policy template
// parameter (CUTLASS-style), and launchers are plain functions shared by
// the torch binding and the C tests.
#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);
// instantiate it with the kernel's T8. Accumulates in-place: callers pass
// the same accumulator array as both `d` and `c`.
// The cp.async pipeline primitives (predicated 16-byte copy, commit_group,
// 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.
// Swizzled address inside a flat [rows * K] staging tile: the 16B chunk
// index is XORed with the row bits at [3, 3+log2(kChunks)) so a warp's
// ldmatrix fragment load (8 consecutive rows x 16B) hits all 32 banks
// exactly once; chunks stay contiguous, so cp.async staging 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-
// contiguous — the only cp.async-able shape for the canonical tile) into the
// flat [rows * K] shared tile via tile_at's swizzle. Crosswise operands go
// through load_crosswise_direct instead.
template <typename T8, int K, int RowsTile, int kThreads>
// Stage-load a CONGRUOUS operand (contract-contiguous storage — the only
// cp.async-able shape) into the flat [rows * K] swizzled tile. kInterior
// drops all predication: valid only for a fully interior CTA (whole rows,
// 16B-aligned base|ld, k_base + K <= contract — the fast_cta peel
// 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
// share the row base's alignment.
// k_base and every c are multiples of 16, so all chunks share the
// 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,55 +92,23 @@ 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] =
row_ok && k_base + c + i < contract ? src[c + i] : T8(0.0f);
}
}
}
}
// 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):
// per-thread (r, c0) of the interior copy — the same mapping
// load_operand_tile_interior uses — with the swizzled stage destination and
// the global source pointer both carried across k-tiles, so each prefetch
// chunk is one LDGSTS at [wr ^ (j << 4)] / [src + j*16] issued straight from
// registers.
//
// 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]].
// Loop-carried prefetch state for one congruous operand ring: per-thread
// (r, c0) mapping with the swizzled stage destination and global source
// pointer carried across k-tiles, so each prefetch chunk is one LDGSTS
// issued straight from registers. The guard is a property of the operand's
// layout, so it lives in the type: the false specialization (crosswise
// operand) is an empty no-op — no dead declarations, no if constexpr at
// the use sites.
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
// tail) zero-fills: src_size=0 reads nothing, and the destination is
// the slot compute(i-1) already released.
// Emit this thread's chunks for the current tile; pf false (loop tail)
// zero-fills into 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
// accumulation, BF16 or FP8 output. The input format follows Traits; the
// tile is compact (row = kK bytes) so MMA fragments read directly — no
// in-kernel transpose of the operands (the binding handles transposes).
// Pre-quantized GEMM kernel: FP8 A/B staged into shared memory, FP32
// accumulation, BF16 output. Operands materialize in the compact canonical
// [rows][kK] tile so MMA fragments read directly — no in-kernel transpose.
// ---------------------------------------------------------------------------
// 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
// for each of 16 rows), so they take this path. A staged variant
// (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.
// canonical tile (a 16B global run holds one contract byte for each of 16
// rows), so they take this path.
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
// each group covers 16 rows), and p*ld keeps the base 16B-aligned
// whenever ld is, so every run of a chunk shares one alignment verdict.
// r0 is a multiple of 16 and p*ld preserves alignment whenever ld has
// it, 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
// lives in its (i>>2)-th 32-bit register.
// [v0.b(i), v1.b(i), v2.b(i), v3.b(i)].
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
// path nor the direct-crosswise path needs a post-compute barrier (one
// __syncthreads per k-tile). (A lean kStages-deep ring traded that barrier
// 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).
// (i-1)%(kStages+1) — already consumed — so neither load path needs a
// post-compute barrier (one __syncthreads per k-tile). The 48KB static
// watermark picks the resident-CTA hint for __launch_bounds__.
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
// shape rides on Fp8GemmTraits and the operand layouts + scheduling knobs
// hang beside them — this is the single template parameter fp8_gemm_kernel
// (and both collectives) take, mirroring CUTLASS's kernel-policy
// consolidation.
// Kernel policy: one type per kernel instantiation (CUTLASS-style
// consolidation) — traits + layout tags + scheduling knobs, the single
// template parameter the kernel and both collectives take.
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 =
Fp8GemmSmem<Traits, LayoutA_, LayoutB_>::kMinCtas;
static constexpr int kMinCtas = Smem::kMinCtas;
static constexpr int kSmemBytes = Smem::kBytes;
};
// LayoutA / LayoutB tag the operands' storage (CUTLASS-style, see common.h):
// A RowMajor = [M][K] / ColMajor = [K][M]; B RowMajor = [K][N] /
// ColMajor = [N][K]. The kernel always computes
// out[m][n] = sum_p tileA[m][p] * tileB[n][p]
// with the tiles materialized in the canonical [M][kK] / [N][kK] layout, so the
// MMA fragments are read identically regardless of layout. The tags only
// 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]
// 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).
// LayoutA / LayoutB tag the operands' storage; the kernel always computes
// out[m][n] = sum_p tileA[m][p] * tileB[n][p] with tiles materialized in
// the canonical [M][kK] / [N][kK] layout, so the tags only change how the
// stage-load gathers from global memory. With p.out_transposed set (the
// swap dispatch for NN problems) the kernel runs the transposed problem
// E = B^T * A^T and the epilogue scatters D[m][n] = E[n][m]; bias then
// indexes D-cols, i.e. the kernel's rows.
//
// 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
// makes consecutive CTAs cover a group of kRasterGroup M-tiles before
// advancing along N, so all CTAs of one group share the same B column
// 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.
// (L2-friendly) raster — consecutive CTAs share one B column stripe — or
// plain N-fastest raster (kRasterGroup=0, the measured best for dX's
// crosswise-B layouts where grouping was neutral).
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,
// each warp computing (WarpM/16) x (WarpN/8) m16n8k32 MMAs (mt x nt).
// The default 128x128 CTA runs 8 warps of 64x32 (mt x nt = 4x4); the
// small-shape path uses 64x64 CTAs of 32x32 warps (cuBLAS-style) so more
// CTAs fit per SM (see launch_fp8_gemm).
// CTA = (BlockM/WarpM) x (BlockN/WarpN) warps, each warp computing
// kMt x kNt m16n8k32 MMAs. Rings rotate kStages+1 buffers (see
// Fp8GemmSmem) — one __syncthreads per k-tile.
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
// are congruous, whole-CTA, 16B-aligned and K has no tail, the mainloop
// runs a compile-time-specialized copy whose loads carry no predication
// — the per-chunk guards cost ~6 of ~100 instructions per warp per
// k-tile, and the small-CTA path is issue-bound there (measured
// +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.
// Interior-CTA peel (kFastLoop instantiations only): whole-CTA,
// 16B-aligned, K without tail — the mainloop then runs a compile-time
// specialized copy with no per-chunk predication (measured +4.5..10% on
// the issue-bound small CTA; the 128x128 CTA regressed, so only the
// small CTA opts in). The verdict is uniform per CTA.
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
// their canonical rings. Called after the post-compute barrier, alongside
// the commit.
// Asynchronous congruous loads for one k-tile: cp.async into the
// canonical rings; kFast selects the predication-free interior copy
// (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
// with full CTA rows, aligned (base | ld), k_base + kK <= k. fast_cta
// admits only congruous operands, so no crosswise fallback is needed.
__device__ __forceinline__ void load_async_fast(T8* a_stage, T8* b_stage,
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.
// Synchronous direct-crosswise loads for one k-tile. In the steady
// state this runs right after barrier 1, so the LDG latency and the
// PRMT transpose overlap the MMA phase instead of stalling the
// inter-barrier window (which dominated the dX/dW stall profile).
__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
// memory stage. The commit is unconditional: when K is shorter than the
// pipeline (tile_count < kStages) the skipped stages commit empty groups
// so the group sequence stays tile-indexed — the steady-state
// wait_group<kStages-1> below is then correct for every iteration and
// 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.
// Prime the pipeline: kStages committed groups, one per stage slot.
// The commit is unconditional — when K is shorter than the pipeline the
// skipped stages commit empty groups, so the group sequence stays
// tile-indexed and the steady-state wait count never needs a runtime
// dispatch.
__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
// copy runs predication-free loads; the generic copy keeps full
// predication. kFastLoop=false instantiates only the generic copy —
// codegen identical to the pre-peel kernel.
// Steady-state mainloop, compile-time specialized on kFast: the fast
// copy runs predication-free loads with loop-carried read/write
// pointers; the generic copy keeps full predication. kFastLoop=false
// 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
// PrefetchCarry; crosswise operands get the empty no-op type).
// Construction targets the first prefetched tile (kStages).
// Fast-path write carries: one per congruous operand (crosswise
// operands get the empty no-op type), targeting the first
// 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
// replaces the per-k-tile (tile % ring) * stage_bytes recomputation —
// its SASS form was a UIMAD.WIDE magic-division ladder, ~10
// uniform-pipe instructions per operand per k-tile (perf 6.2).
// iteration with an equality wrap — replaces the per-k-tile
// (tile % ring) * stage_bytes recomputation (a UIMAD.WIDE
// magic-division ladder in SASS).
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
// keep that invariant true for every iteration.
// when this fires; the tail's unconditional (possibly empty)
// 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 MMA in the 128x64-tile version (the kernel was LSU-issue-bound
// there). B fragments double-buffer across k_segs. kPairB folds the
// 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.
// per k_seg — 0.5 load instructions per MMA. B fragments
// double-buffer across k_segs; kPairB folds the two adjacent nt
// fragments of one pair into a single x4 (see b4_lane_off).
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
// is issued before the MMAs consuming row mt, so the LDS fixed
// latency hides behind tensor-pipe work (cuts the `wait` stall,
// ~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.)
// Software-pipelined A fragments: the ldmatrix.x4 for row mt+1 is
// issued before the MMAs consuming row mt, so the LDS latency hides
// behind tensor-pipe work. Costs 4 extra registers.
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
// emitted inside it.
// loads run after the MMA phase.
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
// predicated-off chunks' zero-fill lands in the slot compute(i-1)
// released — nothing reads it again before the epilogue drain).
// sequence so the fixed wait above stays correct.
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
// the cuBLAS SASS: one base register per operand per k_seg, every
// fragment offset an LDSM immediate — zero address arithmetic inside the
// MMA phase). The closure works because the XOR swizzle's source bits
// come only from the lane's row-within-matrix (r7): the 8- and 16-row
// fragment steps (nt*8, mt*16) never reach them, so
// addr(s, mt) = lane_base + mt*(16*kK) ^ (s<<5) [A, x4 fragment]
// addr(s, nt) = lane_base + nt*(8*kK) ^ (s<<5) [B, x2 fragment]
// 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).
// Per-lane ldmatrix fragment addressing (base-pair scheme, mirrored
// from the cuBLAS SASS): one base register per operand per k_seg,
// every fragment offset an LDSM immediate — zero address arithmetic
// inside the MMA phase. The XOR swizzle's source bits come only from
// the lane's row-within-matrix (r7), so the 8/16-row fragment steps
// never reach them and
// addr(s, mt) = lane_base + mt*(16*kK) ^ (s<<5) [A, x4]
// addr(s, nt) = lane_base + nt*(8*kK) ^ (s<<5) [B, x2 / x4]
__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
// slot's converted base once per k-tile). A's fragment row carries
// the +8-row half (rh8) and the +1-chunk half (rh16) — matching the
// m16n8k32 operand layouts above.
// Stage-relative, loop-invariant per-lane base; A's fragment row
// carries the +8-row (rh8) and +1-chunk (rh16) halves.
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
// the two adjacent nt fragments — 2 x4 instead of 4 x2 per k_seg (12
// LDSM per k-tile instead of 16). Lane contract: lanes 0-7 address
// rows n0..n7 chunk c, lanes 8-15 rows n0..n7 chunk c+1, lanes 16-23
// rows n8..n15 chunk c, lanes 24-31 rows n8..n15 chunk c+1; regs
// {r0,r1} are the even nt's k-halves, {r2,r3} the odd nt's. The +8-row
// 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.
// x4-paired B loads: one ldmatrix.x4 feeds the two adjacent nt
// fragments. Lane contract: lanes 0-7 address rows n0..n7 chunk c,
// lanes 8-15 rows n0..n7 chunk c+1, lanes 16-23 rows n8..n15 chunk c,
// lanes 24-31 rows n8..n15 chunk c+1. The +8-row step never reaches
// the swizzle source bits for kK <= 64; kK=128 swizzles on row[2:0]
// where +8 flips bits, so 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
// kNt ldmatrix.x2 from the seg's base-pair address. frag2/frag4 are the
// flat bases of one b_frag / b_frag4 buffer (the unused one of the pair
// is never touched).
// double-buffer's next-seg fill. frag2/frag4 are the flat bases of one
// b_frag / b_frag4 buffer (the unused one 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.
// Plain orientation: D-local, kBlockM rows of kBlockN elems.
// Out-transposed (swap dispatch): the tile stages D-local rows over the
// swapped problem, so it has kBlockN rows of kBlockM — rows and row
// length trade places. Both row-chunk counts are powers of two, keeping
// the 16B-chunk XOR swizzle well-defined.
// Swizzled address of one 16B chunk (row r, chunk c) of the staged
// tile. Plain orientation: kBlockM rows of kBlockN elems; out-
// transposed (swap dispatch): rows and row length trade places. Both
// row-chunk counts are powers of two, keeping the 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
// epilogue goes through the operand shared memory: the A/B rings are
// dead once the mainloop ends, so their space stages the output tile
// (kBlockM x kBlockN bf16, always <= the ring budget). Threads first
// scatter their accumulators into the tile (STS.32 of bf16x2 pairs), a
// barrier makes the tile coherent, then the whole CTA copies it out in
// 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.
// Scatter the accumulators into the staging tile: the operand rings are
// dead once the mainloop ends, so their space stages the bf16 output
// tile. Threads scatter (STS.32 of bf16x2 pairs), a barrier makes the
// tile coherent, then the whole CTA copies it out in fully-coalesced
// 16B chunks. The 16B-chunk XOR swizzle keeps both the scatter and the
// gather conflict-free.
__device__ __forceinline__ void stage(float acc[kNt][kMt][4]) const {
// Fused bias (idea B): added to the fp32 accumulator before the
// single bf16 rounding — one fewer rounding than the out + bias
// elementwise pass this replaces, and no extra kernel launch / m*n
// round-trip. The per-lane loads (2 per nt, kMt-times re-read) are
// 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.
// Fused bias: added to the fp32 accumulator before the single bf16
// rounding. The per-lane loads are L1 broadcasts; rows past the
// edge skip the load (their smem slots never copy out). Under
// out_transposed the bias indexes D-cols = the kernel's rows.
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
// inside one 16B chunk.
// g+8 of the m16n8 output, columns tig*2/tig*2+1 inside
// 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
// (4 per (nt, mt) vs the packed bf16x2 pair — the swap path is
// the rare NN layout); the row swizzle keeps the quad's stores
// bank-spread. OOB elements store dead lanes of the tile, never
// copied out.
// The acc pair spans two staged rows, so these are scalar
// stores (the swap path is the rare NN layout). 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
// (D[row][col] = out[row * p.m + col]).
// the row length is kernel m', so row/stride flip to the swapped dims.
__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
// L20 squares and -3..4% on rects (the evict-first
// policy hurts more than the L2 B-tile protection helps
// at these sizes); kept as a template knob for other
// SKUs. Default off.
// Evict-first streaming store knob: neutral on L20
// squares, -3..4% on rects; kept for other SKUs.
__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):
// ldmatrix reads whole 16B chunks through the same mapping the staging
// 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.
// Stages live in dynamic shared memory so deep pipelines (> 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
// output staging: the loop's last commits (possibly only zero-filling
// predicated-off chunks) are nobody's wait target anymore.
// Drain the pipeline before the epilogue reclaims the operand rings.
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
// device-adaptive dispatch thresholds.
// every writer stores the same value).
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: stages live
// in dynamic smem, so budgets beyond the 48KB static limit opt in once per
// instantiation via cudaFuncSetAttribute (see AGENTS.md "dynamic shared
// memory"). Templated on the kernel *value* (auto NTTP) so every
// instantiation owns its own armed flag — same-signature kernels must not
// 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.
// Launch one kernel instantiation with its shared-memory budget: budgets
// beyond the 48KB static limit opt in once per instantiation via
// cudaFuncSetAttribute. Templated on the kernel *value* (auto NTTP) so
// every instantiation owns its own armed flag — same-signature kernels
// must not share it. A failed opt-in arms nothing, so the launch below
// fails loudly through the caller's error checks.
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).
// kK selects the K tile (32 / 64 / 128; larger kK halves the __syncthreads
// count per K and doubles the MMA work per stage at more smem per stage).
// Stages is the cp.async pipeline depth (smem = Stages * (BM + BN) * kK
// 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%.
// Padding-driven small-CTA rule: m or n <= 64 wastes half a 128-row CTA's
// MMA work, and a non-128-divisible shape drags its edge tiles through the
// predicated generic path — when 64 divides both dims, the 64x64 CTA tiles
// exactly and wins that band.
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
// comments). Raster order is not a plan field: every canonical layout
// runs grouped raster (see gemm); the plain-raster knob stays available
// through launch_plan's GroupRaster template parameter for experiments.
// without a GPU). Raster order is not a plan field: every canonical layout
// runs grouped raster; the plain-raster knob stays available through
// launch_plan's GroupRaster 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
// ColMajor storage / B RowMajor storage, see Fp8GemmSmem). 0 = the
// dual-congruous NT problem, 1 = TN and the NN swap, 2 = TT. The layout
// shifts the crossovers: the small CTA hides the crosswise LDG+PRMT
// latency far better (more resident CTAs, deeper pipeline), while the big
// CTA's operand reuse mostly buys back load bandwidth the crosswise path
// does not traffic in.
// crosswise_ops counts the operands taking the direct crosswise load
// (A ColMajor / B RowMajor storage): 0 = dual-congruous NT, 1 = TN and the
// NN swap, 2 = TT. The layout shifts the crossovers: the small CTA hides
// the crosswise LDG+PRMT latency far better, while the big CTA's operand
// reuse 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
// narrow CTA never wins — below the wave band the small CTA beats
// it (M=128: 82.8 vs 72.3T), above it the big CTA does. The small
// s3 CTA holds ~3/4 of the big CTA's per-SM throughput but tiles 4x
// 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).
// Crosswise ladder (L20 measured): the small s3 CTA holds ~3/4 of
// the big CTA's per-SM throughput but tiles 4x finer, so it owns
// the whole sub-wave band and past it; the big CTA takes over once
// its grid fills ~1.5 waves.
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
// (L20 measured; integer-scaled by 100 below). This formula
// reproduces every measured crossover: narrow wins the poor-fill
// 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).
// ~94% of its per-SM efficiency (T_narrow ~= 0.53 * T_big,
// integer-scaled by 100 below) — reproduces every measured
// crossover.
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
// full warp depth — measured sm_120, it beats the small CTA by +7..77%
// across the band once the narrow grid passes ~3/8 of a wave (128x4096
// 132 vs 123T, 1024^3 174 vs 131T, 4096x384 242 vs 147T, 8192x128 233
// 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).
// Sub-wave band: the narrow CTA fills the wave with N-tiles at full
// warp depth once its grid passes ~3/8 of a wave; below that the plain
// 64x64 CTA's extra parallelism wins, and past ~5/8 of a wave of
// 128x128 tiles the big CTA's operand reuse wins instead.
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
// structure. Two depths by grid shape: the 24KB 3-slot s2 variant keeps
// 4 CTAs/SM while the whole grid stays resident (<= one 3-CTA wave);
// 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%).
// Full-ring small CTAs: the 24KB s2 variant keeps 4 CTAs/SM while the
// whole grid stays resident; past that the 32KB s3 variant's deeper
// pipeline wins on multi-wave grids.
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
// goes to the wire.
// Grid + launch for one concrete Policy — the only place a GEMM kernel goes
// 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;
constexpr int kBN = Traits::kBlockN;
dim3 grid((p.n + kBN - 1) / kBN, (p.m + kBM - 1) / kBM, p.batch);
dim3 grid((p.n + Traits::kBlockN - 1) / Traits::kBlockN,
(p.m + Traits::kBlockM - 1) / Traits::kBlockM, p.batch);
launch_with_smem<fp8_gemm_kernel<Policy>>(
Fp8GemmSmem<typename Policy::Traits, typename Policy::LayoutTagA,
typename Policy::LayoutTagB>::kBytes,
grid, dim3(Traits::kCtaThreads), stream, p);
Policy::kSmemBytes, 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
// dual-congruous layouts (crosswise instantiations keep the single generic
// body — no dead second loop in their I-cache). Small CTA: 64x64 of 4 warps
// 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.
// x 64x32, kK=64, 2-stage full ring, fast loop only for dual-congruous
// layouts. Narrow: 128x64. Small CTA: 64x64 of 4 warps x 32x32, kK=64,
// kFastLoop always on.
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
// both 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
// M'-contiguous (crosswise load), new B = A^T is K-contiguous (congruous
// load), and p.out_transposed makes the epilogue scatter into the caller's
// Pure problem rewrite: the dual-N-contiguous problem (trans_a/trans_b both
// false) has no dedicated instantiation — it runs as its transpose
// E[N][M] = B^T @ A^T (CUTLASS-sm90's is_swapAB) over swapped operands,
// with p.out_transposed making 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,
// tile-config); the NN path pays a scalar-store scatter, which its rare
// usage (no LLM-linear operand pair is dual-N-contiguous) makes the right
// trade.
// the launcher instantiates; the NN path pays a scalar-store scatter, which
// its rare usage 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
// plain-raster knob stays available through launch_plan for experiments.)
// tags through.
template <FP8Format Fmt>
void gemm(FP8Params p, cudaStream_t stream, bool trans_a, bool trans_b) {
canonicalize_gemm(p, trans_a, trans_b);
csrc/kernels/fp8/ops.cu
+34 -55
View File
@@ -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);
// the storage may independently be a col-major view (.t() of a contiguous
// math (0 = last two dims are [rows][contract], 1 = transposed); the
// 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
// gemm.cuh). Tensors whose inner dims are neither natural layout fall back
// to .contiguous().
// user flag only. Tensors whose inner dims are neither natural layout fall
// back 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
// transpose kernel (out_layout 1/2) over the vectorized elementwise one.
template <bool Tiled, FP8Format Fmt, typename InT>
void launch_one(const FP8QuantizeParams& p, cudaStream_t stream) {
if constexpr (Tiled)
launch_fp8_quantize_tiled<Fmt, InT>(p, stream);
else
launch_fp8_quantize<Fmt, InT>(p, stream);
// Dtype dispatch over the unified quantize launcher.
template <bool Tiled, FP8Format Fmt>
void launch_for_dtype(const torch::Tensor& x, const FP8QuantizeParams& p,
cudaStream_t stream) {
switch (x.scalar_type()) {
case torch::kHalf:
launch_fp8_quantize<Fmt, __half, Tiled>(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);
} 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);
}
launch_for_dtype<Tiled, FP8Format::E4M3>(x, p, stream);
}
} // namespace
// Output-layout dispatch: 0 = [rows][cols] row-major (the historic 2-tuple
// return), 1 = transposed [cols][rows] only (2-tuple), 2 = both orientations
// from a single read of the input (3-tuple). Layouts 1/2 feed the NT GEMM
// fast path from crosswise consumers (backward grad_x / grad_w).
// Output-layout dispatch: 0 = [rows][cols] row-major (2-tuple return),
// 1 = transposed [cols][rows] only (2-tuple), 2 = both orientations from a
// single read (3-tuple). Layouts 1/2 feed the NT GEMM fast path.
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,
m, n, k, a_ld, b_ld);
p.a_ptr = a_st.data_ptr();
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"
// (resolved to an undefined tensor here, so every Python layer can pass its
// bias argument through untouched instead of normalizing it host-side).
// mm_fp8 binding: Python None and an omitted argument both mean "no bias",
// so every Python layer can pass its bias argument through untouched.
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("quantize", &quantize, py::arg("x"), py::arg("scale"),
py::arg("fmt"), py::arg("layout") = 0);
csrc/kernels/fp8/quantize.cuh
+77 -95
View File
@@ -1,10 +1,7 @@
#pragma once
// FP8 quantize device code — pure CUDA, no torch. Any float input element
// type (bf16 / fp16 / fp32) converts to E4M3 or E5M2 with a fused amax over
// the raw (unscaled) values. Mirrors the GEMM file's split: kernels take the
// 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.
// FP8 quantize device code — pure CUDA, no torch: kernels take the
// FP8QuantizeParams POD, format and input type ride on template parameters,
// and the launcher is shared by the torch binding and the 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
// unpack of one 16-byte load into kVecElems floats.
// Input element type traits: one element -> float, and the unpack of one
// 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] =
__bfloat162float(__ushort_as_bfloat16(w[j] & 0xffffu));
f[2 * j + 1] = __bfloat162float(__ushort_as_bfloat16(w[j] >> 16));
const float2 p = __bfloat1622float2(b2[j]);
f[2 * j] = p.x;
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);
f[1] = __uint_as_float(raw.y);
f[2] = __uint_as_float(raw.z);
f[3] = __uint_as_float(raw.w);
const unsigned* w = reinterpret_cast<const unsigned*>(&raw);
#pragma unroll
for (int j = 0; j < 4; ++j) f[j] = __uint_as_float(w[j]);
}
};
// Convert one float pair to one packed fp8 pair. The stored bytes see
// value * mult (round-nearest-even + satfinite).
// One float -> one fp8 byte (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
// values.
// Block-wide amax reduce -> one atomic per block: warp-reduce, park one
// 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;
float* amax = p.amax;
uint8_t* x8 = static_cast<uint8_t*>(p.output_ptr);
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
// elements for 16-bit inputs, 4 for fp32). Torch allocations are >=16B
// aligned and the binding passes freshly allocated contiguous buffers,
// so element 0 keeps the uint4 access natural; a misaligned base
// (contiguous view with an odd storage offset) falls back to the scalar
// loop below via total_vec = 0.
// One 16B load -> kVecElems bytes per step. Torch allocations are >=16B
// aligned, so element 0 keeps the uint4 access natural; a misaligned
// base (odd storage offset view) falls to the scalar tail 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] =
make_uint2(packed[0], packed[1]);
reinterpret_cast<uint2*>(x8)[i] = 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) {
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);
}
x8[i] = cvt_fp8<Fmt>(v * mult);
}
if (p.amax) publish_amax<8>(p.amax, local_amax);
}
template <FP8Format Fmt, typename InT>
void launch_fp8_quantize(const FP8QuantizeParams& p, cudaStream_t stream) {
constexpr int kThreads = 256;
// One block per 256 vectors; at least one block so the scalar tail of a
// tiny / misaligned tensor is still covered.
constexpr int kVecElems = quant_in_traits<InT>::kVecElems;
int64_t blocks = (p.total / kVecElems + kThreads - 1) / kThreads;
if (blocks < 1) blocks = 1;
fp8_quantize_kernel<Fmt, InT><<<blocks, kThreads, 0, stream>>>(p);
}
// 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.
// Tiled transpose quantize (out_layout 1/2): reads the [rows][cols] 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 row-major
// copy too. A 32x32 tile stages through shared memory: loads and 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
// read is a 32-byte broadcast segment. (A 64x64 split-half variant measured
// +21% L2-resident but -3..5% DRAM-bound; the real step mix ties, 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] = __nv_fp8_e5m2(v * mult).__x;
else
q[j] = __nv_fp8_e4m3(v * mult).__x;
q[j] = cvt_fp8<Fmt>(v * mult);
}
}
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
// swaps the staging indices — tile[col][row] was written, so the value
// for input (r0+tx, c0+ty*4+j) sits at tile[ty*4+j][tx].
// tracks threadIdx.x so each warp writes one contiguous run. tile was
// written as tile[col][row], so input (r0+tx, c0+ty*4+j) reads back
// 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) {
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);
}
}
if (p.amax) publish_amax<8>(p.amax, local_amax);
}
template <FP8Format Fmt, typename InT>
void launch_fp8_quantize_tiled(const FP8QuantizeParams& p,
cudaStream_t stream) {
// Unified quantize launcher: Tiled selects the transpose kernel (out_layout
// 1/2) over the vectorized elementwise one.
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