benchmarks/benchmark.py

#!/usr/bin/env python3
"""Benchmark fp8-gemm."""

from __future__ import annotations

import argparse
import importlib
import json
import os
import statistics
import sys
from dataclasses import asdict, dataclass
from pathlib import Path

import torch


ROOT = Path(__file__).resolve().parents[2]
PACKAGE = ROOT / "fp8-gemm"
REGISTRATION_INCLUDE = (
    ROOT.parent
    / "kernels"
    / "kernel-builder"
    / "src"
    / "pyproject"
    / "templates"
    / "torch"
)

SHAPES = {
    "decode_m1_k4096_n2048": (1, 4096, 2048),
    "decode_m1_k4096_n8192": (1, 4096, 8192),
    "small_m16_k4096_n4096": (16, 4096, 4096),
    "small_m32_k4096_n8192": (32, 4096, 8192),
    "small_m64_k512_n1024": (64, 512, 1024),
    "pi05_action_qkv": (51, 2048, 2560),
    "pi05_action_o": (51, 2048, 2048),
    "pi05_action_gate_up": (51, 2048, 16384),
    "pi05_action_down": (51, 8192, 2048),
    "groot_dit_qkv": (51, 1536, 4608),
    "groot_n17_llm_o": (277, 2048, 2048),
    "groot_n17_llm_gate_up": (277, 2048, 16384),
    "groot_n17_llm_down": (277, 8192, 2048),
    "groot_n17_vit_o": (1024, 1024, 1024),
    "cosmos_edge_action": (64, 2048, 9216),
    "lingbot_vision_o": (1024, 1280, 1280),
    "lingbot_action_gate_up": (105, 2048, 16384),
}

MODES = {
    "smoke": ["decode_m1_k4096_n2048", "small_m16_k4096_n4096"],
    "headline": [
        "decode_m1_k4096_n2048",
        "pi05_action_qkv",
        "pi05_action_gate_up",
        "pi05_action_down",
        "groot_n17_llm_o",
        "cosmos_edge_action",
        "lingbot_action_gate_up",
    ],
    "thor-full": list(SHAPES),
}


@dataclass
class Result:
    shape: str
    M: int
    K: int
    N: int
    variant: int
    tile: str
    flashrt_us: float
    flashrt_graph_us: float | None
    package_best_graph_tile: str | None
    package_best_graph_us: float | None
    auto_over_package_best: float | None
    package_tiles_us: dict[str, dict[str, float]] | None
    native_best_tile: str | None
    native_best_us: float | None
    native_best_graph_tile: str | None
    native_best_graph_us: float | None
    native_tiles_us: dict[str, dict[str, float]] | None
    wrapper_vs_native: float | None
    graph_vs_native: float | None
    torch_eager_us: float
    torch_compile_us: float | None
    speedup_vs_eager: float
    speedup_vs_compile: float | None
    max_abs: float
    mean_abs: float
    p99_abs: float
    cosine: float
    status: str


class SourceOps:
    def __init__(self, namespace: str) -> None:
        self._ops = getattr(torch.ops, namespace)

    @staticmethod
    def select_fp8_linear_tile(m: int, n: int, k: int, variant: int = 0) -> str:
        return select_tile(m, n, k, variant)

    def fp8_linear_bf16(self, x, w, alpha=1.0, out=None, variant=0):
        if out is None:
            out = torch.empty((x.shape[0], w.shape[0]), device=x.device, dtype=torch.bfloat16)
        self._ops.fp8_linear_bf16(x, w, float(alpha), int(variant), out)
        return out


def _current_arch_list() -> str:
    major, minor = torch.cuda.get_device_capability(0)
    if (major, minor) == (11, 0):
        return "11.0a"
    if major >= 12:
        return "12.0a"
    return f"{major}.{minor}"


