import argparse
import itertools
import json
import os

import torch
import triton

from aiter.test_common import run_perftest
from aiter.ops.triton.unified_attention import (
    kernel_unified_attention_2d,
    kernel_unified_attention_3d,
    reduce_segments,
    unified_attention as unified_attention_runtime,
)
from op_tests.triton_tests.test_unified_attention import input_helper, ref_paged_attn

os.environ["AMDGCN_USE_BUFFER_OPS"] = "1"


def get_bench_inputs_2d():
    names = [
        "case_name", "seq_lens", "num_heads", "head_size", "block_size", "dtype", "q_dtype", "num_blocks", "sliding_window", "soft_cap", "use_sinks", "num_par_softmax_segments", "seq_threshold_3d"
    ]
    cases = [
        # MIMO-V2
        ("prefill_gqa_sink_swa",    [(49, 4096)], (16, 2), 192, 16, torch.bfloat16, None, 2048, 128,  None, True,  16, 128),
        ("prefill_mqa",             [(49, 4096)], (16, 1), 192, 32, torch.bfloat16, None, 2048, None, None, False, 16, 128),
        ("prefill_gqa_sink_swa_multibatch", [(3, 4096), (49, 4096), (511, 4096)], (16, 2), 192, 16, torch.bfloat16, None, 2048, 128, None, True, 16, 128),
        ("prefill_mqa_multibatch",          [(3, 4096), (49, 4096), (511, 4096)],    (16, 1), 192, 32, torch.bfloat16, None, 2048, None, None, False, 16, 2),
    ]
    return names, cases


def get_bench_inputs_3d():
    names = [
        "case_name", "seq_lens", "num_heads", "head_size", "block_size", "dtype", "q_dtype", "num_blocks", "sliding_window", "soft_cap", "use_sinks", "num_par_softmax_segments", "seq_threshold_3d"
    ]
    cases = [
        ("decode_gqa_sink_swa_3d",  [(1, 4096)], (16, 2), 192, 16, torch.bfloat16, None, 2048, 128,  None, True,  16, 128),
        ("decode_mqa_3d",           [(1, 4096)], (16, 1), 192, 32, torch.bfloat16, None, 2048, None, None, False, 16, 128),
    ]
    return names, cases


def generate_triton_config_combinations(config):
    keys = list(config.keys())
    values = list(config.values())
    configs_list = []
    for combination in itertools.product(*values):
        cfg = dict(zip(keys, combination))
        configs_list.append(cfg)
    return configs_list


def get_unified_attention_2d_triton_configs():

    config = {
        "BLOCK_M": [16, 32, 64],
        "TILE_SIZE": [16, 32, 64],
        "num_warps": [4, 8],
        "num_stages": [1, 2],
    }

    tt_configs = []
    for c in generate_triton_config_combinations(config):
        num_warps = c["num_warps"]
        num_stages = c["num_stages"]
        del c["num_warps"]
        del c["num_stages"]
        tt_configs.append(
            triton.Config(c, num_warps=num_warps, num_stages=num_stages)
        )

    return tt_configs

def prune_unified_attention_2d_configs(configs, nargs, **kwargs):
    def _prune(config, nargs, **kwargs):
        c = config.all_kwargs()
        block_m = c["BLOCK_M"]
        tile_size = c["TILE_SIZE"]

        head_size = kwargs["HEAD_SIZE"]
        num_queries_per_kv = kwargs["num_queries_per_kv"]
        use_fp8 = kwargs["USE_FP8"]

        if block_m < num_queries_per_kv:
            return True
        if head_size > 128 and block_m >= 64:
            return True
        if use_fp8 and tile_size < 32:
            return True

        return False

    res = [c for c in configs if not _prune(c, nargs, **kwargs)]
    return res

unified_attention_2d_tuner = triton.utils.hcutune(
    configs=get_unified_attention_2d_triton_configs(),
    key=[
        "BLOCK_SIZE",
        "HEAD_SIZE",
        "num_queries_per_kv",
        "SLIDING_WINDOW",
        "USE_ALIBI_SLOPES",
        "USE_QQ_BIAS",
        "USE_SOFTCAP",
        "USE_SINKS",
        "USE_MM_PREFIX",
        "USE_FP8",
    ],
    perf_debug=True,
    prune_configs_by={"early_config_prune": prune_unified_attention_2d_configs},
)(kernel_unified_attention_2d)


