perf: optimize fused FP8 GEMM kernel

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);
+    x/w: [..., K] / [N, K] bf16; sx/sw and their inverses control the fused
-    sx_inv/sw_inv: 1/scale; amax_x/amax_w: f32 buffers receiving max-abs.
+    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}")

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

csrc/kernels/fp8_mm.cu

@@ -1,444 +1,583 @@
-// FP8 e4m3 matrix multiply via cuBLASLt (sm89 TN layout).
+// Fused BF16 -> E4M3 MMA -> BF16 matrix multiplication
-//
-// 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.
 
 #include <torch/extension.h>
 #include <ATen/cuda/CUDAContext.h>
 #include <c10/cuda/CUDAGuard.h>
-#include <cublasLt.h>
 #include <cuda_fp8.h>
+#include <cuda_runtime.h>
 #include <cstdint>
 #include <mutex>
 #include <unordered_map>
 
-static std::recursive_mutex g_mutex;
+namespace {
 
-static cublasLtHandle_t g_handle = nullptr;
+constexpr int kMmaM = 16;
-static cublasLtMatmulDesc_t g_desc = nullptr;
+constexpr int kMmaN = 8;
-static cublasLtMatrixLayout_t g_layout_a = nullptr;
+constexpr int kMmaK = 32;
-static cublasLtMatrixLayout_t g_layout_b = nullptr;
+constexpr int kBlockM = 32;
-static cublasLtMatrixLayout_t g_layout_c = nullptr;
+constexpr int kBlockN = 32;
-static cublasLtMatmulPreference_t g_pref = nullptr;
+constexpr int kWarps = 8;
-static void* g_workspace = nullptr;
+constexpr int kForwardBlockM = 64;
-static size_t g_ws_size = 0;
+constexpr int kForwardBlockN = 64;
 
-struct ShapeKey {
+__device__ __forceinline__ unsigned pack_fp8x4_scalar(float x0, float x1,
-    int64_t m;
+                                                       float x2, float x3) {
-    int64_t k;
+    __nv_fp8_e4m3 q0(x0);
-    int64_t n;
+    __nv_fp8_e4m3 q1(x1);
-    bool operator==(const ShapeKey& other) const {
+    __nv_fp8_e4m3 q2(x2);
-        return m == other.m && k == other.k && n == other.n;
+    __nv_fp8_e4m3 q3(x3);
+    return static_cast<unsigned>(q0.__x) |
+           (static_cast<unsigned>(q1.__x) << 8) |
+           (static_cast<unsigned>(q2.__x) << 16) |
+           (static_cast<unsigned>(q3.__x) << 24);
+}
+
+__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<unsigned>(low) | (static_cast<unsigned>(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 <bool Transpose>
+__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<unsigned*>(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<const uint4*>(source);
+    asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;"
+                 :: "r"(shared_address), "l"(source_vec),
+                    "r"(valid ? 16 : 0));
+}
+
+template <bool TrackAmax = true, bool VectorPack = true>
+__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);
+    }
+}
+
+template <bool TransposeA, bool TransposeB>
+__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<TransposeA>(a, row, col, m, k, amax_a) * inv_a;
+    };
+    auto load_b = [&](int64_t col) {
+        return load_bf16<TransposeB>(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));
+}
+
+template <bool TransposeA, bool TransposeB, bool AddBias, int BlockM = kBlockM,
+          int BlockN = kBlockN, bool TrackAmax = true>
+__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<uintptr_t>(a_source) & 15) == 0;
+                const bool aligned_b =
+                    (reinterpret_cast<uintptr_t>(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);
+                    }
+                }
+            }
         }
     };
 
-struct ShapeKeyHash {
+    if constexpr (AsyncContiguous) {
-    size_t operator()(const ShapeKey& s) const {
+        load_bf16_tile(0, 0);
-        size_t h = std::hash<int64_t>()(s.m);
+        asm volatile("cp.async.commit_group;");
-        h ^= std::hash<int64_t>()(s.k) + 0x9e3779b9 + (h << 6) + (h >> 2);
+        asm volatile("cp.async.wait_group 0;");
-        h ^= std::hash<int64_t>()(s.n) + 0x9e3779b9 + (h << 6) + (h >> 2);
+        __syncthreads();
-        return h;
-    }
-};
-
-using AlgoCache = std::unordered_map<ShapeKey, cublasLtMatmulAlgo_t, ShapeKeyHash>;
-
-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);
     }
 
