File size: 4,672 Bytes
fa85b55 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 | #!/usr/bin/env python3
"""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 # noqa: E402
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())
|