From 7adcae9b70362ac18f68f05493a35b97762f34ab Mon Sep 17 00:00:00 2001 From: ESI-SYD Date: Mon, 18 Nov 2024 05:47:56 +0000 Subject: [PATCH] num_stages change --- benchmarks/triton_kernels_benchmark/gemm_benchmark.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/benchmarks/triton_kernels_benchmark/gemm_benchmark.py b/benchmarks/triton_kernels_benchmark/gemm_benchmark.py index 9941b0c5f0..9634b91855 100644 --- 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'},