chore: import upstream snapshot with attribution

python/tvm/relax/frontend/nn/llm/kv_cache.py

@@ -0,0 +1,686 @@
+# 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,
+        )