diff --git a/benchmarks/microbenchmarks/README.md b/benchmarks/microbenchmarks/README.md index 79d5709006..99aba90e77 100644 --- a/benchmarks/microbenchmarks/README.md +++ b/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 diff --git a/benchmarks/microbenchmarks/benchmark_casting.py b/benchmarks/microbenchmarks/benchmark_casting.py index 9ad99db8e6..2221d1736f 100755 --- a/benchmarks/microbenchmarks/benchmark_casting.py +++ b/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) diff --git a/benchmarks/microbenchmarks/benchmark_gemm.py b/benchmarks/microbenchmarks/benchmark_gemm.py index 156dc4bcf3..29445179f9 100755 --- a/benchmarks/microbenchmarks/benchmark_gemm.py +++ b/benchmarks/microbenchmarks/benchmark_gemm.py @@ -4,11 +4,20 @@ # # 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, @@ -16,21 +25,39 @@ 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"], ) diff --git a/benchmarks/microbenchmarks/benchmark_gemm_fp8.py b/benchmarks/microbenchmarks/benchmark_gemm_fp8.py deleted file mode 100755 index cdf17fb4c1..0000000000 --- a/benchmarks/microbenchmarks/benchmark_gemm_fp8.py +++ /dev/null @@ -1,83 +0,0 @@ -#!/usr/bin/env python -############################################################################### -# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. -# -# See LICENSE for license information. -############################################################################### -""" -FP8 GEMM micro-benchmark using te.Linear under fp8_autocast. - -Same model shapes as benchmark_gemm.py. -Output: benchmark_gemm_fp8.csv (written to cwd) -""" - -import torch -import transformer_engine.pytorch as te -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, - make_input, -) - -RECIPES = { - "hybrid": DelayedScaling( - fp8_format=Format.HYBRID, - amax_history_len=16, - amax_compute_algo="max", - ), -} - -FP8_RECIPE = RECIPES["hybrid"] - -BENCHMARK_LABEL = "FP8 GEMM" - - -def bench_fp8_gemm(Case, M, N, K, dtype): - device = "cuda" - - linear = te.Linear(K, N, bias=False).to(device=device, dtype=dtype) - 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(next_x()) - - def fwd_bwd_func(): - xb = next_x() - with te.fp8_autocast(enabled=True, fp8_recipe=FP8_RECIPE): - out = linear(xb) - out.backward(grad_out) - xb.grad = None - linear.weight.grad = None - - 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 - - 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, - ) - - -if __name__ == "__main__": - run_benchmarks( - test_cases=generate_gemm_test_cases(), - bench_fn=bench_fp8_gemm, - param_columns=["Case", "M", "N", "K", "dtype"], - ) diff --git a/benchmarks/microbenchmarks/benchmark_grouped_gemm.py b/benchmarks/microbenchmarks/benchmark_grouped_gemm.py index af79ff697f..d95e184c4c 100755 --- a/benchmarks/microbenchmarks/benchmark_grouped_gemm.py +++ b/benchmarks/microbenchmarks/benchmark_grouped_gemm.py @@ -9,6 +9,7 @@ import transformer_engine.pytorch as te from utils import ( DTYPE_LIST, + build_recipes, time_func, compute_tflops, make_forward_backward_metric_records, @@ -18,6 +19,11 @@ 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")) + def generate_grouped_gemm_group_lens(b, m, balance: bool): if balance: return torch.full((b,), m, dtype=torch.int64) @@ -60,16 +66,18 @@ def _generate_moe_test_cases( for M in GROUPED_GEMM_M_SIZE_LIST: for name, (N, K) in shapes_dict.items(): for dtype in DTYPE_LIST: - test_cases.append( - { - "Case": name, - "B": B, - "M": M, - "N": N, - "K": K, - "dtype": dtype, - } - ) + for recipe in RECIPES: + test_cases.append( + { + "Case": name, + "B": B, + "M": M, + "N": N, + "K": K, + "dtype": dtype, + "recipe": recipe, + } + ) return test_cases @@ -99,9 +107,12 @@ def generate_grok_v2_test_cases(): ) -def bench_grouped_gemm(Case, B, M, N, K, dtype): +def bench_grouped_gemm(Case, B, M, N, K, dtype, recipe): device = "cuda" + fp8_recipe = RECIPES[recipe] + use_fp8 = fp8_recipe is not None + group_lens = generate_grouped_gemm_group_lens(B, M, balance=True) m_splits = [int(v) for v in group_lens.tolist()] m_splits_tensor = torch.tensor(m_splits, dtype=torch.int32, device=device) @@ -120,15 +131,17 @@ def bench_grouped_gemm(Case, B, M, N, K, dtype): next_x = make_input((sum_M, K), dtype, device=device, requires_grad=True) def fwd_func_te(): - return grouped_linear(next_x(), m_splits, m_splits_tensor=m_splits_tensor) + 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() grad_out = torch.randn_like(out_te) def fwd_bwd_func_te(): xb = next_x() - out = grouped_linear(xb, m_splits, m_splits_tensor=m_splits_tensor) - out.backward(grad_out) + with te.autocast(enabled=use_fp8, recipe=fp8_recipe): + out = grouped_linear(xb, m_splits, m_splits_tensor=m_splits_tensor) + out.backward(grad_out) xb.grad = None for param in grouped_linear.parameters(): param.grad = None @@ -169,5 +182,5 @@ def fwd_bwd_func_te(): run_benchmarks( test_cases=test_cases, bench_fn=bench_grouped_gemm, - param_columns=["Case", "B", "M", "N", "K", "dtype"], + param_columns=["Case", "B", "M", "N", "K", "dtype", "recipe"], ) diff --git a/benchmarks/microbenchmarks/benchmark_normalization.py b/benchmarks/microbenchmarks/benchmark_normalization.py index 7ff86cfbfb..ff9f7b28b4 100755 --- a/benchmarks/microbenchmarks/benchmark_normalization.py +++ b/benchmarks/microbenchmarks/benchmark_normalization.py @@ -5,94 +5,113 @@ # See LICENSE for license information. ############################################################################### """ -Normalization micro-benchmark using te.LayerNorm and te.RMSNorm. - -Both LayerNorm and RMSNorm share the same kernel infrastructure. -The M dimension (batch * seq_len) is swept across typical training sizes. - +Normalization micro-benchmark using the fusible-ops LayerNorm / RMSNorm. + +Sweeps BF16 plus the quantized-output precisions (FP8, MXFP8) that TE training +recipes produce. In FP8/FP4 training the norm is fused with the following +Linear (LayerNormLinear / LayerNormMLP) and writes its output already +quantized. That is reproduced here with ``ops.Sequential(Norm, Quantize)``: the +op fuser threads the Quantize op's input quantizer into the norm, so the norm +writes the quantized output directly under ``autocast``. The Quantize op is an +identity outside ``autocast``, so the bf16 baseline uses the same harness. + +Forward only: the quantize epilogue is a forward-pass phenomenon; the norm +backward reads/writes high precision and is precision-independent. NVFP4/MXFP4 +norm outputs are not swept (no validated norm->FP4 cast path). Precisions +unsupported on the current device are skipped. + +These are memory-bound; we report GB/s (input read + output write). Output: benchmark_normalization.csv (written to cwd) """ import torch import transformer_engine.pytorch as te +from transformer_engine.pytorch import ops from utils import ( - DTYPE_LIST, MODEL_HIDDEN_SIZES, M_SIZE_LIST, - time_func, compute_gbps, make_forward_backward_metric_records, run_benchmarks, + MODEL_HIDDEN_SIZES, M_SIZE_LIST, + build_recipes, + time_func, compute_gbps, make_metric_record, run_benchmarks, make_input, ) NORM_TYPES = [ - ("RMSNorm", te.RMSNorm), - ("LayerNorm", te.LayerNorm), + ("RMSNorm", ops.RMSNorm), + ("LayerNorm", ops.LayerNorm), ] -BENCHMARK_LABEL = "Normalization" +BENCHMARK_LABEL = "Normalization Forward" + +# Quantized-output precisions validated for the norm cast; tests/pytorch/ +# triton_kernels/test_norms.py exercises fp8 and mxfp8 norm quantizers. bf16 is +# the plain baseline (Quantize is an identity outside autocast). +RECIPES = build_recipes(names=("bf16", "fp8", "mxfp8")) + +# Forward output bytes/elem by precision. Under autocast the norm quantizes +# through the recipe-created quantizers, which default to columnwise usage on, so +# it writes BOTH rowwise and columnwise data (plus both MXFP8 scale buffers) -- +# the norm->Linear training layout (rowwise feeds fprop, columnwise feeds wgrad). +# Input is always bf16 (2 bytes/elem). +# bf16 : 2.0 (single bf16 output; Quantize is identity) +# fp8 : 2.0 (rowwise + columnwise E4M3/E5M2 data; scale ~ 0) +# mxfp8 : 2.0 + 2/32 = 2 + 1/16 (rowwise + columnwise data + both E8M0 scales) +_FWD_WRITE_BYTES = { + "bf16": 2.0, + "fp8": 2.0, + "mxfp8": 2.0 + 1.0 / 16, +} def _generate_test_cases(): test_cases = [] for model_name, hidden in MODEL_HIDDEN_SIZES: - for norm_name, norm_cls in NORM_TYPES: - for M in M_SIZE_LIST: - for dtype in DTYPE_LIST: + for norm_name, norm_op_cls in NORM_TYPES: + 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_name": norm_name, - "norm_cls": norm_cls, - "dtype": dtype, + "norm_op_cls": norm_op_cls, + "dtype": torch.bfloat16, }) return test_cases -def bench_norm(Case, M, hidden_size, norm_name, norm_cls, dtype): +def bench_norm(Case, Precision, M, hidden_size, norm_op_cls, dtype): device = "cuda" - norm = norm_cls(hidden_size).to(device=device, dtype=dtype) - next_x = make_input((M, hidden_size), dtype, device=device, requires_grad=True) + recipe = RECIPES[Precision] + use_fp8 = recipe is not None - fwd_func = lambda: norm(next_x()) - out = fwd_func() - grad_out = torch.randn_like(out) + # Norm followed by Quantize so the norm writes its output directly in the + # target precision under autocast (identity when use_fp8 is False). + model = ops.Sequential( + norm_op_cls(hidden_size, device=device, dtype=dtype), + ops.Quantize(), + ) + next_x = make_input((M, hidden_size), dtype, device=device, requires_grad=False) - def fwd_bwd_func(): - xb = next_x() - out = norm(xb) - out.backward(grad_out) - xb.grad = None - for p in norm.parameters(): - p.grad = None + def fwd_func(): + with te.autocast(enabled=use_fp8, recipe=recipe): + return model(next_x()) - fwd_bwd_func() + fwd_func() - 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 + # BF16 read + quantized write. + fwd_bytes = int(M * hidden_size * (2 + _FWD_WRITE_BYTES[Precision])) 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_gbps = compute_gbps(fwd_bytes, fwd_ms) - bwd_gbps = compute_gbps(bwd_bytes, bwd_ms) - - return make_forward_backward_metric_records( - BENCHMARK_LABEL, - "GB/s", - fwd_ms, - fwd_gbps, - bwd_ms, - bwd_gbps, - backward_derived=True, - fwd_measurement=fwd_measurement, - fwd_bwd_measurement=fwd_bwd_measurement, - ) + + return [make_metric_record( + BENCHMARK_LABEL, fwd_ms, "GB/s", fwd_gbps, measurement=fwd_measurement, + )] if __name__ == "__main__": run_benchmarks( test_cases=_generate_test_cases(), bench_fn=bench_norm, - param_columns=["Case", "M", "hidden_size", "dtype"], + param_columns=["Case", "Precision", "M", "hidden_size", "dtype"], ) diff --git a/benchmarks/microbenchmarks/utils.py b/benchmarks/microbenchmarks/utils.py index 92bbf66263..45101b60f7 100644 --- a/benchmarks/microbenchmarks/utils.py +++ b/benchmarks/microbenchmarks/utils.py @@ -7,6 +7,7 @@ """Shared utilities for microbenchmarks: model configs, timing, throughput, runner.""" import argparse +import importlib.util import itertools import math import torch @@ -86,6 +87,91 @@ def generate_gemm_test_cases(configs=None, m_sizes=None, dtypes=None): return test_cases +# --------------------------------------------------------------------------- +# Low-precision recipe sweep (shared by the dense and grouped GEMM benchmarks) +# --------------------------------------------------------------------------- +# Transformer Engine imports are deferred into the helpers so importing +# utils.py stays free of a GPU / built TE (keeps the non-GEMM benchmarks and +# offline tooling importable). + + +def _check_mxfp4_support_with_aiter(): + """MXFP4 gate: device support plus the aiter FP4 GEMM backend. + + The MXFP4 GEMM path calls into aiter's a4w4 kernels, so a missing aiter + package would crash at benchmark time even on supported hardware. + """ + from transformer_engine.pytorch.quantization import check_mxfp4_support + + supported, reason = check_mxfp4_support() + if not supported: + return supported, reason + if importlib.util.find_spec("aiter") is None: + return False, "aiter is not installed (required for the MXFP4 GEMM backend)." + return True, "" + + +def _precision_specs(): + """Ordered sweep of (name, recipe factory | None, support check | None). + + A ``None`` factory is the bf16 baseline (no autocast). The fp8 entry uses + HYBRID delayed scaling and is shared by the dense and grouped GEMM + benchmarks so their fp8 numbers stay comparable. + """ + from transformer_engine.common.recipe import ( + DelayedScaling, + Format, + MXFP4BlockScaling, + MXFP8BlockScaling, + NVFP4BlockScaling, + ) + from transformer_engine.pytorch.quantization import ( + check_fp8_support, + check_mxfp8_support, + check_nvfp4_support, + ) + + return ( + ("bf16", None, None), + ( + "fp8", + lambda: DelayedScaling( + fp8_format=Format.HYBRID, + amax_history_len=16, + amax_compute_algo="max", + ), + check_fp8_support, + ), + ("mxfp8", MXFP8BlockScaling, check_mxfp8_support), + ("mxfp4", MXFP4BlockScaling, _check_mxfp4_support_with_aiter), + ("nvfp4", NVFP4BlockScaling, check_nvfp4_support), + ) + + +def build_recipes(names=None): + """Build an ordered ``{name: recipe_or_None}`` sweep of supported precisions. + + ``bf16`` maps to ``None`` (no autocast). Each low-precision entry is + included only when its support check passes on the current device; + unsupported ones are dropped with a short notice. Pass *names* to restrict + and order the sweep, e.g. ``("bf16", "fp8", "mxfp8", "nvfp4")`` for grouped + GEMM, which has no MXFP4 grouped kernel. + """ + specs = _precision_specs() + if names is not None: + by_name = {spec[0]: spec for spec in specs} + specs = tuple(by_name[name] for name in names) + recipes = {} + for name, factory, support_check in specs: + if support_check is not None: + supported, reason = support_check() + if not supported: + print(f"Skipping {name} precision: {reason}") + continue + recipes[name] = factory() if factory is not None else None + return recipes + + # --------------------------------------------------------------------------- # Timing helpers # ---------------------------------------------------------------------------