def load_source_ops() -> SourceOps:
    from torch.utils.cpp_extension import load

    os.environ.setdefault("TORCH_CUDA_ARCH_LIST", _current_arch_list())
    namespace = "fp8_gemm_source_bench"
    capability = torch.cuda.get_device_capability(0)
    cutlass_include = Path(os.environ.get("CUTLASS_INCLUDE", ""))
    if capability == (11, 0):
        if not (cutlass_include / "cutlass" / "cutlass.h").is_file():
            raise RuntimeError("set CUTLASS_INCLUDE for the SM110 source benchmark")
        cuda_sources = [str(PACKAGE / "csrc" / "cutlass_sm110_fp8_gemm.cu")]
        source_define = "-DFLASHRT_FP8_GEMM_SOURCE_SM110_ONLY"
        extra_includes = [
            str(cutlass_include),
            str(cutlass_include.parent / "tools" / "util" / "include"),
        ]
    else:
        cuda_sources = [
            str(PACKAGE / "csrc" / "fp8_gemv_m1_sm120.cu"),
            str(PACKAGE / "csrc" / "fp8_smallM_handtuned_sm120.cu"),
            str(PACKAGE / "csrc" / "fp8_smallM_handtuned_ldmatrix_sm120.cu"),
        ]
        source_define = "-DFLASHRT_FP8_GEMM_SOURCE_SM120_ONLY"
        extra_includes = []
    load(
        name=namespace,
        sources=[str(PACKAGE / "torch-ext" / "torch_binding.cpp"), *cuda_sources],
        extra_include_paths=[
            str(PACKAGE / "csrc"), str(REGISTRATION_INCLUDE), *extra_includes
        ],
        extra_cflags=["-O3", "-DNDEBUG", "-DCUDA_KERNEL", source_define],
        extra_cuda_cflags=[
            "-O3", "-DNDEBUG", "--expt-relaxed-constexpr", "--use_fast_math",
            "-U__CUDA_NO_HALF_OPERATORS__",
            "-U__CUDA_NO_HALF_CONVERSIONS__",
            "-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
            "-U__CUDA_NO_HALF2_OPERATORS__",
            "-DCUDA_KERNEL", source_define
        ],
        verbose=False,
    )
    return SourceOps(namespace)


def load_installed_ops(artifact: str | None):
    if artifact:
        sys.path.insert(0, artifact)
    try:
        return importlib.import_module("fp8_gemm")
    finally:
        if artifact:
            sys.path.remove(artifact)


def select_tile(m: int, n: int, k: int, variant: int = 0) -> str:
    if torch.cuda.get_device_capability(0) == (11, 0):
        forced = {1: "sm110_sq_bf16", 2: "sm110_t1_bf16", 3: "sm110_wide_bf16"}
        if variant not in {0, *forced}:
            raise RuntimeError("SM110 variant must be in [0, 3]")
        if variant:
            return forced[variant]
        if n >= 8 * k:
            return "sm110_wide_bf16"
        if m >= 128 and k >= 4 * n:
            return "sm110_sq_bf16"
        if n == k and m >= 512:
            return "sm110_sq_bf16" if k <= 1024 else "sm110_wide_bf16"
        if n == k and m >= 128:
            return "sm110_wide_bf16"
        return "sm110_t1_bf16"
    if m == 1:
        if variant == 4:
            return "gemv_fp8_m1_w4"
        if variant == 8:
            return "gemv_fp8_m1_w8"
        if variant == 16:
            return "gemv_fp8_m1_w16"
        if n <= 2048:
            return "gemv_fp8_m1_w4"
        if n <= 8192:
            return "gemv_fp8_m1_w8"
        return "gemv_fp8_m1_w16"
    if m <= 16:
        if k % 256 == 0:
            return "ld_fp8_gemm_16x128x256_w4" if n % 128 == 0 else "ld_fp8_gemm_16x64x256_w4"
        if n % 256 == 0:
            return "ld_fp8_gemm_16x256x128_w8"
        if n % 192 == 0:
            return "ld_fp8_gemm_16x192x128_w4"
        if n % 128 == 0:
            return "ld_fp8_gemm_16x128x128_w4"
        return "ld_fp8_gemm_16x64x128_w4"
    if m <= 32:
        if k % 256 == 0:
            return "ld_fp8_gemm_32x128x256_w4" if n % 128 == 0 else "ld_fp8_gemm_32x64x256_w4"
        if n % 192 == 0:
            return "ld_fp8_gemm_32x192x128_w4"
        if n % 128 == 0:
            return "ld_fp8_gemm_32x128x128_w4"
        return "ld_fp8_gemm_32x64x128_w4"
    if m <= 64:
        if k % 256 == 0:
            return "ld_fp8_gemm_64x128x256_w4" if n % 128 == 0 else "ld_fp8_gemm_64x64x256_w4"
        if n % 128 == 0:
            return "ld_fp8_gemm_64x128x128_w4"
        return "ld_fp8_gemm_64x64x128_w4"
    if m <= 64:
        if k % 256 == 0:
            return "ld_fp8_gemm_64x128x256_w4" if n % 128 == 0 else "ld_fp8_gemm_64x64x256_w4"
        if n % 128 == 0:
            return "ld_fp8_gemm_64x128x128_w4"
        return "ld_fp8_gemm_64x64x128_w4"
    raise RuntimeError("unsupported M")


