"""
Scaled Dot Product Attention (SDPA) benchmark

This script benchmarks a single SDPA compute instance.
The SDPA backend can be chosen. Performance is measured using torch profiler.

Can be used as CLI or imported as a module:

    # CLI usage
    python benchmark_single_sdpa.py --batch_size 1 --q_seqlen 8192 ...

    # Module usage
    from benchmark_single_sdpa import run_benchmark
    result = run_benchmark(batch_size=1, q_seqlen=8192, ...)
"""

import argparse
import torch
from torch.nn.attention import SDPBackend, sdpa_kernel
import os
import numpy as np
import functools
import math
import threading
import time
from typing import Optional, Dict, Any

from torch.profiler import profile, record_function, ProfilerActivity


# Dense MMA throughput (FLOPs / clock / SM) for Blackwell datacenter SKUs.
# BF16/FP16 dense = 8192, FP8/MXFP8 dense = 16384 (MXFP8 uses the FP8
# datapath with block scaling).
# Keys match the strings accepted by the --data_type CLI flag.
_BLACKWELL_DC_FLOPS_PER_CLOCK_PER_SM = {
    "bfloat16":  8192,
    "float16":   8192,
    "fp8":      16384,
    "mxfp8":    16384,
}


def _peak_flops_per_clock_per_sm(dtype_str):
    """Return per-SM per-clock dense FLOPs for the current GPU + dtype.
    Returns None on unsupported arch (anything other than Blackwell DC for now)."""
    if not torch.cuda.is_available():
        return None
    props = torch.cuda.get_device_properties(torch.cuda.current_device())
    if props.major != 10:  # only Blackwell DC is in scope
        return None
    return _BLACKWELL_DC_FLOPS_PER_CLOCK_PER_SM.get(dtype_str)


class _SmClockSampler:
    """Background thread that polls SM clock via NVML at ~1 kHz.

    Used to capture the actual boost clock during the benchmark window.
    `nvmlDeviceGetMaxClockInfo` is unreliable on some Blackwell datacenter
    SKUs: it can report a value below the boost the kernel actually runs
    at, producing nonsensical (>100%) SOL numbers downstream.
    """

    def __init__(self):
        self._samples = []
        self._stop = threading.Event()
        self._thread = None
        self._handle = None
        self._pynvml = None

    def start(self):
        try:
            import pynvml
            pynvml.nvmlInit()
            self._pynvml = pynvml
            self._handle = pynvml.nvmlDeviceGetHandleByIndex(torch.cuda.current_device())
        except Exception:
            self._pynvml = None
            return
        self._thread = threading.Thread(target=self._run, daemon=True)
        self._thread.start()

    def _run(self):
        pynvml = self._pynvml
        while not self._stop.is_set():
            try:
                self._samples.append(pynvml.nvmlDeviceGetClockInfo(self._handle, pynvml.NVML_CLOCK_SM))
            except Exception:
                pass
            # Sample at ~1 kHz; kernels run much longer than this in aggregate
            # across warmup + measurement iterations.
            time.sleep(0.001)

    def stop(self):
        if self._thread is None:
            return
        self._stop.set()
        self._thread.join()
        try:
            if self._pynvml is not None:
                self._pynvml.nvmlShutdown()
        except Exception:
            pass

    def peak_mhz(self):
        """Return max sampled SM clock (MHz), or None if no samples."""
        return max(self._samples) if self._samples else None


try:
    import cutlass.cute as cute
    import cutlass
    from cutlass.cute.runtime import from_dlpack

    @cute.jit
    def cvt_sf_MKL_to_M32x4xrm_K4xrk_L(sf_ref_tensor, sf_mma_tensor):
        sf_mma_tensor = cute.group_modes(sf_mma_tensor, 0, 3)
        sf_mma_tensor = cute.group_modes(sf_mma_tensor, 1, 3)
        for i in cutlass.range(cute.size(sf_ref_tensor)):
            mkl_coord = sf_ref_tensor.layout.get_hier_coord(i)
            sf_mma_tensor[mkl_coord] = sf_ref_tensor[mkl_coord]

    HAS_CUTLASS = True
except Exception:
    HAS_CUTLASS = False


def ceil_div(a: int, b: int) -> int:
    return (a + b - 1) // b


def create_scale_factor_tensor_mxfp8(l, mn, k, sf_vec_size):
    """Create MXFP8 scale factor tensor in F8_128x4 layout for cuDNN.
    Returns cute tensor (F8_128x4 layout) on GPU.
    l: batch dimension (b * h), mn: non-contracting dim, k: contracting dim.
    """
    sf_k = ceil_div(k, sf_vec_size)
    atom_m = (32, 4)
    atom_k = 4
    mma_shape = (l, ceil_div(mn, atom_m[0] * atom_m[1]), ceil_div(sf_k, atom_k), atom_m[0], atom_m[1], atom_k)
    mma_permute_order = (3, 4, 1, 5, 2, 0)
    cute_f32_cpu = torch.zeros(mma_shape, dtype=torch.float32).permute(mma_permute_order)

    ref_shape = (l, mn, sf_k)
    ref_permute = (1, 2, 0)
    ref_f32_cpu = torch.empty(ref_shape).uniform_(0.5, 2.0).permute(ref_permute).to(torch.int8).to(torch.float32)

    cvt_sf_MKL_to_M32x4xrm_K4xrk_L(from_dlpack(ref_f32_cpu), from_dlpack(cute_f32_cpu))
    return cute_f32_cpu.to(torch.float8_e8m0fnu).cuda()


def parse_args():
    """Parse command line arguments."""
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("--batch_size", default=1, type=int, help="Batch size to input to the layer")
    parser.add_argument("--q_seqlen", default=8192, type=int, help="Sequence length to input to the layer")
    parser.add_argument("--kv_seqlen", default=8192, type=int, help="Sequence length to input to the layer")
    parser.add_argument(
        "--num_q_heads",
        default=16,
        type=int,
        help="Number of query heads to input to the layer",
    )
    parser.add_argument(
        "--num_kv_heads",
        default=8,
        type=int,
        help="Number of key/value heads to input to the layer",
    )
    parser.add_argument("--head_dim", default=128, type=int, help="Head dimension to input to the layer")
    parser.add_argument(
        "--head_dim_qk",
        default=None,
        type=int,
        help="Optional: head dimension for Q/K. If set, must also set --head_dim_vo",
    )
    parser.add_argument(
        "--head_dim_vo",
        default=None,
        type=int,
        help="Optional: head dimension for V/O. If set, must also set --head_dim_qk",
    )
    parser.add_argument(
        "--data_type",
        default="bfloat16",
        choices=["bfloat16", "float16", "float", "fp8", "mxfp8"],
        help="Data type. Can be bfloat16, float16, fp8, or mxfp8",
    )
    parser.add_argument(
        "--num_iterations",
        default=20,
        type=int,
        help="Number of iterations to run the layer for performance measurement",
    )
    parser.add_argument(
        "--num_warmup_iterations",
        default=0,
        type=int,
        help="Number of warmup iterations to run before measuring performance",
    )
    parser.add_argument("--verbose", action="store_true", help="Verbose output")
    parser.add_argument(
        "--fa4_num_splits",
        default=None,
        type=int,
        help="FlashAttention-4 only: force num_splits (KV split count). " "Default is None (FA4 picks automatically).",
    )
    parser.add_argument(
        "--fwd_bwd",
        action="store_true",
        help="Run both forward and backward pass (fwd only by default)",
    )
    parser.add_argument(
        "--profile_pass",
        default=None,
        type=str,
        choices=["fwd", "bwd", "both"],
        help="Which pass to profile (default: fwd unless --fwd_bwd is set).",
    )
    parser.add_argument(
        "--deterministic_bwd",
        action="store_true",
        help="Use deterministic algorithm for backward pass where supported (cudnn FP16/BF16/FP8, flash_attention, flash_attention_3, flash_attention_4)",
    )
    parser.add_argument(
        "--attn_mask",
        default="no_mask",
        type=str,
        help="Attn mask to use. Can be 'top_left' (causal) or 'no_mask' (bidirectional).",
        choices=["top_left", "bottom_right", "no_mask"],
    )
    parser.add_argument(
        "--sdpa_backend",
        default="pyt_cudnn",
        type=str,
        help="SDPA backend to use",
        choices=[
            "pyt_math",
            "pyt_cudnn",
            "pyt_efficient_attention",
            "pyt_flash_attention",
            "flash_attention",
            "flash_attention_3",
            "flash_attention_4",
            "cudnn",
        ],
    )
    parser.add_argument("--format_output", action="store_true", help="Format output to be used in benchmark")
    parser.add_argument(
        "--case_tag",
        default="",
        type=str,
        help="Tag to identify the case. Not used in calculations. Only for formatted output",
    )
    parser.add_argument(
        "--skip_ref",
        action="store_true",
        help="Skip reference SDPA implementation",
    )
    parser.add_argument(
        "--sliding_window_size",
        default=None,
        type=int,
        help="Sliding window attention size (number of tokens to look back from diagonal)",
    )
    return parser.parse_args()


