Skip to content

Alizaidy/bf16 gemm tuning 081426 #4749

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
junhaha666 merged 17 commits into from
Aug 19, 2026
Merged

Alizaidy/bf16 gemm tuning 081426 #4749

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
d39d8a7
Use triton for DSv4 Pro bf16 GEMM
azaidy Jul 10, 2026
620fa5e
Add bf16 GEMM tuning for DSv4 Pro
azaidy Jul 10, 2026
e45728e
Tune triton BF16 GEMMs more for DSv4
azaidy Jul 10, 2026
d3713ec
tuned m=64,128
ahmed-bsod Jul 13, 2026
3b24f42
tuned m=1,32
ahmed-bsod Jul 13, 2026
264bb9a
tuned m=256
ahmed-bsod Jul 13, 2026
21c0bba
tuned m=16k
ahmed-bsod Jul 13, 2026
86f440b
black formatting fix
ahmed-bsod Jul 13, 2026
04609b4
fixes WAR hazard in kernel rn due to compiler issuing tdm load to sam…
ahmed-bsod Jul 16, 2026
765c020
Add time, tf, BW
azaidy Aug 14, 2026
ce1e638
Resolve conflicts with main
azaidy Aug 14, 2026
614d573
Merge branch 'main' into alizaidy/bf16-gemm-tuning-081426
azaidy Aug 14, 2026
13725ff
Add DSv4 Flash configs
azaidy Aug 14, 2026
91b84a2
Fix configs
azaidy Aug 15, 2026
476e406
Restore opus on shapes where it performs better
azaidy Aug 15, 2026
4a25059
Adjust configs
azaidy Aug 15, 2026
87d2899
Merge branch 'main' into alizaidy/bf16-gemm-tuning-081426
azaidy Aug 15, 2026
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
339 changes: 162 additions & 177 deletions aiter/configs/model_configs/dsv4_bf16_tuned_gemm.csv
View file
Open in desktop

Large diffs are not rendered by default.

