microbenchmarks: extend all to additional low-precision dtypes, restructure GEMM

benchmarks/microbenchmarks/README.md

@@ -5,11 +5,11 @@ Transformer Engine kernels and helper scripts for comparing benchmark CSVs.
 
 ## Benchmarks
 
-- `benchmark_gemm.py`: dense BF16 GEMM benchmark
-- `benchmark_gemm_fp8.py`: dense FP8 GEMM benchmark using `fp8_autocast`
+- `benchmark_gemm.py`: dense GEMM benchmark sweeping BF16 plus the supported
+  low-precision recipes (FP8, MXFP8, MXFP4, NVFP4) via `autocast`
 - `benchmark_grouped_gemm.py`: grouped GEMM benchmark for MoE-style shapes
-- `benchmark_casting.py`: BF16 `<->` FP8 casting benchmark
-- `benchmark_normalization.py`: LayerNorm and RMSNorm benchmark
+- `benchmark_casting.py`: quantize / dequantize benchmark across FP8, MXFP8, NVFP4, and MXFP4
+- `benchmark_normalization.py`: LayerNorm / RMSNorm forward benchmark across BF16 and quantized (FP8, MXFP8) output
 
 Run a benchmark directly from this directory. Pass `--csv` to write results.
 When no filename is provided, `run_benchmarks` derives the CSV name from the

benchmarks/microbenchmarks/benchmark_casting.py

@@ -5,10 +5,21 @@
 # See LICENSE for license information.
 ###############################################################################
 """
-FP8 casting micro-benchmark.
+Low-precision casting micro-benchmark.
 
-Benchmarks quantization (BF16 -> FP8) and dequantization (FP8 -> BF16) for
-both E4M3 (activations/weights) and E5M2 (gradients) formats.
+Benchmarks quantization (BF16 -> low precision) and dequantization
+(low precision -> BF16) for the formats used by TE training recipes:
+
+  * FP8 per-tensor scaling: E4M3 (activations/weights) and E5M2 (gradients)
+  * MXFP8 block scaling:    E4M3 and E5M2 (32-elem blocks, E8M0 scales)
+  * NVFP4 block scaling:    E2M1 (16-elem blocks, E4M3 scales), no RHT
+  * MXFP4 block scaling:    E2M1 (32-elem blocks, E8M0 scales)
+
+Rowwise-only casts are measured (one output tensor), so the numbers reflect the
+core cast kernel; training additionally computes the columnwise/transpose copy.
+The NVFP4 random-Hadamard-transform fused cast is covered separately by
+benchmarks/benchmark_rht_cast.py. Formats unsupported on the current device
+are skipped.
 
 These casts are memory-bound; we report GB/s (input + output bytes).
 Output: benchmark_casting.csv (written to cwd)
@@ -17,7 +28,14 @@
 import torch
 import transformer_engine
 import transformer_engine_torch as tex
-from transformer_engine.pytorch import Float8Quantizer
+from transformer_engine.pytorch import Float8Quantizer, MXFP8Quantizer, NVFP4Quantizer
+from transformer_engine.pytorch.tensor.mxfp4_tensor import MXFP4Quantizer
+from transformer_engine.pytorch.quantization import (
+    check_fp8_support,
+    check_mxfp4_support,
+    check_mxfp8_support,
+    check_nvfp4_support,
+)
 from utils import (
     MODEL_HIDDEN_SIZES, M_SIZE_LIST,
     time_func, compute_gbps, make_metric_record, run_benchmarks,
@@ -26,57 +44,125 @@
 
 TE_FP8_E4M3 = tex.DType.kFloat8E4M3
 TE_FP8_E5M2 = tex.DType.kFloat8E5M2
+TE_FP4_E2M1 = tex.DType.kFloat4E2M1
 
 CAST_LABEL = "Cast"
 
-CAST_CONFIGS = [
-    # (name, direction, fp8_dtype)
-    ("BF16-to-FP8-E4M3", "quantize",   TE_FP8_E4M3),
-    ("FP8-E4M3-to-BF16", "dequantize", TE_FP8_E4M3),
-    ("BF16-to-FP8-E5M2", "quantize",   TE_FP8_E5M2),
-    ("FP8-E5M2-to-BF16", "dequantize", TE_FP8_E5M2),
-]
+
+def _fp8_quantizer(fp8_dtype):
+    """Per-tensor delayed-scaling FP8 quantizer factory (needs scale/amax buffers)."""
+
+    def build():
+        scale = torch.ones(1, dtype=torch.float32, device="cuda")
+        amax = torch.zeros(1, dtype=torch.float32, device="cuda")
+        return Float8Quantizer(scale, amax, fp8_dtype)
+
+    return build
+
+
+# Per-format cast specs:
+#   (name, quantizer factory, quantized bytes/elem, support check, dequant supported).
+# "quantized bytes/elem" = packed data + block-scale bytes per element:
+#   FP8   : 1.0 data, per-tensor scale ~ 0          -> 1.0
+#   MXFP8 : 1.0 data + E8M0 1 byte / 32-elem block  -> 1 + 1/32
+#   NVFP4 : 0.5 data + E4M3 1 byte / 16-elem block  -> 0.5 + 1/16
+#   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),
+    (
+        "MXFP8-E4M3",
+        lambda: MXFP8Quantizer(TE_FP8_E4M3, rowwise=True, columnwise=False),
+        1.0 + 1.0 / 32,
+        check_mxfp8_support,
+        True,
+    ),
+    (
+        "MXFP8-E5M2",
+        lambda: MXFP8Quantizer(TE_FP8_E5M2, rowwise=True, columnwise=False),
+        1.0 + 1.0 / 32,
+        check_mxfp8_support,
+        True,
+    ),
+    (
+        "NVFP4",
+        lambda: NVFP4Quantizer(
+            fp4_dtype=TE_FP4_E2M1, rowwise=True, columnwise=False, with_rht=False
+        ),
+        0.5 + 1.0 / 16,
+        check_nvfp4_support,
+        True,
+    ),
+    (
+        "MXFP4",
+        lambda: MXFP4Quantizer(fp4_dtype=TE_FP4_E2M1, rowwise=True, columnwise=False),
+        0.5 + 1.0 / 32,
+        check_mxfp4_support,
+        False,
+    ),
+)
+
+DIRECTIONS = ("quantize", "dequantize")
+
+
+def _active_formats():
+    """Filter cast formats to those supported on the current device."""
+    formats = []
+    for name, make_quantizer, q_bytes_per_elem, support_check, dequant_supported in _CAST_FORMATS:
+        supported, reason = support_check()
+        if not supported:
+            print(f"Skipping {name} casts: {reason}")
+            continue
+        formats.append((name, make_quantizer, q_bytes_per_elem, dequant_supported))
+    return formats
 
 
 def _generate_test_cases():
     test_cases = []
+    active = _active_formats()
     for model_name, hidden in MODEL_HIDDEN_SIZES:
-        for cast_name, direction, fp8_dtype in CAST_CONFIGS:
-            for M in M_SIZE_LIST:
-                test_cases.append({
-                    "Case": f"{model_name}/{cast_name}",
-                    "M": M,
-                    "hidden_size": hidden,
-                    "direction": direction,
-                    "fp8_dtype": fp8_dtype,
-                    "dtype_str": cast_name,
-                })
+        for fmt_name, make_quantizer, q_bytes_per_elem, dequant_supported in active:
+            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, fp8_dtype, dtype_str):
+def bench_cast(Case, M, hidden_size, direction, make_quantizer, q_bytes_per_elem, dtype_str):
     device = "cuda"
 
     numel = M * hidden_size
-    scale = torch.ones(1, dtype=torch.float32, device=device)
-    amax = torch.zeros(1, dtype=torch.float32, device=device)
-    quantizer = Float8Quantizer(scale, amax, fp8_dtype)
+    quantizer = make_quantizer()
 
     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)
-        total_bytes = numel * (2 + 1)  # BF16 read + FP8 write
+        total_bytes = int(numel * (2 + q_bytes_per_elem))  # BF16 read + quantized write
     else:
-        # Rotate a ring of FP8 tensors (bytes can't be inferred, so hint numel).
-        next_fp8 = rotating(
+        # Rotate a ring of quantized tensors (packed bytes can't be inferred, so hint).
+        next_q = rotating(
             lambda: quantizer(
                 torch.randn(M, hidden_size, dtype=torch.bfloat16, device=device)
             ),
-            bytes_per_buffer=numel,  # FP8 ~ 1 byte/element
+            bytes_per_buffer=int(numel * q_bytes_per_elem),
         )
-        cast_func = lambda: next_fp8().dequantize()
-        total_bytes = numel * (1 + 2)  # FP8 read + BF16 write
+        cast_func = lambda: next_q().dequantize()
+        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)

benchmarks/microbenchmarks/benchmark_gemm.py

@@ -4,33 +4,60 @@
 #
 # See LICENSE for license information.
 ###############################################################################
+"""Dense GEMM micro-benchmark using te.Linear across precisions.
 
