From 40bd5ca00deb649563c6f53734c7e59bda2da937 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Wed, 20 May 2026 19:16:14 +0300 Subject: [PATCH] Make Triton MM block sizes configurable --- modules/sdnq/triton_mm.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/modules/sdnq/triton_mm.py b/modules/sdnq/triton_mm.py index 4cbf01b4e..d2acdb9aa 100644 --- a/modules/sdnq/triton_mm.py +++ b/modules/sdnq/triton_mm.py @@ -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(",")] ]