From 10ae462f18b9e352b210de539146f4657c26062a Mon Sep 17 00:00:00 2001 From: Yu-Yun Chang Date: Wed, 7 Oct 2026 10:31:00 +0000 Subject: [PATCH] perf(gemm): reuse output buffers and large-M A8W8 tuning - Accept caller-owned A8W8 outputs and optionally reuse A4W4 scratch. - Reuse the largest tuned bpreshuffle row only above the table and chunk FlyDSL launches at their row limit. - Keep KDA prefill tuning and dispatch tests with the GEMM implementation. --- .../a8w8_bpreshuffle_tuned_gemm_kimik3.csv | 87 ++++++++++ .../a8w8_bpreshuffle_untuned_gemm_kimik3.csv | 50 ++++++ .../dsv3_a8w8_bpreshuffle_tuned_gemm.csv | 9 - aiter/ops/gemm_op_a4w4.py | 30 +++- aiter/ops/gemm_op_a8w8.py | 154 ++++++++++++++++-- .../test_gemm_a8w8_bpreshuffle_large_m.py | 88 ++++++++++ 6 files changed, 395 insertions(+), 23 deletions(-) create mode 100644 op_tests/test_gemm_a8w8_bpreshuffle_large_m.py 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 0899c7909f6..b2c5afbfa1b 100644 --- a/aiter/configs/model_configs/a8w8_bpreshuffle_tuned_gemm_kimik3.csv +++ b/aiter/configs/model_configs/a8w8_bpreshuffle_tuned_gemm_kimik3.csv @@ -1100,3 +1100,90 @@ gfx950,256,15488,7168,1536,torch.float8_e4m3fn,flydsl,1000034,0,174.6294,flydsl_ gfx950,256,15744,7168,1536,torch.float8_e4m3fn,flydsl,473,0,172.9698,flydsl_bpreshuflle_96x256x128_F8_F8_B16_1x1x4x2_default,2004.31,1508.35,0.0 gfx950,256,16000,7168,1536,torch.float8_e4m3fn,flydsl,1000034,0,175.6231,flydsl_bpreshuffle_8w_256x256x128_F8_F8_B16_4x4,2006.12,1508.7,0.0 gfx950,256,16256,7168,1536,torch.float8_e4m3fn,flydsl,1000031,0,184.4351,flydsl_bpreshuffle_8w_256x256x128_F8_F8_B16_3x4,1940.84,1458.64,0.0 +gfx950,256,16,7168,3584,torch.float8_e4m3fn,ck,10,0,7.7344,a8w8_bpreshuffle_128x16x32x512_16x16_16x16_32x4x1_32x4x1_1x16x1x8_4x4x1_1x1_intrawave_v1,106.29,3358.61,0.0 +gfx950,256,72,6336,7168,torch.float8_e4m3fn,flydsl,2061,0,19.224,flydsl_bpreshuflle_16x64x512_F8_F8_B16_1x3x4x1_default,340.2,2436.79,0.0 +gfx950,256,72,7168,3584,torch.float8_e4m3fn,cktile,78,0,12.512,a8w8_bpreshuffle_cktile_0x0x8x4x1x0x0x0x0x1_48x64x256_1x4x1_16x16x128_default,295.67,2156.36,0.0 +gfx950,256,80,6336,7168,torch.float8_e4m3fn,flydsl,1343,0,20.0478,flydsl_bpreshuflle_16x64x512_F8_F8_B16_1x0x4x1_default,362.47,2344.58,0.0 +gfx950,256,80,7168,3584,torch.float8_e4m3fn,cktile,78,0,12.6711,a8w8_bpreshuffle_cktile_0x0x8x4x1x0x0x0x0x1_48x64x256_1x4x1_16x16x128_default,324.39,2140.6,0.0 +gfx950,256,88,6336,7168,torch.float8_e4m3fn,flydsl,921,0,19.847,flydsl_bpreshuflle_48x64x256_F8_F8_B16_1x3x4x2_default,402.75,2376.3,0.0 +gfx950,256,88,7168,3584,torch.float8_e4m3fn,flydsl,451,0,11.8022,flydsl_bpreshuflle_48x64x256_F8_F8_B16_1x1x4x2_default,383.1,2310.34,0.0 +gfx950,256,96,6336,7168,torch.float8_e4m3fn,flydsl,921,0,20.133,flydsl_bpreshuflle_48x64x256_F8_F8_B16_1x3x4x2_default,433.12,2350.42,0.0 +gfx950,256,96,7168,3584,torch.float8_e4m3fn,flydsl,633,0,12.3565,flydsl_bpreshuflle_48x64x256_F8_F8_B16_1x2x0x2_default,399.18,2218.3,0.0 +gfx950,256,104,6336,7168,torch.float8_e4m3fn,cktile,179,0,20.7096,a8w8_bpreshuffle_cktile_0x0x8x4x1x0x0x0x0x4_64x64x256_1x4x1_16x16x128_default,456.15,2292.65,0.0 +gfx950,256,104,7168,3584,torch.float8_e4m3fn,cktile,179,0,12.9034,a8w8_bpreshuffle_cktile_0x0x8x4x1x0x0x0x0x4_64x64x256_1x4x1_16x16x128_default,414.12,2135.39,0.0 +gfx950,256,112,6336,7168,torch.float8_e4m3fn,cktile,179,0,20.7828,a8w8_bpreshuffle_cktile_0x0x8x4x1x0x0x0x0x4_64x64x256_1x4x1_16x16x128_default,489.5,2292.21,0.0 +gfx950,256,112,7168,3584,torch.float8_e4m3fn,cktile,179,0,13.0583,a8w8_bpreshuffle_cktile_0x0x8x4x1x0x0x0x0x4_64x64x256_1x4x1_16x16x128_default,440.68,2121.04,0.0 +gfx950,256,120,6336,7168,torch.float8_e4m3fn,flydsl,698,0,20.707,flydsl_bpreshuflle_64x64x256_F8_F8_B16_1x2x4x2_default,526.39,2308.27,0.0 +gfx950,256,120,7168,3584,torch.float8_e4m3fn,flydsl,208,0,11.9168,flydsl_bpreshuflle_64x64x256_F8_F8_B16_1x0x4x2_default,517.39,2336.24,0.0 +gfx950,256,128,6336,7168,torch.float8_e4m3fn,ck,114,0,20.8839,a8w8_bpreshuffle_256x64x64x256_16x16_16x16_16x16x1_16x16x1_1x32x1x8_8x8x1_2x1_intrawave_v3,556.73,2296.31,0.0 +gfx950,256,128,7168,3584,torch.float8_e4m3fn,flydsl,1129,0,11.9853,flydsl_bpreshuflle_64x64x256_F8_F8_B16_1x4x4x2_default,548.73,2334.85,0.0 +gfx950,256,136,6336,7168,torch.float8_e4m3fn,cktile,84,0,23.1458,a8w8_bpreshuffle_cktile_0x0x8x4x1x0x0x0x0x1_80x64x256_1x4x1_16x16x128_default,533.72,2078.77,0.0 +gfx950,256,136,7168,3584,torch.float8_e4m3fn,flydsl,452,0,13.2018,flydsl_bpreshuflle_48x128x256_F8_F8_B16_1x1x4x2_default,529.3,2130.56,0.0 +gfx950,256,144,6336,7168,torch.float8_e4m3fn,cktile,84,0,23.3525,a8w8_bpreshuffle_cktile_0x0x8x4x1x0x0x0x0x1_80x64x256_1x4x1_16x16x128_default,560.11,2067.16,0.0 +gfx950,256,144,7168,3584,torch.float8_e4m3fn,flydsl,452,0,13.3876,flydsl_bpreshuflle_48x128x256_F8_F8_B16_1x1x4x2_default,552.66,2111.7,0.0 +gfx950,256,152,6336,7168,torch.float8_e4m3fn,flydsl,1958,0,22.106,flydsl_bpreshuflle_32x64x512_F8_F8_B16_0x3x4x1_default,624.56,2190.9,0.0 +gfx950,256,152,7168,3584,torch.float8_e4m3fn,flydsl,933,0,13.3574,flydsl_bpreshuflle_80x64x256_F8_F8_B16_1x3x4x2_default,584.68,2127.21,0.0 +gfx950,256,160,6336,7168,torch.float8_e4m3fn,flydsl,1958,0,22.5446,flydsl_bpreshuflle_32x64x512_F8_F8_B16_0x3x4x1_default,644.64,2155.32,0.0 +gfx950,256,160,7168,3584,torch.float8_e4m3fn,flydsl,933,0,13.0677,flydsl_bpreshuflle_80x64x256_F8_F8_B16_1x3x4x2_default,629.1,2185.34,0.0 +gfx950,256,168,6336,7168,torch.float8_e4m3fn,cktile,227,0,24.5243,a8w8_bpreshuffle_cktile_0x0x8x4x1x0x0x0x0x4_96x64x256_1x4x1_16x16x128_default,622.24,1987.81,0.0 +gfx950,256,168,7168,3584,torch.float8_e4m3fn,cktile,227,0,15.3486,a8w8_bpreshuffle_cktile_0x0x8x4x1x0x0x0x0x4_96x64x256_1x4x1_16x16x128_default,562.39,1869.92,0.0 +gfx950,256,176,6336,7168,torch.float8_e4m3fn,cktile,227,0,24.7487,a8w8_bpreshuffle_cktile_0x0x8x4x1x0x0x0x0x4_96x64x256_1x4x1_16x16x128_default,645.96,1976.2,0.0 +gfx950,256,176,7168,3584,torch.float8_e4m3fn,cktile,85,0,15.3969,a8w8_bpreshuffle_cktile_0x0x8x4x1x0x0x0x0x1_96x64x256_1x4x1_16x16x128_default,587.32,1873.37,0.0 +gfx950,256,184,6336,7168,torch.float8_e4m3fn,flydsl,710,0,23.3337,flydsl_bpreshuflle_96x64x256_F8_F8_B16_1x2x4x2_default,716.27,2102.84,0.0 +gfx950,256,184,7168,3584,torch.float8_e4m3fn,flydsl,468,0,13.3569,flydsl_bpreshuflle_96x64x256_F8_F8_B16_1x1x4x2_default,707.8,2170.22,0.0 +gfx950,256,192,6336,7168,torch.float8_e4m3fn,flydsl,468,0,23.5475,flydsl_bpreshuflle_96x64x256_F8_F8_B16_1x1x4x2_default,740.63,2090.49,0.0 +gfx950,256,192,7168,3584,torch.float8_e4m3fn,flydsl,938,0,13.4702,flydsl_bpreshuflle_96x64x256_F8_F8_B16_1x3x4x2_default,732.36,2162.61,0.0 +gfx950,256,200,6336,7168,torch.float8_e4m3fn,cktile,231,0,26.5673,a8w8_bpreshuffle_cktile_0x0x8x4x1x0x0x0x0x4_112x64x256_1x4x1_16x16x128_default,683.79,1858.84,0.0 +gfx950,256,200,7168,3584,torch.float8_e4m3fn,cktile,50,0,16.4589,a8w8_bpreshuffle_cktile_0x0x8x4x1x0x0x0x0x2_64x64x256_1x4x1_16x16x128_default,624.35,1778.62,0.0 +gfx950,256,208,6336,7168,torch.float8_e4m3fn,cktile,78,0,26.7745,a8w8_bpreshuffle_cktile_0x0x8x4x1x0x0x0x0x1_48x64x256_1x4x1_16x16x128_default,705.64,1850.39,0.0 +gfx950,256,208,7168,3584,torch.float8_e4m3fn,cktile,178,0,16.5657,a8w8_bpreshuffle_cktile_0x0x8x4x1x0x0x0x0x3_64x64x256_1x4x1_16x16x128_default,645.13,1775.81,0.0 +gfx950,256,216,6336,7168,torch.float8_e4m3fn,flydsl,227,0,25.644,flydsl_bpreshuflle_112x64x256_F8_F8_B16_1x0x4x2_default,765.09,1938.15,0.0 +gfx950,256,216,7168,3584,torch.float8_e4m3fn,flydsl,945,0,14.3439,flydsl_bpreshuflle_112x64x256_F8_F8_B16_1x3x4x2_default,773.72,2060.86,0.0 +gfx950,256,224,6336,7168,torch.float8_e4m3fn,flydsl,475,0,25.8574,flydsl_bpreshuflle_112x64x256_F8_F8_B16_1x1x4x2_default,786.88,1928.29,0.0 +gfx950,256,224,7168,3584,torch.float8_e4m3fn,flydsl,945,0,14.4748,flydsl_bpreshuflle_112x64x256_F8_F8_B16_1x3x4x2_default,795.12,2052.13,0.0 +gfx950,256,232,6336,7168,torch.float8_e4m3fn,flydsl,693,0,25.2873,flydsl_bpreshuflle_48x64x256_F8_F8_B16_1x2x4x2_default,833.35,1978.04,0.0 +gfx950,256,232,7168,3584,torch.float8_e4m3fn,flydsl,464,0,15.1655,flydsl_bpreshuflle_80x128x256_F8_F8_B16_1x1x4x2_default,786.01,1968.12,0.0 +gfx950,256,240,6336,7168,torch.float8_e4m3fn,flydsl,809,0,25.9572,flydsl_bpreshuflle_48x64x256_F8_F8_B16_0x3x4x2_default,839.84,1933.11,0.0 +gfx950,256,240,7168,3584,torch.float8_e4m3fn,flydsl,216,0,15.3646,flydsl_bpreshuflle_80x128x256_F8_F8_B16_1x0x4x2_default,802.58,1951.95,0.0 +gfx950,256,248,6336,7168,torch.float8_e4m3fn,flydsl,480,0,27.0873,flydsl_bpreshuflle_128x64x256_F8_F8_B16_1x1x4x2_default,831.63,1858.32,0.0 +gfx950,256,248,7168,3584,torch.float8_e4m3fn,flydsl,700,0,14.4839,flydsl_bpreshuflle_64x128x256_F8_F8_B16_1x2x4x2_default,879.76,2080.54,0.0 +gfx950,256,256,6336,7168,torch.float8_e4m3fn,ck,65,0,27.0261,a8w8_bpreshuffle_256x128x64x256_16x16_16x16_16x16x1_16x16x1_1x32x1x8_8x8x1_2x1_intrawave_v3,860.4,1868.4,0.0 +gfx950,256,256,7168,3584,torch.float8_e4m3fn,flydsl,210,0,14.8679,flydsl_bpreshuflle_64x128x256_F8_F8_B16_1x0x4x2_default,884.68,2036.44,0.0 +gfx950,256,1,2112,7168,torch.float8_e4m3fn,flydsl,617,7,7.3166,flydsl_bpreshuflle_16x64x512_F8_F8_B16_1x2x0x2_default_ks7,4.14,2070.66,0.0 +gfx950,256,2,2112,7168,torch.float8_e4m3fn,flydsl,1059,7,7.4531,flydsl_bpreshuflle_16x64x512_F8_F8_B16_1x4x0x2_default_ks7,8.12,2034.27,0.0 +gfx950,256,4,2112,7168,torch.float8_e4m3fn,flydsl,249,7,7.0448,flydsl_bpreshuflle_16x64x512_F8_F8_B16_0x1x0x2_default_ks7,17.19,2155.4,0.0 +gfx950,256,8,2112,7168,torch.float8_e4m3fn,flydsl,961,7,7.3166,flydsl_bpreshuflle_16x64x512_F8_F8_B16_0x4x0x2_default_ks7,33.11,2081.56,0.0 +gfx950,256,12,2112,7168,torch.float8_e4m3fn,flydsl,617,7,7.7125,flydsl_bpreshuflle_16x64x512_F8_F8_B16_1x2x0x2_default_ks7,47.11,1980.62,0.0 +gfx950,256,16,2112,7168,torch.float8_e4m3fn,flydsl,1,7,7.4831,flydsl_bpreshuflle_16x64x512_F8_F8_B16_0x0x0x2_default_ks7,64.74,2047.43,0.0 +gfx950,256,24,2112,7168,torch.float8_e4m3fn,flydsl,258,7,8.0538,flydsl_bpreshuflle_32x64x512_F8_F8_B16_0x1x0x2_default_ks7,90.23,1913.66,0.0 +gfx950,256,32,2112,7168,torch.float8_e4m3fn,flydsl,506,7,8.404,flydsl_bpreshuflle_32x64x512_F8_F8_B16_0x2x0x2_default_ks7,115.29,1844.76,0.0 +gfx950,256,40,2112,7168,torch.float8_e4m3fn,flydsl,905,7,9.3632,flydsl_bpreshuflle_16x64x512_F8_F8_B16_1x3x4x2_default_ks7,129.35,1665.51,0.0 +gfx950,256,48,2112,7168,torch.float8_e4m3fn,flydsl,311,7,9.3142,flydsl_bpreshuflle_16x64x512_F8_F8_B16_0x1x4x2_default_ks7,156.03,1684.06,0.0 +gfx950,256,56,2112,7168,torch.float8_e4m3fn,flydsl,686,7,9.7215,flydsl_bpreshuflle_32x64x512_F8_F8_B16_1x2x4x2_default_ks7,174.41,1622.87,0.0 +gfx950,256,64,2112,7168,torch.float8_e4m3fn,flydsl,444,7,10.16,flydsl_bpreshuflle_32x64x512_F8_F8_B16_1x1x4x2_default_ks7,190.73,1561.8,0.0 +gfx950,256,80,2112,7168,torch.float8_e4m3fn,flydsl,1218,4,11.1387,flydsl_bpreshuflle_16x64x256_F8_F8_B16_0x0x4x1_default_ks4,217.46,1440.94,0.0 +gfx950,256,96,2112,7168,torch.float8_e4m3fn,flydsl,1019,2,10.6301,flydsl_bpreshuflle_32x64x512_F8_F8_B16_0x4x4x2_default_ks2,273.44,1527.03,0.0 +gfx950,256,112,2112,7168,torch.float8_e4m3fn,ck,11,0,12.1248,a8w8_bpreshuffle_256x16x64x512_16x16_16x16_32x8x1_32x8x1_1x16x1x16_4x4x1_1x1_intrawave_v1,279.68,1353.81,0.0 +gfx950,256,128,2112,7168,torch.float8_e4m3fn,flydsl,566,2,12.2171,flydsl_bpreshuflle_32x64x512_F8_F8_B16_0x2x4x2_default_ks2,317.22,1358.51,0.0 +gfx950,256,1,6336,7168,torch.float8_e4m3fn,flydsl,125,0,12.0313,flydsl_bpreshuflle_16x64x512_F8_F8_B16_1x0x0x2_default,7.55,3776.51,0.0 +gfx950,256,2,6336,7168,torch.float8_e4m3fn,flydsl,249,2,10.9925,flydsl_bpreshuflle_16x64x512_F8_F8_B16_0x1x0x2_default_ks2,16.53,4135.19,0.0 +gfx950,256,4,6336,7168,torch.float8_e4m3fn,flydsl,249,2,10.5699,flydsl_bpreshuflle_16x64x512_F8_F8_B16_0x1x0x2_default_ks2,34.37,4304.28,0.0 +gfx950,256,8,6336,7168,torch.float8_e4m3fn,flydsl,737,2,10.72,flydsl_bpreshuflle_16x64x512_F8_F8_B16_0x3x0x2_default_ks2,67.79,4251.41,0.0 +gfx950,256,16,6336,7168,torch.float8_e4m3fn,flydsl,737,2,11.2482,flydsl_bpreshuflle_16x64x512_F8_F8_B16_0x3x0x2_default_ks2,129.21,4065.89,0.0 +gfx950,256,24,6336,7168,torch.float8_e4m3fn,flydsl,1108,0,13.4163,flydsl_bpreshuflle_16x64x512_F8_F8_B16_1x4x4x2_default,162.49,3420.66,0.0 +gfx950,256,32,6336,7168,torch.float8_e4m3fn,flydsl,746,2,11.9173,flydsl_bpreshuflle_32x64x512_F8_F8_B16_0x3x0x2_default_ks2,243.9,3864.24,0.0 +gfx950,256,40,6336,7168,torch.float8_e4m3fn,cktile,151,0,14.9131,a8w8_bpreshuffle_cktile_0x0x8x4x1x0x0x0x0x3_32x64x512_1x4x1_16x16x128_default,243.63,3098.62,0.0 +gfx950,256,56,6336,7168,torch.float8_e4m3fn,flydsl,444,0,14.6913,flydsl_bpreshuflle_32x64x512_F8_F8_B16_1x1x4x2_default,346.23,3167.01,0.0 +gfx950,256,64,6336,7168,torch.float8_e4m3fn,flydsl,914,0,15.0527,flydsl_bpreshuflle_32x64x512_F8_F8_B16_1x3x4x2_default,386.2,3101.52,0.0 +gfx950,256,1,7168,3584,torch.float8_e4m3fn,flydsl,125,0,8.0637,flydsl_bpreshuflle_16x64x512_F8_F8_B16_1x0x0x2_default,6.37,3188.12,0.0 +gfx950,256,2,7168,3584,torch.float8_e4m3fn,ck,10,0,7.5534,a8w8_bpreshuffle_128x16x32x512_16x16_16x16_32x4x1_32x4x1_1x16x1x8_4x4x1_1x1_intrawave_v1,13.6,3405.88,0.0 +gfx950,256,4,7168,3584,torch.float8_e4m3fn,flydsl,617,0,7.6211,flydsl_bpreshuflle_16x64x512_F8_F8_B16_1x2x0x2_default,26.97,3380.32,0.0 +gfx950,256,8,7168,3584,torch.float8_e4m3fn,flydsl,849,0,7.8481,flydsl_bpreshuflle_16x64x512_F8_F8_B16_1x3x0x2_default,52.37,3291.68,0.0 +gfx950,256,12,7168,3584,torch.float8_e4m3fn,flydsl,617,0,8.1372,flydsl_bpreshuflle_16x64x512_F8_F8_B16_1x2x0x2_default,75.77,3183.55,0.0 +gfx950,256,24,7168,3584,torch.float8_e4m3fn,flydsl,677,0,8.732,flydsl_bpreshuflle_16x64x512_F8_F8_B16_1x2x4x2_default,141.22,2991.32,0.0 +gfx950,256,32,7168,3584,torch.float8_e4m3fn,flydsl,311,0,8.4693,flydsl_bpreshuflle_16x64x512_F8_F8_B16_0x1x4x2_default,194.13,3101.03,0.0 +gfx950,256,56,7168,3584,torch.float8_e4m3fn,flydsl,1117,0,9.0747,flydsl_bpreshuflle_32x64x512_F8_F8_B16_1x4x4x2_default,317.07,2941.54,0.0 +gfx950,256,64,7168,3584,torch.float8_e4m3fn,flydsl,1019,0,9.2027,flydsl_bpreshuflle_32x64x512_F8_F8_B16_0x4x4x2_default,357.32,2916.21,0.0 +gfx950,256,8192,6144,7168,torch.float8_e4m3fn,flydsl,1000026,0,258.9,flydsl_bpreshuffle_8w_256x256x128_F8_F8_B16_1x8,2787.10,785.72,0.0 +gfx950,256,16384,6144,7168,torch.float8_e4m3fn,flydsl,1000026,0,484.2,flydsl_bpreshuffle_8w_256x256x128_F8_F8_B16_1x8,2980.40,749.29,0.0 +gfx950,256,32768,6144,7168,torch.float8_e4m3fn,flydsl,1000026,0,1001.4,flydsl_bpreshuffle_8w_256x256x128_F8_F8_B16_1x8,2881.70,680.62,0.0 diff --git a/aiter/configs/model_configs/a8w8_bpreshuffle_untuned_gemm_kimik3.csv b/aiter/configs/model_configs/a8w8_bpreshuffle_untuned_gemm_kimik3.csv index dbd846d68bd..3fbaffe5a29 100644 --- a/aiter/configs/model_configs/a8w8_bpreshuffle_untuned_gemm_kimik3.csv +++ b/aiter/configs/model_configs/a8w8_bpreshuffle_untuned_gemm_kimik3.csv @@ -364,3 +364,53 @@ M,N,K,q_dtype_w 8192,3072,128,torch.float8_e4m3fn 16384,3072,128,torch.float8_e4m3fn 32768,3072,128,torch.float8_e4m3fn +16,7168,3584,torch.float8_e4m3fn +32,6336,7168,torch.float8_e4m3fn +72,6336,7168,torch.float8_e4m3fn +72,7168,3584,torch.float8_e4m3fn +80,6336,7168,torch.float8_e4m3fn +80,7168,3584,torch.float8_e4m3fn +88,6336,7168,torch.float8_e4m3fn +88,7168,3584,torch.float8_e4m3fn +96,6336,7168,torch.float8_e4m3fn +96,7168,3584,torch.float8_e4m3fn +104,6336,7168,torch.float8_e4m3fn +104,7168,3584,torch.float8_e4m3fn +112,6336,7168,torch.float8_e4m3fn +112,7168,3584,torch.float8_e4m3fn +120,6336,7168,torch.float8_e4m3fn +120,7168,3584,torch.float8_e4m3fn +128,6336,7168,torch.float8_e4m3fn +128,7168,3584,torch.float8_e4m3fn +136,6336,7168,torch.float8_e4m3fn +136,7168,3584,torch.float8_e4m3fn +144,6336,7168,torch.float8_e4m3fn +144,7168,3584,torch.float8_e4m3fn +152,6336,7168,torch.float8_e4m3fn +152,7168,3584,torch.float8_e4m3fn +160,6336,7168,torch.float8_e4m3fn +160,7168,3584,torch.float8_e4m3fn +168,6336,7168,torch.float8_e4m3fn +168,7168,3584,torch.float8_e4m3fn +176,6336,7168,torch.float8_e4m3fn +176,7168,3584,torch.float8_e4m3fn +184,6336,7168,torch.float8_e4m3fn +184,7168,3584,torch.float8_e4m3fn +192,6336,7168,torch.float8_e4m3fn +192,7168,3584,torch.float8_e4m3fn +200,6336,7168,torch.float8_e4m3fn +200,7168,3584,torch.float8_e4m3fn +208,6336,7168,torch.float8_e4m3fn +208,7168,3584,torch.float8_e4m3fn +216,6336,7168,torch.float8_e4m3fn +216,7168,3584,torch.float8_e4m3fn +224,6336,7168,torch.float8_e4m3fn +224,7168,3584,torch.float8_e4m3fn +232,6336,7168,torch.float8_e4m3fn +232,7168,3584,torch.float8_e4m3fn +240,6336,7168,torch.float8_e4m3fn +240,7168,3584,torch.float8_e4m3fn +248,6336,7168,torch.float8_e4m3fn +248,7168,3584,torch.float8_e4m3fn +256,6336,7168,torch.float8_e4m3fn +256,7168,3584,torch.float8_e4m3fn diff --git a/aiter/configs/model_configs/dsv3_a8w8_bpreshuffle_tuned_gemm.csv b/aiter/configs/model_configs/dsv3_a8w8_bpreshuffle_tuned_gemm.csv index b8a35407090..7f6be003c1c 100644 --- a/aiter/configs/model_configs/dsv3_a8w8_bpreshuffle_tuned_gemm.csv +++ b/aiter/configs/model_configs/dsv3_a8w8_bpreshuffle_tuned_gemm.csv @@ -1404,15 +1404,6 @@ gfx950,256,4096,256,7168,torch.float8_e4m3fn,ck,114,0,20.7918,a8w8_bpreshuffle_2 gfx950,256,8192,256,7168,torch.float8_e4m3fn,flydsl,872,0,26.8278,flydsl_bpreshuflle_64x128x256_F8_F8_B16_1x3x0x2_default,1120.66,2413.53,0.0 gfx950,256,16384,256,7168,torch.float8_e4m3fn,flydsl,1152,0,43.3141,flydsl_bpreshuflle_128x128x256_F8_F8_B16_1x4x4x2_default,1388.22,2947.4,0.0 gfx950,256,32768,256,7168,torch.float8_e4m3fn,flydsl,175,0,68.3212,flydsl_bpreshuflle_128x256x128_F8_F8_B16_1x0x0x2_default,1760.2,3710.32,0.0 -gfx950,256,1,2112,7168,torch.float8_e4m3fn,flydsl,737,0,11.2132,flydsl_bpreshuflle_16x64x512_F8_F8_B16_0x3x0x2_default,2.7,1351.1,0.0 -gfx950,256,2,2112,7168,torch.float8_e4m3fn,flydsl,1059,0,11.2082,flydsl_bpreshuflle_16x64x512_F8_F8_B16_1x4x0x2_default,5.4,1352.72,0.0 -gfx950,256,4,2112,7168,torch.float8_e4m3fn,flydsl,1059,0,11.1888,flydsl_bpreshuflle_16x64x512_F8_F8_B16_1x4x0x2_default,10.82,1357.11,0.0 -gfx950,256,8,2112,7168,torch.float8_e4m3fn,flydsl,1,0,11.4583,flydsl_bpreshuflle_16x64x512_F8_F8_B16_0x0x0x2_default,21.14,1329.16,0.0 -gfx950,256,16,2112,7168,torch.float8_e4m3fn,ck,24,0,11.552,a8w8_bpreshuffle_128x16x32x512_16x16_16x16_32x4x1_32x4x1_1x16x1x8_4x4x1_1x1_intrawave_v2,41.94,1326.27,0.0 -gfx950,256,32,2112,7168,torch.float8_e4m3fn,flydsl,1157,0,11.7147,flydsl_bpreshuflle_16x64x512_F8_F8_B16_0x0x0x1_default,82.71,1323.41,0.0 -gfx950,256,64,2112,7168,torch.float8_e4m3fn,ck,10,0,10.6857,a8w8_bpreshuffle_128x16x32x512_16x16_16x16_32x4x1_32x4x1_1x16x1x8_4x4x1_1x1_intrawave_v1,181.34,1484.97,0.0 -gfx950,256,96,2112,7168,torch.float8_e4m3fn,ck,11,0,11.5132,a8w8_bpreshuffle_256x16x64x512_16x16_16x16_32x8x1_32x8x1_1x16x1x16_4x4x1_1x1_intrawave_v1,252.46,1409.9,0.0 -gfx950,256,128,2112,7168,torch.float8_e4m3fn,flydsl,566,0,12.5191,flydsl_bpreshuflle_32x64x512_F8_F8_B16_0x2x4x2_default,309.57,1325.73,0.0 gfx950,256,256,2112,7168,torch.float8_e4m3fn,flydsl,1476,0,15.9831,flydsl_bpreshuflle_32x64x512_F8_F8_B16_0x1x4x1_default,484.95,1129.64,0.0 gfx950,256,512,2112,7168,torch.float8_e4m3fn,cktile,179,0,22.6088,a8w8_bpreshuffle_cktile_0x0x8x4x1x0x0x0x0x4_64x64x256_1x4x1_16x16x128_default,685.67,927.58,0.0 gfx950,256,1024,2112,7168,torch.float8_e4m3fn,cktile,117,0,32.0104,a8w8_bpreshuffle_cktile_0x0x8x4x1x0x0x0x0x2_80x64x256_1x4x1_16x16x128_default,968.57,837.36,0.0 diff --git a/aiter/ops/gemm_op_a4w4.py b/aiter/ops/gemm_op_a4w4.py index 6be85c0ccda..1cbe3f9b82a 100644 --- a/aiter/ops/gemm_op_a4w4.py +++ b/aiter/ops/gemm_op_a4w4.py @@ -2,6 +2,7 @@ # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. import functools +import os import pandas as pd import torch @@ -17,6 +18,33 @@ from ..ops.gemm_op_common import get_padded_m from ..utility import dtypes +# Latent MXFP4 up-proj at M=16384 allocates a padded [M, N] bf16 output per +# layer (~224 MiB). Reuse one buffer per (device, stream, shape) like MoE scratch. +_GEMM_A4W4_OUT_CACHE: dict[ + tuple[torch.device, int, tuple[int, int], torch.dtype], torch.Tensor +] = {} + + +def _get_gemm_a4w4_out( + m_padded: int, + n: int, + dtype: torch.dtype, + device: torch.device, +) -> torch.Tensor: + shape = (m_padded, n) + if os.environ.get("AITER_GEMM_A4W4_OUT_SCRATCH_REUSE", "0").lower() not in ( + "1", + "true", + ): + return torch.empty(shape, dtype=dtype, device=device) + stream = torch.cuda.current_stream(device=device).cuda_stream + key = (device, stream, shape, dtype) + out = _GEMM_A4W4_OUT_CACHE.get(key) + if out is None: + out = torch.empty(shape, dtype=dtype, device=device) + _GEMM_A4W4_OUT_CACHE[key] = out + return out + @functools.lru_cache(maxsize=1024) def compute_gemm_SplitK(M: int, N: int, K: int, tile_m: int, tile_n: int, tile_k: int): @@ -207,7 +235,7 @@ def gemm_a4w4( beta=beta, ) return out.view(*A.shape[:-1], out.shape[-1]) - out = torch.empty(((m + 31) // 32 * 32, n), dtype=dtype, device=A.device) + out = _get_gemm_a4w4_out((m + 31) // 32 * 32, n, dtype, A.device) if gfx_arch in ["gfx942"]: raise RuntimeError( f"A4W4 GEMM kernel is not supported on gfx942, but got {gfx_arch}!" diff --git a/aiter/ops/gemm_op_a8w8.py b/aiter/ops/gemm_op_a8w8.py index 9f8109643bb..eaabed487cd 100644 --- a/aiter/ops/gemm_op_a8w8.py +++ b/aiter/ops/gemm_op_a8w8.py @@ -604,6 +604,113 @@ def get_GEMM_config_with_quant_type( return config +def _bpreshuffle_group(has_gfx, gfx, cu_num, n, k, q_dtype_w): + quant = str(q_dtype_w) + if has_gfx: + return (gfx, cu_num, n, k, quant) + return (cu_num, n, k, quant) + + +@functools.cache +def _largest_bpreshuffle_rows(tuned_file: str) -> dict: + """Largest tuned M for each device and (N, K, dtype). Misses below that M + stay on the default kernel; only an M above the table reuses the row.""" + if tuned_file not in _GEMM_QUANT_TYPE_CACHE: + get_GEMM_config_with_quant_type(1, 1, 1, dtypes.fp8, tuned_file) + table = _GEMM_QUANT_TYPE_CACHE[tuned_file] + has_gfx = _GEMM_QUANT_TYPE_HAS_GFX[tuned_file] + best: dict = {} + for key, row in table.items(): + if has_gfx: + gfx, cu_num, tuned_m, n, k, quant = key + group = (gfx, cu_num, n, k, quant) + else: + cu_num, tuned_m, n, k, quant = key + group = (cu_num, n, k, quant) + current = best.get(group) + if current is None or tuned_m > current[0]: + best[group] = (tuned_m, row) + return best + + +_REUSED_BPRESUFFLE_SHAPES: set[tuple] = set() + + +def reuse_largest_bpreshuffle_config( + m: int, + n: int, + k: int, + q_dtype_w: torch.dtype, + tuned_file: str, +): + """Tuned row for an M larger than every tuned M of this (N, K), else None.""" + rows = _largest_bpreshuffle_rows(tuned_file) + has_gfx = _GEMM_QUANT_TYPE_HAS_GFX.get(tuned_file, False) + found = rows.get( + _bpreshuffle_group(has_gfx, get_gfx(), get_cu_num(), n, k, q_dtype_w) + ) + if found is None or m <= found[0]: + return None + tuned_m, row = found + token = (n, k, tuned_m, str(q_dtype_w)) + if token not in _REUSED_BPRESUFFLE_SHAPES: + _REUSED_BPRESUFFLE_SHAPES.add(token) + logger.info( + "gemm_a8w8_bpreshuffle M:%s is above the tuned table for N:%s, K:%s; " + "reusing the M:%s row %s", + m, + n, + k, + tuned_m, + row.get("kernelName", row.get("libtype")), + ) + return row + + +def _slice_activation_scale(scale: Tensor, start: int, end: int, rows: int) -> Tensor: + if scale.ndim > 0 and scale.shape[0] == rows: + return scale[start:end] + return scale + + +def _invoke_bpreshuffle_config( + XQ: Tensor, + WQ: Tensor, + x_scale: Tensor, + w_scale: Tensor, + Y: Tensor, + config: dict, +) -> Tensor: + libtype = config["libtype"] + split_k = int(config["splitK"]) + k = XQ.shape[-1] + w_k = WQ.shape[-1] + if libtype == "ck": + return gemm_a8w8_bpreshuffle_ck(XQ, WQ, x_scale, w_scale, Y, split_k) + if libtype == "cktile": + return gemm_a8w8_bpreshuffle_cktile(XQ, WQ, x_scale, w_scale, Y, split_k) + if libtype == "flydsl": + from .flydsl.gemm_kernels import PRESHUFFLE_M_MAX + + rows = XQ.shape[0] + if rows > PRESHUFFLE_M_MAX: + for start in range(0, rows, PRESHUFFLE_M_MAX): + end = min(start + PRESHUFFLE_M_MAX, rows) + _invoke_bpreshuffle_config( + XQ[start:end], + WQ, + _slice_activation_scale(x_scale, start, end, rows), + w_scale, + Y[start:end], + config, + ) + return Y + if w_k > k: + XQ = F.pad(XQ.contiguous(), (0, w_k - k), value=0) + return gemm_a8w8_bpreshuffle_flydsl(XQ, WQ, x_scale, w_scale, Y, config) + raise RuntimeError(f"gemm_a8w8_bpreshuffle has no libtype {libtype}") + + def gemm_a8w8_fake( XQ: Tensor, WQ: Tensor, @@ -738,8 +845,13 @@ def gemm_a8w8_bpreshuffle_fake( bias: Tensor | None = None, dtype: torch.dtype = dtypes.bf16, check: bool = False, + out: Tensor | None = None, ) -> Tensor: - return torch.empty(XQ.shape[0], WQ.shape[0], dtype=dtype, device=XQ.device) + return ( + out + if out is not None + else torch.empty(XQ.shape[0], WQ.shape[0], dtype=dtype, device=XQ.device) + ) @torch_compile_guard(gen_fake=gemm_a8w8_bpreshuffle_fake) @@ -751,6 +863,7 @@ def gemm_a8w8_bpreshuffle( bias: Tensor | None = None, dtype: torch.dtype = dtypes.bf16, check: bool = False, + out: Tensor | None = None, ) -> Tensor: assert dtype in [ torch.bfloat16, @@ -762,7 +875,7 @@ def gemm_a8w8_bpreshuffle( w_k = WQ.shape[-1] if w_k < k: raise RuntimeError( - f"gemm_a8w8_bpreshuffle requires WQ K >= XQ K, got WQ K={w_k}, " f"XQ K={k}" + f"gemm_a8w8_bpreshuffle requires WQ K >= XQ K, got WQ K={w_k}, XQ K={k}" ) # if ( @@ -776,7 +889,15 @@ def gemm_a8w8_bpreshuffle( # return res assert WQ.dtype == dtypes.fp8, "gemm_a8w8_bpreshuffle only support fp8 now" assert bias is None, "gemm_a8w8_bpreshuffle does not support bias now" - Y = torch.empty(m, n, dtype=dtype, device=XQ.device) + if out is not None: + if out.shape != (m, n) or out.dtype != dtype or out.device != XQ.device: + raise ValueError( + f"out must be shape {(m, n)}, dtype {dtype}, device {XQ.device}; " + f"got shape {tuple(out.shape)}, dtype {out.dtype}, device {out.device}" + ) + Y = out + else: + Y = torch.empty(m, n, dtype=dtype, device=XQ.device) # CKTile only supports bf16 dtype config = get_GEMM_config_with_quant_type( @@ -794,17 +915,24 @@ def gemm_a8w8_bpreshuffle( dtypes.fp8, AITER_CONFIGS.AITER_CONFIG_GEMM_A8W8_BPRESHUFFLE_FILE, ) + if config is None: + config = reuse_largest_bpreshuffle_config( + m, + n, + k, + dtypes.fp8, + AITER_CONFIGS.AITER_CONFIG_GEMM_A8W8_BPRESHUFFLE_FILE, + ) + if config is None and w_k > k: + config = reuse_largest_bpreshuffle_config( + m, + n, + w_k, + dtypes.fp8, + AITER_CONFIGS.AITER_CONFIG_GEMM_A8W8_BPRESHUFFLE_FILE, + ) if config is not None: - libtype = config["libtype"] - splitK = int(config["splitK"]) - if libtype == "ck": - return gemm_a8w8_bpreshuffle_ck(XQ, WQ, x_scale, w_scale, Y, splitK) - elif libtype == "cktile": - return gemm_a8w8_bpreshuffle_cktile(XQ, WQ, x_scale, w_scale, Y, splitK) - elif libtype == "flydsl": - if w_k > k: - XQ = F.pad(XQ.contiguous(), (0, w_k - k), value=0) - return gemm_a8w8_bpreshuffle_flydsl(XQ, WQ, x_scale, w_scale, Y, config) + return _invoke_bpreshuffle_config(XQ, WQ, x_scale, w_scale, Y, config) if get_gfx() == "gfx1250": from ..ops.flydsl.gemm_tune.flydsl_gemm_a8w8_bpreshuffle_wmma_common import ( diff --git a/op_tests/test_gemm_a8w8_bpreshuffle_large_m.py b/op_tests/test_gemm_a8w8_bpreshuffle_large_m.py new file mode 100644 index 00000000000..814b69dc7ee --- /dev/null +++ b/op_tests/test_gemm_a8w8_bpreshuffle_large_m.py @@ -0,0 +1,88 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +import pytest +import torch + +from aiter import dtypes +from aiter.ops import gemm_op_a8w8 as gemm_mod + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), + reason="a8w8 bpreshuffle large-M tests require a CUDA/HIP device", +) + +TUNED = "bpreshuffle-test.csv" +KERNEL = "flydsl_bpreshuflle_128x256x128_F8_F8_B16_1x2x4x2_default" + + +def _row(): + return {"libtype": "flydsl", "splitK": 0, "kernelName": KERNEL, "kernelId": 727} + + +@pytest.fixture +def tuned_table(monkeypatch): + gemm_mod._GEMM_QUANT_TYPE_CACHE[TUNED] = { + ("gfx950", 256, 32768, 3072, 512, "torch.float8_e4m3fn"): _row(), + ("gfx950", 256, 16, 3072, 512, "torch.float8_e4m3fn"): _row(), + } + gemm_mod._GEMM_QUANT_TYPE_HAS_GFX[TUNED] = True + gemm_mod._largest_bpreshuffle_rows.cache_clear() + gemm_mod._REUSED_BPRESUFFLE_SHAPES.clear() + monkeypatch.setattr(gemm_mod, "get_gfx", lambda: "gfx950") + monkeypatch.setattr(gemm_mod, "get_cu_num", lambda: 256) + monkeypatch.setattr( + type(gemm_mod.AITER_CONFIGS), + "AITER_CONFIG_GEMM_A8W8_BPRESHUFFLE_FILE", + property(lambda self: TUNED), + ) + monkeypatch.setattr( + gemm_mod, + "get_GEMM_config_with_quant_type", + lambda *args, **kwargs: None, + ) + yield + gemm_mod._GEMM_QUANT_TYPE_CACHE.pop(TUNED, None) + gemm_mod._largest_bpreshuffle_rows.cache_clear() + + +def test_m_above_table_reuses_largest_row_and_chunks(tuned_table, monkeypatch): + seen = [] + + def fake_flydsl(XQ, WQ, x_scale, w_scale, Y, config): + seen.append((XQ.shape[0], x_scale.shape[0], config["kernelName"])) + return Y + + monkeypatch.setattr(gemm_mod, "gemm_a8w8_bpreshuffle_flydsl", fake_flydsl) + rows = 70000 + xq = torch.zeros((rows, 512), device="cuda", dtype=dtypes.fp8) + wq = torch.zeros((3072, 512), device="cuda", dtype=dtypes.fp8) + x_scale = torch.ones((rows, 1), device="cuda", dtype=torch.float32) + w_scale = torch.ones((3072, 1), device="cuda", dtype=torch.float32) + + out = gemm_mod.gemm_a8w8_bpreshuffle(xq, wq, x_scale, w_scale) + + assert out.shape == (rows, 3072) + assert seen == [(65536, 65536, KERNEL), (4464, 4464, KERNEL)] + + +def test_gap_below_largest_tuned_m_stays_on_default_kernel(tuned_table, monkeypatch): + seen = [] + monkeypatch.setattr( + gemm_mod, + "gemm_a8w8_bpreshuffle_ck", + lambda *args: seen.append("ck") or args[4], + ) + monkeypatch.setattr( + gemm_mod, + "gemm_a8w8_bpreshuffle_flydsl", + lambda *args: seen.append("flydsl") or args[4], + ) + xq = torch.zeros((1000, 512), device="cuda", dtype=dtypes.fp8) + wq = torch.zeros((3072, 512), device="cuda", dtype=dtypes.fp8) + x_scale = torch.ones((1000, 1), device="cuda", dtype=torch.float32) + w_scale = torch.ones((3072, 1), device="cuda", dtype=torch.float32) + + gemm_mod.gemm_a8w8_bpreshuffle(xq, wq, x_scale, w_scale) + + assert seen == ["ck"]