# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.
# ruff: noqa: E501, RUF012
# fmt: off

"""Attention KV cache modeling.

This module exposes the public ``PagedKVCache`` classes (``FlashInferPagedKVCache``
and ``TIRPagedKVCache``). The kernel factories that build the underlying TIR
functions are split across sibling private modules:

- ``_kernel_common``: shared helpers (enums, RoPE, mask, tile allocators,
  ``@T.macro`` bundle, tiling config, scheduling).
- ``_page_kernels``: page management (append, debug, copy, compact).
- ``_prefill_kernels``: prefill attention kernels (paged/ragged/MLA/dense).
- ``_decode_kernels``: decode attention kernels and state-merge helpers.

The private-named kernel factories are re-exported from this module so the
test suite can continue to import them via ``tvm.relax.frontend.nn.llm.kv_cache``.
"""

# pylint: disable=too-many-statements,too-many-arguments,invalid-name,line-too-long
import math
from typing import Any, Literal

import tvm
from tvm import relax as rx
from tvm import tirx
from tvm.relax.frontend.nn import Object, Tensor
from tvm.target import Target

# Re-export enums + kernel factories so existing ``from kv_cache import ...``
# users (test suite, tree_attn.py, mlc-llm, etc.) continue to work after the
# split. These names are referenced in ``__all__`` below to signal to linters
# that the imports are intentional public API (not dead code).
from ._decode_kernels import (
    _attention_decode,
    _attention_decode_cpu,
    _merge_state_inplace,
    _merge_state_inplace_cpu,
)
from ._kernel_common import AttnKind, RopeMode
from ._page_kernels import (
    _compact_kv_copy,
    _compact_kv_copy_cpu,
    _copy_single_page,
    _copy_single_page_cpu,
    _copy_single_page_mla,
    _kv_cache_debug_get_kv,
    _kv_cache_debug_get_kv_mla,
    _kv_cache_transpose_append,
    _kv_cache_transpose_append_mla,
)
from ._prefill_kernels import (
    _attention_prefill,
    _attention_prefill_cpu,
    _attention_prefill_mla,
    _attention_prefill_ragged,
    _attention_prefill_ragged_cpu,
    _attention_sequence_prefill,
    _attention_sequence_prefill_with_mask,
)
from .position_embedding import llama_rope_with_position_map
from .tree_attn import (
    tree_attn,
    tree_attn_cpu,
    tree_attn_with_paged_kv_cache,
    tree_attn_with_paged_kv_cache_cpu,
)

__all__ = [
    "AttnKind",
    "FlashInferPagedKVCache",
    "PagedKVCache",
    "RopeMode",
    "TIRPagedKVCache",
    "_attention_decode",
    "_attention_decode_cpu",
    "_attention_prefill",
    "_attention_prefill_cpu",
    "_attention_prefill_mla",
    "_attention_prefill_ragged",
    "_attention_prefill_ragged_cpu",
    "_attention_sequence_prefill",
    "_attention_sequence_prefill_with_mask",
    "_compact_kv_copy",
    "_compact_kv_copy_cpu",
    "_copy_single_page",
    "_copy_single_page_cpu",
    "_copy_single_page_mla",
    "_kv_cache_debug_get_kv",
    "_kv_cache_debug_get_kv_mla",
    "_kv_cache_transpose_append",
    "_kv_cache_transpose_append_mla",
    "_merge_state_inplace",
    "_merge_state_inplace_cpu",
    "llama_rope_with_position_map",
    "tree_attn",
    "tree_attn_cpu",
    "tree_attn_with_paged_kv_cache",
    "tree_attn_with_paged_kv_cache_cpu",
]


