diff --git a/benchmarks/prototype/moe_training/mxfp8/bench_wgrad_dim1.py b/benchmarks/prototype/moe_training/mxfp8/bench_wgrad_dim1.py new file mode 100644 index 00000000000..9fe453f2d4e --- /dev/null +++ b/benchmarks/prototype/moe_training/mxfp8/bench_wgrad_dim1.py @@ -0,0 +1,240 @@ +# 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 + +import torch + +from benchmarks.utils import benchmark_cuda_function_in_microseconds +from torchao.prototype.moe_training.kernels.mxfp8 import ( + mxfp8_quantize_2d_32x1_cutedsl, + triton_mx_block_rearrange_2d_K_groups, +) +from torchao.prototype.moe_training.utils import generate_jagged_offs +from torchao.prototype.mx_formats.config import ( + MXFP8Dim1CastKernelChoice, + ScaleCalculationMode, +) +from torchao.prototype.mx_formats.utils import _to_mxfp8_dim1_kernel_wrapper +from torchao.quantization.quantize_.common import KernelPreference + + +def _quantize_dim1_cuda(x, block_size, scale_mode): + x_mx = _to_mxfp8_dim1_kernel_wrapper( + x, + block_size, + elem_dtype=torch.float8_e4m3fn, + hp_dtype=x.dtype, + kernel_preference=KernelPreference.AUTO, + cast_kernel_choice=MXFP8Dim1CastKernelChoice.CUDA, + scale_calculation_mode=scale_mode, + ) + return x_mx.qdata, x_mx.scale + + +def _quantize_dim1_cutedsl(x, block_size, scale_mode): + qdata, scales = mxfp8_quantize_2d_32x1_cutedsl( + x, + block_size=block_size, + scaling_mode=scale_mode.value, + blocked_scale_output=False, + ) + return qdata.t(), scales + + +def _wgrad(grad_output, input_act, offs, block_size, scale_mode, choice): + """The dim1 (wgrad) path: cast both operands, block the scales, grouped mm.""" + quantize_dim1 = ( + _quantize_dim1_cutedsl + if choice is MXFP8Dim1CastKernelChoice.CUTEDSL + else _quantize_dim1_cuda + ) + + grad_output_t_data, grad_output_t_scales = quantize_dim1( + grad_output, block_size, scale_mode + ) + input_act_t_data, input_act_t_scales = quantize_dim1( + input_act, block_size, scale_mode + ) + + scale_group_offsets = offs // block_size + grad_output_t_scales_blocked = triton_mx_block_rearrange_2d_K_groups( + grad_output_t_scales, + scale_group_offsets, + ) + input_act_t_scales_blocked = triton_mx_block_rearrange_2d_K_groups( + input_act_t_scales, + scale_group_offsets, + ) + + grad_weight = torch._scaled_grouped_mm( + grad_output_t_data, + input_act_t_data.transpose(-2, -1), + grad_output_t_scales_blocked, + input_act_t_scales_blocked, + offs=offs, + out_dtype=torch.bfloat16, + ) + # Grouped output is 3-D (one weight-gradient per expert), hence the transpose + # of the last two dims rather than a plain .t(). + return grad_weight.transpose(-2, -1) + + +def benchmark_cuda_graph_function_in_microseconds(f, *args, iters=1000, **kwargs): + for _ in range(10): + f(*args, **kwargs) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + f(*args, **kwargs) + for _ in range(10): + graph.replay() + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iters): + graph.replay() + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) * 1000.0 / iters + + +def benchmark_function(f, *args, use_cuda_graph=False, graph_iters=1000, **kwargs): + if use_cuda_graph: + return benchmark_cuda_graph_function_in_microseconds( + f, + *args, + iters=graph_iters, + **kwargs, + ) + return benchmark_cuda_function_in_microseconds(f, *args, **kwargs) + + +def _parse_shape(spec: str): + """Parses `MxNxKxGROUPS` or `LABEL:MxNxKxGROUPS`.""" + label, _, dims = spec.rpartition(":") + parts = dims.split("x") + if len(parts) != 4: + raise SystemExit(f"Malformed shape {spec!r}, expected [LABEL:]MxNxKxGROUPS") + m, n, k, groups = (int(p) for p in parts) + return (label or None, m, n, k, groups) + + +def _markdown(rows, columns): + widths = [ + max(len(c), *(len(str(r[c])) for r in rows)) if rows else len(c) + for c in columns + ] + header = "| " + " | ".join(c.ljust(w) for c, w in zip(columns, widths)) + " |" + sep = "|" + "|".join("---:".rjust(w + 2) for w in widths) + "|" + body = [ + "| " + " | ".join(str(r[c]).rjust(w) for c, w in zip(columns, widths)) + " |" + for r in rows + ] + return "\n".join([header, sep, *body]) + + +def main(): + parser = argparse.ArgumentParser( + description="Benchmark the MXFP8 dim1 (wgrad) path with the CUDA vs CuTeDSL " + "dim1 cast kernel." + ) + parser.add_argument( + "shapes", + nargs="+", + metavar="[LABEL:]MxNxKxGROUPS", + help="shapes to benchmark, e.g. 4096x4096x2048x8 or dsv3:4096x4096x2048x8", + ) + parser.add_argument("--multiple-of", type=int, default=128) + parser.add_argument( + "--scale-mode", + choices=("floor", "rceil"), + default="rceil", + ) + parser.add_argument( + "--bench-only", + choices=("cuda", "cutedsl"), + default=None, + ) + parser.add_argument("--deterministic", action="store_true") + parser.add_argument("--cuda-graph-bench", action="store_true") + parser.add_argument("--graph-iters", type=int, default=1000) + args = parser.parse_args() + + if not torch.cuda.is_available(): + raise SystemExit("CUDA is not available") + + shapes = [_parse_shape(s) for s in args.shapes] + labelled = any(label is not None for label, *_ in shapes) + block_size = 32 + scale_mode = ScaleCalculationMode(args.scale_mode) + rows_out = [] + + for label, m, n, k, groups in shapes: + if args.deterministic: + grad_output = ( + torch.arange(m * n, device="cuda", dtype=torch.float32) + .remainder(17) + .sub(8) + .div(16) + .reshape(m, n) + .to(torch.bfloat16) + ) + input_act = ( + torch.arange(m * k, device="cuda", dtype=torch.float32) + .remainder(17) + .sub(8) + .div(16) + .reshape(m, k) + .to(torch.bfloat16) + ) + else: + grad_output = torch.randn(m, n, dtype=torch.bfloat16, device="cuda") + input_act = torch.randn(m, k, dtype=torch.bfloat16, device="cuda") + offsets = generate_jagged_offs( + groups, m, multiple_of=args.multiple_of, device="cuda" + ) + + row = {"Model layer": label or ""} if labelled else {} + row["MxNxK"] = f"{m}x{n}x{k}" + row["Groups"] = groups + + if args.bench_only in (None, "cuda"): + cuda_us = benchmark_function( + _wgrad, + grad_output, + input_act, + offsets, + block_size, + scale_mode, + MXFP8Dim1CastKernelChoice.CUDA, + use_cuda_graph=args.cuda_graph_bench, + graph_iters=args.graph_iters, + ) + row["CUDA us"] = f"{cuda_us:.2f}" + if args.bench_only in (None, "cutedsl"): + cutedsl_us = benchmark_function( + _wgrad, + grad_output, + input_act, + offsets, + block_size, + scale_mode, + MXFP8Dim1CastKernelChoice.CUTEDSL, + use_cuda_graph=args.cuda_graph_bench, + graph_iters=args.graph_iters, + ) + row["CuTeDSL us"] = f"{cutedsl_us:.2f}" + if args.bench_only is None: + row["CuTeDSL speedup"] = f"{cuda_us / cutedsl_us:.2f}" + rows_out.append(row) + + print(_markdown(rows_out, list(rows_out[0].keys()))) + + +if __name__ == "__main__": + main() diff --git a/test/prototype/moe_training/test_kernels.py b/test/prototype/moe_training/test_kernels.py index 077491aa35a..398e9a6dda3 100644 --- a/test/prototype/moe_training/test_kernels.py +++ b/test/prototype/moe_training/test_kernels.py @@ -6,6 +6,7 @@ import pytest import torch +from torch._subclasses.fake_tensor import FakeTensorMode # FP8 MoE kernels require FP8-capable hardware (SM 10.x on CUDA, MI300+ on ROCm) from torchao.utils import is_MI300, is_MI350 @@ -793,6 +794,29 @@ def test_cuda_mx_dim1_2d_numerics_32x1( # Check quantized values - no padding needed for data torch.testing.assert_close(y_d1, y_d1_ref, rtol=0, atol=0) + # The registered fake must report the same metadata as the real op, otherwise + # torch.compile / fake-tensor tracing sees the wrong shapes, strides or dtypes. + # This covers both blocked_scale_output values via the parametrization above. + with FakeTensorMode(): + x_fake = torch.empty_like(x, device="cuda") + y_fake, s_fake = torch.ops.torchao.mxfp8_quantize_2d_32x1_cutedsl.default( + x_fake, + block_size=block_size, + scaling_mode=scaling_mode_str, + blocked_scale_output=blocked_scale_output, + ) + assert y_fake.shape == y_d1.shape, ( + f"fake data shape {y_fake.shape} != real {y_d1.shape}" + ) + assert y_fake.stride() == y_d1.stride(), ( + f"fake data stride {y_fake.stride()} != real {y_d1.stride()}" + ) + assert y_fake.dtype == y_d1.dtype + assert s_fake.shape == s_d1.shape, ( + f"fake scale shape {s_fake.shape} != real {s_d1.shape}" + ) + assert s_fake.dtype == s_d1.dtype + def _make_mxfp8_edge_input(shape, dtype, device, pattern: str) -> torch.Tensor: numel = 1 diff --git a/torchao/prototype/moe_training/kernels/mxfp8/quant.py b/torchao/prototype/moe_training/kernels/mxfp8/quant.py index d5e42e6d74a..ffb0f3ca9ae 100644 --- a/torchao/prototype/moe_training/kernels/mxfp8/quant.py +++ b/torchao/prototype/moe_training/kernels/mxfp8/quant.py @@ -1141,22 +1141,30 @@ def _fake_mxfp8_quantize_2d_32x1_cutedsl_custom_op( assert block_size == 32, "Only block_size=32 is supported" m, k = x.shape - # For 32x1 scaling, output data has same dimensions as input (no padding) + # Output data has the same dimensions as the input, but the 32x1 kernel writes + # it column-major, hence the (1, m) strides. q_data = torch.empty_strided( (m, k), - (k, 1), + (1, m), device=x.device, dtype=torch.float8_e4m3fn, ) - # Only scales tensor needs padding for 4x128 scale factor tiles m_blocks = m // block_size - padded_scale_rows = ceil_div(m_blocks, 4) * 4 - padded_scale_cols = ceil_div(k, 128) * 128 - scales = x.new_empty( - (padded_scale_cols * padded_scale_rows,), - dtype=torch.float8_e8m0fnu, - ) + if blocked_scale_output: + # Only the scales need padding, to whole 128x4 scale-factor tiles. Names + # match cutedsl_quantize_2d_32x1.py: rows come from K, cols from m_blocks. + padded_scale_rows = ceil_div(k, 128) * 128 + padded_scale_cols = ceil_div(m_blocks, 4) * 4 + scales = x.new_empty( + (padded_scale_rows * padded_scale_cols,), + dtype=torch.float8_e8m0fnu, + ) + else: + scales = x.new_empty( + (k, m_blocks), + dtype=torch.float8_e8m0fnu, + ) return q_data, scales