ViperEkura
28d11f1610
- 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/
334 lines
10 KiB
Python
334 lines
10 KiB
Python
"""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()
|