# SPDX-License-Identifier: MIT
 
import math
import random
import pytest
import torch
import triton
import os
from aiter.ops.triton.chunked_pa_prefill import chunked_prefill_paged_decode
from aiter.ops.triton.utils.types import str_to_torch_dtype

NUM_HEADS = [32]
NUM_QUERIES_PER_KV = [1, 4, 8, 16]
# NUM_QUERIES_PER_KV = [4]

HEAD_SIZES = [128, 96, 24]
# HEAD_SIZES = [128]

CACHE_BLOCK_SIZES = [128, 64, 32]
# CACHE_BLOCK_SIZES = [32]

DTYPES = [torch.float16]
CUDA_DEVICES = [f"cuda:{i}" for i in range(1)]

# SLIDING_WINDOW = [0, 16, 64, 128, 256, 512, 2048]
SLIDING_WINDOW = [0]

KV_CACHE_DTYPES = ["auto", "fp8e4m3", "fp8e5m2"]
# KV_CACHE_DTYPES = ["auto"]

OPS = [chunked_prefill_paged_decode]
dump_data = int(os.environ.get("DUMP_DATA", "0"))
run_with_dump_data = int(os.environ.get("RUN_WITH_DUMP", "0"))
print(f"{dump_data=} {run_with_dump_data=} ..................................")

def context_attention_fwd_torch(
    query,  # [num_tokens, H, D]
    k,  # [num_tokens, Hkv, D]
    v,  # [num_tokens, Hkv, D]
    output,  # [num_tokens, H, D]
    k_cache,  # [B, Hkv, D/8, Blk_sz, 8]
    v_cache,  # [B, Hkv, D, Blk_sz]
    b_start_loc,  # [B+1]
    b_seq_len,  # [B]
    k_scale,
    v_scale,
    alibi_slopes=None,
    sliding_window=None,
    sm_scale=None,
):
    # Setup
    num_blocks = b_seq_len.shape[0]
    head_dim = query.shape[-1]
    num_heads = query.shape[1]
    num_kv_heads = k.shape[1]
    num_queries_per_kv = num_heads // num_kv_heads
    device = query.device

    # Softmax scale fallback
    if sm_scale is None:
        sm_scale = 1.0 / (head_dim**0.5)

    is_kv_cache_fp8 = torch.finfo(k_cache.dtype).bits == 8

    # Cast all inputs to float32
    query = query.to(torch.float32)
    k = k.to(torch.float32)
    v = v.to(torch.float32)
    k_cache = k_cache.to(torch.float32)
    v_cache = v_cache.to(torch.float32)

    for b in range(num_blocks):
        q_start = b_start_loc[b]
        q_end = b_start_loc[b + 1]
        q_len = q_end - q_start
        ctx_len = b_seq_len[b] - q_len

        q = query[q_start:q_end]  # [q_len, H, D]
        k_local = k[q_start:q_end]  # [q_len, Hkv, D]
        v_local = v[q_start:q_end]  # [q_len, Hkv, D]

        for h in range(num_heads):
            kv_h = h // num_queries_per_kv

            qh = q[:, h]  # [q_len, D]

            kc = k_cache[:, kv_h]  # [B, D//8, Blk_sz, 8]
            kc = kc.permute(0, 2, 1, 3).reshape(-1, head_dim)  # [B * Blk_sz, D]
            kc = kc[:ctx_len]  # [ctx_len, D]

            vc = v_cache[:, kv_h]  # [B, D, Blk_sz]
            vc = vc.permute(0, 2, 1).reshape(-1, head_dim)  # [B * Blk_sz, D]
            vc = vc[:ctx_len]  # [ctx_len, D]

            if is_kv_cache_fp8:
                kc = kc * k_scale
                vc = vc * v_scale

            # Compute query against context
            qk_ctx = torch.matmul(qh, kc.T)
            qk_ctx *= sm_scale

            if sliding_window and sliding_window > 0:
                q_pos = torch.arange(ctx_len, ctx_len + q_len, device=device)
                k_pos = torch.arange(ctx_len, device=device)
                rel_dist = q_pos[:, None] - k_pos[None, :]
                qk_ctx = qk_ctx.masked_fill(rel_dist >= sliding_window, -1e4)

            elif alibi_slopes is not None:
                alibi_slope = alibi_slopes[h]
                q_pos = torch.arange(ctx_len, ctx_len + q_len, device=device)[:, None]
                k_pos = torch.arange(ctx_len, device=device)[None, :]
                rel_pos = k_pos - q_pos
                alibi_bias = rel_pos.to(torch.float32) * alibi_slope
                mask = (rel_pos <= 0) & (q_pos < b_seq_len[b])
                alibi_bias = torch.where(mask, alibi_bias, float("-inf"))
                qk_ctx += alibi_bias

            p_ctx = torch.softmax(qk_ctx, dim=-1)
            acc = torch.matmul(p_ctx, vc)

            # Compute query against itself (with causal mask)
            kh = k_local[:, kv_h]  # [q_len, D]
            vh = v_local[:, kv_h]  # [q_len, D]

            qk_self = torch.matmul(qh, kh.T)
            qk_self *= sm_scale

            causal_mask = torch.triu(
                torch.ones(q_len, q_len, dtype=torch.bool, device=device), 1
            )
            qk_self = qk_self.masked_fill(causal_mask, float("-inf"))

            if sliding_window and sliding_window > 0:
                q_pos = torch.arange(q_len, device=device)
                k_pos = torch.arange(q_len, device=device)
                rel_dist = q_pos[:, None] - k_pos[None, :]
                qk_self = qk_self.masked_fill(rel_dist >= sliding_window, -10000)

            if alibi_slopes is not None:
                alibi_slope = alibi_slopes[h]
                q_pos = torch.arange(q_len, device=device)[:, None]
                k_pos = torch.arange(q_len, device=device)[None, :]
                rel_pos = k_pos - q_pos
                alibi_bias = rel_pos.to(torch.float32) * alibi_slope
                mask = (rel_pos <= 0) & (q_pos < q_len)
                alibi_bias = torch.where(mask, alibi_bias, float("-inf"))
                qk_self += alibi_bias

            p_self = torch.softmax(qk_self, dim=-1)
            acc_self = torch.matmul(p_self, vh)

            # Output
            acc_total = acc + acc_self
            output[q_start:q_end, h] = acc_total.to(output.dtype)
    return


