From 80286736d361e424e544433e1b9fbb6b869eef68 Mon Sep 17 00:00:00 2001
From: Matthias Diener <matthias.diener@gmail.com>
Date: Mon, 31 Aug 2026 18:09:25 -0500
Subject: [PATCH 1/5] microbenchmarks: use pytest as execution backend

---
 .../microbenchmarks/benchmark_casting.py      | 110 +++++---
 benchmarks/microbenchmarks/benchmark_gemm.py  | 131 ++++++----
 .../microbenchmarks/benchmark_grouped_gemm.py | 137 ++++++----
 .../benchmark_normalization.py                |  82 ++++--
 benchmarks/microbenchmarks/conftest.py        | 112 +++++++++
 benchmarks/microbenchmarks/utils.py           | 236 +++++++++++++++++-
 6 files changed, 658 insertions(+), 150 deletions(-)
 create mode 100644 benchmarks/microbenchmarks/conftest.py

diff --git a/benchmarks/microbenchmarks/benchmark_casting.py b/benchmarks/microbenchmarks/benchmark_casting.py
index 2221d1736..b2a1bb2ac 100755
--- a/benchmarks/microbenchmarks/benchmark_casting.py
+++ b/benchmarks/microbenchmarks/benchmark_casting.py
@@ -25,6 +25,7 @@
 Output: benchmark_casting.csv (written to cwd)
 """
 
+import pytest
 import torch
 import transformer_engine
 import transformer_engine_torch as tex
@@ -38,7 +39,7 @@
 )
 from utils import (
     MODEL_HIDDEN_SIZES, M_SIZE_LIST,
-    time_func, compute_gbps, make_metric_record, run_benchmarks,
+    apply_backend_env, time_func, compute_gbps, make_metric_record,
     make_input, rotating,
 )
 
@@ -118,37 +119,73 @@ def _active_formats():
     return formats
 
 
-def _generate_test_cases():
-    test_cases = []
-    active = _active_formats()
+# Backend axis (None unsets, so "default" is the native path even if the ambient
+# env has a toggle set). "triton" flips the Triton kernel for the op being timed:
+# quantize -> NVTE_USE_CAST_TRANSPOSE_TRITON, dequantize -> NVTE_USE_DEQUANTIZE_TRITON.
+CAST_BACKENDS = {
+    "default": {"NVTE_USE_CAST_TRANSPOSE_TRITON": None, "NVTE_USE_DEQUANTIZE_TRITON": None},
+    "triton": {"NVTE_USE_CAST_TRANSPOSE_TRITON": "1", "NVTE_USE_DEQUANTIZE_TRITON": "1"},
+}
+
+_FORMATS = None
+
+
+def _triton_applies(fmt, direction):
+    # Cast-transpose Triton covers FP8/MXFP8/MXFP4 quantize (not NVFP4); the
+    # dequantize Triton path exists only for MXFP8 (mxfp8_tensor_storage).
+    if direction == "quantize":
+        return fmt != "NVFP4"
+    return fmt.startswith("MXFP8")
+
+
+def _backends_for(fmt, direction):
+    return ["default", "triton"] if _triton_applies(fmt, direction) else ["default"]
+
+
+def _formats():
+    """{format_name: (quantizer_factory, quantized_bytes/elem, dequant_supported)}."""
+    global _FORMATS
+    if _FORMATS is None:
+        _FORMATS = {
+            name: (make_quantizer, q_bytes, dequant_supported)
+            for name, make_quantizer, q_bytes, dequant_supported in _active_formats()
+        }
+    return _FORMATS
+
+
+def generate_cases():
+    """Cross models x cast format x direction x backend x M."""
+    cases = []
     for model_name, hidden in MODEL_HIDDEN_SIZES:
-        for fmt_name, make_quantizer, q_bytes_per_elem, dequant_supported in active:
+        for fmt_name, (_mk, _qb, dequant_supported) in _formats().items():
             for direction in DIRECTIONS:
                 if direction == "dequantize" and not dequant_supported:
                     continue
-                cast_name = (
-                    f"BF16-to-{fmt_name}" if direction == "quantize" else f"{fmt_name}-to-BF16"
-                )
-                for M in M_SIZE_LIST:
-                    test_cases.append({
-                        "Case": f"{model_name}/{cast_name}",
-                        "M": M,
-                        "hidden_size": hidden,
-                        "direction": direction,
-                        "make_quantizer": make_quantizer,
-                        "q_bytes_per_elem": q_bytes_per_elem,
-                        "dtype_str": cast_name,
-                    })
-    return test_cases
-
-
-def bench_cast(Case, M, hidden_size, direction, make_quantizer, q_bytes_per_elem, dtype_str):
+                for backend in _backends_for(fmt_name, direction):
+                    for M in M_SIZE_LIST:
+                        cases.append({
+                            "Case": model_name,
+                            "Format": fmt_name,
+                            "Direction": direction,
+                            "Backend": backend,
+                            "M": M,
+                            "hidden_size": hidden,
+                        })
+    return cases
+
+
+def _case_id(c):
+    return f"{c['Case']}-{c['Format']}-{c['Direction']}-{c['Backend']}-M{c['M']}"
+
+
+def bench_cast(Format, Direction, M, hidden_size):
     device = "cuda"
 
+    make_quantizer, q_bytes_per_elem, _deq = _formats()[Format]
     numel = M * hidden_size
     quantizer = make_quantizer()
 
-    if direction == "quantize":
+    if Direction == "quantize":
         next_x = make_input((M, hidden_size), torch.bfloat16, device=device)
         out = quantizer(next_x())
         cast_func = lambda: quantizer.quantize(next_x(), out=out)
@@ -165,14 +202,27 @@ def bench_cast(Case, M, hidden_size, direction, make_quantizer, q_bytes_per_elem
         total_bytes = int(numel * (q_bytes_per_elem + 2))  # quantized read + BF16 write
 
     ms, measurement = time_func(cast_func, method="blocked")
-    gbps = compute_gbps(total_bytes, ms)
+    return [make_metric_record(
+        CAST_LABEL, ms, "GB/s", compute_gbps(total_bytes, ms), measurement=measurement,
+    )]
 
-    return [make_metric_record(CAST_LABEL, ms, "GB/s", gbps, measurement=measurement)]
 
+def pytest_generate_tests(metafunc):
+    if "case" in metafunc.fixturenames:
+        cases = generate_cases()
+        metafunc.parametrize("case", cases, ids=[_case_id(c) for c in cases])
 
-if __name__ == "__main__":
-    run_benchmarks(
-        test_cases=_generate_test_cases(),
-        bench_fn=bench_cast,
-        param_columns=["Case", "M", "hidden_size", "dtype_str"],
+
+@pytest.mark.benchmark
+def test_cast(microbench, case, monkeypatch):
+    apply_backend_env(monkeypatch, CAST_BACKENDS[case["Backend"]])
+    microbench.run(
+        case,
+        lambda: bench_cast(case["Format"], case["Direction"], case["M"], case["hidden_size"]),
     )
+
+
+if __name__ == "__main__":
+    import sys
+    # Make the file runnable directly: python benchmark_casting.py [--csv -k ...].
+    raise SystemExit(pytest.main([__file__, *sys.argv[1:]]))
diff --git a/benchmarks/microbenchmarks/benchmark_gemm.py b/benchmarks/microbenchmarks/benchmark_gemm.py
index 29445179f..329354841 100755
--- a/benchmarks/microbenchmarks/benchmark_gemm.py
+++ b/benchmarks/microbenchmarks/benchmark_gemm.py
@@ -4,22 +4,30 @@
 #
 # See LICENSE for license information.
 ###############################################################################
-"""Dense GEMM micro-benchmark using te.Linear across precisions.
+"""Dense GEMM micro-benchmark using te.Linear across precisions and backends.
 
-Sweeps the shared model GEMM shapes over BF16 (the high-precision baseline)
-plus every supported low-precision recipe (FP8, MXFP8, MXFP4, NVFP4) via
-te.autocast. Precisions whose hardware/runtime support is unavailable on the
-current device are skipped automatically.
+Runs under pytest (see conftest.py). Sweeps the shared model GEMM shapes over
+BF16 (the high-precision baseline) plus every supported low-precision recipe
+(FP8, MXFP8, MXFP4, NVFP4) via te.autocast, crossed with a selectable kernel
+Backend. Precisions whose hardware/runtime support is unavailable on the current
+device are skipped automatically.
 
-Output: benchmark_gemm.csv (written to cwd)
+Examples::
+
+    pytest benchmark_gemm.py --csv                 # -> benchmark_gemm.csv
+    pytest benchmark_gemm.py -k "bf16 and QKV"     # select shapes/precisions
+    pytest benchmark_gemm.py -k triton             # select the triton backend
+
+Output: benchmark_gemm.csv (written to cwd when --csv is passed).
 """
 
