FP8 GEMM Grok 4.5

"""FP8 e4m3 GEMM via Triton tl.dot (tensor-core fp8 x fp8, fp32 accumulate).

y = (x @ weight.T) * weight_scale  -> bf16
x: (M, K) fp8_e4m3, weight: (N, K) fp8_e4m3, weight_scale: (N,) fp32

Strategy:
  - Real fp8 x fp8 tensor-core MMA via tl.dot (fp32 accumulate), then per-channel scale.
  - Unaligned K is zero-padded to a multiple of BLOCK_K so row-stride stays tile-aligned
    (raw unaligned K is ~80x slower on Hopper).
  - CUDA graph over static workspaces removes Python/launch overhead (critical for skinny M).
  - Weight staging is version-tracked so numeric-stress in-place mutations stay correct.
"""
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,
    EVEN_M: tl.constexpr,
    EVEN_N: tl.constexpr,
):
    pid = tl.program_id(axis=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)

    a_ptrs = a_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak
    # B.T tile (BLOCK_K, BLOCK_N): element (k, n) = B[n, k]
    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, tl.cdiv(K, BLOCK_K)):
        if EVEN_M:
            a = tl.load(a_ptrs)
        else:
            a = tl.load(a_ptrs, mask=offs_m[:, None] < M, other=0.0)
        if EVEN_N:
            b = tl.load(b_ptrs)
        else:
            b = tl.load(b_ptrs, mask=offs_n[None, :] < N, other=0.0)
        acc = tl.dot(a, b, acc)
        a_ptrs += BLOCK_K * stride_ak
        b_ptrs += BLOCK_K * stride_bk

    if EVEN_N:
        scales = tl.load(scale_ptr + offs_n)
    else:
        scales = tl.load(scale_ptr + offs_n, mask=offs_n < N, other=0.0)
    acc = acc * scales[None, :]

    c = acc.to(tl.bfloat16)
    c_ptrs = c_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn
    if EVEN_M and EVEN_N:
        tl.store(c_ptrs, c)
    else:
        tl.store(c_ptrs, c, mask=(offs_m[:, None] < M) & (offs_n[None, :] < N))


def _config_for(M: int, N: int, K: int):
    """BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages, GROUP_M."""
    if M <= 32:
        return 32, 128, 128, 4, 4, 1
    if M <= 64:
        return 64, 128, 128, 4, 4, 4
    return 128, 128, 128, 4, 3, 8


def _launch_gemm(
    x, weight, weight_scale, out,
    M, N, K,
    BLOCK_M, BLOCK_N, BLOCK_K,
    num_warps, num_stages, GROUP_M,
):
    EVEN_M = 1 if (M % BLOCK_M == 0) else 0
    EVEN_N = 1 if (N % BLOCK_N == 0) else 0
    grid = (triton.cdiv(M, BLOCK_M) * triton.cdiv(N, BLOCK_N),)
    _fp8_gemm_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),
        BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K,
        GROUP_M=GROUP_M, EVEN_M=EVEN_M, EVEN_N=EVEN_N,
        num_warps=num_warps, num_stages=num_stages,
    )


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)
        self.register_buffer("weight_scale", s.squeeze(1).to(torch.float32))

        self._BM, self._BN, self._BK, self._nw, self._ns, self._gm = _config_for(M, N, K)
        self._K_work = (K + self._BK - 1) // self._BK * self._BK
        self._need_pad = self._K_work != K

        # Static device workspaces for CUDA-graph capture
        self.register_buffer(
            "_out", torch.empty(M, N, dtype=torch.bfloat16), persistent=False
        )
        self.register_buffer(
            "_x_work",
            torch.zeros(M, self._K_work, dtype=torch.float8_e4m3fn),
            persistent=False,
        )
        self.register_buffer(
            "_w_work",
            torch.zeros(N, self._K_work, dtype=torch.float8_e4m3fn),
            persistent=False,
        )
        self._w_version = -1
        self._graph: torch.cuda.CUDAGraph | None = None
        self._graph_ready = False

    def _sync_w_work(self) -> None:
        """Refresh staged weight when the live buffer is mutated."""
        if self.weight._version == self._w_version:
            return
        if self._need_pad:
            self._w_work[:, : self.K].copy_(self.weight)
        else:
            self._w_work.copy_(self.weight)
        self._w_version = self.weight._version

    def _capture_graph(self) -> None:
        self._sync_w_work()
        for _ in range(3):
            _launch_gemm(
                self._x_work, self._w_work, self.weight_scale, self._out,
                self.M, self.N, self._K_work,
                self._BM, self._BN, self._BK, self._nw, self._ns, self._gm,
            )
        torch.cuda.synchronize()
        g = torch.cuda.CUDAGraph()
        with torch.cuda.graph(g):
            _launch_gemm(
                self._x_work, self._w_work, self.weight_scale, self._out,
                self.M, self.N, self._K_work,
                self._BM, self._BN, self._BK, self._nw, self._ns, self._gm,
            )
        self._graph = g
        self._graph_ready = True

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if not self._graph_ready:
            self._capture_graph()
        self._sync_w_work()
        if self._need_pad:
            self._x_work[:, : self.K].copy_(x)
        else:
            self._x_work.copy_(x)
        assert self._graph is not None
        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]