FP8 GEMM GPT-5.6 Sol

"""Native FP8 GEMM for NVIDIA Blackwell.

The dot operands stay in e4m3 all the way to ``tl.dot``; accumulation and the
per-output-channel epilogue are FP32, followed by a single bf16 conversion.
"""

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

# Warp-specialized Blackwell kernels use a small compiler-managed global
# scratch region.  Triton intentionally leaves its allocation policy to the
# embedding framework.
_triton_scratch = {}
_num_sms = None


def _triton_allocator(size: int, alignment: int, stream):
    del alignment, stream
    buf = _triton_scratch.get(size)
    if buf is None:
        buf = torch.empty(size, device="cuda", dtype=torch.uint8)
        _triton_scratch[size] = buf
    return buf


triton.set_allocator(_triton_allocator)


@triton.jit
def _pad_k_kernel(
    src_ptr,
    dst_ptr,
    ROWS: tl.constexpr,
    K: tl.constexpr,
    K_PAD: tl.constexpr,
    BLOCK_ROWS: tl.constexpr,
    BLOCK_K: tl.constexpr,
):
    rows = tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS)
    cols = tl.program_id(1) * BLOCK_K + tl.arange(0, BLOCK_K)
    values = tl.load(
        src_ptr + rows[:, None] * K + cols[None, :],
        mask=(rows[:, None] < ROWS) & (cols[None, :] < K),
        other=0.0,
    )
    tl.store(
        dst_ptr + rows[:, None] * K_PAD + cols[None, :],
        values,
        mask=(rows[:, None] < ROWS) & (cols[None, :] < K_PAD),
    )


def _pad_k(src: torch.Tensor, k_pad: int, dst: torch.Tensor | None = None) -> torch.Tensor:
    rows, k = src.shape
    if dst is None:
        dst = torch.empty((rows, k_pad), device=src.device, dtype=src.dtype)
    block_rows, block_k = 8, 256
    grid = (triton.cdiv(rows, block_rows), triton.cdiv(k_pad, block_k))
    _pad_k_kernel[grid](
        src,
        dst,
        ROWS=rows,
        K=k,
        K_PAD=k_pad,
        BLOCK_ROWS=block_rows,
        BLOCK_K=block_k,
        num_warps=8,
    )
    return dst


@triton.jit
def _fp8_gemm_kernel(
    x_ptr,
    w_ptr,
    scale_ptr,
    y_ptr,
    M: tl.constexpr,
    N: tl.constexpr,
    K: tl.constexpr,
    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)

    # Grouped column-major ordering keeps an x tile live in L2 across several
    # output-channel tiles without making the wide-N case badly imbalanced.
    num_pid_group = GROUP_M * num_pid_n
    group = pid // num_pid_group
    first_m = group * GROUP_M
    group_m = tl.minimum(num_pid_m - first_m, GROUP_M)
    pid_m = first_m + (pid % group_m)
    pid_n = (pid % num_pid_group) // group_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] * K + offs_k[None, :]
    w_ptrs = w_ptr + offs_n[None, :] * K + offs_k[:, None]
    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)

    for k0 in tl.range(0, K, BLOCK_K, num_stages=4):
        k_mask = (k0 + offs_k) < K
        x = tl.load(
            x_ptrs,
            mask=(offs_m[:, None] < M) & k_mask[None, :],
            other=0.0,
        )
        w = tl.load(
            w_ptrs,
            mask=k_mask[:, None] & (offs_n[None, :] < N),
            other=0.0,
        )
        acc = tl.dot(x, w, acc)
        x_ptrs += BLOCK_K
        w_ptrs += BLOCK_K

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


@triton.jit
def _tile_coords(tile_id, num_pid_m, num_pid_n, GROUP_M: tl.constexpr):
    num_pid_group = GROUP_M * num_pid_n
    group = tile_id // num_pid_group
    first_m = group * GROUP_M
    group_m = tl.minimum(num_pid_m - first_m, GROUP_M)
    pid_m = first_m + (tile_id % group_m)
    pid_n = (tile_id % num_pid_group) // group_m
    return pid_m, pid_n


