"""FP8 e4m3 GEMM solution using Triton with fp8 tensor cores. Genuine fp8 x fp8: both operands are fp8_e4m3 with fp32 accumulation. The per-channel dequant scale (weight_scale) is applied to the fp32 accumulator before casting to bf16. """ import torch import torch.nn as nn import triton import triton.language as tl OP_TYPE = "gemm" SUPPORTED_PRECISIONS = ["fp8_e4m3"] HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"] E4M3_MAX = 448.0 # ----------------------------------------------------------------------------- # Triton kernels # ----------------------------------------------------------------------------- # Autotune configs - common Blackwell-friendly block sizes for FP8. # We provide a few shapes so autotune doesn't take forever at first run. def _get_autotune_configs(): cfgs = [] for bm in (64, 128): for bn in (128, 256): for bk in (64, 128): for nw in (4, 8): for ns in (1, 2, 3): cfgs.append( triton.Config( { "BLOCK_M": bm, "BLOCK_N": bn, "BLOCK_K": bk, "GROUP_M": 8, }, num_warps=nw, num_stages=ns, ) ) return cfgs @triton.autotune(configs=_get_autotune_configs(), key=["M", "N", "K"]) @triton.jit def gemm_fp8_kernel( x_ptr, # (M, K) fp8_e4m3 w_ptr, # (N, K) fp8_e4m3 scale_ptr, # (N,) fp32 out_ptr, # (M, N) bf16 M, N, K, stride_xm, stride_xk, stride_wn, stride_wk, stride_om, stride_on, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, GROUP_M: tl.constexpr, ): pid = tl.program_id(0) num_pid_m = tl.cdiv(M, BLOCK_M) num_pid_n = tl.cdiv(N, BLOCK_N) # L2-friendly grouping 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 * BLOCK_M + tl.arange(0, BLOCK_M) offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) offs_k = tl.arange(0, BLOCK_K) # x: (M, K), w: (N, K) -> w.T: (K, N) x_ptrs = x_ptr + offs_m[:, None] * stride_xm + offs_k[None, :] * stride_xk w_ptrs = w_ptr + offs_n[None, :] * stride_wn + offs_k[:, None] * stride_wk # w.T indexed -> (K, N) acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) for k in range(0, tl.cdiv(K, BLOCK_K)): # Predicated loads on K tail k_mask = offs_k[None, :] + k * BLOCK_K < K x = tl.load( x_ptrs + k * BLOCK_K * stride_xk, mask=(offs_m[:, None] < M) & k_mask, other=0.0, ) w = tl.load( w_ptrs + k * BLOCK_K * stride_wk, mask=k_mask.T & (offs_n[None, :] < N), other=0.0, ) # tl.dot of fp8 x fp8 -> fp32 accumulator acc = tl.dot(x, w, acc=acc, out_dtype=tl.float32) # Apply per-channel scale scale = tl.load(scale_ptr + offs_n, mask=offs_n < N, other=1.0) acc = acc * scale[None, :] # Store as bf16 out_ptrs = out_ptr + offs_m[:, None] * stride_om + offs_n[None, :] * stride_on tl.store( out_ptrs, acc.to(tl.bfloat16), mask=(offs_m[:, None] < M) & (offs_n[None, :] < N), ) # Skinny-M variant: small BLOCK_M (32) and large BLOCK_N / BLOCK_K def _get_skinny_autotune_configs(): cfgs = [] for bn in (128, 256): for bk in (128, 256): for nw in (4, 8): for ns in (2, 3, 4): cfgs.append( triton.Config( { "BLOCK_M": 32, "BLOCK_N": bn, "BLOCK_K": bk, "GROUP_M": 1, }, num_warps=nw, num_stages=ns, ) ) return cfgs @triton.autotune(configs=_get_skinny_autotune_configs(), key=["M", "N", "K"]) @triton.jit def gemm_fp8_kernel_skinny( x_ptr, w_ptr, scale_ptr, out_ptr, M, N, K, stride_xm, stride_xk, stride_wn, stride_wk, stride_om, stride_on, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, GROUP_M: tl.constexpr, ): pid = tl.program_id(0) num_pid_m = tl.cdiv(M, BLOCK_M) num_pid_n = tl.cdiv(N, BLOCK_N) 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 * BLOCK_M + tl.arange(0, BLOCK_M) offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) offs_k = tl.arange(0, BLOCK_K) x_ptrs = x_ptr + offs_m[:, None] * stride_xm + offs_k[None, :] * stride_xk w_ptrs = w_ptr + offs_n[None, :] * stride_wn + offs_k[:, None] * stride_wk acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) for k in range(0, tl.cdiv(K, BLOCK_K)): k_mask = offs_k[None, :] + k * BLOCK_K < K x = tl.load( x_ptrs + k * BLOCK_K * stride_xk, mask=(offs_m[:, None] < M) & k_mask, other=0.0, ) w = tl.load( w_ptrs + k * BLOCK_K * stride_wk, mask=k_mask.T & (offs_n[None, :] < N), other=0.0, ) acc = tl.dot(x, w, acc=acc, out_dtype=tl.float32) scale = tl.load(scale_ptr + offs_n, mask=offs_n < N, other=1.0) acc = acc * scale[None, :] out_ptrs = out_ptr + offs_m[:, None] * stride_om + offs_n[None, :] * stride_on tl.store( out_ptrs, acc.to(tl.bfloat16), mask=(offs_m[:, None] < M) & (offs_n[None, :] < N), ) def _gemm_fp8(x: torch.Tensor, weight: torch.Tensor, weight_scale: torch.Tensor) -> torch.Tensor: M, K = x.shape N = weight.shape[0] out = torch.empty((M, N), dtype=torch.bfloat16, device=x.device) if M >= 64: grid = lambda meta: (triton.cdiv(M, meta["BLOCK_M"]) * triton.cdiv(N, meta["BLOCK_N"]),) gemm_fp8_kernel[grid]( x, weight, weight_scale, out, M, N, K, x.stride(0), x.stride(1), weight.stride(0), weight.stride(1), out.stride(0), out.stride(1), ) else: grid = lambda meta: (triton.cdiv(M, meta["BLOCK_M"]) * triton.cdiv(N, meta["BLOCK_N"]),) gemm_fp8_kernel_skinny[grid]( x, weight, weight_scale, out, M, N, K, x.stride(0), x.stride(1), weight.stride(0), weight.stride(1), out.stride(0), out.stride(1), ) return out # ----------------------------------------------------------------------------- # Model # ----------------------------------------------------------------------------- class Model(nn.Module): """y = ((x @ w.T) * weight_scale).to(bf16). x: fp8_e4m3 (M, K). w: fp8_e4m3 (N, K) normalized to the e4m3 range. weight_scale: (N,) per-output-channel dequant scale. """ def __init__(self, M: int, N: int, K: int): super().__init__() self.M, self.N, self.K = M, N, K w = torch.empty(N, K, dtype=torch.bfloat16) nn.init.normal_(w, std=0.02) s = (w.float().abs().amax(dim=1, keepdim=True) / E4M3_MAX).clamp(min=1e-12) # (N,1) w_fp8 = (w.float() / s).to(torch.float8_e4m3fn) self.register_buffer("weight", w_fp8) # (N, K) fp8 self.register_buffer("weight_scale", s.squeeze(1).to(torch.float32)) # (N,) def forward(self, x: torch.Tensor) -> torch.Tensor: return _gemm_fp8(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]