+import pytest
 import torch
 import transformer_engine.pytorch as te
 from utils import (
     build_recipes,
     generate_gemm_test_cases,
-    time_func, compute_tflops, make_forward_backward_metric_records, run_benchmarks,
+    apply_backend_env, compute_tflops, direction_records,
     make_input,
 )
 
@@ -27,17 +35,56 @@
 
 RECIPES = build_recipes()
 
+# Env recipes to force a dense-GEMM kernel backend (None unsets the var). Per the
+# C++ dispatch: bf16 defaults to hipBLASLt, forced to Triton via NVTE_USE_GEMM_TRITON;
+# mxfp8 defaults to HipKittens, forced to hipBLASLt via NVTE_ROCM_USE_HIPBLASLT_MXFP8
+# (rocm_gemm.cu). fp8 has a single backend.
+_GEMM_TRITON = "NVTE_USE_GEMM_TRITON"
+_HIPBLASLT_MXFP8 = "NVTE_ROCM_USE_HIPBLASLT_MXFP8"
 
-def generate_precision_gemm_test_cases():
-    """Cross the shared dense GEMM shapes with each supported precision."""
-    test_cases = []
-    for base_case in generate_gemm_test_cases():
-        for precision in RECIPES:
-            test_cases.append({**base_case, "Precision": precision})
-    return test_cases
+GEMM_BACKENDS = {
+    "hipblaslt":  {_GEMM_TRITON: None, _HIPBLASLT_MXFP8: "1"},
+    "triton":     {_GEMM_TRITON: "1", _HIPBLASLT_MXFP8: None},
+    "hipkittens": {_GEMM_TRITON: None, _HIPBLASLT_MXFP8: None},
+}
+
+# Backends with a real choice per precision (the supported-backends table).
+_BACKENDS_BY_PRECISION = {
+    "bf16": ["hipblaslt", "triton"],
+    "fp8": ["hipblaslt"],
+    "mxfp8": ["hipblaslt", "hipkittens"],
+}
 
 
-def bench_gemm(Case, Precision, M, N, K, dtype):
+def _backends_for(precision):
+    return _BACKENDS_BY_PRECISION.get(precision, ["hipblaslt"])
+
+
+def generate_cases():
+    """Cross the shared dense GEMM shapes with each precision, backend, direction."""
+    cases = []
+    for base in generate_gemm_test_cases():
+        for precision in RECIPES:
+            for backend in _backends_for(precision):
+                for direction in ("fwd", "bwd"):
+                    cases.append({
+                        "Case": base["Case"],
+                        "Precision": precision,
+                        "Backend": backend,
+                        "Direction": direction,
+                        "M": base["M"],
+                        "N": base["N"],
+                        "K": base["K"],
+                        "dtype": base["dtype"],
+                    })
+    return cases
+
+
+def _case_id(c):
+    return f"{c['Case']}-{c['Precision']}-{c['Backend']}-{c['Direction']}-M{c['M']}"
+
+
+def bench_gemm(Case, Precision, Direction, M, N, K, dtype):
     device = "cuda"
 
     recipe = RECIPES[Precision]
@@ -56,39 +103,37 @@ def fwd_func():
     def fwd_bwd_func():
         xb = next_x()
         with te.autocast(enabled=use_fp8, recipe=recipe):
-            out = linear(xb)
-            out.backward(grad_out)
+            o = linear(xb)
+            o.backward(grad_out)
         xb.grad = None
         linear.weight.grad = None
 
-    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
-
-    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_ms,
-        fwd_tflops,
-        bwd_ms,
-        bwd_tflops,
-        backward_derived=True,
-        fwd_measurement=fwd_measurement,
-        fwd_bwd_measurement=fwd_bwd_measurement,
+    return direction_records(
+        Direction, BENCHMARK_LABEL, "TFLOPS", compute_tflops,
+        fwd_func, fwd_bwd_func, fwd_flops, 2 * fwd_flops,
     )
 
 
