gemm radeon kernel

build.toml

@@ -0,0 +1,39 @@
+[general]
+name = "gemm"
+universal = false
+
+[torch]
+src = [
+  "torch-ext/torch_binding.cpp",
+  "torch-ext/torch_binding.h",
+]
+
+[kernel.gemm]
+backend = "rocm"
+rocm-archs = [
+  #"gfx908",
+  #  "gfx90a",
+  # "gfx940",
+  # "gfx941",
+    "gfx942",
+  # "gfx1100",
+  # "gfx1101",
+  # "gfx1102",
+  # "gfx1200",
+  # "gfx1201",
+]
+depends = ["torch"]
+src = [
+  "include/clangd_workaround.h",
+  "include/gpu_libs.h",
+  "include/gpu_types.h",
+  "include/timer.h",
+  "gemm/gemm_kernel.h",
+  "gemm/gemm_kernel_legacy.h",
+  "gemm/gemm_launcher.hip",
+  "gemm/transpose_kernel.h",
+  "src/utils/arithmetic.h",
+  "src/utils/timer.hip",
+  "tests/checker/metrics.h",
+]
+include = ["include"]

build/torch27-cxx11-rocm63-x86_64-linux/gemm/__init__.py

@@ -0,0 +1,19 @@
+from typing import Optional
+import torch
+from ._ops import ops
+
+def gemm(a: torch.Tensor, b: torch.Tensor, as_: torch.Tensor, bs: torch.Tensor, 
+         out: Optional[torch.Tensor] = None) -> torch.Tensor:
+         
+    if out is None:
+        # Create output tensor with appropriate shape and dtype
+        M, K = a.shape
+        K_b, N = b.shape
+        assert K == K_b, f"Matrix dimension mismatch: A has {K} cols, B has {K_b} rows"
+        
+        # Output should be BF16 type on the same device as inputs
+        out = torch.empty((M, N), dtype=torch.bfloat16, device=a.device)
+    
+    ops.gemm(out, a, b, as_, bs)
+    return out
+

build/torch27-cxx11-rocm63-x86_64-linux/gemm/_gemm_60de0b2.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:abf5894843b26edc6c1894a0e5e9829441567319f446600f1e772fe2621e4faf
+size 2196896

build/torch27-cxx11-rocm63-x86_64-linux/gemm/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _gemm_60de0b2
+ops = torch.ops._gemm_60de0b2
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_gemm_60de0b2::{op_name}"

build/torch28-cxx11-rocm63-x86_64-linux/gemm/__init__.py

@@ -0,0 +1,19 @@
+from typing import Optional
+import torch
+from ._ops import ops
+
+def gemm(a: torch.Tensor, b: torch.Tensor, as_: torch.Tensor, bs: torch.Tensor, 
+         out: Optional[torch.Tensor] = None) -> torch.Tensor:
+         
+    if out is None:
+        # Create output tensor with appropriate shape and dtype
+        M, K = a.shape
+        K_b, N = b.shape
+        assert K == K_b, f"Matrix dimension mismatch: A has {K} cols, B has {K_b} rows"
+        
+        # Output should be BF16 type on the same device as inputs
+        out = torch.empty((M, N), dtype=torch.bfloat16, device=a.device)
+    
+    ops.gemm(out, a, b, as_, bs)
+    return out
+

build/torch28-cxx11-rocm63-x86_64-linux/gemm/_gemm_60de0b2.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:da6a47f577368eda36d189c4840e5a1da512a4a5ed55ac159fb516698f417d49
+size 2196896

build/torch28-cxx11-rocm63-x86_64-linux/gemm/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _gemm_60de0b2
+ops = torch.ops._gemm_60de0b2
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_gemm_60de0b2::{op_name}"

build/torch28-cxx11-rocm64-x86_64-linux/gemm/__init__.py

@@ -0,0 +1,19 @@
+from typing import Optional
+import torch
+from ._ops import ops
+
+def gemm(a: torch.Tensor, b: torch.Tensor, as_: torch.Tensor, bs: torch.Tensor, 
+         out: Optional[torch.Tensor] = None) -> torch.Tensor:
+         
+    if out is None:
+        # Create output tensor with appropriate shape and dtype
+        M, K = a.shape
+        K_b, N = b.shape
+        assert K == K_b, f"Matrix dimension mismatch: A has {K} cols, B has {K_b} rows"
+        
+        # Output should be BF16 type on the same device as inputs
+        out = torch.empty((M, N), dtype=torch.bfloat16, device=a.device)
+    
+    ops.gemm(out, a, b, as_, bs)
+    return out
+

build/torch28-cxx11-rocm64-x86_64-linux/gemm/_gemm_60de0b2.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aaedc163991e19d7ed36c1367fd5627daf741d54c1abfda128d00ea702cb18f6
+size 2195560

build/torch28-cxx11-rocm64-x86_64-linux/gemm/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _gemm_60de0b2
+ops = torch.ops._gemm_60de0b2
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_gemm_60de0b2::{op_name}"

flake.lock

@@ -0,0 +1,169 @@
+{
+  "nodes": {
+    "flake-compat": {
+      "locked": {
+        "lastModified": 1747046372,
+        "narHash": "sha256-CIVLLkVgvHYbgI2UpXvIIBJ12HWgX+fjA8Xf8PUmqCY=",
+        "owner": "edolstra",
+        "repo": "flake-compat",
+        "rev": "9100a0f413b0c601e0533d1d94ffd501ce2e7885",
+        "type": "github"
+      },
+      "original": {
+        "owner": "edolstra",
+        "repo": "flake-compat",
+        "type": "github"
+      }
+    },
+    "flake-compat_2": {
+      "locked": {
+        "lastModified": 1747046372,
+        "narHash": "sha256-CIVLLkVgvHYbgI2UpXvIIBJ12HWgX+fjA8Xf8PUmqCY=",
+        "owner": "edolstra",
+        "repo": "flake-compat",
+        "rev": "9100a0f413b0c601e0533d1d94ffd501ce2e7885",
+        "type": "github"
+      },
+      "original": {
+        "owner": "edolstra",
+        "repo": "flake-compat",
+        "type": "github"
+      }
+    },
+    "flake-utils": {
+      "inputs": {
+        "systems": "systems"
+      },
+      "locked": {
+        "lastModified": 1731533236,
+        "narHash": "sha256-l0KFg5HjrsfsO/JpG+r7fRrqm12kzFHyUHqHCVpMMbI=",
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "rev": "11707dc2f618dd54ca8739b309ec4fc024de578b",
+        "type": "github"
+      },
+      "original": {
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "type": "github"
+      }
+    },
+    "flake-utils_2": {
+      "inputs": {
+        "systems": "systems_2"
+      },
+      "locked": {
+        "lastModified": 1731533236,
+        "narHash": "sha256-l0KFg5HjrsfsO/JpG+r7fRrqm12kzFHyUHqHCVpMMbI=",
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "rev": "11707dc2f618dd54ca8739b309ec4fc024de578b",
+        "type": "github"
+      },
+      "original": {
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "type": "github"
+      }
+    },
+    "hf-nix": {
+      "inputs": {
+        "flake-compat": "flake-compat_2",
+        "flake-utils": "flake-utils_2",
+        "nixpkgs": "nixpkgs"
+      },
+      "locked": {
+        "lastModified": 1759385472,
+        "narHash": "sha256-a1YMZp3Yc1RJfLIObRKBTTbjMKL91IYbzTjG/HNZN+I=",
+        "owner": "huggingface",
+        "repo": "hf-nix",
+        "rev": "050dd78a64cb58fb1f9fb29ca498c73107a9a13e",
+        "type": "github"
+      },
+      "original": {
+        "owner": "huggingface",
+        "repo": "hf-nix",
+        "type": "github"
+      }
+    },
+    "kernel-builder": {
+      "inputs": {
+        "flake-compat": "flake-compat",
+        "flake-utils": "flake-utils",
+        "hf-nix": "hf-nix",
+        "nixpkgs": [
+          "kernel-builder",
+          "hf-nix",
+          "nixpkgs"
+        ]
+      },
+      "locked": {
+        "lastModified": 1759391810,
+        "narHash": "sha256-7thAD4hsNGvDb59rne7Vt0JsdyjK/DaeW3PIXWjueYc=",
+        "owner": "huggingface",
+        "repo": "kernel-builder",
+        "rev": "6029e51ead3bc6328fc817be44a546f366b84d93",
+        "type": "github"
+      },
+      "original": {
+        "owner": "huggingface",
+        "ref": "rocm-per-source-arches",
+        "repo": "kernel-builder",
+        "type": "github"
+      }
+    },
+    "nixpkgs": {
+      "locked": {
+        "lastModified": 1755963616,
+        "narHash": "sha256-6yD0ww/S8n+U2uPYcJZ3DRURP8Kx036GRpR2uPNZroE=",
+        "owner": "nixos",
+        "repo": "nixpkgs",
+        "rev": "73e96df7cff5783f45e21342a75a1540c4eddce4",
+        "type": "github"
+      },
+      "original": {
+        "owner": "nixos",
+        "ref": "nixos-unstable-small",
+        "repo": "nixpkgs",
+        "type": "github"
+      }
+    },
+    "root": {
+      "inputs": {
+        "kernel-builder": "kernel-builder"
+      }
+    },
+    "systems": {
+      "locked": {
+        "lastModified": 1681028828,
+        "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=",
+        "owner": "nix-systems",
+        "repo": "default",
+        "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e",
+        "type": "github"
+      },
+      "original": {
+        "owner": "nix-systems",
+        "repo": "default",
+        "type": "github"
+      }
+    },
+    "systems_2": {
+      "locked": {
+        "lastModified": 1681028828,
+        "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=",
+        "owner": "nix-systems",
+        "repo": "default",
+        "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e",
+        "type": "github"
+      },
+      "original": {
+        "owner": "nix-systems",
+        "repo": "default",
+        "type": "github"
+      }
+    }
+  },
+  "root": "root",
+  "version": 7
+}

flake.nix

@@ -0,0 +1,18 @@
+{
+  description = "Flake for GEMM kernel";
+
+  inputs = {
+    kernel-builder.url = "github:huggingface/kernel-builder/rocm-per-source-arches";
+  };
+
+  outputs =
+    {
+      self,
+      kernel-builder,
+    }:
+
+    kernel-builder.lib.genFlakeOutputs {
+      inherit self;
+      path = ./.;
+    };
+}

gemm/gemm_kernel.h

