| |
| """Benchmark flashrt-vla-residual-gates against PyTorch eager.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import ctypes |
| import ctypes.util |
| import json |
| import math |
| import os |
| import sys |
| from dataclasses import asdict, dataclass |
| from pathlib import Path |
|
|
| import torch |
|
|
|
|
| ROOT = Path(__file__).resolve().parents[2] |
| PACKAGE = ROOT / "flashrt-vla-residual-gates" |
| REGISTRATION_INCLUDE = ( |
| ROOT.parent |
| / "kernels" |
| / "kernel-builder" |
| / "src" |
| / "pyproject" |
| / "templates" |
| / "torch" |
| ) |
|
|
|
|
| SHAPES = { |
| "small": ((64, 8, 4), 1024), |
| "vla_2k": ((2520, 16, 16), 3072), |
| "vla_4k": ((4096, 16, 16), 3072), |
| } |
| SHAPE_GROUPS = { |
| "smoke": ["small"], |
| "headline": ["vla_2k"], |
| "all": list(SHAPES.keys()), |
| } |
|
|
|
|
| @dataclass |
| class Result: |
| shape: str |
| rows: str |
| dim: int |
| flashrt_us: float |
| torch_eager_us: float |
| speedup_vs_eager: float |
| p99_abs: float |
| cosine: float |
| status: str |
|
|
|
|
| class SourceOps: |
| def __init__(self, namespace: str) -> None: |
| self._ops = getattr(torch.ops, namespace) |
|
|
| def joint3_bias_gate_residual_action_nobias_bf16(self, *args): |
| self._ops.joint3_bias_gate_residual_action_nobias_bf16(*args) |
|
|
|
|
| def _preload_cublaslt() -> None: |
| for parent in Path(torch.__file__).resolve().parents: |
| candidate = parent / "nvidia" / "cublas" / "lib" / "libcublasLt.so.12" |
| if candidate.exists(): |
| ctypes.CDLL(str(candidate), mode=ctypes.RTLD_GLOBAL) |
| return |
| library = ctypes.util.find_library("cublasLt") |
| if library: |
| ctypes.CDLL(library, mode=ctypes.RTLD_GLOBAL) |
|
|
|
|
| def _current_arch_list() -> str: |
| major, minor = torch.cuda.get_device_capability(0) |
| return f"{major}.{minor}" |
|
|
|
|
| def load_source_ops() -> SourceOps: |
| from torch.utils.cpp_extension import load |
|
|
| if not REGISTRATION_INCLUDE.is_dir(): |
| raise RuntimeError(f"missing kernel-builder registration include: {REGISTRATION_INCLUDE}") |
| _preload_cublaslt() |
| os.environ.setdefault("TORCH_CUDA_ARCH_LIST", _current_arch_list()) |
| namespace = "flashrt_vla_residual_gates_benchmark" |
| load( |
| name=namespace, |
| sources=[ |
| str(PACKAGE / "torch-ext" / "torch_binding.cpp"), |
| str(PACKAGE / "csrc" / "residual_gates.cu"), |
| ], |
| extra_include_paths=[str(PACKAGE / "csrc"), str(REGISTRATION_INCLUDE)], |
| extra_cflags=["-O3", "-DCUDA_KERNEL"], |
| extra_cuda_cflags=["-O3", "--expt-relaxed-constexpr", "-DCUDA_KERNEL"], |
| verbose=False, |
| ) |
| return SourceOps(namespace) |
|
|
|
|
| def load_installed_ops(artifact: str | None): |
| if artifact: |
| sys.path.insert(0, artifact) |
| try: |
| return importlib.import_module("flashrt_vla_residual_gates") |
| finally: |
| if artifact: |
| sys.path.remove(artifact) |
|
|
|
|
| def make_segment(rows: int, dim: int): |
| residual = torch.randn((rows, dim), device="cuda", dtype=torch.bfloat16) |
| x = torch.randn_like(residual) |
| gate = torch.randn_like(residual) |
| bias = (0.02 * torch.randn((dim,), device="cuda", dtype=torch.bfloat16)).contiguous() |
| out = torch.empty_like(residual) |
| return residual, x, bias, gate, out |
|
|
|
|
| def make_case(rows: tuple[int, int, int], dim: int): |
| v = make_segment(rows[0], dim) |
| a = make_segment(rows[1], dim) |
| u_residual = torch.randn((rows[2], dim), device="cuda", dtype=torch.bfloat16) |
| u_x = torch.randn_like(u_residual) |
| u_out = torch.empty_like(u_residual) |
| return v, a, (u_residual, u_x, u_out) |
|
|
|
|
| def torch_ref(v, a, u): |
| v_ref = (v[0].float() + (v[1].float() + v[2].float().view(1, -1)) * v[3].float()).to(torch.bfloat16) |
| a_ref = (a[0].float() + a[1].float() * a[3].float()).to(torch.bfloat16) |
| u_ref = (u[0].float() + u[1].float()).to(torch.bfloat16) |
| return v_ref, a_ref, u_ref |
|
|
|
|
| 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 percentile(x: torch.Tensor, q: float) -> torch.Tensor: |
| flat = x.flatten() |
| k = max(1, min(flat.numel(), math.ceil(q * flat.numel()))) |
| return flat.kthvalue(k).values |
|
|
|
|
| def metrics(got_parts, expected_parts): |
| got = torch.cat([x.flatten() for x in got_parts]).float() |
| expected = torch.cat([x.flatten() for x in expected_parts]).float() |
| diff = (got - expected).abs() |
| return float(percentile(diff, 0.99).item()), float( |
| torch.nn.functional.cosine_similarity(got, expected, dim=0).item() |
| ) |
|
|
|
|
| def run_one(ops, name: str, rows: tuple[int, int, int], dim: int, args) -> Result: |
| v, a, u = make_case(rows, dim) |
| ops.joint3_bias_gate_residual_action_nobias_bf16( |
| v[0], v[1], v[2], v[3], v[4], |
| a[0], a[1], a[3], a[4], |
| u[0], u[1], u[2], |
| ) |
| expected = torch_ref(v, a, u) |
| p99_abs, cosine = metrics((v[4], a[4], u[2]), expected) |
| flashrt_us = time_us( |
| lambda: ops.joint3_bias_gate_residual_action_nobias_bf16( |
| v[0], v[1], v[2], v[3], v[4], |
| a[0], a[1], a[3], a[4], |
| u[0], u[1], u[2], |
| ), |
| args.warmup, |
| args.iters, |
| ) |
| eager_us = time_us(lambda: torch_ref(v, a, u), args.warmup, args.iters) |
| status = "PASS" if p99_abs <= args.p99_abs_limit else "FAIL" |
| return Result( |
| shape=name, |
| rows=f"{rows[0]},{rows[1]},{rows[2]}", |
| dim=dim, |
| flashrt_us=flashrt_us, |
| torch_eager_us=eager_us, |
| speedup_vs_eager=eager_us / flashrt_us, |
| p99_abs=p99_abs, |
| cosine=cosine, |
| status=status, |
| ) |
|
|
|
|
| def write_markdown(path: Path, results: list[Result]) -> None: |
| lines = [ |
| "# Source Benchmark Results", |
| "", |
| "Environment: NVIDIA GeForce RTX 5090 local source-extension build.", |
| "Baseline: PyTorch eager tensor postprocess chain with matching BF16 math.", |
| "", |
| "| Shape | V,A,U rows | Dim | FlashRT us | Eager us | vs eager | p99 abs | Cosine | Status |", |
| "|---|---:|---:|---:|---:|---:|---:|---:|---|", |
| ] |
| for r in results: |
| lines.append( |
| f"| {r.shape} | {r.rows} | {r.dim} | {r.flashrt_us:.3f} | " |
| f"{r.torch_eager_us:.3f} | {r.speedup_vs_eager:.2f}x | " |
| f"{r.p99_abs:.6f} | {r.cosine:.8f} | {r.status} |" |
| ) |
| path.write_text("\n".join(lines) + "\n") |
|
|
|
|
| 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", choices=sorted(SHAPE_GROUPS), default="smoke") |
| parser.add_argument("--warmup", type=int, default=5) |
| parser.add_argument("--iters", type=int, default=20) |
| parser.add_argument("--p99-abs-limit", type=float, default=0.0) |
| parser.add_argument("--output", default=None) |
| parser.add_argument("--markdown", default=None) |
| args = parser.parse_args() |
|
|
| if not torch.cuda.is_available(): |
| raise SystemExit("CUDA is required") |
| torch.manual_seed(43) |
| ops = load_source_ops() if args.backend == "source" else load_installed_ops(args.artifact) |
| results = [run_one(ops, name, *SHAPES[name], args) for name in SHAPE_GROUPS[args.shapes]] |
| for r in results: |
| print( |
| f"{r.status} {r.shape}: flashrt={r.flashrt_us:.3f}us eager={r.torch_eager_us:.3f}us " |
| f"speedup={r.speedup_vs_eager:.2f}x p99={r.p99_abs:.6f}" |
| ) |
| if args.output: |
| Path(args.output).parent.mkdir(parents=True, exist_ok=True) |
| Path(args.output).write_text(json.dumps([asdict(r) for r in results], indent=2) + "\n") |
| if args.markdown: |
| Path(args.markdown).parent.mkdir(parents=True, exist_ok=True) |
| write_markdown(Path(args.markdown), results) |
| if any(r.status != "PASS" for r in results): |
| raise SystemExit(1) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
