| |
| """Benchmark fp4-gemm.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import importlib.util |
| import json |
| import os |
| import sys |
| from dataclasses import asdict, dataclass |
| from pathlib import Path |
|
|
| import torch |
|
|
|
|
| ROOT = Path(__file__).resolve().parents[2] |
| TEST_FILE = ROOT / "fp4-gemm" / "tests" / "test_fp4_gemm.py" |
|
|
|
|
| @dataclass |
| class BenchResult: |
| shape: str |
| M: int |
| N: int |
| K: int |
| variant: int |
| native_us: float |
| flashrt_us: float |
| torch_eager_us: float |
| torch_compile_us: float |
| speedup_vs_eager: float |
| speedup_vs_compile: float |
| wrapper_over_native: float |
| max_abs: float |
| mean_abs: float |
| p99_abs: float |
| cosine: float |
| status: str |
|
|
|
|
| def load_helpers(): |
| spec = importlib.util.spec_from_file_location("fp4_gemm_test_helpers", TEST_FILE) |
| if spec is None or spec.loader is None: |
| raise RuntimeError(f"cannot load helpers from {TEST_FILE}") |
| module = importlib.util.module_from_spec(spec) |
| sys.modules["fp4_gemm_test_helpers"] = module |
| spec.loader.exec_module(module) |
| return module |
|
|
|
|
| def measure(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 float(start.elapsed_time(end) * 1000.0 / iters) |
|
|
|
|
| def bench_case(helpers, ops, native, name: str, shape: tuple[int, int, int], warmup: int, iters: int) -> list[BenchResult]: |
| m, n, k = shape |
| a_packed, b_packed, sfa, sfb, expected = helpers.prepare_quantized(ops, m, n, k) |
| a_deq = torch.empty((m, k), device="cuda", dtype=torch.float16) |
| b_deq = torch.empty((n, k), device="cuda", dtype=torch.float16) |
| ops.dequantize_fp4_sfa_fp16(a_packed, sfa, a_deq, False) |
| ops.dequantize_fp4_sfa_fp16(b_packed, sfb, b_deq, True) |
| torch.cuda.synchronize() |
|
|
| def torch_ref(): |
| return (a_deq.float() @ b_deq.float().T).to(torch.bfloat16) |
|
|
| torch_eager_us = measure(torch_ref, warmup, iters) |
| compiled_ref = torch.compile(torch_ref, mode="max-autotune-no-cudagraphs") |
| torch_compile_us = measure(compiled_ref, warmup, iters) |
| stream = torch.cuda.current_stream().cuda_stream |
| results: list[BenchResult] = [] |
| variants = (-1, 0, 1, 2) if torch.cuda.get_device_capability(0) == (11, 0) else (0, 1, 2) |
| for variant in variants: |
| out = torch.empty((m, n), device="cuda", dtype=torch.bfloat16) |
| ops.nvfp4_gemm_bf16(a_packed, b_packed, sfa, sfb, out, 1.0, variant) |
| torch.cuda.synchronize() |
| max_abs, mean_abs, p99_abs, cosine = helpers.metrics(out, expected) |
| flashrt_us = measure( |
| lambda: ops.nvfp4_gemm_bf16(a_packed, b_packed, sfa, sfb, out, 1.0, variant), |
| warmup, |
| iters, |
| ) |
| native_variant = variant |
| if native_variant < 0: |
| native_variant = helpers.select_sm110_variant(shape) |
| native_function = ( |
| native.fp4_w4a16_gemm_sm120_bf16out |
| if native_variant == 0 |
| else native.fp4_w4a16_gemm_sm120_bf16out_widen |
| if native_variant == 1 |
| else native.fp4_w4a16_gemm_sm120_bf16out_pingpong |
| ) |
| native_us = measure( |
| lambda: native_function( |
| a_packed.data_ptr(), |
| b_packed.data_ptr(), |
| out.data_ptr(), |
| m, |
| n, |
| k, |
| sfa.data_ptr(), |
| sfb.data_ptr(), |
| 1.0, |
| stream, |
| ), |
| warmup, |
| iters, |
| ) |
| results.append( |
| BenchResult( |
| shape=name, |
| M=m, |
| N=n, |
| K=k, |
| variant=variant, |
| native_us=native_us, |
| flashrt_us=flashrt_us, |
| torch_eager_us=torch_eager_us, |
| torch_compile_us=torch_compile_us, |
| speedup_vs_eager=torch_eager_us / flashrt_us, |
| speedup_vs_compile=torch_compile_us / flashrt_us, |
| wrapper_over_native=flashrt_us / native_us, |
| max_abs=max_abs, |
| mean_abs=mean_abs, |
| p99_abs=p99_abs, |
| cosine=cosine, |
| status="ok", |
| ) |
| ) |
| return results |
|
|
|
|
| def bench_bf16_producer(ops, native, k: int, warmup: int, iters: int): |
| x = torch.randn((1, k), device="cuda", dtype=torch.bfloat16) |
| direct_packed, direct_sfa = ops.alloc_fp4(1, k) |
| compat_packed, compat_sfa = ops.alloc_fp4(1, k) |
| native_packed, native_sfa = ops.alloc_fp4(1, k) |
| stream = torch.cuda.current_stream().cuda_stream |
|
|
| def direct(): |
| ops.quantize_fp4_sfa_bf16( |
| x, direct_packed, direct_sfa, False |
| ) |
|
|
| def compat(): |
| ops.quantize_fp4_sfa_fp16( |
| x.to(torch.float16), compat_packed, compat_sfa, False |
| ) |
|
|
| def native_direct(): |
| native.quantize_bf16_to_nvfp4_swizzled( |
| x.data_ptr(), native_packed.data_ptr(), native_sfa.data_ptr(), |
| 1, k, stream, |
| ) |
|
|
| direct() |
| compat() |
| torch.cuda.synchronize() |
| direct_us = measure(direct, warmup, iters) |
| compat_us = measure(compat, warmup, iters) |
| native_us = measure(native_direct, warmup, iters) |
| return { |
| "M": 1, |
| "K": k, |
| "direct_bf16_us": direct_us, |
| "cast_plus_fp16_us": compat_us, |
| "native_bf16_us": native_us, |
| "speedup_vs_cast_plus_fp16": compat_us / direct_us, |
| "wrapper_over_native": direct_us / native_us, |
| "packed_exact_vs_fp16_contract": bool( |
| torch.equal(direct_packed, compat_packed) |
| ), |
| "note": "native_bf16 uses a distinct FlashRT quantization strategy", |
| } |
|
|
|
|
| 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=["smoke", "headline", "thor-models"], default="headline" |
| ) |
| parser.add_argument("--warmup", type=int, default=20) |
| parser.add_argument("--iterations", type=int, default=100) |
| parser.add_argument("--json-out", default=None) |
| args = parser.parse_args() |
|
|
| helpers = load_helpers() |
| native_root = Path( |
| os.environ.get("FLASHRT_NATIVE_ROOT", str(ROOT.parent / "official" / "FlashRT")) |
| ) |
| sys.path.insert(0, str(native_root)) |
| try: |
| import flash_rt.flash_rt_kernels as native |
| finally: |
| sys.path.pop(0) |
| ops = ( |
| helpers.load_source_ops() |
| if args.backend == "source" |
| else helpers.load_installed_ops(args.artifact) |
| ) |
| shapes = { |
| "small_m16_n128_k128": (16, 128, 128), |
| "small_m32_n256_k256": (32, 256, 256), |
| "mlp_tile_m64_n512_k512": (64, 512, 512), |
| "groot_dit_projection": (51, 1536, 1536), |
| "vla_projection": (105, 2048, 2048), |
| "motus_up": (360, 14336, 3072), |
| "motus_down": (360, 3072, 14336), |
| } |
| if args.mode == "smoke": |
| shapes = {"small_m16_n128_k128": shapes["small_m16_n128_k128"]} |
| elif args.mode == "thor-models": |
| shapes = dict(helpers.SM110_SHAPES) |
| results: list[BenchResult] = [] |
| for name, shape in shapes.items(): |
| results.extend( |
| bench_case( |
| helpers, ops, native, name, shape, args.warmup, args.iterations |
| ) |
| ) |
| producer_results = [ |
| bench_bf16_producer(ops, native, k, args.warmup, args.iterations) |
| for k in (5120, 6144, 17408) |
| ] |
| payload = { |
| "mode": args.mode, |
| "backend": args.backend, |
| "device": torch.cuda.get_device_name(), |
| "torch": torch.__version__, |
| "results": [asdict(item) for item in results], |
| "bf16_producer_results": producer_results, |
| } |
| print(json.dumps(payload, indent=2)) |
| if args.json_out: |
| out = Path(args.json_out) |
| out.parent.mkdir(parents=True, exist_ok=True) |
| out.write_text(json.dumps(payload, indent=2) + "\n") |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
