flash_attn_mma.py

import argparse
import math
import os
import random
import time
from functools import partial
from typing import Optional

import numpy as np
import torch
from flash_attn import flash_attn_func
from torch import Tensor
from torch.nn import functional as F
from torch.nn.attention import SDPBackend, sdpa_kernel
from torch.utils.cpp_extension import load

torch.set_grad_enabled(False)
torch.set_printoptions(
    precision=6, threshold=8, edgeitems=3, linewidth=120, sci_mode=False
)


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--no-rand-q", "--no-rq", action="store_true")
    parser.add_argument("--no-rand-k", "--no-rk", action="store_true")
    parser.add_argument("--no-rand-v", "--no-rv", action="store_true")
    parser.add_argument("--no-rand-qkv", "--no-rqkv", action="store_true")
    parser.add_argument("--run-torch-unfused", "--torch", action="store_true")
    parser.add_argument("--run-torch-sdpa", "--sdpa", action="store_true")
    parser.add_argument("--check", action="store_true")
    parser.add_argument("--check-all", action="store_true")
    parser.add_argument("--show-all", "--show", action="store_true")
    parser.add_argument("--show-matrix", action="store_true")
    parser.add_argument(
        "--only-flops-matmul", "--flops-mm", action="store_true"
    )
    parser.add_argument(
        "--run-acc-f32", "--acc-f32", "--f32", action="store_true"
    )
    parser.add_argument("--B", type=int, default=None)
    parser.add_argument("--H", type=int, default=None)
    parser.add_argument("--N", type=int, default=None)
    parser.add_argument("--D", type=int, default=None)
    parser.add_argument("--seed", type=int, default=None)
    parser.add_argument("--sleep", type=float, default=0.05)
    parser.add_argument("--debug", action="store_true")
    parser.add_argument("--verbose", "--v", action="store_true")
    parser.add_argument("--warmup", "--w", type=int, default=1)
    parser.add_argument("--iters", "--i", type=int, default=5)
    parser.add_argument("--range-k", "--gk", action="store_true")
    parser.add_argument("--build-others", "--others", action="store_true")
    parser.add_argument(
        "--tag-hints", "--tags", "--hints", type=str, default=None
    )
    return parser.parse_args()


args = get_args()


def set_rand_seed(seed: int = 1):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


def get_device_name():
    device_name = torch.cuda.get_device_name(torch.cuda.current_device())
    # since we will run GPU on WSL2, so add WSL2 tag.
    if "Laptop" in device_name:
        device_name += " WSL2"
    return device_name


def get_device_capability():
    return torch.cuda.get_device_capability(torch.cuda.current_device())


def get_build_sources():
    build_sources = []
    # Basic
    build_sources.append("./mma/basic/flash_attn_mma_split_kv.cu")
    build_sources.append("./mma/basic/flash_attn_mma_split_q.cu")
    build_sources.append("./mma/basic/flash_attn_mma_share_kv.cu")
    build_sources.append("./mma/basic/flash_attn_mma_share_qkv.cu")
    build_sources.append("./mma/basic/flash_attn_mma_tiling_qk.cu")
    build_sources.append("./mma/basic/flash_attn_mma_tiling_qkv.cu")
    build_sources.append("./mma/basic/flash_attn_mma_share_kv_F32F16F16F32.cu")
    build_sources.append("./mma/basic/flash_attn_mma_share_qkv_F32F16F16F32.cu")
    build_sources.append("./mma/basic/flash_attn_mma_tiling_qk_F32F16F16F32.cu")
    build_sources.append(
        "./mma/basic/flash_attn_mma_tiling_qkv_F32F16F16F32.cu"
    )
    # Swizzle
    build_sources.append("./mma/swizzle/flash_attn_mma_share_kv_swizzle_q.cu")
    build_sources.append("./mma/swizzle/flash_attn_mma_share_kv_swizzle_qk.cu")
    build_sources.append("./mma/swizzle/flash_attn_mma_share_kv_swizzle_qkv.cu")
    build_sources.append("./mma/swizzle/flash_attn_mma_share_qkv_swizzle_q.cu")
    build_sources.append("./mma/swizzle/flash_attn_mma_share_qkv_swizzle_qk.cu")
    build_sources.append(
        "./mma/swizzle/flash_attn_mma_share_qkv_swizzle_qkv.cu"
    )
    build_sources.append("./mma/swizzle/flash_attn_mma_tiling_qk_swizzle_q.cu")
    build_sources.append("./mma/swizzle/flash_attn_mma_tiling_qk_swizzle_qk.cu")
    build_sources.append(
        "./mma/swizzle/flash_attn_mma_tiling_qk_swizzle_qkv.cu"
    )
    build_sources.append("./mma/swizzle/flash_attn_mma_tiling_qkv_swizzle_q.cu")
    build_sources.append(
        "./mma/swizzle/flash_attn_mma_tiling_qkv_swizzle_qk.cu"
    )
    build_sources.append(
        "./mma/swizzle/flash_attn_mma_tiling_qkv_swizzle_qkv.cu"
    )
    build_sources.append(
        "./mma/swizzle/flash_attn_mma_tiling_qkv_swizzle_q_F32F16F16F32.cu"
    )
    build_sources.append(
        "./mma/swizzle/flash_attn_mma_tiling_qkv_swizzle_qk_F32F16F16F32.cu"
    )
    build_sources.append(
        "./mma/swizzle/flash_attn_mma_tiling_qkv_swizzle_qkv_F32F16F16F32.cu"
    )
    # Others
    if args.build_others:
        build_sources.append("./mma/others/flash_attn_mma_share_qkv_Os2g.cu")
        build_sources.append(
            "./mma/others/flash_attn_mma_share_kv_F32F16F16F32_rr.cu"
        )
        build_sources.append(
            "./mma/others/flash_attn_mma_share_qkv_F32F16F16F32_rr.cu"
        )
    # Pybind
    build_sources.append("./pybind/flash_attn.cc")
    return build_sources


def get_project_dir():
    return os.path.dirname(
        os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    )


project_dir = get_project_dir()


