[Triton/Gluon] : Change RotatingBuffersBench API and update any existing benchmarks

op_tests/op_benchmarks/triton/AGENTS.md

@@ -1,9 +1,9 @@
 # Agent instructions: Triton and Gluon benchmarks
 
-The rules for all Triton and Gluon code, benchmarks included, are in
-`aiter/ops/triton/AGENTS.md`. Read them before you change or review a
-benchmark here.
+The benchmark rules are in `op_tests/op_benchmarks/triton/README.md`. The rules for all 
+Triton and Gluon code are in `aiter/ops/triton/AGENTS.md`. Read both before 
+you change or review a benchmark here.
 
 Claude Code loads them through this import:
 
-@../../../aiter/ops/triton/AGENTS.md
+@../../../aiter/ops/triton/AGENTS.md

op_tests/op_benchmarks/triton/README.md

@@ -0,0 +1,11 @@
+# Triton/Gluon benchmarks
+
+## Rules
+1. Use RotatingBufferBench from `op_tests/op_benchmarks/triton/utils/rotating_buffer_bench.py` for creating rotating buffer pool and benchmarking
+2. Do not use custom benchmarking functions or any triton lang do_bench functions. 
+3. RotatingBufferBench pool creation API requires input generator function. Try to share
+that with the unit test. If it is shared then import the function from unit test file
+4. For ops that are multi-GPU, use bench_synced
+5. No testing for accuracy or assert_close calls in the benchmarks. Do that in the unit
+test.
+6. Benchmark only Triton ops. You may see some current benchmarks that do that, but they will be moved out in future.

op_tests/op_benchmarks/triton/bench_all_reduce.py

@@ -26,10 +26,9 @@
     get_rot_bufs_kwargs,
     override_parser_action,
 )
-from op_tests.op_benchmarks.triton.utils.benchmark_utils import (
-    do_bench_aiter_triton_synced,
-    generate_rotating_buffers_pool,
-    get_caller_name_no_ext,
+from op_tests.op_benchmarks.triton.utils.benchmark_utils import get_caller_name_no_ext
+from op_tests.op_benchmarks.triton.utils.rotating_buffer_bench import (
+    RotatingBufferBench,
 )
 
 WORLD_SIZE = 4
@@ -58,6 +57,10 @@ def benchmark(args, rank, device, tp, inboxes, calls):
         args={},
     )
 
+    bencher = RotatingBufferBench.from_dict(
+        {**get_rot_bufs_kwargs(args), **get_do_bench_graph_kwargs(args)}
+    )
+
     @triton.testing.perf_report([config])
     def _run(tokens, provider):
         numel = tokens * args.hidden
@@ -66,15 +69,8 @@ def gen():
             x = torch.randn(numel, dtype=torch.bfloat16, device=device)
             return x, torch.empty_like(x)
 
-        bufs = generate_rotating_buffers_pool(
-            [gen], [["x", "out"]], **get_rot_bufs_kwargs(args)
-        )
-        ms = do_bench_aiter_triton_synced(
-            fns[provider],
-            bufs,
-            dist_group=tp.cpu_group,
-            **get_do_bench_graph_kwargs(args),
-        )
+        bencher.build_pool([gen], [["x", "out"]])
+        ms = bencher.bench_synced(fns[provider], dist_group=tp.cpu_group)
         per_rank = [None] * WORLD_SIZE
         dist.all_gather_object(per_rank, ms, group=tp.cpu_group)
         ms = max(per_rank)

op_tests/op_benchmarks/triton/bench_all_reduce_rmsnorm_quant.py

@@ -27,10 +27,9 @@
     get_rot_bufs_kwargs,
     override_parser_action,
 )
-from op_tests.op_benchmarks.triton.utils.benchmark_utils import (
-    do_bench_aiter_triton_synced,
-    generate_rotating_buffers_pool,
-    get_caller_name_no_ext,
+from op_tests.op_benchmarks.triton.utils.benchmark_utils import get_caller_name_no_ext
+from op_tests.op_benchmarks.triton.utils.rotating_buffer_bench import (
+    RotatingBufferBench,
 )
 
 WORLD_SIZE = 4
@@ -54,6 +53,10 @@ def benchmark(args, rank, device, tp, inboxes, calls):
         args={},
     )
 
+    bencher = RotatingBufferBench.from_dict(
+        {**get_rot_bufs_kwargs(args), **get_do_bench_graph_kwargs(args)}
+    )
+
     @triton.testing.perf_report([config])
     def _run(tokens, provider):
         shape = (tokens, args.hidden)
