diff --git a/astrai/extension/ops/fp8.py b/astrai/extension/ops/fp8.py index ad34e48..d0e55e1 100644 --- a/astrai/extension/ops/fp8.py +++ b/astrai/extension/ops/fp8.py @@ -127,6 +127,18 @@ def quantize( ``scale`` is the quantization multiplier (device scalar); ``fmt`` selects E4M3 or E5M2. ``amax`` is a fresh 1-element float32 tensor. """ + # Hot-path bypass of the torch.library dispatch (~5us/call, ~40% of a + # 512-wide GEMM): real CUDA tensors of a supported dtype go straight to + # the extension. Fake/subclass tensors and non-CUDA inputs keep the + # custom_op route so torch.compile / meta / fake-tensor tracing and the + # CPU fallback behave exactly as before. + if ( + type(x) is torch.Tensor + and x.is_cuda + and x.dtype in _QUANT_INPUT_DTYPES + and fmt in _FMT_TO_INT + ): + return get_module("fp8_ops").quantize(x, scale, _FMT_TO_INT[fmt]) return fp8_quantize(x, scale, _fmt_int(fmt)) @@ -143,4 +155,12 @@ def mm_fp8( combined dequantization scale. The result is BF16; FP8 output is a separate quantize operation. """ + # Same hot-path bypass as quantize(): the binding's TORCH_CHECKs keep + # validation identical on the direct route. + if ( + type(a) is torch.Tensor + and a.is_cuda + and a.dtype in (torch.float8_e4m3fn, torch.float8_e5m2) + ): + return get_module("fp8_ops").mm_fp8(a, b, scale, int(trans_a), int(trans_b)) return fp8_gemm(a, b, scale, trans_a, trans_b) diff --git a/csrc/kernels/fp8/common.h b/csrc/kernels/fp8/common.h index 68e09b3..0163bc7 100644 --- a/csrc/kernels/fp8/common.h +++ b/csrc/kernels/fp8/common.h @@ -48,25 +48,36 @@ using transpose_layout_t = typename transpose_layout::type; // 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 and -// the cp.async pipeline depth. -template +// 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. +template struct Fp8GemmTraits { static constexpr FP8Format kFormat = Fmt; static constexpr int kBlockM = BlockM; static constexpr int kBlockN = BlockN; static constexpr int kK = K; static constexpr int kStages = Stages; + static constexpr int kWarpM = WarpM; + static constexpr int kWarpN = WarpN; static constexpr bool kIsE5M2 = (Fmt == FP8Format::E5M2); static constexpr __nv_fp8_interpretation_t kNvFormat = kIsE5M2 ? __NV_E5M2 : __NV_E4M3; static constexpr float kFp8Max = kIsE5M2 ? 57344.0f : 448.0f; - // Derived launch geometry: 64x32 warp tiles give the CTA thread count. - // 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__. - static constexpr int kCtaThreads = (BlockM / 64) * (BlockN / 32) * 32; + // 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__. + static constexpr int kWarpsM = BlockM / WarpM; + static constexpr int kWarpsN = BlockN / WarpN; + static constexpr int kCtaThreads = kWarpsM * kWarpsN * 32; + static_assert(kWarpsM * WarpM == BlockM && kWarpsN * WarpN == BlockN, + "warp tiles must exactly tile the CTA"); + static_assert(WarpM % 16 == 0 && WarpN % 8 == 0, + "warp tile must be a multiple of the m16n8 MMA shape"); }; // Quantize-kernel parameter POD: float input (bf16 / fp16 / fp32) -> FP8 diff --git a/csrc/kernels/fp8/gemm.cuh b/csrc/kernels/fp8/gemm.cuh index 9eab35c..2dc604f 100644 --- a/csrc/kernels/fp8/gemm.cuh +++ b/csrc/kernels/fp8/gemm.cuh @@ -290,16 +290,19 @@ transpose_crosswise_region(T8* tile, const T8* staging, int idx, int quad0) { } } -// Layout-aware shared-memory budget and occupancy hint. A congruous operand -// needs its kStages rotating canonical buffers; a staged-crosswise operand -// (crosswise B with kBStaged) needs kStages K-major staging buffers plus ONE -// canonical buffer (rewritten every tile by the in-kernel transpose); a -// direct-crosswise operand rotates kStages+1 canonical buffers so its load -// can run ahead of the compute phase (see the kernel's pipelining note). +// Layout-aware shared-memory budget and occupancy hint. Canonic rings hold +// kStages+1 buffers (LeanRing=false): the load for tile i+kStages targets +// slot (i-1)%(kStages+1) — already consumed — so the pure-congruous path +// needs no post-compute barrier (one __syncthreads per k-tile). LeanRing +// keeps the ring at kStages buffers for small CTAs whose occupancy comes +// from more resident CTAs (less smem) rather than a deeper rotation; it +// brings back barrier 4. A staged-crosswise B always costs kStages K-major +// staging buffers + one canonical buffer. // The 48KB static-smem watermark picks the resident-CTA hint for // __launch_bounds__ (sm_89: 100KB smem per SM, so two CTAs fit while each // stays within the static budget). -template +template struct Fp8GemmSmem { // Crosswise = the stage-load's view: A's tag directly, B's transposed. // A-crosswise always loads direct (L2-typical activations); B-crosswise @@ -309,11 +312,17 @@ struct Fp8GemmSmem { static constexpr bool kBStagePath = kCrossB && StagedB; static constexpr bool kDirectA = kCrossA; static constexpr bool kDirectB = kCrossB && !kBStagePath; + // LeanRing shrinks only the congruous (async) operand rings; a direct + // operand's ring stays kStages+1 deep (see the kernel's ring note). + static constexpr int kARing = kDirectA ? Traits::kStages + 1 + : Traits::kStages + !LeanRing; + static constexpr int kBRing = + kBStagePath ? Traits::kStages + 1 + : (kDirectB ? Traits::kStages + 1 + : Traits::kStages + !LeanRing); static constexpr int kBytes = - (kDirectA ? Traits::kStages + 1 : Traits::kStages) * - Traits::kBlockM * Traits::kK + - (kDirectB || kBStagePath ? Traits::kStages + 1 : Traits::kStages) * - Traits::kBlockN * Traits::kK; + kARing * Traits::kBlockM * Traits::kK + + kBRing * Traits::kBlockN * Traits::kK; static constexpr int kMinCtas = kBytes <= 48 * 1024 ? 2 : 1; }; @@ -331,10 +340,10 @@ struct Fp8GemmSmem { // exists for small-M calls: m <= 64 wastes half of every 128-row CTA, so the // launcher dispatches to it there (see launch_fp8_gemm). template + bool kBStaged = true, bool kLeanRing = false> __global__ void __launch_bounds__(Traits::kCtaThreads, Fp8GemmSmem::kMinCtas) + kBStaged, kLeanRing>::kMinCtas) fp8_gemm_kernel(FP8Params p) { using T8 = std::conditional_t; constexpr int kBlockM = Traits::kBlockM; @@ -357,19 +366,25 @@ __global__ void __launch_bounds__(Traits::kCtaThreads, // shared memory so deep pipelines (kStages * (kBlockM + kBlockN) * kK > // 48KB static limit) opt in via cudaFuncSetAttribute in the launcher. extern __shared__ __align__(16) char fp8_gemm_smem[]; - // Per operand: congruous = kStages rotating canonical buffers; direct- - // crosswise = kStages+1 of them (the load for tile i+kStages targets - // buffer (i-1)%(kStages+1) — the one compute(i-1) finished reading at - // the previous barrier — so it issues right after barrier 1 and its - // global-load latency overlaps the MMA phase below); staged-crosswise - // (B) = kStages K-major staging buffers (filled by cp.async, one per - // tile in flight) followed by one canonical buffer the per-tile - // transpose rewrites. + // Per operand: congruous = kStages+1 rotating canonical buffers — the + // load for tile i+kStages targets slot (i-1)%(kStages+1), which compute + // finished reading before this iteration's barrier 1, so NO post-compute + // barrier is needed on the pure-congruous path (one __syncthreads per + // k-tile, the classic multistage rotation); direct-crosswise rotates the + // same kStages+1 ring for the same reason; staged-crosswise (B) keeps + // kStages K-major staging buffers (filled by cp.async) plus ONE canonical + // buffer the per-tile transpose rewrites (its barrier structure keeps + // barrier 4). constexpr int kAStageBytes = kBlockM * kK; constexpr int kBStageBytes = kBlockN * kK; - constexpr int kARing = kDirectA ? kStages + 1 : kStages; // A canonic ring - constexpr int kBRing = kDirectB ? kStages + 1 : kStages; // B canonic ring - constexpr int kStB = kStages; // B staging ring size (see above) + // Direct-crosswise operands always rotate kStages+1 buffers: their + // prefetch issues right after barrier 1 (targeting the slot compute(i-1) + // released), so a kStages-deep lean ring would race the in-flight MMA + // reads. The lean ring applies only to congruous operands, whose cp.async + // prefetch sits behind the restored barrier 4. + constexpr int kARing = kDirectA ? kStages + 1 : kStages + !kLeanRing; + constexpr int kBRing = kDirectB ? kStages + 1 : kStages + !kLeanRing; + constexpr int kStB = kStages; // B staging ring size T8* const a_base = reinterpret_cast(fp8_gemm_smem); T8* const b_base = reinterpret_cast(fp8_gemm_smem + kARing * kAStageBytes); @@ -410,17 +425,24 @@ __global__ void __launch_bounds__(Traits::kCtaThreads, block_m = blockIdx.y; block_n = blockIdx.x; } - // 128x128 CTA = 8 warps as 2x4 warp tiles of 64x32 (mt x nt = 4x4 MMA). - constexpr int warps_n = kBlockN / 32; - const int warp_m = warp / warps_n; - const int warp_n = warp % warps_n; - const int64_t row_base = (int64_t)block_m * kBlockM + warp_m * 64 + group; - const int64_t output_col = - (int64_t)block_n * kBlockN + warp_n * 32 + thread_in_group * 2; - const int a_row0 = warp_m * 64; // + mt * 16 in the loop - const int b_row0 = warp_n * 32; // + nt * 8 + // 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). + constexpr int kMt = Traits::kWarpM / 16; // 16-row MMA tiles per warp + constexpr int kNt = Traits::kWarpN / 8; // 8-col MMA tiles per warp + const int warp_m = warp / Traits::kWarpsN; + const int warp_n = warp % Traits::kWarpsN; + const int64_t row_base = + (int64_t)block_m * kBlockM + warp_m * Traits::kWarpM + group; + const int64_t output_col = (int64_t)block_n * kBlockN + + warp_n * Traits::kWarpN + + thread_in_group * 2; + const int a_row0 = warp_m * Traits::kWarpM; // + mt * 16 in the loop + const int b_row0 = warp_n * Traits::kWarpN; // + nt * 8 const float scale = *p.scale; - float acc[4][4][4] = {}; // [nt][mt][acc] + float acc[kNt][kMt][4] = {}; // [nt][mt][acc] // Both operands end up in the canonical [M][kK] / [N][kK] shared tiles // the MMA fragments read, regardless of their global layout. A's tag @@ -498,6 +520,45 @@ __global__ void __launch_bounds__(Traits::kCtaThreads, const int rh8 = (lane >> 3) & 1; // +8 rows (A: lanes 8-15, 24-31) const int rh16 = lane >> 4; // +1 chunk (A: lanes 16-31; B uses rh8) + // Precomputed per-lane fragment offsets (stage-relative): the XOR + // swizzle inside tile_at depends only on (row, chunk) — never on the + // ring slot or tile_index — so every lane's ldmatrix address is its + // stage base plus one of these fixed offsets. Building the table once, + // outside the mainloop, removes the per-k_seg swizzle arithmetic + // (IMAD/LOP3 chains) from the innermost loop; the SASS compute window + // was ~36% integer address math before this. + constexpr int kSegs = kK / kMmaK; + unsigned a_off[kSegs][kMt]; // stage-relative byte offsets + unsigned b_off[kSegs][kNt]; + { + // The probe addresses are converted and immediately rebased to the + // stage origin, so the table holds pure offsets to add to any ring + // slot's converted base (double-adding the base was the bug here). + const unsigned a0 = __cvta_generic_to_shared(a_base); +#pragma unroll + for (int s = 0; s < kSegs; ++s) { +#pragma unroll + for (int mt = 0; mt < kMt; ++mt) + a_off[s][mt] = + __cvta_generic_to_shared( + tile_at(a_base, a_row0 + mt * 16 + rh8 * 8 + r7, + (s * 2 + rh16) * 16)) - + a0; + } + const T8* b_probe = kBStagePath ? b_canon : b_base; + const unsigned b0 = __cvta_generic_to_shared(b_probe); +#pragma unroll + for (int s = 0; s < kSegs; ++s) { +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + b_off[s][nt] = + __cvta_generic_to_shared( + tile_at(b_probe, b_row0 + nt * 8 + r7, + (s * 2 + rh8) * 16)) - + b0; + } + } + // Prime the pipeline. Each committed group occupies one circular shared // memory stage; the loop also handles K dimensions smaller than kStages. // Direct loads run synchronously here (back to back with their commit); @@ -532,7 +593,6 @@ __global__ void __launch_bounds__(Traits::kCtaThreads, // a time so each region's transpose overlaps the previous region's // MMA sequence (the transposes are pure shared-memory traffic — B's // global path stayed fully asynchronous above). - constexpr int kSegs = kK / kMmaK; if constexpr (kBStagePath) { transpose_tile(tile_index, 0); // Barrier 2: region 0 visible to every thread before its @@ -544,41 +604,35 @@ __global__ void __launch_bounds__(Traits::kCtaThreads, const T8* b_tile = kBStagePath ? b_canon : b_base + (size_t)(tile_index % kBRing) * kBStageBytes; + const unsigned a_base_addr = __cvta_generic_to_shared(a_tile); + const unsigned b_base_addr = __cvta_generic_to_shared(b_tile); - // 4 ldmatrix.x2 (B) + 4 ldmatrix.x4 (A) feed 16 mma.sync per k_seg — - // 0.5 load instructions per MMA, versus 4.5 scalar LDS per MMA in - // the 128x64-tile version (the kernel was LSU-issue-bound there). - // B fragments double-buffer across k_segs while B is congruous (no - // region writes in flight); a crosswise B reloads per k_seg after - // the region's transpose became visible. - unsigned b_frag[2][4][2]; + // 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 while B is + // congruous (no region writes in flight); a crosswise B reloads per + // k_seg after the region's transpose became visible. + unsigned b_frag[2][kNt][2]; if constexpr (!kBStagePath) { #pragma unroll - for (int nt = 0; nt < 4; ++nt) { - const int row = b_row0 + nt * 8 + r7; + for (int nt = 0; nt < kNt; ++nt) astrai::ldmatrix_x2_lane(b_frag[0][nt], - frag_addr(b_tile, row, rh8)); - } + b_base_addr + b_off[0][nt]); } #pragma unroll for (int k_seg = 0; k_seg < kSegs; ++k_seg) { const int bcur = k_seg & 1, bnext = bcur ^ 1; if constexpr (kBStagePath) { #pragma unroll - for (int nt = 0; nt < 4; ++nt) { - const int row = b_row0 + nt * 8 + r7; + for (int nt = 0; nt < kNt; ++nt) astrai::ldmatrix_x2_lane( - b_frag[bcur][nt], - frag_addr(b_tile, row, k_seg * 2 + rh8)); - } + b_frag[bcur][nt], b_base_addr + b_off[k_seg][nt]); } else if (k_seg + 1 < kSegs) { #pragma unroll - for (int nt = 0; nt < 4; ++nt) { - const int row = b_row0 + nt * 8 + r7; + for (int nt = 0; nt < kNt; ++nt) astrai::ldmatrix_x2_lane( - b_frag[bnext][nt], - frag_addr(b_tile, row, (k_seg + 1) * 2 + rh8)); - } + b_frag[bnext][nt], b_base_addr + b_off[k_seg + 1][nt]); } // Region k_seg+1's transpose overlaps this region's MMA work // (disjoint canonical regions, no race). @@ -590,22 +644,17 @@ __global__ void __launch_bounds__(Traits::kCtaThreads, // 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. - unsigned a_frag[5][4]; - astrai::ldmatrix_x4_lane( - a_frag[0], frag_addr(a_tile, a_row0 + rh8 * 8 + r7, - k_seg * 2 + rh16)); + unsigned a_frag[kMt + 1][4]; + astrai::ldmatrix_x4_lane(a_frag[0], a_base_addr + a_off[k_seg][0]); #pragma unroll - for (int mt = 0; mt < 4; ++mt) { - if (mt < 3) + for (int mt = 0; mt < kMt; ++mt) { + if (mt + 1 < kMt) astrai::ldmatrix_x4_lane( - a_frag[mt + 1], - frag_addr(a_tile, - a_row0 + (mt + 1) * 16 + rh8 * 8 + r7, - k_seg * 2 + rh16)); + a_frag[mt + 1], a_base_addr + a_off[k_seg][mt + 1]); #pragma unroll - for (int nt = 0; nt < 4; ++nt) - astrai::mma_sync(acc[nt][mt], a_frag[mt], b_frag[bcur][nt], - acc[nt][mt]); + for (int nt = 0; nt < kNt; ++nt) + astrai::mma_sync(acc[nt][mt], a_frag[mt], + b_frag[bcur][nt], acc[nt][mt]); } // Barrier 3: region k_seg+1's transposes complete and become // visible before the next k_seg reads them. @@ -613,45 +662,84 @@ __global__ void __launch_bounds__(Traits::kCtaThreads, if (k_seg + 1 < kSegs) __syncthreads(); } } - // Barrier 4: every thread finished reading this stage's tiles before - // the prefetch for the (i+kStages)-th tile overwrites them (and the - // next iteration's transposes rewrite the canonical buffer). - __syncthreads(); + // Barrier 4 (staged-B / lean-ring only): every thread finished + // reading this stage's tiles before the prefetch for the + // (i+kStages)-th tile overwrites them (and the next iteration's + // transposes rewrite the canonical buffer). With the kStages+1 + // canonic rotation the prefetch targets the slot compute(i-1) + // released before barrier 1, so the pure-congruous path skips this + // barrier entirely — one __syncthreads per k-tile. + if constexpr (kBStagePath || kLeanRing) __syncthreads(); if (tile_index + kStages < tile_count) { load_async(tile_index + kStages); astrai::cp_async_commit_group(); } } + // Direct bf16 epilogue 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. const float output_scale = scale; + __nv_bfloat16* tile_out = reinterpret_cast<__nv_bfloat16*>(fp8_gemm_smem); + constexpr int kRowChunks = kBlockN / 8; // 16B chunks per tile row + static_assert(kBlockM * kBlockN * 2 <= + kARing * kBlockM * kK + kBRing * kBlockN * kK, + "output tile must fit the reclaimed operand smem"); + // Swizzled address of one 16B chunk (row r, chunk c) of the tile. + auto out_chunk = [&](int r, int c) -> __nv_bfloat16* { + return tile_out + (size_t)r * kBlockN + + ((c ^ (r & (kRowChunks - 1))) * 8); + }; + const int local_col0 = warp_n * Traits::kWarpN + thread_in_group * 2; #pragma unroll - for (int nt = 0; nt < 4; ++nt) { - const int64_t col = output_col + nt * 8; - // Per-row store: FP8 packs two adjacent columns into one 16-bit - // write, BF16 into one 32-bit __nv_bfloat162 (single cvt+pack - // instruction); boundary or unaligned columns fall back to scalar - // converts so a pack never crosses the row edge or misaligns. - auto store_out = [&](int64_t row, float v0, float v1) { - if (row >= m) return; - const float r0 = v0 * output_scale; - const float r1 = v1 * output_scale; - auto* dst = out_bf16 + row * n + col; - if (col + 1 < n && (reinterpret_cast(dst) & 3) == 0) { - *reinterpret_cast<__nv_bfloat162*>(dst) = - __floats2bfloat162_rn(r0, r1); - } else { - dst[0] = __float2bfloat16(r0); - if (col + 1 < n) dst[1] = __float2bfloat16(r1); - } - }; + for (int nt = 0; nt < kNt; ++nt) { + const int col = local_col0 + nt * 8; #pragma unroll - for (int mt = 0; mt < 4; ++mt) { - const int64_t row0 = row_base + mt * 16; - float* tile_acc = acc[nt][mt]; - if (col < n) { - store_out(row0, tile_acc[0], tile_acc[1]); - store_out(row0 + 8, tile_acc[2], tile_acc[3]); - } + for (int mt = 0; mt < kMt; ++mt) { + 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. + const int off = col & 7; // element offset within the chunk + *reinterpret_cast<__nv_bfloat162*>(out_chunk(r0, col >> 3) + off) = + __floats2bfloat162_rn(tile_acc[0] * output_scale, + tile_acc[1] * output_scale); + *reinterpret_cast<__nv_bfloat162*>(out_chunk(r0 + 8, col >> 3) + + off) = + __floats2bfloat162_rn(tile_acc[2] * output_scale, + tile_acc[3] * output_scale); + } + } + __syncthreads(); + // Coalesced copy-out: thread -> one 16B chunk; consecutive threads walk + // a row so each global transaction covers a full 128B line. + const int64_t row0_global = (int64_t)block_m * kBlockM; + const int64_t col0_global = (int64_t)block_n * kBlockN; + constexpr int kTotalChunks = kBlockM * kRowChunks; + for (int idx = tid; idx < kTotalChunks; idx += kCtaThreads) { + const int r = idx / kRowChunks; + const int c = idx % kRowChunks; + const int64_t row = row0_global + r; + if (row >= m) break; // rows are consecutive: nothing left in range + const int64_t col = col0_global + (int64_t)c * 8; + const uint4 v = *reinterpret_cast(out_chunk(r, c)); + auto* dst = out_bf16 + row * n + col; + if (col + 8 <= n && (reinterpret_cast(dst) & 15) == 0) { + *reinterpret_cast(dst) = v; + } else { + // N-tail chunk or an odd-n row base: spill the elements that + // survive the row edge (and stay aligned). + const __nv_bfloat16* elems = + reinterpret_cast(&v); + for (int e = 0; e < 8 && col + e < n; ++e) dst[e] = elems[e]; } } } @@ -698,37 +786,66 @@ void launch_with_smem(int smem_bytes, dim3 grid, dim3 block, // dX ~39 TF staged vs ~38 direct). constexpr int64_t kCrossStageMinK = 8192; +// Shape-based tile dispatch (grid-searched on the production shapes, see +// csrc/tests/fp8_sweep.cu): small outputs — fewer than ~2 waves of 128x128 +// CTAs on a 24-SM part — take 64x64 CTAs of 32x32 warps with a lean +// (kStages-deep) ring: 24KB of smem keeps 4 CTAs resident, and the extra +// blocks fill the wave quantization gap (512^3: 64 vs 16 CTAs). Everything +// larger takes the 128x128 CTA (8 warps x 64x32) with the kStages+1 ring — +// one __syncthreads per k-tile. m <= 64 keeps the 64x128 CTA so a 128-row +// tile never wastes half its MMA work on predicated-off rows. +constexpr int64_t kSmallShapeMaxTiles = 48; + template || std::is_same_v> + bool GroupRaster = std::is_same_v || + std::is_same_v> void launch_fp8_gemm(const FP8Params& p, cudaStream_t stream) { - dim3 grid((p.n + 127) / 128, (p.m + 127) / 128); const bool b_staged = p.k >= kCrossStageMinK; if (p.m <= 64) { using Traits = Fp8GemmTraits; + dim3 grid((p.n + 127) / 128, (p.m + 63) / 64); if (b_staged) launch_with_smem>( - Fp8GemmSmem::kBytes, grid, - dim3(Traits::kCtaThreads), stream, p); + GroupRaster, true, true>>( + Fp8GemmSmem::kBytes, + grid, dim3(Traits::kCtaThreads), stream, p); else launch_with_smem>( - Fp8GemmSmem::kBytes, grid, - dim3(Traits::kCtaThreads), stream, p); - } else { - using Traits = Fp8GemmTraits; - if (b_staged) - launch_with_smem>( - Fp8GemmSmem::kBytes, grid, - dim3(Traits::kCtaThreads), stream, p); - else - launch_with_smem>( - Fp8GemmSmem::kBytes, grid, - dim3(Traits::kCtaThreads), stream, p); + GroupRaster, false, true>>( + Fp8GemmSmem::kBytes, + grid, dim3(Traits::kCtaThreads), stream, p); + return; } + const int64_t tiles_128 = + ((p.m + 127) / 128) * ((p.n + 127) / 128); + if (tiles_128 < kSmallShapeMaxTiles) { + using Traits = Fp8GemmTraits; + dim3 grid((p.n + 63) / 64, (p.m + 63) / 64); + if (b_staged) + launch_with_smem>( + Fp8GemmSmem::kBytes, + grid, dim3(Traits::kCtaThreads), stream, p); + else + launch_with_smem>( + Fp8GemmSmem::kBytes, + grid, dim3(Traits::kCtaThreads), stream, p); + return; + } + using Traits = Fp8GemmTraits; + dim3 grid((p.n + 127) / 128, (p.m + 127) / 128); + if (b_staged) + launch_with_smem>( + Fp8GemmSmem::kBytes, grid, + dim3(Traits::kCtaThreads), stream, p); + else + launch_with_smem>( + Fp8GemmSmem::kBytes, grid, + dim3(Traits::kCtaThreads), stream, p); } } // namespace fp8 diff --git a/csrc/tests/fp8_sweep.cu b/csrc/tests/fp8_sweep.cu new file mode 100644 index 0000000..1424198 --- /dev/null +++ b/csrc/tests/fp8_sweep.cu @@ -0,0 +1,139 @@ +/* +FP8 GEMM config sweep — pure C, no torch. Times (BM, BN, WarpM, WarpN, kK, +Stages, raster) tile configurations across the production square shapes so +the launcher's shape dispatch table is grounded in measurements. + + nvcc -I csrc -arch=sm_89 -std=c++17 -O3 --use_fast_math \ + csrc/tests/fp8_sweep.cu -o /tmp/fp8_sweep && /tmp/fp8_sweep [iters] [sizes...] +*/ + +#include "test_utils.cuh" + +#include +#include +#include +#include +#include + +#include "../kernels/fp8/gemm.cuh" + +using namespace astrai::fp8; + +namespace { + +struct BenchData { + __nv_fp8_e4m3 *da, *db; + __nv_bfloat16* dout; + float* dscale; +}; + +template +float bench_config(BenchData& d, int m, int n, int k, int iters) { + FP8Params p = {}; + p.a_ptr = d.da; + p.b_ptr = d.db; + p.out_ptr = d.dout; + p.scale = d.dscale; + p.m = m; + p.n = n; + p.k = k; + p.a_ld = k; + p.b_ld = k; + + using Traits = Fp8GemmTraits; + using Smem = Fp8GemmSmem; + dim3 grid((n + BN - 1) / BN, (m + BM - 1) / BM); + dim3 block(Traits::kCtaThreads); + const int smem = Smem::kBytes; + + auto launch = [&] { + launch_with_smem>( + smem, grid, block, 0, p); + }; + launch(); + CUDA_CHECK(cudaGetLastError()); + CUDA_CHECK(cudaDeviceSynchronize()); + + cudaEvent_t start, end; + cudaEventCreate(&start); + cudaEventCreate(&end); + for (int i = 0; i < 3; ++i) launch(); + cudaDeviceSynchronize(); + cudaEventRecord(start); + for (int i = 0; i < iters; ++i) launch(); + cudaEventRecord(end); + cudaEventSynchronize(end); + float ms = 0; + cudaEventElapsedTime(&ms, start, end); + cudaEventDestroy(start); + cudaEventDestroy(end); + return ms / iters; +} + +// One named config column. +struct Col { + const char* name; + float (*fn)(BenchData&, int, int, int, int); +}; + +template +float run(BenchData& d, int m, int n, int k, int iters) { + return bench_config(d, m, n, k, iters); +} + +} // namespace + +int main(int argc, char** argv) { + const int iters = argc > 1 ? atoi(argv[1]) : 50; + int sizes[] = {512, 1024, 2048, 4096, 8192, 0, 0, 0}; + for (int i = 2; i < argc && i < 10; ++i) sizes[i - 2] = atoi(argv[i]); + + Col cols[] = { + {"128ring3", &run<128, 128, 64, 32, 64, 2, true>}, + {"128ring3s3", &run<128, 128, 64, 32, 64, 3, true>}, + {"128lean", &run<128, 128, 64, 32, 64, 2, true, true>}, + {"64r5s4", &run<64, 64, 32, 32, 64, 4, false>}, + {"64lean4", &run<64, 64, 32, 32, 64, 4, false, true>}, + {"64lean5", &run<64, 64, 32, 32, 64, 5, false, true>}, + {"64lean3", &run<64, 64, 32, 32, 64, 3, false, true>}, + }; + const int ncols = sizeof(cols) / sizeof(cols[0]); + + printf("%6s |", "shape"); + for (auto& c : cols) printf(" %11s |", c.name); + printf("\n"); + + for (int s : sizes) { + if (s <= 0) continue; + const int m = s, n = s, k = s; + BenchData d; + std::vector<__nv_fp8_e4m3> a((size_t)m * k), b((size_t)n * k); + for (auto& v : a) v = __nv_fp8_e4m3(randf()); + for (auto& v : b) v = __nv_fp8_e4m3(randf()); + cudaMalloc(&d.da, a.size()); + cudaMalloc(&d.db, b.size()); + cudaMalloc(&d.dout, (size_t)m * n * 2); + cudaMalloc(&d.dscale, 4); + const float one = 1.0f; + cudaMemcpy(d.dscale, &one, 4, cudaMemcpyHostToDevice); + cudaMemcpy(d.da, a.data(), a.size(), cudaMemcpyHostToDevice); + cudaMemcpy(d.db, b.data(), b.size(), cudaMemcpyHostToDevice); + + const double flops = 2.0 * m * n * k; + printf("%6d |", s); + for (auto& c : cols) { + const float ms = c.fn(d, m, n, k, iters); + printf(" %5.2fus %4.1fT |", ms * 1000, + flops / (ms * 1e-3) / 1e12); + } + printf("\n"); + cudaFree(d.da); + cudaFree(d.db); + cudaFree(d.dout); + cudaFree(d.dscale); + } + return 0; +} diff --git a/scripts/tools/bench_fp8_gemm.py b/scripts/tools/bench_fp8_gemm.py new file mode 100644 index 0000000..79a246f --- /dev/null +++ b/scripts/tools/bench_fp8_gemm.py @@ -0,0 +1,88 @@ +#!/usr/bin/env python3 +"""FP8 GEMM micro-benchmark: AstrAI kernel vs cuBLAS (torch._scaled_mm). + +Sizes 512-8192, forward (NT) layout: x8[M,K] @ w8[N,K]^T -> bf16. +cuBLAS reference uses the same pre-quantized fp8 operands and the same +combined scale, so the comparison isolates the GEMM loop itself. + + python scripts/tools/bench_fp8_gemm.py --sizes 512 1024 2048 +""" + +import argparse +import sys + +import torch + +sys.path.insert(0, ".") + +from astrai.extension import loader + + +def bench(fn, iters=50, warmup=10): + for _ in range(warmup): + fn() + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iters): + fn() + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) / iters # ms + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--sizes", type=int, nargs="+", default=[512, 1024, 2048, 4096, 8192] + ) + parser.add_argument( + "--rect", action="store_true", help="also bench M=4096xN=1024 style rectangles" + ) + parser.add_argument("--iters", type=int, default=50) + args = parser.parse_args() + + assert loader.is_available("fp8_ops"), "fp8_ops extension not built" + from astrai.extension.ops.fp8 import mm_fp8 + + dev = torch.device("cuda") + torch.manual_seed(0) + + shapes = [(s, s, s) for s in args.sizes] + if args.rect: + shapes += [(4096, 1024, 4096), (8192, 4096, 8192), (2048, 8192, 2048)] + + print( + f"{'M':>6} {'N':>6} {'K':>6} | {'ours(ms)':>9} {'TFLOPs':>7} | " + f"{'cublas(ms)':>10} {'TFLOPs':>7} | {'ratio':>6}" + ) + print("-" * 72) + for m, n, k in shapes: + x8 = (torch.randn(m, k, device=dev) * 0.05).to(torch.float8_e4m3fn) + w8 = (torch.randn(n, k, device=dev) * 0.05).to(torch.float8_e4m3fn) + scale = torch.ones(1, device=dev, dtype=torch.float32) + + # ours: NT (x8 @ w8^T, LayoutB=ColMajor = weight layout) + t_ours = bench(lambda: mm_fp8(x8, w8, scale, trans_b=True), iters=args.iters) + # cuBLAS: _scaled_mm needs A row-major, B column-major (= w8.t()) + wt = w8.t() + sa = torch.ones(1, device=dev) + sb = torch.ones(1, device=dev) + + def cublas(): + return torch._scaled_mm(x8, wt, sa, sb, out_dtype=torch.bfloat16) + + t_cublas = bench(cublas, iters=args.iters) + flops = 2.0 * m * n * k + tf_ours = flops / (t_ours * 1e-3) / 1e12 + tf_cublas = flops / (t_cublas * 1e-3) / 1e12 + print( + f"{m:>6} {n:>6} {k:>6} | {t_ours:>9.3f} {tf_ours:>7.1f} | " + f"{t_cublas:>10.3f} {tf_cublas:>7.1f} | " + f"{tf_ours / tf_cublas:>5.0%}" + ) + + +if __name__ == "__main__": + main()