tests/test_fp4_gemm.py

#!/usr/bin/env python3
"""Correctness tests for fp4-gemm."""

from __future__ import annotations

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

import torch


ROOT = Path(__file__).resolve().parents[2]
PACKAGE = ROOT / "fp4-gemm"
REGISTRATION_INCLUDE = (
    ROOT.parent
    / "kernels"
    / "kernel-builder"
    / "src"
    / "pyproject"
    / "templates"
    / "torch"
)
DEFAULT_CUTLASS_INCLUDE = (
    ROOT.parent
    / "flashrt_pr31_review"
    / "third_party"
    / "cutlass"
    / "include"
)

SHAPES = {
    "small_m16_n128_k128": (16, 128, 128),
    "small_m32_n256_k256": (32, 256, 256),
    "mlp_tile_m64_n512_k512": (64, 512, 512),
}

SM110_SHAPES = {
    "pi05_action_gate_up": (51, 16384, 2048),
    "pi05_action_down": (51, 2048, 8192),
    "groot_n17_dit_qkv": (41, 4608, 1536),
    "groot_n17_dit_ffn_up": (41, 6144, 1536),
    "groot_n17_dit_ffn_down": (41, 1536, 6144),
    "groot_legacy_dit_qkv": (51, 4608, 1536),
    "groot_backbone_gate_up": (277, 16384, 2048),
    "cosmos_edge_action": (64, 9216, 2048),
    "lingbot_action_gate_up": (105, 16384, 2048),
}

EPILOGUE_SHAPES = {
    "epilogue_tile": (64, 512, 512),
    "motus_up": (360, 14336, 3072),
    "motus_down": (360, 3072, 14336),
}

MODES = {
    "smoke": ["small_m16_n128_k128"],
    "full": list(SHAPES),
    "thor-models": list(SM110_SHAPES),
}


