[FLYDSL] Rebase flydsl hgemm kernels with mixed policies

aiter/aot/flydsl/gemm.py

@@ -49,7 +49,7 @@
 from aiter.jit.core import AITER_CONFIGS
 from aiter.ops.flydsl.gemm_kernels import (
     SPLIT_K_SEMAPHORE_MAX_LEN,
-    get_flydsl_splitk_hgemm_kernel_params,
+    get_flydsl_hgemm_kernel_params,
 )
 from aiter.ops.flydsl.kernels.hgemm_dispatch import compile_flydsl_hgemm_kernel
 from aiter.ops.flydsl.kernels.preshuffle_gemm import compile_preshuffle_gemm_a8
@@ -147,8 +147,8 @@ def parse_csv(csv_path: str):
 
             if kernel_name.startswith("flydsl_bpreshuflle_"):
                 params = _parse_preshuffle_kernel_name(kernel_name)
-            elif kernel_name.startswith("flydsl_gemm"):
-                params = get_flydsl_splitk_hgemm_kernel_params(kernel_name)
+            elif kernel_name.startswith("flydsl_hgemm"):
+                params = get_flydsl_hgemm_kernel_params(kernel_name)
                 if params is not None:
                     params = dict(params)
                     params["kind"] = "hgemm"
@@ -211,25 +211,19 @@ def _compile_hgemm_to_cache(
     k: int,
     dtype: str,
     out_dtype: str,
-    tile_m: int,
-    tile_n: int,
-    tile_k: int,
-    stages: int,
-    split_k: int,
-    block_m_warps: int,
-    block_n_warps: int,
-    block_k_warps: int,
-    n_tile_repeat: int = 1,
-    persistent_n_tiles: int = 1,
-    waves_per_eu: int = 0,
-    b_to_lds_unroll: int = 0,
-    async_copy: bool,
-    b_to_lds: bool,
-    b_preshuffle: bool,
-    c_to_lds: bool,
-    target_gfx: str,
-    kernel_family: str = "hgemm",
-    has_bias: bool = False,
+    has_bias: bool,
+    TILE_M: int,
+    TILE_N: int,
+    TILE_K: int,
+    STAGES: int,
+    SPLIT_K: int,
+    BLOCK_M_WARPS: int,
+    BLOCK_N_WARPS: int,
+    BLOCK_K_WARPS: int,
+    USE_HALF_TILE_INTERLEAVED: bool,
+    GROUP_M: int,
+    HAS_BIAS: bool,
+    HAS_K_TAILL: bool,
     **kwargs,
 ):
     del kwargs, out_dtype
