Skip to content

cktile weight preshuffle test and auto tuning for a8w8 #1400

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 17 commits into from
Nov 21, 2025
Merged

cktile weight preshuffle test and auto tuning for a8w8 #1400

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
17 commits
Select commit Hold shift + click to select a range
e633777
cktile bpreshuffle && tuning code
solinzby Nov 3, 2025
78ddd59
refine code
solinzby Nov 10, 2025
9915dc7
refine code
solinzby Nov 13, 2025
5c6e7db
refine code
solinzby Nov 13, 2025
7d42059
refine
solinzby Nov 13, 2025
848a154
refine
solinzby Nov 13, 2025
96819e0
fix merge conflict
solinzby Nov 14, 2025
305e274
fix conflict
solinzby Nov 14, 2025
1fe33b1
Merge branch 'main' into flatmm_merge
solinzby Nov 14, 2025
829c6aa
fix CI build fail
solinzby Nov 14, 2025
f312c3a
refine code
solinzby Nov 17, 2025
a6ab41b
align aiter interface
solinzby Nov 18, 2025
ccda771
refine
solinzby Nov 19, 2025
eb066a9
add get_padded_M
yzhou103 Nov 19, 2025
c9e8dcc
Merge branch 'main' into flatmm_merge
yzhou103 Nov 19, 2025
4a55093
Merge branch 'main' into flatmm_merge
yzhou103 Nov 20, 2025
fee705a
Merge branch 'main' into flatmm_merge
yzhou103 Nov 20, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
221 changes: 221 additions & 0 deletions aiter/configs/a8w8_bpreshuffle_cktile_tuned_gemm.csv
View file
Open in desktop

Large diffs are not rendered by default.