@@ -89,16 +92,9 @@ def gen():
             x = torch.randn(shape, dtype=torch.bfloat16, device=device)
             return x, torch.randn(shape, dtype=torch.bfloat16, device=device)
 
-        bufs = generate_rotating_buffers_pool(
-            [gen], [["x", "res"]], **get_rot_bufs_kwargs(args)
-        )
+        bencher.build_pool([gen], [["x", "res"]])
         with graph_capture():  # the CAR paths capture as they do in a model
-            ms = do_bench_aiter_triton_synced(
-                fns[provider],
-                bufs,
-                dist_group=tp.cpu_group,
-                **get_do_bench_graph_kwargs(args),
-            )
+            ms = bencher.bench_synced(fns[provider], dist_group=tp.cpu_group)
         per_rank = [None] * WORLD_SIZE
         dist.all_gather_object(per_rank, ms, group=tp.cpu_group)
         return max(per_rank) * 1e3

op_tests/op_benchmarks/triton/bench_batched_gemm_a16wfp4.py

@@ -19,21 +19,21 @@
 )
 from op_tests.op_benchmarks.triton.utils.benchmark_utils import (
     batched_model_benchmark_shapes,
-    do_bench_aiter_triton,
-    generate_rotating_buffers_pool,
     get_caller_name_no_ext,
     get_model_benchmark_object,
     get_shape_benchmark_object,
     print_vgpr,
 )
