chore: relocate kernel benchmarks to csrc/bench
- Move benchmark_gemv.py, benchmark_swiglu.py, and benchmark_gemv_common.py from scripts/tools/ to csrc/bench/ so kernel benchmarks live next to the kernels they measure - Update reproduction commands in decode_linear_benchmark.md, swiglu_benchmark.md, and cuda_kernels.md - Codify the placement convention in AGENTS.md: kernel benchmarks in csrc/bench/, pure-CUDA harnesses in csrc/tests/*.cu, engine and evaluation benchmarks in scripts/
This commit is contained in:
@@ -1,333 +0,0 @@
|
||||
"""Benchmark decode-time linear shapes before enabling custom GEMV dispatch.
|
||||
|
||||
The benchmark deliberately calls ``torch.nn.functional.linear`` directly. It
|
||||
establishes the per-architecture cuBLAS baseline that later GEMV primitives and
|
||||
dispatch decisions must beat.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import math
|
||||
import statistics
|
||||
from dataclasses import asdict, dataclass
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Callable, Iterable
|
||||
|
||||
import click
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class LinearShape:
|
||||
name: str
|
||||
n: int
|
||||
k: int
|
||||
|
||||
|
||||
DEFAULT_SHAPES = (
|
||||
LinearShape("q_proj", 1536, 1536),
|
||||
LinearShape("k_proj", 256, 1536),
|
||||
LinearShape("v_proj", 256, 1536),
|
||||
LinearShape("attn_out", 1536, 1536),
|
||||
LinearShape("mlp_up", 6912, 1536),
|
||||
LinearShape("mlp_gate", 6912, 1536),
|
||||
LinearShape("mlp_down", 1536, 6912),
|
||||
LinearShape("lm_head", 100000, 1536),
|
||||
)
|
||||
DTYPES = {"bfloat16": torch.bfloat16, "float16": torch.float16}
|
||||
|
||||
|
||||
def parse_positive_ints(value: str) -> tuple[int, ...]:
|
||||
"""Parse a comma-separated, duplicate-free list of positive integers."""
|
||||
try:
|
||||
values = tuple(dict.fromkeys(int(item.strip()) for item in value.split(",")))
|
||||
except ValueError as exc:
|
||||
raise click.BadParameter("expected comma-separated integers") from exc
|
||||
if not values or any(item <= 0 for item in values):
|
||||
raise click.BadParameter("values must be positive integers")
|
||||
return values
|
||||
|
||||
|
||||
def parse_shape(value: str) -> LinearShape:
|
||||
"""Parse NAME:N:K into a benchmark shape."""
|
||||
parts = value.split(":")
|
||||
if len(parts) != 3 or not parts[0]:
|
||||
raise click.BadParameter("shape must use NAME:N:K")
|
||||
try:
|
||||
n, k = (int(item) for item in parts[1:])
|
||||
except ValueError as exc:
|
||||
raise click.BadParameter("N and K must be integers") from exc
|
||||
if n <= 0 or k <= 0:
|
||||
raise click.BadParameter("N and K must be positive")
|
||||
return LinearShape(parts[0], n, k)
|
||||
|
||||
|
||||
def estimate_io_bytes(
|
||||
m: int, n: int, k: int, element_size: int, *, has_bias: bool
|
||||
) -> int:
|
||||
"""Estimate bytes touched once by Y[M,N] = X[M,K] @ W[N,K].T."""
|
||||
elements = m * k + n * k + m * n
|
||||
if has_bias:
|
||||
elements += n
|
||||
return elements * element_size
|
||||
|
||||
|
||||
def percentile(values: Iterable[float], quantile: float) -> float:
|
||||
ordered = sorted(values)
|
||||
if not ordered:
|
||||
raise ValueError("percentile requires at least one sample")
|
||||
rank = (len(ordered) - 1) * quantile
|
||||
lower = math.floor(rank)
|
||||
upper = math.ceil(rank)
|
||||
if lower == upper:
|
||||
return ordered[lower]
|
||||
fraction = rank - lower
|
||||
return ordered[lower] * (1 - fraction) + ordered[upper] * fraction
|
||||
|
||||
|
||||
def summarize_latency(samples_ms: list[float]) -> dict[str, float]:
|
||||
return {
|
||||
"median_ms": statistics.median(samples_ms),
|
||||
"p90_ms": percentile(samples_ms, 0.90),
|
||||
"p99_ms": percentile(samples_ms, 0.99),
|
||||
"min_ms": min(samples_ms),
|
||||
"max_ms": max(samples_ms),
|
||||
}
|
||||
|
||||
|
||||
def measure_cuda_ms(
|
||||
operation: Callable[[], torch.Tensor], *, warmup: int, iterations: int, trials: int
|
||||
) -> list[float]:
|
||||
for _ in range(warmup):
|
||||
operation()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
samples = []
|
||||
for _ in range(trials):
|
||||
start = torch.cuda.Event(enable_timing=True)
|
||||
end = torch.cuda.Event(enable_timing=True)
|
||||
start.record()
|
||||
for _ in range(iterations):
|
||||
operation()
|
||||
end.record()
|
||||
end.synchronize()
|
||||
samples.append(start.elapsed_time(end) / iterations)
|
||||
return samples
|
||||
|
||||
|
||||
def count_cuda_kernels(
|
||||
operation: Callable[[], torch.Tensor], repeats: int = 5
|
||||
) -> float:
|
||||
"""Profile a few calls and return the average device events per call."""
|
||||
with torch.profiler.profile(
|
||||
activities=[
|
||||
torch.profiler.ProfilerActivity.CPU,
|
||||
torch.profiler.ProfilerActivity.CUDA,
|
||||
],
|
||||
acc_events=True,
|
||||
) as profile:
|
||||
for _ in range(repeats):
|
||||
operation()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
device_type = torch.autograd.DeviceType.CUDA
|
||||
events = [event for event in profile.events() if event.device_type == device_type]
|
||||
return len(events) / repeats
|
||||
|
||||
|
||||
def capture_linear(
|
||||
x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None
|
||||
) -> tuple[torch.cuda.CUDAGraph, torch.Tensor]:
|
||||
for _ in range(3):
|
||||
F.linear(x, weight, bias)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(graph):
|
||||
output = F.linear(x, weight, bias)
|
||||
return graph, output
|
||||
|
||||
|
||||
def benchmark_case(
|
||||
shape: LinearShape,
|
||||
m: int,
|
||||
*,
|
||||
dtype: torch.dtype,
|
||||
mode: str,
|
||||
bias_enabled: bool,
|
||||
warmup: int,
|
||||
iterations: int,
|
||||
trials: int,
|
||||
) -> dict[str, object]:
|
||||
x = torch.randn((m, shape.k), device="cuda", dtype=dtype)
|
||||
weight = torch.randn((shape.n, shape.k), device="cuda", dtype=dtype)
|
||||
bias = torch.randn(shape.n, device="cuda", dtype=dtype) if bias_enabled else None
|
||||
|
||||
graph = None
|
||||
graph_output = None
|
||||
if mode == "graph":
|
||||
graph, graph_output = capture_linear(x, weight, bias)
|
||||
|
||||
def operation() -> torch.Tensor:
|
||||
graph.replay()
|
||||
return graph_output
|
||||
|
||||
else:
|
||||
|
||||
def operation() -> torch.Tensor:
|
||||
return F.linear(x, weight, bias)
|
||||
|
||||
samples_ms = measure_cuda_ms(
|
||||
operation, warmup=warmup, iterations=iterations, trials=trials
|
||||
)
|
||||
latency = summarize_latency(samples_ms)
|
||||
io_bytes = estimate_io_bytes(
|
||||
m, shape.n, shape.k, x.element_size(), has_bias=bias is not None
|
||||
)
|
||||
median_seconds = latency["median_ms"] / 1000
|
||||
|
||||
result: dict[str, object] = {
|
||||
"name": shape.name,
|
||||
"m": m,
|
||||
"n": shape.n,
|
||||
"k": shape.k,
|
||||
"mode": mode,
|
||||
"bias": bias is not None,
|
||||
"estimated_io_bytes": io_bytes,
|
||||
"effective_bandwidth_gbps": io_bytes / median_seconds / 1e9,
|
||||
"cuda_kernel_launches_per_call": count_cuda_kernels(operation),
|
||||
**latency,
|
||||
"samples_ms": samples_ms,
|
||||
}
|
||||
return result
|
||||
|
||||
|
||||
def render_markdown(payload: dict[str, object]) -> str:
|
||||
metadata = payload["metadata"]
|
||||
assert isinstance(metadata, dict)
|
||||
results = payload["results"]
|
||||
assert isinstance(results, list)
|
||||
|
||||
lines = [
|
||||
"# Decode linear baseline",
|
||||
"",
|
||||
f"- GPU: {metadata['gpu_name']}",
|
||||
f"- Compute capability: {metadata['compute_capability']}",
|
||||
f"- PyTorch / CUDA: {metadata['torch_version']} / {metadata['cuda_version']}",
|
||||
f"- Dtype: {metadata['dtype']}",
|
||||
"",
|
||||
"| Layer | M | N | K | Mode | Median (ms) | p99 (ms) | GB/s | CUDA kernels/call |",
|
||||
"|---|---:|---:|---:|---|---:|---:|---:|---:|",
|
||||
]
|
||||
for item in results:
|
||||
assert isinstance(item, dict)
|
||||
lines.append(
|
||||
"| {name} | {m} | {n} | {k} | {mode} | {median_ms:.4f} | "
|
||||
"{p99_ms:.4f} | {effective_bandwidth_gbps:.1f} | "
|
||||
"{cuda_kernel_launches_per_call:.2f} |".format(**item)
|
||||
)
|
||||
lines.append("")
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def device_metadata(dtype_name: str) -> dict[str, object]:
|
||||
props = torch.cuda.get_device_properties(0)
|
||||
return {
|
||||
"timestamp_utc": datetime.now(timezone.utc).isoformat(),
|
||||
"gpu_name": props.name,
|
||||
"compute_capability": f"{props.major}.{props.minor}",
|
||||
"total_memory_bytes": props.total_memory,
|
||||
"torch_version": torch.__version__,
|
||||
"cuda_version": torch.version.cuda,
|
||||
"dtype": dtype_name,
|
||||
}
|
||||
|
||||
|
||||
@click.command(help=__doc__)
|
||||
@click.option("--output", type=click.Path(path_type=Path), required=True)
|
||||
@click.option("--markdown-output", type=click.Path(path_type=Path))
|
||||
@click.option("--m-values", default="1,2,4,8,16,32", show_default=True)
|
||||
@click.option(
|
||||
"--shape",
|
||||
"shape_values",
|
||||
multiple=True,
|
||||
help="Override defaults with repeatable NAME:N:K shapes.",
|
||||
)
|
||||
@click.option("--dtype", type=click.Choice(tuple(DTYPES)), default="bfloat16")
|
||||
@click.option("--mode", type=click.Choice(("eager", "graph", "both")), default="both")
|
||||
@click.option("--bias/--no-bias", default=False)
|
||||
@click.option("--warmup", type=click.IntRange(min=1), default=10, show_default=True)
|
||||
@click.option(
|
||||
"--iterations", type=click.IntRange(min=1), default=100, show_default=True
|
||||
)
|
||||
@click.option("--trials", type=click.IntRange(min=1), default=20, show_default=True)
|
||||
@click.option("--seed", type=int, default=0, show_default=True)
|
||||
def benchmark_command(
|
||||
output: Path,
|
||||
markdown_output: Path | None,
|
||||
m_values: str,
|
||||
shape_values: tuple[str, ...],
|
||||
dtype: str,
|
||||
mode: str,
|
||||
bias: bool,
|
||||
warmup: int,
|
||||
iterations: int,
|
||||
trials: int,
|
||||
seed: int,
|
||||
) -> None:
|
||||
if not torch.cuda.is_available():
|
||||
raise click.ClickException("CUDA is required")
|
||||
|
||||
parsed_m = parse_positive_ints(m_values)
|
||||
shapes = tuple(parse_shape(item) for item in shape_values) or DEFAULT_SHAPES
|
||||
modes = ("eager", "graph") if mode == "both" else (mode,)
|
||||
torch.manual_seed(seed)
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
|
||||
results = []
|
||||
for shape in shapes:
|
||||
for m in parsed_m:
|
||||
for current_mode in modes:
|
||||
click.echo(
|
||||
f"{shape.name}: M={m} N={shape.n} K={shape.k} {current_mode}"
|
||||
)
|
||||
results.append(
|
||||
benchmark_case(
|
||||
shape,
|
||||
m,
|
||||
dtype=DTYPES[dtype],
|
||||
mode=current_mode,
|
||||
bias_enabled=bias,
|
||||
warmup=warmup,
|
||||
iterations=iterations,
|
||||
trials=trials,
|
||||
)
|
||||
)
|
||||
|
||||
payload: dict[str, object] = {
|
||||
"schema_version": 1,
|
||||
"metadata": device_metadata(dtype),
|
||||
"parameters": {
|
||||
"m_values": list(parsed_m),
|
||||
"shapes": [asdict(shape) for shape in shapes],
|
||||
"modes": list(modes),
|
||||
"bias": bias,
|
||||
"warmup": warmup,
|
||||
"iterations": iterations,
|
||||
"trials": trials,
|
||||
"seed": seed,
|
||||
},
|
||||
"results": results,
|
||||
}
|
||||
output.parent.mkdir(parents=True, exist_ok=True)
|
||||
output.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")
|
||||
if markdown_output is not None:
|
||||
markdown_output.parent.mkdir(parents=True, exist_ok=True)
|
||||
markdown_output.write_text(render_markdown(payload), encoding="utf-8")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
benchmark_command()
|
||||
@@ -1,450 +0,0 @@
|
||||
"""Benchmark the BF16 GEMV primitive and guarded linear dispatcher.
|
||||
|
||||
The kernel suite covers AstrAI's native projections plus common LLaMA and
|
||||
GPT-NeoX matrix shapes. The chain suite is a synthetic projection/MLP chain;
|
||||
it measures dispatcher overhead and dependent MLP work, but is deliberately
|
||||
not presented as a whole-model throughput benchmark.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import gc
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import statistics
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from astrai.extension import bf16_gemv, is_available, linear
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Shape:
|
||||
label: str
|
||||
n: int
|
||||
k: int
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Chain:
|
||||
label: str
|
||||
hidden: int
|
||||
kv: int
|
||||
intermediate: int
|
||||
fused_qkv: bool = False
|
||||
gated_mlp: bool = True
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Timing:
|
||||
median_ms: float
|
||||
p90_ms: float
|
||||
|
||||
|
||||
ASTRAI_SHAPES = (
|
||||
Shape("astrai_qkv", 256, 1536),
|
||||
Shape("astrai_square", 1536, 1536),
|
||||
Shape("astrai_up_gate", 6912, 1536),
|
||||
Shape("astrai_down", 1536, 6912),
|
||||
Shape("astrai_lm_head", 100000, 1536),
|
||||
)
|
||||
|
||||
TRADITIONAL_SHAPES = (
|
||||
Shape("llama2_7b_qo", 4096, 4096),
|
||||
Shape("llama2_7b_up_gate", 11008, 4096),
|
||||
Shape("llama2_7b_down", 4096, 11008),
|
||||
Shape("llama3_8b_kv", 1024, 4096),
|
||||
Shape("llama3_8b_up_gate", 14336, 4096),
|
||||
Shape("llama3_8b_down", 4096, 14336),
|
||||
Shape("llama2_13b_qo", 5120, 5120),
|
||||
Shape("llama2_13b_up_gate", 13824, 5120),
|
||||
Shape("llama2_13b_down", 5120, 13824),
|
||||
Shape("gpt_neox_up", 16384, 4096),
|
||||
Shape("gpt_neox_down", 4096, 16384),
|
||||
Shape("qwen2_7b_kv", 512, 3584),
|
||||
Shape("qwen2_7b_qo", 3584, 3584),
|
||||
Shape("qwen2_7b_up_gate", 18944, 3584),
|
||||
Shape("qwen2_7b_down", 3584, 18944),
|
||||
Shape("llama3_70b_kv", 1024, 8192),
|
||||
Shape("llama3_70b_qo", 8192, 8192),
|
||||
Shape("llama3_70b_up_gate", 28672, 8192),
|
||||
Shape("llama3_70b_down", 8192, 28672),
|
||||
Shape("opt_1_3b_qkvo", 2048, 2048),
|
||||
Shape("opt_1_3b_up", 8192, 2048),
|
||||
Shape("opt_1_3b_down", 2048, 8192),
|
||||
)
|
||||
|
||||
CHAINS = (
|
||||
Chain("llama2_7b", 4096, 4096, 11008),
|
||||
Chain("llama3_8b", 4096, 1024, 14336),
|
||||
Chain("llama2_13b", 5120, 5120, 13824),
|
||||
Chain("gpt_neox_20b", 4096, 4096, 16384, fused_qkv=True),
|
||||
Chain("qwen2_7b", 3584, 512, 18944),
|
||||
Chain("llama3_70b", 8192, 1024, 28672),
|
||||
Chain("opt_1_3b", 2048, 2048, 8192, gated_mlp=False),
|
||||
)
|
||||
|
||||
|
||||
def _elapsed_ms(fn: Callable[[], torch.Tensor], inner: int) -> float:
|
||||
start = torch.cuda.Event(enable_timing=True)
|
||||
end = torch.cuda.Event(enable_timing=True)
|
||||
start.record()
|
||||
for _ in range(inner):
|
||||
fn()
|
||||
end.record()
|
||||
end.synchronize()
|
||||
return start.elapsed_time(end) / inner
|
||||
|
||||
|
||||
def _timing(values: list[float]) -> Timing:
|
||||
ordered = sorted(values)
|
||||
p90_index = max(0, math.ceil(0.9 * len(ordered)) - 1)
|
||||
return Timing(statistics.median(ordered), ordered[p90_index])
|
||||
|
||||
|
||||
def _measure_pair(
|
||||
baseline: Callable[[], torch.Tensor],
|
||||
candidate: Callable[[], torch.Tensor],
|
||||
*,
|
||||
warmup: int,
|
||||
samples: int,
|
||||
inner: int,
|
||||
prepare_baseline: Callable[[], None] = lambda: None,
|
||||
prepare_candidate: Callable[[], None] = lambda: None,
|
||||
) -> tuple[Timing, Timing]:
|
||||
cases = (
|
||||
("baseline", prepare_baseline, baseline),
|
||||
("candidate", prepare_candidate, candidate),
|
||||
)
|
||||
for iteration in range(warmup):
|
||||
_, prepare, fn = cases[iteration % 2]
|
||||
prepare()
|
||||
fn()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
values: dict[str, list[float]] = {"baseline": [], "candidate": []}
|
||||
for sample in range(samples):
|
||||
order = cases if sample % 2 == 0 else tuple(reversed(cases))
|
||||
for label, prepare, fn in order:
|
||||
prepare()
|
||||
values[label].append(_elapsed_ms(fn, inner))
|
||||
return _timing(values["baseline"]), _timing(values["candidate"])
|
||||
|
||||
|
||||
def _print_header() -> None:
|
||||
print(
|
||||
"suite,label,m,n,k,torch_median_ms,torch_p90_ms,"
|
||||
"candidate_median_ms,candidate_p90_ms,speedup_pct,"
|
||||
"max_abs,relative_l2,argmax_equal"
|
||||
)
|
||||
|
||||
|
||||
def _print_result(
|
||||
suite: str,
|
||||
label: str,
|
||||
m: int,
|
||||
n: int,
|
||||
k: int,
|
||||
baseline: Timing,
|
||||
candidate: Timing,
|
||||
reference: torch.Tensor,
|
||||
actual: torch.Tensor,
|
||||
) -> dict[str, object]:
|
||||
difference = actual.float() - reference.float()
|
||||
max_abs = difference.abs().max().item()
|
||||
relative_l2 = difference.norm().item() / max(reference.float().norm().item(), 1e-12)
|
||||
argmax_equal = torch.equal(actual.argmax(dim=-1), reference.argmax(dim=-1))
|
||||
speedup = (baseline.median_ms / candidate.median_ms - 1.0) * 100.0
|
||||
result: dict[str, object] = {
|
||||
"suite": suite,
|
||||
"label": label,
|
||||
"m": m,
|
||||
"n": n,
|
||||
"k": k,
|
||||
"torch_median_ms": baseline.median_ms,
|
||||
"torch_p90_ms": baseline.p90_ms,
|
||||
"candidate_median_ms": candidate.median_ms,
|
||||
"candidate_p90_ms": candidate.p90_ms,
|
||||
"speedup_pct": speedup,
|
||||
"max_abs": max_abs,
|
||||
"relative_l2": relative_l2,
|
||||
"argmax_equal": argmax_equal,
|
||||
}
|
||||
print(
|
||||
f"{suite},{label},{m},{n},{k},"
|
||||
f"{baseline.median_ms:.6f},{baseline.p90_ms:.6f},"
|
||||
f"{candidate.median_ms:.6f},{candidate.p90_ms:.6f},"
|
||||
f"{speedup:+.2f},{max_abs:.6f},{relative_l2:.8f},"
|
||||
f"{str(argmax_equal).lower()}",
|
||||
flush=True,
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
def _weight(n: int, k: int, device: torch.device, std: float) -> torch.Tensor:
|
||||
weight = torch.empty((n, k), device=device, dtype=torch.bfloat16)
|
||||
weight.normal_(mean=0.0, std=std)
|
||||
return weight.requires_grad_(True)
|
||||
|
||||
|
||||
def _kernel_functions(
|
||||
x: torch.Tensor, weight: torch.Tensor
|
||||
) -> tuple[Callable[[], torch.Tensor], Callable[[], torch.Tensor]]:
|
||||
def baseline() -> torch.Tensor:
|
||||
return F.linear(x, weight)
|
||||
|
||||
def candidate() -> torch.Tensor:
|
||||
return bf16_gemv(x, weight.detach())
|
||||
|
||||
return baseline, candidate
|
||||
|
||||
|
||||
def benchmark_kernels(
|
||||
args: argparse.Namespace, device: torch.device
|
||||
) -> list[dict[str, object]]:
|
||||
if args.family == "astrai":
|
||||
shapes = ASTRAI_SHAPES
|
||||
elif args.family == "traditional":
|
||||
shapes = TRADITIONAL_SHAPES
|
||||
else:
|
||||
shapes = ASTRAI_SHAPES + TRADITIONAL_SHAPES
|
||||
if args.shape_label:
|
||||
requested = set(args.shape_label)
|
||||
shapes = tuple(shape for shape in shapes if shape.label in requested)
|
||||
missing = requested - {shape.label for shape in shapes}
|
||||
if missing:
|
||||
raise ValueError(f"unknown shape labels: {', '.join(sorted(missing))}")
|
||||
|
||||
results: list[dict[str, object]] = []
|
||||
for shape in shapes:
|
||||
weight = _weight(shape.n, shape.k, device, args.weight_std)
|
||||
for m in args.m:
|
||||
x = torch.randn((m, shape.k), device=device, dtype=torch.bfloat16)
|
||||
baseline_fn, candidate_fn = _kernel_functions(x, weight)
|
||||
with torch.inference_mode():
|
||||
reference = baseline_fn()
|
||||
actual = candidate_fn()
|
||||
baseline, candidate = _measure_pair(
|
||||
baseline_fn,
|
||||
candidate_fn,
|
||||
warmup=args.warmup,
|
||||
samples=args.samples,
|
||||
inner=args.inner,
|
||||
)
|
||||
results.append(
|
||||
_print_result(
|
||||
"kernel",
|
||||
shape.label,
|
||||
m,
|
||||
shape.n,
|
||||
shape.k,
|
||||
baseline,
|
||||
candidate,
|
||||
reference,
|
||||
actual,
|
||||
)
|
||||
)
|
||||
del baseline_fn, candidate_fn, x, reference, actual
|
||||
del weight
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
return results
|
||||
|
||||
|
||||
def _set_mode(mode: str) -> None:
|
||||
os.environ["ASTRAI_GEMV"] = mode
|
||||
|
||||
|
||||
def _chain_weights(
|
||||
spec: Chain, device: torch.device, std: float
|
||||
) -> dict[str, torch.Tensor]:
|
||||
weights = {
|
||||
"o": _weight(spec.hidden, spec.hidden, device, std),
|
||||
"up": _weight(spec.intermediate, spec.hidden, device, std),
|
||||
"down": _weight(spec.hidden, spec.intermediate, device, std),
|
||||
}
|
||||
if spec.fused_qkv:
|
||||
weights["qkv"] = _weight(3 * spec.hidden, spec.hidden, device, std)
|
||||
else:
|
||||
weights.update(
|
||||
{
|
||||
"q": _weight(spec.hidden, spec.hidden, device, std),
|
||||
"k": _weight(spec.kv, spec.hidden, device, std),
|
||||
"v": _weight(spec.kv, spec.hidden, device, std),
|
||||
}
|
||||
)
|
||||
if spec.gated_mlp:
|
||||
weights["gate"] = _weight(spec.intermediate, spec.hidden, device, std)
|
||||
return weights
|
||||
|
||||
|
||||
def _chain_fn(
|
||||
x: torch.Tensor, weights: dict[str, torch.Tensor], spec: Chain
|
||||
) -> Callable[[], torch.Tensor]:
|
||||
def run() -> torch.Tensor:
|
||||
output_projection = linear(x, weights["o"])
|
||||
up = linear(x, weights["up"])
|
||||
if spec.fused_qkv:
|
||||
attention_projection = linear(x, weights["qkv"])[..., : x.shape[-1]]
|
||||
hidden = F.gelu(up)
|
||||
else:
|
||||
attention_projection = linear(x, weights["q"])
|
||||
linear(x, weights["k"])
|
||||
linear(x, weights["v"])
|
||||
if spec.gated_mlp:
|
||||
gate = linear(x, weights["gate"])
|
||||
hidden = F.silu(gate) * up
|
||||
else:
|
||||
hidden = F.gelu(up)
|
||||
down = linear(hidden, weights["down"])
|
||||
return attention_projection + output_projection + down
|
||||
|
||||
return run
|
||||
|
||||
|
||||
def benchmark_chains(
|
||||
args: argparse.Namespace, device: torch.device
|
||||
) -> list[dict[str, object]]:
|
||||
results: list[dict[str, object]] = []
|
||||
chains = CHAINS
|
||||
if args.chain_label:
|
||||
requested = set(args.chain_label)
|
||||
chains = tuple(chain for chain in chains if chain.label in requested)
|
||||
missing = requested - {chain.label for chain in chains}
|
||||
if missing:
|
||||
raise ValueError(f"unknown chain labels: {', '.join(sorted(missing))}")
|
||||
for spec in chains:
|
||||
weights = _chain_weights(spec, device, args.weight_std)
|
||||
for m in args.m:
|
||||
x = torch.randn((m, spec.hidden), device=device, dtype=torch.bfloat16)
|
||||
run = _chain_fn(x, weights, spec)
|
||||
with torch.inference_mode():
|
||||
_set_mode("0")
|
||||
reference = run()
|
||||
_set_mode(args.candidate_mode)
|
||||
actual = run()
|
||||
baseline, candidate = _measure_pair(
|
||||
run,
|
||||
run,
|
||||
warmup=args.warmup,
|
||||
samples=args.samples,
|
||||
inner=args.chain_inner,
|
||||
prepare_baseline=lambda: _set_mode("0"),
|
||||
prepare_candidate=lambda: _set_mode(args.candidate_mode),
|
||||
)
|
||||
results.append(
|
||||
_print_result(
|
||||
"synthetic_chain",
|
||||
spec.label,
|
||||
m,
|
||||
spec.hidden,
|
||||
spec.intermediate,
|
||||
baseline,
|
||||
candidate,
|
||||
reference,
|
||||
actual,
|
||||
)
|
||||
)
|
||||
del x, reference, actual
|
||||
del weights
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
return results
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--suite", choices=("kernel", "chain", "all"), default="all")
|
||||
parser.add_argument(
|
||||
"--family", choices=("astrai", "traditional", "all"), default="all"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--m", type=int, nargs="+", choices=(1, 2, 4, 8), default=(1, 2, 4, 8)
|
||||
)
|
||||
parser.add_argument(
|
||||
"--shape-label",
|
||||
action="append",
|
||||
help="limit the kernel suite to one or more named shape labels",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--chain-label",
|
||||
action="append",
|
||||
help="limit the chain suite to one or more named model families",
|
||||
)
|
||||
parser.add_argument("--device", type=int, default=0)
|
||||
parser.add_argument("--warmup", type=int, default=20)
|
||||
parser.add_argument("--samples", type=int, default=9)
|
||||
parser.add_argument("--inner", type=int, default=100)
|
||||
parser.add_argument("--chain-inner", type=int, default=20)
|
||||
parser.add_argument(
|
||||
"--candidate-mode",
|
||||
choices=("auto", "1"),
|
||||
default="auto",
|
||||
help="dispatcher mode for the candidate side of the chain suite",
|
||||
)
|
||||
parser.add_argument("--weight-std", type=float, default=0.02)
|
||||
parser.add_argument("--seed", type=int, default=20260902)
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
type=Path,
|
||||
help="optional JSON output; stdout always retains the compact CSV table",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
if not torch.cuda.is_available() or not is_available("bf16_gemv"):
|
||||
raise RuntimeError("benchmark requires CUDA and the built bf16_gemv extension")
|
||||
if args.warmup < 0 or args.samples < 1 or args.inner < 1 or args.chain_inner < 1:
|
||||
raise ValueError("warmup must be non-negative and sample/inner counts positive")
|
||||
|
||||
torch.cuda.set_device(args.device)
|
||||
device = torch.device("cuda", args.device)
|
||||
torch.manual_seed(args.seed)
|
||||
torch.cuda.manual_seed_all(args.seed)
|
||||
properties = torch.cuda.get_device_properties(device)
|
||||
print(
|
||||
f"# device={properties.name}, capability={properties.major}.{properties.minor}, "
|
||||
f"seed={args.seed}, weight_std={args.weight_std}"
|
||||
)
|
||||
_print_header()
|
||||
results: list[dict[str, object]] = []
|
||||
if args.suite in ("kernel", "all"):
|
||||
results.extend(benchmark_kernels(args, device))
|
||||
if args.suite in ("chain", "all"):
|
||||
results.extend(benchmark_chains(args, device))
|
||||
if args.output is not None:
|
||||
payload = {
|
||||
"environment": {
|
||||
"device": properties.name,
|
||||
"capability": f"{properties.major}.{properties.minor}",
|
||||
"torch": torch.__version__,
|
||||
"cuda": torch.version.cuda,
|
||||
},
|
||||
"parameters": {
|
||||
"suite": args.suite,
|
||||
"family": args.family,
|
||||
"m": args.m,
|
||||
"shape_labels": args.shape_label,
|
||||
"chain_labels": args.chain_label,
|
||||
"candidate_mode": args.candidate_mode,
|
||||
"seed": args.seed,
|
||||
"weight_std": args.weight_std,
|
||||
"warmup": args.warmup,
|
||||
"samples": args.samples,
|
||||
"inner": args.inner,
|
||||
"chain_inner": args.chain_inner,
|
||||
},
|
||||
"results": results,
|
||||
}
|
||||
args.output.parent.mkdir(parents=True, exist_ok=True)
|
||||
args.output.write_text(json.dumps(payload, indent=2) + "\n")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,325 +0,0 @@
|
||||
"""Benchmark fused BF16 SwiGLU against torch and unfused GEMV chains."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import math
|
||||
import statistics
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Callable, Iterable
|
||||
|
||||
import click
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from astrai.extension import bf16_gemv, bf16_swiglu, is_available
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SwiGLUShape:
|
||||
name: str
|
||||
n: int
|
||||
k: int
|
||||
|
||||
|
||||
DEFAULT_SHAPES = (
|
||||
SwiGLUShape("astrai_1b", 6912, 1536),
|
||||
SwiGLUShape("llama2_7b", 11008, 4096),
|
||||
SwiGLUShape("llama3_8b", 14336, 4096),
|
||||
SwiGLUShape("llama2_13b", 13824, 5120),
|
||||
SwiGLUShape("gpt_neox_20b", 16384, 6144),
|
||||
)
|
||||
|
||||
|
||||
def parse_positive_ints(value: str) -> tuple[int, ...]:
|
||||
try:
|
||||
values = tuple(dict.fromkeys(int(item.strip()) for item in value.split(",")))
|
||||
except ValueError as exc:
|
||||
raise click.BadParameter("expected comma-separated integers") from exc
|
||||
if not values or any(item <= 0 for item in values):
|
||||
raise click.BadParameter("values must be positive integers")
|
||||
return values
|
||||
|
||||
|
||||
def parse_shape(value: str) -> SwiGLUShape:
|
||||
parts = value.split(":")
|
||||
if len(parts) != 3 or not parts[0]:
|
||||
raise click.BadParameter("shape must use NAME:N:K")
|
||||
try:
|
||||
n, k = (int(item) for item in parts[1:])
|
||||
except ValueError as exc:
|
||||
raise click.BadParameter("N and K must be integers") from exc
|
||||
if n <= 0 or k <= 0 or k % 8:
|
||||
raise click.BadParameter("N must be positive and K positive/divisible by 8")
|
||||
return SwiGLUShape(parts[0], n, k)
|
||||
|
||||
|
||||
def percentile(values: Iterable[float], quantile: float) -> float:
|
||||
ordered = sorted(values)
|
||||
rank = (len(ordered) - 1) * quantile
|
||||
lower = math.floor(rank)
|
||||
upper = math.ceil(rank)
|
||||
if lower == upper:
|
||||
return ordered[lower]
|
||||
fraction = rank - lower
|
||||
return ordered[lower] * (1 - fraction) + ordered[upper] * fraction
|
||||
|
||||
|
||||
def summarize(values: list[float]) -> dict[str, float]:
|
||||
return {
|
||||
"median_ms": statistics.median(values),
|
||||
"p90_ms": percentile(values, 0.90),
|
||||
"p99_ms": percentile(values, 0.99),
|
||||
"min_ms": min(values),
|
||||
"max_ms": max(values),
|
||||
}
|
||||
|
||||
|
||||
def time_operation(operation: Callable[[], torch.Tensor], iterations: int) -> float:
|
||||
start = torch.cuda.Event(enable_timing=True)
|
||||
end = torch.cuda.Event(enable_timing=True)
|
||||
start.record()
|
||||
for _ in range(iterations):
|
||||
operation()
|
||||
end.record()
|
||||
end.synchronize()
|
||||
return start.elapsed_time(end) / iterations
|
||||
|
||||
|
||||
def count_cuda_kernels(
|
||||
operation: Callable[[], torch.Tensor], repeats: int = 5
|
||||
) -> float:
|
||||
with torch.profiler.profile(
|
||||
activities=[
|
||||
torch.profiler.ProfilerActivity.CPU,
|
||||
torch.profiler.ProfilerActivity.CUDA,
|
||||
],
|
||||
acc_events=True,
|
||||
) as profile:
|
||||
for _ in range(repeats):
|
||||
operation()
|
||||
torch.cuda.synchronize()
|
||||
device_type = torch.autograd.DeviceType.CUDA
|
||||
events = [event for event in profile.events() if event.device_type == device_type]
|
||||
return len(events) / repeats
|
||||
|
||||
|
||||
def capture(operation: Callable[[], torch.Tensor]):
|
||||
for _ in range(3):
|
||||
operation()
|
||||
torch.cuda.synchronize()
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(graph):
|
||||
output = operation()
|
||||
|
||||
def replay() -> torch.Tensor:
|
||||
graph.replay()
|
||||
return output
|
||||
|
||||
return replay
|
||||
|
||||
|
||||
def make_operations(x, up_weight, gate_weight, mode: str):
|
||||
operations: dict[str, Callable[[], torch.Tensor]] = {
|
||||
"torch": lambda: F.linear(x, up_weight) * F.silu(F.linear(x, gate_weight)),
|
||||
"gemv_chain": lambda: (
|
||||
bf16_gemv(x, up_weight) * F.silu(bf16_gemv(x, gate_weight))
|
||||
),
|
||||
"fused": lambda: bf16_swiglu(x, up_weight, gate_weight),
|
||||
}
|
||||
if mode == "graph":
|
||||
operations = {name: capture(op) for name, op in operations.items()}
|
||||
return operations
|
||||
|
||||
|
||||
def benchmark_case(
|
||||
shape: SwiGLUShape,
|
||||
m: int,
|
||||
mode: str,
|
||||
*,
|
||||
warmup: int,
|
||||
iterations: int,
|
||||
trials: int,
|
||||
) -> list[dict[str, object]]:
|
||||
x = torch.randn((m, shape.k), device="cuda", dtype=torch.bfloat16) * 0.1
|
||||
scale = shape.k**-0.5
|
||||
up_weight = (
|
||||
torch.randn((shape.n, shape.k), device="cuda", dtype=torch.bfloat16) * scale
|
||||
)
|
||||
gate_weight = (
|
||||
torch.randn((shape.n, shape.k), device="cuda", dtype=torch.bfloat16) * scale
|
||||
)
|
||||
operations = make_operations(x, up_weight, gate_weight, mode)
|
||||
for operation in operations.values():
|
||||
for _ in range(warmup):
|
||||
operation()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
samples = {name: [] for name in operations}
|
||||
forward_order = tuple(operations)
|
||||
# A-B-C-C-B-A order balances cache, clock, and temperature drift.
|
||||
for _ in range(trials):
|
||||
for name in (*forward_order, *reversed(forward_order)):
|
||||
samples[name].append(time_operation(operations[name], iterations))
|
||||
|
||||
with torch.no_grad():
|
||||
expected = operations["torch"]().clone()
|
||||
actual = operations["fused"]().clone()
|
||||
difference = (actual.float() - expected.float()).abs()
|
||||
max_abs_error = float(difference.max())
|
||||
mean_abs_error = float(difference.mean())
|
||||
cosine_similarity = float(
|
||||
F.cosine_similarity(actual.float().flatten(), expected.float().flatten(), dim=0)
|
||||
)
|
||||
|
||||
results = []
|
||||
for name, operation in operations.items():
|
||||
result: dict[str, object] = {
|
||||
"shape": shape.name,
|
||||
"m": m,
|
||||
"n": shape.n,
|
||||
"k": shape.k,
|
||||
"mode": mode,
|
||||
"implementation": name,
|
||||
"cuda_kernel_launches_per_call": count_cuda_kernels(operation),
|
||||
**summarize(samples[name]),
|
||||
}
|
||||
if name == "fused":
|
||||
result.update(
|
||||
max_abs_error=max_abs_error,
|
||||
mean_abs_error=mean_abs_error,
|
||||
cosine_similarity=cosine_similarity,
|
||||
)
|
||||
results.append(result)
|
||||
return results
|
||||
|
||||
|
||||
def device_metadata() -> dict[str, object]:
|
||||
props = torch.cuda.get_device_properties(0)
|
||||
return {
|
||||
"timestamp_utc": datetime.now(timezone.utc).isoformat(),
|
||||
"gpu_name": props.name,
|
||||
"compute_capability": f"{props.major}.{props.minor}",
|
||||
"total_memory_bytes": props.total_memory,
|
||||
"torch_version": torch.__version__,
|
||||
"cuda_version": torch.version.cuda,
|
||||
"dtype": "bfloat16",
|
||||
}
|
||||
|
||||
|
||||
def render_markdown(payload: dict[str, object]) -> str:
|
||||
metadata = payload["metadata"]
|
||||
results = payload["results"]
|
||||
assert isinstance(metadata, dict)
|
||||
assert isinstance(results, list)
|
||||
by_case = {
|
||||
(item["shape"], item["m"], item["mode"], item["implementation"]): item
|
||||
for item in results
|
||||
}
|
||||
cases = sorted({(item["shape"], item["m"], item["mode"]) for item in results})
|
||||
lines = [
|
||||
"# Fused SwiGLU benchmark",
|
||||
"",
|
||||
f"- GPU: {metadata['gpu_name']}",
|
||||
f"- Compute capability: {metadata['compute_capability']}",
|
||||
f"- PyTorch / CUDA: {metadata['torch_version']} / {metadata['cuda_version']}",
|
||||
"",
|
||||
"| Shape | M | Mode | torch ms | GEMV chain ms | fused ms | "
|
||||
"vs best unfused | fused kernels | max abs | cosine |",
|
||||
"|---|---:|---|---:|---:|---:|---:|---:|---:|---:|",
|
||||
]
|
||||
for shape, m, mode in cases:
|
||||
torch_item = by_case[(shape, m, mode, "torch")]
|
||||
gemv_item = by_case[(shape, m, mode, "gemv_chain")]
|
||||
fused_item = by_case[(shape, m, mode, "fused")]
|
||||
best = min(torch_item["median_ms"], gemv_item["median_ms"])
|
||||
improvement = (best / fused_item["median_ms"] - 1) * 100
|
||||
lines.append(
|
||||
f"| {shape} | {m} | {mode} | {torch_item['median_ms']:.5f} | "
|
||||
f"{gemv_item['median_ms']:.5f} | {fused_item['median_ms']:.5f} | "
|
||||
f"{improvement:+.2f}% | "
|
||||
f"{fused_item['cuda_kernel_launches_per_call']:.1f} | "
|
||||
f"{fused_item['max_abs_error']:.5f} | "
|
||||
f"{fused_item['cosine_similarity']:.8f} |"
|
||||
)
|
||||
lines.append("")
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
@click.command(help=__doc__)
|
||||
@click.option("--output", type=click.Path(path_type=Path), required=True)
|
||||
@click.option("--markdown-output", type=click.Path(path_type=Path))
|
||||
@click.option("--m-values", default="1,2,4,8", show_default=True)
|
||||
@click.option("--shape", "shape_values", multiple=True, help="Repeat NAME:N:K.")
|
||||
@click.option("--mode", type=click.Choice(("eager", "graph", "both")), default="both")
|
||||
@click.option("--warmup", type=click.IntRange(min=1), default=10, show_default=True)
|
||||
@click.option(
|
||||
"--iterations", type=click.IntRange(min=1), default=100, show_default=True
|
||||
)
|
||||
@click.option("--trials", type=click.IntRange(min=1), default=10, show_default=True)
|
||||
@click.option("--seed", type=int, default=0, show_default=True)
|
||||
def benchmark_command(
|
||||
output: Path,
|
||||
markdown_output: Path | None,
|
||||
m_values: str,
|
||||
shape_values: tuple[str, ...],
|
||||
mode: str,
|
||||
warmup: int,
|
||||
iterations: int,
|
||||
trials: int,
|
||||
seed: int,
|
||||
) -> None:
|
||||
if not torch.cuda.is_available():
|
||||
raise click.ClickException("CUDA is required")
|
||||
if not is_available("bf16_gemv") or not is_available("bf16_swiglu"):
|
||||
raise click.ClickException("built bf16_gemv and bf16_swiglu are required")
|
||||
shapes = tuple(parse_shape(value) for value in shape_values) or DEFAULT_SHAPES
|
||||
m_values_parsed = parse_positive_ints(m_values)
|
||||
if any(m > 8 for m in m_values_parsed):
|
||||
raise click.BadParameter("fused primitive supports M up to 8")
|
||||
modes = ("eager", "graph") if mode == "both" else (mode,)
|
||||
torch.manual_seed(seed)
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
|
||||
results = []
|
||||
with torch.inference_mode():
|
||||
for shape in shapes:
|
||||
for m in m_values_parsed:
|
||||
for current_mode in modes:
|
||||
click.echo(
|
||||
f"{shape.name}: M={m} N={shape.n} K={shape.k} {current_mode}"
|
||||
)
|
||||
results.extend(
|
||||
benchmark_case(
|
||||
shape,
|
||||
m,
|
||||
current_mode,
|
||||
warmup=warmup,
|
||||
iterations=iterations,
|
||||
trials=trials,
|
||||
)
|
||||
)
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
payload: dict[str, object] = {
|
||||
"metadata": device_metadata(),
|
||||
"settings": {
|
||||
"warmup": warmup,
|
||||
"iterations": iterations,
|
||||
"trials": trials,
|
||||
"seed": seed,
|
||||
"order": "A-B-C-C-B-A",
|
||||
},
|
||||
"results": results,
|
||||
}
|
||||
output.parent.mkdir(parents=True, exist_ok=True)
|
||||
output.write_text(json.dumps(payload, indent=2) + "\n")
|
||||
if markdown_output is not None:
|
||||
markdown_output.parent.mkdir(parents=True, exist_ok=True)
|
||||
markdown_output.write_text(render_markdown(payload))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
benchmark_command()
|
||||
Reference in New Issue
Block a user