-if __name__ == "__main__":
-    run_benchmarks(
-        test_cases=generate_precision_gemm_test_cases(),
-        bench_fn=bench_gemm,
-        param_columns=["Case", "Precision", "M", "N", "K", "dtype"],
+def pytest_generate_tests(metafunc):
+    if "case" in metafunc.fixturenames:
+        cases = generate_cases()
+        metafunc.parametrize("case", cases, ids=[_case_id(c) for c in cases])
+
+
+@pytest.mark.benchmark
+def test_gemm(microbench, case, monkeypatch):
+    apply_backend_env(monkeypatch, GEMM_BACKENDS[case["Backend"]])
+    microbench.run(
+        case,
+        lambda: bench_gemm(
+            case["Case"], case["Precision"], case["Direction"],
+            case["M"], case["N"], case["K"], case["dtype"],
+        ),
     )
+
+
+if __name__ == "__main__":
+    import sys
+    # Make the file runnable directly: python benchmark_gemm.py [--csv -k ...].
+    raise SystemExit(pytest.main([__file__, *sys.argv[1:]]))
diff --git a/benchmarks/microbenchmarks/benchmark_grouped_gemm.py b/benchmarks/microbenchmarks/benchmark_grouped_gemm.py
index d95e184c4..007767ea8 100755
--- a/benchmarks/microbenchmarks/benchmark_grouped_gemm.py
+++ b/benchmarks/microbenchmarks/benchmark_grouped_gemm.py
@@ -4,25 +4,59 @@
 #
 # See LICENSE for license information.
 ###############################################################################
+"""Grouped GEMM micro-benchmark using te.GroupedLinear across precisions and backends.
 
+Runs under pytest (see conftest.py). Sweeps MoE grouped-GEMM shapes over BF16,
+FP8, and MXFP8, crossed with the selectable kernel backend for each precision
+(hipBLASLt / CK_Tile / Triton / HipKittens) and forward/backward direction.
+
+    pytest benchmark_grouped_gemm.py --csv
+    pytest benchmark_grouped_gemm.py -k "mxfp8 and hipkittens"
+"""
+
+import pytest
 import torch
 import transformer_engine.pytorch as te
 from utils import (
     DTYPE_LIST,
+    apply_backend_env,
     build_recipes,
-    time_func,
     compute_tflops,
-    make_forward_backward_metric_records,
-    run_benchmarks,
+    direction_records,
     make_input,
 )
 
 BENCHMARK_LABEL = "Grouped GEMM"
 
-# Same precision sweep as benchmark_gemm.py, minus MXFP4 (GroupedLinear has no
-# MXFP4 grouped kernel). Each test case carries a recipe label; bf16 maps to
-# None (plain path) and unsupported precisions are skipped by build_recipes().
-RECIPES = build_recipes(names=("bf16", "fp8", "mxfp8", "nvfp4"))
+# bf16/fp8/mxfp8 grouped GEMM (no grouped MXFP4 kernel yet -- that's a separate PR).
+RECIPES = build_recipes(names=("bf16", "fp8", "mxfp8"))
+
+# Env recipes to force a grouped-GEMM kernel backend (None unsets the var). Per the
+# C++ dispatch (cublaslt_gemm.cu / rocm_gemm.cu): all-unset -> multi-stream hipBLASLt;
+# CUTLASS+CK -> CK; CUTLASS+HK -> HipKittens; NVTE_USE_GROUPED_GEMM_TRITON routes bf16
+# to the Triton grouped GEMM.
+_CUTLASS = "NVTE_USE_CUTLASS_GROUPED_GEMM"
+_CK = "NVTE_USE_CK_GROUPED_GEMM"
+_HK = "NVTE_USE_HIPKITTENS_GROUPED_GEMM"
+_TRITON = "NVTE_USE_GROUPED_GEMM_TRITON"
+
+GROUPED_BACKENDS = {
+    "hipblaslt":  {_CUTLASS: None, _CK: None, _HK: None, _TRITON: None},
+    "ck_tile":    {_CUTLASS: "1", _CK: "1", _HK: None, _TRITON: None},
+    "hipkittens": {_CUTLASS: "1", _CK: None, _HK: "1", _TRITON: None},
+    "triton":     {_CUTLASS: None, _CK: None, _HK: None, _TRITON: "1"},
+}
+
+# Backends with a real choice per precision (the supported-backends table).
+_BACKENDS_BY_PRECISION = {
+    "bf16": ["hipblaslt", "ck_tile", "triton"],
+    "fp8": ["hipblaslt", "ck_tile"],
+    "mxfp8": ["hipblaslt", "hipkittens"],
+}
+
+
+def _backends_for(recipe):
+    return _BACKENDS_BY_PRECISION.get(recipe, ["hipblaslt"])
 
 def generate_grouped_gemm_group_lens(b, m, balance: bool):
     if balance:
@@ -107,7 +141,7 @@ def generate_grok_v2_test_cases():
     )
 
 
-def bench_grouped_gemm(Case, B, M, N, K, dtype, recipe):
+def bench_grouped_gemm(Case, B, M, N, K, dtype, recipe, Direction):
     device = "cuda"
 
     fp8_recipe = RECIPES[recipe]
@@ -119,25 +153,20 @@ def bench_grouped_gemm(Case, B, M, N, K, dtype, recipe):
     sum_M = sum(m_splits)
 
     grouped_linear = te.GroupedLinear(
-        B,
-        K,
-        N,
-        bias=False,
-        params_dtype=dtype,
-        device=device,
+        B, K, N, bias=False, params_dtype=dtype, device=device,
     )
-    # Rotate the activation buffer (on by default) so back-to-back grouped GEMMs 
+    # 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():
+    def fwd_func():
         with te.autocast(enabled=use_fp8, recipe=fp8_recipe):
             return grouped_linear(next_x(), m_splits, m_splits_tensor=m_splits_tensor)
 
-    out_te = fwd_func_te()
+    out_te = fwd_func()
     grad_out = torch.randn_like(out_te)
 
-    def fwd_bwd_func_te():
+    def fwd_bwd_func():
         xb = next_x()
         with te.autocast(enabled=use_fp8, recipe=fp8_recipe):
             out = grouped_linear(xb, m_splits, m_splits_tensor=m_splits_tensor)
@@ -146,41 +175,57 @@ def fwd_bwd_func_te():
         for param in grouped_linear.parameters():
             param.grad = None
 
-    fwd_bwd_func_te()
-
     fwd_total_flops = 2 * sum_M * N * K
