microbenchmarks: add buffer rotation option

benchmarks/microbenchmarks/README.md

@@ -31,6 +31,28 @@ python benchmark_gemm.py --csv --csv-samples gemm_samples.csv
 The samples CSV contains one row per timing sample with columns for all
 benchmark parameters plus `label`, `sample_idx`, and `time_ms`.
 
+### Rotating input buffers
+
+By default each benchmark cycles its inputs through a ring of buffers whose
+total footprint exceeds the **last-level cache**, so back-to-back kernel
+launches touch different memory (closer to a cold-cache, steady-state workload)
+instead of reading data still resident in cache and reporting optimistic
+numbers. This matches the `--rotating` option of `hipblaslt-bench`, which
+likewise takes a rotating memory budget in MB. Pass `--no-rotating` to instead
+time a single cached input buffer:
+
+```bash
+python benchmark_gemm.py                     # rotate, auto-size the ring past the LLC
+python benchmark_casting.py --rotating 512   # rotate within a 512 MB budget
+python benchmark_gemm.py --no-rotating       # single cached input buffer
+```
+
+Rotation is **on by default**. Passing `--rotating MB` sets the rotating memory
+budget in megabytes (the ring holds enough buffers to span it); omitting the
+value auto-sizes the ring to ~2x a conservative 256 MB last-level cache (the AMD
+Infinity Cache; see `utils.py::_last_level_cache_bytes`). `--no-rotating`
+disables rotation entirely.
+
 ## Shared configuration
 
 Common benchmark settings live in `utils.py`.

benchmarks/microbenchmarks/benchmark_casting.py

@@ -21,6 +21,7 @@
 from utils import (
     MODEL_HIDDEN_SIZES, M_SIZE_LIST,
     time_func, compute_gbps, make_metric_record, run_benchmarks,
+    make_input, rotating,
 )
 
 TE_FP8_E4M3 = tex.DType.kFloat8E4M3
@@ -62,14 +63,19 @@ def bench_cast(Case, M, hidden_size, direction, fp8_dtype, dtype_str):
     quantizer = Float8Quantizer(scale, amax, fp8_dtype)
 
     if direction == "quantize":
-        x = torch.randn(M, hidden_size, dtype=torch.bfloat16, device=device)
-        out = quantizer(x)
-        cast_func = lambda: quantizer.quantize(x, out=out)
+        next_x = make_input((M, hidden_size), torch.bfloat16, device=device)
+        out = quantizer(next_x())
+        cast_func = lambda: quantizer.quantize(next_x(), out=out)
         total_bytes = numel * (2 + 1)  # BF16 read + FP8 write
     else:
-        x = torch.randn(M, hidden_size, dtype=torch.bfloat16, device=device)
-        fp8_tensor = quantizer(x)
-        cast_func = lambda: fp8_tensor.dequantize()
+        # Rotate a ring of FP8 tensors (bytes can't be inferred, so hint numel).
+        next_fp8 = rotating(
+            lambda: quantizer(
+                torch.randn(M, hidden_size, dtype=torch.bfloat16, device=device)
+            ),
+            bytes_per_buffer=numel,  # FP8 ~ 1 byte/element
+        )
+        cast_func = lambda: next_fp8().dequantize()
         total_bytes = numel * (1 + 2)  # FP8 read + BF16 write
 
     ms, measurement = time_func(cast_func, method="blocked")

benchmarks/microbenchmarks/benchmark_gemm.py

@@ -11,6 +11,7 @@
 from utils import (
     generate_gemm_test_cases,
     time_func, compute_tflops, make_forward_backward_metric_records, run_benchmarks,
+    make_input,
 )
 
 BENCHMARK_LABEL = "GEMM"
@@ -20,16 +21,17 @@ def bench_gemm(Case, M, N, K, dtype):
     device = "cuda"
 
     linear = te.Linear(K, N, bias=False).to(device=device, dtype=dtype)
-    x = torch.randn(M, K, dtype=dtype, device=device, requires_grad=True)
+    next_x = make_input((M, K), dtype, device=device, requires_grad=True)
 
-    fwd_func = lambda: linear(x)
+    fwd_func = lambda: linear(next_x())
     out = fwd_func()
     grad_out = torch.randn_like(out)
 
     def fwd_bwd_func():
-        out = linear(x)
+        xb = next_x()
+        out = linear(xb)
         out.backward(grad_out)
