Tune gemm op bf16

aiter/configs/tuned_gemm.csv

@@ -1,11 +1,9 @@
-M,N,K,bias,dtype,outdtype,scaleAB,cu_num,libtype,solidx,soltimes,kernelName,tflops,bw
-64,256,5120,False,torch.bfloat16,torch.float32,False,80,asm,2,0,_ZN5aiter28bf16gemm_outf32_tn_64x64_pf3E,0,0
-80,256,5120,False,torch.bfloat16,torch.float32,False,80,asm,3,0,_ZN5aiter28bf16gemm_outf32_tn_96x64_pf3E,0,0
-128,256,5120,False,torch.bfloat16,torch.float32,False,80,asm,2,0,_ZN5aiter28bf16gemm_outf32_tn_64x64_pf3E,0,0
-150,256,5120,False,torch.bfloat16,torch.float32,False,80,asm,0,0,_ZN5aiter28bf16gemm_outf32_tn_32x64_pf3E,0,0
-192,256,5120,False,torch.bfloat16,torch.float32,False,80,asm,3,0,_ZN5aiter28bf16gemm_outf32_tn_96x64_pf3E,0,0
-220,256,5120,False,torch.bfloat16,torch.float32,False,80,asm,0,0,_ZN5aiter28bf16gemm_outf32_tn_32x64_pf3E,0,0
-256,256,5120,False,torch.bfloat16,torch.float32,False,80,asm,2,0,_ZN5aiter28bf16gemm_outf32_tn_64x64_pf3E,0,0
-384,256,5120,False,torch.bfloat16,torch.float32,False,80,asm,3,0,_ZN5aiter28bf16gemm_outf32_tn_96x64_pf3E,0,0
-448,256,5120,False,torch.bfloat16,torch.float32,False,80,asm,0,0,_ZN5aiter28bf16gemm_outf32_tn_32x64_pf3E,0,0
-512,256,5120,False,torch.bfloat16,torch.float32,False,80,asm,2,0,_ZN5aiter28bf16gemm_outf32_tn_64x64_pf3E,0,0
+M,N,K,bias,dtype,outdtype,scaleAB,cu_num,libtype,solidx,splitK,soltimes,kernelName,err_ratio,tflops,bw
+64,256,5120,False,torch.bfloat16,torch.float32,False,80,asm,1,10,8.3,_ZN5aiter28bf16gemm_outf32_tn_32x64_pf3E,0.0,20.21,402.69
+80,256,5120,False,torch.bfloat16,torch.float32,False,80,asm,2,10,8.8,_ZN5aiter28bf16gemm_outf32_tn_48x64_pf3E,0.0,23.83,400.29
+128,256,5120,False,torch.bfloat16,torch.float32,False,80,asm,3,10,11.2,_ZN5aiter28bf16gemm_outf32_tn_64x64_pf3E,0.0,29.96,362.79
+150,256,5120,False,torch.bfloat16,torch.float32,False,80,asm,2,5,11.7,_ZN5aiter28bf16gemm_outf32_tn_48x64_pf3E,0.0,33.61,368.46
+192,256,5120,False,torch.bfloat16,torch.float32,False,80,asm,5,10,12.8,_ZN5aiter28bf16gemm_outf32_tn_96x64_pf3E,0.0,39.32,373.76
+220,256,5120,False,torch.bfloat16,torch.float32,False,80,asm,3,5,13.1,_ZN5aiter28bf16gemm_outf32_tn_64x64_pf3E,0.0,44.02,389.28
+256,256,5120,False,torch.bfloat16,torch.float32,False,80,asm,3,5,14.4,_ZN5aiter28bf16gemm_outf32_tn_64x64_pf3E,0.0,46.6,382.29
+384,256,5120,False,torch.bfloat16,torch.float32,False,80,asm,5,5,19.7,_ZN5aiter28bf16gemm_outf32_tn_96x64_pf3E,0.0,51.1,352.63

aiter/configs/untuned_gemm.csv

@@ -1 +1,11 @@
 M,N,K,bias,dtype,outdtype,scaleAB
+64,256,5120,False,torch.bfloat16,torch.float32,False
+80,256,5120,False,torch.bfloat16,torch.float32,False
+128,256,5120,False,torch.bfloat16,torch.float32,False
+150,256,5120,False,torch.bfloat16,torch.float32,False
+192,256,5120,False,torch.bfloat16,torch.float32,False
+220,256,5120,False,torch.bfloat16,torch.float32,False
+256,256,5120,False,torch.bfloat16,torch.float32,False
+384,256,5120,False,torch.bfloat16,torch.float32,False
+448,256,5120,False,torch.bfloat16,torch.float32,False
+512,256,5120,False,torch.bfloat16,torch.float32,False

aiter/jit/core.py

@@ -104,6 +104,11 @@ def mp_lock(
     "AITER_CONFIG_BATCHED_GEMM_BF16",
     f"{AITER_ROOT_DIR}/aiter/configs/bf16_tuned_batched_gemm.csv",
 )
+
+AITER_CONFIG_GEMM_BF16 = os.getenv(
+    "AITER_CONFIG_GEMM_BF16",
+    f"{AITER_ROOT_DIR}/aiter/configs/tuned_gemm.csv",
+)
 ## merge config files
 ##example: AITER_CONFIG_GEMM_A4W4="/path1:/path2"
 import pandas as pd
@@ -156,6 +161,9 @@ def update_config_files(file_path: str, merge_name: str):
 AITER_CONFIG_BF16_BATCHED_GEMM_FILE = update_config_files(
     AITER_CONFIG_BF16_BATCHED_GEMM, "bf16_tuned_batched_gemm"
 )
+AITER_CONFIG_GEMM_BF16_FILE = update_config_files(
+    AITER_CONFIG_GEMM_BF16, "bf16_tuned_gemm"
+)
 # config_env end here
 
 find_aiter = importlib.util.find_spec("aiter")

aiter/ops/gemm_op_a4w4.py

@@ -93,8 +93,8 @@ def gemm_a4w4(
         splitK = ck_config.get("splitK", None)
         kernelName = ck_config["kernelName"]
     if (
-        m < 256
-        or (ck_config is not None and kernelName.find("_ZN") == -1)
+        ck_config is not None
+        and kernelName.find("_ZN") == -1
         # or bias is None
     ):
         splitK = 0 if splitK is None else splitK