


# This script provides a template for using load_inline to run a HIP kernel for
import os 
os.environ['PYTORCH_ROCM_ARCH'] = 'gfx942'
import torch
from torch.utils.cpp_extension import load_inline
from task import input_t, output_t
import triton
from copy import deepcopy


CPP_WRAPPER = """
void fp8_mm(torch::Tensor a, torch::Tensor b, torch::Tensor a_scale, torch::Tensor b_scale, torch::Tensor c);
"""

CUDA_SRC = R""" 
// File: fp8_mm_with_smem_vec.cpp

#include <torch/extension.h>
#include <hip/hip_runtime.h>
#include <hip/amd_detail/amd_hip_fp8.h>    // __hip_fp8_e4m3_fnuz, operator float()
#include <hip/amd_detail/amd_hip_bf16.h>   // __hip_bfloat16, __float2bfloat16, __bfloat162float
#include <hip/hip_fp16.h>

// ----------------------------------------
// 1. 辅助宏与模板(与之前保持一致)
// ----------------------------------------
template <int N>
__device__ __host__ inline int cdiv(int x) {
    return (x + N - 1) / N;
}
template <int M, int N>
constexpr __device__ __host__ inline int cdiv() {
    return (M + N - 1) / N;
}

template <int BM, int BN, int mode = 0>
__device__ __host__ inline int2 get_mn_begin(int block_idx, int m, int n) {
    int num_m_block = (m + BM - 1) / BM;
    int num_n_block = (n + BN - 1) / BN;

    if constexpr (mode == 0) {
        int m_idx = block_idx % num_m_block;
        int n_idx = block_idx / num_m_block;
        return make_int2(m_idx * BM, n_idx * BN);
    } else {
        int m_idx = block_idx / num_n_block;
        int n_idx = block_idx % num_n_block;
        return make_int2(m_idx * BM, n_idx * BN);
    }
}

// MFMA wrapper(保持不变)
struct MFMA_FP8_16x16x32 {
    static constexpr int M = 16;
    static constexpr int N = 16;
    static constexpr int K = 32;

    using Areg = __hip_fp8_e4m3_fnuz[8];
    using Breg = __hip_fp8_e4m3_fnuz[8];
    using Creg = float[4];

    __device__ static void load_a(const __hip_fp8_e4m3_fnuz* __restrict__ a_mem,
                                  int stride_m, int stride_k,
                                  Areg& a_reg, int lane_id) {
        constexpr int vec_size = 8;
        int k_i = (lane_id / 16) * vec_size;
        int m_i = (lane_id % 16);
        for (int i = 0; i < vec_size; i++) {
            a_reg[i] = a_mem[m_i * stride_m + (k_i + i) * stride_k];
        }
    }

    __device__ static void load_b(const __hip_fp8_e4m3_fnuz* __restrict__ b_mem,
                                  int stride_n, int stride_k,
                                  Breg& b_reg, int lane_id) {
        constexpr int vec_size = 8;
        int k_i = (lane_id / 16) * vec_size;
        int n_i = (lane_id % 16);
        for (int i = 0; i < vec_size; i++) {
            b_reg[i] = b_mem[n_i * stride_n + (k_i + i) * stride_k];
        }
    }

    __device__ static void mma(const Areg& a_reg,
                               const Breg& b_reg,
                               Creg& c_reg) {
        using floatx4 = __attribute__((__vector_size__(4 * sizeof(float)))) float;
        *(reinterpret_cast<floatx4*>(c_reg)) =
            __builtin_amdgcn_mfma_f32_16x16x32_fp8_fp8(
                *(reinterpret_cast<const long*>(&a_reg)),
                *(reinterpret_cast<const long*>(&b_reg)),
                *(reinterpret_cast<floatx4*>(c_reg)),
                0, 0, 0
            );
    }

    __device__ static void store_c(__hip_bfloat16* __restrict__ c_mem,
                                   int stride_m, int stride_n,
                                   const Creg& c_reg, int lane_id) {
        int n_i = lane_id % 16;
        int m_i = (lane_id / 16) * 4;
        constexpr int vec_size = 4;
        for (int i = 0; i < vec_size; ++i) {
            c_mem[(m_i + i) * stride_m + n_i * stride_n] = __float2bfloat16(c_reg[i]);
        }
    }
};

template <int WM_, int WN_, int WK_, typename INST>
struct CTA {
    static constexpr int WM = WM_;
    static constexpr int WN = WN_;
    static constexpr int WK = WK_;
    static constexpr int INST_M = WM_ * INST::M;
    static constexpr int INST_N = WN_ * INST::N;
    static constexpr int INST_K = WK_ * INST::K;

    __device__ __host__ inline int3 get_offset(int warp_idx,
                                               int inst_m, int inst_n, int inst_k) const {
        int warp_mn    = warp_idx % (WM * WN);
        int warp_m_idx = warp_mn / WN;
        int warp_n_idx = warp_mn % WN;
        int warp_k_idx = warp_idx / (WM * WN);

        int off_m = inst_m + warp_m_idx * INST::M;
        int off_n = inst_n + warp_n_idx * INST::N;
        int off_k = inst_k + warp_k_idx * INST::K;
        return make_int3(off_m, off_n, off_k);
    }
};

constexpr int SCALE_BLOCK_SIZE = 128;
template <int BM, int BK>
__device__ inline void load_a_scale_to_smem(
    const float* __restrict__ a_scale,
    float* __restrict__ smem_a_scale,
    int m_begin,
    int m,
    int base_tile_k
) {
    constexpr int local_num_tile_k = (BK + SCALE_BLOCK_SIZE - 1) / SCALE_BLOCK_SIZE;
    int idx = threadIdx.x;
    int total = BM * local_num_tile_k;
    if (idx < total) {
        int row = idx % BM;
        int p   = idx / BM;
        int global_row = m_begin + row;
        int tile_k = base_tile_k + p;
        smem_a_scale[p * BM + row] =
            a_scale[ global_row + tile_k * m];
    }
}

template <int BN, int BK>
__device__ inline void load_b_scale_to_smem(
    const float* __restrict__ b_scale,
    float* __restrict__ smem_b_scale,
    int tile_n_begin,
    int n_tiles,
    int global_num_tile_n,
    int base_tile_k
) {
    constexpr int local_num_tile_k = (BK + SCALE_BLOCK_SIZE - 1) / SCALE_BLOCK_SIZE;
    int idx = threadIdx.x;
    int total = n_tiles * local_num_tile_k;
    if (idx < total) {
        int ntile = idx / local_num_tile_k;
        int p     = idx % local_num_tile_k;
        int global_ntile = tile_n_begin + ntile;
        int tile_k = base_tile_k + p;
        smem_b_scale[ ntile * local_num_tile_k + p ] =
            b_scale[tile_k * global_num_tile_n + global_ntile];
    }
}
// ------------------------------------------------------------
// 2. 新增模板函数:vectorized load 到 shared memory
//    - BM, BK: tile 尺寸
//    - VEC: 每次循环每个线程最多加载的连续元素个数
//
// 语义:把全局内存中从 (row_begin, col_begin) 开始的 BM×BK 子矩阵，
// “按列主序”(row + col*ld_src) 加载到 shared memory 中 (row + col*ld_smem)。
// 如果越界则补零。
// ------------------------------------------------------------
template <int BM, int BK, int VEC>
__device__ void load_tile_to_smem(
    const __hip_fp8_e4m3_fnuz* __restrict__ src,    // 全局内存起始指针
    __hip_fp8_e4m3_fnuz* __restrict__ smem,          // shared memory 起始指针
    int row_begin, int col_begin,                    // 全局坐标偏移 (row_begin, col_begin)
    int bound_row, int bound_col,                    // 全局矩阵大小 (row_bound, col_bound)
    int ld_src,                                      // 全局内存 row 方向主跨度 (m 或 n)
    int ld_smem                                      // shared memory row 方向跨度 (BM 或 BN)
) {
    int total = BM * BK;                 // tile 内元素总数
    // 每个线程一次尝试加载 VEC 个元素
    int idx = threadIdx.x * VEC;         // 起始 flat index
    while (idx < total) {
        #pragma unroll
        for (int v = 0; v < VEC; ++v) {
            int flat = idx + v;
            if (flat >= total) break;
            int row = flat % BM;         // 在 tile 中的 row 索引
            int col = flat / BM;         // 在 tile 中的 col 索引
            int g_row = row_begin + row;
            int g_col = col_begin + col;
            __hip_fp8_e4m3_fnuz val = (__hip_fp8_e4m3_fnuz)0;
            if (g_row < bound_row && g_col < bound_col) {
                // 在全局内存里 (g_row, g_col) 位置
                val = src[g_row + g_col * ld_src];
            }
            // 存到 shared memory 对应位置 (row + col*ld_smem)
            smem[row + col * ld_smem] = val;
        }
        idx += blockDim.x * VEC;
    }
}

// ------------------------------------------------------------
// 3. 修改后的 custom_kernel:
//    每个 k-block 先把 a/b tile 调用 load_tile_to_smem 一次性 load 到 shared memory。
// ------------------------------------------------------------
template <int BM, int BN, int BK, int WM, int WN, int WK, int VEC>
__global__ void custom_kernel(
    const __hip_fp8_e4m3_fnuz* __restrict__ a,    // (m, k) 列主序
    const __hip_fp8_e4m3_fnuz* __restrict__ b,    // (n, k) 列主序
    const float* __restrict__ a_scale,
    const float* __restrict__ b_scale,
    __hip_bfloat16* __restrict__ c,               // (m, n) 行主序
    int m, int n, int k
) {
    // Layout 注释(与原来保持一致)
    // a: (m, k) : 列主序
    // b: (n, k) : 列主序
    // a_scale: (m, cdiv(k,128))
    // b_scale: (cdiv(n,128), cdiv(k,128))
    // c: (m, n) : 行主序

    using Inst = MFMA_FP8_16x16x32;
    using Cta  = CTA<WM, WN, WK, Inst>;

    constexpr int INST_M = WM * Inst::M;    // 8 * 16 = 128
    constexpr int INST_N = WN * Inst::N;    // 1 * 16 = 16
    constexpr int INST_K = WK * Inst::K;    // 1 * 32 = 32

    constexpr int M_CNT = BM / INST_M;      // 128/128 =1
    constexpr int N_CNT = BN / INST_N;      // 64/16  =4
    constexpr int K_CNT = BK / INST_K;      // 64/32  =2

    int global_num_tile_k = (k + SCALE_BLOCK_SIZE - 1) / SCALE_BLOCK_SIZE;
    int local_num_tile_k  = (BK + SCALE_BLOCK_SIZE - 1) / SCALE_BLOCK_SIZE;  // =1
    int n_tiles_bound     = (BN + SCALE_BLOCK_SIZE - 1) / SCALE_BLOCK_SIZE;  // =1

    // 计算本 block 负责的 (m, n) 小块起始坐标
    int2 mn_begin = get_mn_begin<BM, BN>(blockIdx.x, m, n);
    int m_begin = mn_begin.x;
    int n_begin = mn_begin.y;

    // 线程相关
    int lane_id = threadIdx.x % warpSize;    // warp 内线程 id
    int warp_id = threadIdx.x / warpSize;    // warp id

    // 用来累加 C 的寄存器
    typename Inst::Areg a_reg[M_CNT][K_CNT];
    typename Inst::Breg b_reg[N_CNT][K_CNT];
    typename Inst::Creg c_reg[M_CNT][N_CNT];
    #pragma unroll
    for (int im = 0; im < M_CNT; ++im) {
        #pragma unroll
        for (int in = 0; in < N_CNT; ++in) {
            #pragma unroll
            for (int t = 0; t < 4; ++t) {
                c_reg[im][in][t] = 0.0f;
            }
        }
    }

    // ------------------------------------------------------------
    // 这里声明 shared memory:
    //  0) a_scale: 大小 = BM * local_num_tile_k
    //  1) b_scale: 大小 = n_tiles_bound * local_num_tile_k
    //  2) a_tile:  大小 = BM * BK(FP8 类型，每元素 1 字节)
    //  3) b_tile:  大小 = BN * BK(FP8 类型，每元素 1 字节)
    // ------------------------------------------------------------
    extern __shared__ float shared_mem[]; 
    float* smem_a_scale = shared_mem;                                // size = BM * local_num_tile_k
    float* smem_b_scale = shared_mem + BM * local_num_tile_k;        // size = n_tiles_bound * local_num_tile_k

    // 把 a_tile 放在 float 区之后，需要考虑对齐到 4 字节边界。
    // 我们直接 reinterpret_cast 到 __hip_fp8_e4m3_fnuz*:
    __hip_fp8_e4m3_fnuz* smem_a_tile = reinterpret_cast<__hip_fp8_e4m3_fnuz*>(
            smem_b_scale + (n_tiles_bound * local_num_tile_k)
    );  // 紧接 b_scale 之后

    __hip_fp8_e4m3_fnuz* smem_b_tile = smem_a_tile + (BM * BK);

    // ------------------------------------------------------------
    // 4. 主循环:遍历 k 方向的 tile
    // ------------------------------------------------------------
    for (int k_begin = 0; k_begin < k; k_begin += BK) {
        int base_tile_k = k_begin / SCALE_BLOCK_SIZE;

        // Step A:先把 a_tile 和 b_tile 一次性加载到 shared memory
        // load a: shape [BM, BK]，从全局 (m_begin ... , k_begin ...) 位置开始
        load_tile_to_smem< BM, BK, VEC >(
            a, 
            smem_a_tile,
            m_begin,             // row_begin
            k_begin,             // col_begin
            m,                   // bound_row
            k,                   // bound_col
            m,                   // ld_src = m
            BM                   // ld_smem = BM
        );
        

        // load b: shape [BN, BK]，从全局 (n_begin ... , k_begin ...) 位置开始
        
        load_tile_to_smem< BN, BK, VEC >(
            b,
            smem_b_tile,
            n_begin,
            k_begin,
            n,
            k,
            n,                   // ld_src = n
            BN                   // ld_smem = BN
        );
        

        // Step B:把 a_scale, b_scale 加载进 smem
        int tile_n_begin = n_begin / SCALE_BLOCK_SIZE;
        int tile_n_end   = (n_begin + BN - 1) / SCALE_BLOCK_SIZE;
        int n_tiles      = tile_n_end - tile_n_begin + 1;

        // a_scale -> smem_a_scale
        {
            constexpr int _BM = BM, _BK = BK;
            
            load_a_scale_to_smem< _BM, _BK >(
                a_scale,
                smem_a_scale,
                m_begin,
                m,
                base_tile_k
            );
            
        }
        // b_scale -> smem_b_scale
        {
            constexpr int _BN = BN, _BK = BK;
            
            load_b_scale_to_smem< _BN, _BK >(
                b_scale,
                smem_b_scale,
                tile_n_begin,
                n_tiles,
                (n + SCALE_BLOCK_SIZE - 1) / SCALE_BLOCK_SIZE,
                base_tile_k
            );
            
        }

        // 等待所有线程完成加载
        __syncthreads();

        // Step C:从 shared memory 中读取，执行 MFMA
        #pragma unroll
        for (int ik = 0; ik < K_CNT; ++ik) {
            int sub_k_begin = k_begin + ik * INST_K;  // INST_K=32

            #pragma unroll
            for (int im = 0; im < M_CNT; ++im) {    // M_CNT=1
                #pragma unroll
                for (int in = 0; in < N_CNT; ++in) { // N_CNT=4
                    int3 local_off = Cta().get_offset(
                        warp_id,
                        im * INST_M,
                        in * INST_N,
                        sub_k_begin
                    );
                    int off_m = m_begin + local_off.x;  // 全局 m 坐标
                    int off_n = n_begin + local_off.y;  // 全局 n 坐标
                    int off_k = local_off.z;            // 全局 k 坐标

                    if (off_m >= m || off_n >= n || off_k >= k) {
                        continue;
                    }

                    // 计算在 tile 内部的相对位置
                    int rel_m = off_m - m_begin;  // [0, BM)
                    int rel_n = off_n - n_begin;  // [0, BN)
                    int rel_k = off_k - k_begin;  // [0, BK)

                    // 从 smem_a_tile 读取:按列主序 (row + col*BM)
                    {
                        __hip_fp8_e4m3_fnuz* a_tile_ptr = smem_a_tile + rel_m + rel_k * BM;  
                        
                        Inst::load_a(
                            a_tile_ptr,
                            /*stride_m=*/ 1,
                            /*stride_k=*/ BM,
                            a_reg[im][ik],
                            lane_id
                        );
                        
                    }
                    // 从 smem_b_tile 读取:按列主序 (row + col*BN)
                    {
                        __hip_fp8_e4m3_fnuz* b_tile_ptr = smem_b_tile + rel_n + rel_k * BN; 
                        
                        Inst::load_b(
                            b_tile_ptr,
                            /*stride_n=*/ 1,
                            /*stride_k=*/ BN,
                            b_reg[in][ik],
                            lane_id
                        );
                        
                    }

                    // 原生 MFMA 计算
                    float c_tmp[4] = {0.0f, 0.0f, 0.0f, 0.0f};
                    using floatx4 = __attribute__((__vector_size__(4 * sizeof(float)))) float;
                    *(reinterpret_cast<floatx4*>(c_tmp)) =
                        __builtin_amdgcn_mfma_f32_16x16x32_fp8_fp8(
                            *(reinterpret_cast<const long*>(&a_reg[im][ik])),
                            *(reinterpret_cast<const long*>(&b_reg[in][ik])),
                            *(reinterpret_cast<floatx4*>(c_tmp)),
                            0, 0, 0
                        );

                    // 从 shared memory 中获取 scale
                    int p_local = (off_k - k_begin) / SCALE_BLOCK_SIZE;  // local tile k 索引
                    int row_in_cta = local_off.x + (lane_id / 16) * 4;
                    float a_s[4];
                    a_s[0] = smem_a_scale[ p_local * BM + row_in_cta ];
                    a_s[1] = smem_a_scale[ p_local * BM + row_in_cta + 1 ];
                    a_s[2] = smem_a_scale[ p_local * BM + row_in_cta + 2 ];
                    a_s[3] = smem_a_scale[ p_local * BM + row_in_cta + 3 ];

                    int this_tile_n = off_n / SCALE_BLOCK_SIZE;
                    int ntile_idx   = this_tile_n - tile_n_begin;
                    float b_s = smem_b_scale[ ntile_idx * local_num_tile_k + p_local ];

                    // 累加到 c_reg
                    #pragma unroll
                    for (int t = 0; t < 4; ++t) {
                        c_reg[im][in][t] += c_tmp[t] * a_s[t] * b_s;
                    }
                }
            }
        }

        __syncthreads();
    }

    // Step D:把寄存器结果写回全局内存 c
    #pragma unroll
    for (int im = 0; im < M_CNT; ++im) {
        #pragma unroll
        for (int in = 0; in < N_CNT; ++in) {
            int3 local_off = Cta().get_offset(
                warp_id,
                im * INST_M,
                in * INST_N,
                0
            );
            int off_m = m_begin + local_off.x;
            int off_n = n_begin + local_off.y;
            if (off_m < m && off_n < n) {
                Inst::store_c(
                    c + off_m * n + off_n,
                    n, 1,
                    c_reg[im][in],
                    lane_id
                );
            }
        }
    }
}

// ----------------------------------------
// Host 侧编译与调用
// ----------------------------------------


// Wrapper function
void fp8_mm(torch::Tensor a,
            torch::Tensor b,
            torch::Tensor a_scale,
            torch::Tensor b_scale,
            torch::Tensor c) {
    // Layout
    // a: (m, k) : (1, m)
    // b: (n, k) : (1, m)
    // a_scale: (m, cdiv(k, 128)) : (1, m)
    // b_scale: (cdiv(n, 128), cdiv(k, 128)) : (1, cdiv(n, 128))
    // c: (m, n) : (n, 1)

    int64_t m = a.size(0);
    int64_t k = a.size(1);
    int64_t n = b.size(0);

    int64_t num_tile_k = (k + SCALE_BLOCK_SIZE - 1) / SCALE_BLOCK_SIZE;
    int64_t num_tile_n = (n + SCALE_BLOCK_SIZE - 1) / SCALE_BLOCK_SIZE;

    constexpr int BM_ = 128;
    constexpr int BN_ = 128;
    constexpr int BK_ = 128;
    constexpr int WM_ = 8;
    constexpr int WN_ = 1;
    constexpr int WK_ = 1;
    constexpr int VEC_ = 1; 

    // 计算 grid & block
    int grid_dim_m = ( (int)m + BM_ - 1 ) / BM_;
    int grid_dim_n = ( (int)n + BN_ - 1 ) / BN_;
    dim3 gridDim(grid_dim_m * grid_dim_n, 1, 1);

    // 每个 block 包含 warpSize * WM_ * WN_ * WK_ 个线程
    dim3 blockDim(warpSize * WM_ * WN_ * WK_, 1, 1);

   const int local_num_tile_k = (BK_ + SCALE_BLOCK_SIZE - 1) / SCALE_BLOCK_SIZE;      // = 1
    const int n_tiles_bound    = (BN_ + SCALE_BLOCK_SIZE - 1) / SCALE_BLOCK_SIZE;      // = 1

    // (1)a_scale 部分:BM_ * local_num_tile_k * sizeof(float)
    size_t bytes_a_scale = (size_t)BM_ * (size_t)local_num_tile_k * sizeof(float);
    // (2)b_scale 部分:n_tiles_bound * local_num_tile_k * sizeof(float)
    size_t bytes_b_scale = (size_t)n_tiles_bound * (size_t)local_num_tile_k * sizeof(float);
    // (3)a_tile 部分:BM_ * BK_ * sizeof(__hip_fp8_e4m3_fnuz)  (假设大小为1字节)
    size_t bytes_a_tile  = (size_t)BM_ * (size_t)BK_ * sizeof(__hip_fp8_e4m3_fnuz);
    // (4)b_tile 部分:BN_ * BK_ * sizeof(__hip_fp8_e4m3_fnuz)
    size_t bytes_b_tile  = (size_t)BN_ * (size_t)BK_ * sizeof(__hip_fp8_e4m3_fnuz);

    size_t sharedBytes = bytes_a_scale + bytes_b_scale + bytes_a_tile + bytes_b_tile;

    hipLaunchKernelGGL(
        HIP_KERNEL_NAME(custom_kernel<BM_, BN_, BK_, WM_, WN_, WK_, VEC_>),
        gridDim, blockDim, sharedBytes, 0,
        (__hip_fp8_e4m3_fnuz*)a.data_ptr(), (__hip_fp8_e4m3_fnuz*)b.data_ptr(), 
        a_scale.data_ptr<float>(), b_scale.data_ptr<float>(), (__hip_bfloat16*)c.data_ptr(),
        (int)m, (int)n, (int)k
    );
}
"""

