Skip to content

[Triton] Retune gfx942 a16w4 DSv4.1 tiles for Triton 3.8 #6286

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
frida-andersson wants to merge 2 commits into
base: main
Choose a base branch
from
Draft

[Triton] Retune gfx942 a16w4 DSv4.1 tiles for Triton 3.8 #6286

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
b991bbe
[Triton] Retune gfx942 a16w4 DSv4.1 tiles for Triton 3.8
frida-andersson Oct 8, 2026
ec6debc
[Triton] Keep nonkdim 32 for bm128_n5120_k768 on Triton 3.6
frida-andersson Oct 12, 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
17 changes: 9 additions & 8 deletions aiter/ops/triton/configs/gfx942/triton/moe/a16w4/DEFAULT.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,10 +26,11 @@
"bm128_n5120_k768": {
"BLOCK_SIZE_K": 64,
"BLOCK_SIZE_N": 256,
"matrix_instr_nonkdim": 32,
"num_stages": 2,
"matrix_instr_nonkdim": 16,

@anhminhnguyenhoang anhminhnguyenhoang Oct 16, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

MI300X and MI325X use this same gfx942 table. On MI300X, nonkdim 32 is 10 to 12% faster than 16. Your MI325X data also shows nonkdim 32 faster than #5667, by 0.7 to 2.1%. Could you set "matrix_instr_nonkdim": 32? Then the version check in moe_op_gemm_a16w4.py is not necessary.

"num_stages": 1,
"num_warps": 4,
"waves_per_eu": 2
"waves_per_eu": 0,
"kpack": 2
},
"bm64_n1024_k5120": {
"BLOCK_SIZE_K": 256,
Expand All @@ -43,9 +44,9 @@
"BLOCK_SIZE_K": 128,
"BLOCK_SIZE_N": 128,
"matrix_instr_nonkdim": 16,
"num_stages": 2,
"num_warps": 8,
"waves_per_eu": 2
"num_stages": 1,
"num_warps": 4,
"waves_per_eu": 0
},
"bm64_n5120_k512": {
"BLOCK_SIZE_K": 128,
Expand All @@ -57,9 +58,9 @@
},
"bm64_n5120_k768": {
"BLOCK_SIZE_K": 128,
"BLOCK_SIZE_N": 256,
"BLOCK_SIZE_N": 128,
"matrix_instr_nonkdim": 16,
"num_stages": 2,
"num_stages": 1,
"num_warps": 4,
"waves_per_eu": 0
}
Expand Down
11 changes: 9 additions & 2 deletions aiter/ops/triton/moe/moe_op_gemm_a16w4.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -104,6 +104,13 @@ def get_kernel_config_triton(m, n, k, routing_data):
f"bm{block_m}_n{n}_k{k}"
)
if tuned is not None:
nonkdim = tuned.get("matrix_instr_nonkdim", 16)
# Triton 3.6 fails LLVM lowering of bm128_n5120_k768 at nonkdim 16
# (unrealized_conversion_cast). 3.8 is faster with 16.
if block_m == 128 and n == 5120 and k == 768 and nonkdim == 16:
major, minor = (int(p) for p in triton.__version__.split(".")[:2])
if (major, minor) < (3, 8):
nonkdim = 32
Comment on lines +107 to +113

@anhminhnguyenhoang anhminhnguyenhoang Oct 16, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This stops the crash, but on Triton 3.6.0 and 3.7.1 the row is 17 to 25% slower than base. The cause is num_stages=1, waves_per_eu=0 and kpack=2, which this check does not change. Could you use the base row below Triton 3.8? If the check stays, packaging.Version is what the rest of aiter uses.

return {
"block_m": block_m,
"block_n": tuned["BLOCK_SIZE_N"],
Expand All @@ -115,7 +122,7 @@ def get_kernel_config_triton(m, n, k, routing_data):
"w_cache_modifier": w_cache_modifier,
"split_k": split_k,
"waves_per_eu": tuned.get("waves_per_eu", 0),
"matrix_instr_nonkdim": tuned.get("matrix_instr_nonkdim", 16),
"matrix_instr_nonkdim": nonkdim,
"kpack": tuned.get("kpack", 1),
}

Expand Down Expand Up @@ -162,7 +169,7 @@ def get_kernel_config_triton(m, n, k, routing_data):
"split_k": split_k,
"waves_per_eu": 0,
"matrix_instr_nonkdim": 16,
"kpack": 1,
"kpack": 2 if block_m <= 32 else 1,

@anhminhnguyenhoang anhminhnguyenhoang Oct 16, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

On Triton 3.8.0, kpack=2 at bm32 makes w2 1 to 3% slower and w1 9% faster. Could you note this trade-off in the PR description?

}
return ret

Expand Down