def _get_alibi_slopes(total_num_heads: int) -> torch.Tensor:
    closest_power_of_2 = 2 ** math.floor(math.log2(total_num_heads))
    base = torch.tensor(
        2 ** (-(2 ** -(math.log2(closest_power_of_2) - 3))),
        dtype=torch.float32,
    )
    powers = torch.arange(1, 1 + closest_power_of_2, dtype=torch.int32)
    slopes = torch.pow(base, powers)

    if closest_power_of_2 != total_num_heads:
        extra_base = torch.tensor(
            2 ** (-(2 ** -(math.log2(2 * closest_power_of_2) - 3))),
            dtype=torch.float32,
        )
        num_remaining_heads = min(
            closest_power_of_2, total_num_heads - closest_power_of_2
        )
        extra_powers = torch.arange(
            start=1, end=1 + 2 * num_remaining_heads, step=2, dtype=torch.int32
        )
        slopes = torch.cat([slopes, torch.pow(extra_base, extra_powers)], dim=0)
    return slopes

def seed_everything(seed):
    random.seed(seed)
    torch.manual_seed(seed)

def save_inputs_for_debug(
    query, k, v, key, value, output,
    k_cache, v_cache, block_table,
    b_start_loc, b_seq_len, max_input_len,
    k_scale, v_scale,
    filename="paged_attention_inputs.pt"
):
    # 将所有输入保存到字典中
    inputs_dict = {
        'query': query,
        'k': k,
        'v': v,
        'key':key,
        'value':value,
        'output': output,
        'k_cache': k_cache,
        'v_cache': v_cache,
        'block_table': block_table,
        'b_start_loc': b_start_loc,
        'b_seq_len': b_seq_len,
        'max_input_len': max_input_len,
        'k_scale': k_scale,
        'v_scale': v_scale,
    }

    # 保存到文件
    torch.save(inputs_dict, filename)
    print(f"输入已保存到: {filename}")