import os
os.environ["CXX"] = "clang++"

module = load_inline(
    name='fp8_mm',
    cpp_sources=[CPP_WRAPPER],
    cuda_sources=[CUDA_SRC],
    functions=['fp8_mm'],
    verbose=True,
    extra_cuda_cflags=["--offload-arch=gfx942", "-std=c++20"],
)


def custom_kernel(data: input_t) -> output_t:
    a, b, a_scale, b_scale, c = data
    # Layout
    # a: (m, k) : (1, m)
    # b: (n, k) : (1, m)
    # a_scale: (m, cdiv(k, 128)) : (1, m)
    # b_scale: (cdiv(n, 128), cdiv(k, 128)) : (1, cdiv(n, 128))
    # c: (m, n) : (n, 1)

    module.fp8_mm(a, b, a_scale, b_scale, c)
    return c


def ref_kernel(data: input_t) -> output_t:
    """
    Highly inefficient torch reference implementation of FP8 GEMM.
    You can use this as a reference / starting template for your implementation.
    """
    # c: [m, n] is pre-allocated memory to help remove allocation overhead.
    a, b, a_scale, b_scale, c = data

    # a is M x K in column-major order, we convert here for simplicity.
    a = a.contiguous()
    a_scale = a_scale.contiguous()
    b_scale = b_scale.contiguous()

    # constants
    m = a.shape[0]
    n = b.shape[0]
    k = a.shape[1]
    block_shape_n = 128
    block_shape_k = 128
    scale_n = b_scale.shape[0]
    scale_k = b_scale.shape[1]

    # Apply blockwise scaling to input 'a'
    a_scale = a_scale.unsqueeze(-1).repeat(1, 1, block_shape_k)  # Shape: [m, scale_k, block_shape_k]
    a_scale = a_scale.reshape(m, scale_k * block_shape_k) 
    a_scale = a_scale[:, :k]

    # Dequantize 'a', in your implementation you should do this at the end.
    a = a.to(a_scale.dtype) * a_scale 

    # Apply blockwise scaling to input 'b'
    b_scale = (
        b_scale.view(-1, 1)
        .repeat(1, block_shape_n * block_shape_k)
        .view(scale_n, scale_k, block_shape_n, block_shape_k)
        .permute(0, 2, 1, 3)  # Reorder dimensions: [scale_n, blk_n, scale_k, blk_k]
        .reshape(scale_n * block_shape_n, scale_k * block_shape_k)
    )
    b_scale = b_scale[:n, :k]

    # Dequantize 'b', in your implementation you should do this at the end.
    b = b.to(b_scale.dtype) * b_scale 

    # Compute FP8 GEMM and write to 'c'. 
    c[...] = (a @ b.T).to(torch.bfloat16)
    return c