class PagedKVCache(Object):  # pylint: disable=too-few-public-methods
    """The Paged KV Cache used in LLM batching for efficient attention computation."""

    extern_mods: list[tvm.runtime.Module] = []

    def attention_with_fused_qkv(
        self,
        layer_id: int,
        qkv: Tensor,
        num_qo_heads: int,
        sm_scale: float,
    ) -> Tensor:
        """Compute attention with the given fused q/k/v data and in-cache k/v data
        on the specified layer. Rotary position embeddings are applied to k/v
        within this function.

        - For prefill, the input qkv and output tensor have shape
        (1, total_seq_len) for the first two dimensions.
        - For decode, the input qkv and output tensor have shape
        (batch_size, 1) for the first two dimensions.
        - The input qkv have `2 * num_qo_heads + num_kv_heads` at the third dim.
        - The output tensor have `num_qo_heads` at the third dim.
        - The input qkv and output tensor have `head_dim` at the last dim.
        """
        # pylint: disable=protected-access
        b, s, _, d = qkv._expr.ty.shape
        qkv = qkv.reshape(b * s, qkv.shape[2], d)
        return Tensor(
            _expr=rx.BlockBuilder.current().emit(
                rx.call_dps_packed(
                    "vm.builtin.attention_kv_cache_attention_with_fused_qkv",
                    [
                        self._expr,
                        rx.prim_value(layer_id),  # type: ignore[arg-type]
                        rx.prim_value(sm_scale),
                        qkv._expr,
                    ],
                    out_ty=rx.TensorType((b * s, num_qo_heads, d), qkv.dtype),
                )
            )
        ).reshape(b, s, num_qo_heads, d)

    def self_attention(  # pylint: disable=too-many-locals
        self,
        layer_id: int,
        q: Tensor,
        k: Tensor,
        v: Tensor,
        sm_scale: float,
    ) -> tuple[Tensor, Tensor]:
        """Fine-grained API that computes ragged self attention with Q/K/V data."""
        # pylint: disable=protected-access
        b, s, h_qo, d_qk = q._expr.ty.shape
        _, _, h_kv, d_v = v._expr.ty.shape
        q = q.reshape(b * s, h_qo, d_qk)
        k = k.reshape(b * s, h_kv, d_qk)
        v = v.reshape(b * s, h_kv, d_v)
        bb = rx.BlockBuilder.current()
        attn_results = bb.emit(
            rx.call_dps_packed(
                "vm.builtin.attention_kv_cache_self_attention",
                [
                    self._expr,
                    rx.prim_value(layer_id),  # type: ignore[arg-type]
                    rx.prim_value(sm_scale),
                    q._expr,
                    k._expr,
                    v._expr,
                ],
                out_ty=[
                    rx.TensorType((b * s, h_qo, d_v), q.dtype),
                    rx.TensorType((b * s, h_qo), "float32"),
                ],
            )
        )
        assert isinstance(attn_results.ty, rx.TupleType)
        assert len(attn_results.ty.fields) == 2
        o = Tensor(_expr=bb.emit(rx.TupleGetItem(attn_results, 0))).reshape(b, s, h_qo, d_v)
        lse = Tensor(_expr=bb.emit(rx.TupleGetItem(attn_results, 1))).reshape(b, s, h_qo)
        return o, lse

    def cross_attention(
        self,
        layer_id: int,
        q: Tensor,
        v_head_dim: int,
        sm_scale: float,
    ) -> tuple[Tensor, Tensor]:
        """Fine-grained API that computes paged cross attention with Q and in-cache KV data."""
        # pylint: disable=protected-access
        b, s, h_qo, d_qk = q._expr.ty.shape
        q = q.reshape(b * s, h_qo, d_qk)
        bb = rx.BlockBuilder.current()
        attn_results = bb.emit(
            rx.call_dps_packed(
                "vm.builtin.attention_kv_cache_cross_attention",
                [
                    self._expr,
                    rx.prim_value(layer_id),  # type: ignore[arg-type]
                    rx.prim_value(sm_scale),
                    q._expr,
                ],
                out_ty=[
                    rx.TensorType((b * s, h_qo, v_head_dim), q.dtype),
                    rx.TensorType((b * s, h_qo), "float32"),
                ],
            )
        )
        assert isinstance(attn_results.ty, rx.TupleType)
        assert len(attn_results.ty.fields) == 2
        o = Tensor(_expr=bb.emit(rx.TupleGetItem(attn_results, 0))).reshape(b, s, h_qo, v_head_dim)
        lse = Tensor(_expr=bb.emit(rx.TupleGetItem(attn_results, 1))).reshape(b, s, h_qo)
        return o, lse

    def append_mla_kv(self, layer_id: int, kv: Tensor) -> "PagedKVCache":
        """Fine-grained API that appends the MLA K/V data to KV cache."""
        # pylint: disable=protected-access
        b, s, _, d_qk = kv._expr.ty.shape
        kv = kv.reshape(b * s, d_qk)
        return PagedKVCache(
            _expr=rx.call_pure_packed(
                "vm.builtin.attention_kv_cache_append_mla_kv",
                self._expr,
                rx.prim_value(layer_id),  # type: ignore[arg-type]
                kv._expr,
                ty_args=rx.AnyType(),
            ),
            _name="paged_kv_cache",
        )

    def merge_attn_output_inplace(
        self,
        o_self_attn: Tensor,
        lse_self_attn: Tensor,
        o_cross_attn: Tensor,
        lse_cross_attn: Tensor,
    ) -> tuple[Tensor, Tensor]:
        """Fine-grained API that merges the attention output from two sources.
        The first two tensors will be inplace updated.
        """
        # pylint: disable=protected-access
        b, s, h_qo, d_v = o_self_attn._expr.ty.shape
        o_self_attn = o_self_attn.reshape(b * s, h_qo, d_v)
        lse_self_attn = lse_self_attn.reshape(b * s, h_qo)
        o_cross_attn = o_cross_attn.reshape(b * s, h_qo, d_v)
        lse_cross_attn = lse_cross_attn.reshape(b * s, h_qo)
        bb = rx.BlockBuilder.current()
        merge_results = bb.emit(
            rx.call_pure_packed(
                "vm.builtin.attention_kv_cache_merge_attn_output_inplace",
                self._expr,
                o_self_attn._expr,
                lse_self_attn._expr,
                o_cross_attn._expr,
                lse_cross_attn._expr,
                ty_args=rx.TupleType(
                    [o_self_attn._expr.ty, lse_self_attn._expr.ty]
                ),
            )
        )
        assert isinstance(merge_results.ty, rx.TupleType)
        assert len(merge_results.ty.fields) == 2
        o_self_attn = Tensor(_expr=bb.emit(rx.TupleGetItem(merge_results, 0))).reshape(
            b, s, h_qo, d_v
        )
        lse_self_attn = Tensor(_expr=bb.emit(rx.TupleGetItem(merge_results, 1))).reshape(b, s, h_qo)
        return o_self_attn, lse_self_attn

    def get_query_positions(self, total_length: tirx.Expr) -> Tensor:
        """Get the in-sequence positions of each slot in the query,
        which are needed for applying positional embeddings in some models.

        Parameters
        ----------
        total_length : tirx.Expr
            The summed-up total sequence length of queries in
            the batch being forwarded.

        Returns
        -------
        q_positions : Tensor
            The in-sequence query positions, in shape `(total_length,)`
        """
        return Tensor(
            _expr=rx.BlockBuilder.current().emit(
                rx.call_pure_packed(
                    "vm.builtin.attention_kv_cache_get_query_positions",
                    self._expr,
                    ty_args=rx.TensorType((total_length,), "int32"),
                )
            )
        )

    # pylint: enable=protected-access