@dataclass
class Metrics:
    shape: str
    M: int
    N: int
    K: int
    workload: str
    variant: int | None
    max_abs: float
    mean_abs: float
    p99_abs: float
    cosine: float
    passed: bool


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

    @staticmethod
    def sfa_size_bytes(rows: int, dim: int) -> int:
        n_blocks = dim // 16
        n_row_super = (rows + 127) // 128
        n_col_super = (n_blocks + 3) // 4
        return n_row_super * n_col_super * 512

    def alloc_fp4(self, rows: int, dim: int):
        return (
            torch.empty((rows, dim // 2), device="cuda", dtype=torch.uint8),
            torch.empty((self.sfa_size_bytes(rows, dim),), device="cuda", dtype=torch.uint8),
        )

    def quantize_fp4_sfa_fp16(self, x, packed, sfa, is_sfb=False):
        self._ops.quantize_fp4_sfa_fp16(x, packed, sfa, bool(is_sfb))

    def quantize_fp4_sfa_bf16(self, x, packed, sfa, is_sfb=False):
        self._ops.quantize_fp4_sfa_bf16(x, packed, sfa, bool(is_sfb))

    def dequantize_fp4_sfa_fp16(self, packed, sfa, out, is_sfb=False):
        self._ops.dequantize_fp4_sfa_fp16(packed, sfa, out, bool(is_sfb))

    def nvfp4_gemm_bf16(self, a, b, sfa, sfb, out, alpha=1.0, variant=0):
        self._ops.nvfp4_gemm_bf16(a, b, sfa, sfb, out, float(alpha), int(variant))

    def nvfp4_gemm_bias_bf16(self, a, b, sfa, sfb, bias, out):
        self._ops.nvfp4_gemm_bias_bf16(a, b, sfa, sfb, bias, out)

    def nvfp4_gemm_bias_residual_bf16(
        self, a, b, sfa, sfb, bias, residual, out
    ):
        self._ops.nvfp4_gemm_bias_residual_bf16(
            a, b, sfa, sfb, bias, residual, out
        )

    def nvfp4_gemm_residual_bf16(self, a, b, sfa, sfb, residual, out, alpha=1.0):
        self._ops.nvfp4_gemm_residual_bf16(
            a, b, sfa, sfb, residual, out, float(alpha)
        )

    def nvfp4_gemm_bias_gelu_bf16(self, a, b, sfa, sfb, bias, out, alpha=1.0):
        self._ops.nvfp4_gemm_bias_gelu_bf16(
            a, b, sfa, sfb, bias, out, float(alpha)
        )

    def nvfp4_gemm_bias_gelu_nvfp4(
        self, a, b, sfa, sfb, bias, out_packed, out_sfa, alpha=1.0
    ):
        self._ops.nvfp4_gemm_bias_gelu_nvfp4(
            a, b, sfa, sfb, bias, out_packed, out_sfa, float(alpha)
        )

    def nvfp4_gemm_streamk_bf16(self, a, b, sfa, sfb, out, alpha=1.0):
        self._ops.nvfp4_gemm_streamk_bf16(a, b, sfa, sfb, out, float(alpha))

    def nvfp4_gemm_streamk_bias_bf16(
        self, a, b, sfa, sfb, bias, out, alpha=1.0
    ):
        self._ops.nvfp4_gemm_streamk_bias_bf16(
            a, b, sfa, sfb, bias, out, float(alpha)
        )


class InstalledOps:
    """Adapt the public return-value API to the in-place test interface."""

    def __init__(self, module) -> None:
        self._module = module

    def sfa_size_bytes(self, rows: int, dim: int) -> int:
        return int(self._module.sfa_size_bytes(rows, dim))

    def alloc_fp4(self, rows: int, dim: int):
        return (
            torch.empty((rows, dim // 2), device="cuda", dtype=torch.uint8),
            torch.empty(
                (self._module.sfa_size_bytes(rows, dim),),
                device="cuda",
                dtype=torch.uint8,
            ),
        )

    def quantize_fp4_sfa_fp16(self, x, packed, sfa, is_sfb=False):
        self._module.quantize_fp4_sfa_fp16(
            x, packed=packed, sfa=sfa, is_sfb=bool(is_sfb)
        )

    def quantize_fp4_sfa_bf16(self, x, packed, sfa, is_sfb=False):
        self._module.quantize_fp4_sfa_bf16(
            x, packed=packed, sfa=sfa, is_sfb=bool(is_sfb)
        )

    def dequantize_fp4_sfa_fp16(self, packed, sfa, out, is_sfb=False):
        self._module.dequantize_fp4_sfa_fp16(
            packed, sfa, out=out, is_sfb=bool(is_sfb)
        )

    def nvfp4_gemm_bf16(self, a, b, sfa, sfb, out, alpha=1.0, variant=0):
        self._module.nvfp4_gemm_bf16(
            a,
            b,
            sfa,
            sfb,
            alpha=float(alpha),
            out=out,
            variant=int(variant),
        )

    def nvfp4_gemm_bias_bf16(self, a, b, sfa, sfb, bias, out):
        self._module.nvfp4_gemm_bias_bf16(
            a, b, sfa, sfb, bias, out=out
        )

    def nvfp4_gemm_bias_residual_bf16(
        self, a, b, sfa, sfb, bias, residual, out
    ):
        self._module.nvfp4_gemm_bias_residual_bf16(
            a, b, sfa, sfb, bias, residual, out=out
        )

    def nvfp4_gemm_residual_bf16(self, a, b, sfa, sfb, residual, out, alpha=1.0):
        self._module.nvfp4_gemm_residual_bf16(
            a, b, sfa, sfb, residual, alpha=float(alpha), out=out
        )

    def nvfp4_gemm_bias_gelu_bf16(self, a, b, sfa, sfb, bias, out, alpha=1.0):
        self._module.nvfp4_gemm_bias_gelu_bf16(
            a, b, sfa, sfb, bias, alpha=float(alpha), out=out
        )

    def nvfp4_gemm_bias_gelu_nvfp4(
        self, a, b, sfa, sfb, bias, out_packed, out_sfa, alpha=1.0
    ):
        self._module.nvfp4_gemm_bias_gelu_nvfp4(
            a, b, sfa, sfb, bias, alpha=float(alpha),
            out_packed=out_packed, out_sfa=out_sfa,
        )

    def nvfp4_gemm_streamk_bf16(self, a, b, sfa, sfb, out, alpha=1.0):
        self._module.nvfp4_gemm_streamk_bf16(
            a, b, sfa, sfb, alpha=float(alpha), out=out
        )

    def nvfp4_gemm_streamk_bias_bf16(
        self, a, b, sfa, sfb, bias, out, alpha=1.0
    ):
        self._module.nvfp4_gemm_streamk_bias_bf16(
            a, b, sfa, sfb, bias, alpha=float(alpha), out=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

    cutlass_include = Path(os.environ.get("FLASHRT_CUTLASS_INCLUDE", str(DEFAULT_CUTLASS_INCLUDE)))
    if not REGISTRATION_INCLUDE.is_dir():
        raise RuntimeError(f"missing kernel-builder registration include: {REGISTRATION_INCLUDE}")
    if not cutlass_include.is_dir():
        raise RuntimeError(f"missing CUTLASS include path: {cutlass_include}")
    os.environ.setdefault("TORCH_CUDA_ARCH_LIST", _current_arch_list())
    namespace = "fp4_gemm_source_test"
    capability = torch.cuda.get_device_capability(0)
    if capability == (11, 0):
        gemm_sources = [
            str(PACKAGE / "csrc" / "gemm" / "fp4" / "cutlass_nvfp4_w4a16_gemm_sm100.cu"),
            str(PACKAGE / "csrc" / "gemm" / "fp4" / "cutlass_fp4_gemm_bias_bf16_sm100.cu"),
            str(PACKAGE / "csrc" / "quantize" / "quantize_fp4_sfa_bf16.cu"),
            str(PACKAGE / "csrc" / "gemm" / "fp4" / "sm110_dispatch.cu"),
        ]
        source_define = "-DFLASHRT_FP4_GEMM_SOURCE_SM110_ONLY"
    else:
        gemm_sources = [
            str(PACKAGE / "csrc" / "gemm" / "fp4" / "cutlass_nvfp4_w4a16_gemm_sm120.cu"),
            str(PACKAGE / "csrc" / "gemm" / "fp4" / "fp4_w4a4_mma_warpsplit_sm120.cu"),
            str(PACKAGE / "csrc" / "gemm" / "fp4" / "cutlass_nvfp4_gemm_bias_gelu_bf16out_sm120.cu"),
            str(PACKAGE / "csrc" / "gemm" / "fp4" / "cutlass_nvfp4_gemm_bias_gelu_fp4out_sm120.cu"),
            str(PACKAGE / "csrc" / "gemm" / "fp4" / "cutlass_nvfp4_gemm_dn_streamk_bias_sm120.cu"),
        ]
        source_define = None
    load(
        name=namespace,
        sources=[
            str(PACKAGE / "torch-ext" / "torch_binding.cpp"),
            *gemm_sources,
            str(PACKAGE / "csrc" / "quantize" / "quantize_fp4_sfa.cu"),
            str(PACKAGE / "csrc" / "dequantize_fp4_sfa.cu"),
        ],
        extra_include_paths=[
            str(PACKAGE / "csrc"),
            str(cutlass_include),
            str(REGISTRATION_INCLUDE),
        ],
        extra_cflags=[flag for flag in ["-O3", "-DCUDA_KERNEL", source_define] if flag],
        extra_cuda_cflags=[
            "-O3",
            "--expt-relaxed-constexpr",
            "--expt-extended-lambda",
            "-DCUDA_KERNEL",
            "-DCUTLASS_ARCH_MMA_SM100_SUPPORTED=1",
            *([source_define] if source_define else []),
        ],
        verbose=False,
    )
    return SourceOps(namespace)


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


def make_inputs(m: int, n: int, k: int, seed: int):
    gen = torch.Generator(device="cuda")
    gen.manual_seed(seed)
    a = (torch.randn((m, k), device="cuda", generator=gen) * 0.25).to(torch.float16).contiguous()
    b = (torch.randn((n, k), device="cuda", generator=gen) * 0.25).to(torch.float16).contiguous()
    return a, b


def metrics(got: torch.Tensor, expected: torch.Tensor) -> tuple[float, float, float, float]:
    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 check_bf16_threshold(max_abs: float, mean_abs: float, p99_abs: float, cosine: float) -> bool:
    return max_abs <= 0.125 and mean_abs <= 0.005 and p99_abs <= 0.03125 and cosine >= 0.999


def select_sm110_variant(shape: tuple[int, int, int]) -> int:
    _m, n, k = shape
    if n >= 4 * k:
        return 1
    if n == 3 * k:
        return 2
    return 0


def prepare_quantized(ops: SourceOps, m: int, n: int, k: int):
    a_fp16, b_fp16 = make_inputs(m, n, k, seed=7000 + m + n + k)
    a_packed, sfa = ops.alloc_fp4(m, k)
    b_packed, sfb = ops.alloc_fp4(n, k)
    ops.quantize_fp4_sfa_fp16(a_fp16, a_packed, sfa, False)
    ops.quantize_fp4_sfa_fp16(b_fp16, b_packed, sfb, True)
    a_deq = torch.empty_like(a_fp16)
    b_deq = torch.empty_like(b_fp16)
    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()
    expected = (a_deq.float() @ b_deq.float().T).to(torch.bfloat16)
    return a_packed, b_packed, sfa, sfb, expected


def prepare_quantized_full(ops: SourceOps, m: int, n: int, k: int):
    a_fp16, b_fp16 = make_inputs(m, n, k, seed=9000 + m + n + k)
    a_packed, sfa = ops.alloc_fp4(m, k)
    b_packed, sfb = ops.alloc_fp4(n, k)
    ops.quantize_fp4_sfa_fp16(a_fp16, a_packed, sfa, False)
    ops.quantize_fp4_sfa_fp16(b_fp16, b_packed, sfb, True)
    a_deq = torch.empty_like(a_fp16)
    b_deq = torch.empty_like(b_fp16)
    ops.dequantize_fp4_sfa_fp16(a_packed, sfa, a_deq, False)
    ops.dequantize_fp4_sfa_fp16(b_packed, sfb, b_deq, True)
    return a_packed, b_packed, sfa, sfb, a_deq, b_deq


def run_case(ops: SourceOps, name: str, shape: tuple[int, int, int]) -> list[Metrics]:
    m, n, k = shape
    a_packed, b_packed, sfa, sfb, expected = prepare_quantized(ops, m, n, k)
    results: list[Metrics] = []
    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 = metrics(out, expected)
        results.append(
            Metrics(
                shape=name,
                M=m,
                N=n,
                K=k,
                workload="nvfp4_gemm_bf16",
                variant=variant,
                max_abs=max_abs,
                mean_abs=mean_abs,
                p99_abs=p99_abs,
                cosine=cosine,
                passed=check_bf16_threshold(max_abs, mean_abs, p99_abs, cosine),
            )
        )

    return results


def result_row(name, shape, workload, got, expected, *, fp4_output=False):
    max_abs, mean_abs, p99_abs, cosine = metrics(got, expected)
    if fp4_output:
        mean_magnitude = float(expected.float().abs().mean().item())
        rms = float(expected.float().square().mean().sqrt().item())
        passed = (
            cosine >= 0.9993
            and mean_abs / max(mean_magnitude, 1e-12) <= 0.01
            and p99_abs / max(rms, 1e-12) <= 0.15
        )
    else:
        passed = (
            cosine >= 0.999
            and mean_abs <= 0.008
            and p99_abs <= 0.0625
        )
    return Metrics(
        shape=name,
        M=shape[0],
        N=shape[1],
        K=shape[2],
        workload=workload,
        variant=None,
        max_abs=max_abs,
        mean_abs=mean_abs,
        p99_abs=p99_abs,
        cosine=cosine,
        passed=passed,
    )


def run_epilogue_case(ops, name: str, shape: tuple[int, int, int]):
    m, n, k = shape
    a, b, sfa, sfb, a_deq, b_deq = prepare_quantized_full(ops, m, n, k)
    matmul = a_deq.float() @ b_deq.float().T
    bias = (torch.randn(n, device="cuda") * 0.02).to(torch.bfloat16)
    residual = torch.randn((m, n), device="cuda", dtype=torch.bfloat16)
    rows = []

    out = torch.empty((m, n), device="cuda", dtype=torch.bfloat16)
    ops.nvfp4_gemm_residual_bf16(a, b, sfa, sfb, residual, out)
    expected = (matmul + residual.float()).to(torch.bfloat16)
    rows.append(result_row(name, shape, "nvfp4_gemm_residual_bf16", out, expected))

    ops.nvfp4_gemm_bias_gelu_bf16(a, b, sfa, sfb, bias, out)
    expected_gelu = torch.nn.functional.gelu(
        matmul + bias.float().view(1, -1), approximate="tanh"
    ).to(torch.bfloat16)
    rows.append(result_row(name, shape, "nvfp4_gemm_bias_gelu_bf16", out, expected_gelu))

    out_packed, out_sfa = ops.alloc_fp4(m, n)
    ops.nvfp4_gemm_bias_gelu_nvfp4(
        a, b, sfa, sfb, bias, out_packed, out_sfa
    )
    out_deq = torch.empty((m, n), device="cuda", dtype=torch.float16)
    ops.dequantize_fp4_sfa_fp16(out_packed, out_sfa, out_deq, False)
    staged_packed, staged_sfa = ops.alloc_fp4(m, n)
    ops.quantize_fp4_sfa_fp16(
        expected_gelu.to(torch.float16), staged_packed, staged_sfa, False
    )
    staged_deq = torch.empty_like(out_deq)
    ops.dequantize_fp4_sfa_fp16(
        staged_packed, staged_sfa, staged_deq, False
    )
    rows.append(
        result_row(
            name, shape, "nvfp4_gemm_bias_gelu_nvfp4",
            out_deq, staged_deq, fp4_output=True,
        )
    )

    ops.nvfp4_gemm_streamk_bf16(a, b, sfa, sfb, out)
    expected_linear = matmul.to(torch.bfloat16)
    rows.append(result_row(name, shape, "nvfp4_gemm_streamk_bf16", out, expected_linear))

    ops.nvfp4_gemm_streamk_bias_bf16(a, b, sfa, sfb, bias, out)
    expected_bias = (matmul + bias.float().view(1, -1)).to(torch.bfloat16)
    rows.append(result_row(name, shape, "nvfp4_gemm_streamk_bias_bf16", out, expected_bias))
    return rows


def run_sm110_epilogue_case(ops, name: str, shape: tuple[int, int, int]):
    m, n, k = shape
    a, b, sfa, sfb, a_deq, b_deq = prepare_quantized_full(ops, m, n, k)
    matmul = a_deq.float() @ b_deq.float().T
    bias = (torch.randn(n, device="cuda") * 0.02).to(torch.bfloat16)
    residual = torch.randn((m, n), device="cuda", dtype=torch.bfloat16)
    rows = []

    out = torch.empty_like(residual)
    ops.nvfp4_gemm_bias_bf16(a, b, sfa, sfb, bias, out)
    expected_bias = (matmul + bias.float().view(1, -1)).to(torch.bfloat16)
    rows.append(result_row(
        name, shape, "nvfp4_gemm_bias_bf16", out, expected_bias
    ))

    before = residual.clone()
    ops.nvfp4_gemm_bias_residual_bf16(
        a, b, sfa, sfb, bias, residual, residual
    )
    expected_residual = (
        matmul + bias.float().view(1, -1) + before.float()
    ).to(torch.bfloat16)
    rows.append(result_row(
        name, shape, "nvfp4_gemm_bias_residual_bf16",
        residual, expected_residual,
    ))

    out_packed, out_sfa = ops.alloc_fp4(m, n)
    ops.nvfp4_gemm_bias_gelu_nvfp4(
        a, b, sfa, sfb, bias, out_packed, out_sfa
    )
    out_deq = torch.empty((m, n), device="cuda", dtype=torch.float16)
    ops.dequantize_fp4_sfa_fp16(out_packed, out_sfa, out_deq, False)
    expected_gelu = torch.nn.functional.gelu(
        matmul + bias.float().view(1, -1), approximate="tanh"
    ).to(torch.bfloat16)
    staged_packed, staged_sfa = ops.alloc_fp4(m, n)
    ops.quantize_fp4_sfa_bf16(
        expected_gelu, staged_packed, staged_sfa, False
    )
    staged_deq = torch.empty_like(out_deq)
    ops.dequantize_fp4_sfa_fp16(
        staged_packed, staged_sfa, staged_deq, False
    )
    rows.append(result_row(
        name, shape, "nvfp4_gemm_bias_gelu_nvfp4",
        out_deq, staged_deq, fp4_output=True,
    ))
    return rows


def check_installed_compile(ops: InstalledOps) -> dict[str, object]:
    a_packed, b_packed, sfa, sfb, _ = prepare_quantized(ops, 128, 128, 128)

    def call(a, b, scale_a, scale_b):
        return ops._module.nvfp4_gemm_bf16(a, b, scale_a, scale_b)

    eager = call(a_packed, b_packed, sfa, sfb)
    compiled = torch.compile(call, fullgraph=True)
    got = compiled(a_packed, b_packed, sfa, sfb)
    torch.cuda.synchronize()
    max_abs = float((got.float() - eager.float()).abs().max().item())
    passed = bool(
        got.dtype == torch.bfloat16
        and got.shape == eager.shape
        and torch.equal(got, eager)
    )
    return {
        "fullgraph": True,
        "dtype": str(got.dtype),
        "shape": list(got.shape),
        "max_abs": max_abs,
        "exact": bool(torch.equal(got, eager)),
        "passed": passed,
    }


def check_bf16_quantizer(ops) -> dict[str, object]:
    """Require the direct BF16 producer to preserve the established layout."""
    cases = [
        (1, 5120, False),
        (1, 6144, False),
        (1, 17408, False),
        (16, 2048, False),
        (128, 512, False),
        (64, 1024, True),
    ]
    rows = []
    for case_index, (m, k, is_sfb) in enumerate(cases):
        torch.manual_seed(8100 + case_index)
        x = (torch.randn((m, k), device="cuda") * 1.5).to(torch.bfloat16)
        direct_packed, direct_sfa = ops.alloc_fp4(m, k)
        compat_packed, compat_sfa = ops.alloc_fp4(m, k)
        # CUTLASS SFA/SFB buffers contain alignment padding that producers do
        # not write or consume. Zero it so a full-buffer equality check still
        # proves every mapped scale byte lands at the same address.
        direct_sfa.zero_()
        compat_sfa.zero_()
        ops.quantize_fp4_sfa_bf16(x, direct_packed, direct_sfa, is_sfb)
        ops.quantize_fp4_sfa_fp16(
            x.to(torch.float16), compat_packed, compat_sfa, is_sfb
        )
        torch.cuda.synchronize()
        packed_exact = bool(torch.equal(direct_packed, compat_packed))
        sfa_exact = bool(torch.equal(direct_sfa, compat_sfa))
        direct_deq = torch.empty((m, k), device="cuda", dtype=torch.float16)
        compat_deq = torch.empty_like(direct_deq)
        ops.dequantize_fp4_sfa_fp16(
            direct_packed, direct_sfa, direct_deq, is_sfb
        )
        ops.dequantize_fp4_sfa_fp16(
            compat_packed, compat_sfa, compat_deq, is_sfb
        )
        torch.cuda.synchronize()
        dequant_exact = bool(torch.equal(direct_deq, compat_deq))
        rows.append(
            {
                "shape": [m, k],
                "is_sfb": is_sfb,
                "packed_exact": packed_exact,
                "sfa_exact": sfa_exact,
                "dequant_exact": dequant_exact,
                "passed": packed_exact and sfa_exact and dequant_exact,
            }
        )
    return {"rows": rows, "passed": all(row["passed"] for row in rows)}


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=sorted(MODES), default="smoke")
    parser.add_argument("--json-out", default=None)
    args = parser.parse_args()

    if not torch.cuda.is_available():
        raise RuntimeError("CUDA is required")
    ops = load_source_ops() if args.backend == "source" else load_installed_ops(args.artifact)

    results: list[Metrics] = []
    selected_shapes = SM110_SHAPES if args.mode == "thor-models" else SHAPES
    for name in MODES[args.mode]:
        results.extend(run_case(ops, name, selected_shapes[name]))
    capability = torch.cuda.get_device_capability(0)
    if args.mode == "full" and capability != (11, 0):
        for name, shape in EPILOGUE_SHAPES.items():
            results.extend(run_epilogue_case(ops, name, shape))
    if capability == (11, 0) and args.mode in {"full", "thor-models"}:
        for name in (
            "groot_n17_dit_qkv",
            "groot_n17_dit_ffn_up",
            "groot_n17_dit_ffn_down",
        ):
            results.extend(run_sm110_epilogue_case(
                ops, name, SM110_SHAPES[name]
            ))
    compile_check = None
    bf16_quantizer_check = check_bf16_quantizer(ops)
    if args.backend == "installed" and args.mode == "full":
        compile_check = check_installed_compile(ops)
    passed = sum(1 for item in results if item.passed)
    total = len(results)
    if compile_check is not None:
        total += 1
        passed += int(bool(compile_check["passed"]))
    total += 1
    passed += int(bool(bf16_quantizer_check["passed"]))
    payload = {
        "backend": args.backend,
        "mode": args.mode,
        "device": torch.cuda.get_device_name(),
        "torch": torch.__version__,
        "passed": passed,
        "total": total,
        "results": [asdict(item) for item in results],
        "compile_check": compile_check,
        "bf16_quantizer_check": bf16_quantizer_check,
    }
    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 passed == total else 1


if __name__ == "__main__":
    raise SystemExit(main())