@@ -0,0 +1,896 @@
+// Pipeline GEMM kernel. This version is rushed written and may not applied to all shape.
+// Currently, only selected parameters is tested. (See gemm_launcher )
+#ifndef GEMM_KERNEL
+#define GEMM_KERNEL
+
+#include <cstdio>
+#include <hip/amd_detail/amd_hip_runtime.h>
+#include <hip/amd_detail/amd_warp_functions.h>
+#pragma clang diagnostic push
+#pragma clang diagnostic ignored "-Wunknown-attributes"
+#include "../include/gpu_libs.h"
+#include "../include/gpu_types.h"
+#include "../src/utils/arithmetic.h"
+#include "../include/clangd_workaround.h"
+#include <cstdlib>
+#include <cfloat>
+
+namespace gemm_kernel {
+
+template <typename data_type, int BATCH_SIZE> __device__ inline void read_batch(data_type *dst, const data_type *src) {
+    if constexpr ((sizeof(data_type) * BATCH_SIZE) == 2 * sizeof(ulong4)) {
+        *(reinterpret_cast<ulong4 *>(dst) + 0) = *(reinterpret_cast<const ulong4 *>(src) + 0);
+        *(reinterpret_cast<ulong4 *>(dst) + 1) = *(reinterpret_cast<const ulong4 *>(src) + 1);
+    } else if constexpr ((sizeof(data_type) * BATCH_SIZE) == sizeof(ulong4)) {
+        *reinterpret_cast<ulong4 *>(dst) = *reinterpret_cast<const ulong4 *>(src);
+    } else if constexpr (sizeof(data_type) * BATCH_SIZE == sizeof(ulong2)) {
+        *reinterpret_cast<ulong2 *>(dst) = *reinterpret_cast<const ulong2 *>(src);
+    } else if constexpr (sizeof(data_type) * BATCH_SIZE == sizeof(ulong1)) {
+        *reinterpret_cast<ulong1 *>(dst) = *reinterpret_cast<const ulong1 *>(src);
+    } else if constexpr (sizeof(data_type) * BATCH_SIZE == sizeof(uint1)) {
+        *reinterpret_cast<uint1 *>(dst) = *reinterpret_cast<const uint1 *>(src);
+    } else {
+#pragma unroll
+        for (int b = 0; b < BATCH_SIZE; ++b) {
+            dst[b] = src[b];
+        }
+    }
+}
+
+template <typename data_type, int BATCH_SIZE> __device__ inline void zero_batch(data_type *dst) {
+    if constexpr ((sizeof(data_type) * BATCH_SIZE) == sizeof(ulong4)) {
+        *reinterpret_cast<ulong4 *>(dst) = ulong4{};
+    } else if constexpr (sizeof(data_type) * BATCH_SIZE == sizeof(ulong2)) {
+        *reinterpret_cast<ulong2 *>(dst) = ulong2{};
+    } else if constexpr (sizeof(data_type) * BATCH_SIZE == sizeof(ulong1)) {
+        *reinterpret_cast<ulong1 *>(dst) = ulong1{};
+    } else if constexpr (sizeof(data_type) * BATCH_SIZE == sizeof(uint1)) {
+        *reinterpret_cast<uint *>(dst) = uint{};
+    } else {
+#pragma unroll
+        for (int b = 0; b < BATCH_SIZE; ++b) {
+            dst[b] = 0;
+        }
+    }
+}
+
+template <typename data_type, int DST_Y, int DST_X, int SRC_Y, int SRC_X, int BLOCK_DIM, int BATCH_SIZE>
+__device__ inline void load_input(data_type dst[DST_Y][DST_X], const data_type src[SRC_Y][SRC_X], const int begin_x,
+                                  const int begin_y) {
+    static_assert(BATCH_SIZE > 0);
+    /**
+      Consider (SRC_X % DST_X == 0) && (SRC_Y % DST_Y == 0)
+      Step 1:
+        [   ][***][   ][   ]
+        [   ][   ][   ][   ]
+        [   ][   ][   ][   ]
+        [   ][   ][   ][   ]
+      Step 2:
+        [   ][   ][   ][   ]
+        [   ][***][   ][   ]
+        [   ][   ][   ][   ]
+        [   ][   ][   ][   ]
+    */
+    static_assert((SRC_X % BATCH_SIZE == 0) && (SRC_Y % BATCH_SIZE == 0));
+    static_assert((DST_X % BATCH_SIZE == 0) && (DST_Y % BATCH_SIZE == 0));
+    static_assert(BATCH_SIZE <= DST_X && DST_X % BATCH_SIZE == 0);
+    const int begin_idx = threadIdx.x * BATCH_SIZE;
+    const constexpr int total_elements = DST_X * DST_Y;
+    const constexpr int elements_per_step = BLOCK_DIM * BATCH_SIZE;
+// FIXME: loop unrolling
+#pragma unroll
+    for (int k = begin_idx; k < total_elements; k += elements_per_step) {
+        int l_kx = k % DST_X;
+        int l_ky = k / DST_X;
+        int g_kx = l_kx + begin_x;
+        int g_ky = l_ky + begin_y;
+        auto *dst_flatten = &dst[l_ky][l_kx];
+        // const auto *src_flatten = &src[g_ky][g_kx];
+        // read_batch<data_type, BATCH_SIZE>(dst_flatten, src_flatten);
+        if (((SRC_X % DST_X == 0) || (g_kx < SRC_X)) && ((SRC_Y % DST_Y == 0) || (g_ky < SRC_Y))) {
+            const auto *src_flatten = &src[g_ky][g_kx];
+            read_batch<data_type, BATCH_SIZE>(dst_flatten, src_flatten);
+        } else {
+            zero_batch<data_type, BATCH_SIZE>(dst_flatten);
+        }
+    }
+}
+
+template <int PM, int PN, int QM, int QN, int QK, int QUANT_SIZE, int BLOCK_SIZE, int BATCH_SIZE>
+__device__ void load_scale(float s_s[PM][PN], const float sa[QK][QM], const float sb[QK][QN], const int m, const int n,
+                           const int k) {
+    constexpr int total_elements = PM * PN;
+    constexpr int elements_per_step = BLOCK_SIZE * BATCH_SIZE;
+    // static_assert(PN % BATCH_SIZE)
+
+    const int begin_idx = threadIdx.x * BATCH_SIZE;
+#pragma unroll
+    for (int idx = begin_idx; idx < total_elements; idx += elements_per_step) {
+        static_assert(BATCH_SIZE == 1);
+        int i = idx / PN;
+        int j = idx % PN;
+        if (((QM % PM == 0) || (m + i < QM)) && ((QN % PN == 0) || ((n + j) / QUANT_SIZE < QN))) {
+            s_s[i][j] = sa[k / QUANT_SIZE][(m + i)] * sb[k / QUANT_SIZE][(n) / QUANT_SIZE + j];
+        } else {
+            s_s[i][j] = 1.0f;
+        }
+    }
+}
+
+// don't use __builtin_readcyclecounter(), which would insert waitcnt
+__device__ auto getclock() {
+    uint64_t clk;
+    asm volatile("s_memtime %0" : "=r"(clk));
+    return clk;
+}
+
+
+template <typename Elem> __global__ void check_trans(const Elem *origin, const Elem *tranposed, int m, int n) {
+    auto x = threadIdx.x + blockIdx.x * blockDim.x;
+    auto y = threadIdx.y + blockIdx.y * blockDim.y;
+    if (x < m && y < n) {
+        if (origin[x * n + y] != tranposed[y * m + x]) {
+            printf("Error: %d %d\n", x, y);
+        }
+    }
+}
+
+template <typename in_data_type, typename acc_data_type, typename FragC, typename FragA, typename FragB, int PM, int PN,
+          int BM, int BN, int BK, int FRAG_M, int FRAG_N, int FRAG_K, int WMMA_M, int WMMA_N, int WMMA_K, int WARP_M,
+          int WARP_N, int BLOCK_SIZE, int BATCH_SIZE, int QUANT_SIZE>
+__device__ void wmma_compute(const in_data_type s_a[BM][BK + 8], const in_data_type s_b[BN][BK + 8],
+                             const float s_s[PN][PM], FragC frag_r[FRAG_M][FRAG_N], const int comp_c_frag_m,
+                             const int comp_c_frag_n) {
+    FragC frag_c[FRAG_M][FRAG_N];
+
+#pragma unroll
+    for (int i = 0; i < FRAG_M; i++) {
+#pragma unroll
+        for (int j = 0; j < FRAG_N; j++) {
+            wmma::fill_fragment(frag_c[i][j], 0.0F);
+        }
+    }
+
+#pragma unroll
+    for (int k = 0; k < FRAG_K; ++k) {
+#pragma unroll
+        for (int i = 0; i < FRAG_M; i++) {
+            FragA frag_a;
+            int s_a_row = k * WMMA_K;
+            int s_a_col = (comp_c_frag_m * FRAG_M + i) * WMMA_M;
+            wmma::load_matrix_sync(frag_a, &s_a[s_a_col][s_a_row], BK + 8);
+#pragma unroll
+            for (int j = 0; j < FRAG_N; j++) {
+                FragB frag_b;
+                int s_b_row = k * WMMA_K;
+                int s_b_col = (comp_c_frag_n * FRAG_N + j) * WMMA_N;
+                wmma::load_matrix_sync(frag_b, &s_b[s_b_col][s_b_row], BK + 8);
+
+                wmma::mma_sync(frag_c[i][j], frag_a, frag_b, frag_c[i][j]);
+            }
+        }
+    }
+#pragma unroll
+    for (int i = 0; i < FRAG_M; i++) {
+#pragma unroll
+        for (int j = 0; j < FRAG_N; j++) {
+#pragma unroll
+            for (int k = 0; k < FragC::num_elements; ++k) {
+#ifdef TEST_ON_RDNA4 // RDNA4, WAVE_SIZE = 32
+                int m = ((threadIdx.x & 16) >> 1) | (k & 7) | (comp_c_frag_m * FRAG_M + i) * WMMA_M;
+#else // CDNA3, WAVE_SIZE = 64
+// int m = ((threadIdx.x & 48) >> 2) | (k & 3) | (comp_c_frag_m * FRAG_M + i) * WMMA_M;
+#endif
+                // int n = ((threadIdx.x & 15) | (comp_c_frag_n * FRAG_N + j) * WMMA_N) / QUANT_SIZE;
+                auto lane = threadIdx.x % 64;
+                int m, n;
+                if constexpr (WMMA_M == 32) {
+                    // C or D i: (8 * floor(GPR_num / 4) % 32) + 4 * floor(lane / 32) + (GPR_num % 4)
+                    // C or D j: (lane % 32)
+                    m = (8 * (k / 4) % 32) + 4 * (lane / 32) + (k % 4);
+                    n = lane % 32;
+                } else {
+                    // C or D i: 4 * floor(lane / 16) + (GPR_num % 4)
+                    // C or D j: (lane % 16)
+                    m = 4 * (lane / 16) + (k % 4);
+                    n = lane % 16;
+                }
+                m += (comp_c_frag_m * FRAG_M + i) * WMMA_M;
+                n += (comp_c_frag_n * FRAG_N + j) * WMMA_N;
+                n = n / QUANT_SIZE;
+                // if(threadIdx.x == 192 && blockIdx.x ==0 && blockIdx.y == 0 && blockIdx.z == 0)
+                // printf("m: %d, n: %d\n", m, n);
+                float scale = s_s[n][m];
+                frag_r[i][j].x[k] += (acc_data_type)scale * (acc_data_type)frag_c[i][j].x[k];
+            }
+        }
+    }
+}
+
+__device__ rocwmma::bfloat16_t fast_f32tob16(float f) {
+    union {
+        float fp32;
+        unsigned int u32;
+    } u = {f};
+    u.u32 += 0x7fff + ((u.u32 >> 16) & 1);
+    auto ret = u.u32 >> 16;
+    return reinterpret_cast<rocwmma::bfloat16_t &>(ret);
+}
+
+template <typename acc_data_type, typename out_data_type, typename FragC, typename FragOut, int WMMA_M, int WMMA_N,
+          int BM, int BN, int M, int N, int FRAG_M, int FRAG_N>
+__device__ inline void store_result(out_data_type c[M][N], FragC frag_r[FRAG_M][FRAG_N], const int m, const int n,
+                                    const int comp_c_frag_m, const int comp_c_frag_n) {
+#pragma unroll
+    for (int i = 0; i < FRAG_M; i++) {
+#pragma unroll
+        for (int j = 0; j < FRAG_N; j++) {
+            int frag_m = comp_c_frag_m * FRAG_M + i;
+            int frag_n = comp_c_frag_n * FRAG_N + j;
+            int row = m + frag_m * WMMA_M;
+            int col = n + frag_n * WMMA_N;
+            if (((M % BM == 0) || (row < M)) && ((N % BN == 0) || (col < N))) {
+                out_data_type *c_ptr = &c[row][col];
+                if constexpr (sizeof(acc_data_type) == sizeof(out_data_type)) { // split_k
+                    auto lane = threadIdx.x % 64;
+                    #pragma unroll
+                    for (int k = 0; k < FragC::num_elements; ++k) {
+                        int m, n;
+                        if constexpr (WMMA_M == 32) {
+                            // C or D i: (8 * floor(GPR_num / 4) % 32) + 4 * floor(lane / 32) + (GPR_num % 4)
+                            // C or D j: (lane % 32)
+                            m = (8 * (k / 4) % 32) + 4 * (lane / 32) + (k % 4);
+                            n = lane % 32;
+                        } else {
+                            // C or D i: 4 * floor(lane / 16) + (GPR_num % 4)
+                            // C or D j: (lane % 16)
+                            m = 4 * (lane / 16) + (k % 4);
+                            n = lane % 16;
+                        }
+                        c_ptr[m * N + n] = frag_r[i][j].x[k];;
+                    }
+                    
+                    // wmma::store_matrix_sync(reinterpret_cast<out_data_type *>(c_ptr), frag_r[i][j], N,
+                    //                         wmma::mem_row_major);
+                } else if constexpr (sizeof(out_data_type) == sizeof(half)) {
+                    FragOut frag_out;
+                    static_assert(sizeof(half) == sizeof(out_data_type));
+                    static_assert(FragOut::num_elements == FragC::num_elements);
+                    for (int k = 0; k < FragOut::num_elements; ++k) {
+                        auto reg = fast_f32tob16(frag_r[i][j].x[k]);
+                        frag_out.x[k] = *reinterpret_cast<half *>(&reg);
+                    }
+                    wmma::store_matrix_sync(reinterpret_cast<half *>(c_ptr), frag_out, N, wmma::mem_row_major);
+                } else {
+                    static_assert(0, "Unsupported data type for output");
+                }
+            }
+        }
+    }
+}
+
+// a dummy template to allow inlcuding this file
+template <int Splitk> __global__ void reduce(uint32_t m, uint32_t n, const float *c_splitk, __hip_bfloat16 *c) {
+    auto tid = blockIdx.x * blockDim.x + threadIdx.x;
+    if (tid >= m * n) {
+        return;
+    }
+    float4 sum{};
+#pragma unroll
+    for (auto i = 0; i < Splitk; ++i) {
+        sum += *(float4 *)&c_splitk[i * (m * n) + tid * 4];
+    }
+    auto res =
+        rocwmma::make_vector(fast_f32tob16(sum.x), fast_f32tob16(sum.y), fast_f32tob16(sum.z), fast_f32tob16(sum.w));
+    *(decltype(res) *)&c[tid * 4] = res;
+}
+
+template<int M, int N, int SPLITK_FACTOR, int BLOCK_SIZE>
+__launch_bounds__(BLOCK_SIZE)
+__global__ void reduce_kernel(const float c_splitk[SPLITK_FACTOR][M][N], __hip_bfloat16 c[M][N]) {
+    auto tid = blockIdx.x * blockDim.x + threadIdx.x;
+    if (tid >= M * N) {
+        return;
+    }
+    float4 sum{};
+    #pragma unroll
+    for (auto i = 0; i < SPLITK_FACTOR; ++i) {
+        sum += *(float4 *)&reinterpret_cast<const float*>(c_splitk)[i * (M * N) + tid * 4];
+    }
+    auto res =
+        rocwmma::make_vector(fast_f32tob16(sum.x), fast_f32tob16(sum.y), fast_f32tob16(sum.z), fast_f32tob16(sum.w));
+    *(decltype(res) *)&reinterpret_cast< __BF16_TYPE*>(c)[tid * 4] = res;
+}
+
+
+#ifdef PARAMETERIZE_LIBRARY
+template <typename in_data_type,
+          typename acc_data_type, // Accumulator type (e.g., float)
+          typename out_data_type, // Output type (e.g., __hip_bfloat16)
+          int M, int N, int K,    // Matrix dimensions
+          int BM, int BN, int BK, // Tile dimensions
+          int QUANT_SIZE,         // Quantization block size
+          int BLOCK_SIZE,         // Block size
+          int WARP_M, int WARP_N, // Warp dimensions
+          int LDA, int LDB,
+          int LOAD_BATCH_SIZE>    // Load batch size for vectorized memory operations
+#else
+using in_data_type = __FP8_TYPE;
+using out_data_type = __BF16_TYPE;
+using acc_data_type = float;
+// constexpr int M = 4096, N = 4096, K = 4096;
+constexpr int M = 6144, N = 4608, K = 7168;
+constexpr int LDA = K, LDB = K;
+// constexpr int M = 512, N = 512, K = 512;
+constexpr int BM = 256, BN = 128, BK = 128;
+constexpr int QUANT_SIZE = 128, BLOCK_SIZE = 512;
+constexpr int LOAD_BATCH_SIZE = 16;
+#ifdef TEST_ON_RDNA4 // RDNA4, WAVE_SIZE = 32
+constexpr int WARP_M = 4, WARP_N = 2;
+#else                // CDNA3, WAVE_SIZE = 64
+constexpr int WARP_M = 4, WARP_N = 2;
+#endif
+#endif // End of parameterization
+__global__ __launch_bounds__(BLOCK_SIZE) void gemm_kernel(
+    const in_data_type a[M][LDA], const in_data_type b[N][LDB], out_data_type c[M][N],
+    const float sa[ceil_div(K, QUANT_SIZE)][M / 1], // Assuming M is divisible by 1 (always true)
+    const float sb[ceil_div(K, QUANT_SIZE)][ceil_div(N, QUANT_SIZE)]) {
+    // --- Start: Derived parameters and constants ---
+    constexpr int WMMA_M = 16; // Fixed WMMA dimension M
+    constexpr int WMMA_N = 16; // Fixed WMMA dimension N
+    constexpr int WMMA_K = 32; // Fixed WMMA dimension K (for FP8)
+
+    // WARP_M/N define the 2D arrangement of warps in the block grid.
+    // These might need adjustment based on BLOCK_DIM_X/Y strategy.
+    // Using fixed values based on the non-parameterized version for now.
+    // TODO: Derive WARP_M/N from BLOCK_DIM_X/Y if a flexible strategy is needed.
+    constexpr int WARP_NUM = WARP_M * WARP_N; // Total warps per block
+
+    // Assertion: Check if the assumed warp layout matches the block size
+    static_assert(WARP_NUM * WAVE_SIZE == BLOCK_SIZE, "WARP_M * WARP_N * WAVE_SIZE must equal BLOCK_SIZE");
+
+    // Fragments per warp
+    constexpr int FRAG_M_PER_WARP = BM / WMMA_M / WARP_M;
+    constexpr int FRAG_N_PER_WARP = BN / WMMA_N / WARP_N;
+    constexpr int FRAG_K = BK / WMMA_K; // Fragments along K dimension tile
+
+    static_assert(BM % (WMMA_M * WARP_M) == 0, "BM must be divisible by WMMA_M * WARP_M");
+    static_assert(BN % (WMMA_N * WARP_N) == 0, "BN must be divisible by WMMA_N * WARP_N");
+    static_assert(BK % WMMA_K == 0, "BK must be divisible by WMMA_K");
+    static_assert(BK >= 32, "BK must be at least 32");
+    // --- End: Derived parameters and constants ---
+
+    constexpr int QM = M;                        // Dimension M for scale A
+    constexpr int QN = ceil_div(N, QUANT_SIZE);  // Dimension N for scale B (quantized)
+    constexpr int QK = ceil_div(K, QUANT_SIZE);  // Dimension K for scales (quantized)
+    constexpr int PM = BM;                       // Block size M for scale A * B
+    constexpr int PN = ceil_div(BN, QUANT_SIZE); // Block size N for scale A * B
+
+    // Ensure derived fragment counts are positive
+    static_assert(FRAG_M_PER_WARP > 0, "FRAG_M_PER_WARP must be positive");
+    static_assert(FRAG_N_PER_WARP > 0, "FRAG_N_PER_WARP must be positive");
+    static_assert(FRAG_K > 0, "FRAG_K must be positive");
+
+    using FragA = wmma::fragment<wmma::matrix_a, WMMA_M, WMMA_N, WMMA_K, in_data_type, wmma::row_major>;
+    using FragB = wmma::fragment<wmma::matrix_b, WMMA_M, WMMA_N, WMMA_K, in_data_type, wmma::col_major>;
+    using FragC = wmma::fragment<wmma::accumulator, WMMA_M, WMMA_N, WMMA_K, acc_data_type>;
+    using FragOut = wmma::fragment<wmma::accumulator, WMMA_M, WMMA_N, WMMA_K,
+                                   half>; // Output uses half for storage via bfloat16 reinterpret
+
+    __shared__ in_data_type s_a[BM][BK + 8];
+    __shared__ in_data_type s_b[BN][BK + 8];
+    __shared__ acc_data_type s_s[PN][PM];           // Accumulator type for scales
+    FragC frag_r[FRAG_M_PER_WARP][FRAG_N_PER_WARP]; // Accumulator fragments
+
+    // handle splitk
+    a = (decltype(a))((in_data_type *)a + blockIdx.z * K);
+    b = (decltype(b))((in_data_type *)b + blockIdx.z * K);
+    c += blockIdx.z * M;
+    sa += blockIdx.z * QK;
+    sb += blockIdx.z * QK;
+
+    int tid = threadIdx.x;     // Linear thread ID within the block
+    int wid = tid / WAVE_SIZE; // Warp ID within the block
+
+    // Spilt and compute fragments
+    constexpr int iteration_over_k = ceil_div(K, BK); // Use ceil_div for potentially non-divisible K
+    static_assert(LOAD_BATCH_SIZE > 0, "LOAD_BATCH_SIZE must be positive");
+
+    constexpr auto PIPELINE = true;
+    // using LoadVec = rocwmma::VecT<float, LOAD_BATCH_SIZE / sizeof(float)>;
+    using LoadVec = __attribute__((__vector_size__(LOAD_BATCH_SIZE))) float;
+    static_assert(((BK * BM) % (BLOCK_SIZE * LOAD_BATCH_SIZE)) == 0,
+                  "BK * BM must be divisible by BLOCK_SIZE * LOAD_BATCH_SIZE");
+    static_assert(BK % LOAD_BATCH_SIZE == 0, "BK must be divisible by LOAD_BATCH_SIZE");
+    LoadVec reg_a[BK * BM / BLOCK_SIZE / LOAD_BATCH_SIZE];
+    LoadVec reg_b[BK * BN / BLOCK_SIZE / LOAD_BATCH_SIZE];
+    constexpr auto PK = ceil_div(BK, QUANT_SIZE);
+    static_assert(PK == 1, "PK must be 1 for now");
+    float reg_sa[ceil_div(PM, BLOCK_SIZE)];
+    float reg_sb[ceil_div(PN, BLOCK_SIZE)];
+
+    // threadblock swizzle
+    auto log_tile = 1;
+    auto block_idx_x = blockIdx.x >> log_tile;
+    auto block_idx_y = (blockIdx.y << log_tile) + ((blockIdx.x) & ((1 << (log_tile)) - 1));
+    if (block_idx_x >= ceil_div(N, BN) || block_idx_y >= ceil_div(M, BM)) {
+        return;
+    }
+
+    const int m = block_idx_y * BM;
+    const int n = block_idx_x * BN;
+    int k = 0;
+
+    auto global2reg = [&]() {
+#pragma unroll
+        for (int reg = 0; reg < sizeof(reg_sa) / sizeof(float); reg++) {
+            // NOTE: must iter over reg to make compiler unroll the loop
+            // and thus be able to allocate reg_a on register instead of on scratch memroy
+            int t = tid + reg * BLOCK_SIZE;
+            // NOTE: don't branch here
+            // if (t > PM) {
+            //     break;
+            // }
+            int i = t / PM;
+            int j = t % PM;
+            reg_sa[reg] = sa[k / QUANT_SIZE][(m + j)];
+        }
+#pragma unroll
+        for (int reg = 0; reg < sizeof(reg_sb) / sizeof(float); reg++) {
+            // NOTE: must iter over reg to make compiler unroll the loop
+            // and thus be able to allocate reg_a on register instead of on scratch memroy
+            int t = tid + reg * BLOCK_SIZE;
+            // NOTE: don't branch here
+            // if (t > PN) {
+            //     break;
+            // }
+            int i = t / PN;
+            int j = t % PN;
+            reg_sb[reg] = sb[k / QUANT_SIZE][(n) / QUANT_SIZE + j];
+        }
+#pragma unroll
+        for (int reg = 0; reg < sizeof(reg_a) / sizeof(LoadVec); reg++) {
+            // NOTE: must iter over reg to make compiler unroll the loop
+            // and thus be able to allocate reg_a on register instead of on scratch memroy
+            int t = tid * LOAD_BATCH_SIZE + reg * BLOCK_SIZE * LOAD_BATCH_SIZE;
+            int i = t / BK;
+            int j = t % BK;
+            reg_a[reg] = *(LoadVec *)&a[m + i][k + j];
+        }
+#pragma unroll
+        for (int reg = 0; reg < sizeof(reg_b) / sizeof(LoadVec); reg++) {
+            // NOTE: must iter over reg to make compiler unroll the loop
+            // and thus be able to allocate reg_a on register instead of on scratch memroy
+            int t = tid * LOAD_BATCH_SIZE + reg * BLOCK_SIZE * LOAD_BATCH_SIZE;
+            int i = t / BK;
+            int j = t % BK;
+            reg_b[reg] = *(LoadVec *)&b[n + i][k + j];
+        }
+    };
+
+    auto reg2lds = [&]() {
+#pragma unroll
+        for (int rega = 0; rega < sizeof(reg_sa) / sizeof(float); rega++) {
+            int ta = tid + rega * BLOCK_SIZE;
+            int j = ta % PM;
+#pragma unroll
+            for (int regb = 0; regb < sizeof(reg_sb) / sizeof(float); regb++) {
+                int tb = tid + regb * BLOCK_SIZE;
+                int i = tb % PN;
+                s_s[i][j] = reg_sa[rega] * reg_sb[regb];
+            }
+        }
+#pragma unroll
+        for (int reg = 0; reg < sizeof(reg_a) / sizeof(LoadVec); reg++) {
+            int t = tid * LOAD_BATCH_SIZE + reg * BLOCK_SIZE * LOAD_BATCH_SIZE;
+            int i = t / BK;
+            int j = t % BK;
+            *(LoadVec *)&s_a[i][j] = reg_a[reg];
+        }
+#pragma unroll
+        for (int reg = 0; reg < sizeof(reg_b) / sizeof(LoadVec); reg++) {
+            int t = tid * LOAD_BATCH_SIZE + reg * BLOCK_SIZE * LOAD_BATCH_SIZE;
+            int i = t / BK;
+            int j = t % BK;
+            *(LoadVec *)&s_b[i][j] = reg_b[reg];
+        }
+    };
+
+    if constexpr (PIPELINE) {
+        global2reg();
+    }
+
+// Initialize the output accumulator fragments to zero
+#pragma unroll
+    for (int i = 0; i < FRAG_M_PER_WARP; i++) {
+#pragma unroll
+        for (int j = 0; j < FRAG_N_PER_WARP; j++) {
+            wmma::fill_fragment(frag_r[i][j], 0.0f); // Use float literal
+        }
+    }
+
+    if constexpr (!PIPELINE) {
+        global2reg();
+    }
+
+    reg2lds();
+
+    for (int bk = 1; bk < iteration_over_k; bk++) {
+        k = bk * BK;
+
+        // Calculate remaining K for boundary checks if needed (not currently used by load_input)
+        // const int k_rem = K - k;
+
+        // Load data into shared memory
+        // load_input<in_data_type, BK, BM, K, M, BLOCK_SIZE, 32>(
+        //     s_a, a, m, k);
+        // load_input<in_data_type, BK, BN, K, N, BLOCK_SIZE, 32>(
+        //     s_b, b, n, k);
+        // Load scales into shared memory (using acc_data_type for s_s)
+        // load_scale<PM, PN, QM, QN, QK, QUANT_SIZE, BLOCK_SIZE, 1>(
+        //     s_s, sa, sb, m, n, k);
+
+        if constexpr (PIPELINE) {
+            global2reg();
+        }
+
+        __syncthreads();
+
+        // Perform matrix multiplication using WMMA
+        wmma_compute<in_data_type, acc_data_type, FragC, FragA, FragB, PM, PN, BM, BN, BK, FRAG_M_PER_WARP,
+                     FRAG_N_PER_WARP, FRAG_K, WMMA_M, WMMA_N, WMMA_K, WARP_M, WARP_N, BLOCK_SIZE, LOAD_BATCH_SIZE,
+                     QUANT_SIZE>( // Pass calculated BLOCK_SIZE and LOAD_BATCH_SIZE
+            s_a, s_b, s_s, frag_r, wid / WARP_N, wid % WARP_N);
+        __syncthreads();
+
+        if constexpr (!PIPELINE) {
+            global2reg();
+        }
+
+        // __builtin_amdgcn_sched_barrier(0);
+
+        reg2lds();
+    }
+    __syncthreads();
+    wmma_compute<in_data_type, acc_data_type, FragC, FragA, FragB, PM, PN, BM, BN, BK, FRAG_M_PER_WARP, FRAG_N_PER_WARP,
+                 FRAG_K, WMMA_M, WMMA_N, WMMA_K, WARP_M, WARP_N, BLOCK_SIZE, LOAD_BATCH_SIZE,
+                 QUANT_SIZE>( // Pass calculated BLOCK_SIZE and LOAD_BATCH_SIZE
+        s_a, s_b, s_s, frag_r, wid / WARP_N, wid % WARP_N);
+    // Store results from accumulator fragments to global memory
+    store_result<acc_data_type, out_data_type, FragC, FragOut, WMMA_M, WMMA_N, BM, BN, M, N, FRAG_M_PER_WARP,
+                 FRAG_N_PER_WARP>(c, frag_r, block_idx_y * BM, block_idx_x * BN, wid / WARP_N, wid % WARP_N);
+};
+
+}; // namespace gemm_kernel
+
+HOST_CODE_BELOW
+
+#ifndef PARAMETERIZE_LIBRARY
+// Define type aliases to match those in the namespace
+using fp8_type = gemm_kernel::in_data_type;       // __hip_fp8_e4m3
+using fp16_type = gemm_kernel::out_data_type;     // __hip_bfloat16
+using acc_data_type = gemm_kernel::acc_data_type; // float
+
+// Define constants to match those in the namespace
+constexpr int M = gemm_kernel::M;   // 4096
+constexpr int N = gemm_kernel::N;   // 4096
+constexpr int K = gemm_kernel::K;   // 4096
+constexpr int BM = gemm_kernel::BM; // 256
+constexpr int BN = gemm_kernel::BN; // 128
+constexpr int BK = gemm_kernel::BK; // 32
+constexpr int BLOCK_SIZE = gemm_kernel::BLOCK_SIZE;
+constexpr int QUANT_SIZE = gemm_kernel::QUANT_SIZE; // 128
+
+// Define derived constants for the test
+constexpr int QK = K / QUANT_SIZE;
+constexpr int QM = M;
+constexpr int QN = N / QUANT_SIZE;
+
+// Helper function to check HIP errors
+#define CHECK_HIP_ERROR(val) check((val), #val, __FILE__, __LINE__)
+template <typename T> void check(T err, const char *const func, const char *const file, const int line) {
+    if (err != hipSuccess) {
+        fprintf(stderr, "HIP Runtime Error at: %s:%d\n", file, line);
+        fprintf(stderr, "%s %s\n", hipGetErrorString(err), func);
+        exit(1);
+    }
+}
+
+// Define a macro to check HIP errors
+#define HIP_CALL(call)                                                                                                 \
+    do {                                                                                                               \
+        hipError_t err = call;                                                                                         \
+        if (err != hipSuccess) {                                                                                       \
+            fprintf(stderr, "HIP Error: %s at %s:%d\n", hipGetErrorString(err), __FILE__, __LINE__);                   \
+            exit(EXIT_FAILURE);                                                                                        \
+        }                                                                                                              \
+    } while (0)
+
+// CPU matrix multiplication implementation for result verification
+void cpu_gemm(const fp8_type a[K][M], const fp8_type b[K][N], fp16_type c[M][N], const float sa[QK][QM],
+              const float sb[QK][QN]) {
+    float(*rc)[N] = new float[M][N];
+    for (int m = 0; m < M; ++m) {
+        for (int n = 0; n < N; ++n) {
+            rc[m][n] = 0.0f;
+        }
+    }
+    for (int k = 0; k < K; ++k) {
+        for (int m = 0; m < M; ++m) {
+            for (int n = 0; n < N; ++n) {
+                float scale = sa[k / QUANT_SIZE][m] * sb[k / QUANT_SIZE][n / QUANT_SIZE];
+                rc[m][n] += (scale * (float)a[k][m] * (float)b[k][n]);
+            }
+        }
+    }
+    for (int m = 0; m < M; ++m) {
+        for (int n = 0; n < N; ++n) {
+            c[m][n] = (fp16_type)rc[m][n];
+        }
+    }
+    delete[] rc;
+}
+
+int main() {
+    // Allocate host memory
+    fp8_type(*h_a)[M] = new fp8_type[K][M];
+    fp8_type(*h_b)[N] = new fp8_type[K][N];
+    fp16_type(*h_c)[N] = new fp16_type[M][N];
+    fp16_type(*h_c_ref)[N] = new fp16_type[M][N];
+
+    // Allocate host memory for quantization scale factors
+    float(*h_sa)[QM] = new float[QK][QM];
+    float(*h_sb)[QN] = new float[QK][QN];
+
+    // Initialize input data
+    for (int i = 0; i < K; ++i) {
+        for (int j = 0; j < M; ++j) {
+            h_a[i][j] = (fp8_type)((rand() % 10000) / 10000.0f);
+        }
+    }
+    for (int i = 0; i < K; ++i) {
+        for (int j = 0; j < N; ++j) {
+            h_b[i][j] = (fp8_type)((rand() % 10000) / 10000.0f);
+        }
+    }
+
+    // Initialize quantization scale factors
+    for (int i = 0; i < QK; ++i) {
+        for (int j = 0; j < QM; ++j) {
+            h_sa[i][j] = 1.0f;
+        }
+    }
+    for (int i = 0; i < QK; ++i) {
+        for (int j = 0; j < QN; ++j) {
+            h_sb[i][j] = 1.0f;
+        }
+    }
+
+    // Allocate device memory
+    fp8_type(*d_a)[K];
+    fp8_type(*d_b)[K];
+    fp16_type(*d_c)[N];
+    float(*d_sa)[QM];
+    float(*d_sb)[QN];
+
+    CHECK_HIP_ERROR(hipMalloc(&d_a, K * M * sizeof(fp8_type)));
+    CHECK_HIP_ERROR(hipMalloc(&d_b, K * N * sizeof(fp8_type)));
+    CHECK_HIP_ERROR(hipMalloc(&d_c, M * N * sizeof(fp16_type)));
+    CHECK_HIP_ERROR(hipMalloc(&d_sa, QK * QM * sizeof(float)));
+    CHECK_HIP_ERROR(hipMalloc(&d_sb, QK * QN * sizeof(float)));
+
+    // Copy data from host memory to device memory
+    CHECK_HIP_ERROR(hipMemcpy(d_a, h_a, K * M * sizeof(fp8_type), hipMemcpyHostToDevice));
+    CHECK_HIP_ERROR(hipMemcpy(d_b, h_b, K * N * sizeof(fp8_type), hipMemcpyHostToDevice));
+    CHECK_HIP_ERROR(hipMemcpy(d_sa, h_sa, QK * QM * sizeof(float), hipMemcpyHostToDevice));
+    CHECK_HIP_ERROR(hipMemcpy(d_sb, h_sb, QK * QN * sizeof(float), hipMemcpyHostToDevice));
+
+    // Calculate grid and block sizes - ensure coverage of the entire matrix
+    dim3 grid((N + BN - 1) / BN, (M + BM - 1) / BM);
+    dim3 block(BLOCK_SIZE);
+
+    // Ensure block size is a multiple of 32, since warp size is 32
+    if (BLOCK_SIZE % 32 != 0) {
+        printf("Error: Block size must be a multiple of warp size (32)\n");
+        return 1;
+    }
+
+    // Check if device supports required compute capability
+    int deviceId;
+    HIP_CALL(hipGetDevice(&deviceId));
+    hipDeviceProp_t deviceProp;
+    HIP_CALL(hipGetDeviceProperties(&deviceProp, deviceId));
+
+    if (deviceProp.major < 7) {
+        printf("Error: This kernel requires a GPU with compute capability 7.0 or higher\n");
+        return 1;
+    }
+
+    printf("Running GEMM kernel with grid(%d,%d), block(%d)...\n", grid.x, grid.y, block.x);
+
+    // Query and print kernel and device information
+    printf("Querying kernel and device information...\n");
+
+    // Get device properties
+    HIP_CALL(hipGetDeviceProperties(&deviceProp, deviceId));
+    printf("Device Name: %s\n", deviceProp.name);
+    printf("Total Global Memory: %lu bytes\n", deviceProp.totalGlobalMem);
+    printf("Shared Memory per Block: %lu bytes\n", deviceProp.sharedMemPerBlock);
+    printf("Registers per Block: %d\n", deviceProp.regsPerBlock);
+    printf("Warp Size: %d\n", deviceProp.warpSize);
+    printf("Max Threads per Block: %d\n", deviceProp.maxThreadsPerBlock);
+    printf("Max Threads per Multiprocessor: %d\n", deviceProp.maxThreadsPerMultiProcessor);
+    printf("Number of Multiprocessors: %d\n", deviceProp.multiProcessorCount);
+
+    // Query kernel attributes
+    hipFuncAttributes funcAttr;
+    HIP_CALL(hipFuncGetAttributes(&funcAttr, (const void *)gemm_kernel::gemm_kernel));
+    printf("Kernel Attributes:\n");
+    printf("  Shared Memory Size: %lu bytes\n", funcAttr.sharedSizeBytes);
+    printf("  Number of Registers: %d\n", funcAttr.numRegs);
+    printf("  Max Threads per Block: %d\n", funcAttr.maxThreadsPerBlock);
+    printf("  Local Memory Size: %lu bytes\n", funcAttr.localSizeBytes);
+
+    // Zero the C matrix before launching kernel
+    CHECK_HIP_ERROR(hipMemset(d_c, 0, M * N * sizeof(fp16_type)));
+
+    // Perform warmup runs
+    printf("Performing warmup runs...\n");
+    gemm_kernel::gemm_kernel<<<grid, block>>>(d_a, d_b, d_c, d_sa, d_sb);
+    CHECK_HIP_ERROR(hipDeviceSynchronize());
+    gemm_kernel::gemm_kernel<<<grid, block>>>(d_a, d_b, d_c, d_sa, d_sb);
+    CHECK_HIP_ERROR(hipDeviceSynchronize());
+
+    // Declare and create timing events
+    hipEvent_t start, stop;
+    HIP_CALL(hipEventCreate(&start));
+    HIP_CALL(hipEventCreate(&stop));
+
+    // Ensure device synchronization before formal timing
+    CHECK_HIP_ERROR(hipDeviceSynchronize());
+    HIP_CALL(hipEventRecord(start));
+
+    // Launch kernel
+    printf("Launching kernel...\n");
+    gemm_kernel::gemm_kernel<<<grid, block>>>(d_a, d_b, d_c, d_sa, d_sb);
+
+    // Record end time and calculate execution time
+    HIP_CALL(hipEventRecord(stop));
+
+    // Record end time and calculate execution time
+    HIP_CALL(hipEventSynchronize(stop));
+    float milliseconds = 0;
+    HIP_CALL(hipEventElapsedTime(&milliseconds, start, stop));
+    printf("Kernel execution time: %f ms\n", milliseconds);
+
+    // Check HIP errors
+    CHECK_HIP_ERROR(hipGetLastError());
+
+    // Calculate GPU performance metrics
+    double operations = 2.0 * M * N * K; // Each multiply-add operation counts as 2 floating-point operations
+    double seconds = milliseconds / 1000.0;
+    double tflops = (operations / seconds) / 1e12;
+    printf("GPU Performance: %.2f TFLOPS\n", tflops);
+
+    return 0;
+
+    // Copy results from device memory back to host memory
+    CHECK_HIP_ERROR(hipMemcpy(h_c, d_c, M * N * sizeof(fp16_type), hipMemcpyDeviceToHost));
+
+    // Calculate reference results
+    printf("Computing reference result on CPU...\n");
+    cpu_gemm(h_a, h_b, h_c_ref, h_sa, h_sb);
+
+    // Print the first 10 values for comparison
+    printf("First 10 values (GPU vs CPU):\n");
+    int print_count = 0;
+    for (int i = 0; i < M && print_count < 10; ++i) {
+        for (int j = 0; j < N && print_count < 10; ++j) {
+            printf("  [%d, %d]: GPU=%f, CPU=%f\n", i, j, (float)h_c[i][j], (float)h_c_ref[i][j]);
+            print_count++;
+        }
+    }
+
+    // Verify results
+    printf("Verifying results...\n");
+    int errors = 0;
+    float max_abs_diff = 0.0f;
+    float max_rel_diff = 0.0f;
+    struct ErrorInfo {
+        int row, col;
+        float gpu_val, cpu_val, abs_diff, rel_diff;
+    };
+    ErrorInfo first_10_errors[10];
+    ErrorInfo max_10_errors[10] = {};
+
+    // Add a configurable variable for the number of errors to output
+    int max_errors_to_output = 10; // You can modify this value as needed
+
+    for (int i = 0; i < M; ++i) {
+        for (int j = 0; j < N; ++j) {
+            float gpu_val = (float)h_c[i][j];
+            float cpu_val = (float)h_c_ref[i][j];
+            float abs_diff;
+            float rel_diff;
+
+            if (std::isnan(gpu_val) || std::isnan(cpu_val)) {
+                abs_diff = INFINITY;
+                rel_diff = INFINITY;
+            } else {
+                abs_diff = abs(gpu_val - cpu_val);
+                rel_diff = abs_diff / (abs(cpu_val) + FLT_EPSILON);
+            }
+
+            // Track max absolute and relative differences
+            max_abs_diff = fmaxf(max_abs_diff, abs_diff);
+            max_rel_diff = fmaxf(max_rel_diff, rel_diff);
+
+            // Record first 10 errors
+            if (errors < max_errors_to_output && (rel_diff > 1e-2 || abs_diff > 1e-3)) {
+                first_10_errors[errors] = {i, j, gpu_val, cpu_val, abs_diff, rel_diff};
+            }
+
+            // Track top 10 largest errors
+            if (rel_diff > 1e-2 || abs_diff > 1e-3) {
+                errors++;
+                for (int k = 0; k < max_errors_to_output; ++k) {
+                    if (abs_diff > max_10_errors[k].abs_diff) {
+                        for (int l = max_errors_to_output - 1; l > k; --l) {
+                            max_10_errors[l] = max_10_errors[l - 1];
+                        }
+                        max_10_errors[k] = {i, j, gpu_val, cpu_val, abs_diff, rel_diff};
+                        break;
+                    }
+                }
+            }
+        }
+    }
+
+    // Print first 10 errors
+    printf("First %d errors:\n", max_errors_to_output);
+    for (int i = 0; i < fmin(errors, max_errors_to_output); ++i) {
+        printf("Error at [%d, %d]: GPU=%f, CPU=%f, AbsDiff=%f, RelDiff=%f\n", first_10_errors[i].row,
+               first_10_errors[i].col, first_10_errors[i].gpu_val, first_10_errors[i].cpu_val,
+               first_10_errors[i].abs_diff, first_10_errors[i].rel_diff);
+    }
+
+    // Print top 10 largest errors
+    printf("Top %d largest errors:\n", max_errors_to_output);
+    for (int i = 0; i < max_errors_to_output && max_10_errors[i].abs_diff > 0; ++i) {
+        printf("Error at [%d, %d]: GPU=%f, CPU=%f, AbsDiff=%f, RelDiff=%f\n", max_10_errors[i].row,
+               max_10_errors[i].col, max_10_errors[i].gpu_val, max_10_errors[i].cpu_val, max_10_errors[i].abs_diff,
+               max_10_errors[i].rel_diff);
+    }
+
+    printf("Max abs_diff: %f, Max rel_diff: %f\n", max_abs_diff, max_rel_diff);
+    if (errors == 0) {
+        printf("Test PASSED!\n");
+    } else {
+        printf("Test FAILED with %d errors\n", errors);
+    }
+
+    // Calculate performance
+    double flops = 2.0 * M * N * K;
+    double gflops = (flops * 1e-9) / (milliseconds * 1e-3);
+    printf("Performance: %.2f GFLOPS\n", gflops);
+
+    // Free memory
+    delete[] h_a;
+    delete[] h_b;
+    delete[] h_c;
+    delete[] h_c_ref;
+    delete[] h_sa;
+    delete[] h_sb;
+    HIP_CALL(hipFree(d_a));
+    HIP_CALL(hipFree(d_b));
+    HIP_CALL(hipFree(d_c));
+    HIP_CALL(hipFree(d_sa));
+    HIP_CALL(hipFree(d_sb));
+    HIP_CALL(hipEventDestroy(start));
+    HIP_CALL(hipEventDestroy(stop));
+
+    return 0;
+}
+#endif
+#pragma clang diagnostic pop
+#endif