def _prepare_yarn_rope_scaling(rope_scaling: dict[str, Any] | None, rope_theta: float | None) -> dict[str, Any] | None:
    """Ensure Yarn-specific scaling configs include the theta metadata."""
    if rope_scaling is None:
        return None
    if rope_scaling.get("rope_type") != "yarn":
        return rope_scaling

    rope_scaling_updated = dict(rope_scaling)
    if "inv_theta_log_scale" not in rope_scaling_updated and rope_theta is not None:
        theta_value = float(rope_theta)
        rope_scaling_updated["inv_theta_log_scale"] = 1.0 / (2 * math.log(theta_value))
    return rope_scaling_updated


class FlashInferPagedKVCache(PagedKVCache):  # pylint: disable=too-few-public-methods
    """Paged KV cache using FlashInfer (CUDA) kernels."""

    def __init__(  # pylint: disable=too-many-locals
        self,
        attn_kind: Literal["mha", "mla"] | list[Literal["mha", "mla", "mha_sliding"]],
        max_batch_size: tirx.Var,
        max_total_seq_len: tirx.Var,
        prefill_chunk_size: tirx.Var,
        page_size: tirx.Var,
        support_sliding_window: tirx.Var,
        layer_partition: rx.ShapeExpr,
        num_hidden_layers: int,
        num_attention_heads: int,
        num_key_value_heads: int,
        qk_head_dim: int,
        v_head_dim: int,
        mla_original_qk_head_dim: int,
        mla_original_v_head_dim: int,
        rope_mode: RopeMode,
        rope_scale: int,
        rope_theta: int,
        rope_scaling: dict[str, Any],
        rope_ext_factors: rx.Expr,
        rotary_dim: int,
        enable_disaggregation: bool,
        dtype: str,
        target: Target,
        name: str = "paged_kv_cache",
    ) -> None:
        """Create a paged KV cache object with FlashInfer kernels.

        Parameters
        ----------
        max_batch_size : tirx.Var
            The maximum allowed batch size of the KV cache.
            It is a symbolic variable whose concrete value is specified
            at runtime.
        max_total_seq_len : tirx.Var
            The maximum allowed total sequence length of the KV cache.
            It is a symbolic variable whose concrete value is specified
            at runtime.
        prefill_chunk_size : tirx.Var
            The maximum total sequence length in a prefill.
            It is a symbolic variable whose concrete value is specified
            at runtime.
        page_size : tirx.Var
            The size (a.k.a. number of tokens) of each page.
            It is a symbolic variable whose concrete value is specified
            at runtime.
        support_sliding_window : tirx.Var
            0 or 1, denoting whether the KV cache supports sliding window.
            It is a symbolic variable whose concrete value is specified
            at runtime.
        layer_partition : rx.ShapeExpr
            The KV cache layer partition for pipeline stages.
            It is an indptr array, denoting the starting layer of each pipeline stage.
        rope_mode : RopeMode
            The RoPE mode of the Paged KV cache.
            If it is normal, RoPE will be applied to k before adding k to cache.
            Otherwise, RoPE will be applied to q/k in attention kernel on-the-fly.
        rope_scale : int
            The scale of rotary position embedding.
        rope_theta : int
            The base of rotary position embedding.
        rope_scaling: Dict[str, Any]
            The RoPE scaling information dict.
        rope_ext_factors: rx.Expr
            The RoPE extension factors when "longrope" mode RoPE scaling is enabled.
        rotary_dim : int
            The number of dimensions in the embedding that RoPE is applied to.
        enable_disaggregation : bool
            Whether to enable disaggregation in the KV cache.
        """
        assert rope_mode != RopeMode.INLINE, "FlashInfer RoPE does not support inline mode."
        rope_scaling = _prepare_yarn_rope_scaling(rope_scaling, rope_theta)

        attn_kind_single = attn_kind[0] if isinstance(attn_kind, list) else attn_kind
        if attn_kind_single == "mha_sliding":
            attn_kind_single = "mha"
        flashinfer_prefill_mods = rx.backend.cuda.flashinfer.gen_flashinfer_prefill_module(
            dtype_q=dtype,
            dtype_kv=dtype,
            dtype_o=dtype,
            qk_head_dim=(qk_head_dim if attn_kind_single == "mha" else mla_original_qk_head_dim),
            v_head_dim=(v_head_dim if attn_kind_single == "mha" else mla_original_v_head_dim),
            enable_inline_rope=False,
            return_static_libs=True,
        )
        flashinfer_decode_mods = (
            rx.backend.cuda.flashinfer.gen_flashinfer_decode_module(
                dtype_q=dtype,
                dtype_kv=dtype,
                dtype_o=dtype,
                qk_head_dim=qk_head_dim,
                v_head_dim=v_head_dim,
                enable_inline_rope=False,
                return_static_libs=True,
            )
            if attn_kind_single == "mha"
            else []
        )
        flashinfer_mla_mods = (
            rx.backend.cuda.flashinfer.gen_flashinfer_mla_module(
                dtype_q=dtype,
                dtype_kv=dtype,
                dtype_o=dtype,
                head_dim_ckv=v_head_dim,
                head_dim_kpe=qk_head_dim - v_head_dim,
                return_static_libs=True,
            )
            if attn_kind_single == "mla"
            else []
        )
        self.extern_mods = flashinfer_prefill_mods + flashinfer_decode_mods + flashinfer_mla_mods

        bb = rx.BlockBuilder.current()
        mha_functions = (
            [
                rx.Tuple([rx.StringImm("flashinfer"), rx.ExternFunc("batch_prefill_paged_run"), rx.ExternFunc("batch_prefill_plan")]),
                rx.Tuple([rx.StringImm("flashinfer"), rx.ExternFunc("batch_decode_run"), rx.ExternFunc("batch_decode_plan")]),
                rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_prefill(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, True, rope_scaling, target), "tir_attention_prefill_sliding_window")]),
                rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_decode(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, True, rope_scaling, target), "tir_attention_decode_sliding_window")]),
                rx.Tuple([rx.StringImm("tirx"), bb.add_func(tree_attn_with_paged_kv_cache(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, rope_scaling, target), "tir_attention_prefill_with_tree_mask_with_paged_kv_cache")]),
                rx.Tuple([rx.StringImm("tirx"), bb.add_func(tree_attn(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, rope_scaling, target), "tir_attention_prefill_with_tree_mask")]),
            ]
            if attn_kind_single == "mha"
            else [rx.Tuple([]) for _ in range(6)]
        )
        ragged_prefill_function = rx.Tuple([rx.StringImm("flashinfer"), rx.ExternFunc("batch_prefill_ragged_run"), rx.ExternFunc("batch_prefill_plan")]) if attn_kind_single == "mha" else rx.Tuple([rx.StringImm("flashinfer"), rx.ExternFunc("batch_prefill_ragged_run"), rx.ExternFunc("batch_prefill_plan"), rx.prim_value(mla_original_qk_head_dim), rx.prim_value(mla_original_v_head_dim)])
        mla_function = rx.Tuple([rx.StringImm("flashinfer"), rx.ExternFunc("batch_mla_run"), rx.ExternFunc("batch_mla_plan")] if attn_kind_single == "mla" else [])
        attn_merge_functions = [
            bb.add_func(_merge_state_inplace(num_attention_heads, v_head_dim, dtype, target, "tir_attention_merge_state"), "tir_attention_merge_state"),
        ]
        if attn_kind_single == "mla":
            attn_merge_functions.append(bb.add_func(_merge_state_inplace(num_attention_heads, mla_original_v_head_dim, dtype, target, "tir_attention_merge_state_mla"), "tir_attention_merge_state_mla"))

        if isinstance(attn_kind, list):
            attn_kind = [int(getattr(AttnKind, layer_kind.upper())) for layer_kind in attn_kind]
        else:
            attn_kind = [int(getattr(AttnKind, attn_kind.upper())) for _ in range(num_hidden_layers)]

        args = [
            rx.ShapeExpr(
                [
                    max_batch_size,
                    max_total_seq_len,
                    prefill_chunk_size,
                    page_size,
                    support_sliding_window,
                ]
            ),
            layer_partition,
            rx.prim_value(num_attention_heads),
            rx.prim_value(num_key_value_heads),
            rx.prim_value(qk_head_dim),
            rx.prim_value(v_head_dim),
            rx.ShapeExpr(attn_kind),
            rx.prim_value(enable_disaggregation),
            rx.prim_value(rope_mode),
            rx.prim_value(rope_scale),
            rx.prim_value(rope_theta),
            rope_ext_factors,
            rx.op.zeros((), dtype),
            bb.add_func(_kv_cache_transpose_append(num_key_value_heads, qk_head_dim, dtype), "kv_cache_transpose_append"),
            bb.add_func(_kv_cache_transpose_append_mla(qk_head_dim, dtype), "kv_cache_transpose_append_mla"),
            ragged_prefill_function,
            *mha_functions,
            mla_function,
            rx.Tuple(attn_merge_functions),
            bb.add_func(llama_rope_with_position_map(rope_theta, rope_scale, qk_head_dim, num_attention_heads, num_key_value_heads, dtype, rope_scaling, rotary_dim), "tir_split_rotary"),
            bb.add_func(_copy_single_page(num_key_value_heads, page_size, qk_head_dim, dtype, target) if attn_kind_single == "mha" else _copy_single_page_mla(page_size, qk_head_dim, dtype, target), "kv_cache_copy_single_page"),
            bb.add_func(_kv_cache_debug_get_kv(num_hidden_layers, num_key_value_heads, qk_head_dim, dtype), "kv_cache_debug_get_kv"),
            bb.add_func(_compact_kv_copy(num_key_value_heads, qk_head_dim, dtype, target), "kv_cache_compact_kv_copy"),
        ]
        super().__init__(
            _expr=rx.call_pure_packed(
                "vm.builtin.paged_attention_kv_cache_create",
                *args,
                ty_args=rx.AnyType(),
            ),
            _name=name,
        )