-        x.grad = None
+        xb.grad = None
         linear.weight.grad = None
 
     fwd_bwd_func()

benchmarks/microbenchmarks/benchmark_gemm_fp8.py

@@ -17,6 +17,7 @@
 from utils import (
     generate_gemm_test_cases,
     time_func, compute_tflops, make_forward_backward_metric_records, run_benchmarks,
+    make_input,
 )
 
 RECIPES = {
@@ -36,18 +37,19 @@ def bench_fp8_gemm(Case, M, N, K, dtype):
     device = "cuda"
 
     linear = te.Linear(K, N, bias=False).to(device=device, dtype=dtype)
-    x = torch.randn(M, K, dtype=dtype, device=device, requires_grad=True)
+    next_x = make_input((M, K), dtype, device=device, requires_grad=True)
     grad_out = torch.randn(M, N, dtype=dtype, device=device)
 
     def fwd_func():
         with te.fp8_autocast(enabled=True, fp8_recipe=FP8_RECIPE):
-            return linear(x)
+            return linear(next_x())
 
     def fwd_bwd_func():
+        xb = next_x()
         with te.fp8_autocast(enabled=True, fp8_recipe=FP8_RECIPE):
-            out = linear(x)
+            out = linear(xb)
             out.backward(grad_out)
-        x.grad = None
+        xb.grad = None
         linear.weight.grad = None
 
     fwd_flops = 2 * M * N * K

benchmarks/microbenchmarks/benchmark_grouped_gemm.py

@@ -13,6 +13,7 @@
     compute_tflops,
     make_forward_backward_metric_records,
     run_benchmarks,
+    make_input,
 )
 
 BENCHMARK_LABEL = "Grouped GEMM"
@@ -114,18 +115,21 @@ def bench_grouped_gemm(Case, B, M, N, K, dtype):
         params_dtype=dtype,
         device=device,
     )
-    x = torch.randn((sum_M, K), dtype=dtype, device=device, requires_grad=True)
+    # Rotate the activation buffer (on by default) so back-to-back grouped GEMMs 
+    # read different memory; GroupedLinear splits it internally per m_splits.
+    next_x = make_input((sum_M, K), dtype, device=device, requires_grad=True)
 
     def fwd_func_te():
-        return grouped_linear(x, m_splits, m_splits_tensor=m_splits_tensor)
+        return grouped_linear(next_x(), m_splits, m_splits_tensor=m_splits_tensor)
 
     out_te = fwd_func_te()
     grad_out = torch.randn_like(out_te)
 
     def fwd_bwd_func_te():
-        out = grouped_linear(x, m_splits, m_splits_tensor=m_splits_tensor)
+        xb = next_x()
+        out = grouped_linear(xb, m_splits, m_splits_tensor=m_splits_tensor)
         out.backward(grad_out)
-        x.grad = None
+        xb.grad = None
         for param in grouped_linear.parameters():
             param.grad = None
 

benchmarks/microbenchmarks/benchmark_normalization.py

