# SPDX-License-Identifier: MIT
 
# The kernel in this file is adapted from the VLLM project:
# https://github.com/ROCm/vllm/blob/aiter_integration_final/vllm/attention/ops/chunked_prefill_paged_decode.py

# Authors:
#  - Burkhard Ringlein
#  - Jan van Lunteren
#  - Thomas Parnell

import os
import json
import torch
import functools
from typing import Any, Dict, Optional, List
import triton
import triton.language as tl
from triton.utils import annotate_hint
from aiter.ops.triton.pa_prefill import context_attention_fwd
from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH
import aiter.ops.triton.utils.arch_info as arch_info
from aiter import logger
from ast import literal_eval

NUM_WARPS=4

@triton.jit
def cdiv_fn(x, y):
    return (x + y - 1) // y


@triton.jit
def _kernel_paged_attention_2d(
    output_ptr,  # [num_tokens, num_query_heads, head_size]
    query_ptr,  # [num_tokens, num_query_heads, head_size]
    key_cache_ptr,  # [num_blks, num_kv_heads, head_size // x, blk_size, x]
    value_cache_ptr,  # [num_blks, num_kv_heads, head_size, blk_size]
    block_tables_ptr,  # [num_seqs, max_num_blocks_per_seq]
    seq_lens_ptr,  # [num_seqs]
    alibi_slopes_ptr,  # [num_query_heads]
    scale,  # float32
    k_scale,  # float32
    v_scale,  # float32
    num_query_heads: tl.constexpr,  # int
    num_queries_per_kv: tl.constexpr,  # int
    num_queries_per_kv_padded: tl.constexpr,  # int
    block_table_stride,  # int
    query_stride_0: tl.constexpr,  # int
    query_stride_1: tl.constexpr,  # int, should be equal to head_size
    output_stride_0: tl.constexpr,  # int
    output_stride_1: tl.constexpr,  # int, should be equal to head_size
    BLOCK_SIZE: tl.constexpr,  # int
    CACHE_BLOCK_SIZE: tl.constexpr,  # int
    HEAD_SIZE: tl.constexpr,  # int
    HEAD_SIZE_PADDED: tl.constexpr,  # int, must be power of 2
    USE_ALIBI_SLOPES: tl.constexpr,  # bool
    SLIDING_WINDOW: tl.constexpr,  # int
    x: tl.constexpr,  # int
    stride_k_cache_0: tl.constexpr,  # int
    stride_k_cache_1: tl.constexpr,  # int
    stride_k_cache_2: tl.constexpr,  # int
    stride_k_cache_3: tl.constexpr,  # int
    stride_k_cache_4: tl.constexpr,  # int
    stride_v_cache_0: tl.constexpr,  # int
    stride_v_cache_1: tl.constexpr,  # int
    stride_v_cache_2: tl.constexpr,  # int
    stride_v_cache_3: tl.constexpr,  # int
    SKIP_PREFILL: tl.constexpr,  # bool
    USE_MATRIX_LOAD: tl.constexpr,
    query_start_len_ptr,  # [num_seqs+1]
    HEAD_DIM_PAD_REQ: tl.constexpr, # bool
):
    seq_idx = tl.program_id(0)
    kv_head_idx = tl.program_id(1)

    tl.static_assert(CACHE_BLOCK_SIZE % BLOCK_SIZE == 0, "CACHE_BLOCK_SIZE must be divisible by BLOCK_SIZE")

    if SKIP_PREFILL:
        cur_batch_in_all_start_index = tl.load(query_start_len_ptr + seq_idx)
        cur_batch_in_all_stop_index = tl.load(query_start_len_ptr + seq_idx +
                                              1)
        cur_batch_query_len = cur_batch_in_all_stop_index \
            - cur_batch_in_all_start_index
        if cur_batch_query_len > 1:
            return
    else:
        cur_batch_in_all_start_index = seq_idx

    query_head_idx = kv_head_idx * num_queries_per_kv + tl.arange(
        0, num_queries_per_kv_padded)
    head_mask = query_head_idx < (kv_head_idx + 1) * num_queries_per_kv
    head_mask = head_mask & (query_head_idx < num_query_heads)

    query_offset = (cur_batch_in_all_start_index * query_stride_0 +
                    query_head_idx[:, None] * query_stride_1)
    # Q : (num_queries_per_kv, HEAD_SIZE,)
    if HEAD_DIM_PAD_REQ:
        dim_mask = tl.where(tl.arange(0, HEAD_SIZE_PADDED) < HEAD_SIZE, 1,
                            0).to(tl.int1)
        Q = tl.load(
            query_ptr + query_offset + tl.arange(0, HEAD_SIZE_PADDED)[None, :],
            mask=dim_mask[None, :] & head_mask[:, None],
            other=0.0,
        )
    else:
        Q = tl.load(
            query_ptr + query_offset + tl.arange(0, HEAD_SIZE_PADDED)[None, :],
            mask=head_mask[:, None],
            other=0.0,
        )      

    block_table_offset = seq_idx * block_table_stride

    M = tl.full([num_queries_per_kv_padded], float("-inf"), dtype=tl.float32)
    L = tl.full([num_queries_per_kv_padded], 1.0, dtype=tl.float32)
    acc = tl.zeros([num_queries_per_kv_padded, HEAD_SIZE_PADDED],
                   dtype=tl.float32)

    # sequence len for this particular sequence
    seq_len = tl.load(seq_lens_ptr + seq_idx)
    seq_len = annotate_hint(seq_len, "non-negative")

    # alibi slope for this head
    if USE_ALIBI_SLOPES:
        alibi_slope = tl.load(alibi_slopes_ptr + query_head_idx,
                              mask=head_mask,
                              other=0.0)

    # iterate through tiles
    for start_n in range(0, seq_len, BLOCK_SIZE):
        physical_block_idx = tl.load(block_tables_ptr + block_table_offset + (start_n // CACHE_BLOCK_SIZE))
        physical_block_idx = annotate_hint(physical_block_idx, "non-negative")
        offs_n = tl.arange(0, BLOCK_SIZE)
        offs_d = tl.arange(0, HEAD_SIZE_PADDED)

        if BLOCK_SIZE == CACHE_BLOCK_SIZE:
            if USE_MATRIX_LOAD:
                v_offset = (physical_block_idx * stride_v_cache_0 +
                            kv_head_idx * stride_v_cache_1)
            else:
                v_offset = (physical_block_idx * stride_v_cache_0 +
                            kv_head_idx * stride_v_cache_1 +
                            offs_d[None, :] * stride_v_cache_2 +
                            offs_n[:, None] * stride_v_cache_3)

            k_offset = (physical_block_idx * stride_k_cache_0 +
                        kv_head_idx * stride_k_cache_1 +
                        (offs_d[:, None] // x) * stride_k_cache_2 +
                        offs_n[None, :] * stride_k_cache_3 +
                        (offs_d[:, None] % x) * stride_k_cache_4)
        else:
            if USE_MATRIX_LOAD:
                v_offset = (physical_block_idx * stride_v_cache_0 +
                            kv_head_idx * stride_v_cache_1)
            else:
                v_offset = (physical_block_idx * stride_v_cache_0 +
                            kv_head_idx * stride_v_cache_1 +
                            offs_d[None, :] * stride_v_cache_2 +
                            ((start_n + offs_n[:, None]) % CACHE_BLOCK_SIZE) * stride_v_cache_3)

            k_offset = (physical_block_idx * stride_k_cache_0 +
                        kv_head_idx * stride_k_cache_1 +
                        (offs_d[:, None] // x) * stride_k_cache_2 +
                        ((start_n + offs_n[None, :]) % CACHE_BLOCK_SIZE) * stride_k_cache_3 +
                        (offs_d[:, None] % x) * stride_k_cache_4)
        # K : (HEAD_SIZE, BLOCK_SIZE)
        if HEAD_DIM_PAD_REQ:
            K_load = tl.load(key_cache_ptr + k_offset,
                            mask=dim_mask[:, None],
                            other=0.0)
        else:
            K_load = tl.load(key_cache_ptr + k_offset)           

        if K_load.dtype.is_fp8():
            K = (K_load.to(tl.float32) * tl.load(k_scale)).to(Q.dtype)
        else:
            K = K_load

        # V : (BLOCK_SIZE, HEAD_SIZE)
        if USE_MATRIX_LOAD:
            if HEAD_DIM_PAD_REQ:
                V_load = tl.matrix_load(
                                value_cache_ptr + v_offset,
                                shape=[CACHE_BLOCK_SIZE, HEAD_SIZE],
                                strides=[stride_v_cache_3, stride_v_cache_2],
                                block_shape=[BLOCK_SIZE, HEAD_SIZE_PADDED],
                                offsets=[(start_n % CACHE_BLOCK_SIZE).to(tl.int32), 0],
                                boundary_check=(1,))
            else:
                V_load = tl.matrix_load(
                                value_cache_ptr + v_offset,
                                shape=[CACHE_BLOCK_SIZE, HEAD_SIZE],
                                strides=[stride_v_cache_3, stride_v_cache_2],
                                block_shape=[BLOCK_SIZE, HEAD_SIZE_PADDED],
                                offsets=[(start_n % CACHE_BLOCK_SIZE).to(tl.int32), 0])
        else:
            if HEAD_DIM_PAD_REQ:
                V_load = tl.load(value_cache_ptr + v_offset,
                                mask=dim_mask[None, :],
                                other=0.0)
            else:
                V_load = tl.load(value_cache_ptr + v_offset)

        if V_load.dtype.is_fp8():
            V = (V_load.to(tl.float32) * tl.load(v_scale)).to(Q.dtype)
        else:
            V = V_load

        seq_offset = start_n + tl.arange(0, BLOCK_SIZE)
        boundary = tl.full([BLOCK_SIZE], seq_len, dtype=tl.int32)
        seq_mask = seq_offset < boundary
        # S : (num_queries_per_kv, BLOCK_SIZE,)
        S = tl.where(head_mask[:, None] & seq_mask[None, :], 0.0,
                     float("-inf")).to(tl.float32)
        S += scale * tl.dot(Q, K)

        context_len = seq_len - 1

        if SLIDING_WINDOW > 0:
            S = tl.where((context_len - seq_offset) < SLIDING_WINDOW, S,
                         -10000)

        if USE_ALIBI_SLOPES:
            S += alibi_slope[:, None] * (seq_offset - context_len)

        # compute running maximum
        # m_j : (num_queries_per_kv,)
        m_j = tl.maximum(M, tl.max(S, axis=1))

        # P : (num_queries_per_kv, BLOCK_SIZE,)
        P = tl.exp(S - m_j[:, None])

        # l_j : (num_queries_per_kv,)
        l_j = tl.sum(P, axis=1)

        # alpha : (num_queries_per_kv, )
        alpha = tl.exp(M - m_j)

        # acc : (num_queries_per_kv, BLOCK_SIZE,)
        acc = acc * alpha[:, None]

        # update constants
        L = L * alpha + l_j
        M = m_j

        # acc : (num_queries_per_kv, BLOCK_SIZE,)
        acc += tl.dot(P.to(V.dtype), V)

    # epilogue
    acc = acc / L[:, None]

    output_offset = (cur_batch_in_all_start_index * output_stride_0 +
                     query_head_idx * output_stride_1)

    if HEAD_DIM_PAD_REQ:
        tl.store(
            output_ptr + output_offset[:, None] +
            tl.arange(0, HEAD_SIZE_PADDED)[None, :],
            acc,
            mask=dim_mask[None, :] & head_mask[:, None],
        )
    else:
        tl.store(
            output_ptr + output_offset[:, None] +
            tl.arange(0, HEAD_SIZE_PADDED)[None, :],
            acc,
            mask=head_mask[:, None],
        )

@functools.lru_cache
def find_block(a, b):
    if a < 16 or b < 16:
        return None
    # 找到小于等于b的最大2的幂
    max_power = (b).bit_length() - 1
    # 从大到小检查2的幂
    for k in range(max_power, 3, -1):
        power = 1 << k
        if a % power == 0:
            return power
    return None

@functools.lru_cache
def get_paged_attention_2d_config_filepath(cache_block_size, head_size, slide_window, 
                                           use_alibi_slopes, filter_by_query_len, 
                                           kv_dtype, **kwargs) -> str:
    kv_type = "auto"
    if kv_dtype == torch.float8_e4m3fn or kv_dtype == torch.float8_e5m2:
        kv_type = "fp8"
    device_name = arch_info.get_arch()
    head_size_padded = triton.next_power_of_2(head_size)
    head_size_pad_need = head_size != head_size_padded
    json_file_name = (
        f"paged_attention_2d-device={device_name}"
        f"-CACHE_BLOCK_SIZE={cache_block_size}"
        f"-HEAD_SIZE_PADDED={head_size_padded}"
        f"-SLIDING_WINDOW={slide_window}"
        f"-USE_ALIBI_SLOPES={use_alibi_slopes}"
        f"-HEAD_DIM_PAD_REQ={head_size_pad_need}"
        f"-kv_dtype={kv_type}.json"
    )

    config_file_path = os.path.join(
        f"{AITER_TRITON_CONFIGS_PATH}", "paged_attention_2d", json_file_name
    )
    return config_file_path

@functools.lru_cache
def get_paged_attention_2d_config(
    cache_block_size,
    head_size,
    num_querys_per_kv,
    slide_window,
    use_alibi_slopes,
    filter_by_query_len,
    kv_dtype
) -> Optional[Dict]:
    config_file_path = get_paged_attention_2d_config_filepath(cache_block_size, head_size,
                                                              slide_window, use_alibi_slopes, 
                                                              filter_by_query_len, kv_dtype)
    if os.path.exists(config_file_path):
        with open(config_file_path) as f:
            configs = {int(key): val for key, val in json.load(f)["config"].items()}
            if configs:
                num_querys_per_kv_padded = triton.next_power_of_2(num_querys_per_kv)
                config = configs[min(configs.keys(), key=lambda x: abs(x - num_querys_per_kv_padded))]
                # logger.info(f"paged_attention_2d use kernel config from:{config_file_path}")
                return config
    # If no optimized configuration is available, we will use the default
    logger.warning(
            f"\nUsing default paged_attention_2d kernel config. Performance might "
            f"be sub-optimal! Config not found at {config_file_path}")
    return None

@torch.inference_mode()
def paged_attention_2d(
    query,
    output,
    kv_cache_dtype,
    key_cache,
    value_cache,
    block_table,
    query_start_loc,
    seq_lens,
    k_scale,
    v_scale,
    alibi_slopes=None,
    sliding_window=None,
    sm_scale=None,
    filter_by_query_len=True,
):
    if sm_scale is None:
        sm_scale = 1.0 / (query.shape[1]**0.5)
    use_alibi_slopes = alibi_slopes is not None

    if sliding_window is None or sliding_window <= 0:
        sliding_window = 0

    cache_block_size = value_cache.shape[3]
    head_size = query.shape[2]
    head_size_padded = triton.next_power_of_2(head_size)

    num_seqs = len(seq_lens)
    num_query_heads = query.shape[1]
    num_kv_heads = key_cache.shape[1]
    num_queries_per_kv = query.shape[1] // key_cache.shape[1]
    num_queries_per_kv_padded = max(triton.next_power_of_2(num_queries_per_kv), 16)

    if "fp8" in kv_cache_dtype and (key_cache.dtype == torch.uint8 or value_cache.dtype == torch.uint8):
        # kv_cache may view as uint8
        if kv_cache_dtype in ("fp8", "fp8e4m3"):
            target_dtype = torch.float8_e4m3fn
        elif kv_cache_dtype == "fp8e5m2":
            target_dtype = torch.float8_e5m2
        else:
            raise ValueError("Unsupported FP8 dtype:", kv_cache_dtype)
        key_cache = key_cache.view(target_dtype)
        value_cache = value_cache.view(target_dtype)

    config = get_paged_attention_2d_config(cache_block_size, head_size, num_queries_per_kv,
                                           sliding_window, use_alibi_slopes, filter_by_query_len, key_cache.dtype)
    if not config:
        config = {'num_warps': 4, 'num_stages': 1, 'USE_MATRIX_LOAD': False}
    if 'BLOCK_SIZE' not in config:
        cache_ele_size = key_cache.element_size()
        block_size = cache_block_size
        if block_size * head_size_padded * cache_ele_size > 16384: # 64 * 128 * 2
            block_size = (16384 // (head_size_padded * cache_ele_size))
            block_size = find_block(cache_block_size, block_size)
            assert block_size != None, "can not find suitable block_size size for kernel_paged_attention_2d"
            assert (cache_ele_size >= 2 or block_size >= 32), "Block size must be at least 32 for fp8"
        config['BLOCK_SIZE'] = block_size

    assert cache_block_size % config['BLOCK_SIZE'] == 0, "cache_block_size % block_size need be 0."
    # print(f"kernel_paged_attention_2d: {num_queries_per_kv=} {config=}")

    _kernel_paged_attention_2d[(
        num_seqs,
        num_kv_heads,
    )](
        output_ptr=output,
        query_ptr=query,
        key_cache_ptr=key_cache,
        value_cache_ptr=value_cache,
        block_tables_ptr=block_table,
        seq_lens_ptr=seq_lens,
        alibi_slopes_ptr=alibi_slopes,
        scale=sm_scale,
        k_scale=k_scale,
        v_scale=v_scale,
        num_query_heads=num_query_heads,
        num_queries_per_kv=num_queries_per_kv,
        num_queries_per_kv_padded=num_queries_per_kv_padded,
        block_table_stride=block_table.stride(0),
        query_stride_0=query.stride(0),
        query_stride_1=query.stride(1),
        output_stride_0=output.stride(0),
        output_stride_1=output.stride(1),
        CACHE_BLOCK_SIZE=cache_block_size,
        HEAD_SIZE=head_size,
        HEAD_SIZE_PADDED=head_size_padded,
        USE_ALIBI_SLOPES=use_alibi_slopes,
        SLIDING_WINDOW=sliding_window,
        x=key_cache.shape[4],
        stride_k_cache_0=key_cache.stride(0),
        stride_k_cache_1=key_cache.stride(1),
        stride_k_cache_2=key_cache.stride(2),
        stride_k_cache_3=key_cache.stride(3),
        stride_k_cache_4=key_cache.stride(4),
        stride_v_cache_0=value_cache.stride(0),
        stride_v_cache_1=value_cache.stride(1),
        stride_v_cache_2=value_cache.stride(2),
        stride_v_cache_3=value_cache.stride(3),
        SKIP_PREFILL=filter_by_query_len,
        query_start_len_ptr=query_start_loc,
        HEAD_DIM_PAD_REQ=(head_size != head_size_padded),
        **config,
    )


def chunked_prefill_paged_decode(
    query,
    key,
    value,
    output,
    kv_cache_dtype,
    key_cache,
    value_cache,
    block_table,
    query_start_loc,
    seq_lens,
    max_query_len,
    k_scale,
    v_scale,
    alibi_slopes=None,
    sliding_window=None,
    sm_scale=None,
):
    if sm_scale is None:
        sm_scale = 1.0 / (query.shape[1]**0.5)

    use_alibi_slopes = alibi_slopes is not None

    if sliding_window is None or sliding_window <= 0:
        sliding_window = 0

    if max_query_len > 1:
        context_attention_fwd(
            q=query,
            k=key,
            v=value,
            o=output,
            kv_cache_dtype=kv_cache_dtype,
            k_cache=key_cache,
            v_cache=value_cache,
            b_loc=block_table,
            b_start_loc=query_start_loc,
            b_seq_len=seq_lens,
            max_input_len=max_query_len,
            k_scale=k_scale,
            v_scale=v_scale,
            alibi_slopes=alibi_slopes,
            sliding_window=sliding_window,
            sm_scale=sm_scale,
            skip_decode=True,
        )

    paged_attention_2d(
        query,
        output,
        kv_cache_dtype,
        key_cache,
        value_cache,
        block_table,
        query_start_loc,
        seq_lens,
        k_scale,
        v_scale,
        alibi_slopes=alibi_slopes,
        sliding_window=sliding_window,
        sm_scale=sm_scale,
        filter_by_query_len=True,
    )
