from collections import defaultdict

import numpy as np
import torch
import torch.nn.functional as F
from einops import einsum

import flashinfer
from flashinfer.autotuner import autotune
from flashinfer.fp8_quantization import mxfp8_quantize
from flashinfer.testing.utils import (
    bench_gpu_time,
    dequantize_fp8,
    quantize_fp8,
)

from .flashinfer_benchmark_utils import (
    dtype_str_to_torch_dtype,
    get_device,
    print_perf_metrics,
    is_close_stats,
    filter_backends_by_compute_capability,
    warn_if_pdl_unsupported,
)


def run_gemm_test(args):
    """
    Run a gemm test.

    Args:
        args: Parsed command line arguments containing test configuration

    Returns:
        dict: List of dictionaries containing performance results
    """
    if args.routine == "gemm_fp8_nt_groupwise":
        return testGemmFp8NtGroupwise(args)
    elif args.routine == "group_gemm_fp8_nt_groupwise":
        return testGroupGemmFp8NtGroupwise(args)
    elif args.routine == "bmm_fp8":
        return testBmmFp8(args)
    elif args.routine == "mm_fp8":
        return testMmFp8(args)
    elif args.routine == "bmm_mxfp8":
        return testBmmMxfp8(args)
    elif args.routine == "mm_fp4":
        return testMmFp4(args)
    elif args.routine == "mm_bf16_fp4":
        return testMmBf16Fp4(args)
    elif args.routine == "mm_mxfp8":
        return testMmMxfp8(args)
    elif args.routine == "mm_bf16":
        return testMmBf16(args)
    elif args.routine == "bmm_bf16":
        return testBmmBf16(args)
    elif args.routine == "tinygemm_bf16":
        return testTinygemmBf16(args)
    else:
        raise ValueError(f"Unsupported routine: {args.routine}")


def parse_gemm_args(line, parser):
    """
    Parse command line arguments for gemm test configuration.

    Args:
        line: Command line arguments
        parser: ArgumentParser object already populated with shared arguments

    Returns:
        Parsed argument namespace
    """
    parser.add_argument(
        "--batch_size",
        type=int,
        required=False,
        default=1,
        help="Batch size of test case.",
    )
    parser.add_argument(
        "--m", type=int, required=True, help="Number of rows in the first matrix."
    )
    parser.add_argument(
        "--n", type=int, required=True, help="Number of columns in the second matrix."
    )
    parser.add_argument(
        "--k",
        type=int,
        required=True,
        help="Number of columns in the first matrix and number of rows in the second matrix.",
    )
    parser.add_argument(
        "--tile_size",
        type=int,
        required=False,
        default=128,
        help="Tile size for the gemm operation.",
    )
    parser.add_argument(
        "--group_size",
        type=int,
        required=False,
        default=1,
        help="Group size for the group gemm operation.",
    )
    parser.add_argument(
        "--scale_major_mode",
        type=str,
        required=False,
        default="MN",
        choices=["MN", "K"],
        help="Scale major mode.",
    )
    parser.add_argument(
        "--input_dtype",
        type=str,
        required=False,
        default="fp8_e4m3",
        help="Data type of the input.",
    )
    parser.add_argument(
        "--mat2_dtype",
        type=str,
        required=False,
        default="fp8_e4m3",
        help="Data type of the mat2.",
    )
    parser.add_argument(
        "--out_dtype",
        type=str,
        required=False,
        default="bfloat16",
        help="Data type of the output.",
    )
    parser.add_argument(
        "--mma_sm",
        type=int,
        required=False,
        default=1,
        choices=[1, 2],
        help="How many SMs to use for the MMA operation, must be 1 or 2",
    )
    parser.add_argument(
        "--backends",
        type=str,
        required=False,
        nargs="+",
        default=["cudnn"],
        choices=[
            "cudnn",
            "cublas",
            "trtllm",
            "cutlass",
            "tgv",
            "cublaslt",
            "cute-dsl",
            "b12x",
            "auto",
            "tinygemm",
            "cutile",
            "trtllm_low_latency",
        ],
        help="Kernel backends to test. Default: cudnn",
    )
    parser.add_argument(
        "--use_128x4_sf_layout",
        action="store_true",
        help="Use 128x4 SF layout for the input and mat2.",
    )
    parser.add_argument(
        "--use_nvfp4",
        action="store_true",
        help="In mm_fp4, whether to use nvfp4 quantization or mxfp4 quantization, defaults to False.",
    )
    parser.add_argument(
        "--autotune",
        action="store_true",
        default=False,
        help=(
            "Enable autotuner warmup for supported routines (mm_fp4, bmm_fp8, mm_fp8, bmm_mxfp8, mm_mxfp8, mm_bf16, bmm_bf16)."
        ),
    )
    parser.add_argument(
        "--bias",
        action="store_true",
        default=False,
        help="Use bias (enabled for mm_bf16 with TGV and TinyGEMM backends)",
    )

    args = parser.parse_args(line)
    has_backends_arg = any(
        token == "--backends" or token.startswith("--backends=") for token in line
    )
    has_input_dtype_arg = any(
        token == "--input_dtype" or token.startswith("--input_dtype=") for token in line
    )
    has_mat2_dtype_arg = any(
        token == "--mat2_dtype" or token.startswith("--mat2_dtype=") for token in line
    )
    if args.routine == "tinygemm_bf16":
        if not has_backends_arg:
            args.backends = ["tinygemm"]
        if not has_input_dtype_arg:
            args.input_dtype = "bfloat16"
        if not has_mat2_dtype_arg:
            args.mat2_dtype = "bfloat16"
    if args.routine == "mm_bf16_fp4":
        if not has_backends_arg:
            args.backends = ["cute-dsl"]
        if not has_input_dtype_arg:
            args.input_dtype = "bfloat16"
    if args.routine == "mm_fp8":
        if not has_backends_arg:
            args.backends = ["trtllm_low_latency"]
    if args.verbose >= 1:
        print(f"[INFO] {args = }")
    return args


def to_float8(x, dtype=torch.float8_e4m3fn):
    finfo = torch.finfo(dtype)
    min_val, max_val = x.aminmax()
    amax = torch.maximum(min_val.abs(), max_val.abs()).clamp(min=1e-12)
    scale = finfo.max / amax
    x_scl_sat = (x * scale).clamp(min=finfo.min, max=finfo.max)
    return x_scl_sat.to(dtype), scale.float().reciprocal()