def load_inputs_for_debug(device, filename="paged_attention_inputs_kv_fp8.pt"):
    # 从文件加载
    inputs_dict = torch.load(filename)

    # 返回所有变量
    return (
        inputs_dict['query'].to(device),
        inputs_dict['k'].to(device),
        inputs_dict['v'].to(device),
        inputs_dict['key'].to(device),
        inputs_dict['value'].to(device),
        inputs_dict['output'].to(device),
        inputs_dict['k_cache'].to(device),
        inputs_dict['v_cache'].to(device),
        inputs_dict['block_table'].to(device),
        inputs_dict['b_start_loc'].to(device),
        inputs_dict['b_seq_len'].to(device),
        inputs_dict['max_input_len'],
        inputs_dict['k_scale'].to(device),
        inputs_dict['v_scale'].to(device)
    )

def save_golden_for_debug(
    golden,
    filename="paged_attention_goldens.pt"
):
    # 将所有输入保存到字典中
    goldens_dict = {
        'golden': golden,
    }

    # 保存到文件
    torch.save(goldens_dict, filename)
    print(f"golden已保存到: {filename}")

def load_goldens_for_debug(device, filename="paged_attention_goldens.pt"):
    # 从文件加载
    goldens_dict = torch.load(filename)

    # 返回所有变量
    return (
        goldens_dict['golden'].to(device),
    )

def compute_golden_attention(query, key, value, b_start_loc, b_seq_len, num_heads, num_queries_per_kv, head_size, dtype, device, sliding_window=0):
    """
    计算标准的 attention 作为 golden reference
    使用 causal mask，只允许关注当前位置及之前的位置
    精简参数版本 - 优化计算
    """
    num_tokens = query.shape[0]
    num_kv_heads = num_heads // num_queries_per_kv
    batch_size = len(b_seq_len)

    # 创建输出 tensor
    output = torch.empty(num_tokens, num_heads, head_size, dtype=dtype, device=device)

    # print("开始计算 Golden Attention (Causal Mask)...")

    # 在循环外部预先计算所有必要的信息
    b_seq_start_loc = torch.cumsum(torch.cat([torch.tensor([0], device=device), b_seq_len[:-1]]), dim=0)

    # 预先计算所有 batch 的 query_lens
    query_lens = torch.diff(b_start_loc)
    # 预先计算所有 batch 的 ctx_lens
    ctx_lens = b_seq_len - query_lens
    # print(f"{b_start_loc=} {query_lens=} {b_seq_start_loc=} {ctx_lens=} {b_seq_len=}")

    # 为每个 batch 单独计算
    for batch_idx in range(batch_size):
        seq_start = b_seq_start_loc[batch_idx]
        query_start = b_start_loc[batch_idx]
        query_len_batch = query_lens[batch_idx]
        ctx_len_batch = ctx_lens[batch_idx]

        # 当前序列的完整 key 和 value
        seq_key = key[seq_start:seq_start + b_seq_len[batch_idx]]  # [seq_len, num_kv_heads, head_size]
        seq_value = value[seq_start:seq_start + b_seq_len[batch_idx]]  # [seq_len, num_kv_heads, head_size]

        # 当前序列的 query（new tokens）
        seq_query = query[query_start:query_start + query_len_batch]  # [query_len, num_heads, head_size]

        # 将 key 和 value 扩展到与 query 相同的 head 数量
        key_expanded = seq_key.repeat_interleave(num_queries_per_kv, dim=1)
        value_expanded = seq_value.repeat_interleave(num_queries_per_kv, dim=1)

        # 对于每个 new token 位置（只遍历 query 长度）
        for token_idx in range(query_len_batch):
            # 当前 token 在完整序列中的位置
            current_position = ctx_len_batch + token_idx
            # 获取当前 token 的 query
            current_query = seq_query[token_idx]  # [num_heads, head_size]
            # 计算 attention scores - 使用矩阵运算避免循环
            attention_scores = torch.einsum('hd,shd->sh', current_query, key_expanded) / (head_size ** 0.5)
            # 应用 causal mask
            causal_mask = torch.arange(b_seq_len[batch_idx], device=device) > current_position
            attention_scores = attention_scores.masked_fill(causal_mask.unsqueeze(1), -float('inf'))
            if sliding_window > 0:
                sliding_mask = current_position - torch.arange(b_seq_len[batch_idx], device=device) >= sliding_window
                attention_scores = attention_scores.masked_fill(sliding_mask.unsqueeze(1), -10000)
            # 应用 softmax
            attention_weights = torch.softmax(attention_scores, dim=0)  # [seq_len, num_heads]
            # 计算加权和 - 使用矩阵运算
            weighted_sum = torch.einsum('sh,shd->hd', attention_weights, value_expanded)
            # 保存结果
            output[query_start + token_idx] = weighted_sum

    # print("Golden Attention (Causal Mask) 计算完成!")
    return output

