"""FP8 e4m3 GEMM using Triton fp8 tensor-core MMA.

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.

Notes:
  * A real fp8 x fp8 MMA (fp8 inputs, fp32 accumulate) is run via tl.dot.
  * K is padded to a multiple of 128 so every K-tile is full, keeping the fp8
    MMA codegen on the fast path (partial-K tail tiles generate slow code on
    non-128-multiple K shapes such as 4127).
  * The per-channel dequant scale is fused into the store to avoid a second
    pass over the output.
"""
import torch
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.autotune(
    configs=[
        triton.Config({"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 128}, num_warps=4, num_stages=3),
        triton.Config({"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 256}, num_warps=4, num_stages=3),
        triton.Config({"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 128}, num_warps=4, num_stages=3),
        triton.Config({"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128}, num_warps=4, num_stages=3),
        triton.Config({"BLOCK_M": 256, "BLOCK_N": 256, "BLOCK_K": 128}, num_warps=8, num_stages=3),
        triton.Config({"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 256}, num_warps=8, num_stages=3),
        triton.Config({"BLOCK_M": 64,  "BLOCK_N": 256, "BLOCK_K": 128}, num_warps=4, num_stages=3),
        triton.Config({"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128}, num_warps=4, num_stages=5),
        triton.Config({"BLOCK_M": 64,  "BLOCK_N": 128, "BLOCK_K": 128}, num_warps=4, num_stages=3),
        triton.Config({"BLOCK_M": 256, "BLOCK_N": 256, "BLOCK_K": 256}, num_warps=8, num_stages=4),
        triton.Config({"BLOCK_M": 128, "BLOCK_N": 512, "BLOCK_K": 128}, num_warps=8, num_stages=3),
        triton.Config({"BLOCK_M": 256, "BLOCK_N": 512, "BLOCK_K": 128}, num_warps=8, num_stages=3),
        triton.Config({"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 128}, num_warps=4, num_stages=5),
        triton.Config({"BLOCK_M": 512, "BLOCK_N": 256, "BLOCK_K": 128}, num_warps=8, num_stages=4),
        triton.Config({"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 256}, num_warps=8, num_stages=4),
    ],
    key=["M", "N", "K"],
)
@triton.jit
def _fp8_gemm_kernel(
    x_ptr, w_ptr, y_ptr, scale_ptr,
    M, N, K,
    stride_xm, stride_xk,
    stride_wn, stride_wk,
    stride_ym, stride_yn,
    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)

    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

    # K is padded to a multiple of BLOCK_K by the caller, so every tile is full
    # and the fp8 MMA codegen stays on the fast path (no partial-K tail tile).
    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    for k in range(0, tl.cdiv(K, BLOCK_K)):
        x = tl.load(x_ptrs, mask=offs_m[:, None] < M, other=0.0).to(tl.float8e4nv)
        w = tl.load(w_ptrs, mask=offs_n[:, None] < N, other=0.0).to(tl.float8e4nv)
        acc = tl.dot(x, tl.trans(w), acc, out_dtype=tl.float32)
        x_ptrs += BLOCK_K * stride_xk
        w_ptrs += BLOCK_K * stride_wk

    # Fused per-channel dequant scale.
    scale = tl.load(scale_ptr + offs_n, mask=offs_n < N).to(tl.float32)
    acc = acc * scale[None, :]
    y = acc.to(tl.bfloat16)
    y_ptrs = y_ptr + offs_m[:, None] * stride_ym + offs_n[None, :] * stride_yn
    tl.store(y_ptrs, y, mask=(offs_m[:, None] < M) & (offs_n[None, :] < N))


class Model(torch.nn.Module):
    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)
        torch.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))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x.contiguous()
        w = self.weight  # (N, K) fp8
        M, N, K = self.M, self.N, self.K
        # Pad K to a multiple of 128 so every K-tile is full (avoids the slow
        # partial-K fp8 MMA codegen on non-128-multiple K shapes like 4127).
        Kpad = (K + 127) // 128 * 128
        if Kpad != K:
            x = torch.nn.functional.pad(x, (0, Kpad - K))
            w = torch.nn.functional.pad(w, (0, Kpad - K))
        y = torch.empty(M, N, dtype=torch.bfloat16, device=x.device)
        grid = lambda meta: (
            triton.cdiv(M, meta["BLOCK_M"]),
            triton.cdiv(N, meta["BLOCK_N"]),
        )
        _fp8_gemm_kernel[grid](
            x, w, y, self.weight_scale,
            M, N, Kpad,
            x.stride(0), x.stride(1),
            w.stride(0), w.stride(1),
            y.stride(0), y.stride(1),
        )
        return y


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]