gemm/gemm_kernel_legacy.h

@@ -0,0 +1,377 @@
+// Legacy version of gemm kernel, support all shape and various value of parameters (BM, BN, BK, etc.)
+// It has been replace with faster pipeline version.
+#pragma once
+#include <cstdio>
+#include "../include/gpu_libs.h"
+#include "../include/gpu_types.h"
+#include "../src/utils/arithmetic.h"
+#include "../include/clangd_workaround.h"
+#include <cstdlib>
+#include <cfloat>
+
+DEVICE_CODE_BELOW
+namespace gemm_kernel_legacy {
+
+
+
+template <typename data_type, int BATCH_SIZE> 
+__device__ inline void read_batch(data_type *dst, const data_type *src) {
+    if constexpr ((sizeof(data_type) * BATCH_SIZE) == 2 * sizeof(ulong4)) {
+        *(reinterpret_cast<ulong4 *>(dst) + 0) = *(reinterpret_cast<const ulong4 *>(src) + 0);
+        *(reinterpret_cast<ulong4 *>(dst) + 1) = *(reinterpret_cast<const ulong4 *>(src) + 1);
+    } else if constexpr ((sizeof(data_type) * BATCH_SIZE) == sizeof(ulong4)) {
+        *reinterpret_cast<ulong4 *>(dst) = *reinterpret_cast<const ulong4 *>(src);
+    } else if constexpr (sizeof(data_type) * BATCH_SIZE == sizeof(ulong2)) {
+        *reinterpret_cast<ulong2 *>(dst) = *reinterpret_cast<const ulong2 *>(src);
+    } else if constexpr (sizeof(data_type) * BATCH_SIZE == sizeof(ulong1)) {
+        *reinterpret_cast<ulong1 *>(dst) = *reinterpret_cast<const ulong1 *>(src);
+    } else if constexpr (sizeof(data_type) * BATCH_SIZE == sizeof(uint1)) {
+        *reinterpret_cast<uint1 *>(dst) = *reinterpret_cast<const uint1 *>(src);
+    } else {
+        #pragma unroll
+        for (int b = 0; b < BATCH_SIZE; ++b) {
+            dst[b] = src[b];
+        }
+    }
+}
+
+template <typename data_type, int BATCH_SIZE> 
+__device__ inline void zero_batch(data_type *dst) {
+    if constexpr ((sizeof(data_type) * BATCH_SIZE) == sizeof(ulong4)) {
+        *reinterpret_cast<ulong4 *>(dst) = ulong4{};
+    } else if constexpr (sizeof(data_type) * BATCH_SIZE == sizeof(ulong2)) {
+        *reinterpret_cast<ulong2 *>(dst) = ulong2{};
+    } else if constexpr (sizeof(data_type) * BATCH_SIZE == sizeof(ulong1)) {
+        *reinterpret_cast<ulong1 *>(dst) = ulong1{};
+    } else if constexpr (sizeof(data_type) * BATCH_SIZE == sizeof(uint1)) {
+        *reinterpret_cast<uint *>(dst) = uint{};
+    } else {
+        #pragma unroll
+        for (int b = 0; b < BATCH_SIZE; ++b) {
+            dst[b] = 0;
+        }
+    }
+}
+
+template <typename data_type, int DST_Y, int DST_X, int SRC_Y, int SRC_X, int BLOCK_DIM, int BATCH_SIZE>
+__device__ inline void load_input(data_type dst[DST_Y][DST_X], const data_type src[SRC_Y][SRC_X],
+                                         const int begin_x, const int begin_y) {
+  static_assert(BATCH_SIZE > 0);
+  /**
+    Consider (SRC_X % DST_X == 0) && (SRC_Y % DST_Y == 0)
+    Step 1:
+      [   ][***][   ][   ]
+      [   ][   ][   ][   ]
+      [   ][   ][   ][   ]
+      [   ][   ][   ][   ]
+    Step 2:
+      [   ][   ][   ][   ]
+      [   ][***][   ][   ]
+      [   ][   ][   ][   ]
+      [   ][   ][   ][   ]
+  */
+  static_assert((SRC_X % BATCH_SIZE == 0) && (SRC_Y % BATCH_SIZE == 0));
+  static_assert((DST_X % BATCH_SIZE == 0) && (DST_Y % BATCH_SIZE == 0));
+  static_assert(BATCH_SIZE <= DST_X && DST_X % BATCH_SIZE == 0);
+  const int begin_idx = threadIdx.x * BATCH_SIZE;
+  const constexpr int total_elements = DST_X * DST_Y;
+  const constexpr int elements_per_step = BLOCK_DIM * BATCH_SIZE;
+  // FIXME: loop unrolling
+  #pragma unroll
+  for (int k = begin_idx; k < total_elements; k += elements_per_step) {
+    int l_kx = k % DST_X;
+    int l_ky = k / DST_X;
+    int g_kx = l_kx + begin_x;
+    int g_ky = l_ky + begin_y;
+    auto *dst_flatten = &dst[l_ky][l_kx];
+    // const auto *src_flatten = &src[g_ky][g_kx];
+    // read_batch<data_type, BATCH_SIZE>(dst_flatten, src_flatten);
+    if (((SRC_X % DST_X == 0) || (g_kx < SRC_X)) && ((SRC_Y % DST_Y == 0) || (g_ky < SRC_Y))) {
+      const auto *src_flatten = &src[g_ky][g_kx];
+      read_batch<data_type, BATCH_SIZE>(dst_flatten, src_flatten);
+    } else {
+      zero_batch<data_type, BATCH_SIZE>(dst_flatten);
+    }
+  }
+}
+
+template <int PM, int PN, int QM, int QN, int QK, int QUANT_SIZE, int BLOCK_SIZE, int BATCH_SIZE>
+__device__ void load_scale(float s_s[PM][PN], const float sa[QK][QM], const float sb[QK][QN], 
+                const int m, const int n, const int k) {
+    constexpr int total_elements = PM * PN;
+    constexpr int elements_per_step = BLOCK_SIZE * BATCH_SIZE;
+    // static_assert(PN % BATCH_SIZE)
+    
+    const int begin_idx = threadIdx.x * BATCH_SIZE;
+    #pragma unroll
+    for (int idx = begin_idx; idx < total_elements; idx += elements_per_step) {
+        static_assert(BATCH_SIZE == 1);
+        int i = idx / PN;
+        int j = idx % PN;
+        if (((QM % PM == 0) || (m + i < QM)) && ((QN % PN == 0) || ((n + j) / QUANT_SIZE < QN))) {
+            s_s[i][j] = sa[k / QUANT_SIZE][(m + i)] * sb[k / QUANT_SIZE][(n) / QUANT_SIZE + j];
+        } else {
+            s_s[i][j] = 1.0f;
+        }
+    }
+    
+}
+
+template <typename in_data_type, typename acc_data_type,
+    typename FragC, typename FragA, typename FragB, 
+    int PM, int PN,
+    int BM, int BN, int BK, 
+    int FRAG_M, int FRAG_N, int FRAG_K, 
+    int WMMA_M, int WMMA_N, int WMMA_K,
+    int WARP_M, int WARP_N,
+    int BLOCK_SIZE, int BATCH_SIZE, int QUANT_SIZE>
+__device__ void wmma_compute(
+    const in_data_type s_a[BK][BM],
+    const in_data_type s_b[BK][BN],
+    const float s_s[PM][PN],
+    FragC frag_r[FRAG_M][FRAG_N],
+    const int comp_c_frag_m,
+    const int comp_c_frag_n
+) {
+    FragA frag_a[FRAG_K][FRAG_M]; 
+    FragB frag_b[FRAG_K][FRAG_N];
+
+    // Spilt k over BK
+    for (int k = 0; k < FRAG_K; ++k) {
+        #pragma unroll
+        for (int i = 0; i < FRAG_M; ++i) {
+            int s_a_row = k * WMMA_K;
+            int s_a_col = (comp_c_frag_m * FRAG_M + i) * WMMA_M;
+            wmma::load_matrix_sync(frag_a[k][i], &s_a[s_a_row][s_a_col], BM);
+        }
+        #pragma unroll
+        for (int j = 0; j < FRAG_N; ++j) {
+            int s_b_row = k * WMMA_K;
+            int s_b_col = (comp_c_frag_n * FRAG_N + j) * WMMA_N;
+            wmma::load_matrix_sync(frag_b[k][j], &s_b[s_b_row][s_b_col], BN);
+        }
+    }
+
+    #pragma unroll
+    for (int i = 0; i < FRAG_M; i++) {
+        #pragma unroll
+        for (int j = 0; j < FRAG_N; j++) {
+            FragC frag_c;
+            wmma::fill_fragment(frag_c, 0.0F);
+            #pragma unroll
+            for (int k = 0; k < FRAG_K; ++k) {
+                wmma::mma_sync(frag_c, frag_a[k][i], frag_b[k][j], frag_c);
+            }
+            #pragma unroll
+            for (int k = 0; k < FragC::num_elements; ++k) {
+                #ifdef TEST_ON_RDNA4 // RDNA4, WAVE_SIZE = 32
+                int m = ((threadIdx.x & 16) >> 1) | (k & 7) | (comp_c_frag_m * FRAG_M + i) * WMMA_M;
+                #else // CDNA3, WAVE_SIZE = 64
+                int m = ((threadIdx.x & 48) >> 2) | (k & 3) | (comp_c_frag_m * FRAG_M + i) * WMMA_M;
+                #endif
+                int n = ((threadIdx.x & 15) | (comp_c_frag_n * FRAG_N + j) * WMMA_N) / QUANT_SIZE;
+                float scale = s_s[m][n];
+                frag_r[i][j].x[k] += (acc_data_type)scale * (acc_data_type)frag_c.x[k];
+            }  
+        }
+    }
+}
+
+
+template <typename acc_data_type, typename out_data_type, 
+typename FragC, typename FragOut, int WMMA_M, int WMMA_N, 
+int BM, int BN, int M, int N, int FRAG_M, int FRAG_N>
+__device__ inline void store_result(
+    out_data_type c[M][N], 
+    FragC frag_r[FRAG_M][FRAG_N], 
+    const int m,
+    const int n,
+    const int comp_c_frag_m, 
+    const int comp_c_frag_n
+) {
+    #pragma unroll
+    for (int i = 0; i < FRAG_M; i++) {
+        #pragma unroll
+        for (int j = 0; j < FRAG_N; j++) {
+            int frag_m = comp_c_frag_m * FRAG_M + i;
+            int frag_n = comp_c_frag_n * FRAG_N + j;
+            int row = m + frag_m * WMMA_M;
+            int col = n + frag_n * WMMA_N;
+            if (((M % BM == 0) || (row < M)) && ((N % BN == 0) || (col < N))) {
+                out_data_type *c_ptr = &c[row][col];
+                if constexpr (sizeof(acc_data_type) == sizeof(out_data_type)) {
+                    wmma::store_matrix_sync(reinterpret_cast<out_data_type*>(c_ptr), frag_r[i][j], N, wmma::mem_row_major);
+                } else if constexpr (sizeof(out_data_type) == sizeof(half)) {
+                    FragOut frag_out;
+                    static_assert(sizeof(half) == sizeof(out_data_type));
+                    static_assert(FragOut::num_elements == FragC::num_elements);
+                    for (int k=0;k<FragOut::num_elements;++k) {
+                        __hip_bfloat16 reg = frag_r[i][j].x[k];
+                        frag_out.x[k] = *reinterpret_cast<half*>(&reg);
+                    }
+                    wmma::store_matrix_sync(reinterpret_cast<half*>(c_ptr), frag_out, N, wmma::mem_row_major);
+                } else {
+                    static_assert(0, "Unsupported data type for output");
+                }
+
+            }
+        }
+    }
+}
+
+// a dummy template to allow inlcuding this file
+template<int Dummy=0>
+__global__ void reduce(uint32_t m, uint32_t n, uint32_t splitk, const float *c_splitk, __hip_bfloat16 *c) {
+    auto tid = blockIdx.x * blockDim.x + threadIdx.x;
+    if (tid >= m * n) {
+        return;
+    }
+    float sum = 0;
+    for (auto i = 0; i < splitk; ++i) {
+        sum += c_splitk[i * (m * n) + tid];
+    }
+    c[tid] = sum;
+}
+
+    
+#ifdef PARAMETERIZE_LIBRARY
+template <
+    typename in_data_type,
+    typename acc_data_type, // Accumulator type (e.g., float)
+    typename out_data_type, // Output type (e.g., __hip_bfloat16)
+    int M, int N, int K,    // Matrix dimensions
+    int BM, int BN, int BK, // Tile dimensions
+    int QUANT_SIZE,         // Quantization block size
+    int BLOCK_SIZE,         // Block size
+    int WARP_M, int WARP_N // Warp dimensions
+>
+#else
+using in_data_type = __FP8_TYPE;
+using out_data_type = __BF16_TYPE;
+using acc_data_type = float;
+// constexpr int M = 4096, N = 4096, K = 4096;
+constexpr int M = 96, N = 1024, K = 1024;
+// constexpr int M = 512, N = 512, K = 512;
+constexpr int BM = 64, BN = 256, BK = 32;
+constexpr int QUANT_SIZE = 128, BLOCK_SIZE = 256;
+#ifdef TEST_ON_RDNA4 // RDNA4, WAVE_SIZE = 32
+constexpr int WARP_M = 4, WARP_N = 2;
+#else // CDNA3, WAVE_SIZE = 64
+constexpr int WARP_M = 2, WARP_N = 2;
+#endif
+#endif // End of parameterization
+__global__ void gemm_kernel(
+    const in_data_type a[K][M],
+    const in_data_type b[K][N],
+    out_data_type c[M][N],
+    const float sa[ceil_div(K, QUANT_SIZE)][M / 1        ], // Assuming M is divisible by 1 (always true)
+    const float sb[ceil_div(K, QUANT_SIZE)][ceil_div(N, QUANT_SIZE)]
+) {
+    // --- Start: Derived parameters and constants ---
+    constexpr int WMMA_M = 16; // Fixed WMMA dimension M
+    constexpr int WMMA_N = 16; // Fixed WMMA dimension N
+    constexpr int WMMA_K = 32; // Fixed WMMA dimension K (for FP8)
+
+    // WARP_M/N define the 2D arrangement of warps in the block grid.
+    // These might need adjustment based on BLOCK_DIM_X/Y strategy.
+    // Using fixed values based on the non-parameterized version for now.
+    // TODO: Derive WARP_M/N from BLOCK_DIM_X/Y if a flexible strategy is needed.
+    constexpr int WARP_NUM = WARP_M * WARP_N; // Total warps per block
+
+    // Assertion: Check if the assumed warp layout matches the block size
+    static_assert(WARP_NUM * WAVE_SIZE == BLOCK_SIZE, "WARP_M * WARP_N * WAVE_SIZE must equal BLOCK_SIZE");
+
+    // Fragments per warp
+    constexpr int FRAG_M_PER_WARP = BM / WMMA_M / WARP_M;
+    constexpr int FRAG_N_PER_WARP = BN / WMMA_N / WARP_N;
+    constexpr int FRAG_K = BK / WMMA_K; // Fragments along K dimension tile
+
+    static_assert(BM % (WMMA_M * WARP_M) == 0, "BM must be divisible by WMMA_M * WARP_M");
+    static_assert(BN % (WMMA_N * WARP_N) == 0, "BN must be divisible by WMMA_N * WARP_N");
+    static_assert(BK % WMMA_K == 0, "BK must be divisible by WMMA_K");
+    static_assert(BK >= 32, "BK must be at least 32");
+    // --- End: Derived parameters and constants ---
+
+    constexpr int QM = M; // Dimension M for scale A
+    constexpr int QN = ceil_div(N, QUANT_SIZE); // Dimension N for scale B (quantized)
+    constexpr int QK = ceil_div(K, QUANT_SIZE); // Dimension K for scales (quantized)
+    constexpr int PM = BM; // Block size M for scale A * B
+    constexpr int PN = ceil_div(BN, QUANT_SIZE); // Block size N for scale A * B
+
+    // Ensure derived fragment counts are positive
+    static_assert(FRAG_M_PER_WARP > 0, "FRAG_M_PER_WARP must be positive");
+    static_assert(FRAG_N_PER_WARP > 0, "FRAG_N_PER_WARP must be positive");
+    static_assert(FRAG_K > 0, "FRAG_K must be positive");
+
+    using FragA = wmma::fragment<wmma::matrix_a, WMMA_M, WMMA_N, WMMA_K, in_data_type, wmma::col_major>;
+    using FragB = wmma::fragment<wmma::matrix_b, WMMA_M, WMMA_N, WMMA_K, in_data_type, wmma::row_major>;
+    using FragC = wmma::fragment<wmma::accumulator, WMMA_M, WMMA_N, WMMA_K, acc_data_type>;
+    using FragOut = wmma::fragment<wmma::accumulator, WMMA_M, WMMA_N, WMMA_K, half>; // Output uses half for storage via bfloat16 reinterpret
+
+    __shared__ in_data_type  s_a[BK][BM];
+    __shared__ in_data_type  s_b[BK][BN];
+    __shared__ acc_data_type s_s[PM][PN]; // Accumulator type for scales
+    FragC frag_r[FRAG_M_PER_WARP][FRAG_N_PER_WARP]; // Accumulator fragments
+
+    // handle splitk
+    a += blockIdx.z * K;
+    b += blockIdx.z * K;
+    c += blockIdx.z * M;
+    sa += blockIdx.z * QK;
+    sb += blockIdx.z * QK;
+
+    int tid = threadIdx.x; // Linear thread ID within the block
+    int wid = tid / WAVE_SIZE; // Warp ID within the block
+
+    // Initialize the output accumulator fragments to zero
+    #pragma unroll
+    for (int i = 0; i < FRAG_M_PER_WARP; i++) {
+        #pragma unroll
+        for (int j = 0; j < FRAG_N_PER_WARP; j++) {
+            wmma::fill_fragment(frag_r[i][j], 0.0f); // Use float literal
+        }
+    }
+
+    // Spilt and compute fragments
+    constexpr int iteration_over_k = ceil_div(K, BK); // Use ceil_div for potentially non-divisible K
+    constexpr int LOAD_BATCH_SIZE = (2 * sizeof(float4) / sizeof(in_data_type)) > 0 ? (2 * sizeof(float4) / sizeof(in_data_type)) : 1; // Ensure batch size > 0
+    static_assert(LOAD_BATCH_SIZE > 0, "LOAD_BATCH_SIZE must be positive");
+
+    for (int bk = 0; bk < iteration_over_k; bk++) {
+        const int m = blockIdx.y * BM;
+        const int n = blockIdx.x * BN;
+        const int k = bk * BK;
+
+        // Calculate remaining K for boundary checks if needed (not currently used by load_input)
+        // const int k_rem = K - k;
+
+        // Load data into shared memory
+        load_input<in_data_type, BK, BM, K, M, BLOCK_SIZE, LOAD_BATCH_SIZE>(
+            s_a, a, m, k);
+        load_input<in_data_type, BK, BN, K, N, BLOCK_SIZE, LOAD_BATCH_SIZE>(
+            s_b, b, n, k);
+        // Load scales into shared memory (using acc_data_type for s_s)
+        load_scale<PM, PN, QM, QN, QK, QUANT_SIZE, BLOCK_SIZE, 1>(
+            s_s, sa, sb, m, n, k);
+        __syncthreads();
+
+        // Perform matrix multiplication using WMMA
+        wmma_compute<in_data_type, acc_data_type, FragC, FragA, FragB,
+            PM, PN, BM, BN, BK, FRAG_M_PER_WARP, FRAG_N_PER_WARP, FRAG_K,
+            WMMA_M, WMMA_N, WMMA_K,
+            WARP_M, WARP_N,
+            BLOCK_SIZE, LOAD_BATCH_SIZE, QUANT_SIZE>( // Pass calculated BLOCK_SIZE and LOAD_BATCH_SIZE
+                s_a, s_b, s_s, frag_r, wid / WARP_N, wid % WARP_N);
+        __syncthreads();
+    }
+    // Store results from accumulator fragments to global memory
+    store_result<acc_data_type, out_data_type, FragC, FragOut,
+        WMMA_M, WMMA_N, BM, BN, M, N, FRAG_M_PER_WARP, FRAG_N_PER_WARP>(
+            c, frag_r, blockIdx.y * BM, blockIdx.x * BN,
+            wid / WARP_N, wid % WARP_N);
+
+
+};
+
+
+}; // namespace gemm_kernel_legacy

