KernelBench hard · H100
FP8 GEMM Kimi K3 (1M)
22.9%geomean peak fraction across shapes
manually audited: clean
Clean cell. The submission computes the real scaled FP8 e4m3 GEMM from live x, weight, and per-output-channel scale using Triton tensor-core dot kernels with an optional genuine CUTLASS SM90 path. Odd K is handled by exact zero-padding. Its caches hold compiled code, a shape-selected implementation, a weight-derived padded tensor invalidated by data_ptr/_version, CUTLASS capability metadata, and an L2-flush allocation; no output or input-derived answer is cached. There is no CUDA graph, constant answer, forbidden scaled-mm path, cross-run solution read, grader mutation, tolerance change, or numeric-stress bypass. The unmodified official checker passed, and the four exact benchmark timings independently yield peak_fraction 0.2290.
harnesskinetic-claudeagent session6h 4mtotal wall—check38sbenchmark38soutput tokens198,279cost$327.37gpu-lock wait1h 7mgpu-lock held3mregimecompute
Per-shape vs governing ceilingeach shape graded against whichever binds — fp8 compute or HBM bandwidth
4096×4096×40960.233 ms38.9%589 TFLOPS · 39% of 1,513 TF fp8 peak · also 0.29 TB/s (14% of HBM)
4096×4096×41270.265 ms34.6%523 TFLOPS · 35% of 1,513 TF fp8 peak · also 0.25 TB/s (12% of HBM)
32×8192×81920.057 ms5.0%1.19 TB/s · 59% of 2.0 TB/s HBM · also 76 TFLOPS (5% of compute)
4096×14336×40960.777 ms40.9%619 TFLOPS · 41% of 1,513 TF fp8 peak · also 0.25 TB/s (12% of HBM)
compute-bound memory-bound · bar + right column = official fraction of the ceiling (the geomean input)
geomean(38.9% · 34.6% · 5.0% · 40.9%) = 22.9%
Kernel source (redacted)
"""FP8 e4m3 x e4m3 GEMM with per-output-channel dequant scale.
y = (x @ weight.T) * weight_scale in bf16.
Engines:
1. CUTLASS SM90 warp-specialized FP8 GEMM (TMA + wgmma, FP32 accumulate,
custom EVT epilogue applying the per-N-channel scale) for compute-bound
shapes. Vendored headers live in `cutlass_3.9/` next to this file; the
extension is JIT-built on first use and silently falls back if it can't
be built (the Triton path below handles every shape on its own).
2. Triton tl.dot fp8 kernel (tensor-core MMA, fp32 accumulate) for
skinny-M shapes and as a universal fallback.
3. K not a multiple of 16: operands are zero-padded along K (the pad adds
zeros, which is mathematically exact) and the weight pad is cached
(input-derived, invalidated by tensor version).
Engine choice per (M, N, K) is made by a tiny one-shot micro-tuner on first
encounter of a shape, so any shape gets the best available backend.
"""
import os
import time
import torch
import torch.nn as nn
import triton
import triton.language as tl
E4M3_MAX = 448.0
# ---------------------------------------------------------------------------
# Triton kernels
# ---------------------------------------------------------------------------
def _gen_configs():
cfgs = []
for bm, bn, bk, g, w, s in [
(128, 128, 64, 8, 8, 4),
(128, 128, 128, 8, 8, 3),
(128, 256, 64, 8, 8, 3),
(128, 256, 128, 8, 8, 3),
(64, 128, 128, 4, 4, 4),
(64, 256, 128, 4, 4, 4),
(128, 64, 128, 4, 4, 4),
]:
cfgs.append(
triton.Config(
{"BM": bm, "BN": bn, "BK": bk, "GROUP_M": g}, num_warps=w, num_stages=s
)
)
for bm, bn, bk, g, w, s in [
(16, 128, 128, 1, 4, 4),
(16, 128, 128, 1, 4, 5),
(16, 128, 256, 1, 4, 3),
(32, 128, 128, 1, 4, 5),
]:
cfgs.append(
triton.Config(
{"BM": bm, "BN": bn, "BK": bk, "GROUP_M": g}, num_warps=w, num_stages=s
)
)
return cfgs
@triton.autotune(configs=_gen_configs(), key=["M", "N", "K"])
@triton.jit
def _fp8_gemm_kernel(
a_ptr, b_ptr, c_ptr, s_ptr,
M, N, K,
stride_am, stride_ak,
stride_bn, stride_bk,
stride_cm, stride_cn,
BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr, GROUP_M: tl.constexpr,
):
pid = tl.program_id(0)
num_pid_m = tl.cdiv(M, BM)
num_pid_n = tl.cdiv(N, BN)
num_pid_in_group = GROUP_M * num_pid_n
group_id = pid // num_pid_in_group
first_pid_m = group_id * GROUP_M
group_size_m = min(num_pid_m - first_pid_m, GROUP_M)
pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
pid_n = (pid % num_pid_in_group) // group_size_m
offs_m = pid_m * BM + tl.arange(0, BM)
offs_n = pid_n * BN + tl.arange(0, BN)
offs_k = tl.arange(0, BK)
m_mask = offs_m < M
n_mask = offs_n < N
a_ptrs = a_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak
b_ptrs = b_ptr + offs_n[None, :] * stride_bn + offs_k[:, None] * stride_bk
acc = tl.zeros((BM, BN), dtype=tl.float32)
k_full = K // BK
for _ in range(0, k_full):
a = tl.load(a_ptrs, mask=m_mask[:, None], other=0.0)
b = tl.load(b_ptrs, mask=n_mask[None, :], other=0.0)
# max_num_imprecise_acc=BK: promote from imprecise wgmma accumulation
# into fp32 FFMA every block (correctness for long K; ~5-10% cost).
acc = tl.dot(a, b, acc, max_num_imprecise_acc=BK)
a_ptrs += BK * stride_ak
b_ptrs += BK * stride_bk
if K % BK != 0:
k_rem = K - k_full * BK
a = tl.load(
a_ptrs, mask=m_mask[:, None] & (offs_k[None, :] < k_rem), other=0.0
)
b = tl.load(
b_ptrs, mask=n_mask[None, :] & (offs_k[:, None] < k_rem), other=0.0
)
acc = tl.dot(a, b, acc, max_num_imprecise_acc=BK)
s = tl.load(s_ptr + offs_n, mask=n_mask, other=0.0)
c = acc * s[None, :]
c_ptrs = c_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn
tl.store(c_ptrs, c.to(tl.bfloat16), mask=m_mask[:, None] & n_mask[None, :])
def _triton_gemm(x, w, w_scale, y=None):
M, K = x.shape
N = w.shape[0]
if y is None:
y = torch.empty((M, N), dtype=torch.bfloat16, device=x.device)
grid = lambda meta: ( # noqa: E731
triton.cdiv(M, meta["BM"]) * triton.cdiv(N, meta["BN"]),
)
_fp8_gemm_kernel[grid](
x, w, y, w_scale,
M, N, K,
x.stride(0), x.stride(1),
w.stride(0), w.stride(1),
y.stride(0), y.stride(1),
)
return y
# --- K-padding kernels ------------------------------------------------------
@triton.jit
def _pad_k_kernel(src, dst, R, K, K_PADDED, BK: tl.constexpr):
pid_m = tl.program_id(0)
pid_k = tl.program_id(1)
offs_k = pid_k * BK + tl.arange(0, BK)
mask = offs_k < K
val = tl.load(src + pid_m * K + offs_k, mask=mask, other=0.0)
tl.store(dst + pid_m * K_PADDED + offs_k, val, mask=offs_k < K_PADDED)
@triton.jit
def _pad_k_shl(src_i32, dst_i32, R, K, K_PAD, TOTAL_SRC_DW, BLOCK: tl.constexpr):
# dst-side aligned int32 ops; src bytes gathered via funnel shift of two
# aligned dwords (little-endian). Zeros beyond K.
pid = tl.program_id(0)
offs = pid * BLOCK + tl.arange(0, BLOCK)
kpad_dw = K_PAD // 4
total_dst_dw = R * kpad_dw
valid = offs < total_dst_dw
r = offs // kpad_dw
c = (offs % kpad_dw) * 4
p = r * K + c
inside = c < K
d0 = p // 4
sh8 = (p % 4) * 8
v0 = tl.load(src_i32 + d0, mask=valid & inside, other=0)
v1 = tl.load(
src_i32 + d0 + 1, mask=valid & inside & (sh8 > 0) & (d0 + 1 < TOTAL_SRC_DW),
other=0,
)
sh32 = tl.where(sh8 > 0, 32 - sh8, 0)
u0 = v0.to(tl.uint32, bitcast=True)
u1 = v1.to(tl.uint32, bitcast=True)
val = (u0 >> sh8) | tl.where(sh8 > 0, u1 << sh32, 0)
rem = K - c
keep = tl.where(rem >= 4, -1, (1 << (rem * 8)) - 1).to(tl.uint32, bitcast=True)
val = tl.where(inside & valid, val & keep, 0)
tl.store(dst_i32 + offs, val.to(tl.int32, bitcast=True), mask=valid)
def _pad_k(t, k_pad):
"""(R, K) opaque 1-byte-dtype -> (R, k_pad), zero-filled tail."""
R, K = t.shape
out = torch.empty((R, k_pad), dtype=t.dtype, device=t.device)
if K >= 4 and (R * K) % 4 == 0 and k_pad % 4 == 0:
src = t.view(-1).view(torch.int32)
dst = out.view(-1).view(torch.int32)
BLOCK = 512
grid = (triton.cdiv(R * (k_pad // 4), BLOCK),)
_pad_k_shl[grid](src, dst, R, K, k_pad, (R * K + 3) // 4, BLOCK=BLOCK,
num_warps=4)
return out
BK = 512
grid = (R, triton.cdiv(k_pad, BK))
_pad_k_kernel[grid](t, out, R, K, k_pad, BK=BK)
return out
class _WPadCache:
"""Cached zero-padded-K copy of a weight tensor, invalidated by version."""
def __init__(self):
self._entry = None
def get(self, w, k_pad):
key = (w.data_ptr(), w._version, tuple(w.shape), k_pad, w.device.index)
if self._entry is not None and self._entry[0] == key:
return self._entry[1]
t = _pad_k(w, k_pad)
self._entry = (key, t)
return t
_w_pad_cache = _WPadCache()
# ---------------------------------------------------------------------------
# CUTLASS extension (optional, JIT-built from vendored headers)
# ---------------------------------------------------------------------------
_HERE = os.path.dirname(os.path.abspath(__file__))
_CUDA_SRC = r"""
// FP8 e4m3 x e4m3 GEMM: D = (A @ B^T) * scale_n, bf16 out.
// A: (M,K) fp8 row-major (K-contig). B: (N,K) fp8 row-major (K-contig)
// == CUTLASS LayoutB ColumnMajor (K,N). scale: (N,) fp32. D: (M,N) bf16 row-major.
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include "cutlass/cutlass.h"
#include "cutlass/numeric_types.h"
#include "cute/tensor.hpp"
#include "cutlass/gemm/dispatch_policy.hpp"
#include "cutlass/gemm/collective/collective_builder.hpp"
#include "cutlass/gemm/device/gemm_universal_adapter.h"
#include "cutlass/gemm/kernel/gemm_universal.hpp"
#include "cutlass/epilogue/dispatch_policy.hpp"
#include "cutlass/epilogue/collective/collective_builder.hpp"
#include "cutlass/epilogue/fusion/sm90_callbacks_tma_warpspecialized.hpp"
#include "cutlass/epilogue/fusion/sm90_visitor_tma_warpspecialized.hpp"
#include "cutlass/epilogue/fusion/sm90_visitor_compute_tma_warpspecialized.hpp"
#include "cutlass/epilogue/fusion/sm90_visitor_load_tma_warpspecialized.hpp"
using namespace cute;
namespace fp8gemm {
using ElementA = cutlass::float_e4m3_t;
using LayoutA = cutlass::layout::RowMajor;
constexpr int AlignmentA = 16;
using ElementB = cutlass::float_e4m3_t;
using LayoutB = cutlass::layout::ColumnMajor;
constexpr int AlignmentB = 16;
using ElementC = cutlass::bfloat16_t;
using LayoutC = cutlass::layout::RowMajor;
constexpr int AlignmentC = 8;
using ElementD = cutlass::bfloat16_t;
using LayoutD = cutlass::layout::RowMajor;
constexpr int AlignmentD = 8;
using ElementAccumulator = float;
using ElementCompute = float;
using ArchTag = cutlass::arch::Sm90;
using OperatorClass = cutlass::arch::OpClassTensorOp;
template <class TileShapeMNK, class ClusterShapeMNK, class KernelSched,
class EpiSched = cutlass::epilogue::TmaWarpSpecializedCooperative>
struct GemmTraits {
using TileShape = TileShapeMNK;
using ClusterShape = ClusterShapeMNK;
using KernelSchedule = KernelSched;
using EpilogueSchedule = EpiSched;
static constexpr auto RoundStyle = cutlass::FloatRoundStyle::round_to_nearest;
// D = acc * per-N scale (row broadcast of an (N,) vector, in fp32, cast bf16)
using CustomEVT = cutlass::epilogue::fusion::Sm90EVT<
cutlass::epilogue::fusion::Sm90Compute<cutlass::multiplies, ElementD, ElementCompute, RoundStyle>,
cutlass::epilogue::fusion::Sm90RowBroadcast<
0, TileShapeMNK, float, ElementCompute, Stride<_0, _1, int64_t>>,
cutlass::epilogue::fusion::Sm90AccFetch>;
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
ArchTag, OperatorClass,
TileShape, ClusterShape,
cutlass::epilogue::collective::EpilogueTileAuto,
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
EpilogueSchedule,
CustomEVT>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
ArchTag, OperatorClass,
ElementA, LayoutA, AlignmentA,
ElementB, LayoutB, AlignmentB,
ElementAccumulator,
TileShape, ClusterShape,
cutlass::gemm::collective::StageCountAutoCarveout<
static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
KernelSchedule>::CollectiveOp;
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue,
cutlass::gemm::PersistentScheduler>;
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
};
template <class Traits>
inline cutlass::Status gemm_can_implement(const typename Traits::Gemm::Arguments& args) {
return Traits::Gemm::can_implement(args);
}
template <class Traits>
void run_gemm(const torch::Tensor& a, const torch::Tensor& b,
const torch::Tensor& scale, torch::Tensor& d,
int max_swizzle = 1, int raster_dir = 1, int promo_interval = 4) {
using Gemm = typename Traits::Gemm;
int M = a.size(0), K = a.size(1), N = b.size(0);
typename Gemm::Arguments args{
cutlass::gemm::GemmUniversalMode::kGemm,
{M, N, K, 1},
{(const ElementA*)a.data_ptr(), cute::make_stride(int64_t(K), Int<1>{}, int64_t(0)),
(const ElementB*)b.data_ptr(), cute::make_stride(int64_t(K), Int<1>{}, int64_t(0))},
{// epilogue
{ // fusion args (EVT): {row_broadcast args, acc args(empty), compute args(empty)}
{(const float*)scale.data_ptr(), 0.0f, {_0{}, _1{}, int64_t(0)}},
{},
{},
},
(const ElementC*)d.data_ptr(), // C unused; alias D (beta path never taken)
cute::make_stride(int64_t(N), Int<1>{}, int64_t(0)),
(ElementD*)d.data_ptr(),
cute::make_stride(int64_t(N), Int<1>{}, int64_t(0))}};
args.scheduler.max_swizzle_size = max_swizzle;
args.scheduler.raster_order =
raster_dir == 0
? cutlass::gemm::kernel::detail::RasterOrderOptions::AlongM
: cutlass::gemm::kernel::detail::RasterOrderOptions::AlongN;
// Fast-accum FFMA promotion interval (multiple of TileK/GMMA_K = 4)
args.mainloop.mma_promotion_interval = (uint32_t)promo_interval;
// can_implement + workspace size are pointer-independent; cache per shape/schedule.
struct ShapeCache {
cutlass::Status status = cutlass::Status::kInvalid;
size_t ws_size = 0;
int max_swizzle = -1;
int raster_dir = -1;
int M = -1, N = -1, K = -1;
};
static thread_local ShapeCache cache;
if (cache.M != M || cache.N != N || cache.K != K ||
cache.max_swizzle != max_swizzle || cache.raster_dir != raster_dir) {
cache.status = gemm_can_implement<Traits>(args);
cache.ws_size = cache.status == cutlass::Status::kSuccess
? Gemm::get_workspace_size(args)
: 0;
cache.max_swizzle = max_swizzle;
cache.raster_dir = raster_dir;
cache.M = M; cache.N = N; cache.K = K;
}
TORCH_CHECK(cache.status == cutlass::Status::kSuccess,
"CUTLASS can_implement failed: ", (int)cache.status);
torch::Tensor ws;
void* ws_ptr = nullptr;
if (cache.ws_size > 0) {
ws = torch::empty({(int64_t)cache.ws_size},
torch::TensorOptions().dtype(torch::kUInt8).device(a.device()));
ws_ptr = ws.data_ptr();
}
Gemm gemm;
auto status = gemm.initialize(args, ws_ptr);
TORCH_CHECK(status == cutlass::Status::kSuccess, "CUTLASS initialize failed: ",
(int)status);
status = gemm.run(at::cuda::getCurrentCUDAStream());
TORCH_CHECK(status == cutlass::Status::kSuccess, "CUTLASS run failed: ", (int)status);
}
} // namespace fp8gemm
torch::Tensor fp8_gemm_scaled(torch::Tensor a, torch::Tensor b, torch::Tensor scale,
int64_t variant, int64_t max_swizzle, int64_t raster_dir,
int64_t promo_interval) {
TORCH_CHECK(a.is_cuda() && b.is_cuda() && scale.is_cuda());
TORCH_CHECK(a.dtype() == torch::kFloat8_e4m3fn && b.dtype() == torch::kFloat8_e4m3fn);
TORCH_CHECK(scale.dtype() == torch::kFloat32);
TORCH_CHECK(a.is_contiguous() && b.is_contiguous() && scale.is_contiguous());
const auto& a_c = a;
const auto& b_c = b;
const auto& s_c = scale;
int M = a_c.size(0), N = b_c.size(0);
auto d = torch::empty({M, N}, torch::TensorOptions()
.dtype(torch::kBFloat16)
.device(a_c.device()));
using COOP_FA = cutlass::gemm::KernelTmaWarpSpecializedCooperativeFP8FastAccum;
using T1 = fp8gemm::GemmTraits<Shape<_128, _128, _128>, Shape<_2, _1, _1>, COOP_FA>;
using T2 = fp8gemm::GemmTraits<Shape<_128, _256, _128>, Shape<_2, _2, _1>, COOP_FA>;
using T3 = T1; // was: COOP non-fast-accum (slower)
using T4 = T1; // was: pingpong (broken with custom EVT)
using T5 = fp8gemm::GemmTraits<Shape<_128, _128, _128>, Shape<_1, _1, _1>, COOP_FA>;
using T6 = fp8gemm::GemmTraits<Shape<_256, _128, _128>, Shape<_2, _1, _1>, COOP_FA>;
using T7 = fp8gemm::GemmTraits<Shape<_128, _128, _128>, Shape<_4, _1, _1>, COOP_FA>;
using T8 = T1;
using T9 = T1;
using T10 = T1;
using T11 = fp8gemm::GemmTraits<Shape<_256, _128, _128>, Shape<_2, _2, _1>, COOP_FA>;
using T12 = fp8gemm::GemmTraits<Shape<_256, _128, _128>, Shape<_4, _1, _1>, COOP_FA>;
using T13 = fp8gemm::GemmTraits<Shape<_128, _128, _256>, Shape<_2, _1, _1>, COOP_FA>;
using T14 = fp8gemm::GemmTraits<Shape<_128, _256, _128>, Shape<_2, _1, _1>, COOP_FA>;
using T15 = T1;
using T16 = T1;
using T17 = T1;
using T18 = T1;
using T19 = T1;
switch (variant) {
case 0: fp8gemm::run_gemm<T1>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
case 1: fp8gemm::run_gemm<T2>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
case 2: fp8gemm::run_gemm<T3>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
case 3: fp8gemm::run_gemm<T4>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
case 4: fp8gemm::run_gemm<T5>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
case 5: fp8gemm::run_gemm<T6>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
case 6: fp8gemm::run_gemm<T7>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
case 7: fp8gemm::run_gemm<T8>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
case 8: fp8gemm::run_gemm<T9>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
case 9: fp8gemm::run_gemm<T10>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
case 10: fp8gemm::run_gemm<T11>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
case 11: fp8gemm::run_gemm<T12>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
case 12: fp8gemm::run_gemm<T13>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
case 13: fp8gemm::run_gemm<T14>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
case 14: fp8gemm::run_gemm<T15>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
case 15: fp8gemm::run_gemm<T16>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
case 16: fp8gemm::run_gemm<T17>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
case 17: fp8gemm::run_gemm<T18>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
case 18: fp8gemm::run_gemm<T19>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
default: fp8gemm::run_gemm<T1>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval);
}
return d;
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("fp8_gemm_scaled", &fp8_gemm_scaled, "fp8 gemm with per-n scale");
}
"""
_CUTLASS_DIR = os.path.join(_HERE, "cutlass_3.9")
_EXT = None
_EXT_FAILED = False
def _load_cutlass_ext():
global _EXT, _EXT_FAILED
if _EXT is not None or _EXT_FAILED:
return _EXT
try:
from torch.utils.cpp_extension import load
build_dir = os.path.join(_HERE, "build")
os.makedirs(build_dir, exist_ok=True)
# torch's jit-compile lock persists after killed processes; a lock
# whose mtime is >15 min old is certainly abandoned — remove it.
lock_path = os.path.join(build_dir, "lock")
try:
if os.path.exists(lock_path):
age = time.time() - os.path.getmtime(lock_path)
if age > 900:
os.remove(lock_path)
except OSError:
pass
src_path = os.path.join(_HERE, "fp8_gemm_cutlass.cu")
try:
if not os.path.exists(src_path) or \
open(src_path).read().strip() != _CUDA_SRC.strip():
with open(src_path, "w") as f:
f.write(_CUDA_SRC.strip() + "\n")
except OSError:
pass
build_dir = os.path.join(_HERE, "build")
os.makedirs(build_dir, exist_ok=True)
_EXT = load(
name="fp8_gemm_cutlass",
sources=[os.path.join(_HERE, "fp8_gemm_cutlass.cu")],
extra_include_paths=[os.path.join(_CUTLASS_DIR, "include")],
extra_cuda_cflags=[
"-O3", "-std=c++17", "--expt-relaxed-constexpr",
"--expt-extended-lambda", "-gencode=arch=compute_90a,code=sm_90a",
"--use_fast_math", "-DNDEBUG",
],
extra_cflags=["-O3", "-std=c++17"],
build_directory=build_dir,
verbose=False,
)
except Exception:
_EXT = None
_EXT_FAILED = True
return _EXT
# CUTLASS (variant, max_swizzle, raster_order) candidates per problem size.
def _cutlass_candidates(M, N, K):
return [(0, 1, 1), (4, 1, 1), (5, 4, 1), (5, 8, 1), (1, 1, 1),
(10, 1, 1), (11, 1, 1), (7, 1, 1), (13, 1, 1)]
# ---------------------------------------------------------------------------
# Dispatch with one-shot micro-autotune per shape
# ---------------------------------------------------------------------------
_plan_cache = {}
_l2_scratch = None
def _flush_l2():
global _l2_scratch
if _l2_scratch is None:
_l2_scratch = torch.empty(128 * 1024 * 1024 // 4, dtype=torch.float32, device="cuda")
_l2_scratch.zero_()
def _bench_once(fn, iters=9):
fn() # warm
torch.cuda.synchronize()
times = []
for _ in range(iters):
_flush_l2()
torch.cuda.synchronize()
s = torch.cuda.Event(enable_timing=True)
e = torch.cuda.Event(enable_timing=True)
s.record()
fn()
e.record()
torch.cuda.synchronize()
times.append(s.elapsed_time(e))
times.sort()
return times[len(times) // 2]
def _make_plan(x, w, w_scale):
"""Choose the fastest engine for this (M, N, K). Returns a callable."""
M, K = x.shape
N = w.shape[0]
candidates = []
ext = None
if K % 16 == 0 and (2 * N) % 16 == 0:
ext = _load_cutlass_ext()
if ext is not None:
for v, sw, ro in _cutlass_candidates(M, N, K):
def cutlass_cand(a, b, c, v=v, sw=sw, ro=ro):
return ext.fp8_gemm_scaled(a, b, c, v, sw, ro, 4)
candidates.append(("cutlass", cutlass_cand))
if K % 16 != 0:
# zero-pad operands to a 16B-aligned K (exact: pad contributes 0s)
k_pad = (K + 127) // 128 * 128
ext2 = _load_cutlass_ext()
def triton_padded(a, b, c, k_pad=k_pad):
wp = _w_pad_cache.get(b, k_pad) # re-validated against b version
xp = _pad_k(a, k_pad)
return _triton_gemm(xp, wp, c)
candidates.append(("triton_pad", triton_padded))
if ext2 is not None and M * N >= 1 << 22:
for v, sw, ro in [(0, 1, 1), (10, 1, 1), (5, 4, 1)]:
def cutlass_pad(a, b, c, k_pad=k_pad, v=v, sw=sw, ro=ro):
wp = _w_pad_cache.get(b, k_pad)
xp = _pad_k(a, k_pad)
return ext2.fp8_gemm_scaled(xp, wp, c, v, sw, ro, 4)
candidates.append(("cutlass_pad", cutlass_pad))
# Plain triton always works.
candidates.append(("triton", _triton_gemm))
if len(candidates) == 1:
return candidates[0][1]
best, best_t, best_name = candidates[-1][1], float("inf"), None
for name, cand in candidates:
try:
cand(x, w, w_scale) # correctness smoke (throws on failure)
t = _bench_once(lambda: cand(x, w, w_scale))
except Exception:
continue
if t < best_t:
best, best_t, best_name = cand, t, name
return best
def fp8_gemm(x, w, w_scale):
if not x.is_contiguous():
x = x.contiguous()
if not w.is_contiguous():
w = w.contiguous()
M, K = x.shape
N = w.shape[0]
key = (M, N, K, x.device.index)
plan = _plan_cache.get(key)
if plan is None:
plan = _make_plan(x, w, w_scale)
_plan_cache[key] = plan
return plan(x, w, w_scale)
class Model(nn.Module):
"""Same interface as reference.Model: buffers `weight` (N,K) fp8 and
`weight_scale` (N,) fp32; forward returns y = (x @ w.T) * scale in bf16."""
def __init__(self, M: int, N: int, K: int):
super().__init__()
self.M, self.N, self.K = M, N, K
self.register_buffer("weight", torch.zeros(N, K, dtype=torch.float8_e4m3fn))
self.register_buffer("weight_scale", torch.ones(N, dtype=torch.float32))
def forward(self, x: torch.Tensor) -> torch.Tensor:
return fp8_gemm(x, self.weight, self.weight_scale)
M = 4096
N = 4096
K = 4096
def get_inputs():
x = (torch.rand(M, K) * 8 - 4).to(torch.float8_e4m3fn)
return [x]
def get_init_inputs():
return [M, N, K]
# ==================================================================
# ===== sidecar: fp8_gemm_cutlass.cu (11027 bytes, loaded by solution.py) =====
# ==================================================================
// FP8 e4m3 x e4m3 GEMM: D = (A @ B^T) * scale_n, bf16 out.
// A: (M,K) fp8 row-major (K-contig). B: (N,K) fp8 row-major (K-contig)
// == CUTLASS LayoutB ColumnMajor (K,N). scale: (N,) fp32. D: (M,N) bf16 row-major.
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include "cutlass/cutlass.h"
#include "cutlass/numeric_types.h"
#include "cute/tensor.hpp"
#include "cutlass/gemm/dispatch_policy.hpp"
#include "cutlass/gemm/collective/collective_builder.hpp"
#include "cutlass/gemm/device/gemm_universal_adapter.h"
#include "cutlass/gemm/kernel/gemm_universal.hpp"
#include "cutlass/epilogue/dispatch_policy.hpp"
#include "cutlass/epilogue/collective/collective_builder.hpp"
#include "cutlass/epilogue/fusion/sm90_callbacks_tma_warpspecialized.hpp"
#include "cutlass/epilogue/fusion/sm90_visitor_tma_warpspecialized.hpp"
#include "cutlass/epilogue/fusion/sm90_visitor_compute_tma_warpspecialized.hpp"
#include "cutlass/epilogue/fusion/sm90_visitor_load_tma_warpspecialized.hpp"
using namespace cute;
namespace fp8gemm {
using ElementA = cutlass::float_e4m3_t;
using LayoutA = cutlass::layout::RowMajor;
constexpr int AlignmentA = 16;
using ElementB = cutlass::float_e4m3_t;
using LayoutB = cutlass::layout::ColumnMajor;
constexpr int AlignmentB = 16;
using ElementC = cutlass::bfloat16_t;
using LayoutC = cutlass::layout::RowMajor;
constexpr int AlignmentC = 8;
using ElementD = cutlass::bfloat16_t;
using LayoutD = cutlass::layout::RowMajor;
constexpr int AlignmentD = 8;
using ElementAccumulator = float;
using ElementCompute = float;
using ArchTag = cutlass::arch::Sm90;
using OperatorClass = cutlass::arch::OpClassTensorOp;
template <class TileShapeMNK, class ClusterShapeMNK, class KernelSched,
class EpiSched = cutlass::epilogue::TmaWarpSpecializedCooperative>
struct GemmTraits {
using TileShape = TileShapeMNK;
using ClusterShape = ClusterShapeMNK;
using KernelSchedule = KernelSched;
using EpilogueSchedule = EpiSched;
static constexpr auto RoundStyle = cutlass::FloatRoundStyle::round_to_nearest;
// D = acc * per-N scale (row broadcast of an (N,) vector, in fp32, cast bf16)
using CustomEVT = cutlass::epilogue::fusion::Sm90EVT<
cutlass::epilogue::fusion::Sm90Compute<cutlass::multiplies, ElementD, ElementCompute, RoundStyle>,
cutlass::epilogue::fusion::Sm90RowBroadcast<
0, TileShapeMNK, float, ElementCompute, Stride<_0, _1, int64_t>>,
cutlass::epilogue::fusion::Sm90AccFetch>;
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
ArchTag, OperatorClass,
TileShape, ClusterShape,
cutlass::epilogue::collective::EpilogueTileAuto,
ElementAccumulator, ElementCompute,
ElementC, LayoutC, AlignmentC,
ElementD, LayoutD, AlignmentD,
EpilogueSchedule,
CustomEVT>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
ArchTag, OperatorClass,
ElementA, LayoutA, AlignmentA,
ElementB, LayoutB, AlignmentB,
ElementAccumulator,
TileShape, ClusterShape,
cutlass::gemm::collective::StageCountAutoCarveout<
static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
KernelSchedule>::CollectiveOp;
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue,
cutlass::gemm::PersistentScheduler>;
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
};
template <class Traits>
inline cutlass::Status gemm_can_implement(const typename Traits::Gemm::Arguments& args) {
return Traits::Gemm::can_implement(args);
}
template <class Traits>
void run_gemm(const torch::Tensor& a, const torch::Tensor& b,
const torch::Tensor& scale, torch::Tensor& d,
int max_swizzle = 1, int raster_dir = 1, int promo_interval = 4) {
using Gemm = typename Traits::Gemm;
int M = a.size(0), K = a.size(1), N = b.size(0);
typename Gemm::Arguments args{
cutlass::gemm::GemmUniversalMode::kGemm,
{M, N, K, 1},
{(const ElementA*)a.data_ptr(), cute::make_stride(int64_t(K), Int<1>{}, int64_t(0)),
(const ElementB*)b.data_ptr(), cute::make_stride(int64_t(K), Int<1>{}, int64_t(0))},
{// epilogue
{ // fusion args (EVT): {row_broadcast args, acc args(empty), compute args(empty)}
{(const float*)scale.data_ptr(), 0.0f, {_0{}, _1{}, int64_t(0)}},
{},
{},
},
(const ElementC*)d.data_ptr(), // C unused; alias D (beta path never taken)
cute::make_stride(int64_t(N), Int<1>{}, int64_t(0)),
(ElementD*)d.data_ptr(),
cute::make_stride(int64_t(N), Int<1>{}, int64_t(0))}};
args.scheduler.max_swizzle_size = max_swizzle;
args.scheduler.raster_order =
raster_dir == 0
? cutlass::gemm::kernel::detail::RasterOrderOptions::AlongM
: cutlass::gemm::kernel::detail::RasterOrderOptions::AlongN;
// Fast-accum FFMA promotion interval (multiple of TileK/GMMA_K = 4)
args.mainloop.mma_promotion_interval = (uint32_t)promo_interval;
// can_implement + workspace size are pointer-independent; cache per shape/schedule.
struct ShapeCache {
cutlass::Status status = cutlass::Status::kInvalid;
size_t ws_size = 0;
int max_swizzle = -1;
int raster_dir = -1;
int M = -1, N = -1, K = -1;
};
static thread_local ShapeCache cache;
if (cache.M != M || cache.N != N || cache.K != K ||
cache.max_swizzle != max_swizzle || cache.raster_dir != raster_dir) {
cache.status = gemm_can_implement<Traits>(args);
cache.ws_size = cache.status == cutlass::Status::kSuccess
? Gemm::get_workspace_size(args)
: 0;
cache.max_swizzle = max_swizzle;
cache.raster_dir = raster_dir;
cache.M = M; cache.N = N; cache.K = K;
}
TORCH_CHECK(cache.status == cutlass::Status::kSuccess,
"CUTLASS can_implement failed: ", (int)cache.status);
torch::Tensor ws;
void* ws_ptr = nullptr;
if (cache.ws_size > 0) {
ws = torch::empty({(int64_t)cache.ws_size},
torch::TensorOptions().dtype(torch::kUInt8).device(a.device()));
ws_ptr = ws.data_ptr();
}
Gemm gemm;
auto status = gemm.initialize(args, ws_ptr);
TORCH_CHECK(status == cutlass::Status::kSuccess, "CUTLASS initialize failed: ",
(int)status);
status = gemm.run(at::cuda::getCurrentCUDAStream());
TORCH_CHECK(status == cutlass::Status::kSuccess, "CUTLASS run failed: ", (int)status);
}
} // namespace fp8gemm
torch::Tensor fp8_gemm_scaled(torch::Tensor a, torch::Tensor b, torch::Tensor scale,
int64_t variant, int64_t max_swizzle, int64_t raster_dir,
int64_t promo_interval) {
TORCH_CHECK(a.is_cuda() && b.is_cuda() && scale.is_cuda());
TORCH_CHECK(a.dtype() == torch::kFloat8_e4m3fn && b.dtype() == torch::kFloat8_e4m3fn);
TORCH_CHECK(scale.dtype() == torch::kFloat32);
TORCH_CHECK(a.is_contiguous() && b.is_contiguous() && scale.is_contiguous());
const auto& a_c = a;
const auto& b_c = b;
const auto& s_c = scale;
int M = a_c.size(0), N = b_c.size(0);
auto d = torch::empty({M, N}, torch::TensorOptions()
.dtype(torch::kBFloat16)
.device(a_c.device()));
using COOP_FA = cutlass::gemm::KernelTmaWarpSpecializedCooperativeFP8FastAccum;
using T1 = fp8gemm::GemmTraits<Shape<_128, _128, _128>, Shape<_2, _1, _1>, COOP_FA>;
using T2 = fp8gemm::GemmTraits<Shape<_128, _256, _128>, Shape<_2, _2, _1>, COOP_FA>;
using T3 = T1; // was: COOP non-fast-accum (slower)
using T4 = T1; // was: pingpong (broken with custom EVT)
using T5 = fp8gemm::GemmTraits<Shape<_128, _128, _128>, Shape<_1, _1, _1>, COOP_FA>;
using T6 = fp8gemm::GemmTraits<Shape<_256, _128, _128>, Shape<_2, _1, _1>, COOP_FA>;
using T7 = fp8gemm::GemmTraits<Shape<_128, _128, _128>, Shape<_4, _1, _1>, COOP_FA>;
using T8 = T1;
using T9 = T1;
using T10 = T1;
using T11 = fp8gemm::GemmTraits<Shape<_256, _128, _128>, Shape<_2, _2, _1>, COOP_FA>;
using T12 = fp8gemm::GemmTraits<Shape<_256, _128, _128>, Shape<_4, _1, _1>, COOP_FA>;
using T13 = fp8gemm::GemmTraits<Shape<_128, _128, _256>, Shape<_2, _1, _1>, COOP_FA>;
using T14 = fp8gemm::GemmTraits<Shape<_128, _256, _128>, Shape<_2, _1, _1>, COOP_FA>;
using T15 = T1;
using T16 = T1;
using T17 = T1;
using T18 = T1;
using T19 = T1;
switch (variant) {
case 0: fp8gemm::run_gemm<T1>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
case 1: fp8gemm::run_gemm<T2>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
case 2: fp8gemm::run_gemm<T3>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
case 3: fp8gemm::run_gemm<T4>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
case 4: fp8gemm::run_gemm<T5>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
case 5: fp8gemm::run_gemm<T6>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
case 6: fp8gemm::run_gemm<T7>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
case 7: fp8gemm::run_gemm<T8>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
case 8: fp8gemm::run_gemm<T9>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
case 9: fp8gemm::run_gemm<T10>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
case 10: fp8gemm::run_gemm<T11>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
case 11: fp8gemm::run_gemm<T12>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
case 12: fp8gemm::run_gemm<T13>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
case 13: fp8gemm::run_gemm<T14>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
case 14: fp8gemm::run_gemm<T15>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
case 15: fp8gemm::run_gemm<T16>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
case 16: fp8gemm::run_gemm<T17>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
case 17: fp8gemm::run_gemm<T18>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
case 18: fp8gemm::run_gemm<T19>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
default: fp8gemm::run_gemm<T1>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval);
}
return d;
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("fp8_gemm_scaled", &fp8_gemm_scaled, "fp8 gemm with per-n scale");
}
20260716_145802_kinetic-claude_kinetic-0715_1m__01_fp8_gemm