microbenchmarks: use pytest as execution backend

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,
 )
 
@@ -69,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
         ),
@@ -95,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,
@@ -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:]]))

benchmarks/microbenchmarks/benchmark_gemm.py

@@ -4,40 +4,87 @@
 #
 # 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,
 )
 
 BENCHMARK_LABEL = "GEMM"
 
 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 _backends_for(precision):
+    return _BACKENDS_BY_PRECISION.get(precision, ["hipblaslt"])
 
 
-def bench_gemm(Case, Precision, M, N, K, dtype):
+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,41 @@ 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):
+    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,
+        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:]]))