gemm/gemm_launcher.hip

@@ -0,0 +1,267 @@
+// Wrapped of gemm kernel launcher.
+#include <unistd.h>
+#include <chrono>
+#define PARAMETERIZE_LIBRARY
+#include "gemm_kernel.h"
+#include "gemm_kernel_legacy.h"
+#include "transpose_kernel.h"
+#undef PARAMETERIZE_LIBRARY
+#include "../include/gpu_types.h"
+#include "../include/timer.h"
+#include "../tests/checker/metrics.h"
+#include <iostream>
+
+#include <stdio.h>
+
+HOST_CODE_BELOW
+
+std::vector<std::shared_ptr<KernelTimer>> timers;
+
+using namespace std;
+
+float *c_splitk = nullptr;
+__FP8_TYPE *a_trans = nullptr;
+__FP8_TYPE *b_trans = nullptr;
+constexpr int MAX_MATRIX_M = 6144;
+constexpr int MAX_MATRIX_N = 7168;
+constexpr int MAX_MATRIX_K = 7168;
+constexpr int MAX_SPLITK_FACTOR = 8;
+
+void init_workspace() {
+    LIB_CALL(HOST_TYPE(Malloc)(&c_splitk, MAX_MATRIX_M * MAX_MATRIX_N * sizeof(float) * MAX_SPLITK_FACTOR));
+    LIB_CALL(HOST_TYPE(Malloc)(&a_trans, MAX_MATRIX_M * MAX_MATRIX_K * sizeof(__FP8_TYPE)));
+    LIB_CALL(HOST_TYPE(Malloc)(&b_trans, MAX_MATRIX_N * MAX_MATRIX_K * sizeof(__FP8_TYPE)));
+    // LIB_CALL(HOST_TYPE(StreamCreateWithFlags)(&job_stream0, HOST_TYPE(StreamNonBlocking)));
+    // job_stream0 = 0;
+}
+
+
+// Launch pipeline gemm kernels (most performant).
+// 1. Transpose input A & B.
+// 2. GEMM compute.
+// 3. Reduce (if spilt-k is enable)
+template <int M, int N, int K, int BM, int BN, int BK, int WARP_M, int WARP_N, int BLOCK_SIZE, int QUANT_BLOCK_SIZE,
+          int SPLITK_FACTOR, int LOAD_BATCH_SIZE = 16>
+void launch_gemm(const __FP8_TYPE *a, const __FP8_TYPE *b, __BF16_TYPE *c, const float *as, const float *bs, HOST_TYPE(Stream_t) job_stream0) {
+    static_assert(M <= MAX_MATRIX_M, "M exceeds maximum supported size");
+    static_assert(N <= MAX_MATRIX_N, "N exceeds maximum supported size");
+    static_assert(K <= MAX_MATRIX_K, "K exceeds maximum supported size");
+    static_assert(SPLITK_FACTOR <= MAX_SPLITK_FACTOR, "SPLITK_FACTOR exceeds maximum supported size");
+    if (__builtin_expect(c_splitk == nullptr, 0)) {
+        init_workspace();
+        LIB_CALL(hipDeviceSynchronize());
+    }
+    
+    transpose_kernel::transpose_fp8<K, N>(b_trans, b, job_stream0);
+    transpose_kernel::transpose_fp8<K, M>(a_trans, a, job_stream0);
+    // transpose_kernel::launch_transpose<__FP8_TYPE, K, N, 64, 512, 4>(b_trans, b, job_stream0);
+    // transpose_kernel::launch_transpose<__FP8_TYPE, K, M, 64, 512, 4>(a_trans, a, job_stream0);
+    // Busy wait for 150 microseconds
+    // auto start = std::chrono::high_resolution_clock::now();
+    // while (std::chrono::duration_cast<std::chrono::microseconds>(
+    //     std::chrono::high_resolution_clock::now() - start).count() < 150) {
+    //     // Busy wait
+    // }
+    // be careful that blocksize < 1024, or there's a silent fault
+    // gemm_kernel::check_trans<<<dim3(K / 32, M / 32), dim3(32, 32)>>>(a, a_trans, K, M);
+
+    static_assert(K % SPLITK_FACTOR == 0, "K not divisible by SPLITK_FACTOR");
+    dim3 grid(ceil_div(N, BN) << 1, ceil_div(M, BM) >> 1, SPLITK_FACTOR);
+    static_assert(BLOCK_SIZE >= 32, "BLOCK_SIZE must be at least 32");
+    dim3 block(BLOCK_SIZE);
+    if constexpr (SPLITK_FACTOR == 1) {
+        hipLaunchKernelGGL(
+            HIP_KERNEL_NAME(gemm_kernel::gemm_kernel<__FP8_TYPE, float, __BF16_TYPE, M, N, K, BM, BN, BK, QUANT_BLOCK_SIZE, BLOCK_SIZE, WARP_M, WARP_N, K, K, LOAD_BATCH_SIZE>),
+            grid, block, 0, job_stream0, 
+            reinterpret_cast<const __FP8_TYPE(*)[K]>(a_trans), 
+            reinterpret_cast<const __FP8_TYPE(*)[K]>(b_trans),
+            reinterpret_cast<__BF16_TYPE(*)[N]>(c), reinterpret_cast<const float(*)[M]>(as),
+            reinterpret_cast<const float(*)[ceil_div(N, QUANT_BLOCK_SIZE)]>(bs)
+        );
+    } else {
+        hipLaunchKernelGGL(
+            HIP_KERNEL_NAME(gemm_kernel::gemm_kernel<__FP8_TYPE, float, float, M, N, K / SPLITK_FACTOR, BM, BN, BK, QUANT_BLOCK_SIZE, BLOCK_SIZE, WARP_M, WARP_N, K, K, LOAD_BATCH_SIZE>),
+            grid, block, 0, job_stream0,
+            reinterpret_cast<const __FP8_TYPE(*)[K]>(a_trans), 
+            reinterpret_cast<const __FP8_TYPE(*)[K]>(b_trans),
+            reinterpret_cast<float(*)[N]>(c_splitk), reinterpret_cast<const float(*)[M]>(as),
+            reinterpret_cast<const float(*)[ceil_div(N, QUANT_BLOCK_SIZE)]>(bs));
+        constexpr uint32_t REDUCE_BLOCK = 256;
+        hipLaunchKernelGGL(
+            HIP_KERNEL_NAME(gemm_kernel::reduce_kernel<M, N, SPLITK_FACTOR, REDUCE_BLOCK>),
+            ceil_div(M * N / 4, REDUCE_BLOCK), REDUCE_BLOCK, 0, job_stream0,
+            reinterpret_cast<const float(*)[M][N]>(c_splitk), 
+            reinterpret_cast<__BF16_TYPE(*)[N]>(c)
+        );    }
+    auto err = HOST_TYPE(GetLastError)();
+    if (err != HOST_TYPE(Success)) {
+        std::cerr << "Kernel execution failed.\n" << HOST_TYPE(GetErrorString)(err) << std::endl;
+        abort();
+    }
+}
+
+
+// Launch legacy gemm kernel. (most compellable)
+template <int M, int N, int K, int BM, int BN, int BK, int WARP_M, int WARP_N, int BLOCK_SIZE, int QUANT_BLOCK_SIZE, int SPLITK_FACTOR>
+void launch_gemm_legacy(const __FP8_TYPE *a, const __FP8_TYPE *b, __BF16_TYPE *c, const float *as, const float *bs, HOST_TYPE(Stream_t) job_stream0) {
+    static_assert(K % SPLITK_FACTOR == 0, "K not divisible by SPLITK_FACTOR");
+    dim3 grid(ceil_div(N, BN), ceil_div(M, BM), SPLITK_FACTOR);
+    static_assert(BLOCK_SIZE >= 32, "BLOCK_SIZE must be at least 32");
+    dim3 block(BLOCK_SIZE);
+    if (__builtin_expect(c_splitk == nullptr, 0)) {
+        init_workspace();
+        LIB_CALL(hipDeviceSynchronize());
+    }
+
+    if constexpr (SPLITK_FACTOR == 1) {
+        hipLaunchKernelGGL(
+            HIP_KERNEL_NAME(gemm_kernel_legacy::gemm_kernel<__FP8_TYPE, float, __BF16_TYPE, M, N, K, BM, BN, BK, QUANT_BLOCK_SIZE, BLOCK_SIZE, WARP_M, WARP_N>),
+            grid, block, 0, job_stream0,
+            reinterpret_cast<const __FP8_TYPE (*)[M]>(a),
+            reinterpret_cast<const __FP8_TYPE (*)[N]>(b),
+            reinterpret_cast<__BF16_TYPE (*)[N]>(c),
+            reinterpret_cast<const float (*)[M]>(as),
+            reinterpret_cast<const float (*)[ceil_div(N, QUANT_BLOCK_SIZE)]>(bs)
+        );
+    } else {
+        hipLaunchKernelGGL(
+            HIP_KERNEL_NAME(gemm_kernel_legacy::gemm_kernel<__FP8_TYPE, float, float, M, N, K / SPLITK_FACTOR, BM, BN, BK, QUANT_BLOCK_SIZE, BLOCK_SIZE, WARP_M, WARP_N>),
+            grid, block, 0, job_stream0,
+            reinterpret_cast<const __FP8_TYPE (*)[M]>(a),
+            reinterpret_cast<const __FP8_TYPE (*)[N]>(b),
+            reinterpret_cast<float (*)[N]>(c_splitk),
+            reinterpret_cast<const float (*)[M]>(as),
+            reinterpret_cast<const float (*)[ceil_div(N, QUANT_BLOCK_SIZE)]>(bs)
+        );
+        constexpr uint32_t REDUCE_BLOCK = 256;
+        hipLaunchKernelGGL(
+            HIP_KERNEL_NAME(gemm_kernel_legacy::reduce<0>),
+            ceil_div(M * N, REDUCE_BLOCK), REDUCE_BLOCK, 0, job_stream0,
+            M, N, SPLITK_FACTOR, c_splitk, (__BF16_TYPE *)c
+        );
+    }
+    auto err = HOST_TYPE(GetLastError)();
+    if (err != HOST_TYPE(Success)) {
+        std::cerr << "Kernel execution failed.\n" << HOST_TYPE(GetErrorString)(err) << std::endl;
+        abort();
+    }
+}
+
+constexpr inline uint32_t pack_shape(uint32_t m, uint32_t n, uint32_t k) {
+    // Pack m, n, k into a 32-bit integer
+    // Use 8 bits for each dimension (supports 32-aligned values from 32 to 8192)
+    // Divide each value by 32 to fit into 8 bits
+    return ((m / 32) << 16) | ((n / 32) << 8) | (k / 32);
+}
+// int M, int N, int K, int BM, int BN, int BK, int WARP_M, int WARP_N, int BLOCK_SIZE, int QUANT_BLOCK_SIZE, int
+// SPLITK_FACTOR, int LOAD_BATCH_SIZE
+#define DISPATCH_GEMM(M, N, K, BM, BN, BK, WARP_M, WARP_N, BLOCK_SIZE, SPLITK_FACTOR, LOAD_BATCH_SIZE)                                  \
+    case pack_shape_checked<M, N, K>(): {                                                                              \
+        launch_gemm<M, N, K, BM, BN, BK, WARP_M, WARP_N, BLOCK_SIZE, 128, SPLITK_FACTOR, LOAD_BATCH_SIZE>(a_ptr, b_ptr, c_ptr, as_ptr, bs_ptr, job_stream0);          \
+        break;                                                                                                         \
+    }
+
+#define DISPATCH_GEMM_LEGACY(M, N, K, BM, BN, BK, WARP_M, WARP_N, BLOCK_SIZE, SPLITK_FACTOR)                                                 \
+    case pack_shape_checked<M, N, K>(): {                                                                                \
+        launch_gemm_legacy<M, N, K, BM, BN, BK, WARP_M, WARP_N, BLOCK_SIZE, 128, SPLITK_FACTOR>(a_ptr, b_ptr, c_ptr, as_ptr, bs_ptr, job_stream0); \
+        break;                                                                                                              \
+    }
+
+template <int M, int N, int K> constexpr inline uint32_t pack_shape_checked() {
+    static_assert(M % 32 == 0, "M must be a multiple of 32");
+    static_assert(N % 32 == 0, "N must be a multiple of 32");
+    static_assert(K % 32 == 0, "K must be a multiple of 32");
+    static_assert(M >= 32 && M <= 8192, "M must be between 32 and 8192");
+    static_assert(N >= 32 && N <= 8192, "N must be between 32 and 8192");
+    static_assert(K >= 32 && K <= 8192, "K must be between 32 and 8192");
+    return pack_shape(M, N, K);
+}
+
+
+
+extern "C" {
+// Basically, it dispatch GEMM to fatest implementations according to inputs' shape. 
+void run(void *a, void *b, void *as, void *bs, void *c, int m, int n, int k, PerfMetrics *metrics, hipStream_t job_stream0) {
+    // Cast pointers once
+    const __FP8_TYPE *a_ptr = static_cast<const __FP8_TYPE *>(a);
+    const __FP8_TYPE *b_ptr = static_cast<const __FP8_TYPE *>(b);
+    __BF16_TYPE *c_ptr = static_cast<__BF16_TYPE *>(c);
+    const float *as_ptr = static_cast<const float *>(as);
+    const float *bs_ptr = static_cast<const float *>(bs);
+    KernelTimerScoped timer(timers, 2LL * m * n * k,
+        metrics ? &metrics->entries[0].time : nullptr,
+        metrics ? &metrics->entries[0].gflops : nullptr, job_stream0);
+
+    switch (pack_shape(m, n, k)) {
+#ifdef TEST_ON_RDNA4 // RDNA4, WAVE_SIZE = 32
+        // Test:      M,      N,      K,      BM,    BN,    BK,   WARP_M, WARP_N, BLOCK_SIZE,   SPLITK_FACTOR, LOAD_BATCH_SIZE
+        DISPATCH_GEMM(64, 64, 128, 64, 64, 32, 1, 4, 128, 1, 16);
+        DISPATCH_GEMM(64, 1536, 7168, 64, 128, 64, 4, 2, 256, 1, 16);
+        DISPATCH_GEMM(64, 3072, 1536, 64, 128, 64, 4, 2, 256, 1, 16);
+        DISPATCH_GEMM(64, 576, 7168, 64, 128, 64, 4, 2, 256, 1, 16);
+        DISPATCH_GEMM(96, 7168, 256, 96, 256, 64, 2, 4, 256, 1, 16);
+        DISPATCH_GEMM(96, 7168, 2048, 96, 256, 64, 2, 4, 256, 1, 16);
+        DISPATCH_GEMM(96, 4608, 7168, 96, 256, 64, 2, 4, 256, 1, 16);
+        DISPATCH_GEMM(128, 7168, 2304, 128, 128, 64, 4, 2, 256, 1, 16);
+        DISPATCH_GEMM(128, 512, 7168, 128, 128, 64, 4, 2, 256, 1, 16);
+        DISPATCH_GEMM(512, 4096, 512, 256, 128, 64, 4, 2, 256, 1, 16);
+        DISPATCH_GEMM(512, 1536, 7168, 256, 128, 64, 4, 2, 256, 1, 16);
+        // Benchmark: M,      N,      K,      BM,    BN,    BK,   WARP_M, WARP_N, BLOCK_SIZE,   SPLITK_FACTOR, LOAD_BATCH_SIZE
+        DISPATCH_GEMM(1024, 1536, 7168, 128, 128, 64, 1, 4, 128, 4, 16); // Optimized: 0.49 ms (45.65 TFlops)
+        DISPATCH_GEMM(1024, 3072, 1536, 256, 128, 32, 4, 2, 256, 1, 16); // Optimized: 0.19 ms (51.32 TFlops)
+        DISPATCH_GEMM(1024, 576, 7168, 128, 64, 32, 4, 1, 128, 4, 16);   // Optimized: 0.30 ms (28.16 TFlops)
+        DISPATCH_GEMM(1024, 7168, 256, 256, 128, 32, 4, 2, 256, 1, 16);  // Optimized: 0.08 ms (46.49 TFlops)
+        DISPATCH_GEMM(1024, 7168, 2048, 256, 128, 32, 4, 2, 256, 1, 16); // Optimized: 0.49 ms (61.92 TFlops)
+        DISPATCH_GEMM(1024, 4608, 7168, 128, 128, 32, 2, 2, 128, 1, 16); // Optimized: 0.99 ms (68.16 TFlops)
+        DISPATCH_GEMM(1024, 7168, 2304, 256, 128, 32, 4, 2, 256, 1, 16); // Optimized: 0.51 ms (66.04 TFlops)
+        DISPATCH_GEMM(1024, 512, 7168, 64, 128, 32, 2, 2, 128, 4, 16);   // Optimized: 0.26 ms (28.97 TFlops)
+        DISPATCH_GEMM(1024, 4096, 512, 128, 256, 32, 2, 4, 256, 1, 16);  // Optimized: 0.08 ms (54.27 TFlops)
+        DISPATCH_GEMM(6144, 1536, 7168, 256, 128, 32, 4, 2, 256, 1, 16); // Optimized: 1.76 ms (76.76 TFlops)
+        DISPATCH_GEMM(6144, 3072, 1536, 256, 128, 32, 4, 2, 256, 1, 16); // Optimized: 0.88 ms (66.00 TFlops)
+        DISPATCH_GEMM(6144, 576, 7168, 256, 128, 32, 4, 2, 256, 1, 16);  // Optimized: 0.84 ms (60.68 TFlops)
+        DISPATCH_GEMM(6144, 7168, 256, 256, 128, 32, 4, 2, 256, 1, 16);  // Optimized: 0.49 ms (45.76 TFlops)
+        DISPATCH_GEMM(6144, 7168, 2048, 256, 128, 32, 4, 2, 256, 1, 16); // Optimized: 2.17 ms (83.11 TFlops)
+        DISPATCH_GEMM(6144, 4608, 7168, 256, 128, 32, 4, 2, 256, 1, 16); // Optimized: 4.56 ms (88.99 TFlops)
+        DISPATCH_GEMM(6144, 7168, 2304, 256, 128, 32, 4, 2, 256, 1, 16); // Optimized: 2.41 ms (84.32 TFlops)
+        DISPATCH_GEMM(6144, 512, 7168, 256, 128, 32, 4, 2, 256, 1, 16);  // Optimized: 0.67 ms (67.45 TFlops)
+        DISPATCH_GEMM(6144, 4096, 512, 256, 128, 32, 4, 2, 256, 1, 16);  // Optimized: 0.51 ms (50.79 TFlops)
+#else                                                                // CDNA3, WAVE_SIZE = 64
+      // Benchmark:   M,    N,      K,      BM,    BN,    BK,   WARP_M, WARP_N, BLOCK_SZ, SPLITK_F, LOAD_BS
+        DISPATCH_GEMM(1024, 1536,   7168,   256,   128,   128,   4,      2,      512,      4,        16); // #0
+        DISPATCH_GEMM(1024, 3072,   1536,   256,   128,   128,   4,      2,      512,      2,        16); // #1
+        DISPATCH_GEMM(1024, 576,    7168,   256,   128,   128,   4,      2,      512,      8,        16); // #2
+        DISPATCH_GEMM(1024, 7168,   256,    256,   128,   128,   4,      2,      512,      1,        16); // #3
+        DISPATCH_GEMM(1024, 7168,   2048,   256,   128,   128,   4,      2,      512,      1,        16); // #4
+        DISPATCH_GEMM(1024, 4608,   7168,   256,   128,   128,   4,      2,      512,      2,        16); // #5
+        DISPATCH_GEMM(1024, 7168,   2304,   256,   128,   128,   4,      2,      512,      1,        16); // #6
+        DISPATCH_GEMM(1024, 512,    7168,   256,   128,   128,   4,      2,      512,      8,        16); // #7
+        DISPATCH_GEMM(1024, 4096,   512,    256,   128,   128,   4,      2,      512,      1,        16); // #8
+        DISPATCH_GEMM(6144, 1536,   7168,   256,   128,   128,   4,      2,      512,      1,        16); // #9
+        DISPATCH_GEMM(6144, 3072,   1536,   256,   128,   128,   4,      2,      512,      1,        16); // #10
+        DISPATCH_GEMM(6144, 576,    7168,   256,   128,   128,   4,      2,      512,      2,        16); // #11
+        DISPATCH_GEMM(6144, 7168,   256,    256,   128,   128,   4,      2,      512,      1,        16); // #12
+        DISPATCH_GEMM(6144, 7168,   2048,   256,   128,   128,   4,      2,      512,      1,        16); // #13
+        DISPATCH_GEMM(6144, 4608,   7168,   256,   128,   128,   4,      2,      512,      1,        16); // #14
+        DISPATCH_GEMM(6144, 7168,   2304,   256,   128,   128,   4,      2,      512,      1,        16); // #15
+        DISPATCH_GEMM(6144, 512,    7168,   256,   128,   128,   4,      2,      512,      2,        16); // #16
+        DISPATCH_GEMM(6144, 4096,   512,    256,   128,   128,   4,      2,      512,      1,        16); // #17
+    // Test:                 M,    N,      K,      BM,    BN,    BK,   WARP_M, WARP_N, BLOCK_SZ,    SPLITK_F, 
+        DISPATCH_GEMM_LEGACY(64,   64,     128,    64,    64,    32,   4,      2,      512,          1);
+        DISPATCH_GEMM_LEGACY(64,   1536,   7168,   64,    128,   64,   4,      2,      512,          1);
+        DISPATCH_GEMM_LEGACY(64,   3072,   1536,   64,    128,   64,   4,      2,      512,          1);
+        DISPATCH_GEMM_LEGACY(64,   576,    7168,   64,    128,   64,   4,      2,      512,          1);
+        DISPATCH_GEMM_LEGACY(96,   7168,   256,    96,    256,   64,   2,      4,      512,          1);
+        DISPATCH_GEMM_LEGACY(96,   7168,   2048,   96,    256,   64,   2,      4,      512,          1);
+        DISPATCH_GEMM_LEGACY(96,   4608,   7168,   96,    256,   64,   2,      4,      512,          1);
+        DISPATCH_GEMM_LEGACY(128,  7168,   2304,   128,   128,   64,   4,      2,      512,          1);
+        DISPATCH_GEMM_LEGACY(128,  512,    7168,   128,   128,   64,   4,      2,      512,          1);
+        DISPATCH_GEMM_LEGACY(512,  4096,   512,    256,   128,   64,   4,      2,      512,          1);
+        DISPATCH_GEMM_LEGACY(512,  1536,   7168,   256,   128,   64,   4,      2,      512,          1);
+#endif
+    default: {
+        printf("Error: Unsupported shape M=%d, K=%d, N=%d\n", m, k, n);
+        abort();
+    }
+    }
+}
+} // extern "C"