@torch.inference_mode()
def unified_attention_2d(
    q,
    k,
    v,
    out,
    cu_seqlens_q,
    seqused_k,
    max_seqlen_q,
    max_seqlen_k,
    softmax_scale,
    causal,
    window_size,
    block_table,
    softcap,
    q_descale,
    k_descale,
    v_descale,
    alibi_slopes=None,
    output_scale=None,
    qq_bias=None,
    sinks=None,
    mm_prefix_range=None,
    use_alibi_sqrt=False,
):
    del max_seqlen_q, max_seqlen_k
    assert causal, "Only causal attention is supported"
    assert q_descale is None, "Q scales not supported"

    use_mm_prefix = False
    max_mm_ranges = 0
    if mm_prefix_range is not None:
        if mm_prefix_range.ndim == 3:
            use_mm_prefix = True
            max_mm_ranges = mm_prefix_range.shape[1]
        else:
            raise ValueError(
                f"Unsupported mm_prefix_range shape: {mm_prefix_range.shape}"
            )

    use_alibi_slopes = alibi_slopes is not None
    use_qq_bias = qq_bias is not None

    block_size = v.shape[1]
    num_seqs = len(seqused_k)
    num_query_heads = q.shape[1]
    num_kv_heads = k.shape[2]
    num_queries_per_kv = num_query_heads // num_kv_heads
    head_size = q.shape[2]

    grid = lambda meta: (
        q.shape[0] // (meta["BLOCK_M"] // num_queries_per_kv) + num_seqs,
        num_kv_heads,
    )

    unified_attention_2d_tuner[grid](
        output_ptr=out,
        query_ptr=q,
        key_cache_ptr=k,
        value_cache_ptr=v,
        sink_ptr=sinks,
        block_tables_ptr=block_table,
        seq_lens_ptr=seqused_k,
        alibi_slopes_ptr=alibi_slopes,
        qq_bias_ptr=qq_bias,
        scale=softmax_scale,
        k_scale=k_descale,
        v_scale=v_descale,
        out_scale=1 / output_scale if output_scale is not None else 1.0,
        softcap=softcap,
        num_query_heads=num_query_heads,
        num_queries_per_kv=num_queries_per_kv,
        block_table_stride=block_table.stride(0),
        query_stride_0=q.stride(0),
        query_stride_1=q.stride(1),
        output_stride_0=out.stride(0),
        output_stride_1=out.stride(1),
        qq_bias_stride_0=qq_bias.stride(0) if use_qq_bias else 0,
        BLOCK_SIZE=block_size,
        HEAD_SIZE=head_size,
        HEAD_SIZE_PADDED=triton.next_power_of_2(head_size),
        USE_ALIBI_SLOPES=use_alibi_slopes,
        USE_ALIBI_SQRT=use_alibi_sqrt,
        USE_QQ_BIAS=use_qq_bias,
        USE_SOFTCAP=(softcap > 0),
        USE_SINKS=(sinks is not None),
        USE_MM_PREFIX=use_mm_prefix,
        MAX_MM_RANGES=max_mm_ranges,
        mm_prefix_range_ptr=mm_prefix_range,
        SLIDING_WINDOW=(1 + window_size[0]),
        stride_k_cache_0=k.stride(0),
        stride_k_cache_1=k.stride(1),
        stride_k_cache_2=k.stride(2),
        stride_k_cache_3=k.stride(3),
        stride_v_cache_0=v.stride(0),
        stride_v_cache_1=v.stride(1),
        stride_v_cache_2=v.stride(2),
        stride_v_cache_3=v.stride(3),
        query_start_len_ptr=cu_seqlens_q,
        num_seqs=num_seqs,
        USE_FP8=output_scale is not None,
    )


##################################### 3d kernel #####################################

def get_unified_attention_3d_triton_configs():

    config = {
        "BLOCK_M": [16, 32, 64],
        "TILE_SIZE": [16, 32, 64],
        "num_warps": [4, 8],
        "num_stages": [1, 2],
    }

    tt_configs = []
    for c in generate_triton_config_combinations(config):
        num_warps = c["num_warps"]
        num_stages = c["num_stages"]
        del c["num_warps"]
        del c["num_stages"]
        tt_configs.append(
            triton.Config(c, num_warps=num_warps, num_stages=num_stages)
        )

    return tt_configs


def prune_unified_attention_3d_configs(configs, nargs, **kwargs):
    def _prune(config, nargs, **kwargs):
        c = config.all_kwargs()
        block_m = c["BLOCK_M"]
        tile_size = c["TILE_SIZE"]

        head_size = kwargs["HEAD_SIZE"]
        num_queries_per_kv = kwargs["num_queries_per_kv"]
        num_segments_per_seq = kwargs["NUM_SEGMENTS_PER_SEQ"]
        use_fp8 = "float8" in str(kwargs["key_cache_ptr"].dtype)

        if block_m < num_queries_per_kv:
            return True
        if head_size > 128 and block_m >= 64:
            return True
        if use_fp8 and tile_size < 32:
            return True
        if num_segments_per_seq is not None and num_segments_per_seq <= 0:
            return True

        return False

    res = [c for c in configs if not _prune(c, nargs, **kwargs)]
    return res