+from op_tests.op_benchmarks.triton.utils.rotating_buffer_bench import (
+    RotatingBufferBench,
+)
 from op_tests.triton_tests.gemm.batched.test_batched_gemm_a16wfp4 import (
     generate_batched_gemm_a16wfp4_inputs,
 )
 
 
 def bench_gemm_fn(
-    kwargs_rot_bufs: dict,
-    kwargs_do_bench_graph: dict,
+    bencher: RotatingBufferBench,
     batch: int,
     M: int,
     N: int,
@@ -56,9 +56,7 @@ def bench_gemm_fn(
     # (x, w, x_scales, w_scales, y); kernel's weight-scale kwarg is "w_scales"
     _out_vals_keys = ["x", "w", "_", "w_scales", "y"]
 
-    rot_bufs = generate_rotating_buffers_pool(
-        [_p_gen], [_out_vals_keys], **kwargs_rot_bufs
-    )
+    rot_bufs = bencher.build_pool([_p_gen], [_out_vals_keys])
 
     kernel_fn = functools.partial(batched_gemm_a16wfp4, dtype=c_dtype)
 
@@ -74,7 +72,7 @@ def bench_gemm_fn(
     mem_write = (batch * M * N) * rot_bufs[0]["y"].element_size()
     mem = mem_read + mem_write
 
-    ms = do_bench_aiter_triton(kernel_fn, rot_bufs, **kwargs_do_bench_graph)
+    ms = bencher.bench(kernel_fn)
 
     # Return exactly one scalar depending on which metric is active
     if metric == "time":
@@ -97,8 +95,9 @@ def run_model_benchmark(args):
         model_benchmark_shapes_fn=batched_model_benchmark_shapes,
     )
 
-    kwargs_do_bench_graph = get_do_bench_graph_kwargs(args)
-    kwargs_rot_bufs = get_rot_bufs_kwargs(args)
+    bencher = RotatingBufferBench.from_dict(
+        {**get_rot_bufs_kwargs(args), **get_do_bench_graph_kwargs(args)}
+    )
 
     @triton.testing.perf_report([benchmark])
     def bench_batched_gemm_a16wfp4(
@@ -117,9 +116,7 @@ def bench_batched_gemm_a16wfp4(
             K = math.ceil(K / args.tp)
         # print(f"Layer: {layer}, B: {batch}, M: {M}, N: {N}, K: {K}, hidden_dim: {hidden_dim}, intermediate_dim: {intermediate_dim}")
 
-        return bench_gemm_fn(
-            kwargs_rot_bufs, kwargs_do_bench_graph, batch, M, N, K, metric, args.layout
-        )
+        return bench_gemm_fn(bencher, batch, M, N, K, metric, args.layout)
 
     bench_batched_gemm_a16wfp4.run(save_path="." if args.o else None, print_data=True)
 
@@ -131,8 +128,9 @@ def run_shape_benchmark(args):
         x_names=["batch", "M", "N", "K"],
     )
 
-    kwargs_do_bench_graph = get_do_bench_graph_kwargs(args)
-    kwargs_rot_bufs = get_rot_bufs_kwargs(args)
+    bencher = RotatingBufferBench.from_dict(
+        {**get_rot_bufs_kwargs(args), **get_do_bench_graph_kwargs(args)}
+    )
 
     @triton.testing.perf_report([benchmark])
     def bench_batched_gemm_a16wfp4(
@@ -143,9 +141,7 @@ def bench_batched_gemm_a16wfp4(
         metric,
         **kwargs,
     ):
-        return bench_gemm_fn(
-            kwargs_rot_bufs, kwargs_do_bench_graph, batch, M, N, K, metric, args.layout
-        )
+        return bench_gemm_fn(bencher, batch, M, N, K, metric, args.layout)
 
     bench_batched_gemm_a16wfp4.run(save_path="." if args.o else None, print_data=True)
 

op_tests/op_benchmarks/triton/bench_batched_gemm_a8w8.py

@@ -16,21 +16,21 @@
 )
 from op_tests.op_benchmarks.triton.utils.benchmark_utils import (
     batched_model_benchmark_shapes,
-    do_bench_aiter_triton,
-    generate_rotating_buffers_pool,
     get_caller_name_no_ext,
     get_model_benchmark_object,
     get_shape_benchmark_object,
     print_vgpr,
 )
+from op_tests.op_benchmarks.triton.utils.rotating_buffer_bench import (
+    RotatingBufferBench,
+)
 from op_tests.triton_tests.gemm.batched.test_batched_gemm_a8w8 import (
     generate_batched_gemm_a8w8_inputs,
 )
 
 
 def bench_gemm_fn(
-    kwargs_rot_bufs: dict,
-    kwargs_do_bench_graph: dict,
+    bencher: RotatingBufferBench,
     batch: int,
     M: int,
     N: int,
@@ -53,9 +53,7 @@ def bench_gemm_fn(
     # batched_gemm_a8w8 names its activation/weight/output kwargs XQ/WQ/YQ
     _out_vals_keys = ["XQ", "WQ", "x_scale", "w_scale", "bias", "YQ"]
 
-    rot_bufs = generate_rotating_buffers_pool(
-        [_p_gen], [_out_vals_keys], **kwargs_rot_bufs
-    )
+    rot_bufs = bencher.build_pool([_p_gen], [_out_vals_keys])
 
     kernel_fn = functools.partial(batched_gemm_a8w8, dtype=c_dtype)
 
@@ -73,7 +71,7 @@ def bench_gemm_fn(
     mem_write = (batch * M * N) * rot_bufs[0]["YQ"].element_size()
     mem = mem_read + mem_write
 
-    ms = do_bench_aiter_triton(kernel_fn, rot_bufs, **kwargs_do_bench_graph)
+    ms = bencher.bench(kernel_fn)
 
     # Return exactly one scalar depending on which metric is active
     if metric == "time":
@@ -96,8 +94,9 @@ def run_model_benchmark(args):
         model_benchmark_shapes_fn=batched_model_benchmark_shapes,
     )
 
-    kwargs_do_bench_graph = get_do_bench_graph_kwargs(args)
-    kwargs_rot_bufs = get_rot_bufs_kwargs(args)
+    bencher = RotatingBufferBench.from_dict(
+        {**get_rot_bufs_kwargs(args), **get_do_bench_graph_kwargs(args)}
+    )
 
     @triton.testing.perf_report([benchmark])
     def bench_batched_gemm_a8w8(
@@ -116,9 +115,7 @@ def bench_batched_gemm_a8w8(
             K = math.ceil(K / args.tp)
         # print(f"Layer: {layer}, B: {batch}, M: {M}, N: {N}, K: {K}, hidden_dim: {hidden_dim}, intermediate_dim: {intermediate_dim}")
 
-        return bench_gemm_fn(
-            kwargs_rot_bufs, kwargs_do_bench_graph, batch, M, N, K, metric, args.layout
-        )
+        return bench_gemm_fn(bencher, batch, M, N, K, metric, args.layout)
 
     bench_batched_gemm_a8w8.run(save_path="." if args.o else None, print_data=True)
 
@@ -130,14 +127,13 @@ def run_shape_benchmark(args):
         x_names=["batch", "M", "N", "K"],
     )
 
-    kwargs_do_bench_graph = get_do_bench_graph_kwargs(args)
-    kwargs_rot_bufs = get_rot_bufs_kwargs(args)
+    bencher = RotatingBufferBench.from_dict(
+        {**get_rot_bufs_kwargs(args), **get_do_bench_graph_kwargs(args)}
+    )
 
     @triton.testing.perf_report([benchmark])
     def bench_batched_gemm_a8w8(batch, M, N, K, metric, **kwargs):
-        return bench_gemm_fn(
-            kwargs_rot_bufs, kwargs_do_bench_graph, batch, M, N, K, metric, args.layout
-        )
+        return bench_gemm_fn(bencher, batch, M, N, K, metric, args.layout)
 
     bench_batched_gemm_a8w8.run(save_path="." if args.o else None, print_data=True)
 

op_tests/op_benchmarks/triton/bench_batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant.py

@@ -19,21 +19,21 @@
 )
 from op_tests.op_benchmarks.triton.utils.benchmark_utils import (
     batched_model_benchmark_shapes,
-    do_bench_aiter_triton,
-    generate_rotating_buffers_pool,
     get_caller_name_no_ext,
     get_model_benchmark_object,
     get_shape_benchmark_object,
     print_vgpr,
 )
+from op_tests.op_benchmarks.triton.utils.rotating_buffer_bench import (
+    RotatingBufferBench,
+)
 from op_tests.triton_tests.gemm.batched.test_batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant import (
     generate_batched_gemm_a16w8_inputs as generate_batched_gemm_a8w8_per_token_group_inputs,
 )
 
 
 def bench_gemm_fn(
-    kwargs_rot_bufs: dict,
-    kwargs_do_bench_graph: dict,
+    bencher: RotatingBufferBench,
     batch: int,
     M: int,
     N: int,
@@ -64,7 +64,10 @@ def bench_gemm_fn(
     gens = [_p_gen]
     out_vals_keys = [["X", "WQ", "w_scale", "bias", "YQ"]]
     if y_per_token_quant:
-        # The y_scale_out path writes an FP8 YQ and a per-token scale, not a bf16 YQ.
+        # The y_scale_out path writes an FP8 YQ and a per-token scale, not a
+        # bf16 YQ -- _p_gen's own "YQ" slot comes back None (output=False
+        # above), so drop it ("_") and build the real YQ/y_scale_out pair
+        # here instead, rotating per entry like every other output buffer.
         out_vals_keys = [["X", "WQ", "w_scale", "bias", "_"], ["YQ", "y_scale_out"]]
         gens.append(
             lambda: (
@@ -73,7 +76,7 @@ def bench_gemm_fn(
             )
         )
 
-    rot_bufs = generate_rotating_buffers_pool(gens, out_vals_keys, **kwargs_rot_bufs)
+    rot_bufs = bencher.build_pool(gens, out_vals_keys)
 
     kernel_fn = functools.partial(
         batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant,
@@ -105,7 +108,7 @@ def bench_gemm_fn(
         )
     mem = mem_read + mem_write
 
-    ms = do_bench_aiter_triton(kernel_fn, rot_bufs, **kwargs_do_bench_graph)
+    ms = bencher.bench(kernel_fn)
 
     # Return exactly one scalar depending on which metric is active
     if metric == "time":
@@ -128,8 +131,9 @@ def run_model_benchmark(args):
         model_benchmark_shapes_fn=batched_model_benchmark_shapes,
     )
 
-    kwargs_do_bench_graph = get_do_bench_graph_kwargs(args)
-    kwargs_rot_bufs = get_rot_bufs_kwargs(args)
+    bencher = RotatingBufferBench.from_dict(
+        {**get_rot_bufs_kwargs(args), **get_do_bench_graph_kwargs(args)}
+    )
 
     @triton.testing.perf_report([benchmark])
     def bench_batched_gemm_a8w8_per_token_group_prequant_w_per_batched_tensor_quant(
@@ -148,8 +152,7 @@ def bench_batched_gemm_a8w8_per_token_group_prequant_w_per_batched_tensor_quant(
             raise ValueError(f"Unsupported layer: {layer}")
 
         return bench_gemm_fn(
-            kwargs_rot_bufs,
-            kwargs_do_bench_graph,
+            bencher,
             batch,
             M,
             N,
@@ -173,16 +176,16 @@ def run_shape_benchmark(args):
     x_names = ["batch", "M", "N", "K"]
     benchmark = get_shape_benchmark_object(plot_name, args, x_names=x_names)
 
-    kwargs_do_bench_graph = get_do_bench_graph_kwargs(args)
-    kwargs_rot_bufs = get_rot_bufs_kwargs(args)
+    bencher = RotatingBufferBench.from_dict(
+        {**get_rot_bufs_kwargs(args), **get_do_bench_graph_kwargs(args)}
+    )
 
     @triton.testing.perf_report([benchmark])
     def bench_batched_gemm_a8w8_per_token_group_prequant_w_per_batched_tensor_quant(
         batch, M, N, K, metric, **kwargs
     ):
         return bench_gemm_fn(
-            kwargs_rot_bufs,
-            kwargs_do_bench_graph,
+            bencher,
             batch,
             M,
             N,