Tianxing/moe gemm

.github/workflows/amd_perf_kernel_Integration_tests.yml

@@ -130,6 +130,7 @@ jobs:
           pytest -vvvv ./python/perf-kernels/softmax.py
           pytest -vvv ./python/perf-kernels/rmsnorm.py
           pytest -vvv ./python/perf-kernels/layernorm.py
+          pytest -vvv ./python/perf-kernels/fused_moe/moe-gemm.py
           sh ./python/perf-kernels/streamk/utils/unittest.sh
           pytest -vvv ./python/perf-kernels/multreduce_matmul_kernel.py
       - name: Run Perf Kernels Benchmark

python/perf-kernels/README.md

@@ -99,3 +99,6 @@ Kernel that implements RMS Norm over a row of tensor.
 
 ## `layernorm.py`
 Kernel that implements Layer Normalization over a row on tensor
+
+## `fused_moe/moe-gemm.py`
+Kernel that implements moe gemm.

python/perf-kernels/fused_moe/configs/device_name=AMD_Instinct_MI300X.json

@@ -0,0 +1,35 @@
+{
+  "small_M": {
+    "BLOCK_SIZE_M": 64,
+    "BLOCK_SIZE_N": 64,
+    "BLOCK_SIZE_K": 64,
+    "GROUP_SIZE_M": 4,
+    "num_warps": 8,
+    "num_stages": 2,
+    "waves_per_eu": 0,
+    "matrix_instr_nonkdim": 16,
+    "kpack": 2
+  },
+  "medium_M": {
+    "BLOCK_SIZE_M": 128,
+    "BLOCK_SIZE_N": 128,
+    "BLOCK_SIZE_K": 128,
+    "GROUP_SIZE_M": 1,
+    "num_warps": 8,
+    "num_stages": 2,
+    "waves_per_eu": 0,
+    "matrix_instr_nonkdim": 16,
+    "kpack": 2
+  },
+  "large_M": {
+    "BLOCK_SIZE_M": 256,
+    "BLOCK_SIZE_N": 256,
+    "BLOCK_SIZE_K": 64,
+    "GROUP_SIZE_M": 1,
+    "num_warps": 8,
+    "num_stages": 2,
+    "waves_per_eu": 0,
+    "matrix_instr_nonkdim": 16,
+    "kpack": 2
+  }
+}