kernel_unified_attention_3d_h = triton.heuristics(
    values={"BLOCK_Q": lambda args: args["BLOCK_M"] // args["num_queries_per_kv"]}
)(kernel_unified_attention_3d)

unified_attention_3d_tuner = triton.utils.hcutune(
    configs=get_unified_attention_3d_triton_configs(),
    key=[
        "BLOCK_SIZE",
        "HEAD_SIZE",
        "num_queries_per_kv",
        "SLIDING_WINDOW",
        "USE_ALIBI_SLOPES",
        "USE_QQ_BIAS",
        "USE_SOFTCAP",
        "USE_SINKS",
        "USE_MM_PREFIX",
        "NUM_SEGMENTS_PER_SEQ",
    ],
    perf_debug=True,
    prune_configs_by={"early_config_prune": prune_unified_attention_3d_configs},
)(kernel_unified_attention_3d_h)


def get_reduce_segments_triton_configs():
    config = {
        "num_warps": [1, 2, 4, 8],
        "num_stages": [1, 2],
    }

    tt_configs = []
    for c in generate_triton_config_combinations(config):
        num_warps = c["num_warps"]
        num_stages = c["num_stages"]
        del c["num_warps"]
        del c["num_stages"]
        tt_configs.append(
            triton.Config(c, num_warps=num_warps, num_stages=num_stages)
        )

    return tt_configs


def prune_reduce_segments_configs(configs, nargs, **kwargs):
    return configs

reduce_segments_tuner = triton.utils.hcutune(
    configs=get_reduce_segments_triton_configs(),
    key=[
        "HEAD_SIZE",
        "TILE_SIZE",
        "BLOCK_Q",
        "NUM_SEGMENTS_PER_SEQ",
        "USE_FP8",
    ],
    perf_debug=True,
    prune_configs_by={"early_config_prune": prune_reduce_segments_configs},
)(reduce_segments)


def _get_selected_3d_kernel_config() -> dict:
    best_config = getattr(unified_attention_3d_tuner, "best_config", None)
    if best_config is None:
        raise RuntimeError("unified_attention_3d_tuner did not select a config")

    selected_config = dict(best_config.all_kwargs())
    for key in ("BLOCK_M", "TILE_SIZE"):
        if key not in selected_config:
            raise RuntimeError(
                f"unified_attention_3d_tuner best config is missing required key {key}: "
                f"{selected_config}"
            )
    return selected_config


@torch.inference_mode()
def unified_attention_3d(
    q,
    k,
    v,
    out,
    cu_seqlens_q,
    seqused_k,
    max_seqlen_q,
    max_seqlen_k,
    softmax_scale,
    causal,
    window_size,
    block_table,
    softcap,
    q_descale,
    k_descale,
    v_descale,
    softmax_segm_output,
    softmax_segm_max,
    softmax_segm_expsum,
    alibi_slopes=None,
    output_scale=None,
    qq_bias=None,
    sinks=None,
    mm_prefix_range=None,
    use_alibi_sqrt=False,
    num_par_softmax_segments=None,
):
    del max_seqlen_q, max_seqlen_k
    assert causal, "Only causal attention is supported"
    assert q_descale is None, "Q scales not supported"
    assert num_par_softmax_segments is not None, "3D path requires num_par_softmax_segments"

    use_mm_prefix = False
    max_mm_ranges = 0
    if mm_prefix_range is not None:
        if mm_prefix_range.ndim == 3:
            use_mm_prefix = True
            max_mm_ranges = mm_prefix_range.shape[1]
        else:
            raise ValueError(
                f"Unsupported mm_prefix_range shape: {mm_prefix_range.shape}"
            )

    use_alibi_slopes = alibi_slopes is not None
    use_qq_bias = qq_bias is not None

    block_size = v.shape[1]
    num_seqs = len(seqused_k)
    num_query_heads = q.shape[1]
    num_kv_heads = k.shape[2]
    num_queries_per_kv = num_query_heads // num_kv_heads
    head_size = q.shape[2]

    grid_3d = lambda meta: (
        q.shape[0] // (meta["BLOCK_M"] // num_queries_per_kv) + num_seqs,
        num_kv_heads,
        num_par_softmax_segments,
    )
    unified_attention_3d_tuner[grid_3d](
        segm_output_ptr=softmax_segm_output,
        segm_max_ptr=softmax_segm_max,
        segm_expsum_ptr=softmax_segm_expsum,
        query_ptr=q,
        key_cache_ptr=k,
        value_cache_ptr=v,
        sink_ptr=sinks,
        block_tables_ptr=block_table,
        seq_lens_ptr=seqused_k,
        alibi_slopes_ptr=alibi_slopes,
        qq_bias_ptr=qq_bias,
        scale=softmax_scale,
        k_scale=k_descale,
        v_scale=v_descale,
        softcap=softcap,
        num_query_heads=num_query_heads,
        num_queries_per_kv=num_queries_per_kv,
        block_table_stride=block_table.stride(0),
        query_stride_0=q.stride(0),
        query_stride_1=q.stride(1),
        qq_bias_stride_0=qq_bias.stride(0) if use_qq_bias else 0,
        BLOCK_SIZE=block_size,
        HEAD_SIZE=head_size,
        HEAD_SIZE_PADDED=triton.next_power_of_2(head_size),
        USE_ALIBI_SLOPES=use_alibi_slopes,
        USE_ALIBI_SQRT=use_alibi_sqrt,
        USE_QQ_BIAS=use_qq_bias,
        USE_SOFTCAP=(softcap > 0),
        USE_SINKS=(sinks is not None),
        USE_MM_PREFIX=use_mm_prefix,
        MAX_MM_RANGES=max_mm_ranges,
        mm_prefix_range_ptr=mm_prefix_range,
        SLIDING_WINDOW=(1 + window_size[0]),
        stride_k_cache_0=k.stride(0),
        stride_k_cache_1=k.stride(1),
        stride_k_cache_2=k.stride(2),
        stride_k_cache_3=k.stride(3),
        stride_v_cache_0=v.stride(0),
        stride_v_cache_1=v.stride(1),
        stride_v_cache_2=v.stride(2),
        stride_v_cache_3=v.stride(3),
        query_start_len_ptr=cu_seqlens_q,
        num_seqs=num_seqs,
        NUM_SEGMENTS_PER_SEQ=num_par_softmax_segments,
    )

    selected_config = _get_selected_3d_kernel_config()
    reduce_tile_size = selected_config["TILE_SIZE"]
    reduce_block_q = selected_config["BLOCK_M"] // num_queries_per_kv

    # Keep the reduce stage coupled to the producer kernel's selected meta.
    assert reduce_tile_size == selected_config["TILE_SIZE"]
    assert reduce_block_q == selected_config["BLOCK_M"] // num_queries_per_kv

    grid_reduce = (q.shape[0], num_query_heads)
    reduce_segments_tuner[grid_reduce](
        output_ptr=out,
        segm_output_ptr=softmax_segm_output,
        segm_max_ptr=softmax_segm_max,
        segm_expsum_ptr=softmax_segm_expsum,
        seq_lens_ptr=seqused_k,
        num_seqs=num_seqs,
        num_query_heads=num_query_heads,
        out_scale_inv=1 / output_scale if output_scale is not None else 1.0,
        output_stride_0=out.stride(0),
        output_stride_1=out.stride(1),
        block_table_stride=block_table.stride(0),
        TILE_SIZE=reduce_tile_size,
        HEAD_SIZE=head_size,
        HEAD_SIZE_PADDED=triton.next_power_of_2(head_size),
        query_start_len_ptr=cu_seqlens_q,
        BLOCK_Q=reduce_block_q,
        NUM_SEGMENTS_PER_SEQ=num_par_softmax_segments,
        USE_FP8=output_scale is not None,
    )
    return None


##################################### Benchmark #####################################

def _estimate_unified_attention_total_attended(case, inputs) -> int:
    (
        _name,
        _seq_lens,
        _num_heads,
        _head_size,
        _block_size,
        _dtype,
        _q_dtype,
        _num_blocks,
        sliding_window,
        _soft_cap,
        _use_sinks,
        _num_par_softmax_segments,
        _seq_threshold_3d,
    ) = case
    (
        query_lens,
        kv_lens,
        _query,
        _key_cache,
        _value_cache,
        _maybe_quantized_query,
        _maybe_quantized_key_cache,
        _maybe_quantized_value_cache,
        _output,
        _cu_query_lens,
        _kv_lens_tensor,
        _block_tables,
        _q_descale,
        _k_descale,
        _v_descale,
        _sinks,
        _max_query_len,
        _max_kv_len,
        _softmax_segm_output,
        _softmax_segm_max,
        _softmax_segm_expsum,
    ) = inputs
    total_attended = 0
    for query_len, kv_len in zip(query_lens, kv_lens):
        context_len = kv_len - query_len
        for query_idx in range(query_len):
            attended = context_len + query_idx + 1
            if sliding_window is not None:
                attended = min(attended, sliding_window)
            total_attended += attended
    return total_attended


def _estimate_unified_attention_kv_io_tokens(case, inputs) -> int:
    (
        _name,
        _seq_lens,
        _num_heads,
        _head_size,
        _block_size,
        _dtype,
        _q_dtype,
        _num_blocks,
        sliding_window,
        _soft_cap,
        _use_sinks,
        _num_par_softmax_segments,
        _seq_threshold_3d,
    ) = case
    (
        query_lens,
        kv_lens,
        _query,
        _key_cache,
        _value_cache,
        _maybe_quantized_query,
        _maybe_quantized_key_cache,
        _maybe_quantized_value_cache,
        _output,
        _cu_query_lens,
        _kv_lens_tensor,
        _block_tables,
        _q_descale,
        _k_descale,
        _v_descale,
        _sinks,
        _max_query_len,
        _max_kv_len,
        _softmax_segm_output,
        _softmax_segm_max,
        _softmax_segm_expsum,
    ) = inputs

    total_loaded_kv_tokens = 0
    for query_len, kv_len in zip(query_lens, kv_lens):
        context_len = kv_len - query_len
        for query_idx in range(query_len):
            attended = context_len + query_idx + 1
            if sliding_window is not None:
                attended = min(attended, sliding_window)
            total_loaded_kv_tokens += attended

    return total_loaded_kv_tokens


def _estimate_unified_attention_flops(case, inputs) -> float:
    total_attended = _estimate_unified_attention_total_attended(case, inputs)
    (
        _query_lens,
        _kv_lens,
        query,
        _key_cache,
        _value_cache,
        _maybe_quantized_query,
        _maybe_quantized_key_cache,
        _maybe_quantized_value_cache,
        _output,
        _cu_query_lens,
        _kv_lens_tensor,
        _block_tables,
        _q_descale,
        _k_descale,
        _v_descale,
        _sinks,
        _max_query_len,
        _max_kv_len,
        _softmax_segm_output,
        _softmax_segm_max,
        _softmax_segm_expsum,
    ) = inputs
    num_query_heads = query.shape[1]
    head_size = query.shape[2]
    return 4.0 * num_query_heads * head_size * total_attended


def _estimate_unified_attention_io_bytes(case, inputs) -> int:
    total_loaded_kv_tokens = _estimate_unified_attention_kv_io_tokens(case, inputs)
    (
        _query_lens,
        _kv_lens,
        _query,
        _key_cache,
        _value_cache,
        maybe_quantized_query,
        maybe_quantized_key_cache,
        maybe_quantized_value_cache,
        output,
        _cu_query_lens,
        _kv_lens_tensor,
        _block_tables,
        _q_descale,
        _k_descale,
        _v_descale,
        _sinks,
        _max_query_len,
        _max_kv_len,
        _softmax_segm_output,
        _softmax_segm_max,
        _softmax_segm_expsum,
    ) = inputs
    num_kv_heads = maybe_quantized_key_cache.shape[2]
    head_size = maybe_quantized_key_cache.shape[3]
    q_dtype_size = maybe_quantized_query.element_size()
    k_dtype_size = maybe_quantized_key_cache.element_size()
    v_dtype_size = maybe_quantized_value_cache.element_size()
    out_dtype_size = output.element_size()

    io_bytes = 0
    io_bytes += maybe_quantized_query.numel() * q_dtype_size
    io_bytes += total_loaded_kv_tokens * num_kv_heads * head_size * k_dtype_size
    io_bytes += total_loaded_kv_tokens * num_kv_heads * head_size * v_dtype_size
    io_bytes += output.numel() * out_dtype_size

    return io_bytes


def _validate_runtime_output_with_reference(
    *,
    query_lens,
    kv_lens,
    query,
    key_cache,
    value_cache,
    output,
    block_tables,
    scale,
    sliding_window,
    soft_cap,
    sinks,
    q_dtype,
) -> None:
    ref_output = ref_paged_attn(
        query=query,
        key_cache=key_cache,
        value_cache=value_cache,
        query_lens=query_lens,
        kv_lens=kv_lens,
        block_tables=block_tables,
        scale=scale,
        sliding_window=sliding_window,
        soft_cap=soft_cap,
        sinks=sinks,
    )
    atol, rtol = 1.5e-2, 1e-2
    if q_dtype is not None:
        atol, rtol = 1.5e-1, 1.5e-1
    torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol)


