#!POPCORN leaderboard amd-fp8-mm
#!POPCORN gpus MI300

import os
import sys
from task import input_t, output_t
from torch.utils.cpp_extension import load_inline

os.environ["CXX"] = "clang++"
os.environ["PYTORCH_ROCM_ARCH"] = "gfx942"

cuda_src = r"""
#include <iostream>
#include <hip/hip_runtime.h>
#include <hip/amd_detail/amd_hip_bf16.h>
#include <hip/amd_detail/amd_hip_fp8.h>

__host__ __device__ __forceinline__ constexpr int cdiv(int a, int b) {
  return (a + b - 1) / b;
}

#define BLOCK_DIM 128

template <const uint32_t BN, const uint32_t BK, const uint32_t BM,
          const uint32_t WITERN, const uint32_t WITERM, const uint32_t SM_COUNT>
__global__ void fp8_mm_kernel_streamk(const __hip_fp8_e4m3_fnuz *A,
                                      const __hip_fp8_e4m3_fnuz *B,
                                      const float *A_scale,
                                      const float *B_scale, float *C,
                                      uint32_t N, uint32_t K, uint32_t M) {
  using uint8x16 =
      __attribute__((__vector_size__(16 * sizeof(uint8_t)))) uint8_t;
  using uint8x4 = __attribute__((__vector_size__(4 * sizeof(uint8_t)))) uint8_t;
  using floatx16 = __attribute__((__vector_size__(16 * sizeof(float)))) float;

  static constexpr uint32_t VECTOR_SIZE = 4;
  static constexpr uint32_t WARPSIZE = 64;
  static constexpr uint32_t WN = 32 * WITERN;
  static constexpr uint32_t WM = 32 * WITERM;
  static constexpr uint32_t numThreads = (BN * BM) / (16 * WITERN * WITERM);
  static constexpr uint32_t strideA =
      (numThreads / (BN / VECTOR_SIZE)) * VECTOR_SIZE;
  static constexpr uint32_t strideB =
      (numThreads / (BM / VECTOR_SIZE)) * VECTOR_SIZE;
  static constexpr uint32_t nstridesA = BK >= strideA ? cdiv(BK, strideA) : 1;

  static_assert(numThreads % BN == 0, "BN should be a multiple of numThreads");
  static_assert(numThreads % BM == 0, "BM should be a multiple of numThreads");
  static_assert(BK <= 128 && BM <= 128, "Range above 128 is not supported");

  uint32_t numTiles = cdiv(N, BN) * cdiv(M, BM);
  uint32_t subTilesPerTile = cdiv(K, BK);
  uint32_t totalSubTiles = numTiles * subTilesPerTile;
  uint32_t subTilesToProcess = cdiv(totalSubTiles, SM_COUNT);
  uint32_t startSubTileIdx = subTilesToProcess * blockIdx.x;
  uint32_t startTileIdx = startSubTileIdx / subTilesPerTile;
  uint32_t subTile = startSubTileIdx;

  for (uint32_t tileIdx = startTileIdx, subTilesProcessed = 0;
       tileIdx < numTiles && subTilesProcessed < subTilesToProcess; ++tileIdx) {
    uint32_t rowOffsetC = (tileIdx % cdiv(N, BN)) * BN;
    uint32_t colOffsetC = (tileIdx / cdiv(N, BN)) * BM;
    uint32_t colOffsetA = rowOffsetC;
    uint32_t colOffsetB = colOffsetC;
    uint32_t M_scale = cdiv(M, BLOCK_DIM);

    uint32_t innerColA = (threadIdx.x % (BN / VECTOR_SIZE)) * VECTOR_SIZE;
    uint32_t innerRowA = (threadIdx.x / (BN / VECTOR_SIZE)) * VECTOR_SIZE;
    uint32_t innerColB = (threadIdx.x % (BM / VECTOR_SIZE)) * VECTOR_SIZE;
    uint32_t innerRowB = (threadIdx.x / (BM / VECTOR_SIZE)) * VECTOR_SIZE;

    uint32_t laneIdx = threadIdx.x % WARPSIZE;
    uint32_t warpIdx = threadIdx.x / WARPSIZE;
    uint32_t warpColOffset = (warpIdx % (BM / WM)) * WM;
    uint32_t warpRowOffset = (warpIdx / (BM / WM)) * WN;
    uint32_t warpX = laneIdx % 32;
    uint32_t warpY = laneIdx / 32;

    __shared__ __hip_fp8_e4m3_fnuz As[2][BN][BK + 4], Bs[BM][BK + 4];
    __shared__ float Ws[2][BN + 4];
    int curr = 0, next = 1;

    subTile = subTile % subTilesPerTile;
    uint16_t tileOffset = subTile * BK;

    for (uint32_t innerRowOffsetA = 0; innerRowOffsetA < BK;
         innerRowOffsetA += strideA) {
      if ((tileOffset + innerRowOffsetA + innerRowA) < K &&
          (colOffsetA + innerColA) < N && (innerRowOffsetA + innerRowA) < BK) {
        uint8x4 x = *reinterpret_cast<const uint8x4 *>(
            &A[(tileOffset + innerRowOffsetA + innerRowA) * N +
               (colOffsetA + innerColA)]);
        uint8x4 y = *reinterpret_cast<const uint8x4 *>(
            &A[(tileOffset + innerRowOffsetA + innerRowA + 1) * N +
               (colOffsetA + innerColA)]);
        uint8x4 z = *reinterpret_cast<const uint8x4 *>(
            &A[(tileOffset + innerRowOffsetA + innerRowA + 2) * N +
               (colOffsetA + innerColA)]);
        uint8x4 w = *reinterpret_cast<const uint8x4 *>(
            &A[(tileOffset + innerRowOffsetA + innerRowA + 3) * N +
               (colOffsetA + innerColA)]);

        uint8x4 xt = {x[0], y[0], z[0], w[0]};
        uint8x4 yt = {x[1], y[1], z[1], w[1]};
        uint8x4 zt = {x[2], y[2], z[2], w[2]};
        uint8x4 wt = {x[3], y[3], z[3], w[3]};

        *reinterpret_cast<uint8x4 *>(&As[curr][innerColA + (innerColA / 32) % 4]
                                        [innerRowOffsetA + innerRowA]) = xt;
        *reinterpret_cast<uint8x4 *>(
            &As[curr][innerColA + (innerColA / 32 + 1) % 4]
               [innerRowOffsetA + innerRowA]) = yt;
        *reinterpret_cast<uint8x4 *>(
            &As[curr][innerColA + (innerColA / 32 + 2) % 4]
               [innerRowOffsetA + innerRowA]) = zt;
        *reinterpret_cast<uint8x4 *>(
            &As[curr][innerColA + (innerColA / 32 + 3) % 4]
               [innerRowOffsetA + innerRowA]) = wt;
      } else if ((innerRowOffsetA + innerRowA) < BK) {
        *reinterpret_cast<uint8x4 *>(&As[curr][innerColA + (innerColA / 32) % 4]
                                        [innerRowOffsetA + innerRowA]) = {0};
        *reinterpret_cast<uint8x4 *>(
            &As[curr][innerColA + (innerColA / 32 + 1) % 4]
               [innerRowOffsetA + innerRowA]) = {0};
        *reinterpret_cast<uint8x4 *>(
            &As[curr][innerColA + (innerColA / 32 + 2) % 4]
               [innerRowOffsetA + innerRowA]) = {0};
        *reinterpret_cast<uint8x4 *>(
            &As[curr][innerColA + (innerColA / 32 + 3) % 4]
               [innerRowOffsetA + innerRowA]) = {0};
      }
    }
    if (threadIdx.x < (BN / 4)) {
      *reinterpret_cast<float4 *>(&Ws[curr][threadIdx.x * 4]) =
          *reinterpret_cast<const float4 *>(
              &A_scale[(tileOffset / BLOCK_DIM) * N +
                       (colOffsetA + threadIdx.x * 4)]);
    }
    for (uint32_t innerRowOffsetB = 0; innerRowOffsetB < BK;
         innerRowOffsetB += strideB) {
      if ((tileOffset + innerRowOffsetB + innerRowB) < K &&
          (colOffsetB + innerColB) < M && (innerRowOffsetB + innerRowB) < BK) {
        uint8x4 x = *reinterpret_cast<const uint8x4 *>(
            &B[(tileOffset + innerRowOffsetB + innerRowB) * M +
               (colOffsetB + innerColB)]);
        uint8x4 y = *reinterpret_cast<const uint8x4 *>(
            &B[(tileOffset + innerRowOffsetB + innerRowB + 1) * M +
               (colOffsetB + innerColB)]);
        uint8x4 z = *reinterpret_cast<const uint8x4 *>(
            &B[(tileOffset + innerRowOffsetB + innerRowB + 2) * M +
               (colOffsetB + innerColB)]);
        uint8x4 w = *reinterpret_cast<const uint8x4 *>(
            &B[(tileOffset + innerRowOffsetB + innerRowB + 3) * M +
               (colOffsetB + innerColB)]);

        uint8x4 xt = {x[0], y[0], z[0], w[0]};
        uint8x4 yt = {x[1], y[1], z[1], w[1]};
        uint8x4 zt = {x[2], y[2], z[2], w[2]};
        uint8x4 wt = {x[3], y[3], z[3], w[3]};

        *reinterpret_cast<uint8x4 *>(&Bs[innerColB + (innerColB / 32) % 4]
                                        [innerRowOffsetB + innerRowB]) = xt;
        *reinterpret_cast<uint8x4 *>(&Bs[innerColB + (innerColB / 32 + 1) % 4]
                                        [innerRowOffsetB + innerRowB]) = yt;
        *reinterpret_cast<uint8x4 *>(&Bs[innerColB + (innerColB / 32 + 2) % 4]
                                        [innerRowOffsetB + innerRowB]) = zt;
        *reinterpret_cast<uint8x4 *>(&Bs[innerColB + (innerColB / 32 + 3) % 4]
                                        [innerRowOffsetB + innerRowB]) = wt;
      } else if ((innerRowOffsetB + innerRowB) < BK) {
        *reinterpret_cast<uint8x4 *>(&Bs[innerColB + (innerColB / 32) % 4]
                                        [innerRowOffsetB + innerRowB]) = {0};
        *reinterpret_cast<uint8x4 *>(&Bs[innerColB + (innerColB / 32 + 1) % 4]
                                        [innerRowOffsetB + innerRowB]) = {0};
        *reinterpret_cast<uint8x4 *>(&Bs[innerColB + (innerColB / 32 + 2) % 4]
                                        [innerRowOffsetB + innerRowB]) = {0};
        *reinterpret_cast<uint8x4 *>(&Bs[innerColB + (innerColB / 32 + 3) % 4]
                                        [innerRowOffsetB + innerRowB]) = {0};
      }
    }
    if (threadIdx.x == numThreads - 1) { // load using different warp
      Ws[curr][BN] = B_scale[(tileOffset / BLOCK_DIM) * M_scale +
                             (colOffsetB / BLOCK_DIM)];
    }

    __syncthreads();

    uint32_t A_x_tmp[nstridesA], A_y_tmp[nstridesA], A_z_tmp[nstridesA],
        A_w_tmp[nstridesA];
    float4 A_scale_temp;
    float B_scale_temp;

    floatx16 d[WITERN][WITERM] = {0};

    for (tileOffset += BK;
         subTile < subTilesPerTile && subTilesProcessed < subTilesToProcess;
         ++subTile, ++subTilesProcessed, tileOffset += BK) {
      if (tileOffset < K && subTilesProcessed + 1 < subTilesToProcess) {
        for (uint32_t innerRowOffsetA = 0; innerRowOffsetA < BK;
             innerRowOffsetA += strideA) {
          if ((tileOffset + innerRowOffsetA + innerRowA) < K &&
              (colOffsetA + innerColA) < N &&
              (innerRowOffsetA + innerRowA) < BK) {
            A_x_tmp[innerRowOffsetA / strideA] =
                *reinterpret_cast<const uint32_t *>(
                    &A[(tileOffset + innerRowOffsetA + innerRowA) * N +
                       (colOffsetA + innerColA)]);
            A_y_tmp[innerRowOffsetA / strideA] =
                *reinterpret_cast<const uint32_t *>(
                    &A[(tileOffset + innerRowOffsetA + innerRowA + 1) * N +
                       (colOffsetA + innerColA)]);
            A_z_tmp[innerRowOffsetA / strideA] =
                *reinterpret_cast<const uint32_t *>(
                    &A[(tileOffset + innerRowOffsetA + innerRowA + 2) * N +
                       (colOffsetA + innerColA)]);
            A_w_tmp[innerRowOffsetA / strideA] =
                *reinterpret_cast<const uint32_t *>(
                    &A[(tileOffset + innerRowOffsetA + innerRowA + 3) * N +
                       (colOffsetA + innerColA)]);
          } else if ((innerRowOffsetA + innerRowA) < BK) {
            A_x_tmp[innerRowOffsetA / strideA] = {0};
            A_y_tmp[innerRowOffsetA / strideA] = {0};
            A_z_tmp[innerRowOffsetA / strideA] = {0};
            A_w_tmp[innerRowOffsetA / strideA] = {0};
          }
        }
        if (threadIdx.x < (BN / 4)) {
          A_scale_temp = *reinterpret_cast<const float4 *>(
              &A_scale[(tileOffset / BLOCK_DIM) * N +
                       (colOffsetA + threadIdx.x * 4)]);
        }
        if (threadIdx.x == numThreads - 1) { // load using different warp
          B_scale_temp = B_scale[(tileOffset / BLOCK_DIM) * M_scale +
                                 (colOffsetB / BLOCK_DIM)];
        }
      }

      long a[WITERN], b[WITERM];
      float b_scale = Ws[curr][BN];
      floatx16 c[WITERN][WITERM] = {0};

      for (uint32_t BKOffset = 0; BKOffset < BK; BKOffset += 16) {
        for (uint32_t wn = 0; wn < WITERN; ++wn) {
          int row = warpRowOffset + wn * 32 + warpX;
          row = (row & ~3) + (((row & ~3) / 32 + row % 4) % 4);
          a[wn] =
              *reinterpret_cast<long *>(&As[curr][row][BKOffset + warpY * 8]);
        }
        for (uint32_t wm = 0; wm < WITERM; ++wm) {
          int row = warpColOffset + wm * 32 + warpX;
          row = (row & ~3) + (((row & ~3) / 32 + row % 4) % 4);
          b[wm] = *reinterpret_cast<long *>(&Bs[row][BKOffset + warpY * 8]);
        }
        for (uint32_t wn = 0; wn < WITERN; ++wn) {
          for (uint32_t wm = 0; wm < WITERM; ++wm) {
            c[wn][wm] = __builtin_amdgcn_mfma_f32_32x32x16_fp8_fp8(
                a[wn], b[wm], c[wn][wm], 0, 0, 0);
          }
        }
      }
      for (uint32_t wn = 0; wn < WITERN; ++wn) {
        for (uint32_t wm = 0; wm < WITERM; ++wm) {
          for (uint32_t j = 0; j < 4; ++j) {
            for (uint32_t i = 0; i < 4; ++i) {
              d[wn][wm][i + j * 4] +=
                  c[wn][wm][i + j * 4] *
                  Ws[curr][warpRowOffset + wn * 32 + j * 8 + warpY * 4 + i] *
                  b_scale;
            }
          }
        }
      }

      if (tileOffset < K && subTilesProcessed + 1 < subTilesToProcess) {
        for (uint32_t innerRowOffsetA = 0; innerRowOffsetA < BK;
             innerRowOffsetA += strideA) {
          if ((innerRowOffsetA + innerRowA) < BK) {
            uint8x4 xt = {uint8_t(A_x_tmp[innerRowOffsetA / strideA]),
                          uint8_t(A_y_tmp[innerRowOffsetA / strideA]),
                          uint8_t(A_z_tmp[innerRowOffsetA / strideA]),
                          uint8_t(A_w_tmp[innerRowOffsetA / strideA])};
            uint8x4 yt = {uint8_t(A_x_tmp[innerRowOffsetA / strideA] >> 8),
                          uint8_t(A_y_tmp[innerRowOffsetA / strideA] >> 8),
                          uint8_t(A_z_tmp[innerRowOffsetA / strideA] >> 8),
                          uint8_t(A_w_tmp[innerRowOffsetA / strideA] >> 8)};
            uint8x4 zt = {uint8_t(A_x_tmp[innerRowOffsetA / strideA] >> 16),
                          uint8_t(A_y_tmp[innerRowOffsetA / strideA] >> 16),
                          uint8_t(A_z_tmp[innerRowOffsetA / strideA] >> 16),
                          uint8_t(A_w_tmp[innerRowOffsetA / strideA] >> 16)};
            uint8x4 wt = {uint8_t(A_x_tmp[innerRowOffsetA / strideA] >> 24),
                          uint8_t(A_y_tmp[innerRowOffsetA / strideA] >> 24),
                          uint8_t(A_z_tmp[innerRowOffsetA / strideA] >> 24),
                          uint8_t(A_w_tmp[innerRowOffsetA / strideA] >> 24)};

            *reinterpret_cast<uint8x4 *>(
                &As[next][innerColA + (innerColA / 32) % 4]
                   [innerRowOffsetA + innerRowA]) = xt;
            *reinterpret_cast<uint8x4 *>(
                &As[next][innerColA + (innerColA / 32 + 1) % 4]
                   [innerRowOffsetA + innerRowA]) = yt;
            *reinterpret_cast<uint8x4 *>(
                &As[next][innerColA + (innerColA / 32 + 2) % 4]
                   [innerRowOffsetA + innerRowA]) = zt;
            *reinterpret_cast<uint8x4 *>(
                &As[next][innerColA + (innerColA / 32 + 3) % 4]
                   [innerRowOffsetA + innerRowA]) = wt;
          }
        }
        if (threadIdx.x < (BN / 4)) {
          *reinterpret_cast<float4 *>(&Ws[next][threadIdx.x * 4]) =
              A_scale_temp;
        }
        if (threadIdx.x == numThreads - 1) { // load using different warp
          Ws[next][BN] = B_scale_temp;
        }

        __syncthreads();

        for (uint32_t innerRowOffsetB = 0; innerRowOffsetB < BK;
             innerRowOffsetB += strideB) {
          if ((tileOffset + innerRowOffsetB + innerRowB) < K &&
              (colOffsetB + innerColB) < M &&
              (innerRowOffsetB + innerRowB) < BK) {
            uint8x4 x = *reinterpret_cast<const uint8x4 *>(
                &B[(tileOffset + innerRowOffsetB + innerRowB) * M +
                   (colOffsetB + innerColB)]);
            uint8x4 y = *reinterpret_cast<const uint8x4 *>(
                &B[(tileOffset + innerRowOffsetB + innerRowB + 1) * M +
                   (colOffsetB + innerColB)]);
            uint8x4 z = *reinterpret_cast<const uint8x4 *>(
                &B[(tileOffset + innerRowOffsetB + innerRowB + 2) * M +
                   (colOffsetB + innerColB)]);
            uint8x4 w = *reinterpret_cast<const uint8x4 *>(
                &B[(tileOffset + innerRowOffsetB + innerRowB + 3) * M +
                   (colOffsetB + innerColB)]);

            uint8x4 xt = {x[0], y[0], z[0], w[0]};
            uint8x4 yt = {x[1], y[1], z[1], w[1]};
            uint8x4 zt = {x[2], y[2], z[2], w[2]};
            uint8x4 wt = {x[3], y[3], z[3], w[3]};

            *reinterpret_cast<uint8x4 *>(&Bs[innerColB + (innerColB / 32) % 4]
                                            [innerRowOffsetB + innerRowB]) = xt;
            *reinterpret_cast<uint8x4 *>(
                &Bs[innerColB + (innerColB / 32 + 1) % 4]
                   [innerRowOffsetB + innerRowB]) = yt;
            *reinterpret_cast<uint8x4 *>(
                &Bs[innerColB + (innerColB / 32 + 2) % 4]
                   [innerRowOffsetB + innerRowB]) = zt;
            *reinterpret_cast<uint8x4 *>(
                &Bs[innerColB + (innerColB / 32 + 3) % 4]
                   [innerRowOffsetB + innerRowB]) = wt;
          } else if ((innerRowOffsetB + innerRowB) < BK) {
            *reinterpret_cast<uint8x4 *>(&Bs[innerColB + (innerColB / 32) % 4]
                                            [innerRowOffsetB + innerRowB]) = {
                0};
            *reinterpret_cast<uint8x4 *>(
                &Bs[innerColB + (innerColB / 32 + 1) % 4]
                   [innerRowOffsetB + innerRowB]) = {0};
            *reinterpret_cast<uint8x4 *>(
                &Bs[innerColB + (innerColB / 32 + 2) % 4]
                   [innerRowOffsetB + innerRowB]) = {0};
            *reinterpret_cast<uint8x4 *>(
                &Bs[innerColB + (innerColB / 32 + 3) % 4]
                   [innerRowOffsetB + innerRowB]) = {0};
          }
        }

        __syncthreads();
      }

      curr = 1 - curr;
      next = 1 - next;
    }

    for (uint32_t wn = 0; wn < WITERN; ++wn) {
      for (uint32_t wm = 0; wm < WITERM; ++wm) {
        for (uint32_t j = 0; j < 4; ++j) {
          for (uint32_t i = 0; i < 4; ++i) {
            if ((rowOffsetC + warpRowOffset + wn * 32 + j * 8 + warpY * 4 + i) <
                    N &&
                (colOffsetC + warpColOffset + wm * 32 + warpX) < M) {
              atomicAdd(&C[(rowOffsetC + warpRowOffset + wn * 32 + j * 8 +
                            warpY * 4 + i) *
                               M +
                           (colOffsetC + warpColOffset + wm * 32 + warpX)],
                        d[wn][wm][i + j * 4]);
            }
          }
        }
      }
    }
  }
}

template <const uint32_t BN, const uint32_t BK, const uint32_t BM,
          const uint32_t WITERN, const uint32_t WITERM, const uint32_t SM_COUNT>
__global__ void
fp8_mm_kernel_N(const __hip_fp8_e4m3_fnuz *A, const __hip_fp8_e4m3_fnuz *B,
                const float *A_scale, const float *B_scale, __hip_bfloat16 *C,
                uint32_t N, uint32_t K, uint32_t M) {
  using uint8x16 =
      __attribute__((__vector_size__(16 * sizeof(uint8_t)))) uint8_t;
  using uint8x4 = __attribute__((__vector_size__(4 * sizeof(uint8_t)))) uint8_t;
  using floatx16 = __attribute__((__vector_size__(16 * sizeof(float)))) float;

  static constexpr uint32_t VECTOR_SIZE = 4;
  static constexpr uint32_t WARPSIZE = 64;
  static constexpr uint32_t WN = 32 * WITERN;
  static constexpr uint32_t WM = 32 * WITERM;
  static constexpr uint32_t numThreads = (BN * BM) / (16 * WITERN * WITERM);
  static constexpr uint32_t strideA =
      (numThreads / (BN / VECTOR_SIZE)) * VECTOR_SIZE;
  static constexpr uint32_t strideB =
      (numThreads / (BM / VECTOR_SIZE)) * VECTOR_SIZE;
  static constexpr uint32_t nstridesA = BK >= strideA ? cdiv(BK, strideA) : 1;

  static_assert(numThreads % BN == 0, "BN should be a multiple of numThreads");
  static_assert(numThreads % BM == 0, "BM should be a multiple of numThreads");
  static_assert(BK <= 128 && BM <= 128, "Range above 128 is not supported");

  uint32_t numTiles = cdiv(N, BN) * cdiv(M, BM);

  for (uint32_t tileIdx = blockIdx.x; tileIdx < numTiles; tileIdx += SM_COUNT) {
    uint32_t rowOffsetC = (tileIdx % cdiv(N, BN)) * BN;
    uint32_t colOffsetC = (tileIdx / cdiv(N, BN)) * BM;
    uint32_t colOffsetA = rowOffsetC;
    uint32_t colOffsetB = colOffsetC;
    uint32_t M_scale = cdiv(M, BLOCK_DIM);

    uint32_t innerColA = (threadIdx.x % (BN / VECTOR_SIZE)) * VECTOR_SIZE;
    uint32_t innerRowA = (threadIdx.x / (BN / VECTOR_SIZE)) * VECTOR_SIZE;
    uint32_t innerColB = (threadIdx.x % (BM / VECTOR_SIZE)) * VECTOR_SIZE;
    uint32_t innerRowB = (threadIdx.x / (BM / VECTOR_SIZE)) * VECTOR_SIZE;

    uint32_t laneIdx = threadIdx.x % WARPSIZE;
    uint32_t warpIdx = threadIdx.x / WARPSIZE;
    uint32_t warpColOffset = (warpIdx % (BM / WM)) * WM;
    uint32_t warpRowOffset = (warpIdx / (BM / WM)) * WN;
    uint32_t warpX = laneIdx % 32;
    uint32_t warpY = laneIdx / 32;

    __shared__ __hip_fp8_e4m3_fnuz As[2][BN][BK + 4], Bs[BM][BK + 4];
    __shared__ float Ws[2][BN + 4];
    int curr = 0, next = 1;

    for (uint32_t innerRowOffsetA = 0; innerRowOffsetA < BK;
         innerRowOffsetA += strideA) {
      if ((innerRowOffsetA + innerRowA) < K && (colOffsetA + innerColA) < N &&
          (innerRowOffsetA + innerRowA) < BK) {
        uint8x4 x = *reinterpret_cast<const uint8x4 *>(
            &A[(innerRowOffsetA + innerRowA) * N + (colOffsetA + innerColA)]);
        uint8x4 y = *reinterpret_cast<const uint8x4 *>(
            &A[(innerRowOffsetA + innerRowA + 1) * N +
               (colOffsetA + innerColA)]);
        uint8x4 z = *reinterpret_cast<const uint8x4 *>(
            &A[(innerRowOffsetA + innerRowA + 2) * N +
               (colOffsetA + innerColA)]);
        uint8x4 w = *reinterpret_cast<const uint8x4 *>(
            &A[(innerRowOffsetA + innerRowA + 3) * N +
               (colOffsetA + innerColA)]);

        uint8x4 xt = {x[0], y[0], z[0], w[0]};
        uint8x4 yt = {x[1], y[1], z[1], w[1]};
        uint8x4 zt = {x[2], y[2], z[2], w[2]};
        uint8x4 wt = {x[3], y[3], z[3], w[3]};

        *reinterpret_cast<uint8x4 *>(&As[curr][innerColA + (innerColA / 32) % 4]
                                        [innerRowOffsetA + innerRowA]) = xt;
        *reinterpret_cast<uint8x4 *>(
            &As[curr][innerColA + (innerColA / 32 + 1) % 4]
               [innerRowOffsetA + innerRowA]) = yt;
        *reinterpret_cast<uint8x4 *>(
            &As[curr][innerColA + (innerColA / 32 + 2) % 4]
               [innerRowOffsetA + innerRowA]) = zt;
        *reinterpret_cast<uint8x4 *>(
            &As[curr][innerColA + (innerColA / 32 + 3) % 4]
               [innerRowOffsetA + innerRowA]) = wt;
      } else if ((innerRowOffsetA + innerRowA) < BK) {
        *reinterpret_cast<uint8x4 *>(&As[curr][innerColA + (innerColA / 32) % 4]
                                        [innerRowOffsetA + innerRowA]) = {0};
        *reinterpret_cast<uint8x4 *>(
            &As[curr][innerColA + (innerColA / 32 + 1) % 4]
               [innerRowOffsetA + innerRowA]) = {0};
        *reinterpret_cast<uint8x4 *>(
            &As[curr][innerColA + (innerColA / 32 + 2) % 4]
               [innerRowOffsetA + innerRowA]) = {0};
        *reinterpret_cast<uint8x4 *>(
            &As[curr][innerColA + (innerColA / 32 + 3) % 4]
               [innerRowOffsetA + innerRowA]) = {0};
      }
    }
    if (threadIdx.x < (BN / 4)) {
      *reinterpret_cast<float4 *>(&Ws[curr][threadIdx.x * 4]) =
          *reinterpret_cast<const float4 *>(
              &A_scale[(colOffsetA + threadIdx.x * 4)]);
    }
    for (uint32_t innerRowOffsetB = 0; innerRowOffsetB < BK;
         innerRowOffsetB += strideB) {
      if ((innerRowOffsetB + innerRowB) < K && (colOffsetB + innerColB) < M &&
          (innerRowOffsetB + innerRowB) < BK) {
        uint8x4 x = *reinterpret_cast<const uint8x4 *>(
            &B[(innerRowOffsetB + innerRowB) * M + (colOffsetB + innerColB)]);
        uint8x4 y = *reinterpret_cast<const uint8x4 *>(
            &B[(innerRowOffsetB + innerRowB + 1) * M +
               (colOffsetB + innerColB)]);
        uint8x4 z = *reinterpret_cast<const uint8x4 *>(
            &B[(innerRowOffsetB + innerRowB + 2) * M +
               (colOffsetB + innerColB)]);
        uint8x4 w = *reinterpret_cast<const uint8x4 *>(
            &B[(innerRowOffsetB + innerRowB + 3) * M +
               (colOffsetB + innerColB)]);

        uint8x4 xt = {x[0], y[0], z[0], w[0]};
        uint8x4 yt = {x[1], y[1], z[1], w[1]};
        uint8x4 zt = {x[2], y[2], z[2], w[2]};
        uint8x4 wt = {x[3], y[3], z[3], w[3]};

        *reinterpret_cast<uint8x4 *>(&Bs[innerColB + (innerColB / 32) % 4]
                                        [innerRowOffsetB + innerRowB]) = xt;
        *reinterpret_cast<uint8x4 *>(&Bs[innerColB + (innerColB / 32 + 1) % 4]
                                        [innerRowOffsetB + innerRowB]) = yt;
        *reinterpret_cast<uint8x4 *>(&Bs[innerColB + (innerColB / 32 + 2) % 4]
                                        [innerRowOffsetB + innerRowB]) = zt;
        *reinterpret_cast<uint8x4 *>(&Bs[innerColB + (innerColB / 32 + 3) % 4]
                                        [innerRowOffsetB + innerRowB]) = wt;
      } else if ((innerRowOffsetB + innerRowB) < BK) {
        *reinterpret_cast<uint8x4 *>(&Bs[innerColB + (innerColB / 32) % 4]
                                        [innerRowOffsetB + innerRowB]) = {0};
        *reinterpret_cast<uint8x4 *>(&Bs[innerColB + (innerColB / 32 + 1) % 4]
                                        [innerRowOffsetB + innerRowB]) = {0};
        *reinterpret_cast<uint8x4 *>(&Bs[innerColB + (innerColB / 32 + 2) % 4]
                                        [innerRowOffsetB + innerRowB]) = {0};
        *reinterpret_cast<uint8x4 *>(&Bs[innerColB + (innerColB / 32 + 3) % 4]
                                        [innerRowOffsetB + innerRowB]) = {0};
      }
    }
    if (threadIdx.x == numThreads - 1) { // load using different warp
      Ws[curr][BN] = B_scale[(colOffsetB / BLOCK_DIM)];
    }

    __syncthreads();

    uint32_t A_tmp[4][nstridesA];
    float4 A_scale_temp;
    float B_scale_temp;

    floatx16 d[WITERN][WITERM] = {0};

    for (uint32_t tileOffset = BK; tileOffset < K + BK; tileOffset += BK) {
      if (tileOffset < K) {
        for (uint32_t innerRowOffsetA = 0; innerRowOffsetA < BK;
             innerRowOffsetA += strideA) {
          if ((tileOffset + innerRowOffsetA + innerRowA) < K &&
              (colOffsetA + innerColA) < N &&
              (innerRowOffsetA + innerRowA) < BK) {
            A_tmp[0][innerRowOffsetA / strideA] =
                *reinterpret_cast<const uint32_t *>(
                    &A[(tileOffset + innerRowOffsetA + innerRowA) * N +
                       (colOffsetA + innerColA)]);
            A_tmp[1][innerRowOffsetA / strideA] =
                *reinterpret_cast<const uint32_t *>(
                    &A[(tileOffset + innerRowOffsetA + innerRowA + 1) * N +
                       (colOffsetA + innerColA)]);
            A_tmp[2][innerRowOffsetA / strideA] =
                *reinterpret_cast<const uint32_t *>(
                    &A[(tileOffset + innerRowOffsetA + innerRowA + 2) * N +
                       (colOffsetA + innerColA)]);
            A_tmp[3][innerRowOffsetA / strideA] =
                *reinterpret_cast<const uint32_t *>(
                    &A[(tileOffset + innerRowOffsetA + innerRowA + 3) * N +
                       (colOffsetA + innerColA)]);
          } else if ((innerRowOffsetA + innerRowA) < BK) {
            A_tmp[0][innerRowOffsetA / strideA] = {0};
            A_tmp[1][innerRowOffsetA / strideA] = {0};
            A_tmp[2][innerRowOffsetA / strideA] = {0};
            A_tmp[3][innerRowOffsetA / strideA] = {0};
          }
        }
        if (threadIdx.x < (BN / 4)) {
          A_scale_temp = *reinterpret_cast<const float4 *>(
              &A_scale[(tileOffset / BLOCK_DIM) * N +
                       (colOffsetA + threadIdx.x * 4)]);
        }
        if (threadIdx.x == numThreads - 1) { // load using different warp
          B_scale_temp = B_scale[(tileOffset / BLOCK_DIM) * M_scale +
                                 (colOffsetB / BLOCK_DIM)];
        }
      }

      long a[WITERN], b[WITERM];
      float b_scale = Ws[curr][BN];
      floatx16 c[WITERN][WITERM] = {0};

      for (uint32_t BKOffset = 0; BKOffset < BK; BKOffset += 16) {
        for (uint32_t wn = 0; wn < WITERN; ++wn) {
          int row = warpRowOffset + wn * 32 + warpX;
          row = (row & ~3) + (((row & ~3) / 32 + row % 4) % 4);
          a[wn] =
              *reinterpret_cast<long *>(&As[curr][row][BKOffset + warpY * 8]);
        }
        for (uint32_t wm = 0; wm < WITERM; ++wm) {
          int row = warpColOffset + wm * 32 + warpX;
          row = (row & ~3) + (((row & ~3) / 32 + row % 4) % 4);
          b[wm] = *reinterpret_cast<long *>(
              &Bs[row][BKOffset + warpY * 8]);
        }
        for (uint32_t wn = 0; wn < WITERN; ++wn) {
          for (uint32_t wm = 0; wm < WITERM; ++wm) {
            c[wn][wm] = __builtin_amdgcn_mfma_f32_32x32x16_fp8_fp8(
                a[wn], b[wm], c[wn][wm], 0, 0, 0);
          }
        }
      }
      for (uint32_t wn = 0; wn < WITERN; ++wn) {
        for (uint32_t wm = 0; wm < WITERM; ++wm) {
          for (uint32_t j = 0; j < 4; ++j) {
            for (uint32_t i = 0; i < 4; ++i) {
              d[wn][wm][i + j * 4] +=
                  c[wn][wm][i + j * 4] *
                  Ws[curr][warpRowOffset + wn * 32 + j * 8 + warpY * 4 + i] *
                  b_scale;
            }
          }
        }
      }

      if (tileOffset < K) {
        for (uint32_t innerRowOffsetA = 0; innerRowOffsetA < BK;
             innerRowOffsetA += strideA) {
          if ((innerRowOffsetA + innerRowA) < BK) {
            uint8x4 xt = {uint8_t(A_tmp[0][innerRowOffsetA / strideA]),
                          uint8_t(A_tmp[1][innerRowOffsetA / strideA]),
                          uint8_t(A_tmp[2][innerRowOffsetA / strideA]),
                          uint8_t(A_tmp[3][innerRowOffsetA / strideA])};
            uint8x4 yt = {uint8_t(A_tmp[0][innerRowOffsetA / strideA] >> 8),
                          uint8_t(A_tmp[1][innerRowOffsetA / strideA] >> 8),
                          uint8_t(A_tmp[2][innerRowOffsetA / strideA] >> 8),
                          uint8_t(A_tmp[3][innerRowOffsetA / strideA] >> 8)};
            uint8x4 zt = {uint8_t(A_tmp[0][innerRowOffsetA / strideA] >> 16),
                          uint8_t(A_tmp[1][innerRowOffsetA / strideA] >> 16),
                          uint8_t(A_tmp[2][innerRowOffsetA / strideA] >> 16),
                          uint8_t(A_tmp[3][innerRowOffsetA / strideA] >> 16)};
            uint8x4 wt = {uint8_t(A_tmp[0][innerRowOffsetA / strideA] >> 24),
                          uint8_t(A_tmp[1][innerRowOffsetA / strideA] >> 24),
                          uint8_t(A_tmp[2][innerRowOffsetA / strideA] >> 24),
                          uint8_t(A_tmp[3][innerRowOffsetA / strideA] >> 24)};

            *reinterpret_cast<uint8x4 *>(
                &As[next][innerColA + (innerColA / 32) % 4]
                   [innerRowOffsetA + innerRowA]) = xt;
            *reinterpret_cast<uint8x4 *>(
                &As[next][innerColA + (innerColA / 32 + 1) % 4]
                   [innerRowOffsetA + innerRowA]) = yt;
            *reinterpret_cast<uint8x4 *>(
                &As[next][innerColA + (innerColA / 32 + 2) % 4]
                   [innerRowOffsetA + innerRowA]) = zt;
            *reinterpret_cast<uint8x4 *>(
                &As[next][innerColA + (innerColA / 32 + 3) % 4]
                   [innerRowOffsetA + innerRowA]) = wt;
          }
        }
        if (threadIdx.x < (BN / 4)) {
          *reinterpret_cast<float4 *>(&Ws[next][threadIdx.x * 4]) =
              A_scale_temp;
        }
        if (threadIdx.x == numThreads - 1) { // load using different warp
          Ws[next][BN] = B_scale_temp;
        }

        __syncthreads();

        for (uint32_t innerRowOffsetB = 0; innerRowOffsetB < BK;
             innerRowOffsetB += strideB) {
          if ((tileOffset + innerRowOffsetB + innerRowB) < K &&
              (colOffsetB + innerColB) < M &&
              (innerRowOffsetB + innerRowB) < BK) {
            uint8x4 x = *reinterpret_cast<const uint8x4 *>(
                &B[(tileOffset + innerRowOffsetB + innerRowB) * M +
                   (colOffsetB + innerColB)]);
            uint8x4 y = *reinterpret_cast<const uint8x4 *>(
                &B[(tileOffset + innerRowOffsetB + innerRowB + 1) * M +
                   (colOffsetB + innerColB)]);
            uint8x4 z = *reinterpret_cast<const uint8x4 *>(
                &B[(tileOffset + innerRowOffsetB + innerRowB + 2) * M +
                   (colOffsetB + innerColB)]);
            uint8x4 w = *reinterpret_cast<const uint8x4 *>(
                &B[(tileOffset + innerRowOffsetB + innerRowB + 3) * M +
                   (colOffsetB + innerColB)]);

            uint8x4 xt = {x[0], y[0], z[0], w[0]};
            uint8x4 yt = {x[1], y[1], z[1], w[1]};
            uint8x4 zt = {x[2], y[2], z[2], w[2]};
            uint8x4 wt = {x[3], y[3], z[3], w[3]};

            *reinterpret_cast<uint8x4 *>(&Bs[innerColB + (innerColB / 32) % 4]
                                            [innerRowOffsetB + innerRowB]) = xt;
            *reinterpret_cast<uint8x4 *>(
                &Bs[innerColB + (innerColB / 32 + 1) % 4]
                   [innerRowOffsetB + innerRowB]) = yt;
            *reinterpret_cast<uint8x4 *>(
                &Bs[innerColB + (innerColB / 32 + 2) % 4]
                   [innerRowOffsetB + innerRowB]) = zt;
            *reinterpret_cast<uint8x4 *>(
                &Bs[innerColB + (innerColB / 32 + 3) % 4]
                   [innerRowOffsetB + innerRowB]) = wt;
          } else if ((innerRowOffsetB + innerRowB) < BK) {
            *reinterpret_cast<uint8x4 *>(&Bs[innerColB + (innerColB / 32) % 4]
                                            [innerRowOffsetB + innerRowB]) = {
                0};
            *reinterpret_cast<uint8x4 *>(
                &Bs[innerColB + (innerColB / 32 + 1) % 4]
                   [innerRowOffsetB + innerRowB]) = {0};
            *reinterpret_cast<uint8x4 *>(
                &Bs[innerColB + (innerColB / 32 + 2) % 4]
                   [innerRowOffsetB + innerRowB]) = {0};
            *reinterpret_cast<uint8x4 *>(
                &Bs[innerColB + (innerColB / 32 + 3) % 4]
                   [innerRowOffsetB + innerRowB]) = {0};
          }
        }

        __syncthreads();
      }

      curr = 1 - curr;
      next = 1 - next;
    }

    for (uint32_t wn = 0; wn < WITERN; ++wn) {
      for (uint32_t wm = 0; wm < WITERM; ++wm) {
        for (uint32_t j = 0; j < 4; ++j) {
          for (uint32_t i = 0; i < 4; ++i) {
            if ((rowOffsetC + warpRowOffset + wn * 32 + j * 8 + warpY * 4 + i) <
                    N &&
                (colOffsetC + warpColOffset + wm * 32 + warpX) < M)
              C[(rowOffsetC + warpRowOffset + wn * 32 + j * 8 + warpY * 4 + i) *
                    M +
                (colOffsetC + warpColOffset + wm * 32 + warpX)] =
                  (__hip_bfloat16)d[wn][wm][i + j * 4];
          }
        }
      }
    }
  }
}

at::Tensor fp8_mm(at::Tensor A, at::Tensor B, at::Tensor A_scale,
                  at::Tensor B_scale, at::Tensor C) {
  int N = A.size(0), K = A.size(1), M = B.size(0);

  static float *C_dataptr = NULL;
  if (!C_dataptr) {
    hipMalloc(&C_dataptr, 6144 * 7168 * sizeof(float));
  }

  const uint32_t BK = 128;
  const uint32_t BN = 128;
  const uint32_t BM = 128;
  const uint32_t WITERN = 1;
  const uint32_t WITERM = 1;
  const uint32_t SM_COUNT = 304;
  dim3 numThreads((BN * BM) / (16 * WITERN * WITERM));
  dim3 numBlocks(SM_COUNT);

  if (K >= 7168 && N <= 1024 && M <= 1536) {
    hipMemset(C_dataptr, 0, N * M * sizeof(float));
    at::Tensor Cf32 = at::from_blob(C_dataptr, {N, M},
                  at::TensorOptions().device(at::kCUDA).dtype(at::kFloat));

    fp8_mm_kernel_streamk<BN, BK, BM, WITERN, WITERM, SM_COUNT>
        <<<numBlocks, numThreads>>>(
            (__hip_fp8_e4m3_fnuz *)A.data_ptr(),
            (__hip_fp8_e4m3_fnuz *)B.data_ptr(), A_scale.data_ptr<float>(),
            B_scale.data_ptr<float>(), (float *)Cf32.data_ptr(), N, K, M);
    return Cf32;
  } else {
    fp8_mm_kernel_N<BN, BK, BM, WITERN, WITERM, SM_COUNT>
        <<<numBlocks, numThreads>>>(
            (__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(), N, K, M);
    return C;
  }
}
"""

