import argparse import glob import importlib import itertools import os import torch from common.bench_framework import (make_bwd_benchmark_for_case, make_bwd_benchmark_plot_for_case, make_fwd_benchmark_for_case, make_fwd_benchmark_plot_for_case) from common.diff_engine import DiffCase, calculate_diff def _clean_and_print_csv(csv_path, title): """Remove trailing ' ()' from CSV column names and print the table.""" import pandas as pd df = pd.read_csv(csv_path) df.columns = [c.replace(" ()", "") for c in df.columns] df.to_csv(csv_path, index=False) print(f"{title}:") print(df.to_string(index=False)) print() def make_title_tag(): if torch.cuda.is_available(): dev_name = torch.cuda.get_device_name(0) else: dev_name = "CPU" torch_ver = torch.__version__ return f"[{dev_name} | torch {torch_ver}]" def plot_result(r_path, columns=None): import matplotlib.pyplot as plt import pandas as pd df = pd.read_csv(r_path + ".csv") if columns is None: columns = [ c for c in ["Naive", "Compiled", "Cuda", "Triton"] if c in df.columns ] plt.figure(figsize=(12, 6)) ax = df.plot(x="config", y=columns, kind="bar", ax=plt.gca()) ax.set_title("Speedup over torch (higher is better)\n" + make_title_tag(), fontsize=14, fontweight="bold") ax.set_ylabel("Relative Speedup", fontsize=14) ax.set_xlabel("") plt.xticks(rotation=45, fontsize=12, ha="right", rotation_mode="anchor") for container in ax.containers: labels = [f"x{v.get_height():.2f}" for v in container] ax.bar_label(container, labels=labels, label_type="edge", fontsize=10) plt.tight_layout() plt.savefig(r_path + ".png", bbox_inches="tight") def main(): ap = argparse.ArgumentParser() ap.add_argument( "--case", choices=["rms", "add_rms", "poly", "mul_poly", "grouped_mul_poly", "mla_rope"], required=True) ap.add_argument("--plot", action="store_true") ap.add_argument( "--save-path", type=str, default="./configs/", help="Path to save benchmark results", ) ap.add_argument( "--dtype", choices=["fp16", "bf16", "fp32", "all"], default="bf16", help="Data type for benchmarking (default: bf16)", ) args = ap.parse_args() dtype_map = { "fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32, } if args.dtype == "all": dtypes = [("fp16", torch.float16), ("bf16", torch.bfloat16), ("fp32", torch.float32)] else: dtypes = [(args.dtype, dtype_map[args.dtype])] torch.set_default_device("cuda") mod = importlib.import_module(f"cases.{args.case}") case: DiffCase = mod.CASE # Correctness checks across multiple configs # NOTE: calculate_diff positionally calls build_inputs(hidden_size, bs, sl); # bench framework positionally calls build_inputs(bs, sl, dim). These # disagree — rms-style cases don't care (all 3 axes are flat dims), but # mla_rope does. We match the bench convention in cases/mla_rope.py, so # we swap arg names at the correctness call site below for that case. if args.case == "mla_rope": cfgs = [(1, 1024, 0), (4, 4096, 0), (8, 4096, 0)] # (bs, sl, dummy) else: cfgs = [(2, 128, 4096), (8, 4096, 1280), (1, 32768, 1280)] for bs, sl, hid in cfgs: print( f"Checking correctness: bs={bs}, sl={sl}, D={hid} " f"(N={bs*sl})...", end=" ") if args.case == "mla_rope": # Swap so positional (hidden_size, batch_size, seq_len) maps to # our build_inputs(bs, sl, dim) as (bs, sl, dummy). calculate_diff(case, batch_size=sl, seq_len=hid, hidden_size=bs) else: calculate_diff(case, batch_size=bs, seq_len=sl, hidden_size=hid) print("✅") for dtype_name, dtype in dtypes: print(f"\n{'=' * 60}") print(f" Benchmarking dtype: {dtype_name} ({dtype})") print(f"{'=' * 60}\n") save_dir = os.path.join(args.save_path, args.case, dtype_name) os.makedirs(save_dir, exist_ok=True) is_grouped = args.case == "grouped_mul_poly" if args.plot: batch_size_range = [1] seq_length_range = [4096, 8192, 16384] if is_grouped: dim = [1280] elif "poly" in args.case: dim = [8192, 16384] elif args.case == "mla_rope": dim = [0] # MLA head dims are fixed; dim axis is a dummy else: dim = [2048, 4096] configs = list( itertools.product(batch_size_range, seq_length_range, dim)) if is_grouped: plot_line_vals = ("naive", "compiled", "cuda") plot_line_names = { "naive": "Naive", "compiled": "Compiled", "cuda": "Triton", } else: plot_line_vals = ("naive", "cuda") plot_line_names = { "naive": "Naive", "cuda": "Cuda", } plot_name = f"plot_{args.case}-{dtype_name}-fwd-perf" bench = make_fwd_benchmark_plot_for_case( case=case, configs=configs, plot_name=plot_name, dtype=dtype, line_vals=plot_line_vals, line_names=plot_line_names, ) bench.run(print_data=True, save_path=save_dir) plot_result(os.path.join(save_dir, plot_name)) plot_name = f"plot_{args.case}-{dtype_name}-bwd-perf" bench = make_bwd_benchmark_plot_for_case( case=case, configs=configs, plot_name=plot_name, dtype=dtype, line_vals=plot_line_vals, line_names=plot_line_names, ) bench.run(print_data=True, save_path=save_dir) plot_result(os.path.join(save_dir, plot_name)) for f in glob.glob(os.path.join(save_dir, "*.html")) + \ glob.glob(os.path.join(save_dir, "*.csv")): os.remove(f) else: batch_size_range = [2**i for i in range(0, 4, 1)] seq_length_range = [2**i for i in range(10, 14, 1)] if is_grouped: dim = [1280] elif "poly" in args.case: dim = [8192, 16384] elif args.case == "mla_rope": dim = [0] # MLA head dims are fixed; dim axis is a dummy else: dim = [2048, 4096] configs = list( itertools.product(dim, batch_size_range, seq_length_range)) if is_grouped: fwd_line_vals = ("naive", "naive_bw", "compiled", "compiled_bw", "cuda", "cuda_bw", "speedup") fwd_line_names = { "naive": "Naive (us)", "naive_bw": "Naive (GB/s)", "compiled": "Compiled (us)", "compiled_bw": "Compiled (GB/s)", "cuda": "CUDA (us)", "cuda_bw": "CUDA (GB/s)", "speedup": "SpeedUp (ratio)", } bwd_line_vals = ("naive", "naive_bw", "compiled", "compiled_bw", "compiled_cuda", "compiled_cuda_bw", "speedup") bwd_line_names = { "naive": "Naive (us)", "naive_bw": "Naive (GB/s)", "compiled": "Compiled (us)", "compiled_bw": "Compiled (GB/s)", "compiled_cuda": "CompiledCUDA (us)", "compiled_cuda_bw": "CompiledCUDA (GB/s)", "speedup": "SpeedUp (ratio)", } else: fwd_line_vals = ("naive", "naive_bw", "cuda", "cuda_bw", "speedup") fwd_line_names = { "naive": "Naive (us)", "naive_bw": "Naive (GB/s)", "cuda": "CUDA (us)", "cuda_bw": "CUDA (GB/s)", "speedup": "SpeedUp (ratio)", } bwd_line_vals = fwd_line_vals bwd_line_names = fwd_line_names bench = make_fwd_benchmark_for_case( case=case, configs=configs, plot_name=f"{args.case}-{dtype_name}-fwd-perf", dtype=dtype, line_vals=fwd_line_vals, line_names=fwd_line_names, ) fwd_name = f"{args.case}-{dtype_name}-fwd-perf" bench.run(print_data=False, save_path=save_dir) _clean_and_print_csv(os.path.join(save_dir, fwd_name + ".csv"), fwd_name) bench = make_bwd_benchmark_for_case( case=case, configs=configs, plot_name=f"{args.case}-{dtype_name}-bwd-perf", dtype=dtype, line_vals=bwd_line_vals, line_names=bwd_line_names, ) bwd_name = f"{args.case}-{dtype_name}-bwd-perf" bench.run(print_data=False, save_path=save_dir) _clean_and_print_csv(os.path.join(save_dir, bwd_name + ".csv"), bwd_name) for f in glob.glob(os.path.join(save_dir, "*.html")) + \ glob.glob(os.path.join(save_dir, "*.png")): os.remove(f) if __name__ == "__main__": main()