def _print_benchmark_summary(results: list[dict]) -> None:
    def format_seq_lens(seq_lens: list[tuple[int, int]]) -> str:
        if len(seq_lens) == 1:
            q_len, kv_len = seq_lens[0]
            return f"{q_len}/{kv_len}"
        if len(seq_lens) <= 2:
            return ",".join(f"{q}/{kv}" for q, kv in seq_lens)
        prefix = ",".join(f"{q}/{kv}" for q, kv in seq_lens[:2])
        return f"{prefix},..."

    headers = [
        ("case_name", "case"),
        ("seq_lens", "q/kv"),
        ("num_heads", "hq/hkv"),
        ("head_size", "d"),
        ("block_size", "blk"),
        ("sliding_window", "sw"),
        ("use_sinks", "sink"),
        ("latency_us", "latency_us"),
        ("bandwidth_gbs", "GB/s"),
        ("tflops", "TFLOPS"),
    ]

    rows = []
    for result in results:
        rows.append(
            {
                "case_name": str(result["case_name"]),
                "seq_lens": format_seq_lens(result["seq_lens"]),
                "num_heads": f'{result["num_heads"][0]}/{result["num_heads"][1]}',
                "head_size": str(result["head_size"]),
                "block_size": str(result["block_size"]),
                "sliding_window": (
                    "-" if result["sliding_window"] is None else str(result["sliding_window"])
                ),
                "use_sinks": "Y" if result["use_sinks"] else "N",
                "latency_us": f'{result["latency_us"]:.3f}',
                "bandwidth_gbs": f'{result["bandwidth_gbs"]:.3f}',
                "tflops": f'{result["tflops"]:.3f}',
            }
        )

    widths = {
        key: max(len(title), *(len(row[key]) for row in rows))
        for key, title in headers
    }

    def format_row(row: dict[str, str]) -> str:
        return " | ".join(
            row[key].ljust(widths[key]) if key == "case_name" else row[key].rjust(widths[key])
            for key, _ in headers
        )

    header_row = format_row({key: title for key, title in headers})
    separator = "-+-".join("-" * widths[key] for key, _ in headers)
    print(header_row)
    print(separator)
    for row in rows:
        print(format_row(row))