cpp_src = r"""
at::Tensor fp8_mm(at::Tensor A, at::Tensor B, at::Tensor A_scale,
                  at::Tensor B_scale, at::Tensor C);
"""

if sys.stdout is None:
    sys.stdout = open("/dev/stdout", "w")
if sys.stderr is None:
    sys.stderr = open("/dev/stderr", "w")

module = load_inline(
    name="fp8_mm",
    cpp_sources=[cpp_src],
    cuda_sources=[cuda_src],
    functions=["fp8_mm"],
    verbose=True,
    extra_cuda_cflags=[
        "-Ofast",
        "--offload-arch=gfx942",
        "-std=c++20",
        "-ffp-contract=fast",
        "-lhip_hcc",
         "-mcumode",
    ],
)



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

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

/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/cuda.cu -> /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip [ok]
/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]
Total number of unsupported CUDA function calls: 0


Total number of replaced kernel launches: 2
[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 -Ofast --offload-arch=gfx942 -std=c++20 -ffp-contract=fast -lhip_hcc -mcumode -c /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip -o hip.cuda.o 
clang++: warning: -lhip_hcc: 'linker' input unused [-Wunused-command-line-argument]
/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip:746:5: warning: ignoring return value of function declared with 'nodiscard' attribute [-Wunused-result]
  746 |     hipMalloc(&C_dataptr, 6144 * 7168 * sizeof(float));
      |     ^~~~~~~~~ ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip:759:5: warning: ignoring return value of function declared with 'nodiscard' attribute [-Wunused-result]
  759 |     hipMemset(C_dataptr, 0, N * M * sizeof(float));
      |     ^~~~~~~~~ ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
2 warnings generated when compiling for gfx942.
/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip:746:5: warning: ignoring return value of function declared with 'nodiscard' attribute [-Wunused-result]
  746 |     hipMalloc(&C_dataptr, 6144 * 7168 * sizeof(float));
      |     ^~~~~~~~~ ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip:759:5: warning: ignoring return value of function declared with 'nodiscard' attribute [-Wunused-result]
  759 |     hipMemset(C_dataptr, 0, N * M * sizeof(float));
      |     ^~~~~~~~~ ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
2 warnings generated when compiling for host.
[3/3] clang++ hip.cuda.o main.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/tmpldv1yfnv
[92mSuccessfully preprocessed all matching files.[0m

/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]
/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]
Total number of unsupported CUDA function calls: 0


Total number of replaced kernel launches: 2
[1/1] 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
