| |
| """Benchmark blockwise FP8 producer APIs.""" |
|
|
| 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 / "blockwise-fp8-producers" / "tests")) |
| from test_blockwise_fp8_producers 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("blockwise_fp8_producers") |
| finally: |
| if artifact: |
| sys.path.remove(artifact) |
|
|
|
|
| def reference(kind, x, weight, bias): |
| if kind == "layer_norm": |
| produced = F.layer_norm( |
| x.float(), (x.shape[1],), weight.float(), bias.float(), 1e-6 |
| ) |
| elif kind == "rms_norm": |
| produced = ( |
| x.float() |
| * torch.rsqrt(x.float().square().mean(-1, keepdim=True) + 1e-6) |
| * weight.float() |
| ) |
| elif kind == "gelu_bias": |
| produced = F.gelu(x.float() + bias.float(), approximate="tanh") |
| else: |
| produced = x.float() |
| blocks = produced.reshape(produced.shape[0], produced.shape[1] // 128, 128) |
| scale = torch.clamp(blocks.abs().amax(-1) / 448.0, min=1.0e-12) |
| quantized = torch.clamp( |
| blocks / scale.unsqueeze(-1), -448.0, 448.0 |
| ).to(torch.float8_e4m3fn) |
| return quantized.reshape_as(x), scale |
|
|
|
|
| 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=30) |
| parser.add_argument("--iters", type=int, default=200) |
| args = parser.parse_args() |
| ops = load_ops(args.backend, args.artifact) |
| shapes = [(51, 4096), (277, 9216), (1024, 1152)] |
| if args.mode == "full": |
| shapes = [(1, 4096), (17, 1152), (51, 4096), (65, 4352), (277, 9216), (1024, 1152)] |
|
|
| print("kind,shape,artifact_us,eager_us,compile_us,eager_speedup,compile_speedup") |
| for rows, dim in shapes: |
| x = torch.randn((rows, dim), device="cuda", dtype=torch.bfloat16) |
| weight = torch.randn((dim,), device="cuda", dtype=torch.bfloat16) |
| bias = torch.randn((dim,), device="cuda", dtype=torch.bfloat16) |
| output = torch.empty_like(x, dtype=torch.float8_e4m3fn) |
| scale = torch.empty((rows, dim // 128), device="cuda", dtype=torch.float32) |
| for kind in ("quantize", "layer_norm", "rms_norm", "gelu_bias"): |
| if kind == "quantize": |
| artifact_fn = lambda: ops.quantize_fp8_block128_bf16( |
| x, output=output, scale=scale |
| ) |
| elif kind == "layer_norm": |
| artifact_fn = lambda: ops.layer_norm_fp8_block128_bf16( |
| x, weight, bias, output=output, scale=scale |
| ) |
| elif kind == "rms_norm": |
| artifact_fn = lambda: ops.rms_norm_fp8_block128_bf16( |
| x, weight, output=output, scale=scale |
| ) |
| else: |
| artifact_fn = lambda: ops.gelu_tanh_bias_fp8_block128_bf16( |
| x, bias, output=output, scale=scale |
| ) |
| eager_fn = lambda: reference(kind, x, weight, bias) |
| torch._dynamo.reset() |
| compiled = torch.compile( |
| lambda a, w, b: reference(kind, a, w, b), fullgraph=True |
| ) |
| compiled_fn = lambda: compiled(x, weight, bias) |
| artifact_us = time_us(artifact_fn, args.warmup, args.iters) |
| eager_us = time_us(eager_fn, max(10, args.warmup // 2), max(50, args.iters // 2)) |
| compile_us = time_us(compiled_fn, max(10, args.warmup // 2), max(50, args.iters // 2)) |
| print( |
| f"{kind},{rows}x{dim},{artifact_us:.3f},{eager_us:.3f}," |
| f"{compile_us:.3f},{eager_us/artifact_us:.2f}x," |
| f"{compile_us/artifact_us:.2f}x" |
| ) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
