add-to-benchmarks

benchmarks/float8/bench_matmul.py

@@ -16,6 +16,8 @@
     get_name_to_shapes_iter,
 )
 
+from torchao.ops import mx_fp4_bf16
+from torchao.prototype.mx_formats.mx_tensor import to_mx
 from torchao.testing.training.roofline_utils import get_specs
 
 
@@ -62,29 +64,38 @@ def run(
 ):
     device = "cuda"
     # TODO(future PR): this is ugly
-    assert recipe in ("tensorwise", "rowwise", "mxfp8_cublas"), "unsupported"
+    assert recipe in (
+        "tensorwise",
+        "rowwise",
+        "mxfp8_cublas",
+        "mxfp4_cutlass",
+        "nvfp4",
+    ), "unsupported"
+    use_fp4 = recipe in ("mxfp4_cutlass", "nvfp4")
 
     specs = get_specs()
     bf16_peak_tops = specs["bf16_peak_tops"]
     fp8_peak_tops = specs["fp8_peak_tops"]
+    fp4_peak_tops = specs["fp4_peak_tops"]
     print(f"gpu_name: {torch.cuda.get_device_name(0)}")
-    print(f"peak tops: bf16 {bf16_peak_tops:.2e}, fp8 {fp8_peak_tops:.2e}")
-
+    print(
+        f"peak tops: bf16 {bf16_peak_tops:.2e}, fp8 {fp8_peak_tops:.2e}, fp4 {fp4_peak_tops:.2e}"
+    )
     headers = (
         "fast_accum",
         "name",
         "M",
         "K",
         "N",
-        "ref_time_s",
-        "fp8_time_s",
+        "time_s",
+        "speedup",
         "fp8_speedup",
     )
     results = []
 
     dtype = torch.bfloat16
     name_to_shapes = get_name_to_shapes_iter(shape_gen_name, M, K, N)
