Repository: xlinsist/triton-benchmark
Pull request #22: [benchmarks & cases] Fix bug in Triton matmul benchmark
Author: login: xvyv99
Created: 2025-03-22T09:35:31Z | Merged: 2025-03-23T09:23:09Z



Diff:
diff --git a/benchmarks/patch/triton-cpu-0003-autotuner.patch b/benchmarks/patch/triton-cpu-0003-autotuner.patch
new file mode 100644
index 0000000..a9ec34f
--- /dev/null
+++ b/benchmarks/patch/triton-cpu-0003-autotuner.patch
@@ -0,0 +1,51 @@
+diff --git a/python/triton/runtime/autotuner.py b/python/triton/runtime/autotuner.py
+index c9833c94..64af352f 100644
+--- a/python/triton/runtime/autotuner.py
++++ b/python/triton/runtime/autotuner.py
+@@ -28,6 +28,7 @@ class Autotuner(KernelInterface):
+         rep=None,
+         use_cuda_graph=False,
+         do_bench=None,
++        tuning_time=None
+     ):
+         """
+         :param prune_configs_by: a dict of functions that are used to prune configs, fields:
+@@ -35,6 +36,9 @@ class Autotuner(KernelInterface):
+             'top_k': number of configs to bench
+             'prune_num_stages_by'(optional): a function used to prune num_stages. It takes configs:List[Config] as its input, and returns pruned configs.
+         """
++        if isinstance(tuning_time, dict):
++            self.tuning_time = tuning_time
++
+         if not configs:
+             self.configs = [Config({}, num_warps=4, num_stages=2, num_ctas=1)]
+         else:
+@@ -197,6 +201,10 @@ class Autotuner(KernelInterface):
+         else:
+             config = self.configs[0]
+         self.best_config = config
++        
++        if self.tuning_time != None:
++            self.tuning_time['bench_time'] = self.bench_time
++
+         if os.getenv("TRITON_PRINT_AUTOTUNING", None) == "1" and not used_cached_result:
+             print(f"Triton autotuning for function {self.base_fn.__name__} finished after "
+                   f"{self.bench_time:.2f}s; best config selected: {self.best_config};")
+@@ -301,7 +309,7 @@ class Config:
+ 
+ 
+ def autotune(configs, key, prune_configs_by=None, reset_to_zero=None, restore_value=None, pre_hook=None, post_hook=None,
+-             warmup=None, rep=None, use_cuda_graph=False, do_bench=None):
++             warmup=None, rep=None, use_cuda_graph=False, do_bench=None, tuning_time=None):
+     """
+     Decorator for auto-tuning a :code:`triton.jit`'d function.
+ 
+@@ -360,7 +368,7 @@ def autotune(configs, key, prune_configs_by=None, reset_to_zero=None, restore_va
+     def decorator(fn):
+         return Autotuner(fn, fn.arg_names, configs, key, reset_to_zero, restore_value, pre_hook=pre_hook,
+                          post_hook=post_hook, prune_configs_by=prune_configs_by, warmup=warmup, rep=rep,
+-                         use_cuda_graph=use_cuda_graph)
++                         use_cuda_graph=use_cuda_graph, tuning_time=tuning_time)
+ 
+     return decorator
+ 
diff --git a/cases/matmul/benchmark.py b/cases/matmul/benchmark.py
index 594008b..4028c1e 100644
--- a/cases/matmul/benchmark.py
+++ b/cases/matmul/benchmark.py
@@ -35,6 +35,11 @@ def benchmark_torch(a_np, b_np, num_threads=None):
 
 def run_benchmark(method_name, method_func, shape, a_np, b_np, torch_result):
     """Run a single benchmark and validate results."""
+    
+    # HACK: Refresh threads number, otherwise Triton will default to single-thread mode
+    num_threads = multiprocessing.cpu_count()
+    torch.set_num_threads(num_threads)
+    
     exec_time, result, *rest = method_func(shape, a_np, b_np)
     tune_time = rest[0] if rest else 0.0
     
@@ -45,7 +50,6 @@ def run_benchmark(method_name, method_func, shape, a_np, b_np, torch_result):
         'Shape': shape, 
         'Method': method_name, 
         'Time(s)': exec_time, 
-        # TODO: Capture triton's tuning time
         'TuningTime(s)': tune_time
     }
 
@@ -69,11 +73,12 @@ def main():
         ('hidet', benchmark_hidet), ('tvm', benchmark_tvm),('triton',benchmark_triton), ('autotvm', benchmark_autotvm),('ansor',benchmark_ansor),
         ('hidet_single', benchmark_hidet_single), ('tvm_single', benchmark_tvm_single), ('triton_single', benchmark_triton_single)
     ]
-    
+
     for method, method_func in methods:
         print(f"Running {method} benchmark...")
         records.append(run_benchmark(method, method_func, shape, a_np, b_np, torch_result))
     
+
     df = pd.DataFrame(records)
     df.sort_values(by=['Benchmark', 'Shape'], inplace=True)
 
@@ -89,4 +94,5 @@ def main():
     df.to_csv("./performance_report.csv", index=False)
 
 if __name__ == "__main__":
-    main()
\ No newline at end of file
+    main()
+
diff --git a/cases/matmul/matmul_triton.py b/cases/matmul/matmul_triton.py
index 14dfdef..630941e 100644
--- a/cases/matmul/matmul_triton.py
+++ b/cases/matmul/matmul_triton.py
@@ -7,20 +7,18 @@
 import time
 import os
 
+TUNING_TIME = {}
+triton.runtime.driver.set_active_to_cpu()
+
 # Triton Benchmark
-def get_matmul_kernel_autotune_config():
+def get_matmul_kernel_autotune_config(num_threads = 0):
     configs=[]
     for BLOCK_SIZE_M in [8, 16, 32]:
         for BLOCK_SIZE_N in [8, 16, 32]:
             for BLOCK_SIZE_K in [64]:
-                configs.append(triton.Config({'BLOCK_SIZE_M': BLOCK_SIZE_M, 'BLOCK_SIZE_N': BLOCK_SIZE_N, 'BLOCK_SIZE_K': BLOCK_SIZE_K}))
-    # return configs
-    return [triton.Config({'BLOCK_SIZE_M': 16, 'BLOCK_SIZE_N': 16, 'BLOCK_SIZE_K': 64})]
-@triton.autotune(
-    configs=get_matmul_kernel_autotune_config(),
-    key=[],
-)
-@triton.jit
+                configs.append(triton.Config({'BLOCK_SIZE_M': BLOCK_SIZE_M, 'BLOCK_SIZE_N': BLOCK_SIZE_N, 'BLOCK_SIZE_K': BLOCK_SIZE_K}, num_threads = num_threads))
+    return configs
+
 def triton_matmul_kernel(
         # Pointers to matrices
         a_ptr, b_ptr, c_ptr,
@@ -85,11 +83,16 @@ def triton_matmul_kernel(
     c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
     tl.store(c_ptrs, c, mask=c_mask)
 
-
 def benchmark_triton(shape, a_np, b_np, parallel=True):
-    os.environ["TRITON_CPU_BACKEND"] = "1"
-    os.environ["TRITON_CPU_MAX_THREADS"] = "0" if parallel else "1"
+    fn = triton_matmul_kernel
+    fn_jit = triton.jit(fn)
 
+    fn_jit_tuned = triton.runtime.Autotuner(fn_jit, fn_jit.arg_names, reset_to_zero=None, restore_value=None,
+        configs=get_matmul_kernel_autotune_config(0 if parallel else 1),
+        key=[],
+        tuning_time = TUNING_TIME
+    )
+    
     M, N, K = shape
     a = torch.tensor(a_np, device='cpu', dtype=torch.float32)
     b = torch.tensor(b_np, device='cpu', dtype=torch.float32)
@@ -100,26 +103,27 @@ def benchmark_triton(shape, a_np, b_np, parallel=True):
     # 1D launch kernel where each block gets its own program.
     grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv(N, META['BLOCK_SIZE_N']), )
 
+    # Warm-up
+    fn_jit_tuned[grid](a, b, c, M, N, K, a.stride(0), a.stride(1), b.stride(0), b.stride(1), c.stride(0), c.stride(1))
+
     times = []
     for _ in range(10):
         start = time.perf_counter()
-        triton_matmul_kernel[grid](a, b, c, M, N, K, a.stride(0), a.stride(1), b.stride(0), b.stride(1), c.stride(0), c.stride(1))
+        fn_jit_tuned[grid](a, b, c, M, N, K, a.stride(0), a.stride(1), b.stride(0), b.stride(1), c.stride(0), c.stride(1))
         end = time.perf_counter()
         times.append(end - start)
-    return np.mean(times), c.numpy()
-
+    return np.mean(times), c.numpy(), TUNING_TIME['bench_time']
 
 def benchmark_triton_single(shape, a_np, b_np):
     return benchmark_triton(shape, a_np, b_np, parallel=False)
 
-
 if __name__ == "__main__":
     M = N = K = 512
     a_np = np.random.rand(M, K).astype(np.float32)
     b_np = np.random.rand(K, N).astype(np.float32)
     shape = (M, N, K)
-    time_triton, result_triton = benchmark_triton(shape, a_np, b_np)
-    time_triton_single, result_triton_single = benchmark_triton_single(shape, a_np, b_np)
+    time_triton, result_triton, __ = benchmark_triton(shape, a_np, b_np)
+    time_triton_single, result_triton_single, __ = benchmark_triton_single(shape, a_np, b_np)
     assert np.allclose(result_triton, result_triton_single, atol=1e-3, rtol=1e-3), f"triton result mismatch!"
     print(result_triton)
     print(result_triton_single)


Reviews:
user: login: xlinsist
state: COMMENTED
author_association: OWNER
submitted_at: 2025-03-23T09:40:16Z
user: login: xvyv99
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-24T21:27:30Z

Review comments:
diff_hunk: @@ -35,6 +35,11 @@ def benchmark_torch(a_np, b_np, num_threads=None):
 
 def run_benchmark(method_name, method_func, shape, a_np, b_np, torch_result):
     """Run a single benchmark and validate results."""
+    
+    # HACK: Refresh threads number, otherwise Triton will default to single-thread mode
+    num_threads = multiprocessing.cpu_count()
+    torch.set_num_threads(num_threads)
+    
path: cases/matmul/benchmark.py
user: login: xlinsist
body: Should this be placed in matmul-triton? This run_benchmark function does not explicitly call num_thread.
created_at: 2025-03-23T09:40:16Z
author_association: OWNER
diff_hunk: @@ -35,6 +35,11 @@ def benchmark_torch(a_np, b_np, num_threads=None):
 
 def run_benchmark(method_name, method_func, shape, a_np, b_np, torch_result):
     """Run a single benchmark and validate results."""
+    
+    # HACK: Refresh threads number, otherwise Triton will default to single-thread mode
+    num_threads = multiprocessing.cpu_count()
+    torch.set_num_threads(num_threads)
+    
path: cases/matmul/benchmark.py
user: login: xvyv99
body: Sorry, I forgot to try it before. The code for refreshing the number of threads should indeed be placed in `triton_benchmark`.
created_at: 2025-03-24T21:27:30Z
author_association: CONTRIBUTOR

PR comments:
user: login: xvyv99
created_at: 2025-03-22T12:25:33Z
body: **Result:**
![image](https://github.com/user-attachments/assets/5d91395a-1bd8-4e19-a5ff-b5c4407e245a)
author_association: CONTRIBUTOR
user: login: xlinsist
created_at: 2025-03-22T16:19:27Z
body: Well done! Our next step is to do more observations on the performance to analyze the causes.
author_association: OWNER
user: login: xlinsist
created_at: 2025-03-22T16:28:40Z
body: The patch and the triton implementation look good to me. I am thinking about whether and how to modify other cases by refering to this PR. I'll try to reply on Monday.
author_association: OWNER