def make_inputs(m: int, k: int, n: int, seed: int):
    gen = torch.Generator(device="cuda")
    gen.manual_seed(seed)
    x = (torch.randn((m, k), device="cuda", generator=gen) * 0.25).to(torch.bfloat16).to(torch.float8_e4m3fn)
    w = (torch.randn((n, k), device="cuda", generator=gen) * 0.25).to(torch.bfloat16).to(torch.float8_e4m3fn)
    return x, w


def ref_fn(x, w):
    return (x.float() @ w.float().T).to(torch.bfloat16)


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 measure_median(fn, warmup: int, iters: int, rounds: int = 5) -> float:
    """Reduce clock/order bias without hiding Python-side launch behavior."""
    return float(statistics.median(measure(fn, warmup, iters) for _ in range(rounds)))


def capture_graph(fn, warmup: int) -> torch.cuda.CUDAGraph:
    """Capture one static invocation and retain the graph for paired timing."""
    for _ in range(warmup):
        fn()
    torch.cuda.synchronize()
    graph = torch.cuda.CUDAGraph()
    with torch.cuda.graph(graph):
        fn()
    torch.cuda.synchronize()
    return graph


def measure_group(functions, warmup: int, iters: int, rounds: int = 7):
    """Measure one launch per candidate per round to balance Thor DVFS drift."""
    names = list(functions)
    samples = {name: [] for name in names}
    for _ in range(warmup):
        for fn in functions.values():
            fn()
    torch.cuda.synchronize()
    sample_count = max(iters, rounds * 16)
    event_pairs = {name: [] for name in names}
    for round_index in range(sample_count):
        offset = round_index % len(names)
        for name in names[offset:] + names[:offset]:
            start = torch.cuda.Event(enable_timing=True)
            end = torch.cuda.Event(enable_timing=True)
            start.record()
            functions[name]()
            end.record()
            event_pairs[name].append((start, end))
    torch.cuda.synchronize()
    for name, pairs in event_pairs.items():
        samples[name] = [start.elapsed_time(end) * 1000.0 for start, end in pairs]
    medians = {
        name: float(statistics.median(values)) for name, values in samples.items()
    }
    return medians, samples


def metrics(got, expected):
    diff = (got.float() - expected.float()).abs().flatten()
    return (
        float(diff.max().item()),
        float(diff.mean().item()),
        float(torch.quantile(diff, 0.99).item()),
        float(torch.nn.functional.cosine_similarity(got.float().flatten(), expected.float().flatten(), dim=0).item()),
    )


def load_native_reference():
    root = os.environ.get("FLASHRT_NATIVE_ROOT")
    if not root:
        return None
    sys.path.insert(0, root)
    try:
        return importlib.import_module("flash_rt.flash_rt_kernels")
    finally:
        sys.path.remove(root)