16 changes: 16 additions & 0 deletions aiter/configs/model_configs/dsv4pro_bf16_tuned_gemm.csv
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,16 @@
gfx,cu_num,M,N,K,bias,dtype,outdtype,scaleAB,bpreshuffle,libtype,solidx,splitK,us,kernelName,err_ratio,tflops,bw
gfx1250,256,1,32320,7168,False,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,25.5588,auto,0.0,18.13,18131.47
gfx1250,256,2,32320,7168,False,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,25.6286,auto,0.0,36.16,18085.17
gfx1250,256,4,32320,7168,False,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,25.7178,auto,0.0,72.07,18028.58
gfx1250,256,8,32320,7168,False,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,26.3710,auto,0.0,140.56,17594.0
gfx1250,256,16,32320,7168,False,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,27.1162,auto,0.0,273.39,17133.78
gfx1250,256,32,32320,7168,False,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,28.2982,auto,0.0,523.95,16462.77
gfx1250,256,64,32320,7168,False,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,32.3609,auto,0.0,916.34,14474.07
gfx1250,256,128,32320,7168,False,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,42.7325,auto,0.0,1387.88,11079.35
gfx1250,256,256,32320,7168,False,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,67.0230,auto,0.0,1769.76,7214.8
gfx1250,256,512,32320,7168,False,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,199.3847,auto,0.0,1189.81,2526.65
gfx1250,256,1024,32320,7168,False,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,398.0347,auto,0.0,1192.01,1367.24
gfx1250,256,2048,32320,7168,False,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,768.4669,auto,0.0,1234.82,813.41
gfx1250,256,4096,32320,7168,False,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,1498.5199,auto,0.0,1266.48,525.07
gfx1250,256,8192,32320,7168,False,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,2986.6679,auto,0.0,1270.87,371.76
gfx1250,256,16384,32320,7168,False,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,3336.9446,auto,0.0,2274.94,526.61
16 changes: 8 additions & 8 deletions aiter/ops/triton/_gluon_kernels/gfx1250/gemm/basic/gemm_a16w16.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -500,7 +500,7 @@ def _gemm_a16w16_compute_bound_kernel(
accumulator = gl.zeros((BLOCK_M, BLOCK_N), dtype=gl.float32, layout=WMMA_LAYOUT)

# TDM prologue: fill the pipeline with NUM_BUFFERS-1 tiles
for _ in gl.static_range(NUM_BUFFERS):
for _ in gl.static_range(NUM_BUFFERS - 1):
gl.amd.gfx1250.tdm.async_load(
a_desc, [0, 0], a_buffer.index(load_idx % NUM_BUFFERS)
)
Expand Down Expand Up @@ -534,7 +534,7 @@ def _gemm_a16w16_compute_bound_kernel(
# Register pre-load prologue: wait for tile 0 then read it into cur_a/cur_b.
# After TDM prologue there are (NUM_BUFFERS-1)*2 ops in-flight; waiting for
# (NUM_BUFFERS-2)*2 lets exactly one tile (tile 0) complete.
gl.amd.gfx1250.tdm.async_wait((NUM_BUFFERS - 1) * 2)
gl.amd.gfx1250.tdm.async_wait((NUM_BUFFERS - 2) * 2)

if LAYOUT[0] == "T":
cur_a = gl.amd.cdna4.async_copy.load_shared_relaxed(
Expand Down Expand Up @@ -594,7 +594,7 @@ def _gemm_a16w16_compute_bound_kernel(
# Tighter wait: after issuing the new TDM there are (NUM_BUFFERS-1)*2
# ops in-flight. Waiting for (NUM_BUFFERS-2)*2 guarantees that tile
# compute_idx+1 has landed in LDS.
gl.amd.gfx1250.tdm.async_wait((NUM_BUFFERS - 1) * 2)
gl.amd.gfx1250.tdm.async_wait((NUM_BUFFERS - 2) * 2)

load_idx += 1

Expand Down Expand Up @@ -629,7 +629,7 @@ def _gemm_a16w16_compute_bound_kernel(
# The final K tile is peeled out after this loop so its (possibly partial)
# TDM load can be bounds-checked with set_bounds; the interior iterations
# use the fast add_offsets path that leaves the OOB bound untouched.
for _ in range(num_k_tiles - NUM_BUFFERS - 2):
for _ in range(num_k_tiles - NUM_BUFFERS - 1):

# WMMA for the current tile — uses operands pre-loaded in the
# *previous* iteration so no ds_read stall before the matrix op.
Expand Down Expand Up @@ -665,7 +665,7 @@ def _gemm_a16w16_compute_bound_kernel(
# Tighter wait: after issuing the new TDM there are (NUM_BUFFERS-1)*2
# ops in-flight. Waiting for (NUM_BUFFERS-2)*2 guarantees that tile
# compute_idx+1 has landed in LDS.
gl.amd.gfx1250.tdm.async_wait((NUM_BUFFERS - 1) * 2)
gl.amd.gfx1250.tdm.async_wait((NUM_BUFFERS - 2) * 2)

load_idx += 1

Expand Down Expand Up @@ -729,7 +729,7 @@ def _gemm_a16w16_compute_bound_kernel(
b_desc, [0, 0], b_buffer.index(load_idx % NUM_BUFFERS)
)

gl.amd.gfx1250.tdm.async_wait((NUM_BUFFERS - 1) * 2)
gl.amd.gfx1250.tdm.async_wait((NUM_BUFFERS - 2) * 2)

load_idx += 1

Expand Down Expand Up @@ -759,8 +759,8 @@ def _gemm_a16w16_compute_bound_kernel(

# Epilogue: no more TDM loads; drain the remaining NUM_BUFFERS-1 tiles.
# The first NUM_BUFFERS-2 iterations still use the pre-load / WMMA pattern.
for i in gl.static_range(NUM_BUFFERS - 1):
gl.amd.gfx1250.tdm.async_wait((NUM_BUFFERS - 2 - i) * 2)
for i in gl.static_range(NUM_BUFFERS - 2):
gl.amd.gfx1250.tdm.async_wait((NUM_BUFFERS - 3 - i) * 2)

if LAYOUT[0] == "T":
next_a = gl.amd.cdna4.async_copy.load_shared_relaxed(
Expand Down
138 changes: 113 additions & 25 deletions aiter/ops/triton/configs/gemm/gfx1250-GEMM-A16W16-N=1024-K=4096.json
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,26 +1,114 @@
{
"M_LEQ_64": {
"BLOCK_M": 16,
"BLOCK_N": 16,
"BLOCK_K": 512,
"NUM_BUFFERS": 3,
"num_warps": 2,
"kernel_type": "bandwidth_bound"
},
"any": {
"BLOCK_M": 64,
"BLOCK_N": 64,
"BLOCK_K": 256,
"NUM_BUFFERS": 3,
"num_warps": 4,
"kernel_type": "compute_bound"
},
"M_LEQ_1024": {
"BLOCK_M": 64,
"BLOCK_N": 64,
"BLOCK_K": 256,
"NUM_BUFFERS": 3,
"num_warps": 4,
"kernel_type": "compute_bound"
}
}
"M_LEQ_1": {
"BLOCK_M": 16,
"BLOCK_N": 16,
"BLOCK_K": 512,
"NUM_BUFFERS": 3,
"num_warps": 2,
"kernel_type": "bandwidth_bound"
},
"M_LEQ_4": {
"BLOCK_M": 16,
"BLOCK_N": 16,
"BLOCK_K": 512,
"NUM_BUFFERS": 3,
"num_warps": 2,
"kernel_type": "bandwidth_bound"
},
"M_LEQ_8": {
"BLOCK_M": 16,
"BLOCK_N": 16,
"BLOCK_K": 512,
"NUM_BUFFERS": 3,
"num_warps": 2,
"kernel_type": "bandwidth_bound"
},
"M_LEQ_16": {
"BLOCK_M": 16,
"BLOCK_N": 16,
"BLOCK_K": 512,
"NUM_BUFFERS": 5,
"num_warps": 2,
"kernel_type": "compute_bound"
},
"M_LEQ_32": {
"BLOCK_M": 16,
"BLOCK_N": 16,
"BLOCK_K": 512,
"NUM_BUFFERS": 3,
"num_warps": 2,
"kernel_type": "bandwidth_bound"
},
"M_LEQ_64": {
"BLOCK_M": 16,
"BLOCK_N": 32,
"BLOCK_K": 512,
"NUM_BUFFERS": 4,
"num_warps": 4,
"kernel_type": "bandwidth_bound"
},
"M_LEQ_128": {
"BLOCK_M": 16,
"BLOCK_N": 32,
"BLOCK_K": 512,
"NUM_BUFFERS": 4,
"num_warps": 4,
"kernel_type": "compute_bound"
},
"M_LEQ_256": {
"BLOCK_M": 32,
"BLOCK_N": 32,
"BLOCK_K": 512,
"NUM_BUFFERS": 4,
"num_warps": 4,
"kernel_type": "compute_bound"
},
"M_LEQ_512": {
"BLOCK_M": 64,
"BLOCK_N": 64,
"BLOCK_K": 256,
"NUM_BUFFERS": 4,
"num_warps": 4,
"kernel_type": "compute_bound"
},
"M_LEQ_1024": {
"BLOCK_M": 64,
"BLOCK_N": 64,
"BLOCK_K": 256,
"NUM_BUFFERS": 4,
"num_warps": 4,
"kernel_type": "compute_bound"
},
"M_LEQ_2048": {
"BLOCK_M": 128,
"BLOCK_N": 128,
"BLOCK_K": 128,
"NUM_BUFFERS": 4,
"num_warps": 4,
"kernel_type": "compute_bound"
},
"M_LEQ_4096": {
"BLOCK_M": 128,
"BLOCK_N": 128,
"BLOCK_K": 128,
"NUM_BUFFERS": 4,
"num_warps": 4,
"kernel_type": "compute_bound"
},
"M_LEQ_8192": {
"BLOCK_M": 128,
"BLOCK_N": 128,
"BLOCK_K": 128,
"NUM_BUFFERS": 3,
"num_warps": 4,
"kernel_type": "compute_bound"
},
"any": {
"BLOCK_M": 128,
"BLOCK_N": 128,
"BLOCK_K": 128,
"NUM_BUFFERS": 4,
"num_warps": 4,
"kernel_type": "compute_bound"
}
}
90 changes: 90 additions & 0 deletions aiter/ops/triton/configs/gemm/gfx1250-GEMM-A16W16-N=1024-K=7168.json
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,90 @@
{
"M_LEQ_4": {
"BLOCK_M": 16,
"BLOCK_N": 16,
"BLOCK_K": 512,
"NUM_BUFFERS": 7,
"num_warps": 1,
"kernel_type": "compute_bound"
},
"M_LEQ_8": {
"BLOCK_M": 16,
"BLOCK_N": 16,
"BLOCK_K": 256,
"NUM_BUFFERS": 8,
"num_warps": 1,
"kernel_type": "compute_bound"
},
"M_LEQ_16": {
"BLOCK_M": 16,
"BLOCK_N": 16,
"BLOCK_K": 256,
"NUM_BUFFERS": 7,
"num_warps": 1,
"kernel_type": "compute_bound"
},
"M_LEQ_32": {
"BLOCK_M": 16,
"BLOCK_N": 16,
"BLOCK_K": 512,
"NUM_BUFFERS": 5,
"num_warps": 2,
"kernel_type": "compute_bound"
},
"M_LEQ_64": {
"BLOCK_M": 16,
"BLOCK_N": 32,
"BLOCK_K": 256,
"NUM_BUFFERS": 7,
"num_warps": 2,
"kernel_type": "compute_bound"
},
"M_LEQ_128": {
"BLOCK_M": 32,
"BLOCK_N": 32,
"BLOCK_K": 256,
"NUM_BUFFERS": 6,
"num_warps": 4,
"kernel_type": "compute_bound"
},
"M_LEQ_512": {
"BLOCK_M": 64,
"BLOCK_N": 32,
"BLOCK_K": 256,
"NUM_BUFFERS": 4,
"num_warps": 8,
"kernel_type": "compute_bound"
},
"M_LEQ_2048": {
"BLOCK_M": 64,
"BLOCK_N": 128,
"BLOCK_K": 128,
"NUM_BUFFERS": 4,
"num_warps": 4,
"kernel_type": "compute_bound"
},
"M_LEQ_8192": {
"BLOCK_M": 256,
"BLOCK_N": 128,
"BLOCK_K": 128,
"NUM_BUFFERS": 2,
"num_warps": 8,
"kernel_type": "compute_bound"
},
"any": {
"BLOCK_M": 256,
"BLOCK_N": 256,
"BLOCK_K": 64,
"NUM_BUFFERS": 3,
"num_warps": 4,
"kernel_type": "compute_bound"
},
"M_LEQ_256": {
"BLOCK_M": 32,
"BLOCK_N": 64,
"BLOCK_K": 256,
"NUM_BUFFERS": 6,
"num_warps": 4,
"kernel_type": "compute_bound"
}
}
Loading
Loading