import os
import sys

os.environ["AMDGCN_USE_BUFFER_OPS"] = "1"
# :class:`Hcutuner` reads ``TRITON_HCUTUNE_PERF_MODE`` in ``__init__``. Module-level
# ``fn`` / ``fn_v2`` / ``fn_v2_decode`` = ``triton.utils.hcutune(...)`` run at import time, *before*
# ``if __name__ == "__main__"``, so ``--perf`` must be applied here (or set in the shell).
if "--perf" in sys.argv:
    os.environ["TRITON_HCUTUNE_PERF_MODE"] = "1"

# GPU / ROCm tuning: run **inside** the ``zww_tl_1`` container (not the bare host), e.g.:
#   docker exec zww_tl_1 bash -lc 'cd /data/zhouweiwang/aiter/op_tests/triton_autotune && python tune_extend_attention.py --perf'

# ``do_bench`` timing during tuning (smaller => faster iteration; raise for stable prod numbers).
TUNE_DO_BENCH_WARMUP = 5
TUNE_DO_BENCH_REP = 20

import json
import torch
import triton
import random
import itertools
import argparse

from aiter.ops.triton.extend_attention import (
    _fwd_kernel,
    _fwd_kernel_v2,
    _fwd_kernel_v2_decode,
)

_is_hip = True

# hcutune key for :func:`_fwd_kernel_v2`. JSON block-size lookup uses :func:`_get_config_v2`
# ``want7`` only; these names add kernel constexprs (``SKIP_PREFIX_CUSTOM_MASK``,
# ``xai_temperature_len``) used for autotune but not in the V2 JSON key.
# Log alignment (e.g. ``fp8_dp2_tp8_415_triton_rocm_nomtp.log`` ~2291): ``kv_group_num = q.size(-2)//k.size(-2)``,
# ``Lq``/``Lv`` last dims, ``USE_CUSTOM_MASK = custom_mask is not None``, ``HAS_SINK = sinks is not None``.
HCUTUNE_KEY_V2 = [
    "kv_group_num",
    "Lq",
    "Lv",
    "USE_CUSTOM_MASK",
    "IS_CAUSAL",
    "SKIP_PREFIX_CUSTOM_MASK",
    "HAS_SINK",
    "SLIDING_WINDOW_SIZE",
    "xai_temperature_len",
]

version = triton.__version__.split(".")
major_version, minor_version = eval(version[0]), eval(version[1])


def input_helper(
    B,
    H,
    prefix_length,
    extend_length,
    kv_lora_rank,
    qk_rope_head_dim,
    v_head_dim,
    dtype,
    device,
    attn_impl="normal",
    equal_seqlens=False,
    requires_grad=False,
    kv_num_heads: int = 1,
):
    torch.manual_seed(0)

    if not equal_seqlens:
        max_extend_length = extend_length
        max_prefix_length = prefix_length

        seqlens_extend = torch.randint(
            1, max_extend_length + 1, (B,), dtype=torch.int32
        )
        if prefix_length == 0:
            seqlens_prefix = torch.full((B,), prefix_length, dtype=torch.int32)
        else:
            seqlens_prefix = torch.randint(
                1, max_prefix_length + 1, (B,), dtype=torch.int32
            )

    else:
        seqlens_extend = torch.full((B,), extend_length, dtype=torch.int32)
        seqlens_prefix = torch.full((B,), prefix_length, dtype=torch.int32)

    cu_seqlens_extend = torch.cat(
        [
            torch.tensor([0], dtype=torch.int32),
            seqlens_extend.cumsum(dim=0, dtype=torch.int32),
        ]
    )
    cu_seqlens_prefix = torch.cat(
        [
            torch.tensor([0], dtype=torch.int32),
            seqlens_prefix.cumsum(dim=0, dtype=torch.int32),
        ]
    )

    cu_seqlens_extend = cu_seqlens_extend.to(device="cuda")
    cu_seqlens_prefix = cu_seqlens_prefix.to(device="cuda")

    total_extend = cu_seqlens_extend[-1].item()
    total_prefix = cu_seqlens_prefix[-1].item()

    if attn_impl == "absorb":
        Lq = kv_lora_rank + qk_rope_head_dim
        Lk = kv_lora_rank + qk_rope_head_dim
        Lv = kv_lora_rank
    else:
        Lq = v_head_dim + qk_rope_head_dim
        Lk = v_head_dim + qk_rope_head_dim
        Lv = v_head_dim

    q_extend = torch.randn(
        total_extend, H, Lq, dtype=dtype, device=device
    ).requires_grad_(requires_grad)

    # extend parts (``kv_num_heads`` for GQA: e.g. 2 when q has 16 heads and kv_group_num is 8)
    k_extend = torch.randn(
        total_extend, kv_num_heads, Lk, dtype=dtype, device=device
    ).requires_grad_(requires_grad)
    v_extend = k_extend[..., :Lv]

    # extend indexing
    qo_indptr = cu_seqlens_extend

    # prefix parts
    k_buffer = torch.randn(
        total_prefix, kv_num_heads, Lk, dtype=dtype, device=device
    ).requires_grad_(requires_grad)
    v_buffer = k_buffer[..., :Lv]

    if attn_impl != "absorb":
        # simulate v = kv_latent * w_vc which changes the values compared to k
        v_extend = torch.randn_like(v_extend, dtype=v_extend.dtype)
        v_buffer = torch.randn_like(v_buffer, dtype=v_buffer.dtype)

    # prefix indexing
    kv_indptr = cu_seqlens_prefix
    kv_indices = torch.arange(total_prefix, device=device, dtype=torch.int32)

    custom_mask = None
    mask_indptr = None
    max_len_extend = extend_length

    return (
        q_extend,
        k_extend,
        v_extend,
        k_buffer,
        v_buffer,
        kv_indptr,
        kv_indices,
        qo_indptr,
        custom_mask,
        mask_indptr,
        max_len_extend,
    )