def run_benchmark(
    batch_size: int,
    q_seqlen: int,
    kv_seqlen: int,
    num_q_heads: int,
    num_kv_heads: int,
    head_dim: int = 128,
    head_dim_qk: Optional[int] = None,
    head_dim_vo: Optional[int] = None,
    data_type: str = "bfloat16",
    backend: str = "cudnn",
    attn_mask: str = "no_mask",
    profile_pass: str = "fwd",
    num_iterations: int = 10,
    num_warmup_iterations: int = 0,
    skip_ref: bool = True,
    deterministic_bwd: bool = False,
    sliding_window_size: Optional[int] = None,
    verbose: bool = False,
    fa4_num_splits: Optional[int] = None,
) -> Dict[str, Any]:
    """
    Run a single SDPA benchmark.

    This function can be called directly when using the module as a library.
    Internally uses subprocess to call this script with the appropriate arguments.

    Args:
        batch_size: Batch size
        q_seqlen: Query sequence length
        kv_seqlen: Key/value sequence length
        num_q_heads: Number of query heads
        num_kv_heads: Number of key/value heads
        head_dim: Head dimension (used if head_dim_qk/vo not specified)
        head_dim_qk: Head dimension for Q/K (optional, for asymmetric)
        head_dim_vo: Head dimension for V/O (optional, for asymmetric)
        data_type: Data type ("bfloat16", "float16", "fp8")
        backend: Backend name ("cudnn", "flash_attention_4", etc.)
        attn_mask: Attention mask ("no_mask", "top_left", "bottom_right")
        profile_pass: Which pass to profile ("fwd", "bwd", "both")
        num_iterations: Number of benchmark iterations
        num_warmup_iterations: Warmup iterations before measurement
        skip_ref: Skip reference validation
        deterministic_bwd: Use deterministic backward algorithm
        sliding_window_size: Sliding window attention size (None = no sliding window)
        verbose: Print verbose output

    Returns:
        Dict with keys:
            - time_ms: Median time of the requested pass in milliseconds
            - tflops: TFLOPS for the requested pass
            - max_diff: Maximum difference vs reference
            - gpu_name: GPU name string
            - cudnn_version: cuDNN version (if available)

    Raises:
        RuntimeError: If the benchmark subprocess fails or profile_pass=="both"
            (callers must invoke once per pass so each has independent success)
    """
    if profile_pass == "both":
        raise RuntimeError("run_benchmark no longer accepts profile_pass='both'. Call once " "per pass ('fwd' or 'bwd') so failures remain independent.")
    import subprocess
    import sys

    # Build command
    script_path = os.path.abspath(__file__)
    cmd = [
        sys.executable,
        script_path,
        "--batch_size",
        str(batch_size),
        "--q_seqlen",
        str(q_seqlen),
        "--kv_seqlen",
        str(kv_seqlen),
        "--num_q_heads",
        str(num_q_heads),
        "--num_kv_heads",
        str(num_kv_heads),
        "--data_type",
        data_type,
        "--sdpa_backend",
        backend,
        "--attn_mask",
        attn_mask,
        "--num_iterations",
        str(num_iterations),
        "--num_warmup_iterations",
        str(num_warmup_iterations),
        "--format_output",  # Get CSV-formatted output for parsing
    ]

    # Handle head dimensions
    if head_dim_qk is not None and head_dim_vo is not None:
        cmd.extend(["--head_dim_qk", str(head_dim_qk)])
        cmd.extend(["--head_dim_vo", str(head_dim_vo)])
    else:
        cmd.extend(["--head_dim", str(head_dim)])

    # Handle profile pass (single pass only)
    cmd.extend(["--profile_pass", profile_pass])

    # Handle flags
    if skip_ref:
        cmd.append("--skip_ref")
    if deterministic_bwd:
        cmd.append("--deterministic_bwd")
    if sliding_window_size is not None:
        cmd.extend(["--sliding_window_size", str(sliding_window_size)])
    if verbose:
        cmd.append("--verbose")
    if fa4_num_splits is not None:
        cmd.extend(["--fa4_num_splits", str(fa4_num_splits)])

    # Run benchmark
    result = subprocess.run(
        cmd,
        capture_output=True,
        text=True,
        check=False,
    )

    if result.returncode != 0:
        raise RuntimeError(f"Benchmark failed with return code {result.returncode}.\n" f"stderr: {result.stderr}\n" f"stdout: {result.stdout}")

    # Parse CSV output
    # Format: case_tag,backend,batch_size,q_seqlen,kv_seqlen,num_q_heads,num_kv_heads,head_dim,fwd_time,bwd_time,fwd_tflops,bwd_tflops,max_diff,num_iters
    output_line = result.stdout.strip().split("\n")[-1]
    parts = output_line.split(",")

    if len(parts) < 12:
        raise RuntimeError(f"Unexpected output format: {output_line}")

    # Get GPU name from torch
    gpu_name = torch.cuda.get_device_name(torch.cuda.current_device()) if torch.cuda.is_available() else "Unknown"

    # Try to get cudnn version
    cudnn_version = None
    cudnn_backend_version = None
    try:
        import cudnn

        cudnn_version = cudnn.__version__
        cudnn_backend_version = cudnn.backend_version()
    except ImportError:
        pass

    # Subprocess CSV layout keeps both fwd and bwd columns; pick the one
    # corresponding to the requested pass. The unused pass is 0 when not run.
    if profile_pass == "fwd":
        time_ms = float(parts[8])
        tflops = float(parts[10])
    else:  # "bwd"
        time_ms = float(parts[9])
        tflops = float(parts[11])

    return {
        "time_ms": time_ms,
        "tflops": tflops,
        "max_diff": float(parts[12]) if len(parts) > 12 else 0.0,
        "gpu_name": gpu_name,
        "cudnn_version": cudnn_version,
        "cudnn_backend_version": cudnn_backend_version,
    }


# ============================================================================
# Main benchmark implementation (runs when script is executed directly)
# ============================================================================

# Note: All code below this point is only executed when running as a script.
# When imported as a module, use the run_benchmark() function above.

if __name__ != "__main__":
    # Stop here when imported as module
    pass
