ViperEkura
6709534d64
- delete csrc/kernels/gemm.cu and swiglu.cu and drop their CMake and setup.py registration - remove the ops wrappers plus backend/linear.py and backend/swiglu.py so Linear and MLP call F.linear directly - drop the four gemm and swiglu kernel test files and prune the stale cuda_kernels.md sections - add csrc/bench benchmarks for the remaining kernels: attention decode prefill paged decode paged prefill versus single-launch SDPA references, rotary versus the torch fallback, fp8 quantize and mm_fp8 versus torch baselines - attention, rotary_emb, and fp8_ops kernels are unchanged
253 lines
8.6 KiB
Python
253 lines
8.6 KiB
Python
"""Benchmark the fused rotary-embedding kernel against the torch fallback.
|
|
|
|
The baseline is the complex-multiply fallback from
|
|
``astrai.extension.backend.rotary``. Layouts mirror the production call
|
|
shapes: packed 3D [tokens, n_heads, head_dim] and dense 4D
|
|
[batch, seq_len, n_heads, head_dim]; positions are random integers so every
|
|
row exercises a distinct cos/sin gather.
|
|
"""
|
|
|
|
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
|
|
|
|
import click
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
from astrai.extension import is_available
|
|
from astrai.extension.ops import rotary_emb
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RotaryCase:
|
|
name: str
|
|
layout: str
|
|
batch: int
|
|
seq_len: int
|
|
heads: int
|
|
head_dim: int
|
|
|
|
|
|
DEFAULT_CASES = (
|
|
RotaryCase("decode_bs1", "packed", 1, 1, 32, 128),
|
|
RotaryCase("decode_bs32", "packed", 32, 1, 32, 128),
|
|
RotaryCase("prefill_4k_llama7b", "packed", 1, 4096, 32, 128),
|
|
RotaryCase("prefill_4k_llama70b", "packed", 1, 4096, 64, 128),
|
|
RotaryCase("train_8x2k_llama7b", "dense", 8, 2048, 32, 128),
|
|
RotaryCase("train_4x1k_d64", "dense", 4, 1024, 32, 64),
|
|
)
|
|
|
|
|
|
def parse_case(value: str) -> RotaryCase:
|
|
parts = value.split(":")
|
|
if len(parts) != 6 or not parts[0]:
|
|
raise click.BadParameter("case must use NAME:LAYOUT:BATCH:SEQ:HEADS:HEAD_DIM")
|
|
name, layout, batch, seq_len, heads, head_dim = parts
|
|
try:
|
|
fields = (int(batch), int(seq_len), int(heads), int(head_dim))
|
|
except ValueError as exc:
|
|
raise click.BadParameter("fields must be integers") from exc
|
|
if layout not in ("packed", "dense"):
|
|
raise click.BadParameter("layout must be 'packed' or 'dense'")
|
|
if any(field <= 0 for field in fields) or head_dim % 2:
|
|
raise click.BadParameter("fields must be positive; HEAD_DIM even")
|
|
return RotaryCase(name, layout, fields[0], fields[1], fields[2], fields[3])
|
|
|
|
|
|
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 summarize(values: list[float]) -> dict[str, float]:
|
|
ordered = sorted(values)
|
|
return {
|
|
"median_ms": statistics.median(ordered),
|
|
"p90_ms": ordered[max(0, math.ceil(0.9 * len(ordered)) - 1)],
|
|
}
|
|
|
|
|
|
def torch_apply(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
|
|
"""The complex-multiply fallback (mirrors backend.rotary._torch_apply)."""
|
|
cos, sin = freqs_cis[..., 0], freqs_cis[..., 1]
|
|
dtype = x.dtype
|
|
x_ = x.float().reshape(*x.shape[:-1], -1, 2)
|
|
x_complex = torch.view_as_complex(x_)
|
|
freqs_cis_complex = torch.complex(cos, sin).unsqueeze(-2)
|
|
x_rotated = x_complex * freqs_cis_complex
|
|
return torch.view_as_real(x_rotated).flatten(-2).to(dtype)
|
|
|
|
|
|
def build_freqs(head_dim: int, positions: torch.Tensor) -> torch.Tensor:
|
|
"""[cos, sin] pairs for the given positions, laid out [..., head_dim/2, 2]."""
|
|
theta = 10000.0 ** (
|
|
-torch.arange(0, head_dim, 2, dtype=torch.float64, device=positions.device)
|
|
/ head_dim
|
|
)
|
|
freqs = positions.double().unsqueeze(-1) * theta
|
|
return torch.stack([freqs.cos(), freqs.sin()], dim=-1).float()
|
|
|
|
|
|
def benchmark_case(
|
|
case: RotaryCase, *, warmup: int, iterations: int, trials: int
|
|
) -> dict[str, object]:
|
|
if case.layout == "packed":
|
|
tokens = case.batch * case.seq_len
|
|
x = torch.randn(
|
|
tokens, case.heads, case.head_dim, device="cuda", dtype=torch.bfloat16
|
|
)
|
|
positions = torch.randint(0, 65536, (tokens,), device="cuda")
|
|
else:
|
|
x = torch.randn(
|
|
case.batch,
|
|
case.seq_len,
|
|
case.heads,
|
|
case.head_dim,
|
|
device="cuda",
|
|
dtype=torch.bfloat16,
|
|
)
|
|
positions = torch.randint(0, 65536, (case.batch, case.seq_len), device="cuda")
|
|
freqs_cis = build_freqs(case.head_dim, positions)
|
|
operations: dict[str, Callable[[], torch.Tensor]] = {
|
|
"torch": lambda: torch_apply(x, freqs_cis),
|
|
"cuda": lambda: rotary_emb(x, freqs_cis),
|
|
}
|
|
for operation in operations.values():
|
|
for _ in range(warmup):
|
|
operation()
|
|
torch.cuda.synchronize()
|
|
|
|
samples: dict[str, list[float]] = {name: [] for name in operations}
|
|
order = tuple(operations)
|
|
# A-B-B-A order balances cache, clock, and temperature drift.
|
|
for _ in range(trials):
|
|
for name in (*order, *reversed(order)):
|
|
samples[name].append(time_operation(operations[name], iterations))
|
|
|
|
with torch.no_grad():
|
|
expected = operations["torch"]().float()
|
|
actual = operations["cuda"]().float()
|
|
difference = actual - expected
|
|
io_bytes = (
|
|
2 * x.numel() * x.element_size() + freqs_cis.numel() * freqs_cis.element_size()
|
|
)
|
|
|
|
result: dict[str, object] = {
|
|
"case": case.name,
|
|
"layout": case.layout,
|
|
"heads": case.heads,
|
|
"head_dim": case.head_dim,
|
|
"estimated_io_bytes": io_bytes,
|
|
"max_abs_error": float(difference.abs().max()),
|
|
"cosine_similarity": float(
|
|
F.cosine_similarity(actual.flatten(), expected.flatten(), dim=0)
|
|
),
|
|
}
|
|
for name, samples_ms in samples.items():
|
|
latency = summarize(samples_ms)
|
|
result[name] = {
|
|
"effective_bandwidth_gbps": io_bytes / (latency["median_ms"] / 1000) / 1e9,
|
|
**latency,
|
|
}
|
|
speedup = (result["torch"]["median_ms"] / result["cuda"]["median_ms"] - 1.0) * 100.0
|
|
print(
|
|
f"{case.name},{case.layout},{result['torch']['median_ms']:.5f},"
|
|
f"{result['cuda']['median_ms']:.5f},{speedup:+.1f}%,"
|
|
f"{result['max_abs_error']:.5f}"
|
|
)
|
|
return result
|
|
|
|
|
|
@click.command(help=__doc__)
|
|
@click.option("--output", type=click.Path(path_type=Path), help="Optional JSON output.")
|
|
@click.option(
|
|
"--case",
|
|
"case_values",
|
|
multiple=True,
|
|
help="Filter defaults by bare name, or add/override with "
|
|
"NAME:LAYOUT:BATCH:SEQ:HEADS:HEAD_DIM.",
|
|
)
|
|
@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 | None,
|
|
case_values: tuple[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("rotary_emb"):
|
|
raise click.ClickException("the built rotary_emb kernel is required")
|
|
|
|
# A bare name filters the matching default; a full spec overrides or appends.
|
|
chosen: dict[str, RotaryCase] = {}
|
|
for value in case_values:
|
|
case = (
|
|
next((c for c in DEFAULT_CASES if c.name == value), None)
|
|
if ":" not in value
|
|
else parse_case(value)
|
|
)
|
|
if case is None:
|
|
raise click.BadParameter(f"unknown default case {value!r}")
|
|
chosen[case.name] = case
|
|
cases = tuple(chosen.values()) or DEFAULT_CASES
|
|
|
|
torch.manual_seed(seed)
|
|
torch.cuda.manual_seed_all(seed)
|
|
print("case,layout,torch_ms,cuda_ms,speedup,max_abs")
|
|
results = []
|
|
with torch.inference_mode():
|
|
for case in cases:
|
|
results.append(
|
|
benchmark_case(
|
|
case, warmup=warmup, iterations=iterations, trials=trials
|
|
)
|
|
)
|
|
torch.cuda.empty_cache()
|
|
|
|
if output is not None:
|
|
props = torch.cuda.get_device_properties(0)
|
|
payload = {
|
|
"metadata": {
|
|
"gpu_name": props.name,
|
|
"compute_capability": f"{props.major}.{props.minor}",
|
|
"torch_version": torch.__version__,
|
|
"cuda_version": torch.version.cuda,
|
|
"timestamp_utc": datetime.now(timezone.utc).isoformat(),
|
|
},
|
|
"settings": {
|
|
"warmup": warmup,
|
|
"iterations": iterations,
|
|
"trials": trials,
|
|
"seed": seed,
|
|
"order": "A-B-B-A",
|
|
},
|
|
"results": results,
|
|
}
|
|
output.parent.mkdir(parents=True, exist_ok=True)
|
|
output.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
benchmark_command()
|