# Copyright (C) 2023-2025 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

"""
Memory-efficient attention for prefill.
It supports page size = 1 and prefill with KV cache (i.e. extend).
"""

import functools
import json
from typing import Any, Optional
import torch
import triton
import triton.language as tl

import os
import types

try:
    from triton.knobs import cache as cache_knob
except ImportError:
    # Triton builds without `triton.knobs` (e.g. 3.2.x in some images): disable saved-kernel path.
    cache_knob = types.SimpleNamespace(dir="__triton_knobs_unavailable__")

from aiter.ops.triton.prefill_attention import context_attention_fwd
from aiter.ops.triton.activation import _tanh
import aiter.ops.triton.utils.arch_info as arch_info
from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH

from triton import __version__ as triton_version
triton_minor_version = int(triton_version.split(".")[1])

@triton.jit
def _fwd_kernel(
    Q_Extend,
    K_Extend,
    V_Extend,
    O_Extend,
    K_Buffer,
    V_Buffer,
    qo_indptr,
    kv_indptr,
    kv_indices,
    mask_ptr,
    mask_indptr,
    sm_scale,
    kv_group_num,
    stride_qbs,
    stride_qh,
    stride_kbs,
    stride_kh,
    stride_vbs,
    stride_vh,
    stride_obs,
    stride_oh,
    stride_buf_kbs,
    stride_buf_kh,
    stride_buf_vbs,
    stride_buf_vh,
    logit_cap: tl.constexpr,
    Lq: tl.constexpr,
    Lv: tl.constexpr,
    BLOCK_DMODEL: tl.constexpr,
    BLOCK_DPE: tl.constexpr,
    BLOCK_DV: tl.constexpr,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    USE_CUSTOM_MASK: tl.constexpr,
    IS_CAUSAL: tl.constexpr,
    SKIP_PREFIX_CUSTOM_MASK: tl.constexpr,
    STORE_TRANSPOSE: tl.constexpr,
):
    cur_seq = tl.program_id(0)
    cur_head = tl.program_id(1)
    cur_block_m = tl.program_id(2)
    tl.assume(Q_Extend.to(tl.int64) >= 0)
    tl.assume(K_Extend.to(tl.int64) >= 0)
    tl.assume(V_Extend.to(tl.int64) >= 0)

    cur_kv_head = cur_head // kv_group_num

    cur_seq_extend_start_idx = tl.load(qo_indptr + cur_seq)
    cur_seq_len_extend = tl.load(qo_indptr + cur_seq + 1) - cur_seq_extend_start_idx
    cur_seq_kv_start_idx = tl.load(kv_indptr + cur_seq)
    cur_seq_len_prefix = tl.load(kv_indptr + cur_seq + 1) - cur_seq_kv_start_idx
    cur_seq_len = cur_seq_len_prefix + cur_seq_len_extend

    if USE_CUSTOM_MASK:
        cur_seq_mask_start_idx = tl.load(mask_indptr + cur_seq)

    offs_d = tl.arange(0, BLOCK_DMODEL)
    offs_dv = tl.arange(0, BLOCK_DV)
    offs_m = tl.arange(0, BLOCK_M)
    mask_m = (cur_block_m * BLOCK_M + offs_m) < cur_seq_len_extend

    mask_d = offs_d < Lq
    mask_dv = offs_dv < Lv

    offs_q = (
        (cur_seq_extend_start_idx + cur_block_m * BLOCK_M + offs_m[:, None])
        * stride_qbs
        + cur_head * stride_qh
        + offs_d[None, :]
    )
    q = tl.load(
        Q_Extend + offs_q, mask=(mask_m[:, None]) & (mask_d[None, :]), other=0.0
    )

    if BLOCK_DPE > 0:
        offs_dpe = BLOCK_DMODEL + tl.arange(0, BLOCK_DPE)
        offs_qpe = (
            (cur_seq_extend_start_idx + cur_block_m * BLOCK_M + offs_m[:, None])
            * stride_qbs
            + cur_head * stride_qh
            + offs_dpe[None, :]
        )
        qpe = tl.load(Q_Extend + offs_qpe, mask=mask_m[:, None], other=0.0)

    # stage 1: compute scores with prefix
    offs_n = tl.arange(0, BLOCK_N)

    acc = tl.zeros([BLOCK_M, BLOCK_DV], dtype=tl.float32)
    deno = tl.zeros([BLOCK_M], dtype=tl.float32)
    e_max = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")

    for start_n in range(0, cur_seq_len_prefix, BLOCK_N):
        start_n = tl.multiple_of(start_n, BLOCK_N)
        mask_n = (start_n + offs_n) < cur_seq_len_prefix

        offs_kv_loc = tl.load(
            kv_indices + cur_seq_kv_start_idx + start_n + offs_n, mask=mask_n, other=0
        )

        # load k in transposed way
        offs_buf_k = (
            offs_kv_loc[None, :] * stride_buf_kbs
            + cur_kv_head * stride_buf_kh
            + offs_d[:, None]
        )
        k = tl.load(
            K_Buffer + offs_buf_k, mask=(mask_n[None, :]) & (mask_d[:, None]), other=0.0
        )

        qk = tl.dot(q.to(k.dtype), k)
        if BLOCK_DPE > 0:
            offs_kpe = (
                offs_kv_loc[None, :] * stride_buf_kbs
                + cur_kv_head * stride_buf_kh
                + offs_dpe[:, None]
            )
            kpe = tl.load(
                K_Buffer + offs_kpe,
                mask=mask_n[None, :],
                other=0.0,
            )
            qk += tl.dot(qpe.to(kpe.dtype), kpe)
        qk *= sm_scale

        if logit_cap > 0:
            qk = logit_cap * _tanh(qk / logit_cap)

        if USE_CUSTOM_MASK and not SKIP_PREFIX_CUSTOM_MASK:
            custom_mask = tl.load(
                mask_ptr
                + cur_seq_mask_start_idx
                + (cur_block_m * BLOCK_M + offs_m[:, None]) * cur_seq_len
                + start_n
                + offs_n[None, :],
                mask=(mask_m[:, None] & mask_n[None, :]),
                other=0,
            )
            custom_mask &= mask_m[:, None] & mask_n[None, :]
            qk = tl.where(custom_mask, qk, float("-inf"))
        else:
            qk = tl.where(mask_m[:, None] & mask_n[None, :], qk, float("-inf"))

        n_e_max = tl.maximum(tl.max(qk, 1), e_max)
        re_scale = tl.exp(e_max - n_e_max)
        p = tl.exp(qk - n_e_max[:, None])
        deno = deno * re_scale + tl.sum(p, 1)

        offs_buf_v = (
            offs_kv_loc[:, None] * stride_buf_vbs
            + cur_kv_head * stride_buf_vh
            + offs_dv[None, :]
        )
        v = tl.load(
            V_Buffer + offs_buf_v, mask=mask_n[:, None] & mask_dv[None, :], other=0.0
        )
        p = p.to(v.dtype)
        acc = acc * re_scale[:, None] + tl.dot(p, v)

        e_max = n_e_max

    # stage 2: compute the triangle part

    cur_block_m_end = (
        cur_seq_len_extend
        if not IS_CAUSAL
        else tl.minimum(cur_seq_len_extend, (cur_block_m + 1) * BLOCK_M)
    )
    for start_n in range(0, cur_block_m_end, BLOCK_N):
        start_n = tl.multiple_of(start_n, BLOCK_N)
        mask_n = (start_n + offs_n) < cur_block_m_end

        # load k in transposed way
        offs_k = (
            (cur_seq_extend_start_idx + start_n + offs_n[None, :]) * stride_kbs
            + cur_kv_head * stride_kh
            + offs_d[:, None]
        )
        k = tl.load(
            K_Extend + offs_k, mask=(mask_n[None, :]) & (mask_d[:, None]), other=0.0
        )

        qk = tl.dot(q.to(k.dtype), k, out_dtype=tl.float32)
        if BLOCK_DPE > 0:
            offs_kpe = (
                (cur_seq_extend_start_idx + start_n + offs_n[None, :]) * stride_kbs
                + cur_kv_head * stride_kh
                + offs_dpe[:, None]
            )
            kpe = tl.load(
                K_Extend + offs_kpe,
                mask=mask_n[None, :],
                other=0.0,
            )
            qk += tl.dot(qpe.to(kpe.dtype), kpe)

        qk *= sm_scale

        if logit_cap > 0:
            qk = logit_cap * _tanh(qk / logit_cap)

        if USE_CUSTOM_MASK:
            custom_mask = tl.load(
                mask_ptr
                + cur_seq_mask_start_idx
                + (cur_block_m * BLOCK_M + offs_m[:, None]) * cur_seq_len
                + cur_seq_len_prefix
                + start_n
                + offs_n[None, :],
                mask=(mask_m[:, None] & mask_n[None, :]),
                other=0,
            )
            custom_mask &= mask_m[:, None] & mask_n[None, :]
            qk = tl.where(custom_mask, qk, float("-inf"))
        elif IS_CAUSAL:
            mask_causual = (cur_block_m * BLOCK_M + offs_m[:, None]) >= (
                start_n + offs_n[None, :]
            )
            mask_causual &= mask_m[:, None] & mask_n[None, :]
            qk = tl.where(mask_causual, qk, float("-inf"))
        else:
            mask_non_causal = mask_m[:, None] & mask_n[None, :]
            qk = tl.where(mask_non_causal, qk, float("-inf"))

        n_e_max = tl.maximum(tl.max(qk, 1), e_max)
        re_scale = tl.exp(e_max - n_e_max)
        p = tl.exp(qk - n_e_max[:, None])
        deno = deno * re_scale + tl.sum(p, 1)

        offs_v = (
            (cur_seq_extend_start_idx + start_n + offs_n[:, None]) * stride_vbs
            + cur_kv_head * stride_vh
            + offs_dv[None, :]
        )
        v = tl.load(
            V_Extend + offs_v, mask=mask_n[:, None] & mask_dv[None, :], other=0.0
        )
        p = p.to(v.dtype)
        acc = acc * re_scale[:, None] + tl.dot(p, v)

        e_max = n_e_max

    offs_o = (
        (cur_seq_extend_start_idx + cur_block_m * BLOCK_M + offs_m[:, None])
        * stride_obs
        + cur_head * stride_oh
        + offs_dv[None, :]
    )
    if STORE_TRANSPOSE:
        tl.store(
            O_Extend + offs_o.T,
            (acc / deno[:, None]).T,
            mask=(mask_m[:, None] & mask_dv[None, :]).T,
        )
    else:
        tl.store(
            O_Extend + offs_o,
            acc / deno[:, None],
            mask=mask_m[:, None] & mask_dv[None, :],
        )


