|
| 1 | +"""Benchmark softmax_online op against torch.softmax and torch.compile(torch.softmax).""" |
| 2 | + |
| 3 | +import argparse |
| 4 | +import os |
| 5 | + |
| 6 | +import torch |
| 7 | + |
| 8 | +from forge_cute_py.ops.softmax_online import softmax_fwd, softmax_bwd |
| 9 | +from forge_cute_py.util.bench import do_bench, estimate_bandwidth, summarize_times |
| 10 | + |
| 11 | +SHORT_M = [128, 512, 2048, 8192] |
| 12 | +SHORT_N = [1024, 2048, 4096, 8192] |
| 13 | + |
| 14 | +LONG_M = [64, 128, 256] |
| 15 | +LONG_N = [16384, 32768, 65536, 131072] |
| 16 | + |
| 17 | +DEFAULT_DTYPES = ["float16", "bfloat16", "float32"] |
| 18 | + |
| 19 | + |
| 20 | +def parse_int_list(s: str) -> list[int]: |
| 21 | + return [int(x.strip()) for x in s.split(",")] |
| 22 | + |
| 23 | + |
| 24 | +def parse_str_list(s: str) -> list[str]: |
| 25 | + return [x.strip() for x in s.split(",")] |
| 26 | + |
| 27 | + |
| 28 | +def main(): |
| 29 | + parser = argparse.ArgumentParser(description="Benchmark softmax_online op") |
| 30 | + parser.add_argument( |
| 31 | + "--long", action="store_true", help="Use long-N benchmark suite (small M, large N)" |
| 32 | + ) |
| 33 | + parser.add_argument("--m-sizes", type=parse_int_list, default=None) |
| 34 | + parser.add_argument("--n-sizes", type=parse_int_list, default=None) |
| 35 | + parser.add_argument("--dtypes", type=parse_str_list, default=DEFAULT_DTYPES) |
| 36 | + parser.add_argument("--warmup", type=int, default=20) |
| 37 | + parser.add_argument("--iterations", type=int, default=100) |
| 38 | + parser.add_argument( |
| 39 | + "--impl", |
| 40 | + choices=["auto", "ref", "kernel"], |
| 41 | + default="auto", |
| 42 | + help="softmax_online backend mode (FORGE_SOFTMAX_IMPL)", |
| 43 | + ) |
| 44 | + args = parser.parse_args() |
| 45 | + os.environ["FORGE_SOFTMAX_IMPL"] = args.impl |
| 46 | + |
| 47 | + if args.m_sizes is None: |
| 48 | + args.m_sizes = LONG_M if args.long else SHORT_M |
| 49 | + if args.n_sizes is None: |
| 50 | + args.n_sizes = LONG_N if args.long else SHORT_N |
| 51 | + |
| 52 | + if not torch.cuda.is_available(): |
| 53 | + raise RuntimeError("CUDA required for benchmarking") |
| 54 | + |
| 55 | + gpu_name = torch.cuda.get_device_name(0) |
| 56 | + suite = "long" if args.long else "short" |
| 57 | + print(f"softmax_online benchmarks [{suite}] ({gpu_name}) [impl={args.impl}]") |
| 58 | + print() |
| 59 | + |
| 60 | + header = ( |
| 61 | + f"{'M':>6} {'N':>6} {'Dtype':<10} {'Op':<18} {'Pass':<5} " |
| 62 | + f"{'p50 (ms)':>10} {'BW (GB/s)':>10} {'vs torch':>10}" |
| 63 | + ) |
| 64 | + print(header) |
| 65 | + print("-" * len(header)) |
| 66 | + |
| 67 | + for m in args.m_sizes: |
| 68 | + for n in args.n_sizes: |
| 69 | + for dtype_str in args.dtypes: |
| 70 | + dtype = getattr(torch, dtype_str) |
| 71 | + x = torch.randn(m, n, device="cuda", dtype=dtype) |
| 72 | + elem = x.element_size() |
| 73 | + |
| 74 | + # --- Forward bandwidth: read input + write output --- |
| 75 | + fwd_bytes = 2 * m * n * elem |
| 76 | + |
| 77 | + # --- torch.softmax fwd baseline --- |
| 78 | + torch_fn = lambda: torch.softmax(x, dim=-1) |
| 79 | + torch_times = do_bench(torch_fn, warmup=args.warmup, rep=args.iterations) |
| 80 | + torch_stats = summarize_times(torch_times) |
| 81 | + torch_fwd_p50 = torch_stats["p50_ms"] |
| 82 | + torch_fwd_bw = estimate_bandwidth(fwd_bytes, torch_fwd_p50) |
| 83 | + print( |
| 84 | + f"{m:>6} {n:>6} {dtype_str:<10} {'torch.softmax':<18} {'fwd':<5} " |
| 85 | + f"{torch_fwd_p50:>10.4f} {torch_fwd_bw:>10.2f} {1.0:>10.2f}x" |
| 86 | + ) |
| 87 | + |
| 88 | + # --- torch.compile fwd --- |
| 89 | + try: |
| 90 | + compiled_ref = torch.compile(lambda t: torch.softmax(t, dim=-1)) |
| 91 | + compiled_ref(x) |
| 92 | + fn = lambda: compiled_ref(x) |
| 93 | + compiled_times = do_bench(fn, warmup=args.warmup, rep=args.iterations) |
| 94 | + compiled_stats = summarize_times(compiled_times) |
| 95 | + compiled_p50 = compiled_stats["p50_ms"] |
| 96 | + compiled_bw = estimate_bandwidth(fwd_bytes, compiled_p50) |
| 97 | + ratio = compiled_p50 / torch_fwd_p50 if torch_fwd_p50 > 0 else float("inf") |
| 98 | + print( |
| 99 | + f"{m:>6} {n:>6} {dtype_str:<10} {'torch.compile':<18} {'fwd':<5} " |
| 100 | + f"{compiled_p50:>10.4f} {compiled_bw:>10.2f} {ratio:>10.2f}x" |
| 101 | + ) |
| 102 | + except Exception as e: |
| 103 | + print( |
| 104 | + f"{m:>6} {n:>6} {dtype_str:<10} {'torch.compile':<18} {'fwd':<5} " |
| 105 | + f"{'ERROR':>10} {'':>10} {'':>10} {e}" |
| 106 | + ) |
| 107 | + |
| 108 | + # --- softmax_online fwd --- |
| 109 | + try: |
| 110 | + softmax_fwd(x, dim=-1) |
| 111 | + fn = lambda: softmax_fwd(x, dim=-1) |
| 112 | + times = do_bench(fn, warmup=args.warmup, rep=args.iterations) |
| 113 | + stats = summarize_times(times) |
| 114 | + p50 = stats["p50_ms"] |
| 115 | + bw = estimate_bandwidth(fwd_bytes, p50) |
| 116 | + ratio = p50 / torch_fwd_p50 if torch_fwd_p50 > 0 else float("inf") |
| 117 | + print( |
| 118 | + f"{m:>6} {n:>6} {dtype_str:<10} {'softmax_online':<18} {'fwd':<5} " |
| 119 | + f"{p50:>10.4f} {bw:>10.2f} {ratio:>10.2f}x" |
| 120 | + ) |
| 121 | + except Exception as e: |
| 122 | + print( |
| 123 | + f"{m:>6} {n:>6} {dtype_str:<10} {'softmax_online':<18} {'fwd':<5} " |
| 124 | + f"{'ERROR':>10} {'':>10} {'':>10} {e}" |
| 125 | + ) |
| 126 | + |
| 127 | + # --- Backward pass benchmarks --- |
| 128 | + # Pre-compute softmax output y and fake upstream gradient dy |
| 129 | + y = torch.softmax(x, dim=-1) |
| 130 | + dy = torch.randn_like(y) |
| 131 | + |
| 132 | + # Backward bandwidth: read dy + read y + write dx = 3 * M * N * elem |
| 133 | + bwd_bytes = 3 * m * n * elem |
| 134 | + |
| 135 | + # --- torch backward baseline --- |
| 136 | + torch_bwd_fn = lambda: torch._softmax_backward_data(dy, y, -1, x.dtype) |
| 137 | + torch_bwd_times = do_bench(torch_bwd_fn, warmup=args.warmup, rep=args.iterations) |
| 138 | + torch_bwd_stats = summarize_times(torch_bwd_times) |
| 139 | + torch_bwd_p50 = torch_bwd_stats["p50_ms"] |
| 140 | + torch_bwd_bw = estimate_bandwidth(bwd_bytes, torch_bwd_p50) |
| 141 | + print( |
| 142 | + f"{m:>6} {n:>6} {dtype_str:<10} {'torch.softmax':<18} {'bwd':<5} " |
| 143 | + f"{torch_bwd_p50:>10.4f} {torch_bwd_bw:>10.2f} {1.0:>10.2f}x" |
| 144 | + ) |
| 145 | + |
| 146 | + # --- softmax_online bwd --- |
| 147 | + try: |
| 148 | + y_ours = softmax_fwd(x, dim=-1) |
| 149 | + softmax_bwd(dy, y_ours, dim=-1) |
| 150 | + fn = lambda: softmax_bwd(dy, y_ours, dim=-1) |
| 151 | + times = do_bench(fn, warmup=args.warmup, rep=args.iterations) |
| 152 | + stats = summarize_times(times) |
| 153 | + p50 = stats["p50_ms"] |
| 154 | + bw = estimate_bandwidth(bwd_bytes, p50) |
| 155 | + ratio = p50 / torch_bwd_p50 if torch_bwd_p50 > 0 else float("inf") |
| 156 | + print( |
| 157 | + f"{m:>6} {n:>6} {dtype_str:<10} {'softmax_online':<18} {'bwd':<5} " |
| 158 | + f"{p50:>10.4f} {bw:>10.2f} {ratio:>10.2f}x" |
| 159 | + ) |
| 160 | + except Exception as e: |
| 161 | + print( |
| 162 | + f"{m:>6} {n:>6} {dtype_str:<10} {'softmax_online':<18} {'bwd':<5} " |
| 163 | + f"{'ERROR':>10} {'':>10} {'':>10} {e}" |
| 164 | + ) |
| 165 | + |
| 166 | + print() |
| 167 | + |
| 168 | + |
| 169 | +if __name__ == "__main__": |
| 170 | + main() |
0 commit comments