x_names_2d, x_vals_2d = get_bench_inputs_2d()
configs_2d = [
    triton.testing.Benchmark(
        x_names=x_names_2d,
        x_vals=x_vals_2d,
        line_arg="provider",
        line_vals=["triton"],
        line_names=["triton"],
        styles=[("red", "-")],
        ylabel="ms",
        plot_name="unified_attention_2d",
        args={},
    )
]


@triton.utils.dist_perf_report(configs_2d)
def bench_unified_attention_2d(
    case_name,
    seq_lens,
    num_heads,
    head_size,
    block_size,
    dtype,
    q_dtype,
    num_blocks,
    sliding_window,
    soft_cap,
    use_sinks,
    num_par_softmax_segments,
    seq_threshold_3d,
    provider,
):
    del provider
    device = "cpu" if os.getenv("TRITON_HCUTUNE_COMPILE_ONLY", "") == "1" else "cuda"
    (
        _query_lens,
        _kv_lens,
        _query,
        _key_cache,
        _value_cache,
        maybe_quantized_query,
        maybe_quantized_key_cache,
        maybe_quantized_value_cache,
        output,
        cu_query_lens,
        kv_lens_tensor,
        block_tables,
        q_descale,
        k_descale,
        v_descale,
        sinks,
        max_query_len,
        max_kv_len,
        softmax_segm_output,
        softmax_segm_max,
        softmax_segm_expsum,
    ) = input_helper(
        seq_lens=seq_lens,
        num_heads=num_heads,
        head_size=head_size,
        block_size=block_size,
        num_blocks=num_blocks,
        dtype=dtype,
        q_dtype=q_dtype,
        use_sinks=use_sinks,
        seq_threshold_3D=seq_threshold_3d,
        num_par_softmax_segments=num_par_softmax_segments,
        device=device,
    )
    window_size = (
        (sliding_window - 1, 0) if sliding_window is not None else (-1, -1)
    )
    scale = head_size**-0.5
    tuning_mod = int(os.environ.get("TUNING_MOD", "1"))
    if tuning_mod:
        fn = lambda: unified_attention_2d(
            q=maybe_quantized_query,
            k=maybe_quantized_key_cache,
            v=maybe_quantized_value_cache,
            out=output,
            cu_seqlens_q=cu_query_lens,
            seqused_k=kv_lens_tensor,
            max_seqlen_q=max_query_len,
            max_seqlen_k=max_kv_len,
            softmax_scale=scale,
            causal=True,
            window_size=window_size,
            block_table=block_tables,
            softcap=soft_cap if soft_cap is not None else 0,
            q_descale=q_descale,
            k_descale=k_descale,
            v_descale=v_descale,
            alibi_slopes=None,
            sinks=sinks,
        )
    else:
        fn = lambda: unified_attention_runtime(
            q=maybe_quantized_query,
            k=maybe_quantized_key_cache,
            v=maybe_quantized_value_cache,
            out=output,
            cu_seqlens_q=cu_query_lens,
            seqused_k=kv_lens_tensor,
            max_seqlen_q=max_query_len,
            max_seqlen_k=max_kv_len,
            softmax_scale=scale,
            causal=True,
            window_size=window_size,
            block_table=block_tables,
            softcap=soft_cap if soft_cap is not None else 0,
            q_descale=q_descale,
            k_descale=k_descale,
            v_descale=v_descale,
            sinks=sinks,
            seq_threshold_3D=seq_threshold_3d,
            num_par_softmax_segments=num_par_softmax_segments,
            softmax_segm_output=softmax_segm_output,
            softmax_segm_max=softmax_segm_max,
            softmax_segm_expsum=softmax_segm_expsum,
        )
    return triton.testing.do_bench(fn)


