| |
| """Benchmark sageattention2-blackwell against PyTorch SDPA.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import importlib |
| import sys |
| from pathlib import Path |
|
|
| import torch |
| import torch.nn.functional as F |
|
|
| ROOT = Path(__file__).resolve().parents[2] |
| sys.path.insert(0, str(ROOT / "sageattention2-blackwell" / "tests")) |
| from test_sageattention2_blackwell import load_source_ops, make_inputs, reference, stats |
|
|
|
|
| def load_installed_ops(artifact: str | None): |
| if artifact: |
| sys.path.insert(0, artifact) |
| try: |
| return importlib.import_module("sageattention2_blackwell") |
| finally: |
| if artifact: |
| sys.path.remove(artifact) |
|
|
|
|
| def time_cuda(fn, iters: int, warmup: int) -> float: |
| for _ in range(warmup): |
| fn() |
| torch.cuda.synchronize() |
| start = torch.cuda.Event(enable_timing=True) |
| end = torch.cuda.Event(enable_timing=True) |
| start.record() |
| for _ in range(iters): |
| fn() |
| end.record() |
| torch.cuda.synchronize() |
| return start.elapsed_time(end) * 1000.0 / iters |
|
|
|
|
| def sdpa(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, causal: bool) -> torch.Tensor: |
| q_t = q.transpose(1, 2) |
| if q.shape[2] != k.shape[2]: |
| repeat = q.shape[2] // k.shape[2] |
| k = k.repeat_interleave(repeat, dim=2) |
| v = v.repeat_interleave(repeat, dim=2) |
| return F.scaled_dot_product_attention(q_t, k.transpose(1, 2), v.transpose(1, 2), is_causal=causal).transpose(1, 2) |
|
|
|
|
| def run_case(ops, name: str, seqlen: int, q_heads: int, kv_heads: int, causal: bool, iters: int, warmup: int): |
| q, k, v = make_inputs(1, seqlen, q_heads, kv_heads) |
| out = torch.empty_like(q, dtype=torch.bfloat16) |
| ref = reference(q, k, v, causal) |
|
|
| q_i8, q_scale = ops.quantize_q_bf16_d128(q) |
| k_i8, k_scale = ops.quantize_k_bf16_d128(k) |
| v_half = ops.quantize_v_fp16_bf16_d128(v) |
| torch.cuda.synchronize() |
|
|
| def run_sdpa(): |
| return sdpa(q, k, v, causal) |
|
|
| def run_core(): |
| return ops.sage2_qk_int8_sv_f16_bf16_d128( |
| q_i8, k_i8, v_half, q_scale, k_scale, causal=causal, out=out |
| ) |
|
|
| def run_bf16(): |
| return ops.sage2_prefill_f16_bf16_d128(q, k, v, causal=causal, out=out) |
|
|
| got = run_core() |
| torch.cuda.synchronize() |
| s = stats(got, ref) |
| sdpa_us = time_cuda(run_sdpa, iters, warmup) |
| core_us = time_cuda(run_core, iters, warmup) |
| bf16_us = time_cuda(run_bf16, iters, warmup) |
| return { |
| "name": name, |
| "seqlen": seqlen, |
| "q_heads": q_heads, |
| "kv_heads": kv_heads, |
| "causal": causal, |
| "sdpa_us": sdpa_us, |
| "core_us": core_us, |
| "bf16_us": bf16_us, |
| "core_speedup": sdpa_us / core_us, |
| "bf16_speedup": sdpa_us / bf16_us, |
| **s, |
| } |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--backend", choices=["source", "installed"], default="source") |
| parser.add_argument("--artifact", default=None) |
| parser.add_argument("--mode", choices=["smoke", "full"], default="smoke") |
| parser.add_argument("--iters", type=int, default=100) |
| parser.add_argument("--warmup", type=int, default=20) |
| parser.add_argument("--markdown", default=None) |
| args = parser.parse_args() |
|
|
| if not torch.cuda.is_available(): |
| raise SystemExit("CUDA is required") |
| major, _minor = torch.cuda.get_device_capability(0) |
| if major < 12: |
| raise SystemExit("sageattention2-blackwell requires Blackwell-class CUDA capability") |
| torch.manual_seed(2026) |
| ops = load_source_ops() if args.backend == "source" else load_installed_ops(args.artifact) |
| cases = [ |
| ("qwen3_prefill", 1024, 32, 8, True), |
| ("wan_self_attn", 1024, 24, 24, False), |
| ] |
| if args.mode == "full": |
| cases = [ |
| ("qwen3_prefill", 1024, 32, 8, True), |
| ("qwen3_prefill", 2048, 32, 8, True), |
| ("qwen3_prefill", 4096, 32, 8, True), |
| ("qwen3_prefill", 8192, 32, 8, True), |
| ("wan_self_attn", 1024, 24, 24, False), |
| ("wan_self_attn", 2520, 24, 24, False), |
| ("wan_self_attn", 4096, 24, 24, False), |
| ("wan_self_attn", 5070, 24, 24, False), |
| ] |
| rows = [run_case(ops, *case, args.iters, args.warmup) for case in cases] |
| lines = [ |
| "| Workload | S | Hq/Hkv | Mask | SDPA us | Sage core us | Core speedup | BF16 wrapper us | Wrapper speedup | Cos | p99 abs |", |
| "|---|---:|---:|---|---:|---:|---:|---:|---:|---:|---:|", |
| ] |
| for row in rows: |
| line = ( |
| f"| {row['name']} | {row['seqlen']} | {row['q_heads']}/{row['kv_heads']} | " |
| f"{'causal' if row['causal'] else 'none'} | {row['sdpa_us']:.3f} | {row['core_us']:.3f} | " |
| f"{row['core_speedup']:.2f}x | {row['bf16_us']:.3f} | {row['bf16_speedup']:.2f}x | " |
| f"{row['cos']:.6f} | {row['p99_abs']:.6f} |" |
| ) |
| print(line) |
| lines.append(line) |
| if args.markdown: |
| Path(args.markdown).write_text("\n".join(lines) + "\n", encoding="utf-8") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