@triton.jit
def _fp8_gemm_tma_persistent(
    x_ptr,
    w_ptr,
    scale_ptr,
    y_ptr,
    M: tl.constexpr,
    N: tl.constexpr,
    K: tl.constexpr,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    BLOCK_K: tl.constexpr,
    GROUP_M: tl.constexpr,
    NUM_SMS: tl.constexpr,
    EPILOGUE_SUBTILE: tl.constexpr,
):
    x_desc = tl.make_tensor_descriptor(
        x_ptr,
        shape=[M, K],
        strides=[K, 1],
        block_shape=[BLOCK_M, BLOCK_K],
    )
    w_desc = tl.make_tensor_descriptor(
        w_ptr,
        shape=[N, K],
        strides=[K, 1],
        block_shape=[BLOCK_N, BLOCK_K],
    )
    y_desc = tl.make_tensor_descriptor(
        y_ptr,
        shape=[M, N],
        strides=[N, 1],
        block_shape=[BLOCK_M, BLOCK_N // 2 if EPILOGUE_SUBTILE else BLOCK_N],
    )

    start_pid = tl.program_id(0)
    num_pid_m = tl.cdiv(M, BLOCK_M)
    num_pid_n = tl.cdiv(N, BLOCK_N)
    num_tiles = num_pid_m * num_pid_n

    # Keeping the epilogue's counter independent from the matmul loop counter
    # lets the Blackwell scheduler overlap TMA, tcgen05 MMA, and stores.
    tile_id_out = start_pid - NUM_SMS
    for tile_id in tl.range(
        start_pid,
        num_tiles,
        NUM_SMS,
        flatten=True,
        warp_specialize=True,
    ):
        pid_m, pid_n = _tile_coords(tile_id, num_pid_m, num_pid_n, GROUP_M)
        start_m = pid_m * BLOCK_M
        start_n = pid_n * BLOCK_N

        acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
        for k0 in range(0, K, BLOCK_K):
            x = x_desc.load([start_m, k0])
            w = w_desc.load([start_n, k0])
            acc = tl.dot(x, w.T, acc)

        tile_id_out += NUM_SMS
        out_m, out_n = _tile_coords(tile_id_out, num_pid_m, num_pid_n, GROUP_M)
        start_out_m = out_m * BLOCK_M
        start_out_n = out_n * BLOCK_N

        if EPILOGUE_SUBTILE:
            shaped = tl.reshape(acc, (BLOCK_M, 2, BLOCK_N // 2))
            shaped = tl.permute(shaped, (0, 2, 1))
            acc0, acc1 = tl.split(shaped)
            offs_n = start_out_n + tl.arange(0, BLOCK_N // 2)
            s0 = tl.load(scale_ptr + offs_n)
            s1 = tl.load(scale_ptr + offs_n + BLOCK_N // 2)
            y_desc.store([start_out_m, start_out_n], (acc0 * s0[None, :]).to(tl.bfloat16))
            y_desc.store(
                [start_out_m, start_out_n + BLOCK_N // 2],
                (acc1 * s1[None, :]).to(tl.bfloat16),
            )
        else:
            offs_n = start_out_n + tl.arange(0, BLOCK_N)
            s = tl.load(scale_ptr + offs_n)
            y_desc.store([start_out_m, start_out_n], (acc * s[None, :]).to(tl.bfloat16))


@triton.jit
def _fp8_gemm_tma_skinny(
    x_ptr,
    w_ptr,
    scale_ptr,
    y_ptr,
    M: tl.constexpr,
    N: tl.constexpr,
    K: tl.constexpr,
    BLOCK_N: tl.constexpr,
    BLOCK_K: tl.constexpr,
):
    x_desc = tl.make_tensor_descriptor(
        x_ptr,
        shape=[M, K],
        strides=[K, 1],
        block_shape=[M, BLOCK_K],
    )
    w_desc = tl.make_tensor_descriptor(
        w_ptr,
        shape=[N, K],
        strides=[K, 1],
        block_shape=[BLOCK_N, BLOCK_K],
    )
    y_desc = tl.make_tensor_descriptor(
        y_ptr,
        shape=[M, N],
        strides=[N, 1],
        block_shape=[M, BLOCK_N],
    )

    start_n = tl.program_id(0) * BLOCK_N
    acc = tl.zeros((M, BLOCK_N), dtype=tl.float32)
    for k_tile in tl.range(0, K // BLOCK_K, warp_specialize=True):
        k0 = k_tile * BLOCK_K
        x = x_desc.load([0, k0])
        w = w_desc.load([start_n, k0])
        acc = tl.dot(x, w.T, acc)

    offs_n = start_n + tl.arange(0, BLOCK_N)
    scale = tl.load(scale_ptr + offs_n)
    y_desc.store([0, start_n], (acc * scale[None, :]).to(tl.bfloat16))


def _gemm(
    x: torch.Tensor,
    weight: torch.Tensor,
    scale: torch.Tensor,
    y: torch.Tensor | None = None,
) -> torch.Tensor:
    global _num_sms

    M, K = x.shape
    N = weight.shape[0]
    if y is None:
        y = torch.empty((M, N), device=x.device, dtype=torch.bfloat16)

    if M <= 32 and K % 256 == 0 and N % 64 == 0:
        block_n, block_k = 64, 256
        _fp8_gemm_tma_skinny[(N // block_n,)](
            x,
            weight,
            scale,
            y,
            M=M,
            N=N,
            K=K,
            BLOCK_N=block_n,
            BLOCK_K=block_k,
            num_warps=4,
            num_stages=3,
        )
        return y

    if K % 128 == 0 and M > 32:
        block_m = 128
        block_n = 256
        block_k = 128
        if _num_sms is None:
            _num_sms = torch.cuda.get_device_properties(x.device).multi_processor_count
        num_sms = _num_sms
        grid = (min(num_sms, triton.cdiv(M, block_m) * triton.cdiv(N, block_n)),)
        _fp8_gemm_tma_persistent[grid](
            x,
            weight,
            scale,
            y,
            M=M,
            N=N,
            K=K,
            BLOCK_M=block_m,
            BLOCK_N=block_n,
            BLOCK_K=block_k,
            GROUP_M=8,
            NUM_SMS=num_sms,
            EPILOGUE_SUBTILE=block_n == 256,
            num_warps=8,
            num_stages=4,
        )
        return y

    if M <= 32:
        block_m, block_n, block_k = 32, 32, 128
        num_warps, num_stages = 4, 4
    else:
        block_m, block_n, block_k = 128, 128, 128
        num_warps, num_stages = 8, 4

    grid = (triton.cdiv(M, block_m) * triton.cdiv(N, block_n),)
    _fp8_gemm_kernel[grid](
        x,
        weight,
        scale,
        y,
        M=M,
        N=N,
        K=K,
        BLOCK_M=block_m,
        BLOCK_N=block_n,
        BLOCK_K=block_k,
        GROUP_M=8,
        num_warps=num_warps,
        num_stages=num_stages,
    )
    return y


class Model(nn.Module):
    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.empty((N, K), dtype=torch.float8_e4m3fn))
        self.register_buffer("weight_scale", torch.empty((N,), dtype=torch.float32))
        self._packed_weight = None
        self._packed_weight_version = -1
        self._packed_x = None
        self._output = None

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self._output is None:
            self._output = torch.empty(
                (self.M, self.N), device=x.device, dtype=torch.bfloat16
            )
        if self.K % 128:
            k_pad = triton.cdiv(self.K, 128) * 128
            version = self.weight._version
            if self._packed_weight is None or self._packed_weight_version != version:
                self._packed_weight = _pad_k(self.weight, k_pad)
                self._packed_weight_version = version
            if self._packed_x is None:
                self._packed_x = torch.empty(
                    (self.M, k_pad), device=x.device, dtype=x.dtype
                )
            x = _pad_k(x, k_pad, self._packed_x)
            return _gemm(x, self._packed_weight, self.weight_scale, self._output)
        return _gemm(x, self.weight, self.weight_scale, self._output)


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]