Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-19 09:14:35 +02:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #4778 from resonantsky/dev

Browse Source 
Established rdna2 best triton matmul configs
This commit is contained in:
Disty0
2026-04-17 20:26:53 +03:00
committed by GitHub
parent 3e2c16e8ae cefad871b9
commit 86cd6b500b
1 changed files with 24 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/sdnq/triton_mm.py
+24 -9
View File
@@ -12,16 +12,31 @@ import torch
import triton
import triton.language as tl
try:
from .common import is_rdna2_and_older
except Exception:
is_rdna2_and_older = False
matmul_configs = [
triton.Config({'BLOCK_SIZE_M': BM, 'BLOCK_SIZE_N': BN, "BLOCK_SIZE_K": BK, "GROUP_SIZE_M": GM}, num_warps=w, num_stages=s)
for BM in [32, 64, 128, 256]
for BN in [32, 64, 128, 256]
for BK in [32, 64, 128]
for GM in [4, 8]
for w in [4, 8]
for s in [2]
]
if is_rdna2_and_older:
matmul_configs = [
triton.Config({'BLOCK_SIZE_M': BM, 'BLOCK_SIZE_N': BN, "BLOCK_SIZE_K": BK, "GROUP_SIZE_M": GM}, num_warps=w, num_stages=s)
for BM in [64, 128]
for BN in [64, 128]
for BK in [64]
for GM in [2, 4]
for w in [2, 4]
for s in [2]
]
else:
matmul_configs = [
triton.Config({'BLOCK_SIZE_M': BM, 'BLOCK_SIZE_N': BN, "BLOCK_SIZE_K": BK, "GROUP_SIZE_M": GM}, num_warps=w, num_stages=s)
for BM in [32, 64, 128, 256]
for BN in [32, 64, 128, 256]
for BK in [32, 64, 128]
for GM in [4, 8]
for w in [4, 8]
for s in [2]
]
@triton.autotune(configs=matmul_configs, key=["M", "N", "K", "stride_bk", "ACCUMULATOR_DTYPE"], cache_results=True)