def quantize_to_fp8(tensor, dtype="fp8", quant_method="rms"):
    """
    FP16 到 FP8 量化
    """
    if dtype not in ["fp8", "fp8e4m3", "fp8e5m2"]:
        return tensor

    # 根据不同的 FP8 格式设置范围
    if dtype == "fp8" or dtype == "fp8e4m3":
        # FP8 E4M3 范围: [-448, 448]
        fp8_max = torch.finfo(torch.float8_e4m3fn).max
        # print(f"{fp8_max=}")
    elif dtype == "fp8e5m2":
        # FP8 E5M2 范围: [-57344, 57344]
        fp8_max = torch.finfo(torch.float8_e5m2).max
        # print(f"{fp8_max=}")
    else:
        raise ValueError(f"Unsupported FP8 format: {dtype}")

    tensor_fp32 = tensor.to(torch.float32)
    # 计算 scale
    abs_tensor = torch.abs(tensor_fp32)
    if quant_method == "rms":
        # RMS缩放，减少异常值影响
        rms = torch.sqrt(torch.mean(tensor_fp32 ** 2))
        abs_max = torch.max(abs_tensor)
        # 加权平均：70% RMS + 30% 最大值
        scale_base = 0.7 * rms + 0.3 * abs_max
        scale = scale_base / fp8_max
    elif quant_method == "percentile":
        # 95%分位数，避免极端值
        scale_base = torch.quantile(abs_tensor, 0.95)
        scale = scale_base / fp8_max

    elif quant_method == "adaptive":
        # 自适应方法：根据数据分布选择策略
        abs_mean = torch.mean(abs_tensor)
        abs_max = torch.max(abs_tensor)

        # 如果存在极端值（最大值远大于均值），使用分位数
        if abs_max / abs_mean > 8:
            scale_base = torch.quantile(abs_tensor, 0.95)
        else:
            # 否则使用RMS
            scale_base = torch.sqrt(torch.mean(tensor_fp32 ** 2))

        scale = scale_base / fp8_max
    elif quant_method == "max":
        abs_max = torch.max(abs_tensor)
        scale = abs_max / fp8_max
    else:
        raise ValueError(f"Unsupported quant method: {quant_method}")

    scale = torch.clamp(scale, min=1e-8)
    # 量化
    normalized_tensor = tensor_fp32 / scale
    clipped_tensor = torch.clamp(normalized_tensor, -fp8_max, fp8_max)

    # 转换为 FP8 格式，然后以 uint8 视图输出
    if dtype == "fp8" or dtype == "fp8e4m3":
        # 使用 torch.float8_e4m3fn 如果支持，否则用 uint8 模拟
        if hasattr(torch, 'float8_e4m3fn'):
            tensor_fp8 = clipped_tensor.to(torch.float8_e4m3fn)
        else:
            assert False
    else:  # e5m2
        if hasattr(torch, 'float8_e5m2'):
            tensor_fp8 = clipped_tensor.to(torch.float8_e5m2)
        else:
            assert False
    if (torch.isnan(tensor_fp8).any()):
        assert False, f"tensor_fp8 has nan, {dtype=}"
    return tensor_fp8, scale