gemm/transpose_kernel.h

@@ -0,0 +1,120 @@
+// Implementation of transpose kernel.
+#pragma once
+
+#include <hip/amd_detail/amd_hip_runtime.h>
+#include <hip/amd_detail/amd_warp_functions.h>
+#include "../include/gpu_libs.h"
+#include "../include/gpu_types.h"
+#include "../src/utils/arithmetic.h"
+#include "../include/clangd_workaround.h"
+
+DEVICE_CODE_BELOW
+
+namespace transpose_kernel {
+
+
+
+template <typename Elem, int M, int N, int TILE_DIM, int BLOCK_SIZE, int VEC_SIZE>
+__launch_bounds__(BLOCK_SIZE)
+__global__ void transpose_kernel(Elem *odata, const Elem *idata) {
+    constexpr auto TBLOCK_X = TILE_DIM / VEC_SIZE;
+    constexpr auto TBLOCK_Y = BLOCK_SIZE / TBLOCK_X;
+
+    // avoid read bank conflict
+    // VEC_SIZE * (TILE_DIM + d) * sizeof(Elem) = TBLOCK_Y / (BLOCK_SIZE / WARP_SIZE) * sizeof(Elem) + 128k
+    // each warp read row = TILE_DIM (in VEC_SIZE reads), col = TBLOCK_Y / (BLOCK_SIZE / WARP_SIZE)
+    // warp 0                     warp 1
+    // t0    t16    t32    t48    ...
+    // ...
+    // t1
+    // ...
+    // t15
+    // don't know why padding to d as described above is not working, maybe gpu could merge contigious ds_read_u8 and
+    // cause padding to be TBLOCK_Y / (BLOCK_SIZE / WARP_SIZE)
+    constexpr auto PADDING = TBLOCK_Y / (BLOCK_SIZE / warpSize);
+    __shared__ Elem tile[TILE_DIM][TILE_DIM + PADDING];
+
+    int x = blockIdx.x * TILE_DIM + threadIdx.x * VEC_SIZE;
+    int y = blockIdx.y * TILE_DIM + threadIdx.y;
+
+// Load tile
+#pragma unroll
+    for (int i = 0; i < TILE_DIM; i += TBLOCK_Y) {
+#pragma unroll
+        for (int v = 0; v < VEC_SIZE; v++) {
+            tile[threadIdx.y + i][threadIdx.x * VEC_SIZE + v] = idata[(y + i) * N + x + v];
+        }
+    }
+
+    __syncthreads();
+
+    // Transpose indices
+    x = blockIdx.y * TILE_DIM + threadIdx.x * VEC_SIZE;
+    y = blockIdx.x * TILE_DIM + threadIdx.y;
+
+// Write tile
+#pragma unroll
+    for (int i = 0; i < TILE_DIM; i += TBLOCK_Y) {
+#pragma unroll
+        for (int v = 0; v < VEC_SIZE; v++) {
+            odata[(y + i) * M + x + v] = tile[threadIdx.x * VEC_SIZE + v][threadIdx.y + i];
+        }
+    }
+}
+
+template <typename Elem, int M, int N, int TILE_DIM, int BLOCK_SIZE, int VEC_SIZE>
+void launch_transpose(Elem *out, const Elem *in, hipStream_t stream = 0) {
+    static_assert(TILE_DIM % VEC_SIZE == 0);
+    constexpr auto TBLOCK_X = TILE_DIM / VEC_SIZE;
+    static_assert(BLOCK_SIZE % TBLOCK_X == 0);
+    constexpr auto TBLOCK_Y = BLOCK_SIZE / TBLOCK_X;
+    static_assert(M % TILE_DIM == 0 && N % TILE_DIM == 0);
+    hipLaunchKernelGGL(
+        HIP_KERNEL_NAME(transpose_kernel<Elem, M, N, TILE_DIM, BLOCK_SIZE, VEC_SIZE>),
+        dim3(N / TILE_DIM, M / TILE_DIM), dim3(TBLOCK_X, TBLOCK_Y), 0, stream,
+        out, in);
+}
+
+#define DISPATCH_TRANSPOSE(DIM_0, DIM_1, TILE_DIM, BLOCK_SIZE, VEC_SIZE) else if constexpr(IN_DIM_0 == DIM_0 && IN_DIM_1 == DIM_1) launch_transpose<__FP8_TYPE, IN_DIM_0, IN_DIM_1, TILE_DIM, BLOCK_SIZE, VEC_SIZE>(out, in, stream)
+
+template <int DIM0, int DIM1>
+struct unsupported_config {
+    static_assert(DIM0 == -1, "Unsupported transpose configuration - check template parameters");
+};
+
+// Selecte best parameters for tranpose kernel.
+template <int IN_DIM_0, int IN_DIM_1>
+void transpose_fp8(__FP8_TYPE *out, const __FP8_TYPE *in, hipStream_t stream = 0) {
+    if constexpr (false /* dummy*/ ) static_assert(true);
+    DISPATCH_TRANSPOSE(   256,   1024,     64,    256, 4); // Optimized: 2.71 µs (193.46 GB/s)
+    DISPATCH_TRANSPOSE(   256,   6144,     64,    256, 4); // Optimized: 2.72 µs (1157.37 GB/s)
+    DISPATCH_TRANSPOSE(   256,   7168,     64,    256, 8); // Optimized: 2.99 µs (1225.38 GB/s)
+    DISPATCH_TRANSPOSE(   512,   1024,     64,    512, 4); // Optimized: 2.55 µs (411.21 GB/s)
+    DISPATCH_TRANSPOSE(   512,   4096,     64,    256, 4); // Optimized: 3.01 µs (1394.85 GB/s)
+    DISPATCH_TRANSPOSE(   512,   6144,     64,    512, 4); // Optimized: 3.58 µs (1755.43 GB/s)
+    DISPATCH_TRANSPOSE(  1536,   1024,     64,   1024, 4); // Optimized: 2.78 µs (1130.74 GB/s)
+    DISPATCH_TRANSPOSE(  1536,   3072,     64,    512, 4); // Optimized: 3.57 µs (2641.99 GB/s)
+    DISPATCH_TRANSPOSE(  1536,   6144,    128,   1024, 8); // Optimized: 7.09 µs (2661.36 GB/s)
+    DISPATCH_TRANSPOSE(  2048,   1024,     64,   1024, 4); // Optimized: 2.84 µs (1477.91 GB/s)
+    DISPATCH_TRANSPOSE(  2048,   6144,    128,    512, 8); // Optimized: 8.94 µs (2816.23 GB/s)
+    DISPATCH_TRANSPOSE(  2048,   7168,    128,    512, 8); // Optimized: 9.56 µs (3070.50 GB/s)
+    DISPATCH_TRANSPOSE(  2304,   1024,     64,   1024, 4); // Optimized: 3.08 µs (1532.51 GB/s)
+    DISPATCH_TRANSPOSE(  2304,   6144,    128,    512, 8); // Optimized: 9.30 µs (3043.93 GB/s)
+    DISPATCH_TRANSPOSE(  2304,   7168,    128,    512, 8); // Optimized: 10.39 µs (3179.95 GB/s)
+    DISPATCH_TRANSPOSE(  7168,    512,     64,    512, 4); // Optimized: 3.25 µs (2257.78 GB/s)
+    DISPATCH_TRANSPOSE(  7168,    576,     64,    512, 4); // Optimized: 3.44 µs (2403.24 GB/s)
+    DISPATCH_TRANSPOSE(  7168,   1024,     64,    256, 4); // Optimized: 5.07 µs (2892.62 GB/s)
+    DISPATCH_TRANSPOSE(  7168,   1536,    128,   1024, 8); // Optimized: 7.72 µs (2851.97 GB/s)
+    DISPATCH_TRANSPOSE(  7168,   4608,    128,    512, 8); // Optimized: 16.87 µs (3915.84 GB/s)
+    DISPATCH_TRANSPOSE(  7168,   6144,    128,    256, 8); // Optimized: 21.59 µs (4079.12 GB/s)
+    else static_assert(false);
+}
+
+} // namespace transpose_kernel
+
+
+
+
+#ifndef PARAMETERIZE_LIBRARY
+int main() {}
+#endif

