| |
| """Benchmark generic FP8 linear+bias projection paths.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import importlib |
| import importlib.util |
| import json |
| 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 = { |
| "decode_m1_d2048_o": (1, 2048, 2048), |
| "small_m8_d1536_o": (8, 1536, 1536), |
| "groot_dit_m51_qkv": (51, 1536, 4608), |
| "groot_dit_m51_o": (51, 1536, 1536), |
| "mid_m64_d2048_o": (64, 2048, 2048), |
| "groot_backbone_m105_qkv": (105, 2048, 4096), |
| "groot_backbone_m105_o": (105, 2048, 2048), |
| "mid_m128_siglip_qkv": (128, 1152, 3456), |
| "prefill_m256_d1536_o": (256, 1536, 1536), |
| "prefill_m512_d2048_o": (512, 2048, 2048), |
| } |
|
|
|
|
| def load_test_helpers(): |
| spec = importlib.util.spec_from_file_location("flashrt_fp8_ffn_test_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 load_ops(backend: str, artifact: str | None): |
| helpers = load_test_helpers() |
| if backend == "source": |
| return helpers.load_source_ops(), helpers |
| if artifact: |
| sys.path.insert(0, artifact) |
| try: |
| return importlib.import_module("flashrt_fp8_ffn"), helpers |
| finally: |
| if artifact: |
| sys.path.remove(artifact) |
|
|
|
|
| def cuda_time_us(fn, warmup: int, iters: int, rounds: int) -> 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(iters): |
| fn() |
| end.record() |
| end.synchronize() |
| samples.append(start.elapsed_time(end) * 1000.0 / iters) |
| return statistics.median(samples) |
|
|
|
|
| def graph_runner(fn): |
| fn() |
| torch.cuda.synchronize() |
| graph = torch.cuda.CUDAGraph() |
| with torch.cuda.graph(graph): |
| fn() |
| graph.replay() |
| torch.cuda.synchronize() |
| return graph.replay |
|
|
|
|
| def maybe_fvk_runner(x_fp8, w_fp8, bias, out, x_scale, w_scale, M, N, K): |
| try: |
| from flash_rt import flash_rt_kernels as fvk |
| except ImportError: |
| return None, "flash_rt.flash_rt_kernels is unavailable" |
| gemm = fvk.GemmRunner() |
| alpha = float((x_scale * w_scale).item()) |
| stream = torch.cuda.current_stream().cuda_stream |
|
|
| def run(): |
| gemm.fp8_nn_bias_bf16( |
| x_fp8.data_ptr(), w_fp8.data_ptr(), out.data_ptr(), bias.data_ptr(), |
| M, N, K, alpha, stream, |
| ) |
|
|
| try: |
| run() |
| except RuntimeError as error: |
| return None, str(error) |
| return run, None |
|
|
|
|
| def maybe_fvk_fp16_runners( |
| x_fp8, w_fp8, bias_fp16, out_fp16, x_scale, w_scale, M, N, K |
| ): |
| try: |
| from flash_rt import flash_rt_kernels as fvk |
| except ImportError: |
| return None, None, "flash_rt.flash_rt_kernels is unavailable" |
| gemm = fvk.GemmRunner() |
| alpha = float((x_scale * w_scale).item()) |
| stream = torch.cuda.current_stream().cuda_stream |
|
|
| def fused(): |
| gemm.fp8_nn_bias( |
| x_fp8.data_ptr(), w_fp8.data_ptr(), out_fp16.data_ptr(), |
| bias_fp16.data_ptr(), M, N, K, alpha, stream, |
| ) |
|
|
| def decomposed(): |
| gemm.fp8_descale_fp16( |
| x_fp8.data_ptr(), w_fp8.data_ptr(), out_fp16.data_ptr(), M, N, K, |
| x_scale.data_ptr(), w_scale.data_ptr(), stream, |
| ) |
| fvk.add_bias_fp16( |
| out_fp16.data_ptr(), bias_fp16.data_ptr(), M, N, stream |
| ) |
|
|
| try: |
| fused() |
| fused_runner = fused |
| fused_error = None |
| except RuntimeError as error: |
| fused_runner = None |
| fused_error = str(error) |
| decomposed() |
| return fused_runner, decomposed, fused_error |
|
|
|
|
| def metrics(got: torch.Tensor, expected: torch.Tensor) -> dict[str, float]: |
| got_f = got.float().flatten() |
| exp_f = expected.float().flatten() |
| diff = (got_f - exp_f).abs() |
| return { |
| "max_abs": float(diff.max().item()), |
| "mean_abs": float(diff.mean().item()), |
| "p99_abs": float(torch.quantile(diff, 0.99).item()), |
| "cosine": float(torch.nn.functional.cosine_similarity(got_f, exp_f, dim=0).item()), |
| } |
|
|
|
|
| def run(args) -> dict: |
| if not torch.cuda.is_available(): |
| raise SystemExit("CUDA is required") |
| torch.manual_seed(20260718) |
| ops, helpers = load_ops(args.backend, args.artifact) |
| labels = list(SHAPES) if args.shapes == "all" else args.shapes.split(",") |
| results = [] |
|
|
| for label in labels: |
| if label not in SHAPES: |
| raise ValueError(f"unknown shape {label!r}") |
| M, K, N = SHAPES[label] |
| x_bf16 = torch.randn((M, K), device="cuda", dtype=torch.bfloat16) * 0.25 |
| w_bf16 = torch.randn((N, K), device="cuda", dtype=torch.bfloat16) * (K**-0.5) |
| bias = torch.randn((N,), device="cuda", dtype=torch.bfloat16) * 0.01 |
| x_scale = ( |
| x_bf16.float().abs().max() / (0.9 * helpers.fp8_max()) |
| ).clamp_min(1e-6).reshape(1) |
| w_scale = ( |
| w_bf16.float().abs().max() / (0.9 * helpers.fp8_max()) |
| ).clamp_min(1e-6).reshape(1) |
| x_fp8 = helpers.quantize_fp8_reciprocal(x_bf16, x_scale) |
| w_fp8 = helpers.quantize_fp8(w_bf16, w_scale) |
| out = torch.empty((M, N), device="cuda", dtype=torch.bfloat16) |
| region_out = torch.empty_like(out) |
| input_fp8 = torch.empty_like(x_fp8) |
| fvk_out = torch.empty_like(out) |
| fvk_fp16_out = torch.empty((M, N), device="cuda", dtype=torch.float16) |
| bias_fp16 = bias.to(torch.float16) |
|
|
| def package_fp8(): |
| ops.fp8_linear_bias_bf16( |
| x_fp8, w_fp8, bias, x_scale, w_scale, out=out |
| ) |
|
|
| def package_region(): |
| ops.bf16_fp8_linear_bias_bf16( |
| x_bf16, w_fp8, bias, x_scale, w_scale, |
| input_fp8=input_fp8, out=region_out, pad_to=M, |
| ) |
|
|
| reference = helpers.ref_linear_bias( |
| x_fp8, w_fp8, bias, x_scale, w_scale |
| ) |
| package_fp8() |
| package_metrics = metrics(out, reference) |
| if package_metrics["p99_abs"] > 0.015625 or package_metrics["cosine"] < 0.9999: |
| raise AssertionError(f"{label} package correctness failed: {package_metrics}") |
| package_region() |
| if not torch.equal(region_out, out) or not torch.equal(input_fp8, x_fp8): |
| raise AssertionError(f"{label} BF16 region does not match FP8 entry") |
|
|
| graph = graph_runner(package_region) |
| if args.compare_fvk: |
| fvk, fvk_error = maybe_fvk_runner( |
| x_fp8, w_fp8, bias, fvk_out, x_scale, w_scale, M, N, K |
| ) |
| else: |
| fvk, fvk_error = None, "not requested" |
| if args.compare_fvk: |
| fvk_fp16_fused, fvk_fp16_decomposed, fvk_fp16_error = ( |
| maybe_fvk_fp16_runners( |
| x_fp8, w_fp8, bias_fp16, fvk_fp16_out, x_scale, w_scale, |
| M, N, K, |
| ) |
| ) |
| else: |
| fvk_fp16_fused = None |
| fvk_fp16_decomposed = None |
| fvk_fp16_error = "not requested" |
| if fvk is not None: |
| fvk() |
| fvk_metrics = metrics(fvk_out, reference) |
| if fvk_metrics["p99_abs"] > 0.015625 or fvk_metrics["cosine"] < 0.9999: |
| raise AssertionError(f"{label} FVK correctness failed: {fvk_metrics}") |
| else: |
| fvk_metrics = None |
|
|
| def eager(value): |
| return torch.addmm(bias, value, w_bf16.T) |
|
|
| compiled = torch.compile(eager, fullgraph=True) if args.compile_baseline else None |
| if compiled is not None: |
| compiled_out = compiled(x_bf16) |
| changed_out = compiled(x_bf16 + torch.ones_like(x_bf16) * 0.125) |
| torch.cuda.synchronize() |
| if torch.equal(compiled_out, changed_out): |
| raise AssertionError( |
| f"{label} compiled baseline did not respond to changed input" |
| ) |
|
|
| package_us = cuda_time_us(package_fp8, args.warmup, args.iters, args.rounds) |
| region_us = cuda_time_us(package_region, args.warmup, args.iters, args.rounds) |
| graph_us = cuda_time_us(graph, args.warmup, args.iters, args.rounds) |
| eager_call = lambda: eager(x_bf16) |
| compiled_call = (lambda: compiled(x_bf16)) if compiled is not None else None |
| eager_us = cuda_time_us(eager_call, args.warmup, args.iters, args.rounds) |
| compiled_us = ( |
| cuda_time_us(compiled_call, args.warmup, args.iters, args.rounds) |
| if compiled_call is not None else None |
| ) |
| fvk_us = ( |
| cuda_time_us(fvk, args.warmup, args.iters, args.rounds) |
| if fvk is not None else None |
| ) |
| fvk_fp16_fused_us = ( |
| cuda_time_us( |
| fvk_fp16_fused, args.warmup, args.iters, args.rounds |
| ) |
| if fvk_fp16_fused is not None else None |
| ) |
| fvk_fp16_decomposed_us = ( |
| cuda_time_us( |
| fvk_fp16_decomposed, args.warmup, args.iters, args.rounds |
| ) |
| if fvk_fp16_decomposed is not None else None |
| ) |
| row = { |
| "shape": label, |
| "M": M, |
| "K": K, |
| "N": N, |
| "fp8_linear_bias_us": package_us, |
| "bf16_region_us": region_us, |
| "bf16_region_graph_us": graph_us, |
| "torch_bf16_eager_us": eager_us, |
| "torch_bf16_compiled_us": compiled_us, |
| "fvk_fp8_nn_bias_bf16_us": fvk_us, |
| "fvk_status": "ok" if fvk is not None else "unsupported", |
| "fvk_error": fvk_error, |
| "fvk_fp16_fused_us": fvk_fp16_fused_us, |
| "fvk_fp16_decomposed_us": fvk_fp16_decomposed_us, |
| "fvk_fp16_fused_error": fvk_fp16_error, |
| "package_vs_eager": eager_us / package_us, |
| "region_vs_eager": eager_us / region_us, |
| "graph_vs_eager": eager_us / graph_us, |
| "package_vs_fvk": fvk_us / package_us if fvk_us is not None else None, |
| "package_vs_fvk_fp16_fused": ( |
| fvk_fp16_fused_us / package_us |
| if fvk_fp16_fused_us is not None else None |
| ), |
| "package_vs_fvk_fp16_decomposed": ( |
| fvk_fp16_decomposed_us / package_us |
| if fvk_fp16_decomposed_us is not None else None |
| ), |
| "package_correctness": package_metrics, |
| "fvk_correctness": fvk_metrics, |
| } |
| results.append(row) |
| print(json.dumps(row, sort_keys=True)) |
|
|
| return { |
| "backend": args.backend, |
| "device": torch.cuda.get_device_name(0), |
| "capability": list(torch.cuda.get_device_capability(0)), |
| "torch": torch.__version__, |
| "cuda": torch.version.cuda, |
| "warmup": args.warmup, |
| "iters": args.iters, |
| "rounds": args.rounds, |
| "results": results, |
| } |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--backend", choices=["source", "installed"], default="source") |
| parser.add_argument("--artifact", default=None) |
| parser.add_argument("--shapes", default="all") |
| parser.add_argument("--warmup", type=int, default=20) |
| parser.add_argument("--iters", type=int, default=100) |
| parser.add_argument("--rounds", type=int, default=5) |
| parser.add_argument("--compile-baseline", action="store_true") |
| parser.add_argument("--compare-fvk", action="store_true") |
| parser.add_argument("--output", type=Path, default=None) |
| args = parser.parse_args() |
| payload = run(args) |
| if args.output is not None: |
| args.output.parent.mkdir(parents=True, exist_ok=True) |
| args.output.write_text(json.dumps(payload, indent=2) + "\n") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