+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.
+
+Output: benchmark_gemm.csv (written to cwd)
+"""
 
 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,
     make_input,
 )
 
 BENCHMARK_LABEL = "GEMM"
 
+RECIPES = build_recipes()
+
+
+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
 
-def bench_gemm(Case, M, N, K, dtype):
+
+def bench_gemm(Case, Precision, M, N, K, dtype):
     device = "cuda"
 
+    recipe = RECIPES[Precision]
+    use_fp8 = recipe is not None
+
     linear = te.Linear(K, N, bias=False).to(device=device, dtype=dtype)
     next_x = make_input((M, K), dtype, device=device, requires_grad=True)
 
-    fwd_func = lambda: linear(next_x())
+    def fwd_func():
+        with te.autocast(enabled=use_fp8, recipe=recipe):
+            return linear(next_x())
+
     out = fwd_func()
     grad_out = torch.randn_like(out)
 
     def fwd_bwd_func():
         xb = next_x()
-        out = linear(xb)
-        out.backward(grad_out)
+        with te.autocast(enabled=use_fp8, recipe=recipe):
+            out = linear(xb)
+            out.backward(grad_out)
         xb.grad = None
         linear.weight.grad = None
 
@@ -61,7 +88,7 @@ def fwd_bwd_func():
 
 if __name__ == "__main__":
     run_benchmarks(
-        test_cases=generate_gemm_test_cases(),
+        test_cases=generate_precision_gemm_test_cases(),
         bench_fn=bench_gemm,
-        param_columns=["Case", "M", "N", "K", "dtype"],
+        param_columns=["Case", "Precision", "M", "N", "K", "dtype"],
     )