#!POPCORN leaderboard amd-fp8-mm
import os, torch
from torch.utils.cpp_extension import load_inline
from task import input_t, output_t

os.environ.setdefault("CXX", "clang++")

CPP = r"""
void fp8_mm(torch::Tensor a, torch::Tensor b,
            torch::Tensor as, torch::Tensor bs,
            torch::Tensor c);
"""

HIP = r"""
#include <hip/amd_detail/amd_hip_fp8.h>
#include <hip/amd_detail/amd_hip_bf16.h>

constexpr int BLOCK = 128;

/*  每线程计算 1×1 C 元素，保证正确即可  */
__global__ void custom_kernel(const __hip_fp8_e4m3_fnuz* __restrict__ A,
                              const __hip_fp8_e4m3_fnuz* __restrict__ B,
                              const float*  __restrict__ AS,
                              const float*  __restrict__ BS,
                              __hip_bfloat16*             C,
                              int M, int N, int K)
{
    int row = blockIdx.x * blockDim.x + threadIdx.x;   // M 方向
    int col = blockIdx.y * blockDim.y + threadIdx.y;   // N 方向
    if (row >= M || col >= N) return;

    const int scale_n = (N + BLOCK - 1) / BLOCK;       // = N /128
    float acc = 0.f;

    for (int kb = 0; kb < K; kb += BLOCK) {
        /* ---- 当前 128-k 块的缩放 ---- */
        float sA = AS[row + (kb / BLOCK) * M];               // α(row, kblk)
        float sB = BS[(col / BLOCK) + (kb / BLOCK) * scale_n]; // β(colBlk, kblk)
        float sboth = sA * sB;

        #pragma unroll
        for (int kk = 0; kk < BLOCK; ++kk) {
            int k = kb + kk;

            /*  A、B 都是列主序：addr = row + col*ld  */
            float a = static_cast<float>(A[row + k * M]);
            float b = static_cast<float>(B[col + k * N]);    // ← 修正：k*N 而不是 col*K
            acc += a * b * sboth;
        }
    }
    C[row * N + col] = (__hip_bfloat16)acc;            // C 行主序
}

void fp8_mm(torch::Tensor a, torch::Tensor b,
            torch::Tensor as, torch::Tensor bs,
            torch::Tensor c)
{
    int M = a.size(0);
    int N = b.size(0);
    int K = a.size(1);

    dim3 block(16, 16);
    dim3 grid((M + 15) / 16, (N + 15) / 16);

    hipLaunchKernelGGL(custom_kernel, grid, block, 0, 0,
        reinterpret_cast<const __hip_fp8_e4m3_fnuz*>(a.data_ptr()),
        reinterpret_cast<const __hip_fp8_e4m3_fnuz*>(b.data_ptr()),
        as.data_ptr<float>(),
        bs.data_ptr<float>(),
        reinterpret_cast<__hip_bfloat16*>(c.data_ptr()),
        M, N, K);
}
"""

module = load_inline(
    name         = "fp8_mm_correct",
    cpp_sources  = [CPP],
    cuda_sources = [HIP],
    functions    = ["fp8_mm"],
    verbose      = False,
    extra_cuda_cflags=["--offload-arch=gfx942", "-O3", "-std=c++20"],
)

def custom_kernel(data: input_t) -> output_t:
    a, b, a_scale, b_scale, c = data
    module.fp8_mm(a, b, a_scale, b_scale, c)
    return c

python eval.py leaderboard /tmp/tmp_pnhzpv8
[92mSuccessfully preprocessed all matching files.[0m
