| import collections |
| import math |
| import re |
| from typing import Any, Dict, Sequence |
|
|
| import torch |
| import triton |
|
|
| from .diff_engine import DiffCase |
|
|
|
|
| def make_fwd_key(batch_size, seq_len, dim): |
| return f"forward : ({batch_size}, {seq_len}, {dim})" |
|
|
|
|
| def make_bwd_key(batch_size, seq_len, dim): |
| return f"backward : ({batch_size}, {seq_len}, {dim})" |
|
|
|
|
| def parse_config_string(config_str): |
| match = re.match(r"(\w+)\s*:\s*\(\s*(\d+)\s*,\s*(\d+)\s*,\s*(\d+)\s*\)", |
| config_str) |
| if not match: |
| raise ValueError(f"Invalid config string: {config_str}") |
| _, bs, sl, d = match.groups() |
| return int(bs), int(sl), int(d) |
|
|
|
|
| def make_fwd_benchmark_for_case( |
| *, |
| case: DiffCase, |
| configs: Sequence[tuple[int, int, int]], |
| plot_name: str, |
| ylabel: str = "us", |
| line_vals=("naive", "cuda", "speedup"), |
| line_names: Dict[str, str] | None = None, |
| dtype=torch.bfloat16, |
| eps: float = 1e-6, |
| time_unit_scale: float = 1000, |
| ): |
| timings_ms = collections.defaultdict(dict) |
| line_vals = list(line_vals) |
| line_names = line_names or {v: v.title() for v in line_vals} |
| x_vals = [list(_) for _ in configs] |
|
|
| @triton.testing.perf_report( |
| triton.testing.Benchmark(x_names=["dim", "batch_size", "seq_len"], |
| x_vals=x_vals, |
| line_arg="provider", |
| line_vals=line_vals, |
| line_names=[line_names[v] for v in line_vals], |
| ylabel=ylabel, |
| plot_name=plot_name, |
| args={})) |
| def bench(dim, batch_size, seq_len, provider): |
| key = make_fwd_key(dim, batch_size, seq_len) |
| I = case.build_inputs(batch_size, seq_len, dim, dtype, eps) |
| if provider == "speedup": |
| return timings_ms["naive"][key] / timings_ms["cuda"][key] |
| obj = case.make_naive(I) if provider == "naive" else case.make_cuda(I) |
| run = lambda: case.forward(obj, I) |
| ms = triton.testing.do_bench(run) |
| timings_ms[provider][key] = ms |
| return time_unit_scale * ms |
|
|
| return bench |
|
|
|
|
| def make_fwd_benchmark_plot_for_case( |
| *, |
| case: DiffCase, |
| configs: Sequence[tuple[int, int, int]], |
| plot_name: str, |
| ylabel: str = "Relative Speedup", |
| line_vals=("naive", "cuda"), |
| line_names: Dict[str, str] | None = None, |
| dtype=torch.bfloat16, |
| eps: float = 1e-6, |
| ): |
| timings_ms = collections.defaultdict(dict) |
| spdup_ratio = list() |
| line_vals = list(line_vals) |
| line_names = line_names or {v: v.title() for v in line_vals} |
| x_vals = [make_fwd_key(*_) for _ in configs] |
| x_vals.append("Geometric Mean") |
|
|
| @triton.testing.perf_report( |
| triton.testing.Benchmark(x_names=["config"], |
| x_vals=x_vals, |
| line_arg="provider", |
| line_vals=line_vals, |
| line_names=[line_names[v] for v in line_vals], |
| ylabel=ylabel, |
| plot_name=plot_name, |
| args={})) |
| def bench(config, provider): |
| if config == "Geometric Mean": |
| if provider == "cuda": |
| return round(math.prod(spdup_ratio)**(1 / len(spdup_ratio)), 2) |
| else: |
| return 1.00 |
| batch_size, seq_len, dim = parse_config_string(config) |
| I = case.build_inputs(batch_size, seq_len, dim, dtype, eps) |
| obj = case.make_naive(I) if provider == "naive" else case.make_cuda(I) |
| run = lambda: case.forward(obj, I) |
| ms = triton.testing.do_bench(run) |
| timings_ms[provider][config] = ms |
| if provider == "cuda": |
| ratio = timings_ms["naive"][config] / timings_ms["cuda"][config] |
| spdup_ratio.append(ratio) |
| return round(ratio, 2) |
| else: |
| return 1.00 |
|
|
| return bench |
|
|
|
|
| def make_bwd_benchmark_for_case( |
| *, |
| case: DiffCase, |
| configs: Sequence[tuple[int, int, int]], |
| plot_name: str, |
| ylabel: str = "us", |
| line_vals=("naive", "cuda", "speedup"), |
| line_names: Dict[str, str] | None = None, |
| dtype=torch.bfloat16, |
| eps: float = 1e-6, |
| time_unit_scale: float = 1000, |
| ): |
| timings_ms = collections.defaultdict(dict) |
| line_vals = list(line_vals) |
| line_names = line_names or {v: v.title() for v in line_vals} |
| x_vals = [list(_) for _ in configs] |
|
|
| @triton.testing.perf_report( |
| triton.testing.Benchmark(x_names=["dim", "batch_size", "seq_len"], |
| x_vals=x_vals, |
| line_arg="provider", |
| line_vals=line_vals, |
| line_names=[line_names[v] for v in line_vals], |
| ylabel=ylabel, |
| plot_name=plot_name, |
| args={})) |
| def bench(dim, batch_size, seq_len, provider): |
| key = make_bwd_key(dim, batch_size, seq_len) |
| I = case.build_inputs(batch_size, seq_len, dim, dtype, eps) |
| if provider == "speedup": |
| return timings_ms["naive"][key] / timings_ms["cuda"][key] |
| obj = case.make_naive(I) if provider == "naive" else case.make_cuda(I) |
| y = case.forward(obj, I) |
| gin = list(case.grad_inputs(I)) + list(obj.parameters()) |
| if isinstance(y, torch.Tensor): |
| g = [torch.randn_like(y)] |
| else: |
| g = [torch.randn_like(r) for r in y] |
| run = lambda: torch.autograd.grad(y, |
| gin, |
| g, |
| retain_graph=True, |
| create_graph=False, |
| allow_unused=False) |
| ms = triton.testing.do_bench(run) |
| timings_ms[provider][key] = ms |
| return time_unit_scale * ms |
|
|
| return bench |
|
|
|
|
| def make_bwd_benchmark_plot_for_case( |
| *, |
| case: DiffCase, |
| configs: Sequence[tuple[int, int, int]], |
| plot_name: str, |
| ylabel: str = "Relative Speedup", |
| line_vals=("naive", "cuda"), |
| line_names: Dict[str, str] | None = None, |
| dtype=torch.bfloat16, |
| eps: float = 1e-6, |
| ): |
| timings_ms = collections.defaultdict(dict) |
| spdup_ratio = list() |
| line_vals = list(line_vals) |
| line_names = line_names or {v: v.title() for v in line_vals} |
| x_vals = [make_bwd_key(*_) for _ in configs] |
| x_vals.append("Geometric Mean") |
|
|
| @triton.testing.perf_report( |
| triton.testing.Benchmark(x_names=["config"], |
| x_vals=x_vals, |
| line_arg="provider", |
| line_vals=line_vals, |
| line_names=[line_names[v] for v in line_vals], |
| ylabel=ylabel, |
| plot_name=plot_name, |
| args={})) |
| def bench(config, provider): |
| if config == "Geometric Mean": |
| if provider == "cuda": |
| return round(math.prod(spdup_ratio)**(1 / len(spdup_ratio)), 2) |
| else: |
| return 1.00 |
| batch_size, seq_len, dim = parse_config_string(config) |
| I = case.build_inputs(batch_size, seq_len, dim, dtype, eps) |
| obj = case.make_naive(I) if provider == "naive" else case.make_cuda(I) |
| y = case.forward(obj, I) |
| gin = list(case.grad_inputs(I)) + list(obj.parameters()) |
| if isinstance(y, torch.Tensor): |
| g = [torch.randn_like(y)] |
| else: |
| g = [torch.randn_like(r) for r in y] |
| run = lambda: torch.autograd.grad(y, |
| gin, |
| g, |
| retain_graph=True, |
| create_graph=False, |
| allow_unused=False) |
| ms = triton.testing.do_bench(run) |
| timings_ms[provider][config] = ms |
| if provider == "cuda": |
| ratio = timings_ms["naive"][config] / timings_ms["cuda"][config] |
| spdup_ratio.append(ratio) |
| return round(ratio, 2) |
| else: |
| return 1.00 |
|
|
| return bench |
|
|