def get_build_cuda_cflags(build_pkg: bool = False):
    device_name = get_device_name()
    project_dir = get_project_dir()
    extra_cuda_cflags = []
    extra_cuda_cflags.append("-O3")
    extra_cuda_cflags.append("-std=c++17")
    extra_cuda_cflags.append("-U__CUDA_NO_HALF_OPERATORS__")
    extra_cuda_cflags.append("-U__CUDA_NO_HALF_CONVERSIONS__")
    extra_cuda_cflags.append("-U__CUDA_NO_HALF2_OPERATORS__")
    extra_cuda_cflags.append("-U__CUDA_NO_BFLOAT16_CONVERSIONS__")
    extra_cuda_cflags.append("--expt-relaxed-constexpr")
    extra_cuda_cflags.append("--expt-extended-lambda")
    extra_cuda_cflags.append("--use_fast_math")
    extra_cuda_cflags.append("-DFLASH_ATTN_MMA_DEBUG" if args.debug else "")
    extra_cuda_cflags.append(
        "-DBUILD_FLASH_ATTN_MMA_OTHERS" if args.build_others else ""
    )
    extra_cuda_cflags.append(
        "-DBUILD_FLASH_ATTN_MMA_L20" if "L20" in device_name else ""
    )
    extra_cuda_cflags.append(
        "-DBUILD_FLASH_ATTN_MMA_4090" if "4090" in device_name else ""
    )
    extra_cuda_cflags.append(
        "-DBUILD_FLASH_ATTN_MMA_3080" if "3080" in device_name else ""
    )
    extra_cuda_cflags.append(
        "-diag-suppress 177" if not build_pkg else "--ptxas-options=-v"
    )
    extra_cuda_cflags.append(
        "-Xptxas -v" if not build_pkg else "--ptxas-options=-O3"
    )
    extra_cuda_cflags.append(f"-I {project_dir}/kernels/flash-attn")
    extra_cuda_cflags.append(f"-I {project_dir}/kernels/flash-attn/utils")
    extra_cuda_cflags.append(f"-I {project_dir}/kernels/flash-attn/mma")
    extra_cuda_cflags.append(f"-I {project_dir}/kernels/flash-attn/mma/basic")
    extra_cuda_cflags.append(f"-I {project_dir}/kernels/flash-attn/mma/swizzle")
    extra_cuda_cflags.append(f"-I {project_dir}/kernels/flash-attn/mma/others")
    extra_cuda_cflags.append(f"-I {project_dir}/kernels/flash-attn/cutlass")
    extra_cuda_cflags.append(f"-I {project_dir}/kernels/flash-attn/pybind")
    extra_cuda_cflags.append(f"-I {project_dir}/third-party/cutlass/include")
    extra_cuda_cflags.append(
        f"-I {project_dir}/third-party/cutlass/tools/util/include"
    )
    return extra_cuda_cflags


def get_build_cflags():
    extra_cflags = []
    extra_cflags.append("-std=c++17")
    extra_cflags.append(
        "-DBUILD_FLASH_ATTN_MMA_OTHERS" if args.build_others else ""
    )
    return extra_cflags


def pretty_print_line(m: str = "", sep: str = "-", width: int = 150):
    res_len = width - len(m)
    left_len = int(res_len / 2)
    right_len = res_len - left_len
    pretty_line = sep * left_len + m + sep * right_len
    print(pretty_line)


if args.D and args.D > 256:
    args.run_torch_sdpa = True
pretty_print_line()
print(args)
pretty_print_line()

# Load the CUDA kernel as a python module
lib = load(
    name="flash_attn_lib",
    sources=get_build_sources(),
    extra_cuda_cflags=get_build_cuda_cflags(),
    extra_cflags=get_build_cflags(),
    verbose=args.verbose,
)

if not args.build_others:
    fake_fa_func = lambda q, k, v, o, s: o  # fake FA func
    setattr(lib, "flash_attn_mma_stages_split_q_shared_qkv_Os2g", fake_fa_func)
    setattr(
        lib, "flash_attn_mma_stages_split_q_shared_kv_acc_f32_rr", fake_fa_func
    )
    setattr(
        lib, "flash_attn_mma_stages_split_q_shared_qkv_acc_f32_rr", fake_fa_func
    )


def get_mha_tflops(
    B: int, H: int, N: int, D: int, secs: float = 1.0, only_matmul: bool = False
):
    # Q @ K^T FLOPs
    flops_qk = B * H * N * N * (2 * D - 1)

    # Scaling FLOPs
    flops_scaling = B * H * N * N

    # Safe_Softmax FLOPs
    flops_row_max = B * H * N * (N - 1)  # row max
    flops_subtract_max = B * H * N * N  # sub max
    flops_exp = B * H * N * N  # pointwise exp
    flops_row_sum = B * H * N * (N - 1)  # row sum
    flops_normalization = B * H * N * N  # normalization

    flops_safe_softmax = (
        flops_row_max
        + flops_subtract_max
        + flops_exp
        + flops_row_sum
        + flops_normalization
    )

    # P @ V FLOPs
    flops_pv = B * H * N * D * (2 * N - 1)

    # Total FLOPs
    total_flops = flops_qk + flops_scaling + flops_safe_softmax + flops_pv
    if only_matmul:
        total_flops = flops_qk + flops_pv

    # Convert to TFLOPS
    # 1 TFLOPS = 10^12 FLOPS
    # ref: https://imgtec.eetrend.com/blog/2021/100062210.html.
    tflops = total_flops * 1e-12 / (secs)

    return tflops


MAX_TFLOPS = -1
STATIS_INFO: dict[str, list[float]] = {}
TOATL_TFLOPS: dict[str, float] = {}