@@ -18,6 +18,7 @@
 from utils import (
     DTYPE_LIST, MODEL_HIDDEN_SIZES, M_SIZE_LIST,
     time_func, compute_gbps, make_forward_backward_metric_records, run_benchmarks,
+    make_input,
 )
 
 NORM_TYPES = [
@@ -49,22 +50,23 @@ def bench_norm(Case, M, hidden_size, norm_name, norm_cls, dtype):
     device = "cuda"
 
     norm = norm_cls(hidden_size).to(device=device, dtype=dtype)
-    x = torch.randn(M, hidden_size, dtype=dtype, device=device, requires_grad=True)
+    next_x = make_input((M, hidden_size), dtype, device=device, requires_grad=True)
 
-    fwd_func = lambda: norm(x)
+    fwd_func = lambda: norm(next_x())
     out = fwd_func()
     grad_out = torch.randn_like(out)
 
     def fwd_bwd_func():
-        out = norm(x)
+        xb = next_x()
+        out = norm(xb)
         out.backward(grad_out)
-        x.grad = None
+        xb.grad = None
         for p in norm.parameters():
             p.grad = None
 
     fwd_bwd_func()
 
-    elem_bytes = x.element_size()
+    elem_bytes = torch.empty(0, dtype=dtype).element_size()
     fwd_bytes = 2 * M * hidden_size * elem_bytes   # read x, write y
     bwd_bytes = 4 * M * hidden_size * elem_bytes   # read grad+x+y, write grad_x
 

benchmarks/microbenchmarks/run_benchmarks.sh

@@ -0,0 +1,38 @@
+#!/bin/bash
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+# Run every microbenchmark.
+#
+#   ./run_all.sh                  # defaults (rotation on)
+#   ./run_all.sh --rotating 512   # rotate within a 512 MB budget
+#   ./run_all.sh --no-rotating    # disable input rotation
+#   ./run_all.sh --csv            # also write per-benchmark CSVs
+#
+# Set PYTHON to pick a specific interpreter (default: python).
+
+shopt -s nullglob
+
+cd "$(dirname "$0")"
+PYTHON="${PYTHON:-python}"
+
+failed=()
+for bench in benchmark_*.py; do
+    echo
+    echo "############################################################"
+    echo "# ${bench} $*"
+    echo "############################################################"
+    if ! "$PYTHON" "$bench" "$@"; then
+        echo "!!! ${bench} FAILED" >&2
+        failed+=("$bench")
+    fi
+done
+
+echo
+if (( ${#failed[@]} )); then
+    echo "FAILED: ${failed[*]}" >&2
+    exit 1
+fi
+echo "All benchmarks completed."

benchmarks/microbenchmarks/utils.py

@@ -7,6 +7,8 @@
 """Shared utilities for microbenchmarks: model configs, timing, throughput, runner."""
 
 import argparse
+import itertools
+import math
 import torch
 import torch.utils.benchmark as benchmark
 
@@ -106,6 +108,112 @@ def time_func(fn, method="adaptive", min_run_time=DEFAULT_MIN_RUN_TIME_SECONDS):
     return m.mean * 1e3, m
 
 
+# ---------------------------------------------------------------------------
+# Rotating input buffers (on by default; disable via --no-rotating)
+# ---------------------------------------------------------------------------
+# Benchmark inputs are cycled through a ring of buffers so that back-to-back
+# kernel launches read different input memory and don't benefit from artificial
+# cache residency.  Populated by run_benchmarks() from the parsed CLI args.
+_ROTATE_BUFFERS = True
+_ROTATE_MB = 0  # rotation memory budget in MB; 0 => auto-size to exceed the LLC
+# Ceiling on the rotation ring size. hipBLASLt-bench caps its rotating block
+# count at the iteration count (max(cold_iters, iters)) so it never allocates a
+# buffer it won't revisit. torch.utils.benchmark picks the iteration count
+# adaptively, so there is no fixed value to cap against; we instead bound the
+# ring at a fixed maximum (mirroring hipBLASLt's default cold_iters of 1000).
+# With the auto budget this ceiling is never reached; it only guards a very
+# large explicit --rotating budget on a small buffer, which would otherwise
+# allocate a copy per few MB up to the whole budget.
+_ROTATE_MAX_BUFFERS = 1000
+
+
+def _last_level_cache_bytes():
+    """Bytes of the last-level cache that buffer rotation must exceed.
+
+    HIP reports ``L2_cache_size`` as the small per-XCD L2 (e.g. 4 MB on gfx950),
+    but the real last-level cache is the much larger AMD Infinity Cache.
+
+    Actual last-level/Infinity Cache sizes:
+      - gfx942 / gfx950: 256 MB
+      - gfx1250:         192 MB
+
+    We use 256 MB for all devices: a slightly oversized ring is harmless (it
+    only allocates a little more memory) and avoids per-arch probing.
+    """
+    return 256 * 1024 * 1024
+
+
+def _rotation_count(bytes_per_buffer, cache_mult=2.0, min_buffers=2):
+    """Number of buffers so the rotation ring spans the requested memory budget.
+
+    With an explicit ``--rotating MB`` the budget is that many megabytes; when
+    omitted it is *cache_mult* x the last-level cache (the ~256 MB AMD Infinity
+    Cache), so a buffer is evicted before it is reused.  The ring is floored at
+    *min_buffers* (so enabling rotation always rotates) and capped at
+    ``_ROTATE_MAX_BUFFERS`` (the adaptive-timer analog of hipBLASLt-bench capping
+    its block count at the iteration count, so a huge budget on a small buffer
+    can't allocate an unbounded ring).
+    """
+    if bytes_per_buffer <= 0:
+        return min_buffers
+    if _ROTATE_MB and _ROTATE_MB > 0:
+        budget = _ROTATE_MB * 1024 * 1024
+    else:
+        cache = _last_level_cache_bytes()
+        if not cache:
+            return min_buffers
+        budget = cache_mult * cache
+    count = math.ceil(budget / bytes_per_buffer)
+    if _ROTATE_MAX_BUFFERS and _ROTATE_MAX_BUFFERS > 0:
+        count = min(count, _ROTATE_MAX_BUFFERS)
+    return max(min_buffers, count)
+
+
+def _tensor_nbytes(t):
+    """Byte size of a torch tensor, or 0 if it can't be determined."""
+    numel = getattr(t, "numel", None)
+    element_size = getattr(t, "element_size", None)
+    if callable(numel) and callable(element_size):
+        return int(numel()) * int(element_size())
+    return 0
+
+
+def rotating(build, *, bytes_per_buffer=None):
+    """Return a zero-arg callable yielding an input buffer to time.
+
+    Rotation is on by default: it builds a ring of ``build()`` buffers (spanning
+    the ``--rotating MB`` budget, or ~2x the last-level cache when the size is
+    omitted) and returns the next one on each call.  With ``--no-rotating`` it
+    returns a single cached buffer from ``build()`` on every call, matching the
+    original single-buffer behavior.
+
+    ``build`` is a zero-arg callable returning one fresh buffer.
+    ``bytes_per_buffer`` overrides the auto-sizing hint for buffers whose byte
+    size can't be inferred (e.g. FP8 tensors).
+    """
+    first = build()
+    if not _ROTATE_BUFFERS:
+        return lambda: first
+    nbytes = bytes_per_buffer if bytes_per_buffer is not None else _tensor_nbytes(first)
+    count = _rotation_count(nbytes)
+    buffers = [first] + [build() for _ in range(max(0, count - 1))]
+    ring = itertools.cycle(buffers)
+    return lambda: next(ring)
+
+
+def make_input(shape, dtype, *, device="cuda", requires_grad=False):
+    """Rotation-aware input: a zero-arg callable returning a ``randn`` tensor.
+
+    Honors ``--rotating`` (see :func:`rotating`); on by default, so it returns
+    the next tensor in the ring each call (``--no-rotating`` for a single one).
+    """
+    return rotating(
+        lambda: torch.randn(
+            *shape, dtype=dtype, device=device, requires_grad=requires_grad
+        )
+    )
+
+
 # ---------------------------------------------------------------------------
 # Throughput helpers
 # ---------------------------------------------------------------------------
@@ -275,6 +383,26 @@ def make_parser(**kwargs):
             "--csv-samples is ignored in this mode."
         ),
     )
+    rotating_group = parser.add_mutually_exclusive_group()
+    rotating_group.add_argument(
+        "--rotating", nargs="?", type=int, const=0, default=None, metavar="MB",
+        help=(
+            "Rotate benchmark inputs through a ring of buffers so back-to-back "
+            "launches touch different memory (avoids artificial cache "
+            "residency), like hipBLASLt-bench --rotating. Optionally pass the "
+            "rotating memory budget in MB; omit it to auto-size the ring to "
+            "exceed the last-level cache (the 256 MB Infinity Cache on "
+            "gfx942/gfx950, not just L2). On by default; disable with "
+            "--no-rotating."
+        ),
+    )
+    rotating_group.add_argument(
+        "--no-rotating", action="store_true", default=False,
+        help=(
+            "Disable input buffer rotation (see --rotating) and time a single "
+            "cached input buffer."
+        ),
+    )
     return parser
 
 
@@ -334,6 +462,13 @@ def run_benchmarks(test_cases, bench_fn, param_columns, default_csv=None,
     if args is None:
         args = make_parser().parse_args()
 
+    global _ROTATE_BUFFERS, _ROTATE_MB
+    _rotating = getattr(args, "rotating", None)
+    if _rotating is not None and _rotating < 0:
+        raise ValueError("--rotating expects a non-negative size in MB")
+    _ROTATE_BUFFERS = not getattr(args, "no_rotating", False)
+    _ROTATE_MB = _rotating or 0
+
     if args.kernel_profile:
         from torch.profiler import profile, ProfilerActivity