microbenchmarks: add buffer rotation option

benchmarks/microbenchmarks/README.md

@@ -31,6 +31,24 @@ 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 reuses a single input buffer, so back-to-back kernel
+launches may read data still resident in cache and report optimistic numbers.
+Pass `--rotating-buffers` to instead cycle inputs through a ring of buffers
+sized to exceed the **last-level cache**, so each launch touches different
+memory (closer to a cold-cache, steady-state workload):
+
+```bash
+python benchmark_gemm.py --rotating-buffers        # auto-size the ring past the LLC
+python benchmark_casting.py --rotating-buffers 16  # or fix the ring size (16 buffers)
+```
+
+Passing `--rotating-buffers N` fixes the number of buffers; omitting `N`
+auto-sizes the ring to ~2x a conservative 256 MB last-level cache (the AMD
+Infinity Cache; see `utils.py::_last_level_cache_bytes`). The option is
+**off by default**.
+
 ## 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

@@ -12,6 +12,7 @@
     compute_tflops,
     make_forward_backward_metric_records,
     run_benchmarks,
+    rotating,
 )
 
 BENCHMARK_LABEL = "Grouped GEMM"
@@ -97,7 +98,7 @@ def generate_grok_v2_test_cases():
     )
 
 
-def make_fwd_bwd_funcs_te(x, w, group_lens, activation_dtype):
+def make_fwd_bwd_funcs_te(next_xs, w, group_lens, activation_dtype):
     from transformer_engine.pytorch.cpp_extensions import general_grouped_gemm
 
     B = int(group_lens.numel())
@@ -107,18 +108,19 @@ def make_fwd_bwd_funcs_te(x, w, group_lens, activation_dtype):
     m_splits = [int(v) for v in group_lens.tolist()]
     assert len(m_splits) == B
     sum_M = sum(m_splits)
-    assert x.numel() > 0 and x.shape[0] == sum_M
 
-    x_view = x.reshape(-1, x.shape[-1])
-    xs = list(torch.split(x_view, m_splits))
     weights = [w[i] for i in range(B)]
+    device = w.device
 
-    out = torch.empty((sum_M, N), device=x.device, dtype=activation_dtype)
+    # next_xs() yields the per-expert split of the (current) activation buffer.
+    # The splits are precomputed once per buffer (see bench_grouped_gemm), so the
+    # timed region below only measures general_grouped_gemm, not reshape/split.
+    out = torch.empty((sum_M, N), device=device, dtype=activation_dtype)
 
     def fwd_func_te():
         general_grouped_gemm(
             A=weights,
-            B=xs,
+            B=next_xs(),
             out=[out],
             quantization_params=[None] * B,
             out_dtype=activation_dtype,
@@ -130,10 +132,10 @@ def fwd_func_te():
         )
         return out
 
-    dx = torch.empty((sum_M, K), device=x.device, dtype=activation_dtype)
+    dx = torch.empty((sum_M, K), device=device, dtype=activation_dtype)
     dxs = list(torch.split(dx, m_splits))
 
-    dw_stacked = torch.empty((B, N, K), device=x.device, dtype=activation_dtype)
+    dw_stacked = torch.empty((B, N, K), device=device, dtype=activation_dtype)
     dws = [dw_stacked[i] for i in range(B)]
 
     def bwd_func_te(grad_out):