221 changes: 221 additions & 0 deletions aiter/configs/a8w8_bpreshuffle_cktile_untuned_gemm.csv
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,221 @@
M,N,K,q_dtype_w
1,9216,4096,torch.float8_e4m3fnuz
2,9216,4096,torch.float8_e4m3fnuz
4,9216,4096,torch.float8_e4m3fnuz
8,9216,4096,torch.float8_e4m3fnuz
16,9216,4096,torch.float8_e4m3fnuz
32,9216,4096,torch.float8_e4m3fnuz
64,9216,4096,torch.float8_e4m3fnuz
128,9216,4096,torch.float8_e4m3fnuz
256,9216,4096,torch.float8_e4m3fnuz
1024,9216,4096,torch.float8_e4m3fnuz
2048,9216,4096,torch.float8_e4m3fnuz
4096,9216,4096,torch.float8_e4m3fnuz
4240,9216,4096,torch.float8_e4m3fnuz
16384,9216,4096,torch.float8_e4m3fnuz
32768,9216,4096,torch.float8_e4m3fnuz
1,4608,4096,torch.float8_e4m3fnuz
2,4608,4096,torch.float8_e4m3fnuz
4,4608,4096,torch.float8_e4m3fnuz
8,4608,4096,torch.float8_e4m3fnuz
16,4608,4096,torch.float8_e4m3fnuz
32,4608,4096,torch.float8_e4m3fnuz
64,4608,4096,torch.float8_e4m3fnuz
128,4608,4096,torch.float8_e4m3fnuz
256,4608,4096,torch.float8_e4m3fnuz
1024,4608,4096,torch.float8_e4m3fnuz
2048,4608,4096,torch.float8_e4m3fnuz
4096,4608,4096,torch.float8_e4m3fnuz
16384,4608,4096,torch.float8_e4m3fnuz
32768,4608,4096,torch.float8_e4m3fnuz
1,1280,8192,torch.float8_e4m3fnuz
32,1280,8192,torch.float8_e4m3fnuz
64,1280,8192,torch.float8_e4m3fnuz
128,1280,8192,torch.float8_e4m3fnuz
192,1280,8192,torch.float8_e4m3fnuz
256,1280,8192,torch.float8_e4m3fnuz
320,1280,8192,torch.float8_e4m3fnuz
512,1280,8192,torch.float8_e4m3fnuz
1024,1280,8192,torch.float8_e4m3fnuz
2048,1280,8192,torch.float8_e4m3fnuz
4096,1280,8192,torch.float8_e4m3fnuz
8192,1280,8192,torch.float8_e4m3fnuz
16384,1280,8192,torch.float8_e4m3fnuz
1,8192,1024,torch.float8_e4m3fnuz
32,8192,1024,torch.float8_e4m3fnuz
64,8192,1024,torch.float8_e4m3fnuz
128,8192,1024,torch.float8_e4m3fnuz
192,8192,1024,torch.float8_e4m3fnuz
256,8192,1024,torch.float8_e4m3fnuz
320,8192,1024,torch.float8_e4m3fnuz
512,8192,1024,torch.float8_e4m3fnuz
1024,8192,1024,torch.float8_e4m3fnuz
2048,8192,1024,torch.float8_e4m3fnuz
4096,8192,1024,torch.float8_e4m3fnuz
8192,8192,1024,torch.float8_e4m3fnuz
16384,8192,1024,torch.float8_e4m3fnuz
16,1536,7168,torch.float8_e4m3fnuz
32,1536,7168,torch.float8_e4m3fnuz
64,1536,7168,torch.float8_e4m3fnuz
128,1536,7168,torch.float8_e4m3fnuz
256,1536,7168,torch.float8_e4m3fnuz
512,1536,7168,torch.float8_e4m3fnuz
1024,1536,7168,torch.float8_e4m3fnuz
1536,1536,7168,torch.float8_e4m3fnuz
2048,1536,7168,torch.float8_e4m3fnuz
4096,1536,7168,torch.float8_e4m3fnuz
8192,1536,7168,torch.float8_e4m3fnuz
16384,1536,7168,torch.float8_e4m3fnuz
20480,1536,7168,torch.float8_e4m3fnuz
16,3072,1536,torch.float8_e4m3fnuz
32,3072,1536,torch.float8_e4m3fnuz
64,3072,1536,torch.float8_e4m3fnuz
128,3072,1536,torch.float8_e4m3fnuz
256,3072,1536,torch.float8_e4m3fnuz
512,3072,1536,torch.float8_e4m3fnuz
1024,3072,1536,torch.float8_e4m3fnuz
1536,3072,1536,torch.float8_e4m3fnuz
2048,3072,1536,torch.float8_e4m3fnuz
4096,3072,1536,torch.float8_e4m3fnuz
8192,3072,1536,torch.float8_e4m3fnuz
16384,3072,1536,torch.float8_e4m3fnuz
20480,3072,1536,torch.float8_e4m3fnuz
16,576,7168,torch.float8_e4m3fnuz
32,576,7168,torch.float8_e4m3fnuz
64,576,7168,torch.float8_e4m3fnuz
128,576,7168,torch.float8_e4m3fnuz
256,576,7168,torch.float8_e4m3fnuz
512,576,7168,torch.float8_e4m3fnuz
1024,576,7168,torch.float8_e4m3fnuz
1536,576,7168,torch.float8_e4m3fnuz
2048,576,7168,torch.float8_e4m3fnuz
4096,576,7168,torch.float8_e4m3fnuz
8192,576,7168,torch.float8_e4m3fnuz
16384,576,7168,torch.float8_e4m3fnuz
20480,576,7168,torch.float8_e4m3fnuz
16,7168,2048,torch.float8_e4m3fnuz
32,7168,2048,torch.float8_e4m3fnuz
64,7168,2048,torch.float8_e4m3fnuz
128,7168,2048,torch.float8_e4m3fnuz
256,7168,2048,torch.float8_e4m3fnuz
512,7168,2048,torch.float8_e4m3fnuz
1024,7168,2048,torch.float8_e4m3fnuz
1536,7168,2048,torch.float8_e4m3fnuz
2048,7168,2048,torch.float8_e4m3fnuz
4096,7168,2048,torch.float8_e4m3fnuz
8192,7168,2048,torch.float8_e4m3fnuz
16384,7168,2048,torch.float8_e4m3fnuz
20480,7168,2048,torch.float8_e4m3fnuz
16,4608,7168,torch.float8_e4m3fnuz
32,4608,7168,torch.float8_e4m3fnuz
64,4608,7168,torch.float8_e4m3fnuz
128,4608,7168,torch.float8_e4m3fnuz
256,4608,7168,torch.float8_e4m3fnuz
512,4608,7168,torch.float8_e4m3fnuz
1024,4608,7168,torch.float8_e4m3fnuz
1536,4608,7168,torch.float8_e4m3fnuz
2048,4608,7168,torch.float8_e4m3fnuz
4096,4608,7168,torch.float8_e4m3fnuz
8192,4608,7168,torch.float8_e4m3fnuz
16384,4608,7168,torch.float8_e4m3fnuz
20480,4608,7168,torch.float8_e4m3fnuz
16,7168,2304,torch.float8_e4m3fnuz
32,7168,2304,torch.float8_e4m3fnuz
64,7168,2304,torch.float8_e4m3fnuz
128,7168,2304,torch.float8_e4m3fnuz
256,7168,2304,torch.float8_e4m3fnuz
512,7168,2304,torch.float8_e4m3fnuz
1024,7168,2304,torch.float8_e4m3fnuz
1536,7168,2304,torch.float8_e4m3fnuz
2048,7168,2304,torch.float8_e4m3fnuz
4096,7168,2304,torch.float8_e4m3fnuz
8192,7168,2304,torch.float8_e4m3fnuz
16384,7168,2304,torch.float8_e4m3fnuz
20480,7168,2304,torch.float8_e4m3fnuz
16,512,7168,torch.float8_e4m3fnuz
32,512,7168,torch.float8_e4m3fnuz
64,512,7168,torch.float8_e4m3fnuz
128,512,7168,torch.float8_e4m3fnuz
256,512,7168,torch.float8_e4m3fnuz
512,512,7168,torch.float8_e4m3fnuz
1024,512,7168,torch.float8_e4m3fnuz
1536,512,7168,torch.float8_e4m3fnuz
2048,512,7168,torch.float8_e4m3fnuz
4096,512,7168,torch.float8_e4m3fnuz
8192,512,7168,torch.float8_e4m3fnuz
16384,512,7168,torch.float8_e4m3fnuz
20480,512,7168,torch.float8_e4m3fnuz
16,4096,512,torch.float8_e4m3fnuz
32,4096,512,torch.float8_e4m3fnuz
64,4096,512,torch.float8_e4m3fnuz
128,4096,512,torch.float8_e4m3fnuz
256,4096,512,torch.float8_e4m3fnuz
512,4096,512,torch.float8_e4m3fnuz
1024,4096,512,torch.float8_e4m3fnuz
1536,4096,512,torch.float8_e4m3fnuz
2048,4096,512,torch.float8_e4m3fnuz
4096,4096,512,torch.float8_e4m3fnuz
8192,4096,512,torch.float8_e4m3fnuz
16384,4096,512,torch.float8_e4m3fnuz
20480,4096,512,torch.float8_e4m3fnuz
16,7168,256,torch.float8_e4m3fnuz
32,7168,256,torch.float8_e4m3fnuz
64,7168,256,torch.float8_e4m3fnuz
128,7168,256,torch.float8_e4m3fnuz
256,7168,256,torch.float8_e4m3fnuz
512,7168,256,torch.float8_e4m3fnuz
1024,7168,256,torch.float8_e4m3fnuz
1536,7168,256,torch.float8_e4m3fnuz
2048,7168,256,torch.float8_e4m3fnuz
4096,7168,256,torch.float8_e4m3fnuz
8192,7168,256,torch.float8_e4m3fnuz
16384,7168,256,torch.float8_e4m3fnuz
20480,7168,256,torch.float8_e4m3fnuz
1,4096,512,torch.float8_e4m3fnuz
1,2112,7168,torch.float8_e4m3fnuz
1,4608,7168,torch.float8_e4m3fnuz
1,7168,2304,torch.float8_e4m3fnuz
1,512,7168,torch.float8_e4m3fnuz
1,7168,256,torch.float8_e4m3fnuz
16,2112,7168,torch.float8_e4m3fnuz
32,2112,7168,torch.float8_e4m3fnuz
48,4096,512,torch.float8_e4m3fnuz
48,2112,7168,torch.float8_e4m3fnuz
48,4608,7168,torch.float8_e4m3fnuz
48,7168,2304,torch.float8_e4m3fnuz
48,512,7168,torch.float8_e4m3fnuz
48,7168,256,torch.float8_e4m3fnuz
64,2112,7168,torch.float8_e4m3fnuz
80,4096,512,torch.float8_e4m3fnuz
80,2112,7168,torch.float8_e4m3fnuz
80,4608,7168,torch.float8_e4m3fnuz
80,7168,2304,torch.float8_e4m3fnuz
80,512,7168,torch.float8_e4m3fnuz
80,7168,256,torch.float8_e4m3fnuz
96,4096,512,torch.float8_e4m3fnuz
96,2112,7168,torch.float8_e4m3fnuz
96,4608,7168,torch.float8_e4m3fnuz
96,7168,2304,torch.float8_e4m3fnuz
96,512,7168,torch.float8_e4m3fnuz
96,7168,256,torch.float8_e4m3fnuz
112,4096,512,torch.float8_e4m3fnuz
112,2112,7168,torch.float8_e4m3fnuz
112,4608,7168,torch.float8_e4m3fnuz
112,7168,2304,torch.float8_e4m3fnuz
112,512,7168,torch.float8_e4m3fnuz
112,7168,256,torch.float8_e4m3fnuz
128,2112,7168,torch.float8_e4m3fnuz
256,2112,7168,torch.float8_e4m3fnuz
512,2112,7168,torch.float8_e4m3fnuz
1024,2112,7168,torch.float8_e4m3fnuz
1536,2112,7168,torch.float8_e4m3fnuz
2048,2112,7168,torch.float8_e4m3fnuz
4096,2112,7168,torch.float8_e4m3fnuz
8192,2112,7168,torch.float8_e4m3fnuz
16384,2112,7168,torch.float8_e4m3fnuz
32768,4096,512,torch.float8_e4m3fnuz
32768,2112,7168,torch.float8_e4m3fnuz
32768,4608,7168,torch.float8_e4m3fnuz
32768,7168,2304,torch.float8_e4m3fnuz
32768,512,7168,torch.float8_e4m3fnuz
32768,7168,256,torch.float8_e4m3fnuz
11 changes: 11 additions & 0 deletions aiter/jit/core.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -152,6 +152,12 @@ def get_config_file(env_name, default_file, tuned_file_name):
"AITER_CONFIG_GEMM_A8W8_BPRESHUFFLE",
f"{AITER_ROOT_DIR}/aiter/configs/a8w8_bpreshuffle_tuned_gemm.csv",
)