def bench_case(ops, native, name: str, shape: tuple[int, int, int], variant: int, warmup: int, iters: int, compile_ref: bool):
    m, k, n = shape
    x, w = make_inputs(m, k, n, seed=3000 + m + k + n + variant)
    out = torch.empty((m, n), device="cuda", dtype=torch.bfloat16)
    expected = ref_fn(x, w)
    got = ops.fp8_linear_bf16(x, w, out=out, variant=variant)
    torch.cuda.synchronize()
    max_abs, mean_abs, p99_abs, cos = metrics(got, expected)
    tile = ops.select_fp8_linear_tile(m, n, k, variant)

    wrapper_invoke = lambda: ops.fp8_linear_bf16(x, w, out=out, variant=variant)
    eager_functions = {"wrapper": wrapper_invoke}
    graph_objects = {"wrapper": capture_graph(wrapper_invoke, warmup)}
    package_forced_tiles = {}
    if torch.cuda.get_device_capability(0) == (11, 0):
        for forced_variant, forced_tile in {
            1: "sm110_sq_bf16",
            2: "sm110_t1_bf16",
            3: "sm110_wide_bf16",
        }.items():
            invoke = lambda forced_variant=forced_variant: ops.fp8_linear_bf16(
                x, w, out=out, variant=forced_variant
            )
            invoke()
            torch.cuda.synchronize()
            fmax, fmean, fp99, fcos = metrics(out, expected)
            if fmax > 0.5 or fmean > 0.02 or fp99 > 0.25 or fcos < 0.999:
                raise RuntimeError(
                    f"package {forced_tile} failed correctness for {name}: "
                    f"{fmax=}, {fmean=}, {fp99=}, {fcos=}"
                )
            key = f"package:{forced_tile}"
            package_forced_tiles[forced_tile] = invoke
            eager_functions[key] = invoke
            graph_objects[key] = capture_graph(invoke, warmup)
    package_best_graph_tile = None
    package_best_graph_us = None
    package_tiles_us = None
    native_best_tile = None
    native_best_us = None
    native_best_graph_tile = None
    native_best_graph_us = None
    native_tiles_us = None
    native_invokes = {}
    wrapper_vs_native = None
    graph_vs_native = None
    auto_over_package_best = None
    if native is not None and torch.cuda.get_device_capability(0) == (11, 0):
        native_out = torch.empty_like(out)
        candidates = [
            ("sm110_sq_bf16", native.cutlass_fp8_sq_bf16out),
            ("sm110_t1_bf16", native.cutlass_fp8_t1_bf16out),
            ("sm110_wide_bf16", native.cutlass_fp8_wide_bf16out),
        ]
        for tile_name, fn in candidates:
            invoke = lambda fn=fn: fn(
                x.data_ptr(), w.data_ptr(), native_out.data_ptr(),
                m, n, k, 1.0, 0.0,
                int(torch.cuda.current_stream().cuda_stream),
            )
            rc = invoke()
            if rc != 0:
                continue
            torch.cuda.synchronize()
            nmax, nmean, np99, ncos = metrics(native_out, expected)
            if nmax > 0.5 or nmean > 0.02 or np99 > 0.25 or ncos < 0.999:
                raise RuntimeError(
                    f"native {tile_name} failed correctness for {name}: "
                    f"{nmax=}, {nmean=}, {np99=}, {ncos=}"
                )
            key = f"native:{tile_name}"
            native_invokes[tile_name] = invoke
            eager_functions[key] = invoke
            graph_objects[key] = capture_graph(invoke, warmup)
    if package_forced_tiles:
        eager_times, eager_samples = measure_group(eager_functions, warmup, iters)
        graph_times, graph_samples = measure_group(
            {
                key: graph.replay
                for key, graph in graph_objects.items()
                if key != "wrapper"
            },
            warmup,
            iters,
        )
        flashrt_us = eager_times["wrapper"]
        package_tiles_us = {
            tile_name: {
                "eager": eager_times[f"package:{tile_name}"],
                "graph": graph_times[f"package:{tile_name}"],
            }
            for tile_name in package_forced_tiles
        }
        package_best_graph_us, package_best_graph_tile = min(
            (times["graph"], tile_name)
            for tile_name, times in package_tiles_us.items()
        )
        # Auto and its matching forced variant resolve to the same native
        # function. Use that single graph measurement for the tile gate; timing
        # duplicate graph objects is vulnerable to Thor DVFS order bias.
        flashrt_graph_us = package_tiles_us[tile]["graph"]
        auto_over_package_best = float(statistics.median(
            selected / best
            for selected, best in zip(
                graph_samples[f"package:{tile}"],
                graph_samples[f"package:{package_best_graph_tile}"],
            )
        ))
        if native_invokes:
            native_tiles_us = {
                tile_name: {
                    "eager": eager_times[f"native:{tile_name}"],
                    "graph": graph_times[f"native:{tile_name}"],
                }
                for tile_name in native_invokes
            }
            native_best_us, native_best_tile = min(
                (times["eager"], tile_name)
                for tile_name, times in native_tiles_us.items()
            )
            native_best_graph_us, native_best_graph_tile = min(
                (times["graph"], tile_name)
                for tile_name, times in native_tiles_us.items()
            )
            wrapper_vs_native = float(statistics.median(
                wrapper / native_sample
                for wrapper, native_sample in zip(
                    eager_samples["wrapper"],
                    eager_samples[f"native:{native_best_tile}"],
                )
            ))
            graph_vs_native = float(statistics.median(
                package_sample / native_sample
                for package_sample, native_sample in zip(
                    graph_samples[f"package:{tile}"],
                    graph_samples[f"native:{native_best_graph_tile}"],
                )
            ))
    else:
        flashrt_us = measure_median(wrapper_invoke, warmup, iters)
        flashrt_graph_us = measure_median(
            graph_objects["wrapper"].replay, warmup, iters
        )
    tile_pass = auto_over_package_best is None or auto_over_package_best <= 1.10
    eager_us = measure(lambda: ref_fn(x, w), warmup, iters)
    compile_us = None
    if compile_ref:
        try:
            compiled = torch.compile(ref_fn, fullgraph=True)
            compiled(x, w)
            torch.cuda.synchronize()
            compile_us = measure(lambda: compiled(x, w), warmup, iters)
        except Exception:
            compile_us = None

    return Result(
        shape=name,
        M=m,
        K=k,
        N=n,
        variant=variant,
        tile=tile,
        flashrt_us=flashrt_us,
        flashrt_graph_us=flashrt_graph_us,
        package_best_graph_tile=package_best_graph_tile,
        package_best_graph_us=package_best_graph_us,
        auto_over_package_best=auto_over_package_best,
        package_tiles_us=package_tiles_us,
        native_best_tile=native_best_tile,
        native_best_us=native_best_us,
        native_best_graph_tile=native_best_graph_tile,
        native_best_graph_us=native_best_graph_us,
        native_tiles_us=native_tiles_us,
        wrapper_vs_native=wrapper_vs_native,
        graph_vs_native=graph_vs_native,
        torch_eager_us=eager_us,
        torch_compile_us=compile_us,
        speedup_vs_eager=eager_us / flashrt_us,
        speedup_vs_compile=(compile_us / flashrt_us) if compile_us else None,
        max_abs=max_abs,
        mean_abs=mean_abs,
        p99_abs=p99_abs,
        cosine=cos,
        status=(
            "pass"
            if max_abs <= 0.5 and p99_abs <= 0.25 and cos >= 0.999 and tile_pass
            else "fail"
        ),
    )


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--backend", choices=["source", "installed"], default="source")
    parser.add_argument("--artifact", default=None)
    parser.add_argument("--mode", choices=sorted(MODES), default="smoke")
    parser.add_argument("--warmup", type=int, default=20)
    parser.add_argument("--iterations", type=int, default=100)
    parser.add_argument("--compile-ref", action="store_true")
    parser.add_argument("--json-out", default=None)
    args = parser.parse_args()

    if not torch.cuda.is_available():
        raise SystemExit("CUDA is required")
    capability = torch.cuda.get_device_capability(0)
    if capability not in {(11, 0), (12, 0)}:
        raise SystemExit("fp8-gemm benchmark requires SM110 or SM120")

    ops = load_source_ops() if args.backend == "source" else load_installed_ops(args.artifact)
    native = load_native_reference()
    rows: list[Result] = []
    for name in MODES[args.mode]:
        shape = SHAPES[name]
        variants = [0]
        if shape[0] == 1 and capability == (12, 0):
            variants = [0, 4, 8, 16]
        for variant in variants:
            rows.append(bench_case(
                ops, native, name, shape, variant,
                args.warmup, args.iterations, args.compile_ref,
            ))

    payload = {"rows": [asdict(row) for row in rows]}
    print(json.dumps(payload, indent=2, sort_keys=True))
    if args.json_out:
        Path(args.json_out).write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
    if any(row.status != "pass" for row in rows):
        raise SystemExit(1)


if __name__ == "__main__":
    main()