@triton.jit
def _fwd_kernel_v2(
    Q_Extend,
    K_Extend,
    V_Extend,
    O_Extend,
    K_Buffer,
    V_Buffer,
    qo_indptr,
    kv_indptr,
    kv_indices,
    mask_ptr,
    mask_indptr,
    sink_ptr,
    window_kv_offset_ptr,
    sm_scale,
    k_scale,
    v_scale,
    kv_group_num,
    stride_qbs,
    stride_qh,
    stride_kbs,
    stride_kh,
    stride_vbs,
    stride_vh,
    stride_obs,
    stride_oh,
    stride_buf_kbs,
    stride_buf_kh,
    stride_buf_vbs,
    stride_buf_vh,
    SLIDING_WINDOW_SIZE: tl.constexpr,
    logit_cap: tl.constexpr,
    xai_temperature_len: tl.constexpr,
    Lq: tl.constexpr,
    Lv: tl.constexpr,
    BLOCK_DMODEL: tl.constexpr,
    BLOCK_DPE: tl.constexpr,
    BLOCK_DV: tl.constexpr,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    USE_CUSTOM_MASK: tl.constexpr,
    IS_CAUSAL: tl.constexpr,
    SKIP_PREFIX_CUSTOM_MASK: tl.constexpr,
    STORE_TRANSPOSE: tl.constexpr,
    HAS_SINK: tl.constexpr,
    head_num: tl.constexpr,
    USE_MLS: tl.constexpr,
    batch_size: tl.constexpr,
    # max_len_extend: tl.constexpr,
):
    cur_seq = tl.program_id(0)
    cur_head = tl.program_id(1)
    cur_block_m = tl.program_id(2)
    
    tl.assume(Q_Extend.to(tl.int64) >= 0)
    tl.assume(K_Extend.to(tl.int64) >= 0)
    tl.assume(V_Extend.to(tl.int64) >= 0)

    tl.assume(kv_group_num >= 0)
    tl.assume(stride_qbs >= 0)
    tl.assume(stride_qh >= 0)
    tl.assume(stride_kbs >= 0)
    tl.assume(stride_kh >= 0)
    tl.assume(stride_vbs >= 0)
    tl.assume(stride_vh >= 0)
    tl.assume(stride_obs >= 0)
    tl.assume(stride_oh >= 0)
    tl.assume(stride_buf_kbs >= 0)
    tl.assume(stride_buf_kh >= 0)
    tl.assume(stride_buf_vbs >= 0)
    tl.assume(stride_buf_vh >= 0)
    tl.assume(head_num >= 0)
    tl.assume(batch_size >= 0)
    # tl.assume(max_len_extend >= 0)

    kv_head_num = head_num // kv_group_num
    cur_kv_head = cur_head // kv_group_num

    cur_seq_extend_start_idx = tl.load(qo_indptr + cur_seq)
    cur_seq_len_extend = tl.load(qo_indptr + cur_seq + 1) - cur_seq_extend_start_idx
    cur_seq_kv_start_idx = tl.load(kv_indptr + cur_seq)
    cur_seq_len_prefix = tl.load(kv_indptr + cur_seq + 1) - cur_seq_kv_start_idx
    cur_seq_len = cur_seq_len_prefix + cur_seq_len_extend

    if USE_CUSTOM_MASK:
        cur_seq_mask_start_idx = tl.load(mask_indptr + cur_seq)

    window_kv_offset = 0
    if USE_CUSTOM_MASK and SLIDING_WINDOW_SIZE > 0:
        window_kv_offset = tl.load(window_kv_offset_ptr + cur_seq)

    offs_d = tl.arange(0, BLOCK_DMODEL)
    offs_dv = tl.arange(0, BLOCK_DV)
    offs_m = tl.arange(0, BLOCK_M)
    mask_m = (cur_block_m * BLOCK_M + offs_m) < cur_seq_len_extend

    mask_d = offs_d < Lq
    mask_dv = offs_dv < Lv

    ALL_MASK_M = tl.min(mask_m.to(tl.int32), axis=0) == 1
    ALL_MASK_D = tl.min(mask_d.to(tl.int32), axis=0) == 1
    ALL_MASK_DV = tl.min(mask_dv.to(tl.int32), axis=0) == 1

    if xai_temperature_len > 0:
        offs_qidx = cur_seq_len_prefix + cur_block_m * BLOCK_M + offs_m
        xai_temperature_scale = 1.0 / tl.log2(float(xai_temperature_len))
        xai_temperature_reg = tl.where(
            offs_qidx > xai_temperature_len,
            tl.log2(offs_qidx.to(tl.float32)) * xai_temperature_scale,
            1.0,
        )

    if USE_MLS:
        q = tl.matrix_load(
            Q_Extend + cur_head * stride_qh,
            shape=(head_num, Lq),
            strides=(stride_qbs, 1),
            block_shape=(BLOCK_M, BLOCK_DMODEL),
            offsets=((cur_seq_extend_start_idx + cur_block_m * BLOCK_M).to(tl.int32), 0),
        )
        if not (ALL_MASK_M & ALL_MASK_D):
            q = tl.where((mask_m[:, None]) & (mask_d[None, :]), q, 0.0)
    else:
        offs_q = (
            (cur_seq_extend_start_idx + cur_block_m * BLOCK_M + offs_m[:, None])
            * stride_qbs
            + cur_head * stride_qh
            + offs_d[None, :]
        )
        q = tl.load(
            Q_Extend + offs_q, mask=(mask_m[:, None]) & (mask_d[None, :]), other=0.0
        )

    if BLOCK_DPE > 0:
        offs_dpe = BLOCK_DMODEL + tl.arange(0, BLOCK_DPE)
        if USE_MLS:
            qpe = tl.matrix_load(Q_Extend + cur_head * stride_qh,
                                shape=(head_num, Lq),
                                strides=(stride_qbs, 1),
                                block_shape=(BLOCK_M, BLOCK_DPE),
                                offsets=((cur_seq_extend_start_idx + cur_block_m * BLOCK_M).to(tl.int32),
                                         BLOCK_DMODEL),
                                )
            if not ALL_MASK_M:
                qpe = tl.where(mask_m[:, None], qpe, 0.0)
        else:
            offs_qpe = (
                (cur_seq_extend_start_idx + cur_block_m * BLOCK_M + offs_m[:, None])
                * stride_qbs
                + cur_head * stride_qh
                + offs_dpe[None, :]
            )
            qpe = tl.load(Q_Extend + offs_qpe, mask=mask_m[:, None], other=0.0)

    offs_n = tl.arange(0, BLOCK_N)

    acc = tl.zeros([BLOCK_M, BLOCK_DV], dtype=tl.float32)
    deno = tl.zeros([BLOCK_M], dtype=tl.float32)
    e_max = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")

    for start_n in range(0, cur_seq_len_prefix, BLOCK_N):
        start_n = tl.multiple_of(start_n, BLOCK_N)
        mask_n = (start_n + offs_n) < cur_seq_len_prefix
        ALL_MASK_N = tl.min(mask_n.to(tl.int32), axis=0) == 1

        final_mask = mask_m[:, None] & mask_n[None, :]
        if USE_CUSTOM_MASK and not SKIP_PREFIX_CUSTOM_MASK:
            if USE_MLS:
                custom_mask = tl.matrix_load(
                    mask_ptr + cur_seq_mask_start_idx,
                    shape=(cur_seq_len_extend, cur_seq_len + window_kv_offset),
                    strides=(cur_seq_len + window_kv_offset, 1),
                    block_shape=(BLOCK_M, BLOCK_N),
                    offsets=((cur_block_m * BLOCK_M).to(tl.int32),
                             (window_kv_offset + start_n).to(tl.int32)),
                )
                if not (ALL_MASK_M & ALL_MASK_N):
                    custom_mask = tl.where((mask_m[:, None]) & (mask_n[None, :]), custom_mask, 0)
            else:
                custom_mask = tl.load(
                    mask_ptr
                    + cur_seq_mask_start_idx
                    + (cur_block_m * BLOCK_M + offs_m[:, None])
                    * (cur_seq_len + window_kv_offset)
                    + window_kv_offset
                    + start_n
                    + offs_n[None, :],
                    mask=(mask_m[:, None] & mask_n[None, :]),
                    other=0,
                )
            final_mask &= custom_mask
        if SLIDING_WINDOW_SIZE > 0:
            window_mask = (
                cur_seq_len_prefix + cur_block_m * BLOCK_M + offs_m[:, None]
            ) <= (start_n + offs_n[None, :] + SLIDING_WINDOW_SIZE)
            final_mask &= window_mask

        SKIP_TILE = False
        if (USE_CUSTOM_MASK and not SKIP_PREFIX_CUSTOM_MASK) or SLIDING_WINDOW_SIZE > 0:
            SKIP_TILE = tl.max(tl.max(final_mask.to(tl.int32), axis=1), axis=0) == 0

        if not SKIP_TILE:
            offs_kv_loc = tl.load(
                kv_indices + cur_seq_kv_start_idx + start_n + offs_n,
                mask=mask_n,
                other=0,
            )

            # offs_kv_next = offs_kv_loc[1:]           # [1, 2, ..., N-1]
            # offs_kv_curr = offs_kv_loc[:-1]          # [0, 1, ..., N-2]
            # diff = offs_kv_next - offs_kv_curr       # 长度 N-1
            # is_continuous = tl.all(diff == 1)

            # if USE_MLS and is_continuous:
            #     offs_kv_start_idx = tl.load(
            #         kv_indices + cur_seq_kv_start_idx + start_n,
            #     )

            #     k = tl.matrix_load(
            #         K_Buffer + cur_kv_head * stride_buf_kh,
            #         shape=(Lq, cur_seq_len_prefix.to(tl.int32)),
            #         strides=(1, stride_buf_kbs),
            #         block_shape=(BLOCK_DMODEL, BLOCK_N),
            #         offsets=(0, offs_kv_start_idx.to(tl.int32)),
            #         # mask=(mask_m[:, None] & mask_n[None, :]),
            #         # boundary_check=(0, 1),
            #     )
            #     if not (ALL_MASK_N & ALL_MASK_D):
            #         k = tl.where((mask_n[None, :]) & (mask_d[:, None]), k, 0.0)
            # else:
            #     offs_buf_k = (
            #         offs_kv_loc[None, :] * stride_buf_kbs
            #         + cur_kv_head * stride_buf_kh
            #         + offs_d[:, None]
            #     )
            #     k = tl.load(
            #         K_Buffer + offs_buf_k,
            #         mask=(mask_n[None, :]) & (mask_d[:, None]),
            #         other=0.0,
            #     )
            offs_buf_k = (
                offs_kv_loc[None, :] * stride_buf_kbs
                + cur_kv_head * stride_buf_kh
                + offs_d[:, None]
            )
            k = tl.load(
                K_Buffer + offs_buf_k,
                mask=(mask_n[None, :]) & (mask_d[:, None]),
                other=0.0,
            )
            qk = tl.dot(q.to(k.dtype), k)
            if BLOCK_DPE > 0:
                # if USE_MLS:
                #     kpe = tl.matrix_load(
                #         K_Buffer + cur_kv_head * stride_buf_kh,
                #         shape=(Lq, cur_seq_len_prefix.to(tl.int32)),
                #         strides=(1, stride_buf_kbs),
                #         block_shape=(BLOCK_DPE, BLOCK_N),
                #         offsets=(BLOCK_DMODEL, offs_kv_start_idx.to(tl.int32)),
                #         # mask=(mask_m[:, None] & mask_n[None, :]),
                #         # boundary_check=(0, 1),
                #     )
                #     if not ALL_MASK_N:
                #         kpe = tl.where((mask_n[None, :]), kpe, 0.0)
                # else:
                #     offs_kpe = (
                #         offs_kv_loc[None, :] * stride_buf_kbs
                #         + cur_kv_head * stride_buf_kh
                #         + offs_dpe[:, None]
                #     )
                #     kpe = tl.load(
                #         K_Buffer + offs_kpe,
                #         mask=mask_n[None, :],
                #         other=0.0,
                #     )
                offs_kpe = (
                    offs_kv_loc[None, :] * stride_buf_kbs
                    + cur_kv_head * stride_buf_kh
                    + offs_dpe[:, None]
                )
                kpe = tl.load(
                    K_Buffer + offs_kpe,
                    mask=mask_n[None, :],
                    other=0.0,
                )
                qk += tl.dot(qpe.to(kpe.dtype), kpe)
            qk *= sm_scale * k_scale

            if logit_cap > 0:
                qk = logit_cap * _tanh(qk / logit_cap)

            if xai_temperature_len > 0:
                qk *= xai_temperature_reg[:, None]

            qk = tl.where(final_mask, qk, float("-inf"))

            row_max = tl.max(qk, 1)
            row_max_fixed = tl.where(row_max == float("-inf"), -1e20, row_max)
            n_e_max = tl.maximum(row_max_fixed, e_max)

            re_scale = tl.exp(e_max - n_e_max)
            p = tl.exp(qk - n_e_max[:, None])
            deno = deno * re_scale + tl.sum(p, 1)

            # if USE_MLS:
            #     v = tl.matrix_load(
            #         V_Buffer + cur_kv_head * stride_buf_vh,
            #         shape=(cur_seq_len_prefix.to(tl.int32), Lv),
            #         strides=(stride_buf_kbs, 1),
            #         block_shape=(BLOCK_N, BLOCK_DV),
            #         offsets=(offs_kv_start_idx.to(tl.int32), 0),
            #         # mask=(mask_m[:, None] & mask_n[None, :]),
            #         # boundary_check=(0, 1),
            #     )
            #     if not (ALL_MASK_N & ALL_MASK_DV):
            #         v = tl.where((mask_n[:, None] & mask_dv[None, :]), v, 0.0)
            # else:
            #     offs_buf_v = (
            #         offs_kv_loc[:, None] * stride_buf_vbs
            #         + cur_kv_head * stride_buf_vh
            #         + offs_dv[None, :]
            #     )
            #     v = tl.load(
            #         V_Buffer + offs_buf_v,
            #         mask=mask_n[:, None] & mask_dv[None, :],
            #         other=0.0,
            #     )
            offs_buf_v = (
                offs_kv_loc[:, None] * stride_buf_vbs
                + cur_kv_head * stride_buf_vh
                + offs_dv[None, :]
            )
            v = tl.load(
                V_Buffer + offs_buf_v,
                mask=mask_n[:, None] & mask_dv[None, :],
                other=0.0,
            )
            p = p.to(v.dtype)
            acc = acc * re_scale[:, None] + tl.dot(p, v) * v_scale

            e_max = n_e_max

    cur_block_m_end = (
        cur_seq_len_extend
        if not IS_CAUSAL
        else tl.minimum(cur_seq_len_extend, (cur_block_m + 1) * BLOCK_M)
    )
    for start_n in range(0, cur_block_m_end, BLOCK_N):
        start_n = tl.multiple_of(start_n, BLOCK_N)
        mask_n = (start_n + offs_n) < cur_block_m_end
        ALL_MASK_N = tl.min(mask_n.to(tl.int32), axis=0) == 1

        final_mask = mask_m[:, None] & mask_n[None, :]
        if USE_CUSTOM_MASK:
            if USE_MLS:
                custom_mask = tl.matrix_load(
                    mask_ptr + cur_seq_mask_start_idx,
                    shape=(cur_block_m_end, cur_seq_len + window_kv_offset),
                    strides=(cur_seq_len + window_kv_offset, 1),
                    block_shape=(BLOCK_M, BLOCK_N),
                    offsets=((cur_block_m * BLOCK_M).to(tl.int32),
                             (window_kv_offset + cur_seq_len_prefix + start_n).to(tl.int32)),
                )
                if not (ALL_MASK_M & ALL_MASK_N):
                    custom_mask = tl.where((mask_m[:, None]) & (mask_n[None, :]), custom_mask, 0)
            else:
                custom_mask = tl.load(
                    mask_ptr
                    + cur_seq_mask_start_idx
                    + (cur_block_m * BLOCK_M + offs_m[:, None])
                    * (cur_seq_len + window_kv_offset)
                    + window_kv_offset
                    + cur_seq_len_prefix
                    + start_n
                    + offs_n[None, :],
                    mask=(mask_m[:, None] & mask_n[None, :]),
                    other=0,
                )
            custom_mask &= mask_m[:, None] & mask_n[None, :]
            final_mask &= custom_mask
        elif IS_CAUSAL:
            mask_causual = (cur_block_m * BLOCK_M + offs_m[:, None]) >= (
                start_n + offs_n[None, :]
            )
            mask_causual &= mask_m[:, None] & mask_n[None, :]
            final_mask &= mask_causual
        else:
            mask_non_causal = mask_m[:, None] & mask_n[None, :]
            final_mask &= mask_non_causal

        if SLIDING_WINDOW_SIZE > 0:
            window_mask = (cur_block_m * BLOCK_M + offs_m[:, None]) <= (
                start_n + offs_n[None, :] + SLIDING_WINDOW_SIZE
            )
            final_mask &= window_mask

        SKIP_TILE = False
        if USE_CUSTOM_MASK or SLIDING_WINDOW_SIZE > 0:
            SKIP_TILE = tl.max(tl.max(final_mask.to(tl.int32), axis=1), axis=0) == 0

        if not SKIP_TILE:
            if USE_MLS:
                k = tl.matrix_load(
                    K_Extend + cur_kv_head * stride_kh,
                    shape=(kv_head_num, Lq),
                    strides=(1, stride_kbs),
                    block_shape=(BLOCK_DMODEL, BLOCK_N),
                    offsets=(0, (cur_seq_extend_start_idx + start_n).to(tl.int32)),
                )
                if not (ALL_MASK_N & ALL_MASK_D):
                    k = tl.where((mask_d[:, None]) & (mask_n[None, :]), k, 0.0)
            else:
                offs_k = (
                    (cur_seq_extend_start_idx + start_n + offs_n[None, :]) * stride_kbs
                    + cur_kv_head * stride_kh
                    + offs_d[:, None]
                )
                k = tl.load(
                    K_Extend + offs_k, mask=(mask_n[None, :]) & (mask_d[:, None]), other=0.0
                )

            qk = tl.dot(q.to(k.dtype), k, out_dtype=tl.float32)
            if BLOCK_DPE > 0:
                if USE_MLS:
                    kpe = tl.matrix_load(
                        K_Extend + cur_kv_head * stride_kh,
                        shape=(kv_head_num, Lq),
                        strides=(1, stride_kbs),
                        block_shape=(BLOCK_DPE, BLOCK_N),
                        offsets=(BLOCK_DMODEL, (cur_seq_extend_start_idx + start_n).to(tl.int32)),
                    )
                    if not ALL_MASK_N:
                        kpe = tl.where(mask_n[None, :], kpe, 0.0)
                else:
                    offs_kpe = (
                        (cur_seq_extend_start_idx + start_n + offs_n[None, :]) * stride_kbs
                        + cur_kv_head * stride_kh
                        + offs_dpe[:, None]
                    )
                    kpe = tl.load(
                        K_Extend + offs_kpe,
                        mask=mask_n[None, :],
                        other=0.0,
                    )
                qk += tl.dot(qpe.to(kpe.dtype), kpe)

            qk *= sm_scale

            if logit_cap > 0:
                qk = logit_cap * _tanh(qk / logit_cap)

            if xai_temperature_len > 0:
                qk *= xai_temperature_reg[:, None]

            qk = tl.where(final_mask, qk, float("-inf"))

            row_max = tl.max(qk, 1)
            row_max_fixed = tl.where(row_max == float("-inf"), -1e20, row_max)
            n_e_max = tl.maximum(row_max_fixed, e_max)

            re_scale = tl.exp(e_max - n_e_max)
            p = tl.exp(qk - n_e_max[:, None])
            deno = deno * re_scale + tl.sum(p, 1)

            if USE_MLS:
                v = tl.matrix_load(
                    V_Extend + cur_kv_head * stride_vh,
                    shape=(kv_head_num, Lv),
                    strides=(stride_vbs, 1),
                    block_shape=(BLOCK_N, BLOCK_DV),
                    offsets=((cur_seq_extend_start_idx + start_n).to(tl.int32), 0),
                )
                if not (ALL_MASK_N & ALL_MASK_DV):
                    v = tl.where((mask_n[:, None]) & (mask_dv[None, :]), v, 0.0)
            else:
                offs_v = (
                    (cur_seq_extend_start_idx + start_n + offs_n[:, None]) * stride_vbs
                    + cur_kv_head * stride_vh
                    + offs_dv[None, :]
                )
                v = tl.load(
                    V_Extend + offs_v, mask=mask_n[:, None] & mask_dv[None, :], other=0.0
                )
            p = p.to(v.dtype)
            acc = acc * re_scale[:, None] + tl.dot(p, v)

            e_max = n_e_max

    if HAS_SINK:
        cur_sink = tl.load(sink_ptr + cur_head)
        deno += tl.exp(cur_sink - e_max)

    offs_o = (
        (cur_seq_extend_start_idx + cur_block_m * BLOCK_M + offs_m[:, None])
        * stride_obs
        + cur_head * stride_oh
        + offs_dv[None, :]
    )
    if STORE_TRANSPOSE:
        tl.store(
            O_Extend + offs_o.T,
            (acc / deno[:, None]).T,
            mask=(mask_m[:, None] & mask_dv[None, :]).T,
        )
    else:
        tl.store(
            O_Extend + offs_o,
            acc / deno[:, None],
            mask=mask_m[:, None] & mask_dv[None, :],
        )


