| |
| """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()) |
|
|