@@ -256,27 +250,19 @@ def _compile_hgemm_to_cache(
     stream = fx.Stream(0)
 
     exe = compile_flydsl_hgemm_kernel(
-        dtype,
-        n,
-        k,
-        kernel_family=kernel_family,
-        tile_m=tile_m,
-        tile_n=tile_n,
-        tile_k=tile_k,
-        stages=stages,
-        split_k=split_k,
-        block_m_warps=block_m_warps,
-        block_n_warps=block_n_warps,
-        block_k_warps=block_k_warps,
-        n_tile_repeat=n_tile_repeat,
-        persistent_n_tiles=persistent_n_tiles,
-        waves_per_eu=waves_per_eu,
-        b_to_lds_unroll=b_to_lds_unroll,
-        async_copy=async_copy,
-        b_to_lds=b_to_lds,
-        b_preshuffle=b_preshuffle,
-        c_to_lds=c_to_lds,
-        has_bias=has_bias,
+        dtype=dtype,
+        tile_m=TILE_M,
+        tile_n=TILE_N,
+        tile_k=TILE_K,
+        stages=STAGES,
+        split_k=SPLIT_K,
+        block_m_warps=BLOCK_M_WARPS,
+        block_n_warps=BLOCK_N_WARPS,
+        block_k_warps=BLOCK_K_WARPS,
+        has_bias=HAS_BIAS,
+        has_k_tail=HAS_K_TAILL,
+        group_m=GROUP_M,
+        policy="ht" if USE_HALF_TILE_INTERLEAVED else "ft",
     )
     # FlyDSL JIT does not accept None for tensor slots; pass real buffers for
     # optional bias and split-K sync tensors.
@@ -287,9 +273,11 @@ def _compile_hgemm_to_cache(
         _ptr_view_safe(a),
         _ptr_view_safe(b),
         _ptr_view_safe(launch_bias),
-        m,
         _ptr_view_safe(semaphore),
         _ptr_view_safe(signal),
+        m,
+        n,
+        k,
         stream,
     )
 

aiter/configs/model_configs/kimi_bf16_tuned_gemm.csv

@@ -1,24 +1,44 @@
 gfx,cu_num,M,N,K,bias,dtype,outdtype,scaleAB,bpreshuffle,libtype,solidx,splitK,us,kernelName,err_ratio,tflops,bw
-gfx950,256,1024,384,7168,False,torch.bfloat16,torch.bfloat16,False,False,asm,7,3,20.1108,_ZN5aiter39bf16gemm_fp32bf16_tn_80x64_splitk_cleanE,0.0104,280.3,1042.8
-gfx950,256,96,1024,7168,False,torch.bfloat16,torch.bfloat16,False,False,flydsl,4127,4,11.0603,flydsl_gemm4_abf16_wbf16_bf16_t32x64x128_split_k4_block_m_warp1_block_n_warp4_block_k_warp1_async_copyTrue_b_to_ldsTrue_b_preshuffleFalse_c_to_ldsFalse_gfx950,0.0142,127.42,1469.48
-gfx950,256,32,2112,7168,False,torch.bfloat16,torch.bfloat16,False,False,flydsl,1926,7,10.6796,flydsl_gemm6_abf16_wbf16_bf16_t32x64x64_split_k7_block_m_warp1_block_n_warp4_block_k_warp1_async_copyTrue_b_to_ldsTrue_b_preshuffleFalse_c_to_ldsFalse_gfx950,0.0199,90.72,2890.7
-gfx950,256,64,2112,7168,False,torch.bfloat16,torch.bfloat16,False,False,flydsl,2321,7,13.7074,flydsl_gemm5_abf16_wbf16_bf16_t64x64x64_split_k7_block_m_warp2_block_n_warp2_block_k_warp1_async_copyTrue_b_to_ldsTrue_b_preshuffleFalse_c_to_ldsFalse_gfx950,0.0199,141.37,2295.51
-gfx950,256,1024,2112,7168,False,torch.bfloat16,torch.bfloat16,False,False,flydsl,1542,1,68.5856,flydsl_gemm3_abf16_wbf16_bf16_t64x64x64_split_k1_block_m_warp1_block_n_warp1_block_k_warp1_async_copyTrue_b_to_ldsTrue_b_preshuffleFalse_c_to_ldsFalse_gfx950,0.0,452.05,718.56
-gfx950,256,4096,2112,7168,False,torch.bfloat16,torch.bfloat16,False,False,torch,0,0,117.8273,native,0.0,1052.53,902.16
-gfx950,256,16384,2112,7168,False,torch.bfloat16,torch.bfloat16,False,False,torch,0,0,480.3821,native,0.0,1032.65,696.04
-gfx950,256,1024,3072,1536,False,torch.bfloat16,torch.bfloat16,False,False,flydsl,5322,1,17.6546,flydsl_gemm4_abf16_wbf16_bf16_t128x128x64_split_k1_block_m_warp2_block_n_warp4_block_k_warp1_async_copyTrue_b_to_ldsTrue_b_preshuffleFalse_c_to_ldsFalse_gfx950,0.0,547.37,1069.09
-gfx950,256,4096,3072,1536,False,torch.bfloat16,torch.bfloat16,False,False,torch,0,0,40.3383,native,0.0,958.26,1169.75
-gfx950,256,16384,3072,1536,False,torch.bfloat16,torch.bfloat16,False,False,torch,0,0,130.6942,native,0.0,1183.06,1227.54
-gfx950,256,256,4096,512,False,torch.bfloat16,torch.bfloat16,False,False,flydsl,4620,1,5.2609,flydsl_gemm6_abf16_wbf16_bf16_t64x64x64_split_k1_block_m_warp4_block_n_warp2_block_k_warp1_async_copyTrue_b_to_ldsTrue_b_preshuffleFalse_c_to_ldsFalse_gfx950,0.0,204.1,1245.72
-gfx950,256,1024,4096,512,False,torch.bfloat16,torch.bfloat16,False,False,flydsl,2875,1,9.5558,flydsl_gemm4_abf16_wbf16_bf16_t128x128x64_split_k1_block_m_warp4_block_n_warp2_block_k_warp1_async_copyTrue_b_to_ldsTrue_b_preshuffleFalse_c_to_ldsFalse_gfx950,0.0,449.46,1426.51
-gfx950,256,4096,4096,512,False,torch.bfloat16,torch.bfloat16,False,False,torch,0,0,20.9336,native,0.0,820.68,2003.62
-gfx950,256,16384,4096,512,False,torch.bfloat16,torch.bfloat16,False,False,torch,0,0,80.4652,native,0.0,854.03,1928.65
-gfx950,256,96,7168,512,False,torch.bfloat16,torch.bfloat16,False,False,flydsl,3107,1,5.2439,flydsl_gemm4_abf16_wbf16_bf16_t64x64x128_split_k1_block_m_warp1_block_n_warp4_block_k_warp1_async_copyTrue_b_to_ldsTrue_b_preshuffleFalse_c_to_ldsFalse_gfx950,0.0,134.37,1680.92
-gfx950,256,128,7168,512,False,torch.bfloat16,torch.bfloat16,False,False,flydsl,4843,1,5.4135,flydsl_gemm7_abf16_wbf16_bf16_t64x64x64_split_k1_block_m_warp1_block_n_warp4_block_k_warp1_async_copyTrue_b_to_ldsTrue_b_preshuffleFalse_c_to_ldsFalse_gfx950,0.0,173.55,1719.06
-gfx950,256,256,7168,512,False,torch.bfloat16,torch.bfloat16,False,False,flydsl,2961,1,6.9191,flydsl_gemm4_abf16_wbf16_bf16_t128x64x64_split_k1_block_m_warp4_block_n_warp1_block_k_warp1_async_copyTrue_b_to_ldsTrue_b_preshuffleFalse_c_to_ldsFalse_gfx950,0.0,271.57,1629.14
-gfx950,256,1024,7168,512,False,torch.bfloat16,torch.bfloat16,False,False,flydsl,91,1,13.7628,flydsl_gemm2_abf16_wbf16_bf16_t128x128x64_split_k1_block_m_warp4_block_n_warp4_block_k_warp1_async_copyTrue_b_to_ldsTrue_b_preshuffleFalse_c_to_ldsFalse_gfx950,0.0,546.12,1676.16
-gfx950,256,4096,7168,512,False,torch.bfloat16,torch.bfloat16,False,False,asm,10,1,40.9031,_ZN5aiter24bf16gemm_bf16_tn_256x256E,0.0,735.02,1717.59
-gfx950,256,16384,7168,512,False,torch.bfloat16,torch.bfloat16,False,False,torch,0,0,131.8454,native,0.0,912.12,1964.41
+gfx950,256,32,384,7168,False,torch.bfloat16,torch.bfloat16,False,False,flydsl,319,8,6.9304,flydsl_hgemm_abf16_wbf16_bf16_t16x64x64x7_ks8_w1x2x1_bias0_ktail0_gm0_pft_gfx950,0.0229,25.42,864.07
+gfx950,256,64,384,7168,False,torch.bfloat16,torch.bfloat16,False,False,flydsl,1046,8,7.8039,flydsl_hgemm_abf16_wbf16_bf16_t32x64x64x5_ks8_w1x4x1_bias0_ktail0_gm0_pft_gfx950,0.0212,45.15,829.29
+gfx950,256,96,384,7168,False,torch.bfloat16,torch.bfloat16,False,False,flydsl,1161,8,8.762,flydsl_hgemm_abf16_wbf16_bf16_t32x64x64x8_ks8_w1x4x1_bias0_ktail0_gm0_pft_gfx950,0.0225,60.32,793.77
+gfx950,256,256,384,7168,False,torch.bfloat16,torch.bfloat16,False,False,flydsl,1359,4,11.8424,flydsl_hgemm_abf16_wbf16_bf16_t32x64x128x5_ks4_w1x4x1_bias0_ktail0_gm0_pft_gfx950,0.0143,119.0,791.36
+gfx950,256,512,384,7168,False,torch.bfloat16,torch.bfloat16,False,False,flydsl,2722,2,15.1024,flydsl_hgemm_abf16_wbf16_bf16_t32x64x128x5_ks2_w1x4x1_bias0_ktail0_gm0_pft_gfx950,0.006,186.63,876.57
+gfx950,256,1024,384,7168,False,torch.bfloat16,torch.bfloat16,False,False,asm,7,3,20.1669,_ZN5aiter39bf16gemm_fp32bf16_tn_80x64_splitk_cleanE,0.0104,279.52,1039.9
+gfx950,256,2048,384,7168,False,torch.bfloat16,torch.bfloat16,False,False,torch,0,0,28.2385,native,0.0,399.25,1290.37
+gfx950,256,4096,384,7168,False,torch.bfloat16,torch.bfloat16,False,False,flydsl,2717,1,42.1358,flydsl_hgemm_abf16_wbf16_bf16_t64x96x128x3_ks1_w1x2x2_bias0_ktail0_gm4_pft_gfx950,0.0062,535.14,1598.9
+gfx950,256,8192,384,7168,False,torch.bfloat16,torch.bfloat16,False,False,flydsl,473,1,66.5626,flydsl_hgemm_abf16_wbf16_bf16_t128x96x64x3_ks1_w2x2x1_bias0_ktail0_gm4_pft_gfx950,0.0,677.51,1941.59
+gfx950,256,16384,384,7168,False,torch.bfloat16,torch.bfloat16,False,False,torch,0,0,101.313,native,0.0,890.25,2496.91
+gfx950,256,96,1024,7168,False,torch.bfloat16,torch.bfloat16,False,False,flydsl,755,4,11.2045,flydsl_hgemm_abf16_wbf16_bf16_t32x64x128x4_ks4_w1x4x1_bias0_ktail0_gm0_pft_gfx950,0.014,125.78,1450.57
+gfx950,256,128,1024,7168,False,torch.bfloat16,torch.bfloat16,False,False,flydsl,1342,4,12.3884,flydsl_hgemm_abf16_wbf16_bf16_t32x64x128x5_ks4_w1x4x1_bias0_ktail0_gm0_pft_gfx950,0.0145,151.68,1354.27
+gfx950,256,256,1024,7168,False,torch.bfloat16,torch.bfloat16,False,False,asm,8,5,16.2424,_ZN5aiter39bf16gemm_fp32bf16_tn_96x64_splitk_cleanE,0.0163,231.38,1162.04
+gfx950,256,512,1024,7168,False,torch.bfloat16,torch.bfloat16,False,False,torch,0,0,19.4378,native,0.0,386.68,1186.79
+gfx950,256,1024,1024,7168,False,torch.bfloat16,torch.bfloat16,False,False,torch,0,0,28.2549,native,0.0,532.03,1113.34
+gfx950,256,2048,1024,7168,False,torch.bfloat16,torch.bfloat16,False,False,torch,0,0,43.3843,native,0.0,692.99,1111.8
+gfx950,256,4096,1024,7168,False,torch.bfloat16,torch.bfloat16,False,False,torch,0,0,69.9934,native,0.0,859.07,1168.52
+gfx950,256,8192,1024,7168,False,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,112.7983,auto,0.0,1066.14,1320.04
+gfx950,256,16384,1024,7168,False,torch.bfloat16,torch.bfloat16,False,False,torch,0,0,168.2342,native,0.0,1429.66,1682.87
+gfx950,256,32,2112,7168,False,torch.bfloat16,torch.bfloat16,False,False,flydsl,1073,7,10.8391,flydsl_hgemm_abf16_wbf16_bf16_t32x64x64x5_ks7_w1x4x1_bias0_ktail0_gm0_pft_gfx950,0.0202,89.39,2848.17
+gfx950,256,64,2112,7168,False,torch.bfloat16,torch.bfloat16,False,False,flydsl,2332,7,14.6748,flydsl_hgemm_abf16_wbf16_bf16_t64x64x64x6_ks7_w1x4x1_bias0_ktail0_gm0_pft_gfx950,0.0194,132.05,2144.18
+gfx950,256,256,2112,7168,False,torch.bfloat16,torch.bfloat16,False,False,asm,8,2,28.2104,_ZN5aiter39bf16gemm_fp32bf16_tn_96x64_splitk_cleanE,0.006,274.76,1241.7
+gfx950,256,1024,2112,7168,False,torch.bfloat16,torch.bfloat16,False,False,torch,0,0,51.2476,native,0.0,604.99,961.67
+gfx950,256,4096,2112,7168,False,torch.bfloat16,torch.bfloat16,False,False,torch,0,0,117.5728,native,0.0,1054.81,904.12
+gfx950,256,16384,2112,7168,False,torch.bfloat16,torch.bfloat16,False,False,flydsl,371,1,480.0168,flydsl_hgemm_abf16_wbf16_bf16_t256x128x64x3_ks1_w4x4x1_bias0_ktail0_gm4_pft_gfx950,0.0,1033.44,696.57
+gfx950,256,256,3072,1536,False,torch.bfloat16,torch.bfloat16,False,False,flydsl,1886,1,9.5581,flydsl_hgemm_abf16_wbf16_bf16_t64x64x128x4_ks1_w2x4x1_bias0_ktail0_gm0_pft_gfx950,0.0,252.76,1234.19
+gfx950,256,1024,3072,1536,False,torch.bfloat16,torch.bfloat16,False,False,flydsl,675,1,17.7036,flydsl_hgemm_abf16_wbf16_bf16_t128x96x64x5_ks1_w2x2x1_bias0_ktail0_gm4_pft_gfx950,0.0,545.86,1066.13
+gfx950,256,4096,3072,1536,False,torch.bfloat16,torch.bfloat16,False,False,torch,0,0,41.5199,native,0.0,930.99,1136.47
+gfx950,256,16384,3072,1536,False,torch.bfloat16,torch.bfloat16,False,False,flydsl,393,1,129.8909,flydsl_hgemm_abf16_wbf16_bf16_t256x256x64x2_ks1_w2x4x1_bias0_ktail0_gm4_pht_gfx950,0.0,1190.37,1235.13
+gfx950,256,256,4096,512,False,torch.bfloat16,torch.bfloat16,False,False,flydsl,1267,1,5.6431,flydsl_hgemm_abf16_wbf16_bf16_t64x64x128x4_ks1_w2x4x1_bias0_ktail0_gm0_pft_gfx950,0.0,190.28,1161.35
+gfx950,256,1024,4096,512,False,torch.bfloat16,torch.bfloat16,False,False,flydsl,131,1,10.0214,flydsl_hgemm_abf16_wbf16_bf16_t128x128x64x3_ks1_w2x4x1_bias0_ktail0_gm4_pft_gfx950,0.0,428.58,1360.24
+gfx950,256,4096,4096,512,False,torch.bfloat16,torch.bfloat16,False,False,torch,0,0,20.8466,native,0.0,824.11,2011.98
+gfx950,256,16384,4096,512,False,torch.bfloat16,torch.bfloat16,False,False,torch,0,0,78.7947,native,0.0,872.13,1969.54
+gfx950,256,96,7168,512,False,torch.bfloat16,torch.bfloat16,False,False,flydsl,291,1,5.6879,flydsl_hgemm_abf16_wbf16_bf16_t32x64x128x3_ks1_w1x2x1_bias0_ktail0_gm0_pft_gfx950,0.0,123.88,1549.71
+gfx950,256,128,7168,512,False,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,5.4031,auto,0.0,173.89,1722.37
+gfx950,256,256,7168,512,False,torch.bfloat16,torch.bfloat16,False,False,flydsl,326,1,7.1982,flydsl_hgemm_abf16_wbf16_bf16_t128x64x64x5_ks1_w4x2x1_bias0_ktail0_gm0_pft_gfx950,0.0,261.04,1565.97
+gfx950,256,1024,7168,512,False,torch.bfloat16,torch.bfloat16,False,False,torch,0,0,14.1562,native,0.0,530.95,1629.58
+gfx950,256,4096,7168,512,False,torch.bfloat16,torch.bfloat16,False,False,asm,10,1,39.1359,_ZN5aiter24bf16gemm_bf16_tn_256x256E,0.0,768.21,1795.14
+gfx950,256,16384,7168,512,False,torch.bfloat16,torch.bfloat16,False,False,torch,0,0,132.1562,native,0.0,909.98,1959.79
+gfx950,256,8192,8192,8192,False,torch.bfloat16,torch.bfloat16,False,False,asm,10,1,695.6068,_ZN5aiter24bf16gemm_bf16_tn_256x256E,0.0,1580.65,578.85
 gfx942,80,128,7168,512,False,torch.bfloat16,torch.bfloat16,False,False,opus,10000,0,14.1933,opus_gemm_gfx942_512x128x128x64_2x4_16x16x16_0x0x0,0.0,66.19,655.67
 gfx942,80,256,7168,512,False,torch.bfloat16,torch.bfloat16,False,False,opus,10000,0,27.0618,opus_gemm_gfx942_512x128x128x64_2x4_16x16x16_0x0x0,0.0,69.44,416.54
 gfx942,80,1024,7168,512,False,torch.bfloat16,torch.bfloat16,False,False,opus,10000,0,74.6028,opus_gemm_gfx942_512x128x128x64_2x4_16x16x16_0x0x0,0.0,100.75,309.22