From bf239d194c6166260715bfd5dd22c191af4134e1 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Fri, 28 Aug 2026 16:45:21 +0800 Subject: [PATCH] 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; 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 --- csrc/kernels/fp8/common.h | 94 ++--- csrc/kernels/fp8/gemm.cuh | 660 ++++++++++++---------------------- csrc/kernels/fp8/ops.cu | 93 ++--- csrc/kernels/fp8/quantize.cuh | 180 +++++----- 4 files changed, 376 insertions(+), 651 deletions(-) diff --git a/csrc/kernels/fp8/common.h b/csrc/kernels/fp8/common.h index d09d6f9..342fb2d 100644 --- a/csrc/kernels/fp8/common.h +++ b/csrc/kernels/fp8/common.h @@ -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 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 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; }; diff --git a/csrc/kernels/fp8/gemm.cuh b/csrc/kernels/fp8/gemm.cuh index e3f03f4..2f5be7b 100644 --- a/csrc/kernels/fp8/gemm.cuh +++ b/csrc/kernels/fp8/gemm.cuh @@ -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 #include @@ -21,7 +20,7 @@ namespace fp8 { // m16n8k32 (see astrai::mma_shape::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 struct log2_const : log2_const<(N >> 1), Acc + 1> {}; template @@ -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 __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 +// 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 __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,73 +66,49 @@ 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; - 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. - const auto* src = operand + row * ld + k_base; - const bool chunk_aligned = (reinterpret_cast(src) & 15) == 0; + if constexpr (kInterior) { + const char* src = reinterpret_cast( + operand + (block_row + r) * ld + k_base + c0); + const uintptr_t dst = + reinterpret_cast(tile_at(tile, r, c0)); #pragma unroll - for (int j = 0; j < kCpt; ++j) { - const int c = c0 + j * 16; - T8* dst = tile_at(tile, r, c); - 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. + for (int j = 0; j < kCpt; ++j) + astrai::cp_async_16(reinterpret_cast(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 all chunks share the + // row base's alignment verdict. + const auto* src = operand + row * ld + k_base; + const bool chunk_aligned = (reinterpret_cast(src) & 15) == 0; #pragma unroll - for (int i = 0; i < 16; ++i) - dst[i] = - row_ok && k_base + c + i < contract ? src[c + i] : T8(0.0f); + for (int j = 0; j < kCpt; ++j) { + const int c = c0 + j * 16; + T8* dst = tile_at(tile, r, c); + if (row_ok && chunk_aligned && k_base + c + 15 < contract) { + astrai::cp_async_16(dst, src + c); + } else { + // 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 -__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( - operand + (block_row + r) * ld + k_base + c0); - const uintptr_t dst = reinterpret_cast(tile_at(tile, r, c0)); -#pragma unroll - for (int j = 0; j < kCpt; ++j) - astrai::cp_async_16(reinterpret_cast(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 struct PrefetchCarry; @@ -172,9 +138,8 @@ struct PrefetchCarry { (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 { }; // --------------------------------------------------------------------------- -// 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 __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(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 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 @@ -323,30 +274,22 @@ struct Fp8GemmPolicy { static constexpr int kGroupRaster = GroupRaster_; static constexpr bool kStreamOut = StreamOut_; static constexpr bool kFastLoop = FastLoop_; + using Smem = Fp8GemmSmem; // Flattened for __launch_bounds__, which takes no dependent type names. static constexpr int kCtaThreads = Traits::kCtaThreads; - static constexpr int kMinCtas = - Fp8GemmSmem::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 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(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 __device__ __forceinline__ void load_async(T8* a_stage, T8* b_stage, int64_t k_base) const { if constexpr (!kDirectA) - load_operand_tile( + load_operand_tile( a_stage, a, m, k, a_ld, tid, k_base, block_m * kBlockM); if constexpr (!kDirectB) - load_operand_tile( + load_operand_tile( 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( - a_stage, a, a_ld, tid, k_base, block_m * kBlockM); - if constexpr (!kDirectB) - load_operand_tile_interior( - 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,22 +424,18 @@ 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 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), - (int64_t)stage * kK); + load_async(a_stage_of(stage), b_stage_of(stage), + (int64_t)stage * kK); else load_async(a_stage_of(stage), b_stage_of(stage), (int64_t)stage * kK); @@ -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 __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 carry_a( a_base, kARing, kAStageBytes, a, a_ld, block_m * kBlockM, tid, kStages); PrefetchCarry 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(); // 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::value; // tile_at's shift const unsigned lswz = static_cast((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((a_row0 + rh8 * 8 + r7) * kK + ((rh16 ^ lswz) << 4)); } @@ -734,16 +607,12 @@ struct Fp8CollectiveMainloop { static_cast((r7 >> kShift) & (kChunks - 1)); return static_cast((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(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(dst), v); } else { *reinterpret_cast(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(&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; using Epilogue = Fp8CollectiveEpilogue; - // 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 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<<>>(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 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>( - Fp8GemmSmem::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 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 void gemm(FP8Params p, cudaStream_t stream, bool trans_a, bool trans_b) { canonicalize_gemm(p, trans_a, trans_b); diff --git a/csrc/kernels/fp8/ops.cu b/csrc/kernels/fp8/ops.cu index ca8a2ab..0d2587d 100644 --- a/csrc/kernels/fp8/ops.cu +++ b/csrc/kernels/fp8/ops.cu @@ -6,7 +6,6 @@ #include #include -#include #include #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(); - p.m = static_cast(m); - p.n = static_cast(n); - p.k = static_cast(k); - p.a_ld = static_cast(a_ld); - p.b_ld = static_cast(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 -void launch_one(const FP8QuantizeParams& p, cudaStream_t stream) { - if constexpr (Tiled) - launch_fp8_quantize_tiled(p, stream); - else - launch_fp8_quantize(p, stream); +// Dtype dispatch over the unified quantize launcher. +template +void launch_for_dtype(const torch::Tensor& x, const FP8QuantizeParams& p, + cudaStream_t stream) { + switch (x.scalar_type()) { + case torch::kHalf: + launch_fp8_quantize(p, stream); + break; + case torch::kFloat32: + launch_fp8_quantize(p, stream); + break; + default: + launch_fp8_quantize(p, stream); + } } template 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(p, stream); - else - launch_one(p, stream); - } else if (x.scalar_type() == torch::kFloat32) { - if (e5m2) - launch_one(p, stream); - else - launch_one(p, stream); - } else { - if (e5m2) - launch_one(p, stream); - else - launch_one(p, stream); - } + if (e5m2) + launch_for_dtype(x, p, stream); + else + launch_for_dtype(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(); + p.m = static_cast(m); + p.n = static_cast(n); + p.k = static_cast(k); + p.a_ld = static_cast(a_ld); + p.b_ld = static_cast(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); diff --git a/csrc/kernels/fp8/quantize.cuh b/csrc/kernels/fp8/quantize.cuh index 4642888..1f70159 100644 --- a/csrc/kernels/fp8/quantize.cuh +++ b/csrc/kernels/fp8/quantize.cuh @@ -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 #include @@ -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 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(&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 { 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(&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 +__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 __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 +__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 __global__ void fp8_quantize_kernel(FP8QuantizeParams p) { const float mult = *p.scale; const auto* x = static_cast(p.input_ptr); - void* x8 = p.output_ptr; - float* amax = p.amax; + uint8_t* x8 = static_cast(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::kVecElems; const bool aligned = ((reinterpret_cast(x) | @@ -125,8 +143,7 @@ __global__ void fp8_quantize_kernel(FP8QuantizeParams p) { packed[j] = (lo & 0xffffu) | (hi << 16); } if constexpr (kVecElems == 8) - reinterpret_cast(x8)[i] = - make_uint2(packed[0], packed[1]); + reinterpret_cast(x8)[i] = make_uint2(packed[0], packed[1]); else reinterpret_cast(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::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(v * mult); } + if (p.amax) publish_amax<8>(p.amax, local_amax); } -template -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::kVecElems; - int64_t blocks = (p.total / kVecElems + kThreads - 1) / kThreads; - if (blocks < 1) blocks = 1; - fp8_quantize_kernel<<>>(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 __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::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(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(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 -void launch_fp8_quantize_tiled(const FP8QuantizeParams& p, - cudaStream_t stream) { - const dim3 grid((p.cols + 31) / 32, (p.rows + 31) / 32); - if (grid.x == 0 || grid.y == 0) return; - fp8_quantize_tiled_kernel - <<>>(p); +// Unified quantize launcher: Tiled selects the transpose kernel (out_layout +// 1/2) over the vectorized elementwise one. +template +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 + <<>>(p); + } else { + constexpr int kThreads = 256; + constexpr int kVecElems = quant_in_traits::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<<>>(p); + } } } // namespace fp8