class TIRPagedKVCache(PagedKVCache):  # pylint: disable=too-few-public-methods
    """Paged KV cache using TIR kernels."""

    def __init__(  # pylint: disable=too-many-locals
        self,
        attn_kind: Literal["mha", "mla"] | list[Literal["mha", "mla", "mha_sliding"]],
        max_batch_size: tirx.Var,
        max_total_seq_len: tirx.Var,
        prefill_chunk_size: tirx.Var,
        page_size: tirx.Var,
        support_sliding_window: tirx.Var,
        layer_partition: rx.ShapeExpr,
        num_hidden_layers: int,
        num_attention_heads: int,
        num_key_value_heads: int,
        qk_head_dim: int,
        v_head_dim: int,
        mla_original_qk_head_dim: int,
        mla_original_v_head_dim: int,
        rope_mode: RopeMode,
        rope_scale: int,
        rope_theta: int,
        rope_scaling: dict[str, Any],
        rope_ext_factors: rx.Expr,
        rotary_dim: int,
        enable_disaggregation: bool,
        dtype: str,
        target: Target,
        name: str = "paged_kv_cache",
    ) -> None:
        """Create a paged KV cache object with TIR kernels.

        Parameters
        ----------
        max_batch_size : tirx.Var
            The maximum allowed batch size of the KV cache.
            It is a symbolic variable whose concrete value is specified
            at runtime.
        max_total_seq_len : tirx.Var
            The maximum allowed total sequence length of the KV cache.
            It is a symbolic variable whose concrete value is specified
            at runtime.
        prefill_chunk_size : tirx.Var
            The maximum total sequence length in a prefill.
            It is a symbolic variable whose concrete value is specified
            at runtime.
        page_size : tirx.Var
            The size (a.k.a. number of tokens) of each page.
            It is a symbolic variable whose concrete value is specified
            at runtime.
        support_sliding_window : tirx.Var
            0 or 1, denoting whether the KV cache supports sliding window.
            It is a symbolic variable whose concrete value is specified
            at runtime.
        layer_partition : rx.ShapeExpr
            The KV cache layer partition for pipeline stages.
            It is an indptr array, denoting the starting layer of each pipeline stage.
        rope_mode : RopeMode
            The RoPE mode of the Paged KV cache.
            If it is normal, RoPE will be applied to k before adding k to cache.
            Otherwise, RoPE will be applied to q/k in attention kernel on-the-fly.
        rope_scale : int
            The scale of rotary position embedding.
        rope_theta : int
            The base of rotary position embedding.
        rope_scaling: Dict[str, Any]
            The RoPE scaling information dict.
        rope_ext_factors: rx.Expr
            The RoPE extension factors when "longrope" mode RoPE scaling is enabled.
        rotary_dim : int
            The number of dimensions in the embedding that RoPE is applied to.
        enable_disaggregation : bool
            Whether to enable disaggregation in the KV cache.
        target : Target
            The target to build the model to.
        """
        rope_scaling = _prepare_yarn_rope_scaling(rope_scaling, rope_theta)
        attn_kind_single = attn_kind[0] if isinstance(attn_kind, list) else attn_kind
        if attn_kind_single == "mha_sliding":
            attn_kind_single = "mha"
        if isinstance(attn_kind, list):
            attn_kind = [int(getattr(AttnKind, layer_kind.upper())) for layer_kind in attn_kind]
        else:
            attn_kind = [int(getattr(AttnKind, attn_kind.upper())) for _ in range(num_hidden_layers)]
        bb = rx.BlockBuilder.current()
        args = [
            rx.ShapeExpr(
                [
                    max_batch_size,
                    max_total_seq_len,
                    prefill_chunk_size,
                    page_size,
                    support_sliding_window,
                ]
            ),
            layer_partition,
            rx.prim_value(num_attention_heads),
            rx.prim_value(num_key_value_heads),
            rx.prim_value(qk_head_dim),
            rx.prim_value(v_head_dim),
            rx.ShapeExpr(attn_kind),
            rx.prim_value(enable_disaggregation),
            rx.prim_value(rope_mode),
            rx.prim_value(rope_scale),
            rx.prim_value(rope_theta),
            rope_ext_factors,
            rx.op.zeros((), dtype),
            bb.add_func(_kv_cache_transpose_append(num_key_value_heads, qk_head_dim, dtype), "kv_cache_transpose_append"),
            bb.add_func(_kv_cache_transpose_append_mla(qk_head_dim, dtype), "kv_cache_transpose_append_mla"),
        ]

        if target.kind.name == "llvm":
            if attn_kind_single == "mla":
                raise ValueError("MLA is not supported in TIR kernels for now.")
            args.extend(
                [
                    rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_prefill_ragged_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, v_head_dim, dtype, rope_scaling), "tir_attention_prefill_ragged_cpu")]),
                    rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_prefill_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, False, rope_scaling), "tir_attention_prefill_cpu")]),
                    rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_decode_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, False, rope_scaling), "tir_attention_decode_cpu")]),
                    rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_prefill_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, True, rope_scaling), "tir_attention_prefill_cpu_sliding_window")]),
                    rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_decode_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, True, rope_scaling), "tir_attention_decode_cpu_sliding_window")]),
                    rx.Tuple([rx.StringImm("tirx"), bb.add_func(tree_attn_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, rope_scaling), "tir_attention_prefill_with_tree_mask_cpu")]),
                    rx.Tuple([rx.StringImm("tirx"), bb.add_func(tree_attn_with_paged_kv_cache_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, rope_scaling), "tir_attention_prefill_with_tree_mask_with_paged_kv_cache_cpu")]),
                    rx.Tuple([]),  # f_mla_prefill
                    rx.Tuple([bb.add_func(_merge_state_inplace_cpu(dtype), "tir_attention_merge_state_cpu")]),
                    bb.add_func(llama_rope_with_position_map(rope_theta, rope_scale, qk_head_dim, num_attention_heads, num_key_value_heads, dtype, rope_scaling, rotary_dim), "tir_split_rotary"),
                    bb.add_func(_copy_single_page_cpu(num_key_value_heads, page_size, qk_head_dim, dtype), "kv_cache_copy_single_page_cpu"),
                    bb.add_func(_kv_cache_debug_get_kv(num_hidden_layers, num_key_value_heads, qk_head_dim, dtype), "kv_cache_debug_get_kv"),
                    bb.add_func(_compact_kv_copy_cpu(num_key_value_heads, qk_head_dim, dtype), "kv_cache_compact_kv_copy_cpu"),
                ]
            )
        else:
            ragged_qk_head_dim = qk_head_dim if attn_kind_single == "mha" else mla_original_qk_head_dim
            ragged_v_head_dim = v_head_dim if attn_kind_single == "mha" else mla_original_v_head_dim
            args.append(rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_prefill_ragged(num_key_value_heads if attn_kind_single == "mha" else num_attention_heads, num_attention_heads, ragged_qk_head_dim, ragged_v_head_dim, dtype, rope_scaling, target), "tir_attention_prefill_ragged")]))
            mha_functions = (
                [
                    rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_prefill(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, False, rope_scaling, target), "tir_attention_prefill")]),
                    rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_decode(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, False, rope_scaling, target), "tir_attention_decode")]),
                    rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_prefill(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, True, rope_scaling, target), "tir_attention_prefill_sliding_window")]),
                    rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_decode(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, True, rope_scaling, target), "tir_attention_decode_sliding_window")]),
                    rx.Tuple([rx.StringImm("tirx"), bb.add_func(tree_attn_with_paged_kv_cache(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, rope_scaling, target), "tir_attention_prefill_with_tree_mask_with_paged_kv_cache")]),
                    rx.Tuple([rx.StringImm("tirx"), bb.add_func(tree_attn(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, rope_scaling, target), "tir_attention_prefill_with_tree_mask")]),
                ]
                if attn_kind_single == "mha"
                else [rx.Tuple([]) for _ in range(6)]
            )
            mla_function = rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_prefill_mla(num_attention_heads, v_head_dim, qk_head_dim - v_head_dim, dtype, False, target), "tir_attention_prefill_mla")] if attn_kind_single == "mla" else [])
            attn_merge_functions = [
                bb.add_func(_merge_state_inplace(num_attention_heads, v_head_dim, dtype, target, "tir_attention_merge_state"), "tir_attention_merge_state"),
            ]
            if attn_kind_single == "mla":
                attn_merge_functions.append(bb.add_func(_merge_state_inplace(num_attention_heads, mla_original_v_head_dim, dtype, target, "tir_attention_merge_state_mla"), "tir_attention_merge_state_mla"))
            args.extend(mha_functions)
            args.append(mla_function)
            args.extend(
                [
                    rx.Tuple(attn_merge_functions),
                    bb.add_func(llama_rope_with_position_map(rope_theta, rope_scale, qk_head_dim, num_attention_heads, num_key_value_heads, dtype, rope_scaling, rotary_dim), "tir_split_rotary"),
                    bb.add_func(_copy_single_page(num_key_value_heads, page_size, qk_head_dim, dtype, target) if attn_kind_single == "mha" else _copy_single_page_mla(page_size, qk_head_dim, dtype, target), "kv_cache_copy_single_page"),
                    bb.add_func(_kv_cache_debug_get_kv(num_hidden_layers, num_key_value_heads, qk_head_dim, dtype), "kv_cache_debug_get_kv"),
                    bb.add_func(_compact_kv_copy(num_key_value_heads, qk_head_dim, dtype, target), "kv_cache_compact_kv_copy"),
                ]
            )

        super().__init__(
            _expr=rx.call_pure_packed(
                "vm.builtin.paged_attention_kv_cache_create",
                *args,
                ty_args=rx.AnyType(),
            ),
            _name=name,
        )