# # m,n,k = 256, 256, 128
# m,n,k = 1024, 1024, 1024
# scale_n = triton.cdiv(n, 128)
# scale_k = triton.cdiv(k, 128)
# a = torch.randn((k, m), dtype=torch.bfloat16, device="cuda").to(torch.float8_e4m3fnuz)
# b = torch.randn((k, n), dtype=torch.bfloat16, device="cuda").to(torch.float8_e4m3fnuz)
# a_scale = torch.randn([scale_k, m], dtype=torch.float32, device="cuda")
# b_scale = torch.randn([scale_k, scale_n], dtype=torch.float32, device="cuda")

# c = torch.zeros((m, n), dtype=torch.bfloat16, device="cuda")
# data = (a.T, b.T, a_scale.T, b_scale.T, c)
# copy_data = deepcopy(data)
# out = custom_kernel(data)
# ref_out = ref_kernel(copy_data)
# diff = torch.abs(out - ref_out)
# # print(f"diff shape: {diff.shape}")
# for i in range(m):
#     print(f"i={i}: {diff[i]}")

python submission.py
[92mSuccessfully preprocessed all matching files.[0m

/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/main.cpp -> /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/main.cpp [skipped, no changes]
/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/cuda.cu -> /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip [ok]
Total number of unsupported CUDA function calls: 0


Total number of replaced kernel launches: 0
[1/3] clang++ -MMD -MF main.o.d -DTORCH_EXTENSION_NAME=fp8_mm -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1016\" -isystem /usr/local/lib/python3.10/dist-packages/torch/include -isystem /usr/local/lib/python3.10/dist-packages/torch/include/torch/csrc/api/include -isystem /usr/local/lib/python3.10/dist-packages/torch/include/THH -isystem /opt/rocm-6.3.1/include -isystem /usr/include/python3.10 -fPIC -std=c++17 -c /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/main.cpp -o main.o -D__HIP_PLATFORM_AMD__=1 -DUSE_ROCM=1 -DHIPBLAS_V2 -fPIC
[2/3] /opt/rocm-6.3.1/bin/hipcc  -DWITH_HIP -DTORCH_EXTENSION_NAME=fp8_mm -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1016\" -isystem /usr/local/lib/python3.10/dist-packages/torch/include -isystem /usr/local/lib/python3.10/dist-packages/torch/include/torch/csrc/api/include -isystem /usr/local/lib/python3.10/dist-packages/torch/include/THH -isystem /opt/rocm-6.3.1/include -isystem /usr/include/python3.10 -fPIC -std=c++17 -D__HIP_PLATFORM_AMD__=1 -DUSE_ROCM=1 -DHIPBLAS_V2 -fPIC -DCUDA_HAS_FP16=1 -D__HIP_NO_HALF_OPERATORS__=1 -D__HIP_NO_HALF_CONVERSIONS__=1 -DHIP_ENABLE_WARP_SYNC_BUILTINS=1 --offload-arch=gfx942 -fno-gpu-rdc --offload-arch=gfx942 -std=c++20 -c /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip -o hip.cuda.o 
[3/3] clang++ main.o hip.cuda.o -shared -L/usr/local/lib/python3.10/dist-packages/torch/lib -lc10 -lc10_hip -ltorch_cpu -ltorch_hip -ltorch -ltorch_python -L/opt/rocm-6.3.1/lib -lamdhip64 -o fp8_mm.so

python eval.py test /tmp/tmplpy0049p
[92mSuccessfully preprocessed all matching files.[0m

/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/main.cpp -> /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/main.cpp [skipped, no changes]
/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/cuda.cu -> /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip [skipped, already hipified]
Total number of unsupported CUDA function calls: 0


Total number of replaced kernel launches: 0
ninja: no work to do.