x_names_3d, x_vals_3d = get_bench_inputs_3d()
configs_3d = [
    triton.testing.Benchmark(
        x_names=x_names_3d,
        x_vals=x_vals_3d,
        line_arg="provider",
        line_vals=["triton"],
        line_names=["triton"],
        styles=[("red", "-")],
        ylabel="ms",
        plot_name="unified_attention_3d",
        args={},
    )
]


@triton.utils.dist_perf_report(configs_3d)
def bench_unified_attention_3d(
    case_name,
    seq_lens,
    num_heads,
    head_size,
    block_size,
    dtype,
    q_dtype,
    num_blocks,
    sliding_window,
    soft_cap,
    use_sinks,
    num_par_softmax_segments,
    seq_threshold_3d,
    provider,
):
    del provider, case_name
    device = "cpu" if os.getenv("TRITON_HCUTUNE_COMPILE_ONLY", "") == "1" else "cuda"
    (
        _query_lens,
        _kv_lens,
        _query,
        _key_cache,
        _value_cache,
        maybe_quantized_query,
        maybe_quantized_key_cache,
        maybe_quantized_value_cache,
        output,
        cu_query_lens,
        kv_lens_tensor,
        block_tables,
        q_descale,
        k_descale,
        v_descale,
        sinks,
        max_query_len,
        max_kv_len,
        softmax_segm_output,
        softmax_segm_max,
        softmax_segm_expsum,
    ) = input_helper(
        seq_lens=seq_lens,
        num_heads=num_heads,
        head_size=head_size,
        block_size=block_size,
        num_blocks=num_blocks,
        dtype=dtype,
        q_dtype=q_dtype,
        use_sinks=use_sinks,
        seq_threshold_3D=seq_threshold_3d,
        num_par_softmax_segments=num_par_softmax_segments,
        device=device,
    )
    window_size = (
        (sliding_window - 1, 0) if sliding_window is not None else (-1, -1)
    )
    scale = head_size**-0.5
    num_seqs = len(kv_lens_tensor)
    num_query_heads = maybe_quantized_query.shape[1]
    num_kv_heads = maybe_quantized_key_cache.shape[2]
    num_queries_per_kv = num_query_heads // num_kv_heads

    tuning_mod = int(os.environ.get("TUNING_MOD", "1"))
    if tuning_mod:
        def fn():
            unified_attention_3d(
                q=maybe_quantized_query,
                k=maybe_quantized_key_cache,
                v=maybe_quantized_value_cache,
                out=output,
                cu_seqlens_q=cu_query_lens,
                seqused_k=kv_lens_tensor,
                max_seqlen_q=max_query_len,
                max_seqlen_k=max_kv_len,
                softmax_scale=scale,
                causal=True,
                window_size=window_size,
                block_table=block_tables,
                softcap=soft_cap if soft_cap is not None else 0,
                q_descale=q_descale,
                k_descale=k_descale,
                v_descale=v_descale,
                softmax_segm_output=softmax_segm_output,
                softmax_segm_max=softmax_segm_max,
                softmax_segm_expsum=softmax_segm_expsum,
                sinks=sinks,
                num_par_softmax_segments=num_par_softmax_segments,
            )
    else:
        fn = lambda: unified_attention_runtime(
            q=maybe_quantized_query,
            k=maybe_quantized_key_cache,
            v=maybe_quantized_value_cache,
            out=output,
            cu_seqlens_q=cu_query_lens,
            seqused_k=kv_lens_tensor,
            max_seqlen_q=max_query_len,
            max_seqlen_k=max_kv_len,
            softmax_scale=scale,
            causal=True,
            window_size=window_size,
            block_table=block_tables,
            softcap=soft_cap if soft_cap is not None else 0,
            q_descale=q_descale,
            k_descale=k_descale,
            v_descale=v_descale,
            sinks=sinks,
            seq_threshold_3D=seq_threshold_3d,
            num_par_softmax_segments=num_par_softmax_segments,
            softmax_segm_output=softmax_segm_output,
            softmax_segm_max=softmax_segm_max,
            softmax_segm_expsum=softmax_segm_expsum,
        )
    return triton.testing.do_bench(fn)