@triton.jit
def _fwd_kernel_v2_decode(
    Q_Extend,
    K_Extend,
    V_Extend,
    O_Extend,
    K_Buffer,
    V_Buffer,
    qo_indptr,
    kv_indptr,
    kv_indices,
    mask_ptr,
    mask_indptr,
    sink_ptr,
    window_kv_offset_ptr,
    sm_scale,
    k_scale,
    v_scale,
    stride_qbs,
    stride_qh,
    stride_kbs,
    stride_kh,
    stride_vbs,
    stride_vh,
    stride_obs,
    stride_oh,
    stride_buf_kbs,
    stride_buf_kh,
    stride_buf_vbs,
    stride_buf_vh,
    SLIDING_WINDOW_SIZE: tl.constexpr,
    logit_cap: tl.constexpr,
    xai_temperature_len: tl.constexpr,
    Lq: tl.constexpr,
    Lv: tl.constexpr,
    BLOCK_DMODEL: tl.constexpr,
    BLOCK_DPE: tl.constexpr,
    BLOCK_DV: tl.constexpr,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    USE_CUSTOM_MASK: tl.constexpr,
    IS_CAUSAL: tl.constexpr,
    SKIP_PREFIX_CUSTOM_MASK: tl.constexpr,
    STORE_TRANSPOSE: tl.constexpr,
    HAS_SINK: tl.constexpr,
    kv_group_num: tl.constexpr,
    num_query_heads: tl.constexpr,
    USE_MLS: tl.constexpr,
    batch_size: tl.constexpr,
    # max_len_extend: tl.constexpr,
):
    """
    v2 decode: grid (batch, num_kv_heads, cdiv(max_len_extend, Q_SEQ)) with
    Q_SEQ = BLOCK_M // kv_group_num (same as unified ``BLOCK_Q``; floor, not ceil). If BLOCK_M is not
    a multiple of G, adjacent ``cur_block_m`` may overlap in query_pos like unified, but `mask_m`
    and sequence bounds keep correctness. BLOCK_M is a power of 2 (host). Require BLOCK_M // G >= 1 to launch.
    """
    # Per unified: BLOCK_Q = BLOCK_M // G; stride in query token index for each +1 of cur_block_m.
    Q_SEQ: tl.constexpr = BLOCK_M // kv_group_num
    cur_seq = tl.program_id(0)
    cur_kv_head = tl.program_id(1)
    cur_block_m = tl.program_id(2)
    tl.assume(Q_Extend.to(tl.int64) >= 0)
    tl.assume(K_Extend.to(tl.int64) >= 0)
    tl.assume(V_Extend.to(tl.int64) >= 0)

    tl.assume(stride_qbs >= 0)
    tl.assume(stride_qh >= 0)
    tl.assume(stride_kbs >= 0)
    tl.assume(stride_kh >= 0)
    tl.assume(stride_vbs >= 0)
    tl.assume(stride_vh >= 0)
    tl.assume(stride_obs >= 0)
    tl.assume(stride_oh >= 0)
    tl.assume(stride_buf_kbs >= 0)
    tl.assume(stride_buf_kh >= 0)
    tl.assume(stride_buf_vbs >= 0)
    tl.assume(stride_buf_vh >= 0)
    tl.assume(batch_size >= 0)
    # tl.assume(max_len_extend >= 0)

    kv_head_num = num_query_heads // kv_group_num

    cur_seq_extend_start_idx = tl.load(qo_indptr + cur_seq)
    cur_seq_len_extend = tl.load(qo_indptr + cur_seq + 1) - cur_seq_extend_start_idx
    cur_seq_kv_start_idx = tl.load(kv_indptr + cur_seq)
    cur_seq_len_prefix = tl.load(kv_indptr + cur_seq + 1) - cur_seq_kv_start_idx
    cur_seq_len = cur_seq_len_prefix + cur_seq_len_extend

    if cur_block_m * Q_SEQ >= cur_seq_len_extend:
        return

    if USE_CUSTOM_MASK:
        cur_seq_mask_start_idx = tl.load(mask_indptr + cur_seq)

    window_kv_offset = 0
    if USE_CUSTOM_MASK and SLIDING_WINDOW_SIZE > 0:
        window_kv_offset = tl.load(window_kv_offset_ptr + cur_seq)

    offs_d = tl.arange(0, BLOCK_DMODEL)
    offs_dv = tl.arange(0, BLOCK_DV)
    offs_m = tl.arange(0, BLOCK_M)
    # unified_attention-style: per offs_m, row = offs_m // G, h = offs_m % G (G = kv_group_num)
    query_pos = cur_block_m * Q_SEQ + (offs_m // kv_group_num)
    q_head_in_group = offs_m % kv_group_num
    query_offset_0 = cur_seq_extend_start_idx + query_pos
    query_offset_1 = cur_kv_head * kv_group_num + q_head_in_group
    mask_m = (query_pos < cur_seq_len_extend) & (query_offset_1 < num_query_heads)

    mask_d = offs_d < Lq
    mask_dv = offs_dv < Lv

    ALL_MASK_M = tl.min(mask_m.to(tl.int32), axis=0) == 1
    ALL_MASK_D = tl.min(mask_d.to(tl.int32), axis=0) == 1
    ALL_MASK_DV = tl.min(mask_dv.to(tl.int32), axis=0) == 1

    if xai_temperature_len > 0:
        offs_qidx = cur_seq_len_prefix + query_pos
        xai_temperature_scale = 1.0 / tl.log2(float(xai_temperature_len))
        xai_temperature_reg = tl.where(
            offs_qidx > xai_temperature_len,
            tl.log2(offs_qidx.to(tl.float32)) * xai_temperature_scale,
            1.0,
        )

    offs_q = (
        query_offset_0[:, None] * stride_qbs
        + query_offset_1[:, None] * stride_qh
        + offs_d[None, :]
    )
    q = tl.load(
        Q_Extend + offs_q, mask=(mask_m[:, None]) & (mask_d[None, :]), other=0.0
    )

    if BLOCK_DPE > 0:
        offs_dpe = BLOCK_DMODEL + tl.arange(0, BLOCK_DPE)
        offs_qpe = (
            query_offset_0[:, None] * stride_qbs
            + query_offset_1[:, None] * stride_qh
            + offs_dpe[None, :]
        )
        qpe = tl.load(Q_Extend + offs_qpe, mask=mask_m[:, None], other=0.0)

    offs_n = tl.arange(0, BLOCK_N)

    acc = tl.zeros([BLOCK_M, BLOCK_DV], dtype=tl.float32)
    deno = tl.zeros([BLOCK_M], dtype=tl.float32)
    e_max = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")

    for start_n in range(0, cur_seq_len_prefix, BLOCK_N):
        start_n = tl.multiple_of(start_n, BLOCK_N)
        mask_n = (start_n + offs_n) < cur_seq_len_prefix
        ALL_MASK_N = tl.min(mask_n.to(tl.int32), axis=0) == 1

        final_mask = mask_m[:, None] & mask_n[None, :]
        if USE_CUSTOM_MASK and not SKIP_PREFIX_CUSTOM_MASK:
            # if USE_MLS:
            #     group_id = offs_m // kv_group_num
            #     custom_mask_group = tl.matrix_load(
            #         mask_ptr + cur_seq_mask_start_idx,
            #         shape=(cur_seq_len_prefix, cur_seq_len + window_kv_offset),
            #         strides=(cur_seq_len + window_kv_offset, 1),
            #         block_shape=(BLOCK_M // kv_group_num, BLOCK_N),
            #         offsets=((cur_block_m * Q_SEQ).to(tl.int32),
            #                  (window_kv_offset + start_n).to(tl.int32)),
            #     )
            #     custom_mask = custom_mask_group[group_id[:, None], offs_n[None, :]]
            #     if not (ALL_MASK_M & ALL_MASK_N):
            #         custom_mask = tl.where((mask_m[:, None] & mask_n[None, :]), custom_mask, 0)
            # else:
            #     custom_mask = tl.load(
            #         mask_ptr
            #         + cur_seq_mask_start_idx
            #         + (query_pos[:, None]) * (cur_seq_len + window_kv_offset)
            #         + window_kv_offset
            #         + start_n
            #         + offs_n[None, :],
            #         mask=(mask_m[:, None] & mask_n[None, :]),
            #         other=0,
            #     )
            custom_mask = tl.load(
                mask_ptr
                + cur_seq_mask_start_idx
                + (query_pos[:, None]) * (cur_seq_len + window_kv_offset)
                + window_kv_offset
                + start_n
                + offs_n[None, :],
                mask=(mask_m[:, None] & mask_n[None, :]),
                other=0,
            )
            final_mask &= custom_mask
        if SLIDING_WINDOW_SIZE > 0:
            window_mask = (
                cur_seq_len_prefix + query_pos[:, None]
            ) <= (start_n + offs_n[None, :] + SLIDING_WINDOW_SIZE)
            final_mask &= window_mask

        SKIP_TILE = False
        if (USE_CUSTOM_MASK and not SKIP_PREFIX_CUSTOM_MASK) or SLIDING_WINDOW_SIZE > 0:
            SKIP_TILE = tl.max(tl.max(final_mask.to(tl.int32), axis=1), axis=0) == 0

        if not SKIP_TILE:
            offs_kv_loc = tl.load(
                kv_indices + cur_seq_kv_start_idx + start_n + offs_n,
                mask=mask_n,
                other=0,
            )

            # if USE_MLS:
            #     offs_kv_start_idx = tl.load(
            #         kv_indices + cur_seq_kv_start_idx + start_n,
            #     )

            #     k = tl.matrix_load(
            #         K_Buffer + cur_kv_head * stride_buf_kh,
            #         shape=(Lq, cur_seq_len_prefix.to(tl.int32)),
            #         strides=(1, stride_buf_kbs),
            #         block_shape=(BLOCK_DMODEL, BLOCK_N),
            #         offsets=(0, offs_kv_start_idx.to(tl.int32)),
            #     )
            #     if not (ALL_MASK_N & ALL_MASK_D):
            #         k = tl.where((mask_n[None, :]) & (mask_d[:, None]), k, 0.0)
            # else:
            #     offs_buf_k = (
            #         offs_kv_loc[None, :] * stride_buf_kbs
            #         + cur_kv_head * stride_buf_kh
            #         + offs_d[:, None]
            #     )
            #     k = tl.load(
            #         K_Buffer + offs_buf_k,
            #         mask=(mask_n[None, :]) & (mask_d[:, None]),
            #         other=0.0,
            #     )
            offs_buf_k = (
                offs_kv_loc[None, :] * stride_buf_kbs
                + cur_kv_head * stride_buf_kh
                + offs_d[:, None]
            )
            k = tl.load(
                K_Buffer + offs_buf_k,
                mask=(mask_n[None, :]) & (mask_d[:, None]),
                other=0.0,
            )
            qk = tl.dot(q.to(k.dtype), k)
            if BLOCK_DPE > 0:
                # if USE_MLS:
                #     kpe = tl.matrix_load(
                #         K_Buffer + cur_kv_head * stride_buf_kh,
                #         shape=(Lq, cur_seq_len_prefix.to(tl.int32)),
                #         strides=(1, stride_buf_kbs),
                #         block_shape=(BLOCK_DPE, BLOCK_N),
                #         offsets=(BLOCK_DMODEL, offs_kv_start_idx.to(tl.int32)),
                #     )
                #     if not ALL_MASK_N:
                #         kpe = tl.where((mask_n[None, :]), kpe, 0.0)
                # else:
                #     offs_kpe = (
                #         offs_kv_loc[None, :] * stride_buf_kbs
                #         + cur_kv_head * stride_buf_kh
                #         + offs_dpe[:, None]
                #     )
                #     kpe = tl.load(
                #         K_Buffer + offs_kpe,
                #         mask=mask_n[None, :],
                #         other=0.0,
                #     )
                offs_kpe = (
                    offs_kv_loc[None, :] * stride_buf_kbs
                    + cur_kv_head * stride_buf_kh
                    + offs_dpe[:, None]
                )
                kpe = tl.load(
                    K_Buffer + offs_kpe,
                    mask=mask_n[None, :],
                    other=0.0,
                )
                qk += tl.dot(qpe.to(kpe.dtype), kpe)
            qk *= sm_scale * k_scale

            if logit_cap > 0:
                qk = logit_cap * _tanh(qk / logit_cap)

            if xai_temperature_len > 0:
                qk *= xai_temperature_reg[:, None]

            qk = tl.where(final_mask, qk, float("-inf"))

            row_max = tl.max(qk, 1)
            row_max_fixed = tl.where(row_max == float("-inf"), -1e20, row_max)
            n_e_max = tl.maximum(row_max_fixed, e_max)

            re_scale = tl.exp(e_max - n_e_max)
            p = tl.exp(qk - n_e_max[:, None])
            deno = deno * re_scale + tl.sum(p, 1)

            # if USE_MLS:
            #     v = tl.matrix_load(
            #         V_Buffer + cur_kv_head * stride_buf_vh,
            #         shape=(cur_seq_len_prefix.to(tl.int32), Lv),
            #         strides=(stride_buf_kbs, 1),
            #         block_shape=(BLOCK_N, BLOCK_DV),
            #         offsets=(offs_kv_start_idx.to(tl.int32), 0),
            #     )
            #     if not (ALL_MASK_N & ALL_MASK_DV):
            #         v = tl.where((mask_n[:, None] & mask_dv[None, :]), v, 0.0)
            # else:
            #     offs_buf_v = (
            #         offs_kv_loc[:, None] * stride_buf_vbs
            #         + cur_kv_head * stride_buf_vh
            #         + offs_dv[None, :]
            #     )
            #     v = tl.load(
            #         V_Buffer + offs_buf_v, mask=mask_n[:, None] & mask_dv[None, :], other=0.0
            #     )
            offs_buf_v = (
                offs_kv_loc[:, None] * stride_buf_vbs
                + cur_kv_head * stride_buf_vh
                + offs_dv[None, :]
            )
            v = tl.load(
                V_Buffer + offs_buf_v, mask=mask_n[:, None] & mask_dv[None, :], other=0.0
            )
            p = p.to(v.dtype)
            acc = acc * re_scale[:, None] + tl.dot(p, v) * v_scale

            e_max = n_e_max

    cur_block_m_end = (
        cur_seq_len_extend
        if not IS_CAUSAL
        else tl.minimum(cur_seq_len_extend, (cur_block_m + 1) * Q_SEQ)
    )
    for start_n in range(0, cur_block_m_end, BLOCK_N):
        start_n = tl.multiple_of(start_n, BLOCK_N)
        mask_n = (start_n + offs_n) < cur_block_m_end
        ALL_MASK_N = tl.min(mask_n.to(tl.int32), axis=0) == 1

        final_mask = mask_m[:, None] & mask_n[None, :]
        if USE_CUSTOM_MASK:
            # if USE_MLS:
            #     group_id = offs_m // kv_group_num
            #     custom_mask_group = tl.matrix_load(
            #         mask_ptr + cur_seq_mask_start_idx,
            #         shape=(cur_block_m_end, cur_seq_len + window_kv_offset),
            #         strides=(cur_seq_len + window_kv_offset, 1),
            #         block_shape=(BLOCK_M // kv_group_num, BLOCK_N),
            #         offsets=((cur_block_m * Q_SEQ).to(tl.int32),
            #                  (window_kv_offset + cur_seq_len_prefix + start_n).to(tl.int32)),
            #     )
            #     custom_mask = custom_mask_group[group_id[:, None], offs_n[None, :]]
            #     if not (ALL_MASK_M & ALL_MASK_N):
            #         custom_mask = tl.where((mask_m[:, None] & mask_n[None, :]), custom_mask, 0)
            # else:
            #     custom_mask = tl.load(
            #         mask_ptr
            #         + cur_seq_mask_start_idx
            #         + (query_pos[:, None]) * (cur_seq_len + window_kv_offset)
            #         + window_kv_offset
            #         + cur_seq_len_prefix
            #         + start_n
            #         + offs_n[None, :],
            #         mask=(mask_m[:, None] & mask_n[None, :]),
            #         other=0,
            #     )
            custom_mask = tl.load(
                mask_ptr
                + cur_seq_mask_start_idx
                + (query_pos[:, None]) * (cur_seq_len + window_kv_offset)
                + window_kv_offset
                + cur_seq_len_prefix
                + start_n
                + offs_n[None, :],
                mask=(mask_m[:, None] & mask_n[None, :]),
                other=0,
            )
            custom_mask &= mask_m[:, None] & mask_n[None, :]
            final_mask &= custom_mask
        elif IS_CAUSAL:
            mask_causual = query_pos[:, None] >= (start_n + offs_n[None, :])
            mask_causual &= mask_m[:, None] & mask_n[None, :]
            final_mask &= mask_causual
        else:
            mask_non_causal = mask_m[:, None] & mask_n[None, :]
            final_mask &= mask_non_causal

        if SLIDING_WINDOW_SIZE > 0:
            window_mask = query_pos[:, None] <= (
                start_n + offs_n[None, :] + SLIDING_WINDOW_SIZE
            )
            final_mask &= window_mask

        SKIP_TILE = False
        if USE_CUSTOM_MASK or SLIDING_WINDOW_SIZE > 0:
            SKIP_TILE = tl.max(tl.max(final_mask.to(tl.int32), axis=1), axis=0) == 0

        if not SKIP_TILE:
            if USE_MLS:
                k = tl.matrix_load(
                    K_Extend + cur_kv_head * stride_kh,
                    shape=(kv_head_num, Lq),
                    strides=(1, stride_kbs),
                    block_shape=(BLOCK_DMODEL, BLOCK_N),
                    offsets=(0, (cur_seq_extend_start_idx + start_n).to(tl.int32)),
                )
                if not (ALL_MASK_D & ALL_MASK_N):
                    k = tl.where((mask_d[:, None] & (mask_n[None, :])), k, 0.0)
            else:
                offs_k = (
                    (cur_seq_extend_start_idx + start_n + offs_n[None, :]) * stride_kbs
                    + cur_kv_head * stride_kh
                    + offs_d[:, None]
                )
                k = tl.load(
                    K_Extend + offs_k, mask=(mask_n[None, :]) & (mask_d[:, None]), other=0.0
                )

            qk = tl.dot(q.to(k.dtype), k, out_dtype=tl.float32)
            if BLOCK_DPE > 0:
                if USE_MLS:
                    kpe = tl.matrix_load(
                        K_Extend + cur_kv_head * stride_kh,
                        shape=(kv_head_num, Lq),
                        strides=(1, stride_kbs),
                        block_shape=(BLOCK_DPE, BLOCK_N),
                        offsets=(BLOCK_DMODEL, (cur_seq_extend_start_idx + start_n).to(tl.int32)),
                    )
                    if not ALL_MASK_N:
                        kpe = tl.where(mask_n[None, :], kpe, 0.0)
                else:    
                    offs_kpe = (
                        (cur_seq_extend_start_idx + start_n + offs_n[None, :]) * stride_kbs
                        + cur_kv_head * stride_kh
                        + offs_dpe[:, None]
                    )
                    kpe = tl.load(
                        K_Extend + offs_kpe,
                        mask=mask_n[None, :],
                        other=0.0,
                    )
                qk += tl.dot(qpe.to(kpe.dtype), kpe)

            qk *= sm_scale

            if logit_cap > 0:
                qk = logit_cap * _tanh(qk / logit_cap)

            if xai_temperature_len > 0:
                qk *= xai_temperature_reg[:, None]

            qk = tl.where(final_mask, qk, float("-inf"))

            row_max = tl.max(qk, 1)
            row_max_fixed = tl.where(row_max == float("-inf"), -1e20, row_max)
            n_e_max = tl.maximum(row_max_fixed, e_max)

            re_scale = tl.exp(e_max - n_e_max)
            p = tl.exp(qk - n_e_max[:, None])
            deno = deno * re_scale + tl.sum(p, 1)

            if USE_MLS:
                v = tl.matrix_load(
                    V_Extend + cur_kv_head * stride_vh,
                    shape=(kv_head_num, Lv),
                    strides=(stride_vbs, 1),
                    block_shape=(BLOCK_N, BLOCK_DV),
                    offsets=((cur_seq_extend_start_idx + start_n).to(tl.int32), 0),
                )
                if not (ALL_MASK_N & ALL_MASK_DV):
                    v = tl.where((mask_n[:, None] & mask_dv[None, :]), v, 0.0)
            else:    
                offs_v = (
                    (cur_seq_extend_start_idx + start_n + offs_n[:, None]) * stride_vbs
                    + cur_kv_head * stride_vh
                    + offs_dv[None, :]
                )
                v = tl.load(
                    V_Extend + offs_v, mask=mask_n[:, None] & mask_dv[None, :], other=0.0
                )
            p = p.to(v.dtype)
            acc = acc * re_scale[:, None] + tl.dot(p, v)

            e_max = n_e_max

    if HAS_SINK:
        cur_sink = tl.load(
            sink_ptr + cur_kv_head * kv_group_num + q_head_in_group,
            mask=mask_m,
            other=0.0,
        )
        deno += tl.exp(cur_sink - e_max)

    offs_o = (
        query_offset_0[:, None] * stride_obs
        + query_offset_1[:, None] * stride_oh
        + offs_dv[None, :]
    )
    if STORE_TRANSPOSE:
        tl.store(
            O_Extend + offs_o.T,
            (acc / deno[:, None]).T,
            mask=(mask_m[:, None] & mask_dv[None, :]).T,
        )
    else:
        tl.store(
            O_Extend + offs_o,
            acc / deno[:, None],
            mask=mask_m[:, None] & mask_dv[None, :],
        )


def create_tuple(k):
    if k[0] != '(' and k[-1] != ')':
        return k

    s = k[1:-1]
    entries = s.split(", ")
    ret = []
    for e in entries:
        if e[0] == "'" or e[0] == '"':
            ret.append(e[1:-1])
        else:
            ret.append(eval(e))
    ret_t = tuple(ret)
    return ret_t


def _load_config():
    dev = arch_info.get_device()
    fpath = f"{AITER_TRITON_CONFIGS_PATH}/{dev}-EXTEND_ATTENTION-FP16.json"
    try:
        with open(fpath, "r") as file:
            data = json.load(file)
    except FileNotFoundError:
        return {"config": {}, "path": {}, "key": [], "keys": []}
    res = {}
    res['config'] = data['config']
    res['path'] = data['path']
    res['key'] = list(data['config'].keys())
    res['keys'] = [create_tuple(k) for k in res['key']]
    return res


global_config = _load_config()


def _load_config_v2():
    """Autotuned configs for :func:`_fwd_kernel_v2` (fp8 / sglang-style scale path).

    Each ``config`` entry key must parse to a **7-tuple** via :func:`create_tuple`, matching
    runtime ``want7``; 5-tuple keys are not accepted.
    """
    dev = arch_info.get_device()
    fpath = f"{AITER_TRITON_CONFIGS_PATH}/{dev}-EXTEND_ATTENTION-V2-FP16.json"
    try:
        with open(fpath, "r") as file:
            data = json.load(file)
    except FileNotFoundError:
        return {"config": {}, "path": {}, "key": [], "keys": []}
    res = {}
    res["config"] = data["config"]
    res["path"] = data.get("path", {})
    res["key"] = list(data["config"].keys())
    res["keys"] = []
    for k in res["key"]:
        tup = create_tuple(k)
        if len(tup) != 7:
            raise ValueError(
                f"{dev}-EXTEND_ATTENTION-V2-FP16.json keys must be 7-tuples matching runtime "
                f"want7 (kv_group_num, Lq, Lv, USE_CUSTOM_MASK, IS_CAUSAL, HAS_SINK, "
                f"USE_SLIDING_WINDOW); got length {len(tup)} for {k!r}"
            )
        res["keys"].append(tup)
    return res


def _load_config_v2_decode():
    """Autotuned block sizes for :func:`_fwd_kernel_v2_decode` (short extend path)."""
    dev = arch_info.get_device()
    fpath = f"{AITER_TRITON_CONFIGS_PATH}/{dev}-EXTEND_ATTENTION-V2-DECODE-FP16.json"
    try:
        with open(fpath, "r") as file:
            data = json.load(file)
    except FileNotFoundError:
        return {"config": {}, "path": {}, "key": [], "keys": []}
    res = {}
    res["config"] = data["config"]
    res["path"] = data.get("path", {})
    res["key"] = list(data["config"].keys())
    res["keys"] = []
    for k in res["key"]:
        tup = create_tuple(k)
        if len(tup) != 7:
            raise ValueError(
                f"{dev}-EXTEND_ATTENTION-V2-DECODE-FP16.json keys must be 7-tuples matching runtime "
                f"want7 (kv_group_num, Lq, Lv, USE_CUSTOM_MASK, IS_CAUSAL, HAS_SINK, "
                f"USE_SLIDING_WINDOW); got length {len(tup)} for {k!r}"
            )
        res["keys"].append(tup)
    return res


TORCH_DTYPE_TO_DTYPE = {
    torch.float32: "f32",
    torch.float: "f32",
    torch.float16: "f16",
    torch.half: "f16",
    torch.bfloat16: "bf16",
    torch.float64: "f64",
    torch.double: "f64",
    torch.float8_e4m3fn: "f8_e4m3fn",
    torch.float8_e5m2: "f8_e5m2",
    torch.int8: "i8",
    torch.int16: "i16",
    torch.int32: "i32",
    torch.int64: "i64",
    torch.long: "i64",
    torch.uint8: "u8",
    torch.bool: "bool",
}


@functools.lru_cache
def get_gpu_label():
    target = triton.runtime.driver.active.get_current_target()
    device = torch.cuda.current_device()
    num_cu = torch.cuda.get_device_properties(device).multi_processor_count
    return f"{target.arch}_cu{num_cu}"


def _load_config_v3():
    """Autotuned configs for :func:`_fwd_kernel_v2` (fp8 / sglang-style scale path).

    Each ``config`` entry key must parse to a **7-tuple** via :func:`create_tuple`, matching
    runtime ``want7``; 5-tuple keys are not accepted.
    """
    fpath = f"{AITER_TRITON_CONFIGS_PATH}/extend_attn/_fwd_kernel_v2-device={get_gpu_label()}.json"
    try:
        with open(fpath, "r") as file:
            data = json.load(file)
    except FileNotFoundError:
        return {"key_name": {}, "fpath": {}, "config": {}, "path": {}, "key": [], "keys": []}
    res = {}
    res["key_name"] = data["key"]
    res['fpath'] = fpath
    res["config"] = data["config"]
    res["path"] = data.get("path", {})
    res["key"] = list(data["config"].keys())
    res["keys"] = []
    for k in res["key"]:
        tup = create_tuple(k)
        res["keys"].append(tup)
    return res


def _load_config_v3_decode():
    """Autotuned block sizes for :func:`_fwd_kernel_v2_decode` (short extend path)."""
    dev = arch_info.get_device()
    # fpath = f"{AITER_TRITON_CONFIGS_PATH}/extend_attn/_fwd_kernel_v2_decode-device={get_gpu_label()}-dtype=bf16.json"
    fpath = f"{AITER_TRITON_CONFIGS_PATH}/extend_attn/_fwd_kernel_v2_decode-device={get_gpu_label()}.json"
    try:
        with open(fpath, "r") as file:
            data = json.load(file)
    except FileNotFoundError:
        return {"key_name": {}, "fpath": {}, "config": {}, "path": {}, "key": [], "keys": []}
    res = {}
    res["key_name"] = data["key"]
    res['fpath'] = fpath
    res["config"] = data["config"]
    res["path"] = data.get("path", {})
    res["key"] = list(data["config"].keys())
    res["keys"] = []
    for k in res["key"]:
        tup = create_tuple(k)
        res["keys"].append(tup)
    return res


global_config_v2 = _load_config_v2()
global_config_v2_decode = _load_config_v2_decode()
global_config_v3 = _load_config_v3()
global_config_v3_decode = _load_config_v3_decode()

default_config = {
    "BLOCK_M": 32,
    "BLOCK_N": 32,
    "waves_per_eu": 1,
    "matrix_instr_nonkdim": 16,
    "kpack": 2,
    "num_warps": 4,
    "num_stages": 2,
    "USE_MLS": False,
}


@functools.lru_cache(maxsize=1024)
def _get_config(kv_group_num, Lq, Lv, use_custom_mask, is_causal):
    idx = -1
    for i, keys in enumerate(global_config['keys']):
        if keys[0] == kv_group_num and keys[1] == Lq and keys[2] == Lv \
            and keys[3] == use_custom_mask and keys[4] == is_causal:
            idx = i
            break

    if idx < 0:
        print("WARNING: optimal config not found, just use default config")
        return default_config, None
    else:
        key = global_config['key'][idx]
        return global_config['config'][key], global_config['path'][key]


@functools.lru_cache(maxsize=1024)
def _get_config_v2(
    kv_group_num,
    Lq,
    Lv,
    use_custom_mask,
    is_causal,
    has_sink: bool,
    use_sliding_window: bool,
):
    """
    Lookup order for ``_fwd_kernel_v2`` block sizes:

    1. ``want7 = (kv_group_num, Lq, Lv, use_custom_mask, is_causal, has_sink, USE_SLIDING_WINDOW)``
       against ``{arch}-EXTEND_ATTENTION-V2-FP16.json``. JSON keys must be **7-tuple** strings,
       same shape as ``want7`` (see :func:`_load_config_v2`). The last element is a bool:
       same tuning bucket for any ``sliding_window_size > 0``; use ``False`` when disabled (``<= 0``).
    2. If no V2 entry matches, :data:`default_config` (no fallback to v1 JSON).

    Log field mapping (typical): ``kv_group_num = q_extend.size(-2) // k_extend.size(-2)``,
    ``Lq = q_extend.size(-1)``, ``Lv = v_extend.size(-1)``,
    ``use_custom_mask = custom_mask is not None``, ``is_causal`` as passed,
    ``has_sink = sinks is not None``, ``USE_SLIDING_WINDOW = (sliding_window_size > 0)``.
    """
    want7 = (
        kv_group_num,
        Lq,
        Lv,
        use_custom_mask,
        is_causal,
        has_sink,
        use_sliding_window,
    )
    for i, keys in enumerate(global_config_v2["keys"]):
        if keys == want7:
            key = global_config_v2["key"][i]
            return global_config_v2["config"][key], global_config_v2["path"].get(key)

    print("WARNING: optimal V2 config not found, just use default config")
    return default_config, None


@functools.lru_cache(maxsize=1024)
def _get_config_v2_decode(
    kv_group_num,
    Lq,
    Lv,
    use_custom_mask,
    is_causal,
    has_sink: bool,
    use_sliding_window: bool,
):
    """
    Same ``want7`` as :func:`_get_config_v2`, but loads ``{arch}-EXTEND_ATTENTION-V2-DECODE-FP16.json``
    for :func:`_fwd_kernel_v2_decode`.
    """
    want7 = (
        kv_group_num,
        Lq,
        Lv,
        use_custom_mask,
        is_causal,
        has_sink,
        use_sliding_window,
    )
    for i, keys in enumerate(global_config_v2_decode["keys"]):
        if keys == want7:
            key = global_config_v2_decode["key"][i]
            return global_config_v2_decode["config"][key], global_config_v2_decode[
                "path"
            ].get(key)

    print("WARNING: optimal V2 decode config not found, just use default config")
    return default_config, None


def find_closest_index(target, lst):
    # lst: [(index, batch_size), ...]
    def key(item):
        val = item[1]
        diff = val - target
        if diff >= 0:
            return (0, diff)
        else:
            return (1, -diff)

    best = min(lst, key=key)
    return best[0]


@functools.lru_cache(maxsize=1024)
def _get_config_v3(key):
    _key = str(key)
    _configs = global_config_v3["config"]
    for k, v in _configs.items():
        if k == _key:
            return v

    # find the nearest batch size
    # key = (batch_size, *other)
    bs = []
    _keys = global_config_v3["keys"]
    for i, k in enumerate(_keys):
        if k[1:] == key[1:]:
            bs.append((i, k[0]))

    if bs:
        __key = global_config_v3["key"][find_closest_index(key[0], bs)]
        print(f'WARNING: Not found key {_key} from {global_config_v3["fpath"]}, mapping to key {__key}')
        return _configs[__key]
    else:
        print(f'WARNING: Not found optimal config from {global_config_v3["fpath"]} with key {str(key)}, just use default config')
        return default_config


@functools.lru_cache(maxsize=1024)
def _get_config_v3_decode(key):
    _key = str(key)
    _configs = global_config_v3_decode["config"]
    for k, v in _configs.items():
        if k == _key:
            return v

    # find the nearest batch size
    # key = (batch_size, *other)
    bs = []
    _keys = global_config_v3_decode["keys"]
    for i, k in enumerate(_keys):
        if k[1:] == key[1:]:
            bs.append((i, k[0]))

    if bs:
        __key = global_config_v3_decode["key"][find_closest_index(key[0], bs)]
        print(f'WARNING: Not found key {_key} from {global_config_v3["fpath"]}, mapping to key {__key}')
        return _configs[__key]
    else:
        print(f'WARNING: Not found optimal config from {global_config_v3_decode["fpath"]} with key {str(key)}, just use default config')
        return default_config


def has_kernel_cache(path):
    return False if not path or not os.path.isdir(f'{cache_knob.dir}/{path}') else True


@functools.lru_cache(maxsize=1024)
def get_v2_decode_final_grid(
    max_len_extend: int,
    kv_group_num: int,
    block_m_cfg: int,
    batch_size: int,
    kv_head_num: int,
) -> tuple[int, tuple[int, int, int]]:
    """Decode path: power-of-2 ``block_m_decode``; grid-3 cdiv matches kernel q_seq stride."""
    prod = max_len_extend * kv_group_num
    npo2 = triton.next_power_of_2(prod)
    kv_group_num_align = triton.next_power_of_2(kv_group_num)
    block_m_decode = block_m_cfg
    if prod < 16:
        block_m_decode = 16
    elif block_m_cfg > npo2:
        block_m_decode = npo2
    else:
        block_m_decode = max(block_m_cfg, kv_group_num_align)
    block_count = batch_size * kv_head_num * triton.cdiv(
        max_len_extend, block_m_decode // kv_group_num
    )
    if block_count <= 32:
        block_m_decode = max(max(block_m_decode // 2, 16), kv_group_num_align)
    q_seq = block_m_decode // kv_group_num
    grid = (batch_size, kv_head_num, triton.cdiv(max_len_extend, q_seq))
    return block_m_decode, grid


def to_dtype(torch_dtype):
    if torch_dtype == torch.float32:
        return 'fp32'
    elif torch_dtype == torch.float16:
        return 'fp16'
    elif torch_dtype == torch.bfloat16:
        return 'bf16'
    elif torch_dtype == torch.int32:
        return 'i32'
    else:
        return str(torch_dtype)

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,
    k_scale=None,
    v_scale=None,
    sliding_window_size=-1,
    sinks=None,
    window_kv_offsets=None,
    xai_temperature_len=-1,
    force_v2_prefill: bool = False,
):
    """
    q_extend, k_extend, v_extend, o_extend: contiguous tensors

    k_buffer, v_buffer: (prefix + extend) tensors in mem_manager

    Through ``config`` the signature matches the original aiter API. v2 / sglang
    extensions follow with defaults. ``k_scale`` / ``v_scale`` must both be
    ``None`` or both set (``float`` / ``int`` like sglang, or 1-element
    ``torch.Tensor`` on device); if both are set, :func:`_fwd_kernel_v2` is used.

    If ``force_v2_prefill`` is True and v2 is active, always use :func:`_fwd_kernel_v2`
    even when ``max_len_extend < 32`` (for tests / parity vs :func:`_fwd_kernel_v2_decode`).
    """
    # force_v2_prefill = True
    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_head_num = k_extend.shape[1]
    kv_group_num = head_num // kv_head_num

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

    use_v2 = k_scale is not None or v_scale is not None
    use_v2_decode = (
        use_v2 and max_len_extend < 32 and not force_v2_prefill
    )
    USE_SLIDING_WINDOW = sliding_window_size > 0

    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)
        # set to None to avoid capture cudagraph err
        custom_mask = None
        mask_indptr = None

    if config is None:
        if q_extend.dtype == torch.float16 or q_extend.dtype == torch.bfloat16:
            if use_v2:
                if triton_minor_version >= 5: # >= 3.5
                    key = [
                        batch_size,
                        kv_group_num,
                        Lq,
                        Lv,
                        USE_CUSTOM_MASK,
                        is_causal,
                        skip_prefix_custom_mask,
                        sinks is not None,
                        sliding_window_size,
                        xai_temperature_len,
                        str(q_extend.dtype),
                        str(k_extend.dtype),
                        str(v_extend.dtype),
                        str(o_extend.dtype),
                        str(k_buffer.dtype),
                        str(v_buffer.dtype),
                        str(qo_indptr.dtype),
                        str(kv_indptr.dtype),
                        str(kv_indices.dtype),
                    ]
                    for o in [
                        custom_mask,
                        mask_indptr,
                        sinks,
                        window_kv_offsets
                    ]:
                        if o is not None:
                            if hasattr(o, 'dtype'):
                                key.append(str(o.dtype))

                    key = tuple(key)
                    if use_v2_decode:
                        config = _get_config_v3_decode(key)
                    else:
                        config = _get_config_v3(key)
                else:
                    if use_v2_decode:
                        config, path = _get_config_v2_decode(
                            kv_group_num,
                            Lq,
                            Lv,
                            USE_CUSTOM_MASK,
                            is_causal,
                            sinks is not None,
                            USE_SLIDING_WINDOW,
                        )
                    else:
                        config, path = _get_config_v2(
                            kv_group_num,
                            Lq,
                            Lv,
                            USE_CUSTOM_MASK,
                            is_causal,
                            sinks is not None,
                            USE_SLIDING_WINDOW,
                        )
            else:
                keys = [kv_group_num, Lq, Lv, USE_CUSTOM_MASK, is_causal]
                config, path = _get_config(*keys)
        else:
            config, path = default_config, None
        assert config is not None, "ERROR: optimal config not found"

    block_m_cfg = config["BLOCK_M"]
    # Decode: block_m_decode is power of 2; Q_SEQ = block_m // G (floor), same as unified BLOCK_Q;
    # grid-3 cdiv(max_len_extend, q_seq) matches kernel cur_block_m stride.

    if use_v2_decode:
        block_m_decode, grid = get_v2_decode_final_grid(
            max_len_extend, kv_group_num, block_m_cfg, batch_size, kv_head_num
        )
        # print(f"{max_len_extend=}, {use_v2_decode=}, {block_m_decode=}, {grid=}")
    else:
        grid = (batch_size, head_num, triton.cdiv(max_len_extend, block_m_cfg))
    # num_stages = 1

    # extra_kargs = {}

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

    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),
    )

    block_const = dict(
        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,
    )

    if use_v2:
        HAS_SINK = sinks is not None
        assert k_scale is not None and v_scale is not None, "k_scale and v_scale must both be set"
        # k_scale / v_scale kept in Python API; v2 kernel TEMP omits them for perf vs v1.
        block_const_v2 = {
            **block_const,
            **config,
        }
        if use_v2_decode:
            block_const_v2 = {**block_const_v2, "BLOCK_M": block_m_decode}
            _fwd_kernel_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_const_v2,
                batch_size=batch_size,
                # USE_MLS=True if os.getenv("TRITON_USE_MLS", "0") == "1" and kv_group_num == head_num else False,
            )
        else:
            _fwd_kernel_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_const_v2,
                head_num=head_num,
                batch_size=batch_size,
                # USE_MLS=True if os.getenv("TRITON_USE_MLS", "0") == "1" and kv_group_num == head_num else False,
            )
        return

    fn = (
        _fwd_kernel[grid]
        if not has_kernel_cache(path)
        else functools.partial(
            triton.utils.run_saved_kernel, _fwd_kernel, path, grid=grid
        )
    )

    fn(
        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,
        *stride_args,
        logit_cap=logit_cap,
        **block_const,
        **config,
    )


