|
|
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 |
|
|
|