@torch.inference_mode()
def _benchmark_unified_attention_case(
    case_name,
    seq_lens,
    num_heads,
    head_size,
    block_size,
    dtype,
    q_dtype,
    num_blocks,
    sliding_window,
    soft_cap,
    use_sinks,
    num_par_softmax_segments,
    seq_threshold_3d,
    use_tuned: bool,
    check_correctness: bool = False,
) -> dict:
    (
        _query_lens,
        _kv_lens,
        _query,
        _key_cache,
        _value_cache,
        maybe_quantized_query,
        maybe_quantized_key_cache,
        maybe_quantized_value_cache,
        output,
        cu_query_lens,
        kv_lens_tensor,
        block_tables,
        q_descale,
        k_descale,
        v_descale,
        sinks,
        max_query_len,
        max_kv_len,
        softmax_segm_output,
        softmax_segm_max,
        softmax_segm_expsum,
    ) = input_helper(
        seq_lens=seq_lens,
        num_heads=num_heads,
        head_size=head_size,
        block_size=block_size,
        num_blocks=num_blocks,
        dtype=dtype,
        q_dtype=q_dtype,
        use_sinks=use_sinks,
        seq_threshold_3D=seq_threshold_3d,
        num_par_softmax_segments=num_par_softmax_segments,
        device="cuda",
    )
    window_size = (
        (sliding_window - 1, 0) if sliding_window is not None else (-1, -1)
    )
    scale = head_size**-0.5
    num_seqs = len(kv_lens_tensor)
    num_query_heads = maybe_quantized_query.shape[1]
    num_kv_heads = maybe_quantized_key_cache.shape[2]
    num_queries_per_kv = num_query_heads // num_kv_heads
    use_3d = (
        max_query_len == 1
        and num_par_softmax_segments is not None
        and seq_threshold_3d is not None
        and num_seqs <= seq_threshold_3d
    )

    def _run_tuned_2d():
        unified_attention_2d(
            q=maybe_quantized_query,
            k=maybe_quantized_key_cache,
            v=maybe_quantized_value_cache,
            out=output,
            cu_seqlens_q=cu_query_lens,
            seqused_k=kv_lens_tensor,
            max_seqlen_q=max_query_len,
            max_seqlen_k=max_kv_len,
            softmax_scale=scale,
            causal=True,
            window_size=window_size,
            block_table=block_tables,
            softcap=soft_cap if soft_cap is not None else 0,
            q_descale=q_descale,
            k_descale=k_descale,
            v_descale=v_descale,
            alibi_slopes=None,
            sinks=sinks,
        )
        return output

    def _run_tuned_3d():
        unified_attention_3d(
            q=maybe_quantized_query,
            k=maybe_quantized_key_cache,
            v=maybe_quantized_value_cache,
            out=output,
            cu_seqlens_q=cu_query_lens,
            seqused_k=kv_lens_tensor,
            max_seqlen_q=max_query_len,
            max_seqlen_k=max_kv_len,
            softmax_scale=scale,
            causal=True,
            window_size=window_size,
            block_table=block_tables,
            softcap=soft_cap if soft_cap is not None else 0,
            q_descale=q_descale,
            k_descale=k_descale,
            v_descale=v_descale,
            alibi_slopes=None,
            sinks=sinks,
            softmax_segm_output=softmax_segm_output,
            softmax_segm_max=softmax_segm_max,
            softmax_segm_expsum=softmax_segm_expsum,
            num_par_softmax_segments=num_par_softmax_segments,
        )
        return output

    run_case = _run_tuned_3d if (use_tuned and use_3d) else _run_tuned_2d
    if not use_tuned:
        def _run_runtime():
            unified_attention_runtime(
                q=maybe_quantized_query,
                k=maybe_quantized_key_cache,
                v=maybe_quantized_value_cache,
                out=output,
                cu_seqlens_q=cu_query_lens,
                seqused_k=kv_lens_tensor,
                max_seqlen_q=max_query_len,
                max_seqlen_k=max_kv_len,
                softmax_scale=scale,
                causal=True,
                window_size=window_size,
                block_table=block_tables,
                softcap=soft_cap if soft_cap is not None else 0,
                q_descale=q_descale,
                k_descale=k_descale,
                v_descale=v_descale,
                sinks=sinks,
                seq_threshold_3D=seq_threshold_3d,
                num_par_softmax_segments=num_par_softmax_segments,
                softmax_segm_output=softmax_segm_output,
                softmax_segm_max=softmax_segm_max,
                softmax_segm_expsum=softmax_segm_expsum,
            )
            return output
        run_case = _run_runtime
    _, latency_us = run_perftest(
        run_case,
        num_iters=20,
        num_warmup=5,
    )
    if check_correctness:
        run_case()
        _validate_runtime_output_with_reference(
            query_lens=_query_lens,
            kv_lens=_kv_lens,
            query=_query,
            key_cache=_key_cache,
            value_cache=_value_cache,
            output=output,
            block_tables=block_tables,
            scale=scale,
            sliding_window=sliding_window,
            soft_cap=soft_cap,
            sinks=sinks,
            q_dtype=q_dtype,
        )

    case_tuple = (
        case_name,
        seq_lens,
        num_heads,
        head_size,
        block_size,
        dtype,
        q_dtype,
        num_blocks,
        sliding_window,
        soft_cap,
        use_sinks,
        num_par_softmax_segments,
        seq_threshold_3d,
    )
    inputs_tuple = (
        _query_lens,
        _kv_lens,
        _query,
        _key_cache,
        _value_cache,
        maybe_quantized_query,
        maybe_quantized_key_cache,
        maybe_quantized_value_cache,
        output,
        cu_query_lens,
        kv_lens_tensor,
        block_tables,
        q_descale,
        k_descale,
        v_descale,
        sinks,
        max_query_len,
        max_kv_len,
        softmax_segm_output,
        softmax_segm_max,
        softmax_segm_expsum,
    )
    total_flops = _estimate_unified_attention_flops(case_tuple, inputs_tuple)
    total_io_bytes = _estimate_unified_attention_io_bytes(case_tuple, inputs_tuple)
    return {
        "case_name": case_name,
        "seq_lens": seq_lens,
        "num_heads": num_heads,
        "head_size": head_size,
        "block_size": block_size,
        "sliding_window": sliding_window,
        "use_sinks": use_sinks,
        "latency_us": latency_us,
        "bandwidth_gbs": total_io_bytes / latency_us / 1e3,
        "tflops": total_flops / latency_us / 1e6,
    }