def testGemmFp8NtGroupwise(args):
    """
    Test gemm_fp8_nt_groupwise API.

    This test:
    1. Generates random input tensors
    2. Quantizes input tensors to FP8
    3. Runs gemm_fp8_nt_groupwise
    4. Runs reference check
    5. Measures performance metrics (TFLOPS, TB/sec)

    Args:
        args: Parsed command line arguments containing test configuration

    Returns:
        dict: List of dictionaries containing performance results
    """
    warn_if_pdl_unsupported(args, args.routine)
    if args.verbose >= 1:
        print("[INFO] Running testGemmFp8NtGroupwise")
        print(f"[INFO] FlashInfer version: {flashinfer.__version__}")

    device = get_device(args)
    if args.generate_repro_command:
        print(
            f"[INFO] To reproduce this test case, run the following command: {args.repro_command}"
        )

    ## Parse input arguments
    backends = args.backends
    m = args.m
    n = args.n
    k = args.k
    tile_size = args.tile_size
    scale_major_mode = args.scale_major_mode
    mma_sm = args.mma_sm
    is_cuda_graph_compatible = not args.no_cuda_graph
    run_refcheck = args.refcheck
    res = []

    backends = filter_backends_by_compute_capability(backends, args.routine, device)
    if len(backends) == 0:
        print("[ERROR] No backends to test. Exiting.")
        return res

    out_dtype = dtype_str_to_torch_dtype(args.out_dtype)
    if out_dtype not in [torch.bfloat16, torch.float16]:
        raise ValueError(f"Unsupported output dtype: {args.out_dtype}")
    ## Done parsing input arguments

    if "trtllm" in backends:
        remove_trtllm = False
        if scale_major_mode != "MN":
            print(
                "[INFO] trtllm only supports MN scale_major_mode, removing trtllm from backends"
            )
            remove_trtllm = True
        if k < 256:
            print("[INFO] trtllm only supports k >= 256, removing trtllm from backends")
            remove_trtllm = True
        if remove_trtllm:
            backends.remove("trtllm")

    if len(backends) == 0:
        print("[ERROR] No backends to test. Exiting.")
        return

    ## Prepare input tensors
    a_val = torch.randn((m, k), dtype=torch.float, device=device)
    b_val = torch.randn((n, k), dtype=torch.float, device=device) / np.sqrt(k)

    if args.verbose >= 2:
        print(f"[VVERBOSE] {a_val.shape = }")
        print(f"[VVERBOSE] {b_val.shape = }")

    if scale_major_mode == "K":
        a_scale_shape = (m, k // tile_size)
        b_scale_shape = (n // tile_size, k // tile_size)
    else:
        a_scale_shape = (k // tile_size, m)
        b_scale_shape = (k // tile_size, n // tile_size)

    a_tile_shape = (1, tile_size)
    b_tile_shape = (tile_size, tile_size)

    a_fp8, a_scale = quantize_fp8(a_val, a_scale_shape, a_tile_shape, scale_major_mode)
    b_fp8, b_scale = quantize_fp8(b_val, b_scale_shape, b_tile_shape, scale_major_mode)

    if args.verbose >= 2:
        print(f"[VVERBOSE] {a_fp8.shape = }")
        print(f"[VVERBOSE] {b_fp8.shape = }")
        print(f"[VVERBOSE] {a_scale.shape = }")
        print(f"[VVERBOSE] {b_scale.shape = }")

    a_dequant = dequantize_fp8(a_fp8, a_scale, scale_major_mode)
    b_dequant = dequantize_fp8(b_fp8, b_scale, scale_major_mode)

    def run_backend(backend, a_fp8, b_fp8, a_scale, b_scale):
        if backend in ["cutlass", "trtllm", "cutile"]:
            return flashinfer.gemm.gemm_fp8_nt_groupwise(
                a=a_fp8,
                b=b_fp8,
                a_scale=a_scale,
                b_scale=b_scale,
                scale_major_mode=scale_major_mode,
                out_dtype=out_dtype,
                mma_sm=mma_sm,
                backend=backend,
            )
        else:
            raise ValueError(f"Unsupported backend: {backend}")

    has_reference_output = False
    if run_refcheck:
        reference_output = einsum(a_dequant, b_dequant, "m k, n k -> m n").to(out_dtype)
        has_reference_output = True

    # Storage for timing results and outputs
    backend_times = {backend: [] for backend in backends}
    outputs = {}
    for cur_backend in backends:
        if run_refcheck:
            outputs[cur_backend] = run_backend(
                cur_backend, a_fp8, b_fp8, a_scale, b_scale
            ).detach()
        backend_times[cur_backend] = bench_gpu_time(
            fn=run_backend,
            dry_run_iters=args.dry_run_iters,
            repeat_iters=args.num_iters,
            sleep_after_run=True,  # GEMMs are very MMA-heavy, so prefer sleep to reduce throttling.
            enable_cupti=args.use_cupti,
            use_cuda_graph=is_cuda_graph_compatible,
            cold_l2_cache=True,
            input_args=(cur_backend, a_fp8, b_fp8, a_scale, b_scale),
        )

    tested_backends = list(outputs.keys())
    tested_outputs = list(outputs.values())
    if len(tested_backends) > 0:
        if run_refcheck and has_reference_output:
            for i in range(len(tested_backends)):
                (
                    num_different_elements,
                    num_elements,
                    num_different_elements_percentage,
                ) = is_close_stats(
                    reference_output, tested_outputs[i], rtol=1e-2, atol=1e-2
                )
                if num_different_elements > 0:
                    print(
                        f"[ERROR] Output tensor mismatch from backend {tested_backends[i]}"
                    )
                    if not args.allow_output_mismatch:
                        raise AssertionError(
                            f"[ERROR] Backend {tested_backends[i]} output mismatch with {num_different_elements} elements"
                        )

    for backend in backends:
        if len(backend_times[backend]) > 0:
            median_time = np.median(backend_times[backend])
            std_time = np.std(backend_times[backend])

            problem_flops = 2 * m * n * k
            problem_bytes = (m * k + n * k) * torch.float8_e4m3fn.itemsize + (
                m * n
            ) * out_dtype.itemsize
            tflops = problem_flops / (10**9 * median_time)  # in TFLOPs/sec
            tb_per_sec = problem_bytes / (10**9 * median_time)  # in TB/sec
            print_perf_metrics(backend, median_time, std_time, tflops, tb_per_sec)

        if args.output_path is not None:
            cur_res = defaultdict(str)
            cur_res["routine"] = args.routine
            cur_res["median_time"] = median_time
            cur_res["std_time"] = std_time
            cur_res["tflops"] = tflops
            cur_res["tb_per_sec"] = tb_per_sec
            cur_res["m"] = m
            cur_res["n"] = n
            cur_res["k"] = k
            cur_res["tile_size"] = tile_size
            cur_res["scale_major_mode"] = scale_major_mode
            cur_res["out_dtype"] = out_dtype
            cur_res["mma_sm"] = mma_sm
            cur_res["backend"] = backend
            cur_res["case_tag"] = args.case_tag
            res.append(cur_res)
    return res


def testGroupGemmFp8NtGroupwise(args):
    """
    Test group_gemm_fp8_nt_groupwise API.

    This test:
    1. Generates random input tensors
    2. Quantizes input tensors to FP8
    3. Runs group_gemm_fp8_nt_groupwise
    4. Runs reference check
    5. Measures performance metrics (TFLOPS, TB/sec)

    Args:
        args: Parsed command line arguments containing test configuration

    Returns:
        dict: List of dictionaries containing performance results
    """
    warn_if_pdl_unsupported(args, args.routine)
    if args.verbose >= 1:
        print("[INFO] Running testGroupGemmFp8NtGroupwise")
        print(f"[INFO] FlashInfer version: {flashinfer.__version__}")

    device = get_device(args)
    if args.generate_repro_command:
        print(
            f"[INFO] To reproduce this test case, run the following command: {args.repro_command}"
        )

    ## Parse input arguments
    backends = ["cutlass"]  # Cutlass is currently the only supported backend
    m = args.m
    n = args.n
    k = args.k
    group_size = args.group_size
    tile_size = args.tile_size
    scale_major_mode = args.scale_major_mode
    mma_sm = args.mma_sm
    is_cuda_graph_compatible = not args.no_cuda_graph
    run_refcheck = args.refcheck
    res = []

    backends = filter_backends_by_compute_capability(backends, args.routine, device)
    if len(backends) == 0:
        print("[ERROR] No backends to test. Exiting.")
        return res

    out_dtype = dtype_str_to_torch_dtype(args.out_dtype)
    if out_dtype not in [torch.bfloat16, torch.float16]:
        raise ValueError(f"Unsupported output dtype: {args.out_dtype}")
    ## Done parsing input arguments
    ## Prepare input tensors
    a_val = torch.randn((group_size * m, k), dtype=torch.float, device="cuda")
    b_val = torch.randn((group_size, n, k), dtype=torch.float, device="cuda") / np.sqrt(
        k
    )

    if args.verbose >= 2:
        print(f"[VVERBOSE] {a_val.shape = }")
        print(f"[VVERBOSE] {b_val.shape = }")

    if scale_major_mode == "K":
        a_scale_shape = (group_size * m, k // tile_size)
        b_scale_shape = (group_size, n // tile_size, k // tile_size)
    else:
        a_scale_shape = (k // tile_size, m * group_size)
        b_scale_shape = (group_size, k // tile_size, n // tile_size)

    a_tile_shape = (1, tile_size)
    b_tile_shape = (1, tile_size, tile_size)

    a_fp8, a_scale = quantize_fp8(a_val, a_scale_shape, a_tile_shape, scale_major_mode)
    b_fp8, b_scale = quantize_fp8(b_val, b_scale_shape, b_tile_shape, scale_major_mode)

    a_dequant = dequantize_fp8(a_fp8, a_scale, scale_major_mode)
    b_dequant = dequantize_fp8(b_fp8, b_scale, scale_major_mode)

    m_indptr = torch.arange(0, group_size + 1, dtype=torch.int32, device="cuda") * m

    if args.verbose >= 2:
        print(f"[VVERBOSE] {a_fp8.shape = }")
        print(f"[VVERBOSE] {b_fp8.shape = }")
        print(f"[VVERBOSE] {a_scale.shape = }")
        print(f"[VVERBOSE] {b_scale.shape = }")
        print(f"[VVERBOSE] {m_indptr.shape = }")

    def run_backend(backend, a_fp8, b_fp8, a_scale, b_scale, m_indptr):
        if backend == "cutlass":
            return flashinfer.gemm.group_gemm_fp8_nt_groupwise(
                a=a_fp8,
                b=b_fp8,
                a_scale=a_scale,
                b_scale=b_scale,
                m_indptr=m_indptr,
                scale_major_mode=scale_major_mode,
                out_dtype=out_dtype,
                mma_sm=mma_sm,
            )
        else:
            raise ValueError(f"Unsupported backend: {backend}")

    has_reference_output = False
    if run_refcheck:
        reference_output = (
            einsum(
                a_dequant.view((group_size, m, k)), b_dequant, "b m k, b n k -> b m n"
            )
            .view((group_size * m, n))
            .to(out_dtype)
        )
        has_reference_output = True

    # Storage for timing results and outputs
    backend_times = {backend: [] for backend in backends}
    outputs = {}
    for cur_backend in backends:
        if run_refcheck:
            outputs[cur_backend] = run_backend(
                cur_backend, a_fp8, b_fp8, a_scale, b_scale, m_indptr
            ).detach()
        backend_times[cur_backend] = bench_gpu_time(
            fn=run_backend,
            dry_run_iters=args.dry_run_iters,
            repeat_iters=args.num_iters,
            sleep_after_run=True,  # GEMMs are very MMA-heavy, so prefer sleep to reduce throttling.
            enable_cupti=args.use_cupti,
            use_cuda_graph=is_cuda_graph_compatible,
            cold_l2_cache=True,
            input_args=(cur_backend, a_fp8, b_fp8, a_scale, b_scale, m_indptr),
        )

    tested_backends = list(outputs.keys())
    tested_outputs = list(outputs.values())
    if len(tested_backends) > 0:
        if run_refcheck and has_reference_output:
            for i in range(len(tested_backends)):
                (
                    num_different_elements,
                    num_elements,
                    num_different_elements_percentage,
                ) = is_close_stats(
                    reference_output, tested_outputs[i], rtol=1e-2, atol=1e-2
                )
                if num_different_elements > 0:
                    print(
                        f"[ERROR] Output tensor mismatch from backend {tested_backends[i]}"
                    )
                    if not args.allow_output_mismatch:
                        raise AssertionError(
                            f"[ERROR] Backend {tested_backends[i]} output mismatch with {num_different_elements} elements"
                        )

    for backend in backends:
        if len(backend_times[backend]) > 0:
            median_time = np.median(backend_times[backend])
            std_time = np.std(backend_times[backend])
            problem_flops = 2 * m * n * k * group_size
            problem_bytes = (
                group_size * m * k + group_size * n * k
            ) * torch.float8_e4m3fn.itemsize + (group_size * m * n) * out_dtype.itemsize
            tflops = problem_flops / (10**9 * median_time)  # in TFLOPs/sec
            tb_per_sec = problem_bytes / (10**9 * median_time)  # in TB/sec

            print_perf_metrics(backend, median_time, std_time, tflops, tb_per_sec)

        if args.output_path is not None:
            cur_res = defaultdict(str)
            cur_res["routine"] = args.routine
            cur_res["median_time"] = median_time
            cur_res["std_time"] = std_time
            cur_res["tflops"] = tflops
            cur_res["tb_per_sec"] = tb_per_sec
            cur_res["m"] = m
            cur_res["n"] = n
            cur_res["k"] = k
            cur_res["group_size"] = group_size
            cur_res["tile_size"] = tile_size
            cur_res["scale_major_mode"] = scale_major_mode
            cur_res["out_dtype"] = out_dtype
            cur_res["mma_sm"] = mma_sm
            cur_res["backend"] = backend
            cur_res["case_tag"] = args.case_tag
            res.append(cur_res)
    return res


def testBmmFp8(args):
    """
    Test bmm_fp8 API.

    This test:
    1. Generates random input tensors
    2. Quantizes input tensors to FP8
    3. Runs bmm_fp8
    4. Runs reference check
    5. Measures performance metrics (TFLOPS, TB/sec)

    Args:
        args: Parsed command line arguments containing test configuration

    Returns:
        dict: List of dictionaries containing performance results
    """
    warn_if_pdl_unsupported(args, args.routine)
    if args.verbose >= 1:
        print("[INFO] Running testBmmFp8")
        print(f"[INFO] FlashInfer version: {flashinfer.__version__}")

    device = get_device(args)
    if args.generate_repro_command:
        print(
            f"[INFO] To reproduce this test case, run the following command: {args.repro_command}"
        )

    ## Parse input arguments
    backends = args.backends
    batch_size = args.batch_size
    m = args.m
    n = args.n
    k = args.k
    input_dtype = args.input_dtype
    mat2_dtype = args.mat2_dtype
    res_dtype = args.out_dtype
    backends = args.backends
    is_cuda_graph_compatible = not args.no_cuda_graph
    run_refcheck = args.refcheck
    autotune_supported_backends = [
        "cutlass",
        "cudnn",
        "cublas",
        "auto",
    ]
    res = []

    input_dtype = dtype_str_to_torch_dtype(args.input_dtype)
    if input_dtype not in [torch.float8_e4m3fn, torch.float8_e5m2]:
        raise ValueError(
            f"Unsupported input dtype: {input_dtype}. Supported dtypes are fp8_e4m3 and fp8_e5m2."
        )

    mat2_dtype = dtype_str_to_torch_dtype(args.mat2_dtype)
    if mat2_dtype not in [torch.float8_e4m3fn, torch.float8_e5m2]:
        raise ValueError(
            f"Unsupported mat2 dtype: {mat2_dtype}. Supported dtypes are fp8_e4m3 and fp8_e5m2."
        )

    res_dtype = dtype_str_to_torch_dtype(args.out_dtype)
    if res_dtype not in [torch.bfloat16, torch.float16]:
        raise ValueError(
            f"Unsupported res dtype: {res_dtype}. Supported dtypes are bfloat16 and float16."
        )
    ## Done parsing input arguments

    ## Prepare input tensors
    input = torch.randn([batch_size, m, k], device=device, dtype=torch.bfloat16)
    input_fp8, input_inv_s = to_float8(input, dtype=input_dtype)

    mat2 = torch.randn(
        [batch_size, n, k], device=device, dtype=torch.bfloat16
    ).transpose(-2, -1)
    mat2_fp8, mat2_inv_s = to_float8(mat2, dtype=mat2_dtype)

    if args.verbose >= 2:
        print(f"[VVERBOSE] {input_fp8.shape = }")
        print(f"[VVERBOSE] {input_fp8.dtype = }")
        print(f"[VVERBOSE] {mat2_fp8.shape = }")
        print(f"[VVERBOSE] {mat2_fp8.dtype = }")
        print(f"[VVERBOSE] {input_inv_s = }")
        print(f"[VVERBOSE] {input_inv_s.dtype = }")
        print(f"[VVERBOSE] {mat2_inv_s = }")
        print(f"[VVERBOSE] {mat2_inv_s.dtype = }")

    # Programmatically filter backends: rely on bmm_fp8's @backend_requirement
    # support checks (instead of a hard-coded compute-capability table) by
    # probing each backend with a trial call.
    backends_to_remove = []
    for backend in backends:
        if (
            getattr(args, "autotune", False)
            and backend not in autotune_supported_backends
        ):
            print(f"[INFO] {backend} backend does not support autotune")
            backends_to_remove.append(backend)
            continue

        try:
            flashinfer.gemm.bmm_fp8(
                A=input_fp8,
                B=mat2_fp8,
                A_scale=input_inv_s,
                B_scale=mat2_inv_s,
                dtype=res_dtype,
                backend=backend,
            )
        except Exception as e:
            print(
                f"[INFO] {backend} backend does not support this configuration: {type(e).__name__}: {e}"
            )
            backends_to_remove.append(backend)

    for backend in backends_to_remove:
        backends.remove(backend)

    if len(backends) == 0:
        print("[ERROR] No backends passed validation. Exiting.")
        return res

    def run_backend(backend, input_fp8, mat2_fp8, input_inv_s, mat2_inv_s):
        return flashinfer.gemm.bmm_fp8(
            A=input_fp8,
            B=mat2_fp8,
            A_scale=input_inv_s,
            B_scale=mat2_inv_s,
            dtype=res_dtype,
            backend=backend,
        )

    has_reference_output = False
    if run_refcheck:
        reference_output = torch.bmm(input, mat2)
        has_reference_output = True

    cache_path = getattr(args, "autotune_cache", None)
    if getattr(args, "autotune", False):
        warmup_iters = (
            args.dry_run_iters if args.dry_run_iters and args.dry_run_iters > 0 else 10
        )
        for cur_backend in backends:
            if cur_backend in autotune_supported_backends:
                if args.verbose >= 1:
                    print(f"[INFO] Autotune warmup for bmm_fp8: {warmup_iters} iters")
                with autotune(True, cache=cache_path):
                    for _ in range(warmup_iters):
                        run_backend(
                            cur_backend, input_fp8, mat2_fp8, input_inv_s, mat2_inv_s
                        )
    elif cache_path:
        with autotune(False, cache=cache_path):
            pass

    # Storage for timing results and outputs
    backend_times = {backend: [] for backend in backends}
    outputs = {}
    for cur_backend in backends:
        if run_refcheck:
            outputs[cur_backend] = run_backend(
                cur_backend, input_fp8, mat2_fp8, input_inv_s, mat2_inv_s
            ).detach()
        backend_times[cur_backend] = bench_gpu_time(
            fn=run_backend,
            dry_run_iters=args.dry_run_iters,
            repeat_iters=args.num_iters,
            sleep_after_run=True,
            enable_cupti=args.use_cupti,
            use_cuda_graph=is_cuda_graph_compatible,
            cold_l2_cache=True,
            input_args=(cur_backend, input_fp8, mat2_fp8, input_inv_s, mat2_inv_s),
        )

    tested_backends = list(outputs.keys())
    tested_outputs = list(outputs.values())
    if len(tested_backends) > 0:
        if run_refcheck and has_reference_output:
            if reference_output.dtype in [torch.float8_e4m3fn, torch.float8_e5m2]:
                print(
                    "[INFO] Reference output is FP8. Converting to float32 for reference check."
                )
                reference_output = reference_output.to(torch.float32)
                tested_outputs = [output.to(torch.float32) for output in tested_outputs]
            for i in range(len(tested_backends)):
                cos_sim = F.cosine_similarity(
                    reference_output.reshape(-1),
                    tested_outputs[i].reshape(-1),
                    dim=0,
                )
                if cos_sim < 0.99:
                    print(
                        f"[ERROR] Output tensor mismatch between backends {tested_backends[0]} and {tested_backends[i]}"
                    )
                    if not args.allow_output_mismatch:
                        raise AssertionError(
                            f"[ERROR] Backend {tested_backends[i]} output mismatch with cos_sim={cos_sim}"
                        )

    for backend in backends:
        backend_name = backend + (
            "_autotune"
            if (
                getattr(args, "autotune", False)
                and backend in autotune_supported_backends
            )
            else ""
        )
        if len(backend_times[backend]) > 0:
            median_time = np.median(backend_times[backend])
            std_time = np.std(backend_times[backend])
            problem_flops = 2 * m * n * k * batch_size
            problem_bytes = (
                m * k * input_dtype.itemsize
                + n * k * mat2_dtype.itemsize
                + m * n * res_dtype.itemsize
            )
            tflops = problem_flops / (10**9 * median_time)  # in TFLOPs/sec
            tb_per_sec = problem_bytes / (10**9 * median_time)  # in TB/sec
            print_perf_metrics(backend_name, median_time, std_time, tflops, tb_per_sec)

            if args.output_path is not None:
                cur_res = defaultdict(str)
                cur_res["batch_size"] = batch_size
                cur_res["routine"] = args.routine
                cur_res["median_time"] = median_time
                cur_res["std_time"] = std_time
                cur_res["tflops"] = tflops
                cur_res["tb_per_sec"] = tb_per_sec
                cur_res["m"] = m
                cur_res["n"] = n
                cur_res["k"] = k
                cur_res["input_dtype"] = input_dtype
                cur_res["mat2_dtype"] = mat2_dtype
                cur_res["out_dtype"] = res_dtype
                cur_res["backend"] = backend_name
                cur_res["case_tag"] = args.case_tag
                res.append(cur_res)
    return res


def testMmFp8(args):
    """
    Test mm_fp8 API.

    This test:
    1. Generates random input tensors and quantizes them to FP8
    2. Pre-shuffles the weight tensor into the low-latency block layout
    3. Runs mm_fp8
    4. Runs reference check against the unquantized bf16 GEMM
    5. Measures performance metrics (TFLOPS, TB/sec)

    Args:
        args: Parsed command line arguments containing test configuration

    Returns:
        dict: List of dictionaries containing performance results
    """
    warn_if_pdl_unsupported(args, args.routine)
    if args.verbose >= 1:
        print("[INFO] Running testMmFp8")
        print(f"[INFO] FlashInfer version: {flashinfer.__version__}")

    device = get_device(args)
    if args.generate_repro_command:
        print(
            f"[INFO] To reproduce this test case, run the following command: {args.repro_command}"
        )

    ## Parse input arguments
    backends = list(args.backends)
    m = args.m
    n = args.n
    k = args.k
    is_cuda_graph_compatible = not args.no_cuda_graph
    run_refcheck = args.refcheck
    autotune_supported_backends = ["trtllm_low_latency"]
    res = []

    input_dtype = dtype_str_to_torch_dtype(args.input_dtype)
    if input_dtype != torch.float8_e4m3fn:
        raise ValueError(
            f"Unsupported input dtype: {input_dtype}. mm_fp8 only supports fp8_e4m3."
        )

    mat2_dtype = dtype_str_to_torch_dtype(args.mat2_dtype)
    if mat2_dtype != torch.float8_e4m3fn:
        raise ValueError(
            f"Unsupported mat2 dtype: {mat2_dtype}. mm_fp8 only supports fp8_e4m3."
        )

    res_dtype = dtype_str_to_torch_dtype(args.out_dtype)
    if res_dtype != torch.bfloat16:
        raise ValueError(
            f"Unsupported res dtype: {res_dtype}. mm_fp8 only supports bfloat16."
        )

    if args.batch_size != 1:
        raise ValueError("mm_fp8 is a 2D GEMM; only --batch_size 1 is supported.")
    if k % 128 != 0:
        raise ValueError(f"mm_fp8 requires k % 128 == 0, got k = {k}")
    if n % 32 != 0:
        raise ValueError(f"mm_fp8 requires n % 32 == 0, got n = {n}")
    ## Done parsing input arguments

    ## Prepare input tensors
    input_tensor = torch.randn([m, k], device=device, dtype=torch.bfloat16)
    input_fp8, input_inv_s = to_float8(input_tensor, dtype=input_dtype)

    mat2 = torch.randn([n, k], device=device, dtype=torch.bfloat16)
    mat2_fp8, mat2_inv_s = to_float8(mat2, dtype=mat2_dtype)

    # Pre-shuffle the weights into the (k // 128, n, 128) low-latency block
    # layout. The dict memoizes the (expensive) permutation indices; a fresh one
    # is fine since weight prep runs once here, outside the timed region.
    mat2_prepared = flashinfer.prepare_low_latency_gemm_weights(mat2_fp8, {})
    alpha = (input_inv_s * mat2_inv_s).float()  # a_scale * b_scale

    if args.verbose >= 2:
        print(f"[VVERBOSE] {input_fp8.shape = }")
        print(f"[VVERBOSE] {input_fp8.dtype = }")
        print(f"[VVERBOSE] {mat2_prepared.shape = }")
        print(f"[VVERBOSE] {mat2_prepared.dtype = }")
        print(f"[VVERBOSE] {input_inv_s = }")
        print(f"[VVERBOSE] {mat2_inv_s = }")

    # Programmatically filter backends by probing each backend with a trial
    # call (mm_fp8 has no @backend_requirement decorator; it raises on
    # unsupported backends/architectures).
    # NOTE: unlike some other testers, backends without autotune support are
    # kept when --autotune is set (they just skip the warmup loop below and
    # are reported without the "_autotune" suffix).
    backends_to_remove = []
    for backend in backends:
        try:
            flashinfer.mm_fp8(
                input_fp8,
                mat2_prepared,
                alpha,
                out_dtype=res_dtype,
                backend=backend,
            )
        except Exception as e:
            print(
                f"[INFO] {backend} backend does not support this configuration: {type(e).__name__}: {e}"
            )
            backends_to_remove.append(backend)

    for backend in backends_to_remove:
        backends.remove(backend)

    if len(backends) == 0:
        print("[ERROR] No backends passed validation. Exiting.")
        return res

    def run_backend(backend, input_fp8, mat2_prepared, alpha):
        return flashinfer.mm_fp8(
            input_fp8,
            mat2_prepared,
            alpha,
            out_dtype=res_dtype,
            backend=backend,
        )

    has_reference_output = False
    if run_refcheck:
        reference_output = torch.mm(input_tensor, mat2.T)
        has_reference_output = True

    cache_path = getattr(args, "autotune_cache", None)
    if getattr(args, "autotune", False):
        warmup_iters = (
            args.dry_run_iters if args.dry_run_iters and args.dry_run_iters > 0 else 10
        )
        for cur_backend in backends:
            if cur_backend in autotune_supported_backends:
                if args.verbose >= 1:
                    print(f"[INFO] Autotune warmup for mm_fp8: {warmup_iters} iters")
                with autotune(True, cache=cache_path):
                    for _ in range(warmup_iters):
                        run_backend(cur_backend, input_fp8, mat2_prepared, alpha)
    elif cache_path:
        with autotune(False, cache=cache_path):
            pass

    # Storage for timing results and outputs
    backend_times = {backend: [] for backend in backends}
    outputs = {}
    for cur_backend in backends:
        if run_refcheck:
            outputs[cur_backend] = run_backend(
                cur_backend, input_fp8, mat2_prepared, alpha
            ).detach()
        backend_times[cur_backend] = bench_gpu_time(
            fn=run_backend,
            dry_run_iters=args.dry_run_iters,
            repeat_iters=args.num_iters,
            sleep_after_run=True,
            enable_cupti=args.use_cupti,
            use_cuda_graph=is_cuda_graph_compatible,
            cold_l2_cache=True,
            input_args=(cur_backend, input_fp8, mat2_prepared, alpha),
        )

    tested_backends = list(outputs.keys())
    tested_outputs = list(outputs.values())
    if len(tested_backends) > 0:
        if run_refcheck and has_reference_output:
            for i in range(len(tested_backends)):
                cos_sim = F.cosine_similarity(
                    reference_output.reshape(-1),
                    tested_outputs[i].reshape(-1),
                    dim=0,
                )
                if torch.isnan(cos_sim) or cos_sim < 0.99:
                    print(
                        f"[ERROR] Output tensor mismatch between reference and backend {tested_backends[i]}"
                    )
                    if not args.allow_output_mismatch:
                        raise AssertionError(
                            f"[ERROR] Backend {tested_backends[i]} output mismatch with cos_sim={cos_sim}"
                        )

    for backend in backends:
        backend_name = backend + (
            "_autotune"
            if (
                getattr(args, "autotune", False)
                and backend in autotune_supported_backends
            )
            else ""
        )
        if len(backend_times[backend]) > 0:
            median_time = np.median(backend_times[backend])
            std_time = np.std(backend_times[backend])
            problem_flops = 2 * m * n * k
            problem_bytes = (
                m * k * input_dtype.itemsize
                + n * k * mat2_dtype.itemsize
                + m * n * res_dtype.itemsize
            )
            tflops = problem_flops / (10**9 * median_time)  # in TFLOPs/sec
            tb_per_sec = problem_bytes / (10**9 * median_time)  # in TB/sec
            print_perf_metrics(backend_name, median_time, std_time, tflops, tb_per_sec)

            if args.output_path is not None:
                cur_res = defaultdict(str)
                cur_res["batch_size"] = 1
                cur_res["routine"] = args.routine
                cur_res["median_time"] = median_time
                cur_res["std_time"] = std_time
                cur_res["tflops"] = tflops
                cur_res["tb_per_sec"] = tb_per_sec
                cur_res["m"] = m
                cur_res["n"] = n
                cur_res["k"] = k
                cur_res["input_dtype"] = input_dtype
                cur_res["mat2_dtype"] = mat2_dtype
                cur_res["out_dtype"] = res_dtype
                cur_res["backend"] = backend_name
                cur_res["case_tag"] = args.case_tag
                res.append(cur_res)
    return res


def testBmmMxfp8(args):
    """
    Test bmm_mxfp8 API.

    This test:
    1. Generates random input tensors
    2. Quantizes input tensors to MXFP8
    3. Runs bmm_mxfp8
    4. Runs reference check
    5. Measures performance metrics (TFLOPS, TB/sec)

    Args:
        args: Parsed command line arguments containing test configuration

    Returns:
        dict: List of dictionaries containing performance results
    """
    warn_if_pdl_unsupported(args, args.routine)
    if args.verbose >= 1:
        print("[INFO] Running testBmmMxfp8")
        print(f"[INFO] FlashInfer version: {flashinfer.__version__}")

    device = get_device(args)
    if args.generate_repro_command:
        print(
            f"[INFO] To reproduce this test case, run the following command: {args.repro_command}"
        )

    ## Parse input arguments
    backends = args.backends
    batch_size = args.batch_size
    m = args.m
    n = args.n
    k = args.k
    res_dtype = args.out_dtype
    is_cuda_graph_compatible = not args.no_cuda_graph
    run_refcheck = args.refcheck
    autotune_supported_backends = [
        "cudnn",
    ]
    res = []

    backends = filter_backends_by_compute_capability(backends, args.routine, device)
    if len(backends) == 0:
        print("[ERROR] No backends to test. Exiting.")
        return res

    res_dtype = dtype_str_to_torch_dtype(args.out_dtype)
    if res_dtype not in [torch.bfloat16, torch.float16]:
        raise ValueError(
            f"Unsupported res dtype: {res_dtype}. Supported dtypes are bfloat16 and float16."
        )
    ## Done parsing input arguments

    if getattr(args, "autotune", False):
        backends_to_remove = []
        for cur_backend in backends:
            if cur_backend not in autotune_supported_backends:
                print(f"[INFO] {cur_backend} backend does not support autotune")
                backends_to_remove.append(cur_backend)
        for cur_backend in backends_to_remove:
            backends.remove(cur_backend)

    if len(backends) == 0:
        print("[ERROR] No backends to test. Exiting.")
        return res

    ## Prepare input tensors
    input = torch.randn([batch_size, m, k], device=device, dtype=torch.bfloat16)
    input_mxfp8, input_scale = mxfp8_quantize(input, is_sf_swizzled_layout=True)

    # Quantize [b, n, k] weights, then pass their [b, k, n] transpose as B.
    # Keep the transpose as a view so the K dimension remains contiguous.
    mat2_weight = torch.randn([batch_size, n, k], device=device, dtype=torch.bfloat16)
    mat2_mxfp8, mat2_scale = mxfp8_quantize(mat2_weight, is_sf_swizzled_layout=True)
    # [b, k, n] views for the GEMM call and the reference.
    mat2_mxfp8 = mat2_mxfp8.transpose(-2, -1)
    mat2 = mat2_weight.transpose(-2, -1)

    if args.verbose >= 2:
        print(f"[VVERBOSE] {input_mxfp8.shape = }")
        print(f"[VVERBOSE] {input_mxfp8.dtype = }")
        print(f"[VVERBOSE] {mat2_mxfp8.shape = }")
        print(f"[VVERBOSE] {mat2_mxfp8.dtype = }")
        print(f"[VVERBOSE] {input_scale.shape = }")
        print(f"[VVERBOSE] {input_scale.dtype = }")
        print(f"[VVERBOSE] {mat2_scale.shape = }")
        print(f"[VVERBOSE] {mat2_scale.dtype = }")

    def run_backend(backend, input_mxfp8, mat2_mxfp8, input_scale, mat2_scale):
        if backend == "cudnn":
            return flashinfer.gemm.bmm_mxfp8(
                A=input_mxfp8,
                B=mat2_mxfp8,
                A_scale=input_scale,
                B_scale=mat2_scale,
                dtype=res_dtype,
                backend=backend,
            )
        else:
            raise ValueError(f"Unsupported backend: {backend}")

    has_reference_output = False
    if run_refcheck:
        reference_output = torch.bmm(input, mat2)
        has_reference_output = True

    cache_path = getattr(args, "autotune_cache", None)
    if getattr(args, "autotune", False):
        warmup_iters = (
            args.dry_run_iters if args.dry_run_iters and args.dry_run_iters > 0 else 10
        )
        for cur_backend in backends:
            if cur_backend in autotune_supported_backends:
                if args.verbose >= 1:
                    print(f"[INFO] Autotune warmup for bmm_mxfp8: {warmup_iters} iters")
                with autotune(True, cache=cache_path):
                    for _ in range(warmup_iters):
                        run_backend(
                            cur_backend,
                            input_mxfp8,
                            mat2_mxfp8,
                            input_scale,
                            mat2_scale,
                        )
    elif cache_path:
        with autotune(False, cache=cache_path):
            pass

    # Storage for timing results and outputs
    backend_times = {backend: [] for backend in backends}
    outputs = {}
    for cur_backend in backends:
        if run_refcheck:
            outputs[cur_backend] = run_backend(
                cur_backend, input_mxfp8, mat2_mxfp8, input_scale, mat2_scale
            ).detach()
        backend_times[cur_backend] = bench_gpu_time(
            fn=run_backend,
            dry_run_iters=args.dry_run_iters,
            repeat_iters=args.num_iters,
            sleep_after_run=True,
            enable_cupti=args.use_cupti,
            use_cuda_graph=is_cuda_graph_compatible,
            cold_l2_cache=True,
            input_args=(cur_backend, input_mxfp8, mat2_mxfp8, input_scale, mat2_scale),
        )

    min_cos_sim = 0.9  # TODO: check if can be increased

    tested_backends = list(outputs.keys())
    tested_outputs = list(outputs.values())
    if len(tested_backends) > 0:
        if run_refcheck and has_reference_output:
            for i in range(len(tested_backends)):
                cos_sim = F.cosine_similarity(
                    reference_output.reshape(-1),
                    tested_outputs[i].reshape(-1),
                    dim=0,
                )
                if cos_sim < min_cos_sim:
                    print(
                        f"[ERROR] Output tensor mismatch between backends {tested_backends[0]} and {tested_backends[i]}"
                    )
                    if not args.allow_output_mismatch:
                        raise AssertionError(
                            f"[ERROR] Backend {tested_backends[i]} output mismatch with cos_sim={cos_sim}"
                        )

    for backend in backends:
        backend_name = backend + (
            "_autotune"
            if (
                getattr(args, "autotune", False)
                and backend in autotune_supported_backends
            )
            else ""
        )
        if len(backend_times[backend]) > 0:
            median_time = np.median(backend_times[backend])
            std_time = np.std(backend_times[backend])
            problem_flops = 2 * m * n * k * batch_size
            # MXFP8 uses fp8_e4m3fn for data (1 byte) and uint8 for scales
            # Scale tensors are much smaller, so approximate as 1 byte per element for simplicity
            problem_bytes = (
                m * k * torch.float8_e4m3fn.itemsize
                + n * k * torch.float8_e4m3fn.itemsize
                + m * n * res_dtype.itemsize
            ) * batch_size
            tflops = problem_flops / (10**9 * median_time)  # in TFLOPs/sec
            tb_per_sec = problem_bytes / (10**9 * median_time)  # in TB/sec
            print_perf_metrics(backend_name, median_time, std_time, tflops, tb_per_sec)

            if args.output_path is not None:
                cur_res = defaultdict(str)
                cur_res["batch_size"] = batch_size
                cur_res["routine"] = args.routine
                cur_res["median_time"] = median_time
                cur_res["std_time"] = std_time
                cur_res["tflops"] = tflops
                cur_res["tb_per_sec"] = tb_per_sec
                cur_res["m"] = m
                cur_res["n"] = n
                cur_res["k"] = k
                cur_res["out_dtype"] = res_dtype
                cur_res["backend"] = backend_name
                cur_res["case_tag"] = args.case_tag
                res.append(cur_res)
    return res


def testMmFp4(args):
    """
    Test mm_fp4 API.

    This test:
    1. Generates random input tensors
    2. Quantizes input tensors to FP4
    3. Runs mm_fp4
    4. Runs reference check
    5. Measures performance metrics (TFLOPS, TB/sec)

    Args:
        args: Parsed command line arguments containing test configuration

    Returns:
        dict: List of dictionaries containing performance results
    """
    if args.verbose >= 1:
        print("[INFO] Running testMmFp4")
        print(f"[INFO] FlashInfer version: {flashinfer.__version__}")

    device = get_device(args)
    if args.generate_repro_command:
        print(
            f"[INFO] To reproduce this test case, run the following command: {args.repro_command}"
        )

    ## Parse input arguments
    backends = args.backends
    m = args.m
    n = args.n
    k = args.k
    res_dtype = args.out_dtype
    backends = args.backends
    is_cuda_graph_compatible = not args.no_cuda_graph
    run_refcheck = args.refcheck
    use_128x4_sf_layout = args.use_128x4_sf_layout
    use_nvfp4 = args.use_nvfp4
    autotune_supported_backends = [
        "cudnn",
        "cutlass",
        "trtllm",
        "cute-dsl",
        "b12x",
        "auto",
    ]
    res = []

    res_dtype = dtype_str_to_torch_dtype(args.out_dtype)
    if res_dtype not in [torch.bfloat16, torch.float16]:
        raise ValueError(
            f"Unsupported res dtype: {res_dtype}. Supported dtypes are bfloat16 and float16."
        )

    input = torch.randn([m, k], device=device, dtype=torch.bfloat16)
    mat2 = torch.randn([n, k], device=device, dtype=torch.bfloat16)
    a_sf_layout = (
        flashinfer.SfLayout.layout_128x4
        if use_128x4_sf_layout
        else flashinfer.SfLayout.layout_8x4
    )

    global_sf_input = (448 * 6) / input.float().abs().nan_to_num().max()
    global_sf_mat2 = (448 * 6) / mat2.float().abs().nan_to_num().max()

    if use_nvfp4:
        input_fp4, input_inv_s = flashinfer.nvfp4_quantize(
            input, global_sf_input, sfLayout=a_sf_layout, do_shuffle=False
        )
        mat2_fp4, mat2_inv_s = flashinfer.nvfp4_quantize(
            mat2,
            global_sf_mat2,
            sfLayout=flashinfer.SfLayout.layout_128x4,
            do_shuffle=False,
        )
    else:  # mxfp4
        input_fp4, input_inv_s = flashinfer.mxfp4_quantize(input)
        mat2_fp4, mat2_inv_s = flashinfer.mxfp4_quantize(mat2)

    # The trtllm backend consumes a shuffled weight/scale layout. Prepare it
    # only when that backend is requested: the shuffle requires n to be a
    # multiple of 128, and an unaligned n must not fail the other backends.
    mat2_fp4_trtllm, mat2_inv_s_trtllm = None, None
    if "trtllm" in backends:
        try:
            mat2_fp4_trtllm, mat2_inv_s_trtllm = flashinfer.nvfp4_quantize(
                mat2,
                global_sf_mat2,
                sfLayout=flashinfer.SfLayout.layout_128x4,
                do_shuffle=True,
            )
        except Exception as e:
            print(
                f"[INFO] trtllm backend does not support this configuration: {type(e).__name__}: {e}"
            )
            backends.remove("trtllm")

    if args.verbose >= 2:
        print(f"[VVERBOSE] {input_fp4.shape = }")
        print(f"[VVERBOSE] {input_fp4.dtype = }")
        print(f"[VVERBOSE] {mat2_fp4.shape = }")
        print(f"[VVERBOSE] {mat2_fp4.dtype = }")

    alpha = 1.0 / (global_sf_input * global_sf_mat2) if use_nvfp4 else None
    # TODO: for MXFP4, we don't need a global scale, we should change the compile interface to make
    # alpha optional.
    alpha_for_cute_dsl_mxfp4 = (
        torch.tensor([1.0], dtype=torch.float32, device=device)
        if not use_nvfp4
        else None
    )
    # Completed preparing inputs. Now programmatically filter backends
    block_size = 16 if use_nvfp4 else 32
    backends_to_remove = []

    for backend in backends:
        # Skip autotune check for now (handled separately below)
        if (
            getattr(args, "autotune", False)
            and backend not in autotune_supported_backends
        ):
            print(f"[INFO] {backend} backend does not support autotune")
            backends_to_remove.append(backend)
            continue

        try:
            flashinfer.gemm.mm_fp4(
                a=input_fp4,
                b=mat2_fp4.T if backend != "trtllm" else mat2_fp4_trtllm.T,
                a_descale=input_inv_s,
                b_descale=mat2_inv_s.T if backend != "trtllm" else mat2_inv_s_trtllm.T,
                alpha=(alpha_for_cute_dsl_mxfp4 if (backend == "cute-dsl") else alpha),
                out_dtype=res_dtype,
                block_size=16
                if use_nvfp4
                else 32,  # nvfp4 only supports 16; mxfp4 only supports 32.
                use_8x4_sf_layout=not use_128x4_sf_layout,
                backend=backend,
                use_nvfp4=use_nvfp4,
                enable_pdl=args.enable_pdl,
            )
        except Exception as e:
            print(
                f"[INFO] {backend} backend does not support this configuration: {type(e).__name__}: {e}"
            )
            backends_to_remove.append(backend)

    # Remove unsupported backends
    for backend in backends_to_remove:
        backends.remove(backend)

    if len(backends) == 0:
        print("[ERROR] No backends passed validation. Exiting.")
        return res

    def run_backend(
        backend,
        input_fp4,
        mat2_fp4,
        mat2_fp4_trtllm,
        input_inv_s,
        mat2_inv_s,
        mat2_inv_s_trtllm,
    ):
        return flashinfer.gemm.mm_fp4(
            a=input_fp4,
            b=mat2_fp4.T if backend != "trtllm" else mat2_fp4_trtllm.T,
            a_descale=input_inv_s,
            b_descale=mat2_inv_s.T if backend != "trtllm" else mat2_inv_s_trtllm.T,
            alpha=(alpha_for_cute_dsl_mxfp4 if (backend == "cute-dsl") else alpha),
            out_dtype=res_dtype,
            block_size=block_size,
            use_8x4_sf_layout=not use_128x4_sf_layout,
            backend=backend,
            use_nvfp4=use_nvfp4,
            enable_pdl=args.enable_pdl,
        )

    has_reference_output = False
    if run_refcheck:
        reference_output = torch.mm(input, mat2.T)
        has_reference_output = True

    cache_path = getattr(args, "autotune_cache", None)
    if getattr(args, "autotune", False):
        warmup_iters = (
            args.dry_run_iters if args.dry_run_iters and args.dry_run_iters > 0 else 10
        )
        for cur_backend in backends:
            if args.verbose >= 1:
                print(f"[INFO] Autotune warmup for mm_fp4: {warmup_iters} iters")
            with autotune(True, cache=cache_path):
                for _ in range(warmup_iters):
                    run_backend(
                        cur_backend,
                        input_fp4,
                        mat2_fp4,
                        mat2_fp4_trtllm,
                        input_inv_s,
                        mat2_inv_s,
                        mat2_inv_s_trtllm,
                    )
    elif cache_path:
        with autotune(False, cache=cache_path):
            pass

    # Storage for timing results and outputs
    backend_times = {backend: [] for backend in backends}
    outputs = {}
    for cur_backend in backends:
        if run_refcheck:
            outputs[cur_backend] = run_backend(
                cur_backend,
                input_fp4,
                mat2_fp4,
                mat2_fp4_trtllm,
                input_inv_s,
                mat2_inv_s,
                mat2_inv_s_trtllm,
            ).detach()
        backend_times[cur_backend] = bench_gpu_time(
            fn=run_backend,
            dry_run_iters=args.dry_run_iters,
            repeat_iters=args.num_iters,
            sleep_after_run=True,
            enable_cupti=args.use_cupti,
            use_cuda_graph=is_cuda_graph_compatible,
            cold_l2_cache=True,
            input_args=(
                cur_backend,
                input_fp4,
                mat2_fp4,
                mat2_fp4_trtllm,
                input_inv_s,
                mat2_inv_s,
                mat2_inv_s_trtllm,
            ),
        )

    tested_backends = list(outputs.keys())
    tested_outputs = list(outputs.values())
    if len(tested_backends) > 0:
        if run_refcheck and has_reference_output:
            for i in range(len(tested_backends)):
                cos_sim = F.cosine_similarity(
                    reference_output.reshape(-1),
                    tested_outputs[i].reshape(-1),
                    dim=0,
                )
                if cos_sim < 0.97:
                    print(
                        f"[ERROR] Output tensor mismatch between backends {tested_backends[0]} and {tested_backends[i]}"
                    )
                    if not args.allow_output_mismatch:
                        raise AssertionError(
                            f"[ERROR] Backend {tested_backends[i]} output mismatch with cos_sim={cos_sim}"
                        )

    for backend in backends:
        backend_name = backend + (
            "_autotune"
            if (
                getattr(args, "autotune", False)
                and backend in autotune_supported_backends
            )
            else ""
        )
        if len(backend_times[backend]) > 0:
            median_time = np.median(backend_times[backend])
            std_time = np.std(backend_times[backend])
            problem_flops = 2 * m * n * k
            problem_bytes = (
                m * k * 0.5 + n * k * 0.5 + m * n * res_dtype.itemsize
            )  # 0.5 for fp4
            tflops = problem_flops / (10**9 * median_time)  # in TFLOPs/sec
            tb_per_sec = problem_bytes / (10**9 * median_time)  # in TB/sec
            print_perf_metrics(backend_name, median_time, std_time, tflops, tb_per_sec)

            if args.output_path is not None:
                cur_res = defaultdict(str)
                cur_res["routine"] = args.routine
                cur_res["median_time"] = median_time
                cur_res["std_time"] = std_time
                cur_res["tflops"] = tflops
                cur_res["tb_per_sec"] = tb_per_sec
                cur_res["m"] = m
                cur_res["n"] = n
                cur_res["k"] = k
                cur_res["out_dtype"] = res_dtype
                cur_res["use_128x4_sf_layout"] = use_128x4_sf_layout
                cur_res["backend"] = backend_name
                cur_res["use_nvfp4"] = use_nvfp4
                cur_res["case_tag"] = args.case_tag
                res.append(cur_res)
    return res


# E2M1 (FP4) value table, signed (codes 0-7 positive, 8-15 negative), matching
# ``flashinfer.nvfp4_quantize``.
_E2M1_VALUES_FP32 = (
    0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0,
    -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0,
)  # fmt: skip


def _dequantize_bf16_fp4_ref(b, b_descale, alpha, n, k, block_size):
    """PyTorch implementation of swizzled nvfp4 dequantization to fp32."""
    import torch
    from flashinfer.gemm.gemm_bf16_fp4 import _unswizzle_sf_128x4

    device = b.device
    k_sf = k // block_size
    lut = torch.tensor(_E2M1_VALUES_FP32, dtype=torch.float32, device=device)
    b_int = b.to(torch.int64)
    codes = torch.stack([b_int & 0xF, (b_int >> 4) & 0xF], dim=-1).reshape(n, k)
    values = lut[codes]
    sf = _unswizzle_sf_128x4(b_descale, n, k_sf).view(torch.float8_e4m3fn)
    sf_expanded = sf.to(torch.float32).repeat_interleave(block_size, dim=1)
    weight = values * sf_expanded
    if alpha is not None:
        weight = weight * alpha.to(torch.float32)
    return weight


def testMmBf16Fp4(args):
    """Benchmark mm_bf16_fp4 (bf16 activation x FP4 weight, bf16 x fp4).

    Weights are produced by ``flashinfer.nvfp4_quantize(sfLayout=layout_128x4)``
    -- the same format the new API expects.

    Constraints:
      * N must be divisible by 64 (the weight-prepack / kernel N tile).
        The 128x4 SF swizzle only pads N to 128; prepare_bf16_fp4_weights
        unswizzles the padded tail, so any N % 64 == 0 works.
      * K must be divisible by 16 (FP4 block size).
      * input_dtype must be bfloat16 (the only A dtype currently
        supported; fp16 deferred).

    Refcheck uses an fp32 dequant + matmul.
    """
    if args.verbose >= 1:
        print("[INFO] Running testMmBf16Fp4")
        print(f"[INFO] FlashInfer version: {flashinfer.__version__}")

    device = get_device(args)
    if args.generate_repro_command:
        print(
            f"[INFO] To reproduce this test case, run the following command: {args.repro_command}"
        )

    m, n, k = args.m, args.n, args.k
    input_dtype = dtype_str_to_torch_dtype(args.input_dtype)
    out_dtype = dtype_str_to_torch_dtype(args.out_dtype)
    backends = args.backends
    run_refcheck = args.refcheck
    is_cuda_graph_compatible = not args.no_cuda_graph

    if input_dtype != torch.bfloat16:
        raise ValueError(
            f"mm_bf16_fp4 benchmark requires input_dtype=bfloat16, got {args.input_dtype}"
        )
    if out_dtype not in (torch.bfloat16, torch.float16):
        raise ValueError(
            f"mm_bf16_fp4 benchmark requires out_dtype in (bfloat16, float16), got {args.out_dtype}"
        )
    if n % 64 != 0:
        # N must be a multiple of the 64-wide N tile used by the weight
        # prepack and the cute-dsl kernel.  The 128x4 SF swizzle only *pads*
        # N to 128, and prepare_bf16_fp4_weights unswizzles the padded
        # tail, so any n % 64 == 0 works.
        raise ValueError("mm_bf16_fp4 benchmark requires n % 64 == 0")
    if k % 16 != 0:
        raise ValueError("mm_bf16_fp4 benchmark requires k % 16 == 0 (FP4 block size)")

    torch.manual_seed(args.random_seed)
    a = torch.randn((m, k), device=device, dtype=input_dtype) * 0.5
    w = torch.randn((n, k), device=device, dtype=input_dtype) * 0.1
    g_b = (448 * 6) / w.float().abs().nan_to_num().max()
    b_fp4, b_sf = flashinfer.nvfp4_quantize(
        w,
        g_b,
        sfLayout=flashinfer.SfLayout.layout_128x4,
        do_shuffle=False,
        backend="cute-dsl",
    )
    alpha = torch.tensor([1.0 / g_b.item()], device=device, dtype=torch.float32)

    if args.verbose >= 2:
        print(f"[VVERBOSE] {a.shape = } {a.dtype = }")
        print(f"[VVERBOSE] {b_fp4.shape = } {b_fp4.dtype = }")
        print(f"[VVERBOSE] {b_sf.shape = } {b_sf.dtype = }")

    # Per-backend prep + runner closures.  Prep is one-shot and not timed.
    backend_runners = {}

    def make_runner(b_p, sf_p, alpha_p, backend):
        def run(a):
            # mm_bf16_fp4: bf16 activation a against the prepared FP4
            # weight; b_p/sf_p are prepare_bf16_fp4_weights outputs.
            return flashinfer.mm_bf16_fp4(
                a,
                b_p,
                sf_p,
                alpha_p,
                backend=backend,
                out_dtype=out_dtype,
                block_size=16,
                enable_pdl=args.enable_pdl,
            )

        return run

    backends_to_remove = []
    for backend in backends:
        try:
            b_p, sf_p, alpha_p = flashinfer.prepare_bf16_fp4_weights(
                b_fp4, b_sf, alpha, backend=backend
            )
            runner = make_runner(b_p, sf_p, alpha_p, backend)
            runner(a)
            backend_runners[backend] = runner
        except Exception as e:
            print(
                f"[INFO] {backend} backend does not support this configuration: {type(e).__name__}: {e}"
            )
            backends_to_remove.append(backend)

    for backend in backends_to_remove:
        backends.remove(backend)

    if len(backends) == 0:
        print("[ERROR] No backends passed validation. Exiting.")
        return []

    autotune_supported_backends = ["cudnn", "cute-dsl"]
    cache_path = getattr(args, "autotune_cache", None)
    if getattr(args, "autotune", False):
        warmup_iters = (
            args.dry_run_iters if args.dry_run_iters and args.dry_run_iters > 0 else 10
        )
        for cur_backend in backends:
            if cur_backend in autotune_supported_backends:
                if args.verbose >= 1:
                    print(
                        f"[INFO] Autotune warmup for mm_bf16_fp4 {cur_backend}: "
                        f"{warmup_iters} iters"
                    )
                with autotune(True, cache=cache_path):
                    for _ in range(warmup_iters):
                        backend_runners[cur_backend](a)
    elif cache_path:
        with autotune(False, cache=cache_path):
            pass

    ref = None
    if run_refcheck:
        weight_fp32 = _dequantize_bf16_fp4_ref(b_fp4, b_sf, alpha, n, k, 16)
        ref = (a.float() @ weight_fp32.T).to(out_dtype)

    res = []
    flops = 2 * m * n * k
    bytes_accessed = (
        m * k * input_dtype.itemsize
        + (k // 2) * n  # FP4 weight (uint8, 2 codes / byte)
        + (k // 16) * n  # FP8-E4M3 per-block SF
        + m * n * out_dtype.itemsize
    )

    refcheck_tol = dict(rtol=1.5e-2, atol=1.5e-2)
    for backend in backends:
        runner = backend_runners[backend]
        if run_refcheck:
            out = runner(a)
            try:
                torch.testing.assert_close(out, ref, **refcheck_tol)
            except AssertionError as e:
                if args.allow_output_mismatch:
                    print(f"[WARNING] {backend} output mismatch vs fp32 ref: {e}")
                else:
                    raise

        timing = bench_gpu_time(
            fn=runner,
            dry_run_iters=args.dry_run_iters,
            repeat_iters=args.num_iters,
            sleep_after_run=True,  # GEMMs are very MMA-heavy, so prefer sleep to reduce throttling.
            enable_cupti=args.use_cupti,
            use_cuda_graph=is_cuda_graph_compatible,
            cold_l2_cache=True,
            input_args=(a,),
        )
        median_time = float(np.median(timing))
        std_time = float(np.std(timing))
        tflops = flops / median_time / 1e9
        tb_per_sec = bytes_accessed / median_time / 1e9
        backend_name = backend + (
            "_autotune"
            if (
                getattr(args, "autotune", False)
                and backend in autotune_supported_backends
            )
            else ""
        )
        print_perf_metrics(backend_name, median_time, std_time, tflops, tb_per_sec)
        res.append(
            {
                "routine": args.routine,
                "median_time": median_time,
                "std_time": std_time,
                "tflops": tflops,
                "tb_per_sec": tb_per_sec,
                "backend": backend_name,
                "resolved_backend": backend,
                "m": m,
                "n": n,
                "k": k,
                "input_dtype": str(input_dtype).split(".")[-1],
                "out_dtype": str(out_dtype).split(".")[-1],
                "case_tag": args.case_tag,
            }
        )
    return res


def testMmMxfp8(args):
    """
    Test mm_mxfp8 API.

    This test:
    1. Generates random input tensors
    2. Quantizes input tensors to MXFP8
    3. Runs mm_mxfp8
    4. Runs reference check
    5. Measures performance metrics (TFLOPS, TB/sec)

    Args:
        args: Parsed command line arguments containing test configuration

    Returns:
        dict: List of dictionaries containing performance results
    """
    warn_if_pdl_unsupported(args, args.routine)
    if args.verbose >= 1:
        print("[INFO] Running testMmMxfp8")
        print(f"[INFO] FlashInfer version: {flashinfer.__version__}")

    device = get_device(args)
    if args.generate_repro_command:
        print(
            f"[INFO] To reproduce this test case, run the following command: {args.repro_command}"
        )

    ## Parse input arguments
    backends = args.backends
    m = args.m
    n = args.n
    k = args.k
    res_dtype = args.out_dtype
    is_cuda_graph_compatible = not args.no_cuda_graph
    run_refcheck = args.refcheck
    autotune_supported_backends = [
        "cutlass",
        "cute-dsl",
        "trtllm",
        "cudnn",
        "auto",
    ]
    res = []

    backends = filter_backends_by_compute_capability(backends, args.routine, device)
    if len(backends) == 0:
        print("[ERROR] No backends to test. Exiting.")
        return res

    res_dtype = dtype_str_to_torch_dtype(args.out_dtype)
    if res_dtype not in [torch.bfloat16, torch.float16]:
        raise ValueError(
            f"Unsupported res dtype: {res_dtype}. Supported dtypes are bfloat16 and float16."
        )
    ## Done parsing input arguments

    if getattr(args, "autotune", False):
        backends_to_remove = []
        for cur_backend in backends:
            if cur_backend not in autotune_supported_backends:
                print(f"[INFO] {cur_backend} backend does not support autotune")
                backends_to_remove.append(cur_backend)
        for cur_backend in backends_to_remove:
            backends.remove(cur_backend)

    if len(backends) == 0:
        print("[ERROR] No backends to test. Exiting.")
        return res

    inputs = {}
    input = torch.randn([m, k], device=device, dtype=torch.bfloat16)
    mat2 = torch.randn([n, k], device=device, dtype=torch.bfloat16)
    for backend in backends:
        ## Prepare input tensors
        # Every backend consumes swizzled scales; trtllm optionally uses 8x4 for A.
        if backend == "trtllm" and not args.use_128x4_sf_layout:
            sf_layout_input = flashinfer.SfLayout.layout_8x4
        else:
            sf_layout_input = flashinfer.SfLayout.layout_128x4
        input_mxfp8, input_scale = mxfp8_quantize(
            input, sf_swizzle_layout=sf_layout_input
        )
        # when using trtllm, the shuffle_matrix_sf_a will swizzle the layout.
        mat2_mxfp8, mat2_scale = mxfp8_quantize(
            mat2,
            is_sf_swizzled_layout=backend != "trtllm",
        )

        if backend == "trtllm":
            from flashinfer import shuffle_matrix_a, shuffle_matrix_sf_a

            mat2_mxfp8 = shuffle_matrix_a(mat2_mxfp8, 128).reshape(n, k)
            mat2_scale = shuffle_matrix_sf_a(
                mat2_scale.reshape(n, k // 32),
                128,
                num_elts_per_sf=32,
            )
            mat2_scale = mat2_scale.t()

        if args.verbose >= 2:
            print(f"[VERBOSE] {backend}: {input_mxfp8.shape = }")
            print(f"[VERBOSE] {backend}: {input_mxfp8.dtype = }")
            print(f"[VERBOSE] {backend}: {mat2_mxfp8.shape = }")
            print(f"[VERBOSE] {backend}: {mat2_mxfp8.dtype = }")
            print(f"[VERBOSE] {backend}: {input_scale.shape = }")
            print(f"[VERBOSE] {backend}: {input_scale.dtype = }")
            print(f"[VERBOSE] {backend}: {mat2_scale.shape = }")
            print(f"[VERBOSE] {backend}: {mat2_scale.dtype = }")
        inputs[backend] = (input_mxfp8, mat2_mxfp8, input_scale, mat2_scale)

    def run_backend(
        backend: str,
        inputs: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor],
    ) -> torch.Tensor:
        assert backend in ["cutlass", "trtllm", "cute-dsl", "cudnn", "auto"], (
            f"Unsupported backend: {backend}"
        )
        input_mxfp8, mat2_mxfp8, input_scale, mat2_scale = inputs
        return flashinfer.gemm.mm_mxfp8(
            a=input_mxfp8,
            b=mat2_mxfp8.t(),  # mm_mxfp8 expects b.t()
            a_descale=input_scale,
            b_descale=mat2_scale,
            out_dtype=res_dtype,
            backend=backend,
            use_8x4_sf_layout=backend == "trtllm" and not args.use_128x4_sf_layout,
        )

    has_reference_output = False
    if run_refcheck:
        reference_output = torch.mm(input, mat2.t())
        has_reference_output = True

    cache_path = getattr(args, "autotune_cache", None)
    if getattr(args, "autotune", False):
        warmup_iters = (
            args.dry_run_iters if args.dry_run_iters and args.dry_run_iters > 0 else 10
        )
        for cur_backend in backends:
            if cur_backend in autotune_supported_backends:
                if args.verbose >= 1:
                    print(f"[INFO] Autotune warmup for mm_mxfp8: {warmup_iters} iters")
                with autotune(True, cache=cache_path):
                    for _ in range(warmup_iters):
                        run_backend(
                            cur_backend,
                            inputs[cur_backend],
                        )
    elif cache_path:
        with autotune(False, cache=cache_path):
            pass

    # Storage for timing results and outputs
    backend_times = {backend: [] for backend in backends}
    outputs = {}
    for cur_backend in backends:
        if run_refcheck:
            outputs[cur_backend] = run_backend(
                cur_backend, inputs[cur_backend]
            ).detach()
        backend_times[cur_backend] = bench_gpu_time(
            fn=run_backend,
            dry_run_iters=args.dry_run_iters,
            repeat_iters=args.num_iters,
            sleep_after_run=True,
            enable_cupti=args.use_cupti,
            use_cuda_graph=is_cuda_graph_compatible,
            cold_l2_cache=True,
            input_args=(cur_backend, inputs[cur_backend]),
        )

    # Minimum cosine similarity for swizzled layout
    min_cos_sim = 0.98

    tested_backends = list(outputs.keys())
    tested_outputs = list(outputs.values())
    if len(tested_backends) > 0:
        if run_refcheck and has_reference_output:
            for i in range(len(tested_backends)):
                cos_sim = F.cosine_similarity(
                    reference_output.reshape(-1),
                    tested_outputs[i].reshape(-1),
                    dim=0,
                )
                if cos_sim < min_cos_sim:
                    print(
                        "[ERROR] Output tensor mismatch between reference "
                        f"{tested_backends[0]} and backend {tested_backends[i]}"
                    )
                    if not args.allow_output_mismatch:
                        raise AssertionError(
                            "[ERROR] Output tensor mismatch between reference "
                            f"{tested_backends[0]} and backend {tested_backends[i]} "
                            f"with {cos_sim=} (expected >= {min_cos_sim})"
                        )

    for backend in backends:
        backend_name = backend + (
            "_autotune"
            if (
                getattr(args, "autotune", False)
                and backend in autotune_supported_backends
            )
            else ""
        )
        if len(backend_times[backend]) > 0:
            median_time = np.median(backend_times[backend])
            std_time = np.std(backend_times[backend])
            problem_flops = 2 * m * n * k
            # MXFP8 uses fp8_e4m3fn for data (1 byte) and uint8 for scales
            # Scale tensors are much smaller, so approximate as 1 byte per element for simplicity
            problem_bytes = (
                m * k * torch.float8_e4m3fn.itemsize
                + n * k * torch.float8_e4m3fn.itemsize
                + m * n * res_dtype.itemsize
            )
            tflops = problem_flops / (10**9 * median_time)  # in TFLOPs/sec
            tb_per_sec = problem_bytes / (10**9 * median_time)  # in TB/sec
            print_perf_metrics(backend_name, median_time, std_time, tflops, tb_per_sec)

            if args.output_path is not None:
                cur_res = defaultdict(str)
                cur_res["routine"] = args.routine
                cur_res["median_time"] = median_time
                cur_res["std_time"] = std_time
                cur_res["tflops"] = tflops
                cur_res["tb_per_sec"] = tb_per_sec
                cur_res["m"] = m
                cur_res["n"] = n
                cur_res["k"] = k
                cur_res["out_dtype"] = res_dtype
                cur_res["backend"] = backend_name
                cur_res["case_tag"] = args.case_tag
                res.append(cur_res)
    return res


def testMmBf16(args):
    """
    Test mm_bf16 API.

    This test:
    1. Generates random BF16 input tensors
    2. Runs mm_bf16 with specified backend
    3. Runs reference check (torch.mm)
    4. Measures performance metrics (TFLOPS, TB/sec)

    Args:
        args: Parsed command line arguments containing test configuration

    Returns:
        dict: List of dictionaries containing performance results
    """
    if args.verbose >= 1:
        print("[INFO] Running testMmBf16")
        print(f"[INFO] FlashInfer version: {flashinfer.__version__}")

    device = get_device(args)
    if args.generate_repro_command:
        print(
            f"[INFO] To reproduce this test case, run the following command: {args.repro_command}"
        )

    ## Parse input arguments
    backends = args.backends
    m = args.m
    n = args.n
    k = args.k
    use_bias = getattr(args, "bias", False)
    use_pdl = getattr(args, "enable_pdl", False)
    is_cuda_graph_compatible = not args.no_cuda_graph
    run_refcheck = args.refcheck
    autotune_supported_backends = [
        "cudnn",
        "cutlass",
        "tgv",
        "cublaslt",
        "tinygemm",
        "cute-dsl",
        "auto",
    ]
    res = []

    out_dtype = dtype_str_to_torch_dtype(args.out_dtype)
    if out_dtype not in [torch.bfloat16, torch.float16, torch.float32]:
        raise ValueError(
            f"Unsupported output dtype: {args.out_dtype}. Supported dtypes are bfloat16, float16, and float32."
        )

    ## Prepare input tensors
    # a: (m, k), row-major
    a = torch.randn([m, k], device=device, dtype=torch.bfloat16)
    # b: (n, k) then transpose to (k, n) for column-major layout
    b = torch.randn([n, k], device=device, dtype=torch.bfloat16).transpose(-2, -1)

    bias = None
    if use_bias:
        bias = torch.randn([n], device=device, dtype=torch.bfloat16)

    if args.verbose >= 2:
        print(f"[VVERBOSE] {a.shape = }")
        print(f"[VVERBOSE] {a.dtype = }")
        print(f"[VVERBOSE] {b.shape = }")
        print(f"[VVERBOSE] {b.dtype = }")
        if bias is not None:
            print(f"[VVERBOSE] {bias.shape = }")
            print(f"[VVERBOSE] {bias.dtype = }")
        print(f"[VVERBOSE] {use_pdl = }")

    # Programmatically filter backends
    backends_to_remove = []
    for backend in backends:
        # Skip autotune check for now (handled separately below)
        if (
            getattr(args, "autotune", False)
            and backend not in autotune_supported_backends
        ):
            print(f"[INFO] {backend} backend does not support autotune")
            backends_to_remove.append(backend)
            continue

        try:
            flashinfer.mm_bf16(
                a=a,
                b=b,
                bias=bias,
                pdl=use_pdl,
                out_dtype=out_dtype,
                backend=backend,
            )
        except Exception as e:
            print(
                f"[INFO] {backend} backend does not support this configuration: {type(e).__name__}: {e}"
            )
            backends_to_remove.append(backend)

    # Remove unsupported backends
    for backend in backends_to_remove:
        backends.remove(backend)

    if len(backends) == 0:
        print("[ERROR] No backends passed validation. Exiting.")
        return res

    def run_backend(backend, a, b, bias, use_pdl, out_dtype):
        if backend in [
            "cudnn",
            "cutlass",
            "tgv",
            "cublaslt",
            "tinygemm",
            "cutile",
            "cute-dsl",
            "auto",
        ]:
            return flashinfer.mm_bf16(
                a=a,
                b=b,
                bias=bias,
                pdl=use_pdl,
                out_dtype=out_dtype,
                backend=backend,
            )
        else:
            raise ValueError(f"Unsupported backend: {backend}")

    has_reference_output = False
    reference_output_base = None
    if run_refcheck:
        reference_output_base = torch.mm(a, b).to(out_dtype)
        has_reference_output = True

    cache_path = getattr(args, "autotune_cache", None)
    if getattr(args, "autotune", False):
        warmup_iters = (
            args.dry_run_iters if args.dry_run_iters and args.dry_run_iters > 0 else 10
        )
        for cur_backend in backends:
            if cur_backend in autotune_supported_backends:
                if args.verbose >= 1:
                    print(f"[INFO] Autotune warmup for mm_bf16: {warmup_iters} iters")
                with autotune(True, cache=cache_path):
                    for _ in range(warmup_iters):
                        run_backend(cur_backend, a, b, bias, use_pdl, out_dtype)
    elif cache_path:
        with autotune(False, cache=cache_path):
            pass

    # Storage for timing results and outputs
    backend_times = {backend: [] for backend in backends}
    outputs = {}
    for cur_backend in backends:
        if run_refcheck:
            outputs[cur_backend] = run_backend(
                cur_backend, a, b, bias, use_pdl, out_dtype
            ).detach()
        backend_times[cur_backend] = bench_gpu_time(
            fn=run_backend,
            dry_run_iters=args.dry_run_iters,
            repeat_iters=args.num_iters,
            sleep_after_run=True,
            enable_cupti=args.use_cupti,
            use_cuda_graph=is_cuda_graph_compatible,
            cold_l2_cache=True,
            input_args=(cur_backend, a, b, bias, use_pdl, out_dtype),
        )

    tested_backends = list(outputs.keys())
    tested_outputs = list(outputs.values())
    if len(tested_backends) > 0:
        if run_refcheck and has_reference_output:
            for i in range(len(tested_backends)):
                if (
                    tested_backends[i]
                    in (
                        "cudnn",
                        "tgv",
                        "tinygemm",
                        "cute-dsl",
                    )
                    and bias is not None
                ):
                    reference_output = reference_output_base + bias.unsqueeze(0).to(
                        out_dtype
                    )
                else:
                    reference_output = reference_output_base

                cos_sim = F.cosine_similarity(
                    reference_output.reshape(-1),
                    tested_outputs[i].reshape(-1),
                    dim=0,
                )
                if cos_sim < 0.99:
                    print(
                        f"[ERROR] Output tensor mismatch from backend {tested_backends[i]} with cos_sim={cos_sim}"
                    )
                    if not args.allow_output_mismatch:
                        raise AssertionError(
                            f"[ERROR] Backend {tested_backends[i]} output mismatch with cos_sim={cos_sim}"
                        )

    for backend in backends:
        backend_name = backend + (
            "_autotune"
            if (
                getattr(args, "autotune", False)
                and backend in autotune_supported_backends
            )
            else ""
        )
        if len(backend_times[backend]) > 0:
            median_time = np.median(backend_times[backend])
            std_time = np.std(backend_times[backend])
            problem_flops = 2 * m * n * k
            if use_bias:
                problem_flops += m * n  # bias addition
            problem_bytes = (
                m * k * torch.bfloat16.itemsize
                + k * n * torch.bfloat16.itemsize
                + m * n * out_dtype.itemsize
            )
            if use_bias:
                problem_bytes += n * torch.bfloat16.itemsize
            tflops = problem_flops / (10**9 * median_time)  # in TFLOPs/sec
            tb_per_sec = problem_bytes / (10**9 * median_time)  # in TB/sec
            print_perf_metrics(backend_name, median_time, std_time, tflops, tb_per_sec)

            if args.output_path is not None:
                cur_res = defaultdict(str)
                cur_res["routine"] = args.routine
                cur_res["median_time"] = median_time
                cur_res["std_time"] = std_time
                cur_res["tflops"] = tflops
                cur_res["tb_per_sec"] = tb_per_sec
                cur_res["m"] = m
                cur_res["n"] = n
                cur_res["k"] = k
                cur_res["out_dtype"] = str(out_dtype)
                cur_res["backend"] = backend_name
                cur_res["bias"] = use_bias
                cur_res["enable_pdl"] = use_pdl
                cur_res["case_tag"] = args.case_tag
                res.append(cur_res)
    return res


def testTinygemmBf16(args):
    """
    Test tinygemm_bf16 API.

    This test:
    1. Generates random BF16 input tensors
    2. Runs tinygemm_bf16
    3. Runs reference check (F.linear)
    4. Measures performance metrics (TFLOPS, TB/sec)

    Args:
        args: Parsed command line arguments containing test configuration

    Returns:
        dict: List of dictionaries containing performance results
    """
    if args.verbose >= 1:
        print("[INFO] Running testTinygemmBf16")
        print(f"[INFO] FlashInfer version: {flashinfer.__version__}")

    device = get_device(args)
    if args.generate_repro_command:
        print(
            f"[INFO] To reproduce this test case, run the following command: {args.repro_command}"
        )

    backends = list(args.backends)
    m = args.m
    n = args.n
    k = args.k
    use_bias = getattr(args, "bias", False)
    use_pdl = getattr(args, "enable_pdl", False)
    is_cuda_graph_compatible = not args.no_cuda_graph
    run_refcheck = args.refcheck
    res = []

    input_dtype = dtype_str_to_torch_dtype(args.input_dtype)
    mat2_dtype = dtype_str_to_torch_dtype(args.mat2_dtype)
    out_dtype = dtype_str_to_torch_dtype(args.out_dtype)
    if input_dtype != torch.bfloat16 or mat2_dtype != torch.bfloat16:
        raise ValueError(
            "tinygemm_bf16 only supports bfloat16 input and weight tensors."
        )
    if out_dtype != torch.bfloat16:
        raise ValueError("tinygemm_bf16 only supports bfloat16 outputs.")

    backends = filter_backends_by_compute_capability(backends, args.routine, device)
    if len(backends) == 0:
        print("[ERROR] No backends to test. Exiting.")
        return res

    from flashinfer.gemm import tinygemm_bf16

    input = torch.randn([m, k], device=device, dtype=torch.bfloat16)
    weight = torch.randn([n, k], device=device, dtype=torch.bfloat16)
    bias = torch.randn([n], device=device, dtype=torch.bfloat16) if use_bias else None
    outs = {
        backend: torch.empty([m, n], device=device, dtype=torch.bfloat16)
        for backend in backends
    }

    if args.verbose >= 2:
        print(f"[VVERBOSE] {input.shape = }")
        print(f"[VVERBOSE] {input.dtype = }")
        print(f"[VVERBOSE] {weight.shape = }")
        print(f"[VVERBOSE] {weight.dtype = }")
        if bias is not None:
            print(f"[VVERBOSE] {bias.shape = }")
            print(f"[VVERBOSE] {bias.dtype = }")
        print(f"[VVERBOSE] {use_pdl = }")

    backends_to_remove = []
    for backend in backends:
        if backend != "tinygemm":
            backends_to_remove.append(backend)
            continue
        try:
            tinygemm_bf16(input, weight, outs[backend], bias=bias, use_pdl=use_pdl)
        except Exception as e:
            print(
                f"[INFO] {backend} backend does not support this configuration: {type(e).__name__}: {e}"
            )
            backends_to_remove.append(backend)

    for backend in backends_to_remove:
        backends.remove(backend)
        outs.pop(backend, None)

    if len(backends) == 0:
        print("[ERROR] No backends passed validation. Exiting.")
        return res

    def run_backend(backend, input, weight, bias, out, use_pdl):
        if backend != "tinygemm":
            raise ValueError(f"Unsupported backend: {backend}")
        tinygemm_bf16(input, weight, out, bias=bias, use_pdl=use_pdl)
        return out

    has_reference_output = False
    if run_refcheck:
        reference_output = F.linear(
            input.float(),
            weight.float(),
            bias.float() if bias is not None else None,
        ).bfloat16()
        has_reference_output = True

    backend_times = {backend: [] for backend in backends}
    outputs = {}
    for cur_backend in backends:
        if run_refcheck:
            outputs[cur_backend] = (
                run_backend(
                    cur_backend, input, weight, bias, outs[cur_backend], use_pdl
                )
                .detach()
                .clone()
            )
        backend_times[cur_backend] = bench_gpu_time(
            fn=run_backend,
            dry_run_iters=args.dry_run_iters,
            repeat_iters=args.num_iters,
            sleep_after_run=True,
            enable_cupti=args.use_cupti,
            use_cuda_graph=is_cuda_graph_compatible,
            cold_l2_cache=True,
            input_args=(cur_backend, input, weight, bias, outs[cur_backend], use_pdl),
        )

    tested_backends = list(outputs.keys())
    tested_outputs = list(outputs.values())
    if len(tested_backends) > 0 and run_refcheck and has_reference_output:
        for i in range(len(tested_backends)):
            cos_sim = F.cosine_similarity(
                reference_output.reshape(-1).float(),
                tested_outputs[i].reshape(-1).float(),
                dim=0,
            )
            if cos_sim < 0.99:
                print(
                    f"[ERROR] Output tensor mismatch from backend {tested_backends[i]} with cos_sim={cos_sim}"
                )
                if not args.allow_output_mismatch:
                    raise AssertionError(
                        f"[ERROR] Backend {tested_backends[i]} output mismatch with cos_sim={cos_sim}"
                    )

    for backend in backends:
        if len(backend_times[backend]) > 0:
            median_time = np.median(backend_times[backend])
            std_time = np.std(backend_times[backend])
            problem_flops = 2 * m * n * k
            if use_bias:
                problem_flops += m * n
            problem_bytes = (
                m * k * torch.bfloat16.itemsize
                + n * k * torch.bfloat16.itemsize
                + m * n * torch.bfloat16.itemsize
            )
            if use_bias:
                problem_bytes += n * torch.bfloat16.itemsize
            tflops = problem_flops / (10**9 * median_time)
            tb_per_sec = problem_bytes / (10**9 * median_time)
            print_perf_metrics(backend, median_time, std_time, tflops, tb_per_sec)

            if args.output_path is not None:
                cur_res = defaultdict(str)
                cur_res["routine"] = args.routine
                cur_res["median_time"] = median_time
                cur_res["std_time"] = std_time
                cur_res["tflops"] = tflops
                cur_res["tb_per_sec"] = tb_per_sec
                cur_res["m"] = m
                cur_res["n"] = n
                cur_res["k"] = k
                cur_res["input_dtype"] = input_dtype
                cur_res["mat2_dtype"] = mat2_dtype
                cur_res["out_dtype"] = out_dtype
                cur_res["backend"] = backend
                cur_res["bias"] = use_bias
                cur_res["enable_pdl"] = use_pdl
                cur_res["case_tag"] = args.case_tag
                res.append(cur_res)
    return res


def testBmmBf16(args):
    """
    Test bmm_bf16 API.

    This test:
    1. Generates random BF16 batched input tensors
    2. Runs bmm_bf16 with specified backend
    3. Runs reference check (torch.bmm)
    4. Measures performance metrics (TFLOPS, TB/sec)

    Args:
        args: Parsed command line arguments containing test configuration

    Returns:
        dict: List of dictionaries containing performance results
    """
    warn_if_pdl_unsupported(args, args.routine)
    if args.verbose >= 1:
        print("[INFO] Running testBmmBf16")
        print(f"[INFO] FlashInfer version: {flashinfer.__version__}")

    device = get_device(args)
    if args.generate_repro_command:
        print(
            f"[INFO] To reproduce this test case, run the following command: {args.repro_command}"
        )

    ## Parse input arguments
    backends = args.backends
    batch_size = args.batch_size
    m = args.m
    n = args.n
    k = args.k
    is_cuda_graph_compatible = not args.no_cuda_graph
    run_refcheck = args.refcheck
    autotune_supported_backends = ["cudnn", "cutlass", "auto"]
    res = []

    out_dtype = dtype_str_to_torch_dtype(args.out_dtype)
    if out_dtype not in [torch.bfloat16, torch.float16, torch.float32]:
        raise ValueError(
            f"Unsupported output dtype: {args.out_dtype}. Supported dtypes are bfloat16, float16, and float32."
        )

    ## Prepare input tensors
    # A: (batch_size, m, k), row-major
    A = torch.randn([batch_size, m, k], device=device, dtype=torch.bfloat16)
    # B: (batch_size, n, k) then transpose to (batch_size, k, n) for column-major layout
    B = torch.randn([batch_size, n, k], device=device, dtype=torch.bfloat16).transpose(
        -2, -1
    )

    if args.verbose >= 2:
        print(f"[VVERBOSE] {A.shape = }")
        print(f"[VVERBOSE] {A.dtype = }")
        print(f"[VVERBOSE] {B.shape = }")
        print(f"[VVERBOSE] {B.dtype = }")

    # Programmatically filter backends
    backends_to_remove = []
    for backend in backends:
        # Skip autotune check for now (handled separately below)
        if (
            getattr(args, "autotune", False)
            and backend not in autotune_supported_backends
        ):
            print(f"[INFO] {backend} backend does not support autotune")
            backends_to_remove.append(backend)
            continue

        try:
            flashinfer.bmm_bf16(
                A=A,
                B=B,
                out_dtype=out_dtype,
                backend=backend,
            )
        except Exception as e:
            print(
                f"[INFO] {backend} backend does not support this configuration: {type(e).__name__}: {e}"
            )
            backends_to_remove.append(backend)

    # Remove unsupported backends
    for backend in backends_to_remove:
        backends.remove(backend)

    if len(backends) == 0:
        print("[ERROR] No backends passed validation. Exiting.")
        return res

    def run_backend(backend, A, B, out_dtype):
        if backend in ["cudnn", "cutlass", "auto"]:
            return flashinfer.bmm_bf16(
                A=A,
                B=B,
                out_dtype=out_dtype,
                backend=backend,
            )
        else:
            raise ValueError(f"Unsupported backend: {backend}")

    has_reference_output = False
    if run_refcheck:
        reference_output = torch.bmm(A, B).to(out_dtype)
        has_reference_output = True

    if getattr(args, "autotune", False):
        warmup_iters = (
            args.dry_run_iters if args.dry_run_iters and args.dry_run_iters > 0 else 10
        )
        for cur_backend in backends:
            if cur_backend in autotune_supported_backends:
                if args.verbose >= 1:
                    print(f"[INFO] Autotune warmup for bmm_bf16: {warmup_iters} iters")
                with autotune(True):
                    for _ in range(warmup_iters):
                        run_backend(cur_backend, A, B, out_dtype)

    # Storage for timing results and outputs
    backend_times = {backend: [] for backend in backends}
    outputs = {}
    for cur_backend in backends:
        if run_refcheck:
            outputs[cur_backend] = run_backend(cur_backend, A, B, out_dtype).detach()
        backend_times[cur_backend] = bench_gpu_time(
            fn=run_backend,
            dry_run_iters=args.dry_run_iters,
            repeat_iters=args.num_iters,
            sleep_after_run=True,
            enable_cupti=args.use_cupti,
            use_cuda_graph=is_cuda_graph_compatible,
            cold_l2_cache=True,
            input_args=(cur_backend, A, B, out_dtype),
        )

    tested_backends = list(outputs.keys())
    tested_outputs = list(outputs.values())
    if len(tested_backends) > 0:
        if run_refcheck and has_reference_output:
            for i in range(len(tested_backends)):
                cos_sim = F.cosine_similarity(
                    reference_output.reshape(-1),
                    tested_outputs[i].reshape(-1),
                    dim=0,
                )
                if cos_sim < 0.99:
                    print(
                        f"[ERROR] Output tensor mismatch from backend {tested_backends[i]} with cos_sim={cos_sim}"
                    )
                    if not args.allow_output_mismatch:
                        raise AssertionError(
                            f"[ERROR] Backend {tested_backends[i]} output mismatch with cos_sim={cos_sim}"
                        )

    for backend in backends:
        backend_name = backend + (
            "_autotune"
            if (
                getattr(args, "autotune", False)
                and backend in autotune_supported_backends
            )
            else ""
        )
        if len(backend_times[backend]) > 0:
            median_time = np.median(backend_times[backend])
            std_time = np.std(backend_times[backend])
            problem_flops = 2 * batch_size * m * n * k
            problem_bytes = batch_size * (
                m * k * torch.bfloat16.itemsize
                + k * n * torch.bfloat16.itemsize
                + m * n * out_dtype.itemsize
            )
            tflops = problem_flops / (10**9 * median_time)  # in TFLOPs/sec
            tb_per_sec = problem_bytes / (10**9 * median_time)  # in TB/sec
            print_perf_metrics(backend_name, median_time, std_time, tflops, tb_per_sec)

            if args.output_path is not None:
                cur_res = defaultdict(str)
                cur_res["batch_size"] = batch_size
                cur_res["routine"] = args.routine
                cur_res["median_time"] = median_time
                cur_res["std_time"] = std_time
                cur_res["tflops"] = tflops
                cur_res["tb_per_sec"] = tb_per_sec
                cur_res["m"] = m
                cur_res["n"] = n
                cur_res["k"] = k
                cur_res["out_dtype"] = str(out_dtype)
                cur_res["backend"] = backend_name
                cur_res["case_tag"] = args.case_tag
                res.append(cur_res)
    return res