def run_benchmark(
    perf_func: callable,
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    tag: str,
    out: Optional[torch.Tensor] = None,
    s: Optional[torch.Tensor] = None,  # DEBUG
    stages: int = -1,
    warmup: int = args.warmup,
    iters: int = args.iters,
    show_matrix: bool = args.show_matrix,
    only_show_improved: bool = not args.show_all,
):

    global MAX_TFLOPS
    global MAX_HEADDIM_CFG

    tag_hints: str = args.tag_hints  # e.g "share-qkv,tiling-kv,swizzle"
    if tag_hints:
        tag_hints: list = tag_hints.strip().split(",")
        tag_hints.append("flash")
        tag_hints.append("sdpa")
        tag_hints.append("unfused")
        hit_hints = False
        for hint in tag_hints:
            if hint in tag:
                hit_hints = True
        if not hit_hints:
            return None, None

    if not args.build_others:
        others_tags = ["s2g", "rr"]
        for o_tag in others_tags:
            if o_tag in tag:
                return None, None

    if "sdpa" in tag and (not args.run_torch_sdpa):
        return None, None
    if "unfused" in tag and (not args.run_torch_unfused):
        return None, None
    if "acc-f32" in tag and (not args.run_acc_f32):
        return None, None

    B, H, N, D = q.size()
    if "flash" in tag:
        B, N, H, D = q.size()

    max_supported_D = MAX_HEADDIM_CFG.get(tag, None)
    # skip if headdim not supported.
    if max_supported_D is not None:
        if D > max_supported_D:
            return None, None

    if out is not None:
        out.fill_(0)
    if s is not None:
        s.fill_(0)
    if out is not None:
        for i in range(warmup):
            if stages >= 1:
                if s is not None:
                    perf_func(q, k, v, out, s, stages)
                else:
                    perf_func(q, k, v, out, stages)
            else:
                perf_func(q, k, v, out)
    else:
        for i in range(warmup):
            _ = perf_func(q, k, v)

    torch.cuda.synchronize()
    start = time.time()
    # iters
    if out is not None:
        for i in range(iters):
            if stages >= 1:
                if s is not None:
                    perf_func(q, k, v, out, s, stages)
                else:
                    perf_func(q, k, v, out, stages)
            else:
                perf_func(q, k, v, out)
    else:
        for i in range(iters):
            out = perf_func(q, k, v)
    torch.cuda.synchronize()
    end = time.time()
    total_secs = end - start
    total_time = (end - start) * 1000  # ms
    mean_time = total_time / iters
    mean_secs = total_secs / iters

    TFLOPS = get_mha_tflops(
        B, H, N, D, mean_secs, only_matmul=args.only_flops_matmul
    )
    out_info = f"{tag}"
    out_val_first = out.flatten()[:3].detach().cpu().numpy().tolist()
    out_val_last = out.flatten()[-3:].detach().cpu().numpy().tolist()
    out_val_first = [round(v, 8) for v in out_val_first]
    out_val_last = [round(v, 8) for v in out_val_last]
    out_val = out_val_first[:2]
    out_val.append(out_val_last[-1])
    out_val = [f"{v:<12}" for v in out_val]

    # caculate TFLOPS improved.
    if TFLOPS > MAX_TFLOPS:
        if MAX_TFLOPS > 0:
            improve = ((TFLOPS - MAX_TFLOPS) / MAX_TFLOPS) * 100
            improve = round(improve, 2)
        else:
            improve = 0
        MAX_TFLOPS = TFLOPS
        print(
            f"{out_info:>50}: {out_val}, time:{str(mean_time)[:8]}ms, "
            f"TFLOPS:{TFLOPS:<6.2f}(+{improve:.2f}%)"
        )
    else:
        if (not only_show_improved) or (("flash" in tag) or ("sdpa" in tag)):
            print(
                f"{out_info:>50}: {out_val}, time:{str(mean_time)[:8]}ms, "
                f"TFLOPS:{TFLOPS:<6.2f}"
            )

    if show_matrix:
        print(out)
    time.sleep(args.sleep)
    torch.cuda.synchronize()
    return out.clone(), mean_time


def get_qkvo(B, H, N, D):
    if not (args.no_rand_q or args.no_rand_qkv):
        q = torch.randn((B, H, N, D), dtype=torch.half, device="cuda")
    else:
        q = torch.ones(B, H, N, D, device="cuda", dtype=torch.half).contiguous()
    if not (args.no_rand_k or args.no_rand_qkv):
        k = torch.randn((B, H, N, D), dtype=torch.half, device="cuda")
    else:
        k = torch.ones(B, H, N, D, device="cuda", dtype=torch.half).contiguous()
        if args.range_k:
            for i in range(N):
                k[:, :, i, :] = (i + 1) / N
            k = k.cuda().half().contiguous()
    if not (args.no_rand_v or args.no_rand_qkv):
        v = torch.randn((B, H, N, D), dtype=torch.half, device="cuda")
    else:
        v = torch.ones(B, H, N, D, device="cuda", dtype=torch.half).contiguous()

    o = torch.zeros(B, H, N, D, device="cuda", dtype=torch.half).contiguous()
    # transpose (H,N) -> (N,H) for FA2.
    fq = q.transpose(1, 2).contiguous()
    fk = k.transpose(1, 2).contiguous()
    fv = v.transpose(1, 2).contiguous()
    # transpose (N,D) -> (D,N) for V smem swizzle.
    tk = k.transpose(-2, -1).contiguous()  # [B,H,N,D] -> [B,H,D,N]
    tv = v.transpose(-2, -1).contiguous()  # [B,H,N,D] -> [B,H,D,N]

    return q, k, v, o, fq, fk, fv, tk, tv


# un-fused naive attn
def unfused_standard_attn(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor):
    att = q @ k.transpose(-2, -1) * (1.0 / math.sqrt(k.size(-1)))
    att = F.softmax(att, dim=-1)
    y = att @ v
    return y


def sdpa(q: Tensor, k: Tensor, v: Tensor, use_flash: bool = False):
    if not use_flash:
        with sdpa_kernel(SDPBackend.EFFICIENT_ATTENTION):
            out: Tensor = F.scaled_dot_product_attention(q, k, v)
    else:
        with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
            out: Tensor = F.scaled_dot_product_attention(q, k, v)
    return out


def check_all_close(
    out_flash_or_sdpa: torch.Tensor,
    out_mma: torch.Tensor,
    tag: str = "out_mma",
    check_all: bool = False,
    is_flash: bool = True,
):
    if any((out_flash_or_sdpa is None, out_mma is None)):
        return
    if is_flash:
        true_tag = "out_flash"
        out_flash_or_sdpa = out_flash_or_sdpa.transpose(1, 2)
    else:
        true_tag = "out_sdpa"
    if check_all:
        for i in range(int(N / 8)):
            if i < 4:
                pretty_print_line()
                print(f"{true_tag}[:, :,  {(i*8)}:{(i+1)*8}, :]:\n")
                print(out_flash_or_sdpa[:, :, (i * 8) : (i + 1) * 8, :].float())
                print(f"{tag}[:, :, {(i*8)}:{(i+1)*8}, :]:\n")
                print(out_mma[:, :, (i * 8) : (i + 1) * 8, :].float())
        pretty_print_line()
    diff = torch.abs(out_flash_or_sdpa - out_mma)
    all_close = str(torch.allclose(out_flash_or_sdpa, out_mma, atol=1e-2))
    pretty_print_line(
        f"{true_tag} vs {tag:<25}, all close: {all_close:<6}, "
        f"max diff: {diff.max().item():.6f}, min diff: {diff.min().item():.6f}, "
        f"mean diff: {diff.mean().item():.6f}"
    )