include/clangd_workaround.h

@@ -0,0 +1,9 @@
+#define HOST_CODE_BELOW                                                                                                \
+  extern "C" int printf(const char *fmt, ...);                                                                         \
+  extern void operator delete[](void *ptr) _GLIBCXX_USE_NOEXCEPT;                                                      \
+  extern void *operator new[](__SIZE_TYPE__ size);
+
+#define DEVICE_CODE_BELOW                                                                                                \
+  extern "C" __device__ int printf(const char *fmt, ...);                                                                         \
+  extern __device__ void operator delete[](void *ptr) _GLIBCXX_USE_NOEXCEPT;                                                      \
+  extern __device__ void *operator new[](__SIZE_TYPE__ size);

include/gpu_libs.h

@@ -0,0 +1,42 @@
+#ifdef TEST_ON_CUDA
+#include <mma.h>
+
+#include <cuda_fp16.h>
+#include <cuda_fp8.h>
+
+namespace wmma = nvcuda::wmma;
+
+#define LIB_CALL(call)                                                                                                 \
+  do {                                                                                                                 \
+    cudaError_t err = call;                                                                                            \
+    if (err != cudaSuccess) {                                                                                          \
+      abort();                                                                                                         \
+    }                                                                                                                  \
+  } while (0)
+
+#define HOST_TYPE(x) cuda##x
+
+#else
+
+#ifndef HIP_HEADERS__
+#include <hip/hip_runtime.h>
+#include <hip/hip_fp8.h>
+#include <hip/hip_fp16.h>
+#include <rocwmma/rocwmma.hpp>
+#define HIP_HEADERS__
+#endif
+
+namespace wmma = rocwmma;
+
+#define LIB_CALL(call)                                                                                                 \
+  do {                                                                                                                 \
+    hipError_t err = call;                                                                                             \
+    if (err != hipSuccess) {                                                                                           \
+      abort();                                                                                                         \
+    }                                                                                                                  \
+  } while (0)
+
+#define HOST_TYPE(x) hip##x
+
+#endif
+

