Skip to content

Tune gfx950 batched GEMM configs and benchmarks #3283

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

Draft
nidal567 wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

Tune gfx950 batched GEMM configs and benchmarks #3283

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
2 changes: 1 addition & 1 deletion aiter/ops/triton/_triton_kernels/gemm/batched/batched_gemm_a16wfp4.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -317,4 +317,4 @@ def _get_config(
):
# Note: Config files use K=2*K in their naming because FP4 weights are packed,
# so the actual K dimension in the config file corresponds to 2*K unpacked elements
return get_gemm_config("BATCHED_GEMM_PREQUANT-AFP4WFP4", M, N, 2 * K)
return get_gemm_config("BATCHED_GEMM-A16WFP4", M, N, 2 * K)
Comment thread
brunomazzottiamd marked this conversation as resolved.
1 change: 1 addition & 0 deletions aiter/ops/triton/_triton_kernels/gemm/batched/batched_gemm_bf16.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -179,4 +179,5 @@ def _get_config(
K: int,
):

# BF16 uses the shared 16-bit activation / 16-bit weight batched GEMM config.
return get_gemm_config("BATCHED_GEMM-A16W16", M, N, K)
80 changes: 80 additions & 0 deletions aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM-A16WFP4-N=128-K=512.json
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,80 @@
{
"M_LEQ_16": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 256,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 2,
"waves_per_eu": 6,
"matrix_instr_nonkdim": 16,
"kpack": 1,
"cache_modifier": ".cg",
"NUM_KSPLIT": 1
},
"M_LEQ_32": {
"BLOCK_SIZE_M": 32,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 256,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 2,
"waves_per_eu": 4,
"matrix_instr_nonkdim": 16,
"kpack": 1,
"cache_modifier": ".cg",
"NUM_KSPLIT": 1
},
"M_LEQ_64": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 256,
"GROUP_SIZE_M": 1,
"num_warps": 8,
"num_stages": 2,
"waves_per_eu": 4,
"matrix_instr_nonkdim": 16,
"kpack": 1,
"cache_modifier": ".cg",
"NUM_KSPLIT": 1
},
"M_LEQ_128": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 256,
"GROUP_SIZE_M": 1,
"num_warps": 8,
"num_stages": 2,
"waves_per_eu": 4,
"matrix_instr_nonkdim": 16,
"kpack": 1,
"cache_modifier": ".cg",
"NUM_KSPLIT": 1
},
"M_LEQ_256": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 256,
"GROUP_SIZE_M": 1,
"num_warps": 8,
"num_stages": 2,
"waves_per_eu": 4,
"matrix_instr_nonkdim": 16,
"kpack": 1,
"cache_modifier": ".cg",
"NUM_KSPLIT": 1
},
"any": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 256,
"GROUP_SIZE_M": 1,
"num_warps": 8,
"num_stages": 2,
"waves_per_eu": 4,
"matrix_instr_nonkdim": 16,
"kpack": 1,
"cache_modifier": ".cg",
"NUM_KSPLIT": 1
}
}
80 changes: 80 additions & 0 deletions aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM-A16WFP4.json
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,80 @@
{
"M_LEQ_16": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 256,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 2,
"waves_per_eu": 6,
"matrix_instr_nonkdim": 16,
"kpack": 1,
"cache_modifier": ".cg",
"NUM_KSPLIT": 1
},
"M_LEQ_32": {
"BLOCK_SIZE_M": 32,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 256,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 2,
"waves_per_eu": 4,
"matrix_instr_nonkdim": 16,
"kpack": 1,
"cache_modifier": ".cg",
"NUM_KSPLIT": 1
},
"M_LEQ_64": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 256,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 2,
"waves_per_eu": 4,
"matrix_instr_nonkdim": 16,
"kpack": 1,
"cache_modifier": ".cg",
"NUM_KSPLIT": 1
},
"M_LEQ_128": {
"BLOCK_SIZE_M": 128,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 256,
"GROUP_SIZE_M": 2,
"num_warps": 4,
"num_stages": 2,
"waves_per_eu": 4,
"matrix_instr_nonkdim": 16,
"kpack": 1,
"cache_modifier": ".cg",
"NUM_KSPLIT": 1
},
"M_LEQ_256": {
"BLOCK_SIZE_M": 128,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 256,
"GROUP_SIZE_M": 2,
"num_warps": 4,
"num_stages": 2,
"waves_per_eu": 4,
"matrix_instr_nonkdim": 16,
"kpack": 1,
"cache_modifier": ".cg",
"NUM_KSPLIT": 1
},
"any": {
"BLOCK_SIZE_M": 256,
"BLOCK_SIZE_N": 256,
"BLOCK_SIZE_K": 256,
"GROUP_SIZE_M": 64,
"num_warps": 8,
"num_stages": 2,
"waves_per_eu": 1,
"matrix_instr_nonkdim": 16,
"kpack": 1,
"cache_modifier": null,
"NUM_KSPLIT": 1
}
}
16 changes: 8 additions & 8 deletions ...-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT-N=128-K=512.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,26 +33,26 @@
"cache_modifier": ".cg"
},
"M_LEQ_128": {
"BLOCK_SIZE_M": 32,
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"GROUP_SIZE_M": 1,
"num_warps": 8,
"num_stages": 2,
"waves_per_eu": 2,
"matrix_instr_nonkdim": 16,
"waves_per_eu": 1,
"matrix_instr_nonkdim": 32,
"kpack": 1,
"cache_modifier": ".cg"
"cache_modifier": null
Comment thread
brunomazzottiamd marked this conversation as resolved.
},
"M_LEQ_256": {
"BLOCK_SIZE_M": 32,
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"GROUP_SIZE_M": 1,
"num_warps": 8,
"num_stages": 2,
"waves_per_eu": 2,
"matrix_instr_nonkdim": 16,
"waves_per_eu": 1,
"matrix_instr_nonkdim": 32,
"kpack": 1,
"cache_modifier": ".cg"
"cache_modifier": null
},
"any": {
"BLOCK_SIZE_M": 32,
Expand Down
30 changes: 15 additions & 15 deletions ...-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT-N=512-K=128.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,46 +11,46 @@
"cache_modifier": ".cg"
},
"M_LEQ_32": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_M": 32,
"BLOCK_SIZE_N": 256,
"GROUP_SIZE_M": 1,
"num_warps": 8,
"num_stages": 2,
"num_stages": 1,
"waves_per_eu": 1,
"matrix_instr_nonkdim": 16,
"matrix_instr_nonkdim": 32,
"kpack": 1,
"cache_modifier": ".cg"
},
"M_LEQ_64": {
"BLOCK_SIZE_M": 32,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_N": 256,
"GROUP_SIZE_M": 1,
"num_warps": 8,
"num_stages": 1,
"waves_per_eu": 2,
"matrix_instr_nonkdim": 16,
"waves_per_eu": 1,
"matrix_instr_nonkdim": 32,
"kpack": 1,
"cache_modifier": ".cg"
},
"M_LEQ_128": {
"BLOCK_SIZE_M": 32,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_N": 256,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_warps": 8,
"num_stages": 1,
"waves_per_eu": 6,
"matrix_instr_nonkdim": 16,
"waves_per_eu": 1,
"matrix_instr_nonkdim": 32,
"kpack": 1,
"cache_modifier": ".cg"
},
"M_LEQ_256": {
"BLOCK_SIZE_M": 32,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_N": 256,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_warps": 8,
"num_stages": 1,
"waves_per_eu": 6,
"matrix_instr_nonkdim": 16,
"waves_per_eu": 1,
"matrix_instr_nonkdim": 32,
"kpack": 1,
"cache_modifier": ".cg"
},
Expand Down
18 changes: 9 additions & 9 deletions .../gemm/gfx950-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,34 +22,34 @@
"cache_modifier": ".cg"
},
"M_LEQ_64": {
"BLOCK_SIZE_M": 32,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 256,
"GROUP_SIZE_M": 1,
"num_warps": 8,
"num_stages": 2,
"waves_per_eu": 2,
"waves_per_eu": 1,
"matrix_instr_nonkdim": 16,
"kpack": 1,
"cache_modifier": ".cg"
},
"M_LEQ_128": {
"BLOCK_SIZE_M": 32,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 256,
"GROUP_SIZE_M": 1,
"num_warps": 8,
"num_stages": 2,
"waves_per_eu": 2,
"waves_per_eu": 1,
"matrix_instr_nonkdim": 16,
"kpack": 1,
"cache_modifier": ".cg"
},
"M_LEQ_256": {
"BLOCK_SIZE_M": 32,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 256,
"GROUP_SIZE_M": 1,
"num_warps": 8,
"num_stages": 2,
"waves_per_eu": 2,
"waves_per_eu": 1,
"matrix_instr_nonkdim": 16,
"kpack": 1,
"cache_modifier": ".cg"
Expand Down
14 changes: 7 additions & 7 deletions op_tests/op_benchmarks/triton/bench_batched_gemm_a16wfp4.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,8 @@
import triton
import math
import aiter.ops.triton.utils._triton.arch_info as arch_info
from aiter.ops.triton.gemm.batched.batched_gemm_afp4wfp4_pre_quant import (
batched_gemm_afp4wfp4_pre_quant,
from aiter.ops.triton.gemm.batched.batched_gemm_a16wfp4 import (
batched_gemm_a16wfp4,
)
from op_tests.triton_tests.gemm.batched.test_batched_gemm_a16wfp4 import (
generate_batched_gemm_a16wfp4_inputs,
Expand Down Expand Up @@ -47,7 +47,7 @@ def bench_gemm_fn(
mem = mem_read + mem_write

ms = triton.testing.do_bench(
lambda: batched_gemm_afp4wfp4_pre_quant(x, w, w_scale, c_dtype, y),
lambda: batched_gemm_a16wfp4(x, w, w_scale, c_dtype, y),
warmup=25,
rep=100,
)
Expand All @@ -74,7 +74,7 @@ def run_model_benchmark(args):
)

@triton.testing.perf_report([benchmark])
def bench_batched_gemm_afp4wfp4_pre_quant(
def bench_batched_gemm_a16wfp4(
M, hidden_dim, intermediate_dim, batch, metric, layer, **kwargs
):
if layer == "fc1":
Expand All @@ -92,7 +92,7 @@ def bench_batched_gemm_afp4wfp4_pre_quant(

return bench_gemm_fn(batch, M, N, K, metric, args.layout)

bench_batched_gemm_afp4wfp4_pre_quant.run(
bench_batched_gemm_a16wfp4.run(
save_path="." if args.o else None, print_data=True
)

Expand All @@ -105,7 +105,7 @@ def run_shape_benchmark(args):
)

@triton.testing.perf_report([benchmark])
def bench_batched_gemm_afp4wfp4_pre_quant(
def bench_batched_gemm_a16wfp4(
batch,
M,
N,
Expand All @@ -115,7 +115,7 @@ def bench_batched_gemm_afp4wfp4_pre_quant(
):
return bench_gemm_fn(batch, M, N, K, metric, args.layout)

bench_batched_gemm_afp4wfp4_pre_quant.run(
bench_batched_gemm_a16wfp4.run(
save_path="." if args.o else None, print_data=True
)

Expand Down
Loading
Loading