| |
| """Benchmark fp8-kv-attention against a PyTorch FP8-dequant reference.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import time |
| from pathlib import Path |
|
|
| import torch |
|
|
| import sys |
|
|
| TESTS = Path(__file__).resolve().parents[1] / "tests" |
| sys.path.insert(0, str(TESTS)) |
| from test_fp8_kv_attention import SHAPES, SourceOps, load_installed_ops, load_source_ops, make_inputs, reference |
|
|
|
|
| MODES = { |
| "smoke": ["decode_128"], |
| "headline": ["decode_1024", "verify4_1024", "verify8_4096"], |
| "full": ["decode_128", "decode_1024", "verify4_1024", "verify8_4096"], |
| } |
|
|
|
|
| def time_cuda(fn, warmup: int, iters: 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 float(start.elapsed_time(end) * 1000.0 / iters) |
|
|
|
|
| def main() -> int: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--backend", choices=["source", "installed"], default="source") |
| parser.add_argument("--artifact", default=None) |
| parser.add_argument("--mode", choices=sorted(MODES), default="smoke") |
| parser.add_argument("--warmup", type=int, default=20) |
| parser.add_argument("--iters", type=int, default=100) |
| parser.add_argument("--json-out", default=None) |
| args = parser.parse_args() |
|
|
| ops = load_source_ops() if args.backend == "source" else load_installed_ops(args.artifact) |
| rows = [] |
| for name in MODES[args.mode]: |
| q_seq, kv_seq = SHAPES[name] |
| q, k, v = make_inputs(q_seq, kv_seq, seed=3000 + q_seq * 17 + kv_seq) |
| if isinstance(ops, SourceOps): |
| def kernel_call(): |
| return ops.xqa_bf16_fp8kv(q, k, v, kv_seq) |
| else: |
| pages = k.shape[0] |
| page_table = ops.default_page_table(pages, device=q.device) |
| seq_lens = torch.tensor([[kv_seq]], device=q.device, dtype=torch.int32) |
| mask = ops.causal_spec_mask(q_seq, device=q.device) |
| sem, scratch = ops.allocate_workspace(q_seq=q_seq, device=q.device) |
| out = torch.empty_like(q) |
|
|
| def kernel_call(): |
| return ops.xqa_bf16_fp8kv( |
| q, k, v, page_table, seq_lens, mask, |
| out=out, semaphores=sem, scratch=scratch, |
| max_seq_len=pages * 128, |
| ) |
|
|
| def ref_call(): |
| return reference(q, k, v, kv_seq) |
|
|
| kernel_us = time_cuda(kernel_call, args.warmup, args.iters) |
| ref_us = time_cuda(ref_call, max(2, args.warmup // 5), max(5, args.iters // 10)) |
| rows.append( |
| { |
| "shape": name, |
| "q_seq": q_seq, |
| "kv_seq": kv_seq, |
| "kernel_us": kernel_us, |
| "torch_reference_us": ref_us, |
| "speedup": ref_us / kernel_us, |
| } |
| ) |
| print(f"{name}: kernel={kernel_us:.3f}us ref={ref_us:.3f}us speedup={ref_us / kernel_us:.2f}x") |
| if args.json_out: |
| Path(args.json_out).parent.mkdir(parents=True, exist_ok=True) |
| Path(args.json_out).write_text(json.dumps(rows, indent=2) + "\n") |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
