| |
| """Benchmark INT8 transformer primitives.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import importlib |
| import sys |
| from pathlib import Path |
|
|
| import torch |
|
|
| ROOT = Path(__file__).resolve().parents[2] |
| sys.path.insert(0, str(ROOT / "int8-transformer-primitives" / "tests")) |
| from test_int8_transformer_primitives import load_source_ops |
|
|
|
|
| def load_ops(backend: str, artifact: str | None): |
| if backend == "source": |
| return load_source_ops() |
| if artifact: |
| sys.path.insert(0, artifact) |
| try: |
| return importlib.import_module("int8_transformer_primitives") |
| finally: |
| if artifact: |
| sys.path.remove(artifact) |
|
|
|
|
| def time_us(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 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=["headline", "full"], default="headline") |
| parser.add_argument("--warmup", type=int, default=20) |
| parser.add_argument("--iters", type=int, default=100) |
| args = parser.parse_args() |
| ops = load_ops(args.backend, args.artifact) |
|
|
| shapes = [ |
| ("decode_m8", 8, 1024, 2560), |
| ("small_batch", 64, 2048, 8192), |
| ("vision_prefill", 522, 2048, 2560), |
| ] |
| if args.mode == "full": |
| shapes += [ |
| ("m1", 1, 1024, 1024), |
| ("m17", 17, 256, 256), |
| ("wide_ffn", 257, 2048, 8192), |
| ] |
|
|
| print("workload,M,K,N,op,flashrt_us,torch_eager_us,speedup") |
| for name, m, k, n in shapes: |
| x = (torch.randn((m, k), device="cuda") * 0.5).to(torch.bfloat16) |
| w = (torch.randn((n, k), device="cuda") * 0.5).to(torch.bfloat16) |
| x_i8, x_scale = ops.quantize_int8_rowwise_bf16(x) |
| w_i8, w_scale = ops.quantize_int8_rowwise_bf16(w) |
| torch.cuda.synchronize() |
| out = torch.empty((m, n), device="cuda", dtype=torch.bfloat16) |
|
|
| def flash(): |
| ops.int8_rowwise_linear_bf16(x_i8, w_i8, x_scale, w_scale, out=out) |
|
|
| def eager(): |
| ((x_i8.float() @ w_i8.float().t()) * x_scale[:, None] * w_scale[None, :]).to(torch.bfloat16) |
|
|
| fu = time_us(flash, args.warmup, args.iters) |
| eu = time_us(eager, max(5, args.warmup // 2), max(20, args.iters // 2)) |
| print(f"{name},{m},{k},{n},int8_rowwise_linear_bf16,{fu:.3f},{eu:.3f},{eu/fu:.2f}x") |
|
|
| q = (torch.randn((522, 2048), device="cuda") * 0.5).to(torch.bfloat16) |
| weight = torch.randn((2048,), device="cuda", dtype=torch.bfloat16) |
| out = torch.empty_like(q, dtype=torch.int8) |
| scales = torch.empty((q.shape[0],), device="cuda", dtype=torch.float32) |
|
|
| def flash_rms(): |
| ops.rms_norm_quantize_int8_rowwise_bf16(q, weight, out=out, scales=scales) |
|
|
| def eager_rms(): |
| y = q.float() * torch.rsqrt((q.float() * q.float()).mean(dim=1, keepdim=True) + 1e-6) * weight.float() |
| s = torch.clamp(y.abs().amax(dim=1) / 127.0, min=1e-10) |
| torch.clamp(torch.round(y / s[:, None]), -127, 127).to(torch.int8) |
|
|
| fu = time_us(flash_rms, args.warmup, args.iters) |
| eu = time_us(eager_rms, max(5, args.warmup // 2), max(20, args.iters // 2)) |
| print(f"vision_prefill,522,2048,0,rms_norm_quantize_int8_rowwise_bf16,{fu:.3f},{eu:.3f},{eu/fu:.2f}x") |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
