diff --git a/astrai/extension/ops/fp8.py b/astrai/extension/ops/fp8.py index de0f13b..aa1dfbf 100644 --- a/astrai/extension/ops/fp8.py +++ b/astrai/extension/ops/fp8.py @@ -27,17 +27,19 @@ def _mod(): def fp8_mm( a: torch.Tensor, b: torch.Tensor, sx: torch.Tensor, sw: torch.Tensor ) -> torch.Tensor: - """FP8 e4m3 GEMM: a[M,K] x b[N,K] -> bf16[M,N] (pre-scaled inputs).""" + """BF16 inputs, fused FP8 GEMM with FP32 accumulation and BF16 output.""" @fp8_mm.register_fake def _fp8_mm_fake(a, b, sx, sw): - return torch.empty((a.size(0), b.size(1)), device=a.device, dtype=torch.bfloat16) + return torch.empty((a.size(0), b.size(0)), device=a.device, dtype=torch.bfloat16) @fp8_mm.register_kernel("cuda") def _fp8_mm_cuda(a, b, sx, sw): - return _mod().fp8_mm(a, b) + if not (a.dtype == torch.bfloat16 and b.dtype == torch.bfloat16): + raise TypeError(f"bf16 GEMM requires bf16 inputs, got {a.dtype}/{b.dtype}") + return _mod().fp8_mm(a, b, sx, sw) @fp8_mm.register_kernel("cpu") @@ -46,10 +48,10 @@ def _fp8_mm_cpu(a, b, sx, sw): def linear_forward_scaled(x, w, bias, sx, sw, sx_inv, sw_inv, amax_x, amax_w): - """Quantize x/w with per-tensor scales + cuBLASLt GEMM + bias -> bf16. + """Quantize BF16 inputs to FP8, accumulate in FP32, and return BF16. - x/w: [..., K] / [N, K] bf16; sx/sw: f32 scale tensors (device scalars); - sx_inv/sw_inv: 1/scale; amax_x/amax_w: f32 buffers receiving max-abs. + x/w: [..., K] / [N, K] bf16; sx/sw and their inverses control the fused + E4M3 conversion; amax_x/amax_w receive the input max-abs values. """ if not (x.dtype == torch.bfloat16 and w.dtype == torch.bfloat16): raise TypeError(f"fp8 forward requires bf16 inputs, got {x.dtype}/{w.dtype}") diff --git a/csrc/CMakeLists.txt b/csrc/CMakeLists.txt index 321ee08..3aa63bf 100644 --- a/csrc/CMakeLists.txt +++ b/csrc/CMakeLists.txt @@ -61,9 +61,6 @@ foreach(name ${KERNELS}) "${PYTHON_INCLUDE_DIR}") target_link_libraries(${name} PRIVATE ${TORCH_LIBS}) - if(${name} STREQUAL "fp8_mm") - target_link_libraries(${name} PRIVATE CUDA::cublasLt) - endif() target_link_options(${name} PRIVATE "-Wl,-rpath,${TORCH_LIB_DIR}") target_compile_options(${name} PRIVATE diff --git a/csrc/kernels/fp8_mm.cu b/csrc/kernels/fp8_mm.cu index 581f9f5..3388871 100644 --- a/csrc/kernels/fp8_mm.cu +++ b/csrc/kernels/fp8_mm.cu @@ -1,444 +1,583 @@ -// FP8 e4m3 matrix multiply via cuBLASLt (sm89 TN layout). -// -// cuBLASLt exposes fp8 kernels only for op(A)=T, op(B)=N on Ada; we exploit -// the identity: row-major a[M,K] == A^T as col-major [K,M] (zero copy), and -// row-major wT[N,K] == B as col-major [K,N] (zero copy). The col-major -// result D[M,N] is C^T in row-major terms, so we transpose the output once. -// -// Inputs arrive pre-scaled fp8 e4m3 tensors; output is unscaled fp32. +// Fused BF16 -> E4M3 MMA -> BF16 matrix multiplication #include #include #include -#include #include +#include #include #include #include -static std::recursive_mutex g_mutex; +namespace { -static cublasLtHandle_t g_handle = nullptr; -static cublasLtMatmulDesc_t g_desc = nullptr; -static cublasLtMatrixLayout_t g_layout_a = nullptr; -static cublasLtMatrixLayout_t g_layout_b = nullptr; -static cublasLtMatrixLayout_t g_layout_c = nullptr; -static cublasLtMatmulPreference_t g_pref = nullptr; -static void* g_workspace = nullptr; -static size_t g_ws_size = 0; +constexpr int kMmaM = 16; +constexpr int kMmaN = 8; +constexpr int kMmaK = 32; +constexpr int kBlockM = 32; +constexpr int kBlockN = 32; +constexpr int kWarps = 8; +constexpr int kForwardBlockM = 64; +constexpr int kForwardBlockN = 64; -struct ShapeKey { - int64_t m; - int64_t k; - int64_t n; - bool operator==(const ShapeKey& other) const { - return m == other.m && k == other.k && n == other.n; - } -}; - -struct ShapeKeyHash { - size_t operator()(const ShapeKey& s) const { - size_t h = std::hash()(s.m); - h ^= std::hash()(s.k) + 0x9e3779b9 + (h << 6) + (h >> 2); - h ^= std::hash()(s.n) + 0x9e3779b9 + (h << 6) + (h >> 2); - return h; - } -}; - -using AlgoCache = std::unordered_map; - -static void create_matmul_config(cublasLtMatmulDesc_t* desc, - cublasLtMatrixLayout_t* layout_a, - cublasLtMatrixLayout_t* layout_b, - cublasLtMatrixLayout_t* layout_c) { - cublasOperation_t ta = CUBLAS_OP_T, tb = CUBLAS_OP_N; - TORCH_CHECK(cublasLtMatmulDescCreate(desc, CUBLAS_COMPUTE_32F, CUDA_R_32F) == - CUBLAS_STATUS_SUCCESS); - TORCH_CHECK(cublasLtMatmulDescSetAttribute( - *desc, CUBLASLT_MATMUL_DESC_TRANSA, &ta, sizeof(ta)) == - CUBLAS_STATUS_SUCCESS); - TORCH_CHECK(cublasLtMatmulDescSetAttribute( - *desc, CUBLASLT_MATMUL_DESC_TRANSB, &tb, sizeof(tb)) == - CUBLAS_STATUS_SUCCESS); - TORCH_CHECK(cublasLtMatrixLayoutCreate(layout_a, CUDA_R_8F_E4M3, 1, 1, 1) == - CUBLAS_STATUS_SUCCESS); - TORCH_CHECK(cublasLtMatrixLayoutCreate(layout_b, CUDA_R_8F_E4M3, 1, 1, 1) == - CUBLAS_STATUS_SUCCESS); - TORCH_CHECK(cublasLtMatrixLayoutCreate(layout_c, CUDA_R_16BF, 1, 1, 1) == - CUBLAS_STATUS_SUCCESS); +__device__ __forceinline__ unsigned pack_fp8x4_scalar(float x0, float x1, + float x2, float x3) { + __nv_fp8_e4m3 q0(x0); + __nv_fp8_e4m3 q1(x1); + __nv_fp8_e4m3 q2(x2); + __nv_fp8_e4m3 q3(x3); + return static_cast(q0.__x) | + (static_cast(q1.__x) << 8) | + (static_cast(q2.__x) << 16) | + (static_cast(q3.__x) << 24); } -static void ensure_cublas_lt() { - std::lock_guard lock(g_mutex); - if (g_handle) { - return; +__device__ __forceinline__ unsigned pack_fp8x4_vector(float x0, float x1, + float x2, float x3) { + const auto low = __nv_cvt_float2_to_fp8x2( + make_float2(x0, x1), __NV_SATFINITE, __NV_E4M3); + const auto high = __nv_cvt_float2_to_fp8x2( + make_float2(x2, x3), __NV_SATFINITE, __NV_E4M3); + return static_cast(low) | (static_cast(high) << 16); +} + +__device__ __forceinline__ void mma_fp8_16832(float d[4], + const unsigned a[4], + const unsigned b[2]) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 890 + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 " + "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), + "r"(b[0]), "r"(b[1])); +#endif +} + +template +__device__ __forceinline__ float load_bf16( + const __nv_bfloat16* src, int64_t row, int64_t col, + int64_t rows, int64_t cols, float& amax) { + if (row >= rows || col >= cols) return 0.0f; + int64_t index = Transpose ? col * rows + row : row * cols + col; + float value = __bfloat162float(src[index]); + amax = fmaxf(amax, fabsf(value)); + return value; +} + +__device__ __forceinline__ void atomic_max_float(float* destination, + float value) { + if (destination) + atomicMax(reinterpret_cast(destination), __float_as_uint(value)); +} + +// One thread moves eight BF16 values (16 bytes). The async copy is issued +// through a uint4-shaped pointer so the source and destination are both +// naturally 128-bit aligned for contiguous forward GEMMs. +__device__ __forceinline__ void cp_async_bf16_8( + __nv_bfloat16* destination, const __nv_bfloat16* source, bool valid) { + const unsigned shared_address = __cvta_generic_to_shared(destination); + const uint4* source_vec = reinterpret_cast(source); + asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;" + :: "r"(shared_address), "l"(source_vec), + "r"(valid ? 16 : 0)); +} + +template +__device__ __forceinline__ unsigned load_fp8x4_from_bf16( + const __nv_bfloat16* source, float scale_inv, float& amax, + bool track_amax = true) { + float x0 = __bfloat162float(source[0]); + float x1 = __bfloat162float(source[1]); + float x2 = __bfloat162float(source[2]); + float x3 = __bfloat162float(source[3]); + if constexpr (TrackAmax) { + if (track_amax) { + amax = fmaxf(amax, fmaxf(fabsf(x0), fmaxf(fabsf(x1), + fmaxf(fabsf(x2), fabsf(x3))))); + } + } + if constexpr (VectorPack) { + return pack_fp8x4_vector(x0 * scale_inv, x1 * scale_inv, + x2 * scale_inv, x3 * scale_inv); + } else { + return pack_fp8x4_scalar(x0 * scale_inv, x1 * scale_inv, + x2 * scale_inv, x3 * scale_inv); } - TORCH_CHECK(cublasLtCreate(&g_handle) == CUBLAS_STATUS_SUCCESS); - create_matmul_config(&g_desc, &g_layout_a, &g_layout_b, &g_layout_c); - TORCH_CHECK(cublasLtMatmulPreferenceCreate(&g_pref) == CUBLAS_STATUS_SUCCESS); - size_t ws = 16 * 1024 * 1024; - TORCH_CHECK(cublasLtMatmulPreferenceSetAttribute( - g_pref, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &ws, sizeof(ws)) == - CUBLAS_STATUS_SUCCESS); } -static cublasStatus_t get_algo_cached(int64_t m, int64_t k, int64_t n, - AlgoCache* cache, - cublasLtMatmulAlgo_t* algo); - -static void fp8_gemm_into(torch::Tensor lhs, torch::Tensor rhs, torch::Tensor out, - int64_t m, int64_t k, int64_t n, - const float* a_scale, const float* b_scale, - cudaStream_t stream); - -static const float k_scale_one = 1.0f; - -static void set_layout(cublasLtMatrixLayout_t layout, int64_t rows, int64_t cols, - int64_t ld) { - TORCH_CHECK(cublasLtMatrixLayoutSetAttribute(layout, CUBLASLT_MATRIX_LAYOUT_ROWS, - &rows, sizeof(rows)) == - CUBLAS_STATUS_SUCCESS); - TORCH_CHECK(cublasLtMatrixLayoutSetAttribute(layout, CUBLASLT_MATRIX_LAYOUT_COLS, - &cols, sizeof(cols)) == - CUBLAS_STATUS_SUCCESS); - TORCH_CHECK(cublasLtMatrixLayoutSetAttribute(layout, CUBLASLT_MATRIX_LAYOUT_LD, - &ld, sizeof(ld)) == - CUBLAS_STATUS_SUCCESS); +template +__device__ __forceinline__ void load_direct_fragments( + const __nv_bfloat16* a, const __nv_bfloat16* b, + int64_t row0, int64_t row1, int b_row, int64_t k0, + int64_t m, int64_t n, int64_t k, float inv_a, float inv_b, + float& amax_a, float& amax_b, unsigned a_frag[4], unsigned b_frag[2]) { + auto load_a = [&](int64_t row, int64_t col) { + return load_bf16(a, row, col, m, k, amax_a) * inv_a; + }; + auto load_b = [&](int64_t col) { + return load_bf16(b, b_row, col, n, k, amax_b) * inv_b; + }; + a_frag[0] = pack_fp8x4_scalar(load_a(row0, k0), load_a(row0, k0 + 1), + load_a(row0, k0 + 2), load_a(row0, k0 + 3)); + a_frag[1] = pack_fp8x4_scalar(load_a(row1, k0), load_a(row1, k0 + 1), + load_a(row1, k0 + 2), load_a(row1, k0 + 3)); + a_frag[2] = pack_fp8x4_scalar( + load_a(row0, k0 + 16), load_a(row0, k0 + 17), + load_a(row0, k0 + 18), load_a(row0, k0 + 19)); + a_frag[3] = pack_fp8x4_scalar( + load_a(row1, k0 + 16), load_a(row1, k0 + 17), + load_a(row1, k0 + 18), load_a(row1, k0 + 19)); + b_frag[0] = pack_fp8x4_scalar(load_b(k0), load_b(k0 + 1), load_b(k0 + 2), + load_b(k0 + 3)); + b_frag[1] = pack_fp8x4_scalar(load_b(k0 + 16), load_b(k0 + 17), + load_b(k0 + 18), load_b(k0 + 19)); } -torch::Tensor fp8_mm(torch::Tensor a, torch::Tensor b) { +template +__global__ void fused_fp8_gemm_kernel( + const __nv_bfloat16* __restrict__ a, + const __nv_bfloat16* __restrict__ b, + __nv_bfloat16* __restrict__ out, + const __nv_bfloat16* __restrict__ bias, + const float* __restrict__ scale_a, + const float* __restrict__ scale_b, + float* __restrict__ amax_a, + float* __restrict__ amax_b, + int64_t m, int64_t n, int64_t k) { + __shared__ __align__(16) __nv_fp8_e4m3 a_tile[2][BlockM][kMmaK]; + __shared__ __align__(16) __nv_fp8_e4m3 b_tile[2][BlockN][kMmaK]; + __shared__ __align__(16) __nv_bfloat16 a_bf16[2][BlockM][kMmaK]; + __shared__ __align__(16) __nv_bfloat16 b_bf16[2][BlockN][kMmaK]; + __shared__ float warp_amax_a[kWarps]; + __shared__ float warp_amax_b[kWarps]; + + const int tid = threadIdx.x; + const int warp = tid >> 5; + const int lane = tid & 31; + const int group = lane >> 2; + const int thread_in_group = lane & 3; + constexpr int warps_n = kBlockN / kMmaN; + const int warp_m = warp / warps_n; + const int warp_n = warp % warps_n; + const int64_t row_base = blockIdx.y * BlockM + warp_m * kMmaM + group; + const int64_t output_col = blockIdx.x * BlockN + warp_n * kMmaN + + thread_in_group * 2; + const float sa = *scale_a; + const float sb = *scale_b; + const float inv_a = 1.0f / sa; + const float inv_b = 1.0f / sb; + float local_amax_a = 0.0f; + float local_amax_b = 0.0f; + float acc[4 * (BlockM / kMmaM) * (BlockN / kBlockN)] = {}; + + constexpr bool AsyncContiguous = !TransposeA && !TransposeB; + const bool track_amax_a = + TrackAmax && (!AsyncContiguous || blockIdx.x == 0); + const bool track_amax_b = + TrackAmax && (!AsyncContiguous || blockIdx.y == 0); + auto load_bf16_tile = [&](int buffer, int64_t k_base) { + if constexpr (AsyncContiguous) { + const bool active_loader = true; + const int async_row = tid >> 2; + const int async_col = (tid & 3) * 8; + if (active_loader) { + const int64_t a_row = blockIdx.y * BlockM + async_row; + const int64_t b_row = blockIdx.x * BlockN + async_row; + const auto* a_source = a + a_row * k + k_base + async_col; + const auto* b_source = b + b_row * k + k_base + async_col; + const bool full_chunk = k_base + async_col + 7 < k; + const bool aligned_a = + (reinterpret_cast(a_source) & 15) == 0; + const bool aligned_b = + (reinterpret_cast(b_source) & 15) == 0; + if (a_row < m && full_chunk && aligned_a) { + cp_async_bf16_8( + &a_bf16[buffer][async_row][async_col], a_source, true); + } else { +#pragma unroll + for (int i = 0; i < 8; ++i) { + a_bf16[buffer][async_row][async_col + i] = + a_row < m && k_base + async_col + i < k + ? a_source[i] + : __float2bfloat16(0.0f); + } + } + if (b_row < n && full_chunk && aligned_b) { + cp_async_bf16_8( + &b_bf16[buffer][async_row][async_col], b_source, true); + } else if (async_row < BlockN) { +#pragma unroll + for (int i = 0; i < 8; ++i) { + b_bf16[buffer][async_row][async_col + i] = + b_row < n && k_base + async_col + i < k + ? b_source[i] + : __float2bfloat16(0.0f); + } + } + } + } + }; + + if constexpr (AsyncContiguous) { + load_bf16_tile(0, 0); + asm volatile("cp.async.commit_group;"); + asm volatile("cp.async.wait_group 0;"); + __syncthreads(); + } + + const int64_t tile_count = (k + kMmaK - 1) / kMmaK; + for (int64_t tile_index = 0; tile_index < tile_count; ++tile_index) { + const int buffer = tile_index & 1; + const int64_t k_base = tile_index * kMmaK; + if constexpr (AsyncContiguous) { + if (tile_index + 1 < tile_count) { + load_bf16_tile(buffer ^ 1, k_base + kMmaK); + asm volatile("cp.async.commit_group;"); + } + // Quantization is performed from the prefetched BF16 tile while + // the next tile is in flight. No FP8 global temporary is used. + const int quant_row = tid >> 2; + const int quant_col = (tid & 3) * 8; + const __nv_bfloat16* a_source = + &a_bf16[buffer][quant_row][quant_col]; + *reinterpret_cast(&a_tile[buffer][quant_row][quant_col]) = + load_fp8x4_from_bf16( + a_source, inv_a, local_amax_a, track_amax_a); + *reinterpret_cast(&a_tile[buffer][quant_row][quant_col + 4]) = + load_fp8x4_from_bf16( + a_source + 4, inv_a, local_amax_a, track_amax_a); + if (quant_row < BlockN) { + const __nv_bfloat16* b_source = + &b_bf16[buffer][quant_row][quant_col]; + *reinterpret_cast( + &b_tile[buffer][quant_row][quant_col]) = + load_fp8x4_from_bf16( + b_source, inv_b, local_amax_b, track_amax_b); + *reinterpret_cast( + &b_tile[buffer][quant_row][quant_col + 4]) = + load_fp8x4_from_bf16( + b_source + 4, inv_b, local_amax_b, track_amax_b); + } + } else { + const int64_t k0 = k_base + thread_in_group * 4; + const int b_row = blockIdx.x * BlockN + warp_n * kMmaN + group; + unsigned a_direct[4]; + unsigned b_direct[2]; + load_direct_fragments( + a, b, row_base, row_base + 8, b_row, k0, m, n, k, inv_a, inv_b, + local_amax_a, local_amax_b, a_direct, b_direct); + mma_fp8_16832(acc, a_direct, b_direct); + continue; + } + __syncthreads(); + + const int fragment_col = thread_in_group * 4; + const int a_row0 = warp_m * kMmaM + group; + const int a_row1 = a_row0 + 8; +#pragma unroll + for (int n_tile = 0; n_tile < BlockN / kBlockN; ++n_tile) { + const int b_row = warp_n * kMmaN + group + n_tile * kBlockN; + unsigned b_frag[2]; + b_frag[0] = *reinterpret_cast( + &b_tile[buffer][b_row][fragment_col]); + b_frag[1] = *reinterpret_cast( + &b_tile[buffer][b_row][fragment_col + 16]); +#pragma unroll + for (int m_tile = 0; m_tile < BlockM / kBlockM; ++m_tile) { + const int m_offset = m_tile * kBlockM; + unsigned a_frag[4]; + a_frag[0] = *reinterpret_cast( + &a_tile[buffer][a_row0 + m_offset][fragment_col]); + a_frag[1] = *reinterpret_cast( + &a_tile[buffer][a_row1 + m_offset][fragment_col]); + a_frag[2] = *reinterpret_cast( + &a_tile[buffer][a_row0 + m_offset][fragment_col + 16]); + a_frag[3] = *reinterpret_cast( + &a_tile[buffer][a_row1 + m_offset][fragment_col + 16]); + mma_fp8_16832( + acc + (n_tile * (BlockM / kBlockM) + m_tile) * 4, + a_frag, b_frag); + } + } + if constexpr (AsyncContiguous) { + if (tile_index + 1 < tile_count) { + asm volatile("cp.async.wait_group 0;"); + } + } + __syncthreads(); + } + + if constexpr (TrackAmax) { + for (int offset = 16; offset; offset >>= 1) { + local_amax_a = fmaxf(local_amax_a, + __shfl_xor_sync(0xffffffffu, local_amax_a, offset)); + local_amax_b = fmaxf(local_amax_b, + __shfl_xor_sync(0xffffffffu, local_amax_b, offset)); + } + if (lane == 0) { + warp_amax_a[warp] = local_amax_a; + warp_amax_b[warp] = local_amax_b; + } + __syncthreads(); + if (warp == 0) { + float block_amax_a = lane < kWarps ? warp_amax_a[lane] : 0.0f; + float block_amax_b = lane < kWarps ? warp_amax_b[lane] : 0.0f; + for (int offset = 16; offset; offset >>= 1) { + block_amax_a = fmaxf( + block_amax_a, + __shfl_xor_sync(0xffffffffu, block_amax_a, offset)); + block_amax_b = fmaxf( + block_amax_b, + __shfl_xor_sync(0xffffffffu, block_amax_b, offset)); + } + if (lane == 0) { + if (track_amax_a) atomic_max_float(amax_a, block_amax_a); + if (track_amax_b) atomic_max_float(amax_b, block_amax_b); + } + } + } + + const float output_scale = sa * sb; +#pragma unroll + for (int n_tile = 0; n_tile < BlockN / kBlockN; ++n_tile) { + const int64_t col = output_col + n_tile * kBlockN; +#pragma unroll + for (int m_tile = 0; m_tile < BlockM / kBlockM; ++m_tile) { + const int64_t row0 = row_base + m_tile * kBlockM; + const int64_t row1 = row0 + 8; + float* tile_acc = + acc + (n_tile * (BlockM / kBlockM) + m_tile) * 4; + if (col < n) { + float bias0 = 0.0f; + float bias1 = 0.0f; + if constexpr (AddBias) { + bias0 = __bfloat162float(bias[col]); + if (col + 1 < n) + bias1 = __bfloat162float(bias[col + 1]); + } + if (row0 < m) { + out[row0 * n + col] = + __float2bfloat16(tile_acc[0] * output_scale + bias0); + if (col + 1 < n) + out[row0 * n + col + 1] = __float2bfloat16( + tile_acc[1] * output_scale + bias1); + } + if (row1 < m) { + out[row1 * n + col] = + __float2bfloat16(tile_acc[2] * output_scale + bias0); + if (col + 1 < n) + out[row1 * n + col + 1] = __float2bfloat16( + tile_acc[3] * output_scale + bias1); + } + } + } + } +} + +template +void launch_fused_fp8_gemm( + const torch::Tensor& a, const torch::Tensor& b, torch::Tensor& out, + const torch::Tensor& bias, const torch::Tensor& scale_a, + const torch::Tensor& scale_b, torch::Tensor* amax_a, + torch::Tensor* amax_b, int64_t m, int64_t n, int64_t k, + cudaStream_t stream) { + dim3 grid((n + BlockN - 1) / BlockN, + (m + BlockM - 1) / BlockM); + const auto* bias_ptr = AddBias + ? reinterpret_cast(bias.data_ptr()) + : nullptr; + fused_fp8_gemm_kernel + <<>>( + reinterpret_cast(a.data_ptr()), + reinterpret_cast(b.data_ptr()), + reinterpret_cast<__nv_bfloat16*>(out.data_ptr()), bias_ptr, + scale_a.data_ptr(), scale_b.data_ptr(), + amax_a ? amax_a->data_ptr() : nullptr, + amax_b ? amax_b->data_ptr() : nullptr, m, n, k); +} + +void check_fp8_device(const torch::Tensor& tensor) { + static std::mutex mutex; + static std::unordered_map supported; + const int device = tensor.device().index(); + { + std::lock_guard lock(mutex); + auto cached = supported.find(device); + if (cached != supported.end()) { + TORCH_CHECK(cached->second, + "fused FP8 MMA requires compute capability 8.9 or newer"); + return; + } + } + + const auto* properties = at::cuda::getDeviceProperties(device); + const bool is_supported = properties->major > 8 || + (properties->major == 8 && properties->minor >= 9); + { + std::lock_guard lock(mutex); + supported.emplace(device, is_supported); + } + TORCH_CHECK(is_supported, + "fused FP8 MMA requires compute capability 8.9 or newer"); +} + +void check_scale(const torch::Tensor& scale, const torch::Tensor& input, + const char* name) { + TORCH_CHECK(scale.is_cuda() && scale.device() == input.device() && + scale.scalar_type() == torch::kFloat32 && scale.numel() == 1, + name, " must be a CUDA float32 scalar on the input device"); +} + +} // namespace + +torch::Tensor fp8_mm(torch::Tensor a, torch::Tensor b, torch::Tensor sx, + torch::Tensor sw) { TORCH_CHECK(a.is_cuda() && b.is_cuda(), "CUDA tensors required"); - TORCH_CHECK(a.scalar_type() == torch::kFloat8_e4m3fn, "a must be float8_e4m3fn"); - TORCH_CHECK(b.scalar_type() == torch::kFloat8_e4m3fn, "b must be float8_e4m3fn"); - TORCH_CHECK(a.dim() == 2 && b.dim() == 2, "2D tensors required"); + TORCH_CHECK(a.scalar_type() == torch::kBFloat16 && + b.scalar_type() == torch::kBFloat16, + "a and b must be bf16"); + TORCH_CHECK(a.dim() == 2 && b.dim() == 2, "a and b must be 2D"); + TORCH_CHECK(a.device() == b.device(), "a and b must be on the same device"); + TORCH_CHECK(a.size(1) == b.size(1), "inner dim mismatch"); + check_scale(sx, a, "sx"); + check_scale(sw, a, "sw"); + check_fp8_device(a); const at::cuda::OptionalCUDAGuard guard(a.device()); auto stream = at::cuda::getCurrentCUDAStream(); auto a_c = a.contiguous(); auto b_c = b.contiguous(); - int64_t m = a_c.size(0), k = a_c.size(1), n = b_c.size(0); - TORCH_CHECK(b_c.size(1) == k, "inner dim mismatch"); - - auto buf = torch::empty({m, n}, a_c.options().dtype(torch::kBFloat16)); - ensure_cublas_lt(); - fp8_gemm_into(a_c, b_c, buf, m, k, n, &k_scale_one, &k_scale_one, - stream.stream()); - return buf; + auto out = torch::empty({a_c.size(0), b_c.size(0)}, a_c.options()); + torch::Tensor no_bias; + launch_fused_fp8_gemm( + a_c, b_c, out, no_bias, sx, sw, nullptr, nullptr, + a_c.size(0), b_c.size(0), a_c.size(1), stream.stream()); + C10_CUDA_CHECK(cudaGetLastError()); + return out; } - -// --------------------------------------------------------------------------- -// Quantize: bf16 * scale_inv -> fp8, one atomicMax amax per kernel call. -// amax_ptr must be zeroed before launch; float-bits atomicMax works because -// |v| >= 0 has a monotonic IEEE bit pattern. -// --------------------------------------------------------------------------- - -template -__device__ __forceinline__ T8 cast_fp8(float v); - -template <> -__device__ __forceinline__ __nv_fp8_e4m3 cast_fp8<__nv_fp8_e4m3>(float v) { - return __nv_fp8_e4m3(v); -} - -template <> -__device__ __forceinline__ __nv_fp8_e5m2 cast_fp8<__nv_fp8_e5m2>(float v) { - return __nv_fp8_e5m2(v); -} - -template -__global__ void quantize_kernel(const __nv_bfloat16* __restrict__ src, - const float* __restrict__ scale_inv, - T8* __restrict__ dst, - float* __restrict__ amax_ptr, int64_t n) { - int64_t i = blockIdx.x * (int64_t)blockDim.x + threadIdx.x; - float amax = 0.f; - if (i < n) { - float raw = __bfloat162float(src[i]); - dst[i] = cast_fp8(raw * *scale_inv); - amax = fabsf(raw); - } - for (int off = 16; off; off >>= 1) - amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, off)); - __shared__ float sm[8]; - if ((threadIdx.x & 31) == 0) sm[threadIdx.x >> 5] = amax; - __syncthreads(); - if (threadIdx.x == 0) { - float m = 0.f; - for (int w = 0; w < blockDim.x / 32; ++w) m = fmaxf(m, sm[w]); - atomicMax(reinterpret_cast(amax_ptr), __float_as_uint(m)); - } -} - -// Same but with a transpose (rows x cols bf16 row-major -> fp8 [cols, rows]). -template -__global__ void transpose_quantize_kernel( - const __nv_bfloat16* __restrict__ src, const float* __restrict__ scale_inv, - T8* __restrict__ dst, float* __restrict__ amax_ptr, int64_t rows, - int64_t cols) { - __shared__ T8 tile[32][33]; - int64_t x = blockIdx.x * 32 + threadIdx.x; - int64_t y = blockIdx.y * 32 + threadIdx.y; - float amax = 0.f; - for (int j = 0; j < 32; j += 8) { - if (x < cols && y + j < rows) { - float raw = __bfloat162float(src[(y + j) * cols + x]); - tile[threadIdx.y + j][threadIdx.x] = cast_fp8(raw * *scale_inv); - amax = fmaxf(amax, fabsf(raw)); - } - } - __syncthreads(); - - x = blockIdx.y * 32 + threadIdx.x; - y = blockIdx.x * 32 + threadIdx.y; - for (int j = 0; j < 32; j += 8) { - if (x < rows && y + j < cols) { - dst[(y + j) * rows + x] = tile[threadIdx.x][threadIdx.y + j]; - } - } - for (int off = 16; off; off >>= 1) - amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, off)); - __shared__ float sm[8]; - if ((threadIdx.x & 31) == 0) sm[threadIdx.x >> 5] = amax; - __syncthreads(); - if (threadIdx.x == 0) { - float m = 0.f; - for (int w = 0; w < blockDim.x / 32; ++w) m = fmaxf(m, sm[w]); - atomicMax(reinterpret_cast(amax_ptr), __float_as_uint(m)); - } -} - -__global__ void bias_add_bf16_kernel( - __nv_bfloat16* __restrict__ dst, const __nv_bfloat16* __restrict__ bias, - int64_t total, int64_t n) { - // GEMM and output use the same row-major [M,N] layout. - int64_t idx = blockIdx.x * (int64_t)blockDim.x + threadIdx.x; - if (idx >= total) return; - float v = __bfloat162float(dst[idx]); - dst[idx] = __float2bfloat16(v + __bfloat162float(bias[idx % n])); -} - -static cublasStatus_t get_algo_cached(int64_t m, int64_t k, int64_t n, - AlgoCache* cache, - cublasLtMatmulAlgo_t* algo) { - std::lock_guard lock(g_mutex); - ShapeKey key{m, k, n}; - auto it = cache->find(key); - if (it != cache->end()) { - *algo = it->second; - return CUBLAS_STATUS_SUCCESS; - } - cublasLtMatmulHeuristicResult_t heur; - int returned = 0; - cublasStatus_t st = cublasLtMatmulAlgoGetHeuristic( - g_handle, g_desc, g_layout_a, g_layout_b, g_layout_c, g_layout_c, g_pref, 1, - &heur, &returned); - if (st != CUBLAS_STATUS_SUCCESS || returned == 0) - return CUBLAS_STATUS_NOT_SUPPORTED; - if (heur.workspaceSize > g_ws_size) { - if (g_workspace) cudaFree(g_workspace); - TORCH_CHECK(cudaMalloc(&g_workspace, heur.workspaceSize) == cudaSuccess); - g_ws_size = heur.workspaceSize; - } - cache->emplace(key, heur.algo); - *algo = heur.algo; - return CUBLAS_STATUS_SUCCESS; -} - -static void fp8_gemm_into(torch::Tensor lhs, torch::Tensor rhs, torch::Tensor out, - int64_t m, int64_t k, int64_t n, - const float* a_scale, const float* b_scale, - cudaStream_t stream) { - std::lock_guard lock(g_mutex); - set_layout(g_layout_a, k, n, k); // param A = rhs (op=T -> [N,K]) - set_layout(g_layout_b, k, m, k); // param B = lhs (op=N -> [K,M]) - set_layout(g_layout_c, n, m, n); // col-major [N,M] == row-major [M,N] - // Per-tensor FP32 scales applied inside the GEMM: - // D = alpha * A_SCALE * B_SCALE * A * B (alpha = 1). - TORCH_CHECK(cublasLtMatmulDescSetAttribute( - g_desc, CUBLASLT_MATMUL_DESC_A_SCALE_POINTER, &a_scale, - sizeof(a_scale)) == CUBLAS_STATUS_SUCCESS); - TORCH_CHECK(cublasLtMatmulDescSetAttribute( - g_desc, CUBLASLT_MATMUL_DESC_B_SCALE_POINTER, &b_scale, - sizeof(b_scale)) == CUBLAS_STATUS_SUCCESS); - float alpha = 1.0f, beta = 0.0f; - static AlgoCache cache; - cublasLtMatmulAlgo_t algo; - cublasStatus_t st = get_algo_cached(m, k, n, &cache, &algo); - TORCH_CHECK(st == CUBLAS_STATUS_SUCCESS, - "cublasLtMatmulAlgoGetHeuristic failed: ", cublasLtGetStatusName(st)); - st = cublasLtMatmul(g_handle, g_desc, &alpha, rhs.data_ptr(), g_layout_a, - lhs.data_ptr(), g_layout_b, &beta, out.data_ptr(), g_layout_c, - out.data_ptr(), g_layout_c, &algo, g_workspace, g_ws_size, - stream); - TORCH_CHECK(st == CUBLAS_STATUS_SUCCESS, - "cublasLtMatmul failed: ", cublasLtGetStatusName(st)); -} - -// --------------------------------------------------------------------------- -// Scaled FP8 linear forward: quantize x/w with per-tensor scales -> cublasLt -// GEMM (scales applied inside) -> bias in-place -> bf16 [..., N]. -// sx/sw: f32 scale tensors (device scalars); sx_inv/sw_inv: 1/scale. -// amax_x/amax_w: f32 buffers receiving max-abs of the quantized tensors. -// --------------------------------------------------------------------------- - -torch::Tensor fp8_linear_forward_scaled(torch::Tensor x, torch::Tensor w, - torch::Tensor bias, torch::Tensor sx, - torch::Tensor sw, torch::Tensor sx_inv, - torch::Tensor sw_inv, - torch::Tensor amax_x, - torch::Tensor amax_w) { +torch::Tensor fp8_linear_forward_scaled( + torch::Tensor x, torch::Tensor w, torch::Tensor bias, torch::Tensor sx, + torch::Tensor sw, torch::Tensor sx_inv, torch::Tensor sw_inv, + torch::Tensor amax_x, torch::Tensor amax_w) { TORCH_CHECK(x.is_cuda() && w.is_cuda(), "CUDA tensors required"); - TORCH_CHECK(x.dtype() == torch::kBFloat16 && w.dtype() == torch::kBFloat16, + TORCH_CHECK(x.scalar_type() == torch::kBFloat16 && + w.scalar_type() == torch::kBFloat16, "x and w must be bf16"); + TORCH_CHECK(x.device() == w.device(), "x and w must be on the same device"); + check_scale(sx, x, "sx"); + check_scale(sw, x, "sw"); + check_fp8_device(x); const at::cuda::OptionalCUDAGuard guard(x.device()); auto stream = at::cuda::getCurrentCUDAStream(); auto x_c = x.reshape({-1, w.size(1)}).contiguous(); auto w_c = w.contiguous(); int64_t m = x_c.size(0), k = x_c.size(1), n = w_c.size(0); - TORCH_CHECK(w_c.size(1) == k, "inner dim mismatch"); - ensure_cublas_lt(); - - const float* sx_ptr = sx.data_ptr(); - const float* sw_ptr = sw.data_ptr(); - const float* sxi_ptr = sx_inv.data_ptr(); - const float* swi_ptr = sw_inv.data_ptr(); - float* amax_x_ptr = amax_x.data_ptr(); - float* amax_w_ptr = amax_w.data_ptr(); - C10_CUDA_CHECK(cudaMemsetAsync(amax_x_ptr, 0, sizeof(float), stream.stream())); - C10_CUDA_CHECK(cudaMemsetAsync(amax_w_ptr, 0, sizeof(float), stream.stream())); - - auto x8 = torch::empty({m, k}, x_c.options().dtype(torch::kFloat8_e4m3fn)); - auto w8 = torch::empty({n, k}, w_c.options().dtype(torch::kFloat8_e4m3fn)); - int64_t block = 256; - quantize_kernel<__nv_fp8_e4m3> - <<<(unsigned)((m * k + block - 1) / block), block, 0, stream.stream()>>>( - reinterpret_cast(x_c.data_ptr()), sxi_ptr, - reinterpret_cast<__nv_fp8_e4m3*>(x8.data_ptr()), amax_x_ptr, m * k); - quantize_kernel<__nv_fp8_e4m3> - <<<(unsigned)((n * k + block - 1) / block), block, 0, stream.stream()>>>( - reinterpret_cast(w_c.data_ptr()), swi_ptr, - reinterpret_cast<__nv_fp8_e4m3*>(w8.data_ptr()), amax_w_ptr, n * k); + TORCH_CHECK(w_c.dim() == 2 && w_c.size(1) == k, "inner dim mismatch"); + C10_CUDA_CHECK(cudaMemsetAsync(amax_x.data_ptr(), 0, sizeof(float), + stream.stream())); + C10_CUDA_CHECK(cudaMemsetAsync(amax_w.data_ptr(), 0, sizeof(float), + stream.stream())); + auto out = torch::empty({m, n}, x_c.options()); + if (bias.defined() && bias.numel() > 0) { + TORCH_CHECK(bias.is_cuda() && bias.device() == x.device() && + bias.scalar_type() == torch::kBFloat16 && + bias.numel() == n, + "bias must be CUDA bf16 with shape [N]"); + launch_fused_fp8_gemm( + x_c, w_c, out, bias, sx, sw, &amax_x, &amax_w, + m, n, k, stream.stream()); + } else { + launch_fused_fp8_gemm( + x_c, w_c, out, bias, sx, sw, &amax_x, &amax_w, + m, n, k, stream.stream()); + } C10_CUDA_CHECK(cudaGetLastError()); - auto out = torch::empty({m, n}, x_c.options()); - fp8_gemm_into(x8, w8, out, m, k, n, sw_ptr, sx_ptr, stream.stream()); - - if (bias.defined() && bias.numel() > 0) { - TORCH_CHECK(bias.scalar_type() == torch::kBFloat16 && bias.numel() == n, - "bias must be bf16 with shape [N]"); - bias_add_bf16_kernel<<<(unsigned)((m * n + block - 1) / block), block, 0, stream>>>( - reinterpret_cast<__nv_bfloat16*>(out.data_ptr()), - reinterpret_cast(bias.data_ptr()), m * n, n); - C10_CUDA_CHECK(cudaGetLastError()); - } - + (void)sx_inv; + (void)sw_inv; std::vector shape(x.sizes().begin(), x.sizes().end() - 1); shape.push_back(n); return out.reshape(shape); } -// --------------------------------------------------------------------------- -// Scaled FP8 linear backward: dX = g @ W, dW = g^T @ X, dB = sum(g). -// Scales: g uses sg (immediate), w/x reuse the forward scales. -// --------------------------------------------------------------------------- - std::tuple fp8_linear_backward_scaled( torch::Tensor g, torch::Tensor x, torch::Tensor w, std::vector masks, torch::Tensor sg, torch::Tensor sw, torch::Tensor sx, torch::Tensor sg_inv, torch::Tensor sw_inv, torch::Tensor sx_inv, torch::Tensor amax_g) { - const at::cuda::OptionalCUDAGuard guard(g.device()); - TORCH_CHECK(g.dtype() == torch::kBFloat16 && x.dtype() == torch::kBFloat16 && - w.dtype() == torch::kBFloat16, + TORCH_CHECK(g.is_cuda() && x.is_cuda() && w.is_cuda(), "CUDA tensors required"); + TORCH_CHECK(g.scalar_type() == torch::kBFloat16 && + x.scalar_type() == torch::kBFloat16 && + w.scalar_type() == torch::kBFloat16, "g, x, and w must be bf16"); + TORCH_CHECK(g.device() == x.device() && g.device() == w.device(), + "g, x, and w must be on the same device"); + TORCH_CHECK(masks.size() == 3, "masks must contain three values"); + check_fp8_device(g); + const at::cuda::OptionalCUDAGuard guard(g.device()); auto stream = at::cuda::getCurrentCUDAStream(); + auto g_c = g.reshape({-1, w.size(0)}).contiguous(); auto x_c = x.reshape({-1, x.size(-1)}).contiguous(); auto w_c = w.contiguous(); - int64_t m = g_c.size(0); - int64_t n = w.size(0); - int64_t k = w.size(1); + int64_t m = g_c.size(0), n = w_c.size(0), k = w_c.size(1); TORCH_CHECK(x_c.size(0) == m && x_c.size(1) == k && g_c.size(1) == n, "backward shape mismatch"); auto grad_input = torch::empty_like(x); auto grad_weight = torch::empty_like(w); - auto grad_bias = torch::empty({0}, g_c.options().dtype(g.dtype())); - ensure_cublas_lt(); - - const float* sg_ptr = sg.data_ptr(); - const float* sw_ptr = sw.data_ptr(); - const float* sx_ptr = sx.data_ptr(); - const float* sgi_ptr = sg_inv.data_ptr(); - const float* swi_ptr = sw_inv.data_ptr(); - const float* sxi_ptr = sx_inv.data_ptr(); - float* amax_g_ptr = amax_g.data_ptr(); - C10_CUDA_CHECK(cudaMemsetAsync(amax_g_ptr, 0, sizeof(float), stream.stream())); - - auto fp8_options = g_c.options().dtype(torch::kFloat8_e4m3fn); - auto g8 = torch::empty({m, n}, fp8_options); - auto gt8 = masks[1] ? torch::empty({n, m}, fp8_options) : torch::Tensor(); - auto wt8 = masks[0] ? torch::empty({k, n}, fp8_options) : torch::Tensor(); - auto xt8 = masks[1] ? torch::empty({k, m}, fp8_options) : torch::Tensor(); - // w/x transpose-quantize amax goes to a scratch buffer, NOT amax_g: the - // gradient scale must only see the gradient's own max-abs. - auto amax_t = torch::zeros({1}, g_c.options().dtype(torch::kFloat32)); - - int64_t block = 256; - quantize_kernel<__nv_fp8_e4m3> - <<<(unsigned)((m * n + block - 1) / block), block, 0, stream.stream()>>>( - reinterpret_cast(g_c.data_ptr()), sgi_ptr, - reinterpret_cast<__nv_fp8_e4m3*>(g8.data_ptr()), amax_g_ptr, m * n); - dim3 threads(32, 8); + auto grad_bias = torch::empty({0}, g.options()); + C10_CUDA_CHECK(cudaMemsetAsync(amax_g.data_ptr(), 0, sizeof(float), + stream.stream())); + torch::Tensor no_bias; + bool recorded_amax = false; if (masks[0]) { - dim3 blocks((k + 31) / 32, (n + 31) / 32); - transpose_quantize_kernel<__nv_fp8_e4m3> - <<>>( - reinterpret_cast(w_c.data_ptr()), swi_ptr, - reinterpret_cast<__nv_fp8_e4m3*>(wt8.data_ptr()), - amax_t.data_ptr(), n, k); - fp8_gemm_into(g8, wt8, grad_input.reshape({m, k}), m, n, k, sg_ptr, - sw_ptr, stream.stream()); + auto grad_input_2d = grad_input.reshape({m, k}); + launch_fused_fp8_gemm( + g_c, w_c, grad_input_2d, no_bias, sg, sw, &amax_g, nullptr, + m, k, n, stream.stream()); + recorded_amax = true; } if (masks[1]) { - dim3 g_blocks((n + 31) / 32, (m + 31) / 32); - dim3 x_blocks((k + 31) / 32, (m + 31) / 32); - transpose_quantize_kernel<__nv_fp8_e4m3> - <<>>( - reinterpret_cast(g_c.data_ptr()), sgi_ptr, - reinterpret_cast<__nv_fp8_e4m3*>(gt8.data_ptr()), - amax_t.data_ptr(), m, n); - transpose_quantize_kernel<__nv_fp8_e4m3> - <<>>( - reinterpret_cast(x_c.data_ptr()), sxi_ptr, - reinterpret_cast<__nv_fp8_e4m3*>(xt8.data_ptr()), - amax_t.data_ptr(), m, k); - fp8_gemm_into(gt8, xt8, grad_weight, n, m, k, sg_ptr, sx_ptr, - stream.stream()); + launch_fused_fp8_gemm( + g_c, x_c, grad_weight, no_bias, sg, sx, + recorded_amax ? nullptr : &amax_g, nullptr, + n, k, m, stream.stream()); + recorded_amax = true; + } + if (!recorded_amax) { + amax_g.copy_(g_c.abs().amax().to(torch::kFloat32)); } C10_CUDA_CHECK(cudaGetLastError()); - if (masks[2]) { - grad_bias = g_c.sum(0).to(g.dtype()); - } - return std::tuple( - grad_input, grad_weight, grad_bias); + if (masks[2]) grad_bias = g_c.sum(0).to(g.scalar_type()); + + (void)sg_inv; + (void)sw_inv; + (void)sx_inv; + return {grad_input, grad_weight, grad_bias}; } PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("fp8_mm", &fp8_mm, py::arg("a"), py::arg("b"), - "FP8 e4m3 GEMM: a[M,K] x b[N,K] -> bf16[M,N] (pre-scaled inputs)"); + m.def("fp8_mm", &fp8_mm, py::arg("a"), py::arg("b"), py::arg("sx"), + py::arg("sw"), + "Fused BF16 input, E4M3 MMA, FP32 accumulation, BF16 output GEMM"); m.def("fp8_linear_forward_scaled", &fp8_linear_forward_scaled, py::arg("x"), py::arg("w"), py::arg("bias"), py::arg("sx"), py::arg("sw"), py::arg("sx_inv"), py::arg("sw_inv"), py::arg("amax_x"), py::arg("amax_w"), - "Scaled FP8 linear forward: quantize with per-tensor scales + " - "cublasLt GEMM (scales applied inside) + bias -> bf16"); + "Fused BF16-to-FP8 linear forward with FP32 accumulation"); m.def("fp8_linear_backward_scaled", &fp8_linear_backward_scaled, py::arg("g"), py::arg("x"), py::arg("w"), py::arg("masks"), py::arg("sg"), py::arg("sw"), py::arg("sx"), py::arg("sg_inv"), py::arg("sw_inv"), py::arg("sx_inv"), py::arg("amax_g"), - "Scaled FP8 linear backward: dX = g*sw @ W, dW = (g*sx)^T @ X, " - "dB = sum(g)"); + "Fused BF16-to-FP8 linear backward with FP32 accumulation"); } diff --git a/csrc/tests/fp8_mma_test.cu b/csrc/tests/fp8_mma_test.cu new file mode 100644 index 0000000..c949258 --- /dev/null +++ b/csrc/tests/fp8_mma_test.cu @@ -0,0 +1,155 @@ +/* +Single-kernel BF16 -> FP8 MMA -> BF16 demo for Ada (sm_89). + +nvcc -I csrc -arch=sm_89 -std=c++17 -O3 --use_fast_math \ + --ptxas-options=-O3,-v csrc/tests/fp8_mma_test.cu -o fp8_mma_test \ + && ./fp8_mma_test +*/ + +#include "test_utils.cuh" + +#include + +#include +#include + +constexpr int M = 16; +constexpr int N = 8; +constexpr int K = 32; + +__device__ __forceinline__ unsigned pack_fp8x4(float x0, float x1, float x2, + float x3) { + __nv_fp8_e4m3 q0(x0); + __nv_fp8_e4m3 q1(x1); + __nv_fp8_e4m3 q2(x2); + __nv_fp8_e4m3 q3(x3); + return static_cast(q0.__x) | + (static_cast(q1.__x) << 8) | + (static_cast(q2.__x) << 16) | + (static_cast(q3.__x) << 24); +} + +__device__ __forceinline__ unsigned load_quantize_fp8x4( + const bf16* src, float scale_inv) { + return pack_fp8x4(__bfloat162float(src[0]) * scale_inv, + __bfloat162float(src[1]) * scale_inv, + __bfloat162float(src[2]) * scale_inv, + __bfloat162float(src[3]) * scale_inv); +} + +__device__ __forceinline__ void mma_fp8_16832(float d[4], + const unsigned a[4], + const unsigned b[2]) { + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 " + "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), + "r"(b[0]), "r"(b[1])); +} + +__global__ void fused_bf16_fp8_mma_kernel( + const bf16* __restrict__ a, const bf16* __restrict__ b, + bf16* __restrict__ out, float scale_a, float scale_b) { + const int lane = threadIdx.x; + const int group = lane >> 2; + const int thread_in_group = lane & 3; + const int k0 = thread_in_group * 4; + + // PTX m16n8k32 A fragment: two rows, two 16-column K partitions. + unsigned a_frag[4]; + a_frag[0] = load_quantize_fp8x4(&a[group * K + k0], 1.0f / scale_a); + a_frag[1] = load_quantize_fp8x4(&a[(group + 8) * K + k0], 1.0f / scale_a); + a_frag[2] = load_quantize_fp8x4(&a[group * K + k0 + 16], 1.0f / scale_a); + a_frag[3] = load_quantize_fp8x4(&a[(group + 8) * K + k0 + 16], + 1.0f / scale_a); + + // B is supplied as row-major [N,K], equivalent to the col-major [K,N] + // operand required by the MMA instruction. + unsigned b_frag[2]; + b_frag[0] = load_quantize_fp8x4(&b[group * K + k0], 1.0f / scale_b); + b_frag[1] = load_quantize_fp8x4(&b[group * K + k0 + 16], 1.0f / scale_b); + + float acc[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + mma_fp8_16832(acc, a_frag, b_frag); + + const int col = thread_in_group * 2; + const float output_scale = scale_a * scale_b; + *reinterpret_cast<__nv_bfloat162*>(&out[group * N + col]) = + __floats2bfloat162_rn(acc[0] * output_scale, + acc[1] * output_scale); + *reinterpret_cast<__nv_bfloat162*>(&out[(group + 8) * N + col]) = + __floats2bfloat162_rn(acc[2] * output_scale, + acc[3] * output_scale); +} + +static float quantize_e4m3(float value) { + return static_cast(__nv_fp8_e4m3(value)); +} + +int main() { + srand(0); + std::vector a(M * K), b(N * K), reference(M * N, 0.0f); + std::vector a_bf16(M * K), b_bf16(N * K), output(M * N); + for (float& value : a) value = randf() * 4.0f; + for (float& value : b) value = randf() * 4.0f; + for (int i = 0; i < M * K; ++i) { + a_bf16[i] = f2bf(a[i]); + a[i] = bf2f(a_bf16[i]); + } + for (int i = 0; i < N * K; ++i) { + b_bf16[i] = f2bf(b[i]); + b[i] = bf2f(b_bf16[i]); + } + + const float amax = *std::max_element( + a.begin(), a.end(), [](float x, float y) { return fabsf(x) < fabsf(y); }); + const float bmax = *std::max_element( + b.begin(), b.end(), [](float x, float y) { return fabsf(x) < fabsf(y); }); + const float scale_a = fabsf(amax) / 448.0f; + const float scale_b = fabsf(bmax) / 448.0f; + + for (int row = 0; row < M; ++row) { + for (int col = 0; col < N; ++col) { + float sum = 0.0f; + for (int k = 0; k < K; ++k) { + float qa = quantize_e4m3(a[row * K + k] / scale_a); + float qb = quantize_e4m3(b[col * K + k] / scale_b); + sum = fmaf(qa, qb, sum); + } + reference[row * N + col] = sum * scale_a * scale_b; + } + } + + bf16 *d_a, *d_b, *d_out; + CUDA_CHECK(cudaMalloc(&d_a, a_bf16.size() * sizeof(bf16))); + CUDA_CHECK(cudaMalloc(&d_b, b_bf16.size() * sizeof(bf16))); + CUDA_CHECK(cudaMalloc(&d_out, output.size() * sizeof(bf16))); + CUDA_CHECK(cudaMemcpy(d_a, a_bf16.data(), a_bf16.size() * sizeof(bf16), + cudaMemcpyHostToDevice)); + CUDA_CHECK(cudaMemcpy(d_b, b_bf16.data(), b_bf16.size() * sizeof(bf16), + cudaMemcpyHostToDevice)); + + fused_bf16_fp8_mma_kernel<<<1, 32>>>(d_a, d_b, d_out, scale_a, scale_b); + CUDA_CHECK(cudaDeviceSynchronize()); + CUDA_CHECK(cudaMemcpy(output.data(), d_out, output.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + + float max_abs_error = 0.0f; + float max_rel_error = 0.0f; + for (int i = 0; i < M * N; ++i) { + float error = fabsf(bf2f(output[i]) - reference[i]); + max_abs_error = fmaxf(max_abs_error, error); + max_rel_error = fmaxf(max_rel_error, + error / fmaxf(fabsf(reference[i]), 1e-4f)); + } + const bool pass = max_abs_error < 0.05f; + print_test_header(); + print_test_row("M=16 N=8 K=32 fused BF16->E4M3 MMA", max_abs_error, + max_rel_error, pass); + + cudaFree(d_a); + cudaFree(d_b); + cudaFree(d_out); + return pass ? 0 : 1; +} diff --git a/tests/extension/test_fp8_mma.py b/tests/extension/test_fp8_mma.py new file mode 100644 index 0000000..e1280c3 --- /dev/null +++ b/tests/extension/test_fp8_mma.py @@ -0,0 +1,96 @@ +"""Fused BF16-boundary FP8 MMA kernel tests.""" + +import pytest +import torch + +from astrai.extension.loader import get_module, is_available + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() + or torch.cuda.get_device_capability() < (8, 9) + or not is_available("fp8_mm"), + reason="fused FP8 MMA requires a built kernel and compute capability 8.9+", +) + + +def _scale(tensor): + return (tensor.abs().amax().float() / 448.0).clamp_min(1e-12) + + +def _quantize(tensor, scale): + return (tensor.float() / scale).to(torch.float8_e4m3fn).float() + + +@pytest.mark.parametrize( + ("m", "n", "k"), + [(16, 8, 32), (17, 9, 33), (31, 15, 64), (32, 48, 96)], +) +def test_fused_fp8_mma_matches_explicit_quantization(m, n, k): + torch.manual_seed(m + n + k) + a = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) + b = torch.randn(n, k, device="cuda", dtype=torch.bfloat16) + scale_a = _scale(a) + scale_b = _scale(b) + + out = get_module("fp8_mm").fp8_mm(a, b, scale_a, scale_b) + expected = ( + _quantize(a, scale_a) @ _quantize(b, scale_b).t() * scale_a * scale_b + ).to(torch.bfloat16) + + assert out.dtype == torch.bfloat16 + assert out.shape == (m, n) + torch.testing.assert_close(out, expected, atol=0.125, rtol=0.01) + + +def test_fused_fp8_linear_forward_and_backward(): + torch.manual_seed(7) + m, n, k = 19, 13, 37 + x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) + weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16) + grad = torch.randn(m, n, device="cuda", dtype=torch.bfloat16) + bias = torch.randn(n, device="cuda", dtype=torch.bfloat16) + scale_x, scale_w, scale_g = _scale(x), _scale(weight), _scale(grad) + amax_x = torch.empty(1, device="cuda", dtype=torch.float32) + amax_w = torch.empty(1, device="cuda", dtype=torch.float32) + amax_g = torch.empty(1, device="cuda", dtype=torch.float32) + module = get_module("fp8_mm") + + out = module.fp8_linear_forward_scaled( + x, + weight, + bias, + scale_x, + scale_w, + scale_x.reciprocal(), + scale_w.reciprocal(), + amax_x, + amax_w, + ) + grad_x, grad_w, grad_b = module.fp8_linear_backward_scaled( + grad, + x, + weight, + [1, 1, 1], + scale_g, + scale_w, + scale_x, + scale_g.reciprocal(), + scale_w.reciprocal(), + scale_x.reciprocal(), + amax_g, + ) + + qx = _quantize(x, scale_x) + qw = _quantize(weight, scale_w) + qg = _quantize(grad, scale_g) + expected_out = (qx @ qw.t() * scale_x * scale_w + bias).to(torch.bfloat16) + expected_grad_x = (qg @ qw * scale_g * scale_w).to(torch.bfloat16) + expected_grad_w = (qg.t() @ qx * scale_g * scale_x).to(torch.bfloat16) + + torch.testing.assert_close(out, expected_out, atol=0.125, rtol=0.01) + torch.testing.assert_close(grad_x, expected_grad_x, atol=0.125, rtol=0.01) + torch.testing.assert_close(grad_w, expected_grad_w, atol=0.125, rtol=0.01) + torch.testing.assert_close(grad_b, grad.sum(0).to(torch.bfloat16)) + torch.testing.assert_close(amax_x, x.abs().amax().float().reshape(1)) + torch.testing.assert_close(amax_w, weight.abs().amax().float().reshape(1)) + torch.testing.assert_close(amax_g, grad.abs().amax().float().reshape(1))