power of 2 multiple for caching gemms

mobiusml · Dec 9, 2024 · 5ebcca0 · 5ebcca0
1 parent 245c9bc
commit 5ebcca0
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 5 deletions.
diff --git a/gemlite/triton_kernels/gemm_A16fWnO16f_int32packing.py b/gemlite/triton_kernels/gemm_A16fWnO16f_int32packing.py
@@ -11,20 +11,19 @@
 KEYS        = ['M', 'N', 'K', 'group_size', 'elements_per_sample']
 MATMUL_TYPE = "GEMM"
 
-# code based https://github.com/fpgaminer/GPTQ-triton
 def kernel_config_pruner(configs, nargs, **kwargs):
     global KEYS
     from ..core import GEMLITE_TRITON_CONFIG_CACHE
 
-    m = max(2 ** int(math.ceil(math.log2(nargs['M']))), 16) #Need at least 16 here for tl.dot
+    m = max(2 ** int(math.ceil(math.log2(nargs['M']))), 16)
     n = nargs['N'] 
     k = nargs['K'] 
     g = nargs['group_size']
     e = nargs['elements_per_sample']
 
     #Check cache
     if(MATMUL_TYPE in GEMLITE_TRITON_CONFIG_CACHE):
-        _signature = str(tuple([nargs[i] for i in KEYS]))
+        _signature = str(tuple([m, n, k, g, e]))
         if(_signature in GEMLITE_TRITON_CONFIG_CACHE[MATMUL_TYPE]):
             _config     = copy.deepcopy(GEMLITE_TRITON_CONFIG_CACHE[MATMUL_TYPE][_signature])
             _num_stages = _config.pop('num_stages')

diff --git a/gemlite/triton_kernels/gemm_splitK_A16fWnO16f_int32packing.py b/gemlite/triton_kernels/gemm_splitK_A16fWnO16f_int32packing.py
@@ -15,15 +15,15 @@ def kernel_config_pruner(configs, nargs, **kwargs):
     global KEYS
     from ..core import GEMLITE_TRITON_CONFIG_CACHE
 
-    m = nargs['M'] 
+    m = 2 ** int(math.ceil(math.log2(nargs['M'])))
     n = nargs['N'] 
     k = nargs['K'] 
     g = nargs['group_size']
     e = nargs['elements_per_sample']
 
     #Check cache
     if(MATMUL_TYPE in GEMLITE_TRITON_CONFIG_CACHE):
-        _signature = str(tuple([nargs[i] for i in KEYS]))
+        _signature = str(tuple([m, n, k, g, e]))
         if(_signature in GEMLITE_TRITON_CONFIG_CACHE[MATMUL_TYPE]):
             _config     = copy.deepcopy(GEMLITE_TRITON_CONFIG_CACHE[MATMUL_TYPE][_signature])
             _num_stages = _config.pop('num_stages')