diff --git a/benchmarks/microbenchmarks/README.md b/benchmarks/microbenchmarks/README.md index ba868b0f0..19c7cba26 100644 --- a/benchmarks/microbenchmarks/README.md +++ b/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`. 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 a576568a0..7b0520389 100755 --- a/benchmarks/microbenchmarks/benchmark_grouped_gemm.py +++ b/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() 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/utils.py b/benchmarks/microbenchmarks/utils.py index 4eef29556..ed9adfd2c 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,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