def benchmark_unified_attention_tuned_case(*case) -> dict:
    return _benchmark_unified_attention_case(*case, use_tuned=True)


def benchmark_unified_attention_runtime_case(*case) -> dict:
    return _benchmark_unified_attention_case(
        *case,
        use_tuned=False,
        check_correctness=True,
    )

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--perf",
        action="store_true",
        default=False,
        help="benchmark with hcutuner perf mode",
    )
    return parser.parse_args()


if __name__ == "__main__":
    args = parse_args()
    if args.perf:
        os.environ["TRITON_HCUTUNE_PERF_MODE"] = "1"
    bench_unified_attention_2d.run(
        print_data=True,
    )
    bench_unified_attention_3d.run(
        print_data=True
    )
    results = [
        benchmark_unified_attention_tuned_case(*case)
        for case in (x_vals_2d + x_vals_3d)
    ]
    _print_benchmark_summary(results)
    results = [
        benchmark_unified_attention_runtime_case(*case)
        for case in (x_vals_2d + x_vals_3d)
    ]
    _print_benchmark_summary(results)

# opt_config_cli show --all
# Note: need use ai agent to auto fill the sink_ptr of the config, then use cli or index outrange.
# python -m triton.tools.opt_config_cli export \
#   --kernel kernel_unified_attention_2d \
#   --keep_key num_queries_per_kv \
#   --hoist_key BLOCK_SIZE,HEAD_SIZE,SLIDING_WINDOW,USE_QQ_BIAS,USE_SOFTCAP,USE_SINKS,USE_FP8 \
#   --hoist_dtype \
#   --output_dir ./ua2d_export