def _v2_flat_causal_custom_mask(B, prefix_len, extend_len, device):
    """Row-major causal mask per batch: shape [extend_len, prefix_len + extend_len], k_global <= q_global."""
    total = prefix_len + extend_len
    q_row = torch.arange(extend_len, device=device) + prefix_len
    k_col = torch.arange(total, device=device)
    m = k_col.unsqueeze(0) <= q_row.unsqueeze(1)
    one_batch = m.reshape(-1).contiguous()
    custom_mask = one_batch.repeat(B)
    seg = extend_len * total
    mask_indptr = torch.arange(
        0, (B + 1) * seg, seg, dtype=torch.int32, device=device
    )
    return custom_mask, mask_indptr


def get_bench_inputs():
    names = ["B", "H", "prefix", "extend", "kv_lora_rank", "qk_rope_head_dim",
             "v_head_dim", "causal", "custom_mask"]
    vals = []
    shapes = [
        (2, 4, 0, 512, 32, 16, 32),
        (3, 5, 0, 333, 18, 13, 17),
        (3, 5, 512, 333, 18, 0, 17),
        (3, 5, 110, 333, 18, 0, 19),
        (8, 16, 0, 1024, 128, 0, 128), # this one passes
        # (8, 16, 0, 16324, 128, 0, 128), # this one fails, numeric precision is likely the issue
        (2, 1, 64, 32, 128, 64, 128),
        (2, 1, 64, 32, 128, 64, 128),
        (4, 16, 64, 96, 128, 64, 128),
        (1, 16, 0, 7, 512, 64, 512),
        (1, 16, 7, 4, 512, 64, 512),
        (1, 16, 32, 4, 512, 64, 512),
        (1, 16, 64, 3, 512, 64, 512),
        (1, 16, 127, 4, 512, 64, 512),
        (1, 16, 255, 15, 512, 64, 512),
        (3, 16, 452, 16, 512, 64, 512),
        (4, 16, 512, 14, 512, 64, 512),
        (4, 16, 1024, 16, 512, 64, 512),
        (4, 16, 2048, 13, 512, 64, 512),
    ]

    for is_causal in [True, False]:
        for use_custom_mask in [True, False]:
            for s in shapes:
                vals.append((*s, is_causal, use_custom_mask))

    return names, vals