def redundant_attention(
    q_extend,
    o_extend,
    k_buffer,
    v_buffer,
    b_req_idx,
    b_start_loc,
    b_seq_len,
    b_seq_len_prefix,
    max_len_in_batch,
):
    total_token_num = k_buffer.shape[0]
    B, H_Q, D = b_req_idx.shape[0], q_extend.shape[-2], q_extend.shape[-1]
    q_buffer = torch.empty(
        (total_token_num, H_Q, D), dtype=q_extend.dtype, device=q_extend.device
    )

    pt = 0
    for i in range(B):
        cur_seq_len_extend = b_seq_len[i] - b_seq_len_prefix[i]
        pl, pr = b_start_loc[i] + b_seq_len_prefix[i], b_start_loc[i] + b_seq_len[i]
        q_buffer[pl:pr] = q_extend[pt : pt + cur_seq_len_extend]
        pt += cur_seq_len_extend

    o_buffer = torch.empty_like(q_buffer)
    context_attention_fwd(
        q_buffer, k_buffer, v_buffer, o_buffer, b_start_loc, b_seq_len, max_len_in_batch
    )

    pt = 0
    for i in range(B):
        cur_seq_len_extend = b_seq_len[i] - b_seq_len_prefix[i]
        pl, pr = b_start_loc[i] + b_seq_len_prefix[i], b_start_loc[i] + b_seq_len[i]
        o_extend[pt : pt + cur_seq_len_extend] = o_buffer[pl:pr]
        pt += cur_seq_len_extend
