diff --git a/benchmarks/microbenchmarks/README.md b/benchmarks/microbenchmarks/README.md index ba868b0f0..79d570900 100644 --- a/benchmarks/microbenchmarks/README.md +++ b/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`. diff --git a/benchmarks/microbenchmarks/benchmark_casting.py b/benchmarks/microbenchmarks/benchmark_casting.py index 42f279e72..9ad99db8e 100755 --- a/benchmarks/microbenchmarks/benchmark_casting.py +++ b/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") diff --git a/benchmarks/microbenchmarks/benchmark_gemm.py b/benchmarks/microbenchmarks/benchmark_gemm.py index 2069156b7..156dc4bcf 100755 --- a/benchmarks/microbenchmarks/benchmark_gemm.py +++ b/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() diff --git a/benchmarks/microbenchmarks/benchmark_gemm_fp8.py b/benchmarks/microbenchmarks/benchmark_gemm_fp8.py index 7b037b7a6..cdf17fb4c 100755 --- a/benchmarks/microbenchmarks/benchmark_gemm_fp8.py +++ b/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 diff --git a/benchmarks/microbenchmarks/benchmark_grouped_gemm.py b/benchmarks/microbenchmarks/benchmark_grouped_gemm.py index 3f625a6de..3b8b8877c 100755 --- a/benchmarks/microbenchmarks/benchmark_grouped_gemm.py +++ b/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 (off 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 diff --git a/benchmarks/microbenchmarks/benchmark_normalization.py b/benchmarks/microbenchmarks/benchmark_normalization.py index 577bec3d9..7ff86cfbf 100755 --- a/benchmarks/microbenchmarks/benchmark_normalization.py +++ b/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 diff --git a/benchmarks/microbenchmarks/run_benchmarks.sh b/benchmarks/microbenchmarks/run_benchmarks.sh new file mode 100755 index 000000000..5a185177b --- /dev/null +++ b/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." diff --git a/benchmarks/microbenchmarks/utils.py b/benchmarks/microbenchmarks/utils.py index 4eef29556..92bbf6626 100644 --- a/benchmarks/microbenchmarks/utils.py +++ b/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