def generate_configs(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_triton_configs():
#     config = {
#         "BLOCK_M": [16, 32, 64],
#         "BLOCK_N": [16, 32, 64],
#         "waves_per_eu": [1],
#         "num_warps": [4, 8, 16],
#         # "instruction_sched_variant": ["none", "llvm-iglp-1", "llvm-iglp-8", "local-prefetch"],
#         # "schedule_hint": ["none", "llvm-iglp-8", "local-prefetch"],
#         "matrix_instr_nonkdim": [16],
#         "num_stages": [1, 2, 3],
#         "sched_latency": ["none", "mmac5-ds10"],
#         "kpack": [1, 2],
#     }

#     tt_configs = []
#     for c in generate_configs(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 get_triton_configs():
    config = {
        "BLOCK_M": [16, 32, 64],
        "BLOCK_N": [16, 32, 64],
        "waves_per_eu": [1],
        "num_warps": [4, 8],
        # "instruction_sched_variant": ["none", "llvm-iglp-1", "llvm-iglp-8", "local-prefetch"],
        # "schedule_hint": ["none", "llvm-iglp-8", "local-prefetch"],
        "matrix_instr_nonkdim": [16],
        "num_stages": [1, 2],
        "sched_latency": ["none", "mmac5-ds10"],
        "kpack": [1],
    }

    tt_configs = []
    for c in generate_configs(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_configs(configs, nargs, **kwargs):
    def _prune(config):
        c = config.all_kwargs()
        if c['BLOCK_M'] == 64 and c['num_stages'] >= 2:
            return True
        elif c['BLOCK_M'] >= 64 and c['num_warps'] >= 8:
            return True

    res = [c for c in configs if not _prune(c)]
    return res


key = [
    'kv_group_num',
    'SLIDING_WINDOW_SIZE',
    'Lq',
    'Lv',
    'USE_CUSTOM_MASK',
    'IS_CAUSAL',
    'SKIP_PREFIX_CUSTOM_MASK',
    'STORE_TRANSPOSE',
]
fn = triton.utils.hcutune(
    configs=get_triton_configs(),
    key=key,
    perf_debug=True,
    prune_configs_by={"early_config_prune": prune_configs},
    warmup=TUNE_DO_BENCH_WARMUP,
    rep=TUNE_DO_BENCH_REP,
)(_fwd_kernel)

fn_v2 = triton.utils.hcutune(
    configs=get_triton_configs(),
    key=HCUTUNE_KEY_V2,
    perf_debug=True,
    prune_configs_by={"early_config_prune": prune_configs},
    warmup=TUNE_DO_BENCH_WARMUP,
    rep=TUNE_DO_BENCH_REP,
)(_fwd_kernel_v2)

fn_v2_decode = triton.utils.hcutune(
    configs=get_triton_configs(),
    key=HCUTUNE_KEY_V2,
    perf_debug=True,
    prune_configs_by={"early_config_prune": prune_configs},
    warmup=TUNE_DO_BENCH_WARMUP,
    rep=TUNE_DO_BENCH_REP,
)(_fwd_kernel_v2_decode)


def extend_attention_fwd(
    q_extend,
    k_extend,
    v_extend,
    o_extend,
    k_buffer,
    v_buffer,
    qo_indptr,
    kv_indptr,
    kv_indices,
    custom_mask,
    is_causal,
    mask_indptr,
    max_len_extend,
    sm_scale=None,
    logit_cap=0.0,
    skip_prefix_custom_mask=True,
    # config: Optional[dict[str, any]] = None,
):
    """
    q_extend, k_extend, v_extend, o_extend: contiguous tensors

    k_buffer, v_buffer: (prefix + extend) tensors in mem_manager
    """
    Lq, Lv = (
        q_extend.shape[-1],
        v_extend.shape[-1],
    )

    if Lq == 576:
        BLOCK_DMODEL = 512
        BLOCK_DPE = 64
    elif Lq == 288:
        BLOCK_DMODEL = 256
        BLOCK_DPE = 32
    elif Lq == 192:
        BLOCK_DMODEL = 128
        BLOCK_DPE = 64
    else:
        BLOCK_DMODEL = triton.next_power_of_2(Lq)
        BLOCK_DPE = 0
    BLOCK_DV = triton.next_power_of_2(Lv)

    # BLOCK_M, BLOCK_N = (64, 64)
    # num_warps = 4

    sm_scale = sm_scale or 1.0 / (Lq**0.5)
    batch_size, head_num = qo_indptr.shape[0] - 1, q_extend.shape[1]
    kv_group_num = q_extend.shape[1] // k_extend.shape[1]

    USE_CUSTOM_MASK = custom_mask is not None
    # Skip custom mask for prefix part
    SKIP_PREFIX_CUSTOM_MASK = skip_prefix_custom_mask

    if not USE_CUSTOM_MASK:
        custom_mask = torch.tensor([0], dtype=torch.bool, device=q_extend.device)
        mask_indptr = torch.tensor([0], dtype=torch.int32, device=q_extend.device)

    grid = lambda META: (batch_size, head_num, triton.cdiv(max_len_extend, META["BLOCK_M"]))
    # num_stages = 1

    # extra_kargs = {}

    # extra_kargs = {"waves_per_eu": 1, "matrix_instr_nonkdim": 16, "kpack": 2}

    fn[grid](
        q_extend,
        k_extend,
        v_extend,
        o_extend,
        k_buffer,
        v_buffer,
        qo_indptr,
        kv_indptr,
        kv_indices,
        custom_mask,
        mask_indptr,
        sm_scale,
        kv_group_num,
        q_extend.stride(0),
        q_extend.stride(1),
        k_extend.stride(0),
        k_extend.stride(1),
        v_extend.stride(0),
        v_extend.stride(1),
        o_extend.stride(0),
        o_extend.stride(1),
        k_buffer.stride(0),
        k_buffer.stride(1),
        v_buffer.stride(0),
        v_buffer.stride(1),
        logit_cap=logit_cap,
        BLOCK_DMODEL=BLOCK_DMODEL,
        BLOCK_DPE=BLOCK_DPE,
        BLOCK_DV=BLOCK_DV,
        # BLOCK_M=BLOCK_M,
        # BLOCK_N=BLOCK_N,
        Lq=Lq,
        Lv=Lv,
        USE_CUSTOM_MASK=USE_CUSTOM_MASK,
        IS_CAUSAL=is_causal,
        SKIP_PREFIX_CUSTOM_MASK=SKIP_PREFIX_CUSTOM_MASK,
        STORE_TRANSPOSE=True,
        # num_warps=num_warps,
        # num_stages=num_stages,
        # **config,
    )


def extend_attention_fwd_v2(
    q_extend,
    k_extend,
    v_extend,
    o_extend,
    k_buffer,
    v_buffer,
    qo_indptr,
    kv_indptr,
    kv_indices,
    custom_mask,
    is_causal,
    mask_indptr,
    max_len_extend,
    sm_scale,
    k_scale,
    v_scale,
    sliding_window_size,
    sinks,
    window_kv_offsets,
    xai_temperature_len,
    skip_prefix_custom_mask,
    logit_cap,
):
    """Launch :data:`fn_v2` (hcutune-wrapped :func:`_fwd_kernel_v2`); kwargs align with ``extend_attention_fwd`` v2 path."""
    Lq, Lv = q_extend.shape[-1], v_extend.shape[-1]
    if Lq == 576:
        BLOCK_DMODEL, BLOCK_DPE = 512, 64
    elif Lq == 288:
        BLOCK_DMODEL, BLOCK_DPE = 256, 32
    elif Lq == 192:
        BLOCK_DMODEL, BLOCK_DPE = 128, 64
    else:
        BLOCK_DMODEL = triton.next_power_of_2(Lq)
        BLOCK_DPE = 0
    BLOCK_DV = triton.next_power_of_2(Lv)
    sm_scale = sm_scale or 1.0 / (Lq**0.5)
    batch_size, head_num = qo_indptr.shape[0] - 1, q_extend.shape[1]
    kv_group_num = q_extend.shape[1] // k_extend.shape[1]
    USE_CUSTOM_MASK = custom_mask is not None
    SKIP_PREFIX_CUSTOM_MASK = skip_prefix_custom_mask
    HAS_SINK = sinks is not None
    if not USE_CUSTOM_MASK:
        custom_mask = torch.tensor([0], dtype=torch.bool, device=q_extend.device)
        mask_indptr = torch.tensor([0], dtype=torch.int32, device=q_extend.device)
    elif sliding_window_size > 0 and window_kv_offsets is None:
        window_kv_offsets = torch.zeros(
            batch_size, dtype=torch.int32, device=q_extend.device
        )
    grid = lambda META: (
        batch_size,
        head_num,
        triton.cdiv(max_len_extend, META["BLOCK_M"]),
    )
    stride_args = (
        q_extend.stride(0),
        q_extend.stride(1),
        k_extend.stride(0),
        k_extend.stride(1),
        v_extend.stride(0),
        v_extend.stride(1),
        o_extend.stride(0),
        o_extend.stride(1),
        k_buffer.stride(0),
        k_buffer.stride(1),
        v_buffer.stride(0),
        v_buffer.stride(1),
    )
    fn_v2[grid](
        q_extend,
        k_extend,
        v_extend,
        o_extend,
        k_buffer,
        v_buffer,
        qo_indptr,
        kv_indptr,
        kv_indices,
        custom_mask,
        mask_indptr,
        sinks,
        window_kv_offsets,
        sm_scale,
        k_scale,
        v_scale,
        kv_group_num,
        *stride_args,
        SLIDING_WINDOW_SIZE=sliding_window_size,
        logit_cap=logit_cap,
        xai_temperature_len=xai_temperature_len,
        HAS_SINK=HAS_SINK,
        BLOCK_DMODEL=BLOCK_DMODEL,
        BLOCK_DPE=BLOCK_DPE,
        BLOCK_DV=BLOCK_DV,
        Lq=Lq,
        Lv=Lv,
        USE_CUSTOM_MASK=USE_CUSTOM_MASK,
        IS_CAUSAL=is_causal,
        SKIP_PREFIX_CUSTOM_MASK=SKIP_PREFIX_CUSTOM_MASK,
        STORE_TRANSPOSE=True,
    )


def extend_attention_fwd_v2_decode(
    q_extend,
    k_extend,
    v_extend,
    o_extend,
    k_buffer,
    v_buffer,
    qo_indptr,
    kv_indptr,
    kv_indices,
    custom_mask,
    is_causal,
    mask_indptr,
    max_len_extend,
    sm_scale,
    k_scale,
    v_scale,
    sliding_window_size,
    sinks,
    window_kv_offsets,
    xai_temperature_len,
    skip_prefix_custom_mask,
    logit_cap,
):
    """Launch :data:`fn_v2_decode` (hcutune-wrapped :func:`_fwd_kernel_v2_decode`); grid-2 is KV head."""
    Lq, Lv = q_extend.shape[-1], v_extend.shape[-1]
    if Lq == 576:
        BLOCK_DMODEL, BLOCK_DPE = 512, 64
    elif Lq == 288:
        BLOCK_DMODEL, BLOCK_DPE = 256, 32
    elif Lq == 192:
        BLOCK_DMODEL, BLOCK_DPE = 128, 64
    else:
        BLOCK_DMODEL = triton.next_power_of_2(Lq)
        BLOCK_DPE = 0
    BLOCK_DV = triton.next_power_of_2(Lv)
    sm_scale = sm_scale or 1.0 / (Lq**0.5)
    batch_size, head_num = qo_indptr.shape[0] - 1, q_extend.shape[1]
    kv_head_num = k_extend.shape[1]
    kv_group_num = q_extend.shape[1] // kv_head_num
    USE_CUSTOM_MASK = custom_mask is not None
    SKIP_PREFIX_CUSTOM_MASK = skip_prefix_custom_mask
    HAS_SINK = sinks is not None
    if not USE_CUSTOM_MASK:
        custom_mask = torch.tensor([0], dtype=torch.bool, device=q_extend.device)
        mask_indptr = torch.tensor([0], dtype=torch.int32, device=q_extend.device)
    elif sliding_window_size > 0 and window_kv_offsets is None:
        window_kv_offsets = torch.zeros(
            batch_size, dtype=torch.int32, device=q_extend.device
        )

    def grid(META):
        q_seq = max(META["BLOCK_M"] // kv_group_num, 1)
        return (
            batch_size,
            kv_head_num,
            triton.cdiv(max_len_extend, q_seq),
        )

    stride_args = (
        q_extend.stride(0),
        q_extend.stride(1),
        k_extend.stride(0),
        k_extend.stride(1),
        v_extend.stride(0),
        v_extend.stride(1),
        o_extend.stride(0),
        o_extend.stride(1),
        k_buffer.stride(0),
        k_buffer.stride(1),
        v_buffer.stride(0),
        v_buffer.stride(1),
    )
    fn_v2_decode[grid](
        q_extend,
        k_extend,
        v_extend,
        o_extend,
        k_buffer,
        v_buffer,
        qo_indptr,
        kv_indptr,
        kv_indices,
        custom_mask,
        mask_indptr,
        sinks,
        window_kv_offsets,
        sm_scale,
        k_scale,
        v_scale,
        *stride_args,
        SLIDING_WINDOW_SIZE=sliding_window_size,
        logit_cap=logit_cap,
        xai_temperature_len=xai_temperature_len,
        HAS_SINK=HAS_SINK,
        kv_group_num=kv_group_num,
        num_query_heads=head_num,
        BLOCK_DMODEL=BLOCK_DMODEL,
        BLOCK_DPE=BLOCK_DPE,
        BLOCK_DV=BLOCK_DV,
        Lq=Lq,
        Lv=Lv,
        USE_CUSTOM_MASK=USE_CUSTOM_MASK,
        IS_CAUSAL=is_causal,
        SKIP_PREFIX_CUSTOM_MASK=SKIP_PREFIX_CUSTOM_MASK,
        STORE_TRANSPOSE=True,
    )


def get_bench_inputs_v2():
    """Cases for :func:`_fwd_kernel_v2` / ``_get_config_v2`` keys.

    对 ``fp8_dp2_tp8_415_triton_rocm_nomtp.log`` 逐条核对后，与 extend 相关的**键控组合**只有两类：

    - **MHA**（``k_extend[...,1,...]``）：``sliding_window_size=-1``，``sinks=None``（如 log 中 ``k_extend [1152,1,192]`` / ``[3952,1,192]`` 段）。
    - **GQA**（``k_extend[...,2,...]``）：``sliding_window_size=128``，``sinks shape [16]``（全文件未出现 GQA 与 ``-1``/无 sinks 的同框记录）。

    该 log 中 **未出现** ``sliding_window_size=64``，也未出现「GQA + 无 SWA + 无 sinks」；若需覆盖其它模型再单独加行并注明来源。

    额外两行覆盖 **causal=True + custom_mask=True**（flatten 因果 mask，与 kernel 中 ``USE_CUSTOM_MASK`` 分支一致；want7 第 4 维为 True）。
    """
    names = [
        "B",
        "H",
        "prefix",
        "extend",
        "kv_lora_rank",
        "qk_rope_head_dim",
        "v_head_dim",
        "causal",
        "custom_mask",
        "sliding_window_size",
        "has_sink",
        "kv_num_heads",
    ]
    vals = [
        # (prefix, extend) 只影响访存/grid；want7 由 head/dim/SWA/sinks 决定。prefix=8192、extend=1024 与长 KV bench 习惯一致。
        # (1) GQA Q16/KV2：与 log 中 ``[3952,2,192]`` + ``sliding_window_size=128`` + ``sinks [16]`` 一致；want7 (8,192,128,F,T,T,True)
        (4, 16, 8192, 1024, 128, 64, 128, True, False, 128, True, 2),
        # (2) MHA Q16/KV1：与 log 中 ``[...,1,192]`` + ``sliding_window_size=-1`` + ``sinks=None`` 一致；want7 (16,192,128,F,T,F,False)
        (4, 16, 8192, 1024, 128, 64, 128, True, False, -1, False, 1),
        # (3) GQA + custom causal mask；want7 (8,192,128,T,T,T,True)
        (4, 16, 8192, 1024, 128, 64, 128, True, True, 128, True, 2),
        # (4) MHA + custom causal mask；want7 (16,192,128,T,T,F,False)
        (4, 16, 8192, 1024, 128, 64, 128, True, True, -1, False, 1),
    ]
    return names, vals


def get_bench_inputs_v2_decode():
    """Cases for :func:`_fwd_kernel_v2_decode` (short ``extend`` / decode grid).

    Fixed ``B=64``, ``prefix=256``, ``extend=4``, ``equal_seqlens``; causal / custom_mask /
    sliding_window / sinks / GQA|MHA 与 :func:`get_bench_inputs_v2` 四行一一对应。
    """
    names = [
        "B",
        "H",
        "prefix",
        "extend",
        "kv_lora_rank",
        "qk_rope_head_dim",
        "v_head_dim",
        "causal",
        "custom_mask",
        "sliding_window_size",
        "has_sink",
        "kv_num_heads",
    ]
    vals = [
        (64, 16, 256, 4, 128, 64, 128, True, False, 128, True, 2),
        (64, 16, 256, 4, 128, 64, 128, True, False, -1, False, 1),
        (64, 16, 256, 4, 128, 64, 128, True, True, 128, True, 2),
        (64, 16, 256, 4, 128, 64, 128, True, True, -1, False, 1),
    ]
    return names, vals


x_names, x_vals = get_bench_inputs()
configs = [
    triton.testing.Benchmark(
        x_names=x_names,
        x_vals=x_vals,
        line_arg="provider",
        line_vals=["triton"],
        line_names=["triton"],
        styles=[("red", "-")],
        ylabel="ms",
        plot_name="extend_attention",
        args={'dtype': torch.float16},
    )
]

@triton.utils.dist_perf_report(configs)
def bench_extend_attention(B, H, prefix, extend, kv_lora_rank, qk_rope_head_dim,
                            v_head_dim, causal, custom_mask, provider, dtype):
    torch.manual_seed(42)

    device = 'cpu' if os.getenv("TRITON_HCUTUNE_COMPILE_ONLY", "") == "1" else "cuda"
    ref_attn_impl = "normal"
    sm_scale = 1.0
    logit_cap = 0.0

    (
        q_extend,
        k_extend,
        v_extend,
        k_buffer,
        v_buffer,
        kv_indptr,
        kv_indices,
        qo_indptr,
        custom_mask,
        mask_indptr,
        max_len_extend,
    ) = input_helper(
        B,
        H,
        prefix,
        extend,
        kv_lora_rank,
        qk_rope_head_dim,
        v_head_dim,
        dtype,
        device,
        ref_attn_impl,
    )
    tri_out = torch.empty(
        (*q_extend.shape[:-1], v_extend.shape[-1]),
        dtype=q_extend.dtype,
        device=q_extend.device,
    )

    fn = lambda: extend_attention_fwd(
        q_extend,
        k_extend,
        v_extend,
        tri_out,
        k_buffer,
        v_buffer,
        qo_indptr,
        kv_indptr,
        kv_indices,
        custom_mask,
        causal,
        mask_indptr,
        max_len_extend,
        sm_scale=sm_scale,
        logit_cap=logit_cap,
    )
    return triton.testing.do_bench_cudagraph(fn)


x_names_v2, x_vals_v2 = get_bench_inputs_v2()
configs_v2 = [
    triton.testing.Benchmark(
        x_names=x_names_v2,
        x_vals=x_vals_v2,
        line_arg="provider",
        line_vals=["triton_v2"],
        line_names=["triton_v2"],
        styles=[("blue", "-")],
        ylabel="ms",
        plot_name="extend_attention_v2_hcutune",
        args={"dtype": torch.bfloat16},
    )
]


@triton.utils.dist_perf_report(configs_v2)
def bench_extend_attention_v2(
    B,
    H,
    prefix,
    extend,
    kv_lora_rank,
    qk_rope_head_dim,
    v_head_dim,
    causal,
    custom_mask,
    sliding_window_size,
    has_sink,
    kv_num_heads,
    provider,
    dtype,
):
    torch.manual_seed(42)
    device = "cpu" if os.getenv("TRITON_HCUTUNE_COMPILE_ONLY", "") == "1" else "cuda"
    ref_attn_impl = "normal"
    logit_cap = 0.0
    k_scale = 1.0
    v_scale = 1.0
    xai_temperature_len = -1
    skip_prefix_custom_mask = True

    (
        q_extend,
        k_extend,
        v_extend,
        k_buffer,
        v_buffer,
        kv_indptr,
        kv_indices,
        qo_indptr,
        custom_mask_t,
        mask_indptr,
        max_len_extend,
    ) = input_helper(
        B,
        H,
        prefix,
        extend,
        kv_lora_rank,
        qk_rope_head_dim,
        v_head_dim,
        dtype,
        device,
        ref_attn_impl,
        equal_seqlens=True,
        kv_num_heads=kv_num_heads,
    )
    if custom_mask:
        custom_mask_t, mask_indptr = _v2_flat_causal_custom_mask(
            B, prefix, extend, device
        )

    sm_scale = float(1.0 / (q_extend.shape[-1] ** 0.5))
    sinks = (
        torch.randn(H, dtype=q_extend.dtype, device=device)
        if has_sink
        else None
    )
    window_kv_offsets = None
    if custom_mask and sliding_window_size > 0:
        window_kv_offsets = torch.zeros(B, dtype=torch.int32, device=device)

    tri_out = torch.empty(
        (*q_extend.shape[:-1], v_extend.shape[-1]),
        dtype=q_extend.dtype,
        device=q_extend.device,
    )

    def run_once():
        extend_attention_fwd_v2(
            q_extend,
            k_extend,
            v_extend,
            tri_out,
            k_buffer,
            v_buffer,
            qo_indptr,
            kv_indptr,
            kv_indices,
            custom_mask_t,
            causal,
            mask_indptr,
            max_len_extend,
            sm_scale,
            k_scale,
            v_scale,
            sliding_window_size,
            sinks,
            window_kv_offsets,
            xai_temperature_len,
            skip_prefix_custom_mask,
            logit_cap,
        )

    return triton.testing.do_bench_cudagraph(run_once)


x_names_v2_decode, x_vals_v2_decode = get_bench_inputs_v2_decode()
configs_v2_decode = [
    triton.testing.Benchmark(
        x_names=x_names_v2_decode,
        x_vals=x_vals_v2_decode,
        line_arg="provider",
        line_vals=["triton_v2_decode"],
        line_names=["triton_v2_decode"],
        styles=[("green", "-")],
        ylabel="ms",
        plot_name="extend_attention_v2_decode_hcutune",
        args={"dtype": torch.bfloat16},
    )
]


@triton.utils.dist_perf_report(configs_v2_decode)
def bench_extend_attention_v2_decode(
    B,
    H,
    prefix,
    extend,
    kv_lora_rank,
    qk_rope_head_dim,
    v_head_dim,
    causal,
    custom_mask,
    sliding_window_size,
    has_sink,
    kv_num_heads,
    provider,
    dtype,
):
    torch.manual_seed(42)
    device = "cpu" if os.getenv("TRITON_HCUTUNE_COMPILE_ONLY", "") == "1" else "cuda"
    ref_attn_impl = "normal"
    logit_cap = 0.0
    k_scale = 1.0
    v_scale = 1.0
    xai_temperature_len = -1
    skip_prefix_custom_mask = True

    (
        q_extend,
        k_extend,
        v_extend,
        k_buffer,
        v_buffer,
        kv_indptr,
        kv_indices,
        qo_indptr,
        custom_mask_t,
        mask_indptr,
        max_len_extend,
    ) = input_helper(
        B,
        H,
        prefix,
        extend,
        kv_lora_rank,
        qk_rope_head_dim,
        v_head_dim,
        dtype,
        device,
        ref_attn_impl,
        equal_seqlens=True,
        kv_num_heads=kv_num_heads,
    )
    if custom_mask:
        custom_mask_t, mask_indptr = _v2_flat_causal_custom_mask(
            B, prefix, extend, device
        )

    sm_scale = float(1.0 / (q_extend.shape[-1] ** 0.5))
    sinks = (
        torch.randn(H, dtype=q_extend.dtype, device=device)
        if has_sink
        else None
    )
    window_kv_offsets = None
    if custom_mask and sliding_window_size > 0:
        window_kv_offsets = torch.zeros(B, dtype=torch.int32, device=device)

    tri_out = torch.empty(
        (*q_extend.shape[:-1], v_extend.shape[-1]),
        dtype=q_extend.dtype,
        device=q_extend.device,
    )

    def run_once():
        extend_attention_fwd_v2_decode(
            q_extend,
            k_extend,
            v_extend,
            tri_out,
            k_buffer,
            v_buffer,
            qo_indptr,
            kv_indptr,
            kv_indices,
            custom_mask_t,
            causal,
            mask_indptr,
            max_len_extend,
            sm_scale,
            k_scale,
            v_scale,
            sliding_window_size,
            sinks,
            window_kv_offsets,
            xai_temperature_len,
            skip_prefix_custom_mask,
            logit_cap,
        )

    return triton.testing.do_bench_cudagraph(run_once)


def parse_args():
    parser = argparse.ArgumentParser()

    parser.add_argument("--perf", action="store_true", default=False,
                        help="benchmark with hcutuner perf mode")
    parser.add_argument(
        "--v1",
        action="store_true",
        default=False,
        help="Tune v1 ``_fwd_kernel`` only.",
    )
    parser.add_argument(
        "--v2-decode",
        action="store_true",
        default=False,
        help="Tune v2 decode ``_fwd_kernel_v2_decode`` (B=64, prefix=256, extend=4). Default: v2 prefill ``_fwd_kernel_v2``.",
    )

    return parser.parse_args()


if __name__ == "__main__":
    args = parse_args()

    if args.perf:
        os.environ["TRITON_HCUTUNE_PERF_MODE"] = "1"  # idempotent; real set is at top for Hcutuner init

    if args.v1:
        bench_extend_attention.run(print_data=True, save_path="./tune_extend_attention_out")
    elif args.v2_decode:
        bench_extend_attention_v2_decode.run(
            print_data=True, save_path="./tune_extend_attention_v2_decode_out"
        )
    else:
        bench_extend_attention_v2.run(
            print_data=True, save_path="./tune_extend_attention_v2_out"
        )
