refactor: fp8 gemm policy layering with swap-NN and narrow-N ctas
Kernel restructured CUTLASS-style: Fp8GemmPolicy as the kernel's single template parameter (traits + operand layouts + scheduling knobs), the body split into Fp8GemmTileScheduler / Fp8CollectiveMainloop / Fp8CollectiveEpilogue collectives, and the entry split into canonicalize_gemm -> plan_gemm -> launch_plan behind fp8::gemm. - NN (dual-N-contiguous) problems run as their transpose: the swap in canonicalize_gemm plus an out-transposed epilogue removes one kernel instantiation per (format, tile config) - new 128x64 narrow CTA (8 warps of 32x32) serves the sub-wave band once its grid passes ~3/8 of a wave: +7..77% there (128x4096x4096 116->131T, 1024^3 131->174T, 4096x384x4096 147->242T, 8192x128x4096 131->233T); decode, the padding band and multi-wave shapes unchanged - launch_with_smem no longer swallows cudaFuncSetAttribute failures - fp8_test: GPU-side fp32 reference (O(m*n) compare instead of O(m*n*k) host loop), production-dispatch cases for the NN swap and the plan selection; dead transpose_layout trait removed Device: NVIDIA RTX 6000D (sm_120, 156 SMs), CUDA 13.1, torch 2.11.0+cu130. Kernel-only bench vs CUTLASS 4.8.0 sm120 dense fp8: ahead up to 1.68x below one wave (512^3 44 vs 26T, 64x4096x4096 95 vs 62T), within ~7% in the DRAM-streaming regime (8192^3 248 vs 266T).
This commit is contained in:
@@ -29,23 +29,6 @@ enum class FP8Format : int {
|
||||
struct RowMajor {};
|
||||
struct ColMajor {};
|
||||
|
||||
// Transpose of a layout tag: the same buffer with the rows and contract dims
|
||||
// swapped. B's tag is relative to the canonical [K][N] GEMM matrix, so the
|
||||
// stage-load (which views any operand as [rows][contract]) sees the transposed
|
||||
// tag — this trait makes that inversion explicit.
|
||||
template <typename Layout>
|
||||
struct transpose_layout;
|
||||
template <>
|
||||
struct transpose_layout<RowMajor> {
|
||||
using type = ColMajor;
|
||||
};
|
||||
template <>
|
||||
struct transpose_layout<ColMajor> {
|
||||
using type = RowMajor;
|
||||
};
|
||||
template <typename Layout>
|
||||
using transpose_layout_t = typename transpose_layout<Layout>::type;
|
||||
|
||||
// Compile-time tile configuration, mirroring KernelTraits<HEAD_DIM, BC,
|
||||
// WARPS, STAGES> in the attention kernels. `Fmt` selects the FP8 conversion
|
||||
// and the MMA PTX mnemonic; the remaining parameters shape the CTA tile, the
|
||||
@@ -112,6 +95,13 @@ struct FP8Params {
|
||||
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.
|
||||
int out_transposed = 0;
|
||||
// Shapes. `int` covers every realistic LLM shape; the kernels promote
|
||||
// to int64 for all pointer arithmetic.
|
||||
int m, n, k;
|
||||
|
||||
Reference in New Issue
Block a user