diff --git a/aiter/aot/flydsl/gemm.py b/aiter/aot/flydsl/gemm.py index ad38d95d88..f1759a8168 100755 --- a/aiter/aot/flydsl/gemm.py +++ b/aiter/aot/flydsl/gemm.py @@ -97,7 +97,10 @@ r"(?P\d+)x(?P\d+)x(?P\d+)_" r"(?P[A-Z0-9]+)_(?P[A-Z0-9]+)_(?P[A-Z0-9]+)_" r"(?P\d+)x(?P\d+)(?:x(?P\d+))?(?:x(?P\d+))?_" - r"(?P[A-Za-z][A-Za-z0-9]*)$" + r"(?!ks\d+$)(?P[A-Za-z][A-Za-z0-9]*)" + # Trailing _ksN, emitted only for k_split > 1, so pre-split-K names still + # match. Without it they fail fullmatch and drop out of the AOT build. + r"(?:_ks(?P\d+))?$" ) _SHORT_DTYPE = { "F8": "fp8", @@ -147,6 +150,7 @@ def _parse_preshuffle_kernel_name(name: str) -> dict | None: "xcd_swizzle": int(m.group("xcd_swizzle")) if m.group("xcd_swizzle") else 0, "lds_stage": int(m.group("lds_stage")) if m.group("lds_stage") else 2, "scheduler": m.group("scheduler"), + "k_split": int(m.group("k_split")) if m.group("k_split") else 1, } @@ -354,10 +358,12 @@ def _compile_preshuffle_to_cache( xcd_swizzle: int = 0, lds_stage: int = 2, scheduler: str = "Default", + k_split: int = 1, **kwargs, ): del kwargs enable_scheduler = str(scheduler).lower() != "off" + k_split = int(k_split) import torch @@ -367,7 +373,23 @@ def _compile_preshuffle_to_cache( # FlyDSL preshuffle kernels consume raw quantized bytes for fp8/int8 paths. a = torch.empty((m * k,), device=dev, dtype=torch.int8) b = torch.empty((n * k,), device=dev, dtype=torch.int8) + from aiter.ops.flydsl.gemm_kernels import ( + PRESHUFFLE_SPLIT_K_MAX_TILES, + PRESHUFFLE_SPLIT_K_WORKSPACE_ELEMS, + ) + + # Sized to the same bounds the runtime uses, so the signatures match. out = torch.empty((m * n,), device=dev, dtype=out_torch_dtype) + workspace = ( + torch.empty(PRESHUFFLE_SPLIT_K_WORKSPACE_ELEMS, device=dev, dtype=torch.float32) + if k_split > 1 + else out + ) + semaphore = torch.zeros( + PRESHUFFLE_SPLIT_K_MAX_TILES if k_split > 1 else 0, + device=dev, + dtype=torch.int32, + ) scale_a = torch.empty((max(m, 1),), device=dev, dtype=torch.float32) scale_b = torch.empty((max(n, 1),), device=dev, dtype=torch.float32) bias = torch.empty(0, device=dev, dtype=out_torch_dtype) @@ -386,13 +408,16 @@ def _compile_preshuffle_to_cache( enable_scheduler=enable_scheduler, xcd_swizzle=xcd_swizzle, lds_stage=lds_stage, + split_k=k_split, ) # The layout-API launcher uses fx.Tensor args (it builds views via # fx.get_iter/make_view), so pass flat torch tensors directly rather # than raw pointers (pointer args would fail GetIterOp type checks). _compile_executable_to_cache( exe, + workspace, out, + semaphore, a, b, scale_a, diff --git a/aiter/configs/model_configs/a8w8_bpreshuffle_tuned_gemm_kimik3.csv b/aiter/configs/model_configs/a8w8_bpreshuffle_tuned_gemm_kimik3.csv index 03e49b2d7f..5e322c111e 100644 --- a/aiter/configs/model_configs/a8w8_bpreshuffle_tuned_gemm_kimik3.csv +++ b/aiter/configs/model_configs/a8w8_bpreshuffle_tuned_gemm_kimik3.csv @@ -1,10 +1,10 @@ gfx,cu_num,M,N,K,q_dtype_w,libtype,kernelId,splitK,us,kernelName,tflops,bw,errRatio -gfx950,256,1,576,7168,torch.float8_e4m3fn,flydsl,617,0,11.3863,flydsl_bpreshuflle_16x64x512_F8_F8_B16_1x2x0x2_default,0.73,363.34,0.0 -gfx950,256,2,576,7168,torch.float8_e4m3fn,flydsl,373,0,11.4327,flydsl_bpreshuflle_16x64x512_F8_F8_B16_1x1x0x2_default,1.44,362.59,0.0 -gfx950,256,4,576,7168,torch.float8_e4m3fn,flydsl,1059,0,11.4331,flydsl_bpreshuflle_16x64x512_F8_F8_B16_1x4x0x2_default,2.89,364.03,0.0 -gfx950,256,8,576,7168,torch.float8_e4m3fn,flydsl,1,0,11.3459,flydsl_bpreshuflle_16x64x512_F8_F8_B16_0x0x0x2_default,5.82,369.77,0.0 -gfx950,256,16,576,7168,torch.float8_e4m3fn,ck,10,0,10.7183,a8w8_bpreshuffle_128x16x32x512_16x16_16x16_32x4x1_32x4x1_1x16x1x8_4x4x1_1x1_intrawave_v1,12.33,397.63,0.0 -gfx950,256,32,576,7168,torch.float8_e4m3fn,ck,8,0,10.3684,a8w8_bpreshuffle_128x32x16x512_16x16_16x16_32x4x1_32x4x1_1x32x1x4_4x4x1_1x1_intrawave_v1,25.49,423.88,0.0 +gfx950,256,1,576,7168,torch.float8_e4m3fn,flydsl,961,7,5.0341,flydsl_bpreshuflle_16x64x512_F8_F8_B16_0x4x0x2_default_ks7,1.64,821.81,0.0 +gfx950,256,2,576,7168,torch.float8_e4m3fn,flydsl,961,7,5.0499,flydsl_bpreshuflle_16x64x512_F8_F8_B16_0x4x0x2_default_ks7,3.27,820.89,0.0 +gfx950,256,4,576,7168,torch.float8_e4m3fn,flydsl,2116,14,5.5256,flydsl_bpreshuflle_16x64x256_F8_F8_B16_0x4x0x1_default_ks14,5.98,753.23,0.0 +gfx950,256,8,576,7168,torch.float8_e4m3fn,flydsl,373,7,5.2815,flydsl_bpreshuflle_16x64x512_F8_F8_B16_1x1x0x2_default_ks7,12.51,794.34,0.0 +gfx950,256,16,576,7168,torch.float8_e4m3fn,flydsl,373,7,5.4885,flydsl_bpreshuflle_16x64x512_F8_F8_B16_1x1x0x2_default_ks7,24.07,776.51,0.0 +gfx950,256,32,576,7168,torch.float8_e4m3fn,flydsl,961,7,6.0091,flydsl_bpreshuflle_16x64x512_F8_F8_B16_0x4x0x2_default_ks7,43.97,731.39,0.0 gfx950,256,64,576,7168,torch.float8_e4m3fn,ck,10,0,10.7167,a8w8_bpreshuffle_128x16x32x512_16x16_16x16_32x4x1_32x4x1_1x16x1x8_4x4x1_1x1_intrawave_v1,49.31,434.95,0.0 gfx950,256,128,576,7168,torch.float8_e4m3fn,ck,10,0,9.9632,a8w8_bpreshuffle_128x16x32x512_16x16_16x16_32x4x1_32x4x1_1x16x1x8_4x4x1_1x1_intrawave_v1,106.09,521.29,0.0 gfx950,256,256,576,7168,torch.float8_e4m3fn,ck,11,0,11.4055,a8w8_bpreshuffle_256x16x64x512_16x16_16x16_32x8x1_32x8x1_1x16x1x16_4x4x1_1x1_intrawave_v1,185.34,548.74,0.0 @@ -16,17 +16,17 @@ gfx950,256,8192,576,7168,torch.float8_e4m3fn,cktile,94,0,48.2593,a8w8_bpreshuffl gfx950,256,16384,576,7168,torch.float8_e4m3fn,cktile,232,0,75.7943,a8w8_bpreshuffle_cktile_0x0x8x4x1x0x0x0x0x3_224x192x128_1x4x1_16x16x128_default,1784.98,1852.96,0.0 gfx950,256,32768,576,7168,torch.float8_e4m3fn,ck,68,0,137.1244,a8w8_bpreshuffle_256x128x192x128_16x16_16x16_8x32x1_8x32x1_1x32x1x8_8x8x1_2x1_intrawave_v3,1973.27,2018.3,0.0 gfx950,256,1,1536,128,torch.float8_e4m3fn,cktile,15,0,2.8909,a8w8_bpreshuffle_cktile_0x0x8x4x1x0x0x0x0x1_32x64x128_1x4x1_16x16x128_default,0.14,69.12,0.0 -gfx950,256,1,1536,7168,torch.float8_e4m3fn,flydsl,617,0,11.5477,flydsl_bpreshuflle_16x64x512_F8_F8_B16_1x2x0x2_default,1.91,954.33,0.0 +gfx950,256,1,1536,7168,torch.float8_e4m3fn,flydsl,187,7,5.8969,flydsl_bpreshuflle_16x64x512_F8_F8_B16_1x0x4x2_default_ks7,3.73,1868.83,0.0 gfx950,256,2,1536,128,torch.float8_e4m3fn,cktile,165,0,2.7639,a8w8_bpreshuffle_cktile_0x0x8x4x1x0x0x0x0x4_32x64x128_1x4x1_16x16x128_default,0.28,73.45,0.0 -gfx950,256,2,1536,7168,torch.float8_e4m3fn,flydsl,617,0,11.1346,flydsl_bpreshuflle_16x64x512_F8_F8_B16_1x2x0x2_default,3.96,990.65,0.0 +gfx950,256,2,1536,7168,torch.float8_e4m3fn,flydsl,2117,7,6.0302,flydsl_bpreshuflle_16x64x512_F8_F8_B16_0x4x0x1_default_ks7,7.3,1829.21,0.0 gfx950,256,4,1536,128,torch.float8_e4m3fn,cktile,15,0,2.9594,a8w8_bpreshuffle_cktile_0x0x8x4x1x0x0x0x0x1_32x64x128_1x4x1_16x16x128_default,0.53,70.76,0.0 -gfx950,256,4,1536,7168,torch.float8_e4m3fn,flydsl,617,0,11.2096,flydsl_bpreshuflle_16x64x512_F8_F8_B16_1x2x0x2_default,7.86,985.85,0.0 +gfx950,256,4,1536,7168,torch.float8_e4m3fn,flydsl,435,7,5.9077,flydsl_bpreshuflle_16x64x512_F8_F8_B16_1x1x4x2_default_ks7,14.91,1870.61,0.0 gfx950,256,8,1536,128,torch.float8_e4m3fn,cktile,165,0,2.6116,a8w8_bpreshuffle_cktile_0x0x8x4x1x0x0x0x0x4_32x64x128_1x4x1_16x16x128_default,1.2,85.09,0.0 -gfx950,256,8,1536,7168,torch.float8_e4m3fn,flydsl,737,0,11.9393,flydsl_bpreshuflle_16x64x512_F8_F8_B16_0x3x0x2_default,14.75,929.03,0.0 +gfx950,256,8,1536,7168,torch.float8_e4m3fn,flydsl,677,7,6.0842,flydsl_bpreshuflle_16x64x512_F8_F8_B16_1x2x4x2_default_ks7,28.95,1823.08,0.0 gfx950,256,16,1536,128,torch.float8_e4m3fn,cktile,15,0,3.4698,a8w8_bpreshuffle_cktile_0x0x8x4x1x0x0x0x0x1_32x64x128_1x4x1_16x16x128_default,1.81,71.42,0.0 -gfx950,256,16,1536,7168,torch.float8_e4m3fn,ck,10,0,10.3683,a8w8_bpreshuffle_128x16x32x512_16x16_16x16_32x4x1_32x4x1_1x16x1x8_4x4x1_1x1_intrawave_v1,33.98,1077.7,0.0 +gfx950,256,16,1536,7168,torch.float8_e4m3fn,flydsl,1059,7,6.2957,flydsl_bpreshuflle_16x64x512_F8_F8_B16_1x4x0x2_default_ks7,55.96,1774.84,0.0 gfx950,256,32,1536,128,torch.float8_e4m3fn,flydsl,1412,0,2.7372,flydsl_bpreshuflle_32x64x128_F8_F8_B16_0x1x0x1_default,4.6,109.24,0.0 -gfx950,256,32,1536,7168,torch.float8_e4m3fn,ck,8,0,11.2815,a8w8_bpreshuffle_128x32x16x512_16x16_16x16_32x4x1_32x4x1_1x32x1x4_4x4x1_1x1_intrawave_v1,62.46,1004.98,0.0 +gfx950,256,32,1536,7168,torch.float8_e4m3fn,flydsl,258,7,7.2576,flydsl_bpreshuflle_32x64x512_F8_F8_B16_0x1x0x2_default_ks7,97.09,1562.19,0.0 gfx950,256,64,1536,128,torch.float8_e4m3fn,flydsl,1660,0,3.3145,flydsl_bpreshuflle_32x64x128_F8_F8_B16_0x2x0x1_default,7.59,121.11,0.0 gfx950,256,64,1536,7168,torch.float8_e4m3fn,ck,10,0,11.1903,a8w8_bpreshuffle_128x16x32x512_16x16_16x16_32x4x1_32x4x1_1x16x1x8_4x4x1_1x1_intrawave_v1,125.94,1042.46,0.0 gfx950,256,128,1536,128,torch.float8_e4m3fn,flydsl,1900,0,2.7756,flydsl_bpreshuflle_32x64x128_F8_F8_B16_0x3x0x1_default,18.13,218.41,0.0 @@ -47,12 +47,12 @@ gfx950,256,16384,1536,128,torch.float8_e4m3fn,flydsl,931,0,15.3636,flydsl_bpresh gfx950,256,16384,1536,7168,torch.float8_e4m3fn,flydsl,667,0,156.303,flydsl_bpreshuflle_128x256x128_F8_F8_B16_1x2x0x2_default,2308.19,1143.82,0.0 gfx950,256,32768,1536,128,torch.float8_e4m3fn,flydsl,213,0,25.7514,flydsl_bpreshuflle_64x256x128_F8_F8_B16_1x0x4x2_default,500.36,4079.55,0.0 gfx950,256,32768,1536,7168,torch.float8_e4m3fn,ck,33,0,288.0491,a8w8_bpreshuffle_256x256x256x128_16x16_16x16_8x32x1_8x32x1_1x32x1x8_8x8x1_2x1_intrawave_v3,2504.97,1203.11,0.0 -gfx950,256,1,2176,7168,torch.float8_e4m3fn,flydsl,849,0,11.4632,flydsl_bpreshuflle_16x64x512_F8_F8_B16_1x3x0x2_default,2.72,1361.67,0.0 -gfx950,256,2,2176,7168,torch.float8_e4m3fn,flydsl,125,0,11.0924,flydsl_bpreshuflle_16x64x512_F8_F8_B16_1x0x0x2_default,5.62,1408.23,0.0 -gfx950,256,4,2176,7168,torch.float8_e4m3fn,flydsl,373,0,11.1295,flydsl_bpreshuflle_16x64x512_F8_F8_B16_1x1x0x2_default,11.21,1405.6,0.0 -gfx950,256,8,2176,7168,torch.float8_e4m3fn,flydsl,617,0,11.9394,flydsl_bpreshuflle_16x64x512_F8_F8_B16_1x2x0x2_default,20.9,1314.11,0.0 -gfx950,256,16,2176,7168,torch.float8_e4m3fn,ck,10,0,11.2181,a8w8_bpreshuffle_128x16x32x512_16x16_16x16_32x4x1_32x4x1_1x16x1x8_4x4x1_1x1_intrawave_v1,44.49,1406.82,0.0 -gfx950,256,32,2176,7168,torch.float8_e4m3fn,ck,10,0,11.1468,a8w8_bpreshuffle_128x16x32x512_16x16_16x16_32x4x1_32x4x1_1x16x1x8_4x4x1_1x1_intrawave_v1,89.55,1432.36,0.0 +gfx950,256,1,2176,7168,torch.float8_e4m3fn,flydsl,249,7,6.3084,flydsl_bpreshuflle_16x64x512_F8_F8_B16_0x1x0x2_default_ks7,4.95,2474.33,0.0 +gfx950,256,2,2176,7168,torch.float8_e4m3fn,flydsl,125,7,6.5415,flydsl_bpreshuflle_16x64x512_F8_F8_B16_1x0x0x2_default_ks7,9.54,2387.92,0.0 +gfx950,256,4,2176,7168,torch.float8_e4m3fn,flydsl,497,7,6.3485,flydsl_bpreshuflle_16x64x512_F8_F8_B16_0x2x0x2_default_ks7,19.66,2464.15,0.0 +gfx950,256,8,2176,7168,torch.float8_e4m3fn,flydsl,1,7,6.4615,flydsl_bpreshuflle_16x64x512_F8_F8_B16_0x0x0x2_default_ks7,38.62,2428.19,0.0 +gfx950,256,16,2176,7168,torch.float8_e4m3fn,flydsl,961,7,6.8142,flydsl_bpreshuflle_16x64x512_F8_F8_B16_0x4x0x2_default_ks7,73.25,2316.03,0.0 +gfx950,256,32,2176,7168,torch.float8_e4m3fn,flydsl,1068,7,7.7552,flydsl_bpreshuflle_32x64x512_F8_F8_B16_1x4x0x2_default_ks7,128.72,2058.77,0.0 gfx950,256,64,2176,7168,torch.float8_e4m3fn,ck,10,0,11.4336,a8w8_bpreshuffle_128x16x32x512_16x16_16x16_32x4x1_32x4x1_1x16x1x8_4x4x1_1x1_intrawave_v1,174.62,1428.67,0.0 gfx950,256,128,2176,7168,torch.float8_e4m3fn,flydsl,1713,0,12.5177,flydsl_bpreshuflle_16x64x512_F8_F8_B16_0x2x4x1_default,318.99,1363.84,0.0 gfx950,256,256,2176,7168,torch.float8_e4m3fn,flydsl,2175,0,16.0682,flydsl_bpreshuflle_32x64x512_F8_F8_B16_0x4x4x1_default,497.0,1154.25,0.0 @@ -63,12 +63,12 @@ gfx950,256,4096,2176,7168,torch.float8_e4m3fn,flydsl,1000029,0,73.9852,flydsl_bp gfx950,256,8192,2176,7168,torch.float8_e4m3fn,flydsl,1103,0,134.9245,flydsl_bpreshuflle_128x128x256_F8_F8_B16_1x4x0x2_default,1894.03,815.04,0.0 gfx950,256,16384,2176,7168,torch.float8_e4m3fn,flydsl,1000028,0,222.521,flydsl_bpreshuffle_8w_256x256x128_F8_F8_B16_2x4,2296.87,918.3,0.0 gfx950,256,32768,2176,7168,torch.float8_e4m3fn,flydsl,1000029,0,398.7208,flydsl_bpreshuffle_8w_256x256x128_F8_F8_B16_2x8,2563.7,985.87,0.0 -gfx950,256,1,2304,1536,torch.float8_e4m3fn,flydsl,497,0,4.6403,flydsl_bpreshuflle_16x64x512_F8_F8_B16_0x2x0x2_default,1.53,763.98,0.0 -gfx950,256,2,2304,1536,torch.float8_e4m3fn,flydsl,1,0,4.681,flydsl_bpreshuflle_16x64x512_F8_F8_B16_0x0x0x2_default,3.02,758.65,0.0 -gfx950,256,4,2304,1536,torch.float8_e4m3fn,flydsl,249,0,5.1918,flydsl_bpreshuflle_16x64x512_F8_F8_B16_0x1x0x2_default,5.45,686.37,0.0 -gfx950,256,8,2304,1536,torch.float8_e4m3fn,flydsl,737,0,4.8189,flydsl_bpreshuflle_16x64x512_F8_F8_B16_0x3x0x2_default,11.75,744.59,0.0 +gfx950,256,1,2304,1536,torch.float8_e4m3fn,flydsl,849,0,4.5526,flydsl_bpreshuflle_16x64x512_F8_F8_B16_1x3x0x2_default,1.55,778.7,0.0 +gfx950,256,2,2304,1536,torch.float8_e4m3fn,flydsl,1059,0,4.5926,flydsl_bpreshuflle_16x64x512_F8_F8_B16_1x4x0x2_default,3.08,773.25,0.0 +gfx950,256,4,2304,1536,torch.float8_e4m3fn,flydsl,1059,0,5.0649,flydsl_bpreshuflle_16x64x512_F8_F8_B16_1x4x0x2_default,5.59,703.57,0.0 +gfx950,256,8,2304,1536,torch.float8_e4m3fn,flydsl,1059,0,4.6395,flydsl_bpreshuflle_16x64x512_F8_F8_B16_1x4x0x2_default,12.2,773.38,0.0 gfx950,256,16,2304,1536,torch.float8_e4m3fn,flydsl,737,0,5.2365,flydsl_bpreshuflle_16x64x512_F8_F8_B16_0x3x0x2_default,21.63,694.6,0.0 -gfx950,256,32,2304,1536,torch.float8_e4m3fn,flydsl,249,0,4.642,flydsl_bpreshuflle_16x64x512_F8_F8_B16_0x1x0x2_default,48.79,804.73,0.0 +gfx950,256,32,2304,1536,torch.float8_e4m3fn,flydsl,1059,0,5.0118,flydsl_bpreshuflle_16x64x512_F8_F8_B16_1x4x0x2_default,45.19,745.35,0.0 gfx950,256,64,2304,1536,torch.float8_e4m3fn,flydsl,249,0,4.8142,flydsl_bpreshuflle_16x64x512_F8_F8_B16_0x1x0x2_default,94.09,816.78,0.0 gfx950,256,128,2304,1536,torch.float8_e4m3fn,flydsl,746,0,5.1826,flydsl_bpreshuflle_32x64x512_F8_F8_B16_0x3x0x2_default,174.81,834.6,0.0 gfx950,256,256,2304,1536,torch.float8_e4m3fn,flydsl,802,0,6.3352,flydsl_bpreshuflle_32x64x512_F8_F8_B16_0x3x4x2_default,286.01,806.89,0.0 @@ -99,7 +99,7 @@ gfx950,256,1,6400,7168,torch.float8_e4m3fn,flydsl,1,0,12.163,flydsl_bpreshuflle_ gfx950,256,2,6400,7168,torch.float8_e4m3fn,flydsl,617,0,12.7111,flydsl_bpreshuflle_16x64x512_F8_F8_B16_1x2x0x2_default,14.44,3612.21,0.0 gfx950,256,4,6400,7168,torch.float8_e4m3fn,flydsl,737,0,12.7921,flydsl_bpreshuflle_16x64x512_F8_F8_B16_0x3x0x2_default,28.69,3592.46,0.0 gfx950,256,8,6400,7168,torch.float8_e4m3fn,flydsl,1,0,12.526,flydsl_bpreshuflle_16x64x512_F8_F8_B16_0x0x0x2_default,58.6,3675.15,0.0 -gfx950,256,16,6400,7168,torch.float8_e4m3fn,ck,10,0,11.7238,a8w8_bpreshuffle_128x16x32x512_16x16_16x16_32x4x1_32x4x1_1x16x1x8_4x4x1_1x1_intrawave_v1,125.22,3940.25,0.0 +gfx950,256,16,6400,7168,torch.float8_e4m3fn,flydsl,249,2,11.2092,flydsl_bpreshuflle_16x64x512_F8_F8_B16_0x1x0x2_default_ks2,130.96,4121.14,0.0 gfx950,256,32,6400,7168,torch.float8_e4m3fn,flydsl,311,0,13.2236,flydsl_bpreshuflle_16x64x512_F8_F8_B16_0x1x4x2_default,222.03,3517.51,0.0 gfx950,256,64,6400,7168,torch.float8_e4m3fn,flydsl,1117,0,14.7136,flydsl_bpreshuflle_32x64x512_F8_F8_B16_1x4x4x2_default,399.09,3204.73,0.0 gfx950,256,128,6400,7168,torch.float8_e4m3fn,ck,114,0,21.0686,a8w8_bpreshuffle_256x64x64x256_16x16_16x16_16x16x1_16x16x1_1x32x1x8_8x8x1_2x1_intrawave_v3,557.42,2298.73,0.0 @@ -112,7 +112,7 @@ gfx950,256,8192,6400,7168,torch.float8_e4m3fn,flydsl,1000026,0,294.7446,flydsl_b gfx950,256,16384,6400,7168,torch.float8_e4m3fn,flydsl,1000025,0,542.5117,flydsl_bpreshuffle_8w_256x256x128_F8_F8_B16_1x4,2770.89,687.6,0.0 gfx950,256,32768,6400,7168,torch.float8_e4m3fn,flydsl,1000028,0,1050.5652,flydsl_bpreshuffle_8w_256x256x128_F8_F8_B16_2x4,2861.77,666.49,0.0 gfx950,256,1,7168,768,torch.float8_e4m3fn,flydsl,1652,0,4.3118,flydsl_bpreshuflle_16x64x256_F8_F8_B16_0x2x0x1_default,2.55,1280.24,0.0 -gfx950,256,1,7168,1536,torch.float8_e4m3fn,flydsl,1,0,5.4868,flydsl_bpreshuflle_16x64x512_F8_F8_B16_0x0x0x2_default,4.01,2009.54,0.0 +gfx950,256,1,7168,1536,torch.float8_e4m3fn,flydsl,849,0,5.005,flydsl_bpreshuflle_16x64x512_F8_F8_B16_1x3x0x2_default,4.4,2202.98,0.0 gfx950,256,1,7168,4224,torch.float8_e4m3fn,ck,9,0,17.7187,a8w8_bpreshuffle_128x16x32x128_16x16_16x16_8x16x1_8x16x1_1x16x1x8_4x4x1_1x1_intrawave_v1,3.42,1709.84,0.0 gfx950,256,2,7168,768,torch.float8_e4m3fn,flydsl,1404,0,4.3292,flydsl_bpreshuflle_16x64x256_F8_F8_B16_0x1x0x1_default,5.09,1278.58,0.0 gfx950,256,2,7168,1536,torch.float8_e4m3fn,flydsl,497,0,5.1133,flydsl_bpreshuflle_16x64x512_F8_F8_B16_0x2x0x2_default,8.61,2159.43,0.0 @@ -124,11 +124,11 @@ gfx950,256,8,7168,768,torch.float8_e4m3fn,flydsl,1652,0,4.4118,flydsl_bpreshufll gfx950,256,8,7168,1536,torch.float8_e4m3fn,flydsl,961,0,5.0859,flydsl_bpreshuflle_16x64x512_F8_F8_B16_0x4x0x2_default,34.64,2189.78,0.0 gfx950,256,8,7168,4224,torch.float8_e4m3fn,ck,9,0,17.5696,a8w8_bpreshuffle_128x16x32x128_16x16_16x16_8x16x1_8x16x1_1x16x1x8_4x4x1_1x1_intrawave_v1,27.57,1731.75,0.0 gfx950,256,16,7168,768,torch.float8_e4m3fn,flydsl,1156,0,4.9038,flydsl_bpreshuflle_16x64x256_F8_F8_B16_0x0x0x1_default,35.92,1171.88,0.0 -gfx950,256,16,7168,1536,torch.float8_e4m3fn,flydsl,1,0,5.237,flydsl_bpreshuflle_16x64x512_F8_F8_B16_0x0x0x2_default,67.28,2150.85,0.0 +gfx950,256,16,7168,1536,torch.float8_e4m3fn,flydsl,849,0,5.8241,flydsl_bpreshuflle_16x64x512_F8_F8_B16_1x3x0x2_default,60.49,1934.03,0.0 gfx950,256,16,7168,4224,torch.float8_e4m3fn,ck,9,0,17.963,a8w8_bpreshuffle_128x16x32x128_16x16_16x16_8x16x1_8x16x1_1x16x1x8_4x4x1_1x1_intrawave_v1,53.94,1702.09,0.0 gfx950,256,32,7168,768,torch.float8_e4m3fn,flydsl,257,0,4.4742,flydsl_bpreshuflle_32x64x256_F8_F8_B16_0x1x0x2_default,78.75,1338.42,0.0 -gfx950,256,32,7168,1536,torch.float8_e4m3fn,flydsl,1662,0,5.5894,flydsl_bpreshuflle_32x64x512_F8_F8_B16_0x2x0x1_default,126.07,2060.68,0.0 -gfx950,256,32,7168,4224,torch.float8_e4m3fn,ck,9,0,18.8765,a8w8_bpreshuffle_128x16x32x128_16x16_16x16_8x16x1_8x16x1_1x16x1x8_4x4x1_1x1_intrawave_v1,102.66,1635.45,0.0 +gfx950,256,32,7168,1536,torch.float8_e4m3fn,flydsl,677,0,5.9228,flydsl_bpreshuflle_16x64x512_F8_F8_B16_1x2x4x2_default,118.97,1944.68,0.0 +gfx950,256,32,7168,4224,torch.float8_e4m3fn,flydsl,624,3,11.1259,flydsl_bpreshuflle_32x64x128_F8_F8_B16_1x2x0x2_default_ks3,174.17,2774.75,0.0 gfx950,256,64,7168,768,torch.float8_e4m3fn,flydsl,969,0,4.4907,flydsl_bpreshuflle_32x64x256_F8_F8_B16_0x4x0x2_default,156.91,1441.13,0.0 gfx950,256,64,7168,1536,torch.float8_e4m3fn,flydsl,802,0,5.8575,flydsl_bpreshuflle_32x64x512_F8_F8_B16_0x3x4x2_default,240.6,2053.07,0.0 gfx950,256,64,7168,4224,torch.float8_e4m3fn,ck,76,0,19.4931,a8w8_bpreshuffle_256x32x64x128_16x16_16x16_8x32x1_8x32x1_1x32x1x8_8x8x1_2x1_intrawave_v1,198.82,1614.19,0.0 @@ -164,7 +164,7 @@ gfx950,256,2,8448,7168,torch.float8_e4m3fn,flydsl,497,0,13.5444,flydsl_bpreshufl gfx950,256,4,8448,7168,torch.float8_e4m3fn,flydsl,1,0,13.541,flydsl_bpreshuflle_16x64x512_F8_F8_B16_0x0x0x2_default,35.78,4479.1,0.0 gfx950,256,8,8448,7168,torch.float8_e4m3fn,flydsl,1,0,13.2129,flydsl_bpreshuflle_16x64x512_F8_F8_B16_0x0x0x2_default,73.33,4597.61,0.0 gfx950,256,16,8448,7168,torch.float8_e4m3fn,flydsl,1,0,13.324,flydsl_bpreshuflle_16x64x512_F8_F8_B16_0x0x0x2_default,145.43,4573.72,0.0 -gfx950,256,32,8448,7168,torch.float8_e4m3fn,flydsl,1219,0,15.304,flydsl_bpreshuflle_16x64x512_F8_F8_B16_0x0x4x1_default,253.24,4007.14,0.0 +gfx950,256,32,8448,7168,torch.float8_e4m3fn,flydsl,10,0,15.3828,flydsl_bpreshuflle_32x64x512_F8_F8_B16_0x0x0x2_default,251.94,3986.62,0.0 gfx950,256,64,8448,7168,torch.float8_e4m3fn,flydsl,1228,0,18.7312,flydsl_bpreshuflle_32x64x512_F8_F8_B16_0x0x4x1_default,413.81,3315.08,0.0 gfx950,256,128,8448,7168,torch.float8_e4m3fn,flydsl,1240,0,24.2685,flydsl_bpreshuflle_64x64x256_F8_F8_B16_0x0x4x1_default,638.78,2622.14,0.0 gfx950,256,256,8448,7168,torch.float8_e4m3fn,flydsl,88,0,33.9665,flydsl_bpreshuflle_64x192x256_F8_F8_B16_0x0x4x2_default,912.79,1964.16,0.0 diff --git a/aiter/ops/flydsl/gemm_kernels.py b/aiter/ops/flydsl/gemm_kernels.py index 6a81d1da6f..afbf052606 100644 --- a/aiter/ops/flydsl/gemm_kernels.py +++ b/aiter/ops/flydsl/gemm_kernels.py @@ -971,6 +971,54 @@ def _get_compile_fn(): return _flydsl_compile_fn +# Fixed size rather than one buffer per shape: a shape-keyed cache grows without +# limit and can evict a buffer a captured CUDA graph still points at. The bounds +# come from k_split_candidates, which keeps tile_count under CU_NUM and +# k_split * tile_count at four per CU. +# Mirrors preshuffle_gemm.PRESHUFFLE_M_MAX; duplicated so this module imports +# without FlyDSL present. +PRESHUFFLE_M_MAX = 65536 + +PRESHUFFLE_SPLIT_K_MAX_TILES = 256 +PRESHUFFLE_SPLIT_K_MAX_TILE_ELEMS = 32 * 128 +PRESHUFFLE_SPLIT_K_WORKSPACE_ELEMS = ( + 4 * PRESHUFFLE_SPLIT_K_MAX_TILES * PRESHUFFLE_SPLIT_K_MAX_TILE_ELEMS +) + + +@functools.lru_cache(maxsize=128) +def _get_preshuffle_split_buffers( + device: torch.device, + stream: torch.cuda.Stream, +) -> tuple[Tensor, Tensor]: + # Safe to reuse: launches on a stream are ordered and the reduction hands + # the semaphore back zeroed. + workspace = torch.empty( + PRESHUFFLE_SPLIT_K_WORKSPACE_ELEMS, dtype=torch.float32, device=device + ) + semaphore = torch.zeros( + PRESHUFFLE_SPLIT_K_MAX_TILES, dtype=torch.int32, device=device + ) + return workspace, semaphore + + +def _check_preshuffle_split_capacity( + m: int, n: int, tile_m: int, tile_n: int, split_k: int +) -> None: + tiles = ((m + tile_m - 1) // tile_m) * (n // tile_n) + if tiles > PRESHUFFLE_SPLIT_K_MAX_TILES: + raise RuntimeError( + f"[FlyDSL] split_k needs {tiles} tile semaphores, " + f"more than {PRESHUFFLE_SPLIT_K_MAX_TILES}" + ) + elems = split_k * m * n + if elems > PRESHUFFLE_SPLIT_K_WORKSPACE_ELEMS: + raise RuntimeError( + f"[FlyDSL] split_k needs a {elems}-element fp32 workspace, " + f"more than {PRESHUFFLE_SPLIT_K_WORKSPACE_ELEMS}" + ) + + def flydsl_preshuffle_gemm_a8( XQ: Tensor, WQ: Tensor, @@ -985,8 +1033,9 @@ def flydsl_preshuffle_gemm_a8( xcd_swizzle: int = 0, lds_stage: int = 2, enable_scheduler: bool = True, + split_k: int = 1, ) -> Tensor: - """Compile (cached via lru_cache) and run a FlyDSL preshuffle GEMM kernel.""" + """Compile and run FlyDSL preshuffle GEMM, optionally with fp32 split-K.""" compile_fn = _get_compile_fn() if compile_fn is None: raise RuntimeError("[FlyDSL] compile function not available") @@ -995,14 +1044,24 @@ def flydsl_preshuffle_gemm_a8( m, k = XQ.shape[0], XQ.shape[-1] n = WQ.shape[0] + if m > PRESHUFFLE_M_MAX: + raise RuntimeError( + f"[FlyDSL] M ({m}) exceeds {PRESHUFFLE_M_MAX}; the preshuffle kernel " + f"views A and C through a layout bounded by that many rows." + ) if n % tile_n != 0: raise RuntimeError( f"[FlyDSL] N ({n}) is not a multiple of tile_n ({tile_n}). " f"Arguments not supported! Skipping gemm!" ) - if k % tile_k != 0: + if split_k < 1 or k % split_k != 0: + raise RuntimeError( + f"[FlyDSL] K ({k}) must be divisible by split_k ({split_k})." + ) + if (k // split_k) % tile_k != 0: raise RuntimeError( - f"[FlyDSL] K ({k}) is not a multiple of tile_k ({tile_k}). " + f"[FlyDSL] K/split_k ({k // split_k}) is not a multiple of " + f"tile_k ({tile_k}). " f"Arguments not supported! Skipping gemm!" ) @@ -1037,6 +1096,7 @@ def flydsl_preshuffle_gemm_a8( enable_scheduler=bool(enable_scheduler), xcd_swizzle=int(xcd_swizzle), lds_stage=int(lds_stage), + split_k=int(split_k), ) def _as_i8(t): @@ -1047,12 +1107,24 @@ def _as_i8(t): # epilogue != "none"). Pass an empty tensor as a placeholder for the # default epilogue="none" path. _dummy_bias = torch.empty(0, dtype=Out.dtype, device=Out.device) + if split_k > 1: + _check_preshuffle_split_capacity(m, n, tile_m, tile_n, split_k) + workspace, semaphore = _get_preshuffle_split_buffers( + Out.device, torch.cuda.current_stream(device=Out.device) + ) + else: + workspace = out_contig + # dtype is part of the executable's cache signature, so this must match + # what the AOT pre-compile passes or every non-split-K kernel misses it. + semaphore = torch.empty(0, dtype=torch.int32, device=Out.device) # The layout-API launcher (PR #754) takes fx.Tensor args (it builds views via # fx.get_iter/make_view), so pass flat torch tensors directly rather than raw # pointers. _run_compiled( exe, + workspace.view(-1), out_contig.view(-1), + semaphore, _as_i8(XQ.contiguous()).view(-1), _as_i8(WQ.contiguous()).view(-1), x_scale.contiguous().view(-1), diff --git a/aiter/ops/flydsl/gemm_tune/flydsl_gemm_a8w8_bpreshuffle_common.py b/aiter/ops/flydsl/gemm_tune/flydsl_gemm_a8w8_bpreshuffle_common.py index 619979f47b..ffbb3c5e21 100644 --- a/aiter/ops/flydsl/gemm_tune/flydsl_gemm_a8w8_bpreshuffle_common.py +++ b/aiter/ops/flydsl/gemm_tune/flydsl_gemm_a8w8_bpreshuffle_common.py @@ -67,6 +67,7 @@ class kernelInstance: xcd_swizzle: int # 0=off, >0=group size for XCD remap lds_stage: int = 2 # 2=double-buffer ping-pong, 1=single A-LDS buffer (half LDS) sScheduler: str = "Default" # scheduler hints on; "Off" = compiler default + k_split: int = 1 # >1 splits the K loop over gridDim.z (fp32 workspace + reduce) @property def enable_scheduler(self) -> bool: @@ -99,6 +100,7 @@ def name(self) -> str: ), self.sScheduler.lower(), ] + + ([f"ks{self.k_split}"] if self.k_split > 1 else []) ) @@ -114,6 +116,7 @@ def _ki( q_dtype_w="fp8", dtype="bf16", scheduler="Default", + k_split=1, ): return kernelInstance( tile_m, @@ -127,6 +130,7 @@ def _ki( xcd_swizzle, lds_stage, scheduler, + k_split, ) @@ -227,6 +231,8 @@ def kernel_fits_shape(ki: kernelInstance, M: int, N: int, K: int) -> bool: return False if N % ki.tile_n != 0 or K % ki.tile_k != 0: return False + if ki.k_split > 1 and (K // ki.tile_k) % ki.k_split != 0: + return False if _padded_m(M) % ki.tile_m != 0: return False num_ctas = ((M + ki.tile_m - 1) // ki.tile_m) * (N // ki.tile_n) @@ -329,6 +335,32 @@ def _estimate_max_wpe(tile_m: int, tile_n: int, total_vgpr: int = 512) -> int: return int(total_vgpr / max(est_per_wave, 1)) +# Legal values are the divisors of K//tile_k, which is shape-dependent, so they +# are enumerated rather than hardcoded. +K_SPLIT_MIN_TILES_PER_SLICE = 2 # keep the ping-pong loop fed +K_SPLIT_MAX_CTA_OVERSUBSCRIBE = 4 # no point going far past one CU each + + +def k_split_candidates(ki, M: int, N: int, K: int, cu_num: int = 256) -> list[int]: + """Split-K values worth benchmarking; 1 is excluded, the caller has it. + + Empty once the tile grid already fills the GPU -- splitting would only add + the reduce pass. That bound also caps the fp32 workspace, since for a given + tile grid it caps M. + """ + if ki.k_split != 1 or K % ki.tile_k: + return [] + base_ctas = ((M + ki.tile_m - 1) // ki.tile_m) * (N // ki.tile_n) + if base_ctas >= cu_num: + return [] + n_tiles = K // ki.tile_k + max_split = min( + n_tiles // K_SPLIT_MIN_TILES_PER_SLICE, + max(2, cu_num * K_SPLIT_MAX_CTA_OVERSUBSCRIBE // base_ctas), + ) + return [d for d in range(2, max_split + 1) if n_tiles % d == 0] + + def _build_kernels_list(tiles, total_vgpr=512): kl = {} idx = 0 diff --git a/aiter/ops/flydsl/kernels/preshuffle_gemm.py b/aiter/ops/flydsl/kernels/preshuffle_gemm.py index d099783fe3..cc099fcddb 100644 --- a/aiter/ops/flydsl/kernels/preshuffle_gemm.py +++ b/aiter/ops/flydsl/kernels/preshuffle_gemm.py @@ -25,6 +25,7 @@ from aiter.ops.flydsl.kernels import buffer_ops, vector from .mfma_preshuffle_pipeline import xcd_remap_bx_by +from .splitk_epilogue import CPOL_COHERENT, splitk_reduce_epilogue # (dsrd_preload, dvmem_preload) per (tile_m, tile_n, tile_k). _TILE_PRELOAD_TABLE = { @@ -118,6 +119,10 @@ def _get_preload(tile_m, tile_n, tile_k): ) +# Row bound of the A/C layout views; gemm_kernels mirrors it to guard on the host. +PRESHUFFLE_M_MAX = 65536 + + @functools.lru_cache(maxsize=1024) def compile_preshuffle_gemm( *, @@ -134,16 +139,22 @@ def compile_preshuffle_gemm( use_async_copy: bool = False, xcd_swizzle: int = 0, lds_stage: int = 2, + split_k: int = 1, ): """Compile preshuffle GEMM (fp8/int8/fp16/bf16). - Signature: fn(C, A, B, scale_a, scale_b, bias, M, N, stream). bias is the fused - epilogue bias (per-N, out_dtype); unused when epilogue == "none". + Signature: fn(C, out, semaphore, A, B, scale_a, scale_b, bias, M, N, stream). + For split_k > 1, C is an fp32 partial workspace and the last arriving split + performs the reduction and final conversion inside the same kernel launch. """ if in_dtype not in ("fp8", "int8", "fp16", "bf16"): raise ValueError(f"in_dtype must be fp8/int8/fp16/bf16, got {in_dtype!r}") - if tile_k <= 0 or K % tile_k != 0: + if split_k < 1 or K % split_k != 0: + raise ValueError(f"split_k must divide K; got split_k={split_k}, K={K}") + split_k_extent = K // split_k + if tile_k <= 0 or split_k_extent % tile_k != 0: raise ValueError( - f"tile_k must be a positive divisor of K; got tile_k={tile_k}, K={K}" + "tile_k must be a positive divisor of K/split_k; " + f"got tile_k={tile_k}, K={K}, split_k={split_k}" ) if epilogue not in ("none", "bias", "bias_relu", "bias_silu", "bias_gelu"): raise ValueError( @@ -156,6 +167,8 @@ def compile_preshuffle_gemm( _has_relu = epilogue == "bias_relu" _has_silu = epilogue == "bias_silu" _has_gelu = epilogue == "bias_gelu" + if split_k > 1 and _has_epilogue: + raise ValueError("split_k > 1 does not support fused bias or activation") is_fp8 = in_dtype == "fp8" is_int8 = in_dtype == "int8" @@ -181,12 +194,14 @@ def compile_preshuffle_gemm( layout_elem = Int8 else: layout_elem = Float8E4M3FN if is_gfx950 else Float8E4M3FNUZ - out_elem_cls = BFloat16 if out_dtype == "bf16" else Float16 + final_out_elem_cls = BFloat16 if out_dtype == "bf16" else Float16 + out_elem_cls = Float32 if split_k > 1 else final_out_elem_cls + out_elem_bytes = 4 if split_k > 1 else 2 # Tile geometry (tile_K_perm = K-elements grouped per MMA k-step) tile_K_perm = 128 if use_mfma_scale_128 else (64 if is_8bit else 32) k_iters = tile_k // tile_K_perm - num_tiles = K // tile_k + num_tiles = split_k_extent // tile_k m_repeat = tile_m // 16 num_waves = 4 n_per_wave = tile_n // num_waves @@ -212,11 +227,15 @@ class SharedStorage: a0: fx.Array[layout_elem, a_lds_elems, 16] if lds_stage == 2: a1: fx.Array[layout_elem, a_lds_elems, 16] + if split_k > 1: + split_flag: fx.Array[Int32, 1, 4] # ── Kernel ──────────────────────────────────────────────────────── @flyc.kernel def kernel_gemm( arg_c: fx.Tensor, + arg_out: fx.Tensor, + arg_semaphore: fx.Tensor, arg_a: fx.Tensor, arg_b: fx.Tensor, arg_scale_a: fx.Tensor, @@ -228,7 +247,8 @@ def kernel_gemm( tiled_copy_g2s: fx.TiledCopy, ): tid = fx.thread_idx.x - bid_x, bid_y, _ = fx.block_idx + bid_x, bid_y, bid_z = fx.block_idx + k_off = fx.Int32(bid_z) * num_tiles if const_expr(xcd_swizzle > 0): _bx, _by = xcd_remap_bx_by( @@ -265,10 +285,19 @@ def kernel_gemm( num_records_bytes=fx.Int64(i32_m) * fx.Int64(K) * fx.Int64(elem_bytes), ) gB = fx.rocdl.make_buffer_tensor(arg_b) + c_tensor = arg_c + if const_expr(split_k > 1): + c_split_offset = fx.Int64(bid_z) * fx.Int64(i32_m) * fx.Int64(N) + c_tensor = fx.Tensor( + fx.make_view( + fx.add_offset(fx.get_iter(arg_c), c_split_offset), + fx.make_layout((PRESHUFFLE_M_MAX, N), (N, 1)), + ) + ) gC = fx.rocdl.make_buffer_tensor( - arg_c, + c_tensor, max_size=False, - num_records_bytes=fx.Int64(i32_m) * fx.Int64(N) * fx.Int64(2), + num_records_bytes=fx.Int64(i32_m) * fx.Int64(N) * fx.Int64(out_elem_bytes), ) tA = fx.flat_divide(gA, fx.make_tile(tile_m, tile_k))[None, None, bid_x, None] @@ -328,7 +357,9 @@ def _make_sA(arr): frag_C = thr_mma.make_fragment_C(tC) frag_A_retile = thr_s2r.retile(frag_A) frag_B_retile_stages = [thr_g2r_B.retile(b) for b in frag_B_stages] - buf_copy_out = fx.make_copy_atom(fx.rocdl.BufferCopy16b(), out_elem_cls) + out_cpol = CPOL_COHERENT if split_k > 1 else 0 + copy_op = fx.rocdl.BufferCopy32b if split_k > 1 else fx.rocdl.BufferCopy16b + buf_copy_out = fx.make_copy_atom(copy_op(out_cpol), out_elem_cls) thr_r2g_C = fx.make_tiled_copy_C(buf_copy_out, tiled_mma).get_slice(tid) pC_g = thr_r2g_C.partition_S(tC) frag_C_out = fx.make_fragment_like(frag_C, out_elem_cls.ir_type) @@ -341,7 +372,9 @@ def _make_sA(arr): # OOB rows as 0 instead of faulting past the allocation. gA_flat = fx.rocdl.make_buffer_tensor( fx.Tensor( - fx.make_view(fx.get_iter(arg_a), fx.make_layout(65536 * K, 1)) + fx.make_view( + fx.get_iter(arg_a), fx.make_layout(PRESHUFFLE_M_MAX * K, 1) + ) ), max_size=False, num_records_bytes=fx.Int64(i32_m) * fx.Int64(K) * fx.Int64(elem_bytes), @@ -362,7 +395,7 @@ def dma_a_to_lds(k_tile_val, stage): fx.Int32.ir_type, wave_id * wave_stride_bytes ) lds_ptr = fx.add_offset(sA_i8_ptr[stage], wave_off) - base_k = k_tile_val * tile_k + base_k = (k_off + k_tile_val) * tile_k for i in range_constexpr(num_a_loads): if const_expr(i > 0): lds_ptr = fx.add_offset(lds_ptr, step_bytes) @@ -521,7 +554,7 @@ def pipeline_2stage(read_stage, next_k_val=None, read_next=True): dma_a_to_lds(next_k_val, a_write) fx.copy( buf_copy, - pB_g[None, None, None, next_k_val], + pB_g[None, None, None, k_off + next_k_val], frag_B_retile_stages[write_stage], ) mma_kloop(a_read, cur_frag_B) @@ -532,10 +565,11 @@ def pipeline_2stage(read_stage, next_k_val=None, read_next=True): gpu.barrier() return if const_expr(do_next): - fx.copy(buf_copy, pA_g[None, None, None, next_k_val], frag_copy_A) + global_k_next = k_off + next_k_val + fx.copy(buf_copy, pA_g[None, None, None, global_k_next], frag_copy_A) fx.copy( buf_copy, - pB_g[None, None, None, next_k_val], + pB_g[None, None, None, global_k_next], frag_B_retile_stages[write_stage], ) mma_kloop(a_read, cur_frag_B) @@ -554,13 +588,13 @@ def pipeline_2stage(read_stage, next_k_val=None, read_next=True): ) if const_expr(use_async_copy): dma_a_to_lds(fx.Int32(0), 0) - fx.copy(buf_copy, pB_g[None, None, None, 0], frag_B_retile_stages[0]) + fx.copy(buf_copy, pB_g[None, None, None, k_off], frag_B_retile_stages[0]) frag_C.store(acc_zero) rocdl.s_waitcnt(num_b_loads) gpu.barrier() else: - fx.copy(buf_copy, pA_g[None, None, None, 0], frag_copy_A) - fx.copy(buf_copy, pB_g[None, None, None, 0], frag_B_retile_stages[0]) + fx.copy(buf_copy, pA_g[None, None, None, k_off], frag_copy_A) + fx.copy(buf_copy, pB_g[None, None, None, k_off], frag_B_retile_stages[0]) frag_C.store(acc_zero) fx.copy(uni_copy, frag_copy_A, pA_s_stages[0][None, None, None]) gpu.barrier() @@ -575,8 +609,14 @@ def pipeline_2stage(read_stage, next_k_val=None, read_next=True): k_next = fx.Int32(iv + 1) mma_kloop(0, frag_Bc) if const_expr(not use_async_copy): - fx.copy(buf_copy, pA_g[None, None, None, k_next], frag_copy_A) - fx.copy(buf_copy, pB_g[None, None, None, k_next], frag_Bc_retile) + fx.copy( + buf_copy, pA_g[None, None, None, k_off + k_next], frag_copy_A + ) + fx.copy( + buf_copy, + pB_g[None, None, None, k_off + k_next], + frag_Bc_retile, + ) gpu.barrier() # single buffer: all reads done before overwrite if const_expr(use_async_copy): dma_a_to_lds(k_next, 0) @@ -689,7 +729,6 @@ def load_epi_operands(): # ── Epilogue ───────────────────────────────────────────── if const_expr(not is_8bit and not _has_epilogue): frag_C_out.store(Vec(frag_C.load()).to(out_elem_cls)) - fx.copy(buf_copy_out, frag_C_retile, pC_g) else: if const_expr(not overlap_epi_load): s_a_vals, s_b_vals, bias_vals = load_epi_operands() @@ -738,12 +777,40 @@ def apply_activation(val_s): T.vec(acc_size, out_elem_cls.ir_type), out_elems ) frag_C_out.store(out_vec) - fx.copy(buf_copy_out, frag_C_retile, pC_g) + + fx.copy(buf_copy_out, frag_C_retile, pC_g) + if const_expr(split_k > 1): + # The store carries sc0|sc1, so waiting on it is the whole release. + rocdl.s_waitcnt(0) + gpu.barrier() + splitk_reduce_epilogue( + arg_c, + fx.Tensor( + fx.make_view( + fx.get_iter(arg_out), + fx.make_layout((PRESHUFFLE_M_MAX, N), (N, 1)), + ) + ), + arg_semaphore, + lds.split_flag.ptr, + tile_m, + tile_n, + total_threads, + final_out_elem_cls, + tid, + bid_x, + bid_y, + i32_m, + N, + split_k, + ) # ── Host launcher ───────────────────────────────────────────── @flyc.jit def launch_gemm( arg_c: fx.Tensor, + arg_out: fx.Tensor, + arg_semaphore: fx.Tensor, arg_a: fx.Tensor, arg_b: fx.Tensor, arg_scale_a: fx.Tensor, @@ -809,7 +876,7 @@ def launch_gemm( ) # Reshape A and C to 2D - M_max = 65536 + M_max = PRESHUFFLE_M_MAX arg_a_2d = fx.Tensor( fx.make_view(fx.get_iter(arg_a), fx.make_layout((M_max, K), (K, 1))) ) @@ -822,6 +889,8 @@ def launch_gemm( kernel_gemm( arg_c_2d, + arg_out, + arg_semaphore, arg_a_2d, preshuffle_B, arg_scale_a, @@ -833,7 +902,7 @@ def launch_gemm( tiled_copy_g2s, value_attrs={"rocdl.waves_per_eu": waves_per_eu}, ).launch( - grid=(gx, gy, 1), + grid=(gx, gy, split_k), block=(256, 1, 1), stream=stream, ) diff --git a/aiter/ops/flydsl/kernels/splitk_epilogue.py b/aiter/ops/flydsl/kernels/splitk_epilogue.py new file mode 100644 index 0000000000..e5b984646f --- /dev/null +++ b/aiter/ops/flydsl/kernels/splitk_epilogue.py @@ -0,0 +1,163 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""One-stage split-K epilogue: last arriving split reduces the fp32 partials.""" + +import math + +import flydsl.expr as fx +from flydsl.expr import gpu, range_constexpr +from flydsl.expr.typing import Float32, Int32 +from flydsl.expr.typing import Vector as Vec + +from . import communication_ops_utils as comm_ops + +# Split-K partials cross XCDs, which do not share L2. An agent-scope fence +# would write back the whole L2 per CTA and evict the A/B tiles everyone else +# is reading; sc0|sc1 writes through just these accesses. +CPOL_COHERENT = 0x1 | 0x10 + +VEC = 4 # elements per thread per access: dwordx4 loads, dwordx2 bf16 stores + +_COPY_ATOM = { + 32: fx.rocdl.BufferCopy32b, + 64: fx.rocdl.BufferCopy64b, + 128: fx.rocdl.BufferCopy128b, +} + + +def pairwise_sum(parts): + """Sum a Python list of Vectors as a balanced tree. + + Module level because the kernel AST rewriter would turn this ``while`` into + an ``scf.while``; it is trace-time metaprogramming over a Python list. + """ + while len(parts) > 1: + nxt = [] + for lhs, rhs in zip(parts[0::2], parts[1::2]): + nxt.append(lhs + rhs) + if len(parts) % 2: + nxt.append(parts[-1]) + parts = nxt + return parts[0] + + +def reduce_thread_split(tile_m, tile_n, nthreads): + """Threads along (M, N) for the reduce, VEC contiguous columns per thread. + + The MMA's C layout gives each lane elements that are contiguous in M, whose + memory stride is N, so partitioning the reduce with it would issue one dword + per element. This layout is chosen for the memory instead: one row per + thread where the block divides the tile rows, so the ragged-M skip below is + exact, otherwise as many threads along N as the row splits into. + """ + row_vecs = tile_n // VEC + tm_thr = math.gcd(tile_m, nthreads) + if row_vecs % (nthreads // tm_thr): + tm_thr = nthreads // math.gcd(row_vecs, nthreads) + if tile_n % VEC or tile_m % tm_thr or row_vecs % (nthreads // tm_thr): + raise ValueError(f"tile ({tile_m}, {tile_n}) does not split over {nthreads}") + return tm_thr, nthreads // tm_thr + + +@comm_ops.traced +def splitk_reduce_epilogue( + workspace, + out, + semaphore, + flag_ptr, + tile_m, + tile_n, + nthreads, + out_elem_cls, + tid, + bid_x, + bid_y, + i32_m, + n, + split_k, +): + """Reduce ``split_k`` fp32 planes of ``workspace`` into ``out``, in-kernel. + + ``workspace`` and ``out`` are 2D ``[*, n]`` views; workspace holds the planes + back to back. Every CTA bumps this tile's ``semaphore`` slot and the one that + sees ``split_k - 1`` prior arrivals reduces. The caller must have stored its + partial with ``CPOL_COHERENT`` and waited on it: that is the release, and it + is why no agent fence (which flushes the whole L2) is needed. + + ``out_elem_cls`` is the only dtype knob -- any 16/32-bit output works. + """ + out_bytes = out_elem_cls.width // 8 + tiles_n = n // tile_n + + if tid == fx.Int32(0): + arrival = fx.Int32( + comm_ops.atomic_add_agent( + _semaphore_addr(semaphore, bid_x, bid_y, tiles_n), fx.Int32(1) + ) + ) + is_last = (arrival == fx.Int32(split_k - 1)).select(fx.Int32(1), fx.Int32(0)) + fx.ptr_store(Vec.from_elements([is_last], Int32), flag_ptr) + gpu.barrier() + + is_last = Vec(fx.make_view(flag_ptr, fx.make_layout(1, 1)).load())[0] + if is_last != fx.Int32(0): + plane = fx.Int64(i32_m) * fx.Int64(n) + g_workspace = fx.rocdl.make_buffer_tensor( + workspace, + max_size=False, + num_records_bytes=plane * fx.Int64(4 * split_k), + ) + g_out = fx.rocdl.make_buffer_tensor( + out, max_size=False, num_records_bytes=plane * fx.Int64(out_bytes) + ) + tile = fx.make_tile(tile_m, tile_n) + t_workspace = fx.flat_divide(g_workspace, tile)[None, None, bid_x, bid_y] + t_out = fx.flat_divide(g_out, tile)[None, None, bid_x, bid_y] + + tm_thr, tn_thr = reduce_thread_split(tile_m, tile_n, nthreads) + tiler_mn, tv = fx.make_layout_tv( + fx.make_layout((tm_thr, tn_thr), (tn_thr, 1)), + fx.make_layout((1, VEC), (1, 1)), + ) + load_atom = fx.make_copy_atom(_COPY_ATOM[VEC * 32](CPOL_COHERENT), Float32) + store_atom = fx.make_copy_atom( + _COPY_ATOM[VEC * out_elem_cls.width](), out_elem_cls + ) + p_workspace = ( + fx.make_tiled_copy(load_atom, tv, tiler_mn) + .get_slice(tid) + .partition_S(t_workspace) + ) + p_out = ( + fx.make_tiled_copy(store_atom, tv, tiler_mn) + .get_slice(tid) + .partition_D(t_out) + ) + + # Both descriptors stop at row M, so a ragged-M access is dropped in + # hardware; skipping the threads whose rows are all past M as well keeps + # a short tile from issuing them at all. + stride = fx.Int32(i32_m) * fx.Int32(n) + row0 = fx.Int32(bid_x) * tile_m + fx.Int32(tid) // tn_thr + if row0 < fx.Int32(i32_m): + parts = [] + for s in range_constexpr(split_k): + frag = fx.make_fragment_like(p_workspace) + fx.copy(load_atom, p_workspace, frag, soffset=fx.Int32(s) * stride) + parts.append(Vec(frag.load())) + frag_out = fx.make_fragment_like(p_out) + frag_out.store(pairwise_sum(parts).to(out_elem_cls)) + fx.copy(store_atom, frag_out, p_out) + + if tid == fx.Int32(0): + # Undo this tile's arrivals with the same atomic instead of storing + # 0, so an early increment from the next launch is not clobbered. + comm_ops.atomic_add_agent( + _semaphore_addr(semaphore, bid_x, bid_y, tiles_n), fx.Int32(-split_k) + ) + + +def _semaphore_addr(semaphore, bid_x, bid_y, tiles_n): + idx = fx.Int32(bid_x) * tiles_n + fx.Int32(bid_y) + return fx.Int64(fx.ptrtoint(fx.get_iter(semaphore))) + fx.Int64(idx) * fx.Int64(4) diff --git a/aiter/ops/gemm_op_a8w8.py b/aiter/ops/gemm_op_a8w8.py index ec69fd5b96..bfd03dccb4 100644 --- a/aiter/ops/gemm_op_a8w8.py +++ b/aiter/ops/gemm_op_a8w8.py @@ -139,14 +139,16 @@ def gemm_a8w8_bpreshuffle_cktile( def _parse_flydsl_kernel_name(kernel_name: str): """Parse a flydsl kernelName into ``(tile_m, tile_n, tile_k, async_copy, - waves_per_eu, xcd_swizzle, lds_stage, scheduler)``, or None on failure. - Legacy names lacking the xcd/lds/scheduler tokens default them to - ``0``/``2``/``"Default"``. + waves_per_eu, xcd_swizzle, lds_stage, scheduler, k_split)``, or None on + failure. Legacy names lacking the xcd/lds/scheduler tokens default them to + ``0``/``2``/``"Default"``; the ``_ksN`` split-K suffix is only emitted for + k_split > 1, so every previously tuned name still parses to k_split=1. """ import re m = re.match( - r"flydsl_bpreshuflle_(\d+)x(\d+)x(\d+)_\w+_\w+_\w+_(\d+)x(\d+)(?:x(\d+))?(?:x(\d+))?(?:_([A-Za-z][A-Za-z0-9]*))?$", + r"flydsl_bpreshuflle_(\d+)x(\d+)x(\d+)_\w+_\w+_\w+_(\d+)x(\d+)(?:x(\d+))?(?:x(\d+))?" + r"(?:_(?!ks\d+$)([A-Za-z][A-Za-z0-9]*))?(?:_ks(\d+))?$", kernel_name, ) if m is None: @@ -155,7 +157,8 @@ def _parse_flydsl_kernel_name(kernel_name: str): xcd_swizzle = int(m.group(6)) if m.group(6) else 0 lds_stage = int(m.group(7)) if m.group(7) else 2 scheduler = m.group(8) if m.group(8) else "Default" - return (tm, tn, tk, acp, wpe, xcd_swizzle, lds_stage, scheduler) + k_split = int(m.group(9)) if m.group(9) else 1 + return (tm, tn, tk, acp, wpe, xcd_swizzle, lds_stage, scheduler, k_split) def gemm_a8w8_bpreshuffle_flydsl( @@ -187,7 +190,7 @@ def gemm_a8w8_bpreshuffle_flydsl( parsed = _parse_flydsl_kernel_name(kernel_name) if parsed is None: return gemm_a8w8_bpreshuffle_ck(XQ, WQ, x_scale, w_scale, Out) - tm, tn, tk, acp, wpe, xcd_swizzle, lds_stage, scheduler = parsed + tm, tn, tk, acp, wpe, xcd_swizzle, lds_stage, scheduler, k_split = parsed flydsl_preshuffle_gemm_a8( XQ.contiguous(), @@ -203,6 +206,7 @@ def gemm_a8w8_bpreshuffle_flydsl( xcd_swizzle, lds_stage=lds_stage, enable_scheduler=str(scheduler).lower() != "off", + split_k=k_split, ) return Out diff --git a/csrc/ck_gemm_a8w8_bpreshuffle/gemm_a8w8_bpreshuffle_tune.py b/csrc/ck_gemm_a8w8_bpreshuffle/gemm_a8w8_bpreshuffle_tune.py index 96da0e1449..3e7a18a82b 100644 --- a/csrc/ck_gemm_a8w8_bpreshuffle/gemm_a8w8_bpreshuffle_tune.py +++ b/csrc/ck_gemm_a8w8_bpreshuffle/gemm_a8w8_bpreshuffle_tune.py @@ -37,6 +37,9 @@ from aiter.ops.flydsl.gemm_tune.flydsl_gemm_a8w8_bpreshuffle_common import ( PIPELINES as FLYDSL_PIPELINES, ) + from aiter.ops.flydsl.gemm_tune.flydsl_gemm_a8w8_bpreshuffle_common import ( + k_split_candidates, + ) from aiter.ops.flydsl.gemm_tune.flydsl_gemm_a8w8_bpreshuffle_common import ( kernels_list as kernels_list_flydsl, ) @@ -51,6 +54,9 @@ kernels_list_flydsl = {} kernels_list_flydsl_8wave = {} + def k_split_candidates(*_a, **_kw): + return [] + from aiter.ops.flydsl.utils import is_flydsl_available @@ -131,7 +137,7 @@ def run_gemm_a8w8_asm( ) -def run_gemm_flydsl(x, weight_shuffle, x_scale, w_scale, out, kernel_id): +def run_gemm_flydsl(x, weight_shuffle, x_scale, w_scale, out, kernel_id, k_split=1): ki = kernels_list_flydsl[kernel_id] flydsl_preshuffle_gemm_a8( x, @@ -147,11 +153,14 @@ def run_gemm_flydsl(x, weight_shuffle, x_scale, w_scale, out, kernel_id): ki.xcd_swizzle, ki.lds_stage, ki.enable_scheduler, + split_k=k_split, ) return out -def run_gemm_flydsl_8wave(x, weight_shuffle, x_scale, w_scale, out, kernel_id): +def run_gemm_flydsl_8wave( + x, weight_shuffle, x_scale, w_scale, out, kernel_id, k_split=1 +): from aiter.ops.flydsl.gemm_a8w8_bpreshuffle_8wave import flydsl_8wave_gemm_a8 ki = kernels_list_flydsl_8wave[kernel_id] @@ -411,20 +420,10 @@ def get_cktile_gemm_a8w8_bpreshuffle_tune_task( gemm_keys = ["x", "weight_shuffle", "x_scale", "w_scale", "out"] ref_keys = ["x", "weight", "x_scale", "w_scale", "bias_f32"] tasks_ck = [] - for i, kernel in filtered_cktile.items(): - maxsplitK = ( - aiter.compute_gemm_SplitK( - M, - N, - K, - kernel.MTile, - kernel.NTile, - kernel.KTile, - ) - if useSplitK - else 0 - ) - for splitK in range(maxsplitK + 1): + for i in filtered_cktile: + # cktile's flatmm accumulates into E without clearing it, so k_batch + # stays 1 regardless of --splitK. + for splitK in range(1): info = (info_keys, i, splitK, "", "cktile") tasks_ck.append( ( @@ -523,6 +522,7 @@ def get_ck_gemm_a8w8_bpreshuffle_tune_task( def get_flydsl_gemm_a8w8_bpreshuffle_tune_task( self, info_keys, + useSplitK, seed, ): gfx, _cu_num, M, N, K, q_dtype_w = info_keys @@ -557,34 +557,41 @@ def get_flydsl_gemm_a8w8_bpreshuffle_tune_task( ki = pipe.kernels_list[i] if not pipe.fits(ki, M, N, K): continue - tasks.append( - ( - (info_keys, i, 0, ki.name, "flydsl"), - generate_data, - (M, N, K, seed, dtypes.bf16, q_dtype_eval), - runner, + for ks in [1] + ( + k_split_candidates(ki, M, N, K, cu_num=self.get_cu_num()) + if useSplitK and pipe.name == "preshuffle" + else [] + ): + name = ki.name if ks == 1 else f"{ki.name}_ks{ks}" + tasks.append( ( - gemm_flydsl_keys, - i, - ), - { - "num_warmup": args.warmup, - "num_iters": args.iters, - }, - run_torch, - ( - ref_keys, - dtypes.bf16, - ), - {}, - None, - 1e-2, - 0.01, - None, - None, - ("out",), + (info_keys, i, 0 if ks == 1 else ks, name, "flydsl"), + generate_data, + (M, N, K, seed, dtypes.bf16, q_dtype_eval), + runner, + ( + gemm_flydsl_keys, + i, + ks, + ), + { + "num_warmup": args.warmup, + "num_iters": args.iters, + }, + run_torch, + ( + ref_keys, + dtypes.bf16, + ), + {}, + None, + 1e-2, + 0.01, + None, + None, + ("out",), + ) ) - ) return tasks def _get_flydsl_tune_task_gfx1250(self, info_keys, seed): @@ -691,6 +698,7 @@ def tune( task.extend( self.get_flydsl_gemm_a8w8_bpreshuffle_tune_task( info_keys, + useSplitK, seed, ) )