FP8 GEMM Grok 4.6

"""FP8 e4m3 x e4m3 GEMM with per-output-channel dequant scale.

y = ((x @ weight.T) * weight_scale).to(bf16)

Real fp8 tensor-core MMA via Triton tl.dot on fp8 operands (lowers to
mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 on SM120).
"""
from __future__ import annotations

import torch
import torch.nn as nn
import triton
import triton.language as tl

E4M3_MAX = 448.0


@triton.jit
def _fp8_gemm_kernel(
    a_ptr,
    b_ptr,
    scale_ptr,
    c_ptr,
    M,
    N,
    K,
    stride_am,
    stride_ak,
    stride_bn,
    stride_bk,
    stride_cm,
    stride_cn,
    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 % 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)
    offs_m = tl.max_contiguous(tl.multiple_of(tl.where(offs_m < M, offs_m, 0), BLOCK_M), BLOCK_M)
    offs_n = tl.max_contiguous(tl.multiple_of(tl.where(offs_n < N, offs_n, 0), BLOCK_N), BLOCK_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((BLOCK_M, BLOCK_N), dtype=tl.float32)
    for _k in range(0, K, BLOCK_K):
        a = tl.load(a_ptrs)
        b = tl.load(b_ptrs)
        acc = tl.dot(a, b, acc)
        a_ptrs += BLOCK_K * stride_ak
        b_ptrs += BLOCK_K * stride_bk

    scale = tl.load(scale_ptr + offs_n, mask=offs_n < N, other=0.0)
    acc = acc * scale[None, :]
    mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
    tl.store(
        c_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn,
        acc.to(tl.bfloat16),
        mask=mask,
    )


def _pick(m: int, n: int, k: int) -> tuple[int, int, int, int, int, int]:
    if m <= 32:
        return 32, 64, 128, 8, 4, 4
    if m <= 64:
        return 64, 128, 64, 8, 4, 4
    return 128, 256, 64, 8, 8, 3


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)
        w_fp8 = (w.float() / s).to(torch.float8_e4m3fn)
        self.register_buffer("weight", w_fp8)
        self.register_buffer("weight_scale", s.squeeze(1).to(torch.float32))

        self._bm, self._bn, self._bk, self._group, self._warps, self._stages = _pick(M, N, K)
        self._grid = (triton.cdiv(M, self._bm) * triton.cdiv(N, self._bn),)
        self._need_pad = (K % self._bk) != 0
        self._k_pad = (K + self._bk - 1) // self._bk * self._bk if self._need_pad else K
        self._out: torch.Tensor | None = None
        self._a_pad: torch.Tensor | None = None
        self._b_pad: torch.Tensor | None = None
        self._graph: torch.cuda.CUDAGraph | None = None
        self._cap_a_ptr: int = 0
        self._cap_b_ptr: int = 0

    def _ensure_bufs(self, x: torch.Tensor) -> None:
        device = x.device
        if self._out is None or self._out.device != device:
            self._out = torch.empty((self.M, self.N), device=device, dtype=torch.bfloat16)
            self._graph = None
            self._cap_a_ptr = 0
            self._cap_b_ptr = 0
            if self._need_pad:
                self._a_pad = torch.zeros((self.M, self._k_pad), device=device, dtype=x.dtype)
                self._b_pad = torch.zeros((self.N, self._k_pad), device=device, dtype=self.weight.dtype)

    def _launch(self, a: torch.Tensor, b: torch.Tensor, k: int) -> None:
        _fp8_gemm_kernel[self._grid](
            a,
            b,
            self.weight_scale,
            self._out,
            self.M,
            self.N,
            k,
            a.stride(0),
            a.stride(1),
            b.stride(0),
            b.stride(1),
            self._out.stride(0),
            self._out.stride(1),
            BLOCK_M=self._bm,
            BLOCK_N=self._bn,
            BLOCK_K=self._bk,
            GROUP_M=self._group,
            num_warps=self._warps,
            num_stages=self._stages,
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if not x.is_contiguous():
            x = x.contiguous()
        self._ensure_bufs(x)

        if self._need_pad:
            a, b, k = self._a_pad, self._b_pad, self._k_pad
            a[:, : self.K].copy_(x)
            b[:, : self.K].copy_(self.weight)
        else:
            a, b, k = x, self.weight, self.K

        a_ptr = a.data_ptr()
        b_ptr = b.data_ptr()
        if self._graph is None or a_ptr != self._cap_a_ptr or b_ptr != self._cap_b_ptr:
            self._launch(a, b, k)
            torch.cuda.synchronize()
            g = torch.cuda.CUDAGraph()
            with torch.cuda.graph(g):
                self._launch(a, b, k)
            self._graph = g
            self._cap_a_ptr = a_ptr
            self._cap_b_ptr = b_ptr
        self._graph.replay()
        return self._out


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]