File size: 5,060 Bytes
1db1507 | 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 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 | #!/usr/bin/env python3
"""Compare staged-bias v1 and fused-bias v2 FP8 GELU MLP entries."""
from __future__ import annotations
import argparse
import importlib
import importlib.util
import json
import os
import statistics
import sys
from pathlib import Path
import torch
ROOT = Path(__file__).resolve().parents[2]
TEST_FILE = ROOT / "flashrt-fp8-ffn" / "tests" / "test_fp8_ffn.py"
SHAPES = {
"groot_vit": (128, 1024, 4096, 1024),
"groot_deepstack": (128, 4096, 4096, 2048),
"groot_action_dit": (41, 1536, 6144, 1536),
"pi05_decoder": (10, 1024, 4096, 1024),
}
def helpers_module():
spec = importlib.util.spec_from_file_location("fp8_ffn_helpers", TEST_FILE)
if spec is None or spec.loader is None:
raise RuntimeError(f"cannot load {TEST_FILE}")
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
return module
def time_us(fn, warmup: int, iterations: int, rounds: int = 5) -> float:
samples = []
for _ in range(rounds):
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(iterations):
fn()
end.record()
end.synchronize()
samples.append(start.elapsed_time(end) * 1000.0 / iterations)
return float(statistics.median(samples))
def main() -> int:
parser = argparse.ArgumentParser()
parser.add_argument("--backend", choices=("source", "installed"), default="source")
parser.add_argument("--artifact")
parser.add_argument("--warmup", type=int, default=20)
parser.add_argument("--iterations", type=int, default=100)
parser.add_argument("--json-out")
args = parser.parse_args()
helpers = helpers_module()
if args.backend == "source":
ops = helpers.load_source_ops()
else:
if args.artifact:
sys.path.insert(0, args.artifact)
try:
ops = importlib.import_module("flashrt_fp8_ffn")
finally:
if args.artifact:
sys.path.remove(args.artifact)
rows = []
for name, shape in SHAPES.items():
tensors = helpers.make_case(*shape)
m, _, h, n = shape
hidden_v1 = torch.empty((m, h), device="cuda", dtype=torch.bfloat16)
hidden_fp8_v1 = torch.empty_like(hidden_v1, dtype=helpers.fp8_dtype())
out_v1_buf = torch.empty((m, n), device="cuda", dtype=torch.bfloat16)
hidden_v2 = torch.empty_like(hidden_v1)
hidden_fp8_v2 = torch.empty_like(hidden_fp8_v1)
out_v2_buf = torch.empty_like(out_v1_buf)
if args.backend == "source":
v1 = lambda: ops.fp8_gelu_mlp_bf16(
*tensors, hidden=hidden_v1, hidden_fp8=hidden_fp8_v1,
out=out_v1_buf,
)
v2 = lambda: ops.fp8_gelu_mlp_v2_bf16(
*tensors, hidden=hidden_v2, hidden_fp8=hidden_fp8_v2,
out=out_v2_buf,
)
else:
v1 = lambda: ops.fp8_gelu_mlp_bf16(
*tensors, hidden_bf16=hidden_v1, hidden_fp8=hidden_fp8_v1,
out=out_v1_buf,
)
v2 = lambda: ops.fp8_gelu_mlp_v2_bf16(
*tensors, hidden_bf16=hidden_v2, hidden_fp8=hidden_fp8_v2,
out=out_v2_buf,
)
out_v1 = v1()
out_v2 = v2()
torch.cuda.synchronize()
old = os.environ.get("FLASHRT_FP8_FFN_REQUIRE_BIAS_EPILOGUE")
os.environ["FLASHRT_FP8_FFN_REQUIRE_BIAS_EPILOGUE"] = "1"
try:
v2()
torch.cuda.synchronize()
fused_hit = True
except RuntimeError as error:
if "BIAS epilogue was required" not in str(error):
raise
fused_hit = False
finally:
if old is None:
os.environ.pop("FLASHRT_FP8_FFN_REQUIRE_BIAS_EPILOGUE", None)
else:
os.environ["FLASHRT_FP8_FFN_REQUIRE_BIAS_EPILOGUE"] = old
metrics = helpers.distribution_metrics(out_v2, out_v1)
v1_us = time_us(v1, args.warmup, args.iterations)
v2_us = time_us(v2, args.warmup, args.iterations)
row = {
"shape": name,
"M": shape[0],
"K": shape[1],
"H": shape[2],
"N": shape[3],
"fused_bias_epilogue": fused_hit,
"v1_us": v1_us,
"v2_us": v2_us,
"speedup": v1_us / v2_us,
**metrics,
}
rows.append(row)
print(json.dumps(row, sort_keys=True), flush=True)
result = {
"device": torch.cuda.get_device_name(),
"capability": list(torch.cuda.get_device_capability()),
"torch": torch.__version__,
"cuda": torch.version.cuda,
"rows": rows,
}
if args.json_out:
Path(args.json_out).write_text(json.dumps(result, indent=2) + "\n")
return 0
if __name__ == "__main__":
raise SystemExit(main())
|