else:
    # Parse command line arguments
    args = parse_args()

    if args.data_type == "bfloat16":
        target_dtype = torch.bfloat16
    elif args.data_type == "float16":
        target_dtype = torch.float16
    elif args.data_type == "float":
        target_dtype = torch.float
    elif args.data_type == "fp8":
        target_dtype = torch.float8_e4m3fn
    elif args.data_type == "mxfp8":
        target_dtype = torch.float8_e4m3fn  # Q/K/V input type
        output_dtype = torch.bfloat16  # O output type
    else:
        raise ValueError(f"Invalid data type: {args.data_type}")

    if args.sliding_window_size is not None and args.attn_mask not in ("top_left", "bottom_right"):
        raise ValueError("Sliding window attention requires attn_mask='top_left' or 'bottom_right' (causal)")

    # Parse input arguments
    num_iters = args.num_iterations
    dry_run_iters = args.num_warmup_iterations
    batch_size = args.batch_size
    q_seqlen = args.q_seqlen
    kv_seqlen = args.kv_seqlen
    num_q_heads = args.num_q_heads
    num_kv_heads = args.num_kv_heads
    if args.head_dim_qk is None and args.head_dim_vo is None:
        head_dim_qk = args.head_dim
        head_dim_vo = args.head_dim
    elif args.head_dim_qk is not None and args.head_dim_vo is not None:
        head_dim_qk = args.head_dim_qk
        head_dim_vo = args.head_dim_vo
    else:
        raise ValueError("Both --head_dim_qk and --head_dim_vo must be provided together when using asymmetric head dims.")
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    assert device.type == "cuda", "Requires CUDA device"
    if args.profile_pass is not None:
        run_fwd = args.profile_pass in ("fwd", "both")
        run_bwd = args.profile_pass in ("bwd", "both")
    elif args.fwd_bwd:
        run_fwd = True
        run_bwd = True
    else:
        run_fwd = True
        run_bwd = False
    enable_gqa = num_q_heads != num_kv_heads
    if args.data_type == "mxfp8":
        if args.sdpa_backend != "cudnn":
            raise ValueError("mxfp8 is only supported with the 'cudnn' backend")
        if not HAS_CUTLASS:
            raise RuntimeError("CUTLASS is required for mxfp8 but is not installed")
    # if args.sdpa_backend in ["flash_attention", "flash_attention_3", "pyt_flash_attention"]:
    #     assert args.attn_mask != "top_left", "Flash Attention does not support top left causal mask"

    l2_flush_size_mb = 256
    l2_flush_size = l2_flush_size_mb * 1024 * 1024
    l2_flush_buffer = torch.empty(l2_flush_size, device=device, dtype=torch.int8)

    #############################################################
    ########### Set up SDPA function for each backend ###########

    ## If using cuDNN FE, set up cuDNN graph.
    if args.sdpa_backend == "cudnn":
        is_dropout = False  # Hard coded
        dropout_prob = dropout_p if is_dropout else 0.0  # Hard coded to 0
        is_infer = False  # Hard coded
        attn_scale = head_dim_qk ** (-0.5)

        try:
            import cudnn
        except ImportError:
            cudnn = None
        assert cudnn is not None

        if args.verbose:
            print(f"[INFO] cuDNN Backend Version: {cudnn.backend_version() = }")
            print(f"[INFO] cuDNN Frontend Version: {cudnn.__version__ = }")

        # Helper function: Convert torch type to cuDNN type
        def convert_to_cudnn_type(torch_type):
            if torch_type == torch.float16:
                return cudnn.data_type.HALF
            elif torch_type == torch.bfloat16:
                return cudnn.data_type.BFLOAT16
            elif torch_type == torch.float32:
                return cudnn.data_type.FLOAT
            elif torch_type == torch.int32:
                return cudnn.data_type.INT32
            elif torch_type == torch.int64:
                return cudnn.data_type.INT64
            elif torch_type == torch.float8_e4m3fn:
                return cudnn.data_type.FP8_E4M3
            else:
                raise ValueError("Unsupported tensor data type.")

        ## Will define tensors to set up cuDNN graph once.
        # FP8/MXFP8 needs randn in bfloat16 then convert (randn doesn't support fp8 well)
        randn_dtype = torch.bfloat16 if args.data_type in ("fp8", "mxfp8") else target_dtype
        query = torch.randn(batch_size, q_seqlen, num_q_heads, head_dim_qk, dtype=randn_dtype, device=device).to(target_dtype).transpose(1, 2)
        key = torch.randn(batch_size, kv_seqlen, num_kv_heads, head_dim_qk, dtype=randn_dtype, device=device).to(target_dtype).transpose(1, 2)
        value = torch.randn(batch_size, kv_seqlen, num_kv_heads, head_dim_vo, dtype=randn_dtype, device=device).to(target_dtype).transpose(1, 2)
        if args.data_type == "mxfp8":
            output = torch.empty(batch_size, q_seqlen, num_q_heads, head_dim_vo, dtype=output_dtype, device=device).transpose(1, 2)
        else:
            output = torch.empty(batch_size, q_seqlen, num_q_heads, head_dim_vo, dtype=target_dtype, device=device).transpose(1, 2)

        # FP8-specific descale/scale/amax tensors
        if args.data_type == "fp8":
            descale_q_gpu = torch.ones(1, 1, 1, 1, dtype=torch.float, device=device)
            descale_k_gpu = torch.ones(1, 1, 1, 1, dtype=torch.float, device=device)
            descale_v_gpu = torch.ones(1, 1, 1, 1, dtype=torch.float, device=device)
            descale_s_gpu = torch.ones(1, 1, 1, 1, dtype=torch.float, device=device)
            descale_o_gpu = torch.ones(1, 1, 1, 1, dtype=torch.float, device=device)
            descale_dO_gpu = torch.ones(1, 1, 1, 1, dtype=torch.float, device=device)
            descale_dP_gpu = torch.ones(1, 1, 1, 1, dtype=torch.float, device=device)
            scale_s_gpu = torch.ones(1, 1, 1, 1, dtype=torch.float, device=device)
            scale_o_gpu = torch.ones(1, 1, 1, 1, dtype=torch.float, device=device)
            scale_dQ_gpu = torch.ones(1, 1, 1, 1, dtype=torch.float, device=device)
            scale_dK_gpu = torch.ones(1, 1, 1, 1, dtype=torch.float, device=device)
            scale_dV_gpu = torch.ones(1, 1, 1, 1, dtype=torch.float, device=device)
            scale_dP_gpu = torch.ones(1, 1, 1, 1, dtype=torch.float, device=device)
            amax_s_gpu = torch.zeros(1, 1, 1, 1, dtype=torch.float, device=device)
            amax_o_gpu = torch.zeros(1, 1, 1, 1, dtype=torch.float, device=device)
            amax_dQ_gpu = torch.zeros(1, 1, 1, 1, dtype=torch.float, device=device)
            amax_dK_gpu = torch.zeros(1, 1, 1, 1, dtype=torch.float, device=device)
            amax_dV_gpu = torch.zeros(1, 1, 1, 1, dtype=torch.float, device=device)
            amax_dP_gpu = torch.zeros(1, 1, 1, 1, dtype=torch.float, device=device)

        if args.data_type == "mxfp8":
            block_size = 32
            s_q_padded = ceil_div(q_seqlen, 128) * 128
            s_kv_padded = ceil_div(kv_seqlen, 128) * 128
            d_qk_scale_padded = ceil_div(ceil_div(head_dim_qk, block_size), 4) * 4
            d_vo_scale_padded = ceil_div(ceil_div(head_dim_vo, block_size), 4) * 4
            d_qk_padded = ceil_div(head_dim_qk, 128) * 128
            d_vo_padded = ceil_div(head_dim_vo, 128) * 128
            s_q_scale_padded = ceil_div(ceil_div(q_seqlen, block_size), 4) * 4
            s_kv_scale_padded = ceil_div(ceil_div(kv_seqlen, block_size), 4) * 4
            amax_o_gpu = torch.zeros(1, 1, 1, 1, dtype=torch.float, device=device)
            # Forward scale factors
            sf_q_gpu = create_scale_factor_tensor_mxfp8(batch_size * num_q_heads, s_q_padded, head_dim_qk, block_size)
            sf_k_gpu = create_scale_factor_tensor_mxfp8(batch_size * num_kv_heads, s_kv_padded, head_dim_qk, block_size)
            sf_v_gpu = create_scale_factor_tensor_mxfp8(batch_size * num_kv_heads, d_vo_padded, kv_seqlen, block_size)
            # Backward-specific scale factors: transposed views for Q, K, dO
            # SF_Q_T, SF_K_T: scale along sequence dimension [b, h, s_scale_padded, d_padded]
            sf_q_t_gpu = create_scale_factor_tensor_mxfp8(batch_size * num_q_heads, d_qk_padded, q_seqlen, block_size)
            sf_k_t_gpu = create_scale_factor_tensor_mxfp8(batch_size * num_kv_heads, d_qk_padded, kv_seqlen, block_size)
            # SF_dO: scale along hidden dimension [b, h, s_padded, d_vo_scale_padded]
            sf_dO_gpu = create_scale_factor_tensor_mxfp8(batch_size * num_q_heads, s_q_padded, head_dim_vo, block_size)
            # SF_dO_T: scale along sequence dimension [b, h, s_scale_padded, d_vo_padded]
            sf_dO_t_gpu = create_scale_factor_tensor_mxfp8(batch_size * num_q_heads, d_vo_padded, q_seqlen, block_size)
            # Amax tensors for backward gradients
            amax_dQ_gpu = torch.zeros(1, 1, 1, 1, dtype=torch.float, device=device)
            amax_dK_gpu = torch.zeros(1, 1, 1, 1, dtype=torch.float, device=device)
            amax_dV_gpu = torch.zeros(1, 1, 1, 1, dtype=torch.float, device=device)

        if args.data_type == "mxfp8":
            # MXFP8 backward outputs use the same dtype as forward output
            dQuery = torch.empty(batch_size, num_q_heads, q_seqlen, head_dim_qk, dtype=output_dtype, device=device)
            dKey = torch.empty(batch_size, num_kv_heads, kv_seqlen, head_dim_qk, dtype=output_dtype, device=device)
            dValue = torch.empty(batch_size, num_kv_heads, kv_seqlen, head_dim_vo, dtype=output_dtype, device=device)
        else:
            dQuery = torch.empty_like(query)
            dKey = torch.empty_like(key)
            dValue = torch.empty_like(value)
        dOutput = torch.randn(output.shape, dtype=randn_dtype, device=device).to(target_dtype)
        stats = torch.empty(batch_size, num_q_heads, q_seqlen, 1, dtype=torch.float32, device=device)
        if is_dropout:
            dropout_seed = torch.full((1, 1, 1, 1), 123456, dtype=torch.int64, device="cuda")
            dropout_offset = torch.full((1, 1, 1, 1), 789, dtype=torch.int64, device="cuda")

        # cuDNN graph forward
        graph_fwd = cudnn.pygraph(
            io_data_type=convert_to_cudnn_type(target_dtype),
            intermediate_data_type=cudnn.data_type.FLOAT,
            compute_data_type=cudnn.data_type.FLOAT,
        )

        if is_dropout:
            seed_fwd = graph_fwd.tensor_like(dropout_seed)
            offset_fwd = graph_fwd.tensor_like(dropout_offset)
            dropout_tuple = (dropout_prob, seed_fwd, offset_fwd)

        # Compute diagonal band bounds for masking (shared by FP8 and non-FP8)
        right_bound = None if args.attn_mask == "no_mask" else 0
        left_bound = args.sliding_window_size if args.sliding_window_size else None
        diagonal_align = cudnn.diagonal_alignment.BOTTOM_RIGHT if args.attn_mask == "bottom_right" else cudnn.diagonal_alignment.TOP_LEFT

        if args.data_type == "fp8":
            q_fwd = graph_fwd.tensor_like(query).set_data_type(cudnn.data_type.FP8_E4M3)
            k_fwd = graph_fwd.tensor_like(key).set_data_type(cudnn.data_type.FP8_E4M3)
            v_fwd = graph_fwd.tensor_like(value).set_data_type(cudnn.data_type.FP8_E4M3)

            descale_q_fwd = graph_fwd.tensor_like(descale_q_gpu)
            descale_k_fwd = graph_fwd.tensor_like(descale_k_gpu)
            descale_v_fwd = graph_fwd.tensor_like(descale_v_gpu)
            descale_s_fwd = graph_fwd.tensor_like(descale_s_gpu)
            scale_s_fwd = graph_fwd.tensor_like(scale_s_gpu)
            scale_o_fwd = graph_fwd.tensor_like(scale_o_gpu)

            o_fwd, stats_fwd, amax_s_fwd, amax_o_fwd = graph_fwd.sdpa_fp8(
                q=q_fwd,
                k=k_fwd,
                v=v_fwd,
                descale_q=descale_q_fwd,
                descale_k=descale_k_fwd,
                descale_v=descale_v_fwd,
                descale_s=descale_s_fwd,
                scale_s=scale_s_fwd,
                scale_o=scale_o_fwd,
                # generate_stats=not is_infer,
                is_inference=is_infer,
                attn_scale=attn_scale,
                diagonal_alignment=diagonal_align,
                left_bound=left_bound,
                right_bound=right_bound,
                # dropout=dropout_tuple if is_dropout else None,
            )
        elif args.data_type == "mxfp8":
            q_fwd = graph_fwd.tensor_like(query)
            k_fwd = graph_fwd.tensor_like(key)
            v_fwd = graph_fwd.tensor_like(value)

            sf_q_fwd = graph_fwd.tensor(
                dim=(batch_size, num_q_heads, s_q_padded, d_qk_scale_padded),
                stride=(num_q_heads * s_q_padded * d_qk_scale_padded, s_q_padded * d_qk_scale_padded, d_qk_scale_padded, 1),
                data_type=cudnn.data_type.FP8_E8M0,
                reordering_type=cudnn.tensor_reordering.F8_128x4,
            )
            sf_k_fwd = graph_fwd.tensor(
                dim=(batch_size, num_kv_heads, s_kv_padded, d_qk_scale_padded),
                stride=(num_kv_heads * s_kv_padded * d_qk_scale_padded, s_kv_padded * d_qk_scale_padded, d_qk_scale_padded, 1),
                data_type=cudnn.data_type.FP8_E8M0,
                reordering_type=cudnn.tensor_reordering.F8_128x4,
            )
            sf_v_fwd = graph_fwd.tensor(
                dim=(batch_size, num_kv_heads, s_kv_scale_padded, d_vo_padded),
                stride=(num_kv_heads * s_kv_scale_padded * d_vo_padded, s_kv_scale_padded * d_vo_padded, 1, s_kv_scale_padded),
                data_type=cudnn.data_type.FP8_E8M0,
                reordering_type=cudnn.tensor_reordering.F8_128x4,
            )

            o_fwd, stats_fwd, amax_o_fwd = graph_fwd.sdpa_mxfp8(
                q=q_fwd,
                k=k_fwd,
                v=v_fwd,
                descale_q=sf_q_fwd,
                descale_k=sf_k_fwd,
                descale_v=sf_v_fwd,
                attn_scale=attn_scale,
                diagonal_alignment=diagonal_align,
                diagonal_band_left_bound=left_bound,
                diagonal_band_right_bound=right_bound,
                generate_stats=True,
            )
        else:
            q_fwd = graph_fwd.tensor_like(query)
            k_fwd = graph_fwd.tensor_like(key)
            v_fwd = graph_fwd.tensor_like(value)
            o_fwd, stats_fwd = graph_fwd.sdpa(
                q=q_fwd,
                k=k_fwd,
                v=v_fwd,
                # generate_stats=not is_infer,
                is_inference=is_infer,
                attn_scale=attn_scale,
                diagonal_alignment=diagonal_align,
                diagonal_band_left_bound=left_bound,
                diagonal_band_right_bound=right_bound,
                dropout=dropout_tuple if is_dropout else None,
            )

        if run_bwd:
            if args.data_type == "fp8":
                o_fwd.set_output(True).set_dim(output.size()).set_stride(output.stride()).set_data_type(cudnn.data_type.FP8_E4M3)
                (stats_fwd.set_output(True).set_dim(stats.size()).set_stride(stats.stride()).set_data_type(cudnn.data_type.FLOAT) if not is_infer else None)
            elif args.data_type == "mxfp8":
                o_fwd.set_output(True).set_dim(output.size()).set_stride(output.stride()).set_data_type(convert_to_cudnn_type(output_dtype))
                stats_fwd.set_output(True).set_dim(stats.size()).set_stride(stats.stride()).set_data_type(cudnn.data_type.FLOAT)
            else:
                o_fwd.set_output(True).set_dim(output.size()).set_stride(output.stride())
                (stats_fwd.set_output(True).set_dim(stats.size()).set_stride(stats.stride()).set_data_type(cudnn.data_type.FLOAT) if not is_infer else None)
        else:
            if args.data_type == "fp8":
                o_fwd.set_output(True).set_dim(output.size()).set_stride(output.stride()).set_data_type(cudnn.data_type.FP8_E4M3)
            elif args.data_type == "mxfp8":
                o_fwd.set_output(True).set_dim(output.size()).set_stride(output.stride()).set_data_type(convert_to_cudnn_type(output_dtype))
                stats_fwd.set_output(True).set_dim(stats.size()).set_stride(stats.stride()).set_data_type(cudnn.data_type.FLOAT)
            else:
                o_fwd.set_output(True).set_dim(output.size()).set_stride(output.stride())

        if args.data_type == "fp8":
            amax_s_fwd.set_output(True).set_dim(amax_s_gpu.size()).set_stride(amax_s_gpu.stride()).set_data_type(cudnn.data_type.FLOAT)
            amax_o_fwd.set_output(True).set_dim(amax_o_gpu.size()).set_stride(amax_o_gpu.stride()).set_data_type(cudnn.data_type.FLOAT)
        elif args.data_type == "mxfp8":
            amax_o_fwd.set_output(True).set_dim(amax_o_gpu.size()).set_stride(amax_o_gpu.stride()).set_data_type(cudnn.data_type.FLOAT)
        graph_fwd.validate()
        graph_fwd.build_operation_graph()
        graph_fwd.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK])
        graph_fwd.check_support()
        graph_fwd.build_plans()

        # If backward is requested, set up backward graph.
        if run_bwd:
            graph_bwd = cudnn.pygraph(
                io_data_type=(cudnn.data_type.FP8_E4M3 if args.data_type == "fp8" else convert_to_cudnn_type(target_dtype)),
                intermediate_data_type=cudnn.data_type.FLOAT,
                compute_data_type=cudnn.data_type.FLOAT,
            )

            stats_bwd = graph_bwd.tensor_like(stats)
            if is_dropout:
                seed_bwd = graph_bwd.tensor_like(dropout_seed)
                offset_bwd = graph_bwd.tensor_like(dropout_offset)
                dropout_tuple = (dropout_prob, seed_bwd, offset_bwd)

            if args.data_type == "fp8":
                q_bwd = graph_bwd.tensor_like(query).set_data_type(cudnn.data_type.FP8_E4M3)
                k_bwd = graph_bwd.tensor_like(key).set_data_type(cudnn.data_type.FP8_E4M3)
                v_bwd = graph_bwd.tensor_like(value).set_data_type(cudnn.data_type.FP8_E4M3)
                o_bwd = graph_bwd.tensor_like(output).set_data_type(cudnn.data_type.FP8_E4M3)
                dO_bwd = graph_bwd.tensor_like(dOutput).set_data_type(cudnn.data_type.FP8_E4M3)

                descale_q_bwd = graph_bwd.tensor_like(descale_q_gpu)
                descale_k_bwd = graph_bwd.tensor_like(descale_k_gpu)
                descale_v_bwd = graph_bwd.tensor_like(descale_v_gpu)
                descale_o_bwd = graph_bwd.tensor_like(descale_o_gpu)
                descale_dO_bwd = graph_bwd.tensor_like(descale_dO_gpu)
                descale_s_bwd = graph_bwd.tensor_like(descale_s_gpu)
                descale_dP_bwd = graph_bwd.tensor_like(descale_dP_gpu)
                scale_s_bwd = graph_bwd.tensor_like(scale_s_gpu)
                scale_dQ_bwd = graph_bwd.tensor_like(scale_dQ_gpu)
                scale_dK_bwd = graph_bwd.tensor_like(scale_dK_gpu)
                scale_dV_bwd = graph_bwd.tensor_like(scale_dV_gpu)
                scale_dP_bwd = graph_bwd.tensor_like(scale_dP_gpu)

                (
                    dQ_bwd,
                    dK_bwd,
                    dV_bwd,
                    amax_dQ_bwd,
                    amax_dK_bwd,
                    amax_dV_bwd,
                    amax_dP_bwd,
                ) = graph_bwd.sdpa_fp8_backward(
                    q=q_bwd,
                    k=k_bwd,
                    v=v_bwd,
                    o=o_bwd,
                    dO=dO_bwd,
                    stats=stats_bwd,
                    descale_q=descale_q_bwd,
                    descale_k=descale_k_bwd,
                    descale_v=descale_v_bwd,
                    descale_o=descale_o_bwd,
                    descale_dO=descale_dO_bwd,
                    descale_s=descale_s_bwd,
                    descale_dP=descale_dP_bwd,
                    scale_s=scale_s_bwd,
                    scale_dQ=scale_dQ_bwd,
                    scale_dK=scale_dK_bwd,
                    scale_dV=scale_dV_bwd,
                    scale_dP=scale_dP_bwd,
                    attn_scale=attn_scale,
                    use_causal_mask=(args.attn_mask == "top_left"),
                    use_causal_mask_bottom_right=(args.attn_mask == "bottom_right"),
                    dropout=dropout_tuple if is_dropout else None,
                    use_deterministic_algorithm=args.deterministic_bwd,
                )
            elif args.data_type == "mxfp8":
                # MXFP8 backward requires transposed tensor views and scale factors
                # Q, K, V in FP8_E4M3
                q_bwd = graph_bwd.tensor_like(query)
                k_bwd = graph_bwd.tensor_like(key)
                v_bwd = graph_bwd.tensor_like(value)

                # Create transposed tensor views (separate storage with transposed strides)
                q_t_bwd = graph_bwd.tensor_like(query)
                k_t_bwd = graph_bwd.tensor_like(key)

                # O in BFLOAT16 (forward output)
                o_bwd = graph_bwd.tensor(
                    dim=output.size(),
                    stride=output.stride(),
                    data_type=cudnn.data_type.BFLOAT16,
                )
                # dO in FP8 and BFLOAT16
                dO_bwd = graph_bwd.tensor_like(dOutput)
                dO_t_bwd = graph_bwd.tensor_like(dOutput)
                dO_f16_bwd = graph_bwd.tensor(
                    dim=dOutput.size(),
                    stride=dOutput.stride(),
                    data_type=cudnn.data_type.BFLOAT16,
                )

                # Scale factor tensors with E8M0 dtype and F8_128x4 reordering
                # SF_Q: [b, h_q, s_q_padded, d_qk_scale_padded]
                sf_q_bwd = graph_bwd.tensor(
                    dim=(batch_size, num_q_heads, s_q_padded, d_qk_scale_padded),
                    stride=(num_q_heads * s_q_padded * d_qk_scale_padded, s_q_padded * d_qk_scale_padded, d_qk_scale_padded, 1),
                    data_type=cudnn.data_type.FP8_E8M0,
                    reordering_type=cudnn.tensor_reordering.F8_128x4,
                )
                # SF_Q_T: [b, h_q, s_q_scale_padded, d_qk_padded]
                sf_q_t_bwd = graph_bwd.tensor(
                    dim=(batch_size, num_q_heads, s_q_scale_padded, d_qk_padded),
                    stride=(num_q_heads * s_q_scale_padded * d_qk_padded, s_q_scale_padded * d_qk_padded, 1, s_q_scale_padded),
                    data_type=cudnn.data_type.FP8_E8M0,
                    reordering_type=cudnn.tensor_reordering.F8_128x4,
                )
                # SF_K: [b, h_kv, s_kv_padded, d_qk_scale_padded]
                sf_k_bwd = graph_bwd.tensor(
                    dim=(batch_size, num_kv_heads, s_kv_padded, d_qk_scale_padded),
                    stride=(num_kv_heads * s_kv_padded * d_qk_scale_padded, s_kv_padded * d_qk_scale_padded, d_qk_scale_padded, 1),
                    data_type=cudnn.data_type.FP8_E8M0,
                    reordering_type=cudnn.tensor_reordering.F8_128x4,
                )
                # SF_K_T: [b, h_kv, s_kv_scale_padded, d_qk_padded]
                sf_k_t_bwd = graph_bwd.tensor(
                    dim=(batch_size, num_kv_heads, s_kv_scale_padded, d_qk_padded),
                    stride=(num_kv_heads * s_kv_scale_padded * d_qk_padded, s_kv_scale_padded * d_qk_padded, 1, s_kv_scale_padded),
                    data_type=cudnn.data_type.FP8_E8M0,
                    reordering_type=cudnn.tensor_reordering.F8_128x4,
                )
                # SF_V: [b, h_kv, s_kv_padded, d_vo_scale_padded]
                sf_v_bwd = graph_bwd.tensor(
                    dim=(batch_size, num_kv_heads, s_kv_padded, d_vo_scale_padded),
                    stride=(num_kv_heads * s_kv_padded * d_vo_scale_padded, s_kv_padded * d_vo_scale_padded, d_vo_scale_padded, 1),
                    data_type=cudnn.data_type.FP8_E8M0,
                    reordering_type=cudnn.tensor_reordering.F8_128x4,
                )
                # SF_dO: [b, h_q, s_q_padded, d_vo_scale_padded]
                sf_dO_bwd = graph_bwd.tensor(
                    dim=(batch_size, num_q_heads, s_q_padded, d_vo_scale_padded),
                    stride=(num_q_heads * s_q_padded * d_vo_scale_padded, s_q_padded * d_vo_scale_padded, d_vo_scale_padded, 1),
                    data_type=cudnn.data_type.FP8_E8M0,
                    reordering_type=cudnn.tensor_reordering.F8_128x4,
                )
                # SF_dO_T: [b, h_q, s_q_scale_padded, d_vo_padded]
                sf_dO_t_bwd = graph_bwd.tensor(
                    dim=(batch_size, num_q_heads, s_q_scale_padded, d_vo_padded),
                    stride=(num_q_heads * s_q_scale_padded * d_vo_padded, s_q_scale_padded * d_vo_padded, 1, s_q_scale_padded),
                    data_type=cudnn.data_type.FP8_E8M0,
                    reordering_type=cudnn.tensor_reordering.F8_128x4,
                )

                (
                    dQ_bwd,
                    dK_bwd,
                    dV_bwd,
                    amax_dQ_bwd,
                    amax_dK_bwd,
                    amax_dV_bwd,
                ) = graph_bwd.sdpa_mxfp8_backward(
                    q=q_bwd,
                    q_T=q_t_bwd,
                    k=k_bwd,
                    k_T=k_t_bwd,
                    v=v_bwd,
                    o_f16=o_bwd,
                    dO_f16=dO_f16_bwd,
                    dO=dO_bwd,
                    dO_T=dO_t_bwd,
                    stats=stats_bwd,
                    descale_q=sf_q_bwd,
                    descale_q_T=sf_q_t_bwd,
                    descale_k=sf_k_bwd,
                    descale_k_T=sf_k_t_bwd,
                    descale_v=sf_v_bwd,
                    descale_dO=sf_dO_bwd,
                    descale_dO_T=sf_dO_t_bwd,
                    attn_scale=attn_scale,
                    use_causal_mask=(args.attn_mask == "top_left"),
                    use_causal_mask_bottom_right=(args.attn_mask == "bottom_right"),
                    diagonal_alignment=diagonal_align,
                    use_deterministic_algorithm=args.deterministic_bwd,
                )
            else:
                q_bwd = graph_bwd.tensor_like(query)
                k_bwd = graph_bwd.tensor_like(key)
                v_bwd = graph_bwd.tensor_like(value)
                o_bwd = graph_bwd.tensor_like(output)
                dO_bwd = graph_bwd.tensor_like(dOutput)

                dQ_bwd, dK_bwd, dV_bwd = graph_bwd.sdpa_backward(
                    q=q_bwd,
                    k=k_bwd,
                    v=v_bwd,
                    o=o_bwd,
                    dO=dO_bwd,
                    stats=stats_bwd,
                    attn_scale=attn_scale,
                    diagonal_alignment=diagonal_align,
                    diagonal_band_left_bound=left_bound,
                    diagonal_band_right_bound=right_bound,
                    dropout=dropout_tuple if is_dropout else None,
                    use_deterministic_algorithm=args.deterministic_bwd,
                )

            if args.data_type == "fp8":
                dQ_bwd.set_output(True).set_dim(dQuery.size()).set_stride(dQuery.stride()).set_data_type(cudnn.data_type.FP8_E4M3)
                dK_bwd.set_output(True).set_dim(dKey.size()).set_stride(dKey.stride()).set_data_type(cudnn.data_type.FP8_E4M3)
                dV_bwd.set_output(True).set_dim(dValue.size()).set_stride(dValue.stride()).set_data_type(cudnn.data_type.FP8_E4M3)
                amax_dQ_bwd.set_output(True).set_dim(amax_dQ_gpu.size()).set_stride(amax_dQ_gpu.stride()).set_data_type(cudnn.data_type.FLOAT)
                amax_dK_bwd.set_output(True).set_dim(amax_dK_gpu.size()).set_stride(amax_dK_gpu.stride()).set_data_type(cudnn.data_type.FLOAT)
                amax_dV_bwd.set_output(True).set_dim(amax_dV_gpu.size()).set_stride(amax_dV_gpu.stride()).set_data_type(cudnn.data_type.FLOAT)
                amax_dP_bwd.set_output(True).set_dim(amax_dP_gpu.size()).set_stride(amax_dP_gpu.stride()).set_data_type(cudnn.data_type.FLOAT)
            elif args.data_type == "mxfp8":
                # MXFP8 backward outputs in BFLOAT16
                dQ_bwd.set_output(True).set_dim(dQuery.size()).set_stride(dQuery.stride()).set_data_type(cudnn.data_type.BFLOAT16)
                dK_bwd.set_output(True).set_dim(dKey.size()).set_stride(dKey.stride()).set_data_type(cudnn.data_type.BFLOAT16)
                dV_bwd.set_output(True).set_dim(dValue.size()).set_stride(dValue.stride()).set_data_type(cudnn.data_type.BFLOAT16)
                amax_dQ_bwd.set_output(True).set_dim(amax_dQ_gpu.size()).set_stride(amax_dQ_gpu.stride()).set_data_type(cudnn.data_type.FLOAT)
                amax_dK_bwd.set_output(True).set_dim(amax_dK_gpu.size()).set_stride(amax_dK_gpu.stride()).set_data_type(cudnn.data_type.FLOAT)
                amax_dV_bwd.set_output(True).set_dim(amax_dV_gpu.size()).set_stride(amax_dV_gpu.stride()).set_data_type(cudnn.data_type.FLOAT)
            else:
                dQ_bwd.set_output(True).set_dim(dQuery.size()).set_stride(dQuery.stride())
                dK_bwd.set_output(True).set_dim(dKey.size()).set_stride(dKey.stride())
                dV_bwd.set_output(True).set_dim(dValue.size()).set_stride(dValue.stride())

            graph_bwd.validate()
            graph_bwd.build_operation_graph()
            graph_bwd.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK])
            graph_bwd.check_support()
            graph_bwd.build_plans()

            if args.data_type == "fp8":
                variant_pack_fwd = {
                    q_fwd: query,
                    k_fwd: key,
                    v_fwd: value,
                    o_fwd: output,
                    stats_fwd: stats,
                    descale_q_fwd: descale_q_gpu,
                    descale_k_fwd: descale_k_gpu,
                    descale_v_fwd: descale_v_gpu,
                    descale_s_fwd: descale_s_gpu,
                    scale_s_fwd: scale_s_gpu,
                    scale_o_fwd: scale_o_gpu,
                    amax_s_fwd: amax_s_gpu,
                    amax_o_fwd: amax_o_gpu,
                }

                variant_pack_bwd = {
                    q_bwd: query,
                    k_bwd: key,
                    v_bwd: value,
                    o_bwd: output,
                    dQ_bwd: dQuery,
                    dK_bwd: dKey,
                    dV_bwd: dValue,
                    dO_bwd: dOutput,
                    stats_bwd: stats,
                    descale_q_bwd: descale_q_gpu,
                    descale_k_bwd: descale_k_gpu,
                    descale_v_bwd: descale_v_gpu,
                    descale_o_bwd: descale_o_gpu,
                    descale_s_bwd: descale_s_gpu,
                    descale_dP_bwd: descale_dP_gpu,
                    descale_dO_bwd: descale_dO_gpu,
                    scale_s_bwd: scale_s_gpu,
                    scale_dQ_bwd: scale_dQ_gpu,
                    scale_dK_bwd: scale_dK_gpu,
                    scale_dV_bwd: scale_dV_gpu,
                    scale_dP_bwd: scale_dP_gpu,
                    amax_dQ_bwd: amax_dQ_gpu,
                    amax_dK_bwd: amax_dK_gpu,
                    amax_dV_bwd: amax_dV_gpu,
                    amax_dP_bwd: amax_dP_gpu,
                }

                workspace = torch.empty(
                    max(graph_fwd.get_workspace_size(), graph_bwd.get_workspace_size()),
                    device="cuda",
                    dtype=torch.uint8,
                )
            elif args.data_type == "mxfp8":
                # Create additional tensors needed for MXFP8 backward
                # Output and dOutput in output_dtype format for backward
                output_f16 = output.to(output_dtype)
                dOutput_f16 = dOutput.float().to(output_dtype)
                # Transposed tensors (use same data, just for passing to graph)
                query_t = query.clone()
                key_t = key.clone()
                dOutput_t = dOutput.clone()

                variant_pack_fwd = {
                    q_fwd: query,
                    k_fwd: key,
                    v_fwd: value,
                    o_fwd: output,
                    stats_fwd: stats,
                    sf_q_fwd: sf_q_gpu,
                    sf_k_fwd: sf_k_gpu,
                    sf_v_fwd: sf_v_gpu,
                    amax_o_fwd: amax_o_gpu,
                }
                variant_pack_bwd = {
                    q_bwd: query,
                    q_t_bwd: query_t,
                    k_bwd: key,
                    k_t_bwd: key_t,
                    v_bwd: value,
                    o_bwd: output_f16,
                    dO_bwd: dOutput,
                    dO_t_bwd: dOutput_t,
                    dO_f16_bwd: dOutput_f16,
                    stats_bwd: stats,
                    sf_q_bwd: sf_q_gpu,
                    sf_q_t_bwd: sf_q_t_gpu,
                    sf_k_bwd: sf_k_gpu,
                    sf_k_t_bwd: sf_k_t_gpu,
                    sf_v_bwd: sf_v_gpu,
                    sf_dO_bwd: sf_dO_gpu,
                    sf_dO_t_bwd: sf_dO_t_gpu,
                    dQ_bwd: dQuery,
                    dK_bwd: dKey,
                    dV_bwd: dValue,
                    amax_dQ_bwd: amax_dQ_gpu,
                    amax_dK_bwd: amax_dK_gpu,
                    amax_dV_bwd: amax_dV_gpu,
                }
                workspace = torch.empty(
                    max(graph_fwd.get_workspace_size(), graph_bwd.get_workspace_size()),
                    device="cuda",
                    dtype=torch.uint8,
                )
            else:
                variant_pack_fwd = {
                    q_fwd: query,
                    k_fwd: key,
                    v_fwd: value,
                    o_fwd: output,
                    stats_fwd: stats,
                }
                variant_pack_bwd = {
                    q_bwd: query,
                    k_bwd: key,
                    v_bwd: value,
                    o_bwd: output,
                    dO_bwd: dOutput,
                    stats_bwd: stats,
                    dQ_bwd: dQuery,
                    dK_bwd: dKey,
                    dV_bwd: dValue,
                }
                workspace = torch.empty(
                    max(graph_fwd.get_workspace_size(), graph_bwd.get_workspace_size()),
                    device="cuda",
                    dtype=torch.uint8,
                )
        else:
            if args.data_type == "fp8":
                variant_pack_fwd = {
                    q_fwd: query,
                    k_fwd: key,
                    v_fwd: value,
                    o_fwd: output,
                    stats_fwd: stats,
                    descale_q_fwd: descale_q_gpu,
                    descale_k_fwd: descale_k_gpu,
                    descale_v_fwd: descale_v_gpu,
                    descale_s_fwd: descale_s_gpu,
                    scale_s_fwd: scale_s_gpu,
                    scale_o_fwd: scale_o_gpu,
                    amax_s_fwd: amax_s_gpu,
                    amax_o_fwd: amax_o_gpu,
                }
                workspace = torch.empty(graph_fwd.get_workspace_size(), device="cuda", dtype=torch.uint8)
            elif args.data_type == "mxfp8":
                variant_pack_fwd = {
                    q_fwd: query,
                    k_fwd: key,
                    v_fwd: value,
                    o_fwd: output,
                    stats_fwd: stats,
                    sf_q_fwd: sf_q_gpu,
                    sf_k_fwd: sf_k_gpu,
                    sf_v_fwd: sf_v_gpu,
                    amax_o_fwd: amax_o_gpu,
                }
                workspace = torch.empty(graph_fwd.get_workspace_size(), device="cuda", dtype=torch.uint8)
            else:
                variant_pack_fwd = {
                    q_fwd: query,
                    k_fwd: key,
                    v_fwd: value,
                    o_fwd: output,
                }
                workspace = torch.empty(graph_fwd.get_workspace_size(), device="cuda", dtype=torch.uint8)
        if is_dropout:
            variant_pack_fwd[seed_fwd] = dropout_seed
            variant_pack_fwd[offset_fwd] = dropout_offset
            if run_bwd:
                variant_pack_bwd[seed_bwd] = dropout_seed
                variant_pack_bwd[offset_bwd] = dropout_offset
    ## Done setting up cuDNN graph.

    # For backends MATH, EFFICIENT_ATTENTION, CUDNN_ATTENTION, PYTORCH_FLASH_ATTENTION
    if args.attn_mask == "bottom_right" and args.sdpa_backend.startswith("pyt_"):
        raise ValueError("bottom_right attention mask is only supported with the 'cudnn' backend, not PyTorch backends")

    def pyt_backend_sdpa(query, key, value, backend):
        with sdpa_kernel(backends=[backend]):
            return torch.nn.functional.scaled_dot_product_attention(
                query,
                key,
                value,
                enable_gqa=enable_gqa,
                is_causal=args.attn_mask in ("top_left", "bottom_right"),
            )

    if args.sdpa_backend == "flash_attention":
        import flash_attn
        from flash_attn import flash_attn_func

        # Flash Attention Native
        def flash_attention_sdpa(query, key, value):
            window_size = (args.sliding_window_size, 0) if args.sliding_window_size else (None, None)
            return flash_attn_func(
                query,
                key,
                value,
                causal=args.attn_mask != "no_mask",
                window_size=window_size,
                deterministic=args.deterministic_bwd,
            )

    if args.sdpa_backend == "flash_attention_3":
        import flash_attn_interface

        def flash_attention_3_sdpa(query, key, value):
            window_size = (args.sliding_window_size, 0) if args.sliding_window_size else (None, None)
            output, _ = flash_attn_interface.flash_attn_func(
                query,
                key,
                value,
                causal=args.attn_mask != "no_mask",
                window_size=window_size,
                deterministic=args.deterministic_bwd,
            )
            return output

    if args.sdpa_backend == "flash_attention_4" or (not args.skip_ref):
        import flash_attn.cute.interface as flash_attn_interface

        def flash_attention_4_sdpa(query, key, value):
            window_size = (args.sliding_window_size, 0) if args.sliding_window_size else (None, None)
            kwargs = dict(
                causal=args.attn_mask != "no_mask",
                window_size=window_size,
                deterministic=args.deterministic_bwd,
            )
            if args.fa4_num_splits is not None:
                kwargs["num_splits"] = args.fa4_num_splits
            output, _ = flash_attn_interface.flash_attn_func(
                query,
                key,
                value,
                **kwargs,
            )
            return output

    def get_sdpa_function(backend):
        if backend == "pyt_math":
            return functools.partial(pyt_backend_sdpa, backend=SDPBackend.MATH)
        elif backend == "pyt_efficient_attention":
            return functools.partial(pyt_backend_sdpa, backend=SDPBackend.EFFICIENT_ATTENTION)
        elif backend == "pyt_flash_attention":
            return functools.partial(pyt_backend_sdpa, backend=SDPBackend.FLASH_ATTENTION)
        elif backend == "pyt_cudnn":
            return functools.partial(pyt_backend_sdpa, backend=SDPBackend.CUDNN_ATTENTION)
        elif backend == "flash_attention":
            return flash_attention_sdpa
        elif backend == "flash_attention_3":
            return flash_attention_3_sdpa
        elif backend == "flash_attention_4":
            return flash_attention_4_sdpa
        elif backend == "cudnn":
            return None  # Will be set up separately
        else:
            raise ValueError(f"Invalid backend: {backend}")

    # Util function for addressing different qkv formats for each backend
    def preprocess_qkv(query, key, value, backend):
        if backend.startswith("pyt_") or backend == "cudnn":
            return query, key, value
        elif backend.startswith("flash_attention"):
            query = torch.swapaxes(query, 1, 2)
            key = torch.swapaxes(key, 1, 2)
            value = torch.swapaxes(value, 1, 2)
            return query, key, value
        else:
            raise ValueError(f"Invalid backend: {backend}")

    # Util function addressing different qkvo formats for each backend
    def postprocess_qkvo(query, key, value, output, backend):
        if backend.startswith("pyt_") or backend == "cudnn":
            return query, key, value, output
        elif backend.startswith("flash_attention"):
            output = torch.swapaxes(output, 1, 2)
            query = torch.swapaxes(query, 1, 2)
            key = torch.swapaxes(key, 1, 2)
            value = torch.swapaxes(value, 1, 2)
            return query, key, value, output
        else:
            raise ValueError(f"Invalid backend: {backend}")

    def postprocess_dqdkdvdo(dQuery, dKey, dValue, dOutput, backend):
        if backend.startswith("pyt_") or backend == "cudnn":
            return dQuery, dKey, dValue, dOutput
        elif backend.startswith("flash_attention"):
            dQuery = torch.swapaxes(dQuery, 1, 2)
            dKey = torch.swapaxes(dKey, 1, 2)
            dValue = torch.swapaxes(dValue, 1, 2)
            dOutput = torch.swapaxes(dOutput, 1, 2)
            return dQuery, dKey, dValue, dOutput
        else:
            raise ValueError(f"Invalid backend: {backend}")

    # Util functions for calculating flops and tflops/s achieved
    def flops(
        batch_size,
        q_seqlen,
        kv_seqlen,
        head_dim_qk,
        head_dim_vo,
        num_q_heads,
        attn_mask,
        mode="fwd",
        sliding_window_size=None,
    ):
        assert mode in ["fwd", "bwd", "fwd_bwd"]

        if attn_mask == "no_mask":
            num_nonmasked_elems = q_seqlen * kv_seqlen
        elif attn_mask in ("top_left", "bottom_right"):
            if sliding_window_size is not None:
                # With sliding window: each query attends to at most W keys
                # (left_bound is exclusive, right_bound is inclusive)
                # For positions 0 to W-1: 1, 2, ..., W keys (partial window)
                # For positions W to S-1: W keys each (full window)
                W = sliding_window_size
                S = min(q_seqlen, kv_seqlen)
                if S <= W:
                    # Sequence shorter than window, use full causal
                    num_nonmasked_elems = S * (S + 1) // 2
                else:
                    # Partial window (first W positions) + full window (remaining positions)
                    num_nonmasked_elems = W * (W + 1) // 2 + (S - W) * W
            else:
                num_nonmasked_elems = torch.tril(torch.ones((q_seqlen, kv_seqlen), dtype=torch.bool)).sum()
        else:
            raise ValueError(f"Unknown attn_mask: {attn_mask}")
        # BMM FLOPs: 2 * M * N * K.
        # Here, M*N = num_nonmasked_elems per head; add batch_size * num_q_heads multiplier.
        # Forward: 2 BMMs => (1 x head_dim_qk) + (1 x head_dim_vo)
        # Backward: 5 BMMs => (3 x head_dim_qk) + (2 x head_dim_vo)
        base = batch_size * num_q_heads * num_nonmasked_elems * 2
        if mode == "fwd":
            result = base * (head_dim_qk + head_dim_vo)
        elif mode == "bwd":
            result = base * (3 * head_dim_qk + 2 * head_dim_vo)
        else:  # fwd_bwd
            result = base * (4 * head_dim_qk + 3 * head_dim_vo)
        return result

    def tflops_per_sec(
        batch_size,
        q_seqlen,
        kv_seqlen,
        head_dim_qk,
        head_dim_vo,
        num_q_heads,
        attn_mask,
        time,
        mode="fwd",
        sliding_window_size=None,
    ):
        assert mode in ["fwd", "bwd", "fwd_bwd"]
        f = flops(
            batch_size,
            q_seqlen,
            kv_seqlen,
            head_dim_qk,
            head_dim_vo,
            num_q_heads,
            attn_mask,
            mode,
            sliding_window_size,
        )
        return f / time / 1e9 if not math.isnan(time) else 0.0  # Assume time is in msec

    ###### Done setting up SDPA function for each backend #######
    #############################################################

    ###### SDPA Benchmark -- Run ######
    ## Print System Info
    if args.verbose:
        print(f"[INFO] {torch.__version__ = }")
        print(f"[INFO] {torch.version.cuda = }")
        print(f"[INFO] {torch.cuda.is_available() = }")
        print(f"[INFO] {torch.cuda.device_count() = }")
        print(f"[INFO] {torch.cuda.current_device() = }")
        print(f"[INFO] {torch.cuda.get_device_name(torch.cuda.current_device()) = }")
        if args.sdpa_backend == "pyt_cudnn":
            print(f"[INFO] {torch.backends.cudnn.version() = }")
            print(f"[INFO] {torch.backends.cudnn.enabled = }")
        elif args.sdpa_backend == "flash_attention":
            print(f"[INFO] {flash_attn.__version__ = }")

    forward_times = []
    backward_times = []
    forward_diffs = []

    total_iters = num_iters + dry_run_iters

    first_error = True  # For suppressing error message beyond first error
    sdpa_function = get_sdpa_function(args.sdpa_backend)

    # Sample SM clock throughout the benchmark window so SOL% uses the actual
    # boost clock the kernel ran at rather than nvml's (often-stale) max.
    _clock_sampler = _SmClockSampler()
    _clock_sampler.start()
    for i in range(total_iters):
        # FP8/MXFP8 needs randn in bfloat16 then convert (randn doesn't support fp8 well)
        randn_dtype = torch.bfloat16 if args.data_type in ("fp8", "mxfp8") else target_dtype
        query = (
            torch.randn(batch_size, q_seqlen, num_q_heads, head_dim_qk, dtype=randn_dtype, device=device, requires_grad=True).to(target_dtype).transpose(1, 2)
        )
        key = (
            torch.randn(batch_size, kv_seqlen, num_kv_heads, head_dim_qk, dtype=randn_dtype, device=device, requires_grad=True).to(target_dtype).transpose(1, 2)
        )
        value = (
            torch.randn(batch_size, kv_seqlen, num_kv_heads, head_dim_vo, dtype=randn_dtype, device=device, requires_grad=True).to(target_dtype).transpose(1, 2)
        )

        # cuDNN-specific FP8 descale/scale/amax tensors
        if args.data_type == "fp8" and args.sdpa_backend == "cudnn":
            descale_q_gpu = torch.ones(1, 1, 1, 1, dtype=torch.float, device=device)
            descale_k_gpu = torch.ones(1, 1, 1, 1, dtype=torch.float, device=device)
            descale_v_gpu = torch.ones(1, 1, 1, 1, dtype=torch.float, device=device)
            descale_s_gpu = torch.ones(1, 1, 1, 1, dtype=torch.float, device=device)
            descale_o_gpu = torch.ones(1, 1, 1, 1, dtype=torch.float, device=device)
            descale_dO_gpu = torch.ones(1, 1, 1, 1, dtype=torch.float, device=device)
            descale_dP_gpu = torch.ones(1, 1, 1, 1, dtype=torch.float, device=device)
            scale_s_gpu = torch.ones(1, 1, 1, 1, dtype=torch.float, device=device)
            scale_o_gpu = torch.ones(1, 1, 1, 1, dtype=torch.float, device=device)
            scale_dQ_gpu = torch.ones(1, 1, 1, 1, dtype=torch.float, device=device)
            scale_dK_gpu = torch.ones(1, 1, 1, 1, dtype=torch.float, device=device)
            scale_dV_gpu = torch.ones(1, 1, 1, 1, dtype=torch.float, device=device)
            scale_dP_gpu = torch.ones(1, 1, 1, 1, dtype=torch.float, device=device)
            amax_s_gpu = torch.zeros(1, 1, 1, 1, dtype=torch.float, device=device)
            amax_o_gpu = torch.zeros(1, 1, 1, 1, dtype=torch.float, device=device)
            amax_dQ_gpu = torch.zeros(1, 1, 1, 1, dtype=torch.float, device=device)
            amax_dK_gpu = torch.zeros(1, 1, 1, 1, dtype=torch.float, device=device)
            amax_dV_gpu = torch.zeros(1, 1, 1, 1, dtype=torch.float, device=device)
            amax_dP_gpu = torch.zeros(1, 1, 1, 1, dtype=torch.float, device=device)

        query, key, value = preprocess_qkv(query, key, value, args.sdpa_backend)
        dOutput = torch.randn(*query.shape[:-1], head_dim_vo, dtype=randn_dtype, device=device).to(target_dtype)

        if args.sdpa_backend == "cudnn":
            if args.data_type == "mxfp8":
                output = torch.empty(batch_size, q_seqlen, num_q_heads, head_dim_vo, dtype=output_dtype, device=device).transpose(1, 2)
                # MXFP8 backward outputs use the same dtype as forward output
                dQuery = torch.empty(batch_size, num_q_heads, q_seqlen, head_dim_qk, dtype=output_dtype, device=device)
                dKey = torch.empty(batch_size, num_kv_heads, kv_seqlen, head_dim_qk, dtype=output_dtype, device=device)
                dValue = torch.empty(batch_size, num_kv_heads, kv_seqlen, head_dim_vo, dtype=output_dtype, device=device)
            else:
                output = torch.empty(batch_size, q_seqlen, num_q_heads, head_dim_vo, dtype=target_dtype, device=device).transpose(1, 2)
                dQuery = torch.empty_like(query)
                dKey = torch.empty_like(key)
                dValue = torch.empty_like(value)
            stats = torch.empty(batch_size, num_q_heads, q_seqlen, 1, dtype=torch.float32, device=device)
            if is_dropout:
                dropout_seed = torch.full((1, 1, 1, 1), 123456, dtype=torch.int64, device="cuda")
                dropout_offset = torch.full((1, 1, 1, 1), 789, dtype=torch.int64, device="cuda")

            # Only variant pack and workspace need to be updated for each iteration.
            if run_bwd:
                if args.data_type == "fp8":
                    variant_pack_fwd = {
                        q_fwd: query,
                        k_fwd: key,
                        v_fwd: value,
                        o_fwd: output,
                        stats_fwd: stats,
                        descale_q_fwd: descale_q_gpu,
                        descale_k_fwd: descale_k_gpu,
                        descale_v_fwd: descale_v_gpu,
                        descale_s_fwd: descale_s_gpu,
                        scale_s_fwd: scale_s_gpu,
                        scale_o_fwd: scale_o_gpu,
                        amax_s_fwd: amax_s_gpu,
                        amax_o_fwd: amax_o_gpu,
                    }
                    variant_pack_bwd = {
                        q_bwd: query,
                        k_bwd: key,
                        v_bwd: value,
                        o_bwd: output,
                        dQ_bwd: dQuery,
                        dK_bwd: dKey,
                        dV_bwd: dValue,
                        dO_bwd: dOutput,
                        stats_bwd: stats,
                        descale_q_bwd: descale_q_gpu,
                        descale_k_bwd: descale_k_gpu,
                        descale_v_bwd: descale_v_gpu,
                        descale_o_bwd: descale_o_gpu,
                        descale_s_bwd: descale_s_gpu,
                        descale_dP_bwd: descale_dP_gpu,
                        descale_dO_bwd: descale_dO_gpu,
                        scale_s_bwd: scale_s_gpu,
                        scale_dQ_bwd: scale_dQ_gpu,
                        scale_dK_bwd: scale_dK_gpu,
                        scale_dV_bwd: scale_dV_gpu,
                        scale_dP_bwd: scale_dP_gpu,
                        amax_dQ_bwd: amax_dQ_gpu,
                        amax_dK_bwd: amax_dK_gpu,
                        amax_dV_bwd: amax_dV_gpu,
                        amax_dP_bwd: amax_dP_gpu,
                    }
                elif args.data_type == "mxfp8":
                    # MXFP8 backward needs additional tensors
                    output_f16 = output.to(output_dtype)
                    dOutput_f16 = dOutput.float().to(output_dtype)
                    query_t = query.clone()
                    key_t = key.clone()
                    dOutput_t = dOutput.clone()

                    variant_pack_fwd = {
                        q_fwd: query,
                        k_fwd: key,
                        v_fwd: value,
                        o_fwd: output,
                        stats_fwd: stats,
                        sf_q_fwd: sf_q_gpu,
                        sf_k_fwd: sf_k_gpu,
                        sf_v_fwd: sf_v_gpu,
                        amax_o_fwd: amax_o_gpu,
                    }
                    variant_pack_bwd = {
                        q_bwd: query,
                        q_t_bwd: query_t,
                        k_bwd: key,
                        k_t_bwd: key_t,
                        v_bwd: value,
                        o_bwd: output_f16,
                        dO_bwd: dOutput,
                        dO_t_bwd: dOutput_t,
                        dO_f16_bwd: dOutput_f16,
                        stats_bwd: stats,
                        sf_q_bwd: sf_q_gpu,
                        sf_q_t_bwd: sf_q_t_gpu,
                        sf_k_bwd: sf_k_gpu,
                        sf_k_t_bwd: sf_k_t_gpu,
                        sf_v_bwd: sf_v_gpu,
                        sf_dO_bwd: sf_dO_gpu,
                        sf_dO_t_bwd: sf_dO_t_gpu,
                        dQ_bwd: dQuery,
                        dK_bwd: dKey,
                        dV_bwd: dValue,
                        amax_dQ_bwd: amax_dQ_gpu,
                        amax_dK_bwd: amax_dK_gpu,
                        amax_dV_bwd: amax_dV_gpu,
                    }
                else:
                    variant_pack_fwd = {
                        q_fwd: query,
                        k_fwd: key,
                        v_fwd: value,
                        o_fwd: output,
                        stats_fwd: stats,
                    }
                    variant_pack_bwd = {
                        q_bwd: query,
                        k_bwd: key,
                        v_bwd: value,
                        o_bwd: output,
                        dO_bwd: dOutput,
                        stats_bwd: stats,
                        dQ_bwd: dQuery,
                        dK_bwd: dKey,
                        dV_bwd: dValue,
                    }
                workspace = torch.empty(
                    max(graph_fwd.get_workspace_size(), graph_bwd.get_workspace_size()),
                    device="cuda",
                    dtype=torch.uint8,
                )
            else:
                if args.data_type == "fp8":
                    variant_pack_fwd = {
                        q_fwd: query,
                        k_fwd: key,
                        v_fwd: value,
                        o_fwd: output,
                        stats_fwd: stats,
                        descale_q_fwd: descale_q_gpu,
                        descale_k_fwd: descale_k_gpu,
                        descale_v_fwd: descale_v_gpu,
                        descale_s_fwd: descale_s_gpu,
                        scale_s_fwd: scale_s_gpu,
                        scale_o_fwd: scale_o_gpu,
                        amax_s_fwd: amax_s_gpu,
                        amax_o_fwd: amax_o_gpu,
                    }
                elif args.data_type == "mxfp8":
                    variant_pack_fwd = {
                        q_fwd: query,
                        k_fwd: key,
                        v_fwd: value,
                        o_fwd: output,
                        stats_fwd: stats,
                        sf_q_fwd: sf_q_gpu,
                        sf_k_fwd: sf_k_gpu,
                        sf_v_fwd: sf_v_gpu,
                        amax_o_fwd: amax_o_gpu,
                    }
                else:
                    variant_pack_fwd = {
                        q_fwd: query,
                        k_fwd: key,
                        v_fwd: value,
                        o_fwd: output,
                    }
                workspace = torch.empty(graph_fwd.get_workspace_size(), device="cuda", dtype=torch.uint8)

            if is_dropout:
                variant_pack_fwd[seed_fwd] = dropout_seed
                variant_pack_fwd[offset_fwd] = dropout_offset
                if run_bwd:
                    variant_pack_bwd[seed_bwd] = dropout_seed
                    variant_pack_bwd[offset_bwd] = dropout_offset

        l2_flush_buffer.zero_()

        # Run kernel with profiler for forward if requested, else run unprofiled to prep for backward
        if run_fwd:
            with profile(activities=[ProfilerActivity.CUDA], record_shapes=True) as prof:
                with record_function("sdpa.forward"):  # Custom marker
                    if args.sdpa_backend == "cudnn":
                        graph_fwd.execute(variant_pack_fwd, workspace)
                    else:
                        output = sdpa_function(query, key, value)
                torch.cuda.synchronize()  # Ensure all kernels finish

            # Filter profiler results by kernel name prefix
            matched_kernels = [
                item
                for item in prof.key_averages()
                if item.key.startswith("cudnn")
                or item.key.startswith("kernel_cutlass")
                or "pytorch_flash::" in item.key
                or "flash::" in item.key
                or "at::native::" in item.key
                or "cutlass3x" in item.key
                or "(anonymous namespace)::" in item.key
                or item.key.startswith("fmha_")
            ]
            if len(matched_kernels) >= 1:
                fwd_time = sum(item.device_time for item in matched_kernels) / 1000
                if i >= dry_run_iters:
                    forward_times.append(fwd_time)
        else:
            if args.sdpa_backend == "cudnn":
                graph_fwd.execute(variant_pack_fwd, workspace)
            else:
                output = sdpa_function(query, key, value)
            torch.cuda.synchronize()

        if run_bwd:
            # Run backward pass

            l2_flush_buffer.zero_()

            with profile(activities=[ProfilerActivity.CUDA], record_shapes=True) as prof:
                with record_function("sdpa.backward"):  # Custom marker
                    if args.sdpa_backend == "cudnn":
                        graph_bwd.execute(variant_pack_bwd, workspace)
                    else:
                        query.retain_grad()
                        key.retain_grad()
                        value.retain_grad()
                        output.backward(dOutput)

                        dQuery = query.grad
                        dKey = key.grad
                        dValue = value.grad

                        query.grad = None
                        key.grad = None
                        value.grad = None
                torch.cuda.synchronize()

            matched_kernels = [
                item
                for item in prof.key_averages()
                if "cudnn" in item.key
                or item.key.startswith("kernel_cutlass")
                or "pytorch_flash::" in item.key
                or "flash::" in item.key
                or "at::native::" in item.key
                or "cutlass3x" in item.key
                or "(anonymous namespace)::" in item.key
                or item.key.startswith("fmha_")
            ]
            if len(matched_kernels) >= 1:
                bwd_time = sum(item.device_time for item in matched_kernels) / 1000
                if i >= dry_run_iters:
                    backward_times.append(bwd_time)

            dQuery, dKey, dValue, dOutput = postprocess_dqdkdvdo(dQuery, dKey, dValue, dOutput, args.sdpa_backend)

        (
            query,
            key,
            value,
            output,
        ) = postprocess_qkvo(query, key, value, output, args.sdpa_backend)
        if args.data_type not in ("fp8", "mxfp8") and not args.skip_ref and run_fwd:
            try:
                output_ref = flash_attention_4_sdpa(query, key, value)
                if run_bwd:
                    query.retain_grad()
                    key.retain_grad()
                    value.retain_grad()
                    output_ref.backward(dOutput)

                    torch.testing.assert_close(dQuery, query.grad, rtol=2e-2, atol=2e-2)
                    torch.testing.assert_close(dKey, key.grad, rtol=2e-2, atol=2e-2)
                    torch.testing.assert_close(dValue, value.grad, rtol=2e-2, atol=2e-2)

                torch.testing.assert_close(output, output_ref, rtol=1e-2, atol=1e-2)
                forward_diffs.append(torch.max(torch.abs(output.detach() - output_ref.detach())).item())
            except Exception as e:
                if first_error:
                    print(
                        f"[WARN] Failed reference check. Target backend has been run, but output has not been validated. Failure may be due to incorrect output or reference function failure."
                    )
                    print(f"[WARN] See error message: {e}")
                    first_error = False
                forward_diffs.append(0.0)
        else:
            forward_diffs.append(0.0)

        if args.sdpa_backend == "cudnn":
            del query, key, value, output, dQuery, dKey, dValue, dOutput, stats
        else:
            del query, key, value, output

    _clock_sampler.stop()

    ## print results
    fwd_median_time = (
        np.median(np.array(forward_times[5:])) if len(forward_times) > 5 else (np.median(np.array(forward_times)) if len(forward_times) > 0 else 0.0)
    )
    fwd_tflops = 0.0
    if run_fwd and fwd_median_time > 0:
        fwd_tflops = tflops_per_sec(
            args.batch_size,
            args.q_seqlen,
            args.kv_seqlen,
            head_dim_qk,
            head_dim_vo,
            args.num_q_heads,
            args.attn_mask,
            fwd_median_time,
            "fwd",
            args.sliding_window_size,
        )

    bwd_median_time = (
        np.median(np.array(backward_times[5:])) if len(backward_times) > 5 else (np.median(np.array(backward_times)) if len(backward_times) > 0 else 0.0)
    )
    bwd_tflops = 0.0
    if run_bwd and bwd_median_time > 0:
        bwd_tflops = tflops_per_sec(
            args.batch_size,
            args.q_seqlen,
            args.kv_seqlen,
            head_dim_qk,
            head_dim_vo,
            args.num_q_heads,
            args.attn_mask,
            bwd_median_time,
            "bwd",
            args.sliding_window_size,
        )

    # Compute MMA SOL% using the per-arch FLOPs/clk/SM table and the actual
    # sampled boost clock observed during the benchmark window.
    _peak_mma_tflops = None
    try:
        _flops_per_clk_per_sm = _peak_flops_per_clock_per_sm(args.data_type)
        _num_sms = torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count
        _sampled_mhz = _clock_sampler.peak_mhz()
        if _flops_per_clk_per_sm is not None and _sampled_mhz is not None:
            _peak_mma_tflops = _flops_per_clk_per_sm * _num_sms * _sampled_mhz / 1e6
    except Exception:
        pass

    fwd_sol_str = f", {fwd_tflops / _peak_mma_tflops * 100:.1f}% SOL" if _peak_mma_tflops and fwd_tflops > 0 else ""
    bwd_sol_str = f", {bwd_tflops / _peak_mma_tflops * 100:.1f}% SOL" if _peak_mma_tflops and bwd_tflops > 0 else ""

    if args.format_output:
        print(
            f"{args.case_tag},{args.sdpa_backend},{args.batch_size},{args.q_seqlen},{args.kv_seqlen},{args.num_q_heads},{args.num_kv_heads},{head_dim_qk},{fwd_median_time:.3f},{bwd_median_time:.3f},{fwd_tflops:.0f},{bwd_tflops:.0f},{(np.max(np.array(forward_diffs[5:])) if len(forward_diffs) > 5 else (np.max(np.array(forward_diffs)) if len(forward_diffs) > 0 else 0.0)):.6f},{num_iters}"
        )
    else:
        if run_fwd and run_bwd:
            print(
                f"{args.sdpa_backend}:: Median (fwd, bwd) Execution Times: {fwd_median_time:.3f} ms ({fwd_tflops:.0f} TFLOPS{fwd_sol_str}), {bwd_median_time:.3f} ms ({bwd_tflops:.0f} TFLOPS{bwd_sol_str})"
            )
        elif run_fwd:
            print(f"{args.sdpa_backend}:: Median (fwd) Execution Time: {fwd_median_time:.3f} ms ({fwd_tflops:.0f} TFLOPS{fwd_sol_str})")
        elif run_bwd:
            print(f"{args.sdpa_backend}:: Median (bwd) Execution Time: {bwd_median_time:.3f} ms ({bwd_tflops:.0f} TFLOPS{bwd_sol_str})")
