microbenchmarks: usv implementation

benchmarks/microbenchmarks/README.md

@@ -2,6 +2,8 @@
 
 This directory contains lightweight Python microbenchmarks for selected
 Transformer Engine kernels and helper scripts for comparing benchmark CSVs.
+Timing is powered by [`usv`](https://github.com/matthiasdiener/usv) (install it
+in your environment first, e.g. `pip install -e /path/to/usv`).
 
 ## Benchmarks
 
@@ -31,6 +33,30 @@ 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`.
 
+## Timing options (usv)
+
+Each benchmark prints **one line per run** and finishes with a terminal-only
+min / median / max throughput summary across all cases. Timing is delegated to
+`usv`, so the following optional flags are available on every benchmark (all off
+by default, so the default behavior is unchanged):
+
+| Flag | Meaning |
+| --- | --- |
+| `--interleave` | Sample a run's callables (e.g. fwd / fwd+bwd) round-robin, spreading time-correlated noise across them. |
+| `--warmup N` | Untimed warmup iterations per benchmark. |
+| `--iters N` | Timed samples per benchmark. If unset, sample until ~0.2 s of kernel time elapses (autorange-like). |
+| `--cache-flush` | Flush an L2-sized buffer before each sample (cold-cache timing). |
+| `--cudagraph` | Capture each callable into a CUDA/HIP graph and time replays (removes launch overhead). |
+| `--rotate` | Rotate inputs through an L2-sized ring of buffers (defeats L2 residency). |
+| `--cooldown SECONDS` | Idle sleep after each benchmark to let the GPU cool. |
+| `--monitor` | Sample `rocm-smi` during timing and warn if the GPU clock drifts (AMD). |
+| `--timeout SECONDS` | Abort a benchmark if timing exceeds this many seconds (GPU-hang guard). |
+
+```bash
+python benchmark_gemm.py --interleave --rotate
+python benchmark_gemm.py --iters 200 --cache-flush --monitor
+```
+
 ## Shared configuration
 
 Common benchmark settings live in `utils.py`.
@@ -54,11 +80,16 @@ Use `run_benchmarks(test_cases, bench_fn, param_columns)`.
   `make_metric_record(...)` or `make_forward_backward_metric_records(...)`.
 
 Each metric record represents one benchmark line such as `GEMM Forward`. The
-runner prints that line to stdout and expands it into two CSV columns:
+runner prints all of a run's metrics on a single line and expands each into two
+CSV columns:
 
 - `<label> Time (ms)`
 - `<label> <unit>`
 
+To time the callables, `bench_fn` uses `time_funcs({name: callable})` (which
+honors the usv flags above and enables `--interleave`); `make_input(shape,
+dtype, ...)` returns a rotation-aware input factory that respects `--rotate`.
+
 For example, a `GEMM Forward` metric with unit `TFLOPS` becomes:
 
 - `GEMM Forward Time (ms)`

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,
 )
 
 TE_FP8_E4M3 = tex.DType.kFloat8E4M3
@@ -62,17 +63,18 @@ 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:
+        # Dequantize reads an FP8 tensor; --rotate is a no-op for this direction.
         x = torch.randn(M, hidden_size, dtype=torch.bfloat16, device=device)
         fp8_tensor = quantizer(x)
         cast_func = lambda: fp8_tensor.dequantize()
         total_bytes = numel * (1 + 2)  # FP8 read + BF16 write
 
-    ms, measurement = time_func(cast_func, method="blocked")
+    ms, measurement = time_func(cast_func)
     gbps = compute_gbps(total_bytes, ms)
 
     return [make_metric_record(CAST_LABEL, ms, "GB/s", gbps, measurement=measurement)]

benchmarks/microbenchmarks/benchmark_gemm.py

@@ -10,7 +10,8 @@
 import transformer_engine.pytorch as te
 from utils import (
     generate_gemm_test_cases,
-    time_func, compute_tflops, make_forward_backward_metric_records, run_benchmarks,
+    time_funcs, compute_tflops, make_forward_backward_metric_records, run_benchmarks,
+    make_input,
 )
 
 BENCHMARK_LABEL = "GEMM"
@@ -20,13 +21,16 @@ 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)
+
+    def fwd_func():
+        return linear(next_x())
 
-    fwd_func = lambda: linear(x)
     out = fwd_func()
     grad_out = torch.randn_like(out)
 
     def fwd_bwd_func():
+        x = next_x()
         out = linear(x)
         out.backward(grad_out)
         x.grad = None
@@ -37,9 +41,9 @@ def fwd_bwd_func():
     fwd_flops = 2 * M * N * K
     bwd_flops = 2 * fwd_flops  # dX + dW
 
-    fwd_ms, fwd_measurement = time_func(fwd_func)
-    fwd_bwd_ms, fwd_bwd_measurement = time_func(fwd_bwd_func)
-    bwd_ms = fwd_bwd_ms - fwd_ms
+    ms = time_funcs({"fwd": fwd_func, "fwd_bwd": fwd_bwd_func})
+    fwd_ms = ms["fwd"].median * 1e3
+    bwd_ms = ms["fwd_bwd"].median * 1e3 - fwd_ms
 
     fwd_tflops = compute_tflops(fwd_flops, fwd_ms)
     bwd_tflops = compute_tflops(bwd_flops, bwd_ms)
@@ -52,8 +56,8 @@ def fwd_bwd_func():
         bwd_ms,
         bwd_tflops,
         backward_derived=True,
-        fwd_measurement=fwd_measurement,
-        fwd_bwd_measurement=fwd_bwd_measurement,
+        fwd_measurement=ms["fwd"],
+        fwd_bwd_measurement=ms["fwd_bwd"],
     )
 
 

benchmarks/microbenchmarks/benchmark_gemm_fp8.py

@@ -16,7 +16,8 @@
 from transformer_engine.common.recipe import DelayedScaling, Format
 from utils import (
     generate_gemm_test_cases,
-    time_func, compute_tflops, make_forward_backward_metric_records, run_benchmarks,
+    time_funcs, compute_tflops, make_forward_backward_metric_records, run_benchmarks,
+    make_input,
 )
 
 RECIPES = {
@@ -36,14 +37,15 @@ 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():
+        x = next_x()
         with te.fp8_autocast(enabled=True, fp8_recipe=FP8_RECIPE):
             out = linear(x)
             out.backward(grad_out)
@@ -53,9 +55,9 @@ def fwd_bwd_func():
     fwd_flops = 2 * M * N * K
     bwd_flops = 2 * fwd_flops
 
-    fwd_ms, fwd_measurement = time_func(fwd_func)
-    fwd_bwd_ms, fwd_bwd_measurement = time_func(fwd_bwd_func)
-    bwd_ms = fwd_bwd_ms - fwd_ms
+    ms = time_funcs({"fwd": fwd_func, "fwd_bwd": fwd_bwd_func})
+    fwd_ms = ms["fwd"].median * 1e3
+    bwd_ms = ms["fwd_bwd"].median * 1e3 - fwd_ms
 
     fwd_tflops = compute_tflops(fwd_flops, fwd_ms)
     bwd_tflops = compute_tflops(bwd_flops, bwd_ms)
@@ -68,8 +70,8 @@ def fwd_bwd_func():
         bwd_ms,
         bwd_tflops,
         backward_derived=True,
-        fwd_measurement=fwd_measurement,
-        fwd_bwd_measurement=fwd_bwd_measurement,
+        fwd_measurement=ms["fwd"],
+        fwd_bwd_measurement=ms["fwd_bwd"],
     )
 
 

benchmarks/microbenchmarks/benchmark_grouped_gemm.py

@@ -6,12 +6,14 @@
 ###############################################################################
 
 import torch
+import transformer_engine.pytorch as te
 from utils import (
     DTYPE_LIST,
-    time_func,
+    time_funcs,
     compute_tflops,
     make_forward_backward_metric_records,
     run_benchmarks,
+    make_input,
 )
 
 BENCHMARK_LABEL = "Grouped GEMM"
@@ -97,118 +99,54 @@ def generate_grok_v2_test_cases():
     )
 
 
-def make_fwd_bwd_funcs_te(x, w, group_lens, activation_dtype):
-    from transformer_engine.pytorch.cpp_extensions import general_grouped_gemm
+def bench_grouped_gemm(Case, B, M, N, K, dtype):
+    device = "cuda"
 
-    B = int(group_lens.numel())
-    N = int(w.shape[1])
-    K = int(w.shape[2])
+    # Match the dense GEMM benchmark: drive te.GroupedLinear directly (see #639)
+    # rather than the lower-level general_grouped_gemm, so fwd + dgrad + wgrad and
+    # the framework overhead are all captured.
+    grouped = te.GroupedLinear(B, K, N, bias=False).to(device=device, dtype=dtype)
+    next_x = make_input((B * M, K), dtype, device=device, requires_grad=True)
 
+    group_lens = generate_grouped_gemm_group_lens(B, M, balance=True)
     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)]
-
-    out = torch.empty((sum_M, N), device=x.device, dtype=activation_dtype)
-
-    def fwd_func_te():
-        general_grouped_gemm(
-            A=weights,
-            B=xs,
-            out=[out],
-            quantization_params=[None] * B,
-            out_dtype=activation_dtype,
-            single_output=True,
-            m_splits=m_splits,
-            use_bias=False,
-            bias=None,
-            layout="TN",
-        )
-        return out
-
-    dx = torch.empty((sum_M, K), device=x.device, dtype=activation_dtype)
-    dxs = list(torch.split(dx, m_splits))
-
-    dw_stacked = torch.empty((B, N, K), device=x.device, dtype=activation_dtype)
-    dws = [dw_stacked[i] for i in range(B)]
-
-    def bwd_func_te(grad_out):
-        go = grad_out.view(-1, grad_out.shape[-1])
-        splits = torch.split(go, m_splits)
-
-        general_grouped_gemm(
-            A=weights,
-            B=splits,
-            out=dxs,
-            quantization_params=[None] * B,
-            out_dtype=activation_dtype,
-            single_output=False,
-            layout="NN",
-            m_splits=m_splits,
-            grad=False,
-            use_bias=False,
-            bias=None,
-        )
-
-        general_grouped_gemm(
-            A=xs,
-            B=splits,
-            out=dws,
-            quantization_params=[None] * B,
-            out_dtype=activation_dtype,
-            single_output=False,
-            layout="NT",
-            m_splits=m_splits,
-            grad=False,
-            use_bias=False,
-            bias=None,
-            accumulate=False,
-        )
-
-        return dx, dw_stacked
-
-    return fwd_func_te, bwd_func_te
-
 
-def bench_grouped_gemm(Case, B, M, N, K, dtype):
-    device = "cuda"
+    def fwd_func():
+        return grouped(next_x(), m_splits)
 
-    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)
-    group_lens = generate_grouped_gemm_group_lens(B, M, balance=True).to(device)
+    out = fwd_func()
+    grad_out = torch.randn_like(out)
 
-    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
-    )
+    def fwd_bwd_func():
+        x = next_x()
+        out = grouped(x, m_splits)
+        out.backward(grad_out)
+        x.grad = None
+        for p in grouped.parameters():
+            p.grad = None
 
-    out_te = fwd_func_te()
-    grad_out = torch.randn_like(out_te)
-    bwd_func_te = lambda: bwd_func_te_inner(grad_out)
+    fwd_bwd_func()
 
-    fwd_total_flops = 2 * B * M * N * K
-    bwd_total_flops = 2 * fwd_total_flops
+    fwd_flops = 2 * B * M * N * K
+    bwd_flops = 2 * fwd_flops  # dX + dW
 
-    fwd_te_ms, fwd_measurement = time_func(fwd_func_te)
-    bwd_te_ms, bwd_measurement = time_func(bwd_func_te)
+    ms = time_funcs({"fwd": fwd_func, "fwd_bwd": fwd_bwd_func})
+    fwd_ms = ms["fwd"].median * 1e3
+    bwd_ms = ms["fwd_bwd"].median * 1e3 - fwd_ms
 
-    fwd_te_tflops = compute_tflops(fwd_total_flops, fwd_te_ms)
-    bwd_te_tflops = compute_tflops(bwd_total_flops, bwd_te_ms)
+    fwd_tflops = compute_tflops(fwd_flops, fwd_ms)
+    bwd_tflops = compute_tflops(bwd_flops, bwd_ms)
 
     return make_forward_backward_metric_records(
         BENCHMARK_LABEL,
         "TFLOPS",
-        fwd_te_ms,
-        fwd_te_tflops,
-        bwd_te_ms,
-        bwd_te_tflops,
-        fwd_measurement=fwd_measurement,
-        bwd_measurement=bwd_measurement,
+        fwd_ms,
+        fwd_tflops,
+        bwd_ms,
+        bwd_tflops,
+        backward_derived=True,
+        fwd_measurement=ms["fwd"],
+        fwd_bwd_measurement=ms["fwd_bwd"],
     )