def input_helper(
    BS,
    MAX_SEQ_LEN,
    MAX_CTX_LEN,
    cache_size,
    block_size,
    max_block_per_request,
    num_heads: int,
    head_size: int,
    num_queries_per_kv: int,
    dtype: torch.dtype,
    kv_cache_dtype: str,
    device: str,
    use_alibi_slope: bool,
):
    seed_everything(0)
    torch.set_default_device(device)

    # Need this, otherwise when we capture the graph the process
    # for GPU 1 would run on both GPU0 and GPU1 and things would hang
    #
    # see also similar issue: https://github.com/Dao-AILab/flash-attention/issues/523
    torch.cuda.set_device(device)
    assert block_size * max_block_per_request >= MAX_SEQ_LEN + MAX_CTX_LEN, "cache block count may not enough for each request."
    assert cache_size >= max_block_per_request * BS, "cache block count is not enough for total requests."

    if use_alibi_slope:
        alibi_slopes = _get_alibi_slopes(num_heads).to(device)

    query_lens = [random.randint(16, MAX_SEQ_LEN) for _ in range(BS)]
    if BS > 5:
        # ensure some batch is decode
        query_lens[1] = 1
        query_lens[4] = 1
    query_lens[-1] = 1
    ctx_lens = [random.randint(16, MAX_CTX_LEN) for _ in range(BS)]
    seq_lens = [a + b for a, b in zip(query_lens, ctx_lens)]
    num_kv_heads = num_heads // num_queries_per_kv

    num_tokens = sum(query_lens)
    query = torch.empty(num_tokens, num_heads, head_size, dtype=dtype)
    query.uniform_(-1e-1, 1e-1)
    output = torch.empty(num_tokens, num_heads, head_size, dtype=dtype)

    k_scale = v_scale = torch.tensor(1.0, dtype=torch.float32, device=device)
    kv = torch.empty(sum(seq_lens), 2, num_kv_heads, head_size, dtype=dtype)
    kv.uniform_(-1e-1, 1e-1)
    key, value = kv.unbind(dim=1)
    key_q = key
    value_q = value
    if kv_cache_dtype in ["fp8", "fp8e4m3", "fp8e5m2"]:
        kv_q, scale = quantize_to_fp8(kv, kv_cache_dtype)
        k_scale = v_scale = scale
        key_q, value_q = kv_q.unbind(dim=1)

    if kv_cache_dtype == "auto":
        cache_dtype = dtype
    else:
        cache_dtype = str_to_torch_dtype[kv_cache_dtype]

    k_cache = (torch.rand(
        cache_size, block_size, num_kv_heads, head_size, 
        dtype=dtype
    ) * 20 - 10).to(cache_dtype)

    v_cache = (torch.rand(
        cache_size, block_size, num_kv_heads, head_size, 
        dtype=dtype
    ) * 20 - 10).to(cache_dtype)

    k = torch.zeros(sum(query_lens), num_kv_heads, head_size, dtype=dtype)
    v = torch.zeros(sum(query_lens), num_kv_heads, head_size, dtype=dtype)
    values = torch.arange(0, cache_size, dtype=torch.long)
    values = values[torch.randperm(cache_size)]
    block_table = values[: BS * max_block_per_request].view(BS, max_block_per_request)
    b_seq_len = torch.tensor(seq_lens, dtype=torch.long)
    b_ctx_len = torch.tensor(ctx_lens, dtype=torch.long)
    b_start_loc = torch.cumsum(torch.tensor([0] + query_lens, dtype=torch.long), dim=0)
    max_input_len = MAX_SEQ_LEN
    # copy kv to cache
    b_seq_start_loc = torch.cumsum(
        torch.tensor([0] + seq_lens[:-1], dtype=torch.long), dim=0
    )
    query_begin = 0
    for i in range(BS):
        # print(f"batch={i}, seq_lens={seq_lens[i]}, {query_begin=}, query_len={query_lens[i]}, b_ctx_len={b_ctx_len[i]}")
        query_begin += query_lens[i]
        for j in range(query_lens[i]):
            k[b_start_loc[i] + j].copy_(key[b_seq_start_loc[i] + b_ctx_len[i] + j])
            v[b_start_loc[i] + j].copy_(value[b_seq_start_loc[i] + b_ctx_len[i] + j])
        cur_ctx = 0
        block_id = 0
        ctx_len_batch = seq_lens[i]
        while cur_ctx < ctx_len_batch:
            start_loc = b_seq_start_loc[i] + cur_ctx
            if cur_ctx + block_size > ctx_len_batch:
                end_loc = b_seq_start_loc[i] + ctx_len_batch
            else:
                end_loc = start_loc + block_size
            start_slot = block_table[i, block_id] * block_size
            end_slot = start_slot + end_loc - start_loc
            k_cache.view(-1, num_kv_heads, head_size)[start_slot:end_slot].copy_(
                key_q[start_loc:end_loc]
            )
            v_cache.view(-1, num_kv_heads, head_size)[start_slot:end_slot].copy_(
                value_q[start_loc:end_loc]
            )
            cur_ctx += block_size
            block_id += 1
    # transpose K_cache[num_blocks, block_size, num_kv_heads, head_size]
    # to K_cache[num_blocks, num_kv_heads, head_size/8, block_size, 8]
    k_cache = (
        k_cache.view(-1, block_size, num_kv_heads, head_size // 8, 8)
        .permute(0, 2, 3, 1, 4)
        .contiguous()
    )
    # transpose V_cache[num_blocks, block_size, num_kv_heads, head_size]
    # to V_cache[num_blocks, num_kv_heads, head_size, block_size]
    v_cache = (
        v_cache.view(-1, block_size, num_kv_heads, head_size)
        .permute(0, 2, 3, 1)
        .contiguous()
    )

    if use_alibi_slope:
        return (
            query,
            key,
            value,
            k,
            v,
            output,
            k_cache,
            v_cache,
            block_table,
            b_start_loc,
            b_seq_len,
            max_input_len,
            k_scale,
            v_scale,
            alibi_slopes,
        )
    else:
        return (
            query,
            key,
            value,
            k,
            v,
            output,
            k_cache,
            v_cache,
            block_table,
            b_start_loc,
            b_seq_len,
            max_input_len,
            k_scale,
            v_scale,
            None,
        )

@pytest.mark.parametrize("num_heads", NUM_HEADS)
@pytest.mark.parametrize("num_queries_per_kv", NUM_QUERIES_PER_KV)
@pytest.mark.parametrize("head_size", HEAD_SIZES)
@pytest.mark.parametrize("cache_block_size", CACHE_BLOCK_SIZES)
@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPES)
@pytest.mark.parametrize("device", CUDA_DEVICES)
@pytest.mark.parametrize("sliding_window", SLIDING_WINDOW)
@torch.inference_mode()
def test_contexted_kv_attention(
    num_heads: int,
    num_queries_per_kv: int,
    head_size: int,
    cache_block_size: int,
    sliding_window: int,
    dtype: torch.dtype,
    kv_cache_dtype: str,
    device: str,
) -> None:
    if not run_with_dump_data:
        (
            query,
            key,
            value,
            k,
            v,
            output,
            k_cache,
            v_cache,
            block_table,
            b_start_loc,
            b_seq_len,
            max_input_len,
            k_scale,
            v_scale,
            _,
        ) = input_helper(
            BS=10,
            MAX_SEQ_LEN=512,
            MAX_CTX_LEN=512,
            cache_size=640,
            block_size=cache_block_size,
            max_block_per_request=64,
            num_heads=num_heads,
            head_size=head_size,
            num_queries_per_kv=num_queries_per_kv,
            dtype=dtype,
            kv_cache_dtype=kv_cache_dtype,
            device=device,
            use_alibi_slope=False,
        )
        if dump_data:
            torch.cuda.synchronize()
            save_inputs_for_debug(
                query, k, v, key, value, output,
                k_cache, v_cache, block_table,
                b_start_loc, b_seq_len, max_input_len,
                k_scale, v_scale, filename=f"inputs_kv_{num_heads}_{num_queries_per_kv}_{head_size}_{cache_block_size}_{kv_cache_dtype}.pt"
            )
    else:
        (
            query, k, v, key, value, output,
            k_cache, v_cache, block_table,
            b_start_loc, b_seq_len, max_input_len,
            k_scale, v_scale
        ) = load_inputs_for_debug(device, filename=f"inputs_kv_{num_heads}_{num_queries_per_kv}_{head_size}_{cache_block_size}_{kv_cache_dtype}.pt")
        # print(f"{block_table=}, {b_start_loc=}, {b_seq_len=}")
    output_triton = output
    # Run Triton
    chunked_prefill_paged_decode(
        query,
        k,
        v,
        output_triton,
        kv_cache_dtype,
        k_cache,
        v_cache,
        block_table,
        b_start_loc,
        b_seq_len,
        max_input_len,
        k_scale,
        v_scale,
        sliding_window=sliding_window,
    )
    # Run Torch
    if not run_with_dump_data:
        output_torch = compute_golden_attention(
            query=query,
            key=key,
            value=value,
            b_start_loc=b_start_loc,
            b_seq_len=b_seq_len,
            num_heads=num_heads,
            num_queries_per_kv=num_queries_per_kv,
            head_size=head_size,
            dtype=dtype,
            device=device,
            sliding_window=sliding_window
        )
        if dump_data:
            save_golden_for_debug(output_torch, filename=f"gloden_{num_heads}_{num_queries_per_kv}_{head_size}_{cache_block_size}_{sliding_window}_auto.pt")
    else:
        (output_torch,) = load_goldens_for_debug(device, filename=f"gloden_{num_heads}_{num_queries_per_kv}_{head_size}_{cache_block_size}_{sliding_window}_auto.pt")

    atol = 5e-3 if "fp8" in kv_cache_dtype else 1e-4
    # triton.testing.assert_close(output_triton, output_torch, atol=atol, rtol=1e-2)
    torch.testing.assert_close(output_triton, output_torch, atol=atol, rtol=1e-2)