-    fast_accum_vals = [True, False]
+    fast_accum_vals = [False] if use_fp4 else [True, False]
 
     for idx, (fast_accum, (name, (M, K, N))) in enumerate(
         itertools.product(fast_accum_vals, name_to_shapes)
@@ -107,38 +118,82 @@ def run(
 
         del A
 
-        # raw float8 matmul (upper bound for what we can achive in eager mode)
-        # TODO(future): add e5m2
-        d1, d2, d3 = torch.float8_e4m3fn, torch.float8_e4m3fn, dtype
-        A = torch.zeros(M, K, device=device, dtype=d1)
-        B = torch.zeros(K, N, device=device, dtype=d2).t().contiguous().t()
+        A_hp = torch.randn(M, K, device=device)
+        B_hp_t = torch.randn(N, K, device=device)
+
+        if recipe == "mxfp4_cutlass":
+            _, A = to_mx(A_hp, torch.float4_e2m1fn_x2, 32)
+            _, Bt = to_mx(B_hp_t, torch.float4_e2m1fn_x2, 32)
+            B = Bt.contiguous().T
+            peak_tops = fp4_peak_tops
+        elif recipe == "nvfp4":
+            from torchao.prototype.mx_formats.nvfp4_tensor import nvfp4_quantize
+
+            A_scales, A_data = nvfp4_quantize(A_hp, block_size=16)
+            B_scales, B_data = nvfp4_quantize(B_hp_t, block_size=16)
+            A = A_data.view(torch.float4_e2m1fn_x2)
+            B = B_data.view(torch.float4_e2m1fn_x2).T
+            peak_tops = fp4_peak_tops
+        else:
+            # raw float8 matmul (upper bound for what we can achive in eager mode)
+            # TODO(future): add e5m2
+            d1, d2, d3 = torch.float8_e4m3fn, torch.float8_e4m3fn, dtype
+            A = A_hp.to(d1)
+            B = B_hp_t.to(d2).contiguous().T
+            peak_tops = fp8_peak_tops
+
         if recipe == "tensorwise":
             scale_a = torch.tensor([1.0], device=device)
             scale_b = torch.tensor([1.0], device=device)
         elif recipe == "rowwise":
             scale_a = torch.ones(M, 1, device=device)
             scale_b = torch.ones(1, N, device=device)
-        elif recipe == "mxfp8_cublas":
+        elif recipe in ("mxfp8_cublas", "mxfp4_cutlass"):
             scale_a = torch.ones(M, K // 32, device=device, dtype=torch.float8_e8m0fnu)
             scale_b = torch.ones(N, K // 32, device=device, dtype=torch.float8_e8m0fnu)
+        elif recipe == "nvfp4":
+            # Use the blockwise scales from nvfp4_quantize
+            scale_a = A_scales.view(torch.float8_e4m3fn)
+            scale_b = B_scales.view(torch.float8_e4m3fn)
         else:
             assert False, f"unknown recipe {recipe}"
 
-        def do_matmul(A, B):
+        def do_matmul_fp8(A, B):
             nonlocal scale_a
             nonlocal scale_b
             return torch._scaled_mm(
                 A, B, scale_a, scale_b, out_dtype=d3, use_fast_accum=fast_accum
             )
 
-        fp8_time_sec, fp8_tops_sec, fp8_pct_top_peak = do_benchmarks(
-            tops, fp8_peak_tops, use_gpu_kernel_time, do_matmul, A, B
+        def do_matmul_mxfp4(A, B):
+            nonlocal scale_a
+            nonlocal scale_b
+            return mx_fp4_bf16(A, B, scale_a, scale_b)
+
+        def do_matmul_nvfp4(A, B):
+            nonlocal scale_a
+            nonlocal scale_b
+            return torch._scaled_mm(A, B, scale_a, scale_b, out_dtype=dtype)
+
+        if recipe == "mxfp4_cutlass":
+            do_matmul = do_matmul_mxfp4
+        elif recipe == "nvfp4":
+            do_matmul = do_matmul_nvfp4
+        else:
+            do_matmul = do_matmul_fp8
+
+        time_sec, tops_sec, pct_top_peak = do_benchmarks(
+            tops, peak_tops, use_gpu_kernel_time, do_matmul, A, B
         )
         print(
-            f"fp8 time_sec {fp8_time_sec:.2E}, tops/sec {fp8_tops_sec:.2E}, pct_peak {fp8_pct_top_peak:.3f}"
+            f"time_sec {time_sec:.2E}, tops/sec {tops_sec:.2E}, pct_peak {pct_top_peak:.3f}"
         )
 
-        del A, B, scale_a, scale_b
+        del A, B
+        if scale_a is not None:
+            del scale_a
+        if scale_b is not None:
+            del scale_b
 
         results.append(
             [
@@ -148,8 +203,8 @@ def do_matmul(A, B):
                 K,
                 N,
                 ref_time_sec,
-                fp8_time_sec,
-                ref_time_sec / fp8_time_sec,
+                time_sec,
+                ref_time_sec / time_sec,
             ]
         )
 

benchmarks/float8/utils.py

@@ -352,9 +352,6 @@ def get_gpu_kernel_gemm_time_s(f, *args, **kwargs):
     )
     # there is only 1 key, aten::mm or aten::_scaled_mm, with unit nanoseconds
     assert len(data) == 1
-    if "aten::mm" in data:
-        return data["aten::mm"] / 1e6 / n_iter
-    elif "aten::_scaled_mm" in data:
-        return data["aten::_scaled_mm"] / 1e6 / n_iter
-    else:
-        raise AssertionError("unexpected format of data")
+    key, value = next(iter(data.items()))
+    assert key in ("aten::mm", "aten::_scaled_mm", "torchao::mx_fp4_bf16")
+    return value / 1e6 / n_iter

torchao/testing/training/roofline_utils.py

@@ -54,6 +54,13 @@
         # TODO(future): run measurement on hardware
         "pct_achievable_mem_bw": 0.92,
     },
+    "NVIDIA GeForce RTX 5090": {
+        # https://images.nvidia.com/aem-dam/Solutions/geforce/blackwell/nvidia-rtx-blackwell-gpu-architecture.pdf
+        "bf16_peak_tops": 209.5e12,
+        "fp8_peak_tops": 419e12,
+        "fp4_peak_tops": 1676e12,
+        "peak_mem_bw_bytes_sec": 1.792e15,
+    },
     # TODO(future): more GPU names
 }