Add CuTeDSL kernel for 3D tensor quantization to MXFP8

Author: alexsamardzic

benchmarks/mx_formats/mxfp8_backend_bench.py

@@ -0,0 +1,133 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD 3-Clause license found in the
+# LICENSE file in the root directory of this source tree.
+
+import argparse
+from typing import Dict, Tuple
+
+import torch
+
+from torchao.prototype.moe_training.kernels.mxfp8.quant import mxfp8_quantize_cuda_3d
+
+
+def _parse_backends(s: str) -> list[str]:
+    out = [x.strip() for x in s.split(",") if x.strip()]
+    for b in out:
+        if b not in {"cuda", "cutedsl"}:
+            raise ValueError(f"Unsupported backend={b!r}, expected cuda/cutedsl")
+    return out
+
+
+def _dtype_from_str(s: str) -> torch.dtype:
+    if s == "bf16":
+        return torch.bfloat16
+    if s == "fp32":
+        return torch.float32
+    raise ValueError(f"Unsupported dtype={s}")
+
+
+def _tbps(num_bytes: int, ms: float) -> float:
+    return num_bytes / (ms / 1e3) / 1e12
+
+
+def _benchmark(fn, warmup: int, iters: int) -> float:
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+    start.record()
+    for _ in range(iters):
+        fn()
+    end.record()
+    torch.cuda.synchronize()
+    return start.elapsed_time(end) / iters
+
+
+def _max_abs_diff(a: torch.Tensor, b: torch.Tensor) -> float:
+    return (a.float() - b.float()).abs().max().item()
+
+
+def _run_3d(args) -> None:
+    dtype = _dtype_from_str(args.dtype)
+    backends = _parse_backends(args.backends)
+    E, N, K = args.E, args.N, args.K
+
+    props = torch.cuda.get_device_properties(torch.cuda.current_device())
+    cc = torch.cuda.get_device_capability()
+    print(f"GPU: {props.name}")
+    print(f"CC: {cc}")
+    print(
+        f"shape=(E,N,K)=({E},{N},{K}) dtype={dtype} scaling_mode={args.scaling_mode} "
+        f"stage_count={args.stage_count}"
+    )
+
+    x = torch.randn((E, N, K), device="cuda", dtype=dtype) * 1000
+    bytes_moved = (
+        x.numel() * x.element_size()  # input
+        + x.numel()
+        * torch.tensor([], dtype=torch.float8_e4m3fn).element_size()  # output
+        + (E * (N // 32) * K)
+        * torch.tensor([], dtype=torch.float8_e8m0fnu).element_size()  # scale
+    )
+
+    outs: Dict[str, Tuple[torch.Tensor, torch.Tensor]] = {}
+    for b in backends:
+        try:
+            fn = lambda: mxfp8_quantize_cuda_3d(  # noqa: E731
+                x,
+                block_size=32,
+                scaling_mode=args.scaling_mode,
+                backend=b,
+                stage_count=args.stage_count,
+            )
+            ms = _benchmark(fn, args.warmup, args.iters)
+            y, s = fn()
+            outs[b] = (y, s)
+            print(
+                f"[{b:<10}]    {ms:.3f} ms    {_tbps(bytes_moved, ms):.3f} TB/s   "
+                f"y_stride={tuple(y.stride())} s_shape={tuple(s.shape)}"
+            )
+        except Exception as e:
+            print(f"[{b:<10}] FAILED after 0.00s: {type(e).__name__}: {e}")
+
+    if args.check_results and "cuda" in outs:
+        y_ref, s_ref = outs["cuda"]
+        for b in backends:
+            if b == "cuda" or b not in outs:
+                continue
+            y, s = outs[b]
+            dy = _max_abs_diff(y_ref, y)
+            ds = _max_abs_diff(s_ref, s)
+            print(f"diff(cuda vs {b}): y_max_abs={dy} s_max_abs={ds}")
+            ok = dy <= args.atol and ds <= args.atol
+            print(f"check(cuda vs {b}): {'PASS' if ok else 'FAIL'} (atol={args.atol})")
+            if not ok:
+                raise RuntimeError(
+                    f"Result mismatch for backend={b}: y_diff={dy}, s_diff={ds}, atol={args.atol}"
+                )
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--dtype", choices=("bf16", "fp32"), default="bf16")
+    parser.add_argument("--scaling-mode", choices=("floor", "rceil"), default="floor")
+    parser.add_argument("--backends", default="cuda,cutedsl")
+    parser.add_argument("--warmup", type=int, default=20)
+    parser.add_argument("--iters", type=int, default=100)
+    parser.add_argument("--check-results", action="store_true")
+    parser.add_argument("--atol", type=float, default=0.0)
+    parser.add_argument("--stage-count", type=int, default=2)
+
+    parser.add_argument("--E", type=int, default=8)
+    parser.add_argument("--N", type=int, default=7168)
+    parser.add_argument("--K", type=int, default=2048)
+
+    args = parser.parse_args()
+    _run_3d(args)
+
+
+if __name__ == "__main__":
+    main()

test/prototype/moe_training/test_kernels.py

@@ -28,7 +28,6 @@
     triton_fp8_per_group_rowwise_scales,
 )
 from torchao.prototype.moe_training.kernels.mxfp8 import (
-    _mxfp8_cuda_kernels_available,
     fused_pad_token_groups_cuda,
     fused_unpad_token_groups_cuda,
     mx_block_rearrange_2d_M_groups_cuda,
@@ -42,6 +41,10 @@
     triton_mx_block_rearrange_2d_M_groups,
     triton_mx_block_rearrange_per_group_3d,
 )
+from torchao.prototype.moe_training.kernels.mxfp8.quant import (
+    _mxfp8_cuda_kernels_available,
+    _mxfp8_cutedsl_kernels_available,
+)
 from torchao.prototype.moe_training.utils import (
     _is_column_major,
     generate_jagged_offs,
@@ -375,7 +378,13 @@ def test_triton_mx_block_rearrange_2d_K_groups(
 @pytest.mark.parametrize("K", (32, 1536, 5120, 7168, 8192))
 @pytest.mark.parametrize("input_dtype", (torch.bfloat16,))
 @pytest.mark.parametrize("scaling_mode", (ScaleCalculationMode.FLOOR,))
-def test_cuda_mx_dim1_3d_numerics(E, N, K, input_dtype, scaling_mode):
+@pytest.mark.parametrize("backend", ("cuda", "cutedsl"))
+def test_cuda_mx_dim1_3d_numerics(E, N, K, input_dtype, scaling_mode, backend):
+    if backend == "cuda" and not _mxfp8_cuda_kernels_available:
+        pytest.skip("CUDA C++ mxfp8_quantize_3d backend is unavailable")
+    if backend == "cutedsl" and not _mxfp8_cutedsl_kernels_available:
+        pytest.skip("CuTeDSL mxfp8_quantize_3d backend is unavailable")
+
     scaling_mode_str = (
         "floor" if scaling_mode == ScaleCalculationMode.FLOOR else "rceil"
     )
@@ -402,7 +411,10 @@ def test_cuda_mx_dim1_3d_numerics(E, N, K, input_dtype, scaling_mode):
     s_d1_ref = s_d1_ref.transpose(-2, -1)
 
     y_d1, s_d1 = mxfp8_quantize_cuda_3d(
-        x, block_size=block_size, scaling_mode=scaling_mode_str
+        x,
+        block_size=block_size,
+        scaling_mode=scaling_mode_str,
+        backend=backend,
     )
     # Check scales
     torch.testing.assert_close(s_d1, s_d1_ref, rtol=0, atol=0)