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:
@@ -0,0 +1,450 @@
|
||||
"""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()
|
||||
Reference in New Issue
Block a user