-static void ensure_cublas_lt() {
+    const int64_t tile_count = (k + kMmaK - 1) / kMmaK;
-    std::lock_guard<std::recursive_mutex> lock(g_mutex);
+    for (int64_t tile_index = 0; tile_index < tile_count; ++tile_index) {
-    if (g_handle) {
+        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<unsigned*>(&a_tile[buffer][quant_row][quant_col]) =
+                load_fp8x4_from_bf16<TrackAmax, !TransposeA && !TransposeB>(
+                    a_source, inv_a, local_amax_a, track_amax_a);
+            *reinterpret_cast<unsigned*>(&a_tile[buffer][quant_row][quant_col + 4]) =
+                load_fp8x4_from_bf16<TrackAmax, !TransposeA && !TransposeB>(
+                    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<unsigned*>(
+                    &b_tile[buffer][quant_row][quant_col]) =
+                    load_fp8x4_from_bf16<TrackAmax, !TransposeA && !TransposeB>(
+                        b_source, inv_b, local_amax_b, track_amax_b);
+                *reinterpret_cast<unsigned*>(
+                    &b_tile[buffer][quant_row][quant_col + 4]) =
+                    load_fp8x4_from_bf16<TrackAmax, !TransposeA && !TransposeB>(
+                        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<TransposeA, TransposeB>(
+                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<unsigned*>(
+                &b_tile[buffer][b_row][fragment_col]);
+            b_frag[1] = *reinterpret_cast<unsigned*>(
+                &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<unsigned*>(
+                    &a_tile[buffer][a_row0 + m_offset][fragment_col]);
+                a_frag[1] = *reinterpret_cast<unsigned*>(
+                    &a_tile[buffer][a_row1 + m_offset][fragment_col]);
+                a_frag[2] = *reinterpret_cast<unsigned*>(
+                    &a_tile[buffer][a_row0 + m_offset][fragment_col + 16]);
+                a_frag[3] = *reinterpret_cast<unsigned*>(
+                    &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 <bool TransposeA, bool TransposeB, bool AddBias = false,
+          int BlockM = kBlockM, int BlockN = kBlockN, bool TrackAmax = true>
+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<const __nv_bfloat16*>(bias.data_ptr())
+        : nullptr;
+    fused_fp8_gemm_kernel<TransposeA, TransposeB, AddBias, BlockM, BlockN,
+                          TrackAmax>
+        <<<grid, kWarps * 32, 0, stream>>>(
+        reinterpret_cast<const __nv_bfloat16*>(a.data_ptr()),
+        reinterpret_cast<const __nv_bfloat16*>(b.data_ptr()),
+        reinterpret_cast<__nv_bfloat16*>(out.data_ptr()), bias_ptr,
+        scale_a.data_ptr<float>(), scale_b.data_ptr<float>(),
+        amax_a ? amax_a->data_ptr<float>() : nullptr,
+        amax_b ? amax_b->data_ptr<float>() : nullptr, m, n, k);
+}
+
+void check_fp8_device(const torch::Tensor& tensor) {
+    static std::mutex mutex;
+    static std::unordered_map<int, bool> supported;
+    const int device = tensor.device().index();
+    {
+        std::lock_guard<std::mutex> 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;
         }
-    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,
+    const auto* properties = at::cuda::getDeviceProperties(device);
-                                      AlgoCache* cache,
+    const bool is_supported = properties->major > 8 ||
-                                      cublasLtMatmulAlgo_t* algo);
+                              (properties->major == 8 && properties->minor >= 9);
-
+    {
-static void fp8_gemm_into(torch::Tensor lhs, torch::Tensor rhs, torch::Tensor out,
+        std::lock_guard<std::mutex> lock(mutex);
-                          int64_t m, int64_t k, int64_t n,
+        supported.emplace(device, is_supported);
-                          const float* a_scale, const float* b_scale,
+    }
-                          cudaStream_t stream);
+    TORCH_CHECK(is_supported,
-
+                "fused FP8 MMA requires compute capability 8.9 or newer");
-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);
 }
 
-torch::Tensor fp8_mm(torch::Tensor a, torch::Tensor b) {
+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(a.scalar_type() == torch::kBFloat16 &&
-    TORCH_CHECK(b.scalar_type() == torch::kFloat8_e4m3fn, "b must be float8_e4m3fn");
+                    b.scalar_type() == torch::kBFloat16,
-    TORCH_CHECK(a.dim() == 2 && b.dim() == 2, "2D tensors required");
+                "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);
+    auto out = torch::empty({a_c.size(0), b_c.size(0)}, a_c.options());
-    TORCH_CHECK(b_c.size(1) == k, "inner dim mismatch");
+    torch::Tensor no_bias;
-
+    launch_fused_fp8_gemm<false, false, false, kForwardBlockM,
-    auto buf = torch::empty({m, n}, a_c.options().dtype(torch::kBFloat16));
+                          kForwardBlockN, false>(
-    ensure_cublas_lt();
+        a_c, b_c, out, no_bias, sx, sw, nullptr, nullptr,
-    fp8_gemm_into(a_c, b_c, buf, m, k, n, &k_scale_one, &k_scale_one,
+        a_c.size(0), b_c.size(0), a_c.size(1), stream.stream());
-                  stream.stream());
+    C10_CUDA_CHECK(cudaGetLastError());
-    return buf;
+    return out;
 }
 
-
+torch::Tensor fp8_linear_forward_scaled(
-// ---------------------------------------------------------------------------
+    torch::Tensor x, torch::Tensor w, torch::Tensor bias, torch::Tensor sx,
-// Quantize: bf16 * scale_inv -> fp8, one atomicMax amax per kernel call.
+    torch::Tensor sw, torch::Tensor sx_inv, torch::Tensor sw_inv,
-// amax_ptr must be zeroed before launch; float-bits atomicMax works because
+    torch::Tensor amax_x, torch::Tensor amax_w) {
-// |v| >= 0 has a monotonic IEEE bit pattern.
-// ---------------------------------------------------------------------------
-
-template <typename T8>
-__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 <typename T8>
-__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<T8>(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<unsigned*>(amax_ptr), __float_as_uint(m));
-    }
-}
-
-// Same but with a transpose (rows x cols bf16 row-major -> fp8 [cols, rows]).
-template <typename T8>
-__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<T8>(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<unsigned*>(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<std::recursive_mutex> 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<std::recursive_mutex> 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_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");
+    TORCH_CHECK(w_c.dim() == 2 && w_c.size(1) == k, "inner dim mismatch");
-    ensure_cublas_lt();
+    C10_CUDA_CHECK(cudaMemsetAsync(amax_x.data_ptr<float>(), 0, sizeof(float),
-
+                                   stream.stream()));
-    const float* sx_ptr = sx.data_ptr<float>();
+    C10_CUDA_CHECK(cudaMemsetAsync(amax_w.data_ptr<float>(), 0, sizeof(float),
-    const float* sw_ptr = sw.data_ptr<float>();
+                                   stream.stream()));
-    const float* sxi_ptr = sx_inv.data_ptr<float>();
-    const float* swi_ptr = sw_inv.data_ptr<float>();
-    float* amax_x_ptr = amax_x.data_ptr<float>();
-    float* amax_w_ptr = amax_w.data_ptr<float>();
-    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<const __nv_bfloat16*>(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<const __nv_bfloat16*>(w_c.data_ptr()), swi_ptr,
-            reinterpret_cast<__nv_fp8_e4m3*>(w8.data_ptr()), amax_w_ptr, n * k);
-    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,
+        TORCH_CHECK(bias.is_cuda() && bias.device() == x.device() &&
-                    "bias must be bf16 with shape [N]");
+                        bias.scalar_type() == torch::kBFloat16 &&
-        bias_add_bf16_kernel<<<(unsigned)((m * n + block - 1) / block), block, 0, stream>>>(
+                        bias.numel() == n,
-            reinterpret_cast<__nv_bfloat16*>(out.data_ptr()),
+                    "bias must be CUDA bf16 with shape [N]");
-            reinterpret_cast<const __nv_bfloat16*>(bias.data_ptr()), m * n, n);
+        launch_fused_fp8_gemm<false, false, true, kForwardBlockM,
-        C10_CUDA_CHECK(cudaGetLastError());
+                              kForwardBlockN>(
+            x_c, w_c, out, bias, sx, sw, &amax_x, &amax_w,
+            m, n, k, stream.stream());
+    } else {
+        launch_fused_fp8_gemm<false, false, false, kForwardBlockM,
+                              kForwardBlockN>(
+            x_c, w_c, out, bias, sx, sw, &amax_x, &amax_w,
+            m, n, k, stream.stream());
     }
+    C10_CUDA_CHECK(cudaGetLastError());
 
+    (void)sx_inv;
+    (void)sw_inv;
     std::vector<int64_t> 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<torch::Tensor, torch::Tensor, torch::Tensor> fp8_linear_backward_scaled(
     torch::Tensor g, torch::Tensor x, torch::Tensor w,
     std::vector<int64_t> 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.is_cuda() && x.is_cuda() && w.is_cuda(), "CUDA tensors required");
-    TORCH_CHECK(g.dtype() == torch::kBFloat16 && x.dtype() == torch::kBFloat16 &&
+    TORCH_CHECK(g.scalar_type() == torch::kBFloat16 &&
-                    w.dtype() == 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 m = g_c.size(0), n = w_c.size(0), k = w_c.size(1);
-    int64_t n = w.size(0);
-    int64_t k = w.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()));
+    auto grad_bias = torch::empty({0}, g.options());
-    ensure_cublas_lt();
+    C10_CUDA_CHECK(cudaMemsetAsync(amax_g.data_ptr<float>(), 0, sizeof(float),
-
+                                   stream.stream()));
-    const float* sg_ptr = sg.data_ptr<float>();
+    torch::Tensor no_bias;
-    const float* sw_ptr = sw.data_ptr<float>();
+    bool recorded_amax = false;
-    const float* sx_ptr = sx.data_ptr<float>();
-    const float* sgi_ptr = sg_inv.data_ptr<float>();
-    const float* swi_ptr = sw_inv.data_ptr<float>();
-    const float* sxi_ptr = sx_inv.data_ptr<float>();
-    float* amax_g_ptr = amax_g.data_ptr<float>();
-    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<const __nv_bfloat16*>(g_c.data_ptr()), sgi_ptr,
-            reinterpret_cast<__nv_fp8_e4m3*>(g8.data_ptr()), amax_g_ptr, m * n);
-    dim3 threads(32, 8);
     if (masks[0]) {
-        dim3 blocks((k + 31) / 32, (n + 31) / 32);
+        auto grad_input_2d = grad_input.reshape({m, k});
-        transpose_quantize_kernel<__nv_fp8_e4m3>
+        launch_fused_fp8_gemm<false, true>(
-            <<<blocks, threads, 0, stream.stream()>>>(
+            g_c, w_c, grad_input_2d, no_bias, sg, sw, &amax_g, nullptr,
-                reinterpret_cast<const __nv_bfloat16*>(w_c.data_ptr()), swi_ptr,
+            m, k, n, stream.stream());
-                reinterpret_cast<__nv_fp8_e4m3*>(wt8.data_ptr()),
+        recorded_amax = true;
-                amax_t.data_ptr<float>(), n, k);
-        fp8_gemm_into(g8, wt8, grad_input.reshape({m, k}), m, n, k, sg_ptr,
-                      sw_ptr, stream.stream());
     }
     if (masks[1]) {
-        dim3 g_blocks((n + 31) / 32, (m + 31) / 32);
+        launch_fused_fp8_gemm<true, true>(
-        dim3 x_blocks((k + 31) / 32, (m + 31) / 32);
+            g_c, x_c, grad_weight, no_bias, sg, sx,
-        transpose_quantize_kernel<__nv_fp8_e4m3>
+            recorded_amax ? nullptr : &amax_g, nullptr,
-            <<<g_blocks, threads, 0, stream.stream()>>>(
+            n, k, m, stream.stream());
-                reinterpret_cast<const __nv_bfloat16*>(g_c.data_ptr()), sgi_ptr,
+        recorded_amax = true;
-                reinterpret_cast<__nv_fp8_e4m3*>(gt8.data_ptr()),
+    }
-                amax_t.data_ptr<float>(), m, n);
+    if (!recorded_amax) {
-        transpose_quantize_kernel<__nv_fp8_e4m3>
+        amax_g.copy_(g_c.abs().amax().to(torch::kFloat32));
-            <<<x_blocks, threads, 0, stream.stream()>>>(
-                reinterpret_cast<const __nv_bfloat16*>(x_c.data_ptr()), sxi_ptr,
-                reinterpret_cast<__nv_fp8_e4m3*>(xt8.data_ptr()),
-                amax_t.data_ptr<float>(), m, k);
-        fp8_gemm_into(gt8, xt8, grad_weight, n, m, k, sg_ptr, sx_ptr,
-                      stream.stream());
     }
     C10_CUDA_CHECK(cudaGetLastError());
-    if (masks[2]) {
+    if (masks[2]) grad_bias = g_c.sum(0).to(g.scalar_type());
-        grad_bias = g_c.sum(0).to(g.dtype());
+
-    }
+    (void)sg_inv;
-    return std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>(
+    (void)sw_inv;
-        grad_input, grad_weight, grad_bias);
+    (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"),
+    m.def("fp8_mm", &fp8_mm, py::arg("a"), py::arg("b"), py::arg("sx"),
-          "FP8 e4m3 GEMM: a[M,K] x b[N,K] -> bf16[M,N] (pre-scaled inputs)");
+          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 + "
+          "Fused BF16-to-FP8 linear forward with FP32 accumulation");
-          "cublasLt GEMM (scales applied inside) + bias -> bf16");
     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, "
+          "Fused BF16-to-FP8 linear backward with FP32 accumulation");
-          "dB = sum(g)");
 }

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 <cuda_fp8.h>
+
+#include <algorithm>
+#include <vector>
+
+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<unsigned>(q0.__x) |
+           (static_cast<unsigned>(q1.__x) << 8) |
+           (static_cast<unsigned>(q2.__x) << 16) |
+           (static_cast<unsigned>(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<float>(__nv_fp8_e4m3(value));
+}
+
+int main() {
+    srand(0);
+    std::vector<float> a(M * K), b(N * K), reference(M * N, 0.0f);
+    std::vector<bf16> 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;
+}

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))