num_stages change

intel · Nov 18, 2024 · 7adcae9 · 7adcae9
1 parent 8382e76
commit 7adcae9
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/benchmarks/triton_kernels_benchmark/gemm_benchmark.py b/benchmarks/triton_kernels_benchmark/gemm_benchmark.py
@@ -23,12 +23,14 @@
 TRANSPOSE_B = os.getenv('TRANSPOSE_B', '0') == '1'
 use_xetla = not (TRANSPOSE_A or TRANSPOSE_B)
 
+num_stages_range = [4] if os.getenv('TRITON_INTEL_ADVANCED_PATH', '0') == '0' else [1, 2, 3]
+
 
 @triton.autotune(
     configs=[
         triton.Config(
             {'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 4, 'grf_mode': 'large'},
-            num_stages=s, num_warps=32) for s in [1, 2, 3]
+            num_stages=s, num_warps=32) for s in num_stages_range
     ] + [
         triton.Config(
             {'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 4, 'grf_mode': 'large'},