Bs = [1, 4, 8] if not args.B else [args.B]
Hs = [1, 4, 8] if not args.H else [args.H]
Ns = [1024, 2048, 4096, 8192] if not args.N else [args.N]
Ds = [64, 128, 256, 512] if not args.D else [args.D]
# batch_size, n_head, seq_len, head_dim (B,H,N,D)
BHNDs = [(B, H, N, D) for B in Bs for H in Hs for N in Ns for D in Ds]
# max headdim supported for different methods. skip if D > max_D.
MAX_HEADDIM_CFG: dict[str, int] = {
    # FA2, SDPA, Naive MHA.
    "(flash)": 256,
    "(sdpa)": 4096,  # may no limit
    "(unfused)": 4096,  # may no limit
    # Split-KV
    "mma(split-kv+stage1)": 128,
    "mma(split-kv+stage2)": 128,
    # Split-Q
    "mma(split-q+stage1)": 128,
    "mma(split-q+stage2)": 128,
    # Split-Q + Shared KV SMEM
    "mma(split-q+share-kv+stage1)": 256,
    "mma(split-q+share-kv+stage2)": 128,
    "mma(split-q+share-kv+swizzle-q+stage1)": 256,
    "mma(split-q+share-kv+swizzle-q+stage2)": 128,
    "mma(split-q+share-kv+swizzle-qk+stage1)": 256,
    "mma(split-q+share-kv+swizzle-qk+stage2)": 128,
    "mma(split-q+share-kv+swizzle-qkv+stage1)": 256,
    "mma(split-q+share-kv+swizzle-qkv+stage2)": 128,
    "mma(split-q+share-kv+acc-f32+stage1)": 256,
    "mma(split-q+share-kv+acc-f32+stage2)": 128,
    # Split-Q + Fully Shared QKV SMEM
    "mma(split-q+share-qkv+stage1)": 256,
    "mma(split-q+share-qkv+stage2)": 128,
    "mma(split-q+share-qkv+swizzle-q+stage1)": 256,
    "mma(split-q+share-qkv+swizzle-q+stage2)": 128,
    "mma(split-q+share-qkv+swizzle-qk+stage1)": 256,
    "mma(split-q+share-qkv+swizzle-qk+stage2)": 128,
    "mma(split-q+share-qkv+swizzle-qkv+stage1)": 256,
    "mma(split-q+share-qkv+swizzle-qkv+stage2)": 128,
    "mma(split-q+share-qkv+acc-f32+stage1)": 256,
    "mma(split-q+share-qkv+acc-f32+stage2)": 128,
    # Split-Q + QK Fine-grained Tiling
    "mma(split-q+tiling-qk+stage1)": 1024,
    "mma(split-q+tiling-qk+stage2)": 1024,
    "mma(split-q+tiling-qk+swizzle-q+stage1)": 1024,
    "mma(split-q+tiling-qk+swizzle-q+stage2)": 1024,
    "mma(split-q+tiling-qk+swizzle-qk+stage1)": 1024,
    "mma(split-q+tiling-qk+swizzle-qk+stage2)": 1024,
    "mma(split-q+tiling-qk+swizzle-qkv+stage1)": 256,
    "mma(split-q+tiling-qk+swizzle-qkv+stage2)": 256,
    "mma(split-q+tiling-qk+acc-f32+stage1)": 1024,
    "mma(split-q+tiling-qk+acc-f32+stage2)": 1024,
    # Split-Q + Fully QKV Fine-grained Tiling
    "mma(split-q+tiling-qkv+stage1)": 1024,
    "mma(split-q+tiling-qkv+stage2)": 1024,
    "mma(split-q+tiling-qkv+acc-f32+stage1)": 1024,
    "mma(split-q+tiling-qkv+acc-f32+stage2)": 1024,
    "mma(split-q+tiling-qkv+swizzle-q+stage1)": 1024,
    "mma(split-q+tiling-qkv+swizzle-q+stage2)": 1024,
    "mma(split-q+tiling-qkv+swizzle-qk+stage1)": 1024,
    "mma(split-q+tiling-qkv+swizzle-qk+stage2)": 1024,
    "mma(split-q+tiling-qkv+swizzle-qkv+stage1)": 1024,
    "mma(split-q+tiling-qkv+swizzle-qkv+stage2)": 1024,
    "mma(split-q+tiling-qkv+acc-f32+swizzle-q+stage1)": 1024,
    "mma(split-q+tiling-qkv+acc-f32+swizzle-q+stage2)": 1024,
    "mma(split-q+tiling-qkv+acc-f32+swizzle-qk+stage1)": 1024,
    "mma(split-q+tiling-qkv+acc-f32+swizzle-qk+stage2)": 1024,
    "mma(split-q+tiling-qkv+acc-f32+swizzle-qkv+stage1)": 1024,
    "mma(split-q+tiling-qkv+acc-f32+swizzle-qkv+stage2)": 1024,
    # Others, O s2g, etc.
    "mma(split-q+share-qkv+o-s2g+stage1)": 256,
    "mma(split-q+share-qkv+o-s2g+stage2)": 128,
    "mma(split-q+share-kv+acc-f32+rr+stage1)": 256,
    "mma(split-q+share-kv+acc-f32+rr+stage2)": 128,
    "mma(split-q+share-qkv+acc-f32+rr+stage1)": 256,
    "mma(split-q+share-qkv+acc-f32+rr+stage2)": 256,
    "mma(split-q+tiling-qk+acc-f32+rr+stage1)": 1024,
    "mma(split-q+tiling-qk+acc-f32+rr+stage2)": 1024,
}

seed = args.seed if args.seed else random.choice(range(10000))
set_rand_seed(seed)
pretty_print_line()
pretty_print_line(
    f"B: batch_size, H: n_head, N: seq_len, D: head_dim, "
    f"seed: {seed}, Warmup: {args.warmup}, Iters: {args.iters}"
)