-    bwd_total_flops = 2 * fwd_total_flops
-
-    fwd_te_ms, fwd_measurement = time_func(fwd_func_te)
-    fwd_bwd_te_ms, fwd_bwd_measurement = time_func(fwd_bwd_func_te)
-    bwd_te_ms = fwd_bwd_te_ms - fwd_te_ms
-
-    fwd_te_tflops = compute_tflops(fwd_total_flops, fwd_te_ms)
-    bwd_te_tflops = compute_tflops(bwd_total_flops, bwd_te_ms)
-
-    return make_forward_backward_metric_records(
-        BENCHMARK_LABEL,
-        "TFLOPS",
-        fwd_te_ms,
-        fwd_te_tflops,
-        bwd_te_ms,
-        bwd_te_tflops,
-        backward_derived=True,
-        fwd_measurement=fwd_measurement,
-        fwd_bwd_measurement=fwd_bwd_measurement,
+    return direction_records(
+        Direction, BENCHMARK_LABEL, "TFLOPS", compute_tflops,
+        fwd_func, fwd_bwd_func, fwd_total_flops, 2 * fwd_total_flops,
     )
 
 
-if __name__ == "__main__":
-    test_cases = (
+def generate_cases():
+    """MoE grouped-GEMM cases crossed with per-precision backend and direction."""
+    base = (
         generate_deepseekv2_lite_test_cases()
         + generate_deepseekv2_test_cases()
         + generate_deepseekv3_test_cases()
         + generate_grok_v2_test_cases()
     )
-
-    run_benchmarks(
-        test_cases=test_cases,
-        bench_fn=bench_grouped_gemm,
-        param_columns=["Case", "B", "M", "N", "K", "dtype", "recipe"],
+    cases = []
+    for b in base:
+        for backend in _backends_for(b["recipe"]):
+            for direction in ("fwd", "bwd"):
+                cases.append({**b, "Backend": backend, "Direction": direction})
+    return cases
+
+
+def _case_id(c):
+    return f"{c['Case']}-{c['recipe']}-{c['Backend']}-{c['Direction']}-B{c['B']}-M{c['M']}"
+
+
+def pytest_generate_tests(metafunc):
+    if "case" in metafunc.fixturenames:
+        cases = generate_cases()
+        metafunc.parametrize("case", cases, ids=[_case_id(c) for c in cases])
+
+
+@pytest.mark.benchmark
+def test_grouped_gemm(microbench, case, monkeypatch):
+    backend = case["Backend"]
+    if backend in ("ck_tile", "hipkittens") and case["B"] <= 1:
+        pytest.skip(f"{backend} grouped GEMM needs num_groups > 1")
+    if backend == "hipkittens" and (case["N"] % 256 or case["K"] % 256):
+        pytest.skip("HipKittens grouped GEMM needs 256-aligned expert dims")
+    apply_backend_env(monkeypatch, GROUPED_BACKENDS[backend])
+    microbench.run(
+        case,
+        lambda: bench_grouped_gemm(
+            case["Case"], case["B"], case["M"], case["N"], case["K"],
+            case["dtype"], case["recipe"], case["Direction"],
+        ),
     )
+
+
+if __name__ == "__main__":
+    import sys
+    # Make the file runnable directly: python benchmark_grouped_gemm.py [--csv -k ...].
+    raise SystemExit(pytest.main([__file__, *sys.argv[1:]]))
diff --git a/benchmarks/microbenchmarks/benchmark_normalization.py b/benchmarks/microbenchmarks/benchmark_normalization.py
index ff9f7b28b..135ceea37 100755
--- a/benchmarks/microbenchmarks/benchmark_normalization.py
+++ b/benchmarks/microbenchmarks/benchmark_normalization.py
@@ -24,13 +24,14 @@
 Output: benchmark_normalization.csv (written to cwd)
 """
 
+import pytest
 import torch
 import transformer_engine.pytorch as te
 from transformer_engine.pytorch import ops
 from utils import (
     MODEL_HIDDEN_SIZES, M_SIZE_LIST,
     build_recipes,
-    time_func, compute_gbps, make_metric_record, run_benchmarks,
+    apply_backend_env, time_func, compute_gbps, make_metric_record,
     make_input,
 )
 
@@ -61,28 +62,47 @@
 }
 
 
-def _generate_test_cases():
-    test_cases = []
+# Backend axis (None unsets, so "default" is the native path even if the ambient
+# env has a toggle set). "triton" flips the Triton RMSNorm/LayerNorm kernel (each
+# NormType reads only its own toggle).
+NORM_BACKENDS = {
+    "default": {"NVTE_USE_RMSNORM_TRITON": None, "NVTE_USE_LAYERNORM_TRITON": None},
+    "triton": {"NVTE_USE_RMSNORM_TRITON": "1", "NVTE_USE_LAYERNORM_TRITON": "1"},
+}
+
+_NORM_CLS = {name: cls for name, cls in NORM_TYPES}
+
+
+def generate_cases():
+    """Cross models x norm type x precision x backend x M (forward only)."""
+    cases = []
     for model_name, hidden in MODEL_HIDDEN_SIZES:
-        for norm_name, norm_op_cls in NORM_TYPES:
+        for norm_name in _NORM_CLS:
             for precision in RECIPES:
-                for M in M_SIZE_LIST:
-                    test_cases.append({
-                        "Case": f"{model_name}/{norm_name}",
-                        "Precision": precision,
-                        "M": M,
-                        "hidden_size": hidden,
-                        "norm_op_cls": norm_op_cls,
-                        "dtype": torch.bfloat16,
-                    })
-    return test_cases
-
-
-def bench_norm(Case, Precision, M, hidden_size, norm_op_cls, dtype):
+                for backend in NORM_BACKENDS:
+                    for M in M_SIZE_LIST:
+                        cases.append({
+                            "Case": model_name,
+                            "NormType": norm_name,
+                            "Precision": precision,
+                            "Backend": backend,
+                            "M": M,
+                            "hidden_size": hidden,
+                        })
+    return cases
+
+
+def _case_id(c):
+    return f"{c['Case']}-{c['NormType']}-{c['Precision']}-{c['Backend']}-M{c['M']}"
+
+
+def bench_norm(NormType, Precision, M, hidden_size):
     device = "cuda"
+    dtype = torch.bfloat16
 
     recipe = RECIPES[Precision]
     use_fp8 = recipe is not None
+    norm_op_cls = _NORM_CLS[NormType]
 
     # Norm followed by Quantize so the norm writes its output directly in the
     # target precision under autocast (identity when use_fp8 is False).
@@ -102,16 +122,28 @@ def fwd_func():
     fwd_bytes = int(M * hidden_size * (2 + _FWD_WRITE_BYTES[Precision]))
 
     fwd_ms, fwd_measurement = time_func(fwd_func)
-    fwd_gbps = compute_gbps(fwd_bytes, fwd_ms)
-
     return [make_metric_record(
-        BENCHMARK_LABEL, fwd_ms, "GB/s", fwd_gbps, measurement=fwd_measurement,
+        BENCHMARK_LABEL, fwd_ms, "GB/s", compute_gbps(fwd_bytes, fwd_ms),
+        measurement=fwd_measurement,
     )]
 
 
-if __name__ == "__main__":
-    run_benchmarks(
-        test_cases=_generate_test_cases(),
-        bench_fn=bench_norm,
-        param_columns=["Case", "Precision", "M", "hidden_size", "dtype"],
+def pytest_generate_tests(metafunc):
+    if "case" in metafunc.fixturenames:
+        cases = generate_cases()
+        metafunc.parametrize("case", cases, ids=[_case_id(c) for c in cases])
+
+
+@pytest.mark.benchmark
+def test_norm(microbench, case, monkeypatch):
+    apply_backend_env(monkeypatch, NORM_BACKENDS[case["Backend"]])
+    microbench.run(
+        case,
+        lambda: bench_norm(case["NormType"], case["Precision"], case["M"], case["hidden_size"]),
     )
+
+
+if __name__ == "__main__":
+    import sys
+    # Make the file runnable directly: python benchmark_normalization.py [--csv -k ...].
+    raise SystemExit(pytest.main([__file__, *sys.argv[1:]]))
diff --git a/benchmarks/microbenchmarks/conftest.py b/benchmarks/microbenchmarks/conftest.py
new file mode 100644
index 000000000..802f9d0e4
--- /dev/null
+++ b/benchmarks/microbenchmarks/conftest.py
@@ -0,0 +1,112 @@
+#!/usr/bin/env python
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+"""pytest glue for the microbenchmarks.
+
+Thin shim: the option/CSV/timing logic lives in utils.py; this file only wires
+those helpers into pytest hooks and exposes the ``microbench`` fixture.  Run a
+family with, e.g.::
+
+    pytest benchmark_gemm.py --csv                 # write benchmark_gemm.csv
+    pytest benchmark_gemm.py -k "bf16 and QKV"     # select by parametrize id
+    pytest benchmark_gemm.py -k triton             # select the triton backend
+"""
+
+from pathlib import Path
+
+import pytest
+
+from utils import (
+    collect_kernel_rows,
+    configure_rotating,
+    format_results_table,
+    print_case,
+    record_bench,
+    write_bench_outputs,
+)
+
+
+def pytest_addoption(parser):
+    group = parser.getgroup("microbench", "TE GPU microbenchmarks")
+    group.addoption(
+        "--csv", nargs="?", const=True, default=None, metavar="FILE",
+        help="Write results to CSV (one per family; default name from the module).",
+    )
+    group.addoption(
+        "--csv-samples", nargs="?", const=True, default=None, metavar="FILE",
+        help="Write per-sample timing data to a CSV.",
+    )
+    group.addoption(
+        "--kernel-profile", action="store_true", default=False,
+        help="Also profile GPU kernels via torch.profiler and write a _kernel_profile CSV.",
+    )
+    group.addoption(
+        "--rotating", nargs="?", type=int, const=0, default=None, metavar="MB",
+        help="Rotate inputs through a ring of buffers (optional MB budget). On by default.",
+    )
+    group.addoption(
+        "--no-rotating", action="store_true", default=False,
+        help="Disable input buffer rotation.",
+    )
+
+
+def pytest_configure(config):
+    config.addinivalue_line("markers", "benchmark: TE GPU microbenchmark")
+    configure_rotating(config.getoption("--rotating"), config.getoption("--no-rotating"))
+    config._microbench_store = {}
+
+
+class _MicroBench:
+    """Handed to each test as the ``microbench`` fixture."""
+
+    def __init__(self, request):
+        self._request = request
+        self._config = request.config
+
+    def run(self, case, bench_callable):
+        """Time *bench_callable* (returns metric records) and record it under *case*."""
+        records = bench_callable()
+        print_case(case, records)
+        family = Path(self._request.module.__file__).stem
+        kernel_rows = None
+        if self._config.getoption("--kernel-profile"):
+            kernel_rows = collect_kernel_rows(bench_callable, case)
+        record_bench(
+            self._config._microbench_store, family, case, records,
+            kernel_rows, self._request.node.name,
+        )
+        return records
+
+
+@pytest.fixture
+def microbench(request):
+    return _MicroBench(request)
+
+
+def pytest_terminal_summary(terminalreporter, exitstatus, config):
+    store = getattr(config, "_microbench_store", None)
+    if not store:
+        return
+    table = format_results_table(store)
+    if table:
+        terminalreporter.write_line("")
+        for row in table.splitlines():
+            terminalreporter.write_line(row)
+
+
+def pytest_sessionfinish(session, exitstatus):
+    config = session.config
+    store = getattr(config, "_microbench_store", None)
+    if not store:
+        return
+    written = write_bench_outputs(
+        store,
+        csv=config.getoption("--csv"),
+        csv_samples=config.getoption("--csv-samples"),
+        kernel_profile=config.getoption("--kernel-profile"),
+    )
+    for path in written:
+        print(f"microbench: wrote {path}")
diff --git a/benchmarks/microbenchmarks/utils.py b/benchmarks/microbenchmarks/utils.py
index 45101b60f..e7073c948 100644
--- a/benchmarks/microbenchmarks/utils.py
+++ b/benchmarks/microbenchmarks/utils.py
@@ -10,6 +10,7 @@
 import importlib.util
 import itertools
 import math
+from types import SimpleNamespace
 import torch
 import torch.utils.benchmark as benchmark
 
@@ -389,6 +390,35 @@ def make_forward_backward_metric_records(label_prefix, unit,
     return records
 
 
+def direction_records(direction, label, unit, throughput,
+                      fwd_func, fwd_bwd_func, fwd_work, bwd_work):
+    """Metric records for a forward-only or a derived-backward timing.
+
+    *direction* is ``"fwd"`` or ``"bwd"``. *throughput* is ``compute_tflops`` or
+    ``compute_gbps`` and *fwd_work* / *bwd_work* the matching flops / bytes.
+    Backward is ``(fwd+bwd) - fwd``; its per-sample distribution is each fwd+bwd
+    sample shifted by the fwd mean (fwd and fwd+bwd are timed separately, so the
+    spread is inherited from fwd+bwd).
+    """
+    if direction == "fwd":
+        fwd_ms, fwd_measurement = time_func(fwd_func)
+        return [make_metric_record(
+            label, fwd_ms, unit, throughput(fwd_work, fwd_ms), measurement=fwd_measurement,
+        )]
+    fwd_bwd_func()  # warm the backward graph
+    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
+    fwd_mean_s = fwd_measurement.mean
+    bwd_measurement = SimpleNamespace(
+        times=[t - fwd_mean_s for t in fwd_bwd_measurement.times]
+    )
+    return [make_metric_record(
+        label, bwd_ms, unit, throughput(bwd_work, bwd_ms),
+        derived=True, measurement=bwd_measurement,
+    )]
+
+
 def _metric_time_key(metric):
     return f"{metric['label']} Time (ms)"
 
@@ -548,12 +578,7 @@ 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
+    configure_rotating(getattr(args, "rotating", None), getattr(args, "no_rotating", False))
 
     if args.kernel_profile:
         from torch.profiler import profile, ProfilerActivity
@@ -696,3 +721,202 @@ def run_benchmarks(test_cases, bench_fn, param_columns, default_csv=None,
             )
             df.to_csv(samples_csv, index=False)
             print(f"Samples saved to {samples_csv}")
+
+
+# ---------------------------------------------------------------------------
+# pytest-based execution support
+# ---------------------------------------------------------------------------
+# The microbenchmarks can also run under pytest; conftest.py is a thin shim over
+# the framework-agnostic helpers below (no pytest import here, so importing
+# utils.py never requires pytest). Results are collected per family (test module)
+# and written with the same CSV / samples / kernel-profile schema run_benchmarks
+# produces, so downstream tooling (e.g. the dashboard ingest) is unaffected.
+
+def configure_rotating(rotating, no_rotating):
+    """Set module-level input-rotation state from parsed options."""
+    global _ROTATE_BUFFERS, _ROTATE_MB
+    if rotating is not None and rotating < 0:
+        raise ValueError("--rotating expects a non-negative size in MB")
+    _ROTATE_BUFFERS = not no_rotating
+    _ROTATE_MB = rotating or 0
+
+
+def apply_backend_env(monkeypatch, env):
+    """Force a kernel backend for one test by setting/unsetting env vars.
+
+    A ``None`` value unsets the var, so forcing one backend cleanly clears the
+    toggles that would select a competing one; pytest restores them afterwards.
+    """
+    for key, value in env.items():
+        if value is None:
+            monkeypatch.delenv(key, raising=False)
+        else:
+            monkeypatch.setenv(key, value)
+
+
+class _FamilyResults:
+    """Accumulated rows / samples / kernel rows for one benchmark family."""
+
+    def __init__(self):
+        self.param_columns = None
+        self.metric_columns = None
+        self.rows = []
+        self.case_metrics = []
+        self.kernel_rows = []
+
+
+def _stringify_params(case_params):
+    return {k: (str(v) if isinstance(v, torch.dtype) else v) for k, v in case_params.items()}
+
+
+def record_bench(store, family, case_params, metric_records, kernel_rows=None, node_name=""):
+    """Record one benchmark case into *store* (a dict keyed by *family*)."""
+    fam = store.setdefault(family, _FamilyResults())
+    metric_row = _metric_row_from_records(metric_records)
+    metric_columns = list(metric_row.keys())
+    if fam.param_columns is None:
+        fam.param_columns = list(case_params.keys())
+        fam.metric_columns = metric_columns
+    elif metric_columns != fam.metric_columns:
+        raise ValueError(
+            f"Inconsistent metric columns for {family}: "
+            f"expected {fam.metric_columns}, got {metric_columns}"
+        )
+    row = _stringify_params(case_params)
+    row.update(metric_row)
+    fam.rows.append(row)
+    fam.case_metrics.append((_stringify_params(case_params), metric_records, node_name))
+    if kernel_rows:
+        fam.kernel_rows.extend(kernel_rows)
+
+
+def print_case(case_params, metric_records):
+    """Print a case header and its metric lines (reused stdout format)."""
+    label = "  ".join(f"{k}={v}" for k, v in case_params.items())
+    print(f"\n{'='*60}\nTesting: {label}\n{'='*60}")
+    _print_metric_records(metric_records)
+
+
+def collect_kernel_rows(bench_callable, case_params):
+    """Re-run *bench_callable* under torch.profiler and return per-kernel rows."""
+    from torch.profiler import profile, ProfilerActivity
+    with profile(activities=[ProfilerActivity.CUDA]) as prof:
+        bench_callable()
+        torch.cuda.synchronize()
+    events = [e for e in prof.key_averages() if e.self_device_time_total > 0]
+    events.sort(key=lambda e: e.self_device_time_total, reverse=True)
+    params = _stringify_params(case_params)
+    rows = []
+    for e in events:
+        kr = dict(params)
+        kr["kernel_name"] = e.key
+        kr["cuda_time_total_us"] = round(e.self_device_time_total, 1)
+        kr["num_calls"] = e.count
+        kr["cuda_time_avg_us"] = round(e.self_device_time_total / e.count, 2) if e.count else 0
+        rows.append(kr)
+    return rows
+
+
+def write_bench_outputs(store, *, csv=None, csv_samples=None, kernel_profile=False):
+    """Write per-family CSV / samples / kernel-profile outputs; return paths written."""
+    import pandas as pd
+    from pathlib import Path
+
+    written = []
+    for family, fam in store.items():
+        if not fam.rows:
+            continue
+        if csv is not None:
+            out = csv if isinstance(csv, str) else f"{family}.csv"
+            pd.DataFrame(fam.rows, columns=fam.param_columns + fam.metric_columns).to_csv(
+                out, index=False
+            )
+            written.append(out)
+            if kernel_profile and fam.kernel_rows:
+                kout = f"{Path(out).stem}_kernel_profile.csv"
+                cols = fam.param_columns + [
+                    "kernel_name", "cuda_time_total_us", "num_calls", "cuda_time_avg_us",
+                ]
+                pd.DataFrame(fam.kernel_rows, columns=cols).to_csv(kout, index=False)
+                written.append(kout)
+        if csv_samples is not None:
+            sout = csv_samples if isinstance(csv_samples, str) else f"{family}_samples.csv"
+            sample_rows = []
+            for case_params, records, _node in fam.case_metrics:
+                for metric in records:
+                    m = metric.get("measurement")
+                    if m is None:
+                        continue
+                    for i, t in enumerate(m.times):
+                        sr = dict(case_params)
+                        sr["label"] = metric["label"]
+                        sr["sample_idx"] = i
+                        sr["time_ms"] = t * 1e3
+                        sample_rows.append(sr)
+            if sample_rows:
+                pd.DataFrame(
+                    sample_rows,
+                    columns=fam.param_columns + ["label", "sample_idx", "time_ms"],
+                ).to_csv(sout, index=False)
+                written.append(sout)
+    return written
+
+
+def _times_ms(measurement):
+    if measurement is None:
+        return []
+    return [float(t) * 1e3 for t in getattr(measurement, "times", [])]
+
+
+def _result_rows(store):
+    """Flatten *store* into (name, stats_ms, throughput, unit) rows for the summary."""
+    import numpy as np
+
+    rows = []
+    for fam in store.values():
+        for _case_params, records, node_name in fam.case_metrics:
+            base = node_name
+            if base.endswith("]") and "[" in base:
+                base = base[base.index("[") + 1 : -1]  # keep the parametrize id
+            visible = [m for m in records if not m.get("samples_only")]
+            for m in visible:
+                name = base if len(visible) == 1 else f"{base} {m['label']}"
+                times = _times_ms(m.get("measurement"))
+                if times:
+                    a = np.asarray(times)
+                    stats = {
+                        "min": float(a.min()), "median": float(np.median(a)),
+                        "mean": float(a.mean()), "max": float(a.max()), "std": float(a.std()),
+                    }
+                else:  # derived metric: only the mean is meaningful
+                    stats = {"min": None, "median": None, "mean": m["ms"], "max": None, "std": None}
+                rows.append((name, stats, m["throughput"], m["unit"]))
+    return rows
+
+
+def format_results_table(store):
+    """Render a pytest-benchmark-style results table from *store* (times in ms)."""
+    rows = _result_rows(store)
+    if not rows:
+        return ""
+
+    def cell(v):
+        return "-" if v is None else f"{v:.4f}"
+
+    headers = ["Name (time in ms)", "Min", "Median", "Mean", "Max", "StdDev", "Throughput"]
+    body = []
+    for name, s, thr, unit in sorted(rows, key=lambda r: r[0]):
+        body.append([
+            name, cell(s["min"]), cell(s["median"]), cell(s["mean"]),
+            cell(s["max"]), cell(s["std"]), f"{thr:.2f} {unit}",
+        ])
+    widths = [max(len(headers[i]), *(len(r[i]) for r in body)) for i in range(len(headers))]
+
+    def line(cells):
+        return "  ".join(
+            c.ljust(widths[i]) if i == 0 else c.rjust(widths[i]) for i, c in enumerate(cells)
+        )
+
+    sep = "-" * (sum(widths) + 2 * (len(widths) - 1))
+    title = f" benchmark: {len(body)} tests ".center(len(sep), "-")
+    return "\n".join([title, line(headers), sep, *(line(r) for r in body), sep])

From 0a548ee1555909822c041510bb75555b5d1119d1 Mon Sep 17 00:00:00 2001
From: Matthias Diener <matthias.diener@gmail.com>
Date: Tue, 1 Sep 2026 15:36:56 -0500
Subject: [PATCH 2/5] improve formatting

---
 benchmarks/microbenchmarks/utils.py | 31 ++++++++++++++++++-----------
 1 file changed, 19 insertions(+), 12 deletions(-)

diff --git a/benchmarks/microbenchmarks/utils.py b/benchmarks/microbenchmarks/utils.py
index e7073c948..3e059c356 100644
--- a/benchmarks/microbenchmarks/utils.py
+++ b/benchmarks/microbenchmarks/utils.py
@@ -886,16 +886,16 @@ def _result_rows(store):
                     a = np.asarray(times)
                     stats = {
                         "min": float(a.min()), "median": float(np.median(a)),
-                        "mean": float(a.mean()), "max": float(a.max()), "std": float(a.std()),
+                        "max": float(a.max()), "std": float(a.std()),
                     }
-                else:  # derived metric: only the mean is meaningful
-                    stats = {"min": None, "median": None, "mean": m["ms"], "max": None, "std": None}
+                else:  # no per-sample distribution: show the single value as the median
+                    stats = {"min": None, "median": m["ms"], "max": None, "std": None}
                 rows.append((name, stats, m["throughput"], m["unit"]))
     return rows
 
 
 def format_results_table(store):
-    """Render a pytest-benchmark-style results table from *store* (times in ms)."""
+    """Render the results as a Markdown table (times in ms) with a caption line."""
     rows = _result_rows(store)
     if not rows:
         return ""
@@ -903,20 +903,27 @@ def format_results_table(store):
     def cell(v):
         return "-" if v is None else f"{v:.4f}"
 
-    headers = ["Name (time in ms)", "Min", "Median", "Mean", "Max", "StdDev", "Throughput"]
+    headers = ["Name", "Min (ms)", "Median (ms)", "Max (ms)", "StdDev (ms)", "Throughput"]
     body = []
     for name, s, thr, unit in sorted(rows, key=lambda r: r[0]):
         body.append([
-            name, cell(s["min"]), cell(s["median"]), cell(s["mean"]),
+            name, cell(s["min"]), cell(s["median"]),
             cell(s["max"]), cell(s["std"]), f"{thr:.2f} {unit}",
         ])
     widths = [max(len(headers[i]), *(len(r[i]) for r in body)) for i in range(len(headers))]
 
-    def line(cells):
-        return "  ".join(
+    def row(cells):
+        padded = [
             c.ljust(widths[i]) if i == 0 else c.rjust(widths[i]) for i, c in enumerate(cells)
-        )
+        ]
+        return "| " + " | ".join(padded) + " |"
 
-    sep = "-" * (sum(widths) + 2 * (len(widths) - 1))
-    title = f" benchmark: {len(body)} tests ".center(len(sep), "-")
-    return "\n".join([title, line(headers), sep, *(line(r) for r in body), sep])
+    # Markdown delimiter row: left-align Name, right-align the numeric columns.
+    align = [
+        ":" + "-" * (widths[i] - 1) if i == 0 else "-" * (widths[i] - 1) + ":"
+        for i in range(len(headers))
+    ]
+    caption = f"benchmark: {len(body)} tests"
+    return "\n".join(
+        [caption, "", row(headers), "| " + " | ".join(align) + " |", *(row(r) for r in body)]
+    )

From e2109faa528f2dbef025f86e7dd3992ede72f621 Mon Sep 17 00:00:00 2001
From: Matthias Diener <matthias.diener@gmail.com>
Date: Tue, 1 Sep 2026 15:48:11 -0500
Subject: [PATCH 3/5] autodetect benchmark files

---
 benchmarks/microbenchmarks/conftest.py | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/benchmarks/microbenchmarks/conftest.py b/benchmarks/microbenchmarks/conftest.py
index 802f9d0e4..7bd3ba536 100644
--- a/benchmarks/microbenchmarks/conftest.py
+++ b/benchmarks/microbenchmarks/conftest.py
@@ -59,6 +59,18 @@ def pytest_configure(config):
     config._microbench_store = {}
 
 
+def pytest_collect_file(parent, file_path):
+    # Collect benchmark_*.py like test files so `pytest .` finds them without a
+    # rename; skip init paths so an explicitly-passed file isn't double-collected.
+    if (
+        file_path.suffix == ".py"
+        and file_path.name.startswith("benchmark_")
+        and not parent.session.isinitpath(file_path)
+    ):
+        return pytest.Module.from_parent(parent, path=file_path)
+    return None
+
+
 class _MicroBench:
     """Handed to each test as the ``microbench`` fixture."""
 

From eb856279b9884f7d22fd56eaa8622d7f4a62e45e Mon Sep 17 00:00:00 2001
From: Matthias Diener <matthias.diener@gmail.com>
Date: Tue, 1 Sep 2026 16:16:20 -0500
Subject: [PATCH 4/5] fix file naming

---
 benchmarks/microbenchmarks/utils.py | 16 ++++++++++++++--
 1 file changed, 14 insertions(+), 2 deletions(-)

diff --git a/benchmarks/microbenchmarks/utils.py b/benchmarks/microbenchmarks/utils.py
index 3e059c356..3a50c3d95 100644
--- a/benchmarks/microbenchmarks/utils.py
+++ b/benchmarks/microbenchmarks/utils.py
@@ -822,12 +822,24 @@ def write_bench_outputs(store, *, csv=None, csv_samples=None, kernel_profile=Fal
     import pandas as pd
     from pathlib import Path
 
+    # When an explicit filename is given but several families run in one session
+    # (e.g. `pytest .`), insert the family name so they don't overwrite each other.
+    multi = sum(1 for fam in store.values() if fam.rows) > 1
+
+    def _dest(explicit, family, default_name):
+        if not isinstance(explicit, str):
+            return default_name
+        if not multi:
+            return explicit
+        p = Path(explicit)
+        return str(p.with_name(f"{p.stem}-{family}{p.suffix}"))
+
     written = []
     for family, fam in store.items():
         if not fam.rows:
             continue
         if csv is not None:
-            out = csv if isinstance(csv, str) else f"{family}.csv"
+            out = _dest(csv, family, f"{family}.csv")
             pd.DataFrame(fam.rows, columns=fam.param_columns + fam.metric_columns).to_csv(
                 out, index=False
             )
@@ -840,7 +852,7 @@ def write_bench_outputs(store, *, csv=None, csv_samples=None, kernel_profile=Fal
                 pd.DataFrame(fam.kernel_rows, columns=cols).to_csv(kout, index=False)
                 written.append(kout)
         if csv_samples is not None:
-            sout = csv_samples if isinstance(csv_samples, str) else f"{family}_samples.csv"
+            sout = _dest(csv_samples, family, f"{family}_samples.csv")
             sample_rows = []
             for case_params, records, _node in fam.case_metrics:
                 for metric in records:

From 71b939934b7b20664d8b4d0f6fd1acf5e6a40fc1 Mon Sep 17 00:00:00 2001
From: Matthias Diener <matthias.diener@gmail.com>
Date: Tue, 1 Sep 2026 16:55:52 -0500
Subject: [PATCH 5/5] formatting fixes

---
 .../microbenchmarks/benchmark_casting.py      | 16 +++++------
 benchmarks/microbenchmarks/benchmark_gemm.py  |  4 +++
 benchmarks/microbenchmarks/utils.py           | 27 ++++++++++---------
 3 files changed, 27 insertions(+), 20 deletions(-)

diff --git a/benchmarks/microbenchmarks/benchmark_casting.py b/benchmarks/microbenchmarks/benchmark_casting.py
index b2a1bb2ac..555a5704e 100755
--- a/benchmarks/microbenchmarks/benchmark_casting.py
+++ b/benchmarks/microbenchmarks/benchmark_casting.py
@@ -70,24 +70,24 @@ def build():
 #   MXFP4 : 0.5 data + E8M0 1 byte / 32-elem block  -> 0.5 + 1/32
 # MXFP4 has no packed-FP4 dequantize kernel yet, so it runs the quantize direction only.
 _CAST_FORMATS = (
-    ("FP8-E4M3", _fp8_quantizer(TE_FP8_E4M3), 1.0, check_fp8_support, True),
-    ("FP8-E5M2", _fp8_quantizer(TE_FP8_E5M2), 1.0, check_fp8_support, True),
+    ("fp8-e4m3", _fp8_quantizer(TE_FP8_E4M3), 1.0, check_fp8_support, True),
+    ("fp8-e5m2", _fp8_quantizer(TE_FP8_E5M2), 1.0, check_fp8_support, True),
     (
-        "MXFP8-E4M3",
+        "mxfp8-e4m3",
         lambda: MXFP8Quantizer(TE_FP8_E4M3, rowwise=True, columnwise=False),
         1.0 + 1.0 / 32,
         check_mxfp8_support,
         True,
     ),
     (
-        "MXFP8-E5M2",
+        "mxfp8-e5m2",
         lambda: MXFP8Quantizer(TE_FP8_E5M2, rowwise=True, columnwise=False),
         1.0 + 1.0 / 32,
         check_mxfp8_support,
         True,
     ),
     (
-        "NVFP4",
+        "nvfp4",
         lambda: NVFP4Quantizer(
             fp4_dtype=TE_FP4_E2M1, rowwise=True, columnwise=False, with_rht=False
         ),
@@ -96,7 +96,7 @@ def build():
         True,
     ),
     (
-        "MXFP4",
+        "mxfp4",
         lambda: MXFP4Quantizer(fp4_dtype=TE_FP4_E2M1, rowwise=True, columnwise=False),
         0.5 + 1.0 / 32,
         check_mxfp4_support,
@@ -134,8 +134,8 @@ def _triton_applies(fmt, direction):
     # Cast-transpose Triton covers FP8/MXFP8/MXFP4 quantize (not NVFP4); the
     # dequantize Triton path exists only for MXFP8 (mxfp8_tensor_storage).
     if direction == "quantize":
-        return fmt != "NVFP4"
-    return fmt.startswith("MXFP8")
+        return fmt != "nvfp4"
+    return fmt.startswith("mxfp8")
 
 
 def _backends_for(fmt, direction):
diff --git a/benchmarks/microbenchmarks/benchmark_gemm.py b/benchmarks/microbenchmarks/benchmark_gemm.py
index 329354841..ab5884fb1 100755
--- a/benchmarks/microbenchmarks/benchmark_gemm.py
+++ b/benchmarks/microbenchmarks/benchmark_gemm.py
@@ -123,6 +123,10 @@ def pytest_generate_tests(metafunc):
 
 @pytest.mark.benchmark
 def test_gemm(microbench, case, monkeypatch):
+    if case["Precision"] == "mxfp4" and any(
+        dim % 32 for dim in (case["M"], case["N"], case["K"])
+    ):
+        pytest.skip("MXFP4 GEMM needs M/N/K divisible by 32")
     apply_backend_env(monkeypatch, GEMM_BACKENDS[case["Backend"]])
     microbench.run(
         case,
diff --git a/benchmarks/microbenchmarks/utils.py b/benchmarks/microbenchmarks/utils.py
index 3a50c3d95..341f250be 100644
--- a/benchmarks/microbenchmarks/utils.py
+++ b/benchmarks/microbenchmarks/utils.py
@@ -48,9 +48,9 @@
 
 # Unique (model_name, hidden_size) pairs for element-wise benchmarks
 MODEL_HIDDEN_SIZES = [
-    ("Llama3-8B",   4096),
-    ("Llama3-70B",  8192),
-    ("Llama3-405B", 16384),
+    ("Llama3.1-8B",   4096),
+    ("Llama3.1-70B",  8192),
+    ("Llama3.1-405B", 16384),
     ("Qwen2.5-7B",  3584),
     ("Qwen2.5-72B", 8192),
 ]
@@ -881,11 +881,12 @@ def _times_ms(measurement):
 
 
 def _result_rows(store):
-    """Flatten *store* into (name, stats_ms, throughput, unit) rows for the summary."""
+    """Flatten *store* into (suite, name, stats_ms, throughput, unit) rows for the summary."""
     import numpy as np
 
     rows = []
-    for fam in store.values():
+    for family, fam in store.items():
+        suite = family[len("benchmark_"):] if family.startswith("benchmark_") else family
         for _case_params, records, node_name in fam.case_metrics:
             base = node_name
             if base.endswith("]") and "[" in base:
@@ -902,7 +903,7 @@ def _result_rows(store):
                     }
                 else:  # no per-sample distribution: show the single value as the median
                     stats = {"min": None, "median": m["ms"], "max": None, "std": None}
-                rows.append((name, stats, m["throughput"], m["unit"]))
+                rows.append((suite, name, stats, m["throughput"], m["unit"]))
     return rows
 
 
@@ -915,24 +916,26 @@ def format_results_table(store):
     def cell(v):
         return "-" if v is None else f"{v:.4f}"
 
-    headers = ["Name", "Min (ms)", "Median (ms)", "Max (ms)", "StdDev (ms)", "Throughput"]
+    headers = [
+        "Benchmark", "Config", "Min (ms)", "Median (ms)", "Max (ms)", "StdDev (ms)", "Throughput",
+    ]
     body = []
-    for name, s, thr, unit in sorted(rows, key=lambda r: r[0]):
+    for suite, name, s, thr, unit in sorted(rows, key=lambda r: (r[0], r[1])):
         body.append([
-            name, cell(s["min"]), cell(s["median"]),
+            suite, name, cell(s["min"]), cell(s["median"]),
             cell(s["max"]), cell(s["std"]), f"{thr:.2f} {unit}",
         ])
     widths = [max(len(headers[i]), *(len(r[i]) for r in body)) for i in range(len(headers))]
 
     def row(cells):
+        # Left-align the text columns (Benchmark, Config); right-align the numerics.
         padded = [
-            c.ljust(widths[i]) if i == 0 else c.rjust(widths[i]) for i, c in enumerate(cells)
+            c.ljust(widths[i]) if i <= 1 else c.rjust(widths[i]) for i, c in enumerate(cells)
         ]
         return "| " + " | ".join(padded) + " |"
 
-    # Markdown delimiter row: left-align Name, right-align the numeric columns.
     align = [
-        ":" + "-" * (widths[i] - 1) if i == 0 else "-" * (widths[i] - 1) + ":"
+        ":" + "-" * (widths[i] - 1) if i <= 1 else "-" * (widths[i] - 1) + ":"
         for i in range(len(headers))
     ]
     caption = f"benchmark: {len(body)} tests"