AITER_CONFIG_GEMM_A8W8_BPRESHUFFLE_CKTILE = os.getenv(
"AITER_CONFIG_GEMM_A8W8_BPRESHUFFLE_CKTILE",
f"{AITER_ROOT_DIR}/aiter/configs/a8w8_bpreshuffle_cktile_tuned_gemm.csv",
)

AITER_CONFIG_GEMM_A8W8_BLOCKSCALE = os.getenv(
"AITER_CONFIG_GEMM_A8W8_BLOCKSCALE",
f"{AITER_ROOT_DIR}/aiter/configs/a8w8_blockscale_tuned_gemm.csv",
Expand Down Expand Up @@ -192,6 +198,11 @@ def get_config_file(env_name, default_file, tuned_file_name):
AITER_CONFIG_GEMM_A8W8_BPRESHUFFLE,
"a8w8_bpreshuffle_tuned_gemm",
)
AITER_CONFIG_GEMM_A8W8_BPRESHUFFLE_CKTILE_FILE = get_config_file(
"AITER_CONFIG_GEMM_A8W8_BPRESHUFFLE_CKTILE",
AITER_CONFIG_GEMM_A8W8_BPRESHUFFLE_CKTILE,
"a8w8_bpreshuffle_cktile_tuned_gemm",
)
AITER_CONFIG_GEMM_A8W8_BLOCKSCALE_FILE = get_config_file(
"AITER_CONFIG_GEMM_A8W8_BLOCKSCALE",
AITER_CONFIG_GEMM_A8W8_BLOCKSCALE,
Expand Down
37 changes: 37 additions & 0 deletions aiter/jit/optCompilerConfig.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -288,6 +288,25 @@
"hip_clang_path": "os.environ.get('FLATMM_HIP_CLANG_PATH')",
"blob_gen_cmd": "f'{AITER_CSRC_DIR}/ck_deepgemm/gen_instances.py --working_path {{}}'"
},
"module_gemm_a8w8_bpreshuffle_cktile": {
"srcs": [
"f'{AITER_CSRC_DIR}/pybind/gemm_a8w8_bpreshuffle_cktile_pybind.cu'",
"f'{AITER_CSRC_DIR}/py_itfs_cu/gemm_common.cu'",
"f'{AITER_CSRC_DIR}/cktile_gemm_a8w8_bpreshuffle/gemm_a8w8_bpreshuffle_cktile.cu'"
],
"flags_extra_cc": [],
"flags_extra_hip": [],
"extra_ldflags": "None",
"extra_include": [
"f'{AITER_CSRC_DIR}/cktile_gemm_a8w8_bpreshuffle/include'",
"f'{CK_DIR}/example/ck_tile/18_flatmm'"
],
"is_python_module": "True",
"is_standalone": "False",
"verbose": "False",
"hip_clang_path": "os.environ.get('FLATMM_HIP_CLANG_PATH')",
"blob_gen_cmd": "f'{AITER_CSRC_DIR}/cktile_gemm_a8w8_bpreshuffle/gen_instances.py --working_path {{}} --tune_file {AITER_CONFIG_GEMM_A8W8_BPRESHUFFLE_CKTILE_FILE}'"
},
"module_gemm_a8w8_asm": {
"srcs": [
"f'{AITER_CSRC_DIR}/pybind/gemm_a8w8_asm_pybind.cu'",
Expand Down Expand Up @@ -582,6 +601,24 @@
"is_standalone": "False",
"blob_gen_cmd": "f'{AITER_CSRC_DIR}/ck_gemm_a8w8_bpreshuffle/gen_instances.py --working_path {{}} --tune'"
},
"module_gemm_a8w8_bpreshuffle_cktile_tune": {
"srcs": [
"f'{AITER_CSRC_DIR}/pybind/gemm_a8w8_bpreshuffle_cktile_tune_pybind.cu'",
"f'{AITER_CSRC_DIR}/cktile_gemm_a8w8_bpreshuffle/gemm_a8w8_bpreshuffle_cktile_tune.cu'"
],
"flags_extra_cc": [],
"flags_extra_hip": [],
"extra_ldflags": "None",
"extra_include": [
"f'{AITER_CSRC_DIR}/cktile_gemm_a8w8_bpreshuffle/include'",
"f'{CK_DIR}/example/ck_tile/18_flatmm'"
],
"verbose": "False",
"hip_clang_path": "os.environ.get('FLATMM_HIP_CLANG_PATH')",
"is_python_module": "True",
"is_standalone": "False",
"blob_gen_cmd": "f'{AITER_CSRC_DIR}/cktile_gemm_a8w8_bpreshuffle/gen_instances.py --working_path {{}} --tune'"
},
"module_aiter_operator": {
"srcs": [
"f'{AITER_CSRC_DIR}/pybind/aiter_operator_pybind.cu'",
Expand Down
Loading
Loading