# @pytest.mark.parametrize("num_heads", NUM_HEADS)
# @pytest.mark.parametrize("num_queries_per_kv", NUM_QUERIES_PER_KV)
# @pytest.mark.parametrize("head_size", HEAD_SIZES)
# @pytest.mark.parametrize("dtype", DTYPES)
# @pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPES)
# @pytest.mark.parametrize("device", CUDA_DEVICES)
# @torch.inference_mode()
# def test_contexted_kv_attention_alibi(
#     num_heads: int,
#     num_queries_per_kv: int,
#     head_size: int,
#     dtype: torch.dtype,
#     kv_cache_dtype: str,
#     device: str,
# ) -> None:
#     (
#         query,
#         key,
#         value,
#         k,
#         v,
#         output,
#         k_cache,
#         v_cache,
#         block_table,
#         b_start_loc,
#         b_seq_len,
#         max_input_len,
#         k_scale,
#         v_scale,
#         alibi_slopes,
#     ) = input_helper(
#         BS=10,
#         MAX_SEQ_LEN=1024,
#         MAX_CTX_LEN=1024,
#         cache_size=640,
#         block_size=32,
#         max_block_per_request=64,
#         num_heads=num_heads,
#         head_size=head_size,
#         num_queries_per_kv=num_queries_per_kv,
#         dtype=dtype,
#         kv_cache_dtype=kv_cache_dtype,
#         device=device,
#         use_alibi_slope=True,
#     )
#     output_torch = torch.empty_like(output)
#     output_triton = output

#     # Run Triton
#     chunked_prefill_paged_decode(
#         query,
#         k,
#         v,
#         output_triton,
#         kv_cache_dtype,
#         k_cache,
#         v_cache,
#         block_table,
#         b_start_loc,
#         b_seq_len,
#         max_input_len,
#         k_scale,
#         v_scale,
#         alibi_slopes=alibi_slopes,
#     )
#     # Run Torch
#     context_attention_fwd_torch(
#         query,
#         k,
#         v,
#         output_torch,
#         k_cache,
#         v_cache,
#         b_start_loc,
#         b_seq_len,
#         k_scale,
#         v_scale,
#         alibi_slopes=alibi_slopes,
#     )

#     triton.testing.assert_close(output_triton, output_torch, atol=1e-2, rtol=1e-2)