run_torch_sdpa = args.run_torch_sdpa
for B, H, N, D in BHNDs:
    MAX_TFLOPS = -1
    q, k, v, o, fq, fk, fv, tk, tv = get_qkvo(B, H, N, D)
    if D > 256:
        args.run_torch_sdpa = True
    else:
        args.run_torch_sdpa = run_torch_sdpa
    torch.cuda.synchronize()
    pretty_print_line()
    pretty_print_line(
        f"B={B}, H={H}, N={N}, D={D}, Warmup: {args.warmup}, Iters: {args.iters}"
    )
    # Naive MHA.
    out_unfused, _ = run_benchmark(unfused_standard_attn, q, k, v, "(unfused)")
    # Split-KV
    out_mma_split_kv1, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_kv,
        q,
        k,
        v,
        "mma(split-kv+stage1)",
        o,
        stages=1,
    )
    out_mma_split_kv2, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_kv,
        q,
        k,
        v,
        "mma(split-kv+stage2)",
        o,
        stages=2,
    )
    # Split-Q
    out_mma_split_q1, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q,
        q,
        k,
        v,
        "mma(split-q+stage1)",
        o,
        stages=1,
    )
    out_mma_split_q2, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q,
        q,
        k,
        v,
        "mma(split-q+stage2)",
        o,
        stages=2,
    )
    # Split-Q + Shared KV SMEM + Swizzle
    out_mma_share_kv1, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_shared_kv,
        q,
        k,
        v,
        "mma(split-q+share-kv+stage1)",
        o,
        stages=1,
    )
    out_mma_share_kv2, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_shared_kv,
        q,
        k,
        v,
        "mma(split-q+share-kv+stage2)",
        o,
        stages=2,
    )
    out_mma_share_kv_f321, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_shared_kv_acc_f32,
        q,
        k,
        v,
        "mma(split-q+share-kv+acc-f32+stage1)",
        o,
        stages=1,
    )
    out_mma_share_kv_f322, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_shared_kv_acc_f32,
        q,
        k,
        v,
        "mma(split-q+share-kv+acc-f32+stage2)",
        o,
        stages=2,
    )
    out_mma_share_kv_sq1, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_shared_kv_swizzle_q,
        q,
        k,
        v,
        "mma(split-q+share-kv+swizzle-q+stage1)",
        o,
        stages=1,
    )
    out_mma_share_kv_sq2, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_shared_kv_swizzle_q,
        q,
        k,
        v,
        "mma(split-q+share-kv+swizzle-q+stage2)",
        o,
        stages=2,
    )
    out_mma_share_kv_sqk1, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_shared_kv_swizzle_qk,
        q,
        k,
        v,
        "mma(split-q+share-kv+swizzle-qk+stage1)",
        o,
        stages=1,
    )
    out_mma_share_kv_sqk2, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_shared_kv_swizzle_qk,
        q,
        k,
        v,
        "mma(split-q+share-kv+swizzle-qk+stage2)",
        o,
        stages=2,
    )
    out_mma_share_kv_sqkv1, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_shared_kv_swizzle_qkv,
        q,
        k,
        tv,
        "mma(split-q+share-kv+swizzle-qkv+stage1)",
        o,
        stages=1,
    )
    out_mma_share_kv_sqkv2, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_shared_kv_swizzle_qkv,
        q,
        k,
        tv,
        "mma(split-q+share-kv+swizzle-qkv+stage2)",
        o,
        stages=2,
    )
    # Split-Q + Fully Shared QKV SMEM + Swizzle
    out_mma_share_qkv1, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_shared_qkv,
        q,
        k,
        v,
        "mma(split-q+share-qkv+stage1)",
        o,
        stages=1,
    )
    out_mma_share_qkv2, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_shared_qkv,
        q,
        k,
        v,
        "mma(split-q+share-qkv+stage2)",
        o,
        stages=2,
    )
    out_mma_share_qkv_f321, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_shared_qkv_acc_f32,
        q,
        k,
        v,
        "mma(split-q+share-qkv+acc-f32+stage1)",
        o,
        stages=1,
    )
    out_mma_share_qkv_f322, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_shared_qkv_acc_f32,
        q,
        k,
        v,
        "mma(split-q+share-qkv+acc-f32+stage2)",
        o,
        stages=2,
    )
    out_mma_share_qkv_sq1, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_shared_qkv_swizzle_q,
        q,
        k,
        v,
        "mma(split-q+share-qkv+swizzle-q+stage1)",
        o,
        stages=1,
    )
    out_mma_share_qkv_sq2, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_shared_qkv_swizzle_q,
        q,
        k,
        v,
        "mma(split-q+share-qkv+swizzle-q+stage2)",
        o,
        stages=2,
    )
    out_mma_share_qkv_sqk1, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_shared_qkv_swizzle_qk,
        q,
        k,
        v,
        "mma(split-q+share-qkv+swizzle-qk+stage1)",
        o,
        stages=1,
    )
    out_mma_share_qkv_sqk2, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_shared_qkv_swizzle_qk,
        q,
        k,
        v,
        "mma(split-q+share-qkv+swizzle-qk+stage2)",
        o,
        stages=2,
    )
    out_mma_share_qkv_sqkv1, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_shared_qkv_swizzle_qkv,
        q,
        k,
        tv,
        "mma(split-q+share-qkv+swizzle-qkv+stage1)",
        o,
        stages=1,
    )
    out_mma_share_qkv_sqkv2, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_shared_qkv_swizzle_qkv,
        q,
        k,
        tv,
        "mma(split-q+share-qkv+swizzle-qkv+stage2)",
        o,
        stages=2,
    )
    # Split-Q + QK Fine-grained Tiling + Swizzle
    out_mma_tiling_qk1, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_tiling_qk,
        q,
        k,
        v,
        "mma(split-q+tiling-qk+stage1)",
        o,
        stages=1,
    )
    out_mma_tiling_qk2, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_tiling_qk,
        q,
        k,
        v,
        "mma(split-q+tiling-qk+stage2)",
        o,
        stages=2,
    )
    out_mma_tiling_qk_f321, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_tiling_qk_acc_f32,
        q,
        k,
        v,
        "mma(split-q+tiling-qk+acc-f32+stage1)",
        o,
        stages=1,
    )
    out_mma_tiling_qk_f322, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_tiling_qk_acc_f32,
        q,
        k,
        v,
        "mma(split-q+tiling-qk+acc-f32+stage2)",
        o,
        stages=2,
    )
    out_mma_tiling_qk_sq1, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_tiling_qk_swizzle_q,
        q,
        k,
        v,
        "mma(split-q+tiling-qk+swizzle-q+stage1)",
        o,
        stages=1,
    )
    out_mma_tiling_qk_sq2, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_tiling_qk_swizzle_q,
        q,
        k,
        v,
        "mma(split-q+tiling-qk+swizzle-q+stage2)",
        o,
        stages=2,
    )
    out_mma_tiling_qk_sqk1, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_tiling_qk_swizzle_qk,
        q,
        k,
        v,
        "mma(split-q+tiling-qk+swizzle-qk+stage1)",
        o,
        stages=1,
    )
    out_mma_tiling_qk_sqk2, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_tiling_qk_swizzle_qk,
        q,
        k,
        v,
        "mma(split-q+tiling-qk+swizzle-qk+stage2)",
        o,
        stages=2,
    )
    out_mma_tiling_qk_sqkv1, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_tiling_qk_swizzle_qkv,
        q,
        k,
        tv,
        "mma(split-q+tiling-qk+swizzle-qkv+stage1)",
        o,
        stages=1,
    )
    out_mma_tiling_qk_sqkv2, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_tiling_qk_swizzle_qkv,
        q,
        k,
        tv,
        "mma(split-q+tiling-qk+swizzle-qkv+stage2)",
        o,
        stages=2,
    )
    # Split-Q + QKV Fully Fine-grained Tiling
    out_mma_tiling_qkv1, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_tiling_qkv,
        q,
        k,
        v,
        "mma(split-q+tiling-qkv+stage1)",
        o,
        stages=1,
    )
    out_mma_tiling_qkv2, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_tiling_qkv,
        q,
        k,
        v,
        "mma(split-q+tiling-qkv+stage2)",
        o,
        stages=2,
    )
    out_mma_tiling_qkv_sq1, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_tiling_qkv_swizzle_q,
        q,
        k,
        v,
        "mma(split-q+tiling-qkv+swizzle-q+stage1)",
        o,
        stages=1,
    )
    out_mma_tiling_qkv_sq2, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_tiling_qkv_swizzle_q,
        q,
        k,
        v,
        "mma(split-q+tiling-qkv+swizzle-q+stage2)",
        o,
        stages=2,
    )
    out_mma_tiling_qkv_sqk1, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_tiling_qkv_swizzle_qk,
        q,
        k,
        v,
        "mma(split-q+tiling-qkv+swizzle-qk+stage1)",
        o,
        stages=1,
    )
    out_mma_tiling_qkv_sqk2, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_tiling_qkv_swizzle_qk,
        q,
        k,
        v,
        "mma(split-q+tiling-qkv+swizzle-qk+stage2)",
        o,
        stages=2,
    )
    out_mma_tiling_qkv_sqkv1, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_tiling_qkv_swizzle_qkv,
        q,
        k,
        v,
        "mma(split-q+tiling-qkv+swizzle-qkv+stage1)",
        o,
        stages=1,
    )
    out_mma_tiling_qkv_sqkv2, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_tiling_qkv_swizzle_qkv,
        q,
        k,
        v,
        "mma(split-q+tiling-qkv+swizzle-qkv+stage2)",
        o,
        stages=2,
    )
    out_mma_tiling_qkv_f321, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_tiling_qkv_acc_f32,
        q,
        k,
        v,
        "mma(split-q+tiling-qkv+acc-f32+stage1)",
        o,
        stages=1,
    )
    out_mma_tiling_qkv_f322, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_tiling_qkv_acc_f32,
        q,
        k,
        v,
        "mma(split-q+tiling-qkv+acc-f32+stage2)",
        o,
        stages=2,
    )
    out_mma_tiling_qkv_fsq1, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_tiling_qkv_acc_f32_swizzle_q,
        q,
        k,
        v,
        "mma(split-q+tiling-qkv+acc-f32+swizzle-q+stage1)",
        o,
        stages=1,
    )
    out_mma_tiling_qkv_fsq2, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_tiling_qkv_acc_f32_swizzle_q,
        q,
        k,
        v,
        "mma(split-q+tiling-qkv+acc-f32+swizzle-q+stage2)",
        o,
        stages=2,
    )
    out_mma_tiling_qkv_fsqk1, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_tiling_qkv_acc_f32_swizzle_qk,
        q,
        k,
        v,
        "mma(split-q+tiling-qkv+acc-f32+swizzle-qk+stage1)",
        o,
        stages=1,
    )
    out_mma_tiling_qkv_fsqk2, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_tiling_qkv_acc_f32_swizzle_qk,
        q,
        k,
        v,
        "mma(split-q+tiling-qkv+acc-f32+swizzle-qk+stage2)",
        o,
        stages=2,
    )
    out_mma_tiling_qkv_fsqkv1, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_tiling_qkv_acc_f32_swizzle_qkv,
        q,
        k,
        v,
        "mma(split-q+tiling-qkv+acc-f32+swizzle-qkv+stage1)",
        o,
        stages=1,
    )
    out_mma_tiling_qkv_fsqkv2, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_tiling_qkv_acc_f32_swizzle_qkv,
        q,
        k,
        v,
        "mma(split-q+tiling-qkv+acc-f32+swizzle-qkv+stage2)",
        o,
        stages=2,
    )
    # Others, O s2g, etc.
    out_mma_share_kv_rr1, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_shared_kv_acc_f32_rr,
        q,
        k,
        v,
        "mma(split-q+share-kv+acc-f32+rr+stage1)",
        o,
        stages=1,
    )
    out_mma_share_kv_rr2, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_shared_kv_acc_f32_rr,
        q,
        k,
        v,
        "mma(split-q+share-kv+acc-f32+rr+stage2)",
        o,
        stages=2,
    )
    out_mma_share_qkv_s2g1, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_shared_qkv_Os2g,
        q,
        k,
        v,
        "mma(split-q+share-qkv+o-s2g+stage1)",
        o,
        stages=1,
    )
    out_mma_share_qkv_s2g2, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_shared_qkv_Os2g,
        q,
        k,
        v,
        "mma(split-q+share-qkv+o-s2g+stage2)",
        o,
        stages=2,
    )
    out_mma_share_qkv_rr1, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_shared_qkv_acc_f32_rr,
        q,
        k,
        v,
        "mma(split-q+share-qkv+acc-f32+rr+stage1)",
        o,
        stages=1,
    )
    out_mma_share_qkv_rr2, _ = run_benchmark(
        lib.flash_attn_mma_stages_split_q_shared_qkv_acc_f32_rr,
        q,
        k,
        v,
        "mma(split-q+share-qkv+acc-f32+rr+stage2)",
        o,
        stages=2,
    )
    # FA2, SDPA official
    out_flash, _ = run_benchmark(flash_attn_func, fq, fk, fv, "(flash)")
    out_sdpa, _ = run_benchmark(
        partial(sdpa, use_flash=(D <= 256)), q, k, v, "(sdpa)"
    )
    pretty_print_line()

    torch.cuda.synchronize()
    if args.check:
        if D <= 256:
            pretty_print_line()
            # Split-KV
            check_all_close(
                out_flash,
                out_mma_split_kv1,
                "out_mma_split_kv1",
                args.check_all,
            )
            check_all_close(
                out_flash,
                out_mma_split_kv2,
                "out_mma_split_kv2",
                args.check_all,
            )
            # Split-Q
            check_all_close(
                out_flash, out_mma_split_q1, "out_mma_split_q1", args.check_all
            )
            check_all_close(
                out_flash, out_mma_split_q2, "out_mma_split_q2", args.check_all
            )
            # Split-Q + Shared KV SMEM
            check_all_close(
                out_flash,
                out_mma_share_kv1,
                "out_mma_share_kv1",
                args.check_all,
            )
            check_all_close(
                out_flash,
                out_mma_share_kv2,
                "out_mma_share_kv2",
                args.check_all,
            )
            check_all_close(
                out_flash,
                out_mma_share_kv_f321,
                "out_mma_share_kv_f321",
                args.check_all,
            )
            check_all_close(
                out_flash,
                out_mma_share_kv_f322,
                "out_mma_share_kv_f322",
                args.check_all,
            )
            check_all_close(
                out_flash,
                out_mma_share_kv_sq1,
                "out_mma_share_kv_sq1",
                args.check_all,
            )
            check_all_close(
                out_flash,
                out_mma_share_kv_sq2,
                "out_mma_share_kv_sq2",
                args.check_all,
            )
            check_all_close(
                out_flash,
                out_mma_share_kv_sqk1,
                "out_mma_share_kv_sqk1",
                args.check_all,
            )
            check_all_close(
                out_flash,
                out_mma_share_kv_sqk2,
                "out_mma_share_kv_sqk2",
                args.check_all,
            )
            check_all_close(
                out_flash,
                out_mma_share_kv_sqkv1,
                "out_mma_share_kv_sqkv1",
                args.check_all,
            )
            check_all_close(
                out_flash,
                out_mma_share_kv_sqkv2,
                "out_mma_share_kv_sqkv2",
                args.check_all,
            )
            # Split-Q + Fully Shared QKV SMEM
            check_all_close(
                out_flash,
                out_mma_share_qkv1,
                "out_mma_share_qkv1",
                args.check_all,
            )
            check_all_close(
                out_flash,
                out_mma_share_qkv2,
                "out_mma_share_qkv2",
                args.check_all,
            )
            check_all_close(
                out_flash,
                out_mma_share_qkv_f321,
                "out_mma_share_qkv_f321",
                args.check_all,
            )
            check_all_close(
                out_flash,
                out_mma_share_qkv_f322,
                "out_mma_share_qkv_f322",
                args.check_all,
            )
            check_all_close(
                out_flash,
                out_mma_share_qkv_sq1,
                "out_mma_share_qkv_sq1",
                args.check_all,
            )
            check_all_close(
                out_flash,
                out_mma_share_qkv_sq2,
                "out_mma_share_qkv_sq2",
                args.check_all,
            )
            check_all_close(
                out_flash,
                out_mma_share_qkv_sqk1,
                "out_mma_share_qkv_sqk1",
                args.check_all,
            )
            check_all_close(
                out_flash,
                out_mma_share_qkv_sqk2,
                "out_mma_share_qkv_sqk2",
                args.check_all,
            )
            check_all_close(
                out_flash,
                out_mma_share_qkv_sqkv1,
                "out_mma_share_qkv_sqkv1",
                args.check_all,
            )
            check_all_close(
                out_flash,
                out_mma_share_qkv_sqkv2,
                "out_mma_share_qkv_sqkv2",
                args.check_all,
            )
            # Split-Q + QK Fine-grained Tiling
            check_all_close(
                out_flash,
                out_mma_tiling_qk1,
                "out_mma_tiling_qk1",
                args.check_all,
            )
            check_all_close(
                out_flash,
                out_mma_tiling_qk2,
                "out_mma_tiling_qk2",
                args.check_all,
            )
            check_all_close(
                out_flash,
                out_mma_tiling_qk_f321,
                "out_mma_tiling_qk_f321",
                args.check_all,
            )
            check_all_close(
                out_flash,
                out_mma_tiling_qk_f322,
                "out_mma_tiling_qk_f322",
                args.check_all,
            )
            check_all_close(
                out_flash,
                out_mma_tiling_qk_sq1,
                "out_mma_tiling_qk_sq1",
                args.check_all,
            )
            check_all_close(
                out_flash,
                out_mma_tiling_qk_sq2,
                "out_mma_tiling_qk_sq2",
                args.check_all,
            )
            check_all_close(
                out_flash,
                out_mma_tiling_qk_sqk1,
                "out_mma_tiling_qk_sqk1",
                args.check_all,
            )
            check_all_close(
                out_flash,
                out_mma_tiling_qk_sqk2,
                "out_mma_tiling_qk_sqk2",
                args.check_all,
            )
            check_all_close(
                out_flash,
                out_mma_tiling_qk_sqkv1,
                "out_mma_tiling_qk_sqkv1",
                args.check_all,
            )
            check_all_close(
                out_flash,
                out_mma_tiling_qk_sqkv2,
                "out_mma_tiling_qk_sqkv2",
                args.check_all,
            )
            # Split-Q + Fully QKV Fine-grained Tiling
            check_all_close(
                out_flash,
                out_mma_tiling_qkv1,
                "out_mma_tiling_qkv1",
                args.check_all,
            )
            check_all_close(
                out_flash,
                out_mma_tiling_qkv2,
                "out_mma_tiling_qkv2",
                args.check_all,
            )
            check_all_close(
                out_flash,
                out_mma_tiling_qkv_sq1,
                "out_mma_tiling_qkv_sq1",
                args.check_all,
            )
            check_all_close(
                out_flash,
                out_mma_tiling_qkv_sq2,
                "out_mma_tiling_qkv_sq2",
                args.check_all,
            )
            check_all_close(
                out_flash,
                out_mma_tiling_qkv_sqk1,
                "out_mma_tiling_qkv_sqk1",
                args.check_all,
            )
            check_all_close(
                out_flash,
                out_mma_tiling_qkv_sqk2,
                "out_mma_tiling_qkv_sqk2",
                args.check_all,
            )
            check_all_close(
                out_flash,
                out_mma_tiling_qkv_sqkv1,
                "out_mma_tiling_qkv_sqkv1",
                args.check_all,
            )
            check_all_close(
                out_flash,
                out_mma_tiling_qkv_sqkv2,
                "out_mma_tiling_qkv_sqkv2",
                args.check_all,
            )
            check_all_close(
                out_flash,
                out_mma_tiling_qkv_f321,
                "out_mma_tiling_qkv_f321",
                args.check_all,
            )
            check_all_close(
                out_flash,
                out_mma_tiling_qkv_f322,
                "out_mma_tiling_qkv_f322",
                args.check_all,
            )
            check_all_close(
                out_flash,
                out_mma_tiling_qkv_fsq1,
                "out_mma_tiling_qkv_fsq1",
                args.check_all,
            )
            check_all_close(
                out_flash,
                out_mma_tiling_qkv_fsq2,
                "out_mma_tiling_qkv_fsq2",
                args.check_all,
            )
            check_all_close(
                out_flash,
                out_mma_tiling_qkv_fsqk1,
                "out_mma_tiling_qkv_fsqk1",
                args.check_all,
            )
            check_all_close(
                out_flash,
                out_mma_tiling_qkv_fsqk2,
                "out_mma_tiling_qkv_fsqk2",
                args.check_all,
            )
            check_all_close(
                out_flash,
                out_mma_tiling_qkv_fsqkv1,
                "out_mma_tiling_qkv_fsqkv1",
                args.check_all,
            )
            check_all_close(
                out_flash,
                out_mma_tiling_qkv_fsqkv2,
                "out_mma_tiling_qkv_fsqkv2",
                args.check_all,
            )
            # Others, O s2g, etc.
            check_all_close(
                out_flash,
                out_mma_share_kv_rr1,
                "out_mma_share_kv_rr1",
                args.check_all,
            )
            check_all_close(
                out_flash,
                out_mma_share_kv_rr2,
                "out_mma_share_kv_rr2",
                args.check_all,
            )
            check_all_close(
                out_flash,
                out_mma_share_qkv_s2g1,
                "out_mma_share_qkv_s2g1",
                args.check_all,
            )
            check_all_close(
                out_flash,
                out_mma_share_qkv_s2g2,
                "out_mma_share_qkv_s2g2",
                args.check_all,
            )
            check_all_close(
                out_flash,
                out_mma_share_qkv_rr1,
                "out_mma_share_qkv_rr1",
                args.check_all,
            )
            check_all_close(
                out_flash,
                out_mma_share_qkv_rr2,
                "out_mma_share_qkv_rr2",
                args.check_all,
            )
            pretty_print_line()
        elif args.run_torch_sdpa:
            pretty_print_line()
            # Split-Q + Fully Shared QKV SMEM
            check_all_close(
                out_sdpa,
                out_mma_share_qkv1,
                "out_mma_share_qkv1",
                args.check_all,
                False,
            )
            check_all_close(
                out_sdpa,
                out_mma_share_qkv2,
                "out_mma_share_qkv2",
                args.check_all,
                False,
            )
            check_all_close(
                out_sdpa,
                out_mma_share_qkv_f321,
                "out_mma_share_qkv_f321",
                args.check_all,
                False,
            )
            check_all_close(
                out_sdpa,
                out_mma_share_qkv_f322,
                "out_mma_share_qkv_f322",
                args.check_all,
                False,
            )
            # Split-Q + QK Fine-grained Tiling
            check_all_close(
                out_sdpa,
                out_mma_tiling_qk1,
                "out_mma_tiling_qk1",
                args.check_all,
                False,
            )
            check_all_close(
                out_sdpa,
                out_mma_tiling_qk2,
                "out_mma_tiling_qk2",
                args.check_all,
                False,
            )
            check_all_close(
                out_sdpa,
                out_mma_tiling_qk_f321,
                "out_mma_tiling_qk_f321",
                args.check_all,
                False,
            )
            check_all_close(
                out_sdpa,
                out_mma_tiling_qk_f322,
                "out_mma_tiling_qk_f322",
                args.check_all,
                False,
            )
            check_all_close(
                out_sdpa,
                out_mma_tiling_qk_sq1,
                "out_mma_tiling_qk_sq1",
                args.check_all,
                False,
            )
            check_all_close(
                out_sdpa,
                out_mma_tiling_qk_sq2,
                "out_mma_tiling_qk_sq2",
                args.check_all,
                False,
            )
            check_all_close(
                out_sdpa,
                out_mma_tiling_qk_sqk1,
                "out_mma_tiling_qk_sqk1",
                args.check_all,
                False,
            )
            check_all_close(
                out_sdpa,
                out_mma_tiling_qk_sqk2,
                "out_mma_tiling_qk_sqk2",
                args.check_all,
                False,
            )
            check_all_close(
                out_sdpa,
                out_mma_tiling_qk_sqkv1,
                "out_mma_tiling_qk_sqkv1",
                args.check_all,
                False,
            )
            check_all_close(
                out_sdpa,
                out_mma_tiling_qk_sqkv2,
                "out_mma_tiling_qk_sqkv2",
                args.check_all,
                False,
            )
            # Split-Q + Fully QKV Fine-grained Tiling
            check_all_close(
                out_sdpa,
                out_mma_tiling_qkv1,
                "out_mma_tiling_qkv1",
                args.check_all,
                False,
            )
            check_all_close(
                out_sdpa,
                out_mma_tiling_qkv2,
                "out_mma_tiling_qkv2",
                args.check_all,
                False,
            )
            check_all_close(
                out_sdpa,
                out_mma_tiling_qkv_sq1,
                "out_mma_tiling_qkv_sq1",
                args.check_all,
                False,
            )
            check_all_close(
                out_sdpa,
                out_mma_tiling_qkv_sq2,
                "out_mma_tiling_qkv_sq2",
                args.check_all,
                False,
            )
            check_all_close(
                out_sdpa,
                out_mma_tiling_qkv_sqk1,
                "out_mma_tiling_qkv_sqk1",
                args.check_all,
                False,
            )
            check_all_close(
                out_sdpa,
                out_mma_tiling_qkv_sqk2,
                "out_mma_tiling_qkv_sqk2",
                args.check_all,
                False,
            )
            check_all_close(
                out_sdpa,
                out_mma_tiling_qkv_sqkv1,
                "out_mma_tiling_qkv_sqkv1",
                args.check_all,
                False,
            )
            check_all_close(
                out_sdpa,
                out_mma_tiling_qkv_sqkv2,
                "out_mma_tiling_qkv_sqkv2",
                args.check_all,
                False,
            )
            check_all_close(
                out_sdpa,
                out_mma_tiling_qkv_f321,
                "out_mma_tiling_qkv_f321",
                args.check_all,
                False,
            )
            check_all_close(
                out_sdpa,
                out_mma_tiling_qkv_f322,
                "out_mma_tiling_qkv_f322",
                args.check_all,
                False,
            )
            check_all_close(
                out_sdpa,
                out_mma_tiling_qkv_fsq1,
                "out_mma_tiling_qkv_fsq1",
                args.check_all,
                False,
            )
            check_all_close(
                out_sdpa,
                out_mma_tiling_qkv_fsq2,
                "out_mma_tiling_qkv_fsq2",
                args.check_all,
                False,
            )
            check_all_close(
                out_sdpa,
                out_mma_tiling_qkv_fsqk1,
                "out_mma_tiling_qkv_fsqk1",
                args.check_all,
                False,
            )
            check_all_close(
                out_sdpa,
                out_mma_tiling_qkv_fsqk2,
                "out_mma_tiling_qkv_fsqk2",
                args.check_all,
                False,
            )
            check_all_close(
                out_sdpa,
                out_mma_tiling_qkv_fsqkv1,
                "out_mma_tiling_qkv_fsqkv1",
                args.check_all,
                False,
            )
            check_all_close(
                out_sdpa,
                out_mma_tiling_qkv_fsqkv2,
                "out_mma_tiling_qkv_fsqkv2",
                args.check_all,
                False,
            )
            # Others, O s2g, etc.
            check_all_close(
                out_sdpa,
                out_mma_share_qkv_rr1,
                "out_mma_share_qkv_rr1",
                args.check_all,
                False,
            )
            check_all_close(
                out_sdpa,
                out_mma_share_qkv_rr2,
                "out_mma_share_qkv_rr2",
                args.check_all,
                False,
            )
            pretty_print_line()