"""FP8 e4m3 GEMM — Triton kernel for B200 (SM100 Blackwell). Uses tl.dot with fp8_e4m3fn inputs for genuine fp8 x fp8 tensor-core MMA (fp32 accumulate), then applies the per-output-channel weight_scale and stores 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 # --------------------------------------------------------------------------- # Regular FP8 GEMM kernel # --------------------------------------------------------------------------- @triton.jit def _fp8_gemm_kernel( x_ptr, w_ptr, s_ptr, y_ptr, M, N, K, stride_xm, stride_wn, stride_ym, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, GROUP_M: tl.constexpr, ): """fp8 x fp8 GEMM: y = (x @ w.T) * weight_scale, stored as bf16. x: fp8_e4m3 (M, K) row-major w: fp8_e4m3 (N, K) row-major — indexed as w[n, k] s: float32 (N,) per-output-channel dequant scale y: bf16 (M, N) row-major """ # ---- grouped launch order for L2 reuse -------------------------------- 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 = tl.minimum(num_pid_m - first_pid_m, GROUP_M) pid_m = first_pid_m + (pid % group_size_m) pid_n = (pid % num_pid_in_group) // group_size_m # ---- tile offsets ----------------------------------------------------- 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) # ---- pointer arithmetic ----------------------------------------------- # x[m, k] x_ptrs = x_ptr + offs_m[:, None] * stride_xm + offs_k[None, :] # w[n, k] accessed as b[k, n] for tl.dot (K,N) layout w_ptrs = w_ptr + offs_n[None, :] * K + offs_k[:, None] # ---- masks ------------------------------------------------------------ mask_m = offs_m < M mask_n = offs_n < N # ---- accumulator (fp32) ----------------------------------------------- acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) # ---- main K loop ------------------------------------------------------ for _ in range(0, K, BLOCK_K): mask_k = offs_k < K a = tl.load(x_ptrs, mask=mask_m[:, None] & mask_k[None, :], other=0.0) b = tl.load(w_ptrs, mask=mask_k[:, None] & mask_n[None, :], other=0.0) acc = tl.dot(a, b, acc, input_precision="ieee") x_ptrs += BLOCK_K w_ptrs += BLOCK_K offs_k += BLOCK_K # ---- epilogue: per-channel scale -------------------------------------- s_ptrs = s_ptr + offs_n scales = tl.load(s_ptrs, mask=mask_n, other=1.0) acc = acc * scales[None, :] # ---- store bf16 ------------------------------------------------------- y_ptrs = y_ptr + offs_m[:, None] * stride_ym + offs_n[None, :] tl.store(y_ptrs, acc.to(tl.bfloat16), mask=mask_m[:, None] & mask_n[None, :]) # --------------------------------------------------------------------------- # Split-K FP8 GEMM kernel (for skinny M) # --------------------------------------------------------------------------- @triton.jit def _fp8_gemm_splitk_kernel( x_ptr, w_ptr, s_ptr, y_ptr, M, N, K, stride_xm, stride_wn, stride_ym, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, GROUP_M: tl.constexpr, SPLIT_K: tl.constexpr, ): """Split-K fp8 GEMM. Each block computes a partial dot-product over K / SPLIT_K elements, then atomically adds the scaled result to the output. """ pid = tl.program_id(0) pid_sk = tl.program_id(1) # split-K index 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 = tl.minimum(num_pid_m - first_pid_m, GROUP_M) pid_m = first_pid_m + (pid % 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) # ---- K range for this split ------------------------------------------- k_per_split = tl.cdiv(K, SPLIT_K) k_start = pid_sk * k_per_split mask_m = offs_m < M mask_n = offs_n < N acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) for k_off in range(0, k_per_split, BLOCK_K): cur_offs_k = k_off + tl.arange(0, BLOCK_K) global_k = k_start + cur_offs_k mask_k = global_k < K x_ptrs = x_ptr + offs_m[:, None] * stride_xm + global_k[None, :] w_ptrs = w_ptr + offs_n[None, :] * K + global_k[:, None] a = tl.load(x_ptrs, mask=mask_m[:, None] & mask_k[None, :], other=0.0) b = tl.load(w_ptrs, mask=mask_k[:, None] & mask_n[None, :], other=0.0) acc = tl.dot(a, b, acc, input_precision="ieee") # ---- epilogue --------------------------------------------------------- s_ptrs = s_ptr + offs_n scales = tl.load(s_ptrs, mask=mask_n, other=1.0) acc = acc * scales[None, :] # Atomic add into output y_ptrs = y_ptr + offs_m[:, None] * stride_ym + offs_n[None, :] tl.atomic_add(y_ptrs, acc.to(tl.float32), mask=mask_m[:, None] & mask_n[None, :], sem="relaxed", scope="gpu") # --------------------------------------------------------------------------- # Config selection: pick tile sizes based on problem shape # --------------------------------------------------------------------------- def _pick_config(M: int, N: int, K: int): """Return (kernel, kwargs, grid_fn) tuned for the given shape.""" # Split-K for skinny M: increase occupancy by splitting K among more blocks if M <= 64: kw = { "BLOCK_M": 32, "BLOCK_N": 128, "BLOCK_K": 128, "GROUP_M": 1, "SPLIT_K": 8, } grid = lambda: ( triton.cdiv(M, kw["BLOCK_M"]) * triton.cdiv(N, kw["BLOCK_N"]), kw["SPLIT_K"], ) return _fp8_gemm_splitk_kernel, kw, grid, 4, 3 # num_warps, num_stages # Large square or rectangular — maximise tensor-core throughput if M >= 1024 and N >= 1024 and K >= 1024: kw = { "BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 128, "GROUP_M": 4, } grid = lambda: (triton.cdiv(M, kw["BLOCK_M"]) * triton.cdiv(N, kw["BLOCK_N"]),) return _fp8_gemm_kernel, kw, grid, 8, 4 # num_warps, num_stages # Default / medium shapes kw = { "BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "GROUP_M": 4, } grid = lambda: (triton.cdiv(M, kw["BLOCK_M"]) * triton.cdiv(N, kw["BLOCK_N"]),) return _fp8_gemm_kernel, kw, grid, 4, 4 # num_warps, num_stages # --------------------------------------------------------------------------- # Model # --------------------------------------------------------------------------- class Model(nn.Module): """y = ((x @ w.T) * weight_scale).to(bf16).""" 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) 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: M_act, K_act = x.shape assert K_act == self.K, f"K mismatch: input {K_act} vs weight {self.K}" M, N, K = M_act, self.N, self.K y = torch.empty(M, N, dtype=torch.bfloat16, device=x.device) kernel_fn, kw, grid_fn, num_warps, num_stages = _pick_config(M, N, K) if kernel_fn is _fp8_gemm_splitk_kernel: y.zero_() kernel_fn[grid_fn()]( x, self.weight, self.weight_scale, y, M, N, K, x.stride(0), self.weight.stride(0), y.stride(0), num_warps=num_warps, num_stages=num_stages, **kw, ) return y # --------------------------------------------------------------------------- # Module-level shims for compatibility with check.py / benchmark.py # --------------------------------------------------------------------------- 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]