@@ -155,7 +157,7 @@ def bwd_func_te(grad_out):
         )
 
         general_grouped_gemm(
-            A=xs,
+            A=next_xs(),
             B=splits,
             out=dws,
             quantization_params=[None] * B,
@@ -177,14 +179,23 @@ def bwd_func_te(grad_out):
 def bench_grouped_gemm(Case, B, M, N, K, dtype):
     device = "cuda"
 
-    x = torch.randn((B * M, K), dtype=dtype, device=device, requires_grad=True)
-    w = torch.randn((B, N, K), dtype=dtype, device=device, requires_grad=True)
+    w = torch.randn((B, N, K), dtype=dtype, device=device)
     group_lens = generate_grouped_gemm_group_lens(B, M, balance=True).to(device)
+    m_splits = [int(v) for v in group_lens.tolist()]
+
+    # Rotation ring of *pre-split* activation buffers, built once at setup, so
+    # the timed region only measures the grouped GEMM (not reshape/split). With
+    # rotation off this is a single cached split, matching the original behavior.
+    elem_bytes = torch.empty(0, dtype=dtype).element_size()
+
+    def _build_xs():
+        x = torch.randn((B * M, K), dtype=dtype, device=device)
+        return list(torch.split(x.reshape(-1, x.shape[-1]), m_splits))
+
+    next_xs = rotating(_build_xs, bytes_per_buffer=B * M * K * elem_bytes)
 
-    x_te = x.clone().detach()
-    w_te = w.clone().detach()
     fwd_func_te, bwd_func_te_inner = make_fwd_bwd_funcs_te(
-        x_te, w_te, group_lens, activation_dtype=dtype
+        next_xs, w, group_lens, activation_dtype=dtype
     )
 
     out_te = fwd_func_te()

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/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,96 @@ def time_func(fn, method="adaptive", min_run_time=DEFAULT_MIN_RUN_TIME_SECONDS):
     return m.mean * 1e3, m
 
 
+# ---------------------------------------------------------------------------
+# Rotating input buffers (opt-in via --rotating-buffers; default off)
+# ---------------------------------------------------------------------------
+# When enabled, 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; the defaults below preserve the original single-buffer
+# behavior.
+_ROTATE_BUFFERS = False
+_ROTATE_COUNT = 0  # 0 => auto-size the ring to exceed the last-level cache
+
+
+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):
+    """Ring size: ``--rotating-buffers N`` if given, else enough to exceed the
+    last-level cache.
+
+    Sizing the ring to at least *cache_mult* x the last-level cache (the ~256 MB
+    AMD Infinity Cache) ensures a buffer is evicted before it is reused.
+    """
+    if _ROTATE_COUNT and _ROTATE_COUNT > 0:
+        return max(1, int(_ROTATE_COUNT))
+    if bytes_per_buffer <= 0:
+        return min_buffers
+    cache = _last_level_cache_bytes()
+    if not cache:
+        return min_buffers
+    return max(min_buffers, math.ceil(cache_mult * cache / bytes_per_buffer))
+
+
+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.
+
+    With ``--rotating-buffers`` off (default) this returns a single cached
+    buffer from ``build()`` on every call, exactly matching the
+    single-buffer behavior.  With the flag on it builds a ring of ``build()``
+    buffers (sized to exceed the last-level cache, or ``--rotating-buffers N``)
+    and returns the next one on each call.
+
+    ``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-buffers`` (see :func:`rotating`); off by default, so it
+    returns the same tensor each call.
+    """
+    return rotating(
+        lambda: torch.randn(
+            *shape, dtype=dtype, device=device, requires_grad=requires_grad
+        )
+    )
+
+
 # ---------------------------------------------------------------------------
 # Throughput helpers
 # ---------------------------------------------------------------------------
@@ -275,6 +367,16 @@ def make_parser(**kwargs):
             "--csv-samples is ignored in this mode."
         ),
     )
+    parser.add_argument(
+        "--rotating-buffers", nargs="?", type=int, const=0, default=None, metavar="N",
+        help=(
+            "Rotate benchmark inputs through a ring of buffers so back-to-back "
+            "launches touch different memory (avoids artificial cache "
+            "residency). Optionally pass N to set the ring size; omit N to "
+            "auto-size the ring to exceed the last-level cache (the 256 MB "
+            "Infinity Cache on gfx942/gfx950, not just L2). Off by default."
+        ),
+    )
     return parser
 
 
@@ -334,6 +436,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_COUNT
+    _rotating = getattr(args, "rotating_buffers", None)
+    if _rotating is not None and _rotating < 0:
+        raise ValueError("--rotating-buffers expects N >= 0")
+    _ROTATE_BUFFERS = _rotating is not None
+    _ROTATE_COUNT = _rotating or 0
+
     if args.kernel_profile:
         from torch.profiler import profile, ProfilerActivity