diff --git a/benchmarks/prototype/moe_training/mxfp8/bench_pad_token_groups.py b/benchmarks/prototype/moe_training/mxfp8/bench_pad_token_groups.py deleted file mode 100644 index 5fca5bd4fa8..00000000000 --- a/benchmarks/prototype/moe_training/mxfp8/bench_pad_token_groups.py +++ /dev/null @@ -1,237 +0,0 @@ -# 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. -# this benchmarking script is a modified version of the original script from: https://github.com/drisspg/transformer_nuggets/blob/main/transformer_nuggets/utils/benchmark.py - -import argparse -import itertools -from dataclasses import dataclass -from typing import List - -import torch -from tabulate import tabulate -from tqdm import tqdm - -from benchmarks.utils import benchmark_cuda_function_in_microseconds, profile_fn -from torchao.prototype.moe_training.kernels.mxfp8 import ( - _mxfp8_cuda_kernels_available, - fused_pad_token_groups_cuda, - torch_pad_token_groups, -) -from torchao.prototype.moe_training.utils import generate_jagged_offs - -device = torch.device("cuda") - -# Needed since changing args to function causes recompiles -torch._dynamo.config.cache_size_limit = 1000 - - -@dataclass(frozen=True) -class ExperimentConfig: - num_tokens: int - dim: int - num_groups: int - alignment_size: int - - -@dataclass(frozen=True) -class ExperimentResult: - torch_eager_time_us: float - cuda_time_us: float - torch_mem_bw_gbps: float - cuda_mem_bw_gbps: float - - -@dataclass(frozen=True) -class Experiment: - config: ExperimentConfig - result: ExperimentResult - - -def get_configs() -> List[ExperimentConfig]: - # Various token group sizes and dimensions - num_tokens_list = [16384] - dim_list = [1536, 2048, 5120, 7168] - num_groups_list = [1, 4, 8, 16] - alignment_size_list = [32] - - configs = [] - for num_tokens, dim, num_groups, alignment_size in itertools.product( - num_tokens_list, dim_list, num_groups_list, alignment_size_list - ): - configs.append( - ExperimentConfig( - num_tokens=num_tokens, - dim=dim, - num_groups=num_groups, - alignment_size=alignment_size, - ) - ) - return configs - - -def run_experiment( - config: ExperimentConfig, args: argparse.Namespace -) -> ExperimentResult: - num_tokens, dim, num_groups, alignment_size = ( - config.num_tokens, - config.dim, - config.num_groups, - config.alignment_size, - ) - - inputs = torch.randn(num_tokens, dim, dtype=torch.bfloat16, device=device) - - def torch_eager_with_offsets(): - group_offsets = generate_jagged_offs( - num_groups, num_tokens, multiple_of=1, device=device - ) - return torch_pad_token_groups( - inputs, group_offsets, alignment_size - ) # Returns 3 values - - def warmup(fn): - for _ in range(5): - fn() - - # bench torch eager (includes buffer allocation overhead) - warmup(torch_eager_with_offsets) - torch_eager_time_us = benchmark_cuda_function_in_microseconds( - torch_eager_with_offsets - ) - if args.profile: - group_offsets = generate_jagged_offs( - num_groups, num_tokens, multiple_of=1, device=device - ) - profile_fn( - torch_pad_token_groups, - inputs, - group_offsets, - alignment_size, - profile_name="torch_pad_token_groups_eager", - ) - - # bench CUDA kernel if available - if _mxfp8_cuda_kernels_available: - - def cuda_with_offsets(): - group_offsets = generate_jagged_offs( - num_groups, num_tokens, multiple_of=1, device=device - ) - return fused_pad_token_groups_cuda(inputs, group_offsets, alignment_size) - - warmup(cuda_with_offsets) - cuda_time_us = benchmark_cuda_function_in_microseconds(cuda_with_offsets) - if args.profile: - group_offsets = generate_jagged_offs( - num_groups, num_tokens, multiple_of=1, device=device - ) - profile_fn( - fused_pad_token_groups_cuda, - inputs, - group_offsets, - alignment_size, - profile_name="fused_pad_token_groups_cuda", - ) - else: - cuda_time_us = float("inf") # Not available - - # mem bw calculations - run once to get output sizes - group_offsets = generate_jagged_offs( - num_groups, num_tokens, multiple_of=1, device=device - ) - torch_padded_tokens, torch_padded_start_offsets, torch_padded_offsets = ( - torch_pad_token_groups(inputs, group_offsets, alignment_size) - ) - - bytes_per_el = torch.finfo(torch.bfloat16).bits / 8 - - read_bytes = ( - inputs.numel() * bytes_per_el # Read input tokens - + group_offsets.numel() * 4 # Read group offsets (int32) - ) - - write_bytes = ( - torch_padded_tokens.numel() * bytes_per_el # Write zeros (entire buffer) - + inputs.numel() * bytes_per_el # Write actual data (overwrites part of zeros) - + torch_padded_offsets.numel() * 4 # Write output offsets (int32) - ) - - total_bytes = read_bytes + write_bytes - - torch_mem_bw_gbps = (total_bytes / 1e9) / (torch_eager_time_us / 1e6) - - if _mxfp8_cuda_kernels_available and cuda_time_us != float("inf"): - cuda_mem_bw_gbps = (total_bytes / 1e9) / (cuda_time_us / 1e6) - else: - cuda_mem_bw_gbps = 0.0 - - return ExperimentResult( - torch_eager_time_us=torch_eager_time_us, - cuda_time_us=cuda_time_us, - torch_mem_bw_gbps=torch_mem_bw_gbps, - cuda_mem_bw_gbps=cuda_mem_bw_gbps, - ) - - -def print_results(experiments: List[Experiment]): - headers = [ - "num_tokens", - "dim", - "num_groups", - "torch_us", - "cuda_us", - "torch_mem_bw_gbps", - "cuda_mem_bw_gbps", - "cuda_vs_torch", - ] - rows = [] - for experiment in experiments: - cuda_time = experiment.result.cuda_time_us - cuda_vs_torch = ( - f"{experiment.result.torch_eager_time_us / cuda_time:.2f}x" - if cuda_time != float("inf") and cuda_time > 0 - else "N/A" - ) - cuda_bw_str = ( - f"{experiment.result.cuda_mem_bw_gbps:.2f}" - if experiment.result.cuda_mem_bw_gbps > 0 - else "N/A" - ) - - rows.append( - [ - experiment.config.num_tokens, - experiment.config.dim, - experiment.config.num_groups, - experiment.result.torch_eager_time_us, - experiment.result.cuda_time_us, - f"{experiment.result.torch_mem_bw_gbps:.2f}", - cuda_bw_str, - cuda_vs_torch, - ] - ) - print(tabulate(rows, headers=headers)) - - -def main(args: argparse.Namespace): - torch.random.manual_seed(123) - configs = get_configs() - results = [] - for config in tqdm(configs): - result = run_experiment(config, args) - results.append(Experiment(config=config, result=result)) - - # Use Tabulate to print results - print_results(results) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument( - "--profile", action="store_true", help="Enable profiling with PyTorch profiler" - ) - args = parser.parse_args() - main(args) diff --git a/benchmarks/prototype/moe_training/mxfp8/bench_pad_unpad_token_groups.py b/benchmarks/prototype/moe_training/mxfp8/bench_pad_unpad_token_groups.py new file mode 100644 index 00000000000..47142090793 --- /dev/null +++ b/benchmarks/prototype/moe_training/mxfp8/bench_pad_unpad_token_groups.py @@ -0,0 +1,380 @@ +# 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 multiprocessing +import sys + +import torch + +from benchmarks.utils import benchmark_cuda_function_in_microseconds +from torchao.prototype.moe_training.kernels.mxfp8 import ( + _pad_token_groups_cutedsl, + _unpad_token_groups_cutedsl, + fused_pad_token_groups_cuda, + fused_unpad_token_groups_cuda, + torch_pad_token_groups, + torch_unpad_token_groups, +) +from torchao.prototype.moe_training.utils import generate_jagged_offs + +# fused_{pad,unpad}_token_groups_cuda can hit an illegal memory access for some +# large shapes, which poisons the CUDA context; not fixing since this legacy +# kernel will be deleted soon, so run it in a subprocess to isolate it instead. +_CUDA_WORKER_TIMEOUT_S = 300 + + +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() + result = start.elapsed_time(end) * 1000.0 / iters + # Tear the graph (and its private memory pool) down before the caller + # captures another one for a different shape. + del graph + torch.cuda.synchronize() + return result + + +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 _make_inputs(tokens, dim, dtype, deterministic): + if deterministic: + return ( + torch.arange(tokens * dim, device="cuda", dtype=torch.float32) + .remainder(17) + .sub(8) + .div(16) + .reshape(tokens, dim) + .to(dtype) + ) + return torch.randn(tokens, dim, dtype=dtype, device="cuda") + + +def _pad_cuda_worker( + tokens, + dim, + groups, + alignment_size, + dtype_name, + deterministic, + use_cuda_graph, + graph_iters, + multiple_of, + out_queue, +): + dtype = torch.bfloat16 if dtype_name == "bf16" else torch.float32 + inputs = _make_inputs(tokens, dim, dtype, deterministic) + offsets = generate_jagged_offs( + groups, tokens, multiple_of=multiple_of, device="cuda" + ) + us = benchmark_function( + fused_pad_token_groups_cuda, + inputs, + offsets, + alignment_size, + use_cuda_graph=use_cuda_graph, + graph_iters=graph_iters, + ) + out_queue.put(us) + + +def _unpad_cuda_worker( + tokens, + dim, + groups, + alignment_size, + dtype_name, + deterministic, + use_cuda_graph, + graph_iters, + multiple_of, + out_queue, +): + dtype = torch.bfloat16 if dtype_name == "bf16" else torch.float32 + inputs = _make_inputs(tokens, dim, dtype, deterministic) + offsets = generate_jagged_offs( + groups, tokens, multiple_of=multiple_of, device="cuda" + ) + padded_inputs, padded_start_offsets, _ = torch_pad_token_groups( + inputs, offsets, alignment_size + ) + us = benchmark_function( + fused_unpad_token_groups_cuda, + padded_inputs, + offsets, + padded_start_offsets, + tokens, + alignment_size, + use_cuda_graph=use_cuda_graph, + graph_iters=graph_iters, + ) + out_queue.put(us) + + +def _bench_cuda_isolated( + worker, + op_name, + tokens, + dim, + groups, + alignment_size, + dtype_name, + deterministic, + use_cuda_graph, + graph_iters, + multiple_of, +): + ctx = multiprocessing.get_context("spawn") + out_queue = ctx.Queue() + proc = ctx.Process( + target=worker, + args=( + tokens, + dim, + groups, + alignment_size, + dtype_name, + deterministic, + use_cuda_graph, + graph_iters, + multiple_of, + out_queue, + ), + ) + proc.start() + proc.join(_CUDA_WORKER_TIMEOUT_S) + if proc.exitcode == 0 and not out_queue.empty(): + return out_queue.get() + if proc.is_alive(): + proc.terminate() + print( + f"warning: cuda {op_name} crashed for shape {tokens}x{dim}x{groups}, marking N/A", + file=sys.stderr, + ) + return None + + +def _parse_shape(spec: str): + """Parses `TOKENSxDIMxGROUPS` or `LABEL:TOKENSxDIMxGROUPS`.""" + label, _, dims = spec.rpartition(":") + parts = dims.split("x") + if len(parts) != 3: + raise SystemExit( + f"Malformed shape {spec!r}, expected [LABEL:]TOKENSxDIMxGROUPS" + ) + tokens, dim, groups = (int(p) for p in parts) + return (label or None, tokens, dim, 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 torch / cuda / CuTeDSL MXFP8 token-group padding and unpadding." + ) + parser.add_argument( + "shapes", + nargs="+", + metavar="[LABEL:]TOKENSxDIMxGROUPS", + help="shapes to benchmark, e.g. 16384x7168x8 or dsv3:16384x7168x8", + ) + parser.add_argument("--mode", choices=("pad", "unpad", "both"), default="both") + parser.add_argument("--alignment-size", type=int, default=32) + parser.add_argument("--multiple-of", type=int, default=1) + parser.add_argument("--dtype", choices=("bf16", "fp32"), default="bf16") + parser.add_argument( + "--bench-only", + nargs="+", + choices=("torch", "cuda", "cutedsl"), + default=None, + help="restrict to one or more backends; default is all", + ) + 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") + + dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float32 + shapes = [_parse_shape(s) for s in args.shapes] + labelled = any(label is not None for label, _, _, _ in shapes) + selected = set(args.bench_only) if args.bench_only else {"torch", "cuda", "cutedsl"} + if args.cuda_graph_bench and "torch" in selected: + # torch_{pad,unpad}_token_groups call .tolist(), a synchronous D2H copy + # that CUDA graph capture never tolerates, for any shape. + if selected == {"torch"}: + raise SystemExit("torch cannot be benchmarked under --cuda-graph-bench") + print( + "warning: excluding torch from --cuda-graph-bench (unpinned D2H copy)", + file=sys.stderr, + ) + selected.discard("torch") + + do_pad = args.mode in ("pad", "both") + do_unpad = args.mode in ("unpad", "both") + rows_out = [] + + for label, tokens, dim, groups in shapes: + inputs = _make_inputs(tokens, dim, dtype, args.deterministic) + offsets = generate_jagged_offs( + groups, tokens, multiple_of=args.multiple_of, device="cuda" + ) + + row = {"Model layer": label or ""} if labelled else {} + row["Tokens x Dim"] = f"{tokens}x{dim}" + row["Groups"] = groups + + if do_pad: + pad_cuda_us = None + call_args = (inputs, offsets, args.alignment_size) + if "torch" in selected: + us = benchmark_function( + torch_pad_token_groups, + *call_args, + use_cuda_graph=args.cuda_graph_bench, + graph_iters=args.graph_iters, + ) + row["Pad Torch us"] = f"{us:.2f}" + if "cuda" in selected: + pad_cuda_us = _bench_cuda_isolated( + _pad_cuda_worker, + "pad", + tokens, + dim, + groups, + args.alignment_size, + args.dtype, + args.deterministic, + args.cuda_graph_bench, + args.graph_iters, + args.multiple_of, + ) + row["Pad CUDA us"] = ( + f"{pad_cuda_us:.2f}" if pad_cuda_us is not None else "N/A" + ) + if "cutedsl" in selected: + out, _, _ = _pad_token_groups_cutedsl(*call_args) + pad_cutedsl_us = benchmark_function( + _pad_token_groups_cutedsl, + *call_args, + use_cuda_graph=args.cuda_graph_bench, + graph_iters=args.graph_iters, + ) + row["Pad CuTeDSL us"] = f"{pad_cutedsl_us:.2f}" + moved_bytes = (inputs.numel() + out.numel()) * inputs.element_size() + row["Pad CuTeDSL GB/s"] = ( + f"{(moved_bytes / 1e9) / (pad_cutedsl_us / 1e6):.1f}" + ) + if "cuda" in selected and "cutedsl" in selected: + row["Pad CuTeDSL speedup"] = ( + f"{pad_cuda_us / pad_cutedsl_us:.2f}" + if pad_cuda_us is not None + else "N/A" + ) + + if do_unpad: + unpad_cuda_us = None + padded_inputs, padded_start_offsets, _ = torch_pad_token_groups( + inputs, offsets, args.alignment_size + ) + call_args = ( + padded_inputs, + offsets, + padded_start_offsets, + tokens, + args.alignment_size, + ) + if "torch" in selected: + us = benchmark_function( + torch_unpad_token_groups, + *call_args, + use_cuda_graph=args.cuda_graph_bench, + graph_iters=args.graph_iters, + ) + row["Unpad Torch us"] = f"{us:.2f}" + if "cuda" in selected: + unpad_cuda_us = _bench_cuda_isolated( + _unpad_cuda_worker, + "unpad", + tokens, + dim, + groups, + args.alignment_size, + args.dtype, + args.deterministic, + args.cuda_graph_bench, + args.graph_iters, + args.multiple_of, + ) + row["Unpad CUDA us"] = ( + f"{unpad_cuda_us:.2f}" if unpad_cuda_us is not None else "N/A" + ) + if "cutedsl" in selected: + out = _unpad_token_groups_cutedsl(*call_args) + unpad_cutedsl_us = benchmark_function( + _unpad_token_groups_cutedsl, + *call_args, + use_cuda_graph=args.cuda_graph_bench, + graph_iters=args.graph_iters, + ) + row["Unpad CuTeDSL us"] = f"{unpad_cutedsl_us:.2f}" + moved_bytes = ( + padded_inputs.numel() + out.numel() + ) * inputs.element_size() + row["Unpad CuTeDSL GB/s"] = ( + f"{(moved_bytes / 1e9) / (unpad_cutedsl_us / 1e6):.1f}" + ) + if "cuda" in selected and "cutedsl" in selected: + row["Unpad CuTeDSL speedup"] = ( + f"{unpad_cuda_us / unpad_cutedsl_us:.2f}" + if unpad_cuda_us is not None + else "N/A" + ) + + rows_out.append(row) + + print(_markdown(rows_out, list(rows_out[0].keys()))) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/prototype/moe_training/mxfp8/bench_unpad_token_groups.py b/benchmarks/prototype/moe_training/mxfp8/bench_unpad_token_groups.py deleted file mode 100644 index e1d2da82b3f..00000000000 --- a/benchmarks/prototype/moe_training/mxfp8/bench_unpad_token_groups.py +++ /dev/null @@ -1,240 +0,0 @@ -# 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. -# this benchmarking script is a modified version of the original script from: https://github.com/drisspg/transformer_nuggets/blob/main/transformer_nuggets/utils/benchmark.py - -import argparse -import itertools -from dataclasses import dataclass -from typing import List - -import torch -from tabulate import tabulate -from tqdm import tqdm - -from benchmarks.utils import benchmark_cuda_function_in_microseconds, profile_fn -from torchao.prototype.moe_training.kernels.mxfp8 import ( - _mxfp8_cuda_kernels_available, - fused_unpad_token_groups_cuda, - torch_pad_token_groups, - torch_unpad_token_groups, -) -from torchao.prototype.moe_training.utils import generate_jagged_offs - -device = torch.device("cuda") - -# Needed since changing args to function causes recompiles -torch._dynamo.config.cache_size_limit = 1000 - - -@dataclass(frozen=True) -class ExperimentConfig: - num_tokens: int - dim: int - num_groups: int - alignment_size: int - - -@dataclass(frozen=True) -class ExperimentResult: - torch_eager_time_us: float - cuda_time_us: float - torch_mem_bw_gbps: float - cuda_mem_bw_gbps: float - - -@dataclass(frozen=True) -class Experiment: - config: ExperimentConfig - result: ExperimentResult - - -def get_configs() -> List[ExperimentConfig]: - # Various token group sizes and dimensions - num_tokens_list = [16384] - dim_list = [1536, 2048, 5120, 7168] - num_groups_list = [1, 4, 8, 16] - alignment_size_list = [32] - - configs = [] - for num_tokens, dim, num_groups, alignment_size in itertools.product( - num_tokens_list, dim_list, num_groups_list, alignment_size_list - ): - configs.append( - ExperimentConfig( - num_tokens=num_tokens, - dim=dim, - num_groups=num_groups, - alignment_size=alignment_size, - ) - ) - return configs - - -def run_experiment( - config: ExperimentConfig, args: argparse.Namespace -) -> ExperimentResult: - num_tokens, dim, num_groups, alignment_size = ( - config.num_tokens, - config.dim, - config.num_groups, - config.alignment_size, - ) - - # Create inputs and pad them first - inputs = torch.randn(num_tokens, dim, dtype=torch.bfloat16, device=device) - group_offsets = generate_jagged_offs( - num_groups, num_tokens, multiple_of=1, device=device - ) - - # Pad the inputs to get padded tensors for unpad benchmark - padded_inputs, padded_group_start_offsets, padded_group_end_offsets = ( - torch_pad_token_groups(inputs, group_offsets, alignment_size) - ) - - def torch_eager_with_offsets(): - return torch_unpad_token_groups( - padded_inputs, - group_offsets, - padded_group_start_offsets, - num_tokens, - alignment_size, - ) - - def warmup(fn): - for _ in range(5): - fn() - - # bench torch eager (includes buffer allocation overhead) - warmup(torch_eager_with_offsets) - torch_eager_time_us = benchmark_cuda_function_in_microseconds( - torch_eager_with_offsets - ) - if args.profile: - profile_fn( - torch_unpad_token_groups, - padded_inputs, - group_offsets, - padded_group_start_offsets, - alignment_size, - profile_name="torch_unpad_token_groups_eager", - ) - - # bench CUDA kernel if available - if _mxfp8_cuda_kernels_available: - - def cuda_with_offsets(): - return fused_unpad_token_groups_cuda( - padded_inputs, - group_offsets, - padded_group_start_offsets, - num_tokens, - alignment_size, - ) - - warmup(cuda_with_offsets) - cuda_time_us = benchmark_cuda_function_in_microseconds(cuda_with_offsets) - if args.profile: - profile_fn( - fused_unpad_token_groups_cuda, - padded_inputs, - group_offsets, - padded_group_start_offsets, - num_tokens, - alignment_size, - profile_name="fused_unpad_token_groups_cuda", - ) - else: - cuda_time_us = float("inf") # Not available - - # mem bw calculations - bytes_per_el = torch.finfo(torch.bfloat16).bits / 8 - - read_bytes = ( - padded_inputs.numel() * bytes_per_el # Read padded input tokens - + group_offsets.numel() * 4 # Read group offsets (int32) - + padded_group_start_offsets.numel() * 4 # Read padded start offsets (int32) - ) - - write_bytes = ( - inputs.numel() * bytes_per_el # Write unpadded data - ) - - total_bytes = read_bytes + write_bytes - - torch_mem_bw_gbps = (total_bytes / 1e9) / (torch_eager_time_us / 1e6) - - if _mxfp8_cuda_kernels_available and cuda_time_us != float("inf"): - cuda_mem_bw_gbps = (total_bytes / 1e9) / (cuda_time_us / 1e6) - else: - cuda_mem_bw_gbps = 0.0 - - return ExperimentResult( - torch_eager_time_us=torch_eager_time_us, - cuda_time_us=cuda_time_us, - torch_mem_bw_gbps=torch_mem_bw_gbps, - cuda_mem_bw_gbps=cuda_mem_bw_gbps, - ) - - -def print_results(experiments: List[Experiment]): - headers = [ - "num_tokens", - "dim", - "num_groups", - "torch_us", - "cuda_us", - "torch_mem_bw_gbps", - "cuda_mem_bw_gbps", - "cuda_vs_torch", - ] - rows = [] - for experiment in experiments: - cuda_time = experiment.result.cuda_time_us - cuda_vs_torch = ( - f"{experiment.result.torch_eager_time_us / cuda_time:.2f}x" - if cuda_time != float("inf") and cuda_time > 0 - else "N/A" - ) - cuda_bw_str = ( - f"{experiment.result.cuda_mem_bw_gbps:.2f}" - if experiment.result.cuda_mem_bw_gbps > 0 - else "N/A" - ) - - rows.append( - [ - experiment.config.num_tokens, - experiment.config.dim, - experiment.config.num_groups, - experiment.result.torch_eager_time_us, - experiment.result.cuda_time_us, - f"{experiment.result.torch_mem_bw_gbps:.2f}", - cuda_bw_str, - cuda_vs_torch, - ] - ) - print(tabulate(rows, headers=headers)) - - -def main(args: argparse.Namespace): - torch.random.manual_seed(123) - configs = get_configs() - results = [] - for config in tqdm(configs): - result = run_experiment(config, args) - results.append(Experiment(config=config, result=result)) - - # Use Tabulate to print results - print_results(results) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument( - "--profile", action="store_true", help="Enable profiling with PyTorch profiler" - ) - args = parser.parse_args() - main(args) diff --git a/test/prototype/moe_training/test_kernels.py b/test/prototype/moe_training/test_kernels.py index 398e9a6dda3..582b6e6ba73 100644 --- a/test/prototype/moe_training/test_kernels.py +++ b/test/prototype/moe_training/test_kernels.py @@ -55,6 +55,10 @@ def _is_sm_10x() -> bool: triton_mx_block_rearrange_2d_M_groups, triton_mx_block_rearrange_per_group_3d, ) +from torchao.prototype.moe_training.kernels.mxfp8.cutedsl_pad_token_groups import ( + _pad_token_groups_cutedsl, + _unpad_token_groups_cutedsl, +) from torchao.prototype.moe_training.kernels.mxfp8.cutedsl_rearrange_2d_m_groups import ( _mx_block_rearrange_2d_m_groups_cutedsl, ) @@ -1225,109 +1229,127 @@ def test_cuda_mx_3d_cutedsl_edge_inputs( assert y.stride() == y_ref.stride() -@pytest.mark.skipif( - not _mxfp8_cuda_kernels_available, - reason="CUDA kernel requires sm_100 and CUDA 12.8+", -) +# Cases are listed per impl rather than crossed, so merging the cuda and cutedsl +# tests preserves exactly the coverage each had before. +_PAD_CASES = [ + *(("cuda", t, 7168, g, 32) for t in (128, 157, 4096, 16392) for g in (1, 2, 4, 8)), + *( + ("cutedsl", t, d, g, 32) + for t in (128, 157, 4096) + for d in (32, 7168) + for g in (1, 8) + ), +] + +_UNPAD_CASES = [ + *(("cuda", t, 7168, g, 32) for t in (128, 157, 4096) for g in (1, 2, 4, 8)), + *( + ("cutedsl", t, d, g, 32) + for t in (128, 157, 4096) + for d in (32, 7168) + for g in (1, 8) + ), +] + + +def _skip_if_impl_unavailable(impl: str): + if impl == "cuda" and not _mxfp8_cuda_kernels_available: + pytest.skip("CUDA kernel requires sm_100 and CUDA 12.8+") + if impl == "cutedsl" and not _mxfp8_cutedsl_kernels_available: + pytest.skip("MXFP8 cutedsl kernels not available") + + @skip_if_rocm("ROCm enablement in progress") -@pytest.mark.parametrize("num_tokens", [128, 157, 4096, 16392]) -@pytest.mark.parametrize("dim", [7168]) -@pytest.mark.parametrize("num_groups", [1, 2, 4, 8]) -@pytest.mark.parametrize("alignment_size", [32]) +@pytest.mark.parametrize("impl,num_tokens,dim,num_groups,alignment_size", _PAD_CASES) @pytest.mark.parametrize("dtype", [torch.bfloat16]) -def test_cuda_fused_pad_token_groups( - num_tokens: int, dim: int, num_groups: int, alignment_size: int, dtype: torch.dtype +def test_fused_pad_token_groups( + impl: str, + num_tokens: int, + dim: int, + num_groups: int, + alignment_size: int, + dtype: torch.dtype, ): - """Test fused_pad_token_groups_cuda kernel for padding token groups to alignment.""" + """Pad token groups to alignment, cuda and cutedsl against the torch reference.""" + _skip_if_impl_unavailable(impl) device = "cuda" - - # Create input activations inputs = torch.randn(num_tokens, dim, dtype=dtype, device=device) - - # Generate group offsets (end indices for each group) group_offsets = generate_jagged_offs( num_groups, num_tokens, multiple_of=1, device=device ) - # Get reference output - ref_padded_tokens, ref_padded_start_offsets, ref_padded_offsets = ( - torch_pad_token_groups(inputs, group_offsets, alignment_size) + ref_padded_tokens, ref_start_offsets, ref_end_offsets = torch_pad_token_groups( + inputs, group_offsets, alignment_size ) - - # Run CUDA kernel - kernel_padded_tokens, kernel_padded_start_offsets, kernel_padded_end_offsets = ( - fused_pad_token_groups_cuda(inputs, group_offsets, alignment_size) + kernel = ( + _pad_token_groups_cutedsl if impl == "cutedsl" else fused_pad_token_groups_cuda ) - - # All implementations now use the same upper bound output size - # Verify outputs match - assert torch.allclose(ref_padded_tokens, kernel_padded_tokens, rtol=0, atol=1e-5), ( - "Padded tokens do not match" + padded_tokens, start_offsets, end_offsets = kernel( + inputs, group_offsets, alignment_size ) - assert torch.equal(ref_padded_start_offsets, kernel_padded_start_offsets), ( + + if impl == "cutedsl": + assert torch.equal(ref_padded_tokens, padded_tokens) + else: + assert torch.allclose(ref_padded_tokens, padded_tokens, rtol=0, atol=1e-5), ( + "Padded tokens do not match" + ) + assert torch.equal(ref_start_offsets, start_offsets), ( "Padded group start offsets do not match" ) - assert torch.equal(ref_padded_offsets, kernel_padded_end_offsets), ( + assert torch.equal(ref_end_offsets, end_offsets), ( "Padded group end offsets do not match" ) -@pytest.mark.skipif( - not _mxfp8_cuda_kernels_available, - reason="CUDA kernel requires sm_100 and CUDA 12.8+", -) @skip_if_rocm("ROCm enablement in progress") -@pytest.mark.parametrize("num_tokens", [128, 157, 4096]) -@pytest.mark.parametrize("dim", [7168]) -@pytest.mark.parametrize("num_groups", [1, 2, 4, 8]) -@pytest.mark.parametrize("alignment_size", [32]) +@pytest.mark.parametrize("impl,num_tokens,dim,num_groups,alignment_size", _UNPAD_CASES) @pytest.mark.parametrize("dtype", [torch.bfloat16]) -def test_cuda_fused_unpad_token_groups( - num_tokens: int, dim: int, num_groups: int, alignment_size: int, dtype: torch.dtype +def test_fused_unpad_token_groups( + impl: str, + num_tokens: int, + dim: int, + num_groups: int, + alignment_size: int, + dtype: torch.dtype, ): - """Test fused_unpad_token_groups_cuda kernel for removing padding from token groups.""" + """Remove padding from token groups, cuda and cutedsl against the torch reference.""" + _skip_if_impl_unavailable(impl) device = "cuda" - - # Create input activations inputs = torch.randn(num_tokens, dim, dtype=dtype, device=device) - - # Generate group offsets (end indices for each group) group_offsets = generate_jagged_offs( num_groups, num_tokens, multiple_of=1, device=device ) - - # First pad the tokens to create padded inputs - padded_tokens, padded_group_start_offsets, padded_offsets = torch_pad_token_groups( + padded_tokens, padded_group_start_offsets, _ = torch_pad_token_groups( inputs, group_offsets, alignment_size ) - # Get reference output using torch implementation - ref_unpadded_tokens = torch_unpad_token_groups( + call_args = ( padded_tokens, group_offsets, padded_group_start_offsets, num_tokens, alignment_size, ) - - # Run CUDA kernel - kernel_unpadded_tokens = fused_unpad_token_groups_cuda( - padded_tokens, - group_offsets, - padded_group_start_offsets, - num_tokens, - alignment_size, + ref_unpadded_tokens = torch_unpad_token_groups(*call_args) + kernel = ( + _unpad_token_groups_cutedsl + if impl == "cutedsl" + else fused_unpad_token_groups_cuda ) + unpadded_tokens = kernel(*call_args) - # Verify outputs match - assert torch.allclose( - ref_unpadded_tokens, kernel_unpadded_tokens, rtol=0, atol=1e-5 - ), "Unpadded tokens do not match" - - # Verify that unpad correctly reverses pad operation - assert torch.allclose(inputs, kernel_unpadded_tokens, rtol=0, atol=1e-5), ( - "Unpadded tokens should match original inputs" - ) + if impl == "cutedsl": + assert torch.equal(ref_unpadded_tokens, unpadded_tokens) + # unpad must exactly reverse pad + assert torch.equal(inputs, unpadded_tokens) + else: + assert torch.allclose( + ref_unpadded_tokens, unpadded_tokens, rtol=0, atol=1e-5 + ), "Unpadded tokens do not match" + assert torch.allclose(inputs, unpadded_tokens, rtol=0, atol=1e-5), ( + "unpad did not reverse pad" + ) @pytest.mark.parametrize("round_scales_to_power_of_2", [True, False]) diff --git a/torchao/prototype/moe_training/kernels/mxfp8/__init__.py b/torchao/prototype/moe_training/kernels/mxfp8/__init__.py index 0ed217b43a4..ceb0c1459e2 100644 --- a/torchao/prototype/moe_training/kernels/mxfp8/__init__.py +++ b/torchao/prototype/moe_training/kernels/mxfp8/__init__.py @@ -1,3 +1,7 @@ +from torchao.prototype.moe_training.kernels.mxfp8.cutedsl_pad_token_groups import ( # noqa: F401 + _pad_token_groups_cutedsl, + _unpad_token_groups_cutedsl, +) from torchao.prototype.moe_training.kernels.mxfp8.quant import ( _mxfp8_cuda_kernels_available, # noqa: F401 _mxfp8_flydsl_kernels_available, # noqa: F401 diff --git a/torchao/prototype/moe_training/kernels/mxfp8/cutedsl_pad_token_groups.py b/torchao/prototype/moe_training/kernels/mxfp8/cutedsl_pad_token_groups.py new file mode 100644 index 00000000000..ba4612c0a00 --- /dev/null +++ b/torchao/prototype/moe_training/kernels/mxfp8/cutedsl_pad_token_groups.py @@ -0,0 +1,414 @@ +# 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 functools + +import torch + +from torchao.utils import ceil_div, round_up + +from .cute_utils import current_cuda_stream + + +@functools.cache +def _compile_copy_token_groups_cutedsl( + num_groups: int, + element_size: int, + is_pad: bool, + row_bytes_aligned: bool, + alignment_size: int, +): + import cuda.bindings.driver as cuda + import cutlass + import cutlass.cute as cute + from cutlass.cute.runtime import make_fake_stream, make_fake_tensor + + NUM_GROUPS = num_groups + ELEMENT_SIZE = element_size + IS_PAD = is_pad + ROW_BYTES_ALIGNED = row_bytes_aligned + ALIGNMENT_SIZE = alignment_size + VEC_ELEMS = 16 // ELEMENT_SIZE + WARP_SIZE = 32 + WARPS_PER_BLOCK = 8 + THREADS_PER_BLOCK = WARP_SIZE * WARPS_PER_BLOCK + UNROLL = 4 + LANE_STRIDE = WARP_SIZE * VEC_ELEMS + UNROLL_STRIDE = UNROLL * LANE_STRIDE + ELEMENT_DTYPE = cutlass.BFloat16 if ELEMENT_SIZE == 2 else cutlass.Float32 + + class CopyTokenGroups: + @cute.jit + def _find_group( + self, + row: cutlass.Int32, + group_end_offsets: cute.Tensor, + ): + group = cutlass.Int32(0) + keep_looking = cutlass.Int32(1) + while group < NUM_GROUPS - 1 and keep_looking != 0: + end = cutlass.Int32((group_end_offsets.iterator + group).load()) + if row >= end: + group = group + cutlass.Int32(1) + else: + keep_looking = cutlass.Int32(0) + return group + + @cute.jit + def _group_start( + self, + group: cutlass.Int32, + group_end_offsets: cute.Tensor, + ): + group_start = cutlass.Int32(0) + if group > 0: + group_start = cutlass.Int32( + (group_end_offsets.iterator + group - 1).load() + ) + return group_start + + @cute.jit + def _padded_start( + self, + group: cutlass.Int32, + group_end_offsets: cute.Tensor, + ): + cumulative = cutlass.Int32(0) + previous = cutlass.Int32(0) + for g in range(NUM_GROUPS): + group_end = cutlass.Int32((group_end_offsets.iterator + g).load()) + group_size = group_end - previous + padded_size = round_up(group_size, ALIGNMENT_SIZE) + if g < group: + cumulative = cumulative + padded_size + previous = group_end + return cumulative + + @cute.jit + def _fill_padded_offsets( + self, + tidx: cutlass.Int32, + group_end_offsets: cute.Tensor, + padded_group_start_offsets: cute.Tensor, + padded_group_end_offsets: cute.Tensor, + ): + if tidx < NUM_GROUPS: + group = cutlass.Int32(tidx) + group_start = self._group_start(group, group_end_offsets) + group_end = cutlass.Int32((group_end_offsets.iterator + group).load()) + group_size = group_end - group_start + padded_size = round_up(group_size, ALIGNMENT_SIZE) + padded_start = self._padded_start(group, group_end_offsets) + (padded_group_start_offsets.iterator + group).store(padded_start) + (padded_group_end_offsets.iterator + group).store( + padded_start + padded_size + ) + + @cute.jit + def _copy_vec_or_tail( + self, + inputs: cute.Tensor, + input_offset: cutlass.Int64, + elem_col: cutlass.Int32, + output: cute.Tensor, + output_offset: cutlass.Int64, + dim: cutlass.Int32, + ): + if cutlass.const_expr(ROW_BYTES_ALIGNED): + value = cute.recast_ptr( + (inputs.iterator + input_offset).align(16), + dtype=cutlass.Int128, + ).load() + cute.recast_ptr( + (output.iterator + output_offset).align(16), + dtype=cutlass.Int128, + ).store(value) + else: + for i in range(VEC_ELEMS): + if elem_col + i < dim: + value = (inputs.iterator + input_offset + i).load() + (output.iterator + output_offset + i).store(value) + + @cute.kernel + def kernel( + self, + inputs: cute.Tensor, + group_end_offsets: cute.Tensor, + padded_group_start_offsets: cute.Tensor, + padded_group_end_offsets: cute.Tensor, + output: cute.Tensor, + num_tokens: cutlass.Int32, + dim: cutlass.Int32, + ): + tidx, _, _ = cute.arch.thread_idx() + block_idx, _, _ = cute.arch.block_idx() + lane_id = cutlass.Int32(tidx % WARP_SIZE) + warp_id = cutlass.Int32(tidx // WARP_SIZE) + if cutlass.const_expr(IS_PAD): + if block_idx == 0: + self._fill_padded_offsets( + tidx, + group_end_offsets, + padded_group_start_offsets, + padded_group_end_offsets, + ) + row = cutlass.Int32(block_idx * WARPS_PER_BLOCK + warp_id) + if row < num_tokens: + group_start = cutlass.Int32(0) + padded_start = cutlass.Int32(0) + if lane_id == 0: + if cutlass.const_expr(IS_PAD): + group = cutlass.Int32(0) + previous = cutlass.Int32(0) + cumulative = cutlass.Int32(0) + keep_looking = cutlass.Int32(1) + while group < NUM_GROUPS and keep_looking != 0: + group_end = cutlass.Int32( + (group_end_offsets.iterator + group).load() + ) + group_size = group_end - previous + padded_size = round_up(group_size, ALIGNMENT_SIZE) + if row < group_end: + group_start = previous + padded_start = cumulative + keep_looking = cutlass.Int32(0) + else: + cumulative = cumulative + padded_size + previous = group_end + group = group + cutlass.Int32(1) + else: + group = self._find_group(row, group_end_offsets) + group_start = self._group_start(group, group_end_offsets) + padded_start = cutlass.Int32( + (padded_group_start_offsets.iterator + group).load() + ) + group_start = cute.arch.shuffle_sync(group_start, 0) + padded_start = cute.arch.shuffle_sync(padded_start, 0) + offset_in_group = row - group_start + padded_row = padded_start + offset_in_group + input_row = row + output_row = padded_row + if cutlass.const_expr(not IS_PAD): + input_row = padded_row + output_row = row + elem_col = lane_id * VEC_ELEMS + # row*dim overflows Int32 once num_tokens*dim exceeds ~2**31. + input_offset = cutlass.Int64(input_row) * dim + elem_col + output_offset = cutlass.Int64(output_row) * dim + elem_col + if cutlass.const_expr(ROW_BYTES_ALIGNED): + while elem_col + (UNROLL - 1) * LANE_STRIDE + VEC_ELEMS <= dim: + vals = [] + for k in cutlass.range_constexpr(UNROLL): + vals.append( + cute.recast_ptr( + ( + inputs.iterator + input_offset + k * LANE_STRIDE + ).align(16), + dtype=cutlass.Int128, + ).load() + ) + for k in cutlass.range_constexpr(UNROLL): + cute.recast_ptr( + ( + output.iterator + output_offset + k * LANE_STRIDE + ).align(16), + dtype=cutlass.Int128, + ).store(vals[k]) + elem_col = elem_col + UNROLL_STRIDE + input_offset = input_offset + UNROLL_STRIDE + output_offset = output_offset + UNROLL_STRIDE + while elem_col < dim: + self._copy_vec_or_tail( + inputs, + input_offset, + elem_col, + output, + output_offset, + dim, + ) + elem_col = elem_col + WARP_SIZE * VEC_ELEMS + input_offset = input_offset + WARP_SIZE * VEC_ELEMS + output_offset = output_offset + WARP_SIZE * VEC_ELEMS + + @cute.jit + def __call__( + self, + inputs: cute.Tensor, + group_end_offsets: cute.Tensor, + padded_group_start_offsets: cute.Tensor, + padded_group_end_offsets: cute.Tensor, + output: cute.Tensor, + num_tokens: cutlass.Int32, + dim: cutlass.Int32, + num_blocks: cutlass.Int32, + stream: cuda.CUstream, + ): + self.kernel( + inputs, + group_end_offsets, + padded_group_start_offsets, + padded_group_end_offsets, + output, + num_tokens, + dim, + ).launch( + grid=(num_blocks, 1, 1), + block=(THREADS_PER_BLOCK, 1, 1), + cluster=(1, 1, 1), + stream=stream, + ) + + kernel = CopyTokenGroups() + + n = cute.sym_int() + d = cute.sym_int() + out_n = cute.sym_int() + offs_stride = cute.sym_int() + assumed_align = 16 if ROW_BYTES_ALIGNED else 1 + fake_input = make_fake_tensor( + ELEMENT_DTYPE, + (n, d), + stride=(d, 1), + assumed_align=assumed_align, + ) + fake_output = make_fake_tensor( + ELEMENT_DTYPE, + (out_n, d), + stride=(d, 1), + assumed_align=assumed_align, + ) + fake_group_end_offsets = make_fake_tensor( + cutlass.Int32, + (NUM_GROUPS,), + stride=(offs_stride,), + ) + fake_padded_group_start_offsets = make_fake_tensor( + cutlass.Int32, + (NUM_GROUPS,), + stride=(offs_stride,), + ) + fake_padded_group_end_offsets = make_fake_tensor( + cutlass.Int32, + (NUM_GROUPS,), + stride=(offs_stride,), + ) + fake_stream = make_fake_stream() + + return cute.compile( + kernel, + inputs=fake_input, + group_end_offsets=fake_group_end_offsets, + padded_group_start_offsets=fake_padded_group_start_offsets, + padded_group_end_offsets=fake_padded_group_end_offsets, + output=fake_output, + num_tokens=0, + dim=0, + num_blocks=1, + stream=fake_stream, + options="--enable-tvm-ffi", + ) + + +def _pad_token_groups_cutedsl( + inputs: torch.Tensor, + group_end_offsets: torch.Tensor, + alignment_size: int = 32, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + assert inputs.ndim == 2 + assert inputs.is_cuda + assert inputs.is_contiguous() + assert inputs.dtype in (torch.float32, torch.bfloat16) + assert group_end_offsets.ndim == 1 + assert group_end_offsets.is_cuda + assert group_end_offsets.dtype == torch.int32 + assert alignment_size in (32, 128) + num_tokens, dim = inputs.shape + num_groups = group_end_offsets.shape[0] + assert num_groups <= 32 + padded_group_start_offsets = torch.empty( + (num_groups,), + dtype=group_end_offsets.dtype, + device=group_end_offsets.device, + ) + padded_group_end_offsets = torch.empty( + (num_groups,), + dtype=group_end_offsets.dtype, + device=group_end_offsets.device, + ) + output_rows = round_up(num_tokens + num_groups * alignment_size, alignment_size) + output = torch.zeros( + (output_rows, dim), + dtype=inputs.dtype, + device=inputs.device, + ) + compiled = _compile_copy_token_groups_cutedsl( + num_groups, + inputs.element_size(), + True, + dim * inputs.element_size() % 16 == 0, + alignment_size, + ) + stream = current_cuda_stream() + compiled( + inputs, + group_end_offsets, + padded_group_start_offsets, + padded_group_end_offsets, + output, + int(num_tokens), + int(dim), + int(ceil_div(num_tokens, 8)), + stream, + ) + return output, padded_group_start_offsets, padded_group_end_offsets + + +def _unpad_token_groups_cutedsl( + padded_inputs: torch.Tensor, + group_end_offsets: torch.Tensor, + padded_group_start_offsets: torch.Tensor, + num_tokens: int, + alignment_size: int = 32, +) -> torch.Tensor: + assert padded_inputs.ndim == 2 + assert padded_inputs.is_cuda + assert padded_inputs.is_contiguous() + assert padded_inputs.dtype in (torch.float32, torch.bfloat16) + assert group_end_offsets.ndim == 1 + assert group_end_offsets.is_cuda + assert group_end_offsets.dtype == torch.int32 + assert padded_group_start_offsets.ndim == 1 + assert padded_group_start_offsets.is_cuda + assert padded_group_start_offsets.dtype == torch.int32 + assert alignment_size in (32, 128) + dim = padded_inputs.shape[1] + num_groups = group_end_offsets.shape[0] + assert num_groups <= 32 + output = torch.empty( + (num_tokens, dim), + dtype=padded_inputs.dtype, + device=padded_inputs.device, + ) + compiled = _compile_copy_token_groups_cutedsl( + num_groups, + padded_inputs.element_size(), + False, + dim * padded_inputs.element_size() % 16 == 0, + alignment_size, + ) + stream = current_cuda_stream() + compiled( + padded_inputs, + group_end_offsets, + padded_group_start_offsets, + padded_group_start_offsets, + output, + int(num_tokens), + int(dim), + int(ceil_div(num_tokens, 8)), + stream, + ) + return output