include/gpu_types.h

@@ -0,0 +1,28 @@
+#pragma once
+#ifdef TEST_ON_CUDA
+#define __FP8_TYPE __nv_fp8_e4m3
+#define __FP8x4_TYPE __nv_fp8x4_e4m3
+#define __BF16_TYPE __nv_bfloat16
+#define __BF16x2_TYPE __nv_bfloat162
+
+#else
+#ifdef TEST_ON_RDNA4
+#define __FP8_TYPE __hip_fp8_e4m3
+#define __FP8x4_TYPE __hip_fp8x4_e4m3
+constexpr const inline int WAVE_SIZE = 32;
+constexpr const inline int XCD_SWIZZLE = 1;
+#else
+#define __FP8_TYPE __hip_fp8_e4m3_fnuz
+#define __FP8x4_TYPE __hip_fp8x4_e4m3_fnuz
+constexpr const inline int WAVE_SIZE = 64;
+constexpr const inline int XCD_SWIZZLE = 8;
+#endif
+
+#define __BF16_TYPE __hip_bfloat16
+#define __BF16x2_TYPE __hip_bfloat162
+#define __FP16_TYPE __half
+#define __INT16_TYPE int16_t
+
+#endif
+
+#define __FP32_TYPE float

include/timer.h

@@ -0,0 +1,57 @@
+#pragma once
+#include "gpu_types.h"
+#include "gpu_libs.h"
+#include <memory>
+#include <vector>
+
+typedef void (*TimerCompletionCallback)(float elapsed_time, size_t calc_ops, float *time_ptr, float *gflops_ptr,
+                                        void *user_data);
+
+class KernelTimer {
+  private:
+    size_t calc_ops;
+    HOST_TYPE(Event_t) start, stop;
+    float *time_ptr;
+    float *gflops_ptr;
+    void *user_data;
+    TimerCompletionCallback callback;
+    bool callback_executed;
+
+  public:
+    KernelTimer(size_t calc_ops, float *time, float *gflops);
+
+    void start_timer(hipStream_t stream = 0);
+    void stop_timer(hipStream_t stream = 0);
+    void set_callback(TimerCompletionCallback cb, void *data = nullptr);
+
+    // Wait for the timer to complete and execute the callback if set
+    void synchronize();
+
+    // Getter methods for the callback
+    HOST_TYPE(Event_t) get_start_event() const { return start; }
+    HOST_TYPE(Event_t) get_stop_event() const { return stop; }
+    size_t get_calc_ops() const { return calc_ops; }
+    float *get_time_ptr() const { return time_ptr; }
+    float *get_gflops_ptr() const { return gflops_ptr; }
+    void execute_callback(float elapsed_time);
+    void set_callback_executed(bool executed) { callback_executed = executed; }
+    bool is_callback_executed() const { return callback_executed; }
+
+    ~KernelTimer();
+};
+
+class KernelTimerScoped {
+  private:
+    std::shared_ptr<KernelTimer> timer;
+    hipStream_t stream;
+
+  public:
+    KernelTimerScoped(std::vector<std::shared_ptr<KernelTimer>> &timers, size_t calc_ops, float *time, float *gflops,
+                      hipStream_t stream = 0)
+        : timer(std::make_shared<KernelTimer>(calc_ops, time, gflops)), stream(stream) {
+        timers.push_back(timer);
+        timer->start_timer(stream);
+    }
+
+    ~KernelTimerScoped() { timer->stop_timer(stream); }
+};

src/utils/arithmetic.h

@@ -0,0 +1,12 @@
+#ifndef ARITHMETIC_H
+#define ARITHMETIC_H
+
+template <int x, int y> constexpr __device__ __host__ inline int exact_div() {
+  static_assert(x % y == 0);
+  static_assert(x >= y);
+  return x / y;
+}
+
+constexpr __device__ __host__ inline int ceil_div(int x, int y) { return (x + y - 1) / y; }
+
+#endif

src/utils/timer.hip

@@ -0,0 +1,104 @@
+#include "../../include/timer.h"
+
+// Define HIPRT_CB if not already defined
+#ifndef HIPRT_CB
+#define HIPRT_CB
+#endif
+
+// Forward declaration of KernelTimer
+class KernelTimer;
+
+// Static callback function for hipStreamAddCallback
+static void HIPRT_CB eventCallback(hipStream_t stream, hipError_t status, void* userData) {
+  if (status != hipSuccess) return;
+  
+  KernelTimer* timer = static_cast<KernelTimer*>(userData);
+  float elapsed_time;
+  
+  // Use the getter methods to access the private members
+  HOST_TYPE(Event_t) start_event = timer->get_start_event();
+  HOST_TYPE(Event_t) stop_event = timer->get_stop_event();
+  
+  LIB_CALL(HOST_TYPE(EventElapsedTime)(&elapsed_time, start_event, stop_event));
+  
+  size_t calc_ops = timer->get_calc_ops();
+  double flops = static_cast<double>(calc_ops);
+  double gflops_val = (flops / (elapsed_time * 1e-3)) / 1e9;
+
+  // Store results in the provided pointers
+  float* time_ptr = timer->get_time_ptr();
+  float* gflops_ptr = timer->get_gflops_ptr();
+  
+  if (time_ptr != nullptr) {
+    *time_ptr = elapsed_time;
+  }
+  if (gflops_ptr != nullptr) {
+    *gflops_ptr = static_cast<float>(gflops_val);
+  }
+  
+  // Call user callback if provided
+  timer->execute_callback(elapsed_time);
+  timer->set_callback_executed(true);
+}
+
+KernelTimer::KernelTimer(size_t calc_ops, float *time, float *gflops)
+    : calc_ops(calc_ops), time_ptr(time), gflops_ptr(gflops), user_data(nullptr), 
+      callback(nullptr), callback_executed(false) {
+  LIB_CALL(HOST_TYPE(EventCreate)(&start));
+  LIB_CALL(HOST_TYPE(EventCreate)(&stop));
+}
+
+void KernelTimer::start_timer(hipStream_t stream) { 
+  LIB_CALL(HOST_TYPE(EventRecord)(start, stream));
+  callback_executed = false;
+}
+
+void KernelTimer::stop_timer(hipStream_t stream) {
+  LIB_CALL(HOST_TYPE(EventRecord)(stop, stream));
+  // Instead of synchronizing, add a callback to the stream that will be called when the event completes
+  LIB_CALL(hipStreamAddCallback(stream, eventCallback, this, 0));
+}
+
+void KernelTimer::set_callback(TimerCompletionCallback cb, void* data) {
+  callback = cb;
+  user_data = data;
+}
+
+void KernelTimer::execute_callback(float elapsed_time) {
+  if (callback && !callback_executed) {
+    callback(elapsed_time, calc_ops, time_ptr, gflops_ptr, user_data);
+  }
+}
+
+void KernelTimer::synchronize() {
+  // If callback hasn't been executed yet, synchronize and wait for event completion, then manually execute callback
+  if (!callback_executed) {
+    LIB_CALL(HOST_TYPE(EventSynchronize)(stop));
+    float elapsed_time;
+    LIB_CALL(HOST_TYPE(EventElapsedTime)(&elapsed_time, start, stop));
+    
+    double flops = static_cast<double>(calc_ops);
+    double gflops_val = (flops / (elapsed_time * 1e-3)) / 1e9;
+
+    // Store results in the provided pointers
+    if (time_ptr != nullptr) {
+      *time_ptr = elapsed_time;
+    }
+    if (gflops_ptr != nullptr) {
+      *gflops_ptr = static_cast<float>(gflops_val);
+    }
+    
+    // Execute callback
+    if (callback) {
+      callback(elapsed_time, calc_ops, time_ptr, gflops_ptr, user_data);
+    }
+    callback_executed = true;
+  }
+}
+
+KernelTimer::~KernelTimer() {
+  // Synchronize during destruction to ensure callback is executed
+  synchronize();
+  LIB_CALL(HOST_TYPE(EventDestroy)(start));
+  LIB_CALL(HOST_TYPE(EventDestroy)(stop));
+}

tests/checker/checker.cpp

