# SPDX-License-Identifier: MIT
 
# SPDX-License-Identifier: MIT
 
# The kernels in this file are adapted from LightLLM's context_attention_fwd:
# https://github.com/ModelTC/lightllm/blob/main/lightllm/models/llama/triton_kernel/context_flashattention_nopad.py

import os

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

BASE_BLOCK = 128
NUM_WARPS = 4

if triton.__version__ >= "2.1.0":

    @triton.jit
    def _fwd_kernel(
        Q,
        K,
        V,
        K_cache,
        V_cache,
        B_Loc,
        sm_scale,
        k_scale,
        v_scale,
        B_Start_Loc,
        B_Seqlen,
        block_size: tl.constexpr,
        x: tl.constexpr,
        Out,
        stride_b_loc_b,
        stride_b_loc_s,
        stride_qbs: tl.constexpr,
        stride_qh: tl.constexpr,
        stride_qd: tl.constexpr,
        stride_kbs: tl.constexpr,
        stride_kh: tl.constexpr,
        stride_kd: tl.constexpr,
        stride_vbs: tl.constexpr,
        stride_vh: tl.constexpr,
        stride_vd: tl.constexpr,
        stride_obs: tl.constexpr,
        stride_oh: tl.constexpr,
        stride_od: tl.constexpr,
        stride_k_cache_bs: tl.constexpr,
        stride_k_cache_h: tl.constexpr,
        stride_k_cache_d: tl.constexpr,
        stride_k_cache_bl: tl.constexpr,
        stride_k_cache_x: tl.constexpr,
        stride_v_cache_bs: tl.constexpr,
        stride_v_cache_h: tl.constexpr,
        stride_v_cache_d: tl.constexpr,
        stride_v_cache_bl: tl.constexpr,
        num_queries_per_kv: int,
        IN_PRECISION: tl.constexpr,
        BLOCK_M: tl.constexpr,
        BLOCK_DMODEL: tl.constexpr,  # head size
        BLOCK_DMODEL_PADDED: tl.constexpr,  # head size padded to a power of 2
        BLOCK_N: tl.constexpr,
        SLIDING_WINDOW: tl.constexpr,
        SKIP_DECODE: tl.constexpr,
        USE_MATRIX_LOAD: tl.constexpr,  # bool
        HEAD_DIM_PAD_REQ: tl.constexpr,  # bool
        max_input_len,
    ):
        if USE_MATRIX_LOAD:
            # if use matrix load, need make sure all cache tokens in BLOCK_N are in the same cache block
            tl.static_assert(BLOCK_N <= block_size and block_size % BLOCK_N == 0)
        cur_batch = tl.program_id(0)
        cur_head = tl.program_id(1)
        start_m = tl.program_id(2)

        cur_kv_head = cur_head // num_queries_per_kv

        cur_batch_seq_len = tl.load(B_Seqlen + cur_batch)
        cur_batch_in_all_start_index = tl.load(B_Start_Loc + cur_batch)
        cur_batch_in_all_stop_index = tl.load(B_Start_Loc + cur_batch + 1)
        cur_batch_query_len = (cur_batch_in_all_stop_index -
                               cur_batch_in_all_start_index)
        cur_batch_ctx_len = cur_batch_seq_len - cur_batch_query_len
        cur_batch_ctx_len = annotate_hint(cur_batch_ctx_len, "non-negative")

        if SKIP_DECODE and cur_batch_query_len == 1:
            return

        # start position inside of the query
        # generally, N goes over kv, while M goes over query_len
        block_start_loc = BLOCK_M * start_m

        # initialize offsets
        # [N]; starts at 0
        offs_n = tl.arange(0, BLOCK_N)
        # [D]; starts at 0
        offs_d = tl.arange(0, BLOCK_DMODEL_PADDED)
        # [M]; starts at current position in query
        offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
        # [M,D]
        off_q = (
            (cur_batch_in_all_start_index + offs_m[:, None]) * stride_qbs +
            cur_head * stride_qh + offs_d[None, :] * stride_qd)

        if HEAD_DIM_PAD_REQ:
            dim_mask = tl.where(
                tl.arange(0, BLOCK_DMODEL_PADDED) < BLOCK_DMODEL, 1,
                0).to(tl.int1)  # [D]

            q = tl.load(Q + off_q,
                        mask=dim_mask[None, :] &
                        (offs_m[:, None] < cur_batch_query_len),
                        other=0.0)  # [M,D]
        else:
            q = tl.load(Q + off_q,
                        mask=(offs_m[:, None] < cur_batch_query_len),
                        other=0.0)  # [M,D]

        # initialize pointer to m and l
        m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")  # [M]
        l_i = tl.full([BLOCK_M], 1.0, dtype=tl.float32)  # [M]
        acc = tl.zeros([BLOCK_M, BLOCK_DMODEL_PADDED],
                       dtype=tl.float32)  # [M,D]

        # compute query against context (no causal mask here)
        for start_n in range(0, cur_batch_ctx_len, BLOCK_N):
            start_n = tl.multiple_of(start_n, BLOCK_N)
            # -- compute qk ----
            if USE_MATRIX_LOAD:
                # all cache tokens in BLOCK_N are in the same cache block
                bn = tl.load(B_Loc + cur_batch * stride_b_loc_b +
                            (start_n // block_size) * stride_b_loc_s)
                # [D,N]
                off_k = (bn * stride_k_cache_bs +
                        cur_kv_head * stride_k_cache_h +
                        (offs_d[:, None] // x) * stride_k_cache_d +
                        ((start_n + offs_n[None, :]) % block_size) *
                        stride_k_cache_bl +
                        (offs_d[:, None] % x) * stride_k_cache_x)
                # [N,D]
                off_v = (
                    bn * stride_v_cache_bs +
                    cur_kv_head * stride_v_cache_h)
            else:
                bn = tl.load(B_Loc + cur_batch * stride_b_loc_b +
                            ((start_n + offs_n) // block_size) * stride_b_loc_s,
                            mask=(start_n + offs_n) < cur_batch_ctx_len,
                            other=0)  # [N]
                # we explicit tell compiler bn is non-negative
                bn = annotate_hint(bn, "non-negative")
                # set constancy to 16 for v cache load using load_dwordx4
                bn = tl.max_constancy(bn, [16])
                # [D,N]
                off_k = (bn[None, :] * stride_k_cache_bs +
                        cur_kv_head * stride_k_cache_h +
                        (offs_d[:, None] // x) * stride_k_cache_d +
                        ((start_n + offs_n[None, :]) % block_size) *
                        stride_k_cache_bl +
                        (offs_d[:, None] % x) * stride_k_cache_x)
                # [N,D]
                off_v = (
                    bn[:, None] * stride_v_cache_bs +
                    cur_kv_head * stride_v_cache_h +
                    offs_d[None, :] * stride_v_cache_d +
                    ((start_n + offs_n[:, None]) % block_size) * stride_v_cache_bl)

            if HEAD_DIM_PAD_REQ:
                if block_size % BLOCK_N == 0:
                    # block_size % BLOCK_N == 0, seq will never meet memory boundray out
                    k_load = tl.load(K_cache + off_k,
                                    mask=dim_mask[:, None],
                                    other=0.0)  # [D,N]
                else:
                    k_load = tl.load(K_cache + off_k,
                                    mask=dim_mask[:, None] &
                                    ((start_n + offs_n[None, :]) < cur_batch_ctx_len),
                                    other=0.0)  # [D,N]
            else:
                if block_size % BLOCK_N == 0:
                    k_load = tl.load(K_cache + off_k)  # [D,N]
                else:
                    k_load = tl.load(K_cache + off_k,
                                    mask=(start_n + offs_n[None, :]) < cur_batch_ctx_len,
                                    other=0.0)  # [D,N]

            if k_load.dtype.is_fp8():
                k = (k_load.to(tl.float32) * tl.load(k_scale)).to(q.dtype)
            else:
                k = k_load

            qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)  # [M,N]
            qk = tl.dot(q, k, acc=qk, input_precision=IN_PRECISION)
            qk = tl.where((start_n + offs_n[None, :]) < cur_batch_ctx_len, qk,
                          float("-inf"))
            qk *= sm_scale
            if SLIDING_WINDOW > 0:
                qk = tl.where((cur_batch_ctx_len + offs_m[:, None]) -
                              (start_n + offs_n[None, :]) < SLIDING_WINDOW, qk,
                              float("-inf"))

            # -- compute m_ij, p, l_ij
            m_j = tl.maximum(m_i, tl.max(qk, 1))
            if SLIDING_WINDOW > 0:
                # For sliding window there's a chance the max is -inf due to masking of
                # the entire row. In this case we need to set m_j 0 to avoid NaN
                m_j = tl.where(m_j > float("-inf"), m_j, 0.0)

            # P : (BLOCK_M, BLOCK_SIZE)
            p = tl.exp(qk - m_j[:, None])
             # l_j : (BLOCK_M,)
            l_j = tl.sum(p, 1)
            # alpha : (BLOCK_M, )
            alpha = tl.exp(m_i - m_j)
            # scale acc
            acc = acc * alpha[:, None]
            # update acc
            if USE_MATRIX_LOAD:
                if HEAD_DIM_PAD_REQ:
                    v_load = tl.matrix_load(
                                    V_cache + off_v,
                                    shape=[block_size, BLOCK_DMODEL],
                                    strides=[stride_v_cache_bl, stride_v_cache_d],
                                    block_shape=[BLOCK_N, BLOCK_DMODEL_PADDED],
                                    offsets=[(start_n % block_size).to(tl.int32), 0],
                                    boundary_check=(1,)) # [N,D]
                else:
                    v_load = tl.matrix_load(
                                    V_cache + off_v,
                                    shape=[block_size, BLOCK_DMODEL],
                                    strides=[stride_v_cache_bl, stride_v_cache_d],
                                    block_shape=[BLOCK_N, BLOCK_DMODEL_PADDED],
                                    offsets=[(start_n % block_size).to(tl.int32), 0]) # [N,D]
            else:
                seq_mask = (start_n + offs_n[:, None]) < cur_batch_ctx_len
                # set constancy of seq_mask to 8 for using vector load_dwordx4
                seq_mask = tl.max_constancy(seq_mask, [8, BLOCK_DMODEL_PADDED])
                if HEAD_DIM_PAD_REQ:
                    if block_size % BLOCK_N == 0:
                        # block_size % BLOCK_N == 0, seq will never meet memory boundray out
                        v_load = tl.load(V_cache + off_v,
                                        mask=dim_mask[None, :],
                                        other=0.0)  # [N,D]
                    else:
                        v_load = tl.load(V_cache + off_v,
                                        mask=dim_mask[None, :] & seq_mask,
                                        other=0.0)  # [N,D]
                else:
                    if block_size % BLOCK_N == 0:
                        v_load = tl.load(V_cache + off_v)  # [N,D]
                    else:

                        v_load = tl.load(V_cache + off_v,
                                        mask=seq_mask,
                                        other=0.0)  # [N,D]

            if v_load.dtype.is_fp8():
                v = (v_load.to(tl.float32) * tl.load(v_scale)).to(q.dtype)
            else:
                v = v_load
            p = p.to(v.dtype)

            acc = tl.dot(p, v, acc=acc, input_precision=IN_PRECISION)
            # # update m_i and l_i
            l_i = l_i * alpha + l_j
            m_i = m_j

        off_k = (offs_n[None, :] * stride_kbs + cur_kv_head * stride_kh +
                 offs_d[:, None] * stride_kd)
        off_v = (offs_n[:, None] * stride_vbs + cur_kv_head * stride_vh +
                 offs_d[None, :] * stride_vd)
        k_ptrs = K + off_k
        v_ptrs = V + off_v

        # block_mask is 0 when we're already past the current query length
        block_mask = tl.where(block_start_loc < cur_batch_query_len, 1, 0)

        # compute query against itself (with causal mask)
        for start_n in range(0, block_mask * (start_m + 1) * BLOCK_M, BLOCK_N):
            start_n = tl.multiple_of(start_n, BLOCK_N)
            # -- compute qk ----
            if HEAD_DIM_PAD_REQ:
                k = tl.load(k_ptrs +
                            (cur_batch_in_all_start_index + start_n) * stride_kbs,
                            mask=dim_mask[:, None] &
                            ((start_n + offs_n[None, :]) < cur_batch_query_len),
                            other=0.0)
            else:
                k = tl.load(k_ptrs +
                            (cur_batch_in_all_start_index + start_n) * stride_kbs,
                            mask=((start_n + offs_n[None, :]) < cur_batch_query_len),
                            other=0.0)
            qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
            qk = tl.dot(q, k, acc=qk, input_precision=IN_PRECISION)
            qk *= sm_scale
            # apply causal mask
            qk = tl.where(offs_m[:, None] >= (start_n + offs_n[None, :]), qk,
                          float("-inf"))
            if SLIDING_WINDOW > 0:
                qk = tl.where(
                    offs_m[:, None] - (start_n + offs_n[None, :])
                    < SLIDING_WINDOW, qk, float("-inf"))

            m_j = tl.maximum(m_i, tl.max(qk, 1))
            if SLIDING_WINDOW > 0:
                # For sliding window there's a chance the max is -inf due to masking of
                # the entire row. In this case we need to set m_j 0 to avoid NaN
                m_j = tl.where(m_j > float("-inf"), m_j, 0.0)

            # P : (BLOCK_M, BLOCK_SIZE)
            p = tl.exp(qk - m_j[:, None])
             # l_j : (BLOCK_M,)
            l_j = tl.sum(p, 1)
            # alpha : (BLOCK_M, )
            alpha = tl.exp(m_i - m_j)
            # scale acc
            acc = acc * alpha[:, None]
            # update acc
            if HEAD_DIM_PAD_REQ:
                v = tl.load(v_ptrs +
                            (cur_batch_in_all_start_index + start_n) * stride_vbs,
                            mask=dim_mask[None, :] &
                            ((start_n + offs_n[:, None]) < cur_batch_query_len),
                            other=0.0)
            else:
                v = tl.load(v_ptrs +
                            (cur_batch_in_all_start_index + start_n) * stride_vbs,
                            mask=((start_n + offs_n[:, None]) < cur_batch_query_len),
                            other=0.0)
            p = p.to(v.dtype)

            acc = tl.dot(p, v, acc=acc, input_precision=IN_PRECISION)
            # update m_i and l_i
            l_i = alpha * l_i + l_j
            m_i = m_j

        acc = acc / l_i[:, None]
        # initialize pointers to output
        off_o = (
            (cur_batch_in_all_start_index + offs_m[:, None]) * stride_obs +
            cur_head * stride_oh + offs_d[None, :] * stride_od)
        out_ptrs = Out + off_o

        if HEAD_DIM_PAD_REQ:
            tl.store(out_ptrs,
                    acc,
                    mask=dim_mask[None, :] &
                    (offs_m[:, None] < cur_batch_query_len))
        else:
            tl.store(out_ptrs,
                    acc,
                    mask=(offs_m[:, None] < cur_batch_query_len))
        return

    @triton.jit
    def _fwd_kernel_alibi(
        Q,
        K,
        V,
        K_cache,
        V_cache,
        B_Loc,
        sm_scale,
        k_scale,
        v_scale,
        B_Start_Loc,
        B_Seqlen,
        Alibi_slopes,
        block_size,
        x,
        Out,
        stride_b_loc_b,
        stride_b_loc_s,
        stride_qbs,
        stride_qh,
        stride_qd,
        stride_kbs,
        stride_kh,
        stride_kd,
        stride_vbs,
        stride_vh,
        stride_vd,
        stride_obs,
        stride_oh,
        stride_od,
        stride_k_cache_bs,
        stride_k_cache_h,
        stride_k_cache_d,
        stride_k_cache_bl,
        stride_k_cache_x,
        stride_v_cache_bs,
        stride_v_cache_h,
        stride_v_cache_d,
        stride_v_cache_bl,
        num_queries_per_kv: int,
        IN_PRECISION: tl.constexpr,
        BLOCK_M: tl.constexpr,
        BLOCK_DMODEL: tl.constexpr,  # head size
        BLOCK_DMODEL_PADDED: tl.constexpr,  # head size padded to a power of 2
        BLOCK_N: tl.constexpr,
        SKIP_DECODE: tl.constexpr,
    ):
        # attn_bias[]
        cur_batch = tl.program_id(0)
        cur_head = tl.program_id(1)
        start_m = tl.program_id(2)

        cur_kv_head = cur_head // num_queries_per_kv

        # cur_batch_seq_len: the length of prompts
        # cur_batch_ctx_len: the length of prefix
        # cur_batch_in_all_start_index: the start id of the dim=0
        cur_batch_seq_len = tl.load(B_Seqlen + cur_batch)
        cur_batch_in_all_start_index = tl.load(B_Start_Loc + cur_batch)
        cur_batch_in_all_stop_index = tl.load(B_Start_Loc + cur_batch + 1)
        cur_batch_query_len = (cur_batch_in_all_stop_index -
                               cur_batch_in_all_start_index)
        cur_batch_ctx_len = cur_batch_seq_len - cur_batch_query_len

        if SKIP_DECODE and cur_batch_query_len == 1:
            return

        block_start_loc = BLOCK_M * start_m

        # initialize offsets
        offs_n = tl.arange(0, BLOCK_N)
        offs_d = tl.arange(0, BLOCK_DMODEL_PADDED)
        offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
        off_q = (
            (cur_batch_in_all_start_index + offs_m[:, None]) * stride_qbs +
            cur_head * stride_qh + offs_d[None, :] * stride_qd)

        dim_mask = tl.where(
            tl.arange(0, BLOCK_DMODEL_PADDED) < BLOCK_DMODEL, 1, 0).to(tl.int1)

        q = tl.load(Q + off_q,
                    mask=dim_mask[None, :] &
                    (offs_m[:, None] < cur_batch_seq_len - cur_batch_ctx_len),
                    other=0.0)

        # # initialize pointer to m and l
        m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
        l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
        acc = tl.zeros([BLOCK_M, BLOCK_DMODEL_PADDED], dtype=tl.float32)

        alibi_slope = tl.load(Alibi_slopes + cur_head)
        alibi_start_q = tl.arange(
            0, BLOCK_M) + block_start_loc + cur_batch_ctx_len
        alibi_start_k = 0
        for start_n in range(0, cur_batch_ctx_len, BLOCK_N):
            start_n = tl.multiple_of(start_n, BLOCK_N)
            # -- compute qk ----
            bn = tl.load(B_Loc + cur_batch * stride_b_loc_b +
                         ((start_n + offs_n) // block_size) * stride_b_loc_s,
                         mask=(start_n + offs_n) < cur_batch_ctx_len,
                         other=0)
            off_k = (bn[None, :] * stride_k_cache_bs +
                     cur_kv_head * stride_k_cache_h +
                     (offs_d[:, None] // x) * stride_k_cache_d +
                     ((start_n + offs_n[None, :]) % block_size) *
                     stride_k_cache_bl +
                     (offs_d[:, None] % x) * stride_k_cache_x)
            off_v = (
                bn[:, None] * stride_v_cache_bs +
                cur_kv_head * stride_v_cache_h +
                offs_d[None, :] * stride_v_cache_d +
                (start_n + offs_n[:, None]) % block_size * stride_v_cache_bl)
            k_load = tl.load(K_cache + off_k,
                             mask=dim_mask[:, None] &
                             ((start_n + offs_n[None, :]) < cur_batch_ctx_len),
                             other=0.0)  # [D,N]

            if k_load.dtype.is_fp8():
                k = (k_load.to(tl.float32) * tl.load(k_scale)).to(q.dtype)
            else:
                k = k_load

            qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
            qk = tl.dot(q, k, acc=qk, input_precision=IN_PRECISION)
            qk = tl.where((start_n + offs_n[None, :]) < cur_batch_ctx_len, qk,
                          float("-inf"))
            qk *= sm_scale

            # load alibi
            alibi = (tl.arange(0, BLOCK_N)[None, :] + alibi_start_k -
                     alibi_start_q[:, None]) * alibi_slope
            alibi = tl.where(
                (alibi <= 0) & (alibi_start_q[:, None] < cur_batch_seq_len),
                alibi, float("-inf"))
            qk += alibi
            alibi_start_k += BLOCK_N

            # -- compute m_ij, p, l_ij
            m_ij = tl.max(qk, 1)
            m_i_new = tl.maximum(m_i, m_ij)
            p = tl.math.exp(qk - m_i_new[:, None])
            l_ij = tl.sum(p, 1)
            # -- update m_i and l_i

            alpha = tl.math.exp(m_i - m_i_new)
            l_i_new = alpha * l_i + l_ij
            # -- update output accumulator --
            # scale p
            # scale acc
            acc_scale = alpha
            # acc_scale = l_i / l_i_new * alpha
            acc = acc * acc_scale[:, None]
            # update acc
            v_load = tl.load(V_cache + off_v,
                             mask=dim_mask[None, :] &
                             ((start_n + offs_n[:, None]) < cur_batch_ctx_len),
                             other=0.0)
            if v_load.dtype.is_fp8():
                v = (v_load.to(tl.float32) * tl.load(v_scale)).to(q.dtype)
            else:
                v = v_load
            p = p.to(v.dtype)

            acc = tl.dot(p, v, acc=acc, input_precision='ieee')
            # update m_i and l_i
            l_i = l_i_new
            m_i = m_i_new

        off_k = (offs_n[None, :] * stride_kbs + cur_kv_head * stride_kh +
                 offs_d[:, None] * stride_kd)
        off_v = (offs_n[:, None] * stride_vbs + cur_kv_head * stride_vh +
                 offs_d[None, :] * stride_vd)
        k_ptrs = K + off_k
        v_ptrs = V + off_v

        block_mask = tl.where(
            block_start_loc < cur_batch_seq_len - cur_batch_ctx_len, 1, 0)

        # init alibi
        alibi_slope = tl.load(Alibi_slopes + cur_head)
        alibi_start_q = tl.arange(
            0, BLOCK_M) + block_start_loc + cur_batch_ctx_len
        alibi_start_k = cur_batch_ctx_len
        # # init debugger
        # offset_db_q = tl.arange(0, BLOCK_M) + block_start_loc
        # offset_db_k = tl.arange(0, BLOCK_N)
        # calc q[BLOCK_M, BLOCK_MODEL] mul k[prefix_len: , BLOCK_DMODEL]
        for start_n in range(0, block_mask * (start_m + 1) * BLOCK_M, BLOCK_N):
            start_n = tl.multiple_of(start_n, BLOCK_N)
            # -- compute qk ----
            k = tl.load(k_ptrs +
                        (cur_batch_in_all_start_index + start_n) * stride_kbs,
                        mask=dim_mask[:, None] &
                        ((start_n + offs_n[None, :])
                         < cur_batch_seq_len - cur_batch_ctx_len),
                        other=0.0)

            qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
            qk = tl.dot(q, k, acc=qk, input_precision='ieee')
            qk *= sm_scale
            qk = tl.where(offs_m[:, None] >= (start_n + offs_n[None, :]), qk,
                          float("-inf"))

            # load alibi
            alibi = (tl.arange(0, BLOCK_N)[None, :] + alibi_start_k -
                     alibi_start_q[:, None]) * alibi_slope
            alibi = tl.where(
                (alibi <= 0) & (alibi_start_q[:, None] < cur_batch_seq_len),
                alibi, float("-inf"))
            qk += alibi
            alibi_start_k += BLOCK_N

            # -- compute m_ij, p, l_ij
            m_ij = tl.max(qk, 1)
            m_i_new = tl.maximum(m_i, m_ij)
            p = tl.math.exp(qk - m_i_new[:, None])
            l_ij = tl.sum(p, 1)
            # -- update m_i and l_i

            alpha = tl.math.exp(m_i - m_i_new)
            l_i_new = alpha * l_i + l_ij
            # -- update output accumulator --
            # scale p
            # scale acc
            acc_scale = alpha
            # acc_scale = l_i / l_i_new * alpha
            acc = acc * acc_scale[:, None]
            # update acc
            v = tl.load(v_ptrs +
                        (cur_batch_in_all_start_index + start_n) * stride_vbs,
                        mask=dim_mask[None, :] &
                        ((start_n + offs_n[:, None])
                         < cur_batch_seq_len - cur_batch_ctx_len),
                        other=0.0)
            p = p.to(v.dtype)

            acc = tl.dot(p, v, acc=acc, input_precision='ieee')
            # update m_i and l_i
            l_i = l_i_new
            m_i = m_i_new

        acc = acc / l_i[:, None]

        # initialize pointers to output
        off_o = (
            (cur_batch_in_all_start_index + offs_m[:, None]) * stride_obs +
            cur_head * stride_oh + offs_d[None, :] * stride_od)
        out_ptrs = Out + off_o
        tl.store(out_ptrs,
                 acc,
                 mask=dim_mask[None, :] &
                 (offs_m[:, None] < cur_batch_seq_len - cur_batch_ctx_len))
        return

    @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_context_attention_fwd_config_filepath(cache_block_size, head_size, slide_window,
                                            use_alibi_slopes, skip_decode,
                                            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

        kernel_name = "context_attention_fwd_alibi" if use_alibi_slopes else "context_attention_fwd"
        json_file_name = (
            f"{kernel_name}-device={device_name}"
            f"-block_size={cache_block_size}"
            f"-BLOCK_DMODEL_PADDED={head_size_padded}"
            f"-SLIDING_WINDOW={slide_window}"
            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}", "context_attention_fwd", json_file_name
        )
        return config_file_path

    @functools.lru_cache
    def get_context_attention_fwd_config(
        cache_block_size,
        head_size,
        max_input_len,
        slide_window,
        use_alibi_slopes,
        skip_decode,
        kv_dtype
    ) -> Optional[Dict]:
        config_file_path = get_context_attention_fwd_config_filepath(cache_block_size, head_size,
                                                                slide_window, use_alibi_slopes,
                                                                skip_decode, 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:
                    config = configs[min(configs.keys(), key=lambda x: abs(x - max_input_len))]
                    # logger.info(f"context_attention_fwd 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 context_attention_fwd kernel config. Performance might "
                f"be sub-optimal! Config not found at {config_file_path}")
        return None

    @torch.inference_mode()
    def context_attention_fwd(q,
                              k,
                              v,
                              o,
                              kv_cache_dtype: str,
                              k_cache,
                              v_cache,
                              b_loc,
                              b_start_loc,
                              b_seq_len,
                              max_input_len,
                              k_scale: torch.Tensor,
                              v_scale: torch.Tensor,
                              alibi_slopes=None,
                              sliding_window=None,
                              sm_scale=None,
                              skip_decode=False):

        q_dtype_is_f32 = q.dtype is torch.float32
        # shape constraints
        Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1]
        assert Lq == Lk and Lk == Lv
        # round up Lk to a power of 2 - this is required for Triton block size
        Lk_padded = triton.next_power_of_2(Lk)

        # 0 means "disable"
        if sliding_window is None or sliding_window <= 0:
            sliding_window = 0
        use_alibi_slopes = False if alibi_slopes is None else True
        # open when kernel tuned
        # config = get_context_attention_fwd_config(v_cache.shape[3], Lk, max_input_len,
        #                                           sliding_window, use_alibi_slopes,
        #                                           skip_decode, k_cache.dtype)
        config = None
        if not config:
            config = ({'num_warps': NUM_WARPS, 'num_stages': 1}
                       if use_alibi_slopes else
                       {'num_warps': NUM_WARPS, 'num_stages': 1, 'USE_MATRIX_LOAD': False})

        if 'BLOCK_N' not in config or 'BLOCK_M' not in config:
            BLOCK = BASE_BLOCK // 2 if q_dtype_is_f32 else BASE_BLOCK
            config['BLOCK_M'] = BLOCK

            BLOCK = 64
            cache_ele_size = v_cache.element_size()
            if BLOCK * Lk_padded * cache_ele_size > 16384: # 64 * 128 * 2
                BLOCK = 1 << ((16384 // (Lk_padded * cache_ele_size)).bit_length() - 1)
            config['BLOCK_N'] = BLOCK
        # print(f"context_attention_fwd: {config=}")
        # Turing does have tensor core for float32 multiplication
        # use ieee as fallback for triton kernels work. There is also
        # warning on vllm/config.py to inform users this fallback
        # implementation
        IN_PRECISION = None
        if sm_scale is None:
            sm_scale = 1.0 / (Lq**0.5)
        batch, head = b_seq_len.shape[0], q.shape[1]
        num_queries_per_kv = q.shape[1] // k.shape[1]

        assert batch + 1 == len(b_start_loc)
        grid = (batch, head, triton.cdiv(max_input_len, config['BLOCK_M']))  # batch, head,

        if "fp8" in kv_cache_dtype and (k_cache.dtype == torch.uint8 or v_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)
            k_cache = k_cache.view(target_dtype)
            v_cache = v_cache.view(target_dtype)

        if alibi_slopes is not None:
            _fwd_kernel_alibi[grid](
                q,
                k,
                v,
                k_cache,
                v_cache,
                b_loc,
                sm_scale,
                k_scale,
                v_scale,
                b_start_loc,
                b_seq_len,
                alibi_slopes,
                v_cache.shape[3],
                k_cache.shape[4],
                o,
                b_loc.stride(0),
                b_loc.stride(1),
                q.stride(0),
                q.stride(1),
                q.stride(2),
                k.stride(0),
                k.stride(1),
                k.stride(2),
                v.stride(0),
                v.stride(1),
                v.stride(2),
                o.stride(0),
                o.stride(1),
                o.stride(2),
                k_cache.stride(0),
                k_cache.stride(1),
                k_cache.stride(2),
                k_cache.stride(3),
                k_cache.stride(
                    4
                ),  #[num_blocks, num_kv_heads, head_size/x, block_size, x]
                v_cache.stride(0),
                v_cache.stride(1),
                v_cache.stride(2),
                v_cache.stride(
                    3),  #[num_blocks, num_kv_heads, head_size, block_size]
                num_queries_per_kv=num_queries_per_kv,
                IN_PRECISION=IN_PRECISION,
                BLOCK_DMODEL=Lk,
                BLOCK_DMODEL_PADDED=Lk_padded,
                SKIP_DECODE=skip_decode,
                **config,
            )
            return

        _fwd_kernel[grid](
            q,
            k,
            v,
            k_cache,
            v_cache,
            b_loc,
            sm_scale,
            k_scale,
            v_scale,
            b_start_loc,
            b_seq_len,
            v_cache.shape[3],
            k_cache.shape[4],
            o,
            b_loc.stride(0),
            b_loc.stride(1),
            q.stride(0),
            q.stride(1),
            q.stride(2),
            k.stride(0),
            k.stride(1),
            k.stride(2),
            v.stride(0),
            v.stride(1),
            v.stride(2),
            o.stride(0),
            o.stride(1),
            o.stride(2),
            k_cache.stride(0),
            k_cache.stride(1),
            k_cache.stride(2),
            k_cache.stride(3),
            k_cache.stride(
                4),  #[num_blocks, num_kv_heads, head_size/x, block_size, x]
            v_cache.stride(0),
            v_cache.stride(1),
            v_cache.stride(2),
            v_cache.stride(
                3),  #[num_blocks, num_kv_heads, head_size, block_size]
            num_queries_per_kv=num_queries_per_kv,
            IN_PRECISION=IN_PRECISION,
            BLOCK_DMODEL=Lk,
            BLOCK_DMODEL_PADDED=Lk_padded,
            SLIDING_WINDOW=sliding_window,
            SKIP_DECODE=skip_decode,
            HEAD_DIM_PAD_REQ=(Lk != Lk_padded),
            max_input_len=max_input_len,
            **config,
        )
        return
