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

Make Triton MM block sizes configurable

Browse Source
This commit is contained in:
Disty0
2026-05-20 19:16:14 +03:00
parent 4b2f38ab7f
commit 40bd5ca00d
1 changed files with 8 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/sdnq/triton_mm.py
+8 -7
View File
@@ -7,6 +7,7 @@ matmul_configs we use takes AMD and Intel into consideration too.
SDNQ Triton configs can outperform RocBLAS and OneDNN.
"""
import os
import math
import torch
@@ -14,15 +15,15 @@ import triton
import triton.language as tl
min_block_size = 64
min_block_size = int(os.environ.get("SDNQ_TRITON_MM_MIN_BLOCK_SIZE", "64"))
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, 256]
for BN in [64, 128, 256]
for BK in [64, 128]
for GM in [2, 4, 8]
for w in [2, 4, 8]
for s in [2]
for BM in [int(BM) for BM in os.environ.get("SDNQ_TRITON_MM_BLOCK_SIZE_M_LIST", "64,128,256").replace(" ","").split(",")]
for BN in [int(BN) for BN in os.environ.get("SDNQ_TRITON_MM_BLOCK_SIZE_N_LIST", "64,128,256").replace(" ","").split(",")]
for BK in [int(BK) for BK in os.environ.get("SDNQ_TRITON_MM_BLOCK_SIZE_K_LIST", "64,128").replace(" ","").split(",")]
for GM in [int(GM) for GM in os.environ.get("SDNQ_TRITON_MM_GROUP_SIZE_M_LIST", "2,4,8").replace(" ","").split(",")]
for w in [int(w) for w in os.environ.get("SDNQ_TRITON_MM_NUM_WARPS_LIST", "2,4,8").replace(" ","").split(",")]
for s in [int(s) for s in os.environ.get("SDNQ_TRITON_MM_NUM_STAGES_LIST", "2").replace(" ","").split(",")]
]