@@ -0,0 +1,469 @@
+#include "checker.h"
+#include <dlfcn.h>
+#include <sstream>
+#include <fstream>
+#include <iomanip>
+#include <limits>
+#include <getopt.h>
+#include <unistd.h>
+
+std::pair<bool, std::string> verbose_allclose(const torch::Tensor &received, const torch::Tensor &expected,
+                                              float rtol = 1e-05, float atol = 1e-08, int max_print = 5) {
+    // Check if the shapes of the tensors match
+    if (received.sizes() != expected.sizes()) {
+        std::string expected_shape_str = "[";
+        std::string received_shape_str = "[";
+        auto expected_sizes = expected.sizes();
+        auto received_sizes = received.sizes();
+
+        for (int i = 0; i < expected_sizes.size(); i++) {
+            expected_shape_str += std::to_string(expected_sizes[i]);
+            if (i < expected_sizes.size() - 1)
+                expected_shape_str += ", ";
+        }
+        expected_shape_str += "]";
+
+        for (int i = 0; i < received_sizes.size(); i++) {
+            received_shape_str += std::to_string(received_sizes[i]);
+            if (i < received_sizes.size() - 1)
+                received_shape_str += ", ";
+        }
+        received_shape_str += "]";
+
+        return {false, "SIZE MISMATCH: expected " + expected_shape_str + " but got " + received_shape_str};
+    }
+
+    auto diff = torch::abs(received.to(torch::kFloat32) - expected.to(torch::kFloat32));
+
+    auto tolerance = atol + rtol * torch::abs(expected);
+
+    auto tol_mismatched = diff > tolerance;
+    auto nan_mismatched = torch::logical_xor(torch::isnan(received), torch::isnan(expected));
+    auto posinf_mismatched = torch::logical_xor(torch::isposinf(received), torch::isposinf(expected));
+    auto neginf_mismatched = torch::logical_xor(torch::isneginf(received), torch::isneginf(expected));
+
+    auto mismatched = torch::logical_or(torch::logical_or(tol_mismatched, nan_mismatched),
+                                        torch::logical_or(posinf_mismatched, neginf_mismatched));
+
+    auto mismatched_indices = torch::nonzero(mismatched);
+
+    // Count the number of mismatched elements
+    int64_t num_mismatched = mismatched.sum().item<int64_t>();
+
+    // Generate detailed information if there are mismatches
+    if (num_mismatched >= 1) {
+        std::stringstream mismatch_details;
+        auto sizes = received.sizes();
+        mismatch_details << "Mismatch found in tensors with shape [";
+        for (int i = 0; i < sizes.size(); i++) {
+            mismatch_details << sizes[i];
+            if (i < sizes.size() - 1)
+                mismatch_details << ", ";
+        }
+        mismatch_details << "]:\n";
+        mismatch_details << "Number of mismatched elements: " << num_mismatched << "\n";
+
+        for (int i = 0; i < std::min(max_print, (int)mismatched_indices.size(0)); i++) {
+            auto index = mismatched_indices[i];
+            std::vector<int64_t> idx_vec;
+            for (int j = 0; j < index.size(0); j++) {
+                idx_vec.push_back(index[j].item<int64_t>());
+            }
+
+            // Format the index as a string
+            std::string idx_str = "(";
+            for (size_t j = 0; j < idx_vec.size(); j++) {
+                idx_str += std::to_string(idx_vec[j]);
+                if (j < idx_vec.size() - 1)
+                    idx_str += ", ";
+            }
+            idx_str += ")";
+
+            float received_val, expected_val;
+            torch::Tensor received_elem = received;
+            torch::Tensor expected_elem = expected;
+
+            for (size_t j = 0; j < idx_vec.size(); j++) {
+                received_elem = received_elem[idx_vec[j]];
+                expected_elem = expected_elem[idx_vec[j]];
+            }
+
+            received_val = received_elem.item<float>();
+            expected_val = expected_elem.item<float>();
+
+            mismatch_details << "ERROR at " << idx_str << ": " << received_val << " " << expected_val << "\n";
+        }
+
+        if (num_mismatched > max_print) {
+            mismatch_details << "... and " << (num_mismatched - max_print) << " more mismatched elements.";
+        }
+
+        return {false, mismatch_details.str()};
+    }
+
+    return {true, "Maximum error: " + std::to_string(diff.max().item<float>())};
+}
+
+// Check if implementation matches reference within tolerance
+std::pair<bool, std::string> check_implementation(std::ofstream &fout, const torch::Tensor &output,
+                                                  const torch::Tensor &expected, float rtol = 2e-02, float atol = 1e-03,
+                                                  CheckerMode mode = CheckerMode::kElementWise) {
+    if (mode == CheckerMode::kRowIndex) {
+        // For row index mode, we need to sort each row before comparison
+        // since the order of indices with the same values might differ
+        auto sorted_output = output.clone();
+        auto sorted_expected = expected.clone();
+
+        sorted_output = std::get<0>(torch::sort(output, 1));
+        sorted_expected = std::get<0>(torch::sort(expected, 1));
+
+        return verbose_allclose(sorted_output, sorted_expected, rtol, atol);
+    } else if (mode == CheckerMode::kJustDump) {
+        // Dump output and expected tensors to file
+        {
+            fout << "=====OUTPUT=====" << std::endl;
+            fout << output.sizes() << std::endl;
+
+            // Manually print the full tensor to avoid truncation
+            auto sizes = output.sizes();
+            if (sizes.size() == 2) {
+                // For 2D tensors (matrices)
+                for (int64_t i = 0; i < sizes[0]; i++) {
+                    for (int64_t j = 0; j < sizes[1]; j++) {
+                        fout << std::setw(12) << std::setprecision(6) << output[i][j].item<float>() << " ";
+                    }
+                    fout << std::endl;
+                }
+            } else {
+                // Fallback for other tensor dimensions
+                fout << output << std::endl;
+            }
+        }
+
+        {
+            fout << "=====EXPECTED=====" << std::endl;
+            fout << expected.sizes() << std::endl;
+
+            // Manually print the full tensor to avoid truncation
+            auto sizes = output.sizes();
+            if (sizes.size() == 2) {
+                // For 2D tensors (matrices)
+                for (int64_t i = 0; i < sizes[0]; i++) {
+                    for (int64_t j = 0; j < sizes[1]; j++) {
+                        fout << std::setw(12) << std::setprecision(6) << expected[i][j].item<float>() << " ";
+                    }
+                    fout << std::endl;
+                }
+            } else {
+                // Fallback for other tensor dimensions
+                fout << output << std::endl;
+            }
+        }
+
+        return {true, ""};
+    }
+    return verbose_allclose(output, expected, rtol, atol);
+}
+
+constexpr int BENCHMARK_ITERS = 5;
+
+void preload() {
+void *handle_rocblas = dlopen("/usr/local/lib/python3.10/dist-packages/torch/lib/librocblas.so", RTLD_NOW | RTLD_GLOBAL);
+void *handle_hipblas = dlopen("/usr/local/lib/python3.10/dist-packages/torch/lib/libhipblas.so", RTLD_NOW | RTLD_GLOBAL);
+void *handle_hipblaslt = dlopen("/usr/local/lib/python3.10/dist-packages/torch/lib/libhipblaslt.so", RTLD_NOW | RTLD_GLOBAL);
+
+if (!handle_rocblas || !handle_hipblas || !handle_hipblaslt) {
+    fprintf(stderr, "Failed to load required libraries: %s\n", dlerror());
+    exit(1);
+}
+}
+
+int main(int argc, char **argv) {
+    // preload();
+    // bool benchmark = false;
+    bool benchmark = true;
+    bool profile_mode = false;
+    int target_test_case = -1;
+    int target_sub_case = -1;
+    int opt;
+
+    while ((opt = getopt(argc, argv, "bpt:c:")) != -1) {
+        switch (opt) {
+        case 'b':
+            benchmark = false;
+            break;
+        case 'p':
+            profile_mode = true;
+            break;
+        case 't':
+            target_sub_case = std::stoi(optarg);
+            break;
+        case 'c':
+            target_test_case = std::stoi(optarg);
+            break;
+        default:
+            fprintf(stderr, "Usage: %s [-b] [-p] [-t subcase_index] [-c test_case_index]\n", argv[0]);
+            fprintf(stderr, "  -b: Disable benchmark mode\n");
+            fprintf(stderr, "  -p: Enable profile mode (skips reference kernel and comparison)\n");
+            fprintf(stderr, "  -t: Run only the specified subcase index\n");
+            fprintf(stderr, "  -c: Run only the specified test case index\n");
+            exit(EXIT_FAILURE);
+        }
+    }
+
+    case_initialize();
+    int num_params, passed_cases = 0;
+    num_params = get_params_count();
+
+    // Validate test case index if specified
+    if (target_test_case >= 0) {
+        if (target_test_case >= num_params) {
+            std::cerr << "Error: Test case index " << target_test_case << " is out of range (0-" << (num_params - 1)
+                      << ")" << std::endl;
+            exit(EXIT_FAILURE);
+        }
+    }
+
+    std::vector<std::vector<PerfMetrics>> run_times(num_params);
+    std::vector<std::tuple<bool, std::string, std::vector<std::pair<float, float>>>> results;
+
+    // If targeting specific test case and subcase, run multiple times and output only the best time
+    if (target_test_case >= 0 && target_sub_case >= 0) {
+        void *input = case_get_input(target_test_case);
+        std::vector<Checkee> output;
+        float best_time = std::numeric_limits<float>::max();
+
+        for (int j = 0; j < BENCHMARK_ITERS; j++) {
+            PerfMetrics metrics;
+            output = case_run_kernel(input, &metrics);
+
+            if (metrics.count <= target_sub_case) {
+                std::cerr << "Error: Subcase index " << target_sub_case << " is out of range (0-" << (metrics.count - 1)
+                          << ")" << std::endl;
+                exit(EXIT_FAILURE);
+            }
+
+            best_time = std::min(best_time, metrics.entries[target_sub_case].time);
+        }
+
+        std::cout << std::fixed << std::setprecision(6) << best_time * 1e3 << std::endl;
+        case_destroy(input);
+        return 0;
+    }
+
+    // Normal execution path
+    if (!profile_mode && target_test_case < 0) {
+        std::cout << "Found " << num_params << " test cases for " << case_get_name() << '\n';
+    }
+    if (benchmark) {
+        std::cout << "Benchmark mode enabled\n";
+    }
+    if (profile_mode) {
+        std::cout << "Profile mode enabled (skipping reference kernels and comparison)\n";
+    }
+
+    // Determine which test cases to run
+    std::vector<int> test_cases_to_run;
+    if (target_test_case >= 0) {
+        test_cases_to_run.push_back(target_test_case);
+    } else {
+        for (int i = 0; i < num_params; i++) {
+            test_cases_to_run.push_back(i);
+        }
+    }
+
+    for (int i : test_cases_to_run) {
+        std::ofstream *fout = nullptr;
+        void *input = case_get_input(i);
+        if (!profile_mode && target_test_case < 0) {
+            std::cerr << "Running test case " << i << std::flush;
+        }
+        std::vector<Checkee> reference;
+        if (!profile_mode) {
+            reference = case_run_ref_kernel(input);
+        }
+        std::vector<Checkee> output;
+        for (int j = 0; j < (benchmark ? BENCHMARK_ITERS : 1); j++) {
+            PerfMetrics metrics;
+            output = case_run_kernel(input, &metrics);
+            run_times[i].push_back(metrics);
+        }
+
+        bool match = true;
+        std::string case_message;
+
+        if (!profile_mode) {
+            if (reference.size() != output.size()) {
+                std::cerr << "Wrong test definition: reference and output have different sizes" << '\n';
+                abort();
+            }
+
+            for (int j = 0; j < reference.size(); j++) {
+                float rtol, atol;
+                get_error_tolerance(&rtol, &atol);
+                if (output[j].mode == CheckerMode::kJustDump) {
+                    if (!fout) {
+                        fout = new std::ofstream(std::string("case_") + std::to_string(i) + ".txt");
+                    }
+                    *fout << "===== SUBCASE " << output[j].name << "=====" << std::endl;
+                }
+                auto [match_sub, message_sub] =
+                    check_implementation(*fout, *output[j].tensor, *reference[j].tensor, rtol, atol, output[j].mode);
+                if (!match_sub) {
+                    case_message += "Err on sub case " + std::to_string(j) + ": " + message_sub + "\n";
+                    match = false;
+                }
+            }
+            if (match) {
+                passed_cases++;
+            }
+        } else {
+            match = true;
+            passed_cases++;
+        }
+
+        std::vector<std::pair<float, float>> case_metrics;
+
+        // Process metrics for each run
+        for (const auto &run : run_times[i]) {
+            if (run.count == 1) {
+                // Backward compatibility: single metric case
+                case_metrics.push_back({run.entries[0].time, run.entries[0].gflops});
+            } else {
+                // Multiple metrics case - first entry is the total result
+                case_metrics.push_back({run.entries[0].time, run.entries[0].gflops});
+            }
+        }
+
+        results.push_back(std::make_tuple(match, case_message, case_metrics));
+        case_destroy(input);
+        if (!profile_mode && target_test_case < 0) {
+            std::cout << "\033[2K\r" << std::flush;
+        }
+    }
+
+    // Only show detailed output if not in single test case mode
+    if (target_test_case < 0) {
+        std::cout << "=======================" << '\n';
+        if (!profile_mode) {
+            if (passed_cases == num_params) {
+                std::cout << "✅ All " << num_params << " test cases passed!" << '\n';
+            } else {
+                std::cout << "❌ [" << num_params - passed_cases << "/" << num_params << "] test cases failed!" << '\n';
+            }
+        } else {
+            std::cout << "Profile mode: results comparison skipped" << '\n';
+        }
+        std::cout << "-----------------------" << '\n';
+
+        for (int i = 0; i < num_params; i++) {
+            auto [match, message, metrics] = results[i];
+
+            // Calculate best and worst metrics
+            float best_time = std::numeric_limits<float>::max();
+            float best_gflops = 0.0f;
+            float worst_time = 0.0f;
+            float worst_gflops = std::numeric_limits<float>::max();
+
+            for (const auto &[time, gflops] : metrics) {
+                best_time = std::min(best_time, time);
+                best_gflops = std::max(best_gflops, gflops);
+                worst_time = std::max(worst_time, time);
+                worst_gflops = std::min(worst_gflops, gflops);
+            }
+
+            std::string timing_info;
+            if (benchmark) {
+                std::stringstream ss;
+                ss << std::fixed << std::setprecision(2);
+                ss << "Best: [\033[1m" << best_time * 1e3 << "\033[0m us, \033[1m" << best_gflops / 1e3
+                   << "\033[0m TFLOPS], "
+                   << "\033[2mSlowest: [" << worst_time * 1e3 << " us, " << worst_gflops / 1e3 << " TFLOPS]\033[0m";
+                timing_info = ss.str();
+            } else {
+                std::stringstream ss;
+                ss << std::fixed << std::setprecision(2);
+                ss << "Time: " << best_time * 1e3 << " us, TFLOPS: " << best_gflops / 1e3;
+                timing_info = ss.str();
+            }
+
+            if (!profile_mode && !match) {
+                std::cout << "❌ Test case " << i << ": " << timing_info << "\n" << message << '\n';
+            } else {
+                std::cout << "✅ Test case " << i << ": " << timing_info << "\n";
+            }
+
+            // Print sub-results if there are multiple metrics
+            if (run_times[i][0].count > 1) {
+                for (int j = 1; j < run_times[i][0].count; j++) {
+                    std::stringstream ss;
+                    ss << std::fixed << std::setprecision(2);
+                    ss << "    - Sub-case " << run_times[i][0].entries[j].name << ": ";
+
+                    if (benchmark) {
+                        float sub_best_time = std::numeric_limits<float>::max();
+                        float sub_best_gflops = 0.0f;
+                        float sub_worst_time = 0.0f;
+                        float sub_worst_gflops = std::numeric_limits<float>::max();
+
+                        for (const auto &run : run_times[i]) {
+                            sub_best_time = std::min(sub_best_time, run.entries[j].time);
+                            sub_best_gflops = std::max(sub_best_gflops, run.entries[j].gflops);
+                            sub_worst_time = std::max(sub_worst_time, run.entries[j].time);
+                            sub_worst_gflops = std::min(sub_worst_gflops, run.entries[j].gflops);
+                        }
+
+                        ss << "Best: [\033[1m" << sub_best_time * 1e3 << "\033[0m us, \033[1m" << sub_best_gflops / 1e3
+                           << "\033[0m TFLOPS], "
+                           << "\033[2mSlowest: [" << sub_worst_time * 1e3 << " us, " << sub_worst_gflops / 1e3
+                           << " TFLOPS]\033[0m";
+                    } else {
+                        ss << "Time: " << run_times[i][0].entries[j].time * 1e3
+                           << " us, TFLOPS: " << run_times[i][0].entries[j].gflops / 1e3;
+                    }
+
+                    std::cout << ss.str() << std::endl;
+                }
+            }
+        }
+        std::cout << "-----------------------" << '\n';
+
+        // Calculate geometric mean of time and GFLOPS
+        double geo_mean_time = 1.0;
+        double geo_mean_gflops = 1.0;
+
+        for (int i = 0; i < num_params; i++) {
+            auto [match, message, metrics] = results[i];
+            // Always use the best performance metrics for geometric mean
+            float best_time = std::numeric_limits<float>::max();
+            float best_gflops = 0.0f;
+
+            for (const auto &[time, gflops] : metrics) {
+                best_time = std::min(best_time, time);
+                best_gflops = std::max(best_gflops, gflops);
+            }
+
+            geo_mean_time *= best_time;
+            geo_mean_gflops *= best_gflops;
+        }
+
+        geo_mean_time = std::pow(geo_mean_time, 1.0 / num_params);
+        geo_mean_gflops = std::pow(geo_mean_gflops, 1.0 / num_params);
+
+        if (benchmark) {
+            std::stringstream ss;
+            ss << std::fixed << std::setprecision(2);
+            ss << "GeoMean - Best Time: \033[1m" << geo_mean_time * 1e3 << "\033[0m us, Best TFLOPS: \033[1m"
+               << geo_mean_gflops / 1e3 << "\033[0m";
+            std::cout << ss.str() << std::endl;
+        } else {
+            std::stringstream ss;
+            ss << std::fixed << std::setprecision(2);
+            ss << "GeoMean - Time: " << geo_mean_time * 1e3 << " us, TFLOPS: " << geo_mean_gflops / 1e3;
+            std::cout << ss.str() << std::endl;
+        }
+        std::cout << "=======================" << '\n';
+    }
+
+    return 0;
+}

tests/checker/checker.h

@@ -0,0 +1,33 @@
+
+#pragma once
+#include <iostream>
+#include <vector>
+#include "torch/torch.h"
+#include "metrics.h"
+
+enum class CheckerMode {
+  kElementWise,
+  kRowIndex,
+  kJustDump,
+};
+
+struct Checkee {
+  torch::Tensor *tensor;
+  CheckerMode mode;
+  std::string name;
+};
+
+void case_initialize();
+int get_params_count();
+void *case_get_input(int index);
+std::vector<Checkee> case_run_kernel(void *input, PerfMetrics* metrics);
+std::vector<Checkee> case_run_ref_kernel(void *input);
+const char *case_get_name();
+void get_error_tolerance(float *rtol, float *atol);
+void case_destroy(void *input);
+CheckerMode get_checker_mode();
+
+
+// using OutputData = torch::Tensor;
+// void ref_kernel(const BlockwiseMatmulInputs &data);
+// BlockwiseMatmulInputs generate_input(int m, int n, int k, int seed);

tests/checker/metrics.h

@@ -0,0 +1,13 @@
+#pragma once
+
+
+struct PerfMetricEntry {
+    char name[20];
+    float time;
+    float gflops;
+};
+
+struct PerfMetrics {
+    int count = 0;
+    PerfMetricEntry entries[20];
+};

torch-ext/gemm/__init__.py

@@ -0,0 +1,19 @@
+from typing import Optional
+import torch
+from ._ops import ops
+
+def gemm(a: torch.Tensor, b: torch.Tensor, as_: torch.Tensor, bs: torch.Tensor, 
+         out: Optional[torch.Tensor] = None) -> torch.Tensor:
+         
+    if out is None:
+        # Create output tensor with appropriate shape and dtype
+        M, K = a.shape
+        K_b, N = b.shape
+        assert K == K_b, f"Matrix dimension mismatch: A has {K} cols, B has {K_b} rows"
+        
+        # Output should be BF16 type on the same device as inputs
+        out = torch.empty((M, N), dtype=torch.bfloat16, device=a.device)
+    
+    ops.gemm(out, a, b, as_, bs)
+    return out
+

torch-ext/torch_binding.cpp

@@ -0,0 +1,53 @@
+#include <torch/all.h>
+#include <torch/library.h>
+#include <hip/hip_runtime.h>
+
+#include "registration.h"
+#include "torch_binding.h"
+
+// Forward declaration of the C function from gemm_launcher.hip
+extern "C" {
+    struct PerfMetrics;
+    void run(void *a, void *b, void *as, void *bs, void *c, int m, int n, int k, PerfMetrics *metrics, hipStream_t job_stream0);
+}
+
+void gemm(torch::Tensor &out, torch::Tensor const &a, torch::Tensor const &b, 
+          torch::Tensor const &as, torch::Tensor const &bs) {
+    
+    // Validate tensor properties
+    TORCH_CHECK(a.device().is_cuda(), "Input tensor a must be on GPU device");
+    TORCH_CHECK(b.device().is_cuda(), "Input tensor b must be on GPU device");
+    TORCH_CHECK(as.device().is_cuda(), "Scale tensor as must be on GPU device");
+    TORCH_CHECK(bs.device().is_cuda(), "Scale tensor bs must be on GPU device");
+    TORCH_CHECK(out.device().is_cuda(), "Output tensor out must be on GPU device");
+    
+    TORCH_CHECK(a.is_contiguous(), "Input tensor a must be contiguous");
+    TORCH_CHECK(b.is_contiguous(), "Input tensor b must be contiguous");
+    TORCH_CHECK(as.is_contiguous(), "Scale tensor as must be contiguous");
+    TORCH_CHECK(bs.is_contiguous(), "Scale tensor bs must be contiguous");
+    TORCH_CHECK(out.is_contiguous(), "Output tensor out must be contiguous");
+    
+    // Get matrix dimensions from tensor shapes
+    // Assuming a is [M, K], b is [K, N], out is [M, N]
+    int M = a.size(0);
+    int K = a.size(1);
+    int N = b.size(1);
+    
+    TORCH_CHECK(b.size(0) == K, "Matrix dimensions mismatch: a.size(1) != b.size(0)");
+    TORCH_CHECK(out.size(0) == M, "Output tensor dimension mismatch: out.size(0) != M");
+    TORCH_CHECK(out.size(1) == N, "Output tensor dimension mismatch: out.size(1) != N");
+    
+    // Use default HIP stream (stream 0)
+    const hipStream_t stream = 0;
+    
+    // Call the C function
+    run(a.data_ptr(), b.data_ptr(), as.data_ptr(), bs.data_ptr(), out.data_ptr(),
+        M, N, K, nullptr, stream);
+}
+
+TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
+  ops.def("gemm(Tensor! out, Tensor a, Tensor b, Tensor a_scale, Tensor b_scale) -> ()");
+  ops.impl("gemm", torch::kCUDA, &gemm);
+}
+
+REGISTER_EXTENSION(TORCH_EXTENSION_NAME)

torch-ext/torch_binding.h

@@ -0,0 +1,6 @@
+#pragma once
+
+#include <torch/torch.h>
+
+void gemm(torch::Tensor &out, torch::Tensor const &a, torch::Tensor const &b, 
+          torch::Tensor const &as, torch::Tensor const &bs);